Compat: DCP resharding for FSDP2 Gefen state - #83
Merged
Conversation
|
Caution Review failedAn error occurred during the review process. Please try again later. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Address the PR review on GefenDCPState. - Load path re-blocks instead of collapsing to period one. After a DCP reshard, re-derive Gefen's compact per-rank state from the resharded dense fp32 momentum: re-run the block-variance period search on the dense second moment (a grad^2 proxy), relearn the exact codebook on the new local shard, and re-quantize into blocks (per-block max-abs magnitude, nearest codeword indices). vmean is reconstructed as the per-block mean of the dense second moment for the new blocking. This restores the ~1 byte/param footprint that period-one re-encoding destroyed (e.g. a 2048x vmean/m_magnitude blow-up). Re-blocking happens before optimizer.state.clear(), preserving the fail-atomic load. - Fail closed on capturable=True and factored_v_2d=True: capturable (and compiled) counter/seed semantics are not host-serializable, and the factored row/column second moment is not shard-addressable. - Replace the tautological period-one tests with real ones: a genuine Gefen.step learns the codebook and picks period > 1; save, reshard (2->2, 2->4, 4->2, GPU/NCCL 2->4, fully_shard), and assert the restored state stays compact and continues within quantization-noise tolerance of a native run at the target topology. Add capturable/factored guard tests. - Docs (COMPATIBILITY, README, CHANGELOG): drop the "exact period-one" overclaim; state that GefenDCPState is for resharding, that resume is a correct continuation within quantization noise (not bit-exact), steer same-topology users to the native path, and note the compact re-block and the factored_v_2d/capturable requirements.
thad0ctor
changed the title
This was referenced Jul 16, 2026
GefenDCPState only reads optimizer state (m_codebook / m_magnitude / vmean / automatic_period), which is layout-identical whether the step ran through the fused CUDA kernels or the decomposed path, so the adapter is fused-agnostic -- but every test built with fused=False. Thread a fused flag through the test harness (default False, so existing tests are unchanged) and add a GPU/NCCL test that trains with fused=True and asserts the fused-produced state saves, reshards, stays compact, and continues within tolerance.
thad0ctor
marked this pull request as ready for review
July 16, 2026 20:35
Owner
Author
Review summary (findings + resolution)Reviewed the initial adapter ( [HIGH] Load collapsed the block period P→1 — memory blow-up + non-bit-exact resume. The original load path re-encoded the resharded dense momentum as period-1 ( [MEDIUM] Tests were tautological. The originals hand-injected period-1 synthetic state into both the save side and the reference and never called a real [fail-closed] [docs] Overclaim. "projects the restored state into an exact period-one representation" mislead about resume fidelity, and the memory implication was undisclosed. Fused path. The adapter is fused-agnostic (it touches only state, which is layout-identical fused vs decomposed), but every test used Documented tradeoff (not a silent gap)Reshard resume is a correct continuation within ~256-level quantization noise, not bit-exact, even same-topology — inherent to routing momentum through a dense reshard + re-block against a freshly learned per-shard codebook. Same-topology users should use the native ValidationAll DCP tests green on idle RTX 3090s: 9 passed — same-topology + 2→4 + 4→2 reshards (gloo), a real 4-GPU NCCL 2→4 reshard, a ScopePlain |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bd7cfc60ab
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@CHANGELOG.md`:
- Line 9: Remove the “no public API changes” claim from the release summary near
the GefenDCPState entry, while preserving the existing description of the new
public GefenDCPState API and its resharding behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: c95918e4-2af2-4d74-932f-10efd89673eb
📒 Files selected for processing (6)
CHANGELOG.mdCOMPATIBILITY.mdREADME.mdsrc/gefen/__init__.pysrc/gefen/dcp.pytests/test_dcp_resharding.py
Address six correctness findings on the DCP resharding adapter and bump the on-disk schema to v2 (fail-closed on version/identity/topology mismatch): - Serialize and restore per-group hyperparameters (lr/betas/eps/weight_decay) so a resume after an LR-schedule or hyperparameter change adopts the checkpoint's values instead of silently keeping the freshly-constructed optimizer's, matching the native full-state path. - Persist stable per-slot parameter identities (name + group + global shape) and validate them before accepting a load, so a target that registered same-shaped parameters in a different order is rejected instead of cross-assigning each parameter the other's momentum/second-moment. - Reject a deterministic-policy mismatch before mutating live state, matching native Gefen.load_state_dict, instead of silently overwriting the flag. - Treat an initialized slot that reshards to an empty local shard (N->M where dim-0 < the target world) as unmaterialized, name-only state and leave it out of codebook learning, so it no longer feeds a None per-rank codebook into the re-quantize and dereferences codebook.device. - Route the restore period search through the same force_1d_period_one / force_2d_period_one / period_one_substrings gates as _resolve_automatic_period, so an explicit period-one config is honored on resume instead of restoring a frozen period>1 codebook. - Reject a slot carrying only some required optimizer fields instead of treating it as uninitialized and zeroing the momentum/second-moment history. Extend the resharding tests with per-fix coverage (changed-lr restore, swapped-identity rejection, deterministic-mismatch rejection, force-period-one restore, and an N->M reshard to empty target shards) and correct the 0.4.1 changelog summary, which claimed no public API changes despite adding GefenDCPState.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/gefen/dcp.py`:
- Around line 295-310: Update _identities and the DCP setup validation to
require caller-provided, stable parameter names before DCP is enabled; reject
missing or empty synthesized names and duplicate name/group/shape identity
triples instead of deriving identities from slot order. Preserve identity
serialization for valid parameters, and add regression coverage for unnamed and
duplicate identities to ensure loading fails rather than cross-assigning
optimizer history.
- Around line 445-449: Update the checkpoint-loading flow around saved_hypers
and the optimizer-state replacement to fully validate every required
hyperparameter, including presence, numeric type, finite values, and native
range constraints, before mutating live optimizer state or parameter groups.
Reuse the existing hyperparameter parsing/validation rules where available,
commit changes only after all groups pass, and add a corruption test verifying
rejection leaves both optimizer state and groups unchanged.
🪄 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: 4ae0860c-ee7e-4013-b5db-afd29c0328c2
📒 Files selected for processing (3)
CHANGELOG.mdsrc/gefen/dcp.pytests/test_dcp_resharding.py
🚧 Files skipped from review as they are similar to previous changes (1)
- CHANGELOG.md
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7ea3f52e02
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Address the CodeRabbit + Codex critiques of the v2 identity/hyper/init additions: 1. Require caller-stable, unique parameter identities. Reject at construction when param_names is missing/short, matches Gefen's synthesized positional pattern (group_N_param_M / param_N), or yields a duplicate (name, group, shape) triple, instead of silently cross-assigning momentum on a reordered load. 2. Parse and range-validate every group hyperparameter (finite, lr>=0, 0<=betas<1, eps>0, wd>=0) BEFORE any mutation, so a missing/non-numeric key or a poisoning value (lr=NaN, beta1=1, negative eps/wd) is rejected fail-atomically rather than committed. 3. Recompute the slot layout (re-run the fail-closed layout gate) at save and load time so a retained wrapper reflects add_param_group instead of a stale construction-time snapshot. 4. Co-locate the learned codebook with each slot's operand device inside the re-block so mixed CPU/CUDA shards do not abort on a device mismatch. 5. Synchronize staging/validation/re-block success across the process group before committing (all_reduce of a local-ok flag via _synchronize_step_failure) so either every rank commits or every rank raises; gloo-safe for CPU-resident state. 6. Reject incoherent initialization metadata: an initialized slot must carry step>=1 and vmean_step>=1; an uninitialized slot must carry no counters and no dense history that would be silently discarded. Serialized schema is unchanged, so _FORMAT_VERSION stays at 2. Extends tests/test_dcp_resharding.py with per-finding coverage (unnamed/duplicate rejection, fail-atomic corrupt-hyper load, add_param_group save/load, mixed-device codebook, one-rank-failure abort, incoherent-counter rejection).
The PR states that multidimensional meshes, subgroups, and non-Shard(0) placements fail closed, but only the non-DTensor / capturable / factored_v_2d rejections had tests. Add reject tests for a non-Shard(0) (Replicate) placement, a multidimensional DeviceMesh, and a group carrying a Muon sharded_mode.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/test_dcp_resharding.py (1)
943-954: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCompare the complete optimizer state for fail-atomicity.
This signature can miss mutations to
m_codebook,vmean, codebooks, global step, orm_magnitudevalues whose sum is unchanged. Consequently, every downstreamstate_unchangedassertion can pass after a partial commit.Snapshot all state tensors and optimizer-level metadata, then compare them exactly after the expected failure.
🤖 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_dcp_resharding.py` around lines 943 - 954, Update _live_state_signature to snapshot the complete optimizer state, including all per-parameter state tensors and optimizer-level metadata such as global step, m_codebook, vmean, codebooks, and exact m_magnitude values. Use exact tensor comparisons rather than aggregate sums, and ensure downstream state_unchanged assertions compare the full snapshots after expected failures.
🤖 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 `@tests/test_dcp_resharding.py`:
- Line 509: Add the existing Gloo availability skipif guard to
test_rejects_non_shard0_placement and the two immediately following tests.
Ensure each test skips before calling dist.init_process_group("gloo") when Gloo
is unsupported, matching the guard used by other Gloo-only tests.
---
Nitpick comments:
In `@tests/test_dcp_resharding.py`:
- Around line 943-954: Update _live_state_signature to snapshot the complete
optimizer state, including all per-parameter state tensors and optimizer-level
metadata such as global step, m_codebook, vmean, codebooks, and exact
m_magnitude values. Use exact tensor comparisons rather than aggregate sums, and
ensure downstream state_unchanged assertions compare the full snapshots after
expected failures.
🪄 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: 364ee3e1-d28a-41bc-a242-da3f7238fac6
📒 Files selected for processing (2)
src/gefen/dcp.pytests/test_dcp_resharding.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/gefen/dcp.py
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ea89e85a2f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Codex: `_validate_layout` checked mesh size but not rank order, so a reordered
full-world 1-D mesh (e.g. ranks [1, 0]) passed while permuting shard->rank
ownership -- misaligning the replicated initialized/counter metadata with the
sharded momentum. Now reject any mesh whose ranks are not the canonical
0..N-1 order. Test: test_rejects_reordered_full_world_mesh (2-rank).
CodeRabbit: the layout-reject tests called init_process_group("gloo")
unconditionally; add the same `skipif(not gloo available)` guard the reshard
tests use so unsupported builds skip instead of failing.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bd47e80b6a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ) | ||
| ) | ||
| for name, parameter in zip(names, group["params"]): | ||
| if _SYNTHESIZED_NAME_RE.match(str(name)): |
There was a problem hiding this comment.
When a model genuinely declares a parameter named param_0 or group_0_param_0, model.named_parameters() supplies a caller-stable name, but this spelling-only check rejects it as synthesized. Such an optimizer has a valid unique identity yet cannot use GefenDCPState; track whether Gefen actually generated the name rather than rejecting every explicit name matching this pattern.
Useful? React with 👍 / 👎.
Comment on lines
+351
to
+357
| _REQUIRED_STATE_FIELDS = ( | ||
| "automatic_period", | ||
| "step", | ||
| "m_codebook", | ||
| "m_magnitude", | ||
| "vmean", | ||
| ) |
There was a problem hiding this comment.
If a transformed or corrupted slot contains only name and vmean_step, this required-field list sees no initialized fields, so _slot_initialized() classifies it as fresh and state_dict() successfully writes an uninitialized slot with a nonzero vmean_step; every subsequent load then rejects that checkpoint as incoherent. Fresh evidence beyond the earlier partial-state finding is that vmean_step is still omitted from the fields used to detect partial state, so the save-side fix remains incomplete.
Useful? React with 👍 / 👎.
Comment on lines
+805
to
+806
| if torch.is_tensor(current): | ||
| current.fill_(value) |
There was a problem hiding this comment.
When the destination optimizer uses an accepted one-element integral LR tensor, loading a normal fractional checkpoint LR such as 1e-3 executes current.fill_(1e-3) and silently truncates it to zero, freezing subsequent updates. This occurs after the optimizer state has already been replaced, so reject an unrepresentable destination tensor dtype or stage an exact representation before committing.
Useful? React with 👍 / 👎.
Comment on lines
+70
to
+72
| def _dense_momentum(optimizer, parameter, state) -> torch.Tensor: | ||
| indices = state["m_codebook"].reshape(-1).long() | ||
| magnitude = state["m_magnitude"].reshape(-1).float() |
There was a problem hiding this comment.
For a large local shard, converting the entire uint8 m_codebook buffer to long allocates an additional 8 bytes per parameter while the original indices and resulting dense fp32 momentum are also live; a one-billion-parameter shard therefore needs roughly 8 GiB of avoidable transient memory before DCP can write anything and can OOM an otherwise viable save. Use the existing chunked gefen_dequantize_unpacked_indices path, or equivalent bounded chunks, before applying the block magnitudes.
Useful? React with 👍 / 👎.
| | Megatron-LM | All optimizers, including checkpoint resume — [scope](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#megatron-lm-integration-scope) | | ||
|
|
||
| > **FSDP2 checkpoint scope.** Plain Gefen and Muon `approx` save and resume exactly through PyTorch's standard full-state checkpoint calls, as long as the GPU count and sharding layout are unchanged and every GPU joins the save. Anything outside that scope refuses to load instead of corrupting state — [full details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). | ||
| > **FSDP2 checkpoint scope.** `GefenDCPState` is for **resharding** plain Gefen: wrap the optimizer with it for `torch.distributed.checkpoint`, and a checkpoint saved on N ranks loads on M ranks for a one-dimensional default-world `Shard(0)` mesh. Load re-blocks the resharded state back to Gefen's compact ~1 byte/param form, so resume is a correct continuation within quantization noise rather than a bit-exact restore; same-topology resumes should use the bit-exact native full-state path (Muon `approx` too). Requires `factored_v_2d=False` and `capturable=False` — [full details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). |
There was a problem hiding this comment.
When source and target block boundaries differ, the saved dense second moment is only each source block's vmean repeated per element, and load then averages those repeated values into the target blocks; the original per-element second-moment history is unavailable. Consequently continuation error includes irreversible second-moment aggregation error that can be much larger than 256-level momentum quantization noise, so this user-facing guarantee should describe both sources of approximation rather than bounding resume as being within quantization noise.
Useful? React with 👍 / 👎.
This was referenced Jul 16, 2026
Closed
thad0ctor
added a commit
that referenced
this pull request
Jul 16, 2026
Addresses the Codex review on #88. All three findings were verified against the code and tests before fixing; all three were valid. - Restore the [0.4.1] section to its as-released text. #83 wrote its GefenDCPState entry into the already-published 0.4.1 section, which shipped on 2026-07-15 from e986963 -- a commit that contains no src/gefen/dcp.py. gefen-x==0.4.1 on PyPI has no GefenDCPState, so the attribution was wrong and my 0.5.0 entry duplicated it. GefenDCPState is now credited to 0.5.0 only, and the 0.4.1 section is byte-identical to e986963:CHANGELOG.md. - Drop the "resume across a device move is bit-exact" claim. The cross-device test (test_cpu_resident_step.py) asserts state co-location and a finite continuation, not bit-exactness, and test_cpu_vs_cuda_nonfused_parity documents that CPU/CUDA bit-exactness is not achievable because the grad^2 reductions use different accumulation orders. Now described as a correct, numerically close continuation. - Qualify "the CUDA same-device step path is unchanged". Values stay bit-identical, but #80 scalarizes a tensor lr on the non-fused, non-capturable path, and _lr_scalar re-reads via .item() whenever the lr tensor's identity or _version changes. An in-place LR scheduler on a CUDA tensor lr therefore now costs one D2H sync per step on that path. Documented as an explicit behavior change with the float-lr workaround.
This was referenced Jul 16, 2026
thad0ctor
added a commit
that referenced
this pull request
Jul 17, 2026
…off test (#90) * Fix DCP resharding review findings: bounded save transient, honest resume error Post-merge review follow-up to #83. Bound the save-side dequantize transient. _dense_momentum built a full-size int64 copy of the uint8 indices plus a full-size fp32 gather on top of the fp32 result -- a measured 16.00 bytes/param of transient (~16 GiB on a 1B-param local shard), enough to OOM an otherwise-viable save in the one optimizer whose premise is memory efficiency. Route it through the existing chunked gefen_dequantize_unpacked_indices and scale the magnitudes in place: measured 6.00 bytes/param at 64M elements, and the excess over the 4 bytes/param result is now a bounded constant rather than proportional to shard size. The saved momentum is bit-identical, asserted against the previous implementation as an oracle. Reject an orphaned vmean_step at save. vmean_step was omitted from partial-state detection, so a slot carrying only name + vmean_step was classified as wholly fresh and written as an "uninitialized" slot with a nonzero counter -- which every later load rejects as incoherent, i.e. a silently dead checkpoint. It now marks a materialize without being required, so such a slot is rejected at save while a legacy state that carries vmean without the counter still saves. Reject an unrepresentable tensor hyperparameter during staging. Restoring a fractional checkpoint lr into a one-element integral lr tensor truncated it to 0 via fill_, silently freezing every later update -- and it ran after the optimizer state was already replaced, breaking fail-atomic. Representability is now checked before anything is committed. Decide synthesized names by provenance, not spelling. The guard pattern-matched param_N / group_N_param_M, so a model that genuinely declares a parameter named "param_0" was rejected and could not use GefenDCPState at all. Gefen now records at registration whether it generated a name (the only place the distinction is knowable) and the guard asks that instead; a name of unknown provenance still falls back to the conservative spelling check. Document the resume error honestly. The docs bounded it by 256-level quantization noise, which covers only the momentum. The saved dense second moment is each source block's vmean repeated per element, and load averages those into the target blocks: exact to fp32 round-off while the target blocking matches or refines the source's, but irreversible aggregation once a target block spans several source blocks (measured ~29% relative at 2x coarser, vs ~2e-4 for momentum re-quantization). Every existing reshard test re-derives the same period (2048), so they only ever exercise the aligned regime -- _CONTINUE_TOL is empirically fine but its stated reason was wrong; the comment now says what it actually covers, and the coarsening regime is pinned by a new test. * Make the DCP handoff workers report failures instead of a bare assert (#89) test_full_dcp_handoff_is_exact_on_same_dtensor_topology flakes on CI and, when it does, reports only `assert all(process.exitcode == 0)` -> `assert False` with no indication of what went wrong. Two concrete reasons: * `dist.init_process_group` ran OUTSIDE the worker's try, so a rendezvous / init failure -- a real possibility on a loaded CI runner -- produced a nonzero exit with no traceback at all. It now runs inside the try (and the finally only tears down a group that was actually initialized). * Only rank 0 put a result, and the parent drained exactly one item. A failure on any other rank therefore put its traceback on the queue where nobody read it: the parent took rank 0's success instead and surfaced the failure as an unexplained nonzero exit code. Every rank now reports, the parent drains one result per rank, and any worker traceback is raised in the assertion. Also distinguishes "terminated because slow" from "crashed" in the timeout and exitcode assertion messages. Verified by injecting a rank-1 failure: the test now fails in ~4s with the worker's real traceback instead of a bare assert. * Synchronize rank-local DCP validation; correct the re-blocking regimes Codex [P2]: _validate_hyper_destinations reads the LIVE optimizer, so it can fail on one rank and pass on the others, but it ran before the synchronized staging block -- the failing rank raised and left the collective its peers were entering. Rank-local validation (destination representability + the deterministic flag) now runs inside that block as _validate_rank_local, so a one-rank failure becomes a group-wide abort. All checks stay ahead of the commit, preserving fail-atomic. Rank 1 previously escaped the collective only because rank 0 tore its process group down on the way out, surfacing a gloo "Connection closed by peer" instead of the real cause; a failing rank that keeps its group alive leaves peers blocked. The new test therefore asserts which error each rank gets -- asserting only "every rank raised" passes with the bug present. Also corrects the second-moment re-blocking story, which was backwards. It was measured against the SOURCE values rather than a native run at the target blocking, which credits "gave back the number we started with" as fidelity and so rewards the lossy direction; the fixture was constant within each source block, hiding the effect. Against a native target-blocking run: aligned 1.2e-07, coarsen 4x on source boundaries 2.5e-07 (the EMA and block mean are linear, so averaging equal-size block means is the mean over their union), refine 4x 1.6e-01. Refining/straddling loses history, not coarsening; the ~29% figure was never a resume error. README, COMPATIBILITY.md, _CONTINUE_TOL and the regime test are corrected. * Bound the DCP save's peak memory to one slot's dense form Resharding forces the save to expand Gefen's ~1 byte/param block state into dense fp32 (4 B/param momentum + 4 B/param second moment), because the momentum codebook is learned per rank and the indices are meaningless off their owning rank. That expansion is the format's price. Holding it for every slot at once was not: state_dict() built every dense pair up front and handed DCP a finished dict, so a save peaked HIGHER than a training step -- a model that trains fine could OOM on save, in the one optimizer whose premise is memory efficiency. state_dict() now returns zero-storage stand-ins (_LazyDenseShard, a wrapper subclass reporting real shape/dtype/device but owning no storage) and expands exactly one at a time, when the writer asks for that write item. Peak over resident state, 67.1M-param local shard: 12.00 -> 6.00 B/param for one 64M tensor, 8.25 -> 0.50 B/param for 32 tensors. The bound is per slot, so the win grows as state is split across more parameters. dcp.load also calls state_dict() to build its destination, so a stand-in must work as a load destination too: __get_tensor_shard__ retains its buffer, since returning a fresh tensor per call would let load copy the checkpoint into a temporary and drop it -- a silent no-op restore. Only GefenSavePlanner bypasses the cache, which is why the bound is opt-in. Without the planner the save is still correct, at the old cost; a loud failure there is impossible because the hook cannot tell an unplanned save from a load. Also fixes an independent 2x bug this uncovered: _dense_second_moment did expand(...).reshape(...).clone(), and reshaping a stride-0 expand already allocates a full copy that the clone then copied again. The clone was not pointless -- at period == 1 the reshape returns a view aliasing live vmean state -- so it is now a preallocated broadcast fill, single-copy and never-aliasing. That alone is 12.0 -> 8.0. Saved checkpoint bytes are unchanged (payload md5-identical to the previous implementation, verified against it as an oracle). From adversarial review of the above: - dcp.async_save WORKED before this change and now fails. That is a deliberate capability regression (its CPU staging copies the whole state dict up front, which is the materialization the bound exists to prevent), but it was undocumented, and the error told users to pass GefenSavePlanner -- which the repro already did. Message rewritten to name the real constraint; documented in COMPATIBILITY.md, README, and the class docstring; restoring it is #91. - _validate_dense_geometry never checked m_codebook.numel() == local_numel, which _dense_momentum's final reshape requires, so a malformed state failed mid-write rather than before the first byte as its docstring claimed. - repr/str/print on a stand-in raised, so any logger printing the dict would mask the real error with a secondary one. __repr__ no longer dispatches, and pickling raises a TypeError naming the constraint. - The dequantize scratch is a ~128 MiB ceiling, not a constant; as written it contradicted the 0.50 B/param figure in its own sentence. * Scope the DCP async_save claim to where torch's staging actually copies CI on the declared torch floor (py3.10 / torch 2.5.0) failed test_dcp_async_save_is_refused_and_writes_nothing with DID NOT RAISE: on 2.5, dcp.async_save does not refuse a GefenDCPState state dict. The question that mattered was not the test but whether 2.5 users get a silently corrupt checkpoint from a call that appears to succeed -- staging zero-storage stand-ins into empty CPU buffers and writing those. So it was measured, not reasoned about: a real Gefen + DTensor on a single-rank Gloo group, saved both ways and loaded back. async_save's checkpoint restores bit-identically to the synchronous save's (automatic_period, step, m_codebook, m_magnitude, vmean all equal). There is no corruption on the floor, and nothing to fail loudly about. torch 2.5 stages through _offload_state_dict_to_cpu -> _iterate_state_dict, whose entire offload step is `ret = ret.to(cpu_device)`. For state already on the CPU that is a no-op torch short-circuits in C++ to the same object, so it never reaches __torch_dispatch__: nothing is copied, the stand-ins survive staging by reference, and the writer resolves them one at a time through the ordinary save protocol -- which is exactly what the synchronous save does. That is not stand-in-specific; 2.5 aliases any already-CPU state dict. Newer torch stages via _create_cpu_state_dict -> zeros_like, which a stand-in refuses, and so does 2.5 for non-CPU state, where .to(cpu) is a real copy. The refusal therefore still holds for every CUDA-resident save on every supported version -- i.e. every real FSDP2 deployment this adapter exists for. The gap is CPU-resident state on 2.5 alone, which is the test's own toy configuration. So the test now pins the invariant that holds on every version instead of a refusal only some versions produce: async_save either fails loudly and writes nothing, or the checkpoint it leaves restores exactly like the synchronous save. A checkpoint that lands and restores wrong fails on any torch, which is the property the old test was reaching for. It branches on the outcome rather than on a version number, because the deciding factor is whether staging copies -- which depends on the device as much as the version. The docs claimed flatly that async_save "fails loudly with a RuntimeError, writing nothing", which is false on the floor. That claim is now scoped to where it is true, and async_save stays unsupported: the 2.5/CPU pass stages nothing, so the write is not asynchronous with respect to the optimizer, and the dense expansion runs in the writer thread against live state -- resuming training while it runs races the writer against your own step. It "works" only because the test is quiescent. No behaviour change: src/ is docstring-only, so the sync save's bytes remain md5-identical to the pre-#90 implementation. * Reject positional collision names and overflowing hyper destinations A caller-named single-parameter group whose name collided kept synthesized=False, but _unique_name's suffix is decided by registration order alone: two groups both named "weight" become weight/weight_1, and building them in the other order swaps which parameter owns which name. DCP validated that on spelling and cross-assigned the two parameters' momentum silently. Mark the name synthesized when the collision suffix changed it, so DCP refuses an identity it cannot rely on. A group name that needed no suffix is untouched and still reshards. _validate_hyper_destinations rejected only fractional truncation, so an integral value that overflows its destination (lr=128.0 into int8) passed staging and left fill_ to raise -- after the state was swapped in, which is the half-applied restore the check exists to prevent. Range-check integral destinations during staging instead. bool is special-cased to {0, 1}: torch.iinfo has no bool entry, and it is the one integral dtype whose fill_ never raises, so an out-of-range value would coerce to True and commit a silently wrong lr. Complex destinations join the floating-point skip, since they represent fractional values. Correct the two checkpoint summaries: a GefenMuon(sharded_mode= "distributed") checkpoint consolidates every owner's complete state onto every rank and does resume under a different world size, contradicting the blanket same-topology claim. Both cross-world paths are covered (2->4 in test_muon_fsdp2_distributed_checkpoint_restore, 2->1 in test_consolidated_v2_checkpoint_loads_into_single_process_optimizer). Also align README with COMPATIBILITY on the DCP resume: only the momentum is bounded by quantization noise, not the second-moment re-aggregation. Poll the FSDP2 checkpoint drain against liveness so a worker that dies without reporting short-circuits instead of burning the whole timeout. Only a nonzero exit aborts, and only after a flush grace period: a rank that exits 0 has reported or is still flushing its feeder thread, and a failing rank puts its traceback immediately before exiting. * Synchronize the DCP save's rank-local validation; pin name provenance Codex [P2]: state_dict()'s pre-save checks read rank-local state -- the live codebook, each slot's block geometry -- so they can fail on ONE rank while the others pass, and dcp.save calls state_dict() while converting Stateful objects, BEFORE it enters its own planning collectives. A bare raise dropped the failing rank out while its peers walked into collectives it would never join. In a 2-rank repro (rank 0 missing its codebook) the healthy rank sat blocked for 35s and escaped only when rank 0's process exited and tore the transport down; a training loop that catches the error and keeps its group alive blocks its peers for good. Rank-local failures are now agreed across the WORLD group before any rank returns, as load_state_dict already does. That makes state_dict() collective, which the docstring now states plainly. dcp.save and dcp.load both call it on every rank, so documented usage is unaffected; a lone rank-0 call would block. The trade is worth it: the dict is full of stand-ins that carry no data and cannot be pickled, so inspecting it outside a collective save/load yields nothing usable anyway. A conditional collective was considered and rejected -- a rank cannot detect locally whether its peers are also calling, so there is no safe no-op to fall back to. Codex [P2]: GefenSavePlanner inherited flatten_state_dict from DefaultSavePlanner, but a GefenDCPState cannot be saved unflattened. DCP only descends into a nested mapping when it flattens it: unflattened, the whole {"optimizer": ...} mapping becomes a SINGLE BYTE_IO write item that DCP pickles wholesale, and the stand-ins refuse. The per-slot write items are never created, so teaching resolution to walk the mapping could not have fixed it -- there is nothing to resolve. Rejected at construction, naming the real constraint, rather than surfacing as an opaque "cannot be pickled" TypeError mid-write. Codex [P2]: a native load that CHANGES a parameter's name dropped the flag recorded at registration (it describes the OLD name) and re-derived provenance from the new name's spelling -- which cannot see a positional collision. A checkpoint from two unnamed single-parameter groups both called "weight" holds "weight"/"weight_1"; loaded into an optimizer named alpha/beta, neither matches the group_N_param_N / param_N pattern, so both were taken as caller-provided. GefenDCPState then accepted the order-dependent "weight_1", and a later DCP load with those groups rebuilt in the other order could validate while cross-assigning momentum. Widening the pattern is not available: no regex separates a positional "weight_1" from a legitimate "layer_1". Provenance is therefore serialized alongside param_names, so it travels with the name it describes and survives a rename exactly -- "weight" stays resharding-safe while only "weight_1" is refused. The key is additive; a checkpoint written without it leaves provenance unknowable, so a CHANGED name fails closed instead of falling back to spelling. That costs a legacy checkpoint whose names legitimately changed its ability to reshard until re-saved, which is recoverable; cross-assigned momentum is not. An unchanged name still keeps its registration-time flag, so ordinary legacy resumes are unaffected. The reachable path is narrower than reported and the tests pin it: a checkpoint saved from DTensor parameters carries a rank-local topology marker and load rejects any rename against it long before provenance is consulted. It is a checkpoint WITHOUT that marker -- a single-device run resumed into a sharded one, exactly when a caller then reaches for resharding -- that reaches this code. Saved DCP checkpoint bytes are unchanged (payload md5-identical pre/post). * Support async_save with a staging writer that keeps the save bounded async_save returns once the state dict is staged to the CPU and writes from a background thread, so training may resume immediately while the checkpoint still reflects the moment of the call. Staging is what makes that true, and the default writer cannot stage a GefenDCPState: it asks each entry for a CPU copy of itself, and the stand-ins that keep the save memory-bounded have nothing to copy. GefenFileSystemWriter stages those slots itself, one at a time -- expanding a slot's dense form, copying it to the CPU, releasing it before the next -- so the device holds a single slot's dense form during staging, the same bound GefenSavePlanner gives the synchronous save. Measured GPU peak during staging is 0.500 B/param on a 32-tensor shard, identical to the sync path; the CPU then holds the full snapshot, which is what async staging costs for any optimizer. FileSystemWriter already satisfies AsyncStager on both torch 2.5 and 2.12, and both versions route through it, so this is one API with no version-conditional caller code. The default writer is now refused rather than half-working. Torch 2.5 stages with .to(cpu_device), which short-circuits to the same object for CPU-resident state: the stand-ins entered the staged dict by reference, nothing was copied, and the writer thread then expanded them against LIVE optimizer state while training continued -- writing save-time counters beside post-step momentum, silently. It looked correct only while the state was quiescent, which is the condition async_save exists to lift. __torch_dispatch__ cannot see a same-device .to(), so the refusal lives in __torch_function__ and names the writer that stages correctly. Retires the test that pinned async as unsupported. Its premise -- that the 2.5 CPU pass-through "lands a correct checkpoint, so there is nothing to refuse" -- is what the tear above disproves; its assertions could have been reworded to pass, but not over a false premise. It now pins the real contract: the default writer refuses loudly and writes nothing, and GefenFileSystemWriter lands the checkpoint the synchronous save would. The snapshot test gates write_data so the step provably lands before the write, rather than racing it -- an unstaged save cannot accidentally pass. * Hold the DCP save's one-slot bound under a multi-threaded writer Codex [P2]: thread_count is a public FileSystemWriter option, and it splits the plan into one file bucket per thread, so N threads call resolve_data at once. Resolving is what expands a slot, so the advertised "single slot's dense form" was one slot PER THREAD. Measured on the 32-tensor bench, the peak over resident state scaled with the option: 0.500 B/param at thread_count=1, 0.625 at 2, 1.250 at 4, 2.375 at 8. On the 8-slot test model thread_count=8 reached ~8 B/param -- the whole eager cost the planner exists to avoid, silently restored by a throughput knob. torch's per_thread_copy_ahead does not bound this: it feeds _OverlappingCpuLoader, which torch only uses at thread_count == 1, and it bounds CPU bytes in flight rather than the device-side expansion. The planner cannot detect the writer. DCP hands the writer the planner, not the reverse -- set_up_planner receives only (state_dict, storage_meta, is_coordinator), and StorageMeta carries checkpoint_id/save_id/load_id/modules. So rejecting thread_count > 1 is not enforceable where it matters: the synchronous path takes a plain FileSystemWriter that Gefen never constructs. Serializing alone would not have bounded anything either. The writer holds what resolve_data returns until it has finished writing that item, so a device-side return outlives any lock held across the call, and N threads still hold N dense shards. resolve_data therefore lands the shard on the CPU before returning, which makes the device-side expansion dead before the lock is released -- and only then does a lock around it bound one slot. The peak is now flat at the advertised numbers for every thread count measured: 0.500 B/param on the 32-tensor bench and 6.000 on a single 64M tensor at thread_count 1, 2, 4 and 8. The copy is necessarily synchronous -- the writer may write the tensor the moment it is returned -- which gives up the copy/expand overlap torch's single-threaded loader gets from its own non_blocking copy. Pinning the destination buys that bandwidth back: the default thread_count=1 save costs +9.6% wall time (273 -> 299 ms on the 32-tensor bench), while thread_count=4 is 46% faster than before (217 -> 117 ms) and is now bounded, so the knob that exposed this more than pays back the default's cost. Async staging expands through the same helper but unpinned: it keeps every slot's snapshot until the background write drains it, and page-locking the entire dense state dict for a whole write is a cost the sync path's one-slot-at-a-time lifetime does not have. Saved DCP checkpoint bytes are unchanged (payload md5-identical pre/post). empty_like, not empty, is what keeps that true for uninitialized slots: they expand with zeros_like, which preserves a channels_last parameter's format, and the writer serializes strides. * Page-lock one slot at a time on the bounded DCP save path CodeRabbit [Major]: the lock only bounds device-side expansion, so N writer threads can keep N pinned CPU shards alive across the write. Measured, and it is worse than N: torch's FileSystemWriter retains every tensor it has written in a dict scoped to the whole file (filesystem.py populates tensor_dict unconditionally, then reads it only for SAFETENSORS), so at thread_count=8 the page-locked high-water was the entire dense state dict -- 8 bytes/param, 536.87 MB on the 32x2M bench, 66 cudaHostAllocs that the caching host allocator then never returns to the OS. Neutralizing that retention drops the live set to exactly thread_count shards, which is the mechanism the finding names; the writer's own retention is what scales it to all of them. What the finding does not change: the advertised bound is on VRAM, and it holds. GPU peak stays flat at 0.500 B/param (32x2M) and 6.000 (1x64M) for thread_count 1, 2, 4 and 8. Host memory was never bounded to one slot, in any version. The writer holds the full dense state dict by the end of a save whatever the planner returns -- 8 bytes/param of RSS pre-change and post-change alike, at every thread count. That is torch's, not ours, and no permit can take it back: a permit released at the write/consumption boundary still could not free a tensor tensor_dict holds until the file closes. Gating to thread_count=1 is not enforceable either -- set_up_planner receives only (state_dict, storage_meta, is_coordinator), and StorageMeta carries checkpoint_id/save_id/load_id/modules, so the planner never sees the writer, and the sync path takes a plain FileSystemWriter Gefen never constructs. What our pinning did change is the page-locked share. At thread_count=1 it changed nothing: torch's _OverlappingCpuLoader copies with `.to("cpu", non_blocking=True)`, whose destination is page-locked, so the default path already held 536.87 MB pinned with the identical 66 cudaHostAllocs before this branch. Above 1 torch uses _SerialCpuLoader, whose `.cpu()` is pageable, and there returning pinned shards took the page-locked footprint from 0 to 8 bytes/param. So copies now pass through one reusable page-locked buffer and are handed to the writer pageable. The pinned high-water becomes a single slot's dense form at every thread count -- 8.39 MB (0.125 B/param) on the 32x2M bench, 268.44 MB (4.000 B/param, one 64M slot) on the single-tensor bench -- from 536.87 MB, and 3 cudaHostAllocs rather than 66. That is below torch's own default-path footprint, not merely back to it, and the shards the writer accumulates are now memory the host can swap. Held under the existing lock, which every expansion already serializes on: what comes back is a private tensor, so the next expansion may overwrite the buffer immediately. The extra host-to-host copy is not what it costs. Isolated, the per-slot copy goes from 10.9 ms to 19.5 ms per 256 MB of dense state (torch's own pageable `.cpu()` is 26.3 ms), but a save is dominated by writing the payload, and skipping 63 cudaHostAllocs -- milliseconds each -- pays for the copy: the default thread_count=1 save measures at or below the current path end-to-end on the 32x2M bench, against disk-write noise too large to separate them. Saved DCP checkpoint bytes are unchanged: payload __0_0.distcp md5 is identical to the pre-change implementation, to this branch's previous commit, and to this one. Allocating the pageable result with empty_like first is what keeps that true -- it fixes the strides the staging view must match, preserves a channels_last parameter's format, and guarantees the view spans exactly numel elements. * Scope the planner-omission and host-memory claims precisely Omitting GefenSavePlanner is harmless only for the synchronous dcp.save; async_save separately requires GefenFileSystemWriter and fails without it, because staging cannot resolve a stand-in through the planner at all. The two sat in different sections, so the omission claim read as covering both. Pageable is not the same as reclaimable in practice: a host with no swap, or a cgroup memory limit, still has to hold what the writer accumulates. Says the shards may be reclaimed or swapped rather than that the host is free to, and states plainly that a save needs room for the full dense state dict on the host either way. * Hold the DCP save's fail-before-write and one-slot page-locked bounds Two review findings, both cases where a bound the save advertises did not actually hold. Validate that the momentum indices can address the codebook. The existing checks confirm the block arrays tile the shard; they say nothing about whether an index in the shard exists in the codebook. _dense_momentum gathers with index_select, and the expansion is deferred to write time, so a codebook shorter than an index present raised IndexError from inside the write -- 49185 bytes of __0_0.distcp on disk and no .metadata, where the eager expansion this path replaced had failed before the writer started. The new checks are O(1) metadata, deliberately not a max() over the indices: they are uint8, so a 1-D codebook of at least 256 rows covers every index that can exist. A reduction would instead cost a full pass plus a device sync per slot on every save, and would still pass a truncated codebook whose live indices happened to be small, leaving the failure latent for the next step. Size the page-locked staging buffer from the largest slot, up front, instead of growing it as bigger slots arrive. Replacing the buffer does not give the old block back: torch's CachingHostAllocator keeps freed pinned blocks in its own cache and never returns them to the OS -- a free does not reach cudaHostFree at all -- so each new maximum added a page-locked block and the footprint climbed to the sum of the successive maxima. Measured on six ascending slots: 504.0 MiB page-locked against a 256.0 MiB largest slot (1.97x), six pinned blocks, zero driver frees; 256.0 MiB and one block after. That ordering is not a corner case. FileSystemWriter keys its write items by size and sorts them ascending, so at the default thread_count=1 a save with N distinct slot sizes walks the maxima in exactly the order that allocates all N. The existing bench missed it twice over: uniform slots (one distinct size) and thread_count=8, where _split_by_size_and_type sorts descending into per-thread buckets and the largest slot sizes the buffer immediately. set_up_planner names its arguments in full rather than taking *args: dcp.save inspects the signature and, finding no storage_meta, falls back to the pre-2.3 set_up_planner(state_dict, is_coordinator) call, which would bind is_coordinator into storage_meta's slot and leave is_coordinator False on every rank -- including the one that builds the global plan. Payload bytes are unchanged (md5-identical to the pre-change implementation on both the CPU and the CUDA staging path), and the GPU peak stays flat at 0.500 B/param (32x2M) and 6.000 (1x64M) for thread_count 1/2/4/8. * Await the async save through the version-tolerant response Newer torch returns an AsyncSaveResponse holding the future in upload_completion rather than a bare Future, so the documented example's future.result() raises AttributeError instead of waiting -- and a background write failure would surface nowhere. Uses the same accessor the tests use, and says why waiting matters. * Await the async response everywhere, and share the drain's reap budget The README fix for AsyncSaveResponse missed the COMPATIBILITY.md example and the GefenFileSystemWriter docstring, which is the same defect twice more: .result() on the response raises AttributeError on newer torch instead of waiting, and the docstring never awaited at all, so a background write failure surfaced nowhere. All three now use the accessor the tests use. The drain's reap was spent per worker, not shared: after the 180s deadline it joined each worker for 60s in turn, so two hung ranks cost ~300s of CI and _DRAIN_TIMEOUT was not the overall budget it advertises. The joins now share one grace period, and stragglers are all signalled before any is joined, so the world size cannot multiply the wall clock. By that point every worker has either reported -- and exits promptly -- or the drain already spent its deadline on workers that will not. * Reject block state that straddles devices before the save writes The preflight checked counts and dtypes but not where the block state lives, and the expansion is deferred to write time: _dense_momentum scales the gathered values by m_magnitude in place, and _dense_second_moment fills a buffer on vmean's device that is then wrapped against the parameter's mesh. A slot whose tensors straddle devices therefore raised from inside the write, leaving a torn partial checkpoint -- the fail-before-write property this preflight exists to hold. The check cannot cost a save that would otherwise have worked: those same operations fail on any mismatch it rejects, so it only moves the failure ahead of the first byte. Compares device identity rather than type, since cuda:0 and cuda:1 straddle as surely as cuda and cpu. CPU-resident state stays valid -- it is co-located with a CPU-resident shard, not mismatched with it. * Reject one name shared by two parameters before the save The saved identity is (name, group, shape) with a positional group index, so two same-shaped parameters sharing a name produce an identity list that a group reorder leaves unchanged: validation passes and each slot's state lands on the other parameter. Require names to be unique across groups. Covered in both directions -- the collision is rejected, and distinct names in separate groups still round-trip. * Condense the DCP adapter's comments and docstrings The prose added on this branch ran about 2.4 lines per line of code, against 0.29 for the file it was added to. Cut it back: keep the constraints that would otherwise be re-broken -- the stand-in caches because dcp.load also calls state_dict(), the clone at period == 1 aliased live vmean, the staging buffer is sized up front because the host allocator never returns pinned blocks and DCP writes ascending, the device check cannot reject a valid save -- and drop restatements of the code, narrative, measured numbers that belong in the PR, and text written to justify a decision rather than inform a reader. Comments and docstrings only; no logic, error messages, or behavior change. * Pin the async and unplanned save memory bounds Correctness tests would pass even if async staging or the planner's bound regressed, so the two paths measured only by hand now have peak guards. test_dcp_async_save_peak_is_bounded_by_one_slot: GefenFileSystemWriter staging must stay near one slot on a multi-slot model. Fails at ~9 B/param if staging retains every slot on the device. test_dcp_save_without_planner_pays_the_eager_cost: the mirror -- omitting the planner is correct but costs the full eager footprint, which is what makes the one-slot tests a measurement of the planner rather than of something else. Also fixes a latent flake the second test exposed: the page-locked test read the absolute host high-water, but torch's writer retains pinned buffers in a global pool across tests, so any earlier pinning save contaminated it. It now measures the growth its own save adds. * Import fully_shard tolerantly across the torch 2.5 floor (#92) FSDP2's public torch.distributed.fsdp.fully_shard arrived in 2.6, but the package floor is 2.5, where it lives at torch.distributed._composable.fsdp. Six sites imported the public path unconditionally, so on a 2.5 environment with a GPU visible the fully_shard tests raised ImportError instead of running (CI missed it because they are CUDA-gated and CI is CPU-only), and the README example failed on the declared floor. Each site now falls back to the private path on 2.5; the README notes the 2.6 requirement rather than a private import. * Fix two DCP resharding fail-atomic/sync gaps from PR #90 review Finding A: reject float lr destinations that overflow to inf. _validate_hyper_destinations skipped every floating-point destination, so a finite checkpoint lr past the destination dtype's range (e.g. 1e10 into a float16 tensor) reached the in-place fill_ and cast to inf -- only after the optimizer state was already swapped in, defeating fail-atomic. Range-check float destinations against torch.finfo(dtype).max, matching the integral branch; ordinary rounding (fp64 -> fp32/bf16) is still accepted. Finding B: run the duplicate-caller-name check inside the synchronized region. The new duplicate-name/identity guard lived in _validate_layout, which state_dict()/load_state_dict() run BEFORE their _synchronize_step_failure block (on purpose: it establishes the process group and its other checks are replicated). But duplicate names depend on what each rank's optimizer holds, so a one-rank add_param_group/rename desyncs it: the offending rank raised before the sync while its peers walked into dcp.save/load's planning collectives and blocked. Extract the duplicate detection into _reject_duplicate_slots and call it eagerly at construction (replicated -> clean error) and again inside the synchronized regions (_build_state_dict for save, _validate_rank_local for load), so a one-rank duplicate becomes a group-wide abort instead of a hang. Saved checkpoint payload bytes are unchanged (both fixes are validation-only). * Cover the reshard levers left implicit: 2->1, CPU-offload, mixed precision The reshard tests exercised 2->4 and 4->2 but not the world-size-1 collapse the "load on a single GPU" claim rests on, and fully_shard was only tested with the default policy though the docs imply CPUOffloadPolicy and MixedPrecisionPolicy work through the same path. Adds 2->1 (gloo/CPU) and a parametrized fully_shard reshard under both policies (2-GPU NCCL): params staged on the host, and a reduced compute/reduce dtype, must each still round-trip compact and continue. * Run the whole of _validate_layout inside the DCP sync region _validate_layout reads rank-local state a one-rank add_param_group can desync (short param_names, non-DTensor, bad mesh), but state_dict and load_state_dict called it BEFORE their _synchronize_step_failure block. A one-rank layout failure raised before the sync, stranding healthy peers in the following planning collective -- the same class already fixed only for the duplicate-name check. Split out _require_process_group (the one check that must precede the sync, since the sync needs the group it proves exists) and move all remaining layout/metadata validation inside the sync try on both paths. __init__ keeps eager validation (construction is replicated; no collective yet). Payload bytes unchanged. * Harden DCP spawn drains; defer world>1 construct-time validation; bound pinned staging across saves FIX 1 (CI-red flake): factor the hardened spawn-drain from the FSDP2 checkpoint tests into a shared `_drain_workers` helper and route every resharding runner through it. The old per-runner drain took one blocking `result_queue.get(timeout=T)` then joined+terminated in `finally`, so a worker that finished its collective but was slow to flush its result on a loaded runner got reaped before it reported (len(results)==1, exitcodes [-15, 0]). The helper polls a shared deadline, only short-circuits on a NONZERO early exit (a still-running worker keeps the full budget), shares one reap/terminate budget, and tolerates a reported-but-slow-to-exit worker in the exitcode assertion. Desync runner deadlines raised 90->180s. FIX 2 (dcp.py __init__): construction-time `_validate_layout` / `_reject_duplicate_slots` are only strand-safe with no peers -- under rank-conditional construction one rank raises in __init__ while its peers block in state_dict()'s WORLD sync. Run them eagerly only at world size 1; at world > 1 defer to the synchronized save/load paths, which already re-derive and abort the whole group together. The world=2 reordered-mesh rejection now surfaces from the collective state_dict(). FIX 3 (_PinnedStaging): torch's CachingHostAllocator never returns freed pinned blocks to the OS, so successive saves whose largest slot grows each pin a new maximum and the page-locked footprint climbs toward the sum of the historical maxima (measured 240 MiB vs a 128 MiB largest slot). Add a process-wide pin bound fixed by the first staged save; a later slot past the bound stages pageable (a slower copy only on growth). Bounds the page-locked peak to one slot (240 -> 16 MiB in the repro) with the save payload bytes unchanged. New regression test fails without the fix. * Revert the ineffective across-saves pinned bound The process-wide pin cap added in cef24d3 does not work and cannot at this layer: a D2H copy into pageable memory itself uses a pinned bounce buffer that torch's CachingHostAllocator retains, so page-locked memory still climbs toward the sum of the historical maxima across growing saves even when GefenDCPState's own staging buffer is bounded (measured: cap fixed at one slot, pinned peak still 8->24->56->120 MiB across four growing saves). Its regression test passed only by order-dependent allocator state in a full run and failed in isolation. _PinnedStaging is restored to its per-save one-slot form; the across-saves accumulation is documented as torch host-allocator behavior below this layer. The drain-harness hardening and the world>1 construct-time deferral from cef24d3 are kept. * Reject a float hyperparameter that underflows the destination to zero The overflow guard only checked finfo.max, so a nonzero value below the destination dtype's smallest subnormal (lr=1e-10 into float16) passed staging and then fill_ rounded it to 0 after the state was swapped in -- freezing updates as silently as an integral truncation. Rejected during staging now, the mirror of the overflow case; ordinary rounding and representable subnormals still pass. * Reject beta round-to-one, non-finite save state, and cross-rank layout drift Three P2 preflight/validation gaps in the DCP reshard adapter: - A valid beta near 1 can round onto the excluded upper bound under a narrow float destination (0.9999 into float16 -> 1.0), which the overflow/underflow checks pass; the next bias correction 1 - beta**t then divides by zero. Reject a beta whose destination cast leaves [0, 1) during staging. - Non-finite live block state (NaN/Inf in codebook/m_magnitude/vmean) passed the geometry preflight, so the deferred dense expansion wrote it out and completed a checkpoint the load side's torch.isfinite later rejects. Add a cheap isfinite check on the compact block state so the save fails before the first byte. - Two ranks holding DIFFERENT but individually-valid layouts for the same positional slot each pass per-rank validation, so the boolean failure sync reports success and DCP merges two different parameters into one global tensor. Add a cross-rank identity-digest comparison (single min-reduce of a 128-bit hash) inside the synchronized save region; every rank sees the same verdict, so a divergent layout aborts the whole group instead of corrupting silently. * State the torch 2.5 fully_shard import in the resharding example The resharding path works on the declared 2.5 floor -- the code already falls back to torch.distributed._composable.fsdp there -- so "needs 2.6" wrongly excluded a supported version. The example now imports fully_shard with the same fallback, and the note says 2.5+ with the 2.6 public-path caveat.
thad0ctor
added a commit
that referenced
this pull request
Jul 17, 2026
* release: v0.5.0 — version bump + changelog Two user-facing features landed since v0.4.1, so this is a minor bump: - #80 CPU-offloaded training for plain Gefen (CPU-resident stepping, DeepSpeed ZeRO-2/3 CPU-offload, FSDP2 CPUOffloadPolicy). - #83 GefenDCPState, a new public export making plain Gefen's FSDP2 optimizer state reshardable across world sizes (closes #81). Both are additive; the CUDA same-device step path and the native state_dict/load_state_dict path are unchanged. * release: correct changelog attribution and overstated claims Addresses the Codex review on #88. All three findings were verified against the code and tests before fixing; all three were valid. - Restore the [0.4.1] section to its as-released text. #83 wrote its GefenDCPState entry into the already-published 0.4.1 section, which shipped on 2026-07-15 from e986963 -- a commit that contains no src/gefen/dcp.py. gefen-x==0.4.1 on PyPI has no GefenDCPState, so the attribution was wrong and my 0.5.0 entry duplicated it. GefenDCPState is now credited to 0.5.0 only, and the 0.4.1 section is byte-identical to e986963:CHANGELOG.md. - Drop the "resume across a device move is bit-exact" claim. The cross-device test (test_cpu_resident_step.py) asserts state co-location and a finite continuation, not bit-exactness, and test_cpu_vs_cuda_nonfused_parity documents that CPU/CUDA bit-exactness is not achievable because the grad^2 reductions use different accumulation orders. Now described as a correct, numerically close continuation. - Qualify "the CUDA same-device step path is unchanged". Values stay bit-identical, but #80 scalarizes a tensor lr on the non-fused, non-capturable path, and _lr_scalar re-reads via .item() whenever the lr tensor's identity or _version changes. An in-place LR scheduler on a CUDA tensor lr therefore now costs one D2H sync per step on that path. Documented as an explicit behavior change with the float-lr workaround. * changelog: note the #90 bounded/async save and resharding hardening * release: date 0.5.0 to 2026-07-17 * changelog: correct the tensor-LR bit-exactness claim (#80) With nonzero weight decay the non-fused tensor-LR path computes 1 - lr*wd on the host in double precision rather than as tensor arithmetic in the LR dtype, so the update is numerically equivalent, not bit-identical (a float16 LR rounds 1 - tiny to 1.0 on the old device path). The LR value itself is still exact.
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
Adds
GefenDCPState, a standalonetorch.distributed.checkpointadapter that makes plain Gefen's FSDP2 optimizer state reshardable — save on N ranks, load/continue on M — closing #81. The nativestate_dict/load_state_dictpath is unchanged.Shard(0)DTensors, so DCP can reshard them across world sizes.Fidelity / semantics
Resharding routes momentum through a dense reshard and re-blocks it against a freshly learned per-shard codebook, so the restored optimizer is a correct continuation within quantization noise (~256-level), not bit-exact — even same-topology. For unchanged topology, use the native
get_state_dict/set_state_dictpath, which is bit-exact.GefenDCPStateis for when the topology changes.Scope (fail-closed)
Plain
Gefenonly, on one-dimensional default-worldShard(0)DTensors. The following fail closed at construction/validation with clear messages:GefenMuon/GefenMuonHybrid(type(optimizer) is not Gefen),capturable=True,factored_v_2d=True, Muonsharded_mode, non-DTensor params, multidimensional or non-default-world meshes, subgroups, more than one placement, and non-Shard(0)placements. Reshardable DCP for the Muon family is tracked as follow-ups: #84 (GefenMuon) and #85 (GefenMuonHybrid).This is the issue-81 extraction of the ideas proven in #68, implemented directly against current
mainwithout pulling in #67's optimizer-contract convention. Supersedes the incorrectly broad draft #82.Usage
Validation
All on idle RTX 3090s (
CUDA_DEVICE_ORDER=PCI_BUS_ID, never the RTX PRO 6000s):tests/test_dcp_resharding.py) — each builds a realGefen, takes real steps (learned codebook, block period > 1), and asserts the restored state stays compact (vmean/m_magnitude= 1 value/block, not per-element), finite, and continues within tolerance of a native reference at the target topology:fully_shardnn.Linearmodel,fused=True) GPU test,capturable/factored_v_2d/ non-DTensor reject tests.See the review-summary comment on this PR for the findings from review and how each was addressed (the load path originally collapsed to per-element period-1 — a 2048× memory blow-up — which the re-block fix resolves).
Summary by CodeRabbit
New Features
GefenDCPState, an opt-in wrapper enabling FSDP2 optimizer checkpoint resharding across different rank counts without rank-0 full-state gather.GefenDCPStateat the package level for easier adoption.Documentation
Tests