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 14, 05:30 UTC
[BugFix] Fix OOM in vLLM replicas by ensuring consistent NCCL memory accounting (#25359)
Signed-off-by: Kourosh Hakhamaneshi <kourosh@anyscale.com>
kourosh hakhamaneshi committed
Sep 23, 2025 at 15:49 UTC
abad204be6eaaae3630db4287ebdca004197f819
3 files changed
+191
-7
.buildkite/test-pipeline.yaml
+3
@@ -887,6 +887,8 @@ steps:
887
- tests/v1/test_external_lb_dp.py
888
- tests/v1/entrypoints/openai/test_multi_api_servers.py
889
- vllm/v1/engine/
890
+ - vllm/v1/worker/
891
+ - tests/v1/worker/test_worker_memory_snapshot.py
892
commands:
893
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
894
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/test_external_lb_dp.py
@@ -908,6 +910,7 @@ steps:
910
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s test_sharded_state_loader.py
911
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
912
- pytest -v -s models/multimodal/generation/test_maverick.py
913
+ - pytest -v -s v1/worker/test_worker_memory_snapshot.py
914
915
- label: Plugin Tests (2 GPUs) # 40min
916
timeout_in_minutes: 60
tests/v1/worker/test_worker_memory_snapshot.py
new
+174
@@ -0,0 +1,174 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import multiprocessing as mp
5
+import os
6
+import tempfile
7
+from multiprocessing import Queue
8
+from typing import Optional
9
+from unittest.mock import patch
10
+
11
+import pytest
12
+import torch
13
+
14
+from vllm.engine.arg_utils import EngineArgs
15
+from vllm.utils import MemorySnapshot
16
+from vllm.v1.worker.gpu_worker import (Worker,
17
+ init_worker_distributed_environment)
18
+
19
+# Global queue to track operation order across processes
20
+_QUEUE: Optional[Queue] = None
21
+
22
+
23
+def track_operation(operation: str, rank: int):
24
+ """Track when an operation happens and its rank."""
25
+ if _QUEUE is not None:
26
+ _QUEUE.put((operation, rank))
27
+
28
+
29
+def make_operation_tracker(operation_name: str, original_func):
30
+ """Create a mock function that tracks when an operation is called.
31
+
32
+ Args:
33
+ operation_name: Name to use when tracking this operation
34
+ original_func: The original function to wrap
35
+
36
+ Returns:
37
+ A wrapper function that tracks the operation and calls the original
38
+ """
39
+
40
+ def wrapper(*args, **kwargs):
41
+ rank = int(os.environ.get("RANK", "-1"))
42
+ track_operation(operation_name, rank)
43
+ return original_func(*args, **kwargs)
44
+
45
+ return wrapper
46
+
47
+
48
+def worker_process(rank: int, world_size: int, distributed_init_method: str,
49
+ queue: Queue, error_queue: Queue):
50
+ """Worker process that initializes a GPU worker with proper tracking."""
51
+ global _QUEUE
52
+ _QUEUE = queue
53
+
54
+ try:
55
+ # Set environment variables
56
+ os.environ["RANK"] = str(rank)
57
+ os.environ["LOCAL_RANK"] = str(rank)
58
+ os.environ["WORLD_SIZE"] = str(world_size)
59
+
60
+ # Create vLLM config with small model
61
+ vllm_config = EngineArgs(model="facebook/opt-125m",
62
+ tensor_parallel_size=2,
63
+ load_format="dummy").create_engine_config()
64
+
65
+ # Create worker
66
+ worker = Worker(
67
+ vllm_config=vllm_config,
68
+ local_rank=rank,
69
+ rank=rank,
70
+ distributed_init_method=distributed_init_method,
71
+ )
72
+
73
+ # Get original functions before patching
74
+ original_init_worker = init_worker_distributed_environment
75
+ original_memory_snapshot_init = MemorySnapshot.__init__
76
+ original_all_reduce = torch.distributed.all_reduce
77
+
78
+ # Apply minimal patches to track operation order
79
+ init_patch = patch(
80
+ 'vllm.v1.worker.gpu_worker.init_worker_distributed_environment',
81
+ side_effect=make_operation_tracker("init_distributed",
82
+ original_init_worker))
83
+ memory_patch = patch.object(
84
+ MemorySnapshot, '__init__',
85
+ make_operation_tracker("memory_snapshot",
86
+ original_memory_snapshot_init))
87
+ all_reduce_patch = patch('torch.distributed.all_reduce',
88
+ side_effect=make_operation_tracker(
89
+ "nccl_all_reduce", original_all_reduce))
90
+
91
+ with init_patch, memory_patch, all_reduce_patch:
92
+
93
+ # Initialize device (this is where we test the order)
94
+ worker.init_device()
95
+
96
+ # Load model to ensure everything works
97
+ worker.load_model()
98
+
99
+ # Signal success
100
+ queue.put(("success", rank))
101
+
102
+ except Exception as e:
103
+ error_queue.put((rank, str(e), type(e).__name__))
104
+ raise
105
+
106
+
107
+@pytest.mark.skipif(torch.cuda.device_count() < 2,
108
+ reason="Need at least 2 GPUs for tensor parallelism")
109
+def test_init_distributed_is_called_before_memory_snapshot():
110
+ """Test that distributed env is setup before memory snapshot.
111
+
112
+ This test makes sure during worker initialization, the initial memory
113
+ snapshot is taken after distributed env is setup to include all the buffers
114
+ allocated by distributed env.
115
+ """
116
+ world_size = 2
117
+
118
+ # Create a temporary file for distributed init
119
+ with tempfile.NamedTemporaryFile(delete=False) as f:
120
+ distributed_init_method = f"file://{f.name}"
121
+
122
+ # Create queues for inter-process communication
123
+ ctx = mp.get_context("spawn")
124
+ operation_queue = ctx.Queue()
125
+ error_queue = ctx.Queue()
126
+
127
+ # Start worker processes
128
+ processes = []
129
+ for rank in range(world_size):
130
+ p = ctx.Process(target=worker_process,
131
+ args=(rank, world_size, distributed_init_method,
132
+ operation_queue, error_queue))
133
+ p.start()
134
+ processes.append(p)
135
+
136
+ # Wait for all processes to complete
137
+ for p in processes:
138
+ p.join(timeout=60) # 60 second timeout
139
+
140
+ # Check for errors
141
+ errors = []
142
+ while not error_queue.empty():
143
+ rank, error_msg, error_type = error_queue.get()
144
+ errors.append(f"Rank {rank}: {error_type}: {error_msg}")
145
+
146
+ if errors:
147
+ pytest.fail("Worker processes failed:\n" + "\n".join(errors))
148
+
149
+ # Collect all operations from the queue
150
+ operations = []
151
+ while not operation_queue.empty():
152
+ operations.append(operation_queue.get())
153
+
154
+ # Verify we got operations from both ranks
155
+ print(f"Collected operations: {operations}")
156
+
157
+ # Check operations for each rank
158
+ for rank in range(world_size):
159
+ rank_ops = [op for op, r in operations if r == rank]
160
+ print(f"\nRank {rank} operations: {rank_ops}")
161
+
162
+ # Raises ValueError if the operation is not found
163
+ init_distributed = rank_ops.index("init_distributed")
164
+ nccl_all_reduce = rank_ops.index("nccl_all_reduce")
165
+ memory_snapshot = rank_ops.index("memory_snapshot")
166
+
167
+ # Verify order: init_distributed should happen before memory_snapshot
168
+ assert init_distributed < nccl_all_reduce < memory_snapshot, (
169
+ f"Rank {rank}: init_distributed (index {init_distributed}) "
170
+ f"must happen before nccl_all_reduce (index {nccl_all_reduce}) "
171
+ f"and memory_snapshot (index {memory_snapshot})")
172
+
173
+ # Clean up
174
+ os.unlink(distributed_init_method.replace("file://", ""))
vllm/v1/worker/gpu_worker.py
+14
-7
@@ -169,6 +169,20 @@ class Worker(WorkerBase):
169
current_platform.set_device(self.device)
170
171
current_platform.check_if_supports_dtype(self.model_config.dtype)
172
+
173
+ # Initialize the distributed environment BEFORE taking
174
+ # memory snapshot
175
+ # This ensures NCCL buffers are allocated before we measure
176
+ # available memory
177
+ init_worker_distributed_environment(self.vllm_config, self.rank,
178
+ self.distributed_init_method,
179
+ self.local_rank,
180
+ current_platform.dist_backend)
181
+
182
+ # Set random seed.
183
+ set_random_seed(self.model_config.seed)
184
+
185
+ # Now take memory snapshot after NCCL is initialized
186
gc.collect()
187
torch.cuda.empty_cache()
188
@@ -190,13 +204,6 @@ class Worker(WorkerBase):
204
else:
205
raise RuntimeError(
206
f"Not support device type: {self.device_config.device}")
193
- # Initialize the distributed environment.
194
- init_worker_distributed_environment(self.vllm_config, self.rank,
195
- self.distributed_init_method,
196
- self.local_rank,
197
- current_platform.dist_backend)
198
- # Set random seed.
199
- set_random_seed(self.model_config.seed)
207
208
# Construct the model runner
209
self.model_runner: GPUModelRunner = GPUModelRunner(