Conversation
Introduce src/gefen/offload.py (StateOffloadMixin) and wire it into Gefen and GefenMuon with a minimal, fully flag-gated core integration: - move_state_(device=None): fail-atomic movement of authoritative optimizer state to the live parameter devices or an explicit co-located device. - offload_state_() / restore_state_(): synchronous CPU-authoritative optimizer state for plain replicated-CUDA Gefen; each step stages state to the parameter device, runs the ordinary (non-fused, non-batched) update, and copies it back to CPU while the shared codebook stays CUDA-resident. All new behavior is gated on state_offload_active (default off), so existing single-GPU and distributed (FSDP2/ZeRO/DDP) paths are behaviorally unchanged. Native offload supports plain Gefen with ordinary replicated CUDA parameters and fails closed under DTensor/FSDP sharding, rank-local sharded state, capturable/CUDA-graph execution, custom tensor-valued state, and torch.compile. Activation, restore, move, and checkpoint load are staged and validated (device, dtype, shape, tight contiguous storage, no aliasing) before the live mapping is swapped, so a copy failure is a true no-op. A failed post-update CPU copyback marks the optimizer poisoned and refuses further step/export; loading a known-good checkpoint clears the poison unconditionally (regardless of whether offload is still active) and co-locates a map_location='cpu' codebook back onto the parameter device. tests/test_state_offload.py adds 17 tests; full suite green on RTX 3090.
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 2 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughGefen adds atomic optimizer-state movement and synchronous CPU offload for replicated CUDA parameters. Offload-aware stepping, checkpoint reconciliation, poison recovery, codebook placement, documentation, and CUDA/DDP validation are included. ChangesOptimizer state movement
CPU-authoritative stepping
Checkpoint and offload lifecycle
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant GefenStep
participant StateOffloadMixin
participant CUDAParameter
participant CPUOptimizerState
GefenStep->>StateOffloadMixin: validate offload readiness
StateOffloadMixin->>CPUOptimizerState: stage state for update
CPUOptimizerState->>CUDAParameter: copy state to parameter device
GefenStep->>CUDAParameter: perform optimizer update
CUDAParameter->>CPUOptimizerState: copy updated state back
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![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 `@src/gefen/offload.py`:
- Around line 451-463: Update the offload activation method containing
_stage_all_parameter_state_to_cpu so the codebook is moved to CUDA within the
same atomic state update before marking offload active. Clear or refresh the
codebook’s device caches as part of that transaction, preserving rollback
behavior on failure, and add coverage for an inactive CPU-mapped load followed
by offload_state_() and a successful step().
🪄 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: 43cc43ff-cd13-49d7-8269-84a29fd87b37
📒 Files selected for processing (5)
CHANGELOG.mdREADME.mdsrc/gefen/gefen.pysrc/gefen/offload.pytests/test_state_offload.py
![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: 5e4ef1c6af
ℹ️ 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".
…pped load offload_state_() validated with require_cpu_state=False, which permits a CPU-resident codebook left behind by an inactive map_location='cpu' load, but then marked offload active without relocating it. The next step() tripped the "codebook must remain CUDA-resident" guard, bricking the common resume-then-enable-offload flow. Stage the codebook onto its parameter device inside offload_state_'s atomic update (rollback-safe: staging precedes mutation) and clear the stale per-device caches. Factor the staging into _stage_offload_codebook_on_parameter_device and reuse it from the load-path co-location helper. Adds test_offload_activation_colocates_cpu_mapped_codebook.
tests/test_state_offload_distributed.py: a two-rank NCCL DDP test asserting the offloaded run is bit-identical to a non-offloaded run stepping the same all-reduced gradients, with per-parameter state staying CPU-resident across steps (block and factored second-moment paths). Gated on two CUDA GPUs + NCCL. tests/test_state_offload.py: test_offload_end_to_end_training_loss_matches_baseline drives a real forward/backward/step loop over a small MLP and asserts the offloaded loss trajectory matches the non-offloaded baseline bit-for-bit. README notes that offload is validated bit-identical under DistributedDataParallel.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/test_state_offload_distributed.py (1)
44-108: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winSetup/init failures bypass the error-reporting path.
torch.cuda.set_device(rank)anddist.init_process_group(...)(Lines 44-45) run before thetryat Line 46, so a failure there is never written toresult_queueandfinally'sdist.destroy_process_group()never executes for that rank. If one rank fails here while the other blocks inside the same collectiveinit_process_groupcall, the parent's diagnostic ends up as a bare timeout ("DDP offload workers timed out") instead of the actual traceback, and the failure takes the fullresult_queue.get(timeout=180)+join(timeout=180)budget to surface instead of failing fast.Widen the try/except/finally to cover setup, and guard the cleanup call:
♻️ Proposed fix
os.environ["WORLD_SIZE"] = str(world) - torch.cuda.set_device(rank) - dist.init_process_group("nccl", rank=rank, world_size=world) try: + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world) reference_model = _build_model() @@ finally: - dist.destroy_process_group() + if dist.is_initialized(): + dist.destroy_process_group()🤖 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_state_offload_distributed.py` around lines 44 - 108, Move the try block in the distributed worker to encompass torch.cuda.set_device and dist.init_process_group, so setup failures reach the existing BaseException reporting path. Guard dist.destroy_process_group in the finally block with a check that initialization succeeded before calling it, while preserving traceback reporting through result_queue for rank 0.
🤖 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.
Nitpick comments:
In `@tests/test_state_offload_distributed.py`:
- Around line 44-108: Move the try block in the distributed worker to encompass
torch.cuda.set_device and dist.init_process_group, so setup failures reach the
existing BaseException reporting path. Guard dist.destroy_process_group in the
finally block with a check that initialization succeeded before calling it,
while preserving traceback reporting through result_queue for rank 0.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: be011f8f-6d0d-4204-b56d-34461809ed3a
📒 Files selected for processing (3)
README.mdtests/test_state_offload.pytests/test_state_offload_distributed.py
🚧 Files skipped from review as they are similar to previous changes (2)
- README.md
- tests/test_state_offload.py
Move torch.cuda.set_device and dist.init_process_group inside the worker's try/except so a setup or process-group init failure reaches the result_queue reporting path (surfacing the real traceback) rather than manifesting as a bare parent-side timeout. Guard destroy_process_group with dist.is_initialized() so cleanup does not itself raise when init never completed.
|
Fixed in d39070a. Moved |
|
What is this PR intended to achieve? Offloading in inference terms usually means supplementing VRAM with ordinary RAM at the cost of speed to reduce or avoid OOM errors. Is that true for this PR too? If not, you may wish to reconsider the specific terminology to avoid confusion. |
More or less, I was looking at ways to the implement optimizer-state offload similar to paged AdamW, it reduces persistent VRAM usage by keeping Gefen’s optimizer state in ordinary CPU RAM. was going to make branches for a few different approaches to test if any are worth keeping for an even more aggressive VRAM savings. |
2 participants
Summary
Adds a modular optimizer-state movement engine (
src/gefen/offload.py,StateOffloadMixin) and wires it intoGefen/GefenMuonwith a minimal, fully flag-gated core integration:move_state_(device=None)— fail-atomic movement of authoritative optimizer state (momentum codebook indices, magnitudes, second moments, counters) to the live parameter devices, or to an explicit co-locateddevice. Available on bothGefenandGefenMuon.offload_state_()/restore_state_()— synchronous CPU-authoritative optimizer state for plain replicated-CUDAGefen: persistent per-parameter state lives on CPU between steps; each step stages it to the parameter device, runs the ordinary (non-fused, non-batched) update, and copies it back to CPU. The shared codebook stays CUDA-resident.Everything is opt-in and gated on
state_offload_active(default off), so existing single-GPU and distributed (FSDP2 / ZeRO / DDP) training paths are behaviorally unchanged — the coregefen.pydiff is 32 lines.Scope and safety
Native CPU offload supports plain
Gefenwith ordinary replicated CUDA parameters (single-GPU and DDP-style data parallelism). It fails closed (raises before mutating anything) under:torch.compile, and custom tensor-valued state extensions.Sharded offload requires a separate paging policy and is left as future work.
Robustness
load_state_dict()clears the copyback-poison flag unconditionally — recovery no longer depends on offload still being active at load time (previously,restore_state_()before loading left the optimizer permanently bricked).map_location='cpu'while offload is active now co-locates the shared codebook back onto the parameter device, so the next step keeps the codebook CUDA-resident (previously it raised "codebook must remain CUDA-resident").Testing
tests/test_state_offload.py— 17 tests (bit-exact offloaded stepping across factored + block second-moment paths, activate-before-first-step, periodic codebook refresh, activation/restore/move failure atomicity, cross-device movement forGefen/GefenMuon, active-offload checkpoint load, copyback poison stickiness + recovery, CPU-mapped codebook co-location, persistent-storage alias rejection).Validated on an RTX 3090 (CUDA 12.x):
test_state_offload.py: 17 passed.test_native_load_atomicity,test_hybrid_load_atomicity,test_cpu_step_checkpoint,test_gefen_fsdp2_checkpoint,test_muon_distributed_checkpoint_safety,test_state_offload): 107 passed, 3 skipped.Summary by CodeRabbit
move_state_().offload_state_()/restore_state_(), with per-step readiness/validation.map_location='cpu'checkpoint behavior to restore CUDA codebook placement when offload is enabled later.