@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 655 lines 22.3 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # This file is a pure Python wrapper for the NCCL library.
5 # The main purpose is to use NCCL combined with CUDA graph.
6 # Before writing this script, we tried the following approach:
7 # 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself
8 # often gets stuck when initializing the NCCL communicator.
9 # 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce`
10 # contains many other potential cuda APIs, that are not allowed during
11 # capturing the CUDA graph. For further details, please check
12 # https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ .
13 #
14 # Another rejected idea is to write a C/C++ binding for NCCL. It is usually
15 # doable, but we often encounter issues related with nccl versions, and need
16 # to switch between different versions of NCCL. See
17 # https://github.com/NVIDIA/nccl/issues/1234 for more details.
18 # A C/C++ binding is not flexible enough to handle this. It requires
19 # recompilation of the code every time we want to switch between different
20 # versions. This current implementation, with a **pure** Python wrapper, is
21 # more flexible. We can easily switch between different versions of NCCL by
22 # changing the environment variable `VLLM_NCCL_SO_PATH`, or the `so_file`
23 # variable in the code.
24
25 import ctypes
26 import functools
27 import platform
28 from dataclasses import dataclass
29 from typing import Any
30
31 import torch
32 from torch.distributed import ReduceOp
33
34 from vllm import envs
35 from vllm.logger import init_logger
36 from vllm.platforms import current_platform
37 from vllm.utils.nccl import find_nccl_library
38
39 logger = init_logger(__name__)
40
41 # === export types and functions from nccl to Python ===
42 # for the original nccl definition, please check
43 # https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in
44
45 ncclResult_t = ctypes.c_int
46 ncclComm_t = ctypes.c_void_p
47 ncclWindow_t = ctypes.c_void_p
48
49
50 class ncclUniqueId(ctypes.Structure):
51 _fields_ = [("internal", ctypes.c_byte * 128)]
52
53
54 NCCL_UNIQUE_ID_BYTES = ctypes.sizeof(ncclUniqueId)
55
56
57 # Mirror of NCCL's versioned ncclCommProperties_t (nccl_device/core.h,
58 # v2.31.2-1, 144 bytes). ncclCommQueryProperties fills fields gated by the
59 # version the caller declares in props.version, not by props.size, so never
60 # declare a version newer than this layout: clamp to
61 # NCCL_COMM_PROPERTIES_LAYOUT_VERSION. To use fields added in a newer NCCL,
62 # extend the layout and bump the constant.
63 NCCL_COMM_PROPERTIES_LAYOUT_VERSION = 23102 # NCCL_VERSION(2, 31, 2)
64
65
66 class ncclCommProperties(ctypes.Structure):
67 _fields_ = [
68 ("size", ctypes.c_size_t),
69 ("magic", ctypes.c_uint),
70 ("version", ctypes.c_uint),
71 ("rank", ctypes.c_int),
72 ("nRanks", ctypes.c_int),
73 ("cudaDev", ctypes.c_int),
74 ("nvmlDev", ctypes.c_int),
75 ("deviceApiSupport", ctypes.c_bool),
76 ("multimemSupport", ctypes.c_bool),
77 ("ginType", ctypes.c_int),
78 ("nLsaTeams", ctypes.c_int),
79 ("hostRmaSupport", ctypes.c_bool),
80 ("railedGinType", ctypes.c_int),
81 # Filled only when the declared version is >= NCCL_VERSION(2, 31, 0).
82 ("commHash", ctypes.c_uint64),
83 ("ginMinStride", ctypes.c_int),
84 ("ginConnectionType", ctypes.c_int),
85 ("ginSupport", ctypes.c_bool * 64),
86 ("devCommRuntimeVersionSize", ctypes.c_size_t),
87 ]
88
89
90 cudaStream_t = ctypes.c_void_p
91 buffer_type = ctypes.c_void_p
92
93 ncclDataType_t = ctypes.c_int
94
95
96 class ncclDataTypeEnum:
97 ncclInt8 = 0
98 ncclChar = 0
99 ncclUint8 = 1
100 ncclInt32 = 2
101 ncclInt = 2
102 ncclUint32 = 3
103 ncclInt64 = 4
104 ncclUint64 = 5
105 ncclFloat16 = 6
106 ncclHalf = 6
107 ncclFloat32 = 7
108 ncclFloat = 7
109 ncclFloat64 = 8
110 ncclDouble = 8
111 ncclBfloat16 = 9
112 ncclFloat8e4m3 = 10
113 ncclNumTypes = 11
114
115 @classmethod
116 @functools.lru_cache(maxsize=1)
117 def _torch_to_nccl_map(cls) -> dict[torch.dtype, int]:
118 return {
119 torch.int8: cls.ncclInt8,
120 torch.uint8: cls.ncclUint8,
121 torch.int32: cls.ncclInt32,
122 torch.int64: cls.ncclInt64,
123 torch.float16: cls.ncclFloat16,
124 torch.float32: cls.ncclFloat32,
125 torch.float64: cls.ncclFloat64,
126 torch.bfloat16: cls.ncclBfloat16,
127 current_platform.fp8_dtype(): cls.ncclFloat8e4m3,
128 }
129
130 @classmethod
131 def supports_torch_dtype(cls, dtype: torch.dtype) -> bool:
132 return dtype in cls._torch_to_nccl_map()
133
134 @classmethod
135 def try_from_torch(cls, dtype: torch.dtype) -> int | None:
136 return cls._torch_to_nccl_map().get(dtype)
137
138 @classmethod
139 def from_torch(cls, dtype: torch.dtype) -> int:
140 nccl_dtype = cls.try_from_torch(dtype)
141 if nccl_dtype is not None:
142 return nccl_dtype
143 raise ValueError(
144 f"Unsupported dtype {dtype}: should be one of "
145 f"int8, uint8, int32, int64, float16, float32, float64, bfloat16,"
146 " float8e4m3."
147 )
148
149
150 ncclRedOp_t = ctypes.c_int
151
152
153 class ncclRedOpTypeEnum:
154 ncclSum = 0
155 ncclProd = 1
156 ncclMax = 2
157 ncclMin = 3
158 ncclAvg = 4
159 ncclNumOps = 5
160
161 @classmethod
162 def from_torch(cls, op: ReduceOp) -> int:
163 if op == ReduceOp.SUM:
164 return cls.ncclSum
165 if op == ReduceOp.PRODUCT:
166 return cls.ncclProd
167 if op == ReduceOp.MAX:
168 return cls.ncclMax
169 if op == ReduceOp.MIN:
170 return cls.ncclMin
171 if op == ReduceOp.AVG:
172 return cls.ncclAvg
173 raise ValueError(f"Unsupported op: {op}")
174
175
176 @dataclass
177 class Function:
178 name: str
179 restype: Any
180 argtypes: list[Any]
181
182
183 class NCCLLibrary:
184 exported_functions = [
185 # const char* ncclGetErrorString(ncclResult_t result)
186 Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]),
187 # ncclResult_t ncclGetVersion(int *version);
188 Function("ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)]),
189 # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId);
190 Function("ncclGetUniqueId", ncclResult_t, [ctypes.POINTER(ncclUniqueId)]),
191 # ncclResult_t ncclCommInitRank(
192 # ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
193 # note that ncclComm_t is a pointer type, so the first argument
194 # is a pointer to a pointer
195 Function(
196 "ncclCommInitRank",
197 ncclResult_t,
198 [ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int],
199 ),
200 # ncclResult_t ncclAllReduce(
201 # const void* sendbuff, void* recvbuff, size_t count,
202 # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
203 # cudaStream_t stream);
204 # note that cudaStream_t is a pointer type, so the last argument
205 # is a pointer
206 Function(
207 "ncclAllReduce",
208 ncclResult_t,
209 [
210 buffer_type,
211 buffer_type,
212 ctypes.c_size_t,
213 ncclDataType_t,
214 ncclRedOp_t,
215 ncclComm_t,
216 cudaStream_t,
217 ],
218 ),
219 # ncclResult_t ncclReduce(
220 # const void* sendbuff, void* recvbuff, size_t count,
221 # ncclDataType_t datatype, ncclRedOp_t op, int root,
222 # ncclComm_t comm, cudaStream_t stream);
223 # note that cudaStream_t is a pointer type, so the last argument
224 # is a pointer
225 Function(
226 "ncclReduce",
227 ncclResult_t,
228 [
229 buffer_type,
230 buffer_type,
231 ctypes.c_size_t,
232 ncclDataType_t,
233 ncclRedOp_t,
234 ctypes.c_int,
235 ncclComm_t,
236 cudaStream_t,
237 ],
238 ),
239 # ncclResult_t ncclAllGather(
240 # const void* sendbuff, void* recvbuff, size_t count,
241 # ncclDataType_t datatype, ncclComm_t comm,
242 # cudaStream_t stream);
243 # note that cudaStream_t is a pointer type, so the last argument
244 # is a pointer
245 Function(
246 "ncclAllGather",
247 ncclResult_t,
248 [
249 buffer_type,
250 buffer_type,
251 ctypes.c_size_t,
252 ncclDataType_t,
253 ncclComm_t,
254 cudaStream_t,
255 ],
256 ),
257 # ncclResult_t ncclReduceScatter(
258 # const void* sendbuff, void* recvbuff, size_t count,
259 # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
260 # cudaStream_t stream);
261 # note that cudaStream_t is a pointer type, so the last argument
262 # is a pointer
263 Function(
264 "ncclReduceScatter",
265 ncclResult_t,
266 [
267 buffer_type,
268 buffer_type,
269 ctypes.c_size_t,
270 ncclDataType_t,
271 ncclRedOp_t,
272 ncclComm_t,
273 cudaStream_t,
274 ],
275 ),
276 # ncclResult_t ncclSend(
277 # const void* sendbuff, size_t count, ncclDataType_t datatype,
278 # int dest, ncclComm_t comm, cudaStream_t stream);
279 Function(
280 "ncclSend",
281 ncclResult_t,
282 [
283 buffer_type,
284 ctypes.c_size_t,
285 ncclDataType_t,
286 ctypes.c_int,
287 ncclComm_t,
288 cudaStream_t,
289 ],
290 ),
291 # ncclResult_t ncclRecv(
292 # void* recvbuff, size_t count, ncclDataType_t datatype,
293 # int src, ncclComm_t comm, cudaStream_t stream);
294 Function(
295 "ncclRecv",
296 ncclResult_t,
297 [
298 buffer_type,
299 ctypes.c_size_t,
300 ncclDataType_t,
301 ctypes.c_int,
302 ncclComm_t,
303 cudaStream_t,
304 ],
305 ),
306 # ncclResult_t ncclBroadcast(
307 # const void* sendbuff, void* recvbuff, size_t count,
308 # ncclDataType_t datatype, int root, ncclComm_t comm,
309 # cudaStream_t stream);
310 Function(
311 "ncclBroadcast",
312 ncclResult_t,
313 [
314 buffer_type,
315 buffer_type,
316 ctypes.c_size_t,
317 ncclDataType_t,
318 ctypes.c_int,
319 ncclComm_t,
320 cudaStream_t,
321 ],
322 ),
323 # be cautious! this is a collective call, it will block until all
324 # processes in the communicator have called this function.
325 # because Python object destruction can happen in random order,
326 # it is better not to call it at all.
327 # ncclResult_t ncclCommDestroy(ncclComm_t comm);
328 Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]),
329 # ncclCommAbort frees resources associated with the communicator
330 # without requiring a collective synchronization. Unlike
331 # ncclCommDestroy, it is safe to call during an uncoordinated
332 # shutdown when peer ranks may already be gone.
333 # ncclResult_t ncclCommAbort(ncclComm_t comm);
334 Function("ncclCommAbort", ncclResult_t, [ncclComm_t]),
335 # ncclResult_t ncclCommSuspend(ncclComm_t comm, int flags);
336 Function("ncclCommSuspend", ncclResult_t, [ncclComm_t, ctypes.c_int]),
337 # ncclResult_t ncclCommResume(ncclComm_t comm);
338 Function("ncclCommResume", ncclResult_t, [ncclComm_t]),
339 # ncclResult_t ncclGroupStart();
340 Function("ncclGroupStart", ncclResult_t, []),
341 # ncclResult_t ncclGroupEnd();
342 Function("ncclGroupEnd", ncclResult_t, []),
343 # ncclResult_t ncclCommWindowRegister(
344 # ncclComm_t comm, void* buff, size_t size,
345 # ncclWindow_t* win, int winFlags);
346 Function(
347 "ncclCommWindowRegister",
348 ncclResult_t,
349 [
350 ncclComm_t,
351 buffer_type,
352 ctypes.c_size_t,
353 ctypes.POINTER(ncclWindow_t),
354 ctypes.c_int,
355 ],
356 ),
357 # ncclResult_t ncclCommWindowDeregister(
358 # ncclComm_t comm, ncclWindow_t win);
359 Function("ncclCommWindowDeregister", ncclResult_t, [ncclComm_t, ncclWindow_t]),
360 # Query runtime properties of a specific initialized communicator.
361 Function(
362 "ncclCommQueryProperties",
363 ncclResult_t,
364 [ncclComm_t, ctypes.POINTER(ncclCommProperties)],
365 ),
366 ]
367
368 # class attribute to store the mapping from the path to the library
369 # to avoid loading the same library multiple times
370 path_to_library_cache: dict[str, Any] = {}
371
372 # class attribute to store the mapping from library path
373 # to the corresponding dictionary
374 path_to_dict_mapping: dict[str, dict[str, Any]] = {}
375
376 def __init__(self, so_file: str | None = None):
377 so_file = so_file or find_nccl_library()
378
379 try:
380 if so_file not in NCCLLibrary.path_to_dict_mapping:
381 lib = ctypes.CDLL(so_file)
382 NCCLLibrary.path_to_library_cache[so_file] = lib
383 self.lib = NCCLLibrary.path_to_library_cache[so_file]
384 except Exception as e:
385 logger.error(
386 "Failed to load NCCL library from %s. "
387 "It is expected if you are not running on NVIDIA/AMD GPUs."
388 "Otherwise, the nccl library might not exist, be corrupted "
389 "or it does not support the current platform %s. "
390 "If you already have the library, please set the "
391 "environment variable VLLM_NCCL_SO_PATH"
392 " to point to the correct nccl library path.",
393 so_file,
394 platform.platform(),
395 )
396 raise e
397
398 if so_file not in NCCLLibrary.path_to_dict_mapping:
399 _funcs: dict[str, Any] = {}
400 for func in NCCLLibrary.exported_functions:
401 try:
402 f = getattr(self.lib, func.name)
403 f.restype = func.restype
404 f.argtypes = func.argtypes
405 _funcs[func.name] = f
406 except AttributeError:
407 if func.name in [
408 "ncclCommWindowRegister",
409 "ncclCommWindowDeregister",
410 ]:
411 if envs.VLLM_USE_NCCL_SYMM_MEM:
412 logger.warning_once(
413 "The symbol %s is not found in the NCCL "
414 "library %s. To enable VLLM_USE_NCCL_SYMM_MEM "
415 " please update your NCCL version to >= "
416 "2.27.03.",
417 func.name,
418 so_file,
419 )
420 if current_platform.is_rocm():
421 # Having an exception here on ROCm platform is
422 # not allowed during graph capturing
423 continue
424 elif func.name == "ncclCommQueryProperties":
425 # Optional on NCCL versions older than 2.29.
426 continue
427 elif func.name in ("ncclCommSuspend", "ncclCommResume"):
428 # RCCL doesn't export these; NCCL >= 2.29.7 does, and
429 # vLLM's CUDA path already requires that version.
430 # PyNcclCommunicator checks has_symbol() before
431 # calling either.
432 if current_platform.is_rocm():
433 continue
434 raise
435 NCCLLibrary.path_to_dict_mapping[so_file] = _funcs
436 self._funcs = NCCLLibrary.path_to_dict_mapping[so_file]
437
438 def has_symbol(self, name: str) -> bool:
439 """Whether the loaded NCCL/RCCL library exports the given symbol."""
440 return name in self._funcs
441
442 def ncclGetErrorString(self, result: ncclResult_t) -> str:
443 return self._funcs["ncclGetErrorString"](result).decode("utf-8")
444
445 def NCCL_CHECK(self, result: ncclResult_t) -> None:
446 if result != 0:
447 error_str = self.ncclGetErrorString(result)
448 raise RuntimeError(f"NCCL error: {error_str}")
449
450 def ncclGetRawVersion(self) -> int:
451 version = ctypes.c_int()
452 self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version)))
453 # something like 21903
454 return version.value
455
456 def ncclGetVersion(self) -> str:
457 version_str = str(self.ncclGetRawVersion())
458 # something like 21903 --> "2.19.3"
459 major = version_str[0].lstrip("0")
460 minor = version_str[1:3].lstrip("0")
461 patch = version_str[3:].lstrip("0")
462 return f"{major}.{minor}.{patch}"
463
464 def ncclGetUniqueId(self) -> ncclUniqueId:
465 unique_id = ncclUniqueId()
466 self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id)))
467 return unique_id
468
469 def unique_id_from_bytes(self, data: bytes) -> ncclUniqueId:
470 if len(data) != 128:
471 raise ValueError(
472 f"Expected 128 bytes for ncclUniqueId, got {len(data)} bytes"
473 )
474 unique_id = ncclUniqueId()
475 ctypes.memmove(ctypes.addressof(unique_id.internal), data, 128)
476 return unique_id
477
478 def ncclCommInitRank(
479 self, world_size: int, unique_id: ncclUniqueId, rank: int
480 ) -> ncclComm_t:
481 comm = ncclComm_t()
482 self.NCCL_CHECK(
483 self._funcs["ncclCommInitRank"](
484 ctypes.byref(comm), world_size, unique_id, rank
485 )
486 )
487 return comm
488
489 def ncclAllReduce(
490 self,
491 sendbuff: buffer_type,
492 recvbuff: buffer_type,
493 count: int,
494 datatype: int,
495 op: int,
496 comm: ncclComm_t,
497 stream: cudaStream_t,
498 ) -> None:
499 # `datatype` actually should be `ncclDataType_t`
500 # and `op` should be `ncclRedOp_t`
501 # both are aliases of `ctypes.c_int`
502 # when we pass int to a function, it will be converted to `ctypes.c_int`
503 # by ctypes automatically
504 self.NCCL_CHECK(
505 self._funcs["ncclAllReduce"](
506 sendbuff, recvbuff, count, datatype, op, comm, stream
507 )
508 )
509
510 def ncclReduce(
511 self,
512 sendbuff: buffer_type,
513 recvbuff: buffer_type,
514 count: int,
515 datatype: int,
516 op: int,
517 root: int,
518 comm: ncclComm_t,
519 stream: cudaStream_t,
520 ) -> None:
521 # `datatype` actually should be `ncclDataType_t`
522 # and `op` should be `ncclRedOp_t`
523 # both are aliases of `ctypes.c_int`
524 # when we pass int to a function, it will be converted to `ctypes.c_int`
525 # by ctypes automatically
526 self.NCCL_CHECK(
527 self._funcs["ncclReduce"](
528 sendbuff, recvbuff, count, datatype, op, root, comm, stream
529 )
530 )
531
532 def ncclReduceScatter(
533 self,
534 sendbuff: buffer_type,
535 recvbuff: buffer_type,
536 count: int,
537 datatype: int,
538 op: int,
539 comm: ncclComm_t,
540 stream: cudaStream_t,
541 ) -> None:
542 # `datatype` actually should be `ncclDataType_t`
543 # and `op` should be `ncclRedOp_t`
544 # both are aliases of `ctypes.c_int`
545 # when we pass int to a function, it will be converted to `ctypes.c_int`
546 # by ctypes automatically
547 self.NCCL_CHECK(
548 self._funcs["ncclReduceScatter"](
549 sendbuff, recvbuff, count, datatype, op, comm, stream
550 )
551 )
552
553 def ncclAllGather(
554 self,
555 sendbuff: buffer_type,
556 recvbuff: buffer_type,
557 count: int,
558 datatype: int,
559 comm: ncclComm_t,
560 stream: cudaStream_t,
561 ) -> None:
562 # `datatype` actually should be `ncclDataType_t`
563 # which is an aliases of `ctypes.c_int`
564 # when we pass int to a function, it will be converted to `ctypes.c_int`
565 # by ctypes automatically
566 self.NCCL_CHECK(
567 self._funcs["ncclAllGather"](
568 sendbuff, recvbuff, count, datatype, comm, stream
569 )
570 )
571
572 def ncclSend(
573 self,
574 sendbuff: buffer_type,
575 count: int,
576 datatype: int,
577 dest: int,
578 comm: ncclComm_t,
579 stream: cudaStream_t,
580 ) -> None:
581 self.NCCL_CHECK(
582 self._funcs["ncclSend"](sendbuff, count, datatype, dest, comm, stream)
583 )
584
585 def ncclRecv(
586 self,
587 recvbuff: buffer_type,
588 count: int,
589 datatype: int,
590 src: int,
591 comm: ncclComm_t,
592 stream: cudaStream_t,
593 ) -> None:
594 self.NCCL_CHECK(
595 self._funcs["ncclRecv"](recvbuff, count, datatype, src, comm, stream)
596 )
597
598 def ncclBroadcast(
599 self,
600 sendbuff: buffer_type,
601 recvbuff: buffer_type,
602 count: int,
603 datatype: int,
604 root: int,
605 comm: ncclComm_t,
606 stream: cudaStream_t,
607 ) -> None:
608 self.NCCL_CHECK(
609 self._funcs["ncclBroadcast"](
610 sendbuff, recvbuff, count, datatype, root, comm, stream
611 )
612 )
613
614 def ncclCommDestroy(self, comm: ncclComm_t) -> None:
615 self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm))
616
617 def ncclCommAbort(self, comm: ncclComm_t) -> None:
618 self.NCCL_CHECK(self._funcs["ncclCommAbort"](comm))
619
620 def ncclCommSuspend(self, comm: ncclComm_t, flags: int) -> None:
621 self.NCCL_CHECK(self._funcs["ncclCommSuspend"](comm, flags))
622
623 def ncclCommResume(self, comm: ncclComm_t) -> None:
624 self.NCCL_CHECK(self._funcs["ncclCommResume"](comm))
625
626 def ncclGroupStart(self) -> None:
627 self.NCCL_CHECK(self._funcs["ncclGroupStart"]())
628
629 def ncclGroupEnd(self) -> None:
630 self.NCCL_CHECK(self._funcs["ncclGroupEnd"]())
631
632 def ncclCommWindowRegister(
633 self, comm: ncclComm_t, buff: buffer_type, size: int, win_flags: int
634 ) -> ncclWindow_t:
635 window = ncclWindow_t()
636 self.NCCL_CHECK(
637 self._funcs["ncclCommWindowRegister"](
638 comm, buff, size, ctypes.byref(window), win_flags
639 )
640 )
641 return window
642
643 def ncclCommWindowDeregister(self, comm: ncclComm_t, window: ncclWindow_t) -> None:
644 self.NCCL_CHECK(self._funcs["ncclCommWindowDeregister"](comm, window))
645
646
647 __all__ = [
648 "NCCLLibrary",
649 "ncclDataTypeEnum",
650 "ncclRedOpTypeEnum",
651 "ncclUniqueId",
652 "ncclComm_t",
653 "cudaStream_t",
654 "buffer_type",
655 ]