Repository navigation
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
Make staged AdamW and Muon first-class Megatron configuration while simplifying checkpoint failure handling to rely on process recovery. Preserve AWEX ordinary optimizer migration for non-staged colocation. Key changes: - Add core CLI configuration for staged AdamW and staged Muon - Remove rollback, retry, snapshot, and managed async checkpoint state - Support managed and ordinary optimizer residency in AWEX v1 and v2
AWEX 0.8.1 provides the Qwen3 converter directly, so remove the local attention-name monkey patch and its compatibility-only tests.
Separate optimizer algorithm selection from CPU-staged residency so native and staged AdamW and distributed Muon share one configuration model. Key changes: - Add first-class distributed Muon optimizer configuration - Preserve Megatron ownership while staging AdamW and Muon state - Integrate checkpoint, AWEX, documentation, and dependency support
Lock CPU-staged AdamW and Muon behavior to their native numerical\noracles before consolidating the staging runtime. Cover bucket slicing,\ncheckpoint continuation, and real tensor-parallel Muon collectives.
Consolidate bounded CUDA slots, stream scheduling, residency, and\noffload lifecycle without changing AdamW or Muon update math. Keep thin\noptimizer-specific facades for bucket slicing and matrix ownership.
Delegate AdamW math to PyTorch's functional implementation and bind\nslot-backed Muon proxies to the official TensorParallelMuon step. Preserve\nMegatron group-step and checkpoint semantics around CPU-authoritative state.
Use fused AdamW for small slot workloads and foreach for units with many\nparameter slices. Keep explicit backend modes covered against native\nPyTorch while defaulting to the measured adaptive path.
Exercise the optimizer-specific parallel axes that were missing from the initial numerical parity suite. Key changes: - Validate AdamW through real MCore DP=2 shard ownership and all-gather - Cover dense TP axes, expert TP, and QKV-split Muon updates - Exercise layer-wise Muon owners under DP=2 and DP=2 x TP=2
Refresh native Muon settings before each update so scheduler changes take effect. Select checkpoint counters from nonempty owners and provide a buffer device when binding padding-only AdamW shards. Add numerical schedule parity and empty ownership regression coverage.
Resolve the squashed MOPD history against main while retaining the CPU-staged optimizer changes. Integrate managed optimizer state with shared Megatron residency and preserve sequence layout and AWEX updates.
Use official Muon updates throughout and remove the unused workspace, brittle source checks, redundant residency bookkeeping, and duplicate tests. Keep numerical parity and checkpoint recovery coverage. Publish the tested Qwen3.5 Base configuration and document the measured actor and whole-GPU memory peaks separately. Refs: #1613
Account for aligned batch size times maximum sequence length when splitting padded Megatron inputs so padding cannot exceed the token budget. Preserve this cost model when synchronizing microbatch counts. Enable chunked LM-head loss in the Qwen3.5 staged-offload example to bound logits memory with the normal training batch. Validation: 53 microbatch tests passed; Qwen3.5 PPO profiling completed GPU work with the normal batch and chunked loss enabled.
Merge the dense AWEX transfer-rank fix and lockfile cleanup from the optimizer offload base while retaining only the Muon dependency additions in the stacked lockfiles.
yulangz
had a problem deploying
to
AReaL-unittests
September 22, 2026 04:59 — with 
GitHub Actions
Failure
GitHub Actions
Failure
yulangz
had a problem deploying
to
AReaL-unittests
September 22, 2026 04:59 — with
GitHub Actions
Failure
GitHub Actions
Failure
yulangz
had a problem deploying
to
AReaL-unittests
September 22, 2026 06:55 — with
GitHub Actions
Error
GitHub Actions
Error
added 13 commits
September 22, 2026 13:22
Keep staged AdamW compatible with Megatron Core 0.19 and reconcile AWEX optimizer residency with the latest main changes.
Merge the updated AdamW offload base and follow Megatron Core 0.19 Muon APIs so native and staged layer-wise optimizers retain their ownership, numerical parity, and checkpoint behavior.
Use the native layer-wise Muon builder for optimizer classification and DP ownership. This removes the 0.17 compatibility path and the manual scalar-only LayerWise assembly while preserving the staged checkpoint leaf layout. Key changes: - Map AReaL dist_muon to MCore muon and enable layer-wise ownership - Build native and staged leaves through get_megatron_optimizer - Cover scalar-only and all-expert distributed edge cases Refs: #1743
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
Carry forward the Megatron 0.19 test fixtures and main integration fixes from the AdamW branch. Preserve model-specific optimizer group overrides through both native and staged Muon construction. Verify no-decay groups after scheduler initialization with independent native and staged configurations. Refs: #1743, #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
Merge the validated AdamW branch to use main AWEX 0.9.0 and its CI runtime dependency alignment while retaining the Muon dependency install step. Validate 42 Muon, optimizer config, checkpoint and AWEX tests locally. Refs: #1743
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.
Carry HF model-only state reset and AWEX transfer-rank fixes into the stacked Muon PR. Remove activation tracking after splitting it into #1779. Validate shared HF/checkpoint/config/Muon regressions: 61 passed, 2 existing async-save fixture skips.
Le8r0nJames
reviewed
Oct 13, 2026
sitabulaixizawaluduo
requested review from
TaoZex and
daihaowz
as code owners
October 16, 2026 02:15
Resolve conflicts after the AdamW offload PR merged into main. Enable native Muon checkpoint construction and allow storage allocation addresses to change across AWEX release/resume while retaining logical ownership checks. Exercise native and staged checkpoint round trips and post-resume updates under real DP/TP process groups. Refs: #1743
yulangz
added
safe-to-test
This branch was successfully deployed
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
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.
Description
Add
optimizer.type=dist_muonon Megatron-Core 0.19.0, with native GPU state or CPU-staged state. This builds on the AdamW offload infrastructure merged in #1613 and now targetsmain.emerging-optimizers==0.3.0.Merged current
mainand removed inherited AWEX/runtime changes already handled there. CI installs the new Muon dependency into an isolated directory used by both installation validation and tests; it does not resync AWEX.Validation
pre-commit run --all-files: passed.Multi-node tests were not run because this machine is a single node. A staged DCP topology-change restore was not exercised in this validation run.