perf: accelerate colocated and non-colocated vLLM weight updates

[aoshen02]

Summary

• make packed transfer an internal invariant: one CUDA IPC allocation per colocated chunk and vLLM packed NCCL for remote engines
• coalesce TP gathers and EP broadcasts while bypassing TP=1 / EP=1 no-op collectives
• keep PP/EP source selection topology-correct without adding a narrow PP-only fast path or new CLI controls
• enable vLLM packed transfer for non-colocated MoE and compressed-tensors models, including expert weights
• keep each routed-expert layer atomic during conversion

This branch is based on current main, including the vLLM 0.25.1 / CUDA 13 image update in #353, and contains one focused commit. The stale rollout naming and Qwen3-Next quantization changes previously carried by this PR have been removed.

Design

For colocated CUDA IPC, a vLLM worker selects its handle by physical GPU UUID. Every trainer rank in a colocated engine slot therefore needs to expose the same complete HF metadata stream; sending only the local PP stage is not correct unless inference PP is explicitly aligned with training PP.

The raw iterator preserves the general path: PP broadcasts restore stage-owned parameters, EP broadcasts restore expert shards, and TP gathers reconstruct tensor shards before HF conversion. TP tensors are gathered in group/dtype buckets. EP tensors are coalesced by source, with the mapping from original global source rank to the current EP group's local root computed once from PP group membership. This supports mixed TP/PP/EP/DP topologies without topology-specific branches in the hot path.

Each converted chunk is packed into one byte tensor and one CUDA IPC handle per trainer GPU. The historical packed/non-packed CLI switch and legacy IPC payload are removed; transport format is now an internal invariant, matching Slime’s unconditional FlattenedTensorBucket colocated path. Up to four chunk updates remain in flight internally; the window is implementation backpressure rather than a user-facing tuning parameter. Tensors remain alive until vLLM has closed the IPC handles.

The non-colocated path keeps its existing per-PP-stage ownership and now uses packed vLLM transfer for both dense and expert chunks. Expert parameters are bucketed by parent layer so converter calls remain atomic.

H200 validation

All tests used vllm/vime:nightly-dev-20260715c (vLLM 0.25.1, PyTorch 2.11.0+cu130) on one 8xH200 node.

Performance A/B

Moonlight-16B-A3B FP8, trainer TP4/EP2, four rollout engines at TP2/EP2:

Revision	Initial weight update
Current main (ee16693cd, code-equivalent baseline at 6cefd8463)	20.3 s
This PR	5.2 s

All 240 /update_weights requests returned HTTP 200, rollout completed, and training completed. The optimized timing matches the original PR implementation while using a smaller, general design.

The baseline and this PR both later hit the same Transformer Engine FP8 checkpoint-save issue (extra_fp8_variables=None), after training and independently of weight transfer. FP8 train/rollout logprob differences are also large on both revisions, so BF16 was used for the correctness checks below.

Correctness and topology

• Qwen3-0.6B BF16, trainer TP2/PP2, four TP2 rollout engines: both updates completed in 2.1-2.7 s; train_rollout_logprob_abs_diff = 0.013537; job succeeded.
• Qwen3-30B-A3B BF16, trainer TP2/PP2/EP2, four TP2/EP2 rollout engines: both 114-chunk updates completed in 11.1 s and 10.8 s; all update requests returned HTTP 200; train_rollout_logprob_abs_diff = 0.0188007; job succeeded.

Both logprob differences are below the CI threshold of 0.1. The 30B run specifically exercises PP-to-EP source remapping with real MoE shards.

Buildkite A/B

This compares the vLLM 0.25.1 image PR (Buildkite #669, fcdd0db8) with this PR (Buildkite #685, 93c589ade). Update timings use only Timer update_weights end, excluding queueing, model downloads, and initialization. Tests with multiple training steps report their mean.

Across the 56 directly comparable weight updates, total time fell from 751.2 s to 308.4 s (-58.9%, 2.44x faster). All 19 common tests that emit an update timer were non-regressing. Across 33 common logprob samples, mean train_rollout_logprob_abs_diff stayed effectively unchanged (0.0161984 to 0.0161861), while the maximum decreased from 0.0526013 to 0.0496524.

Test	Update #669	Update #685	Speedup	Logprob #669	Logprob #685	Logprob delta
full disk update	1.700 s	1.650 s	2.9%	0.007583	0.008097	+0.000514
GLM4 9B quick start	1.400 s	1.333 s	4.8%	0.010685	0.010255	-0.000430
GLM4.7 30B Mooncake	22.367 s	7.700 s	65.6%	0.051289	0.048423	-0.002865
Qwen3 30B DeepEP/FP8	51.900 s	18.067 s	65.2%	0.037622	0.037401	-0.000221
Qwen3.6 35B Mooncake/DeepEP	50.567 s	13.200 s	73.9%	0.014023	0.013863	-0.000160
Qwen3 30B R3 DeepEP/FP8	54.100 s	18.500 s	65.8%	0.031398	0.030788	-0.000610
Qwen3 30B R3	20.767 s	8.700 s	58.1%	0.013967	0.013402	-0.000564
Qwen3 4B PPO	2.800 s	1.333 s	52.4%	0.015052	0.014850	-0.000202
Qwen3 4B disaggregate	6.933 s	6.567 s	5.3%	0.015282	0.015385	+0.000103
Qwen3 4B critic-only	1.867 s	1.500 s	19.7%	n/a	n/a	n/a
logprob/entropy GPU	n/a	n/a	n/a	n/a	n/a	n/a
release train	1.900 s	1.700 s	10.5%	0.007875	0.008871	+0.000996
streaming partial rollout	2.467 s	1.500 s	39.2%	0.011763	0.013260	+0.001497
Moonlight 16B	10.133 s	3.900 s	61.5%	0.014266	0.013506	-0.000760
Moonlight 16B R3	8.533 s	5.900 s	30.9%	0.008999	0.009599	+0.000601
MIMO 7B MTP	2.700 s	2.150 s	20.4%	0.008805	0.010588	+0.001784
debug rollout/train	0.000 s	0.000 s	same	0.008558	0.007695	-0.000863
OPD vLLM	0.267 s	0.233 s	12.7%	0.007635	0.007783	+0.000149
external PD	8.350 s	7.375 s	11.7%	0.008674	0.008716	+0.000042
fanout short	0.333 s	0.300 s	9.9%	0.012825	0.014692	+0.001867

Two suites are new in #685 and have no #669 baseline:

Test	Update time	Logprob diff	Result
non-colocate PP	1.560 s mean	0.006034 mean / 0.006196 max	passed
precision parallel matrix	0.010 s mean / 0.400 s max across 128 updates	0.016896 mean / 0.017665 max across 64 layouts	passed

Every logprob result remains well below the CI threshold of 0.1.

Validation

• focused updated suites: 91 passed
• complete tests/utils: 165 passed
• four-rank Gloo subgroup probe confirmed that torch.distributed._broadcast_coalesced expects a group-local root rank
• full-repository pre-commit: passed
• git diff --check: passed
• H200 A/B, BF16 PP, and BF16 PP+EP integration runs described above
• Buildkite #685, Read the Docs, and DCO: passed for 93c589ade

Diff reduction

• before cleanup: 6 commits, 16 files, +910/-126
• final: 1 commit, 11 files, +560/-362
• production code: +189/-209; remaining additions are regression and contract tests in existing test files

Related work

• [Perf][Model Loader] Reload layout-preserving weights directly vllm#48382 optimizes receiver reload and supplies the WNA16 runtime-view path.
• [bugfix] Fix distributed update weights for pipeline parallel #329 fixes Bridge PP collective semantics; this PR is rebased on and preserves that behavior.
• Support online MTP draft weight updates #351 adds MTP draft-model synchronization; the unified sender preserves it.

AI assistance was used. The submitter reviewed every changed line and ran the listed tests and model evaluations.

[gemini-code-assist]

Code Review

This pull request introduces a pipeline parallel (PP) broadcast fast path for weight updates, allowing raw-mode pipeline stages to be converted in parallel and broadcast as packed HF chunks across colocated trainer ranks before CUDA IPC transfer. It adds new CLI arguments (--vllm-weight-update-max-inflight and --vllm-colocate-pp-broadcast), updates the weight update logic to support packed metadata, and renames variables in rollout.py for clarity. Feedback on the changes highlights a critical issue in the PP broadcast implementation where dist.broadcast is called with a global rank instead of the relative group rank for the src parameter, which could lead to out-of-bounds errors or incorrect broadcasting.

[gemini-code-assist]

In PyTorch distributed, when calling collective operations like dist.broadcast with a custom process group (group=pp_group), the src argument must be the relative rank (group rank) of the source process within that process group, not its global rank.

Using the global rank (source_rank) will cause out-of-bounds rank errors or incorrect broadcasting behavior if the pipeline parallel group ranks do not start at 0 or are not contiguous global ranks.

Since pp_ranks is the list of global ranks ordered by their rank within pp_group, the index of source_rank in pp_ranks is exactly its group rank. We can use enumerate(pp_ranks) to get both the group rank (src_group_rank) and the global rank (source_rank), and pass src=src_group_rank to dist.broadcast.

Suggested change

	for source_rank in pp_ranks:
	for chunk_index, metadata in enumerate(metadata_by_global_rank[source_rank]):
	if rank == source_rank:
	_, packed_tensor = local_chunks[chunk_index]
	local_chunks[chunk_index] = None
	else:
	packed_tensor = torch.empty(
	sum(metadata["tensor_sizes"]),
	dtype=torch.uint8,
	device=torch.cuda.current_device(),
	)
	dist.broadcast(packed_tensor, src=source_rank, group=pp_group)
	for src_group_rank, source_rank in enumerate(pp_ranks):
	for chunk_index, metadata in enumerate(metadata_by_global_rank[source_rank]):
	if rank == source_rank:
	_, packed_tensor = local_chunks[chunk_index]
	local_chunks[chunk_index] = None
	else:
	packed_tensor = torch.empty(
	sum(metadata["tensor_sizes"]),
	dtype=torch.uint8,
	device=torch.cuda.current_device(),
	)
	dist.broadcast(packed_tensor, src=src_group_rank, group=pp_group)

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #33610022 | 📁 Comparing 93c589a against latest (b929921)

  🔍 Preview build  

44 files changed · ± 44 modified

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 34 more...