feat(RL): NeMo-RL × Megatron ModelExpress v2 weight-refit (trainer path) by KavinKrishnan · Pull Request #3068 · NVIDIA-NeMo/RL

KavinKrishnan · 2026-07-06T04:31:48Z

NeMo-RL × Megatron trainer/orchestration integration for ModelExpress native weight refit. Pairs with ai-dynamo/modelexpress#482.

What's here

Per-rank Megatron native-shard publishing with TP/PP/EP geometry and Bridge name-map sidecars; no trainer-side full-model gather.
Native vLLM MX updater integration, TP-aware MDL, EP filtering, and durable receiver lifecycle.
EP8→TP2 and TP2+EP2 recipes; GB200 RDMA resources, baked images, and fail-fast preflight.
Native-vLLM PyNccl sender/controller harness for a valid static baseline.

Validated target

Qwen3-30B-A3B, Megatron EP8/TP1 on 2 × 4-GPU GB200 nodes → vLLM TP2 on 2 GPUs.
Fully baked three-step GRPO: KL 0.0399–0.0643.
Steady MX+MDL refit: 10.25–11.94s.
MDL warm load: 2.7–4.1s; 291 direct + 144 fused + 18,432 expert writes, zero fallback.
TP2+EP2: 64 experts/rank, 9,651/9,651 coverage, 61.1→32.1 GB/rank (47.5% fewer bytes), KL 0.0586.
Native PyNccl baseline with GPU-preloaded sender: 5.17s update, 61.06 GB/rank, 94.5–96.6 Gbps/rank. Sender HBM preload (14.01s) is reported separately and excluded.

Interpretation

For this fixed full-model topology, NCCL is about 2× faster end-to-end. MX's raw RDMA wire is faster, but its current critical path is Megatron→HF translation/staging plus load. MX's differentiators remain expert filtering, heterogeneous layouts, elasticity, straggler isolation, and fan-out.

Artifacts

infra/nrl_k8s/examples/k8s_exemplars/V1/grpo_moe_qwen3_30b_ep8_tp2_dynamo_mx.yaml
...ep8_tp2_dynamo_mx.gb200.infra.yaml
...ep8_tp2ep2_dynamo_mx.yaml
infra/nrl_k8s/dynamo_mx/bench/preflight_ep8_tp2.sh
infra/nrl_k8s/dynamo_mx/bench/native_nccl_refit_bench.py

copy-pr-bot · 2026-07-06T04:31:51Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Trainer-side MX publish path for the Megatron policy worker (per-rank native shard + layout sidecar, stream_weights_via_mx), MX helpers (role detection, publish-set collection, MxConfig/builders), and reproducible K8s deploy + smoke artifacts (dense, mixed-TP, PP, MoE, multi-cycle, fan-out). Pairs with the MX client umbrella (modelexpress) and the native vLLM "mx" weight-transfer backend. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

…GD template Adds a self-contained guide to stand up the full NeMo-RL x Megatron x ModelExpress deployment (MX server + Megatron trainer + Dynamo/vLLM rollouts) with placeholders instead of cluster/namespace/personal specifics, and a generalized rollout DGD template showing both the native "mx" weight-transfer backend (going-forward) and the current --load-format mx + DYN_MX_REFIT_ENABLED flow. Points at the existing MX server manifest, V2 Megatron exemplar, and smoke harnesses. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

Scrubs personal image / secrets / cache paths / cluster selectors from the V2 rollout exemplar and replaces them with documented <PLACEHOLDER>s. Adds a header explaining what users must create from their own accounts/keys (worker image, a Secret with their HF token, a registry pull secret, PVC, RDMA claim), with kubectl create-secret examples. Adds IPC_LOCK + full UCX aux transports. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

…arks - examples/ep8_rollout_dgd.example.yaml: TP1 x DP8 x EP8 inference DGD (--enable-expert-parallel + --moe-backend triton + native "mx" backend) that makes EP byte-pruning measurable (current deploys are EP=1, where the filter is a no-op). Placeholder-only. - bench/: push-button differentiator orchestrator (run_differentiator_bench.sh) covering transport baseline, refit-phase breakdown, elastic-join, straggler, and partial/EP byte-pruning, with graceful skips when a live publisher / NCCL path / EP>1 deploy is absent; plus mx_vs_nccl_refit_bench.py, the synthetic EP>1 byte-pruning proof (ep_gt1_byte_pruning.py, runnable now, 8.0x), and the MDL partial-update smoke. README documents scenarios + prereqs. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

…stic, reg, fp8) Single-purpose harnesses that produced the 2026-07-08 GB200 numbers: mdl_correctness (corrupt->warm-reload byte-identity), ep_live_test (live EP placement parity + refit), wire_bench (cross-host 900 Gbps + EP byte-pruning), elastic_bench (decoupled/straggler timelines), reg_bench (arena vs per-tensor), fp8_h1_test (fp8 warm-path diagnosis). Kept in-tree so they survive /tmp wipes.

…it run

…DISABLE) - EP4 publishes over RDMA

…per-source pull balance, GT byte-identity)

…32+local) to gather all experts

…T (role-based shard axis)

… per-TP at load); full byte-identity check

…current) — true GPU->GPU transfer time

…oncurrent); fixes use_arena non-reentrancy

…rent multi-arena path)

…acked when no 4-free node)

…ing down under rank-packing; each publishes independently)

- detect_megatron_role/collect_megatron_publish_set: advertise GLOBAL expert id (expert_id + owned_expert_ids = ep_rank*num_local + local; keep local_expert_id) so an EP-trainer -> non-EP/lower-EP rollout gathers experts without collision. - ep_publisher: pass num_local_experts (num_experts // EP). - ep_e2e_rollout.py: live e2e that drives the PRODUCTION MxVllmWeightUpdater in a real vLLM engine (corrupt->refit->generate). NOTE: manual post-hoc refit fights vLLM's pre-owned NIXL/UCX in-process (degraded transport); the real proof needs the NeMo-RL framework wiring the mx backend at engine init.

update_weights_via_mx (megatron branch) now routes through MxVllmWeightUpdater (_update_weights_via_mx_native), which delivers full HF weights (vLLM slices per-TP at load) and supports EP-gather (EP trainer -> lower-EP/non-EP rollout) with per-source global expert remap, reusing the worker's existing receiver (one NIXL agent). This makes the communicated EP-trainer -> TP-rollout config work through the production path (TP1 trainer dense -> TP2 rollout handled by vLLM's loader; EP experts gathered+globalized). MX_MEGATRON_LEGACY_RECEIVER=1 falls back to the bespoke planner path. Full-HF gather is correctness-first; EP byte-pruning is a follow-up. Requires GRPO-loop validation (refit is orchestrated by the training loop).

… EP layout _update_weights_via_mx_native no longer hardcodes moe_expert_filter=False. It now reads cluster.weight_sync.mx_config.moe_expert_filter and introspects the live vLLM FusedMoE layout (ep_world_size/ep_rank/num_experts) to drive EP byte-pruning. Safe: filter is a no-op at EP1 (rollout owns all experts -> full gather, the proven path) and only prunes when the rollout is expert-parallel, using the model's ground-truth layout so we never pull the wrong expert subset.

Capture the exact GB200 EP4-to-TP1 topology and transport settings so the NIXL team can reproduce the TCP-to-RDMA lane selection issue without cluster-specific credentials.

Update the shared GB200 snapshot with the clean three-step result showing that excluding TCP keeps all source pulls on RDMA and removes the dominant refit stall.

Pass each vLLM worker's tensor-parallel identity into the converged MX updater so TP2 refits can validate rank-local behavior.

Provide the two-node GB200 Megatron EP8 and native-MX TP2 GRPO deployment used for the final correctness and performance gate.

Add fail-fast GPU/RDMA/backend checks, size rollout staging memory, and provide a TP2+EP2 variant for live expert-byte pruning validation.

Hold model, topology, batches, and step count constant while swapping only the weight-sync backend for the final MX comparison.

Use published trainer/rollout images, disable unavailable GB200 MNNVL, and size the GRPO batch for eight policy workers.

Use aligned trainer/MDL images, an explicit Dynamo command, namespace-local MX DNS, and a GRPO import preflight so no runtime code overlay is required.

Use vLLM PyNccl on both sender and rollout to measure packed checkpoint update time without NeMo communicator incompatibilities.

Standardize JSON metrics and assertions for filtering, TP slicing, partial refit, elastic joins, stragglers, tree fan-out, and trainer egress balance.

Make BF16/FP8 stage comparisons and MX differentiator runs reproducible with coordinated publishers, normalized timing output, correctness gates, and GB200 benchmark manifests.

Add a standalone two-rank P6 broadcast benchmark and label P5 preconsolidated versus unsupported true-EP semantics so comparisons cannot overclaim topology parity. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

Fall back to torch.distributed (TCPStore + NCCL) when nemo_rl is absent so the P6 wire-only baseline runs in the model-express-dev sender/receiver images; fix RDMA network annotation keys in the receiver manifest. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

Pull only tensors required by non-contiguous mixed-TP plans while retaining full-contiguous transfers, and add a real FP8 partial-refit behavioral correctness gate plus EP publisher byte accounting. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

Pass registry shapes with the pruned tensor-name set so row-parallel fallback receives shaped tensors for local axis-1 slicing instead of flat scratch buffers. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

github-actions Bot added Documentation Improvements or additions to documentation community-request labels Jul 6, 2026

KavinKrishnan force-pushed the kavink/megatron-mx-perf branch from 76b10c4 to 3a1cc9d Compare July 6, 2026 04:41

KavinKrishnan added 4 commits July 7, 2026 09:05

KavinKrishnan force-pushed the kavink/megatron-mx-perf branch from 6334212 to ab6c05a Compare July 7, 2026 22:01

KavinKrishnan added 21 commits July 8, 2026 09:00

bench: EP-configurable Megatron publisher for EP->TP2 first-party ref…

9fab978

…it run

bench: working EP publisher launch.sh (venv UCX/NIXL paths + NCCL_IB_…

95f1107

…DISABLE) - EP4 publishes over RDMA

bench: EP->TP receiver harness (multi-source EP-gather + TP reshard, …

50e8947

…per-source pull balance, GT byte-identity)

bench: EP->TP receiver — globalize expert ids per EP source (ep_rank*…

259701d

…32+local) to gather all experts

bench: EP->TP receiver — add TP>1 byte-identity check vs TP-sharded G…

7de0b85

…T (role-based shard axis)

bench: EP->TP receiver — receiver yields full HF weights (vLLM slices…

9624bc7

… per-TP at load); full byte-identity check

bench: EP4 device-RDMA wire harness (arena + stripe, sequential + con…

6e73720

…current) — true GPU->GPU transfer time

bench: EP4 wire — separate arena setup (sequential) from RDMA read (c…

576b10b

…oncurrent); fixes use_arena non-reentrancy

bench: EP4 wire — per-source isolated arena pull (drop fragile concur…

0881e9f

…rent multi-arena path)

bench: EP publisher — device-pack ranks onto visible GPUs (EP8 rank-p…

f1ba7ee

…acked when no 4-free node)

bench: EP publisher — drop post-publish barrier (kept ranks from tear…

cec646f

…ing down under rank-packing; each publishes independently)

bench: EP wire summary label uses EP_SIZE (was hardcoded EP4)

2492d0a

docs(bench): add sanitized NIXL RC debug config

c512031

Capture the exact GB200 EP4-to-TP1 topology and transport settings so the NIXL team can reproduce the TCP-to-RDMA lane selection issue without cluster-specific credentials.

docs(bench): record validated no-TCP RDMA fix

cc1e58c

Update the shared GB200 snapshot with the clean three-step result showing that excluding TCP keeps all source pulls on RDMA and removes the dominant refit stall.

feat(vllm_backend): propagate live rollout TP layout

9013278

Pass each vLLM worker's tensor-parallel identity into the converged MX updater so TP2 refits can validate rank-local behavior.

feat(k8s): add EP8 trainer to TP2 rollout target

6eb0359

Provide the two-node GB200 Megatron EP8 and native-MX TP2 GRPO deployment used for the final correctness and performance gate.

feat(k8s): gate EP8 TP2 deployment and EP filtering

f21ae57

Add fail-fast GPU/RDMA/backend checks, size rollout staging memory, and provide a TP2+EP2 variant for live expert-byte pruning validation.

KavinKrishnan added 2 commits July 10, 2026 18:31

bench(k8s): add EP8 TP2 NCCL comparison recipe

073d1cb

Hold model, topology, batches, and step count constant while swapping only the weight-sync backend for the final MX comparison.

fix(k8s): finalize baked EP8 TP2 deployment

9124f4c

Use published trainer/rollout images, disable unavailable GB200 MNNVL, and size the GRPO batch for eight policy workers.

KavinKrishnan mentioned this pull request Jul 11, 2026

feat(rl): NeMo-RL × Megatron × ModelExpress v2 weight-refit integration (umbrella) ai-dynamo/modelexpress#482

Draft

KavinKrishnan added 8 commits July 10, 2026 23:25

fix(k8s): make baked EP8 TP2 launch self-contained

b248bc1

Use aligned trainer/MDL images, an explicit Dynamo command, namespace-local MX DNS, and a GRPO import preflight so no runtime code overlay is required.

bench(nccl): add native TP2 refit baseline harness

8cb3822

Use vLLM PyNccl on both sender and rollout to measure packed checkpoint update time without NeMo communicator incompatibilities.

bench(mx): add seven-scenario differentiator suite

b81d549

Standardize JSON metrics and assertions for filtering, TP slicing, partial refit, elastic joins, stragglers, tree fan-out, and trainer egress balance.

bench(refit): add real-30B MX and NCCL measurement harnesses

4ef6208

Make BF16/FP8 stage comparisons and MX differentiator runs reproducible with coordinated publishers, normalized timing output, correctness gates, and GB200 benchmark manifests.

bench(nccl): add wire-only and EP source-layout baselines

62d784c

Add a standalone two-rank P6 broadcast benchmark and label P5 preconsolidated versus unsupported true-EP semantics so comparisons cannot overclaim topology parity. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

fix(refit): preserve shapes for strided scratch pulls

89c543a

Pass registry shapes with the pruned tensor-name set so row-parallel fallback receives shaped tensors for local axis-1 slicing instead of flat scratch buffers. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(RL): NeMo-RL × Megatron ModelExpress v2 weight-refit (trainer path)#3068

feat(RL): NeMo-RL × Megatron ModelExpress v2 weight-refit (trainer path)#3068
KavinKrishnan wants to merge 35 commits into
NVIDIA-NeMo:mainfrom
KavinKrishnan:kavink/megatron-mx-perf

KavinKrishnan commented Jul 6, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

KavinKrishnan commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What's here

Validated target

Interpretation

Artifacts

Uh oh!

copy-pr-bot Bot commented Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

KavinKrishnan commented Jul 6, 2026 •

edited

Loading