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 44experimental support for data-parallel inference with torchrun
55Note the data load balancing and distribution is done out of the vllm engine,
66no 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+ ```
712"""
813
914from vllm import LLM , SamplingParams
1419 "The president of the United States is" ,
1520 "The capital of France is" ,
1621 "The future of AI is" ,
17- ] * 50
22+ ]
1823
1924# Create sampling parameters, the same across all ranks
2025sampling_params = SamplingParams (temperature = 0.8 , top_p = 0.95 )
4550
4651outputs = llm .generate (prompts , sampling_params )
4752
48-
49- # all ranks will have the same outputs
50- print ("-" * 50 )
5153for output in outputs :
5254 prompt = output .prompt
5355 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+
5660"""
5761Further tips:
5862
You can’t perform that action at this time.
0 commit comments