Deliverable 4: the slimmed platform-optimizer convention (contracts, portable state, DCP) - #67
Draft
thad0ctor wants to merge 57 commits into
Draft
Deliverable 4: the slimmed platform-optimizer convention (contracts, portable state, DCP)#67thad0ctor wants to merge 57 commits into
thad0ctor wants to merge 57 commits into
Conversation
OptimizerCapabilities accepted arbitrary objects in training, checkpoints, and precisions and non-bool capability flags, and OptimizerChildContract accepted any object as its child contract, so an invalid declaration constructed silently and failed far from where it was built. Check element and flag types in __post_init__ like every sibling descriptor.
_hybrid_contract copied the muon child's training claims verbatim and fell back to the static base tuple when muon was absent, never consulting the backup child. A Gefen-backed backup-only hybrid therefore omitted the flattened element-shard layout its backup validates, while an AdamW-backed hybrid claimed DTensor training no code validates for that child. Declare the ordered union of the present children's own claims instead; a backup without a contract contributes only plain replicated training.
The canonical-import and hybrid-rebinding atomicity tests asserted only top-level object identity, so a regression that published staged state by mutating live containers in place (state[param].update, copy_() into an existing state tensor, group option or group['params'] element edits) would pass every assertion. Add tests/_state_snapshot.py, a shared deep snapshot helper that captures container identities plus bitwise clones of every reachable tensor, per-parameter state dict, counter, and param-group entry, and rewire both files' snapshot helpers onto it. Verified that all eight simulated in-place regressions now fail the assertions while the current implementation still passes.
test_activation_and_restore_copy_failures_are_atomic only checked that optimizer.state, the per-parameter dict, and the codebook were the same objects after an injected copy failure; an offload/restore regression that partially overwrote m_codebook/m_magnitude/vmean in place before failing would keep every identity intact and pass. Snapshot the persistent tensors and counters with the file's existing _persistent_snapshot helper before each injected failure and assert bitwise equality (plus codebook value and global step) afterwards. Verified the new assertions catch a simulated in-place corruption that the old identity checks missed.
Every flattened shard in the scoped-collective tests had nonzero length, leaving the dedicated empty-shard logic (nonempty_activity filtering of the gradient-presence consensus and the empty-slice no-op paths) entirely unexercised. Add a two-member gloo test where one member binds a legal zero-length flattened slice to a live numel-0 tensor and assert it joins initialization, step, refresh, and failure-sync collectives symmetrically: codebooks agree with the nonempty-only oracle on both members, presence asymmetry on the empty member is accepted, the empty member's state stays inert, a nonempty-member exact-DP failure is synchronized to the empty member and leaves both atomic and retryable.
The reuse_existing_periods flag that gates the conditional "initialize" operation header in _prepare_gefen_exact_codebook was derived from the rank-local _resuming_from_checkpoint() predicate on both hot paths (_maybe_refresh_gefen_codebook and initialize_codebook). Per-parameter state presence is legitimately rank-asymmetric under an explicit scope (empty non-owner slots and zero-length flattened shards never create state), so after a consolidation-style resume that strips the common codebook, state-bearing members computed reuse=True and skipped the header all_gather while empty members computed reuse=False and issued it, mismatching the group's collective schedules. Resolve one group-wide decision before any member branches on it: _scope_agreed_resuming_from_checkpoint all_reduces (MAX) the local predicate over the codebook scope, so any member that restored periods makes the whole scope reuse them. Both call sites are reached by every member of a multi-member scope in the same order, and empty members are unaffected behaviorally by reuse=True because parameters without gradients or elements never enter the period iterator.
_ensure_codebook_scope_agreement early-returned on the rank-local _gefen_codebook_scope_validated flag. Collective-free rank-local operations (move_state_, offload_state_, staged and native checkpoint loads) reset that flag on the member that ran them only, while every value the step operation header fingerprints stays identical (the codebook fingerprint is computed on CPU bytes and is device independent). One member then entered the agreement all_gather while the others early-returned, diverging the scoped collective schedule directly behind a passing header exchange. Fold the flag into the always-exchanged operation header as a trailing decision bit that is excluded from the equality check, and lower the flag on every member when any member reports it cleared, so the group re-validates together and then proceeds with identical collectives. This resolution was chosen over the two alternatives deliberately: raising on a flag mismatch would turn the advertised collective-free rank-local operations into scope-wide faults on legitimately asymmetric use, and dropping the reset for movement/offload would keep the divergence for staged and native loads, which genuinely can change scope-agreement-relevant state and must keep resetting the flag.
_canonical_import_live_token identified locally bound parameters by id() only, so commit_canonical_state_import could not detect that a parameter's storage was retargeted (or mutated in place) between prepare and commit, even though the staged shadow was device-cast and geometry-validated against the parameters as they existed at prepare time. The sibling portable path already folds _parameter_storage_token (device, dtype, layout, shape, stride, storage pointer/size, version) into its live token for every local binding; mirror it here so a stale prepared import is refused with the existing freshness error.
Cover the three review findings: a consolidation-style resume with a whole-parameter owner and an empty non-owner must gate the conditional initialize header on one group-wide reuse decision and then complete the resume step collectively with restored periods intact; a collective-free move_state_ on one member only must lead every member to the same scope re-validation decision on the next step; and a prepared canonical import must go stale when a bound parameter's storage is retargeted or mutated between prepare and commit, while an undisturbed prepare/commit round still succeeds.
…k scope GefenMuonHybrid.step decided the GradScaler overflow skip with the rank-local _amp_prepare_optimizer_step and raised structural gradient preflight errors rank-locally, even when post_sharding installed one multi-member codebook process-group binding shared by both children. A rank whose found_inf was set (or whose local gradients were malformed) then skipped or raised alone while its peers entered the children's scoped step collectives, hanging the group. Route both decisions through the children's scoped protocol before any child steps: the composite runs Gefen._prepare_scoped_amp_optimizer_step on itself (found_inf/grad_scale agreement is validated collectively and an overflow skip is a group-wide decision, entered and exited symmetrically on every member) and synchronizes the structural preflight through the children's _synchronize_codebook_scope_failure so every scope member raises together, keeping the atomic both-children-skip semantics. Without a binding the local behavior is unchanged.
step() re-ran the complete O(params x world) finalized-layout forensic rebuild on every guard call (2 passes per unscoped step, up to 7 under a multi-member codebook scope) and recomputed the manifest sha256 fingerprint inside every scoped operation header, costing seconds of host time per step at large scale. Cache one forensic verdict as an O(local params) identity-token snapshot (finalized registries by object identity, every live group container, parameter and compatibility name, plus a version counter bumped by every legitimate mutating API), and compute the manifest shard set and digest once per finalized manifest at post_sharding. Steady-state step guards now reuse the verdict; checkpoint prepare/commit, canonical export/import, rebinding, state movement/offload, codebook initialize/refresh, scope re-validation, and contract readiness still run the full forensic rebuild. The offload step-readiness scan is deduped under the same scheme. Caches live in __slots__ so staged __dict__ copies and fail-before-mutation snapshots never see them. Public-container tampering (group params/param_names slots, state names, the name cache), including closure-time mutation, still fails the step guard before any state mutation; in-place edits inside the private registries move to detection at the next full-forensics boundary, as now documented in docs/optimizer_contracts.md.
The deep snapshot helper only identity-checked top-level optimizer __dict__ attributes, so a post_sharding regression that populated a live child's per-parameter registries (_param_names, _gefen_shard_bindings) in place before a later child raised would leave a half-populated dict that kept its object identity and slipped past every atomicity test. Capture the key set (by parameter identity) and cloned values of those registries in deep_state_snapshot and value-compare them in assert_deep_state_snapshot, handling optimizer types that do not stage them. Add hybrid-rebinding tests covering injected in-place adds and value changes to both registries.
Under active state offload the per-step readiness guard (_assert_state_offload_step_ready, run ~2x/step) reached _state_movement_rejection_reason, which always consulted the finalized layout with full=True. That forced an uncached full layout rebuild plus a manifest sha256 recompute on every step (two full passes and two digest recomputes per step), bypassing the per-step layout-forensics memoization the non-offload guards already use. The finalized binding layout is immutable across steps, so thread a require_full_layout flag through _state_offload_rejection_reason and _state_movement_rejection_reason. The per-step offload readiness path uses the memoized fast path (full=False, reusing _gefen_layout_forensics_verdict), while the move_state_ / _atomic_state_movement_supported boundaries, activation, load, and contract readiness keep the full rebuild. The per-tensor offload state scan (CPU tightness/dtype, pairwise disjointness, native-schema validation) is untouched and still runs on every step, so a token-preserving corruption of a later parameter's offloaded state is still caught before any earlier parameter mutates.
The fail-before-mutation snapshot compared raw bytes via view(torch.uint8),
which PyTorch rejects on a 0-dim tensor ("self.dim() cannot be 0 to view
Float as Byte"). Capturable device-resident scalar counters hit this, failing
the CUDA-only capturable canonical-import tests (which the CPU CI skips).
Flatten to 1-D before reinterpreting; shapes are already checked equal, so
both sides flatten identically.
…ed load atomicity) main now carries the fail-before-mutation load_state_dict work (Tier-1, #69), so it no longer needs to live in this convention branch. Merging main in subtracts that shared base from this PR's diff: the two load-atomicity test files become identical to main, and the native/hybrid staging split shows up only as this branch's convention-only decorations rather than as net-new code. Reconcile the atomicity methods, which had diverged across the two branches: * main has the review-fixed staging -- isolated foreign-AdamW backup staging (_stage_foreign_backup_load / _commit_foreign_backup_load) because torch AdamW.load_state_dict is not itself fail-before-mutation, plus the _run_load_state_dict_post_hooks split so a composite owner commits every child's raw state before dispatching any child's post-hook -- but none of the convention decorations; * this branch has the decorations (the _assert_finalized_binding_layout guards, the _stage_load_state_dict offload block, the _invalidate_layout_forensics_caches commit hook, and the allow_preinitialized_periods canonical validation) layered on the older, pre-review hybrid load. Keep both: this branch's decorations on top of main's review-fixed base. The foreign-backup helpers and the three-phase hybrid load body are byte-identical to main; gefen.py differs from this branch only by the post-hook-split addition. tests/test_native_load_atomicity.py + tests/test_hybrid_load_atomicity.py: 28 passed against the reconciled tree.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c467cfb28b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…ier-2 convention side) (#72) * Preserve legacy vmean loads on the rank-local sharded path (#70) * Preserve legacy vmean loads on the rank-local sharded path _unwrap_rank_local_sharded_checkpoint validated the payload with the strict default (allow_legacy_vmean_counter=False), so a rank-local (DTensor/FSDP) checkpoint carrying vmean without the separate vmean_step counter -- a pre-counter state the step-time resume path backfills from step -- was rejected at load. The native load path already opts into that tolerance (allow_legacy_vmean_counter=True); the rank-local path did not, so the two disagreed and an otherwise valid legacy resume failed only on the sharded path. This regressed against pre-load-atomicity behavior, which had no such check. Pass allow_legacy_vmean_counter=True at the rank-local call site to match the native path. This only relaxes the vmean-without-vmean_step case; a current checkpoint (which carries vmean_step) is unaffected, and the inverse vmean_step-without-vmean corruption check is unchanged. tests/test_cpu_step_checkpoint.py::test_rank_local_validator_tolerates_legacy_vmean_without_step pins the validator tolerance both ways (the strict default rejects the pre-counter payload; the tolerance the rank-local path now opts into accepts it). * Guard the rank-local unwrap path directly, not just the validator Add test_rank_local_unwrap_tolerates_and_backfills_legacy_vmean: a single-rank gloo world builds a rank-local (rank_local_dtensor_v2) checkpoint, drops the separate vmean_step counter to simulate a pre-counter state, loads it through _unwrap_rank_local_sharded_checkpoint (the call site the fix touches), and asserts the first resumed step backfills vmean_step from step. The existing validator test calls _validate_rank_local_states directly, so it stays green if the unwrap call site regresses to strict validation. This test fails without allow_legacy_vmean_counter=True on the rank-local path (verified: reverting the fix raises "block second moment is missing vmean_step" at load). CPU-only, no multiprocessing (~2s). * Synchronize pre-collective Muon step failures * Disable Dynamo tracing on the pre-collective sync wrappers Match the mainline fix (#71): the sharded sync wrappers (_synchronize_sharded_step_flag/_error/_control_range, _prepare_synchronized_amp_step) and their codebook-scope analogs (_synchronize_codebook_scope_failure, _prepare_scoped_amp_optimizer_step) all branch/raise on dist.all_reduce results, so they carry @torch._dynamo.disable like the sibling _step_failure_process_groups and _assert_sharded_grad_presence_consistent. Tracing-only; no runtime change (precollective + scoped-agreement tests still 8/8). * Close pre-collective sync gaps on the convention side Four distributed-correctness fixes (found in review), on top of re-pointing the codebook scope onto the mainline primitive: * Overlapping-mesh deadlock + wrong flag device: same _step_failure_process_- groups rewrite as the mainline PR -- sync over every participated mesh (deduped, ordered) with the shard device, mirroring _assert_sharded_grad_presence_consistent. Threads collective_device through _synchronize_step_control_range too. * Scoped Gefen AMP presence hang: plain Gefen.step gated _prepare_scoped_amp_optimizer_step() on local found_inf/grad_scale, so with a multi-member codebook scope a rank with those attributes entered the presence all-gather while a rank without them skipped it -- deadlock. Run the scoped AMP agreement on EVERY member of a multi-member scope, matching the unconditional Muon/Hybrid preflight. * Dropped post-closure binding recheck (the CI failure): the merged Tier-2 step rewrite lost the convention's _assert_finalized_binding_layout() recheck after closure() on the mesh (non-scoped) branch, so a closure that rebinds param_groups was not caught -- test_finalized_layout_guard_rechecks_after_- closure[muon] failed. Re-assert it inside the preflight try (caught + synchronized) so a closure rebind fails before mutation and on every member. Full convention suite: 1095 passed, 0 failed. * Preserve mesh dimension order + close hybrid preflight gaps Three more review fixes: * DeviceMesh dimension order (same as the mainline PR): mirror _assert_sharded_grad_presence_consistent -- dedup meshes by content key, sorted-key order, get_all_groups() dimension order per mesh -- instead of flattening and sorting all groups by name, which could reorder a 2-D HSDP/TP mesh's row/column groups per rank and lock-order deadlock the preflight. * Hybrid post-closure binding recheck (unscoped path): the binding-is-None branch stepped children without re-asserting _assert_finalized_binding_layout() after the closure/pre-hooks, so a rank-local closure swapping a same-shaped backup parameter passed the gradient scan and stepped under stale routing. The scoped path already rechecked; re-assert in the unscoped preflight try (caught + synchronized). * Hybrid capture readiness: the _gefen_hybrid_precollective_preflight marker suppresses the muon child's own _assert_codebook_capture_ready() guard, and the hybrid preamble only checked devices, so a hybrid capturing before its Muon codebook initialized could run host-driven codebook init during capture. Call the muon capture/codebook readiness guard in the hybrid preflight. Full convention suite with GPU (2x3090 NCCL): 1420 passed.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a114fcb0b6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/gefen/gefen.py (1)
9506-9532: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winClosure/preamble failures in
Gefen.step()are not synchronized before entering the scoped codebook collective — this can deadlock a multi-member scope.
_assert_capturable_if_capturing(),_assert_codebook_capture_ready(), andclosure()run unguarded here. If any of these raises on only one member of a multi-member_gefen_codebook_process_group(e.g., the closure itself fails asymmetrically, which the docstring explicitly says this API supports for codebook learning), that rank raises immediately while peer ranks proceed into the collective_validate_codebook_scope_operation_header("step")a few lines later — the peers block forever waiting for the failed rank to join.
GefenMuon.step()(this same PR,gefen_muon.pylines 3066-3102) already fixes exactly this pattern for its own scoped branch by wrapping the equivalent calls in try/except and synchronizing via_synchronize_codebook_scope_failure(local_preamble_error, "step preamble")before any collective.Gefen.step()was not updated to match, and no test exercises this scenario for plain Gefen (only for GefenMuon, pertests/test_codebook_scope_distributed.py's_closure_preamble_worker).🔒️ Proposed fix mirroring GefenMuon's synchronized preamble
self._assert_state_offload_step_ready() self._assert_finalized_binding_layout() self._assert_runtime_codebook_process_group() - self._assert_capturable_if_capturing() - self._assert_codebook_capture_ready() - - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() + + loss = None + try: + self._assert_capturable_if_capturing() + self._assert_codebook_capture_ready() + if closure is not None: + with torch.enable_grad(): + loss = closure() + local_preamble_error = None + except Exception as exc: + loss = None + local_preamble_error = exc + if self._gefen_codebook_process_group is not None: + self._synchronize_codebook_scope_failure( + local_preamble_error, "step preamble" + ) + elif local_preamble_error is not None: + raise local_preamble_error self._assert_state_offload_step_ready() self._assert_finalized_binding_layout() self._assert_runtime_codebook_process_group() try: _assert_optimizer_gradients_structurally_valid(self)Do you want me to open an issue to track adding an asymmetric-closure-failure regression test for plain Gefen (mirroring the existing GefenMuon
_closure_preamble_workertest)?🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/gefen/gefen.py` around lines 9506 - 9532, Synchronize all Gefen.step() preamble failures before any scoped codebook collective. Wrap the preamble assertions and optional closure execution in try/except, retain the exception as local_preamble_error, and call _synchronize_codebook_scope_failure(local_preamble_error, "step preamble") for scoped process groups; re-raise locally when no scoped group exists. Mirror GefenMuon.step() while preserving loss assignment and subsequent gradient preflight behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/gefen/hybrid.py`:
- Around line 1252-1259: Update hybrid.step around
_hybrid_codebook_process_group so failure-sync process_groups includes sharded
meshes owned by both the primary Muon child and the backup child, not only
self.muon. Preserve the existing binding guard and ensure any backup-side
preflight failure participates in the same collective synchronization scope.
---
Outside diff comments:
In `@src/gefen/gefen.py`:
- Around line 9506-9532: Synchronize all Gefen.step() preamble failures before
any scoped codebook collective. Wrap the preamble assertions and optional
closure execution in try/except, retain the exception as local_preamble_error,
and call _synchronize_codebook_scope_failure(local_preamble_error, "step
preamble") for scoped process groups; re-raise locally when no scoped group
exists. Mirror GefenMuon.step() while preserving loss assignment and subsequent
gradient preflight behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 683a82f4-ef20-429f-b682-4d65b61f3c34
📒 Files selected for processing (7)
src/gefen/gefen.pysrc/gefen/gefen_muon.pysrc/gefen/hybrid.pytests/test_codebook_scope_distributed.pytests/test_cpu_step_checkpoint.pytests/test_gefen_fsdp2_checkpoint.pytests/test_precollective_failure_sync.py
thad0ctor
changed the title
This was referenced Jul 15, 2026
* Harden convention codebook-scope preambles and contract validation Address a round of Codex + CodeRabbit findings on the slimmed optimizer convention (PR #67). All findings are convention-specific. Codebook-scope collective symmetry: - initialize_codebook() now captures the failure-vote binding and synchronizes a one-sided preamble failure (finalized layout / runtime binding / capture-readiness) through the scope before the scoped "initialize" operation header, so a member that fails its preamble no longer strands peers inside the header all_gather. - plain Gefen.step() rejects a closure that replaces or clears the runtime codebook binding between capture and the header: the captured-binding recheck raises and is synchronized through the captured scope, so every rank follows the same header collective instead of one rank skipping it. Checkpoint and contract validation: - CheckpointProcessGroupBinding no longer rejects a CUDA collective device on Gloo/MPI (both support CUDA tensors); such a device now falls through to the existing device-availability check. - Contract dataclasses enforce exact runtime types: StateField.checkpointed, StateVariant.initialized/migration_only, and CheckpointSupport.requires_collective/atomic_load must be real bools, and OptimizerContract.schema_version must be a real int, so truthy strings or bool/float look-alikes can no longer advertise false guarantees. - The hybrid finalized-layout fast token folds in the _state_param_owner registry contents (keys and values), so an in-place owner replacement that preserves the dict identity and length invalidates the cached verdict. Tests: - Snapshot per-device cache membership so a cache clear/removal is detected. - Verify the backup parameter is untouched in hybrid scoped-failure tests. - Guard the remaining Gloo-only distributed tests with the dist/Gloo skipif. - Documented that only ordinary nested GefenMuonHybrid.load_state_dict() reports atomic_load=False; portable composite imports keep the guarantee. - New gloo regression tests for the initialize-preamble and step group-swap synchronization, plus unit tests for the checkpoint device, contract type, and hybrid owner-token fixes. * Keep MPI checkpoint bindings CPU-only in collective-device validation Narrow the F4 relaxation: Gloo genuinely supports CUDA tensors, but MPI moves GPU tensors only when built CUDA-aware, which PyTorch cannot reliably detect. Reject a CUDA collective device on the MPI backend during validate_runtime() instead of deferring a backend error to the portable collective's all_gather/broadcast (Codex P2 on #76).
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/optimizer_contracts.md`:
- Line 44: Update the private-registry exception in the optimizer contract to
state that in-place replacements are deferred only for registries not included
in the relevant optimizer’s content token. Explicitly exclude
_state_param_owner, whose contents participate in the fast token and therefore
are detected on the next step; preserve the existing boundary behavior for other
unchanged-token private registries.
In `@src/gefen/gefen_muon.py`:
- Around line 3098-3110: Update the synchronized gradient preflight around
_assert_runtime_codebook_process_group to revalidate the current
_gefen_codebook_process_group against the captured scope_binding before
proceeding, matching plain Gefen behavior so closure-driven group replacement or
clearing is reported collectively before
_validate_codebook_scope_operation_header("step"). Extend the Muon group-swap
regression coverage to exercise this closure scenario.
In `@tests/test_hybrid_layout_cache.py`:
- Around line 174-180: Update the assertions in the test around
optimizer._state_param_owner to validate the replacement parameter and owner
separately: assert the tuple’s first member is rogue and its second member
remains child. Remove the tuple-versus-parameter identity check, while
preserving the mapping-length assertion.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: a2b4fb0a-aa89-46e9-95e6-99e8b8d33dc0
📒 Files selected for processing (20)
README.mdbenchmarks/microbench/bench_layout_guard.pydocs/optimizer_contracts.mdsrc/gefen/__init__.pysrc/gefen/checkpoint.pysrc/gefen/contracts.pysrc/gefen/gefen.pysrc/gefen/gefen_muon.pysrc/gefen/hybrid.pysrc/gefen/portable_runtime.pytests/_state_snapshot.pytests/test_canonical_state_cpu.pytests/test_checkpoint_binding.pytests/test_codebook_scope_cpu.pytests/test_codebook_scope_distributed.pytests/test_hybrid_layout_cache.pytests/test_hybrid_scoped_failure_protocol.pytests/test_layout_guard_cost.pytests/test_optimizer_contracts.pytests/test_scoped_collective_agreement_fixes.py
💤 Files with no reviewable changes (4)
- README.md
- src/gefen/init.py
- src/gefen/portable_runtime.py
- tests/test_canonical_state_cpu.py
🚧 Files skipped from review as they are similar to previous changes (8)
- tests/test_checkpoint_binding.py
- benchmarks/microbench/bench_layout_guard.py
- tests/test_scoped_collective_agreement_fixes.py
- src/gefen/checkpoint.py
- tests/test_hybrid_scoped_failure_protocol.py
- tests/test_codebook_scope_cpu.py
- src/gefen/hybrid.py
- src/gefen/gefen.py
Reconcile the hybrid pre-collective failure sync with #74 (merged to main): gefen_muon.py auto-merges (the sync helpers are now @staticmethod and the mesh scope is collected via _collect_sharded_failure_groups). In GefenMuonHybrid.step, keep the convention's codebook-scope binding branch (_synchronize_prevalidated_codebook_scope_failure) unchanged, and in the no-binding branch adopt #74's static UNION-scope sync (GefenMuon._synchronize_sharded_step_error / _prepare_synchronized_amp_step over self._step_failure_process_groups()), which also covers the muon=None backup-only case. Also picks up the v0.4.1 release bump.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 509f929267
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Mirror the plain-Gefen scoped step-preamble hardening into GefenMuon.step: recheck the runtime codebook process-group against the captured binding in the gradient preflight so a closure that clears or replaces it is caught and synchronized before any peer enters the scoped operation-header collective, instead of one rank silently skipping the header and stranding its peers. Add a GefenMuon group-swap regression (fails without the recheck: the swapping rank runs the step while its peer hangs in the all_gather). Fix a tautological assertion in the composite owner in-place-swap layout cache test: the registry value is a (parameter, child) tuple, so unpack and assert the replaced members instead of comparing the tuple to a parameter. Qualify the optimizer-contracts private-registry deferral note: in-place value replacement is deferred to a boundary only for registries not folded into the content/fast token; _state_param_owner contents do participate, so such a replacement is rejected by the next step guard.
Owner
Author
…-contracts # Conflicts: # README.md # src/gefen/__init__.py # src/gefen/gefen.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Deliverable 4 of the upstream-extraction plan: the remaining
model-support-pipelineconvention PR, rebased on top of the merged Tier 1 (load atomicity, #69/#70) and Tier 2 (pre-collective failure sync, #71/#72) and slimmed so offload/movement is out of its review scope.This is the additive convention core going to
main— Optimizer Contracts + Portable State I/O + DCP persistence (+ codebook scopes, rebinding, layout caching, numerical fixes). Per the plan these are inextricable from the contract and are not candidates for further extraction.What this adds over
maincontracts.py) and capability reporting.portable_*modules)._synchronize_step_failureprimitive.What this intentionally does not include (deferred to Tier 3)
move_state_,offload_state_,restore_state_,StateMovementProvider,StateOffloadProvider), docs, release-gate coverage, and tests.OptimizerCapabilities.atomic_state_movement/state_offloadfields are retained and hard-wiredFalse, so serialized contracts and older checkpoints still round-trip unchanged.Validation
portable_*,canonical,checkpoint, andrebindingmodules are byte-identical to the full convention..step(), andstate_dict()/load_state_dict()round-trip on CPU.GPU parity/distributed jobs have no CI runner and show as skipped, consistent with the other PRs in this series.
Summary by CodeRabbit
New Features
Bug Fixes