[megatron] feat: add DeepSeek-V4-Flash SFT example (Megatron-Bridge backend) by Meirtz · Pull Request #6603 · verl-project/verl

Meirtz · 2026-06-04T04:29:25Z

Summary

Adds a GSM8K SFT example for DeepSeek-V4-Flash via the Megatron-Bridge backend, mirroring examples/sft/gsm8k/run_mimo_7b_mtp_megatron.sh. Runs at TP1/PP4/EP8 (32 GPU), THD packed sequences, MTP enabled, unfused mHC/RoPE (apply_rope_fusion=False, use_fused_mhc=False), use_mbridge=True, vanilla_mbridge=False.

Folds in a 2-line fix to verl/models/mcore/mtp_patch.py so the MTP postprocess monkeypatch tolerates extra kwargs (padding_mask) that current Megatron-Core's GPTModel.forward passes — required by model.mtp.enable=True. Supersedes the standalone #6596 (closed).

Environment / image — please read first

DeepSeek-V4-Flash SFT needs a specific stack and will not run on a stock verl install. Required environment:

Base image: nvcr.io/nvidia/pytorch:26.04-py3
verl — this PR (includes the mtp_patch.py fix)
Megatron-Bridge (megatron.bridge) + Megatron-Core with the DSv4 hybrid attention and THD support — Megatron-Core PR [dev] [DeepSeek-v4] Packed Sequence (THD) support for DSv4 Hybrid Attention NVIDIA/Megatron-LM#5011 (open). The example script fetches both onto PYTHONPATH.
DSv4 kernels: fast_hadamard_transform + nvidia-cudnn-frontend — the example script installs these.

Reference Dockerfile that bakes the full environment: https://gist.github.com/Meirtz/7e240c5489851ebf79a6edfd357a40b7

The example script also bootstraps the Megatron-side pieces and kernels at runtime, so a plain nvcr.io/nvidia/pytorch:26.04-py3 container is sufficient to run it.

Not a duplicate

The DeepSeek-V4 GRPO path is #6473 (examples/grpo_trainer/run_deepseek_v4_flash_megatron.sh); this is the SFT counterpart under examples/sft/gsm8k/. Same Megatron-Bridge backend, different trainer.

Test

The DSv4 SFT path through Megatron-Bridge (THD, MTP, the apply_rope_fusion=False / use_fused_mhc=False config) was validated end-to-end on 8×H100 with a 4-layer DSv4 proxy: training and validation loss decrease, no NaN. The full model (TP1/PP4/EP8) and this exact GSM8K configuration have not been run by me; the script mirrors the validated config and the existing megatron SFT example structure. CI cannot exercise it until the DSv4 kernels are in verl's CI image.

AI assistance

Prepared with AI assistance (Claude) and reviewed by the submitter.

CLAassistant · 2026-06-04T04:29:37Z

All committers have signed the CLA.

gemini-code-assist

Code Review

This pull request introduces a new shell script run_deepseek_v4_flash_sft_megatron.sh to run DeepSeek-V4-Flash SFT on GSM8K using the Megatron-Bridge backend. Feedback on the script points out that since set -u is enabled, referencing $PYTHONPATH directly when it is unset will cause an unbound variable error, and suggests using ${PYTHONPATH:-} instead.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-04T04:30:40Z

+# Megatron-Bridge: provides megatron.bridge (the DSv4 AutoBridge / recipes).
+[ -d Megatron-Bridge ] || git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge
+cd -
+export PYTHONPATH=$PYTHONPATH:$PYPATH/Megatron-Bridge/src:$PYPATH/Megatron-LM


Since set -u (nounset) is enabled on line 5, referencing $PYTHONPATH directly when it is not already set in the environment will cause the script to crash with an unbound variable error. Use ${PYTHONPATH:-} to safely default to an empty string if PYTHONPATH is unset.

Suggested change

export PYTHONPATH=$PYTHONPATH:$PYPATH/Megatron-Bridge/src:$PYPATH/Megatron-LM

export PYTHONPATH=${PYTHONPATH:-}:$PYPATH/Megatron-Bridge/src:$PYPATH/Megatron-LM

References

Enable set -xeuo pipefail in shell scripts to ensure that the script exits on errors, treats unset variables as errors, and pipelines fail correctly.

Good catch — fixed, now uses ${PYTHONPATH:-} so it's safe under set -u when PYTHONPATH is unset.

ISEEKYAN · 2026-06-04T08:00:48Z

+#!/usr/bin/env bash
+# DeepSeek-V4-Flash SFT on GSM8K via the Megatron-Bridge backend.
+# DSv4 hybrid attention requires TP=1; scale with PP and EP -> TP1/PP4/EP8 = 32 GPUs
+# (e.g. 4 nodes x 8 GPU on H100/H200, 8 nodes x 4 GPU on GB200/GB300).


32H100 or H200 would meet OOM, better give a more realistic setting for H100.
128H100 with fully recompute and PP16EP8 would be a choice

Right, 32×H100 OOMs (recompute already full). Both 128-GPU options work layout-wise now — NVIDIA-NeMo/Megatron-Bridge#4131 handles PP4 and PP16.

ISEEKYAN · 2026-06-04T08:14:58Z

+
+NUM_GPUS=${NUM_GPUS:-8}          # GPUs per node
+TP_SIZE=${TP_SIZE:-1}
+PP_SIZE=${PP_SIZE:-4}


how to set PP4 or PP16 with flash config:
"num_hidden_layers": 43,
"num_hash_layers": 3,
"num_nextn_predict_layers": 1,

No manual layout — NVIDIA-NeMo/Megatron-Bridge#4131 auto-sets it for DSv4 at PP>1. Hash layers are first, so the even split keeps them on stage 0 for both PP4 and PP16. Depends on that PR; once it lands, just set PP=4 or 16.

ISEEKYAN · 2026-06-04T08:15:44Z

+    engine.vanilla_mbridge=False \
+    engine.use_remove_padding=${USE_REMOVE_PADDING} \
+    engine.use_dist_checkpointing=False \
+    +engine.override_transformer_config.apply_rope_fusion=False \


can we enable rope fusion? it is critical for varlen perf

~~(Earlier attributed this to an upstream kernel bug — wrong.)~~ Root cause was a Megatron-Bridge config-mapping issue (partial_rotary_factor double-applied → the fused rope kernel read its cos/sin cache out of bounds), fixed in NVIDIA-NeMo/Megatron-Bridge#4271 and validated on GB300 (fused now matches unfused to 4 decimals). apply_rope_fusion=True will be enabled here once the example's Megatron-Bridge checkout includes that fix.

Done — rope fusion is now ON in this example. The NaN root cause was a Megatron-Bridge config-mapping bug (fixed in NVIDIA-NeMo/Megatron-Bridge#4271, merged); full-model DSv4-Flash SFT on 8×GB300 with apply_rope_fusion=True now matches the unfused control (iter-10 lm loss 2.3319 vs 2.3306, 0 NaN over 20 iters). Varlen perf path restored.

Meirtz · 2026-06-04T09:54:24Z

Folded the 2-line MTP postprocess fix (previously #6596) into this PR so the example is self-contained for model.mtp.enable=True. Closing #6596 in favor of this.

GSM8K SFT example for DeepSeek-V4-Flash via the Megatron-Bridge backend, mirroring run_mimo_7b_mtp_megatron.sh. TP1/PP4/EP8 (32 GPU), THD packed sequences, MTP enabled, unfused mHC/RoPE (apply_rope_fusion=False, use_fused_mhc=False), vanilla_mbridge=False. Fetches the DSv4 deps not in verl's pinned Megatron-Core (NVIDIA/Megatron-LM#5011 + NVIDIA-NeMo/Megatron-Bridge). Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

GPTModel.forward in current Megatron-Core passes extra keyword arguments (e.g. padding_mask for MoE routing) into self._postprocess. The MTP postprocess monkeypatch reimplemented _postprocess with a fixed signature and no **kwargs, so enabling MTP raised TypeError on padding_mask. Accept padding_mask and add **kwargs (consumed in the decoder, not in postprocess, so ignoring them preserves behavior). Needed by the DSv4-Flash SFT example in this PR (model.mtp.enable=True). Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

…V4-Flash SFT script Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

The example runs at PP4 by default; the DSv4 pipeline-layout auto-set lives in NVIDIA-NeMo/Megatron-Bridge#4131 (not yet merged). Fetch its head the same way the script already fetches Megatron-LM#5011, so PP>1 works. Revert to a plain clone once verl-project#4131 merges. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

pull/5011/head is rebased over time; its current head sits on pure dev and lacks main-side symbols that megatron.bridge imports (megatron.core._rank_utils.safe_get_world_size), breaking the example with an ImportError. Pin the validated main2dev commit ed6b1f65 instead, which carries both the DSv4 hybrid attention + THD support and the main-side helpers. Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

Meirtz · 2026-06-10T10:51:25Z

Heads-up for anyone running this example: Megatron-LM is now pinned to the validated commit ed6b1f65 (commit 8de857f) instead of tracking pull/5011/head. The #5011 head was rebased onto pure dev and dropped main-side symbols that megatron.bridge imports — if you cloned earlier and hit ImportError: cannot import name 'safe_get_world_size' from 'megatron.core._rank_utils', fix it with: cd Megatron-LM && git fetch origin ed6b1f65502aec7f2fe27e14a1245c29e435c2a6 && git checkout ed6b1f65502aec7f2fe27e14a1245c29e435c2a6.

…mplate Two issues reported by a user running this example: (1) the example pointed at the RL preprocess (gsm8k.py, 'prompt' column, no assistant turns) while SFT should use gsm8k_multiturn_sft.py ('messages' column) — mismatched messages_key raised KeyError: 'prompt'; (2) DeepSeek-V4-Flash ships no chat_template in its HF tokenizer, so message rendering needs model.custom_chat_template. The script now generates a minimal template (only bos/eos exist in the tokenizer; no role tokens) and passes it via @file. Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

Meirtz · 2026-06-11T09:46:38Z

Two fixes from a user running this example end-to-end (thanks!), now smoke-validated on a GB300 node against this PR's head:

Data: the example now uses the SFT preprocess examples/data_preprocess/gsm8k_multiturn_sft.py (messages column, includes assistant turns) — the previous messages_key=prompt raised KeyError: 'prompt' against SFT-preprocessed parquet. Validated: parquet schema ✓.
Chat template: DeepSeek-V4-Flash ships no chat_template in its HF tokenizer (and no role tokens — only bos/eos), so the script generates a minimal template and passes it via model.custom_chat_template. Note: upstream verl consumes this value verbatim as the template string (verl/workers/config/model.py) — there is no @file indirection here (some downstream forks have one), so the script inlines the file content: model.custom_chat_template="\"$(cat $CHAT_TEMPLATE_FILE)\"". Validated: tokenizer has no builtin template ✓, MultiTurnSFTDataset sample renders with a non-empty assistant loss mask ✓, the override survives hydra's grammar verbatim ✓.

Commits 640f0ed + 9823b0b. Replace the generated template via CHAT_TEMPLATE_FILE=... if you have an official one.

…has no @file syntax) model.custom_chat_template is consumed verbatim as the template string upstream (verl/workers/config/model.py); the @file indirection only exists in downstream forks. Pass the file content via command substitution, hydra-quoted. Validated end-to-end on a GB300 node: gsm8k_multiturn_sft preprocess -> messages column, DeepSeek-V4-Flash tokenizer (no builtin chat_template) + this generated template -> MultiTurnSFTDataset sample with non-empty assistant loss mask, and the override string survives hydra's grammar verbatim. Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

… plain clone Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

The fused-rope SFT NaN was a Megatron-Bridge config-mapping bug fixed in NVIDIA-NeMo/Megatron-Bridge#4271 (already on Megatron-Bridge main, which this example clones). Full-model DSv4-Flash SFT on 8xGB300 with rope fusion now matches the unfused control (iter-10 lm loss 2.3319 vs 2.3306, 0 NaN). Restores the varlen perf path. Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

shyoshyo

It appears that sparse attention fusion is not enabled by default unless explicitly specified. This can consume significant amounts of GPU memory in long-context scenarios.

shyoshyo · 2026-06-12T10:40:30Z

+    +engine.override_transformer_config.recompute_granularity=full \
+    +engine.override_transformer_config.recompute_method=uniform \
+    +engine.override_transformer_config.recompute_num_layers=1 \
+    "


It appears that sparse attention fusion is not enabled by default unless explicitly specified as +engine.override_transformer_config.apply_dsa_kernel_fusion=True

Good point, thanks for raising it. Two reasons this example keeps it off by default:

The pinned mcore asserts SM100+ for apply_dsa_kernel_fusion (transformer_config.py), so enabling it would break the H100 path this example also targets — and it additionally needs flash_mla from FlashMLA@nv_dev, which the setup here doesn't install. Added it as a commented Blackwell-only switch next to use_fused_mhc (6475db0).

Long context is currently bounded by CP support anyway (mcore ExecuteModelRequest cannot be found #5087 still open), so at the sequence lengths trainable today the unfused fallback is workable.

FWIW, the SM100 gate looks relaxable upstream now that cudnn-frontend v1.25.0 ships the SM90 indexer kernels and FlashMLA's sparse fwd supports SM90a — but that's an mcore-side change.

apply_dsa_kernel_fusion cuts memory via fused DSA kernels but the pinned mcore asserts SM100+ and it needs flash_mla; keep it off for the H100 path and document the Blackwell switch next to use_fused_mhc. Co-authored-by: Claude Signed-off-by: Lingrui Mei <lmei@nvidia.com>

gemini-code-assist Bot reviewed Jun 4, 2026

View reviewed changes

Meirtz mentioned this pull request Jun 4, 2026

[megatron] fix: tolerate extra kwargs in MTP postprocess patch #6596

Closed

Meirtz force-pushed the example/deepseek-v4-flash-sft branch 2 times, most recently from 28478a3 to 34824b1 Compare June 4, 2026 05:11

ISEEKYAN reviewed Jun 4, 2026

View reviewed changes

Meirtz force-pushed the example/deepseek-v4-flash-sft branch from 34824b1 to 433b544 Compare June 4, 2026 09:53

Meirtz requested a review from vermouth1992 as a code owner June 4, 2026 09:53

Meirtz and others added 2 commits June 4, 2026 18:25

Meirtz force-pushed the example/deepseek-v4-flash-sft branch from 433b544 to ac73b51 Compare June 4, 2026 10:25

[megatron] example: drop the environment comment block from DeepSeek-…

e4dfb11

…V4-Flash SFT script Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

Meirtz force-pushed the example/deepseek-v4-flash-sft branch from 9cd93be to e4dfb11 Compare June 4, 2026 12:55

This was referenced Jun 4, 2026

[recipe,docs,test] feat: add DeepSeek-V4-Flash SFT recipes, launcher and tests NVIDIA-NeMo/Megatron-Bridge#4131

Merged

DeepSeek-V4 training support NVIDIA/Megatron-LM#4468

Open

Meirtz mentioned this pull request Jun 10, 2026

fix(model): keep full rotary_percent for DeepSeek-V4 MLA rope NVIDIA-NeMo/Megatron-Bridge#4271

Merged

Meirtz and others added 2 commits June 11, 2026 18:15

[megatron] example: Megatron-Bridge PR verl-project#4131 merged — use…

4c77b6b

… plain clone Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Lingrui Mei <lmei@nvidia.com>

Meirtz mentioned this pull request Jun 11, 2026

feat(recipe): enable rope fusion in DeepSeek-V4-Flash SFT recipes NVIDIA-NeMo/Megatron-Bridge#4305

Merged

shyoshyo reviewed Jun 12, 2026

View reviewed changes

spacegoing mentioned this pull request Jun 27, 2026

[docker] Recommended image / Dockerfile for DeepSeek-V4 + DSA on the Megatron-Lite backend (B200 / H200)? #6864

Open

shyoshyo mentioned this pull request Jul 7, 2026

fix: add label shifting to mlite THD packing to align SFT loss with M… NVIDIA/Megatron-LM#5682

Merged

6 tasks

	export PYTHONPATH=$PYTHONPATH:$PYPATH/Megatron-Bridge/src:$PYPATH/Megatron-LM
	export PYTHONPATH=${PYTHONPATH:-}:$PYPATH/Megatron-Bridge/src:$PYPATH/Megatron-LM

Conversation

Meirtz commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Environment / image — please read first

Not a duplicate

Test

AI assistance

Uh oh!

CLAassistant commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 4, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Meirtz Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Meirtz commented Jun 4, 2026

Uh oh!

Meirtz commented Jun 10, 2026

Uh oh!

Meirtz commented Jun 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shyoshyo left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Meirtz commented Jun 4, 2026 •

edited

Loading

CLAassistant commented Jun 4, 2026 •

edited

Loading

Meirtz Jun 4, 2026 •

edited

Loading

Meirtz commented Jun 11, 2026 •

edited

Loading

shyoshyo left a comment •

edited

Loading