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

Launch 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 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). 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:

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:

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:

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:

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:

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:

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:

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:

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) – 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) – Implements the Transformer class and forward-pass logic used by each rank.
• inference/configs/*.json – Model hyper-parameter definitions (e.g., config_16B.json, config_236B.json).
• [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 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 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}