@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 223 lines 8.37 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import random
5 import typing
6
7 import pytest
8 import torch
9 import torch.distributed as dist
10 import torch.multiprocessing as mp
11
12 import vllm.envs as envs
13 from tests.utils import ensure_current_vllm_config
14 from vllm.distributed import cleanup_dist_env_and_memory
15 from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
16 from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
17 from vllm.distributed.device_communicators.pynccl_allocator import (
18 get_nccl_mem_pool,
19 is_symmetric_memory_enabled,
20 )
21 from vllm.distributed.parallel_state import (
22 get_tp_group,
23 init_distributed_environment,
24 initialize_model_parallel,
25 )
26 from vllm.platforms import current_platform
27 from vllm.utils.system_utils import update_environment_variables
28
29 torch.manual_seed(42)
30 random.seed(44)
31
32 test_size_elements = 4 * 1024 * 1024
33
34
35 def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
36 monkeypatch = pytest.MonkeyPatch()
37 with monkeypatch.context() as m:
38 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
39 dtype = torch.bfloat16
40 device = torch.device(f"cuda:{local_rank}")
41 torch.accelerator.set_device_index(device)
42 torch.set_default_device(device)
43 torch.set_default_dtype(dtype)
44 update_environment_variables(
45 {
46 "RANK": str(local_rank),
47 "LOCAL_RANK": str(local_rank),
48 "WORLD_SIZE": str(world_size),
49 "MASTER_ADDR": "localhost",
50 "MASTER_PORT": "12345",
51 }
52 )
53
54 init_distributed_environment()
55 with ensure_current_vllm_config():
56 initialize_model_parallel(tensor_model_parallel_size=world_size)
57
58 cuda_communicator = typing.cast(
59 CudaCommunicator, get_tp_group().device_communicator
60 )
61 pynccl_comm = cuda_communicator.pynccl_comm
62 if get_nccl_mem_pool() is None:
63 pytest.skip(
64 "NCCL allocator compilation failed (probably missing NCCL headers)."
65 )
66 if not is_symmetric_memory_enabled():
67 pytest.skip("NCCL symmetric memory allreduce is disabled.")
68
69 register_nccl_symmetric_ops(pynccl_comm)
70 input = torch.randint(1, 23, (test_size_elements,), dtype=dtype, device=device)
71 input_clone = input.clone()
72 output = torch.ops.vllm.all_reduce_symmetric_with_copy(input)
73 assert output is not None
74
75 group = get_tp_group().device_group
76 dist.all_reduce(input_clone, group=group)
77 torch.testing.assert_close(output, input_clone, atol=2.5, rtol=0.1)
78
79
80 @pytest.mark.skipif(
81 not current_platform.is_cuda(),
82 reason="NCCLSymmMemAllreduce is only available for CUDA platforms.",
83 )
84 @pytest.mark.parametrize("world_size", [2])
85 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
86 def test_nccl_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, world_size):
87 if world_size > torch.accelerator.device_count():
88 pytest.skip("Not enough GPUs to run the test.")
89
90 # Enable SymmMemCommunicator
91 monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
92 monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
93 monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
94
95 mp.spawn(nccl_symm_mem_allreduce_worker, args=(world_size,), nprocs=world_size)
96 cleanup_dist_env_and_memory()
97
98
99 def nccl_symm_mem_allgather_worker(local_rank: int, world_size: int):
100 monkeypatch = pytest.MonkeyPatch()
101 with monkeypatch.context() as m:
102 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
103 dtype = torch.bfloat16
104 device = torch.device(f"cuda:{local_rank}")
105 torch.accelerator.set_device_index(device)
106 torch.set_default_device(device)
107 torch.set_default_dtype(dtype)
108 update_environment_variables(
109 {
110 "RANK": str(local_rank),
111 "LOCAL_RANK": str(local_rank),
112 "WORLD_SIZE": str(world_size),
113 "MASTER_ADDR": "localhost",
114 "MASTER_PORT": "12346",
115 }
116 )
117
118 init_distributed_environment()
119 with ensure_current_vllm_config():
120 initialize_model_parallel(tensor_model_parallel_size=world_size)
121
122 cuda_communicator = typing.cast(
123 CudaCommunicator, get_tp_group().device_communicator
124 )
125 if get_nccl_mem_pool() is None:
126 pytest.skip(
127 "NCCL allocator compilation failed (probably missing NCCL headers)."
128 )
129 if not is_symmetric_memory_enabled():
130 pytest.skip("NCCL symmetric memory is disabled.")
131
132 per_rank_size = test_size_elements // world_size
133 input_tensor = torch.randint(
134 1, 23, (per_rank_size,), dtype=dtype, device=device
135 )
136 output = cuda_communicator.all_gatherv(input_tensor, dim=0)
137
138 group = get_tp_group().device_group
139 expected = torch.empty(test_size_elements, dtype=dtype, device=device)
140 dist.all_gather_into_tensor(expected, input_tensor, group=group)
141 torch.testing.assert_close(output, expected, atol=0.0, rtol=0.0)
142
143
144 @pytest.mark.skipif(
145 not current_platform.is_cuda(),
146 reason="NCCL symmetric memory is only available for CUDA platforms.",
147 )
148 @pytest.mark.parametrize("world_size", [2])
149 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
150 def test_nccl_symm_mem_allgather(monkeypatch: pytest.MonkeyPatch, world_size):
151 if world_size > torch.accelerator.device_count():
152 pytest.skip("Not enough GPUs to run the test.")
153
154 monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
155 monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
156 monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
157
158 mp.spawn(nccl_symm_mem_allgather_worker, args=(world_size,), nprocs=world_size)
159 cleanup_dist_env_and_memory()
160
161
162 def nccl_symm_mem_reduce_scatter_worker(local_rank: int, world_size: int):
163 monkeypatch = pytest.MonkeyPatch()
164 with monkeypatch.context() as m:
165 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
166 dtype = torch.bfloat16
167 device = torch.device(f"cuda:{local_rank}")
168 torch.accelerator.set_device_index(device)
169 torch.set_default_device(device)
170 torch.set_default_dtype(dtype)
171 update_environment_variables(
172 {
173 "RANK": str(local_rank),
174 "LOCAL_RANK": str(local_rank),
175 "WORLD_SIZE": str(world_size),
176 "MASTER_ADDR": "localhost",
177 "MASTER_PORT": "12347",
178 }
179 )
180
181 init_distributed_environment()
182 with ensure_current_vllm_config():
183 initialize_model_parallel(tensor_model_parallel_size=world_size)
184
185 cuda_communicator = typing.cast(
186 CudaCommunicator, get_tp_group().device_communicator
187 )
188 if get_nccl_mem_pool() is None:
189 pytest.skip(
190 "NCCL allocator compilation failed (probably missing NCCL headers)."
191 )
192 if not is_symmetric_memory_enabled():
193 pytest.skip("NCCL symmetric memory is disabled.")
194
195 per_rank_size = test_size_elements // world_size
196 input_tensor = torch.randint(
197 1, 23, (test_size_elements,), dtype=dtype, device=device
198 )
199 input_clone = input_tensor.clone()
200 output = cuda_communicator.reduce_scatter(input_tensor, dim=0)
201
202 group = get_tp_group().device_group
203 expected = torch.empty(per_rank_size, dtype=dtype, device=device)
204 dist.reduce_scatter_tensor(expected, input_clone, group=group)
205 torch.testing.assert_close(output, expected, atol=2.5, rtol=0.1)
206
207
208 @pytest.mark.skipif(
209 not current_platform.is_cuda(),
210 reason="NCCL symmetric memory is only available for CUDA platforms.",
211 )
212 @pytest.mark.parametrize("world_size", [2])
213 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
214 def test_nccl_symm_mem_reduce_scatter(monkeypatch: pytest.MonkeyPatch, world_size):
215 if world_size > torch.accelerator.device_count():
216 pytest.skip("Not enough GPUs to run the test.")
217
218 monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
219 monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
220 monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
221
222 mp.spawn(nccl_symm_mem_reduce_scatter_worker, args=(world_size,), nprocs=world_size)
223 cleanup_dist_env_and_memory()