feat(refit): add a ModelExpress-brokered nccl_reshard transport

[yixinh-nv]

What does this PR do?

Adds refit_transport=mx_nccl_reshard: the existing nccl_reshard data path, with NCCL communicator bootstrap brokered by ModelExpress instead of a TCPStore.

Below communicator creation, the path is unchanged: nccl.m2n.reshard moves the sharded bulk tensors and the packed NCCL broadcast moves the miscellaneous tensors. MX coordinates admission, immutable membership, worker-generation fencing, plan agreement, per-lane UID publication, and readiness; it is not the tensor-transfer engine.

Design boundary

• Data plane: existing nccl.m2n.reshard plus packed NCCL broadcast.
• Control plane: expected roster, epoch, plan digest, registration lease, and lane bootstrap ID.
• Communicator lifecycle: the MX path creates a bounded nonblocking bootstrap communicator, waits for every asynchronous setup operation, then splits it into a blocking child communicator. M2N receives that child because its current native implementation treats ncclInProgress as an error.
• Lane ordering: rendezvous, all-ranks broadcast lane, then each PP-local reshard lane, with a cluster-wide barrier between phases. Overlapping communicators are never initialized concurrently.

Dependencies

PR	Role
ModelExpress #661	Collective server: admission, TTL/liveness, epochs and report fencing
ModelExpress #665	Client: plan digest and registration lease renewal
ModelExpress #666	Collective execution/data-plane hardening

Matched H200 end-to-end results

Both runs used Nebius 8xH200 (two 4-GPU pods), Qwen3-30B-A3B BF16, Megatron TP2/PP2/EP2 to vLLM TP4, the same 240-prompt file (SHA-256 8fffc35d...be9da1), five GRPO steps, and the same 56.87 GiB refit payload: 54.00 GiB M2N bulk plus 2.87 GiB misc broadcast. The only transport selection was nccl_reshard versus mx_nccl_reshard.

Field	Native, no MX (nccl_reshard)	MX (mx_nccl_reshard)
First refit	6.63 s	7.52 s
Steady refits	0.62, 0.64, 0.92, 0.64 s	0.61, 0.86, 0.74, 0.64 s
Steady mean / median	0.71 / 0.64 s	0.71 / 0.69 s
Five KL errors	0.0029, 0.0032, 0.0034, 0.0025, 0.0029	0.0029, 0.0032, 0.0034, 0.0025, 0.0029
Native M2N proof	All eight workers logged real nccl.m2n.reshard	All eight workers logged real nccl.m2n.reshard
Workload result	5/5 GRPO steps completed	5/5 GRPO steps completed
Driver teardown	Old launcher hit the known Ray finalizer exit bug after completion	Exit status 0; zero live Ray actors/GPU use afterward

For the MX leg, the control plane admitted 8 workers, reached epoch 1 READY, agreed on the plan digest, and created all three lanes before the first refit.

These are enclosing prepare_for_generation/transfer_and_update_weights timings: they include Ray/actor dispatch and synchronization in addition to device work. They are not wire-bandwidth measurements. All samples are retained, including the cold first refit and the 0.92 s native outlier.

The validation used the pinned #2971 runtime base (0c595799) plus the validated MX integration hunks and the current PR's bootstrap/cleanup fixes. The run provenance records the exact source hashes, ModelExpress server 6f74f08d, and ModelExpress client/data 4f85110e.

The samples establish that MX preserves the native M2N execution and correctness path. They are too few, and the native leg predates the final Ray-cleanup fix, to claim an MX performance delta. A new native rerun at the final exact heads would make the performance comparison publication-ready.

Key fixes discovered during H200 validation

• Poll nonblocking communicator progress before reading the bootstrap broadcast buffer; otherwise a lazily connected collective can leave a peer's verification tensor unchanged and trigger a cascading abort.
• Promote the bounded bootstrap communicator to a blocking child before handing it to M2N, whose current native calls reject ncclInProgress.
• Explicitly drain generation, policy, teacher, and placement-group owners before Ray teardown so late destructors cannot reinitialize a second CoreWorker during Python finalization.

Usage

uv run ./examples/run_grpo.py \
  --config <your_config>.yaml \
  policy.generation.colocated.enabled=false \
  policy.generation.refit_transport=mx_nccl_reshard \
  policy.generation.mx_server_url=<modelexpress-server>:8001

Every worker on both sides must receive the same ModelExpress server address. NCCL_COMM_ID must be unset for this path.

Validation status

• Focused MX bootstrap/plan tests: 30 passed in the H200 runtime environment.
• ModelExpress fix: allow last batch less than microbatch_size. #661/feat: introduce megatron checkpoint dir precedence #665/fix: Tie weights after set_model_state_dict if required #666 configured CI and DCO checks are green.
• NeMo RL configured PR checks and DCO are green.
• Real 8xH200 native and MX M2N legs completed all five application steps.
• Native rerun at the final exact heads with the new Ray-cleanup path, if a publication-grade A/B claim is needed.

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