@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 137 lines 3.97 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4
5 from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm
6 from vllm.distributed.parallel_state import (
7 _init_stateless_group,
8 _node_count,
9 get_pcp_group,
10 get_pp_group,
11 get_tp_group,
12 get_world_group,
13 )
14 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
15
16 _STANDBY_WORLD: StatelessGroupCoordinator | None = None
17 _STANDBY_WORLD_NODE_COUNT: int | None = None
18 _STANDBY_DP: StatelessGroupCoordinator | None = None
19 _STANDBY_EP: StatelessGroupCoordinator | None = None
20 _STANDBY_EPLB: StatelessGroupCoordinator | None = None
21
22
23 def get_standby_dp_group() -> StatelessGroupCoordinator | None:
24 return _STANDBY_DP
25
26
27 def get_standby_ep_group() -> StatelessGroupCoordinator | None:
28 return _STANDBY_EP
29
30
31 def get_standby_eplb_group() -> StatelessGroupCoordinator | None:
32 return _STANDBY_EPLB
33
34
35 def get_standby_world_group() -> StatelessGroupCoordinator | None:
36 return _STANDBY_WORLD
37
38
39 def create_standby_groups(
40 new_dp_size: int,
41 new_world_size_across_dp: int,
42 master_ip: str,
43 coord_store_port: int,
44 use_all2all: bool,
45 enable_eplb: bool = True,
46 backend: str | None = None,
47 ) -> None:
48 global \
49 _STANDBY_WORLD, \
50 _STANDBY_WORLD_NODE_COUNT, \
51 _STANDBY_DP, \
52 _STANDBY_EP, \
53 _STANDBY_EPLB
54
55 from vllm.distributed.utils import get_cached_tcp_store_client
56
57 assert new_world_size_across_dp == torch.distributed.get_world_size() * new_dp_size
58 world_group = get_world_group()
59 assert isinstance(world_group, StatelessGroupCoordinator)
60 backend = backend or world_group.backend
61
62 coord_store = get_cached_tcp_store_client(master_ip, coord_store_port)
63
64 standby_world_ranks = [list(range(new_world_size_across_dp))]
65 _STANDBY_WORLD = _init_stateless_group(
66 standby_world_ranks,
67 "world",
68 master_ip,
69 backend,
70 use_device_communicator=False,
71 coord_store=coord_store,
72 )
73 _STANDBY_WORLD_NODE_COUNT = _node_count(_STANDBY_WORLD.tcp_store_group)
74
75 tp_size = get_tp_group().world_size
76 pcp_size = get_pcp_group().world_size
77 pp_size = get_pp_group().world_size
78
79 all_ranks = torch.arange(new_world_size_across_dp).reshape(
80 -1, new_dp_size, pp_size, pcp_size, tp_size
81 )
82 standby_dp_ranks = all_ranks.transpose(1, 4).reshape(-1, new_dp_size).unbind(0)
83 standby_dp_ranks = [x.tolist() for x in standby_dp_ranks]
84
85 # Deferred to commit so the warm-up runs while the engine is paused.
86 with defer_comm_warmup_on_rocm():
87 _STANDBY_DP = _init_stateless_group(
88 standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store
89 )
90
91 standby_ep_ranks = (
92 all_ranks.transpose(1, 2)
93 .reshape(-1, new_dp_size * pcp_size * tp_size)
94 .unbind(0)
95 )
96 standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
97 _STANDBY_EP = _init_stateless_group(
98 standby_ep_ranks,
99 "ep",
100 master_ip,
101 backend,
102 coord_store,
103 use_all2all=use_all2all,
104 )
105
106 if enable_eplb:
107 _STANDBY_EPLB = _init_stateless_group(
108 standby_ep_ranks,
109 "eplb",
110 master_ip,
111 backend,
112 coord_store=coord_store,
113 )
114
115
116 def pop_standby_groups() -> dict:
117 """Return all standby groups and clear the standby state."""
118 global \
119 _STANDBY_WORLD, \
120 _STANDBY_WORLD_NODE_COUNT, \
121 _STANDBY_DP, \
122 _STANDBY_EP, \
123 _STANDBY_EPLB
124
125 result = dict(
126 world=_STANDBY_WORLD,
127 dp=_STANDBY_DP,
128 ep=_STANDBY_EP,
129 eplb=_STANDBY_EPLB,
130 node_count=_STANDBY_WORLD_NODE_COUNT,
131 )
132 _STANDBY_WORLD = None
133 _STANDBY_WORLD_NODE_COUNT = None
134 _STANDBY_DP = None
135 _STANDBY_EP = None
136 _STANDBY_EPLB = None
137 return result