Skip to content

Commit 8eb0a1d

Browse files
authored
[Doc] Polish example for torchrun dp (vllm-project#25899)
1 parent fea3e47 commit 8eb0a1d

File tree

1 file changed

+10
-6
lines changed

1 file changed

+10
-6
lines changed

examples/offline_inference/torchrun_dp_example.py

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,11 @@
44
experimental support for data-parallel inference with torchrun
55
Note the data load balancing and distribution is done out of the vllm engine,
66
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+
```
712
"""
813

914
from vllm import LLM, SamplingParams
@@ -14,7 +19,7 @@
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
2025
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
@@ -45,14 +50,13 @@
4550

4651
outputs = llm.generate(prompts, sampling_params)
4752

48-
49-
# all ranks will have the same outputs
50-
print("-" * 50)
5153
for output in outputs:
5254
prompt = output.prompt
5355
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+
5660
"""
5761
Further tips:
5862

0 commit comments

Comments
 (0)