refactor(vllm): use native IPC engine for colocated weight sync by miracle0517 · Pull Request #345 · vllm-project/vime

miracle0517 · 2026-07-13T03:38:35Z

Summary

1.Replace VIME's custom colocated IPC weight-sync implementation with the native weight transfer engines provided by vLLM and vLLM-Ascend.
This removes the manual IPC handle gathering, serialization, and worker-side chunk loading path while aligning colocated weight updates with vLLM's native /update_weights lifecycle.

2.Fix fused MoE weight shape mismatches during layerwise reload on Ascend NPU.

What changed

Use NPUIPCWeightTransferEngine on Ascend to send colocated weights.
Route native update_info payloads through VLLMEngine.update_weights and vLLM's /update_weights endpoint.
Initialize the native weight transfer engine once when colocated rollout engines are connected.
Use the native start_weight_update / finish_weight_update state machine with checkpoint-format loading.
Remove the custom:
- trainer-rank IPC gather groups
- IPC handle merging and serialization helpers
- update_weights_from_tensor / update_weights_chunk RPC path
- colocated worker-side tensor reconstruction logic
Retain the Ascend worker compatibility patches for MoE expert weight loaders and rotary embeddings.
Update unit tests to cover the native CUDA/NPU IPC transfer flow and isolate optional dependency stubs between test modules.

Why

1.The previous implementation duplicated functionality already supported by vLLM's native IPC transfer engines and required VIME to maintain custom GPU-slot gathering, IPC handle reconstruction, and worker RPC code.
Using the native transfer path reduces VIME-specific integration code and keeps colocated weight synchronization aligned with upstream vLLM and vLLM-Ascend behavior.

2.vLLM-Ascend may reload fused MoE weights with a transposed layout. Direct copying causes npu_grouped_matmul shape errors, while skip wake up transpose can solve it.

Testing

pytest tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py -q
pytest tests/unit/backends/vllm_utils/test_vllm_engine.py -q
Pre-commit formatting and lint checks
End-to-end Qwen3-4B colocated training on Ascend NPU
pytest tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py -q
End-to-end Qwen3-30B-A3B colocated training on Ascend NPU

Comparison of reward and train_rollout_logprob_abs_diff curves between colocated and non-colocated training：

Replace the custom colocated IPC handle gathering and worker chunk RPC path with vLLM/vLLM-Ascend native IPC weight transfer engines. Route generated update_info through VLLMEngine.update_weights and vLLM's /update_weights endpoint, initialize colocated transfer engines once during connection, and update unit tests for the native colocated transfer flow. Signed-off-by: wuxiang <498160096@qq.com>

read-the-docs-community · 2026-07-13T03:39:39Z

Documentation build overview

📚 vime | 🛠️ Build #33580940 | 📁 Comparing 92a4fe9 against latest (36a84ad)

🔍 Preview build

45 files changed · + 2 added · ± 27 modified · - 16 deleted

+ Added

± Modified

- Deleted

gemini-code-assist

Code Review

This pull request refactors the colocated vLLM weight synchronization mechanism to leverage native IPC transfer engines (NPUIPCWeightTransferEngine and IPCWeightTransferEngine) for both NPU and CUDA platforms, removing custom IPC gathering and serialization logic. It also introduces shape-changing MoE restore patches and updates associated unit tests and scripts. Feedback on the changes suggests calling .contiguous() on transposed tensors in the NPU worker's wake_up patch to prevent potential runtime crashes or silent correctness issues caused by non-contiguous memory layouts on NPU platforms.

gemini-code-assist · 2026-07-13T03:40:14Z

+                for name, param in model.named_parameters():
+                    if "w2_weight" in name and param.shape[2] == hidden_size:
+                        transposed = param.transpose(1, 2)
+                    elif "w13_weight" in name and param.shape[1] == hidden_size:
+                        transposed = param.transpose(1, 2)
+                    else:
+                        continue
+                    parent_name, param_name = name.rsplit(".", 1)
+                    parent = model.get_submodule(parent_name)
+                    setattr(parent, param_name, torch.nn.Parameter(transposed, requires_grad=False))


On NPU platforms, custom operators (such as those in Ascend/ATB) are highly sensitive to tensor memory layout and typically require contiguous tensors. Transposing a tensor via param.transpose(1, 2) returns a non-contiguous view, which can lead to silent correctness issues, runtime crashes, or memory access violations during execution. Calling .contiguous() on the transposed tensor ensures it is safely laid out in memory before wrapping it in a Parameter.

Suggested change

for name, param in model.named_parameters():

if "w2_weight" in name and param.shape[2] == hidden_size:

transposed = param.transpose(1, 2)

elif "w13_weight" in name and param.shape[1] == hidden_size:

transposed = param.transpose(1, 2)

else:

continue

parent_name, param_name = name.rsplit(".", 1)

parent = model.get_submodule(parent_name)

setattr(parent, param_name, torch.nn.Parameter(transposed, requires_grad=False))

for name, param in model.named_parameters():

if "w2_weight" in name and param.shape[2] == hidden_size:

transposed = param.transpose(1, 2).contiguous()

elif "w13_weight" in name and param.shape[1] == hidden_size:

transposed = param.transpose(1, 2).contiguous()

else:

continue

parent_name, param_name = name.rsplit(".", 1)

parent = model.get_submodule(parent_name)

setattr(parent, param_name, torch.nn.Parameter(transposed, requires_grad=False))

CalvinXKY · 2026-07-14T02:45:24Z

Title/body only cover the MoE transpose, but commit 38f431a is a large native-IPC refactor. Please expand the description so reviewers don’t miss the weight-sync rewrite.
And close #332.

CalvinXKY · 2026-07-14T02:51:11Z



-def _patch_npu_colocate_worker() -> None:
-    """Skip layerwise_reload in NPUWorker weight-update hooks (colocate OOM fix)."""


Please call out in the PR that this intentionally re-enables native layerwise on NPU (previously skipped for OOM). Any memory regress vs the old skip path?

Please call out in the PR that this intentionally re-enables native layerwise on NPU (previously skipped for OOM). Any memory regress vs the old skip path?

We conducted a comparative experiment on Qwen3-30B. The peak memory usage of the new version compared to the old version is 35.51 GB vs. 34.89 GB, with a difference within 1.78%

CalvinXKY · 2026-07-14T02:51:39Z

With native IPC, rank0 now fans out one merged-UUID payload to all colocated engines. Confirmed OK for multi-engine colocate, not only single-engine 30B?

Signed-off-by: wuxiang <498160096@qq.com>

miracle0517 · 2026-07-15T07:43:14Z

With native IPC, rank0 now fans out one merged-UUID payload to all colocated engines. Confirmed OK for multi-engine colocate, not only single-engine 30B?

uuid is generated in npu_ipc_engine( physical_npu_id = npu_generate_uuid() = f"{get_ip()}-{physical_device}"), so it is unique. And it is tested ok in multi-engine colocate。

CalvinXKY · 2026-07-16T06:45:44Z

@miracle0517 @momo609

gemini-code-assist Bot reviewed Jul 13, 2026

View reviewed changes

miracle0517 force-pushed the fix/moe_colocate branch 3 times, most recently from ad53378 to c5446b7 Compare July 13, 2026 08:57

miracle0517 marked this pull request as ready for review July 14, 2026 01:35

CalvinXKY reviewed Jul 14, 2026

View reviewed changes

fix(npu):skip moe model transport in wake up

92a4fe9

Signed-off-by: wuxiang <498160096@qq.com>

miracle0517 force-pushed the fix/moe_colocate branch from c5446b7 to 92a4fe9 Compare July 14, 2026 10:54

miracle0517 changed the title ~~fix(npu): transpose fused MoE weights during layerwise reload~~ refactor(vllm): use native IPC engine for colocated weight sync Jul 15, 2026

miracle0517 mentioned this pull request Jul 15, 2026

refactor(vllm): use native IPC engine for colocated weight sync #332

Closed

4 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

refactor(vllm): use native IPC engine for colocated weight sync#345

refactor(vllm): use native IPC engine for colocated weight sync#345
miracle0517 wants to merge 2 commits into
vllm-project:ascendfrom
miracle0517:fix/moe_colocate

miracle0517 commented Jul 13, 2026 •

edited

Loading

Uh oh!

read-the-docs-community Bot commented Jul 13, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 13, 2026

Uh oh!

CalvinXKY commented Jul 14, 2026

Uh oh!

CalvinXKY Jul 14, 2026

Uh oh!

miracle0517 Jul 15, 2026

Uh oh!

CalvinXKY commented Jul 14, 2026

Uh oh!

miracle0517 commented Jul 15, 2026

Uh oh!

CalvinXKY commented Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		def _patch_npu_colocate_worker() -> None:
		"""Skip layerwise_reload in NPUWorker weight-update hooks (colocate OOM fix)."""

Uh oh!

Conversation

miracle0517 commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changed

Why

Testing

Uh oh!

read-the-docs-community Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

CalvinXKY commented Jul 14, 2026

Uh oh!

CalvinXKY Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

miracle0517 Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

CalvinXKY commented Jul 14, 2026

Uh oh!

miracle0517 commented Jul 15, 2026

Uh oh!

CalvinXKY commented Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

miracle0517 commented Jul 13, 2026 •

edited

Loading

read-the-docs-community Bot commented Jul 13, 2026 •

edited

Loading