feat(refit): add ModelExpress Megatron-to-vLLM reshard path

[KavinKrishnan]

What this PR does

This PR adds the complete NeMo-RL side of the ModelExpress mx_reshard refit path:

Each part has a narrow job in making trainer-to-generator refits usable from NeMo-RL:

• Describe what every Megatron rank owns. Megatron stores model weights already split across TP, EP, and ETP ranks. The publisher records each rank's real piece and its position in the full tensor, so ModelExpress can read it directly instead of first gathering a complete model on every trainer.
• Translate Megatron names into the names vLLM expects. The trainer and generator use different names for the same weights, and fused tensors may map to several HF tensors. Megatron-Bridge supplies that relationship; this adapter preserves the required Q/K/V and gate/up order and fails when a name is ambiguous.
• Connect ModelExpress to NeMo-RL's normal refit lifecycle. The WeightSynchronizer now starts publishers and receivers, waits for all required ranks, runs the transfer, commits the new version, and shuts both sides down in a safe order. Callers do not need a separate ModelExpress-only control path.
• Prevent partial models from becoming active. A refit proceeds only after every expected publisher is ready, and one receiver failure prevents the version from being committed. This avoids serving a model assembled from a mixture of old and new weights.
• Make correctness and delays visible. Parameter verification catches transfers that report success but install unchanged or wrong bytes. Separate stage records show whether time was spent waiting for publishers, moving bytes, or installing tensors, which keeps performance investigations from guessing.
• Provide a functional GRPO gate that cannot pass by accident. The test checks that the recipe launches, the model actually changes, and trainer/generator probabilities remain aligned. Portable recipes and analysis tools keep the same checks reproducible without upstreaming private cluster names or images.
• Publish global QKV geometry from each live layer. Nemotron-style layers can have fewer KV heads than TP ranks, and different layers may use different attention shapes. Reading each owning linear_qkv module avoids 2 // 8 = 0 and prevents one root-model setting from being incorrectly reused for every layer.

This PR depends on ModelExpress #635, which carries the global QKV interval alias builder the new QKV contract requires. Neither half is useful alone: NeMo-RL publishes the geometry, ModelExpress maps the intervals.

Architecture

flowchart LR
    M[Live Megatron layer] --> G[Per-layer Q/KV geometry resolver]
    G --> C[Parameter classifier]
    C --> B[Megatron-Bridge name map]
    B --> P[ModelExpress publisher]
    P --> Q[Publisher quorum]
    Q --> R[vLLM reshard receiver]
    R --> I[Install and verify]
    I --> V[Commit one version]

    W[WeightSynchronizer] --> P
    W --> Q
    W --> R
    W --> V
    W --> S[Ordered receiver/publisher shutdown]

Loading

For Q=64, KV=2 and trainer TP8, NeMo-RL now publishes global num_heads, num_kv_heads, head_dim, and qkv_interleave from each live linear_qkv.config. It does not publish num_kv_heads_local=0, and it does not pretend every rank owns a replicated KV head. ModelExpress receives each rank's real fused-row range and performs the interval mapping.

Why per-layer geometry matters

Nemotron-family models may use different attention geometry in different layers. Reading one root model config would stamp the same Q/KV shape onto every QKV tensor. The resolver therefore walks to the owning live linear_qkv module first and uses the root values only as a compatibility fallback. The chosen geometry must match the tensor's global fused-row count or publication fails closed.

Review guide

1. QKV/KV-heads-below-TP contract

Commit: 83c3d1324

Review first:

• nemo_rl/distributed/mx_megatron_helpers.py
• nemo_rl/models/policy/workers/megatron_policy_worker.py
• nemo_rl/weight_sync/mx_reshard_weight_synchronizer.py
• the matching helper, publisher, policy-worker, and synchronizer tests

Please check that layer-local config wins, local head fields are emitted only when both counts divide by TP, fused rows are validated, and only the obsolete num_query_groups % TP guard was relaxed.

2. Megatron publisher foundation

Commits: 23b90cf8, eea2ad6f, 537fa3ad

Review parameter role classification, global expert IDs, Megatron-Bridge name order, TP/ETP ranges, and fail-closed handling of unknown or ambiguous fused layouts.

3. Weight-sync and vLLM lifecycle

Review:

• mx_reshard_weight_synchronizer.py
• mx_reshard_publisher.py
• mx_vllm_reshard_receiver.py
• mx_reshard_config.py
• vLLM worker/config/backend wiring

The important ordering is publish all trainer ranks, establish receiver quorum, pull/install on all receiver ranks, then commit. Any failed rank prevents version commit. Shutdown releases receivers before publishers so NIXL registrations remain valid until reads finish.

4. Reliability and observability

Review cleanup after exceptions, telemetry isolation, separate quorum timing, and parameter verification in mx_refit_verify.py. Telemetry must never turn a successful refit into a failed one.

5. Functional gate and portable tools

Review tests/functional/grpo_mx_reshard_refit.sh, the three portable recipe configs, summarize_refit_stages.py, and attribute_variance.py.

The gate intentionally uses js_divergence_error rather than an absolute gen_kl_error ceiling; Megatron and vLLM kernels have a non-zero baseline KL even when every parameter is identical. It also requires a non-zero gradient so an unchanged model cannot pass vacuously.

Private namespace/image-specific cluster manifests are intentionally not included.

Evidence and scope

The pre-GQA integration was validated at 32 GPUs (16 Megatron trainer GPUs and 16 vLLM receiver GPUs) together with the original six ModelExpress #635 commits:

• dense and MoE refits reached complete coverage with zero fallback;
• same-checkpoint verification found zero changed parameters on every receiver rank;
• moving-model probability, JS-divergence, and non-zero-gradient gates passed;
• repeated refits and ordered worker teardown completed cleanly.

The new KV<TP contract is currently unit-qualified and representative CUDA-tensor-qualified with ModelExpress #635. Q=64/KV=2/head_dim=128 logical TP8 passed exact Q/K/V reconstruction, same-weight parity, changed-weight parity, and projection-output equality. This is not yet a full Megatron-to-vLLM TP8 E2E result and is not full Nemotron Ultra qualification.

Test plan

• NeMo helper/publisher suites — 45 passed
• NeMo mx_reshard synchronizer suite — 14 passed
• policy-worker resolver wiring check
• ModelExpress companion focused suite — 266 passed
• existing real 32-GPU BF16 dense/MoE validation for the pre-GQA integration
• CI on head 6a4978963
• real reduced Megatron TP8 -> vLLM TP8 Q=64/KV=2 E2E
• full 520-GPU Nemotron Ultra qualification

[copy-pr-bot]

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.