Skip to content

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

Draft
KavinKrishnan wants to merge 30 commits into
mainfrom
kavink/megatron-mx-main
Draft

feat(rl): NeMo-RL × Megatron × ModelExpress v2 weight-refit integration (umbrella)#482
KavinKrishnan wants to merge 30 commits into
mainfrom
kavink/megatron-mx-main

Conversation

@KavinKrishnan

@KavinKrishnan KavinKrishnan commented Jul 5, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Umbrella integration branch for the ModelExpress v2 RL weight-refit path used by NeMo-RL × Megatron on vLLM.

What's in it

  • Persistent scratch, metadata-safe agent reuse, stale-generation pruning, wire pruning, and coverage guards.
  • Native vLLM MxWeightTransferEngine, spawned-TP registration, and Dynamo load aliases.
  • Merged EP Bridge sidecars, local→global expert recovery, non-expert dedup, and live EP filtering.
  • TP-aware MDL with rank-local direct/fused/expert writes.
  • Multi-rail NIXL/RDMA, DMA-BUF, slice planning, and fan-out substrate.

Live validation

  • Fully baked GB200 EP8→TP2 three-step GRPO: KL 0.0399–0.0643.
  • Steady MX+MDL refit: 10.25–11.94s.
  • TP2+EP2: 47.5% fewer bytes, 64 experts/rank, KL 0.0586.
  • Native PyNccl static baseline: 5.17s packed update for 61.06 GB/rank (94.5–96.6 Gbps/rank); sender HBM preload excluded.

For this fixed full-model topology NCCL is ~2× faster end-to-end. MX is currently dominated by receiver translation/staging and load; its system advantages are filtering, heterogeneous layouts, elasticity, straggler isolation, and fan-out.

Pairs with NVIDIA-NeMo/RL#3068 and the Dynamo native-engine wiring branch.

@copy-pr-bot

copy-pr-bot Bot commented Jul 5, 2026

Copy link
Copy Markdown

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.

This was referenced Jul 6, 2026
Closed
Closed
Closed
Closed
@KavinKrishnan
KavinKrishnan force-pushed the kavink/megatron-mx-main branch from 8338f23 to 25d1169 Compare July 6, 2026 04:29
@KavinKrishnan KavinKrishnan changed the title NeMo-RL × Megatron × ModelExpress v2 refit integration (WIP — needs rebase onto main) feat(rl): NeMo-RL × Megatron × ModelExpress v2 weight-refit integration (umbrella) Jul 6, 2026
@github-actions github-actions Bot added the feat label Jul 6, 2026
@KavinKrishnan
KavinKrishnan force-pushed the kavink/megatron-mx-main branch from 25d1169 to 87e3c23 Compare July 6, 2026 04:36
@KavinKrishnan
KavinKrishnan force-pushed the kavink/megatron-mx-main branch from 87e3c23 to d7c199d Compare July 6, 2026 06:20
@KavinKrishnan
feat(rl): NeMo-RL × Megatron × ModelExpress v2 weight-refit integration
d19bcb4 
Squashed integration of the MX v2 RL weight-refit path (previously 53 commits;
see closed PRs #349/#465/#450/#470 for the incremental history).

Client (modelexpress_client/python/modelexpress):
- nemo_rl_v2.py: MxV2TrainingPublisher / MxV2RefitReceiver, multi-source
  pick_megatron_slice_plans, tree fan-out (publish_self_as_source replica_uid
  unique-sid + discover_v2_sources prefer_replicas).
- megatron_translator.py: Megatron→HF translate (qkv un-interleave, gated-MLP
  split, per-expert unstack).
- nixl_transfer.py: sliced pull (receive_sliced_from_source), host/device
  memtype for pinned-CPU staging, rebind_tensors buffer caching; reconciled
  with the accelerator-backend abstraction on main.
- ucx_utils.py: MX_RDMA_NIC_PIN=stripe multi-NIC parallelism.
- rl_expert_layout.py: MoE expert-parallel wire filtering (compute_local_expert_ids).
- vmm/: VMM-arena single-region registration.
- engines/vllm/weight_transfer.py: MxWeightTransferEngine — native vLLM
  WeightTransferEngine backend registered as "mx".
- engines/vllm/mdl.py: MDL (Mapped Direct Load) decoupled load-side fast-loader.

Server (modelexpress_server/src/p2p) + proto:
- SourceIdentity.extra_parameters round-trip + revision, additive with main's
  version/arch fields.

Docs: MX-RL design + integration/benchmark context under docs/RL, docs/slides.
Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
KavinKrishnan force-pushed the kavink/megatron-mx-main branch from d7c199d to d19bcb4 Compare July 6, 2026 06:30
@KavinKrishnan
feat(engines/vllm): MDL incremental dest-map for partial/subset/delta…
00eb94f 
… updates

Builds the MDL destination map incrementally — each tensor is classified the
first time it's seen (in whatever update it arrives), rather than only in a full
cold cycle. So update_weights(subset=...) (by layer/layer-group/param list),
layerwise reload, or delta updates get the warm in-place fast path instead of
permanently falling back to the stock loader for tensors absent from cycle 1.
Refactors _build_dest_map -> additive _extend_dest_map + a _try_write helper.
Supports the vLLM-layer update(subset) API being aligned with the NIXL team.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
refactor(engines/vllm): factor vLLM weight-update logic into a tier-2…
02d01c8 
… layer

Splits the vLLM refit path into explicit tiers so the WeightTransferEngine
backend is a thin adapter and the generic receiver stays vLLM-free:

- NEW weight_update.py (tier 2): MxVllmWeightUpdater owns the vLLM-specific
  logic — geometry discovery, receive/translate (Megatron + DTensor), buffer
  alloc/registration (host/device/arena), EP filter, byte-verify — behind the
  base lifecycle initialize_weight_update_setup / was_weight_update_setup_
  initialized / start_weight_update / update_weights(subset) / finish_weight_
  update. Info dataclasses (MxInitInfo/MxUpdateInfo, subclassing vLLM's ABC
  bases) + WeightSubset live here. Adds param-name subset scoping.
- weight_transfer.py (tier 3): MxWeightTransferEngine reduces to an adapter
  mapping the ABC's init/receive/shutdown onto the tier-2 lifecycle; re-exports
  the dataclasses for back-compat.
- MxV2RefitReceiver (tier 1) unchanged — stays a generic receiver.

Matches the layering agreed with the NIXL team; a framework can drive tier 2
directly via RPC instead of the ABC. Validated on a live vLLM 0.20.1 pod:
'mx' auto-registers, lifecycle present, shim resolves to the same class.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
feat(engines/vllm): wire layer / layer-group subset selection in upda…
9776b5b 
…te_weights

WeightSubset now supports three selectors (a tensor is kept if it matches ANY):
param_names (exact), layers (transformer indices -> `.layers.<n>.`), and
layer_groups (name fragments, e.g. self_attn / mlp / experts). Replaces the
NotImplementedError stubs on both the Megatron (_apply_subset over receive
specs) and DTensor (scratch-yield filter) paths. Completes the update_weights
(subset) surface aligned with the NIXL-team API discussion. Matching semantics
verified (param/layers/groups/multi-name/empty).

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
refactor(rl): remove deprecated expert shim, planner is strictly rang…
6fb7870 
…e-only

SliceOwnership/SliceRequest drop is_expert/expert_axis/owned_expert_ids/
required_experts; the reshard planner drops the expert special-case. An MoE
expert shard is a plain SHARD range on the expert axis (via
expert_ids_to_contiguous_ranges); ordinary range intersection yields the same
EP coverage (proven equivalent). rank_local_publisher drops the dead expert
params (verl used defaults). Tests: delete superseded ep test, convert ep_ranges
to the canonical ranges-only coverage test, rewrite expert-layout round-trip.
@KavinKrishnan
feat(mdl): loaderless cold path for quant models whose load_weights i…
8230723 
…sn't re-entrant

Block-fp8 MoE replaces params in process_weights_after_loading, so vLLM's
load_weights drops output_dim and a 2nd load raises. When it fails, MDL builds
the map + applies weights via its own in-place writes (needs no output_dim),
remaps checkpoint->param names via the model's hf_to_vllm_mapper, and fails loud
if <90% of params map (won't serve a half-written model). bf16 path unchanged
(loaderless only triggers on load_weights failure).
@KavinKrishnan
vllm weight_update: EP-gather path for EP-trainer -> non-EP/lower-EP …
97edc32 
…rollout

_receive_megatron branches to a new _receive_megatron_ep_gather when sources
span >1 ep_rank: process each EP source with its expert HF names remapped
experts.<local> -> experts.<global> (publisher advertises the global id in
extras), dedupe stale sources by ep_rank (newest), apply the EP filter to keep
only the rollout rank's wanted global experts, and gather. The single-EP /
TP-assembly path is unchanged. Algorithm validated byte-identical at EP4/EP8 in
the first-party harness.
@KavinKrishnan
fix(vllm-wt): complete native EP source translation
22c3c93 
Merge rank-local Bridge name maps, recover global expert identities from legacy publishers, order checkpoint tensors by layer, and register the mx backend deterministically in spawned TP workers.
@KavinKrishnan
test(vllm-wt): fail closed on refit byte drift
cc4801f 
Make ground-truth verification report missing, extra, and changed tensors and cover merged expert maps from multiple EP source sidecars.
@KavinKrishnan
fix(vllm-wt): preserve translated weights across scratch reuse
f63d53b 
Take ownership of each translated tensor before the next EP source overwrites persistent receive buffers, preventing silent cross-expert corruption.
@KavinKrishnan
feat(mdl): support TP-local warm writes
30175fa 
Slice global checkpoint tensors to each vLLM TP rank and derive fused/expert offsets from live local parameter shapes before in-place MDL copies.
@KavinKrishnan KavinKrishnan mentioned this pull request Jul 11, 2026
Draft
@KavinKrishnan
chore(refit): log source identity for egress metrics
3510ba7 
Include source rank and source ID on each RDMA completion so differentiator harnesses can quantify trainer balance reliably.
@KavinKrishnan
feat(refit): add stage timing and harden native updates
32eee74 
Measure normalized refit stages, bound source discovery, persist receive state, and extend TP/FP8/subset safeguards so live NeMo-RL comparisons are attributable and fail closed.
@KavinKrishnan
feat(refit): prune partial updates on the wire
be061fe 
Pass subset tensor names into the existing scratch receive filter so warm partial refits reduce NIXL bytes without invalidating persistent registration.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
fix(mdl): map TP2 expert FP8 shards correctly
3fdab97 
Split fused W13 gate/up destinations before applying TP-local source slicing so weights and block scales do not collide on equal global/local shapes.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
chore(ci): normalize refit source headers
58b1ca0 
Add repository-standard SPDX headers to the refit instrumentation, tests, and existing RL utility scripts so copyright validation can evaluate the branch cleanly.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
feat(discovery): expose version and freshness in source listings
2f8adc3 
Add lightweight training-step and update-time fields so clients can discard stale or old READY rows before fetching full tensor metadata, reducing warm refit discovery fan-out.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
feat(mdl): default FP8 loaderless refit at TP>1
63d580f 
Enable the corrected loaderless install path automatically for FP8 at tensor-parallel width above one, now that TP1 and TP2 corrupt/refit/generate validation passes; keep MX_FP8_LOADERLESS as an override.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
feat(discovery): filter and cache refit source metadata
5aa9ae3 
Push model, rank, version, and freshness predicates into batched Redis discovery and reuse stable worker metadata across training steps, reducing live M5 discovery latency by 4.5x. Add EP transfer descriptor accounting to make full-contiguous fallback behavior measurable.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
feat(discovery): invalidate metadata cache on layout changes
a43825e 
Publish a stable topology and tensor-registry signature with lightweight source rows, key receiver metadata templates by that signature, and refresh safely when TP, PP, EP, expert ownership, or translation metadata changes in place.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan
feat(refit): add persistent device staging for EP gather
b478709 
Allow high-memory receivers to retain translated EP tensors on GPU across source pulls and refit cycles, removing the host round trip while preserving host staging as the compatibility default.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
This was referenced Jul 15, 2026
Open
Open
Merged
KavinKrishnan added a commit that referenced this pull request Jul 17, 2026
@KavinKrishnan
feat(server): add filtered versioned source discovery
6f1120b 
Extract the catalog portion of #482 so receivers can filter lightweight source rows by model, rank, version, and freshness while safely caching topology-stable metadata.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
@KavinKrishnan KavinKrishnan mentioned this pull request Jul 17, 2026
Merged
KavinKrishnan added a commit that referenced this pull request Jul 20, 2026
@KavinKrishnan
feat(vllm): add mapped direct load installer
ccd864a 
Extract the target-side installer from #482 so translated weights can be written into validated direct, fused, expert, TP-local, and FP8 scale destinations without coupling MDL to discovery or transport planning.

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

feat size/XXL

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@KavinKrishnan