Fix post-merge DCP resharding review findings + harden the flaky handoff test - #90
Merged
Conversation
…sume 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.
…#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.
![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: be25ef1a04
ℹ️ 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 [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.
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.
Owner
Author
|
@coderabbitai review |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughGefenDCPState now uses bounded lazy dense-state saving, adds planner and filesystem-writer support for synchronous and asynchronous checkpoints, validates restores collectively, tracks synthesized parameter-name provenance, and expands compatibility documentation and distributed test coverage. ChangesDCP resharding and checkpoint lifecycle
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant GefenDCPState
participant GefenSavePlanner
participant GefenFileSystemWriter
participant DCPStorage
GefenDCPState->>GefenSavePlanner: expose lazy dense optimizer slots
GefenSavePlanner->>GefenFileSystemWriter: materialize one slot
GefenFileSystemWriter->>DCPStorage: stage and write checkpoint data
DCPStorage-->>GefenDCPState: provide checkpoint data for restore
Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
tests/test_dcp_resharding.py (1)
1761-1789: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCompare the restored codebook in planner-equivalence coverage.
m_codebookcontains indices whose meaning depends ontarget._gefen_codebook. Include the codebook—and preferablyvmean_step/global step—in the signatures so semantically different restores cannot compare equal.🤖 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 1761 - 1789, Update the checkpoint signatures in the resharding equivalence test to include target._gefen_codebook alongside m_codebook, and include vmean_step/global-step metadata when available. Compare these restored codebook and step values between planned and plain checkpoints so semantically different optimizer restores cannot pass as equal.
🤖 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 `@COMPATIBILITY.md`:
- Line 115: Update the checkpoint compatibility summaries in COMPATIBILITY.md
(line 115) and README.md (line 639) to explicitly carve out the documented
cross-world GefenMuon(sharded_mode="distributed") path, rather than grouping all
Muon or sharded modes as same-topology. In README.md, also replace “within
quantization noise” with wording that acknowledges data-dependent error from
refining or off-boundary second-moment reblocking can be materially larger.
In `@src/gefen/dcp.py`:
- Around line 811-831: The integral hyperparameter validation in
src/gefen/dcp.py:811-831 must reject values outside the destination tensor dtype
range before optimizer state is committed; update the logic around
_GROUP_HYPER_KEYS to use torch.iinfo(current.dtype), special-casing torch.bool
to 0/1 when supported, while preserving the existing fractional-value rejection.
Add an out-of-range regression in tests/test_dcp_resharding.py:2046-2075, such
as 128.0 for an int8 lr tensor, and assert the live optimizer state remains
unchanged after staging fails.
In `@src/gefen/gefen.py`:
- Around line 1446-1451: The unique naming logic in _unique_name must mark the
parameter identity as synthesized when collision suffixing changes the
normalized caller group name, or reject the collision; preserve
synthesized=false only when the name remains unchanged. Update the
duplicate-group coverage in tests/test_dcp_resharding.py lines 2149-2193 to
exercise two unnamed single-parameter groups sharing a group name and verify DCP
identity validation prevents cross-assigned momentum.
In `@tests/test_gefen_fsdp2_checkpoint.py`:
- Around line 585-588: Update the queue-draining logic around result_queue.get
in the test’s process-waiting flow to poll at short intervals until an overall
deadline, while monitoring worker exit codes. When any worker exits nonzero,
stop waiting after a brief queue-flush grace period instead of allowing another
rank to block for the full 180-second timeout; preserve collection of results
already available.
---
Nitpick comments:
In `@tests/test_dcp_resharding.py`:
- Around line 1761-1789: Update the checkpoint signatures in the resharding
equivalence test to include target._gefen_codebook alongside m_codebook, and
include vmean_step/global-step metadata when available. Compare these restored
codebook and step values between planned and plain checkpoints so semantically
different optimizer restores cannot pass as equal.
🪄 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: 15bf89a2-737d-483a-b5f4-aa8f54e8276d
📒 Files selected for processing (7)
COMPATIBILITY.mdREADME.mdsrc/gefen/__init__.pysrc/gefen/dcp.pysrc/gefen/gefen.pytests/test_dcp_resharding.pytests/test_gefen_fsdp2_checkpoint.py
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.
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.
Owner
Author
|
All four addressed in Collision-suffixed names (Major) — confirmed, fixed. Verified the exact mechanism: two groups named Integral overflow (Major) — confirmed, fixed. Verified both halves of your claim before fixing: Your Muon carve-out (Minor) — you were right, I was wrong. My initial read was that the sentence was scoped to Carved out in both docs, with the distinction made explicit so it doesn't overclaim the other way: that portability is replication, not resharding — each rank materializes full state instead of addressing its own shard, so it doesn't buy what Drain loop (Minor) — implemented. Polls against a deadline; only a nonzero early exit short-circuits, and only after a grace period to flush the queue, since a rank can legitimately exit 0 right after reporting. Validated across four worker behaviors — the load-bearing one being a slow-but-passing rank exiting 0 at 6s, which is not aborted (7.0s, 2/2 results), while a crash + hung peer now surfaces the traceback in ~11s instead of burning the full 180s. Suite: 847 passed / 2 skipped / 0 failed on 4 GPUs; DCP file 53 passed. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f4913ed27a
ℹ️ 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 [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).
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.
Owner
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0440664ddf
ℹ️ 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 [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.
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 `@src/gefen/dcp.py`:
- Around line 176-180: The save path’s one-slot memory assumption is invalid
because FileSystemWriter can retain multiple CPU shards concurrently. Update
GefenSavePlanner.resolve_data or its surrounding save flow to either enforce
FileSystemWriter(thread_count=1) or acquire/release the permit around writer
consumption so only one shard remains live at a time.
🪄 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: 63863270-a990-491e-a093-fd77109af6f0
📒 Files selected for processing (3)
COMPATIBILITY.mdsrc/gefen/dcp.pytests/test_dcp_resharding.py
🚧 Files skipped from review as they are similar to previous changes (2)
- COMPATIBILITY.md
- tests/test_dcp_resharding.py
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b7a160ad75
ℹ️ 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".
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.
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 `@COMPATIBILITY.md`:
- Line 117: Update the planner-omission claim in the documentation to explicitly
limit it to synchronous dcp.save calls using a supported writer. Clarify that
this does not apply to dcp.async_save, whose async section requires
GefenFileSystemWriter.
- Line 138: Update the compatibility documentation paragraph describing
`torch`’s `FileSystemWriter` and planner behavior: replace the guarantee that
host memory “stays free to swap” with wording that pageable staging may be
swappable but remains subject to host-memory limits, including environments
without swap or with cgroup limits. Preserve the distinction between bounded
page-locked memory and the writer retaining the full dense state.
🪄 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: 40042d35-9d90-435f-abe8-13fdb5a847c9
📒 Files selected for processing (3)
COMPATIBILITY.mdsrc/gefen/dcp.pytests/test_dcp_resharding.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/test_dcp_resharding.py
- src/gefen/dcp.py
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e5d736a6de
ℹ️ 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".
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.
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c78aba58be
ℹ️ 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".
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 85c9d9f932
ℹ️ 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".
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b635f6ea7e
ℹ️ 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".
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.
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 35e75d7404
ℹ️ 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".
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.
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9de5c186c9
ℹ️ 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".
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).
…ision 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.
There was a problem hiding this comment.
💡 Codex Review
Lines 1277 to 1278 in 4f1d721
Update differentiable tensor hyperparameters without autograd
When a live one-element tensor hyperparameter has requires_grad=True (which _validate_group_options currently accepts), this in-place fill_ runs with grad tracking and raises because the tensor is a leaf requiring gradients. The exception occurs after optimizer.state has already been cleared and replaced, leaving the restore partially committed; perform the fill under torch.no_grad() or reject this destination during staging.
ℹ️ 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".
_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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bb0edcda52
ℹ️ 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".
…nd 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.
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.
There was a problem hiding this comment.
💡 Codex Review
Line 1304 in c5e3121
Restore leaf tensor hyperparameters without autograd tracking
When the supported live tensor LR is a leaf with requires_grad=True, calling GefenDCPState.load_state_dict() outside a no_grad context makes this in-place fill_ raise an autograd error. Because the optimizer state has already been cleared and replaced at lines 1294–1295, the advertised fail-atomic restore leaves the optimizer partially loaded; perform the tensor hyperparameter commit under torch.no_grad() or reject this destination during staging.
Line 1304 in c5e3121
Reject aliased tensor destinations with different values
When multiple parameter groups share the same live tensor LR—a supported idiom exercised in tests/test_lr_item_no_sync.py—but the checkpoint contains different LRs for those groups, these sequential fill_ calls overwrite the same tensor and leave every aliased group with the last saved value. The load reports success with incorrect group hyperparameters after replacing optimizer state; detect aliases whose staged values differ and reject them before commit, or break the alias while restoring.
ℹ️ 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".
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.
There was a problem hiding this comment.
💡 Codex Review
Lines 1320 to 1321 in 4509add
Avoid post-swap fills on grad-tracked hyperparameters
When a live tensor hyperparameter is a leaf with requires_grad=True (for example, a learnable scalar LR accepted by the current constructor validation), this in-place fill_ runs outside torch.no_grad() and raises. At that point optimizer.state has already been cleared and replaced, so the failed load leaves new per-parameter state combined with the old hyperparameters and optimizer metadata, violating the advertised fail-atomic restore. Reject such destinations during staging or complete their fills safely before publishing the staged state.
ℹ️ 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".
…t 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.
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
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
Follow-up to #83. Makes
GefenDCPState's reshardable save memory-bounded, addsasync_savesupport at the same bound, and hardens the resharding path against a batch of silent-failure and fail-atomic bugs found in review. Saved checkpoint bytes are unchanged throughout (payload md5-identical to the pre-change implementation).1. The save is memory-bounded
Resharding must expand Gefen's ~1 byte/param block state to dense fp32 (4 B/param momentum + 4 B/param second moment), because the momentum codebook is learned per rank. The pre-change
state_dict()built every slot's dense pair up front, holding 8 B/param for the whole write — a higher peak than a training step, enough to OOM a model that trains fine.GefenDCPStatenow hands DCP zero-storage stand-ins and expands one slot at a time, when the writer asks. Peak over resident state, 67.1M-param shard:The bound is per slot, so the win grows as state splits across more parameters — every real model. It holds across
FileSystemWriter(thread_count=N), and page-locked memory stays at one slot regardless of thread count.GefenSavePlanneris new public API. Omitting it is not a correctness bug — the save writes the same checkpoint at the old 8 B/param. (It can't fail loudly instead:dcp.loadalso callsstate_dict(), so the hook can't tell an unplanned save from a load.)2.
async_savesupportGefenFileSystemWriterstages each slot to the CPU one at a time, sodcp.async_savekeeps the same 0.5 B/param bound rather than materializing everything. It snapshots at stage time, so a step resumed while the background write runs can't tear the checkpoint (pinned by a gated-writer test). The default writer is refused with a message namingGefenFileSystemWriter.3. Correctness fixes
Silent-failure and fail-atomic bugs, each with a regression test proven to fail without its fix:
_validate_layoutnow runs inside the synchronized region; only the process-group-exists check precedes it.all_reduceover a per-rank identity hash catches it.fully_shardimports fall back to the torch 2.5 path (_composable.fsdp), so the declared 2.5 floor is actually supported (FSDP2 fully_shard import breaks GefenDCPState test and README example on the declared torch 2.5 floor #92).4. Test coverage
New GPU/CPU coverage where the PR made a claim: async staging peak and the unplanned-path cost; page-locked bound across thread counts and ascending slot sizes; the reshard levers left implicit (2→1 collapse,
CPUOffloadPolicy,MixedPrecisionPolicy); and multi-rank abort tests for every synchronized-validation path. The spawn-based drain harness was hardened to distinguish a slow worker from a broken one.Known limit
dcp.async_save's host memory can still grow across successive saves whose largest slot grows — a D2H copy into pageable memory uses a pinned bounce buffer that torch's caching host allocator retains, one layer below this adapter. Documented; not bounded here.Validation
axolotl_gefen(torch 2.12+cu133), idle RTX 3090s (PRO 6000s never touched). Full suite 881 passed / 2 skipped / 0 failed on 4 GPUs; DCP file green on a real torch 2.5.0 floor venv. Payload__0_0.distcpmd5-identical to pre-change (verified against the merged base).ruffclean; no absolute paths committed.