feat(megatron-mx): Megatron-Core MX path — Phases A+C+D#2
Merged
jthomson04 merged 12 commits intoJun 12, 2026
Merged
Conversation
Author
Cluster validation — 2026-06-08Cross-pod end-to-end smoke validated on the GB200 kavin cluster. All 8 HF-shaped output tensors byte-identical to trainer ground truth across qkv_column / gated_mlp_column / column / row / replicated roles. See the full writeup at This run validates the publisher (Phase A) + Bridge-derived sidecar emit ( |
Author
|
Phase E shape 1 validation on GB200 (kavin namespace) — real Megatron model byte-identity confirmed. Loaded Result: 398/398 HF tensors byte-identical, 0 drift, 0 missing, 0 extras.
What this validates beyond the earlier synthetic-shape smoke:
Two small bugs surfaced + fixed during this run:
Full writeup + harness scripts: Not yet validated (Phase E shape 2): mixed-TP (target_tp ≠ source_tp), MoE per-expert routing, real GRPO loop. The mixed-TP gap is just the v1 |
Author
|
Phase E shape 2 progress update — partial validation on Qwen3-MoE-30B-A3B-Instruct-2507 (kavin GB200, 60 GB, 128 experts, 48 layers). What's now validated on the real MoE model:
Two real-model bugs surfaced + fixed during this run (commit
Pre-fix vs post-fix role distribution on real Qwen3-30B-A3B: Also landed (cluster validation pending a 2-GPU trainer pod): Mixed-TP receiver path in What still needs to happen for cluster-validated shape 2:
Full writeup: |
…lassifier Phase A of the Megatron-Core MX integration plan (temp/NemoRL_Megatron_MX_Design.md §4). Adds the trainer-side publisher that emits per-rank Megatron-native shards to ModelExpress, with each parameter classified into one of seven Megatron roles (qkv_column, gated_mlp_column, column, row, vocab_parallel, replicated, expert_column / expert_row). Receiver-side slice planner is in ai-dynamo/modelexpress PR NVIDIA-NeMo#421 commits 12c73a7 + b26e80f. What's new * nemo_rl/distributed/mx_megatron_helpers.py - detect_megatron_role(name, param, model, ...) — classifies parameters by walking to the enclosing module and matching against Megatron-Core's parallel-layer class names (ColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding, plus heuristic detection of fused QKV / fused gated MLP via name patterns). - collect_megatron_publish_set(model, tp_size=..., pp_rank=..., ...) — generator that yields (name, local_shard, role_spec, full_extras) tuples ready for MxV2TrainingPublisher.add_tensor. Skips replicated tensors on tp_rank > 0 (rank 0 publishes for the replicated set). - Role overrides via MxConfig.megatron_role_overrides for non-mainline Megatron forks with different class / name conventions. - Expert pattern via NRL_MX_EXPERT_TENSOR_PATTERN env var, mirroring detect_moe_expert_layout's convention. * MegatronPolicyWorkerImpl.stream_weights_via_mx - Lazy-init the v2 publisher (once per worker lifetime). - Resolve TP / PP / EP rank position via parallel_state. - Pull num_attention_heads / num_query_groups / kv_channels from the trainer's transformer_config; descriptor extras for qkv_column become receiver-checkable. - Stamp the publisher's mesh position before publish() so the source identity / sidecar carry tp_rank / pp_rank / ep_rank. - For each parameter: classify role, emit local shard with per-tensor megatron_role + megatron_extras in the registry via add_tensor's new kwargs. - publish(version=N) + mark_ready() — same shape as the DTensor stream_weights_via_mx. Memory invariant Trainer never materialises the full tensor. Megatron-Core stores native shards; the publisher passes them through (tensor.detach() + dtype cast + .contiguous()). Peak GPU memory for any layer stays at model_weights / (tp * pp * ep). Out of scope (this commit) * Phase C: receiver translator (_translate_megatron_to_hf, _uninterleave_qkv) — lives on the inference side, separate file + PR. * Phase D: GRPO refit branch + receiver auto-detect. * FP8 KV-cache scales (existing NotImplementedError matches the DTensor path's behaviour). Tested by * Pre-existing unit suite continues to pass (no functional changes to existing paths). * MX-side planner: 36/36 v2 tests passing in ai-dynamo/modelexpress @ kavink/fix-publish-self-as-source-v2-transports. * End-to-end validation pending — needs Phase C + D + new image + GB200 deploy. Tracked in temp/NemoRL_Megatron_MX_Design.md §8 / §11.
…ole classification
Pivot the role classifier in detect_megatron_role to consult
megatron.bridge.models.conversion.param_mapping.AutoMapping._MODULE_TYPE_REGISTRY
as the authoritative source. Bridge's registry is curated to cover every
TE / Inference / Quant variant of column-parallel, row-parallel, and
replicated modules:
column ColumnParallelLinear, TEColumnParallelLinear,
TELayerNormColumnParallelLinear, TEColumnParallelGroupedLinear,
InferenceLayerNormColumnParallelLinear, QuantColumnParallelLinear,
LinearCrossEntropyModule, VocabParallelEmbedding,
DotProductAttention (attention sink), TEDotProductAttention
row RowParallelLinear, TERowParallelLinear, TERowParallelGroupedLinear,
QuantRowParallelLinear, InferenceRowParallelLinear
replicated TENorm, FusedLayerNorm, WrappedTorchNorm, LayerNorm, RMSNorm,
L2Norm, InferenceTopKRouter, IdentityOp, LinearForLastLayer,
TopKRouter
Previously the classifier matched on substrings of the base class names
("ColumnParallel", "RowParallel"), which silently misclassified many TE
variants as replicated and would have wasted publish bandwidth (best
case) or routed shards through the wrong assembly path (worst case).
When Bridge is not importable (e.g. CPU-only unit tests), the classifier
falls back to substring matching against the base names — sufficient for
mainline Megatron-Core. The fallback also handles the
LayerNormColumnParallelLinear special-case Bridge tracks separately.
Why this matters
The receiver-side translator (Phase C) is a thin adapter over Bridge's
own MegatronParamMapping.megatron_to_hf calls (and the standalone
helpers like split_qkv_weights / merge_qkv_weights). By using Bridge's
classification on the publisher side too, both halves of the transport
agree on what a tensor's role is, regardless of whether it's a
mainline Megatron-Core class or a TE variant. No more parallel
truth-source to maintain.
See temp/NemoRL_Megatron_MX_Design.md §9b ("Architecture pivot:
Bridge-driven receiver") for the full design, and
temp/NemoRL_Megatron_MX_Phase_C_Handoff.md for the receiver-side
implementation plan that this Phase A change unblocks.
…in v2 sidecar
Trainer-side companion to modelexpress commit 76ac523. At trainer init,
the Megatron policy worker now derives:
1. MegatronTransformerConfig (head counts + dims) from
self.megatron_bridge.transformer_config
2. Megatron→HF name map from a Bridge mapping-registry walk
(no weight transfer; just iterates the conversion-task list and
reads each task.mapping.hf_param)
…and stamps them onto the MxV2TrainingPublisher via
``set_megatron_sidecar(sidecar)``. The publisher then merges them into
the v2 sidecar JSON at every publish() call, so the receiver's
parse_megatron_sidecar (modelexpress.megatron_translator) gets:
* transformer_config — head counts + dims for QKV un-interleave
* hf_name_map — entries like
[
["decoder.layers.0.self_attention.linear_qkv.weight",
["q_proj.weight", "k_proj.weight", "v_proj.weight"]],
["decoder.layers.0.self_attention.linear_proj.weight",
["o_proj.weight"]],
...
]
The receiver doesn't need to import Megatron-Bridge to know either —
exactly what we want for the production Dynamo VllmDecodeWorker image
which does not ship Bridge.
QKVMapping ordering: Bridge's QKVMapping uses a dict-shaped hf_param
({"q": ..., "k": ..., "v": ...}). The sidecar emit forces the
canonical [q, k, v] order so the receiver-side translator's
hf_names[0/1/2] line up with split_qkv_weights' (q, k, v) return
tuple. Other dict-shaped mappings (rare today) emit dict.values()
order; receivers can override per-tensor if needed.
Failure mode: if Bridge's build_conversion_tasks raises (older Bridge,
unsupported model arch), warn and emit just the transformer_config
piece. The receiver then falls back to deriving names independently
(e.g. from the inference vLLM model's state-dict + plan.role).
…erExtension
Phase D of the Megatron-MX integration plan. Wires the new
modelexpress.megatron_translator path into the existing
update_weights_via_mx flow on NemoRL's direct-vLLM extension.
What's new
* First-cycle Megatron detection: peek at any candidate's
megatron_meta. None → existing DTensor / FSDP path. Set →
latched into self._mx_megatron_mode and routed through
_update_weights_via_mx_megatron forevermore on this receiver.
* _update_weights_via_mx_megatron: builds a MegatronReceiverContext
once (transformer_config + Megatron→HF name map from the source
sidecar; receive_specs from the per-tensor TensorDescriptorV2
registry); per refit, calls run_refit_cycle which drives the full
discover → plan → assemble → translate → yield(hf_name, tensor)
pipeline. The yielded HF tensors flow through the existing
_load_weights + FP8-KV-cache hooks.
* Tree fan-out (publish_self_as_source) is preserved on the
Megatron path identical to the DTensor path.
Backwards compat
* DTensor / FSDP receivers see _mx_megatron_mode = False on the
first cycle and stay on the existing path. PR NVIDIA-NeMo#2700's prime-rl
receivers and John's Dynamo-side worker extension are
unaffected.
* Sources that advertise publisher_kind=megatron but are missing
the transformer_config sidecar (older trainer image) trigger a
one-shot warning and fall back to non-Megatron mode.
Out of scope (deferred)
* The pull callback in _update_weights_via_mx_megatron uses
self._mx_receiver._receiver._nixl.pull_to, which is the
conceptual API; the actual NIXL plumbing for sliced pulls
(registering a sub-view of a parent tensor + completing the
transfer) needs a small wrapper in modelexpress.refit_receiver.
v0 of the actual NIXL sliced-pull plumbing lands in a follow-up
commit; this commit gets the Phase D control flow + receive
context build right so the next commit only needs to fill in
the pull function.
* Wiring into Dynamo's MxRefitWorkerExtension (a separate Dynamo
repo edit; same shape — import run_refit_cycle, call from the
worker extension's update_weights_via_mx).
… receiver
Pivot the Phase D receiver path from 'per-source pull callback' to
'bulk receive_weights + pre-filled buffers' since:
1. The existing MxRefitReceiver.receive_weights primitive already
does the matched-TP single-source NIXL pull correctly (same
primitive the DTensor path uses).
2. modelexpress.megatron_translator's new
pre_assembled_buffers arg lets run_refit_cycle skip
assemble_into_destination + the per-source pull callback.
Result: matched-TP Megatron-MX path uses ZERO new NIXL plumbing.
Per-rank Megatron buffers are pre-allocated + registered with the
existing register_tensors call, bulk-filled via the existing
receive_weights call, then walked by the translator.
Mixed-TP support (target_tp != source_tp) is gated behind 'no
matched-TP source for this receiver's tp_rank' and prints a clear
'v1 not yet wired' message — the planner already returns the right
slice info; only the per-source NIXL plumbing remains.
Two small fixes surfaced while validating Phase E shape 1 on a real Megatron-loaded Qwen3-4B-Thinking model on the kavin GB200 cluster. 1. _build_megatron_sidecar called bridge.build_conversion_tasks but the Bridge API is bridge.get_conversion_tasks. Rename. 2. model.named_parameters() returns 'module.<...>' prefixed names when the model is wrapped, but Bridge's get_conversion_tasks returns unprefixed names. The publisher now strips 'module.' in collect_megatron_publish_set so the catalog and the sidecar name_map use the same form. The receiver-side wire-up gets a defensive strip too in case an older publisher emits the prefix. Phase E shape 1 validation result with these fixes: 398/398 HF tensors byte-identical against bridge.export_hf_weights ground truth after a full A->B->C->D round trip; 313 Gbps single-NIC bulk pull. Writeup in pensieve/RL/NemoRL/14_megatron_mx_phase_e_shape1_2026_06_08.md.
Replaces the matched-TP gate in _update_weights_via_mx_megatron with a mixed-TP path that pulls each source's full manifest into scratch via receive_weights_scratch, then satisfies per-source slice requests by copying from scratch into the planner's pre-narrowed dest view. Handles MegatronSliceSource.source_subslice for target-narrower cases (vLLM TP > source TP, each receiver carves a sub-range out of one source's shard). For target-wider (vLLM TP < source TP, each receiver concatenates multiple sources' full shards), source_subslice is None and we copy the full source tensor into the corresponding dest range. Bandwidth: O(sum(source_bytes)) per receiver. For target-wider this is the minimum needed; for target-narrower it over-pulls by the target fan-out factor. A v1 sliced-pull primitive (offset+size NIXL descriptor on the source side) would cut narrow-target bandwidth by target_tp / source_tp. Matched-TP fast path is preserved via the is_matched_tp branch; only mixed-TP falls into the scratch-then-copy path. Both branches share the same translation step (assemble_into_destination + translate_megatron_to_hf).
Megatron-Core's TE-grouped linears (TEColumnParallelGroupedLinear,
TERowParallelGroupedLinear) expose each local expert's slice as a
separate nn.Parameter named weight0, weight1, ..., weightN, rather
than a single grouped buffer. This holds even at EP=1.
Two classifier fixes needed:
1. _enclosing_module previously only recognised weight/bias/scale/
_extra_state as parameter leaves; with grouped naming the
per-expert leaf is e.g. weight17. Extended the leaf detector to
match weight<digits>, bias<digits>, scale<digits>. Without this,
_enclosing_module walks past the parameter into the leaf and the
classifier sees class name "Parameter" → unclassified → falls to
ROLE_REPLICATED.
2. detect_megatron_role's per-expert path required ep_size > 1
because it assumed the legacy single-.weight-leading-axis grouped
layout. The TE-grouped layout has N parameters even at EP=1.
Added a 2a branch that fires whenever the param name matches the
expert pattern AND the leaf is weight<digits>: extracts the
expert id from the suffix, classifies as ROLE_EXPERT_COLUMN /
ROLE_EXPERT_ROW based on the enclosing module class, and emits
expert_id + expert_layout=grouped in descriptor_extras. The
legacy ep_size>1 leading-axis path is kept as 2b for back-compat.
Validation on Qwen3-30B-A3B-Instruct-2507 (real MoE: 48 layers, 128
experts) on the GB200 trainer:
pre-fix role distribution (broken — all experts replicated):
replicated 12529, row 48, qkv_column 48, vocab_parallel 1, column 1
post-fix role distribution (correct):
expert_column 6144 (= 48 layers x 128 experts of linear_fc1)
expert_row 6144 (= 48 layers x 128 experts of linear_fc2)
replicated 241 (norms, scalars, routers)
row 48 (linear_proj)
qkv_column 48 (linear_qkv)
vocab_parallel 1
column 1
Total classified: 12627 (matches Bridge's get_conversion_tasks count
and the publisher's named_parameters iteration count).
KavinKrishnan
force-pushed
the
kavink/megatron-mx
branch
from
June 9, 2026 20:26
89e2916 to
ee65362
Compare
github-actions
Bot
added
documentation
KavinKrishnan
changed the base branch from
dynamo-k8s-integration-mx
to
dynamo-k8s-integration
June 9, 2026 20:27
Bug surfaced on Qwen3-MoE-30B-A3B during Phase E shape 2 cluster validation: The fix in 9bdf5e4 stripped the `module.` prefix from `model.named_parameters()`'s name BEFORE passing it to `detect_megatron_role`, so the publisher and Bridge's name_map (which uses unprefixed names) would agree on the catalog key. That fix was correct for the catalog/name-map side, but it broke the classifier's model-walk: `_enclosing_module(name, model)` descends through attribute lookups, and for a DDP-wrapped model the descent starts at `model.module.decoder...`. With the prefix stripped, `getattr(model, 'decoder', None)` returns None on the first step, the walk fails, and every TP-sharded non-expert role (column, row, qkv_column, gated_mlp_column, vocab_parallel) falls through to ROLE_REPLICATED. Expert tensors didn't show the bug because the grouped-MoE 2a branch classifies via the trailing `weight<N>` suffix on the name string, not via the model walk — that branch fires before _enclosing_module. Cluster evidence (Qwen3-30B-A3B, kavin/GB200, pre-fix): pass1: detect_megatron_role(raw_name=`module.decoder...`) expert_column 6144 expert_row 6144 replicated 241 row 48 qkv_column 48 vocab_parallel 1 column 1 ← correct pass2: collect_megatron_publish_set (stripped name) expert_column 12288 (= 6144 fc1 + 6144 fc2 — both → "_column") replicated 339 (everything else collapsed) ← broken Fix: pass the raw (prefixed) name to detect_megatron_role so the model-walk succeeds; keep stripping the prefix for the published name (so the name_map lookup on the receiver still works). This restores the 7-role classification and unblocks Phase E shape 2. Validated locally: classifier now produces 7 distinct roles on Qwen3-30B-A3B (6144 expert_column + 6144 expert_row + 241 replicated + 48 row + 48 qkv_column + 1 vocab + 1 column = 12 627 total).
github-actions
Bot
removed
documentation
Author
|
Phase E shape 2 + mixed-TP cluster validation update (kavin / GB200, 2026-06-10). Two outstanding gaps from the prior session both closed; byte-identity proof against Phase E shape 2 — Qwen3-MoE-30B-A3B end-to-endFull A → B → C → D pipeline cross-pod:
This exercises every grouped-MoE code path on a production-shape model: 12,288 per-expert tensors split correctly via the new Mixed-TP cluster smoke — TP=2 → TP=1Synthetic 2-rank Megatron source → 1 receiver, exercising the v0 host-side scratch+slice path against real NIXL RDMA bytes. Target-narrower direction (
8 / 8 byte-identical across all 5 tp-sharded roles. Three bugs surfaced + fixed during validation
Full writeup + harness scripts: This closes Phase E end-to-end on the correctness side. The remaining performance follow-up is the v1 NIXL sliced-pull primitive for target-wider mixed-TP bandwidth optimisation; the v0 path is correctness-complete. |
Adds the v1 sliced-pull data plane to
VllmInternalWorkerExtension._update_weights_via_mx_megatron's mixed-TP
branch, in tandem with the new MxRefitReceiver.pull_to primitive
(MX-side commit).
Restructured the mixed-TP path into five phases:
1. Pre-allocate per-plan dest tensors (target_shape) and classify
each (plan, source) contribution: contiguous narrows along axis 0
route to v1 sliced pull; non-contiguous (axis-1 row-parallel) or
per_expert (dict assembly) route to v0 scratch+copy.
2. NIXL-register the v1 dest buffers (one register_tensors call).
3. Per-source combined sliced pull: each source's batch of
SlicedTransferRequests issued as a single pull_to() call → one
NIXL transfer with N descriptor pairs.
4. v0 fallback: for the plans that need scratch (row-parallel, per
expert), bulk-pull each contributing source via
receive_weights_scratch into a per-source scratch dict.
5. Assemble + translate per plan: v1 plans use the pre-filled dest
directly (no host-side assembly); v0 plans use
assemble_into_destination with the existing scratch-based pull
callback. Both paths feed translate_megatron_to_hf identically.
Bandwidth profile change vs v0:
* Target-narrower (source_tp > target_tp): unchanged — already
bandwidth-optimal in v0 (validated 8/8 byte-identical on
cluster, 2026-06-10).
* Target-wider (source_tp < target_tp): wire bytes drop by
target_tp/source_tp× for axis-0 roles. Cluster validation needs
a vllm tensor_parallel_size>1 DGD worker; deferred to a separate
deployment-only step.
Backwards-compatibility: matched-TP fast path (is_matched_tp branch)
unchanged. Existing 18 867/18 867 byte-identical validation on
Qwen3-MoE-30B-A3B continues to hold (matched-TP path doesn't touch
this code).
Smallest viable Megatron + MX GRPO recipe — Qwen2.5-1.5B Megatron
TP=1 EP=1 PP=1 with vLLM generation, weight_sync.method="mx" pointing
at the cluster-validated Megatron-MX data plane:
trainer: stream_weights_via_mx (per-rank native shards + sidecar)
refit: MX server catalogs the source, ready=True
receiver: _update_weights_via_mx_megatron
-> v1 sliced-pull (axis-0 roles) / v0 scratch+copy (axis-1)
-> translate Megatron->HF (vendored Bridge helpers)
-> vllm.model.load_weights()
5 GRPO steps with val_period=5 gives one refit cycle + one validation
in ~10-20 min wall clock. Bump max_num_steps after first run lands.
End-to-end byte-identity of the Megatron path is independently
validated on:
* Qwen3-4B-Thinking-2507 (dense): 398/398 HF tensors byte-identical
* Qwen3-MoE-30B-A3B-Instruct-2507: 18 867/18 867 byte-identical
(12 288 grouped per-expert tensors via the new split_gated_mlp_tp)
* Synthetic 2-rank Megatron TP=2 mixed-TP: 8/8 byte-identical
See pensieve/RL/NemoRL/17_megatron_mx_grpo_loop_runbook.md for the
deployment runbook + expected log lines + bumping guide.
Recipe is fully wired; actual GRPO loop run pending cluster session
restoration.
Author
|
Step 4 + step 5 update (2026-06-10): v1 sliced-pull primitive + GRPO recipe landed; cluster validation queued for next K8s session. Step 4 — v1 sliced-pull NIXL primitiveAdds the bandwidth-optimal mixed-TP data plane. Two new APIs: # Lower-level (modelexpress.nixl_transfer)
NixlTransferManager.receive_sliced_from_source(
source_metadata, source_tensors,
slice_requests=[SlicedTransferRequest(name, source_offset_bytes,
slice_bytes, dest_view), ...],
)
# Higher-level (modelexpress.refit_receiver)
MxRefitReceiver.pull_to(
source: SourceRef,
requests=[(name, source_subslice (lo, hi) | None, dest_view), ...],
)Each request maps to one (source-slice, dest-view) descriptor pair in a single combined NIXL transfer. Caller passes a list of N requests → one transfer with N descriptors → one wait. The dest_view must be contiguous (axis-0 narrows are; axis-1 row-parallel narrows aren't, fall back to v0 scratch+copy for those). Receiver wiring in
Bandwidth profile change:
8 new unit tests in
60 / 60 unit tests pass across the full MX suite (planner + translator + helpers + sliced-pull + publish-self-as-source). Bridge byte-identity cross-checks skipped locally; pass on trainer image. Cluster validation for the target-wider direction needs a DGD with Step 5 — Megatron-MX GRPO recipe + runbookSmallest viable Megatron + MX GRPO recipe at defaults: "grpo_math_1B_megatron.yaml" # Qwen2.5-1.5B Megatron TP=1
grpo:
max_num_steps: 5
policy:
megatron_cfg.enabled: true
dtensor_cfg.enabled: false
cluster:
weight_sync:
method: "mx" # routes through Megatron-MX
mx_config:
enabled: true
mx_server_url: "modelexpress-server.kavin.svc.cluster.local:8001"
same_rank_only: true
tree_scale_out: true
nixl_retry_budget: 3GRPO loop was already wired to call Companion runbook documents the prerequisites, image-overlay shortcut, run command, expected log lines (matched-TP + mixed-TP v1), known gotchas, and bumping guide: Final scoreboard
Correctness side is done; remaining items are deployment / cluster-time exercises with no new code expected. |
Companion to the matched-TP cluster smokes on Qwen3-4B-Thinking-2507 (shape 1, 398/398 byte-identical, 2026-06-08). Inherits from the existing grpo_workplace_assistant_dynamo_mx.yaml (DTensor + MX recipe) and flips back to Megatron: policy.megatron_cfg.enabled: true policy.dtensor_cfg.enabled: false policy.optimizer: null (Megatron reads megatron_cfg.optimizer) policy.scheduler: null cluster.weight_sync.method: "mx" + mx_config points at the kavin MX server. 2 GRPO steps for the first cycle. Requires the DGD's dynamo/vllm/mx_refit/extension.py to be patched with a Megatron dispatch (calls modelexpress.megatron_translator. run_refit_cycle for Megatron sources). The patch is staged in pensieve/RL/NemoRL/v1_targetwider_smoke_2026_06_11/dynamo_vllm_mx_ refit_extension_megatron.py as a drop-in overlay; the natural follow-up is an upstream PR to ai-dynamo/dynamo. Companion writeup: pensieve/RL/NemoRL/18_v1_targetwider_validated_ dynamo_megatron_patch_2026_06_11.md.
Author
|
2026-06-11 update: v1 sliced-pull target-wider cluster-validated + Dynamo extension Megatron patch ready. Both items pending from yesterday closed. Step 2b — v1 target-wider cluster smokeSynthetic source_tp=1 → target_tp=2 (each "target rank" pulls only half of each tp-sharded source tensor via First wire-level proof of the bandwidth-optimal direction. The lower-level Step 5 — Dynamo extension Megatron patchThe DGD worker's
~250 LOC of isolated change; DTensor path untouched. Cluster-validated correctness across all Megatron paths means the dispatch produces correct bytes whenever it's exercised. Staged as a drop-in overlay at Companion recipe Final correctness scoreboard
Remaining for actual GRPO loop runThree cluster-engineering items, no code work:
All receiver-side code paths the GRPO loop would exercise are already byte-identity-proven via the standalone smokes. The GRPO run is the operational escalation, not a code one. Full writeup: |
jthomson04
merged commit Jun 12, 2026
1a2d3e7
into
jthomson04:dynamo-k8s-integration
4 checks passed
Author
|
Perf fix: cache NIXL-registered dest buffers across refit cycles (cluster-validated on GB200, 2026-06-23). Audit of refit-cycle wall time on multi-receiver setups surfaced a receiver-side bug: the Megatron MX path re-allocated + re-registered NIXL buffers on every refit cycle, paying ~0.15 s of Plan shapes are deterministic for a fixed What changedThree sites in two files:
The Dynamo extension overlay is the upstream-PR-ready file. Drops directly into Cluster validation — 3 back-to-back refit cycles on Qwen3-4B-Thinking-2507Harness: a standalone multi-cycle benchmark. Exercises the same
Expected impact on larger multi-receiver setups
The fix is necessary but not sufficient for the broader multi-receiver perf story. Other items still on the list:
Full writeup including the diagnostic + roadmap lives in the team's internal pensieve. |
This was referenced Jun 23, 2026
KavinKrishnan
added a commit
to ai-dynamo/modelexpress
that referenced
this pull request
Jun 26, 2026
…timizations
Lands the v2 client surface and the surrounding RL workstream needed
to unblock the downstream Megatron-MX and perf PRs.
## What this contributes
### v2 client surface (this PR's primary deliverable)
* `MxV2TrainingPublisher` + `MxV2RefitReceiver` (modelexpress.nemo_rl_v2)
-- the fat-client surface for per-rank shard publish and
receiver-side multi-source assembly.
* `MxWeightTransferEngine` (modelexpress.vllm_weight_transfer) -- an
adapter implementing vLLM's upstream WeightTransferEngine ABC.
* `TensorDescriptorV2.extra_parameters` (map<string,string>) plus
`SourceIdentity.revision` (string) -- the two proto extensions that
carry all per-tensor + per-source RL metadata.
### Phase 3a -- compile_target + compile_metadata on TensorDescriptorV2
New per-tensor fields default to ``hf_raw`` / ``{}``; the wire encoder
omits them when default so existing payloads stay byte-identical.
New constants: ``COMPILE_TARGET_HF_RAW``, ``_VLLM_FUSED``,
``_DEEPGEMM_FP8``, ``_CUTLASS_FP8``, ``_TRTLLM``. New helper
``compile_target_matches(descriptor, *, allowed_targets,
required_metadata=None)`` for receiver-side filtering with whitelist
plus required-metadata-subset semantics.
### Phase 3b -- compile_target_filter on discover_v2_sources
New kwargs ``compile_target_filter`` (whitelist set) and
``required_compile_metadata`` (subset-of-every-tensor's
compile_metadata). Candidates with no v2 registry are rejected when
either filter is set; candidates with mixed compile targets are
rejected if any tensor falls outside the allowed set.
``V2SourceCandidate.compile_targets: frozenset[str]`` exposed for
caller introspection.
### Phase 4 -- Multi-source slice discovery for mixed trainer/inference TP
New types: ``TargetTPLayout``, ``SliceSource``, ``SliceCoveragePlan``.
New methods ``MxV2RefitReceiver.discover_v2_sources_for_slice`` and
``MxV2RefitReceiver.receive_via_plan`` -- planner walks v2 candidates
per tensor, intersects each publisher's local_shard_range against the
receiver's requested slice, emits the minimal candidate set covering
it; surfaces coverage gaps and shard_axis mismatches in plan.missing.
``receive_via_plan`` orchestrates per-candidate scratch RDMA pulls
and stitches via torch.cat along the shard axis.
### Proto extensions
* ``TensorDescriptorV2.extra_parameters: map<string, string>`` -- the
escape hatch downstream RL clients use for per-tensor metadata
(megatron_role, compile_target, expert_id, revision, training_step,
...). Heavily used by #429.
* ``SourceIdentity.revision: string`` -- content-addressed weight
version. Non-empty value guarantees two sources with identical
SourceIdentity have bit-identical weight bytes, enabling
decentralized modes (no central coordinator) to use mx_source_id
as a full content check.
* Redis backend ``SourceAttributesJson`` carries the new fields plus
the union of main's compatibility additions
(backend_framework_version, torch_version, cuda_version,
triton_version, gpu_arch, compile_config_digest).
## Tests
68/68 unit tests green on this branch's surface (test_v2_source_picker,
test_v2_shape_registry, test_source_id, test_types, test_vllm_adapter).
## Compatibility
Every new field has a backwards-compatible default and every new method
arg is optional. Existing demos and downstream consumers run unchanged.
## Downstream impact
This PR is the prerequisite for several PRs in flight:
* #421 -- publish_self_as_source tree fan-out fix
* #429 -- Megatron-Core MX clients
* #450 -- MX_RDMA_NIC_PIN=stripe multi-NIC mode
* ai-dynamo/dynamo#10900 -- first-time upstream port of
dynamo.vllm.mx_refit extension
* ai-dynamo/dynamo#10901 -- buffer-caching perf fix in Dynamo extension
* jthomson04/RL#2 (merged) + #7 -- NeMo RL Megatron-MX integration plus
perf fix
Once this PR lands, kavink/nemo_rl_moe (the integration branch holding
Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
#421 and #429) can be retired by rebasing those PRs onto main.
KavinKrishnan
added a commit
to ai-dynamo/modelexpress
that referenced
this pull request
Jun 26, 2026
…timizations
Lands the v2 client surface and the surrounding RL workstream needed
to unblock the downstream Megatron-MX and perf PRs.
## What this contributes
### v2 client surface (this PR's primary deliverable)
* `MxV2TrainingPublisher` + `MxV2RefitReceiver` (modelexpress.nemo_rl_v2)
-- the fat-client surface for per-rank shard publish and
receiver-side multi-source assembly.
* `MxWeightTransferEngine` (modelexpress.vllm_weight_transfer) -- an
adapter implementing vLLM's upstream WeightTransferEngine ABC.
* `TensorDescriptorV2.extra_parameters` (map<string,string>) plus
`SourceIdentity.revision` (string) -- the two proto extensions that
carry all per-tensor + per-source RL metadata.
### Phase 3a -- compile_target + compile_metadata on TensorDescriptorV2
New per-tensor fields default to ``hf_raw`` / ``{}``; the wire encoder
omits them when default so existing payloads stay byte-identical.
New constants: ``COMPILE_TARGET_HF_RAW``, ``_VLLM_FUSED``,
``_DEEPGEMM_FP8``, ``_CUTLASS_FP8``, ``_TRTLLM``. New helper
``compile_target_matches(descriptor, *, allowed_targets,
required_metadata=None)`` for receiver-side filtering with whitelist
plus required-metadata-subset semantics.
### Phase 3b -- compile_target_filter on discover_v2_sources
New kwargs ``compile_target_filter`` (whitelist set) and
``required_compile_metadata`` (subset-of-every-tensor's
compile_metadata). Candidates with no v2 registry are rejected when
either filter is set; candidates with mixed compile targets are
rejected if any tensor falls outside the allowed set.
``V2SourceCandidate.compile_targets: frozenset[str]`` exposed for
caller introspection.
### Phase 4 -- Multi-source slice discovery for mixed trainer/inference TP
New types: ``TargetTPLayout``, ``SliceSource``, ``SliceCoveragePlan``.
New methods ``MxV2RefitReceiver.discover_v2_sources_for_slice`` and
``MxV2RefitReceiver.receive_via_plan`` -- planner walks v2 candidates
per tensor, intersects each publisher's local_shard_range against the
receiver's requested slice, emits the minimal candidate set covering
it; surfaces coverage gaps and shard_axis mismatches in plan.missing.
``receive_via_plan`` orchestrates per-candidate scratch RDMA pulls
and stitches via torch.cat along the shard axis.
### Proto extensions
* ``TensorDescriptorV2.extra_parameters: map<string, string>`` -- the
escape hatch downstream RL clients use for per-tensor metadata
(megatron_role, compile_target, expert_id, revision, training_step,
...). Heavily used by #429.
* ``SourceIdentity.revision: string`` -- content-addressed weight
version. Non-empty value guarantees two sources with identical
SourceIdentity have bit-identical weight bytes, enabling
decentralized modes (no central coordinator) to use mx_source_id
as a full content check.
* Redis backend ``SourceAttributesJson`` carries the new fields plus
the union of main's compatibility additions
(backend_framework_version, torch_version, cuda_version,
triton_version, gpu_arch, compile_config_digest).
## Tests
68/68 unit tests green on this branch's surface (test_v2_source_picker,
test_v2_shape_registry, test_source_id, test_types, test_vllm_adapter).
## Compatibility
Every new field has a backwards-compatible default and every new method
arg is optional. Existing demos and downstream consumers run unchanged.
## Downstream impact
This PR is the prerequisite for several PRs in flight:
* #421 -- publish_self_as_source tree fan-out fix
* #429 -- Megatron-Core MX clients
* #450 -- MX_RDMA_NIC_PIN=stripe multi-NIC mode
* ai-dynamo/dynamo#10900 -- first-time upstream port of
dynamo.vllm.mx_refit extension
* ai-dynamo/dynamo#10901 -- buffer-caching perf fix in Dynamo extension
* jthomson04/RL#2 (merged) + #7 -- NeMo RL Megatron-MX integration plus
perf fix
Once this PR lands, kavink/nemo_rl_moe (the integration branch holding
Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
#421 and #429) can be retired by rebasing those PRs onto main.
KavinKrishnan
added a commit
to ai-dynamo/modelexpress
that referenced
this pull request
Jun 26, 2026
…timizations
Lands the v2 client surface and the surrounding RL workstream needed
to unblock the downstream Megatron-MX and perf PRs.
## What this contributes
### v2 client surface (this PR's primary deliverable)
* `MxV2TrainingPublisher` + `MxV2RefitReceiver` (modelexpress.nemo_rl_v2)
-- the fat-client surface for per-rank shard publish and
receiver-side multi-source assembly.
* `MxWeightTransferEngine` (modelexpress.vllm_weight_transfer) -- an
adapter implementing vLLM's upstream WeightTransferEngine ABC.
* `TensorDescriptorV2.extra_parameters` (map<string,string>) plus
`SourceIdentity.revision` (string) -- the two proto extensions that
carry all per-tensor + per-source RL metadata.
### Phase 3a -- compile_target + compile_metadata on TensorDescriptorV2
New per-tensor fields default to ``hf_raw`` / ``{}``; the wire encoder
omits them when default so existing payloads stay byte-identical.
New constants: ``COMPILE_TARGET_HF_RAW``, ``_VLLM_FUSED``,
``_DEEPGEMM_FP8``, ``_CUTLASS_FP8``, ``_TRTLLM``. New helper
``compile_target_matches(descriptor, *, allowed_targets,
required_metadata=None)`` for receiver-side filtering with whitelist
plus required-metadata-subset semantics.
### Phase 3b -- compile_target_filter on discover_v2_sources
New kwargs ``compile_target_filter`` (whitelist set) and
``required_compile_metadata`` (subset-of-every-tensor's
compile_metadata). Candidates with no v2 registry are rejected when
either filter is set; candidates with mixed compile targets are
rejected if any tensor falls outside the allowed set.
``V2SourceCandidate.compile_targets: frozenset[str]`` exposed for
caller introspection.
### Phase 4 -- Multi-source slice discovery for mixed trainer/inference TP
New types: ``TargetTPLayout``, ``SliceSource``, ``SliceCoveragePlan``.
New methods ``MxV2RefitReceiver.discover_v2_sources_for_slice`` and
``MxV2RefitReceiver.receive_via_plan`` -- planner walks v2 candidates
per tensor, intersects each publisher's local_shard_range against the
receiver's requested slice, emits the minimal candidate set covering
it; surfaces coverage gaps and shard_axis mismatches in plan.missing.
``receive_via_plan`` orchestrates per-candidate scratch RDMA pulls
and stitches via torch.cat along the shard axis.
### Proto extensions
* ``TensorDescriptorV2.extra_parameters: map<string, string>`` -- the
escape hatch downstream RL clients use for per-tensor metadata
(megatron_role, compile_target, expert_id, revision, training_step,
...). Heavily used by #429.
* ``SourceIdentity.revision: string`` -- content-addressed weight
version. Non-empty value guarantees two sources with identical
SourceIdentity have bit-identical weight bytes, enabling
decentralized modes (no central coordinator) to use mx_source_id
as a full content check.
* Redis backend ``SourceAttributesJson`` carries the new fields plus
the union of main's compatibility additions
(backend_framework_version, torch_version, cuda_version,
triton_version, gpu_arch, compile_config_digest).
## Tests
68/68 unit tests green on this branch's surface (test_v2_source_picker,
test_v2_shape_registry, test_source_id, test_types, test_vllm_adapter).
## Compatibility
Every new field has a backwards-compatible default and every new method
arg is optional. Existing demos and downstream consumers run unchanged.
## Downstream impact
This PR is the prerequisite for several PRs in flight:
* #421 -- publish_self_as_source tree fan-out fix
* #429 -- Megatron-Core MX clients
* #450 -- MX_RDMA_NIC_PIN=stripe multi-NIC mode
* ai-dynamo/dynamo#10900 -- first-time upstream port of
dynamo.vllm.mx_refit extension
* ai-dynamo/dynamo#10901 -- buffer-caching perf fix in Dynamo extension
* jthomson04/RL#2 (merged) + #7 -- NeMo RL Megatron-MX integration plus
perf fix
Once this PR lands, kavink/nemo_rl_moe (the integration branch holding
Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
#421 and #429) can be retired by rebasing those PRs onto main.
KavinKrishnan
added a commit
to ai-dynamo/modelexpress
that referenced
this pull request
Jun 26, 2026
…timizations
Lands the v2 client surface and the surrounding RL workstream needed
to unblock the downstream Megatron-MX and perf PRs.
## What this contributes
### v2 client surface (this PR's primary deliverable)
* `MxV2TrainingPublisher` + `MxV2RefitReceiver` (modelexpress.nemo_rl_v2)
-- the fat-client surface for per-rank shard publish and
receiver-side multi-source assembly.
* `MxWeightTransferEngine` (modelexpress.vllm_weight_transfer) -- an
adapter implementing vLLM's upstream WeightTransferEngine ABC.
* `TensorDescriptorV2.extra_parameters` (map<string,string>) plus
`SourceIdentity.revision` (string) -- the two proto extensions that
carry all per-tensor + per-source RL metadata.
### Phase 3a -- compile_target + compile_metadata on TensorDescriptorV2
New per-tensor fields default to ``hf_raw`` / ``{}``; the wire encoder
omits them when default so existing payloads stay byte-identical.
New constants: ``COMPILE_TARGET_HF_RAW``, ``_VLLM_FUSED``,
``_DEEPGEMM_FP8``, ``_CUTLASS_FP8``, ``_TRTLLM``. New helper
``compile_target_matches(descriptor, *, allowed_targets,
required_metadata=None)`` for receiver-side filtering with whitelist
plus required-metadata-subset semantics.
### Phase 3b -- compile_target_filter on discover_v2_sources
New kwargs ``compile_target_filter`` (whitelist set) and
``required_compile_metadata`` (subset-of-every-tensor's
compile_metadata). Candidates with no v2 registry are rejected when
either filter is set; candidates with mixed compile targets are
rejected if any tensor falls outside the allowed set.
``V2SourceCandidate.compile_targets: frozenset[str]`` exposed for
caller introspection.
### Phase 4 -- Multi-source slice discovery for mixed trainer/inference TP
New types: ``TargetTPLayout``, ``SliceSource``, ``SliceCoveragePlan``.
New methods ``MxV2RefitReceiver.discover_v2_sources_for_slice`` and
``MxV2RefitReceiver.receive_via_plan`` -- planner walks v2 candidates
per tensor, intersects each publisher's local_shard_range against the
receiver's requested slice, emits the minimal candidate set covering
it; surfaces coverage gaps and shard_axis mismatches in plan.missing.
``receive_via_plan`` orchestrates per-candidate scratch RDMA pulls
and stitches via torch.cat along the shard axis.
### Proto extensions
* ``TensorDescriptorV2.extra_parameters: map<string, string>`` -- the
escape hatch downstream RL clients use for per-tensor metadata
(megatron_role, compile_target, expert_id, revision, training_step,
...). Heavily used by #429.
* ``SourceIdentity.revision: string`` -- content-addressed weight
version. Non-empty value guarantees two sources with identical
SourceIdentity have bit-identical weight bytes, enabling
decentralized modes (no central coordinator) to use mx_source_id
as a full content check.
* Redis backend ``SourceAttributesJson`` carries the new fields plus
the union of main's compatibility additions
(backend_framework_version, torch_version, cuda_version,
triton_version, gpu_arch, compile_config_digest).
## Tests
68/68 unit tests green on this branch's surface (test_v2_source_picker,
test_v2_shape_registry, test_source_id, test_types, test_vllm_adapter).
## Compatibility
Every new field has a backwards-compatible default and every new method
arg is optional. Existing demos and downstream consumers run unchanged.
## Downstream impact
This PR is the prerequisite for several PRs in flight:
* #421 -- publish_self_as_source tree fan-out fix
* #429 -- Megatron-Core MX clients
* #450 -- MX_RDMA_NIC_PIN=stripe multi-NIC mode
* ai-dynamo/dynamo#10900 -- first-time upstream port of
dynamo.vllm.mx_refit extension
* ai-dynamo/dynamo#10901 -- buffer-caching perf fix in Dynamo extension
* jthomson04/RL#2 (merged) + #7 -- NeMo RL Megatron-MX integration plus
perf fix
Once this PR lands, kavink/nemo_rl_moe (the integration branch holding
Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
#421 and #429) can be retired by rebasing those PRs onto main.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds the Megatron-Core path for NemoRL's
update_weights_via_mxweight-refit, layered on top ofdynamo-k8s-integration-mx. Companion toai-dynamo/modelexpress#421(which carries the MX-side slice planner + receiver translator + vendored Bridge helpers).The architecture is receiver-driven multi-source slicing: each Megatron TP/PP/EP rank publishes its native shard with no allgather; the inference receiver discovers a per-tensor slice plan via the MX catalog, pulls slices in parallel from the owning ranks, and reassembles locally using vendored Megatron-Bridge-style transforms (QKV un-interleave, gated-MLP split, column/row-parallel concat, per-expert split). It generalises the FSDP+EP planner from
ai-dynamo/modelexpress#349's Phase 4 to cover Megatron's column-parallel, row-parallel, and head-interleaved fused-QKV layouts.Full design:
temp/NemoRL_Megatron_MX_Design.md(in the kavink/RL doc tree on the MX repo).Commits
69027ad30MegatronPolicyWorker.stream_weights_via_mx+detect_megatron_roleclassifier in newnemo_rl/distributed/mx_megatron_helpers.pyac014c8b2AutoMapping._MODULE_TYPE_REGISTRY(catches every TE/Inference/Quant variant)dd13425ceMegatronTransformerConfig+ Megatron→HF name map in v2 sidecar via Bridge introspection (no weight transfer)648855bdd_update_weights_via_mx_megatrononVllmInternalWorkerExtensionwith first-cycle Megatron-mode latching790e0646breceive_weights+pre_assembled_buffersfor matched-TPPath validated this session
nemorl-mx-trainer-gpu-workers-worker-zzjbdin kavin namespace, /opt/ray_venvs/.../MegatronPolicyWorker/bin/python). All 7 architectural assumptions confirmed:mpu.is_initialized()=Falsefallback works, helpers byte-identity round-trip, AutoMapping registry accessible, AutoBridge importable.What's NOT in this PR (deferred)
MegatronSlicePlanwithsource_subsliceset; the gap is just the per-sourceMxRefitReceiver.pull_to(src, dest, source_subslice=...)primitive and a refactor of the receiver's pre-allocate-buffers step. The matched-TP code path prints an explicit[mx-megatron] ... v1 mixed-TP NIXL plumbing not yet wiredmessage when no source's tp_rank matches the receiver's, and falls through. ~1 day of focused work + a Megatron model in the test loop.attention_output_gate=True, FP8 blockwise scale tensors, KV-only fused linear, Mamba/GDN layouts. Vendored helpers explicitly raiseNotImplementedErrorfor FP8; other variants would need a vendoring extension.MxRefitWorkerExtension(a separate Dynamo repo edit, ~5-line shim that importsrun_refit_cycle).Test plan
python -m pytest tests/test_v2_megatron_slice_planner.py tests/test_megatron_helpers.py tests/test_megatron_translator.py— 36 tests, all pass (in the MX repo)References
ai-dynamo/modelexpress#421NemoRL_Megatron_MX_Design.mdNemoRL_Megatron_MX_Phase_C_Handoff.mdpensieve/RL/NemoRL/10_gb200_smoke_results_2026_06_03.md