diff --git a/examples/offline_inference/torchrun_example.py b/examples/offline_inference/torchrun_example.py index 3d3d7946c..2960d3299 100644 --- a/examples/offline_inference/torchrun_example.py +++ b/examples/offline_inference/torchrun_example.py @@ -4,9 +4,10 @@ experimental support for tensor-parallel inference with torchrun, see https://github.com/vllm-project/vllm/issues/11400 for the motivation and use case for this example. -run the script with `torchrun --nproc-per-node=2 torchrun_example.py`, -the argument 2 should match the `tensor_parallel_size` below. -see `tests/distributed/test_torchrun_example.py` for the unit test. +run the script with `torchrun --nproc-per-node=4 torchrun_example.py`, +the argument `4` should match the product of `tensor_parallel_size` and +`pipeline_parallel_size` below. see `tests/distributed/test_torchrun_example.py` +for the unit test. """ import torch.distributed as dist