perf(nixl): pipeline reusable weight pulls

[S1ro1]

Summary

• use ModelExpress only for startup peer discovery and the trainer tensor table, then use exact versioned NIXL notifications for all policy transfers, including startup
• prepare static NIXL READ requests once, repost them after completion, and post every trainer-rank pull in a group before waiting
• attach trainer arena credits to READ completion instead of performing a separate acknowledgement handshake
• add explicit, opt-in two-arena overlap through weight_broadcast.overlap_transfer_and_replay

The path remains receiver-driven PULL. The default remains one arena because the second arena has a model-dependent GPU-memory cost. This PR adds no fallback protocol, runtime sizing heuristic, profiling instrumentation, or inference-scheduler change.

Protocol

ModelExpress is the one-time startup discovery plane. It exchanges the trainer tensor table and NIXL agent metadata, after which the fixed peers establish NIXL/UCX connections. The shared WeightSender/WeightReceiver sentinel handshake offers and acknowledges every policy version before the NIXL transport runs; this is the same outer lifecycle used by the other transports. MX source statuses are not used. Policy 0 and later policies use exact NIXL ready and complete notifications for the transfer itself.

sequenceDiagram
    participant T as Trainer ranks
    participant B as Broadcast directory
    participant R as NIXL weight receiver
    participant MX as ModelExpress
    participant I as Inference ranks

    Note over T,I: One-time setup
    I->>MX: Publish inference NIXL metadata
    R->>MX: Publish receiver NIXL metadata

    loop Each policy version
        T->>B: Offer version with sender ready marker
        R->>B: Acknowledge with receiver ready marker
        T->>B: Mark transfer started
        alt First version
            T->>MX: Publish tensor table and trainer metadata
            T->>MX: Fetch inference and receiver metadata
            R->>MX: Fetch trainer table and metadata
            Note over T,I: Establish NIXL and UCX connections once
        end
        T-->>R: NIXL policy ready notification
        R->>I: Update weights
        loop Each transfer group
            T-->>I: NIXL group ready notification
            I->>T: Post READs to every trainer rank
            T-->>I: RDMA payload
            I-->>T: READ completion returns arena credit
        end
        I-->>R: Update complete
        R-->>T: NIXL policy complete notification
        T->>B: Commit version with finished marker
    end

Loading

Each group signal contains (group index, generation), and each policy signal contains (step, event), so it is consumed exactly once. A trainer staging slot is reused only after every inference rank completes its READ from that slot.

NIXLWeightReceiver owns the ModelExpress discovery session, local NIXL agent, and typed trainer peer in the orchestrator process. The generic orchestrator and WeightWatcher only use the WeightReceiver API and contain no NIXL-specific state. Transfer plans and the trainer broadcaster store ordered NixlPeer lists, so positional trainer-agent indices stay direct list lookups and raw NIXL agent names remain inside the adapter. On the trainer, finish_transfer_group waits for every inference rank to complete a group before advancing its generation and allowing its arena slot to be reused.

Tensor addresses, shard routes, peers, and descriptor lists remain static for the worker lifetime. Requests are therefore prepared once and reposted after completion. For each group, inference waits for every serving trainer rank, posts all trainer-rank READs, and then waits for completion. Inference ranks rotate their trainer order, while #3340 pins each process to the active InfiniBand port nearest its GPU.

With overlap enabled, trainer and inference allocate exactly two largest-group-sized arenas. Inference replays and quantizes group N while receiving group N+1. Allocation is explicit: there is no free-memory probe, automatic downgrade, or hidden retry.

Measured results

Every populated cell below is the median of the positive framework time/broadcast_weights samples; terminal zero entries are excluded. This is the only timing definition used in the description. Parentheses contain Slurm job IDs, and blank cells mean that configuration was not measured.

Model	Old NIXL	New NIXL	New NIXL + overlap	NCCL
Qwen3-30B-A3B	4.651 s (2260)	1.898 s (2278)
GLM-4.5-Air	8.676 s (2286)	4.125 s (2288)	3.043 s (2305)	7.616 s (2289)
GLM-5.2		12.513 s (2297)	8.577 s (2317, EP32)

All NIXL baselines use the correct rail pinning, and router replay is disabled. The GLM-4.5-Air old, new, and overlap measurements are 20-step Hendrycks-math runs with 4096 maximum sampling tokens; the NCCL comparison uses the same model and node roles over three positive broadcasts. NCCL uses BF16 checkpoint-format transfer with quantize_in_weight_transfer = false.

GLM-5.2 job 2317 completed 20 rollout steps and 18 steady trainer broadcasts on 64 trainer and 32 inference H200s. Framework time/broadcast_weights was 8.577s median and 12.987s p95. At the inference client, end-to-end synchronization was 7.899s median and 12.577s p95: pause was 5.204s median / 9.882s p95, while the weight update itself was 2.842s median / 2.960s p95. Each inference GPU received 82.501 GB per policy; across all 576 rank/version samples, RDMA was 1.864s median / 2.027s p95, replay was 1.954s median / 1.993s p95, and their overlapped process time was 2.682s median / 2.836s p95. Median payload was 44.26 GB/s per GPU, or approximately 354 GB/s across an eight-GPU inference node. The W&B run contains the completed training metrics.

On this single framework metric, Qwen improves by 59.2% from old to new. GLM-4.5-Air improves by 52.5% from old to new and 64.9% from old to the final overlapped path; the final path is 60.1% faster than NCCL. GLM-5.2 improves by 31.5% from the measured EP16 one-arena path to the EP32 overlap run. Holding overlap constant, EP32 reduces median framework broadcast time by 11.4% from the EP16 result (9.683s, job 2301), while median client update time falls by 19.4% from 3.524s to 2.842s. Framework timing includes the unchanged vLLM pause/resume interval, whose variance dominates the EP32 tail, so it should not be interpreted as raw fabric bandwidth.

Memory and scope

Overlap adds one largest-group arena rather than another model copy. For GLM-5.2, that is 0.308 GB per trainer GPU and 3.807 GB per inference GPU. The EP32 run therefore used a 7.613 GB two-buffer arena per inference GPU and added 141.5 GB across the tested 64-trainer/32-inference topology. The option is disabled by default; models without that headroom retain the one-arena path.

GLM-5.2 uses 64 trainer H200s with EP=8 and full offload, plus DP/EP inference with online blockwise FP8 and DeepGEMM. EP16 receives 127.800 GB per inference GPU and retains 5,056 host-side handles per rank; EP32 receives 82.501 GB per GPU and prepares 72,819 descriptors across 79 transfer groups per rank. These request structures add no model-sized GPU allocation.

Every serving trainer rank connects to every inference rank so completion returns exact source credits. Static membership is the expected hot path; elastic scaling and fault tolerance remain out of scope.

Numerical validation

Jobs 2286 and 2305 completed the configured 20 GLM-4.5-Air rollout steps and finalized metrics. Step 20 of job 2305 contained 0/8 trainable traces, so it correctly emitted no optimizer or mismatch-KL row. Across the 19 steps with measurements in both runs, mean mismatch KL was 0.002602538 before and 0.002608075 after (+0.2%); median was 0.002786405 before and 0.002702605 after (-3.0%). Router replay was deliberately disabled, so individual steps include normal sampling variation.

Validation

• Ruff 0.15.14: uv run ruff format --check and uv run ruff check
• uv run python -m compileall on changed Python modules
• uv run pytest -q tests/unit/test_configs.py
• ported onto main at 21a4a3245 through feat(slurm): launch ModelExpress for NIXL transfers #3340 at 0d809b1e4; the generic orchestrator and WeightWatcher are unchanged
• two-node, three-step Qwen3-0.6B end-to-end NIXL smoke (Slurm job 2451): startup policy v0 and policies v1-v3 synchronized; trainer and orchestrator completed
• Qwen3-30B-A3B and GLM-4.5-Air eight-rail transfer benchmarks
• GLM-5.2 ten-node EP16 and twelve-node EP32 runs with full trainer offload, online blockwise FP8 inference, and DeepGEMM
• matched NCCL comparison on GLM-4.5-Air
• final 20-step GLM-4.5-Air Hendrycks-math run with 4096 sampling tokens and router replay disabled

---

Note

Medium Risk
Changes core distributed weight-sync choreography between trainer, orchestrator, and vLLM; mis-sync would cause stale weights or hangs, though scope is limited to the NIXL transport path.

Overview
NIXL weight broadcast is reworked so ModelExpress is only used for one-time peer discovery (trainer tensor table and NIXL agent metadata). Per-group and per-policy coordination now uses versioned NIXL notifications instead of ModelExpress READY/INITIALIZING buffer sessions.

The NIXL adapter gains PreparedRead / NixlPeer: READ descriptor lists are prepared once, reposted each update, and all trainer-rank pulls in a group are posted before waiting. Arena reuse is tied to READ completion via group (index, generation) signals; the orchestrator participates with policy ready / complete notifications over NIXL.

Inference follows the trainer’s staging_buffer_count (no GPU free-memory sizing heuristic). Optional weight_broadcast.overlap_transfer_and_replay allocates exactly two largest-group arenas on trainer and inference so the next group can be received while the current one is replayed; default remains one arena.

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

[cursor]

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

There are 2 total unresolved issues (including 1 from previous review).

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

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

[mikasenghaas]

trusss