A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 00:30 UTC
Add `model` to `WeightTransferEngine.__init__` (#42922)
Signed-off-by: SumanthRH <sumanthrh99@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Sumanth R Hegde committed
May 22, 2026 at 17:52 UTC
3cb83c95922f34d0a3dc5d7795ab7b4c6ef6e54b
7 files changed
+78
-29
tests/distributed/test_weight_transfer.py
+39
-13
@@ -100,7 +100,9 @@ class TestNCCLEngineParsing:
100
"""Test parsing valid init info dict."""
101
config = WeightTransferConfig(backend="nccl")
102
parallel_config = create_mock_parallel_config()
103
- engine = NCCLWeightTransferEngine(config, parallel_config)
103
+ engine = NCCLWeightTransferEngine(
104
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
105
+ )
106
107
init_info = engine.parse_init_info(
108
{
@@ -121,7 +123,9 @@ class TestNCCLEngineParsing:
123
"""Test parsing init info with missing required field."""
124
config = WeightTransferConfig(backend="nccl")
125
parallel_config = create_mock_parallel_config()
124
- engine = NCCLWeightTransferEngine(config, parallel_config)
126
+ engine = NCCLWeightTransferEngine(
127
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
128
+ )
129
130
with pytest.raises(ValueError, match="Invalid init_info"):
131
engine.parse_init_info(
@@ -135,7 +139,9 @@ class TestNCCLEngineParsing:
139
"""Test parsing valid update info dict."""
140
config = WeightTransferConfig(backend="nccl")
141
parallel_config = create_mock_parallel_config()
138
- engine = NCCLWeightTransferEngine(config, parallel_config)
142
+ engine = NCCLWeightTransferEngine(
143
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
144
+ )
145
146
update_info = engine.parse_update_info(
147
{
@@ -161,14 +167,18 @@ class TestEngineRegistry:
167
"""Test factory creates NCCL engine."""
168
config = WeightTransferConfig(backend="nccl")
169
parallel_config = create_mock_parallel_config()
164
- engine = WeightTransferEngineFactory.create_engine(config, parallel_config)
170
+ engine = WeightTransferEngineFactory.create_engine(
171
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
172
+ )
173
assert isinstance(engine, NCCLWeightTransferEngine)
174
175
def test_create_engine_ipc(self):
176
"""Test factory creates IPC engine."""
177
config = WeightTransferConfig(backend="ipc")
178
parallel_config = create_mock_parallel_config()
171
- engine = WeightTransferEngineFactory.create_engine(config, parallel_config)
179
+ engine = WeightTransferEngineFactory.create_engine(
180
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
181
+ )
182
assert isinstance(engine, IPCWeightTransferEngine)
183
184
def test_create_engine_invalid_backend(self):
@@ -176,7 +186,9 @@ class TestEngineRegistry:
186
config = WeightTransferConfig(backend="invalid")
187
parallel_config = create_mock_parallel_config()
188
with pytest.raises(ValueError, match="Invalid weight transfer backend"):
179
- WeightTransferEngineFactory.create_engine(config, parallel_config)
189
+ WeightTransferEngineFactory.create_engine(
190
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
191
+ )
192
193
def test_register_duplicate_raises(self):
194
"""Test registering duplicate engine name raises."""
@@ -196,7 +208,9 @@ def test_nccl_receive_weights_without_init_raises():
208
209
config = WeightTransferConfig(backend="nccl")
210
parallel_config = create_mock_parallel_config()
199
- engine = NCCLWeightTransferEngine(config, parallel_config)
211
+ engine = NCCLWeightTransferEngine(
212
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
213
+ )
214
215
update_info = NCCLWeightTransferUpdateInfo(
216
names=["w"],
@@ -273,7 +287,9 @@ def inference_receive_tensor(
287
parallel_config.data_parallel_rank = 0
288
parallel_config.data_parallel_index = 0
289
276
- engine = NCCLWeightTransferEngine(config, parallel_config)
290
+ engine = NCCLWeightTransferEngine(
291
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
292
+ )
293
294
# Initialize the engine (joins as rank 1)
295
init_info = NCCLWeightTransferInitInfo(
@@ -478,7 +494,9 @@ class TestIPCEngineParsing:
494
495
config = WeightTransferConfig(backend="ipc")
496
parallel_config = create_mock_parallel_config()
481
- engine = IPCWeightTransferEngine(config, parallel_config)
497
+ engine = IPCWeightTransferEngine(
498
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
499
+ )
500
501
# Create dummy IPC handles
502
dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
@@ -512,7 +530,9 @@ class TestIPCEngineParsing:
530
531
config = WeightTransferConfig(backend="ipc")
532
parallel_config = create_mock_parallel_config()
515
- engine = IPCWeightTransferEngine(config, parallel_config)
533
+ engine = IPCWeightTransferEngine(
534
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
535
+ )
536
537
dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
538
dummy_tensor2 = torch.ones(50, device="cuda:0")
@@ -545,7 +565,9 @@ class TestIPCEngineParsing:
565
566
config = WeightTransferConfig(backend="ipc")
567
parallel_config = create_mock_parallel_config()
548
- engine = IPCWeightTransferEngine(config, parallel_config)
568
+ engine = IPCWeightTransferEngine(
569
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
570
+ )
571
572
dummy_tensor = torch.ones(10, 10, device="cuda:0")
573
_, ipc_handle = reduce_tensor(dummy_tensor)
@@ -635,7 +657,9 @@ def inference_receive_ipc_tensor(
657
parallel_config.data_parallel_rank = 0
658
parallel_config.data_parallel_index = 0
659
638
- engine = IPCWeightTransferEngine(config, parallel_config)
660
+ engine = IPCWeightTransferEngine(
661
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
662
+ )
663
664
# Initialize the engine (no-op for IPC)
665
init_info = IPCWeightTransferInitInfo()
@@ -763,7 +787,9 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():
787
788
config = WeightTransferConfig(backend="ipc")
789
parallel_config = create_mock_parallel_config()
766
- engine = IPCWeightTransferEngine(config, parallel_config)
790
+ engine = IPCWeightTransferEngine(
791
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
792
+ )
793
794
# Create IPC handle with wrong GPU UUID
795
dummy_tensor = torch.ones(10, 10, device="cuda:0")
vllm/distributed/weight_transfer/__init__.py
+2
@@ -5,8 +5,10 @@ Weight transfer engines for syncing model weights from trainers
5
to inference workers.
6
"""
7
8
+from vllm.distributed.weight_transfer.base import WeightTransferEngine
9
from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
10
11
__all__ = [
12
+ "WeightTransferEngine",
13
"WeightTransferEngineFactory",
14
]
vllm/distributed/weight_transfer/base.py
+6
-1
@@ -65,7 +65,10 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
65
update_info_cls: type[TUpdateInfo]
66
67
def __init__(
68
- self, config: WeightTransferConfig, parallel_config: ParallelConfig
68
+ self,
69
+ config: WeightTransferConfig,
70
+ parallel_config: ParallelConfig,
71
+ model: torch.nn.Module,
72
) -> None:
73
"""
74
Initialize the weight transfer engine.
@@ -73,9 +76,11 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
76
Args:
77
config: The configuration for the weight transfer engine
78
parallel_config: The configuration for the parallel setup
79
+ model: The local model instance which will receive the weights
80
"""
81
self.config = config
82
self.parallel_config = parallel_config
83
+ self.model = model
84
85
def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo:
86
"""
vllm/distributed/weight_transfer/factory.py
+5
-1
@@ -10,6 +10,8 @@ from vllm.distributed.weight_transfer.base import WeightTransferEngine
10
from vllm.logger import init_logger
11
12
if TYPE_CHECKING:
13
+ import torch
14
+
15
from vllm.config.parallel import ParallelConfig
16
from vllm.config.weight_transfer import WeightTransferConfig
17
@@ -75,12 +77,14 @@ class WeightTransferEngineFactory:
77
cls,
78
config: "WeightTransferConfig",
79
parallel_config: "ParallelConfig",
80
+ model: "torch.nn.Module",
81
) -> WeightTransferEngine:
82
"""Create a weight transfer engine instance.
83
84
Args:
85
config: Weight transfer configuration containing the backend name
86
parallel_config: Parallel configuration for the engine
87
+ model: The local model instance which will receive the weights
88
89
Returns:
90
An initialized weight transfer engine instance
@@ -102,7 +106,7 @@ class WeightTransferEngineFactory:
106
engine_cls.__name__,
107
)
108
105
- return engine_cls(config, parallel_config)
109
+ return engine_cls(config, parallel_config, model)
110
111
112
# Register built-in weight transfer engines here.
vllm/distributed/weight_transfer/ipc_engine.py
+6
-2
@@ -125,7 +125,10 @@ class IPCWeightTransferEngine(
125
update_info_cls = IPCWeightTransferUpdateInfo
126
127
def __init__(
128
- self, config: WeightTransferConfig, parallel_config: ParallelConfig
128
+ self,
129
+ config: WeightTransferConfig,
130
+ parallel_config: ParallelConfig,
131
+ model: torch.nn.Module,
132
) -> None:
133
"""
134
Initialize the IPC weight transfer engine.
@@ -133,8 +136,9 @@ class IPCWeightTransferEngine(
136
Args:
137
config: The configuration for the weight transfer engine
138
parallel_config: The configuration for the parallel setup
139
+ model: The local model instance which will receive the weights
140
"""
137
- super().__init__(config, parallel_config)
141
+ super().__init__(config, parallel_config, model)
142
143
def parse_update_info(
144
self, update_dict: dict[str, Any]
vllm/distributed/weight_transfer/nccl_engine.py
+6
-2
@@ -109,7 +109,10 @@ class NCCLWeightTransferEngine(
109
update_info_cls = NCCLWeightTransferUpdateInfo
110
111
def __init__(
112
- self, config: WeightTransferConfig, parallel_config: ParallelConfig
112
+ self,
113
+ config: WeightTransferConfig,
114
+ parallel_config: ParallelConfig,
115
+ model: torch.nn.Module,
116
) -> None:
117
"""
118
Initialize the NCCL weight transfer engine.
@@ -117,8 +120,9 @@ class NCCLWeightTransferEngine(
120
Args:
121
config: The configuration for the weight transfer engine
122
parallel_config: The configuration for the parallel setup
123
+ model: The local model instance which will receive the weights
124
"""
121
- super().__init__(config, parallel_config)
125
+ super().__init__(config, parallel_config, model)
126
self.model_update_group: PyNcclCommunicator | None = None
127
128
def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None:
vllm/v1/worker/gpu_worker.py
+14
-10
@@ -36,7 +36,10 @@ from vllm.distributed.parallel_state import (
36
get_pp_group,
37
get_tp_group,
38
)
39
-from vllm.distributed.weight_transfer import WeightTransferEngineFactory
39
+from vllm.distributed.weight_transfer import (
40
+ WeightTransferEngine,
41
+ WeightTransferEngineFactory,
42
+)
43
from vllm.logger import init_logger
44
from vllm.lora.request import LoRARequest
45
from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
@@ -131,15 +134,9 @@ class Worker(WorkerBase):
134
# Buffers saved before sleep
135
self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
136
134
- # Weight transfer engine (initialized on-demand)
135
- self.weight_transfer_engine = (
136
- WeightTransferEngineFactory.create_engine(
137
- self.vllm_config.weight_transfer_config,
138
- self.vllm_config.parallel_config,
139
- )
140
- if self.vllm_config.weight_transfer_config is not None
141
- else None
142
- )
137
+ # Weight transfer engine is created in `load_model` once the model
138
+ # is available, since the engine needs a reference to the model.
139
+ self.weight_transfer_engine: WeightTransferEngine | None = None
140
self._weight_update_active = False
141
self._is_checkpoint_format = True
142
@@ -344,6 +341,13 @@ class Worker(WorkerBase):
341
):
342
self.model_runner.load_model(load_dummy_weights=load_dummy_weights)
343
344
+ if self.vllm_config.weight_transfer_config is not None:
345
+ self.weight_transfer_engine = WeightTransferEngineFactory.create_engine(
346
+ self.vllm_config.weight_transfer_config,
347
+ self.vllm_config.parallel_config,
348
+ self.model_runner.get_model(),
349
+ )
350
+
351
def update_config(self, overrides: dict[str, Any]) -> None:
352
self.model_runner.update_config(overrides)
353