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 2, 03:30 UTC
Fix sparse NCCL weight transfer test construction (#44345)
Signed-off-by: Siddharth Bedekar <bedeksid@gmail.com>
Siddharth Bedekar committed
Jun 2, 2026 at 14:51 UTC
0917a009d3fc0ff6f9e1a8bca1b10c6c1821d018
2 files changed
+7
-2
docs/training/weight_transfer/base.md
+1
@@ -156,5 +156,6 @@ from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
156
engine = WeightTransferEngineFactory.create_engine(
157
config=weight_transfer_config,
158
parallel_config=parallel_config,
159
+ model=model,
160
)
161
```
tests/distributed/test_weight_transfer.py
+6
-2
@@ -291,7 +291,9 @@ def test_nccl_receive_sparse_weights_without_init_raises():
291
292
config = WeightTransferConfig(backend="nccl")
293
parallel_config = create_mock_parallel_config()
294
- engine = NCCLWeightTransferEngine(config, parallel_config)
294
+ engine = NCCLWeightTransferEngine(
295
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
296
+ )
297
298
update_info = NCCLWeightTransferUpdateInfo(
299
names=["w"],
@@ -526,7 +528,9 @@ def inference_receive_sparse_tensor(
528
parallel_config.data_parallel_rank = 0
529
parallel_config.data_parallel_index = 0
530
529
- engine = NCCLWeightTransferEngine(config, parallel_config)
531
+ engine = NCCLWeightTransferEngine(
532
+ config, parallel_config, MagicMock(spec=torch.nn.Module)
533
+ )
534
engine.init_transfer_engine(
535
NCCLWeightTransferInitInfo(
536
master_address=master_address,