"tests/compile/fullgraph/test_simple.py" did not exist on "74f441f4b517a895ad12afd314a6f40caf657c4e"
parallel_state.py 59.1 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

Woosuk Kwon's avatar
Woosuk Kwon committed
4
# Copyright 2023 The vLLM team.
5
6
# Adapted from
# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py
Zhuohan Li's avatar
Zhuohan Li committed
7
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
8
9
10
11
12
"""vLLM distributed state.
It takes over the control of the distributed environment from PyTorch.
The typical workflow is:

- call `init_distributed_environment` to initialize the distributed environment.
13
- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to
14
15
16
17
18
19
20
21
22
23
24
 initialize the model parallel groups.

- any code dealing with the distributed stuff

- call `destroy_model_parallel` to destroy the model parallel groups.
- call `destroy_distributed_environment` to destroy the distributed environment.

If you only need to use the distributed environment without model/pipeline
 parallelism, you can skip the model parallel initialization and destruction
 steps.
"""
25
import contextlib
26
import gc
27
import pickle
28
import weakref
29
30
31
from collections import namedtuple
from contextlib import contextmanager, nullcontext
from dataclasses import dataclass
32
from datetime import timedelta
33
from multiprocessing import shared_memory
34
from typing import Any, Callable, Optional, Union
35
from unittest.mock import patch
Zhuohan Li's avatar
Zhuohan Li committed
36
37

import torch
38
import torch.distributed
39
from torch.distributed import Backend, ProcessGroup
40
from typing_extensions import deprecated
Zhuohan Li's avatar
Zhuohan Li committed
41

42
import vllm.envs as envs
43
44
from vllm.distributed.device_communicators.base_device_communicator import (
    DeviceCommunicatorBase)
45
from vllm.distributed.utils import StatelessProcessGroup
46
from vllm.logger import init_logger
47
48
from vllm.utils import (direct_register_custom_op, get_distributed_init_method,
                        resolve_obj_by_qualname, supports_custom_op)
49
50


51
52
53
@dataclass
class GraphCaptureContext:
    stream: torch.cuda.Stream
54

55

56
TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
57

58
59

def _split_tensor_dict(
60
61
    tensor_dict: dict[str, Union[torch.Tensor, Any]]
) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:
62
63
64
65
66
    """Split the tensor dictionary into two parts:
    1. A list of (key, value) pairs. If the value is a tensor, it is replaced
         by its metadata.
    2. A list of tensors.
    """
67
68
    metadata_list: list[tuple[str, Any]] = []
    tensor_list: list[torch.Tensor] = []
69
70
71
72
73
74
    for key, value in tensor_dict.items():
        if isinstance(value, torch.Tensor):
            # Note: we cannot use `value.device` here,
            # because it contains not only the device type but also the device
            # index (e.g. "cuda:0"). We only need the device type.
            # receiving side will set the device index.
75
            device = value.device.type
76
            metadata_list.append(
77
                (key, TensorMetadata(device, value.dtype, value.size())))
78
79
            tensor_list.append(value)
        else:
80
            metadata_list.append((key, value))
81
82
83
    return metadata_list, tensor_list


84
_group_name_counter: dict[str, int] = {}
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99


def _get_unique_name(name: str) -> str:
    """Get a unique name for the group.
    Example:
    _get_unique_name("tp") -> "tp:0"
    _get_unique_name("tp") -> "tp:1"
    """
    if name not in _group_name_counter:
        _group_name_counter[name] = 0
    newname = f"{name}:{_group_name_counter[name]}"
    _group_name_counter[name] += 1
    return newname


100
_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}
101
102
103


def _register_group(group: "GroupCoordinator") -> None:
104
    _groups[group.unique_name] = weakref.ref(group)
105
106


107
108
109
110
111
112
def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
    assert group_name in _groups, f"Group {group_name} is not found."
    group = _groups[group_name]()
    if group is None:
        raise ValueError(f"Group {group_name} is destroyed.")
    return group._all_reduce_out_place(tensor)
113
114


115
116
def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
    return torch.empty_like(tensor)
117

118

119
120
121
122
123
124
def reduce_scatter(tensor: torch.Tensor, dim: int, world_size: int,
                   group_name: str) -> torch.Tensor:
    assert group_name in _groups, f"Group {group_name} is not found."
    group = _groups[group_name]()
    if group is None:
        raise ValueError(f"Group {group_name} is destroyed.")
125
    return group._reduce_scatter_out_place(tensor, dim)
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140


def reduce_scatter_fake(tensor: torch.Tensor, dim: int, world_size: int,
                        group_name: str) -> torch.Tensor:
    new_shape = list(tensor.shape)
    new_shape[dim] = tensor.shape[dim] // world_size
    return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device)


def all_gather(tensor: torch.Tensor, dim: int, world_size: int,
               group_name: str) -> torch.Tensor:
    assert group_name in _groups, f"Group {group_name} is not found."
    group = _groups[group_name]()
    if group is None:
        raise ValueError(f"Group {group_name} is destroyed.")
141
    return group._all_gather_out_place(tensor, dim)
142
143
144
145
146
147
148
149
150


def all_gather_fake(tensor: torch.Tensor, dim: int, world_size: int,
                    group_name: str) -> torch.Tensor:
    new_shape = list(tensor.shape)
    new_shape[dim] = tensor.shape[dim] * world_size
    return torch.empty(new_shape, dtype=tensor.dtype, device=tensor.device)


151
if supports_custom_op():
152
    direct_register_custom_op(
153
154
155
        op_name="all_reduce",
        op_func=all_reduce,
        fake_impl=all_reduce_fake,
156
157
    )

158
159
160
161
162
163
164
165
166
167
168
169
    direct_register_custom_op(
        op_name="reduce_scatter",
        op_func=reduce_scatter,
        fake_impl=reduce_scatter_fake,
    )

    direct_register_custom_op(
        op_name="all_gather",
        op_func=all_gather,
        fake_impl=all_gather_fake,
    )

170

171
172
173
174
175
176
class GroupCoordinator:
    """
    PyTorch ProcessGroup wrapper for a group of processes.
    PyTorch ProcessGroup is bound to one specific communication backend,
        e.g. NCCL, Gloo, MPI, etc.
    GroupCoordinator takes charge of all the communication operations among
177
178
        the processes in the group. It manages both CPU and device
        communication.
179
180
181
182
    """

    # available attributes:
    rank: int  # global rank
183
    ranks: list[int]  # global ranks in the group
184
185
186
187
188
189
190
191
192
193
    world_size: int  # size of the group
    # difference between `local_rank` and `rank_in_group`:
    # if we have a group of size 4 across two nodes:
    # Process | Node | Rank | Local Rank | Rank in Group
    #   0     |   0  |  0   |     0      |       0
    #   1     |   0  |  1   |     1      |       1
    #   2     |   1  |  2   |     0      |       2
    #   3     |   1  |  3   |     1      |       3
    local_rank: int  # local rank used to assign devices
    rank_in_group: int  # rank inside the group
194
195
196
197
    cpu_group: ProcessGroup  # group for CPU communication
    device_group: ProcessGroup  # group for device communication
    # device communicator (if use_device_communicator=True)
    device_communicator: Optional[DeviceCommunicatorBase]
198
    mq_broadcaster: Optional[Any]  # shared memory broadcaster
199
200
201

    def __init__(
        self,
202
        group_ranks: list[list[int]],
203
204
        local_rank: int,
        torch_distributed_backend: Union[str, Backend],
205
        use_device_communicator: bool,  # whether to use device communicator
206
        use_message_queue_broadcaster: bool = False,
207
        group_name: Optional[str] = None,
208
    ):
209
210
211
        group_name = group_name or "anonymous"
        self.unique_name = _get_unique_name(group_name)
        _register_group(self)
212
213
214

        self.rank = torch.distributed.get_rank()
        self.local_rank = local_rank
215
216
217

        self_device_group = None
        self_cpu_group = None
218
219
220
221
222
223
224
225
226
227
228

        for ranks in group_ranks:
            device_group = torch.distributed.new_group(
                ranks, backend=torch_distributed_backend)
            # a group with `gloo` backend, to allow direct coordination between
            # processes through the CPU.
            cpu_group = torch.distributed.new_group(ranks, backend="gloo")
            if self.rank in ranks:
                self.ranks = ranks
                self.world_size = len(ranks)
                self.rank_in_group = ranks.index(self.rank)
229
230
231
232
233
                self_device_group = device_group
                self_cpu_group = cpu_group

        assert self_cpu_group is not None
        assert self_device_group is not None
234

235
236
        self.cpu_group = self_cpu_group
        self.device_group = self_device_group
237

238
        from vllm.platforms import current_platform
239

240
        if current_platform.is_cuda_alike():
241
            self.device = torch.device(f"cuda:{local_rank}")
242
243
        elif current_platform.is_xpu():
            self.device = torch.device(f"xpu:{local_rank}")
244
245
246
        elif current_platform.is_out_of_tree():
            self.device = torch.device(
                f"{current_platform.device_name}:{local_rank}")
247
248
249
        else:
            self.device = torch.device("cpu")

250
        self.use_device_communicator = use_device_communicator
251
        self.device_communicator = None
252
253
254
255
256
        if use_device_communicator and self.world_size > 1:
            device_comm_cls = resolve_obj_by_qualname(
                current_platform.get_device_communicator_cls())
            self.device_communicator = device_comm_cls(
                cpu_group=self.cpu_group,
257
                device=self.device,
258
259
                device_group=self.device_group,
                unique_name=self.unique_name,
260
261
            )

262
        from vllm.distributed.device_communicators.shm_broadcast import (
263
264
265
266
            MessageQueue)
        self.mq_broadcaster: Optional[MessageQueue] = None
        if use_message_queue_broadcaster and self.world_size > 1:
            self.mq_broadcaster = MessageQueue.create_from_process_group(
267
                self.cpu_group, 1 << 22, 6)
268

269
        from vllm.platforms import current_platform
270
271
        self.use_custom_op_call = (current_platform.is_cuda_alike()
                                   or current_platform.is_tpu())
272

273
274
275
        self.use_cpu_custom_send_recv = (current_platform.is_cpu() and hasattr(
            torch.ops._C, "init_shm_manager"))

276
277
278
279
280
281
282
283
284
285
    @property
    def first_rank(self):
        """Return the global rank of the first process in the group"""
        return self.ranks[0]

    @property
    def last_rank(self):
        """Return the global rank of the last process in the group"""
        return self.ranks[-1]

286
287
288
289
290
291
292
293
294
295
    @property
    def is_first_rank(self):
        """Return whether the caller is the first process in the group"""
        return self.rank == self.first_rank

    @property
    def is_last_rank(self):
        """Return whether the caller is the last process in the group"""
        return self.rank == self.last_rank

296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
    @property
    def next_rank(self):
        """Return the global rank of the process that follows the caller"""
        rank_in_group = self.rank_in_group
        world_size = self.world_size
        return self.ranks[(rank_in_group + 1) % world_size]

    @property
    def prev_rank(self):
        """Return the global rank of the process that precedes the caller"""
        rank_in_group = self.rank_in_group
        world_size = self.world_size
        return self.ranks[(rank_in_group - 1) % world_size]

    @contextmanager
    def graph_capture(
            self, graph_capture_context: Optional[GraphCaptureContext] = None):
        if graph_capture_context is None:
            stream = torch.cuda.Stream()
            graph_capture_context = GraphCaptureContext(stream)
        else:
            stream = graph_capture_context.stream

319
320
321
322
323
324
325
326
327
328
        # only cuda uses this function,
        # so we don't abstract it into the base class
        maybe_ca_context = nullcontext()
        from vllm.distributed.device_communicators.cuda_communicator import (
            CudaCommunicator)
        if self.device_communicator is not None:
            assert isinstance(self.device_communicator, CudaCommunicator)
            ca_comm = self.device_communicator.ca_comm
            if ca_comm is not None:
                maybe_ca_context = ca_comm.capture()  # type: ignore
329
330
331
332
333
334
335

        # ensure all initialization operations complete before attempting to
        # capture the graph on another stream
        curr_stream = torch.cuda.current_stream()
        if curr_stream != stream:
            stream.wait_stream(curr_stream)

336
        with torch.cuda.stream(stream), maybe_ca_context:
337
            yield graph_capture_context
338
339
340

    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
        """
341
342
343
344
345
346
347
348
349
350
        User-facing all-reduce function before we actually call the
        all-reduce operation.

        We need this because Dynamo does not support passing an arbitrary
        object (`self` in this case) to a custom op. We need to pass the
         group name as a string, and then look up the group coordinator from
         the group name, dispatch the all-reduce operation to the group
         coordinator.

        In addition, PyTorch custom ops do not support mutation or returning
351
352
        a new tensor in the same op. So we always make the all-reduce operation
        out-of-place.
353
354
355
356
357
        """
        # Bypass the function if we are using only 1 GPU.
        if self.world_size == 1:
            return input_

358
359
360
361
362
        if self.use_custom_op_call:
            return torch.ops.vllm.all_reduce(input_,
                                             group_name=self.unique_name)
        else:
            return self._all_reduce_out_place(input_)
363

364
    def _all_reduce_out_place(self, input_: torch.Tensor) -> torch.Tensor:
365
366
        if self.device_communicator is None:
            raise ValueError("No device communicator found")
367
        return self.device_communicator.all_reduce(input_)
368
369
370
371
372
373
374
375

    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
        world_size = self.world_size
        # Bypass the function if we are using only 1 GPU.
        if world_size == 1:
            return input_
        assert -input_.dim() <= dim < input_.dim(), (
            f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
376

377
378
379
380
381
382
383
384
385
386
        if self.use_custom_op_call:
            return torch.ops.vllm.all_gather(input_,
                                             dim,
                                             world_size,
                                             group_name=self.unique_name)
        else:
            return self._all_gather_out_place(input_, dim)

    def _all_gather_out_place(self, input_: torch.Tensor,
                              dim: int) -> torch.Tensor:
387
388
        if self.device_communicator is None:
            raise ValueError("No device communicator found")
389
        return self.device_communicator.all_gather(input_, dim)
390

391
392
393
394
    def all_gatherv(self,
                    input_: Union[torch.Tensor, list[torch.Tensor]],
                    dim: int = 0,
                    sizes: Optional[list[int]] = None):
395
396
        if self.device_communicator is None:
            raise ValueError("No device communicator found")
397
398
        return self.device_communicator.all_gatherv(input_, dim, sizes)

399
400
401
402
403
404
405
406
407
408
    def reduce_scatter(self,
                       input_: torch.Tensor,
                       dim: int = -1) -> torch.Tensor:
        world_size = self.world_size
        # Bypass the function if we are using only 1 GPU.
        if world_size == 1:
            return input_
        assert -input_.dim() <= dim < input_.dim(), (
            f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")

409
410
411
412
413
414
415
416
        if self.use_custom_op_call:
            return torch.ops.vllm.reduce_scatter(input_,
                                                 dim,
                                                 world_size,
                                                 group_name=self.unique_name)
        else:
            return self._reduce_scatter_out_place(input_, dim)

417
418
419
420
    def reduce_scatterv(self,
                        input_: torch.Tensor,
                        dim: int = -1,
                        sizes: Optional[list[int]] = None) -> torch.Tensor:
421
422
        if self.device_communicator is None:
            raise ValueError("No device communicator found")
423
424
        return self.device_communicator.reduce_scatterv(input_, dim, sizes)

425
426
    def _reduce_scatter_out_place(self, input_: torch.Tensor,
                                  dim: int) -> torch.Tensor:
427
428
        if self.device_communicator is None:
            raise ValueError("No device communicator found")
429
430
        return self.device_communicator.reduce_scatter(input_, dim)

431
432
433
    def gather(self,
               input_: torch.Tensor,
               dst: int = 0,
434
               dim: int = -1) -> Optional[torch.Tensor]:
435
436
437
438
439
440
441
442
443
        """
        NOTE: We assume that the input tensor is on the same device across
        all the ranks.
        NOTE: `dst` is the local rank of the destination rank.
        """
        world_size = self.world_size
        # Bypass the function if we are using only 1 GPU.
        if world_size == 1:
            return input_
444
445
        if self.device_communicator is None:
            raise ValueError("No device communicator found")
446
        return self.device_communicator.gather(input_, dst, dim)
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462

    def broadcast(self, input_: torch.Tensor, src: int = 0):
        """Broadcast the input tensor.
        NOTE: `src` is the local rank of the source rank.
        """
        assert src < self.world_size, f"Invalid src rank ({src})"

        # Bypass the function if we are using only 1 GPU.
        if self.world_size == 1:
            return input_
        # Broadcast.
        torch.distributed.broadcast(input_,
                                    src=self.ranks[src],
                                    group=self.device_group)
        return input_

463
464
465
466
467
468
469
470
471
    def broadcast_object(self, obj: Optional[Any] = None, src: int = 0):
        """Broadcast the input object.
        NOTE: `src` is the local rank of the source rank.
        """
        assert src < self.world_size, f"Invalid src rank ({src})"

        # Bypass the function if we are using only 1 GPU.
        if self.world_size == 1:
            return obj
472
473
474
        if self.mq_broadcaster is not None:
            assert src == 0, "Message queue broadcaster only supports src=0"
            return self.mq_broadcaster.broadcast_object(obj)
475
476
477
478
479
480
481
482
483
484
485
486
        if self.rank_in_group == src:
            torch.distributed.broadcast_object_list([obj],
                                                    src=self.ranks[src],
                                                    group=self.cpu_group)
            return obj
        else:
            recv = [None]
            torch.distributed.broadcast_object_list(recv,
                                                    src=self.ranks[src],
                                                    group=self.cpu_group)
            return recv[0]

487
    def broadcast_object_list(self,
488
                              obj_list: list[Any],
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
                              src: int = 0,
                              group: Optional[ProcessGroup] = None):
        """Broadcast the input object list.
        NOTE: `src` is the local rank of the source rank.
        """
        assert src < self.world_size, f"Invalid src rank ({src})"

        # Bypass the function if we are using only 1 GPU.
        if self.world_size == 1:
            return obj_list
        # Broadcast.
        torch.distributed.broadcast_object_list(obj_list,
                                                src=self.ranks[src],
                                                group=self.device_group)
        return obj_list

505
506
507
508
509
510
    def send_object(self, obj: Any, dst: int) -> None:
        """Send the input object list to the destination rank."""
        """NOTE: `dst` is the local rank of the destination rank."""

        assert dst < self.world_size, f"Invalid dst rank ({dst})"

511
        assert dst != self.rank_in_group, (
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
            "Invalid destination rank. Destination rank is the same "
            "as the current rank.")

        # Serialize object to tensor and get the size as well
        object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8)

        size_tensor = torch.tensor([object_tensor.numel()],
                                   dtype=torch.long,
                                   device="cpu")

        # Send object size

        torch.distributed.send(size_tensor,
                               dst=self.ranks[dst],
                               group=self.cpu_group)

        # Send object
        torch.distributed.send(object_tensor,
                               dst=self.ranks[dst],
                               group=self.cpu_group)

        return None

    def recv_object(self, src: int) -> Any:
        """Receive the input object list from the source rank."""
        """NOTE: `src` is the local rank of the source rank."""

        assert src < self.world_size, f"Invalid src rank ({src})"

541
        assert src != self.rank_in_group, (
542
543
544
545
546
547
548
            "Invalid source rank. Source rank is the same as the current rank."
        )

        size_tensor = torch.empty(1, dtype=torch.long, device="cpu")

        # Receive object size
        rank_size = torch.distributed.recv(size_tensor,
549
                                           src=self.ranks[src],
550
551
552
553
554
555
556
557
558
                                           group=self.cpu_group)

        # Tensor to receive serialized objects into.
        object_tensor = torch.empty(  # type: ignore[call-overload]
            size_tensor.item(),  # type: ignore[arg-type]
            dtype=torch.uint8,
            device="cpu")

        rank_object = torch.distributed.recv(object_tensor,
559
                                             src=self.ranks[src],
560
561
562
563
564
565
566
567
568
                                             group=self.cpu_group)

        assert rank_object == rank_size, (
            "Received object sender rank does not match the size sender rank.")

        obj = pickle.loads(object_tensor.numpy().tobytes())

        return obj

569
570
    def broadcast_tensor_dict(
        self,
571
        tensor_dict: Optional[dict[str, Union[torch.Tensor, Any]]] = None,
572
573
574
        src: int = 0,
        group: Optional[ProcessGroup] = None,
        metadata_group: Optional[ProcessGroup] = None
575
    ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:
576
577
578
579
580
581
582
583
584
585
586
        """Broadcast the input tensor dictionary.
        NOTE: `src` is the local rank of the source rank.
        """
        # Bypass the function if we are using only 1 GPU.
        if (not torch.distributed.is_initialized() or self.world_size == 1):
            return tensor_dict

        group = self.device_group
        metadata_group = self.cpu_group
        assert src < self.world_size, f"Invalid src rank ({src})"

587
588
        rank_in_group = self.rank_in_group
        if rank_in_group == src:
589
            metadata_list: list[tuple[Any, Any]] = []
590
591
592
593
594
595
596
            assert isinstance(
                tensor_dict,
                dict), (f"Expecting a dictionary, got {type(tensor_dict)}")
            metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
            # `metadata_list` lives in CPU memory.
            # `broadcast_object_list` has serialization & deserialization,
            # all happening on CPU. Therefore, we can use the CPU group.
597
            self.broadcast_object(metadata_list, src=src)
598
599
600
601
602
603
604
605
            async_handles = []
            for tensor in tensor_list:
                if tensor.numel() == 0:
                    # Skip broadcasting empty tensors.
                    continue
                if tensor.is_cpu:
                    # use metadata_group for CPU tensors
                    handle = torch.distributed.broadcast(tensor,
606
                                                         src=self.ranks[src],
607
608
609
610
611
                                                         group=metadata_group,
                                                         async_op=True)
                else:
                    # use group for GPU tensors
                    handle = torch.distributed.broadcast(tensor,
612
                                                         src=self.ranks[src],
613
614
615
616
617
618
619
                                                         group=group,
                                                         async_op=True)
                async_handles.append(handle)
            for async_handle in async_handles:
                async_handle.wait()

        else:
620
            metadata_list = self.broadcast_object(None, src=src)
621
622
            tensor_dict = {}
            async_handles = []
623
            for key, value in metadata_list:
624
625
626
627
628
629
                if isinstance(value, TensorMetadata):
                    tensor = torch.empty(value.size,
                                         dtype=value.dtype,
                                         device=value.device)
                    if tensor.numel() == 0:
                        # Skip broadcasting empty tensors.
630
                        tensor_dict[key] = tensor
631
632
633
634
635
                        continue
                    if tensor.is_cpu:
                        # use metadata_group for CPU tensors
                        handle = torch.distributed.broadcast(
                            tensor,
636
                            src=self.ranks[src],
637
638
639
640
                            group=metadata_group,
                            async_op=True)
                    else:
                        # use group for GPU tensors
641
642
643
644
645
                        handle = torch.distributed.broadcast(
                            tensor,
                            src=self.ranks[src],
                            group=group,
                            async_op=True)
646
                    async_handles.append(handle)
647
                    tensor_dict[key] = tensor
648
                else:
649
                    tensor_dict[key] = value
650
651
652
653
            for async_handle in async_handles:
                async_handle.wait()
        return tensor_dict

654
655
    def send_tensor_dict(
        self,
656
        tensor_dict: dict[str, Union[torch.Tensor, Any]],
657
658
        dst: Optional[int] = None,
        all_gather_group: Optional["GroupCoordinator"] = None,
659
        all_gather_tensors: Optional[dict[str, bool]] = None,
660
    ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:
661
662
        """Send the input tensor dictionary.
        NOTE: `dst` is the local rank of the source rank.
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677

        all_gather_group: The group for the all-gather operation. If provided,
            an optimization is enabled where each rank in the group sends a
            slice of a tensor and the receiver reconstructs it using an
            all-gather, which can improve performance. This is typically the
            tensor-parallel group.
        all_gather_tensors: A dictionary to specify which tensors should use
            the all-gather optimization, which is only effective when
            `all_gather_group` is provided. By default, this optimization is
            on for any tensor whose size is divisible by the
            `all_gather_group`'s world size. However, it should be disabled
            for tensors that are not fully replicated across the group (e.g.,
            the residual tensor when sequence parallelism is enabled). This
            dictionary allows overriding the default behavior on a per-tensor
            basis.
678
679
680
681
        """
        # Bypass the function if we are using only 1 GPU.
        if not torch.distributed.is_initialized() or self.world_size == 1:
            return tensor_dict
682
683
684
685
686
        all_gather_size = (1 if all_gather_group is None else
                           all_gather_group.world_size)
        all_gather_rank = (0 if all_gather_group is None else
                           all_gather_group.rank_in_group)

687
688
689
690
        group = self.device_group
        metadata_group = self.cpu_group

        if dst is None:
691
            dst = (self.rank_in_group + 1) % self.world_size
692
693
        assert dst < self.world_size, f"Invalid dst rank ({dst})"

694
        if self.use_cpu_custom_send_recv:
695
696
            if self.device_communicator is None:
                raise ValueError("No device communicator found")
697
698
699
700
            self.device_communicator.send_tensor_dict(  # type: ignore
                tensor_dict, dst)
            return None

701
        metadata_list: list[tuple[Any, Any]] = []
702
703
704
705
706
707
708
709
        assert isinstance(
            tensor_dict,
            dict), f"Expecting a dictionary, got {type(tensor_dict)}"
        metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
        # `metadata_list` lives in CPU memory.
        # `send_object_list` has serialization & deserialization,
        # all happening on CPU. Therefore, we can use the CPU group.
        self.send_object(metadata_list, dst=dst)
710
711
712
713
714
715
716

        tensor_keys = [
            k for k, v in tensor_dict.items() if isinstance(v, torch.Tensor)
        ]
        assert len(tensor_keys) == len(tensor_list)

        for key, tensor in zip(tensor_keys, tensor_list):
717
718
719
            if tensor.numel() == 0:
                # Skip sending empty tensors.
                continue
720
721

            # send-allgather: send only a slice, then do allgather.
722
723
724
725
726
            use_all_gather = (all_gather_group is not None
                              and tensor.numel() % all_gather_size == 0)
            use_all_gather = all_gather_tensors.get(key, use_all_gather) \
                if all_gather_tensors else use_all_gather
            if use_all_gather:
727
728
                tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank]

729
730
            if tensor.is_cpu:
                # use metadata_group for CPU tensors
731
732
733
                torch.distributed.send(tensor,
                                       dst=self.ranks[dst],
                                       group=metadata_group)
734
735
            else:
                # use group for GPU tensors
736
737
738
                torch.distributed.send(tensor,
                                       dst=self.ranks[dst],
                                       group=group)
739
740
741
742
        return None

    def recv_tensor_dict(
        self,
743
744
        src: Optional[int] = None,
        all_gather_group: Optional["GroupCoordinator"] = None,
745
        all_gather_tensors: Optional[dict[str, bool]] = None,
746
    ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:
747
748
        """Recv the input tensor dictionary.
        NOTE: `src` is the local rank of the source rank.
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763

        all_gather_group: The group for the all-gather operation. If provided,
            an optimization is enabled where each rank in the group sends a
            slice of a tensor and the receiver reconstructs it using an
            all-gather, which can improve performance. This is typically the
            tensor-parallel group.
        all_gather_tensors: A dictionary to specify which tensors should use
            the all-gather optimization, which is only effective when
            `all_gather_group` is provided. By default, this optimization is
            on for any tensor whose size is divisible by the
            `all_gather_group`'s world size. However, it should be disabled
            for tensors that are not fully replicated across the group (e.g.,
            the residual tensor when sequence parallelism is enabled). This
            dictionary allows overriding the default behavior on a per-tensor
            basis.
764
765
766
767
        """
        # Bypass the function if we are using only 1 GPU.
        if not torch.distributed.is_initialized() or self.world_size == 1:
            return None
768
769
770
771
772
        all_gather_size = (1 if all_gather_group is None else
                           all_gather_group.world_size)
        all_gather_rank = (0 if all_gather_group is None else
                           all_gather_group.rank_in_group)

773
774
775
776
        group = self.device_group
        metadata_group = self.cpu_group

        if src is None:
777
            src = (self.rank_in_group - 1) % self.world_size
778
779
        assert src < self.world_size, f"Invalid src rank ({src})"

780
        if self.use_cpu_custom_send_recv:
781
782
            if self.device_communicator is None:
                raise ValueError("No device communicator found")
783
784
785
            return self.device_communicator.recv_tensor_dict(  # type: ignore
                src)

786
        recv_metadata_list = self.recv_object(src=src)
787
        tensor_dict: dict[str, Any] = {}
788
789
790
791
792
793
794
        for key, value in recv_metadata_list:
            if isinstance(value, TensorMetadata):
                tensor = torch.empty(value.size,
                                     dtype=value.dtype,
                                     device=value.device)
                if tensor.numel() == 0:
                    # Skip broadcasting empty tensors.
795
                    tensor_dict[key] = tensor
796
                    continue
797
798
799
800

                # send-allgather: send only a slice, then do allgather.
                use_all_gather = (all_gather_group is not None
                                  and tensor.numel() % all_gather_size == 0)
801
802
                use_all_gather = all_gather_tensors.get(key, use_all_gather) \
                    if all_gather_tensors else use_all_gather
803
804
805
806
807
808

                if use_all_gather:
                    orig_shape = tensor.shape
                    tensor = tensor.reshape(all_gather_size,
                                            -1)[all_gather_rank]

809
810
811
                if tensor.is_cpu:
                    # use metadata_group for CPU tensors
                    torch.distributed.recv(tensor,
812
                                           src=self.ranks[src],
813
814
815
                                           group=metadata_group)
                else:
                    # use group for GPU tensors
816
817
818
                    torch.distributed.recv(tensor,
                                           src=self.ranks[src],
                                           group=group)
819
820
821
822
823
824
                if use_all_gather:
                    # do the allgather
                    tensor = all_gather_group.all_gather(  # type: ignore
                        tensor, dim=0)
                    tensor = tensor.reshape(orig_shape)

825
                tensor_dict[key] = tensor
826
            else:
827
                tensor_dict[key] = value
828
829
        return tensor_dict

830
831
832
833
834
835
836
837
838
    def barrier(self):
        """Barrier synchronization among the group.
        NOTE: don't use `device_group` here! `barrier` in NCCL is
        terrible because it is internally a broadcast operation with
        secretly created GPU tensors. It is easy to mess up the current
        device. Use the CPU group instead.
        """
        torch.distributed.barrier(group=self.cpu_group)

839
    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
840
        """Sends a tensor to the destination rank in a blocking way"""
841
        """NOTE: `dst` is the local rank of the destination rank."""
842
843
        if self.device_communicator is None:
            raise ValueError("No device communicator found")
844
        self.device_communicator.send(tensor, dst)
845
846
847
848
849

    def recv(self,
             size: torch.Size,
             dtype: torch.dtype,
             src: Optional[int] = None) -> torch.Tensor:
850
851
        """Receives a tensor from the source rank."""
        """NOTE: `src` is the local rank of the source rank."""
852
853
        if self.device_communicator is None:
            raise ValueError("No device communicator found")
854
        return self.device_communicator.recv(size, dtype, src)
855

856
    def destroy(self):
857
        if hasattr(self, "device_group"):
858
            torch.distributed.destroy_process_group(self.device_group)
859
860
            del self.device_group
        if hasattr(self, "cpu_group"):
861
            torch.distributed.destroy_process_group(self.cpu_group)
862
            del self.cpu_group
863
864
        if self.device_communicator is not None:
            self.device_communicator.destroy()
865
866
        if self.mq_broadcaster is not None:
            self.mq_broadcaster = None
867

868
869
870
871
872
873
    def prepare_communication_buffer_for_model(self, model: torch.nn.Module):
        if self.device_communicator is not None:
            self.device_communicator.prepare_communication_buffer_for_model(
                model)

    def dispatch(
874
875
876
877
878
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
        is_sequence_parallel: bool = False
    ) -> tuple[torch.Tensor, torch.Tensor]:
879
880
        if self.device_communicator is not None:
            return self.device_communicator.dispatch(hidden_states,
881
882
                                                     router_logits,
                                                     is_sequence_parallel)
883
884
        else:
            return hidden_states, router_logits
885

886
887
888
    def combine(self,
                hidden_states,
                is_sequence_parallel: bool = False) -> torch.Tensor:
889
        if self.device_communicator is not None:
890
891
            return self.device_communicator.combine(hidden_states,
                                                    is_sequence_parallel)
892
893
        else:
            return hidden_states
894

895
896

_WORLD: Optional[GroupCoordinator] = None
897
_NODE_COUNT: Optional[int] = None
898
899
900
901
902
903
904


def get_world_group() -> GroupCoordinator:
    assert _WORLD is not None, ("world group is not initialized")
    return _WORLD


905
def init_world_group(ranks: list[int], local_rank: int,
906
907
908
909
910
                     backend: str) -> GroupCoordinator:
    return GroupCoordinator(
        group_ranks=[ranks],
        local_rank=local_rank,
        torch_distributed_backend=backend,
911
        use_device_communicator=False,
912
        group_name="world",
913
914
915
    )


916
def init_model_parallel_group(
917
    group_ranks: list[list[int]],
918
919
920
    local_rank: int,
    backend: str,
    use_message_queue_broadcaster: bool = False,
921
    group_name: Optional[str] = None,
922
) -> GroupCoordinator:
923

924
925
926
927
    return GroupCoordinator(
        group_ranks=group_ranks,
        local_rank=local_rank,
        torch_distributed_backend=backend,
928
        use_device_communicator=True,
929
        use_message_queue_broadcaster=use_message_queue_broadcaster,
930
        group_name=group_name,
931
932
933
    )


934
935
936
937
938
939
940
941
_TP: Optional[GroupCoordinator] = None


def get_tp_group() -> GroupCoordinator:
    assert _TP is not None, ("tensor model parallel group is not initialized")
    return _TP


942
943
944
945
946
947
@deprecated("`get_tensor_model_parallel_group` has been replaced with "
            "`get_tp_group` and may be removed after v0.12. Please use "
            "`get_tp_group` instead.")
def get_tensor_model_parallel_group():
    return get_tp_group()

948

949
950
951
952
953
954
955
956
957
958
959
960
_DCP: Optional[GroupCoordinator] = None


def get_dcp_group() -> GroupCoordinator:
    assert _DCP is not None, (
        "decode context model parallel group is not initialized")
    return _DCP


# kept for backward compatibility
get_context_model_parallel_group = get_dcp_group

961
962
_PP: Optional[GroupCoordinator] = None

963
964
965
966
967
968
969
_DP: Optional[GroupCoordinator] = None


def get_dp_group() -> GroupCoordinator:
    assert _DP is not None, ("data parallel group is not initialized")
    return _DP

970

971
972
973
974
975
976
977
978
_EP: Optional[GroupCoordinator] = None


def get_ep_group() -> GroupCoordinator:
    assert _EP is not None, ("expert parallel group is not initialized")
    return _EP


979
980
981
982
def get_pp_group() -> GroupCoordinator:
    assert _PP is not None, (
        "pipeline model parallel group is not initialized")
    return _PP
983
984


985
986
987
988
989
@deprecated("`get_pipeline_model_parallel_group` has been replaced with "
            "`get_pp_group` and may be removed in v0.12. Please use "
            "`get_pp_group` instead.")
def get_pipeline_model_parallel_group():
    return get_pp_group()
990
991


992
@contextmanager
993
def graph_capture(device: torch.device):
994
995
    """
    `graph_capture` is a context manager which should surround the code that
996
997
    is capturing the CUDA graph. Its main purpose is to ensure that some
    operations will be run after the graph is captured, before the graph
998
999
1000
1001
1002
1003
1004
1005
1006
    is replayed. It returns a `GraphCaptureContext` object which contains the
    necessary data for the graph capture. Currently, it only contains the
    stream that the graph capture is running on. This stream is set to the
    current CUDA stream when the context manager is entered and reset to the
    default stream when the context manager is exited. This is to ensure that
    the graph capture is running on a separate stream from the default stream,
    in order to explicitly distinguish the kernels to capture
    from other kernels possibly launched on background in the default stream.
    """
1007
1008
1009
    context = GraphCaptureContext(torch.cuda.Stream(device=device))
    with get_tp_group().graph_capture(context), get_pp_group().graph_capture(
            context):
1010
1011
        yield context

1012

1013
logger = init_logger(__name__)
1014

1015
_ENABLE_CUSTOM_ALL_REDUCE = True
1016
1017


1018
1019
1020
def set_custom_all_reduce(enable: bool):
    global _ENABLE_CUSTOM_ALL_REDUCE
    _ENABLE_CUSTOM_ALL_REDUCE = enable
1021

Zhuohan Li's avatar
Zhuohan Li committed
1022

1023
1024
1025
1026
1027
1028
def init_distributed_environment(world_size: int = -1,
                                 rank: int = -1,
                                 distributed_init_method: str = "env://",
                                 local_rank: int = -1,
                                 backend: str = "nccl",
                                 timeout: Optional[timedelta] = None):
1029
1030
1031
1032
    logger.debug(
        "world_size=%d rank=%d local_rank=%d "
        "distributed_init_method=%s backend=%s", world_size, rank, local_rank,
        distributed_init_method, backend)
1033
1034
    from vllm.config import get_current_vllm_config
    config = get_current_vllm_config()
1035
1036
1037
    if config is not None and config.parallel_config.data_parallel_size > 1 \
        and config.parallel_config.distributed_executor_backend \
        != "external_launcher":
1038
1039
1040
1041
1042
1043
1044
1045
        parallel_config = config.parallel_config
        # adjust to take into account data parallelism
        # offset the rank by the data parallel rank
        rank = parallel_config.data_parallel_rank * world_size + rank
        # adjust the world size to take into account data parallelism
        world_size = parallel_config.world_size_across_dp
        ip = parallel_config.data_parallel_master_ip
        port = parallel_config.get_next_dp_init_port()
1046
        distributed_init_method = get_distributed_init_method(ip, port)
1047
1048
1049
        logger.info(
            "Adjusting world_size=%d rank=%d distributed_init_method=%s for DP",
            world_size, rank, distributed_init_method)
1050
1051
1052
1053
    if not torch.distributed.is_initialized():
        assert distributed_init_method is not None, (
            "distributed_init_method must be provided when initializing "
            "distributed environment")
1054
1055
1056
1057
1058
1059
1060
        if not torch.distributed.is_backend_available(backend):
            logger.warning(
                "Distributed backend %s is not available; "
                "falling back to gloo.", backend)
            assert torch.distributed.is_gloo_available(), (
                "Fallback Gloo backend is not available.")
            backend = "gloo"
1061
1062
1063
1064
1065
        # this backend is used for WORLD
        torch.distributed.init_process_group(
            backend=backend,
            init_method=distributed_init_method,
            world_size=world_size,
1066
1067
            rank=rank,
            timeout=timeout)
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
    # set the local rank
    # local_rank is not available in torch ProcessGroup,
    # see https://github.com/pytorch/pytorch/issues/122816
    if local_rank == -1:
        # local rank not set, this usually happens in single-node
        # setting, where we can use rank as local rank
        if distributed_init_method == "env://":
            local_rank = envs.LOCAL_RANK
        else:
            local_rank = rank
1078
    global _WORLD, _NODE_COUNT
1079
    if _WORLD is None:
1080
        ranks = list(range(torch.distributed.get_world_size()))
1081
        _WORLD = init_world_group(ranks, local_rank, backend)
1082
1083
1084
        _NODE_COUNT = _node_count(_WORLD.cpu_group)
        logger.debug("Detected %d nodes in the distributed environment",
                     _NODE_COUNT)
1085
1086
1087
    else:
        assert _WORLD.world_size == torch.distributed.get_world_size(), (
            "world group already initialized with a different world size")
1088
1089


Zhuohan Li's avatar
Zhuohan Li committed
1090
1091
1092
def initialize_model_parallel(
    tensor_model_parallel_size: int = 1,
    pipeline_model_parallel_size: int = 1,
1093
    decode_context_model_parallel_size: Optional[int] = 1,
1094
    backend: Optional[str] = None,
Zhuohan Li's avatar
Zhuohan Li committed
1095
1096
) -> None:
    """
1097
    Initialize model parallel groups.
Zhuohan Li's avatar
Zhuohan Li committed
1098
1099

    Arguments:
1100
1101
1102
1103
        tensor_model_parallel_size: number of GPUs used for tensor model
            parallelism.
        pipeline_model_parallel_size: number of GPUs used for pipeline model
            parallelism.
1104
        backend: name of torch distributed communication backend.
1105
1106

    Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
Zhuohan Li's avatar
Zhuohan Li committed
1107
1108
    use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
    the model pipeline. The present function will
1109
1110
1111
1112
1113
    create 4 tensor model-parallel groups and 2 pipeline model-parallel groups:
        4 tensor model-parallel groups:
            [g0, g1], [g2, g3], [g4, g5], [g6, g7]
        2 pipeline model-parallel groups:
            [g0, g2, g4, g6], [g1, g3, g5, g7]
Zhuohan Li's avatar
Zhuohan Li committed
1114
1115
1116
1117
1118
1119
1120
1121
    Note that for efficiency, the caller should make sure adjacent ranks
    are on the same DGX box. For example if we are using 2 DGX-1 boxes
    with a total of 16 GPUs, rank 0 to 7 belong to the first box and
    ranks 8 to 15 belong to the second box.
    """
    # Get world size and rank. Ensure some consistencies.
    assert torch.distributed.is_initialized()
    world_size: int = torch.distributed.get_world_size()
1122
    rank = torch.distributed.get_rank()
1123
1124
    backend = backend or torch.distributed.get_backend(
        get_world_group().device_group)
Zhuohan Li's avatar
Zhuohan Li committed
1125

1126
1127
1128
1129
    data_parallel_size = 1
    from vllm.config import get_current_vllm_config
    config = get_current_vllm_config()
    if config is not None:
1130
1131
1132
1133
1134
1135
1136
1137
1138
        data_parallel_size = config.parallel_config.data_parallel_size

    # the layout order is: ExternalDP x DP x PP x TP
    # ExternalDP is the data parallel group that is not part of the model,
    # every dp rank can generate independently (in verl integration).
    # DP is the data parallel group that is part of the model,
    # all the ranks in the same DP group should generate simultaneously,
    # i.e. the `generate` call in the same DP group should be called together,
    # otherwise it will cause deadlock.
1139
1140
1141
    # to get group_ranks for each dimension, transpose that dimension to the
    # last dimension, then reshape to 2D, then unbind the last dimension
    all_ranks = torch.arange(world_size).reshape(
1142
        -1, data_parallel_size, pipeline_model_parallel_size,
1143
1144
        tensor_model_parallel_size)  # noqa

1145
1146
1147
    # Build the tensor model-parallel groups.
    global _TP
    assert _TP is None, ("tensor model parallel group is already initialized")
1148
1149
    group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0)
    group_ranks = [x.tolist() for x in group_ranks]
1150
1151

    # message queue broadcaster is only used in tensor model parallel group
1152
    _TP = init_model_parallel_group(group_ranks,
1153
1154
                                    get_world_group().local_rank,
                                    backend,
1155
1156
                                    use_message_queue_broadcaster=True,
                                    group_name="tp")
1157

1158
1159
1160
1161
1162
1163
    # Build the DCP model-parallel groups.
    global _DCP
    assert _DCP is None, (
        "decode context model parallel group is already initialized")
    # Note(hc): In the current implementation of decode context parallel,
    # dcp_size must not exceed tp_size, because the world size does not
1164
    # change by DCP, it simply reuses the GPUs of TP group, and split one
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
    # TP group into tp_size//dcp_size DCP groups.
    group_ranks = all_ranks.reshape(
        -1, decode_context_model_parallel_size).unbind(0)
    group_ranks = [x.tolist() for x in group_ranks]
    _DCP = init_model_parallel_group(group_ranks,
                                     get_world_group().local_rank,
                                     backend,
                                     use_message_queue_broadcaster=True,
                                     group_name="dcp")

1175
    # Build the pipeline model-parallel groups.
1176
1177
    global _PP
    assert _PP is None, (
1178
        "pipeline model parallel group is already initialized")
1179
    group_ranks = all_ranks.transpose(2, 3).reshape(
1180
1181
        -1, pipeline_model_parallel_size).unbind(0)
    group_ranks = [x.tolist() for x in group_ranks]
1182
    _PP = init_model_parallel_group(group_ranks,
1183
1184
                                    get_world_group().local_rank,
                                    backend,
1185
                                    group_name="pp")
1186

1187
1188
    global _DP
    assert _DP is None, ("data parallel group is already initialized")
1189
1190
    group_ranks = all_ranks.transpose(1,
                                      3).reshape(-1,
1191
1192
1193
1194
1195
1196
1197
                                                 data_parallel_size).unbind(0)
    group_ranks = [x.tolist() for x in group_ranks]
    _DP = init_model_parallel_group(group_ranks,
                                    get_world_group().local_rank,
                                    backend,
                                    group_name="dp")

1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
    global _EP
    assert _EP is None, ("expert parallel group is already initialized")
    group_ranks = all_ranks.transpose(1, 2).reshape(
        -1, data_parallel_size * tensor_model_parallel_size).unbind(0)
    group_ranks = [x.tolist() for x in group_ranks]
    _EP = init_model_parallel_group(group_ranks,
                                    get_world_group().local_rank,
                                    backend,
                                    group_name="ep")

1208
1209
    logger.info(
        "rank %s in world size %s is assigned as "
1210
1211
1212
        "DP rank %s, PP rank %s, TP rank %s, EP rank %s", rank, world_size,
        _DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group,
        _EP.rank_in_group)
1213

Zhuohan Li's avatar
Zhuohan Li committed
1214

1215
1216
1217
def ensure_model_parallel_initialized(
    tensor_model_parallel_size: int,
    pipeline_model_parallel_size: int,
1218
    decode_context_model_parallel_size: Optional[int] = 1,
1219
    backend: Optional[str] = None,
1220
1221
1222
1223
1224
) -> None:
    """Helper to initialize model parallel groups if they are not initialized,
    or ensure tensor-parallel and pipeline-parallel sizes are equal to expected
    values if the model parallel groups are initialized.
    """
1225
1226
    backend = backend or torch.distributed.get_backend(
        get_world_group().device_group)
1227
1228
    if not model_parallel_is_initialized():
        initialize_model_parallel(tensor_model_parallel_size,
1229
1230
                                  pipeline_model_parallel_size,
                                  decode_context_model_parallel_size, backend)
1231
1232
1233
1234
        return

    assert (
        get_tensor_model_parallel_world_size() == tensor_model_parallel_size
1235
1236
1237
    ), ("tensor parallel group already initialized, but of unexpected size. "
        f"got: {get_tensor_model_parallel_world_size()=} vs. "
        f"wanted: {tensor_model_parallel_size=}")
1238
1239
    pp_world_size = get_pp_group().world_size
    assert (pp_world_size == pipeline_model_parallel_size), (
1240
1241
1242
        "pipeline parallel group already initialized, but of unexpected size. "
        f"got: {pp_world_size=} vs. "
        f"wanted: {pipeline_model_parallel_size=}")
1243
1244


1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
def prepare_communication_buffer_for_model(model: torch.nn.Module):
    """Prepare the communication buffer for the model.
    Traditional communication libraries like NCCL are almost
    model agnostic. However, emerging new communication libraries like
    MoE all2all (DeepEP) usually allocate the communication buffer
    based on the model shape for optimal performance.
    """
    if _TP is not None:
        _TP.prepare_communication_buffer_for_model(model)
    if _PP is not None:
        _PP.prepare_communication_buffer_for_model(model)
    if _DP is not None:
        _DP.prepare_communication_buffer_for_model(model)
    if _EP is not None:
        _EP.prepare_communication_buffer_for_model(model)


Zhuohan Li's avatar
Zhuohan Li committed
1262
def model_parallel_is_initialized():
1263
    """Check if tensor and pipeline parallel groups are initialized."""
1264
    return (_TP is not None and _PP is not None)
1265
1266


1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
_TP_STATE_PATCHED = False


@contextmanager
def patch_tensor_parallel_group(tp_group: GroupCoordinator):
    """Patch the tp group temporarily until this function ends.

    This method is for draft workers of speculative decoding to run draft model
    with different tp degree from that of target model workers.

    Args:
        tp_group (GroupCoordinator): the tp group coordinator
    """
    global _TP_STATE_PATCHED
    assert not _TP_STATE_PATCHED, "Should not call when it's already patched"

    _TP_STATE_PATCHED = True
    old_tp_group = get_tp_group()
    global _TP
    _TP = tp_group
    try:
        yield
    finally:
        # restore the original state
        _TP_STATE_PATCHED = False
        _TP = old_tp_group


Zhuohan Li's avatar
Zhuohan Li committed
1295
1296
def get_tensor_model_parallel_world_size():
    """Return world size for the tensor model parallel group."""
1297
    return get_tp_group().world_size
Zhuohan Li's avatar
Zhuohan Li committed
1298
1299
1300
1301


def get_tensor_model_parallel_rank():
    """Return my rank for the tensor model parallel group."""
1302
    return get_tp_group().rank_in_group
Zhuohan Li's avatar
Zhuohan Li committed
1303
1304


1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
def get_decode_context_model_parallel_world_size():
    """Return world size for the decode context model parallel group."""
    return get_dcp_group().world_size


def get_decode_context_model_parallel_rank():
    """Return my rank for the decode context model parallel group."""
    return get_dcp_group().rank_in_group


1315
1316
1317
1318
1319
1320
1321
def get_node_count() -> int:
    """Return the total number of nodes in the distributed environment. """
    assert _NODE_COUNT is not None, (
        "distributed environment is not initialized")
    return _NODE_COUNT


Zhuohan Li's avatar
Zhuohan Li committed
1322
def destroy_model_parallel():
1323
    """Set the groups to none and destroy them."""
1324
    global _TP
1325

1326
1327
1328
1329
1330
1331
1332
1333
1334
    if _TP:
        _TP.destroy()
    _TP = None

    global _PP
    if _PP:
        _PP.destroy()
    _PP = None

1335
1336
1337
1338
1339
    global _DCP
    if _DCP:
        _DCP.destroy()
    _DCP = None

1340
1341
1342
1343
1344
    global _DP
    if _DP:
        _DP.destroy()
    _DP = None

1345
1346
1347
1348
1349
    global _EP
    if _EP:
        _EP.destroy()
    _EP = None

1350
1351

def destroy_distributed_environment():
1352
    global _WORLD, _NODE_COUNT
1353
1354
1355
    if _WORLD:
        _WORLD.destroy()
    _WORLD = None
1356
    _NODE_COUNT = None
1357
1358
    if torch.distributed.is_initialized():
        torch.distributed.destroy_process_group()
1359
1360


1361
1362
1363
1364
1365
1366
1367
def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
    destroy_model_parallel()
    destroy_distributed_environment()
    if shutdown_ray:
        import ray  # Lazy import Ray
        ray.shutdown()
    gc.collect()
1368
    from vllm.platforms import current_platform
1369
1370
1371
    empty_cache = current_platform.empty_cache
    if empty_cache is not None:
        empty_cache()
1372
    try:
1373
1374
        if not current_platform.is_cpu():
            torch._C._host_emptyCache()
1375
1376
1377
    except AttributeError:
        logger.warning(
            "torch._C._host_emptyCache() only available in Pytorch >=2.5")
1378
1379


1380
def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
1381
                        source_rank: int = 0) -> list[bool]:
1382
    """
1383
1384
    This is a collective operation that returns if each rank is in the same node
    as the source rank. It tests if processes are attached to the same
1385
1386
    memory system (shared access to shared memory).
    """
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
    if isinstance(pg, ProcessGroup):
        assert torch.distributed.get_backend(
            pg) != torch.distributed.Backend.NCCL, (
                "in_the_same_node_as should be tested with a non-NCCL group.")
        # local rank inside the group
        rank = torch.distributed.get_rank(group=pg)
        world_size = torch.distributed.get_world_size(group=pg)

        # global ranks of the processes in the group
        ranks = torch.distributed.get_process_group_ranks(pg)
    else:
        rank = pg.rank
        world_size = pg.world_size
        ranks = list(range(world_size))
1401
1402
1403
1404
1405
1406
1407
1408
1409

    # local tensor in each process to store the result
    is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32)

    magic_message = b"magic_message"
    shm = None

    try:
        with contextlib.suppress(OSError):
1410
            if rank == source_rank:
1411
1412
1413
                # create a shared memory segment
                shm = shared_memory.SharedMemory(create=True, size=128)
                shm.buf[:len(magic_message)] = magic_message
1414
1415
1416
1417
1418
                if isinstance(pg, ProcessGroup):
                    torch.distributed.broadcast_object_list(
                        [shm.name], src=ranks[source_rank], group=pg)
                else:
                    pg.broadcast_obj(shm.name, src=source_rank)
1419
                is_in_the_same_node[rank] = 1
1420
1421
            else:
                # try to open the shared memory segment
1422
1423
1424
1425
1426
1427
1428
                if isinstance(pg, ProcessGroup):
                    recv = [None]
                    torch.distributed.broadcast_object_list(
                        recv, src=ranks[source_rank], group=pg)
                    name = recv[0]
                else:
                    name = pg.broadcast_obj(None, src=source_rank)
1429
1430
1431
1432
1433
1434
                # fix to https://stackoverflow.com/q/62748654/9191338
                # Python incorrectly tracks shared memory even if it is not
                # created by the process. The following patch is a workaround.
                with patch("multiprocessing.resource_tracker.register",
                           lambda *args, **kwargs: None):
                    shm = shared_memory.SharedMemory(name=name)
1435
1436
1437
1438
1439
1440
1441
1442
                if shm.buf[:len(magic_message)] == magic_message:
                    is_in_the_same_node[rank] = 1
    except Exception as e:
        logger.error("Error ignored in is_in_the_same_node: %s", e)
    finally:
        if shm:
            shm.close()

1443
1444
1445
1446
    if isinstance(pg, ProcessGroup):
        torch.distributed.barrier(group=pg)
    else:
        pg.barrier()
1447
1448
1449

    # clean up the shared memory segment
    with contextlib.suppress(OSError):
1450
        if rank == source_rank and shm:
1451
            shm.unlink()
1452

1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
    if isinstance(pg, ProcessGroup):
        torch.distributed.all_reduce(is_in_the_same_node, group=pg)
        aggregated_data = is_in_the_same_node
    else:
        aggregated_data = torch.zeros_like(is_in_the_same_node)
        for i in range(world_size):
            rank_data = pg.broadcast_obj(is_in_the_same_node, src=i)
            aggregated_data += rank_data

    return [x == 1 for x in aggregated_data.tolist()]
1463
1464


1465
1466
def is_global_first_rank() -> bool:
    """
1467
    Check if the current process is the first rank globally across all
1468
    parallelism strategies (PP, TP, DP, EP, etc.).
1469

1470
1471
1472
    Unlike group-specific checks like `get_tensor_model_parallel_rank() == 0`
    or `get_pp_group().is_first_rank`, this function checks the global rank
    across all parallelism dimensions.
1473

1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
    Returns:
        bool: True if this is the global first rank (rank 0), False otherwise.
              Returns True if distributed is not initialized (single process).
    """
    try:
        # If world group is available, use it for the most accurate check
        global _WORLD
        if _WORLD is not None:
            return _WORLD.is_first_rank

        # If torch distributed is not initialized, assume single process
        if not torch.distributed.is_initialized():
            return True

        # Fallback to torch's global rank
        return torch.distributed.get_rank() == 0

    except Exception:
        # If anything goes wrong, assume this is the first rank
        return True


1496
1497
1498
1499
1500
1501
def _node_count(pg: Union[ProcessGroup, StatelessProcessGroup]) -> int:
    """
    Returns the total number of nodes in the process group.

    Args:
        pg: The process group to analyze
1502

1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
    Returns:
        int: The total number of nodes
    """
    if isinstance(pg, ProcessGroup):
        world_size = torch.distributed.get_world_size(group=pg)
    else:
        world_size = pg.world_size

    if world_size == 1:
        return 1

    # Build node assignment map
    node_assignment = [0] * world_size  # rank -> node_id
    next_node_id = 0

    for current_rank in range(world_size):
        if node_assignment[current_rank] != 0:
            continue  # Already assigned to a node

        # Assign current rank to a new node
        next_node_id += 1
        node_assignment[current_rank] = next_node_id

        # Find all ranks on the same node as current_rank
        same_node_flags = in_the_same_node_as(pg, current_rank)
        for other_rank, is_same_node in enumerate(same_node_flags):
            if is_same_node and node_assignment[other_rank] == 0:
                node_assignment[other_rank] = next_node_id

    return next_node_id