[Tracking] Large-scale MoE-VLM training improvements

[yuhezhang-ai]

Summary

We have been running large-scale MoE-VLM training (Qwen3.5-VL-scale, EP/TP/CP/FSDP2 mixes up to 128 GPUs) on top of Automodel and accumulated a set of fixes and features that are mature enough to upstream. This issue tracks the themed PR series; each PR ships with a CI-runnable (CPU / 1-GPU) test as its proof, with large-scale runs cited as supporting evidence.

Grouping is by coherent subsystem (one area, one reviewer, shared tests) rather than many atomic PRs, and risky/large changes are kept standalone so they cannot block the safe wins.

PR queue

#	Themed PR	Contents	Proof	Status	PR
1	MoE TP/EP correctness	TP-safe MoE parallelization; EP+TP gradient correctness; TE grouped-expert zero-grad handling	264 unit tests (passing), validated in large-scale campaign	merged	#2995
2	megatron-fsdp compatibility	compat fixes for megatron-fsdp 0.5 API (flattened DP-CP, TP DTensor local shape, API drift)	149 unit tests (passing)	ready	#2986
3	Vision-tower AC on the EP-MoE path	follow-up to #2840: the MoE AC path (ep_size>1) does not cover the vision tower; also a math-SDPA recompute context_fn for full-block vision AC	AC-scoping + recompute-parity unit tests (181 passing incl. regressions)	merged	#2993 a relevant PR landed too: #3025
4	CP vision sharding (frame-level)	frame-level context-parallel sharding of vision inputs; docs + 40 unit tests. Sharding scope by design: frozen ViT splits frames across CP × TP; trainable ViT splits within CP only (TP ranks recompute the ViT) to keep TP-replica gradients correct; pure-TP (cp=1) sharding not enabled	unit tests (40, passing), loss parity + throughput data from validation runs	ready for main but we will stack on top of #2937 and fix conflicts	#2990
5	Setup prewarms	fla / cuBLAS / NCCL communicator prewarms (dropped dry_run warmup RFC)	19 unit tests (passing)	merged	#2992
6	Loss & packing	memory-bounded backward for the existing chunked CE; preserve indexed packing masks under FA2; scoped AC (non_moe / non_moe_no_attn) dropped after the scope decision	331 unit tests passing (2 skipped), including forward/gradient parity and packing-mask preservation	merged	#2996
7	CP block-diagonal varlen attention	block-diagonal varlen attention for context parallelism (the large/risky one — intentionally standalone)	60 parity/safety tests (passing)	merged	#2989
8	async-TP native-linear graph shaping	emit the reshape-mm-reshape pattern Inductor's _micro_pipeline_tp fusion recognizes (the enable flag from #1711 exists, but TPLinear's bmm path defeats the fusion)	13 unit tests (passing)	ready	#2987
9	Zero-numel-shard-safe fused Adam for VLM	guard against zero-numel FSDP shards hitting TE fused Adam (small dense ViT on a wide mesh); Automodel-side fix independent of any TransformerEngine timeline (separate TE issue to be filed with the deterministic repro)	9 unit tests incl. real-DTensor empty-local-shard proof (passing)	merged	#2997
10	Qwen3.5 MoE-VLM long-context integration	integrate the CP block-diagonal packing/GDN path from #2989 and frame-level CP vision sharding from #2990 into Qwen3.5 MoE-VLM; add the Qwen3.5-122B-A10B 128K CP+EP+packing recipe as the first large-scale target	focused integration/parity tests + 100-step 128-GPU Qwen3.5-122B 128K validation	draft; stacked on #2937, #2989, and #2990	#3186

Deferred (not in this series): native Muon optimizer integration — still under active development, will be proposed once it stabilizes.

Already upstreamed from this effort directly

Contribution of this effort merged before this series opened:

• fix(distributed): omit unset FSDP reshard argument #2926 (@zhiqi-li) — omit unset FSDP reshard argument

Process

Each themed PR is staged on a branch off clean main with its ported test run locally before opening. PRs will be linked in the table as they open. cc @yuhezhang-ai