@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 268 lines 7.04 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """DeepEP test utilities."""
4
5 import dataclasses
6 import os
7 import traceback
8 from collections.abc import Callable
9 from typing import Concatenate
10
11 import torch
12 from torch.distributed import ProcessGroup
13 from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage]
14 from typing_extensions import ParamSpec
15
16 from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2
17 from vllm.utils.network_utils import get_open_port
18
19 if has_deep_ep():
20 from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
21 DeepEPHTPrepareAndFinalize,
22 )
23 from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
24 DeepEPLLPrepareAndFinalize,
25 )
26
27 if has_deep_ep_v2():
28 from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
29 DeepEPV2PrepareAndFinalize,
30 )
31
32 ## Parallel Processes Utils
33
34 P = ParamSpec("P")
35
36
37 class GINNotAvailableError(RuntimeError):
38 pass
39
40
41 @dataclasses.dataclass
42 class ProcessGroupInfo:
43 world_size: int
44 world_local_size: int
45 rank: int
46 node_rank: int
47 local_rank: int
48 device: torch.device
49
50
51 def _worker_parallel_launch(
52 local_rank: int,
53 world_size: int,
54 world_local_size: int,
55 node_rank: int,
56 init_method: str,
57 worker: Callable[Concatenate[ProcessGroupInfo, P], None],
58 *args: P.args,
59 **kwargs: P.kwargs,
60 ) -> None:
61 rank = node_rank * world_local_size + local_rank
62 torch.accelerator.set_device_index(local_rank)
63 device = torch.device("cuda", local_rank)
64 torch.distributed.init_process_group(
65 backend="nccl",
66 init_method=init_method,
67 rank=rank,
68 world_size=world_size,
69 )
70 barrier = torch.tensor([rank], device=device)
71 torch.distributed.all_reduce(barrier)
72
73 try:
74 worker(
75 ProcessGroupInfo(
76 world_size=world_size,
77 world_local_size=world_local_size,
78 rank=rank,
79 node_rank=node_rank,
80 local_rank=local_rank,
81 device=device,
82 ),
83 *args,
84 **kwargs,
85 )
86 except Exception as ex:
87 print(ex)
88 traceback.print_exc()
89 raise
90 finally:
91 torch.distributed.destroy_process_group()
92
93
94 def parallel_launch(
95 world_size: int,
96 worker: Callable[Concatenate[ProcessGroupInfo, P], None],
97 *args: P.args,
98 **kwargs: P.kwargs,
99 ) -> None:
100 assert not kwargs
101 try:
102 spawn(
103 _worker_parallel_launch,
104 args=(
105 world_size,
106 world_size,
107 0,
108 f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}",
109 worker,
110 )
111 + args,
112 nprocs=world_size,
113 join=True,
114 )
115 except Exception as exc:
116 # pytest.skip cannot propagate directly through torch.multiprocessing.
117 if "GINNotAvailableError" in str(exc):
118 import pytest
119
120 pytest.skip("NCCL GIN not available (no IBGDA-capable hardware)")
121 raise
122
123
124 ## DeepEP specific utils
125
126
127 @dataclasses.dataclass
128 class DeepEPHTArgs:
129 num_local_experts: int
130
131
132 @dataclasses.dataclass
133 class DeepEPLLArgs:
134 max_tokens_per_rank: int
135 hidden_size: int
136 num_experts: int
137 use_fp8_dispatch: bool
138
139
140 def make_deepep_ht_a2a(
141 pg: ProcessGroup,
142 pgi: ProcessGroupInfo,
143 dp_size: int,
144 ht_args: DeepEPHTArgs,
145 q_dtype: torch.dtype | None = None,
146 block_shape: list[int] | None = None,
147 ):
148 import deep_ep
149
150 # high throughput a2a
151 num_nvl_bytes = 1024 * 1024 * 1024 # 1GB
152 num_rdma_bytes, low_latency_mode, num_qps_per_rank = 0, False, 1
153 buffer = deep_ep.Buffer(
154 group=pg,
155 num_nvl_bytes=num_nvl_bytes,
156 num_rdma_bytes=num_rdma_bytes,
157 low_latency_mode=low_latency_mode,
158 num_qps_per_rank=num_qps_per_rank,
159 )
160 return DeepEPHTPrepareAndFinalize(
161 buffer=buffer,
162 num_dispatchers=pgi.world_size,
163 dp_size=dp_size,
164 rank_expert_offset=pgi.rank * ht_args.num_local_experts,
165 )
166
167
168 def make_deepep_ll_a2a(
169 pg: ProcessGroup,
170 pgi: ProcessGroupInfo,
171 deepep_ll_args: DeepEPLLArgs,
172 q_dtype: torch.dtype | None = None,
173 block_shape: list[int] | None = None,
174 ):
175 import deep_ep
176
177 # low-latency a2a
178 num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
179 deepep_ll_args.max_tokens_per_rank,
180 deepep_ll_args.hidden_size,
181 pgi.world_size,
182 deepep_ll_args.num_experts,
183 )
184
185 buffer = deep_ep.Buffer(
186 group=pg,
187 num_rdma_bytes=num_rdma_bytes,
188 low_latency_mode=True,
189 num_qps_per_rank=deepep_ll_args.num_experts // pgi.world_size,
190 )
191
192 return DeepEPLLPrepareAndFinalize(
193 buffer=buffer,
194 num_dispatchers=pgi.world_size,
195 max_tokens_per_rank=deepep_ll_args.max_tokens_per_rank,
196 use_fp8_dispatch=deepep_ll_args.use_fp8_dispatch,
197 )
198
199
200 def make_deepep_a2a(
201 pg: ProcessGroup,
202 pgi: ProcessGroupInfo,
203 dp_size: int,
204 deepep_ht_args: DeepEPHTArgs | None,
205 deepep_ll_args: DeepEPLLArgs | None,
206 q_dtype: torch.dtype | None = None,
207 block_shape: list[int] | None = None,
208 ):
209 if deepep_ht_args is not None:
210 assert deepep_ll_args is None
211 return make_deepep_ht_a2a(
212 pg, pgi, dp_size, deepep_ht_args, q_dtype, block_shape
213 )
214
215 assert deepep_ll_args is not None
216 return make_deepep_ll_a2a(pg, pgi, deepep_ll_args, q_dtype, block_shape)
217
218
219 @dataclasses.dataclass
220 class DeepEPV2Args:
221 num_local_experts: int
222 num_experts: int
223 num_topk: int
224 hidden_size: int
225 max_tokens_per_rank: int
226 use_fp8_dispatch: bool
227
228
229 def make_deepep_v2_a2a(
230 pg: ProcessGroup,
231 pgi: ProcessGroupInfo,
232 dp_size: int,
233 v2_args: DeepEPV2Args,
234 use_cudagraph: bool = False,
235 ):
236 import deep_ep
237
238 from vllm.utils.nccl import query_nccl_gin_type
239
240 # ElasticBuffer can segfault when GIN is unavailable. Initialize the
241 # lazy communicator and reject unsupported systems before entering DeepEP.
242 probe = torch.zeros(1, device=pgi.device)
243 torch.distributed.all_reduce(probe, group=pg)
244 gin_type = query_nccl_gin_type(pg)
245 if gin_type is None:
246 raise RuntimeError("Failed to determine NCCL GIN support")
247 if gin_type == 0:
248 raise GINNotAvailableError("NCCL GIN not available")
249
250 buffer = deep_ep.ElasticBuffer(
251 group=pg,
252 num_max_tokens_per_rank=v2_args.max_tokens_per_rank,
253 hidden=v2_args.hidden_size,
254 num_topk=v2_args.num_topk,
255 use_fp8_dispatch=v2_args.use_fp8_dispatch,
256 allow_hybrid_mode=False,
257 explicitly_destroy=True,
258 )
259 return DeepEPV2PrepareAndFinalize(
260 buffer=buffer,
261 num_dispatchers=pgi.world_size,
262 dp_size=dp_size,
263 rank_expert_offset=pgi.rank * v2_args.num_local_experts,
264 num_experts=v2_args.num_experts,
265 num_topk=v2_args.num_topk,
266 use_fp8_dispatch=v2_args.use_fp8_dispatch,
267 use_cudagraph=use_cudagraph,
268 )