[megatron] feat: add Qwen3.5 THD (packed sequence) Megatron SFT example

[gaohongkui]

What does this PR do?

examples/sft/gsm8k/ has no Megatron 3D-parallel SFT example that uses packed sequences (THD). Every megatron example there runs bshd, and the only one that sets the flag explicitly — run_qwen3_5_397b_a17b_megatron.sh — sets use_remove_padding=False with this note:

Qwen3.5 uses Gated Delta Net (GDN) linear attention which currently does NOT support packed sequences (THD format) in Megatron-LM. [...] Once NVIDIA/Megatron-LM#2644 is merged, THD format will be supported.

That note was accurate when it was written (#5381, 2026-03-13) but is no longer true:

• Megatron-LM#2644 "feat(moe): Support packed sequence for gated delta net (GDN)" merged 2026-04-07.
• The THD branch in megatron/core/ssm/gated_delta_net.py first ships in a megatron-core release in 0.18.0. Verified by content inspection across release tags — git show <tag>:megatron/core/ssm/gated_delta_net.py | grep "qkv_format == 'thd'" finds nothing in core_v0.16.0 / 0.16.1 / 0.17.0 / 0.17.1, and finds it in core_v0.18.0 / 0.18.2. (git merge-base --is-ancestor is not usable here: NVIDIA's public mirror does not carry the internal merge commit as an ancestor of the release tags, so it reports a false negative.)
• megatron-core 0.18.2 also parametrizes its own GDN test over sequence_packing (tests/unit_tests/ssm/test_gated_delta_net.py).

So THD works for this architecture family on megatron-core >= 0.18.0, and the padding no longer has to be computed.

Checklist Before Starting

• Search for similar PRs. Queries run: GDN THD sequence packing, use_remove_padding megatron sft, attention_backend flash fused NaN. No open PR covers GDN+THD or a Megatron THD SFT example. The nearest related one is [fsdp] fix: make GDN work correctly when using ulysses and varlen sequence #5346 ([fsdp] fix: make GDN work correctly when using ulysses and varlen sequence), which is the FSDP path and does not overlap.
• Format the PR title as [{modules}] {type}: {description}

Test

Ran the new example end-to-end on 2 nodes / 16 GPUs (H20-141G), megatron-core 0.18.2, torch 2.11.0+cu128, Qwen3.5-35B-A3B, gsm8k:

python3 examples/data_preprocess/gsm8k_multiturn_sft.py --local_save_dir ~/dataset
NNODES=2 NUM_GPUS=8 MODEL_PATH=<Qwen3.5-35B-A3B> \
  bash examples/sft/gsm8k/run_qwen3_5_35b_a3b_megatron.sh

Result — completed the full epoch, 116/116 steps:

step	1	13	25	37	49	61	73	85	97	109	116
train/loss	0.7730	0.3938	0.3717	0.3110	0.3379	0.2985	0.3214	0.3267	0.2901	0.2984	0.3116

• 3.76 s/it steady state (9.22 s/it for the first few steps, including warmup)
• max_memory_allocated 42.93 GB, max_memory_reserved 51.85 GB
• Zero NaN / errors — found NaN, loss:nan, Traceback, RuntimeError, AssertionError all have zero occurrences in the logs
• THD was confirmed genuinely active (not silently falling back to bshd): model.use_remove_padding=True, engine.use_remove_padding=True and data.use_dynamic_bsz=True all present in the launched torchrun arguments, with no attention_backend override — which is also what the header claims (the default Megatron-Bridge path already pins AttnBackend.flash).

One deviation to disclose: the cluster used for this run has no external network access, so instead of downloading gsm8k, the run converted an already-present RL-format gsm8k parquet into the SFT messages schema. This is field-for-field equivalent to gsm8k_multiturn_sft.py, because gsm8k.py and gsm8k_multiturn_sft.py build the user turn from the same literal instruction string, so prompt[0].content equals the SFT user content and extra_info.answer equals answer_raw (7473 rows, no rows dropped).

API and Usage Example

No API change. New example script only:

# THD (packed sequences) on megatron-core >= 0.18.0
NNODES=2 NUM_GPUS=8 bash examples/sft/gsm8k/run_qwen3_5_35b_a3b_megatron.sh

# knobs (defaults shown)
TP_SIZE=2 PP_SIZE=2 CP_SIZE=2 EP_SIZE=8 \
MAX_LENGTH=2048 MAX_TOKEN_LEN_PER_GPU=2048 \
  bash examples/sft/gsm8k/run_qwen3_5_35b_a3b_megatron.sh

Design & Code Changes

1. Add examples/sft/gsm8k/run_qwen3_5_35b_a3b_megatron.sh — Qwen3.5-35B-A3B on gsm8k, TP2/PP2/CP2/EP8, with the three settings THD needs (model.use_remove_padding, engine.use_remove_padding, data.use_dynamic_bsz) and a note on how data.max_token_len_per_gpu interacts with CP_SIZE (the effective micro-batch budget is max_token_len_per_gpu * CP_SIZE).

   attention_backend is deliberately left unset, because the default (Megatron-Bridge) path already pins AttnBackend.flash. The header records why the deprecated engine.vanilla_mbridge=True path needs it set explicitly: there attention_backend falls back to megatron's auto, which prefers cuDNN fused attention on Hopper+, and cuDNN's SDPA backward has an uninitialized-workspace defect on Hopper (fixed in cuDNN 9.18) that makes THD backward produce huge dK/dV and then NaN gradients — see Fused attention with THD format + CP may cause a bug in backward pass, leading to NaN values during training LLM. NVIDIA/TransformerEngine#2186.

2. Update the stale THD note in run_qwen3_5_397b_a17b_megatron.sh to point at the new example. Comment-only; its behaviour is deliberately unchanged, since it is pinned to Megatron-LM 0.16.0 and its parallelism config was validated that way, and I cannot test a 128-GPU 397B configuration.

Checklist Before Submitting

• Read the Contribute Guide.
• Apply pre-commit checks — pre-commit was not installable in my environment, so I ran the repo's own sanity checks directly instead: python3 tests/special_sanity/check_example_naming.py --root examples (✅ 101/101), python3 tests/special_sanity/check_pr_title.py with this title (✅), and bash -n on both scripts. Please let me know if you want a pre-commit run before merge.
• Add / Update the documentation — not applicable; this is an example script with an explanatory header.
• Add unit or end-to-end test(s) to the CI workflow — not feasible: this needs 16 GPUs, a 35B MoE checkpoint and megatron-core >= 0.18.0. Validated by the experiment above instead.

---

AI assistance disclosure: this change was prepared with AI assistance (Claude). I reviewed every changed line, ran the verification run described above myself, and can defend the change end-to-end.