feat(refit): integrate ModelExpress weight synchronization

[zhengluo-nv]

Summary

• add model_express as a NeMo-RL weight synchronization transport using the public ModelExpress control, trainer, and generator clients
• integrate Megatron publication and vLLM staged installation through existing framework-native worker RPCs
• keep transport ownership, manifests, tensor selection, NIXL staging, verification, and vLLM installation inside ModelExpress
• remove NeMo-only trainer/generator wrappers; policy and vLLM workers now call the public MX clients directly

Architecture

NeMo-RL owns orchestration and safe points. ModelExpress owns rank-local transport resources and engine-specific transfer/install behavior.

flowchart LR
    subgraph Driver["NeMo-RL driver"]
        WS["ModelExpressWeightSynchronizer"]
        CC["ModelExpressControlClient"]
        WS --> CC
    end

    CP["ModelExpress RefitService<br/>version, shard, worker, lease metadata"]

    subgraph Policy["Megatron policy worker"]
        PW["Framework-native policy RPCs"]
        TC["ModelExpressTrainerClient"]
        TS["MX Megatron tensor selection"]
        SB["Registered trainer buffers"]
        MS["Trainer-local manifest service"]

        PW --> TC
        TC --> TS --> SB
        TC --> MS
    end

    subgraph Generator["vLLM generation worker"]
        GW["Framework-native vLLM RPCs"]
        GC["ModelExpressGeneratorClient"]
        ST["Staged NIXL transfer<br/>plan, pull, verify"]
        VI["vLLM graph-safe installer"]

        GW --> GC --> ST --> VI
    end

    CC -->|"create / inspect / retire version"| CP
    TC -->|"register worker + publish shard"| CP
    GC -->|"discover shards + hold lease"| CP
    GC -->|"fetch exact-version manifest"| MS
    SB -->|"NIXL reads"| ST
    WS -->|"publish / update / release RPCs"| PW
    WS -->|"initialize / update RPCs"| GW

Loading

Lifecycle:

1. Each policy actor constructs ModelExpressTrainerClient before distributed setup so NIXL is initialized before NCCL.
2. After Megatron setup, the actor builds MX-owned Megatron tensor specs and binds them once to the trainer client.
3. The synchronizer creates a version and invokes policy RPCs; each participating trainer client publishes its bound shard.
4. The synchronizer invokes vLLM RPCs; each generator client fetches the exact manifests, holds a lease, stages and verifies the NIXL transfer, and applies at the framework safe point.
5. The synchronizer retires the version and invokes trainer release before training mutates the in-place source buffers.

Integration boundary

The NeMo-RL orchestrator owns the version lifecycle through ModelExpressControlClient. Existing framework-native worker RPCs invoke ModelExpressTrainerClient inside Megatron policy actors and ModelExpressGeneratorClient inside vLLM generator actors. NeMo-RL does not construct ModelExpress adapters or own transfer state.

ModelExpress owns:

• trainer NIXL and manifest-service resources
• Megatron tensor selection and publication metadata
• exact-version source discovery and leases
• NIXL transfer planning, staging, and digest verification
• graph-safe vLLM installation

Dependencies

• ModelExpress RL client and staged generator implementation: ai-dynamo/modelexpress#648
• distributed reshard/QKV reliability prerequisite: ai-dynamo/modelexpress#635, merged

Current integration heads:

• NeMo-RL: 7e5982212
• ModelExpress PR Sequence packing + no dynamic batching failure on 8b #648: 83a446a6e

The latest GPU E2E validated the compatible pair NeMo-RL 51f7e051d and ModelExpress fec7d02d using the image below. The NeMo-RL delta to the current head only removes factory-level backend/topology gates and updates focused tests. ModelExpress has additional client, adapter, protobuf, and compatibility changes after the tested SHA, so a repeat GPU E2E on both current heads remains pending.

Validation

Current-head automated checks

• ModelExpress focused refit suite: 76 passed
• NeMo-RL focused ModelExpress synchronizer/factory suite: 9 passed
• NeMo-RL Ruff formatting and lint checks passed for the changed worker integrations
• NeMo-RL Python compilation checks passed
• git diff --check passed in both repositories
• both current commits are GPG-signed and DCO-signed off

The full NeMo-RL monorepo test suite has not been rerun for this branch.

Latest B200 GPU E2E

Model: meta-llama/Llama-3.1-8B-Instruct, BF16. Both runs used image nvcr.io/0980761089281446/model-express-dev-containers:nemo-rl-mx-fec7d02d-51f7e051d-b200-20260819-r2 (sha256:0f84554615c7bc1e9fa1ba90c8ba0437572be68a647976361847a9a161e7cb4d) in namespace zheng.

TP1 trainer → TP1 generator: lifecycle and redundant publishers

• topology: 4×B200 across two GPU nodes; two Megatron TP1 workers form DP2, and two independent vLLM TP1 generators run non-colocated
• both DP replicas published the same logical source slot as redundant physical candidates
• each generator captured 291 copies, 0 unsupported, 0 unattributed
• each generator loaded the exact manifest-carried NIXL metadata and staged 16.06 GB
• generation, log-prob computation, policy training, and Step 1/1 completed; driver exit code 0
• transfer/update: 15.04 s; total step: 33.65 s

This run validates the version, lease, manifest, transfer, install, redundant-source, and training lifecycle. Because source and destination are both TP1, it does not exercise shard-boundary resharding.

TP2 trainer → TP1 generator: resharding

• topology: 3×B200 on w4xnn; one Megatron TP2 group on two GPUs feeds one non-colocated vLLM TP1 generator on the third GPU, with a separate CPU Ray head
• the generator selected both distinct TP source slots and merged their shard boundaries into complete TP1 destination tensors
• capture: 291 copies, 0 unsupported, 0 unattributed
• both trainer NIXL agents loaded; 16.06 GB across 226 regions staged and installed
• generation and TP2 policy training completed Step 1/1; final head exit code 0
• transfer/update: 14.01 s; total step: 28.84 s

The first TP2 attempt completed training but the 64 GiB CPU head was OOM-killed during teardown. The identical rerun with a 128 GiB CPU-head limit completed cleanly.

The TP1 and TP2 timings are individual functional-run observations, not an A/B performance comparison: generator count, node placement, cache warmth, and head resources differed. Explicit post-transfer parameter equality and generation-parity checks remain pending.

Earlier two-node TP1 smoke

An earlier compatible pair was validated with image nvcr.io/0980761089281446/model-express-dev-containers:nemo-rl-mx-3cb98ac7-d26b92c02-b200-20260819-r3 (sha256:fc14b6f451e4668ca51ef0b4e267e4d714290cb2855fd825906c50cd856fb2a3).

• topology: two B200 nodes, one GPU per node; non-colocated Megatron TP1 trainer and vLLM TP1 generator
• transport: direct trainer-to-generator NIXL over UCX (rc,cuda_copy)
• strict capture: 291 copies, 0 unsupported sources, 0 unattributed copies
• staged transfer completed with full source-digest verification before vLLM apply
• training completed successfully: total step 27.73 s, transfer/update 12.51 s, average reward 1.0
• lifecycle audit: the completed version reached RELEASING with no remaining shard or lease records
• driver exit code 0

A non-fatal interpreter-shutdown warning remained after successful completion; it occurred after the training result and outside the refit data path.

[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.

[KavinKrishnan]

Two things I hit reading through this, neither an argument against the direction.

The first is concrete. ai-dynamo/modelexpress#648 deletes modelexpress/engines/vllm/refit/receiver.py, and that's the class the mx_reshard path in #3632 imports. So whichever of these lands first, the other needs a real fix rather than a rebase, and it'd be good to settle which receiver we're keeping before either merges. On top of that we both add a transport to the same VllmRefitSelector literal, the same VllmRefitConfig, and the same create_weight_synchronizer dispatch, so there's some mechanical overlap in ~11 files either way.

The second is about the E2E. TP1 trainer to TP1 generator doesn't move any shard boundaries, so it isn't really exercising the reshard. The two geometries that caused us trouble were the MoE one (Megatron TP2/EP4 into vLLM TP4) and models where KV heads are fewer than TP -- there the per-rank KV count floors to zero, and most ranks legitimately own no K or V rows at all, so the sparse merge is the whole game. Since this moves Megatron tensor selection into MX, those are the two I'd want green before we swap out a path that's already been through them. Happy to point you at the configs and the reduced 64Q/2KV model if that's useful.

[zhengluo-nv]

deletes modelexpress/engines/vllm/refit/receiver.py, and that's the class the mx_reshard path in #3632 imports

Yeah, I want to only expose 3 public clients to external, and the receiver will become an internal class. A modularized modelexpress-rl with a handful of simple public APIs should make future framework integration much easier. We will make these public APIs stable during MX version upgrade.

The two geometries that caused us trouble were the MoE one (Megatron TP2/EP4 into vLLM TP4) and models where KV heads are fewer than TP

I believe I copied all your code into the new folder. So it should not cause regression on the resharding