Repository navigation
feat(engine): add CPU-staged Megatron AdamW - #1613
Conversation
Enable policy distillation from multiple teacher objectives while sharing colocated training and rollout resources. Key changes: - Add MOPD objectives, routing, losses, and teacher lifecycle management - Support forked colocated workers across local, Ray, and Slurm schedulers - Extend AWEX weight synchronization and memory residency handling - Add bilingual documentation, a local GSM8K example, and focused tests
Forked colocated roles inherited only the owner environment, so rollout allocator settings could not differ from actor and teacher settings. Pass per-task overrides through every scheduler fork path and expose the active allocator configuration in residency logs.
Free CUDA graph allocations alongside weights and KV cache while colocated training owns the GPU. Enable SGLang's graph memory saver before capture and restore graphs after AWEX weight publication. Key changes: - preserve the preload hook required for pauseable CUDA graphs - pass the cuda_graph tag through rollout offload and restore - cover child environment and restore ordering with unit tests
Gate rollout recovery on the v1 AWEX colocate path and align regression tests with the current SGLang pause and Guard lifecycle contracts.
Regenerate the CLI references from the merged SchedulingSpec so the MOPD GPU GRES setting and the new Slurm reservation options are documented together.
Model the v1 actor configuration explicitly and isolate Guard port allocation tests from node-wide abstract socket leases.
Resolve SGLang base GPU IDs in the worker's logical device namespace so Local and Slurm colocation select non-overlapping devices without scheduler-specific overrides. Preserve registered AWEX converter fusion policies to keep dense and Qwen3-MoE train/inference metadata aligned.
Keep AdamW and Muon optimizer state in pinned CPU slabs and stream bounded GPU updates. Add transactional synchronous checkpointing with disk-backed rollback, managed AdamW asynchronous save and finalization, and Muon DP resharding. Integrate AWEX residency handling and comprehensive unit and torchrun coverage.
Provide a reproducible eight-GPU AWEX setup for comparing standard and CPU-staged Megatron optimizers without changing precision-aware optimizer semantics. Key changes: - Add a Qwen3-30B-A3B Attention-DP8 and Expert-EP8 GRPO example - Add an example-local CPU-staged toggle and worker configuration - Capture step-2 CUDA allocator snapshots with tests and usage docs
Keep representative numerical, residency, checkpoint, AWEX, Muon, and async-save paths while removing exhaustive fault matrices and unused torchrun workers. Reuse the existing multi-turn GSM8K workflow in the CPU-staged example so the PR stays below the automated review size limit.
Make the CPU-staged Qwen3 MoE recipe reusable with a local DAPO-Math loader and proxy-based rollout flow. Ensure AWEX releases precision-aware optimizer master state before colocated inference resumes. Key changes: - Add and register a DAPO-Math-17k RL dataset loader - Rename and sanitize the CPU-staged example configuration - Preserve optimizer-state tensor identity during AWEX migration
Route DAPO-Math prompt filtering through the shared helper so base tokenizers without a native chat template use the configured compatibility path.
Prevent Megatron activations, optimizer staging buffers, statistics, and CUDA IPC exports from overlapping the resumed rollout engine. Account for sequence alignment before packing to keep micro-batches under their token cap. Key changes: - release retained activation and optimizer staging storage - defer CUDA IPC exporter cleanup until all readers finish - flush actor statistics before the colocation handoff - account for sequence alignment during micro-batch allocation
Preserve main training metrics, Qwen4 optimizer group overrides and CUDA IPC allocator handling alongside staged optimizer residency. Update the DP1 and DP2 fixtures for Megatron Core 0.19 buffer metadata and bucket-group parameter synchronization. Preserve the parameterized per-token loss configuration instead of overwriting it. Refs: #1613
Merge main to adopt AWEX 0.9.0 in both dependency manifests and lockfiles. Install the checkout-declared AWEX version in GPU CI so older runtime images can collect and run the updated tests. Validate 5325 SGLang CI test collections, 174 AWEX/offload tests and DP2 staged AdamW parity locally. Optional DTE cases remain skipped. Refs: #1613
Le8r0nJames
left a comment
Le8r0nJames
left a comment
There was a problem hiding this comment.
In addition to the points above, could you report optimizer-step and full training-step timings with CPU staging enabled versus disabled under the same workload? The memory savings are clear, but the CPU–GPU transfer overhead and throughput trade-off are not yet quantified.
Reuse the checkpoint load lifecycle to refresh CPU masters and clear moments after a model-only HF load. Prevent updates after a partial load failure and preserve scheduler progress. Validate the next AdamW update against a fresh optimizer with staging slots both resident and released.
Subtract the SGLang server base from scheduler GPU ordinals before adding the dense transfer-rank base. Cover isolated and fully visible GPU layouts across two nodes.
Remove activation storage patching and release from the optimizer-offload PR. Preserve the feature in its own branch for separate ownership and memory-lifecycle review.
The red is baseline and the green is with optimizer-offload.
|
|
Updated the code and addressed the correctness feedback:
Independent activation-storage tracking/reclamation has been removed from this PR and moved to PR #1779, based directly on main, for separate ownership/safety review. Muon support remains in the stacked PR #1743. The redundant AWEX packing clone and post-insertion size-limit check are reported upstream in inclusionAI/Awex issue #128, with reproducible CPU/CUDA cases and proposed fixes. A 20-byte limit currently yields a 24-byte group; packing 256 MiB has a 512 MiB extra allocation peak with Latest local validation: 109 related tests passed, including ACK success/failure/timeout lifetime checks; 10 consecutive exchanges on two GPUs passed with actual MetaServer/CUDA IPC/native AWEX reader/bounded NCCL and exact received weights, while checking exports were reclaimed before final engine completion; full pre-commit passed. The two-GPU exchange validation does not replace full SGLang training or multi-node convergence testing. |
|
Thanks for addressing the correctness issues and splitting out activation tracking. Could we also remove the local packer workaround from this PR and help land the fix upstream in inclusionAI/Awex#128? |
|
Addressed the latest review feedback and pushed through f4a707c:
The AWEX fix is now PR #129 for issue #128. Upstream local CPU CI: 279 passed, 1 skipped; tensor/CUDA IPC tests: 31 passed. GitHub requires maintainer approval before the upstream PR's CI runs. Activation-storage tracking remains in separate formal PR #1779; Muon remains in stacked PR #1743. Updated this PR description to reflect that split and the native-AWEX result. CI for the latest head was retriggered through the PR's |
|
The failure in run 38379205319 is a runner ownership race, rather than a reported pytest failure. The SGLang suite completed successfully: 5309 passed, 187 skipped, 397 deselected. The vLLM job was assigned to Pushed isolated CI fix fdb3955. Fix: scope the provisioned runner's machine label and the test job's selector to Validation: actionlint 1.7.12 passed with ShellCheck disabled; YAML/routing checks cover correct matching, cross-run/prior-attempt/backend exclusion, exclusion of old generic selectors, and identical create/cleanup names. Full pre-commit passed. Also merged latest main CI was retriggered through the PR's Evidence: current vLLM job, older run's VM teardown, successful SGLang job. |
|
Resolved the main conflict and fixed the actual failures from run 38392961547 in two separate commits:
Both previous GPU jobs ran on the correct isolated runners and completed pytest. They had the same five parameterized recovery failures ( Local validation after the merge/fix:
The PR is now conflict-free and head 99e9507 is pushed. CI was triggered through the PR's safe-to-test label; the new GPU run is visible and in progress: 38677858683. It has no final result yet. |
Use the recovery fixture from main now that R3 includes the fix. Initialize R3 state in the existing chunked-forward test fixture.
|
Merged R3 main The conflict was the duplicate Existing Qwen3.5 chunked-forward tests now initialize the three new R3 state fields, matching the existing upstream fixture pattern. No new production workaround was added. Validation: 259 R3/staged-optimizer/chunked-forward/recovery tests passed; full pre-commit passed. No fresh multi-node run was performed. Both GPU variants of the preceding head's CI had passed (38677858683). Head 59f5705 is pushed and GitHub reports no conflicts. Triggered CI via the PR's |
Description
Keep authoritative FP32 AdamW master parameters and moments in pinned CPU memory and stream updates through bounded reusable CUDA buffers. The staged optimizer integrates with Megatron checkpointing and AWEX colocation while preserving ordinary optimizer residency behavior.
Muon support is split into stacked #1743, based on this branch. Activation-storage tracking is split into independent #1779; ordinary Megatron training no longer installs its patches or mutates activation storage in this PR.
Examples
examples/cpu_staged_offload/contains DAPO-Math configurations for Qwen3-30B-A3B-Base and Qwen3.5-35B-A3B-Base, sharing one training entry point. The Qwen3.5 configuration defaults to three training steps on eight GPUs with attention TP2/DP4, MoE EP8, and four TP2 rollout servers.Repeat with
trial_name=baseline actor.megatron.cpu_staged_offload.enabled=falsefor the baseline. Precision-aware mode stays enabled in both runs.Validation
Current review fixes
Local validation:
pre-commit run --all-files: passed.The v1 Qwen3.5-35B-A3B-Base default example completed three PPO updates and three AWEX updates (versions 1–3), exit code 0, on eight L20X GPUs using the unmodified installed AWEX 0.9.0 packer. No local workaround or upstream patch was loaded: the old two
clone()calls and post-insertion size check remain in the tested dependency. Each update completed on all eight writers/readers.Configuration: batch size 8, two samples/prompt, max generation 128, max tokens 2047, microbatch token budget 4096, chunked logits enabled; attention TP2/DP4, MoE EP8, four TP2 rollout servers. Trial:
qwen35-native-awex09-no-packing-fix-r3-20261013. Whole-GPU peak observed by 1-second NVML sampling: 81.32 GiB. No OOM occurred. This default workload does not require the upstream packing fix to run; tighter memory budgets and larger workloads were not validated. This is a short single-node execution check, not a convergence or v2 model-residency validation.Latest completed remote validation: both SGLang and vLLM GPU CI, pre-commit, and package installation passed at
99e95070(run 38677858683). The isolated runner labels/VM identities prevent cross-run cleanup from deleting another PR's test worker.After merging native R3 from main
6cad8710, this PR uses main's Flash Next recovery fixture unchanged; the earlier duplicate fixture fix is no longer part of the PR diff. Existing Qwen3.5 chunked-forward tests initialize the new R3 state in their fake engines. Local R3/staged-optimizer/chunked-forward/recovery regressions: 259 passed; full pre-commit passed. No fresh multi-node integration run was performed for this merge.Historical integration measurements
The following smoke measurements precede the activation-tracker split and describe the earlier integration snapshot.
Qwen3-30B-A3B and Qwen3.5-35B-A3B-Base smoke runs completed three PPO updates and three AWEX weight updates on one eight-L20X node.
Actor allocated memory decreased by 52.4%. Whole-GPU peaks include rollout and weight exchange and were essentially unchanged in these short execution checks.
After splitting Muon:
pytest -q tests/test_gpu_staged_optimizer.py tests/test_gpu_staged_optimizer_awex.py tests/test_megatron_optimizer_config.py tests/test_megatron_async_save.py tests/test_cpu_staged_offload_example.py: 53 passed, 2 skipped (CUDA-only AWEX residency)pre-commit run --all-files: passedThe full multi-node suite was not run because validation used a single node.