[Bugfix] Keep batch-invariance NCCL pins out of the weight-transfer group (#59500)
Signed-off-by: aoshen524 <aoshen524@gmail.com> Co-authored-by: aoshen524 <aoshen524@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>
Guanxin Li committed
Oct 1, 2026 at 20:19 UTC
a810316c0894b8a43e53458ae5b14c9ad0b18b5d
4 files changed
+140
-27
tests/distributed/test_weight_transfer.py
+43
-5
@@ -68,6 +68,7 @@ from vllm.distributed.weight_transfer.sparse_nccl_engine import (
68
SparseWeightPatch,
69
)
70
from vllm.platforms import current_platform
71
+from vllm.utils.nccl import _nccl_has_no_cache
72
from vllm.utils.network_utils import get_open_port
73
74
@@ -433,13 +434,15 @@ def trainer_broadcast_tensor(
434
return True
435
436
436
-@ray.remote(num_gpus=1)
437
+# max_calls=1: a batch-invariant run leaves NCCL pins in the worker process.
438
+@ray.remote(num_gpus=1, max_calls=1)
439
def inference_receive_tensor(
440
master_address: str,
441
master_port: int,
442
world_size: int,
443
tensor_shape: list[int],
444
tensor_dtype: str,
445
+ batch_invariant: bool = False,
446
) -> dict:
447
"""Inference task that receives tensor via NCCLWeightTransferEngine."""
448
import contextlib
@@ -448,6 +451,21 @@ def inference_receive_tensor(
451
import torch
452
453
_set_ray_assigned_device()
454
+ if batch_invariant:
455
+ from vllm.model_executor.determinism.batch_invariant import (
456
+ override_envs_for_invariance,
457
+ )
458
+
459
+ override_envs_for_invariance()
460
+ # Like vLLM's own groups, a pinned communicator makes NCCL read the
461
+ # pins before the transfer group exists.
462
+ torch.distributed.init_process_group(
463
+ "nccl",
464
+ init_method=f"tcp://127.0.0.1:{get_open_port()}",
465
+ rank=0,
466
+ world_size=1,
467
+ )
468
+ torch.distributed.all_reduce(torch.ones(1, device="cuda"))
469
470
from vllm.config.parallel import ParallelConfig
471
from vllm.config.weight_transfer import WeightTransferConfig
@@ -534,11 +552,25 @@ def inference_receive_tensor(
552
torch.accelerator.device_count() < 2,
553
reason="Need at least 2 GPUs to run NCCL weight transfer test.",
554
)
537
-def test_nccl_weight_transfer_between_processes():
555
+@pytest.mark.parametrize(
556
+ "batch_invariant",
557
+ [
558
+ pytest.param(False, id="default"),
559
+ pytest.param(
560
+ True,
561
+ id="batch-invariant-worker",
562
+ marks=pytest.mark.skipif(
563
+ not _nccl_has_no_cache(), reason="Needs CUDA NCCL >= 2.29.7."
564
+ ),
565
+ ),
566
+ ],
567
+)
568
+def test_nccl_weight_transfer_between_processes(batch_invariant):
569
"""Test NCCL weight transfer from trainer to inference process using Ray.
570
571
This test verifies that the NCCLWeightTransferEngine can receive
541
- tensors broadcast by a trainer process via NCCL.
572
+ tensors broadcast by a trainer process via NCCL, including when the
573
+ worker runs with batch-invariance NCCL pins the trainer does not have.
574
"""
575
_init_ray_for_weight_transfer()
576
@@ -550,13 +582,19 @@ def test_nccl_weight_transfer_between_processes():
582
tensor_dtype = "float32"
583
584
inference_future = inference_receive_tensor.remote(
553
- master_address, master_port, world_size, tensor_shape, tensor_dtype
585
+ master_address,
586
+ master_port,
587
+ world_size,
588
+ tensor_shape,
589
+ tensor_dtype,
590
+ batch_invariant,
591
)
592
trainer_future = trainer_broadcast_tensor.remote(
593
master_address, master_port, world_size, tensor_shape, tensor_dtype
594
)
595
559
- trainer_result, result = ray.get([trainer_future, inference_future])
596
+ # A mismatched NCCL config deadlocks instead of failing.
597
+ trainer_result, result = ray.get([trainer_future, inference_future], timeout=300)
598
599
assert trainer_result, "Trainer should complete successfully"
600
assert result["success"], (
vllm/distributed/weight_transfer/nccl_common.py
+10
-7
@@ -26,6 +26,7 @@ from vllm.distributed.weight_transfer.packed_tensor import (
26
DEFAULT_PACKED_BUFFER_SIZE_BYTES,
27
DEFAULT_PACKED_NUM_BUFFERS,
28
)
29
+from vllm.utils.nccl import unpinned_nccl_env
30
31
32
def decode_nccl_unique_id(
@@ -154,7 +155,8 @@ def stateless_init_process_group(
155
pg = StatelessProcessGroup.create(
156
host=master_address, port=master_port, rank=rank, world_size=world_size
157
)
157
- return PyNcclCommunicator(pg, device=device)
158
+ with unpinned_nccl_env():
159
+ return PyNcclCommunicator(pg, device=device)
160
161
162
def uid_init_process_group(
@@ -170,12 +172,13 @@ def uid_init_process_group(
172
"""
173
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
174
173
- return PyNcclCommunicator.from_unique_id_bytes(
174
- nccl_unique_id_bytes,
175
- rank=rank,
176
- world_size=world_size,
177
- device=device,
178
- )
175
+ with unpinned_nccl_env():
176
+ return PyNcclCommunicator.from_unique_id_bytes(
177
+ nccl_unique_id_bytes,
178
+ rank=rank,
179
+ world_size=world_size,
180
+ device=device,
181
+ )
182
183
184
def worker_init_process_group(
vllm/model_executor/determinism/batch_invariant.py
+22
-15
@@ -16,6 +16,7 @@ from vllm.model_executor.determinism.batch_invariant_configs import (
16
from vllm.platforms import current_platform
17
from vllm.triton_utils import tl, triton
18
from vllm.utils.mem_utils import get_max_shared_memory_bytes
19
+from vllm.utils.nccl import pin_nccl_env
20
from vllm.utils.platform_utils import num_compute_units
21
from vllm.utils.torch_utils import is_torch_equal_or_newer
22
@@ -1144,22 +1145,28 @@ def override_envs_for_invariance():
1145
1146
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
1147
1147
- # NCCL determinism settings
1148
+ # NCCL determinism settings. NCCL keeps the launch mode process-wide, so,
1149
+ # unlike the pins below, it also applies to communicators shared with other
1150
+ # processes.
1151
os.environ["NCCL_LAUNCH_MODE"] = "GROUP"
1149
- os.environ["NCCL_COLLNET_ENABLE"] = "0"
1150
- os.environ["NCCL_NVLS_ENABLE"] = "0"
1151
- os.environ["NCCL_P2P_NET_DISABLE"] = "1"
1152
- os.environ["NCCL_MIN_NCHANNELS"] = "1"
1153
- os.environ["NCCL_MAX_NCHANNELS"] = "1"
1154
- os.environ["NCCL_PROTO"] = "Simple"
1155
- # NCCL >= 2.31 zero-fills the algorithm table of every collective when
1156
- # NCCL_ALGO is set and re-enables only the named ones; together with the
1157
- # NCCL_PROTO above, collectives not named here end up with no algorithm
1158
- # and fail with ncclInvalidUsage. Re-enable Ring and Tree for all
1159
- # collectives, then pin AllReduce to Tree for determinism.
1160
- os.environ["NCCL_ALGO"] = "ring,tree;allreduce:tree"
1161
- os.environ["NCCL_NTHREADS"] = "1"
1162
- os.environ["NCCL_SOCKET_NTHREADS"] = "1"
1152
+ pin_nccl_env(
1153
+ {
1154
+ "NCCL_COLLNET_ENABLE": "0",
1155
+ "NCCL_NVLS_ENABLE": "0",
1156
+ "NCCL_P2P_NET_DISABLE": "1",
1157
+ "NCCL_MIN_NCHANNELS": "1",
1158
+ "NCCL_MAX_NCHANNELS": "1",
1159
+ "NCCL_PROTO": "Simple",
1160
+ # NCCL >= 2.31 zero-fills the algorithm table of every collective
1161
+ # when NCCL_ALGO is set and re-enables only the named ones; together
1162
+ # with the NCCL_PROTO above, collectives not named here end up with
1163
+ # no algorithm and fail with ncclInvalidUsage. Re-enable Ring and
1164
+ # Tree for all collectives, then pin AllReduce to Tree.
1165
+ "NCCL_ALGO": "ring,tree;allreduce:tree",
1166
+ "NCCL_NTHREADS": "1",
1167
+ "NCCL_SOCKET_NTHREADS": "1",
1168
+ }
1169
+ )
1170
1171
# torch.compile settings
1172
os.environ["VLLM_USE_AOT_COMPILE"] = "0"
vllm/utils/nccl.py
+65
@@ -3,9 +3,11 @@
3
4
from __future__ import annotations
5
6
+import contextlib
7
import ctypes
8
import importlib.util
9
import os
10
+from collections.abc import Iterator, Mapping
11
12
import torch
13
@@ -129,3 +131,66 @@ def query_nccl_gin_type(
131
logger.warning("ncclCommQueryProperties returned error %d", result)
132
return None
133
return props.railedGinType if railed else props.ginType
134
+
135
+
136
+# Values the variables set by `pin_nccl_env` had before it: None until something
137
+# is pinned, empty if NCCL cannot re-read them.
138
+_unpinned_env: dict[str, str | None] | None = None
139
+
140
+
141
+def pin_nccl_env(pins: dict[str, str]) -> None:
142
+ """Set NCCL variables for this process's own communicators.
143
+
144
+ NCCL never reconciles these across ranks, so communicators shared with a
145
+ process that does not set them are created without them, see
146
+ `unpinned_nccl_env`. NCCL caches most variables on first read; listing the
147
+ pins in NCCL_NO_CACHE makes it re-read them per communicator. NCCL parses
148
+ that list once, so this must run before the process's first NCCL call.
149
+ """
150
+ global _unpinned_env
151
+ launch_env = _swap_env(pins)
152
+ if _unpinned_env is not None:
153
+ return
154
+ _unpinned_env = launch_env if _nccl_has_no_cache() else {}
155
+ if _unpinned_env:
156
+ no_cache = os.environ.get("NCCL_NO_CACHE")
157
+ os.environ["NCCL_NO_CACHE"] = ",".join(filter(None, [no_cache, *pins]))
158
+
159
+
160
+@contextlib.contextmanager
161
+def unpinned_nccl_env() -> Iterator[None]:
162
+ """Create communicators shared with other processes without the pins."""
163
+ if _unpinned_env == {}:
164
+ logger.warning_once(
165
+ "NCCL_NO_CACHE needs CUDA NCCL >= 2.29.7, so the NCCL settings vLLM "
166
+ "pinned also apply here; a process sharing this communicator must "
167
+ "set the same NCCL_* variables or it will hang."
168
+ )
169
+ pinned = _swap_env(_unpinned_env or {})
170
+ try:
171
+ yield
172
+ finally:
173
+ _swap_env(pinned)
174
+
175
+
176
+def _nccl_has_no_cache() -> bool:
177
+ # NCCL_NO_CACHE arrived in NCCL 2.29.7; ask the library vLLM loads.
178
+ if torch.version.cuda is None:
179
+ return False
180
+ version = ctypes.c_int()
181
+ try:
182
+ ctypes.CDLL(find_nccl_library()).ncclGetVersion(ctypes.byref(version))
183
+ except Exception:
184
+ return False
185
+ return version.value >= 22907
186
+
187
+
188
+def _swap_env(values: Mapping[str, str | None]) -> dict[str, str | None]:
189
+ """Set (None: unset) the given variables and return their previous values."""
190
+ old = {name: os.environ.get(name) for name in values}
191
+ for name, value in values.items():
192
+ if value is None:
193
+ os.environ.pop(name, None)
194
+ else:
195
+ os.environ[name] = value
196
+ return old