feat(weight-transfer): add NIXL + ModelExpress synchronization

[S1ro1]

Summary

Adds the initial production path for direct trainer-to-vLLM weight synchronization over NIXL and ModelExpress:

trainer-native weights -> declarative trainer-to-HF graph -> typed BF16/FP32 wire arenas -> recorded HF-to-vLLM loader graph -> final vLLM kernel tensors

This composes the trainer conversion graph from #2797 with vLLM loader recording from #2779. #2598 was used only as a reference for the ModelExpress transport shape.

Key properties:

• declarative trainer-to-HF conversion chains shared by every registered PrimeRL custom model family
• BF16 wire tensors by default, with model-declared FP32 preservation for tensors stored as FP32 in the source checkpoint
• reusable CUDA staging arenas registered with NIXL/UCX on both trainer and inference ranks
• lazy recording and replay of vLLM loaders, including view, split, stack, concatenate, and dtype-conversion paths
• bounded group-by-group transfer with independent trainer and inference sharding topologies
• ModelExpress coordination for startup and per-policy synchronization
• no vLLM source changes
• NIXL-specific transport code is isolated under prime_rl.nixl_weight_transfer

Correctness and failure handling

• Any trainer, inference, transport, or coordination failure fails the synchronization; there is no fallback path.
• Typed trainer tables carry exact tensor shapes, wire dtypes, and shard address ranges.
• Each policy version becomes live only after all participating inference ranks complete the update.
• The initial broadcast and cleanup lifecycle match the existing NCCL behavior.
• FP32 wire declarations were audited against the actual Hugging Face safetensors headers rather than runtime module allocation dtypes.
• NIXL inference identities use the local CUDA device index, preserving unique ranks when one admin endpoint owns multiple workers.
• The ModelExpress watcher stops advancing after the final trainer-published in-memory policy version.

Validation

Ruff 0.13.0 check, Ruff format check, and git diff --check pass after rebasing onto current main. CPU and GPU workflows are delegated to PR CI.
Existing Qwen3.5-MoE, Qwen3.5-MoE VLM, and Nemotron-H tests use the instance conversion API that owns the model config.

WhiteFiber Qwen3-MoE E2E

Validated Qwen3-30B-A3B-Instruct-2507 in nixl-modelexpress-nvl72-e2e:

• trainer: 2x4 GPUs, FSDP + EP8, FP32 master weights, FA4
• inference: 1x4 GPUs, TP4 with expert parallelism, eager mode
• max_off_policy_steps=1
• unquantized BF16/FP32 path
• 50 synchronized steps for both NIXL and matched NCCL controls on the same GPU nodes

Backend	Successful updates	KL mean	KL median	KL p95	KL max	Rank-0 sync mean
NIXL + ModelExpress	50/50	0.01742160	0.01398600	0.03868628	0.08084155	4.02599 s
NCCL control	50/50	0.02322200	0.01630000	0.06676000	0.09850000	8.140388 s

NIXL was 2.02196x faster for end-to-end rank-0 synchronization.

The NCCL-minus-NIXL mean KL difference was +0.0058004, with an approximate 95% interval of [-0.0009297, +0.0125306]; the reverse-text workload did not resolve a backend-related KL regression.

The NIXL run had zero rollout errors and zero pod restarts. NCCL had three expected stale-rollout cancellations at off-policy level 1 due to slower synchronization, with no transfer or inference failures.

W&B:

• NIXL trainer
• NIXL orchestrator
• NCCL trainer
• NCCL orchestrator

Current scope

• NIXL uses the UCX backend.
• A ModelExpress control-plane deployment is required externally.
• Inference ranks currently receive the complete model weight set, including all experts.
• vLLM private loading APIs are intentionally used.

---

Note

High Risk
Large new distributed weight-sync path on the critical RL training loop (trainer, orchestrator, vLLM workers) with NIXL/UCX and ModelExpress coordination; conversion refactors touch every custom model family used for broadcast.

Overview
Introduces weight_broadcast.type = "nixl" as a third in-memory transport (with NCCL and filesystem), wired through shared RL/trainer/orchestrator/inference configs, modelexpress==0.3.0 as a core dependency, and trainer NIXLWeightBroadcast plus vLLM NIXLWeightUpdateWorker for UCX/NIXL pulls coordinated by ModelExpress (session_id, gRPC port, READY/INITIALIZING rendezvous).

The orchestrator WeightWatcher can advance policy versions from ModelExpress status instead of waiting on filesystem STABLE markers; startup and per-step sync paths treat NIXL like NCCL for dispatch gating and worker-count auto-setup. LoRA remains filesystem-only for in-memory transports.

Model weight conversion is refactored: imperative per-family converters become conversion_ops chains (Rename, Stack, routed_experts_op, etc.) with instance convert_to_hf / convert_to_prime on PreTrainedModelPrimeRL, plus keep_in_fp32_for_weight_transfer hooks where wire dtype must stay FP32.

Smaller changes: drop the optional modelexpress uv extra from Docker/entrypoint sync, relax inference quantization typing and remove the mxfp8 SM100 config validator, and pass session_id through /init_broadcaster for NCCL compatibility.

^{Reviewed by Cursor Bugbot for commit 4484cb1. Bugbot is set up for automated code reviews on this repo. Configure here.}

[samsja]

why str ?

[samsja]

can we add docs on:

• the new conversion format for moe
• the way nixl.py file work and record the transformation

I will review the pr in depth, doing step by step

[samsja]

how much did we tested rl run ? I am worried about the potential change in bf16/fp32 for some of the layers, I know that before it was not super clearly definw who was doing what in dtype

[samsja]

can we rename the ht_to_tt function to hf_to_prl ? kinda not tt anymore no ?

[mikasenghaas]

would it make sens to have this in a prime_rl.broadcast.nixl submodule. next to fileystem and nccl ideally

[mikasenghaas]

why is this snake case when the rest is pascal case?

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 4484cb1. Configure here.}