Skip to content

[megatron] Qwen3.8: row-shard the PLE table over EP (any EP size) - #2533

Closed
xinyuangui2 wants to merge 50 commits into
NovaSky-AI:mainfrom
xinyuangui2:qwen38-ple-ep-rowshard
Closed

[megatron] Qwen3.8: row-shard the PLE table over EP (any EP size)#2533
xinyuangui2 wants to merge 50 commits into
NovaSky-AI:mainfrom
xinyuangui2:qwen38-ple-ep-rowshard

Conversation

@xinyuangui2

@xinyuangui2 xinyuangui2 commented Oct 25, 2026

Copy link
Copy Markdown
Collaborator

Stacked on #2513. Only the last commit (9323dd5) is new; the rest of the diff is #2513.

Qwen3.8 PLE n-gram table: split it across the EP group by contiguous rows instead of by n-gram head, so EP no longer has to divide the 16 heads (EP32/64 for the 512 experts).

  • load_table: each EP rank loads [N·r/EP, N·(r+1)/EP) of the frozen table into pinned host memory (~100 GB / EP per rank).
  • _lookup: each id goes to the rank that owns its row by all-to-all, and the rows come back the same way (as EPShardedEmbeddingTable in Megatron-LM#7553). The table is frozen, so there is no backward pass.
  • New CPU test: 3 gloo ranks, with a different number of tokens on each, match the unsharded PLE.

🤖 Generated with Claude Code

https://claude.ai/code/session_015m9Uy57LhVrqUZwKFH2VdR

xinyuangui2 and others added 30 commits October 17, 2026 09:15
@xinyuangui2 @claude
[megatron] Qwen3.8-Flash-Next (qwen4_exp) bridge WIP
3f6dd1f 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[megatron] qwen4_exp: accept transformers' qwen_sparse_attention laye…
c1ebb5b 
…r type; test health timeout

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[megatron] qwen4_exp: resolve GDN gate like HF; CPU test for the prov…
9fd0465 
…ider

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[megatron] qwen4_exp: register Qwen4ExpGatedDeltaNet parallelism type
6105c0e 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[megatron] qwen4_exp: keep SiLU on the GDN conv, sigmoid only on the …
1850880 
…output gate

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[vllm] Preserve model_runner.model_state across level-2 sleep
b236f99 
Qwen4Exp's ngram_context_offsets lives in the weights pool and came back as
zeros after a level-2 sleep, corrupting PLE n-gram contexts after every
weight sync.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[test] qwen3.8 parity row: compare common prefixes; document frozen P…
0137cb0 
…LE in recipe

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[megatron] qwen4_exp: keep the frozen PLE table shard in pinned host …
c41c05b 
…memory

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[examples] qwen3.8 recipe: adapter-only LoRA sync, rank 16, vLLM gpu_…
6ae2f43 
…memory_utilization 0.6

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[examples] qwen3.8 recipe: cap max_num_seqs at 512
7eac53f 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
Forward NCCL_* overrides to Ray workers; qwen3.8 recipe disables NVLS
cba683a 
NVLS buffers stay resident through offload/sleep (~24 -> ~12.5 GiB/GPU trainer
residual at EP16), which blocked the adapter load at step 2.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
Format (ruff, black)
356e54e 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[megatron] qwen4_exp: correct upstream PR references
a04f498 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Hp5LxobgMKPQJkiDxEXcef
@xinyuangui2 @claude
[test] qwen3.8 slice: make_slice.py with a small random PLE table; pa…
a39ee46 
…rity row on 4xH100 (EP4, vLLM TP4)

make_slice.py replaces the ~51B-parameter PLE n-gram table with a small random one
(ngram_vocab_size_base 2048, head sizes/offsets recomputed with the HF prime formula), so the
4-layer slice is 23 GB and fits the 4xH100 CI shape. Parity on 4xH100: forward |dlogp| 0.051,
post-sync common prefix 0.021. The gsm8k recipe is renamed to the _lora_2node variant.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
[test] qwen3.8 slice: LoRA adapter-sync row on 4xH100 (EP4)
ec5e8e1 
Zero adapter 0.052, stale sampler 0.239, updated adapter 0.056 (threshold 0.1).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
[megatron] qwen4_exp: QSA sparse indexer (sequences past 2051 tokens)…
7cfaa00 
…, FlexAttention fast path

Megatron previously ran the QSA indexer only in its exactly-dense regime (<= indexer_budget +
compress_ratio - 1 = 2051 tokens) and raised beyond. Qwen4ExpQSAIndexer now mirrors HF's
selection: per query, the top indexer_budget/r of the complete r-token blocks (fp32-mean pooled
keys, k_layernorm, RoPE at the block start, sum_h relu(q_h.k)/sqrt(d)) plus the incomplete tail,
per packed sequence. qsa_sparse_attention attends over the selection: FlexAttention on CUDA
(~3x faster than SDPA's math kernel, the only one it has for head_dim 256 with a bool mask),
chunked SDPA otherwise. Short sequences keep the TE dense kernel.

Tests: CPU selection/attention/grad parity vs HF (26 pass); GPU vLLM-vs-Megatron logprobs on
~7.7K-token prompts (test_qwen38_qsa_long_context.py): 0.056 mean |dlogp| (short dense row: 0.051).
8K core attention fwd+bwd on one H100: TE dense 6.7 ms, flex 44 ms, SDPA math 138 ms.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
Forward WANDB_ENTITY / WANDB_BASE_URL to Ray workers, like WANDB_API_KEY
c3dcdf9 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
[megatron] Full FT with optimizer CPU offload: FP32 masters only on t…
275b786 
…he host

With optimizer_cpu_offload and offload_fraction 1.0, megatron-core still builds FP32 main-param
shards on the GPU (DistributedOptimizer._build_model_and_main_param_groups) and the
HybridDeviceOptimizer only copies them to the host; at expert-DP 1 (Qwen3.8-Flash-Next, EP16)
that is ~29 GiB/GPU and init OOMed. patch_hybrid_optimizer_cpu_masters routes full-FT runs to
megatron-core's precision-aware CPU-offload path (no GPU masters, BF16 decoupled grads), builds
each master as one pinned host tensor, and skips the throwaway HDO over the full model. LoRA
keeps the default (the adapter store reads the FP32 shards). Also: QSA entries in the patches
README.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
[weight sync] merge_lora: optional sync of LoRA-touched blocks only (…
25976ff 
…sync_lora_targets_only)

merge_lora=true re-sent the whole model each sync. With
megatron_config.lora_config.sync_lora_targets_only=true, MegatronWeightSource gets a
param_filter keeping only the blocks that hold a LoRA-wrapped (or trainable) parameter, as whole
blocks because vLLM's layerwise reload rebuilds a module from all of its tensors. vLLM's copies
of the frozen rest must stay the checkpoint's, so vLLM sleeps at level 1 for weights in this mode.
Rejected with dummy load format, FP8 weight sync, and the sharded_rdt/delta backends. Off by
default.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
[megatron] mHC chunked MLP: keep router replay (R3) per chunk
4b8e417 
chunked_mlp (mlp_chunks_for_training on hyper-connection layers: qwen4_exp, glm5_next) calls the
MoE router once per sequence chunk, but megatron-core's RouterReplay holds one [tokens, topk]
tensor per call: replay handed every chunk the whole micro-batch's indices (Qwen3.8 DAPO + R3:
'expected index [10634, 10] to be no larger than self [1330, 512]'), and recording kept only the
last chunk. Each chunk now gets its slice of the forward target and of the recompute FIFO entry,
and a recording is reassembled from the chunks.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
[megatron] mHC chunked MLP: same number of MoE calls on every rank
fa5f65d 
chunked_mlp ran min(num_chunks, tokens) MoE calls via torch.chunk, which also returns fewer
chunks than asked for short inputs. Each MoE call runs expert-parallel collectives, so a rank on
a short micro-batch (e.g. the one-token padding micro-batches TokenBasedBatchIterator adds to
even out micro-batch counts across DP ranks) made fewer calls than its EP peers and the group
deadlocked: Qwen3.8 DAPO at EP16 hung in the forward pass until the 30 min NCCL watchdog fired
(15 ranks in the MoE dispatch all-gather, one in the PLE lookup). Always make exactly num_chunks
non-empty calls (tensor_split), padding short sequences with masked dummy tokens whose outputs
(and router-replay rows) are dropped. Affects glm5_next too.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
[megatron] qwen4_exp: QSA FlexAttention mask without per-sequence Pyt…
416849e 
…hon ints

mask_mod captured the sequence length and block count as ints, which torch.compile guards on,
recompiling per distinct packed length. Pass them as 0-d tensors.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
[examples] qwen3.8 DAPO LoRA recipe (2 nodes, R3, 2k + 8k); checkpoin…
bf9a7b3 
…ting test rows

The recipe runs the QSA sparse regime and rollout router replay at full size (step 1: 966 s,
rollout/train |dlogp| 0.011). test_trainer_full_checkpointing gains qwen3.8 slice rows for LoRA
and full fine-tuning (EP4, optimizer on the host for full FT).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
Format (ruff, black)
4a4095e 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QepVYWMiu9bi2gEh5omjVi
@xinyuangui2 @claude
[trainer] Advantage-split log-ratio metric for multi-epoch direction …
c735fad 
…checks; qwen3.8 full-FT recipes

compute_advantage_split_logratio_metrics logs the mean logp - old_logp over positive- and
negative-advantage tokens (policy/adv_pos_logratio, policy/adv_neg_logratio) from the Megatron and
FSDP policy losses. With update_epochs_per_batch > 1 the trainer also logs them per epoch, so a few
epochs on one batch show whether the update moves logprobs in the advantage's direction.

Adds run_gsm8k_qwen3_8_flash_next_fullft_2node.sh and run_dapo_qwen3_8_flash_next_fullft_sync_2node.sh
(8xB200 nodes, host-resident optimizer, bf16 grad reduce).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01YbshX4DEqk3oji2hEQBhvz
@xinyuangui2 @claude
[examples] qwen3.8 full FT: pageable host optimizer buffers
c8131bb 
torch's pinned-host allocator rounds each block up to a power of two (a 2.2 GiB pinned tensor costs
4.3 GiB of host RAM). With the HybridDeviceOptimizer's FP32 masters and grad copies pinned, the
first optimizer step pushed a 1.76 TB B200 node past Ray's 95% memory-monitor threshold. With
pin_cpu_grads/pin_cpu_params=false, GSM8K full FT peaks at ~1.62 TB per node and runs 5 steps
(rollout/train |dlogp| 0.020-0.022, ~300 s/step).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01YbshX4DEqk3oji2hEQBhvz
@xinyuangui2 @claude
[vllm] Level-2 sleep model_state restore: Qwen4Exp only
07d4180 
Carry only Qwen4ExpModelState.ngram_context_offsets across a level-2 sleep,
keyed by the model_state class name, instead of every model_state tensor.
Other models' model_state is no longer copied or rewritten. vLLM #60009
lists this tensor as not fixed upstream.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01FT1xhhWniiXEKhg52FSrMc
@xinyuangui2 @claude
[vllm] Level-2 sleep model_state restore: key on the Qwen4Exp attribute
6fc0e7e 
Drop the class-name allowlist: save ngram_context_offsets only when the
runner's model_state has it, matching how the other vLLM patches stay a
no-op for models they don't apply to.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01FT1xhhWniiXEKhg52FSrMc
@xinyuangui2 @claude
Drop changes not needed for Qwen3.8-Flash-Next
68abb93 
Removed (generic; to go in separate PRs if wanted):
- merge_lora sync_lora_targets_only filter (merged_lora_filter.py, config,
  validation, client keep_weights_on_sleep, docs, tests); it saved ~1.5% for
  this model and did not fix the merged-LoRA OOM.
- Advantage-split logratio metrics and the per-epoch log line (direction-test
  diagnostics).
- WANDB_ENTITY / WANDB_BASE_URL forwarding.
- mtp/grad_clip.py shard_float16_groups pair: Qwen3.8 trains no MTP head.
  Instead, the host-masters optimizer routing now skips MTP models, so their
  clipping keeps finding FP32 shards.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01FT1xhhWniiXEKhg52FSrMc
@xinyuangui2 @claude
[megatron] CPU-offloaded optimizer: create the host Adam state at con…
7f5e245 
…struction

Saving a full-FT distributed checkpoint before the first optimizer step failed with
KeyError: 'master_param' in DistributedOptimizer._get_main_param_and_optimizer_states (precision-aware
branch): the HybridDeviceOptimizer's state, which carries master_param and the Adam moments, only
exists after the first step, because torch's Adam/AdamW allocate it lazily. init_cpu_optimizer_state
builds that lazy state eagerly after the HDO's sub-optimizers are created (zero moments, step 0,
with torch's own device/dtype placement, including fused=True, which megatron-core uses for the CPU
AdamW). The first step is bit-identical; the moments are only allocated earlier.

Found by the full-FT row of test_trainer_full_checkpointing.py (Qwen3.8-Flash-Next 4-layer slice,
4x B200), which saves right after build. It now passes (save, resume, second save).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01YbshX4DEqk3oji2hEQBhvz
xinyuangui2 and others added 20 commits October 20, 2026 21:03
@xinyuangui2 @claude
Drop NCCL_* forwarding to Ray workers
168c9f3 
Keep the PR to Qwen3.8 changes: forwarding every NCCL_* variable changes
behavior for all models. The recipes no longer export NCCL_NVLS_ENABLE;
the GSM8K LoRA recipe documents setting it in the worker environment on
80 GB GPUs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014jADLGSS2eEaKAzLuAL7PK
@xinyuangui2 @claude
[examples] qwen3_8_flash_next: complete_export.py makes full-FT HF ex…
9a44003 
…ports loadable

The full-FT HF export (trainer.hf_save_interval) only contains what the Megatron model owns. For
Qwen3.8-Flash-Next that leaves out the frozen PLE n-gram tables, the vision tower
(language_model_only) and the MTP head (mtp_num_layers=0), so vLLM cannot load it. The script
copies those tensors from the base checkpoint (whole PLE-only shards, the rest packed into one
extra file), rewrites the index and verifies every tensor's name/dtype/shape against the base.
Raw safetensors bytes, no torch.

Tested on the step-10 and step-20 exports of a 20-step full-FT DAPO run on 16xB200: 1658/1658
tensors verified, both load in vLLM (TP8) and reproduce the in-training AIME-2024 avg@8
(0.887 / 0.817 vs 0.858 / 0.812). Idempotent on an already completed export.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01S14hc47TJCuJMDwU84K7HW
@xinyuangui2 @claude
[examples] qwen3_8_flash_next DAPO recipes: validated LR and length p…
a6fceec 
…enalty, evals every 5 steps

Defaults now match the 16xB200 runs that learned (held-out AIME-2024, 30 problems):
- full FT: LR 1e-6 -> 3e-6, overlong penalty factor 1.0 -> 0.5. avg@8 0.467 -> 0.829 (step 10) ->
  0.871 (step 20), accuracy on finished responses 0.991. At factor 1.0 full FT learned to cut its
  reasoning short (finished accuracy 0.855 at step 20); at LR 1e-5 entropy fell to 0.62x LoRA's.
- LoRA: LR 1e-5 -> 3e-4, alpha 16 -> 32 (LR 1e-5 / 3e-5 runs stayed flat; 3e-4 reaches 0.85 by step 10),
  same penalty factor 0.5 as full FT.
- both: 32 prompts x 8, mini 16 (2 optimizer steps per batch), 30 steps, AIME eval before training and
  every 5 steps with 12 samples per problem (as in the GLM-5.3-Flash recipe). All overridable via env.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01S14hc47TJCuJMDwU84K7HW
@xinyuangui2 @claude
[vllm] Restore all model_runner.model_state tensors across any worker…
d9402fb 
… sleep

vLLM's wake path restores only the model's parameters and buffers, so any
tensor on Model Runner V2's model_state can come back wrong; Qwen4Exp's
ngram_context_offsets (vllm#60009) is one case. Save every model_state
tensor on each sleep, at any level, and restore it when the weights wake.
patch_vllm_level2_sleep_model_state -> patch_vllm_sleep_model_state.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014jADLGSS2eEaKAzLuAL7PK
@xinyuangui2 @claude
[vllm] Carry only the model_state tensors known to be lost on sleep
3d29bb9 
Copying every model_state tensor on each sleep is wasted work. Keep an
explicit list, MODEL_STATE_TENSORS_LOST_ON_SLEEP, currently Qwen4Exp's
ngram_context_offsets (vllm#60009); the patch still covers any sleep level.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014jADLGSS2eEaKAzLuAL7PK
@xinyuangui2 @claude
[megatron] Drop patch_hybrid_optimizer_cpu_masters
da1c426 
The patch kept FP32 master params off the GPU when the optimizer was fully
offloaded to the host, which full fine-tuning of Qwen3.8-Flash-Next needed on
16 GPUs (expert-DP 1). On 4x8 B200 (EP16, expert-DP 2) the distributed
optimizer's state fits on the GPU, so the recipes no longer offload it and
the patch, its CPU test and its README section go.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LT2b2yhk8WTqxCaKwZD1ML
@xinyuangui2 @claude
[examples] qwen3_8_flash_next: recipes target 4 nodes of 8xB200
c053b6d 
Rename the four recipes *_2node.sh -> *_4node.sh: NUM_NODES=4, four vLLM
engines at TP8, Megatron TP=CP=PP=1 EP16 (expert-DP 2), and no optimizer CPU
offload (drops optimizer_offload_fraction, pin_cpu_* and the related knobs).
Header comments updated for the 32-GPU memory budget. Training settings
(LR, LoRA rank/alpha, overlong penalty, batch sizes, R3, evals) unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LT2b2yhk8WTqxCaKwZD1ML
@xinyuangui2 @claude
[examples] qwen3.8: disable vLLM custom all-reduce from the recipes; …
ed037a9 
…drop the model-type check in inference_servers/utils.py

vLLM's custom all-reduce registers CUDA-graph buffers over CUDA IPC, which
fails on sleep mode's cumem allocator (colocate_all=true, TP>1). Instead of a
qwen4_exp model-type check in build_vllm_cli_args, the four 4-node recipes and
the qwen3.8 GPU test rows pass disable_custom_all_reduce through
engine_init_kwargs. inference_servers/utils.py is back to main.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LT2b2yhk8WTqxCaKwZD1ML
@xinyuangui2 @claude
[megatron] mHC chunked MLP: simpler _ChunkedRouterReplay docstring
4ddb220 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014jADLGSS2eEaKAzLuAL7PK
@xinyuangui2 @claude
[tests] qwen3.8 full-FT checkpoint row: GPU-resident optimizer
a24ba50 
The row offloaded the optimizer to the host, which only worked with
patch_hybrid_optimizer_cpu_masters: stock megatron-core's HybridDeviceOptimizer
creates its Adam state lazily, so a checkpoint saved before the first
optimizer step can't be resumed (KeyError ...exp_avg from model not in state
dict). The 4-node recipes keep the optimizer on the GPU; so does this row now
(~44 GB/GPU on the slice). Passes save, resume and second save.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LT2b2yhk8WTqxCaKwZD1ML
@xinyuangui2 @claude
[megatron] qwen4_exp: TODO to swap in megatron-core's QSA kernel
ab8cfbf 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014jADLGSS2eEaKAzLuAL7PK
@xinyuangui2 @claude
[ci] Run the qwen3.8 checkpoint and QSA long-context tests on H100
a55a36b 
The qwen3.8 rows of test_trainer_full_checkpointing.py were marked only
`megatron`, so SkyRL-GPU-Megatron would run them on its L4 compute config;
the full-FT row now keeps the optimizer on the GPU (~44 GB/GPU). Mark both
qwen3.8 rows `h100` (auto-skipped elsewhere) and add the file to
ci/gpu_ci_run_h100.sh. test_qwen38_qsa_long_context.py was already marked
`h100` but listed in no runner, so it never ran; add it too.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LT2b2yhk8WTqxCaKwZD1ML
@xinyuangui2 @claude
[tests] qwen3.8 checkpoint test: LoRA row only, like GLM-5.3-Flash
aeeb72d 
Drop the qwen3.8 full-FT row of test_trainer_full_checkpointing.py: with the
optimizer on the GPU (as in the 4-node recipes) and expert-DP 1 on 4 GPUs it
peaks at ~85 GB/GPU (measured on B200), more than the CI GPUs have. GLM-5.3-Flash
also checkpoints only its LoRA row here. The qwen3.8 LoRA row goes back to the
`megatron` suite, and the file leaves ci/gpu_ci_run_h100.sh;
test_qwen38_qsa_long_context.py stays there.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LT2b2yhk8WTqxCaKwZD1ML
@xinyuangui2 @claude
[tests] qwen3.8: point the 4-layer slice default at xguihugging/Qwen3…
3cc628f 
….8-Flash-Next-4layer

The slice is now public on HF. Also note that a local SKYRL_QWEN38_4LAYER_PATH must keep
"Qwen3.8-Flash-Next" in its path (the tests key the qwen3.8 settings on the model name), and add
numpy to make_slice.py's usage line.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014jADLGSS2eEaKAzLuAL7PK
@xinyuangui2 @claude
[megatron] qwen4_exp: load the PLE table from a hub repo id
8470878 
Qwen4ExpPLE.load_table joined qwen4_hf_path (the bridge's model_name_or_path)
with file names, which only works for a local directory. With the tests'
default slice now a hub id (xguihugging/Qwen3.8-Flash-Next-4layer), Megatron
init_model failed with FileNotFoundError on model.safetensors.index.json: the
table is read while the model is built, before megatron_worker's
snapshot_download. Resolve a non-directory path with hf_hub_download, which
fetches only the index and the PLE files through the HF cache.

Verified on 4xH100 with SKYRL_QWEN38_4LAYER_PATH unset (slice pulled from HF):
megatron_models qwen3.8 (forward 0.053, post-sync 0.026), megatron_lora_models
qwen3.8 (zero 0.052 / stale 0.243 / synced 0.059), qwen38_qsa_long_context
(0.055), and trainer_full_checkpointing megatron_qwen3_8_flash_next_lora all pass.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Y9h6d6JKmNWdFLNEZHE165
@xinyuangui2 @claude
Merge upstream/main into qwen38-flash-next
9642733 
Conflict in ci/gpu_ci_run_h100.sh: keep both main's test_megatron_vlm_init.py and this
branch's test_qwen38_qsa_long_context.py.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014jADLGSS2eEaKAzLuAL7PK
@xinyuangui2 @claude
Move the chunked_mlp fix (equal MoE calls, per-chunk R3) to NovaSky-A…
3723e76 
…I#2532

This PR keeps mhc_transformer_layer.py and its test as on main. The qwen3.8 recipes
with R3 (DAPO) and multi-rank EP need NovaSky-AI#2532.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014jADLGSS2eEaKAzLuAL7PK
@xinyuangui2 @claude
[megatron] patches README: the qwen4_exp upstream PRs are independent…
35931ae 
… trials; switch to whichever lands

Megatron-LM#7393, Megatron-Bridge#6123, Megatron-LM#7553 and Megatron-Bridge#6354 are separate
implementations (only MB#6123 builds on MLM#7393). List what each needs to land, what to check
before switching (construction path, recompute, PLE placement, packing, LoRA, R3), and verify
with the checkpoint row and a short DAPO run.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014jADLGSS2eEaKAzLuAL7PK
@xinyuangui2 @claude
[vllm] cite vllm-project/vllm#63548 as the upstream fix for model sta…
c38ea29 
…te lost on sleep

vllm#60009 explicitly leaves Qwen4Exp's ngram_context_offsets out of scope (it is held by
the model runner's model state, not a module). vllm#63548 builds the MRV2 model state in the
runtime pool under sleep mode; patch_vllm_sleep_model_state can go once it is in our vLLM pin.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014jADLGSS2eEaKAzLuAL7PK
@xinyuangui2 @claude
[megatron] qwen4_exp: row-shard the PLE table over EP (any EP size)
9323dd5 
The n-gram table was split by head, so EP had to divide the 16 heads.
Split it by contiguous rows instead (as Megatron-LM#7553) and route each
id to its owner rank with a variable-split all-to-all.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015m9Uy57LhVrqUZwKFH2VdR
@xinyuangui2

xinyuangui2 commented Oct 26, 2026

Copy link
Copy Markdown
Collaborator Author

Folded into #2513 (commit 41b0b3b).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@xinyuangui2