parallel_state.py 12.1 KB
Newer Older
Woosuk Kwon's avatar
Woosuk Kwon committed
1
# Copyright 2023 The vLLM team.
2
3
# Adapted from
# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py
Zhuohan Li's avatar
Zhuohan Li committed
4
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
5
"""Tensor and pipeline parallel groups."""
Woosuk Kwon's avatar
Woosuk Kwon committed
6
import contextlib
7
from typing import Optional
Zhuohan Li's avatar
Zhuohan Li committed
8
9
10

import torch

11
12
13
14
from vllm.logger import init_logger

logger = init_logger(__name__)

15
# Tensor model parallel group that the current rank belongs to.
Zhuohan Li's avatar
Zhuohan Li committed
16
_TENSOR_MODEL_PARALLEL_GROUP = None
17
# Pipeline model parallel group that the current rank belongs to.
Zhuohan Li's avatar
Zhuohan Li committed
18
19
_PIPELINE_MODEL_PARALLEL_GROUP = None

20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
# when people blindly call `torch.distributed.all_reduce` etc,
# it will use this group. It is initialized with the `backend`
# parameter of `init_distributed_environment` below.
# Essentially, this is `torch.distributed.group.WORLD`.
# We leave a line here to note that this is device-specific.
# Note that this variable is not safe to use, because when users
# call `init_distributed_environment` first, and then destroy
# the process group themselves, this variable will keep a reference to the
# destroyed process group, which is not useful.
_DEVICE_WORLD_GROUP = None

# duing `init_distributed_environment`, we will also initialize a
# group with `gloo` backend, to allow direct coordination between
# processes through the CPU.
_CPU_WORLD_GROUP = None

# In summary, after calling `init_distributed_environment`, we will
# always have two groups: one for device-specific (and is the default)
# and one for CPU. All processes will be part of both groups.

40
41
42
# A list of global ranks for each pipeline group to ease calculation of the
# source rank when broadcasting from the first or last pipeline stage.
_PIPELINE_GLOBAL_RANKS = None
Zhuohan Li's avatar
Zhuohan Li committed
43
44


45
def init_distributed_environment(
46
47
48
    world_size: int = -1,
    rank: int = -1,
    distributed_init_method: str = "env://",
49
50
51
    local_rank: int = -1,
    backend: str = "nccl",
):
52
53
    logger.debug(f"{world_size=} {rank=} {local_rank=} "
                 f"{distributed_init_method=} {backend=}")
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
    if not torch.distributed.is_initialized():
        assert distributed_init_method is not None, (
            "distributed_init_method must be provided when initializing "
            "distributed environment")
        # this backend is used for WORLD
        torch.distributed.init_process_group(
            backend=backend,
            init_method=distributed_init_method,
            world_size=world_size,
            rank=rank)
        global _DEVICE_WORLD_GROUP, _CPU_WORLD_GROUP
        _DEVICE_WORLD_GROUP = torch.distributed.group.WORLD
        ranks = list(range(torch.distributed.get_world_size()))
        _CPU_WORLD_GROUP = torch.distributed.new_group(ranks=ranks,
                                                       backend="gloo")


Zhuohan Li's avatar
Zhuohan Li committed
71
72
73
def initialize_model_parallel(
    tensor_model_parallel_size: int = 1,
    pipeline_model_parallel_size: int = 1,
74
    backend: Optional[str] = None,
Zhuohan Li's avatar
Zhuohan Li committed
75
76
) -> None:
    """
77
    Initialize model parallel groups.
Zhuohan Li's avatar
Zhuohan Li committed
78
79

    Arguments:
80
81
82
83
84
85
        tensor_model_parallel_size: number of GPUs used for tensor model
            parallelism.
        pipeline_model_parallel_size: number of GPUs used for pipeline model
            parallelism.

    Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
Zhuohan Li's avatar
Zhuohan Li committed
86
87
    use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
    the model pipeline. The present function will
88
89
90
91
92
    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
93
94
95
96
97
98
99
100
    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()
101
102
    # get the backend of _DEVICE_WORLD_GROUP
    backend = backend or torch.distributed.get_backend()
Zhuohan Li's avatar
Zhuohan Li committed
103

104
105
    if (world_size !=
            tensor_model_parallel_size * pipeline_model_parallel_size):
Zhuohan Li's avatar
Zhuohan Li committed
106
        raise RuntimeError(
107
108
109
110
111
112
113
114
            f"world_size ({world_size}) is not equal to "
            f"tensor_model_parallel_size ({tensor_model_parallel_size}) x "
            f"pipeline_model_parallel_size ({pipeline_model_parallel_size})")

    num_tensor_model_parallel_groups: int = (world_size //
                                             tensor_model_parallel_size)
    num_pipeline_model_parallel_groups: int = (world_size //
                                               pipeline_model_parallel_size)
Zhuohan Li's avatar
Zhuohan Li committed
115
116
117
118
    rank = torch.distributed.get_rank()

    # Build the tensor model-parallel groups.
    global _TENSOR_MODEL_PARALLEL_GROUP
119
120
    assert _TENSOR_MODEL_PARALLEL_GROUP is None, (
        "tensor model parallel group is already initialized")
Zhuohan Li's avatar
Zhuohan Li committed
121
122
123
    for i in range(num_tensor_model_parallel_groups):
        ranks = range(i * tensor_model_parallel_size,
                      (i + 1) * tensor_model_parallel_size)
124
        group = torch.distributed.new_group(ranks, backend=backend)
Zhuohan Li's avatar
Zhuohan Li committed
125
126
127
        if rank in ranks:
            _TENSOR_MODEL_PARALLEL_GROUP = group

128
    # Build the pipeline model-parallel groups.
Zhuohan Li's avatar
Zhuohan Li committed
129
130
    global _PIPELINE_MODEL_PARALLEL_GROUP
    global _PIPELINE_GLOBAL_RANKS
131
132
    assert _PIPELINE_MODEL_PARALLEL_GROUP is None, (
        "pipeline model parallel group is already initialized")
Zhuohan Li's avatar
Zhuohan Li committed
133
134
    for i in range(num_pipeline_model_parallel_groups):
        ranks = range(i, world_size, num_pipeline_model_parallel_groups)
135
        group = torch.distributed.new_group(ranks, backend=backend)
Zhuohan Li's avatar
Zhuohan Li committed
136
137
138
139
140
        if rank in ranks:
            _PIPELINE_MODEL_PARALLEL_GROUP = group
            _PIPELINE_GLOBAL_RANKS = ranks


141
142
143
def ensure_model_parallel_initialized(
    tensor_model_parallel_size: int,
    pipeline_model_parallel_size: int,
144
    backend: Optional[str] = None,
145
146
147
148
149
) -> 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.
    """
150
151
    # get the backend of _DEVICE_WORLD_GROUP
    backend = backend or torch.distributed.get_backend()
152
153
    if not model_parallel_is_initialized():
        initialize_model_parallel(tensor_model_parallel_size,
154
                                  pipeline_model_parallel_size, backend)
155
156
157
158
159
160
161
162
163
164
165
166
167
168
        return

    assert (
        get_tensor_model_parallel_world_size() == tensor_model_parallel_size
    ), ("tensor parallel group already initialized, but of unexpected size: "
        f"{get_tensor_model_parallel_world_size()=} vs. "
        f"{tensor_model_parallel_size=}")
    assert (get_pipeline_model_parallel_world_size(
    ) == pipeline_model_parallel_size), (
        "pipeline parallel group already initialized, but of unexpected size: "
        f"{get_pipeline_model_parallel_world_size()=} vs. "
        f"{pipeline_model_parallel_size=}")


Zhuohan Li's avatar
Zhuohan Li committed
169
def model_parallel_is_initialized():
170
    """Check if tensor and pipeline parallel groups are initialized."""
171
172
    return (_TENSOR_MODEL_PARALLEL_GROUP is not None
            and _PIPELINE_MODEL_PARALLEL_GROUP is not None)
Zhuohan Li's avatar
Zhuohan Li committed
173
174


175
176
177
178
179
180
def get_cpu_world_group():
    """Get the CPU world group."""
    assert _CPU_WORLD_GROUP is not None, ("CPU world group is not initialized")
    return _CPU_WORLD_GROUP


Zhuohan Li's avatar
Zhuohan Li committed
181
182
def get_tensor_model_parallel_group():
    """Get the tensor model parallel group the caller rank belongs to."""
183
    assert _TENSOR_MODEL_PARALLEL_GROUP is not None, (
184
        "tensor model parallel group is not initialized")
Zhuohan Li's avatar
Zhuohan Li committed
185
186
187
188
189
    return _TENSOR_MODEL_PARALLEL_GROUP


def get_pipeline_model_parallel_group():
    """Get the pipeline model parallel group the caller rank belongs to."""
190
191
    assert _PIPELINE_MODEL_PARALLEL_GROUP is not None, (
        "pipeline model parallel group is not initialized")
Zhuohan Li's avatar
Zhuohan Li committed
192
193
194
195
196
    return _PIPELINE_MODEL_PARALLEL_GROUP


def get_tensor_model_parallel_world_size():
    """Return world size for the tensor model parallel group."""
197
198
    return torch.distributed.get_world_size(
        group=get_tensor_model_parallel_group())
Zhuohan Li's avatar
Zhuohan Li committed
199
200
201
202


def get_pipeline_model_parallel_world_size():
    """Return world size for the pipeline model parallel group."""
203
204
    return torch.distributed.get_world_size(
        group=get_pipeline_model_parallel_group())
Zhuohan Li's avatar
Zhuohan Li committed
205
206
207
208
209
210
211
212
213


def get_tensor_model_parallel_rank():
    """Return my rank for the tensor model parallel group."""
    return torch.distributed.get_rank(group=get_tensor_model_parallel_group())


def get_pipeline_model_parallel_rank():
    """Return my rank for the pipeline model parallel group."""
214
215
    return torch.distributed.get_rank(
        group=get_pipeline_model_parallel_group())
Zhuohan Li's avatar
Zhuohan Li committed
216
217
218
219
220
221
222
223
224
225
226
227
228


def get_tensor_model_parallel_src_rank():
    """Calculate the global rank corresponding to the first local rank
    in the tensor model parallel group."""
    global_rank = torch.distributed.get_rank()
    local_world_size = get_tensor_model_parallel_world_size()
    return (global_rank // local_world_size) * local_world_size


def get_pipeline_model_parallel_first_rank():
    """Return the global rank of the first process in the pipeline for the
    current tensor parallel group"""
229
230
    assert _PIPELINE_GLOBAL_RANKS is not None, (
        "Pipeline parallel group is not initialized")
Zhuohan Li's avatar
Zhuohan Li committed
231
232
233
234
235
236
    return _PIPELINE_GLOBAL_RANKS[0]


def get_pipeline_model_parallel_last_rank():
    """Return the global rank of the last process in the pipeline for the
    current tensor parallel group"""
237
238
    assert _PIPELINE_GLOBAL_RANKS is not None, (
        "Pipeline parallel group is not initialized")
Zhuohan Li's avatar
Zhuohan Li committed
239
240
241
    last_rank_local = get_pipeline_model_parallel_world_size() - 1
    return _PIPELINE_GLOBAL_RANKS[last_rank_local]

Zhuohan Li's avatar
Zhuohan Li committed
242

Zhuohan Li's avatar
Zhuohan Li committed
243
244
def get_pipeline_model_parallel_next_rank():
    """Return the global rank that follows the caller in the pipeline"""
245
246
    assert _PIPELINE_GLOBAL_RANKS is not None, (
        "Pipeline parallel group is not initialized")
Zhuohan Li's avatar
Zhuohan Li committed
247
248
249
250
251
252
    rank_in_pipeline = get_pipeline_model_parallel_rank()
    world_size = get_pipeline_model_parallel_world_size()
    return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline + 1) % world_size]


def get_pipeline_model_parallel_prev_rank():
253
    """Return the global rank that precedes the caller in the pipeline"""
254
255
    assert _PIPELINE_GLOBAL_RANKS is not None, (
        "Pipeline parallel group is not initialized")
Zhuohan Li's avatar
Zhuohan Li committed
256
257
258
259
260
261
    rank_in_pipeline = get_pipeline_model_parallel_rank()
    world_size = get_pipeline_model_parallel_world_size()
    return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline - 1) % world_size]


def destroy_model_parallel():
262
    """Set the groups to none and destroy them."""
Zhuohan Li's avatar
Zhuohan Li committed
263
    global _TENSOR_MODEL_PARALLEL_GROUP
264
265
    if _TENSOR_MODEL_PARALLEL_GROUP:
        torch.distributed.destroy_process_group(_TENSOR_MODEL_PARALLEL_GROUP)
Zhuohan Li's avatar
Zhuohan Li committed
266
267
    _TENSOR_MODEL_PARALLEL_GROUP = None
    global _PIPELINE_MODEL_PARALLEL_GROUP
268
269
    if _PIPELINE_MODEL_PARALLEL_GROUP:
        torch.distributed.destroy_process_group(_PIPELINE_MODEL_PARALLEL_GROUP)
Zhuohan Li's avatar
Zhuohan Li committed
270
    _PIPELINE_MODEL_PARALLEL_GROUP = None
271
272
    global _PIPELINE_GLOBAL_RANKS
    _PIPELINE_GLOBAL_RANKS = None
273
    from vllm.distributed.device_communicators import pynccl_utils
Woosuk Kwon's avatar
Woosuk Kwon committed
274

275
276
    # Destroy the pynccl states if any.
    pynccl_utils.destroy_process_group()
Woosuk Kwon's avatar
Woosuk Kwon committed
277
278


279
280
# Whether to use pynccl for nccl all reduce.
# We use pynccl for all reduce when using CUDA graph, because torch.distributed
Woosuk Kwon's avatar
Woosuk Kwon committed
281
# is not well supported by CUDA graph.
282
_ENABLE_PYNCCL_FOR_ALL_REDUCE = False
Woosuk Kwon's avatar
Woosuk Kwon committed
283
284
285


@contextlib.contextmanager
286
def with_pynccl_for_all_reduce():
287
    from vllm.distributed.device_communicators import pynccl_utils
288
    """use pynccl instead of torch.distributed for all reduce"""
Woosuk Kwon's avatar
Woosuk Kwon committed
289
290
291
    tp_size = get_tensor_model_parallel_world_size()
    if tp_size == 1:
        # No-op.
292
        # NOTE(woosuk): We don't initialize pynccl when tp_size is 1.
Woosuk Kwon's avatar
Woosuk Kwon committed
293
294
        yield
    else:
295
296
297
        global _ENABLE_PYNCCL_FOR_ALL_REDUCE
        old = _ENABLE_PYNCCL_FOR_ALL_REDUCE
        _ENABLE_PYNCCL_FOR_ALL_REDUCE = True
Woosuk Kwon's avatar
Woosuk Kwon committed
298
299

        stream = torch.cuda.current_stream()
300
        with pynccl_utils.set_pynccl_stream(stream):
Woosuk Kwon's avatar
Woosuk Kwon committed
301
            yield
302
        _ENABLE_PYNCCL_FOR_ALL_REDUCE = old
Woosuk Kwon's avatar
Woosuk Kwon committed
303
304


305
306
307
308
def is_pynccl_enabled_for_all_reduce():
    """check if pynccl is enabled for all reduce"""
    global _ENABLE_PYNCCL_FOR_ALL_REDUCE
    return _ENABLE_PYNCCL_FOR_ALL_REDUCE