Harden distributed optimizer integration and release gates - #65
Merged
Conversation
📝 WalkthroughWalkthroughThis release adds deterministic and AMP-safe optimizer stepping, stricter distributed checkpoint validation, CUDA graph counter handling, a Trainer resume harness, reproducible packaging gates, expanded CI matrices, and extensive CPU/GPU regression coverage. ChangesOptimizer correctness and validation
Release and packaging
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Trainer
participant Accelerate
participant GefenOptimizer
participant CheckpointStore
Trainer->>Accelerate: construct and wrap optimizer
Accelerate->>GefenOptimizer: execute staged training
GefenOptimizer->>CheckpointStore: save optimizer and trainer state
Trainer->>CheckpointStore: resume from staged checkpoint
CheckpointStore->>GefenOptimizer: restore optimizer state
GefenOptimizer-->>Trainer: continue training and report digests
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
thad0ctor
changed the title
thad0ctor
changed the title
thad0ctor
changed the title
Consolidated version-2 Parallel-Muon checkpoints refused to load into an optimizer with no eligible distributed process group (single-process resume/eval), even though released-v1 and markerless loads of the same state succeed; rebind the saved owner manifest to live parameters by saved id and run the same per-parameter proof. Harden the Muon grad-presence preflight the way the AMP DTensor preflight already is: content-keyed meshes, sorted mesh and item order so rank-divergent parameter-group order cannot falsely pass the positional activity vector or interleave per-mesh collectives, and no CUDA query for CPU-only mesh optimizers. Refresh the stale codebook-fallback comment and disclose the fused=False CUDA backend change, the eps>0 constructor requirement, and the quantized-momentum codebook refusal in the changelog. Both new gloo regression tests fail on the previous code.
Tag pushes no longer require a registered self-hosted GPU runner: the gpu_release_tests job is replaced by scripts/release_gpu_gate.sh, which downloads the release run's dist artifact via gh, installs the wheel into a cached --system-site-packages venv instead of mutating the host Python, and runs the identical preflight, mandatory GPU test list, and zero-skip enforcement. Approving the testpypi environment is the release manager's attestation that the local gate passed; the hosted CPU and Transformers Trainer wheel gates still block both publish jobs mechanically. The workflow filename, tag trigger, both environment gates, prerelease-skips-PyPI logic, and OIDC permission scoping are unchanged.
thad0ctor
marked this pull request as ready for review
July 12, 2026 22:26
Owner
Author
|
@coderabbitai full review |
✅ Action performedFull review finished. You're currently rate limited under our Fair Usage Limits Policy. Your recent PR review activity is in the 95th percentile or higher among CodeRabbit users, so adaptive limits apply. Your next review will be available in 59 minutes. |
Owner
Author
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/test_muon_grad_presence.py (1)
731-731: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
rank * 0is dead arithmetic — make the intent explicit.
7300 + rank * 0always evaluates to7300, so every rank seeds identically. That is exactly what's required here (both ranks must build the same full tensor beforedistribute_tensor), but therank * 0reads like an incomplete edit and invites a "fix" torank * N, which would silently break collective safety.Suggested clarification
- generator = torch.Generator(device="cpu").manual_seed(7300 + rank * 0) + # Identical seed on every rank: distribute_tensor needs the same full + # tensor across ranks to stay collective-safe. + generator = torch.Generator(device="cpu").manual_seed(7300)🤖 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_muon_grad_presence.py` at line 731, Update the generator seed expression in the relevant test to use the explicit constant seed 7300, removing the dead rank arithmetic while preserving identical tensor construction across all ranks before distribute_tensor.
🤖 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_muon_distributed_checkpoint_safety.py`:
- Around line 587-588: Add a distributed/Gloo availability skipif guard to
test_mixed_parallel_fallback_and_rank_local_composition_cpu, matching the guard
pattern used by sibling distributed tests such as
test_consolidated_v2_checkpoint_loads_into_single_process_optimizer. Ensure the
test skips when torch.distributed or the Gloo backend is unavailable, while
preserving its existing behavior when supported.
---
Nitpick comments:
In `@tests/test_muon_grad_presence.py`:
- Line 731: Update the generator seed expression in the relevant test to use the
explicit constant seed 7300, removing the dead rank arithmetic while preserving
identical tensor construction across all ranks before distribute_tensor.
🪄 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: 7bc28d60-c710-4472-ba88-3eb19dbb7602
📒 Files selected for processing (30)
.github/workflows/ci.yml.github/workflows/release.yml.gitignoreCHANGELOG.mdCOMPATIBILITY.mdCONTRIBUTING.mdMANIFEST.inREADME.mdbenchmarks/README.mdbenchmarks/trainer_resume/README.mdbenchmarks/trainer_resume/__init__.pybenchmarks/trainer_resume/run.pypyproject.tomlscripts/release_gpu_gate.shsrc/gefen/gefen.pysrc/gefen/gefen_muon.pysrc/gefen/hybrid.pytests/test_amp_grad_scaler.pytests/test_capturable.pytests/test_cpu_step_checkpoint.pytests/test_deterministic_mode.pytests/test_dispatch_gating_cpu.pytests/test_epsilon_muon_param_group_cpu.pytests/test_gefen_fsdp2_checkpoint.pytests/test_muon_distributed_checkpoint_safety.pytests/test_muon_grad_presence.pytests/test_step_preflight_atomicity.pytests/test_training_matrix_harness.pytests/test_transformers_trainer_resume.pytests/test_validation_cpu.py
thad0ctor
force-pushed
the
agent/optimizer-production-hardening
branch
from
July 12, 2026 22:42
f82e2ad to
24ec6ec
Compare
Owner
Author
|
Codex Review: Didn't find any major issues. Keep it up! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
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.
What this PR is
The 0.4.0 production-hardening release. Nothing here makes Gefen faster — it makes it dependable enough to bet a long training run on: checkpoints either resume correctly or refuse loudly, multi-GPU jobs fail fast instead of hanging or silently drifting, and a broken build can no longer reach PyPI. The guiding contract: it either works correctly or it stops with an error before touching state — nothing fails silently.
What changed
Determinism (
deterministic=True) — opt-in bit-identical updates on matching GPUs for all three optimizers. This is how you prove a resumed run equals an uninterrupted one, and how you debug: if two runs differ, something real changed. Works everywhere the optimizer is constructed (PyTorch, HF Trainer, axolotloptim_args: { deterministic: true }).Checkpoints can't lie —
distributedcheckpoints carry a versioned owner manifest, validate fully before any mutation, and resume across world sizes — including into a single-process optimizer (fixed during review; the initial format refused that load while v1 allowed it).approxoptimizer state round-trips through PyTorch full-state DCP under FSDP2, same topology, all-ranks-collective. Out-of-scope loads (resharding, changed world size, stripped/corrupted payloads, quantized momentum that lost its codebook) are refused up front — the state is shard-relative, so any "successful" cross-topology load would be fabricating optimizer state.Multi-GPU jobs fail fast instead of hanging — rank-divergent
grad is Nonepatterns are rejected before any Muon collective, order-insensitively (content-keyed meshes, sorted items — hardened during review after an adversarially-verified false-pass/deadlock scenario). AMP protocol selection is rank-synchronized; overflow steps mutate nothing, atomically across both hybrid children.Releases are gated — publishing requires the actual built wheel to pass the CPU matrix (torch 2.5.0 floor + latest), a real Transformers Trainer save/resume gate, and the two-GPU CUDA/JIT/distributed gate with zero skipped tests. The GPU gate runs locally via
scripts/release_gpu_gate.sh <tag>against the release run's artifact (no self-hosted runner exposure); approving thetestpypienvironment is the release manager's attestation it passed. Build tooling is pinned with byte-reproducible rebuild verification.Review round
The branch went through a four-agent review (core optimizer, Muon/distributed, test quality, CI/release) plus an adversarial verification pass; every confirmed finding was fixed with a regression test that fails on the pre-fix code:
deterministic: False-tagged checkpoint into adeterministic=Trueoptimizer is refused (both now tested).fused=FalseCUDA period-search backend,eps > 0, quantized-momentum codebook refusal).Validation
benchmarks/trainer_resume/).v0.4.0.dev1): tag push → build + all hosted gates green → local GPU gate against the downloaded run artifact passed → pipeline paused at the TestPyPI approval exactly as designed. The dry run itself caught and fixed two harness bugs (mp.Queue tensor transport flake,ghrepo resolution).Known scope
FSDP2 full-state optimizer checkpoints are same-topology by design, not caution: plain-Gefen/
approxstate is quantized against per-rank codebooks and per-shard block geometry, so no rearrangement of the saved bytes is valid on other shard boundaries (an approximate dequantize-and-rewarm migration tool is possible future work). Model-only DCP and Muondistributedcheckpoints are unaffected. Megatron's Adam-specific distributed optimizer, FSDP, fp16, and CPU offload remain guarded rather than silently accepted.Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Chores