@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 213 lines 6.36 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Usage:
4 Single node:
5 python examples/features/data_parallel/data_parallel_offline.py \
6 --model="ibm-research/PowerMoE-3b" \
7 -dp=2 \
8 -tp=2
9
10 Multi-node:
11 Node 0 (assume the node has ip of 10.99.48.128):
12 python examples/features/data_parallel/data_parallel_offline.py \
13 --model="ibm-research/PowerMoE-3b" \
14 -dp=2 \
15 -tp=2 \
16 --dp-num-nodes=2 \
17 --dp-node-rank=0 \
18 --dp-master-addr=10.99.48.128 \
19 --dp-master-port=13345
20 Node 1:
21 python examples/features/data_parallel/data_parallel_offline.py \
22 --model="ibm-research/PowerMoE-3b" \
23 -dp=2 \
24 -tp=2 \
25 --dp-num-nodes=2 \
26 --dp-node-rank=1 \
27 --dp-master-addr=10.99.48.128 \
28 --dp-master-port=13345
29 """
30
31 import os
32 from time import sleep
33
34 from vllm import LLM, EngineArgs, SamplingParams
35 from vllm.platforms import current_platform
36 from vllm.utils.argparse_utils import FlexibleArgumentParser
37 from vllm.utils.network_utils import get_open_port
38
39
40 def create_parser():
41 parser = FlexibleArgumentParser(description="Data Parallel Inference")
42
43 # Add all engine args
44 EngineArgs.add_cli_args(parser)
45 parser.set_defaults(
46 model="ibm-research/PowerMoE-3b",
47 enable_expert_parallel=True,
48 )
49
50 # Add DP-specific args (separate from engine args to avoid conflicts)
51 parser.add_argument(
52 "--dp-num-nodes",
53 type=int,
54 default=1,
55 help="Total number of nodes for data parallel.",
56 )
57 parser.add_argument(
58 "--dp-node-rank",
59 type=int,
60 default=0,
61 help="Rank of the current node for data parallel.",
62 )
63 parser.add_argument(
64 "--dp-master-addr",
65 type=str,
66 default="",
67 help="Master node IP address for DP coordination.",
68 )
69 parser.add_argument(
70 "--dp-master-port",
71 type=int,
72 default=0,
73 help="Master node port for DP coordination.",
74 )
75 parser.add_argument(
76 "--timeout",
77 type=int,
78 default=300,
79 help="Number of seconds before unresponsive process is killed.",
80 )
81
82 return parser
83
84
85 def main(
86 dp_size,
87 local_dp_rank,
88 global_dp_rank,
89 dp_master_ip,
90 dp_master_port,
91 engine_args,
92 ):
93 os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
94 os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
95 os.environ["VLLM_DP_SIZE"] = str(dp_size)
96 os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
97 os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
98
99 # CUDA_VISIBLE_DEVICES for each DP rank is set automatically inside the
100 # engine processes.
101
102 # Sample prompts.
103 prompts = [
104 "Hello, my name is",
105 "The president of the United States is",
106 "The capital of France is",
107 "The future of AI is",
108 ] * 100
109
110 # with DP, each rank should process different prompts.
111 # usually all the DP ranks process a full dataset,
112 # and each rank processes a different part of the dataset.
113 floor = len(prompts) // dp_size
114 remainder = len(prompts) % dp_size
115
116 # Distribute prompts into even groups.
117 def start(rank):
118 return rank * floor + min(rank, remainder)
119
120 prompts = prompts[start(global_dp_rank) : start(global_dp_rank + 1)]
121 if len(prompts) == 0:
122 # if any rank has no prompts to process,
123 # we need to set a placeholder prompt
124 prompts = ["Placeholder"]
125 print(f"DP rank {global_dp_rank} needs to process {len(prompts)} prompts")
126
127 # Create a sampling params object.
128 # since we are doing data parallel, every rank can have different
129 # sampling params. here we set different max_tokens for different
130 # ranks for demonstration.
131 sampling_params = SamplingParams(
132 temperature=0.8, top_p=0.95, max_tokens=[16, 20][global_dp_rank % 2]
133 )
134
135 # Create an LLM.
136 llm = LLM(**engine_args)
137 outputs = llm.generate(prompts, sampling_params)
138 # Print the outputs.
139 for i, output in enumerate(outputs):
140 if i >= 5:
141 # print only 5 outputs
142 break
143 prompt = output.prompt
144 generated_text = output.outputs[0].text
145 print(
146 f"DP rank {global_dp_rank}, Prompt: {prompt!r}, "
147 f"Generated text: {generated_text!r}"
148 )
149
150 # Give engines time to pause their processing loops before exiting.
151 sleep(1)
152
153
154 if __name__ == "__main__":
155 parser = create_parser()
156 args = vars(parser.parse_args())
157
158 # Extract DP-specific args (pop to remove from engine_args)
159 dp_size = args.pop("data_parallel_size")
160 dp_num_nodes = args.pop("dp_num_nodes")
161 dp_node_rank = args.pop("dp_node_rank")
162 dp_master_addr = args.pop("dp_master_addr")
163 dp_master_port = args.pop("dp_master_port")
164 timeout = args.pop("timeout")
165
166 # Remaining args are engine args
167 engine_args = args
168
169 if dp_num_nodes == 1:
170 dp_master_ip = "127.0.0.1"
171 dp_master_port_val = get_open_port()
172 else:
173 dp_master_ip = dp_master_addr
174 dp_master_port_val = dp_master_port
175
176 assert dp_size % dp_num_nodes == 0, "dp_size should be divisible by dp_num_nodes"
177 dp_per_node = dp_size // dp_num_nodes
178
179 from multiprocessing import Process
180
181 if current_platform.is_rocm():
182 from multiprocessing import set_start_method
183
184 set_start_method("spawn", force=True)
185
186 procs = []
187 for local_dp_rank, global_dp_rank in enumerate(
188 range(dp_node_rank * dp_per_node, (dp_node_rank + 1) * dp_per_node)
189 ):
190 proc = Process(
191 target=main,
192 args=(
193 dp_size,
194 local_dp_rank,
195 global_dp_rank,
196 dp_master_ip,
197 dp_master_port_val,
198 engine_args,
199 ),
200 )
201 proc.start()
202 procs.append(proc)
203 exit_code = 0
204 for proc in procs:
205 proc.join(timeout=timeout)
206 if proc.exitcode is None:
207 print(f"Killing process {proc.pid} that didn't stop within 5 minutes.")
208 proc.kill()
209 exit_code = 1
210 elif proc.exitcode:
211 exit_code = proc.exitcode
212
213 exit(exit_code)