Slim the platform-optimizer convention to contracts, portable state, and DCP - #73
Closed
thad0ctor wants to merge 52 commits into
Closed
Slim the platform-optimizer convention to contracts, portable state, and DCP#73thad0ctor wants to merge 52 commits into
thad0ctor wants to merge 52 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.
tests/test_layout_guard_cost.py pins the guard-cost contract: zero full forensic passes in steady-state steps, exactly one after finalization or any mutating API, the manifest digest computed once per finalized manifest, warm-verdict detection of closure/public-container tampering, boundary detection of private-registry in-place tampering, offload scan dedupe, and the caches staying out of the public attribute namespace. benchmarks/microbench/bench_layout_guard.py builds a synthetic 512 member x 300 parameter (153,600 shard) manifest with no process group and compares old per-step guard cost (7 forensic passes + 2 digest computes scoped; 2 passes unscoped) against the warm guard sequence: 42.6s -> 171us scoped (~250,000x), 9.6s -> 170us unscoped, far above the required 100x.
The perf caching applied a cached identity-token verdict to the per-step offload readiness check. Its token captured the state and param-group containers by identity but not the per-parameter offloaded state tensors, which are legitimately replaced every step. A token-preserving in-place corruption of a later parameter's offloaded state therefore slipped past step entry and was only caught mid-step, after earlier parameters had already been updated and committed, breaking fail-before-mutation. The readiness scan is O(local params) and was never the layout-forensics cost the cache targeted, so drop the offload verdict cache and run the full scan before any parameter is staged. The layout-manifest digest cache is unchanged. Update the layout-guard tests to assert every-step scanning and add a CUDA regression that corrupts a later parameter and checks the first is left byte-for-byte untouched.
Split the two run-on paragraphs into readable sentences, state the step-vs-boundary detection split as an explicit consequence for integrators, and list the full-rebuild boundaries plainly. Correct the offload paragraph: the readiness scan now runs in full on every step rather than reusing a cached verdict, and it does not run at state movement (movement disables offload).
… step tokens GefenMuonHybrid.step() calls its finalized-layout guard twice per step and each call ran the full O(params) composite forensic rebuild -- routing, ownership, per-child manifest partitioning, and local-binding sort -- so a finalized hybrid paid the reconstruction Gefen/GefenMuon already avoid. Memoize one composite verdict as an O(local params) identity-token snapshot, mirroring the base scheme. The token captures every field the rebuild reads by identity (the finalized flags, the composite manifest/roles/local bindings/finalized slots/owner registry/codebook binding, the subopt list and each child's private manifest, local bindings, codebook binding and defaults, and every live child group['params'] container) and folds in each child's own already-cached fast-path verdict, so any child-level change the children can detect propagates automatically and any composite container replacement is caught directly. post_sharding -- the only hybrid API that reassigns the composite fields, and the one every rebind helper routes through -- bumps a version counter and clears the verdict. Contract readiness still runs the full rebuild (full=True), and the hybrid guards now accept the same full= kwarg the base scoped path passes through _assert_runtime_codebook_process_group. Detection stays before any mutation: an in-place child param slot swap, a swapped child binding, closure-time public-container tampering, and composite registry replacement all still raise at the step guard.
The import smoke derived optimizer names from __all__ and read __name__ on each, skipping only kernels and __version__. The contracts layer adds public integer schema constants (CONTRACT_SCHEMA_VERSION, CANONICAL_STATE_FORMAT_VERSION, PORTABLE_STATE_FORMAT_VERSION, IDENTITY_SCHEMA_VERSION), which have no __name__, so the step raised AttributeError. Resolve every export instead — which also exercises the lazy __getattr__ hooks — and fall back to repr for value exports.
The composite native load committed the muon child
(self.muon.load_state_dict) before the backup half was validated, then
recovered a backup rejection with a second full
self.muon.load_state_dict(snapshot) reload. That rollback could itself
raise (e.g. CUDA OOM re-staging every muon state tensor), masking the
real backup error and leaving a half-loaded hybrid (new muon, old
backup) -- unlike every other mutating entry point on this branch, which
stages fully then publishes through non-throwing swaps.
Give the composite genuine two-phase semantics: validate and stage BOTH
halves before publishing EITHER, so a rejection on either half leaves
both live children byte-for-byte untouched. Each Gefen child already
loads atomically; expose that as a reusable primitive by splitting
Gefen.load_state_dict into _prepare_load_state_dict (runs the load
pre-hooks and stages an isolated shadow, mutating nothing) and
_publish_load_state_dict (commits through the existing non-throwing dict
swaps and runs post-hooks). The hybrid stages the muon child, then:
* Gefen backup: stages the backup too (all rejection happens here,
before any mutation), then publishes both through non-throwing swaps;
* AdamW backup: torch's Optimizer.load_state_dict validates group
structure and casts every tensor before its single __setstate__, so
it is itself fail-before-mutation -- commit it first while the muon
child is only staged, then publish the muon swap (which cannot fail).
Chosen mechanism is process-group-safe: staging goes through
Gefen._stage_load_state_dict, which shallow-copies the child's __dict__
and swaps in fresh state containers -- it never deepcopies the live
optimizer, so a codebook process-group handle (or any live ProcessGroup)
is shared by reference, never duplicated. The removed snapshot path's
copy.deepcopy is gone (import dropped); we deliberately avoid deepcopy of
a child or its process groups.
tests/test_hybrid_load_atomicity.py pins the contract: a backup
rejection (mismatched param-group count / foreign schema) must not reload
the muon child (byte-for-byte and object-identity unchanged), and an
original backup error must surface over a failing rollback reload rather
than be masked. Covers Gefen and AdamW backups; all three fail on the
pre-fix code.
This reverts commit 993ec96.
The state-offload and state-movement features are exercised only by CUDA-gated tests, which skip in the CPU-only push/PR CI. Their fail-before- mutation and device-transfer invariants were therefore verified by no gate. Add tests/test_state_offload.py and tests/test_state_movement.py to the mandatory two-GPU release gate list (72 passed, 0 skips on two devices), so the release gate — not a hosted CI runner — covers them.
GefenMuon.step and GefenMuonHybrid.step ran the capturable-readiness checks, the instance step pre-hooks, and the closure BEFORE any scope synchronization. Because capture readiness depends on rank-local gradients and the closure is user code, any of these can fail on only a subset of scope members. When it did, the failing member exited step() while its peers proceeded into the scoped operation-header all_gather, the composite preflight synchronization, or a child's scoped step collectives -- and hung. Capture these preamble failures and route them through _synchronize_codebook_scope_failure on the shared codebook binding before any member enters a later scoped collective, so the failure raises on every member together (mirroring the existing gradient-preflight synchronization). The finalized-layout and runtime-process-group guards stay local: they establish the very binding used to synchronize. Add asymmetric two-rank gloo regressions (a rank-0 closure failure) for both GefenMuon and the hybrid; without the fix the peer strands in a scoped collective until the gloo application timeout. Also fail the hybrid scoped-failure harness when a worker hangs or exits nonzero so an unclean distributed shutdown can no longer pass silently.
A composite canonical-global save/load processes EVERY child, so the Hybrid may advertise a checkpoint topology/exactness guarantee only when both children support it. _hybrid_portable_contract_support unioned the children's same_topology, topology_changing, and topology_change_kinds, which could advertise a guarantee that one child does not honor. Intersect them instead. (This is the opposite of the per-routed-parameter TRAINING claims, which remain legitimately unioned.) Add a regression that a Hybrid with children of differing canonical-global support advertises only the intersection.
_portable_tensor_chunk_elements sized chunks on the element size alone, but the strided read path materializes an int64 linear index, its running quotient, and one int64 coordinate tensor per dimension. On a high-rank non-contiguous tensor the scratch oversubscribed the fixed clone budget and could OOM. Account for that per-element scratch in the chunk size, mirroring portable_wire._clone_tensor. Add a high-rank regression asserting the chunk stays within the budget.
The shared deep-snapshot helper compared tensors with torch.equal, which reports +0.0/-0.0 as equal and every NaN as unequal; compare dtype/layout/shape then raw bytes so a rejected transaction that flips a sign bit is caught and an unchanged NaN payload is not a false failure. Several atomicity snapshots retained mutable containers only by reference, so a failed transaction that mutated contents in place (preserving object identity) would still pass. Deep-snapshot and compare the contents: optimizer state / param groups / registries in the codebook-scope atomicity check, the hybrid _state_param_owner mapping, and the Gefen _gefen_shard_bindings registry and per-device caches. Also assert the rejected canonical-import commits are no-ops immediately after each RuntimeError -- before the recovery import -- so a fresh import can no longer mask a fail-before-mutation regression. The real implementation passes every strengthened assertion.
The manifest-digest recompute assertion was trivially satisfied by the deliberate compute earlier in the test, so a per-step recompute regression would still pass. Warm past the first step (which validates the layout fully once, a separately guarded behavior) and then require the computation count to stay fixed across further steps and a cache-backed fingerprint read.
…test _free_port() released the probed socket before the workers bound it, so another process could steal the port and flake CI. Rendezvous over a temporary file:// init_method that stays valid until every rank initializes, mirroring the DCP distributed test.
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.
…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.
📝 WalkthroughWalkthroughChangesThe PR adds immutable optimizer contracts and sharding identities, canonical local state handling, topology-neutral portable state serialization, tensor-only DCP persistence, hybrid composite checkpoint transactions, scoped collective validation, finalized-layout caching, and extensive CPU/distributed test coverage. Portable optimizer integration
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
thad0ctor
marked this pull request as draft
July 15, 2026 18:50
Owner
Author
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 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 133: Clarify the “Hybrid composite loads” statement in the atomic-load
discussion to distinguish ordinary nested native Hybrid loading from portable
composite imports. Preserve the portable Hybrid guarantee described earlier, and
limit atomic_load=False to the ordinary nested native loader if that is the
intended scope.
In `@src/gefen/checkpoint.py`:
- Around line 118-129: The _validate_collective_device method should retain the
NCCL→CUDA validation but remove the Gloo/MPI CPU-only restriction, allowing Gloo
with CUDA bindings to reach the later CUDA-device validation. Update tests in
test_checkpoint_binding.py to cover rejection when the configured CUDA device is
unavailable.
In `@src/gefen/contracts.py`:
- Around line 778-802: Enforce exact runtime types for contract flags and schema
versions across StateField.__post_init__ and the related contract validation
blocks: validate boolean flags with type(value) is bool, and validate
schema_version with type(value) is int so booleans and floats are rejected.
Preserve the existing validation and normalization behavior for valid values.
In `@src/gefen/gefen.py`:
- Around line 8562-8575: Validate that the closure-time
_gefen_codebook_process_group still matches the captured scope_binding before
invoking _validate_codebook_scope_operation_header("step"). Reject replacement
with None or a different binding through the existing synchronized scope-failure
path, ensuring every rank follows the same header collective while preserving
the current preflight checks.
In `@src/gefen/hybrid.py`:
- Around line 986-1003: Update the fast-path token construction in
_hybrid_child_param_group_tokens to include the _state_param_owner mapping’s
ownership entries, not only its identity and length, so in-place replacements
invalidate the cached verdict. Add a regression test covering replacement of one
registry entry while preserving the dictionary length and verify the full
validation path rejects stale ownership.
In `@tests/_state_snapshot.py`:
- Around line 85-112: Update _tensor_value_pairs and the related
snapshot/comparison paths to record mutable cache membership, not just tensor
values. Specifically include the contents and keys of _gefen_codebook_by_device
(and any corresponding cache structures covered at the noted locations) so
clear() and entry removal produce a mismatch even when retained tensor
references and clones are unchanged; preserve existing tensor snapshot behavior.
In `@tests/test_codebook_scope_distributed.py`:
- Around line 786-787: Guard
test_explicit_gloo_scope_aggregates_logical_state_and_fails_atomically and the
related subgroup and zero-length-shard Gloo tests with the existing dist/Gloo
skipif marker used by the closure tests. Ensure these tests skip when Gloo
support is unavailable rather than spawning workers, while preserving their
current behavior when supported.
In `@tests/test_hybrid_scoped_failure_protocol.py`:
- Around line 258-270: Update the _untouched helper and its callers in the
failure tests to accept or provide backup_shard, then verify backup_parameter
matches the corresponding slice of _backup_initial() using _bits_equal. Preserve
the existing checks for optimizer state and Muon parameters while ensuring
backup mutations cannot pass unnoticed.
🪄 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: 57a3b501-3cee-47a2-a080-30da3cf7911e
📒 Files selected for processing (58)
.github/workflows/ci.ymlCOMPATIBILITY.mdREADME.mdbenchmarks/microbench/bench_layout_guard.pydocs/optimizer_contracts.mdsrc/gefen/__init__.pysrc/gefen/canonical.pysrc/gefen/checkpoint.pysrc/gefen/codebook.pysrc/gefen/contracts.pysrc/gefen/gefen.pysrc/gefen/gefen_muon.pysrc/gefen/hybrid.pysrc/gefen/portable.pysrc/gefen/portable_collective.pysrc/gefen/portable_dcp.pysrc/gefen/portable_fields.pysrc/gefen/portable_hybrid.pysrc/gefen/portable_identity.pysrc/gefen/portable_runtime.pysrc/gefen/portable_schema.pysrc/gefen/portable_state.pysrc/gefen/portable_wire.pysrc/gefen/rebinding.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_rebinding.pytests/test_hybrid_scoped_failure_protocol.pytests/test_layout_guard_cost.pytests/test_muon_distributed_checkpoint_safety.pytests/test_muon_fsdp2_approx.pytests/test_muon_fsdp2_parity.pytests/test_optimizer_contracts.pytests/test_portable_collective.pytests/test_portable_dcp.pytests/test_portable_dcp_distributed.pytests/test_portable_dcp_nccl.pytests/test_portable_dcp_topologies.pytests/test_portable_fields.pytests/test_portable_hybrid.pytests/test_portable_hybrid_distributed.pytests/test_portable_hybrid_runtime.pytests/test_portable_identity.pytests/test_portable_runtime.pytests/test_portable_runtime_consensus.pytests/test_portable_runtime_distributed.pytests/test_portable_runtime_hardening.pytests/test_portable_schema.pytests/test_portable_state.pytests/test_portable_state_math.pytests/test_portable_wire.pytests/test_rebinding_cpu.pytests/test_scoped_collective_agreement_fixes.pytests/test_shard_identity_contracts.py
|
|
||
| Plain Gefen declares replicated, flattened element-shard, and the narrow DTensor training layouts. Its PyTorch rank-local checkpoint transport is same-topology only. GefenMuon declares replicated and narrow DTensor training, with mode-specific state extents: `approx` state is local, `exact` state is logically global, and `distributed` momentum is held by the parameter owner while non-owners retain metadata only. Native Parallel-Muon checkpoints separately declare world-size owner redistribution, not placement-changing resharding. `GefenMuonHybrid` retains its nested child namespaces and does not flatten AdamW or Gefen child state into a fabricated common schema. | ||
|
|
||
| Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Distinguish native Hybrid loads from portable composite imports.
Line 133 currently contradicts the portable Hybrid guarantee on Line 97 unless “Hybrid composite loads” is explicitly limited to the ordinary nested native loader.
Proposed clarification
-Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`.
+Ordinary nested `GefenMuonHybrid.load_state_dict()` loads do not yet provide the same guarantee and report `atomic_load=False`; portable composite imports retain the separate fail-before-mutation guarantee described above.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`. | |
| Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Ordinary nested `GefenMuonHybrid.load_state_dict()` loads do not yet provide the same guarantee and report `atomic_load=False`; portable composite imports retain the separate fail-before-mutation guarantee described above. |
🤖 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 `@docs/optimizer_contracts.md` at line 133, Clarify the “Hybrid composite
loads” statement in the atomic-load discussion to distinguish ordinary nested
native Hybrid loading from portable composite imports. Preserve the portable
Hybrid guarantee described earlier, and limit atomic_load=False to the ordinary
nested native loader if that is the intended scope.
Comment on lines
+118
to
+129
| def _validate_collective_device(self, backend: object) -> None: | ||
| backend_name = str(backend).lower() | ||
| if "nccl" in backend_name: | ||
| if self.collective_device.type != "cuda": | ||
| raise ValueError( | ||
| "checkpoint collective device is incompatible with the runtime backend" | ||
| ) | ||
| elif "gloo" in backend_name or "mpi" in backend_name: | ||
| if self.collective_device.type != "cpu": | ||
| raise ValueError( | ||
| "checkpoint collective device is incompatible with the runtime backend" | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== checkpoint.py excerpt ==\n'
sed -n '1,220p' src/gefen/checkpoint.py | cat -n
printf '\n== test_checkpoint_binding.py excerpt ==\n'
sed -n '1,220p' tests/test_checkpoint_binding.py | cat -n
printf '\n== search for collective_device validation and backend usage ==\n'
rg -n "collective_device|_validate_collective_device|gloo|nccl|mpi|broadcast|all_gather" src testsRepository: thad0ctor/Gefen-X
Length of output: 50373
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== PyTorch distributed docs snippet search (local repo references only) ==\n'
rg -n "Gloo|gloo|NCCL|nccl|all_gather|broadcast" README* docs src tests || trueRepository: thad0ctor/Gefen-X
Length of output: 24808
🌐 Web query:
PyTorch distributed Gloo GPU tensors all_gather broadcast documentation
💡 Result:
In PyTorch, the Gloo backend is primarily designed and optimized for CPU-based distributed training [1][2]. While Gloo does have support for GPU tensors, its performance is generally lower than the NCCL backend [1][3], which is the recommended choice for distributed training on NVIDIA GPUs [1][2]. Regarding the specific collective operations you mentioned: 1. Broadcast: The Gloo backend supports the broadcast operation for both CPU and GPU tensors [1][4]. 2. All_gather: The Gloo backend supports the all_gather operation for both CPU and GPU tensors [1][4]. Key Considerations for Using Gloo with GPUs: - Performance: NCCL is significantly more efficient for GPU-to-GPU communication, particularly on high-bandwidth interconnects like NVLink or InfiniBand [3]. Gloo is typically used as a fallback or in environments where NCCL is not available or supported (e.g., heterogeneous clusters or non-NVIDIA GPUs) [1][3][2]. - Use Case: If you are performing distributed training on NVIDIA GPUs, you should prioritize the NCCL backend [1][2]. If you encounter issues with NCCL, Gloo serves as a functional, though slower, alternative for most standard collective operations [1][4]. - Documentation Accuracy: Be aware that official documentation tables may occasionally contain inaccuracies regarding backend support for certain niche operations (e.g., all_to_all) [5][6]. Always verify specific, complex collective behaviors with your current PyTorch version, as support can evolve [5][6]. For the most reliable GPU distributed performance, standard practice is to initialize your process group with the NCCL backend when CUDA is available [2]. Source Summary: The technical capabilities of backends are outlined in the official PyTorch distributed documentation [7][1]. Performance benchmarks and comparative analysis emphasize the efficiency gap between NCCL and Gloo for GPU workloads [3]. Tracking of specific feature support for backends is maintained via official PyTorch documentation and community issue reports [5][6].
Citations:
- 1: https://docs.pytorch.org/docs/stable/distributed.md
- 2: https://runebook.dev/en/docs/pytorch/distributed/torch.distributed.Backend
- 3: https://markaicode.com/benchmarks/pytorch-distributed-training-benchmark/
- 4: https://docs.pytorch.org/docs/stable/distributed
- 5: Torch Distributed documentation falsely claims that GLOO supports all_to_all pytorch/pytorch#162248
- 6: Add
gloosupport forall_to_allpytorch/pytorch#88245 - 7: https://docs.pytorch.org/docs/stable/distributed.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("src/gefen/checkpoint.py")
print(p.exists(), p)
PYRepository: thad0ctor/Gefen-X
Length of output: 184
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "checkpoint.py"
nl -ba src/gefen/checkpoint.py | sed -n '90,170p'
echo
echo "test_checkpoint_binding.py"
nl -ba tests/test_checkpoint_binding.py | sed -n '90,160p'Repository: thad0ctor/Gefen-X
Length of output: 209
🌐 Web query:
site:pytorch.org/docs/stable/distributed gloo cuda tensors all_gather broadcast
💡 Result:
In PyTorch, the Gloo backend is designed primarily for CPU-based distributed operations and is not the recommended or optimal choice for CUDA tensor collective communications like all_gather or broadcast [1]. While PyTorch's distributed framework may allow the Gloo backend to handle operations with CUDA tensors in specific configurations or as a fallback, it is generally intended for CPU tensors [1][2]. The NCCL backend is the standard and highly optimized choice for distributed training and collective communications using CUDA tensors [1]. According to official PyTorch documentation: 1. The Gloo backend is specifically recommended for CPU-based distributed training [1]. 2. For CUDA GPU-based operations, the NCCL backend provides significantly better performance and is the standard practice [1]. 3. In modern PyTorch distributed configurations (specifically when using default settings), the system often automatically routes collectives with CPU tensors to Gloo and collectives with CUDA tensors to NCCL [1][2]. If you attempt to use Gloo for CUDA tensors, you may encounter performance limitations or functional restrictions, as NCCL is purpose-built to leverage NVIDIA GPU interconnects (such as NVLink and InfiniBand) effectively [1].
Citations:
- 1: https://docs.pytorch.org/docs/stable/distributed
- 2: https://docs.pytorch.org/docs/stable/distributed.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== src/gefen/portable_dcp.py validation =="
sed -n '1,160p' src/gefen/portable_dcp.py | cat -n
echo
echo "== src/gefen/portable_collective.py device transfer / transport =="
sed -n '150,220p' src/gefen/portable_collective.py | cat -n
sed -n '380,420p' src/gefen/portable_collective.py | cat -n
echo
echo "== src/gefen/codebook.py binding init =="
sed -n '1,120p' src/gefen/codebook.py | cat -nRepository: thad0ctor/Gefen-X
Length of output: 14211
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== checkpoint-binding related tests and gloo/cuda expectations =="
python3 - <<'PY'
from pathlib import Path
for path in [Path("tests/test_checkpoint_binding.py"), Path("tests/test_gefen_fsdp2_checkpoint.py"), Path("tests/test_portable_dcp_nccl.py"), Path("tests/test_codebook_scope_distributed.py")]:
if path.exists():
text = path.read_text()
print(f"\n## {path}")
for needle in ["cuda:0", "collective_device", "runtime backend", "gloo", "cpu"]:
if needle in text:
print(f"contains: {needle}")
PYRepository: thad0ctor/Gefen-X
Length of output: 580
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for p in ["src/gefen/portable_dcp.py", "src/gefen/codebook.py"]:
print(f"\n## {p}")
text = Path(p).read_text().splitlines()
for i, line in enumerate(text, 1):
if "collective_device" in line or "backend" in line or "nccl" in line or "gloo" in line or "mpi" in line:
start = max(1, i-8)
end = min(len(text), i+12)
for j in range(start, end+1):
print(f"{j:4d}: {text[j-1]}")
print("----")
break
PYRepository: thad0ctor/Gefen-X
Length of output: 2694
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for p in ["src/gefen/portable_dcp.py", "src/gefen/codebook.py"]:
lines = Path(p).read_text().splitlines()
print(f"\n## {p}")
for idx, line in enumerate(lines, 1):
if "if \"nccl\" in backend" in line or "elif \"gloo\" in backend" in line or "collective_device" in line:
for j in range(max(1, idx-5), min(len(lines), idx+15)+1):
print(f"{j:4d}: {lines[j-1]}")
print("----")
PYRepository: thad0ctor/Gefen-X
Length of output: 10471
Allow CUDA bindings on Gloo. This guard rejects valid Gloo/CUDA checkpoint bindings before the later CUDA-device check runs. Keep NCCL→CUDA, but drop the Gloo CPU-only restriction and update tests/test_checkpoint_binding.py to cover the unavailable CUDA-device case instead.
🤖 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/checkpoint.py` around lines 118 - 129, The
_validate_collective_device method should retain the NCCL→CUDA validation but
remove the Gloo/MPI CPU-only restriction, allowing Gloo with CUDA bindings to
reach the later CUDA-device validation. Update tests in
test_checkpoint_binding.py to cover rejection when the configured CUDA device is
unavailable.
Comment on lines
+778
to
+802
| name: str | ||
| scope: StateScope | ||
| geometry: StateGeometry | ||
| checkpointed: bool | ||
| key_match: StateKeyMatch = StateKeyMatch.EXACT | ||
| applicable_sharded_modes: AbstractSet[str] = frozenset() | ||
| description: str = "" | ||
| optional: bool = False | ||
|
|
||
| def __post_init__(self) -> None: | ||
| object.__setattr__( | ||
| self, | ||
| "applicable_sharded_modes", | ||
| _frozenset(self.applicable_sharded_modes), | ||
| ) | ||
| if not self.name: | ||
| raise ValueError("StateField.name must be non-empty") | ||
| if not isinstance(self.scope, StateScope): | ||
| raise TypeError("StateField.scope must be a StateScope") | ||
| if not isinstance(self.geometry, StateGeometry): | ||
| raise TypeError("StateField.geometry must be a StateGeometry") | ||
| if not isinstance(self.key_match, StateKeyMatch): | ||
| raise TypeError("StateField.key_match must be a StateKeyMatch") | ||
| if type(self.optional) is not bool: | ||
| raise TypeError("StateField.optional must be a bool") |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Enforce exact types for contract flags and schema versions.
Values such as atomic_load="false" remain truthy and falsely advertise restore guarantees. Likewise, schema_version=True or 1.0 is currently accepted as version 1.
Proposed fix
class StateField:
def __post_init__(self) -> None:
+ if type(self.checkpointed) is not bool:
+ raise TypeError("StateField.checkpointed must be a bool")
...
if type(self.optional) is not bool:
raise TypeError("StateField.optional must be a bool") class StateVariant:
def __post_init__(self) -> None:
...
+ for name in ("initialized", "migration_only"):
+ if type(getattr(self, name)) is not bool:
+ raise TypeError("StateVariant.{} must be a bool".format(name)) class CheckpointSupport:
def __post_init__(self) -> None:
...
+ for name in ("requires_collective", "atomic_load"):
+ if type(getattr(self, name)) is not bool:
+ raise TypeError("CheckpointSupport.{} must be a bool".format(name)) class OptimizerContract:
def __post_init__(self) -> None:
...
- if self.schema_version != CONTRACT_SCHEMA_VERSION:
+ if (
+ type(self.schema_version) is not int
+ or self.schema_version != CONTRACT_SCHEMA_VERSION
+ ):Also applies to: 822-870, 962-997, 1093-1102
🤖 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/contracts.py` around lines 778 - 802, Enforce exact runtime types
for contract flags and schema versions across StateField.__post_init__ and the
related contract validation blocks: validate boolean flags with type(value) is
bool, and validate schema_version with type(value) is int so booleans and floats
are rejected. Preserve the existing validation and normalization behavior for
valid values.
Comment on lines
+8562
to
+8575
| try: | ||
| self._assert_finalized_binding_layout() | ||
| self._assert_runtime_codebook_process_group() | ||
| _assert_optimizer_gradients_structurally_valid(self) | ||
| local_preflight_error = None | ||
| except Exception as exc: | ||
| local_preflight_error = exc | ||
| if scope_binding is not None: | ||
| self._synchronize_prevalidated_codebook_scope_failure( | ||
| local_preflight_error, "gradient preflight", scope_binding | ||
| ) | ||
| elif local_preflight_error is not None: | ||
| raise local_preflight_error | ||
| self._validate_codebook_scope_operation_header("step") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Reject closure-time scope-binding replacement before the operation header.
A closure can replace _gefen_codebook_process_group on one rank with None or another valid binding. Local preflight then succeeds, but that rank skips or enters a different header collective while peers use the captured group, causing a deadlock.
Proposed fix
try:
+ if self._gefen_codebook_process_group is not scope_binding:
+ raise RuntimeError(
+ "Gefen codebook process-group binding changed during the step preamble"
+ )
self._assert_finalized_binding_layout()
self._assert_runtime_codebook_process_group()
_assert_optimizer_gradients_structurally_valid(self)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| try: | |
| self._assert_finalized_binding_layout() | |
| self._assert_runtime_codebook_process_group() | |
| _assert_optimizer_gradients_structurally_valid(self) | |
| local_preflight_error = None | |
| except Exception as exc: | |
| local_preflight_error = exc | |
| if scope_binding is not None: | |
| self._synchronize_prevalidated_codebook_scope_failure( | |
| local_preflight_error, "gradient preflight", scope_binding | |
| ) | |
| elif local_preflight_error is not None: | |
| raise local_preflight_error | |
| self._validate_codebook_scope_operation_header("step") | |
| try: | |
| if self._gefen_codebook_process_group is not scope_binding: | |
| raise RuntimeError( | |
| "Gefen codebook process-group binding changed during the step preamble" | |
| ) | |
| self._assert_finalized_binding_layout() | |
| self._assert_runtime_codebook_process_group() | |
| _assert_optimizer_gradients_structurally_valid(self) | |
| local_preflight_error = None | |
| except Exception as exc: | |
| local_preflight_error = exc | |
| if scope_binding is not None: | |
| self._synchronize_prevalidated_codebook_scope_failure( | |
| local_preflight_error, "gradient preflight", scope_binding | |
| ) | |
| elif local_preflight_error is not None: | |
| raise local_preflight_error | |
| self._validate_codebook_scope_operation_header("step") |
🤖 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 8562 - 8575, Validate that the closure-time
_gefen_codebook_process_group still matches the captured scope_binding before
invoking _validate_codebook_scope_operation_header("step"). Reject replacement
with None or a different binding through the existing synchronized scope-failure
path, ensuring every rank follows the same header collective while preserving
the current preflight checks.
Comment on lines
+986
to
+1003
| live = [ | ||
| getattr(self, "_hybrid_layout_version", 0), | ||
| self._hybrid_post_sharding_finalized, | ||
| self._hybrid_sharding_manifest, | ||
| self._hybrid_local_shard_bindings, | ||
| self._hybrid_fqn_roles, | ||
| self._hybrid_finalized_slots, | ||
| self._hybrid_codebook_process_group, | ||
| self._state_param_owner, | ||
| len(self._state_param_owner), | ||
| self.defaults, | ||
| self._subopts, | ||
| len(self._subopts), | ||
| ] | ||
| for child in self._subopts: | ||
| live.append(child._finalized_binding_layout_matches()) | ||
| GefenMuonHybrid._hybrid_child_param_group_tokens(child, live) | ||
| return tuple(live) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Include ownership entries in the fast-path token.
An in-place _state_param_owner replacement that preserves the dictionary identity and length reuses the cached True verdict. The full rebuild would reject it, but is never reached, allowing stale or incorrect state routing.
Proposed fix
live = [
...
self._state_param_owner,
len(self._state_param_owner),
...
]
+ for parameter_id in sorted(self._state_param_owner):
+ live.extend(
+ (parameter_id, self._state_param_owner[parameter_id])
+ )
for child in self._subopts:Add a warm-cache regression test that replaces one registry entry in place without changing its length.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| live = [ | |
| getattr(self, "_hybrid_layout_version", 0), | |
| self._hybrid_post_sharding_finalized, | |
| self._hybrid_sharding_manifest, | |
| self._hybrid_local_shard_bindings, | |
| self._hybrid_fqn_roles, | |
| self._hybrid_finalized_slots, | |
| self._hybrid_codebook_process_group, | |
| self._state_param_owner, | |
| len(self._state_param_owner), | |
| self.defaults, | |
| self._subopts, | |
| len(self._subopts), | |
| ] | |
| for child in self._subopts: | |
| live.append(child._finalized_binding_layout_matches()) | |
| GefenMuonHybrid._hybrid_child_param_group_tokens(child, live) | |
| return tuple(live) | |
| live = [ | |
| getattr(self, "_hybrid_layout_version", 0), | |
| self._hybrid_post_sharding_finalized, | |
| self._hybrid_sharding_manifest, | |
| self._hybrid_local_shard_bindings, | |
| self._hybrid_fqn_roles, | |
| self._hybrid_finalized_slots, | |
| self._hybrid_codebook_process_group, | |
| self._state_param_owner, | |
| len(self._state_param_owner), | |
| self.defaults, | |
| self._subopts, | |
| len(self._subopts), | |
| ] | |
| for parameter_id in sorted(self._state_param_owner): | |
| live.extend( | |
| (parameter_id, self._state_param_owner[parameter_id]) | |
| ) | |
| for child in self._subopts: | |
| live.append(child._finalized_binding_layout_matches()) | |
| GefenMuonHybrid._hybrid_child_param_group_tokens(child, live) | |
| return tuple(live) |
🤖 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/hybrid.py` around lines 986 - 1003, Update the fast-path token
construction in _hybrid_child_param_group_tokens to include the
_state_param_owner mapping’s ownership entries, not only its identity and
length, so in-place replacements invalidate the cached verdict. Add a regression
test covering replacement of one registry entry while preserving the dictionary
length and verify the full validation path rejects stale ownership.
Comment on lines
+85
to
+112
| def _tensor_value_pairs(optimizer): | ||
| """Every tensor reachable from ``optimizer.__dict__`` with a clone of it.""" | ||
|
|
||
| pairs = [] | ||
| visited = set() | ||
|
|
||
| def visit(value): | ||
| if torch.is_tensor(value): | ||
| if id(value) not in visited: | ||
| visited.add(id(value)) | ||
| pairs.append((value, value.detach().clone())) | ||
| elif isinstance(value, dict): | ||
| if id(value) in visited: | ||
| return | ||
| visited.add(id(value)) | ||
| for key, item in value.items(): | ||
| visit(key) | ||
| visit(item) | ||
| elif isinstance(value, (list, tuple, set, frozenset)): | ||
| if id(value) in visited: | ||
| return | ||
| visited.add(id(value)) | ||
| for item in value: | ||
| visit(item) | ||
|
|
||
| for value in optimizer.__dict__.values(): | ||
| visit(value) | ||
| return tuple(pairs) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Snapshot cache membership, not only reachable tensor bytes.
An in-place cache.clear() or entry removal from _gefen_codebook_by_device preserves attribute identity, while the retained tensor reference still matches its clone. The atomicity assertion therefore false-passes. Snapshot and compare mutable cache contents explicitly.
Proposed fix
_CHILD_REGISTRY_ATTRS = ("_param_names", "_gefen_shard_bindings")
+_CHILD_CACHE_ATTRS = (
+ "_gefen_codebook_by_device",
+ "_gefen_codebook_lut_by_device",
+ "_sr_seed_by_device",
+ "_gefen_global_step_by_device",
+)
+def _cache_snapshot(optimizer):
+ caches = {}
+ for name in _CHILD_CACHE_ATTRS:
+ cache = getattr(optimizer, name, None)
+ if type(cache) is dict:
+ caches[name] = (cache, _cloned(cache))
+ return caches
+
def deep_state_snapshot(optimizer):
return {
"attributes": optimizer.__dict__.copy(),
+ "caches": _cache_snapshot(optimizer),
"registries": _registry_snapshot(optimizer),
...
}
def assert_deep_state_snapshot(optimizer, snapshot):
...
+ for name, (expected_ref, expected_contents) in snapshot["caches"].items():
+ live = getattr(optimizer, name)
+ assert live is expected_ref
+ _nested_equal(live, expected_contents)Also applies to: 136-156, 185-186
🤖 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 `@tests/_state_snapshot.py` around lines 85 - 112, Update _tensor_value_pairs
and the related snapshot/comparison paths to record mutable cache membership,
not just tensor values. Specifically include the contents and keys of
_gefen_codebook_by_device (and any corresponding cache structures covered at the
noted locations) so clear() and entry removal produce a mismatch even when
retained tensor references and clones are unchanged; preserve existing tensor
snapshot behavior.
Comment on lines
+786
to
+787
| def test_explicit_gloo_scope_aggregates_logical_state_and_fails_atomically(): | ||
| results = _run_workers() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant test file and nearby guard patterns.
git ls-files tests/test_codebook_scope_distributed.py
echo "--- outline ---"
ast-grep outline tests/test_codebook_scope_distributed.py --view expanded || true
echo "--- guard/search ---"
rg -n "skipif|is_gloo_available|is_available|closure" tests/test_codebook_scope_distributed.py
echo "--- relevant ranges ---"
sed -n '740,820p' tests/test_codebook_scope_distributed.py
echo "--- more ranges ---"
sed -n '900,960p' tests/test_codebook_scope_distributed.py
echo "--- more ranges 2 ---"
sed -n '1090,1145p' tests/test_codebook_scope_distributed.pyRepository: thad0ctor/Gefen-X
Length of output: 11706
Guard the Gloo-only tests with skipif. These entry points spawn Gloo workers unconditionally, so builds without dist/Gloo support fail instead of skipping. Add the same skipif used by the closure tests to the subgroup and zero-length-shard cases as well.
🤖 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 `@tests/test_codebook_scope_distributed.py` around lines 786 - 787, Guard
test_explicit_gloo_scope_aggregates_logical_state_and_fails_atomically and the
related subgroup and zero-length-shard Gloo tests with the existing dist/Gloo
skipif marker used by the closure tests. Ensure these tests skip when Gloo
support is unavailable rather than spawning workers, while preserving their
current behavior when supported.
Comment on lines
+258
to
+270
| def _untouched(optimizer, muon_parameter, backup_parameter): | ||
| return ( | ||
| optimizer.muon._gefen_global_step == 0 | ||
| and optimizer.backup._gefen_global_step == 0 | ||
| and optimizer.muon._gefen_codebook is None | ||
| and optimizer.backup._gefen_codebook is None | ||
| and (muon_parameter is None or _bits_equal(muon_parameter, _muon_initial())) | ||
| and ( | ||
| muon_parameter is None | ||
| or optimizer.muon.state[muon_parameter] == {"name": "matrix"} | ||
| ) | ||
| and optimizer.backup.state[backup_parameter] == {"name": "vector"} | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Verify the backup parameter remains untouched.
backup_parameter is never inspected, so these failure tests can pass even if the backup child mutates its weights. Compare it with the corresponding slice of _backup_initial() by passing backup_shard, or snapshot it before step().
Proposed fix
-def _untouched(optimizer, muon_parameter, backup_parameter):
+def _untouched(optimizer, muon_parameter, backup_parameter, backup_shard):
+ start = backup_shard.logical_slice.flat_offset
+ stop = start + backup_shard.logical_slice.length
return (
...
+ and _bits_equal(backup_parameter, _backup_initial()[start:stop])
)Update the callers to pass backup_shard.
🤖 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 `@tests/test_hybrid_scoped_failure_protocol.py` around lines 258 - 270, Update
the _untouched helper and its callers in the failure tests to accept or provide
backup_shard, then verify backup_parameter matches the corresponding slice of
_backup_initial() using _bits_equal. Preserve the existing checks for optimizer
state and Muon parameters while ensuring backup mutations cannot pass unnoticed.
This was referenced Jul 15, 2026
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 slimmed platform-optimizer convention, rebased on top of the merged Tier 1 (load atomicity, #69/#70) and Tier 2 (pre-collective failure sync, #71/#72).
This supersedes #67. It lands the full optimizer-integration convention minus the state-movement/offload feature, which is deferred to Tier 3.
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
Documentation
Performance