Nemotron RL support by Zhichenzzz · Pull Request #1284 · radixark/miles

Zhichenzzz · 2026-06-03T01:18:54Z

ci-sglang-pr: #27110
ci-megatron-pr: #51

sgl-project/sglang#27110
radixark/Megatron-LM#51

gemini-code-assist

Code Review

This pull request introduces a new offline conversion tool convert_hf_to_torch_dist_bridge.py for Nemotron-H models, simplifies conditional checks in model.py using walrus operators, and updates the MambaModel.forward shim to route MTP labels. The review feedback highlights two issues: a potential TypeError in the MambaModel.forward shim due to passing an unsupported loss_mask argument to the original forward method, and a potential directory nesting issue in the conversion tool if the target directory already exists during shutil.move.

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-03T01:20:17Z

+    def forward(self, *args, loss_mask=None, mtp_kwargs=None, **kwargs):
+        # Route miles' GPT-style mtp_kwargs['mtp_labels'] into MambaModel.forward's
+        # mtp_labels arg so MTP actually trains in RL: the MTP loss is gated on
+        # mtp_labels (mirroring gpt_model.forward), while the main path keeps
+        # labels=None and returns logits for the policy-gradient loss.
+        mtp_labels = (mtp_kwargs or {}).get('mtp_labels')
+        return _orig_forward(self, *args, loss_mask=loss_mask, mtp_labels=mtp_labels, **kwargs)


The original MambaModel.forward (captured in _orig_forward) does not accept loss_mask as an argument (which is the entire reason this shim exists, as documented in the docstring). Passing loss_mask=loss_mask to _orig_forward will raise a TypeError: forward() got an unexpected keyword argument 'loss_mask'. You should drop loss_mask when calling _orig_forward just like the original shim did.

Suggested change

def forward(self, *args, loss_mask=None, mtp_kwargs=None, **kwargs):

# Route miles' GPT-style mtp_kwargs['mtp_labels'] into MambaModel.forward's

# mtp_labels arg so MTP actually trains in RL: the MTP loss is gated on

# mtp_labels (mirroring gpt_model.forward), while the main path keeps

# labels=None and returns logits for the policy-gradient loss.

mtp_labels = (mtp_kwargs or {}).get('mtp_labels')

return _orig_forward(self, *args, loss_mask=loss_mask, mtp_labels=mtp_labels, **kwargs)

def forward(self, *args, loss_mask=None, mtp_kwargs=None, **kwargs):

# Route miles' GPT-style mtp_kwargs['mtp_labels'] into MambaModel.forward's

# mtp_labels arg so MTP actually trains in RL: the MTP loss is gated on

# mtp_labels (mirroring gpt_model.forward), while the main path keeps

# labels=None and returns logits for the policy-gradient loss.

mtp_labels = (mtp_kwargs or {}).get('mtp_labels')

return _orig_forward(self, *args, mtp_labels=mtp_labels, **kwargs)

gemini-code-assist · 2026-06-03T01:20:18Z

+    if dist.get_rank() == 0:
+        source_dir = get_checkpoint_name(args.save, 1, False, return_base_dir=True)
+        target_dir = get_checkpoint_name(args.save, -1, True, return_base_dir=True)
+        shutil.move(source_dir, target_dir)


If target_dir already exists (e.g., from a previous run of this conversion tool), shutil.move(source_dir, target_dir) will move source_dir inside target_dir as a subdirectory (e.g., target_dir/iter_0000001), rather than replacing it. To ensure consistent behavior across multiple runs, you should remove target_dir if it already exists before calling shutil.move.

Suggested change

if dist.get_rank() == 0:

source_dir = get_checkpoint_name(args.save, 1, False, return_base_dir=True)

target_dir = get_checkpoint_name(args.save, -1, True, return_base_dir=True)

shutil.move(source_dir, target_dir)

if dist.get_rank() == 0:

source_dir = get_checkpoint_name(args.save, 1, False, return_base_dir=True)

target_dir = get_checkpoint_name(args.save, -1, True, return_base_dir=True)

if os.path.exists(target_dir):

shutil.rmtree(target_dir)

shutil.move(source_dir, target_dir)

Code: minimal MTP fixes only (+21/-4): - forward_step: pre-shift mtp_labels (next-token) so MTP depth-k targets token[i+k+2] as the pretrained head expects (raw tokens were off-by-one). - train(): init mtp_losses=None and gate the metric (tracker can be empty). - bridge: clear mtp_hybrid_override_pattern when dropping MTP (finalize() would hit the nonexistent Symbols.MTP_SEPARATOR), and route miles mtp_kwargs[mtp_labels] into MambaModel.forward(mtp_labels=...). Scripts follow the existing per-model convention (sh+py runner, models/ config, bridge converter tool). The old branch walrus-style churn is gone (main already adopted that style). (cherry picked from commit 2ea1e9b224a45831e595b4e74e032cd59132ea7a)

Colocated weight sync IPC-shares the megatron->HF conversion transients with the engines; after del they sit in the producer IPC pending-free list, and under torch_memory_saver offload nothing triggers the lazy reap -- the freed pages accumulate as unpausable memory (+~full export size per training round) until rollout KV onload OOMs. Reap explicitly: once after update_weights() in the actor (covers every updater path, incl. the new atomic-group/direct-iterator ones) and after each tensor bucket sync (bounds peak pinning). Pairs with the engine-side ipc_collect in sglang (#27110).

# Conflicts: # miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py

Slices the 550B checkpoint down to source layers 0,1,7,8 (renumbered 0..3), which is the cheapest selection covering every block type the 108-layer model has: mamba, moe, attention, moe -> 'ME*E'. A prefix cut would need 8 layers to reach the first attention layer and pull in 4 MoE layers (~44B params) instead of 2. Everything else (512 experts, top-22, moe_latent_size=2048, sigmoid router + expert bias, AutoBridge path) matches the full model, so the runner exercises the same weight-conversion path on one node and asserts it with --check-weight-update-equal.

yueming-yuan · 2026-08-04T22:16:29Z

+                # Reap CUDA-IPC pending frees; under torch_memory_saver
+                # offload nothing else triggers the lazy reap (leaks ~full
+                # export size per round otherwise).
+                torch.cuda.ipc_collect()


is this needed? in my understanding ipc should be automatically collected in weight update (so other models work well)

repeated ipc collect, will delete this and that one in sglang

…550b_a55b.py The three .sh launchers (full run, 4-layer run, offline convert) duplicated the same arg blocks with the differences hardcoded. Merge them into the existing Python launcher using the ScriptArgs/typer convention the other model scripts use (run_qwen3_30b_a3b.py, run_glm5_744b_a40b.py), so every setting is a flag with MILES_SCRIPT_* env binding instead of an edit-the-file constant. Settings are unchanged: parallelism (TP8/PP4/EP32 for the full model on 128 GPUs, TP1/PP1/EP8 for the pruned slice), SGLang DP-attention sizing, GRPO, optimizer, and rollout args all reproduce the shell scripts verbatim. Two deliberate deltas, both to match the other launchers: --save goes to {output_dir}/{run_id}/checkpoints, and --dump-details is now passed. scripts/models/*.sh stay -- execute_train sources them for MODEL_ARGS. Also add the 4-layer slice as a model-scripts CI test on stage-c-8-gpu-h200. 8 GPUs rather than 4 because SGLang DP-attention needs attn_tp to divide Mamba n_groups=8, and --ci-test turns on the Megatron -> SGLang weight equality check that this test exists to guard.

yueming-yuan · 2026-08-04T22:34:52Z

qq: why do we need this script? given that we usually just load hf weight directly when we use megatron bridge

(I removed this file, but correct me if it is needed)

should remove it

Zero bytes (the canonical empty git blob), not on main, nothing references it. Swept into 2e63629 alongside the nemotron work -- 'core' is what Linux names a coredump in cwd.

…nvert miles' load_checkpoint dispatches on what --load points at: an HF directory goes to _load_checkpoint_hf, which is the same megatron.bridge path nemotron_h is mapped through, and args.load falls back to ref_load/hf_checkpoint. So bridge mode already loads HF directly and the offline conversion was never needed for correctness -- only to amortize the bridge mapping across runs. That optimization also came with a constraint the direct path does not have: the converted layout has to match the run's TP/PP/EP or the load reshards. Drop --use-torch-dist-ckpt and _prepare_torch_dist_ckpt, and remove tools/convert_hf_to_torch_dist_bridge.py, which nothing references now. It was a copy of tools/convert_hf_to_torch_dist.py with the mbridge backend swapped for megatron.bridge, minus the ROCm writer patch and --custom-model-provider-path. If the 550B's startup cost turns out to matter, the right shape is a backend branch on the existing tool's --megatron-to-hf-mode flag, which it already parses and currently ignores.

The test itself is correct but hf download of the pruned slice 404s, so it fails in 40s on every run. Re-enable by publishing the checkpoint and dropping the disabled= flag.

…ne reap point --use-kl-loss (coef 0) loads a ref model, which the noop weights backuper rejects under --disable-weights-backuper, so the CI test config could never run; run_deepseek.py already dropped the flag for the same reason. The per-chunk ipc_collect in the base sync loop is redundant (the next chunk's sends reap the previous chunk's frees); keep a single reap after update_weights(), the one point where nothing else runs before the trainer pauses for the rollout.

The pruned checkpoint is published at CharyZeng/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-4layer (layers 0,1,7,8 of the 550B BF16 renumbered 0..3, MTP head kept).

The 4-layer A/B (all ipc_collect calls removed, 6 rounds, offload on, with and without --ci-test) shows no memory growth; the automatic pending-free collection on the next round's IPC activity is sufficient. actor.py is now identical to main.

Zhichenzzz requested review from fzyzcjy, maocheng23, yueming-yuan and yushengsu-thu as code owners June 3, 2026 01:18

gemini-code-assist Bot reviewed Jun 3, 2026

View reviewed changes

Zhichenzzz mentioned this pull request Jun 3, 2026

Megatron: Nemotron RL support radixark/Megatron-LM#51

Merged

Zhichenzzz changed the title ~~feat(nemotron_h): MTP training for Mamba-hybrid path + offline torch_dist convert~~ Nemotron RL support Jun 3, 2026

Zhichenzzz mentioned this pull request Jun 3, 2026

[RL] Nemotron support on sglang-miles sgl-project/sglang#27110

Merged

Zhichenzzz added 2 commits June 11, 2026 01:09

Zhichenzzz force-pushed the nemotron-h-mtp branch from 16ddc6f to 3d811f4 Compare June 11, 2026 01:20

fzyzcjy requested a review from Shi-Dong as a code owner July 25, 2026 23:42

Zhichenzzz requested a review from guapisolo July 29, 2026 01:33

Zhichenzzz and others added 3 commits July 28, 2026 18:33

Merge branch 'main' into nemotron-h-mtp

28553be

Merge remote-tracking branch 'origin/main' into nemotron-h-mtp

918fd62

# Conflicts: # miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py

yueming-yuan reviewed Aug 4, 2026

View reviewed changes

Comment thread miles/backends/megatron_utils/actor.py Outdated

yueming-yuan reviewed Aug 4, 2026

View reviewed changes

yueming-yuan added 2 commits August 4, 2026 15:35

Remove stray empty 'core' file committed at the repo root

b5cd674

Zero bytes (the canonical empty git blob), not on main, nothing references it. Swept into 2e63629 alongside the nemotron work -- 'core' is what Linux names a coredump in cwd.

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

yueming-yuan and others added 6 commits August 4, 2026 15:58

Apply black formatting

80250a1

ci: disable the nemotron 4-layer test until its checkpoint is published

11a1a7e

The test itself is correct but hf download of the pruned slice 404s, so it fails in 40s on every run. Re-enable by publishing the checkpoint and dropping the disabled= flag.

Merge remote-tracking branch 'origin/main' into nemotron-h-mtp

09c8fc2

Enable the nemotron 4-layer CI test

af22b57

The pruned checkpoint is published at CharyZeng/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-4layer (layers 0,1,7,8 of the 550B BF16 renumbered 0..3, MTP head kept).

yueming-yuan approved these changes Aug 5, 2026

View reviewed changes

Zhichenzzz merged commit 98eec31 into main Aug 5, 2026
41 checks passed

Zhichenzzz deleted the nemotron-h-mtp branch August 5, 2026 05:49

Zhichenzzz mentioned this pull request Aug 5, 2026

fix(mtp): double-shift GPT-path MTP labels #2215

Merged

Navigation Menu

Uh oh!

Nemotron RL support - #1284

Nemotron RL support#1284
Zhichenzzz merged 14 commits into
mainfrom
nemotron-h-mtp

Zhichenzzz commented Jun 3, 2026 •

edited by yueming-yuan

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jun 3, 2026

Uh oh!

gemini-code-assist Bot Jun 3, 2026

Uh oh!

yueming-yuan Aug 4, 2026

Uh oh!

Zhichenzzz Aug 5, 2026

Uh oh!

Uh oh!

yueming-yuan Aug 4, 2026

Uh oh!

yueming-yuan Aug 4, 2026

Uh oh!

Zhichenzzz Aug 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

Zhichenzzz commented Jun 3, 2026 • edited by yueming-yuan 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 3, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

yueming-yuan Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

Zhichenzzz Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yueming-yuan Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

yueming-yuan Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

Zhichenzzz Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Zhichenzzz commented Jun 3, 2026 •

edited by yueming-yuan

Loading