fix(reshard): make distributed refit reliable and faster - #635

What this PR does

This PR completes the ModelExpress half of the NeMo-RL reshard-refit path. Each change below removes a specific source of wasted work, incorrect weight placement, or unsafe teardown:

• Collapse duplicate replica offers before planning. Data-parallel and expert-data-parallel ranks often advertise identical copies of the same shard. Keeping every copy made the receiver read the same bytes repeatedly and broke large reads into hundreds of thousands of small ones. This keeps one verified copy and leaves genuinely different shards untouched.
• Build NIXL read descriptors once per plan. Turning a transfer plan into low-level read requests is setup work. Rebuilding it every refit re-derived an identical list of ~413k objects in Python even when neither the layout nor the selected sources changed. It was also outside every timed stage, so it appeared only as unattributed time — which is what kept stage attribution unreportable at 60-86%.
• Install reconstructed tensors in one batch. Full-pull tensors were copied into live parameters through thousands of individual operations. Batching preserves the same destinations while removing repeated launch and interpreter overhead.
• Recognize vLLM's fused MoE expert views. vLLM reaches its per-expert loader path via unsqueeze(0).unbind(). That was not on the geometry allowlist, so every expert source was classified unsupported and MoE refits failed closed at ~5% coverage — 18,432 sources on Qwen3-30B-A3B. Capture now records the real expert storage without copying it.
• Release registered memory before shutting NIXL down. Destroying a NIXL agent while GPU memory was still registered aborted the process inside ucp_worker_destroy, taking down the whole Ray worker. This hit at the end of every GRPO run on this path, so a run that had actually succeeded still exited looking like a crash.
• Check publisher readiness without decoding the whole model. A quorum check needs only the published step and whether the table is non-empty. Parsing every tensor and shard to answer that delayed each refit without improving correctness.
• Map QKV using global row ranges when KV heads are below TP. num_kv_heads // tp_size gives zero for Nemotron Ultra's 2 KV heads at TP8. Megatron actually splits the fused QKV tensor by raw rows, so each rank's real row interval is mapped into Q, K and V instead of inventing a local KV-head count.

The NeMo-RL publisher and lifecycle integration is in NVIDIA-NeMo/RL#3632. The KV<TP fix needs both PRs: NeMo-RL publishes global per-layer QKV geometry, ModelExpress maps the raw fused-row interval.

Data flow

flowchart LR
    L[Live Megatron layer config] --> N[Global per-layer Q/KV descriptor]
    N --> P[Raw TP fused-row shard]
    P --> I[Global interval intersection]
    I --> Q[q_proj shards]
    I --> K[k_proj shards]
    I --> V[v_proj shards]

    A[Trainer shard tables] --> D[Deduplicate identical replicas]
    Q --> D
    K --> D
    V --> D
    D --> C[Build and cache read descriptors]
    C --> R[Batched NIXL reads]
    R --> B[Batched install]

    S[Publisher step stamps] --> U[Cheap quorum check]
    B --> X[Deregister memory]
    X --> Y[Destroy NIXL agent]

Measured effect of the six pre-GQA changes

The performance evidence below predates the GQA commit and does not apply to it.

On real Qwen3-30B BF16 tensors, Megatron EP8 publishers to a vLLM TP2 receiver, per receiver rank:

Also tested at 32 GPUs (16 trainer, 16 receiver) over NIXL/RDMA. Three independent dense cold starts per arm gave a 7.57x receiver speedup. MoE reached 100% coverage across 11 consecutive refits with 0 unsupported and 0 fallback. Exact verification found 0 of 435 parameters changed on all 16 receiver ranks for a same-checkpoint refit, and the moving-model GRPO probability / JS-divergence / non-zero-gradient gates passed.

GQA / KV-heads-below-TP evidence

This arm is unit-qualified plus representative CUDA-tensor-qualified, which is a lower bar than the above, and I'd rather say so than blur the two:

• Q=64, KV=2, head_dim=128 at logical trainer TP8 reconstructs byte-exact Q/K/V with no gaps or overlaps
• ranks without K/V rows publish Q only, and sparse tables merge into complete sources with zero fallback
• divisible layouts remain byte-for-byte identical to the legacy local-head path
• 24Q/6KV/TP4 and heterogeneous per-layer geometry are covered
• missing or invalid global geometry fails closed
• BF16 CUDA same-weight and changed-weight refits agree on parameters and projection outputs

Not yet done: real Megatron TP8 to vLLM TP8 end-to-end for Q=64/KV=2, and full Nemotron Ultra qualification.

Review guide

Roughly in order of how much scrutiny each deserves.

1. Global QKV interval mapping — newest correctness fix

modelexpress_rl/train/engines/megatron/aliases.py, tests/test_reshard_megatron_gqa.py

The code is arranged so the layout is a named thing rather than arithmetic inside a loop:

• _QkvLayout — holds the head counts, derives the group geometry.
• bands() — states Megatron's row order exactly once, as a sequence of runs. This is the single place encoding our assumption about Megatron's internal layout; if upstream changes it, this is the one function to fix.
• _read_qkv_layout — validates published metadata; half-specified metadata fails closed rather than guessing.
• _build_global_qkv_aliases — a plain interval intersection against those runs.

Worth checking: the per-group overlap math, that every source row is mapped exactly once (there is an explicit coverage assertion), that empty K/V tensors are omitted rather than published as zero-row tensors, and that divisible layouts still take the legacy path unchanged.

Known limitation: this reverse-engineers Megatron's fused-QKV row order rather than getting it from a supported interface, so it is sensitive to upstream change. Having Megatron expose the layout is the right long-term fix and is worth a follow-up.

2. Replica deduplication — main performance change

refit/reshard/rendezvous.py, tests/test_reshard_refit_replica_merge.py

Only offers with identical geometry and digest should collapse. This is the change most capable of silently dropping a genuinely distinct shard, so the merge predicate is the thing to check.

3. Fused MoE capture — main pre-GQA correctness change

geometry.py, receiver.py, types.py, tests/test_reshard_refit_moe_experts.py

The geometry.py diff is small (+28): one allowlist entry for unbind, a signature-transparent loader stamp (vLLM invokes the expert loader entirely by keyword, so a stamp naming its first parameter positionally raised TypeError on every expert), and retaining why a source was unsupported rather than only its name.

4. NIXL lifecycle

nixl_transfer.py, tests/test_nixl_peer_lifecycle.py

Every registered region must be released before agent destruction, including under partial setup and repeated shutdown.

5. Batched install and descriptor caching

receiver.py, tests/test_reshard_refit_batch_install.py, tests/test_reshard_refit_descriptor_cache.py

Cache invalidation: _plan is assigned in exactly two places, and _cached_descriptors is cleared immediately after, before any buffer registration. Batched copies: _foreach_copy_ has no defined ordering for overlapping destinations, so update_weights checks that destination views occupy disjoint storage and falls back to the captured per-view order when they don't.

6. Quorum parse skip

rendezvous.py and its tests. The fetch stays serial on purpose — concurrent metadata reads measured slower (median quorum 4.02s to 7.56s) by loading the shared server, and a test pins that.

Scope still pending

• real Megatron TP8 to vLLM TP8 E2E for Q=64/KV=2, with exact receiver parameters and generation agreement
• full 520-GPU Nemotron Ultra qualification
• FP8 installation (separate meta-tensor issue)
• restart and elastic-scale qualification
• the future O(1) quorum protocol change, which would carry publisher_step in the ListSources record and remove the serial fetch entirely

Test plan

• PYTHONPATH=. python3 -m pytest tests/test_*reshard*.py tests/test_*nixl*.py tests/test_envs.py -q — 278 passed
• CPU TP8 reconstruction, rendezvous merge, planner, full-pull and negative controls
• BF16 CUDA same-weight and moving-weight parameter/projection parity
• Real 32-GPU BF16 dense and MoE gates for the six pre-GQA changes
• Companion NeMo-RL #3632 CI on the global per-layer descriptor contract
• Real Megatron TP8 to vLLM TP8 Q=64/KV=2 E2E
• Full Nemotron Ultra qualification