Add DCP-shardable FSDP2 optimizer state and portable resharding - #82
Closed
thad0ctor wants to merge 57 commits into
Closed
Add DCP-shardable FSDP2 optimizer state and portable resharding#82thad0ctor 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.
… 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.
* 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).
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.
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.
…-resharding-main # Conflicts: # src/gefen/gefen.py
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
This was referenced Jul 16, 2026
Owner
Author
|
Closing this draft: it incorrectly bundled the full #67 convention stack and is far larger than the intended standalone issue-81 extraction. I am replacing it with a minimal main-based PR containing only the FSDP2 DCP-shardable/resharding implementation and directly required tests/docs. |
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.
Summary
main, including the reviewed optimizer-contract convention and the DTensor/DCP resharding data plane.torch.distributed.checkpointsave/load test that saves on two ranks and resumes on four ranks without gathering the complete optimizer state onto rank 0.Why
Gefen's existing FSDP2 optimizer checkpoint format stores variable-size rank-local blobs. DCP plans loads from tensor shapes, so a fresh optimizer template cannot load a warmed blob and fails with a size mismatch. The old format also cannot reshard optimizer state across a changed world size.
This change adds a shard-addressable portable representation while retaining the existing native same-topology checkpoint path and failing closed for unsupported layouts and transitions.
PR lineage
This PR deliberately supersedes:
The implementation was rebuilt on current
mainafter #80, then reconciled with the later scoped and unscoped collective-failure hardening.Closes #81.
Validation
git diff --check: clean.GPU allocation used only the permitted idle devices: two-GPU tests on PCI indices
2,4; four-GPU test on2,4,5,1.