|
1 | | -# |
2 | | -# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. |
3 | | -# This file is a part of the vllm-ascend project. |
4 | | -# Adapted from vllm-project/vllm/examples/offline_inference/data_parallel.py |
5 | 1 | # SPDX-License-Identifier: Apache-2.0 |
6 | | -# usage: |
7 | | -# python examples/offline_inference_data_parallel.py |
8 | | -# we need to have a launcher to create multiple data parallel |
9 | | -# ranks. And each rank will create a vLLM instance to process its own prompts. |
| 2 | +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | +""" |
| 4 | +Usage: |
| 5 | +Single node: |
| 6 | + python examples/offline_inference/data_parallel.py \ |
| 7 | + --model="ibm-research/PowerMoE-3b" \ |
| 8 | + --dp-size=2 \ |
| 9 | + --tp-size=2 |
| 10 | +
|
| 11 | +Multi-node: |
| 12 | + Node 0 (assume the node has ip of 10.99.48.128): |
| 13 | + python examples/offline_inference/data_parallel.py \ |
| 14 | + --model="ibm-research/PowerMoE-3b" \ |
| 15 | + --dp-size=2 \ |
| 16 | + --tp-size=2 \ |
| 17 | + --node-size=2 \ |
| 18 | + --node-rank=0 \ |
| 19 | + --master-addr=10.99.48.128 \ |
| 20 | + --master-port=13345 |
| 21 | + Node 1: |
| 22 | + python examples/offline_inference/data_parallel.py \ |
| 23 | + --model="ibm-research/PowerMoE-3b" \ |
| 24 | + --dp-size=2 \ |
| 25 | + --tp-size=2 \ |
| 26 | + --node-size=2 \ |
| 27 | + --node-rank=1 \ |
| 28 | + --master-addr=10.99.48.128 \ |
| 29 | + --master-port=13345 |
| 30 | +""" |
10 | 31 |
|
11 | | -import gc |
12 | 32 | import os |
| 33 | +from time import sleep |
| 34 | + |
| 35 | +from vllm import LLM, SamplingParams |
| 36 | +from vllm.utils import get_open_port |
| 37 | + |
| 38 | + |
| 39 | +def parse_args(): |
| 40 | + import argparse |
13 | 41 |
|
| 42 | + parser = argparse.ArgumentParser(description="Data Parallel Inference") |
| 43 | + parser.add_argument( |
| 44 | + "--model", |
| 45 | + type=str, |
| 46 | + default="ibm-research/PowerMoE-3b", |
| 47 | + help="Model name or path", |
| 48 | + ) |
| 49 | + parser.add_argument("--dp-size", |
| 50 | + type=int, |
| 51 | + default=2, |
| 52 | + help="Data parallel size") |
| 53 | + parser.add_argument("--tp-size", |
| 54 | + type=int, |
| 55 | + default=2, |
| 56 | + help="Tensor parallel size") |
| 57 | + parser.add_argument("--node-size", |
| 58 | + type=int, |
| 59 | + default=1, |
| 60 | + help="Total number of nodes") |
| 61 | + parser.add_argument("--node-rank", |
| 62 | + type=int, |
| 63 | + default=0, |
| 64 | + help="Rank of the current node") |
| 65 | + parser.add_argument("--master-addr", |
| 66 | + type=str, |
| 67 | + default="", |
| 68 | + help="Master node IP address") |
| 69 | + parser.add_argument("--master-port", |
| 70 | + type=int, |
| 71 | + default=0, |
| 72 | + help="Master node port") |
| 73 | + parser.add_argument("--enforce-eager", |
| 74 | + action="store_true", |
| 75 | + help="Enforce eager mode execution.") |
| 76 | + parser.add_argument("--trust-remote-code", |
| 77 | + action="store_true", |
| 78 | + help="Trust remote code.") |
| 79 | + return parser.parse_args() |
14 | 80 |
|
15 | | -def main(): |
16 | | - dp_rank = int(os.environ['RANK']) |
17 | | - local_rank = int(os.environ['LOCAL_RANK']) |
18 | | - dp_size = int(os.environ['WORLD_SIZE']) |
19 | | - master_addr = os.environ['MASTER_ADDR'] |
20 | | - master_port = os.environ['MASTER_PORT'] |
21 | | - tp_size = 1 |
22 | | - etp_size = 1 |
23 | 81 |
|
24 | | - os.environ["VLLM_DP_RANK"] = str(dp_rank) |
| 82 | +def main( |
| 83 | + model, |
| 84 | + dp_size, |
| 85 | + local_dp_rank, |
| 86 | + global_dp_rank, |
| 87 | + dp_master_ip, |
| 88 | + dp_master_port, |
| 89 | + GPUs_per_dp_rank, |
| 90 | + enforce_eager, |
| 91 | + trust_remote_code, |
| 92 | +): |
| 93 | + os.environ["VLLM_DP_RANK"] = str(global_dp_rank) |
| 94 | + os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank) |
25 | 95 | os.environ["VLLM_DP_SIZE"] = str(dp_size) |
26 | | - os.environ["VLLM_DP_MASTER_IP"] = master_addr |
27 | | - os.environ["VLLM_DP_MASTER_PORT"] = master_port |
28 | | - os.environ["ASCEND_RT_VISIBLE_DEVICES"] = ",".join( |
29 | | - str(i) |
30 | | - for i in range(local_rank * tp_size, (local_rank + 1) * tp_size)) |
| 96 | + os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip |
| 97 | + os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port) |
31 | 98 |
|
32 | | - import torch |
33 | | - from vllm import LLM, SamplingParams |
34 | | - from vllm.distributed.parallel_state import ( |
35 | | - destroy_distributed_environment, destroy_model_parallel) |
| 99 | + # CUDA_VISIBLE_DEVICES for each DP rank is set automatically inside the |
| 100 | + # engine processes. |
36 | 101 |
|
| 102 | + # Sample prompts. |
37 | 103 | prompts = [ |
38 | 104 | "Hello, my name is", |
39 | 105 | "The president of the United States is", |
40 | 106 | "The capital of France is", |
41 | 107 | "The future of AI is", |
42 | | - ] * 4 |
| 108 | + ] * 100 |
43 | 109 |
|
44 | | - promts_per_rank = len(prompts) // dp_size |
45 | | - start = dp_rank * promts_per_rank |
46 | | - end = start + promts_per_rank |
47 | | - prompts = prompts[start:end] |
| 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)] |
48 | 121 | if len(prompts) == 0: |
| 122 | + # if any rank has no prompts to process, |
| 123 | + # we need to set a placeholder prompt |
49 | 124 | prompts = ["Placeholder"] |
50 | | - print(f"DP rank {dp_rank} needs to process {len(prompts)} prompts") |
51 | | - num_seqs = len(prompts) |
| 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.0, |
| 133 | + max_tokens=32, |
| 134 | + ) |
52 | 135 |
|
53 | | - sampling_params = SamplingParams(temperature=0.8, |
54 | | - top_p=0.95, |
55 | | - max_tokens=4, |
56 | | - min_tokens=4) |
57 | 136 | # Create an LLM. |
58 | | - llm = LLM(model="deepseek-ai/DeepSeek-V2-Lite-Chat", |
59 | | - tensor_parallel_size=tp_size, |
60 | | - trust_remote_code=True, |
61 | | - max_model_len=4096, |
62 | | - max_num_seqs=num_seqs, |
63 | | - additional_config={ |
64 | | - 'expert_tensor_parallel_size': etp_size, |
65 | | - 'torchair_graph_config': { |
66 | | - 'enabled': False, |
67 | | - }, |
68 | | - }) |
| 137 | + llm = LLM( |
| 138 | + model=model, |
| 139 | + tensor_parallel_size=GPUs_per_dp_rank, |
| 140 | + enforce_eager=enforce_eager, |
| 141 | + trust_remote_code=trust_remote_code, |
| 142 | + distributed_executor_backend="mp", |
| 143 | + max_model_len=2048, |
| 144 | + max_num_batched_tokens=2048, |
| 145 | + max_num_seqs=16, |
| 146 | + enable_prefix_caching=False, |
| 147 | + enable_expert_parallel=True, |
| 148 | + gpu_memory_utilization=0.9, |
| 149 | + additional_config={ |
| 150 | + "ascend_scheduler_config": { |
| 151 | + "enabled": True |
| 152 | + }, |
| 153 | + "torchair_graph_config": { |
| 154 | + "enabled": False, |
| 155 | + "enable_multistream_shared_expert": False |
| 156 | + }, |
| 157 | + }, |
| 158 | + ) |
69 | 159 |
|
70 | 160 | outputs = llm.generate(prompts, sampling_params) |
71 | | - for output in outputs: |
| 161 | + # Print the outputs. |
| 162 | + for i, output in enumerate(outputs): |
| 163 | + if i >= 5: |
| 164 | + # print only 5 outputs |
| 165 | + break |
72 | 166 | prompt = output.prompt |
73 | 167 | generated_text = output.outputs[0].text |
74 | | - print(f"DP rank {dp_rank}, Prompt: {prompt!r}, " |
| 168 | + print(f"DP rank {global_dp_rank}, Prompt: {prompt!r}, " |
75 | 169 | f"Generated text: {generated_text!r}") |
76 | 170 |
|
77 | | - del llm |
78 | | - destroy_model_parallel() |
79 | | - destroy_distributed_environment() |
80 | | - gc.collect() |
81 | | - torch.npu.empty_cache() |
| 171 | + # Give engines time to pause their processing loops before exiting. |
| 172 | + sleep(1) |
82 | 173 |
|
83 | 174 |
|
84 | 175 | if __name__ == "__main__": |
85 | | - main() |
| 176 | + args = parse_args() |
| 177 | + |
| 178 | + dp_size = args.dp_size |
| 179 | + tp_size = args.tp_size |
| 180 | + node_size = args.node_size |
| 181 | + node_rank = args.node_rank |
| 182 | + |
| 183 | + if node_size == 1: |
| 184 | + dp_master_ip = "127.0.0.1" |
| 185 | + dp_master_port = get_open_port() |
| 186 | + else: |
| 187 | + dp_master_ip = args.master_addr |
| 188 | + dp_master_port = args.master_port |
| 189 | + |
| 190 | + assert dp_size % node_size == 0, "dp_size should be divisible by node_size" |
| 191 | + dp_per_node = dp_size // node_size |
| 192 | + |
| 193 | + from multiprocessing import Process |
| 194 | + |
| 195 | + procs = [] |
| 196 | + for local_dp_rank, global_dp_rank in enumerate( |
| 197 | + range(node_rank * dp_per_node, (node_rank + 1) * dp_per_node)): |
| 198 | + proc = Process( |
| 199 | + target=main, |
| 200 | + args=( |
| 201 | + args.model, |
| 202 | + dp_size, |
| 203 | + local_dp_rank, |
| 204 | + global_dp_rank, |
| 205 | + dp_master_ip, |
| 206 | + dp_master_port, |
| 207 | + tp_size, |
| 208 | + args.enforce_eager, |
| 209 | + args.trust_remote_code, |
| 210 | + ), |
| 211 | + ) |
| 212 | + proc.start() |
| 213 | + procs.append(proc) |
| 214 | + exit_code = 0 |
| 215 | + for proc in procs: |
| 216 | + proc.join(timeout=3000) |
| 217 | + if proc.exitcode is None: |
| 218 | + print( |
| 219 | + f"Killing process {proc.pid} that didn't stop within 5 minutes." |
| 220 | + ) |
| 221 | + proc.kill() |
| 222 | + exit_code = 1 |
| 223 | + elif proc.exitcode: |
| 224 | + exit_code = proc.exitcode |
| 225 | + |
| 226 | + exit(exit_code) |
0 commit comments