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.
added 4 commits
September 29, 2026 13:17
Pack state-only prefix exports when splitting boundary states would exceed the minimum-call plan. Preserve the no-replay split for shallow forks and validate nested forward and backward equivalence.
Record the three-step comparison for the corrected nested-state runtime and distinguish actor-core gains from AWEX exchange variance. Document the extra core prefix work required by the public FLA interface.
Plan compact tree microbatches globally and exchange CPU rows before Megatron training. This aligns DP slots while preserving row-associated PPO and chunked-loss tensors. Key changes: - Add exact and heuristic prefix-aware tree scheduling - Feed planned groups into Megatron tree packing - Cover row reordering, capacity, and fallback behavior
Replace pre-scheduler performance measurements with the completed three-step Qwen3.5 and DAPO Math run. Record actor gains, AWEX variance, W&B run, and the text-only scope.
yulangz
marked this pull request as ready for review
September 30, 2026 04:45
yulangz
requested review from
CormickKneey,
HwVanICI,
Le8r0nJames,
PrometheusComing,
fishcrap,
garrett4wade,
geshi001,
guozhihao-224,
nuzant,
rchardx and
sitabulaixizawaluduo
as code owners
September 30, 2026 04:45
Record the matched 20-step DAPO Math comparison at the example token limits so reviewers can assess throughput, memory, and timeout behavior alongside the short smoke test.
sitabulaixizawaluduo
added
safe-to-test
Avoid quadratic attention-mask temporaries and repeated tree work on long shared-prefix batches. Keep the 1024-token chunked LM-head loss enabled in the actor benchmark. Key changes: - Build checkpointable compact tree masks and direct trie position IDs - Bound and refine large-batch scheduling and DP group assignment - Fuse tree convolution and request GDN states only when needed - Cover masks, scheduling, states, and CUDA gradients with tests Refs: #1765
Move CPU-heavy prompt preparation and trajectory export off the event loop so pause and destroy requests can complete during large tree rollouts. Bound the destroy RPC timeout to avoid indefinite cleanup. Key changes: - Offload tokenizer and serialization work to worker threads - Add responsiveness and proxy cleanup regression tests Refs: #1765
Exercise multi-turn agent training with a simulated user under the same long-output pressure as the HARTS comparison. Keep external simulator credentials in a file and retain 1024-token chunked LM-head loss. Key changes: - Add the tau3 workflow, task-ID loader, and reward bridge - Provide smoke, performance, and 10-step pressure configurations - Document native-log results and historical W&B run IDs - Test task repetition and simulated-user credential handling Refs: #1765
sitabulaixizawaluduo
requested review from
TaoZex and
daihaowz
as code owners
October 16, 2026 02:15
|
This pull request has been automatically marked as stale because it has not had recent activity within the last 14 days. Please add a comment or push new commits to keep it active. Thank you for your contribution! |
This branch has not been 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 a Qwen3.5-35B-A3B-Base adaptation of HARTS tree execution on the Megatron Bridge path:
enable_chunked_logits=true,lm_head_loss_chunk_size=1024). Shared compact rows accumulate gradients from distinct fork labels.Milestone commits keep planning, chunked loss, model integration, AWEX rank repair, nested-state correction, and benchmark documentation separately reviewable.
Related Issue
No issue specified. Paper: https://arxiv.org/abs/2608.28158v1.
Type of Change
Checklist
pre-commit run --all-files)feature/optimizer_offload/review-prcommand/create-prBreaking Change Details (if applicable): None.
Additional Context
The baseline and the new globally scheduled HARTS run each completed three steps on one eight-L20X node with Qwen3.5-35B-A3B-Base, DAPO Math, DP4/TP2/EP8, four TP2 SGLang servers, two samples per prompt, 128 maximum generated tokens, and chunked LM-head loss enabled. CPU staged offload was disabled in both arms because the installed Megatron-Core 0.17.0 does not meet this branch's staged optimizer compatibility guard.
Training steps are 1.76× faster and actor input-token throughput is 1.63× higher on the two post-compilation steps. End-to-end time is 1.8% higher, while AWEX weight exchange took 91.90 seconds longer. A previous HARTS run before global scheduling measured 3.76 seconds per actor step and 936.20 seconds end to end; these separate short runs have varying rollout samples and AWEX times. Generated rewards differed, so this is an execution/performance smoke test, not a convergence claim. The memory readings are checkpoints, not all-phase peaks.
W&B: baseline and HARTS with global scheduling.
Validation: 15 HARTS core/planner/mapping/convolution tests, 12 chunked/semantic LM-head tests, 3 AWEX controller tests, 46 sequence-packing tests, 3 targeted Megatron VLM tests, direct CUDA FLA gradient checks, 200 randomized trie forward/gradient checks, and a real one-step AWEX smoke, followed by both three-step runs.
mdformat --checkand Ruff passed. PR pre-commit CI passed. Localpre-commit run --all-fileswas run; itsuv-lockhook rewrote both existing lockfiles under local uv 0.11.8 without dependency changes, so those unrelated rewrites were restored. Multi-node integration was not run because no multi-node allocation was available.The new global scheduler passed a two-rank Gloo row-exchange check and 70 targeted CUDA tests. Its three-step GPU benchmark completed and synced to W&B. Public FLA only exports per-sequence final states, so nested forks may pack an additional state-only prefix into the same call. This preserves minimum call depth but can repeat linear-attention core work beyond the paper's selective-state kernel bound. The tested path is text-only Qwen3.5 without context parallelism.
Long-output 20-step comparison
Both Qwen3.5-35B-A3B-Base DAPO Math runs completed 20 PPO/AWEX updates on eight L20X GPUs. Each step trained on 32 prompts x 16 samples, with rollout concurrency 32, queue 512, input limit 2048, generation limit 20480, and request limit 32767. Both retained chunked LM-head loss. SGLang was capped at 32 running requests per TP2 server to avoid the OOM seen in an uncapped attempt.
HARTS was 38.2% slower end to end on these long outputs. Each completed update still contained 512 trajectories. HARTS requires
pad_to_maximum=true, while the baseline usedfalse, so this run does not isolate padding cost from tree execution. Both runs exited successfully and synced to W&B. The targeted CUDA test for shared-row chunked LM-head gradients passed. The example README records the full setup and limitations.W&B: baseline, HARTS.
Follow-up: tree performance and τ³ airline stress
Three additional commits reduce HARTS tree overhead, keep proxy control RPCs responsive during large rollouts, and add a non-solo τ³-Bench airline workflow with smoke and 10-step pressure configurations. The fixed-input actor benchmark keeps the 1,024-token chunked LM-head loss enabled in both modes. Its three steady-state updates on eight L20X GPUs measured median actor steps of 8.855 s baseline and 4.579 s HARTS (1.93×); this excludes rollout and AWEX.
The separate 10-step airline runs used 32 prompts × 16 samples per step, rollout concurrency 32, queue 512, and input/generation/request caps of 2,048/20,480/32,767 tokens. From the native StatsLogger tables, excluding compilation step 1:
HARTS throughput was 1.1835× baseline. Both runs completed 10 updates with chunked LM-head loss enabled; their native logs record 14 and 15 resampled trajectory timeouts, respectively. These were separate live rollouts and used 32,768 versus 65,536 compact-token microbatch caps, so this is an actor-throughput comparison, not a paired replay or end-to-end speedup. W&B baseline and HARTS contain historical replay of numeric native-log metrics, not live logging.
Validation for these commits: 81 focused non-slow tests passed (one optional τ³ test skipped without its checkout), the three τ³ adapter tests passed with the v1.0.1 checkout on
PYTHONPATH, and the compiled CUDA tree-convolution output/gradient test passed. Full pre-commit was run; local uv 0.11.8 regenerated both lockfiles despite no dependency change, so the unrelated lock changes were restored. All other all-file hooks and each commit's installed hooks passed. Multi-node integration remains untested because no multi-node allocation was available.