[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

SkyRL's own implementation; the open upstream PRs are independent trials (see Upstream PRs to track). 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 row-sharded over the EP group (any EP size; ids are routed to the owning rank by all-to-all) 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 cuDNN frontend's QSA kernels on Blackwell when the installed cudnn-frontend has them (cudnn.ops.qsa_forward/qsa_backward, Add packed QSA training kernels for Blackwell and Rubin NVIDIA/cudnn-frontend#1765; dev wheel via the recipes' QSA_CUDNN_WHEEL, SKYRL_QWEN4_QSA_CUDNN=0 to disable), else a FlexAttention path (44 ms per layer fwd+bwd at 8K, vs 138 ms for SDPA's math kernel). B200 results: comment.
• 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 CP=PP=1 (TP > 1 needs sequence parallelism).
• 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.

chunked_mlp fixes (#2532, merged)

The Qwen3.8 layer reuses chunked_mlp from mcore_ext/mhc_transformer_layer.py. Its two fixes, equal MoE call counts per rank (an EP deadlock otherwise) and per-chunk router replay (R3 crashed at DAPO step 1 otherwise), landed in #2532 and are merged in here.

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. Upstream fix: vllm-project/vllm#63548.

• 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), 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).

Progress vs .agents/docs/new_model_support.md

Stage 0: Megatron port and roundtrip logprobs

• Megatron-Core modules and Megatron-Bridge mapping with LoRA, all under patches/megatron/
• vLLM TP serving, layerwise weight reload and LoRA mapping on main's vLLM 0.31.0 (no custom wheel)
• test_megatron_models.py / test_megatron_lora_models.py rows on the 4-layer slice pass on 4×H100 (re-run 2026-10-22 on aeeb72d)
• Full-model roundtrip; upstream issue list below
• Slice published on HF: xguihugging/Qwen3.8-Flash-Next-4layer (the rows' default)

Stage 1: GSM8K, LoRA and full fine-tuning

• Both run E2E with BF16, vLLM TP only and EP-only training; 4-node recipes committed

Stage 2: DAPO, LoRA and full fine-tuning

• ~8K context, R3 on, reward signal over ~20 steps for both (AIME-2024 avg@12: LoRA 0.489 → 0.886, full FT 0.467 → 0.878; see Validation)
• test_trainer_full_checkpointing.py LoRA row: save, resume, second save (re-run 2026-10-22 on 4×H100)

Stage 3: performance and quantization (follow-up)

• QSA: block-sparse kernel on Blackwell (cuDNN frontend QSA; FlexAttention stays the fallback)
• TP with sequence parallelism (PLE runs on the TP-gathered sequence; TP2+EP4 matches EP4 on B200)
• CP for PLE and the QSA indexer, to train longer context
• Megatron perf features on this architecture (overlap grad reduce, full activation recompute, optimizer offload)
• MXFP8: a qwen4_exp spec in weight_sync/fp8/models and an MXFP8 run matching the BF16 curve

Validation

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

DAPO 2k/8k with the committed 4-node recipes (4×8 B200 with EFA, EP16 / DP32, vLLM 4 × TP8, optimizer on GPU): 30 steps, R3, 32 prompts × 8 samples, 2 optimizer steps per batch. Held-out AIME-2024 avg@12 (30 problems × 12 samples, temperature 1.0, top-p 0.7):

step	0	5	10	15	20	25	30
LoRA (r16 α32, LR 3e-4, wandb)	0.489	0.678	0.808	0.850	0.886	0.867	0.856
full fine-tuning (LR 3e-6, wandb)	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).
• 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 row on the slice (save, resume, second save). Like GLM-5.3-Flash, it has no full-FT row: with the optimizer on the GPU and expert-DP 1 on 4 GPUs, full FT of the slice peaks at ~85 GB/GPU.

Known limitations

• Megatron parallelism is limited to CP=PP=1 (PLE needs each packed sequence whole on one rank); TP > 1 runs with sequence parallelism. 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 without the cuDNN wheel (FlexAttention, e.g. on H100) is ~4× the cost of dense attention at 8K. The cuDNN QSA kernels cover 12 query heads per KV head (24/2 per rank at TP1, 12/1 at TP2); at TP4 (6/1 per rank) QSA falls back to FlexAttention without a log line. On sm100, vLLM falls back from FA4 to FA2 for head_size 256 with GDN's KV block size of 400.
• The PLE gate (sqrt(|key·query|)·sign) is steep near zero, so per-parameter grads of key_proj/norm_key/norm_query are sensitive to tiny input differences (also in HF fp32). Comparing these params across parallel layouts or against HF shows large per-param scatter even though overall grads match.
• A local slice passed via SKYRL_QWEN38_4LAYER_PATH must keep Qwen3.8-Flash-Next in its path: the tests key the qwen3.8 settings on the model name.

Upstream PRs to track

All open. The Megatron PRs are independent trials at qwen4_exp, not one stack: different model classes, config fields and PLE designs. Whichever lands first (and SkyRL bumps the pin to it), switch to it and delete patches/megatron/qwen4_exp/ as a whole. patches/megatron/README.md lists what to check before switching (construction path, activation recompute, PLE placement next to vLLM, packing, LoRA, R3).

• Megatron-LM#7393 + Megatron-Bridge#6123 (community): GPTModel, TP + SP. Both have to land.
• Megatron-Bridge#6354 (NVIDIA): HybridModel with its own modules inside Bridge; lands on its own.
• Megatron-LM#7553 (NVIDIA): core only (HybridModel, PP/VPP, CP, MTP, Triton QSA). Needs a Bridge bridge on top before qwen4_exp/ can go.
• vllm-project/vllm#63548 (build the MRV2 model state in the runtime pool under sleep mode): drop patch_vllm_sleep_model_state in vllm_compat.py.

Upstream issues

• vLLM:
  • model-state tensors lost on sleep (vllm#63548)
  • "Qwen4Exp HC does not support sequence-parallel MoE" blocks DP > 1 + EP + TP > 1 rollout
• megatron-core:
  • GDN uses one activation for the conv and the output gate
  • no hyper-connection, PLE or QSA modules (independent trials: Megatron-LM#7393, Megatron-LM#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 (independent trials: Megatron-Bridge#6123, Megatron-Bridge#6354)

🤖 Generated with Claude Code

https://claude.ai/code/session_01LT2b2yhk8WTqxCaKwZD1ML

---

Note

Medium Risk
Large new model path touches Megatron forward, MoE routing replay, and colocated vLLM weight sync/sleep; mistakes could cause silent logprob drift or training failures, but changes are mostly additive patches and examples.

Overview
Adds Qwen3.8-Flash-Next (qwen4_exp) to the Megatron training path via a new vendored stack under patches/megatron/qwen4_exp/ (bridge, provider, layer specs, and modules for gated hyper-connections, EP-sharded frozen PLE, sigmoid GDN, and QSA indexer plus sparse attention with cuDNN/FlexAttention/SDPA fallbacks). The bridge registers on import through model_bridges.py, and megatron_worker.py now requires text-only training for this VLM checkpoint.

Fixes colocated vLLM sleep/wake corrupting Qwen4Exp rollout state by saving and restoring listed model_runner.model_state tensors (e.g. ngram_context_offsets) in vllm_compat.py and the weight-sync sleep path in new_inference_worker_wrap.py.

Ships 4-node B200 recipes (GSM8K GRPO and DAPO sync for LoRA/full FT), make_slice.py / complete_export.py for CI-sized checkpoints and loadable HF exports, and extends GPU CI with 4-layer parity rows, ~8K QSA long-context logprob tests, LoRA checkpointing, and CPU tests for modules vs HF and the sleep patch. Megatron patch README documents upstream removal criteria for qwen4_exp/.

^{Reviewed by Cursor Bugbot for commit f0edea6. Bugbot is set up for automated code reviews on this repo. Configure here.}

[xinyuangui2]

Let's check the status upstream PR, they should include this. If they don't include, let's comment in megatron/megatron bridge.

[xinyuangui2]

QSA sparse attention on cuDNN frontend (5ca3778, b0339cf): findings on B200

Setup: nvidia-cudnn-frontend 1.32.0.dev76547535 (GitHub release wheel, via uv run --with), CuTeDSL 4.8.0 (cu13), 8×B200 nodes. FlexAttention is the existing path (SKYRL_QWEN4_QSA_CUDNN=0). All model-level A/Bs alternate the two modes on the same batch.

1. Kernel benchmark (one layer, 1×B200, random top-512 selection, ms)

Pack	FlexAttention fwd / fwd+bwd	cuDNN fwd / fwd+bwd	Dense TE (FA2) fwd+bwd
24Q/2KV 8K	14.7 / 40.1	1.64 / 13.3	9.8
24Q/2KV 10K	21.2 / 62.4	2.67 / 17.7	14.8
24Q/2KV 32K	202 / 613	14.1 / 73.3	142
12Q/1KV 8K	9.7 / 22.8	0.91 / 6.5	5.2
12Q/1KV 32K	118 / 324	7.3 / 36.9	71.7

• With a recency-biased selection, FlexAttention gets ~2× faster and cuDNN is unchanged (8K fwd+bwd 22.7 vs 13.3 ms).
• Max error vs the reference: ≤0.6% of the max value for out / dq / dk / dv. Metadata build (cudnn_qsa_metadata): 0.35–1.05 ms per call.

2. Correctness in the model

• test_qsa_cudnn (24Q/2KV, 12Q/1KV): pass. CPU tests: pass.
• Qwen3.8 4-layer slice vs HF: 8K long-context logprob diff 0.0528 with cuDNN vs 0.0523 with FlexAttention; dense rows unchanged.
• 2-node LoRA DAPO (EP16, R3): rollout-vs-train |dlogp| 0.011–0.012 in both modes. Ray workers pick up the --with wheel.
• Same-batch logprob diff cuDNN vs FlexAttention: mean ~0.011 (8K), 0.002–0.026 (14K–62K).

3. 8K, 2-node DAPO (EP16, LoRA, R3)

• Separate 3-step runs looked slower with cuDNN (fwd_logprobs 46–50 vs 32–35 s). A same-batch A/B showed this was batch variance: the cuDNN run's batch put 2.3–2.8× more EP collectives on the profiled ranks.
• Same batch: cuDNN is faster by ~2.2 s (~8%) per fwd_logprobs. That is all attention: 0.2–0.4 s vs 1.0–2.4 s per pass.
• QSA kernels are 1–3% of the step at 8K. They never overlap NCCL kernels. Host time per call: cuDNN 1.5 ms vs FlexAttention 26 ms.

4. Long context (1 node, EP8, LoRA, full recompute, one sequence per rank, warm steps)

Tokens per rank	fwd_logprobs	train fwd+bwd	Peak memory, fwd / train
14K	−10% (3.1 vs 3.5 s)	even (~15.4 s both)	41.0 vs 42.8 / 60.5 GB both
30K	−29% (4.4 vs 6.1 s)	−30% (22.4 vs 32.2 s)	44.5 vs 57.4 / 81.4 vs 87.1 GB
62K	−50% (7.2 vs 14.5 s)	cuDNN 36.8 s; FlexAttention OOM	52.5 vs 114.4 / 128 GB vs OOM

• Attention backward alone: cuDNN is 3.6× faster at 14K (0.25 vs 0.90 s) and 7.6× faster at 30K (0.57 vs 4.35 s).
• The 14K training tie is not cuDNN overhead. That step is host/launch-bound: the GPU is idle ~65% in both modes, with ~263K kernels per step, mostly small MoE ones. All non-attention kernels match, so the saved attention time becomes idle. From ~30K, attention is on the critical path.
• FlexAttention's 62K OOM happens in create_block_mask, which materializes a T×T mask in the recompute forward.
• Cold start: cuDNN's first pass of a run takes ~30 s (kernel compile and cache load), and FlexAttention's first train step takes +10–15 s (compile).

Takeaway: at the 8K DAPO recipe, the cuDNN path is a small win (~8% of the logprob pass), and training there is bound by MoE/host time. From ~30K it is a large win (~30% faster, up to 62 GB less memory), and at 62K it is the only path that trains in this setup.

[xinyuangui2]

TP2 + sequence parallelism, end to end: 30-step full-FT DAPO on 4×8 B200

Committed recipe run_dapo_qwen3_8_flash_next_fullft_sync_4node.sh at bb9c205, with only tensor_model_parallel_size=2 changed (EP16, ETP1, R3 on, LR 3e-6, AIME-2024 eval every 5 steps with 12 samples). Compared with the TP1 run of the same recipe from Oct 21 (older code, a24ba50).

• TP2: wandb 79b1louq
• TP1 reference: wandb u0ccg560

AIME-2024 avg@12:

step	0	5	10	15	20	25	30
TP2	0.467	0.822	0.881	0.875	0.864	0.856	0.811
TP1	0.467	0.839	0.858	0.872	0.878	0.819	0.844

Means over steps 1-30 (TP2 / TP1):

	TP2	TP1
rollout vs trainer |Δlogp|	0.0110	0.0100
grad norm	0.304	0.313
pass@8	0.894	0.901
response length (tokens)	2673	2677
peak allocated GPU memory / rank	93 GB	110 GB
policy_train / step	114 s	85 s
step time	249 s	240 s

The run finished with no errors and learns like TP1. TP2 saves ~17 GB per GPU, at the cost of a slower train phase; total step time is about the same because generation dominates. Policy entropy ran higher than TP1 for steps 7-19 (peak 0.93 at step 12 vs 0.31-0.53), then the two converged again from step 20 (steps 1-6 match, so the metric is not biased by TP). This could be ordinary run-to-run divergence, but the reference ran on older code, so it isn't a like-for-like comparison. A same-batch gradient replay comparison (TP1 vs TP2/TP4/ETP2) is in progress.

[gemini-code-assist]

Code Review

This pull request introduces support for the Qwen3.8-Flash-Next (qwen4_exp) model in the Megatron backend, including custom modules for gated low-rank hyper-connections, PLE n-gram embedding, and QSA indexer/sparse attention. It also adds scripts for model slicing and export completion, training recipes for DAPO and GRPO, and patches vLLM sleep/wake states to preserve model runner tensors. Feedback on the changes suggests using context managers for file operations in complete_export.py and make_slice.py to prevent unclosed file descriptors, and using getattr when accessing checkpoint_core_attention in layer_specs.py to avoid potential AttributeErrors.

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The PR adds a specialized Megatron training path and changes vLLM sleep/wake behavior.

Fix the B200 recipes’ GDN backend choice before merging.

Findings

1.  B200 training uses broken backward ▶
2.  Relative source breaks table links ▶
3.  FlexAttention gradients lack GPU coverage ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  HF[HF checkpoint] --> Bridge[Megatron bridge]
  HF --> PLE[Frozen host PLE table]
  Bridge --> Train[Megatron training]
  PLE --> Train
  Train --> Sync[Weights or LoRA sync]
  Sync --> VLLM[vLLM rollouts]
  VLLM --> Train
  VLLM --> Sleep[Save n-gram constants]
  Sleep --> Wake[Wake weights and restore constants]
  Wake --> VLLM

Loading

_{Reviews (1) · Last reviewed commit: "[examples] qwen3.8: trainer parallelism ..." · Reviewed by Greptile}

[greptile-apps]

FlexAttention gradients lack GPU coverage

This long-context test calls only forward. The cuDNN gradient test explicitly disables FlexAttention for its reference run, so neither checks gradients through the default GPU sparse path when the optional cuDNN wheel is absent.

Add a GPU test that forces FlexAttention and compares outputs and query/key/value gradients with SDPA on a pack containing sparse sequences and padding. Otherwise, wrong gradients or backward-only failures can pass the added tests.

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Reviewed by Cursor Bugbot for commit 3c63288. Configure here.}