File tree Expand file tree Collapse file tree 1 file changed +10
-6
lines changed
examples/offline_inference Expand file tree Collapse file tree 1 file changed +10
-6
lines changed Original file line number Diff line number Diff line change 4
4
experimental support for data-parallel inference with torchrun
5
5
Note the data load balancing and distribution is done out of the vllm engine,
6
6
no internal lb supported in external_launcher mode.
7
+
8
+ To run this example:
9
+ ```bash
10
+ $ torchrun --nproc-per-node=2 examples/offline_inference/torchrun_dp_example.py
11
+ ```
7
12
"""
8
13
9
14
from vllm import LLM , SamplingParams
14
19
"The president of the United States is" ,
15
20
"The capital of France is" ,
16
21
"The future of AI is" ,
17
- ] * 50
22
+ ]
18
23
19
24
# Create sampling parameters, the same across all ranks
20
25
sampling_params = SamplingParams (temperature = 0.8 , top_p = 0.95 )
45
50
46
51
outputs = llm .generate (prompts , sampling_params )
47
52
48
-
49
- # all ranks will have the same outputs
50
- print ("-" * 50 )
51
53
for output in outputs :
52
54
prompt = output .prompt
53
55
generated_text = output .outputs [0 ].text
54
- print (f"Prompt: { prompt !r} \n Generated text: { generated_text !r} \n " )
55
- print ("-" * 50 )
56
+ print (
57
+ f"DP Rank: { dp_rank } Prompt: { prompt !r} \n Generated text: { generated_text !r} \n "
58
+ )
59
+
56
60
"""
57
61
Further tips:
58
62
You can’t perform that action at this time.
0 commit comments