[xinyuangui2]

What

Adds Qwen3.8-Flash-Next (HF qwen4_exp, Qwen4ExpForConditionalGeneration) to the Megatron backend, for both full fine-tuning and LoRA, with vLLM rollouts.

The architecture is new to both upstreams. It has 48 layers of 3:1 Gated DeltaNet / gated full attention, and the full-attention layers use a QSA sparse indexer (top-512 blocks of 4 tokens). The MoE has 512 routed experts (top-10) and a shared expert. Every block has gated low-rank hyper-connections over a 4-stream residual, and a frozen 51B-parameter n-gram embedding (PLE) is injected at layer 1. It ships as a VL checkpoint, and SkyRL trains the language model only (language_model_only=True).

No dependency pins change: this runs on main's vLLM 0.31.0, megatron-core and Megatron-Bridge. What the pinned upstreams don't have yet is under patches/.

Components

patches/megatron/qwen4_exp/: the Megatron-Bridge model

Upstream home: NVIDIA-NeMo/Megatron-Bridge#6123 / #6354, on NVIDIA/Megatron-LM#7393 / #7553 (both open). The modules here were written against HF transformers' modeling_qwen4_exp.py, and their parameter names mirror the HF classes. The bridge builds on Megatron-Bridge's Qwen35MoEBridge mappings.

• modules.py:
  • Qwen4ExpGatedResidual / Qwen4ExpFinalMixer: the gated low-rank hyper-connection and the final stream mixer, which takes the block's final_layernorm slot.
  • Qwen4ExpGatedDeltaNet: megatron-core's GDN with a sigmoid output gate and SiLU on the conv. megatron-core uses one activation for both.
  • Qwen4ExpPLE: the PLE n-gram embedding.
    • Hashing is packed-sequence aware.
    • The table is sharded by n-gram head over the EP group and kept frozen in pinned host memory.
    • Lookup is all-gather ids → local CPU embed → all-to-all.
  • Qwen4ExpQSAIndexer and qsa_sparse_attention: HF's top-k block selection. At up to indexer_budget + compress_ratio - 1 = 2051 tokens every block is kept, so attention stays on the fused TE kernel. Past that it runs a FlexAttention path (44 ms per layer fwd+bwd at 8K, vs 138 ms for SDPA's math kernel).
• layer_specs.py: Qwen4ExpTransformerLayer (expand → PLE → HC-wrapped mixer → HC-wrapped MoE), and Qwen4ExpCoreAttention, which switches between TE dense and the QSA sparse path.
• provider.py: Qwen4ExpModelProvider. It stashes input_ids / cu_seqlens for PLE through a forward pre-hook, and requires TP=CP=PP=1.
• bridge.py: Qwen4ExpBridge. It uses the Qwen3.5 MoE mappings minus the fused pre-norms, plus replicated mappings for the HC, PLE projection and indexer weights. The 51B PLE table, the vision tower and the MTP head aren't bridged.

patches/megatron/README.md is the removal plan: what retires each piece upstream, and what to compare before deleting it.

patches/megatron/mcore_ext/mhc_transformer_layer.py: router replay for chunked_mlp

This fixes shared code that glm5_next also uses:

• RouterReplay replays one top-k tensor per router call, but chunked_mlp calls the router once per chunk. _ChunkedRouterReplay hands each chunk its slice of the forward target and of the recompute queue, and reassembles recorded indices. Before this, DAPO with R3 crashed at step 1 with an index-size error.
• Every rank now makes exactly mlp_chunks_for_training MoE calls, padding short micro-batches with masked tokens. A rank on a one-token padding micro-batch used to make fewer EP collectives and deadlock the group.

inference_servers/vllm_compat.py + new_inference_worker_wrap.py: model state lost on sleep

vLLM's wake path restores only the model's parameters and buffers. Tensors on Model Runner V2's model_state that live in the weights pool can therefore come back wrong. Qwen4Exp's constant n-gram offsets (arange(-2, 0)) come back as zeros, which corrupts every later PLE n-gram context (vllm-project/vllm#60009).

• patch_vllm_sleep_model_state saves the tensors listed in MODEL_STATE_TENSORS_LOST_ON_SLEEP before any sleep, and restores them when the weights wake.
• Before this, rollout↔train |Δlogp| jumped from 0.07 to 1.15 after the first weight sync, even with bit-exact weights.

Other SkyRL changes

• megatron_worker.py: qwen4_exp is text-only on Megatron, so it rejects language_model_only=false with a clear message. model_bridges.py registers the bridge.
• Recipes: examples/train/qwen3_8_flash_next/, 4 nodes of 8×B200, EP16, 4 vLLM engines × TP8, colocated.
  • run_gsm8k_qwen3_8_flash_next_{lora,fullft}_4node.sh
  • run_dapo_qwen3_8_flash_next_{lora,fullft}_sync_4node.sh: 2k/8k, R3 on, evals every 5 steps.
    • Full FT: LR 3e-6.
    • LoRA: r16 α32, LR 3e-4.
    • Both: overlong penalty factor 0.5.
  • They set disable_custom_all_reduce in engine_init_kwargs. vLLM's custom all-reduce registers CUDA-graph buffers over CUDA IPC, which fails on sleep mode's cumem allocator, so every TP worker dies at engine start.
  • make_slice.py: the 4-layer CI slice with a small random PLE table (23 GB).
  • complete_export.py: makes a full-FT HF export loadable (see Known limitations).
• Tests:
  • CPU: qwen4_exp modules and the QSA indexer vs HF (outputs and grads, packed), chunked-MoE R3 and equal call counts, the sleep model-state patch.
  • GPU: qwen3.8 4-layer rows in gpu_ci/megatron/test_megatron_models.py and test_megatron_lora_models.py, and test_qwen38_qsa_long_context.py (sparse regime).

Validation

Megatron vs vLLM on the 4-layer slice, 4×H100 (EP4, vLLM TP4), threshold 0.1:

row	Megatron vs vLLM	post weight sync (common prefix)
qwen3.8-flash-next-4layer_h100_ep4	0.051	0.021
test_qwen38_qsa_long_context.py (~7.7K-token prompts, sparse QSA)	0.056

LoRA roundtrip (..._ep4_adapter):

phase	diff	check
zero adapter	0.052	< 0.1
perturbed trainer vs stale sampler	0.239	> 0.1
synced adapter	0.056	< 0.1

Full model, DAPO 2k/8k, R3, 32 prompts × 8 samples, 2 optimizer steps per batch, on B200 with EFA. Held-out AIME-2024 avg@8 (30 problems × 8 samples, temperature 1.0, top-p 0.7):

step	0	10	20
full fine-tuning (LR 3e-6, overlong penalty 0.5)	0.467	0.829	0.871
full fine-tuning, overlong penalty 1.0	0.463	0.858	0.812
LoRA (r16 α32, LR 3e-4)	0.446	0.846	0.821

• The model learns to finish its reasoning within the 8K budget: truncation goes from ~54% to 5–17%.
• Full fine-tuning at penalty 0.5:
  • Accuracy on finished responses is 0.991 at step 20 (LoRA: 0.985).
  • The rollout↔train |Δlogp| stays at 0.009–0.012 throughout.
  • The completed step-20 export serves in standalone vLLM at 0.883.
• Penalty 1.0 over-compresses full fine-tuning: accuracy on finished responses falls to 0.855 by step 20.
• LR 1e-5 for full fine-tuning sharpens too fast: entropy fell to 0.62× LoRA's on the same batches by step 7.

These runs used 2×8 B200 (EP16, vLLM 2 × TP8) with the optimizer offloaded to the host.

Reruns of the committed 4-node recipes (4×8 B200, EP16 / DP32, vLLM 4 × TP8, optimizer on GPU), 30 steps, R3, 32 × 8, held-out AIME-2024 avg@12 (30 problems × 12 samples):

step	0	5	10	15	20	25	30
LoRA (r16 α32, LR 3e-4)	0.489	0.678	0.808	0.850	0.886	0.867	0.856
full fine-tuning (LR 3e-6)	0.467	0.839	0.858	0.872	0.878	0.819	0.844

• Both peak at step 20. Accuracy on finished responses there is 0.981 for both.
• Full FT learns faster early, then drifts after step 20: finished-but-wrong answers go 6 → 31 → 23, and they are long, coherent guesses near the 8K budget. LoRA holds its level (finished accuracy 0.972 at step 30).
• On the same score scale as GLM-5.3-Flash in [megatron] GLM-5.3-Flash (glm5_next) support: KDA + k-pool DSA + mHC MoE, full fine-tuning and LoRA sync to vLLM (merge_lora=false) #2179 (avg_score 0.072 → 0.711 over 30 steps, i.e. accuracy 0.54 → 0.86), Qwen3.8 reaches +0.772 (LoRA) / +0.756 (full FT) at step 20.
• Rollout↔train |Δlogp| stays at 0.008–0.012 in both runs, and no stop rule fired.
• Steady step time: full FT 180–230 s, LoRA 250–340 s. LoRA is slower because vLLM serves 512-expert LoRA adapters.

Also passing: test_trainer_full_checkpointing.py LoRA and full-FT rows on the slice (save, resume, second save). The full-FT row uses the GPU-resident optimizer, matching the 4-node recipes.

Known limitations

• Megatron parallelism is limited to TP=CP=PP=1 (PLE needs each packed sequence whole on one rank) and EP ≤ 16 (the PLE table is sharded by its 16 heads over the EP group). The recipes use EP16 with DP32 on 32 GPUs, so each expert is held by 2 ranks. Text only; no MTP training.
• The full-FT HF export omits the frozen PLE tables, the vision tower and the MTP head. Run examples/train/qwen3_8_flash_next/complete_export.py <export_dir> <base_dir> to merge them from the base. The MTP drafter is then the base's, so speculative-decoding acceptance may drop.
• QSA sparse attention is ~6× the cost of dense attention at 8K; a block-sparse kernel would close that. On sm100, vLLM falls back from FA4 to FA2 for head_size 256 with GDN's KV block size of 400.
• The 4-layer CI slice isn't published on HF yet. The rows default to a placeholder id; override it with SKYRL_QWEN38_4LAYER_PATH.
• Multi-node runs need EFA/RDMA. Over TCP, NCCL took 97% of a full-FT step.
• With stock megatron-core's CPU-offloaded optimizer (HybridDeviceOptimizer), a checkpoint saved before the first optimizer step can't be resumed (KeyError ...exp_avg from model not in state dict), because its Adam state is created lazily. The recipes don't offload the optimizer.

Upstream issues

• vLLM:
  • model-state tensors lost on sleep (vllm#60009)
  • custom all-reduce fails on cumem memory under sleep mode
  • "Qwen4Exp HC does not support sequence-parallel MoE" blocks DP > 1 + EP + TP > 1 rollout
  • reload_weights() materializes the PLE table on GPU (vllm#62063)
• megatron-core:
  • GDN uses one activation for the conv and the output gate
  • no hyper-connection, PLE or QSA modules (Megatron-LM#7393 / #7553)
• Megatron-Bridge:
  • the Qwen3.5 layer-type helper rejects transformers ≥ 5.16's qwen_sparse_attention
  • AutoMapping detects parallelism by exact class name, so subclasses fail
  • no qwen4_exp bridge yet (#6123 / #6354)

🤖 Generated with Claude Code