# How to Set Up Distributed Inference with torchrun for Multi-Node Deployment

> Learn how to set up distributed inference with torchrun for multi-node deployment. Effortlessly launch generate.py using torchrun flags for efficient model scaling and sharded checkpoint loading across nodes.

- Repository: [DeepSeek/DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3)
- Tags: how-to-guide
- Published: 2026-02-26

---

**Launch [`inference/generate.py`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/generate.py) via `torchrun` with `--nnodes` and `--nproc_per_node` flags; the script automatically detects `WORLD_SIZE` and `RANK` to initialize NCCL process groups and assigns each process to its local GPU for sharded checkpoint loading.**

The DeepSeek-V3 repository provides a pure-PyTorch inference pipeline that scales seamlessly from a single GPU to multi-node clusters. Setting up distributed inference with torchrun for multi-node deployment requires no code modifications—the [`inference/generate.py`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/generate.py) script reads launcher-injected environment variables to coordinate prompt broadcasting and model parallelism across GPUs.

## How Distributed Inference Works in DeepSeek-V3

The core distributed logic resides in [[`inference/generate.py`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/generate.py)](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/generate.py). The script detects the runtime topology, initializes inter-process communication, and ensures every rank generates from identical inputs.

### Rank and World Size Detection

When `torchrun` spawns processes, it injects `WORLD_SIZE`, `RANK`, and `LOCAL_RANK` into the environment. The entry point reads these variables to determine the global topology:

```python
world_size = int(os.getenv("WORLD_SIZE", "1"))
rank       = int(os.getenv("RANK", "0"))
local_rank = int(os.getenv("LOCAL_RANK", "0"))

```

`world_size` represents the total number of GPUs across all nodes, while `local_rank` identifies the specific GPU index on the current machine.

### NCCL Process Group Initialization

When `world_size` exceeds one, the script initializes a distributed backend so that ranks can communicate during inference:

```python
if world_size > 1:
    dist.init_process_group("nccl")

```

This creates an NCCL-backed process group that enables high-speed GPU-to-GPU communication over the network fabric.

### GPU Device Assignment and Checkpoint Sharding

Each process pins itself to its designated GPU to ensure deterministic device placement:

```python
torch.cuda.set_device(local_rank)

```

The DeepSeek-V3 checkpoint is sharded per rank using the naming convention `model{rank}-mp{world_size}.safetensors`. After a warm-up forward pass, the script loads the correct shard for the current process:

```python
load_model(model,
           os.path.join(ckpt_path,
                        f"model{rank}-mp{world_size}.safetensors"))

```

This **tensor-parallel** sharding allows the model weights to be distributed across GPUs, reducing per-device memory requirements.

### Synchronized Prompt Broadcasting

To ensure all nodes generate from the same input in interactive mode, rank 0 broadcasts the user prompt via `dist.broadcast_object_list`:

```python
if world_size == 1:
    prompt = input(">>> ")
elif rank == 0:
    prompt = input(">>> ")
    objects = [prompt]
    dist.broadcast_object_list(objects, 0)
else:
    objects = [None]
    dist.broadcast_object_list(objects, 0)
    prompt = objects[0]

```

Non-master ranks receive the prompt through this broadcast, guaranteeing synchronized generation across the cluster.

## Launching Multi-Node Inference with torchrun

`torchrun` (recommended for PyTorch ≥ 1.10) automates environment variable injection. Use the following command template to deploy across multiple physical machines:

```bash
torchrun \
  --nnodes=2 \
  --nproc_per_node=8 \
  --node_rank=$NODE_RANK \
  --master_addr=$MASTER_ADDR \
  --master_port=29500 \
  inference/generate.py \
    --ckpt-path /path/to/checkpoint \
    --config inference/configs/config_236B.json \
    --interactive

```

Parameter reference:

- **`--nnodes`** – Total number of physical machines participating in the job.
- **`--nproc_per_node`** – Number of GPU processes to spawn per machine (typically one per GPU).
- **`--node_rank`** – Unique integer identifier for each node (0 for the first node, 1 for the second, etc.).
- **`--master_addr`** – IP address or hostname of the rank-0 node that coordinates the cluster.
- **`--master_port`** – Free TCP port on the master node for rendezvous traffic.

For single-node, multi-GPU runs, omit `--nnodes` and specify only the local GPU count:

```bash
torchrun --nproc_per_node=4 inference/generate.py \
  --ckpt-path /path/to/checkpoint \
  --config inference/configs/config_236B.json \
  --interactive

```

## Running Batch Inference Across Nodes

For non-interactive, file-based generation, replace `--interactive` with `--input-file`. Rank 0 reads the file and distributes tokenized inputs across all ranks:

```bash
torchrun --nnodes=2 --nproc_per_node=8 --node_rank=$NODE_RANK \
  --master_addr=$MASTER_ADDR --master_port=29500 \
  inference/generate.py \
    --ckpt-path /path/to/checkpoint \
    --config inference/configs/config_236B.json \
    --input-file prompts.txt

```

The script broadcasts prompts once at startup, then processes the batch in parallel across the cluster.

## Key Files and Architecture

Understanding the repository structure helps debug distributed runs:

- **[[`inference/generate.py`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/generate.py)](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/generate.py)** – Entry point containing all distributed-setup logic, prompt broadcasting, and checkpoint loading.
- **[[`inference/model.py`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py)](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py)** – Implements the `Transformer` class and forward-pass logic used by each rank.
- **[`inference/configs/*.json`](https://github.com/deepseek-ai/DeepSeek-V3/tree/main/inference/configs)** – Model hyper-parameter definitions (e.g., [`config_16B.json`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/config_16B.json), [`config_236B.json`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/config_236B.json)).
- **[[`inference/convert.py`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/convert.py)](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/convert.py)** – Utility for converting consolidated checkpoints into the sharded `safetensors` format required for multi-GPU inference.

## Summary

- **Environment auto-detection** – [`inference/generate.py`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/generate.py) reads `WORLD_SIZE`, `RANK`, and `LOCAL_RANK` injected by `torchrun` to configure the distributed topology.
- **NCCL backend** – Process groups initialize automatically when `world_size > 1`, enabling efficient multi-node GPU communication.
- **Sharded checkpoints** – Each rank loads `model{rank}-mp{world_size}.safetensors`, implementing tensor parallelism without manual partitioning.
- **Prompt synchronization** – `dist.broadcast_object_list` ensures all ranks generate from identical inputs in interactive mode.
- **Launcher simplicity** – No code changes are required; simply adjust `torchrun` flags (`--nnodes`, `--nproc_per_node`, `--node_rank`) to scale from single-node to multi-node clusters.

## Frequently Asked Questions

### What environment variables does torchrun set for distributed inference?

`torchrun` automatically exports `WORLD_SIZE` (total GPU count), `RANK` (global process ID), and `LOCAL_RANK` (GPU index on the current node). The DeepSeek-V3 inference script reads these via `os.getenv()` to initialize the NCCL process group and assign devices.

### How does DeepSeek-V3 handle checkpoint loading in multi-GPU setups?

The implementation uses **tensor parallelism** via file sharding. Each rank loads a specific shard named `model{rank}-mp{world_size}.safetensors` from the checkpoint directory. The `load_model()` function in [`inference/generate.py`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/generate.py) constructs the shard path dynamically based on the detected rank and world size.

### Can I run distributed inference on a single node with multiple GPUs?

Yes. Omit the `--nnodes` flag and set `--nproc_per_node` to the number of local GPUs. The script detects `world_size` from the environment and initializes the NCCL backend for intra-node communication without requiring network configuration.

### What is the purpose of prompt broadcasting in generate.py?

Prompt broadcasting ensures **deterministic synchronization** across all ranks. In interactive mode, only rank 0 reads user input from `stdin`; it then uses `dist.broadcast_object_list()` to transmit the prompt string to all other ranks. This guarantees that every GPU in the cluster begins generation from the exact same token sequence.