A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Doc] Polish example for torchrun dp (#25899)
Zhuohan Li committed
Sep 29, 2025 at 14:31 UTC
8eb0a1d90621927538697f75b4e17c6f79153b4d
1 file changed
+10
-6
examples/offline_inference/torchrun_dp_example.py
+10
-6
@@ -4,6 +4,11 @@
4
experimental support for data-parallel inference with torchrun
5
Note the data load balancing and distribution is done out of the vllm engine,
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
+```
12
"""
13
14
from vllm import LLM, SamplingParams
@@ -14,7 +19,7 @@ prompts = [
19
"The president of the United States is",
20
"The capital of France is",
21
"The future of AI is",
17
-] * 50
22
+]
23
24
# Create sampling parameters, the same across all ranks
25
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
@@ -45,14 +50,13 @@ prompts = [
50
51
outputs = llm.generate(prompts, sampling_params)
52
48
-
49
-# all ranks will have the same outputs
50
-print("-" * 50)
53
for output in outputs:
54
prompt = output.prompt
55
generated_text = output.outputs[0].text
54
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}\n")
55
- print("-" * 50)
56
+ print(
57
+ f"DP Rank: {dp_rank} Prompt: {prompt!r}\nGenerated text: {generated_text!r}\n"
58
+ )
59
+
60
"""
61
Further tips:
62