[yulangz]

Superseded by two focused PRs:

• feat(engine): integrate WeftMoE training and inference with AdamW #1891 targets main: common WeftMoE integration, staged AdamW, AWEX/SGLang, and experiment screenshots.
• feat(engine): adapt CPU-staged Muon to WeftMoE #1892 targets feature/muon-optimizer-offload: only the optional Muon writer/master hooks and associated tests/documentation.

The replacement branches merge automatically and passed single-node 8-GPU AdamW/Muon BF16/FP8 adapter checks. Original description and review history are retained below.

---

Description

Integrate WeftMoE into AReaL's v1 colocated Megatron/SGLang path. Add a DAPO-Math launch script and configuration for baseline, WeftMoE BF16, and WeftMoE blockwise FP8.

• Install Megatron and SGLang adapters and connect optimizer-update, checkpoint, offload, and teardown hooks.
• Publish staged Muon/AdamW FP32 masters into FP8 W/W-transpose; transfer W bytes and block scales through AWEX's existing canonical mapping. Rebind SGLang receive views after offload and publish the received version after transport completes.
• Preserve non-flat expert storage identities during offload, including master-backed BF16 restoration.
• Honor Megatron-Bridge's configured Flex/DeepEP dispatcher for the baseline; publish the actual attention TP size to AWEX when DP attention is enabled.

The supported FP8 configuration uses CPU-staged Muon or AdamW, expert DP=1, and expert TP=1. AdamW requires distributed and precision-aware optimization and the matching WeftMoE #6. Each expert matrix must fit in one staging slot; the writer consumes the updated complete GPU FP32 matrix before that slot is recycled. DTE and v2 FP8 weight transfer are outside this PR.

AdamW support adds a master-writer callback, preserves whole selected matrices when building staging units, and skips BF16 parameter offload/reload for FP8-managed expert buffers in both extended and fallback buffer APIs. The optimizer update formula is unchanged. The example README includes the AdamW overrides; the supplied screenshots remain Muon runs.

Related Issue

Depends on #1743 (CPU-staged Muon), whose branch now includes main at de8ae341.
This PR targets feature/muon-optimizer-offload and includes its updated tip
386337df as an ancestor. The Files changed tab now contains only the integration:
22 files, 1,447 additions, 34 deletions, including the three supplied screenshots.
The lockfiles match the Muon base; unrelated package-index and dependency changes
have been removed.

After #1743 merges into main, rebase/clean and retarget this PR to main.
The Muon branch itself was not changed by this cleanup.

Type of Change

• ✨ New feature

Validation

The following tests ran with the existing Muon-capable Megatron environment: 247 passed, 2 skipped. The two skips are pre-existing async-save metric tests marked skipped for fixture isolation; CUDA IPC and GPU Muon tests ran.

python -m pytest -q \
  tests/test_gpu_staged_optimizer.py \
  tests/test_gpu_staged_muon.py \
  tests/test_gpu_staged_muon_checkpoint.py \
  tests/test_megatron_optimizer_config.py \
  tests/test_awex_weftmoe_fp8.py \
  tests/test_megatron_bridge_deterministic.py \
  tests/test_awex_sglang_plugin.py \
  tests/test_megatron_weight_storage.py \
  tests/test_megatron_async_save.py \
  tests/test_binary_kl_divergence.py
UV_DEFAULT_INDEX=https://pypi.org/simple UV_INDEX_URL=https://pypi.org/simple pre-commit run --all-files
bash -n examples/weftmoe/run.sh

Also loaded all three variant configurations and the 8K/65536-token override, checking actor/rollout flags and baseline dispatcher selection. Tests cover FP32 master callbacks with released BF16 storage, weight/scale mapping and IPC, receive-buffer rebinding across repeated updates, failed-transfer publication ordering, and non-flat storage restoration.

Also validated all three AdamW launch configurations. With the matching WeftMoE #6 code, single-node 8-GPU Megatron tests passed for staged AdamW BF16, staged AdamW FP8, staged AdamW FP8 with fallback buffer residency, and the Muon FP8 regression (TP2/EP8; two layers; hidden 3072, intermediate 1024, 256 experts, top-k 8, 128 tokens/rank; two optimizer updates and full recomputation). The tests check zero BF16 expert storage, exact FP8 W/W-transpose bytes, offload/reload, and ordinary BF16 updates after adapter close. WeftMoE lifecycle/validation unit tests: 22 passed.

Checkpoint checks cover inner AdamW master/m/v state restoration and BF16 model materialization, not a full Megatron distributed-checkpoint roundtrip. No new full-model or multi-node RL run was launched for the rebased PR. The screenshots below are from the previously recorded experiments, not a new run of this commit.

Checklist

• Pre-commit hooks pass (pre-commit run --all-files)
• Relevant tests pass; new tests added for new functionality
• Example documentation updated and Markdown formatted (no docs-site changes)
• Branch contains the updated Muon base 386337df and main at de8ae341
• This PR was created by a coding agent via create-pr
• This PR is a breaking change

Experiment screenshots

All comparisons use DAPO-Math-17k, train split, GRPO, 32 prompts × 16 completions, temperature 1.0, top-p 1.0, and Muon CPU offload with learning rate 1e-6. Each run uses one node with 8× NVIDIA L20X (SM90, approximately 140 GiB/card). Training uses TP2 × DP4, EP8, ETP1, CP1, PP1, BF16 model parameters, and full-layer recomputation. No speculative decoding or MTP training is enabled.

The plots show token counters, task reward, update log-probability differences, and rollout/training/weight-update timers. The two 16K screenshots also include estimated training FLOPs/s. Baseline uses native Megatron Flex/DeepEP and native SGLang MoE; the other curves use WeftMoE BF16 or FP8 routed experts.

1. Qwen3.5-35B-A3B-Base — 8K

• Generation total-token cap: 8191; SGLang context length: 8192.
• Training microbatch token limit: 65536.
• Rollout: DP8/EP8 with DP attention enabled for all three runs.
• Configured for three epochs with checkpoint saving disabled; the recorded runs ended at 111/124/67 completed steps for baseline/BF16/FP8.
• Green = baseline; purple = WeftMoE BF16; pink = WeftMoE FP8.

2. Qwen3.5-35B-A3B-Base — 16K

• Generation total-token cap: 16383; SGLang context length: 16384.
• Training microbatch token limit: 16384; configured step cap: 20.
• Baseline rollout: TP8/EP8, DP attention disabled. WeftMoE BF16/FP8 rollout: DP8/EP8, DP attention enabled.
• Pink = baseline; orange = WeftMoE BF16; green = WeftMoE FP8.

3. Qwen3.5-122B-A10B — 16K

• Model: Qwen3.5-122B-A10B, non-Base checkpoint.
• Generation total-token cap: 16383; SGLang context length: 16384.
• Training microbatch token limit: 16384; configured step cap: 20.
• Rollout: DP8/EP8 with DP attention enabled for all three runs.
• Blue = baseline; light yellow/orange = WeftMoE BF16; light blue = WeftMoE FP8.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.