DeepSeek-V4 support-v2: native megatron implementation as option by yueming-yuan · Pull Request #2706 · radixark/miles

yueming-yuan · 2026-08-21T22:47:41Z

ci-megatron-pr: dsv4-dual-backend

What

Bring DeepSeek-V4 back on the bumped Megatron, and make the training implementation an explicit
choice. Second of the three bump PRs; stacked on #2673 (base bump), which disabled DSv4.

--dsv4-impl {miles,megatron} selects a whole path, at config+spec build time:

	`miles` (default)	`megatron`
attention variant	`dsv4` (ported PR #28 slice)	`dsv4_hybrid` (dev-native)
mHC	miles hyper-connections	native `enable_hyper_connections`
kernels	tilelang, BSHD + sparse CP	cudnn / none
tensor parallelism	yes	TP must be 1

Unsupported combinations raise, they do not fall back — June's MILES_DSV4_ATTENTION_BACKEND
silently forced miles on TP>1 and that behaviour is retired.

Megatron side — `radixark/Megatron-LM:dsv4-dual-backend`

miles-main-20260819 + 6 commits: the miles DSv4 attention variant, hash-routing under TP + MoE
freeze knobs, and four fixes found by running it:

tid2eid was a buffer, not a parameter — miles' weight sync materializes full tensors from the
parameter buffer, so it pushed an all-zero routing table into sglang and the engine became a
different model. The weight checker caught this on the first run.
the shared-expert SwiGLU clamp opt-out (activation_func_clamp_shared_expert) was lost when the
mHC revert landed, so training clamped what DeepSeek-V4 and the engine leave unclamped.
rollout routing replay registered on MTP routers, which desyncs the replay streams.
three sites in csa.py read Optional dsa_indexer_loss_coeff with
getattr(..., 0.0) — whose default only fires when the attribute is missing — so None
reached kl_div * loss_coeff. Two other sites in the same file already did it correctly.

Miles side (this PR)

--dsv4-impl, with the arg translation and the three hard-errors, plus its unit test.
The plugin's weights now carry Megatron's names and layout, so one HF checkpoint loads into both
impls; mbridge, the quantizers, the weight checker and the converters follow the rename.
fp32-only params use Megatron's mark_keep_in_fp32.
The DSv4 4-layer e2e is re-enabled (it was disabled= in Bump Megatron-LM to miles-main-20260819 (latest NVIDIA dev) #2673).
4-layer conversion runs at PP1 — the bumped Megatron asserts hash-MoE + PP>1 needs an explicit
pipeline_model_parallel_layout.

Verification (devbox, 8xH200)

4-layer RL e2e, default thresholds, no skip list, no relaxed gate — matches the pre-bump stack
(miles 49cac7b52 + megatron 2e64d788b):

	pre-bump	this PR
rollout 0 / 1 train-rollout gap	0.0103 / 0.0094	0.0103 / 0.0099
`train_rollout_logprob_abs_diff` step 0 / 1	0.0833	0.0858 / 0.0836
`train_rollout_kl` step 0 / 1	0.00927	0.00973 / 0.00942
`found NaN`	0	0
R3 mismatch tokens	0	0
job	succeeded	succeeded

Both impls, same recorded rollout, same args, TP=1:

	miles impl	megatron impl
rollout logprob (input)	-1.6760895252227783	-1.6760895252227783
train-forward logprob	-1.6852522790431976	-1.6873207092285156
`train_rollout_logprob_abs_diff`	0.0837	0.0885
job	succeeded	succeeded

Native mHC and miles mHC agree to 2.1e-3. Converting with --dsv4-impl megatron gives
27,390,526,557 parameters, identical to the miles impl.

Not covered: every sample in that rollout has reward 0, so GRPO advantage and grad_norm are 0
on both impls — the backward ran clean but its numerics are untested. Also untested: TP>1 (rejected
by design), the cudnn kernel backend (needs flash_mla, absent from the image), multi-step runs of
the megatron impl.

Full investigation, including the wrong turns, in bump_docs/05-dsv4-dual-backend.md.

50-step A/B on 8×GB300 (32 GPU, sm103)

Identical config for both runs — TP1 · PP8 (first 4 / last 3 layers) · CP1 · EP4 · DP4,
max-tokens-per-gpu 4096, FP8 blockwise, fp32 optimizer, GRPO on DAPO-math-17k (32 prompts × 8
samples, max response 4096, lr 1e-6, deterministic mode), cudnn kernel backend for conversion and
training — only --dsv4-impl differs. megatron = job 2695 (50 steps), miles = job 2703 (54 steps).

Timing, p50 over steps 2+ (steps 0–1 excluded: per-shape JIT warm-up). Full cycle = wall-clock
between consecutive rollout completions.

phase	megatron impl	miles impl	Δ
train	82.6 s	79.0 s	−4.4%
log_probs	23.2 s	22.9 s	−1.3%
update_weights	38.4 s	38.3 s	−0.3%
full cycle	200.5 s	197.0 s	−1.7%

	megatron impl · 2695	miles impl · 2703
steps	50	54
reward, mean	0.763	0.765
`train_rollout_kl`, mean	0.0270	0.0254
logprob diff (train − rollout), mean	−0.0313	−0.0295
logprob diff, range	−0.058 … −0.011	−0.066 … −0.012

Megatron dev reimplemented PR #6's MTP-in-RL support natively: - process_mtp_loss derives MTP labels from input_ids when labels is None (RL). - config.mtp_detach_heads detaches output head + MTP embedding gradients. So on the miles side: set config.mtp_detach_heads=True when enable_mtp_training, and stop passing the now-unsupported mtp_kwargs to GPTModel.forward (labels=None + input_ids derivation is equivalent to mtp_labels=batch['tokens']).

… lora test mock Matches the dev Megatron arg rename used in model.py (use_gloo_process_groups).

…put_with_bias New dev split the MLP into _forward_mlp_output_with_bias + _forward_post_mlp and dropped _sglang_pre_mlp_residual (true-on-policy was deferred). Anchor pre_mlp_residual on the unconditional nvtx_range_push("mlp") (residual set in all branches); keep pre_mlp_layernorm_output and mlp_output_with_bias anchors. Will re-verify after the true-on-policy migration.

…bias The megatron true-on-policy merge relocated the _sglang_pre_mlp_residual extraction out of _forward_mlp into _forward_mlp_output_with_bias, so the BSHD source-patcher anchor 'residual = getattr(self, "_sglang_pre_mlp_residual", hidden_states)' no longer exists. Mirror the THD config: target _forward_mlp_output_with_bias and prepend the pre_mlp_residual dump at the mlp nvtx push (matches the new structure). Fixes test_run_megatron + test_miles_dumper source-patch failures.

…_module) New dev's build_module passes name= to every spec-built module (megatron base Attention.__init__ accepts it). miles' custom attention classes overriding __init__ must accept + thread it: DSAMultiLatentAttention/DSAMLASelfAttention (glm5, DeepSeek-V3.2), HuggingfaceAttention base, qwen3_5/qwen3_next Attention. Mirrors the DeepSeekV4Attention name fix. Fixes 'DSAMLASelfAttention.__init__() got an unexpected keyword argument name'.

…lders New dev removed the moe_use_legacy_grouped_gemm argument and the legacy grouped-GEMM concept (grouped_mlp_modules now takes only moe_use_grouped_gemm). miles' model_provider and glm4 spec builders read args.moe_use_legacy_grouped_gemm, which no longer exists -> AttributeError on the Namespace. Drop the kwarg (dev spec fns take **kwargs / the new grouped path is the only one).

New Megatron renamed MultiTokenPredictionLayerSubmodules.transformer_layer -> mtp_model_layer, so MTP weight names are now mtp.layers.{i}.mtp_model_layer.*. miles' MTP weight converters (mcore<->hf) hardcoded 'transformer_layer' -> 'Unknown MTP component' on rollout weight export (test_mimo_7B_mtp_only_grad). Accept both names across export (mimo/qwen3_next/deepseekv3/glm4moe/qwen3_5 + fp8/ mxfp8/nvfp4 quantizers) and mbridge import (mimo/qwen3_next/qwen3_5/glm4moe_lite/glm4moe).

New dev rewrote MTPLossLoggingHelper: it stores loss_sums/num_tokens (or loss_values) and only computes tracker['values'] after reduce_loss_in_tracker() (which also does the cross-rank all-reduce). miles read tracker['values'] directly (never populated) -> 'values' not in tracker -> mtp_losses unbound -> UnboundLocalError at extra_metrics. Call reduce_loss_in_tracker(), read values/loss_values, init mtp_losses=None + guard the metric. Fixes test_mimo_7B_mtp_only_grad past weight-export.

New Megatron reads disable_{param,grad}_buffers_cpu_backup from DistributedDataParallelConfig; _ParamAndGradBuffer no longer takes them as ctor kwargs.

…ollow-up PR

…rd-only run-megatron skips DDP wrap

Rebuilds were driven by whether the PR diff touched docker paths, so every rerun of a Dockerfile-touching PR paid for an identical multi-arch build. Hash the build inputs, stamp the hash on the published tag as a label, and compare against it: a PR keeps one pr-<num> tag and rebuilds it only when the inputs actually move, when the tag is missing, or on the one-shot rebuild-ci-image label. resolve-ci-image now selects the PR image whenever that tag is current rather than only when this run built it, so reusing an image does not silently fall back to dev.

Also scope pull-requests: write to the docker-build call instead of the whole workflow; only that job removes the one-shot label.

…' into megatron-bump-20260819

New dev returns functools.partial builders for the mlp and shared-experts submodules instead of ModuleSpec, and MoELayer now passes name= to the shared-experts builder.

Megatron now rejects the combination: fake QAT swaps in straight-through weight tensors while TE's fused wgrad accumulation writes main_grad onto the original ones, so the quantized weights' gradients would be dropped.

Megatron's dsa variant now builds AbsorbedMLASelfAttention, which reads linear_kv_up_proj.weight directly; the pinned Bridge's LoRALinear has no such property. Upstream Bridge added one that returns the LoRA-effective weight.

Megatron's Mamba mixer exposes conv1d as conv1d_weight/conv1d_bias parameters; the pinned Bridge maps only the conv1d.* module names, so the weight sync skips them silently. Upstream Bridge accepts both spellings.

New Megatron renamed the MTP submodule to mtp_model_layer, so MTP expert weights missed the expert-offset rewrite and ranks broadcast differing tensor sets during weight sync, deadlocking the EP collectives.

Megatron now gates the HF tokenizer on args.trust_remote_code, which miles never set, so checkpoints shipping custom tokenizer code (Kimi) failed to load. miles already passes trust_remote_code=True everywhere it loads a checkpoint itself.

The private _keep_fp32 attribute was only honoured by the fp32 snapshot that rode in on the DSv4 megatron commit, which the base bump defers. Upstream now owns this: mark_keep_in_fp32 exempts the tensor from Float16Module's cast, so the scale survives the bf16 round trip through the weight sync.

TopKRouter.forward gained packed_seq_params upstream, which MoELayer passes positionally. Separately, the plain bridge provider never forwarded gradient_accumulation_fusion, so --no-gradient-accumulation-fusion could not reach models built through it (the LoRA bridge path already synced it).

# Conflicts: # .github/workflows/_build-pr-ci-image.yml

Megatron now resolves te_general_gemm, which it could not on TE 2.17 before the bump, and TE multiplies in the operand dtype -- so --moe-router-dtype fp32 became a bf16 multiply. The last ulp is enough for top-k to pick a different expert, and on qwen3.5 that flipped one token in layer 0 and cascaded to a 15% logprob drift. 25 model scripts ask for fp32 routing, so default it here rather than per model.

claude

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

#2673 was squash-merged, so main's base bump and this branch's copy of it are unrelated histories. Three resolutions git could not get right on its own: - docker/Dockerfile: main is net-unchanged from the merge base (bump set miles-main-20260819, #2734 set it back), so git kept this branch's miles-main-20260819 — a branch that no longer exists on the Megatron remote. Take miles-main. - test_deepseek_v4_flash_4layer_ci.py: the inverse — this branch is net-unchanged, so git kept main's disabled=. Re-enable it; that is this PR. - model.py mtp_loss_scale: this branch predates 59cc3c1, so its side is the unfixed line. Take main's.

Local devbox sync config generated by `rx devbox sync start`; swept into cb252c2 by accident.

…recipe --dsv4-impl reaches both the conversion (the spec is built there too) and the training args. The 8-node GB300 recipe gains a megatron-impl branch: TP must be 1 there, and CP would need qkv_format=thd, so both go to DP and the per-rank token budget doubles to keep max_tokens_per_gpu * cp_size unchanged.

…raining The conversion builds the model spec too, so it hits the same backend validation; --extra-args only reaches _train, which left the conversion on the default (cuDNN for --dsv4-impl megatron, whose flash_mla is not in the image).

Same recipe as the miles-impl test with the megatron implementation selected; the PyTorch DSA fallback is used until the CI image ships flash_mla and cudnn-frontend, so the fused-kernel path stays uncovered.

The --dsv4-* aliases duplicated flags Megatron already exposes (--csa-*, --o-groups, --o-lora-rank, --num-residual-streams, --mhc-sinkhorn-iterations, --moe-n-hash-layers); the plugin now only declares --dsv4-impl. Also regenerates the model_args snapshots, stale since 43c5f1c.

…kend-drop-colocate

…ic batching

…is thd-only

…e lock

Zhichenzzz

Most looks good! just small comments

Zhichenzzz · 2026-08-31T05:12:40Z

+        # RL rollout-routing-replay (R3) seam for the sparse-attention indexer topk: lets the miles
+        # indexer_replay_manager record (on rollout) / replay (on train forward) the top-k KV picks,
+        # mirroring the MoE routing-replay seam. No-op unless the manager is enabled.
+        indexer_replay_manager.register_to_module(self, "indexer_replay", stream_idx=layer_id)


(maybe later we should opt this indexer replay, since now it is easy to cause host oom / and long-latency. and this will be very help for the following models, e.g., glm, qwen, ...

Zhichenzzz · 2026-08-31T05:41:37Z

@@ -1 +1 @@
 import torch


rename the file into deepseek_v4?

Zhichenzzz · 2026-08-31T05:44:27Z



-def convert_deepseekv4_to_hf(args, name, param):
+def _packed_alphas(name: str, param, bucket):


claude code says this will harm async train, but i dont agree (maybe quick check

just checked, this should not harm the training

Zhichenzzz · 2026-08-31T05:45:31Z

        # Non-interleaved indexers keep wk as a standalone FP8 parameter in SGLang.
        fp8_param_names.extend(
            [
                "self_attention.wk.weight",


Zhichenzzz · 2026-08-31T05:47:09Z

@@ -1,2 +1,2 @@
 """
 DeepSeek V4 training script.


do we also need to modify amd script into the latest impl?

…zers

…pshots

yueming-yuan added 30 commits August 19, 2026 20:33

test: rename enable_gloo_process_groups -> use_gloo_process_groups in…

59393ba

… lora test mock Matches the dev Megatron arg rename used in model.py (use_gloo_process_groups).

fix: read renamed use_gloo_process_groups megatron arg

08af1ca

fix: set TMS cpu-backup flags on ddp_config in LoRA buffer patch

a265bf7

New Megatron reads disable_{param,grad}_buffers_cpu_backup from DistributedDataParallelConfig; _ParamAndGradBuffer no longer takes them as ctor kwargs.

megatron-bump: disable true-on-policy on megatron backend until the f…

f175d92

…ollow-up PR

megatron-bump: disable DSv4 4-layer e2e until the DSv4 follow-up PR

2d5b1e9

docker: bump MEGATRON_BRANCH to miles-main-20260819

8bc58b9

fix: vocab_size_with_padding moved to megatron.core.tokenizers; forwa…

3e16b8e

…rd-only run-megatron skips DDP wrap

cleanup: trim comments

ab39831

ci: update workflow structure tests for the rebuild-decision job

a8d68a6

Also scope pull-requests: write to the docker-build call instead of the whole workflow; only that job removes the one-shot label.

Merge remote-tracking branch 'origin/ci-image-rebuild-on-input-change…

316dab0

…' into megatron-bump-20260819

fix(inkling): adapt layer specs to megatron's MlpBuilder partials

6643598

New dev returns functools.partial builders for the mlp and shared-experts submodules instead of ModuleSpec, and MoELayer now passes name= to the shared-experts builder.

fix: disable gradient accumulation fusion under INT4 fake QAT

2921423

Megatron now rejects the combination: fake QAT swaps in straight-through weight tensors while TE's fused wgrad accumulation writes main_grad onto the original ones, so the quantized weights' gradients would be dropped.

test: disable the Nemotron-3 Ultra e2e until Megatron-Bridge is bumped

9f5f079

Megatron's Mamba mixer exposes conv1d as conv1d_weight/conv1d_bias parameters; the pinned Bridge maps only the conv1d.* module names, so the weight sync skips them silently. Upstream Bridge accepts both spellings.

fix: accept the renamed MTP submodule in expert weight-sync names

ddf69ce

New Megatron renamed the MTP submodule to mtp_model_layer, so MTP expert weights missed the expert-offset rewrite and ranks broadcast differing tensor sets during weight sync, deadlocking the EP collectives.

Merge remote-tracking branch 'origin/main' into megatron-bump-20260819

3dcb90d

# Conflicts: # .github/workflows/_build-pr-ci-image.yml

chore: trim comments in megatron backend defaults

6d46875

Format the MTP name-mapping edits

6ddbd76

claude Bot reviewed Aug 21, 2026

View reviewed changes

Base automatically changed from megatron-bump-20260819 to main August 24, 2026 19:37

yueming-yuan and others added 5 commits August 24, 2026 13:07

Drop .rxignore

5356fc1

Local devbox sync config generated by `rx devbox sync start`; swept into cb252c2 by accident.

Merge branch 'main' into dsv4-dual-backend

25a1171

zhoutong-hai mentioned this pull request Aug 26, 2026

feat(deepseek-v4): add native LoRA construction and conversion bridge #2772

Draft

e2e: cover --dsv4-impl megatron on the 4-layer Flash recipe

4b651a4

Same recipe as the miles-impl test with the megatron implementation selected; the PyTorch DSA fallback is used until the CI image ships flash_mla and cudnn-frontend, so the fused-kernel path stays uncovered.

yueming-yuan force-pushed the dsv4-dual-backend branch from 2bb54d6 to 4b651a4 Compare August 27, 2026 07:17

yueming-yuan mentioned this pull request Aug 27, 2026

Add --colocate-memory-peak-device: overlap the trainer/rollout handoff on the GPU #2781

Merged

yueming-yuan added 3 commits August 27, 2026 00:33

dsv4: simplify the stale-checkpoint check

1a454db

Merge remote-tracking branch 'origin/main' into yueming/dsv4-dual-bac…

4f98c7d

…kend-drop-colocate

yueming-yuan added the run-ci-model-scripts Run model script smoke tests label Aug 27, 2026

fix: single-node smoke config selects TP1 for --dsv4-impl megatron

03856d9

yueming-yuan force-pushed the dsv4-dual-backend branch from ea4c4bb to 03856d9 Compare August 27, 2026 21:29

yueming-yuan changed the title ~~DeepSeek-V4 on the bumped Megatron, with a selectable training implementation~~ DeepSeek-V4 support-v2: native megatron implementation as option Aug 27, 2026

yueming-yuan added 4 commits August 27, 2026 14:37

Default DeepSeek-V4 to --dsv4-impl megatron with thd layout and dynam…

bed7093

…ic batching

Keep --micro-batch-size 1 on the bshd (miles) path; dynamic batching …

7ed69ca

…is thd-only

test: pin --dsv4-impl in cases that depended on the old default

f913b76

fix: serialize prepare steps racing on a host-shared cache with a fil…

8a19557

…e lock

Zhichenzzz approved these changes Aug 31, 2026

View reviewed changes

yueming-yuan added 5 commits August 31, 2026 12:55

fix: update v4 indexer wk to its renamed path in the fp8/mxfp8 quanti…

3d71ce0

…zers

rename mbridge/deepseekv4.py to deepseek_v4.py

11b5421

fix: normalize trailing-slash path_dst in the prepare-step lock

00f10e5

amd: pin --dsv4-impl miles (no cudnn path on ROCm); refresh stale sna…

19e3f8d

…pshots

format: isort import order after mbridge rename

745153c

yueming-yuan merged commit 6d0cd93 into main Aug 31, 2026
22 of 31 checks passed

yueming-yuan deleted the dsv4-dual-backend branch August 31, 2026 20:21

Navigation Menu

Uh oh!

DeepSeek-V4 support-v2: native megatron implementation as option - #2706

DeepSeek-V4 support-v2: native megatron implementation as option#2706
yueming-yuan merged 66 commits into
mainfrom
dsv4-dual-backend

yueming-yuan commented Aug 21, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

Zhichenzzz left a comment

Uh oh!

Zhichenzzz Aug 31, 2026

Uh oh!

Zhichenzzz Aug 31, 2026

Uh oh!

Zhichenzzz Aug 31, 2026

Uh oh!

yueming-yuan Aug 31, 2026

Uh oh!

Zhichenzzz Aug 31, 2026

Uh oh!

yueming-yuan Aug 31, 2026

Uh oh!

Zhichenzzz Aug 31, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		def convert_deepseekv4_to_hf(args, name, param):
		def _packed_alphas(name: str, param, bucket):

Uh oh!

Conversation

yueming-yuan commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Megatron side — radixark/Megatron-LM:dsv4-dual-backend

Miles side (this PR)

Verification (devbox, 8xH200)

50-step A/B on 8×GB300 (32 GPU, sm103)

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

Zhichenzzz left a comment

Choose a reason for hiding this comment

Uh oh!

Zhichenzzz Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Zhichenzzz Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Zhichenzzz Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

yueming-yuan Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Zhichenzzz Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

yueming-yuan Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Zhichenzzz Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yueming-yuan commented Aug 21, 2026 •

edited

Loading

Megatron side — `radixark/Megatron-LM:dsv4-dual-backend`