Set up a multi-GPU training environment with uv

Set up a PyTorch environment with uv, then prove that two processes use two GPUs and communicate. A small synthetic optimizer step checks the distributed environment without downloading a model or dataset.

Prerequisites

• uv installed
• A Linux host with two or more NVIDIA GPUs supported by the selected CUDA wheel
• A compatible NVIDIA driver and, for containers, GPU passthrough enabled

This tutorial uses NCCL, PyTorch’s NVIDIA GPU communication backend. Native Windows does not provide NCCL; macOS does not provide CUDA. Treat WSL2 as a separate deployment and verify its GPU and communication support rather than assuming native Linux behavior.

Check the visible hardware

nvidia-smi --query-gpu=index,name,driver_version --format=csv,noheader

Expect two lines, one per GPU, with their driver versions. This proves the driver sees the devices, not that PyTorch can execute on them. A CUDA toolkit installation is unnecessary for this wheel-based environment.

Create the project

uv init --bare --python 3.12 distributed_training
cd distributed_training
uv python pin 3.12

uv reports the created project and pinned interpreter. --bare creates a minimal pyproject.tomlThe standard configuration file for Python projects. Declares the project name, version, dependencies, build system, and tool settings in one place. Learn more → . Python 3.12 is a conservative choice for combining compiled ML packages; a Python version satisfying torch’s metadata may still lack wheels for an extension.

Add these sections to pyproject.toml:

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

[tool.uv.sources]
torch = { index = "pytorch-cu128" }

uv add torch accelerate
uv run python -c "import torch; print(torch.__version__, torch.version.cuda, torch.cuda.device_count())"

The installed torch version includes +cu128, the runtime is 12.8, and the device count should match nvidia-smi. uv run executes the command inside the project’s virtual environment, ensuring all dependencies are installed first.

A zero count can mean a CPU wheel, missing device access, or a driver/library problem. Inspect torch.version.cuda before diagnosing it as a CPU-only installation. The PyTorch installation guide covers other backends.

Check execution and synchronization

Save check_distributed.py:

import argparse

import torch
from accelerate import Accelerator
from accelerate.utils import set_seed

parser = argparse.ArgumentParser()
parser.add_argument("--expected-processes", type=int, required=True)
args = parser.parse_args()
set_seed(42)
accelerator = Accelerator()
assert accelerator.device.type == "cuda", "CUDA is required"
assert accelerator.num_processes == args.expected_processes

# Each process must execute on its assigned GPU and participate in a reduction.
x = torch.tensor([accelerator.process_index + 1.0], device=accelerator.device)
total = accelerator.reduce(x, reduction="sum")
n = accelerator.num_processes
assert total.item() == n * (n + 1) / 2

model = torch.nn.Linear(4, 1).to(accelerator.device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
model, optimizer = accelerator.prepare(model, optimizer)
before = next(model.parameters()).detach().clone()
inputs = torch.ones(2, 4, device=accelerator.device) * (accelerator.process_index + 1)
loss = model(inputs).square().mean()
accelerator.backward(loss)
optimizer.step()
assert not torch.equal(before, next(model.parameters()))

# Different rank inputs produce gradients that DDP must synchronize.
weights = accelerator.gather(next(model.parameters()).detach().reshape(1, -1))
assert torch.allclose(weights, weights[0].expand_as(weights))
print(f"rank={accelerator.process_index} device={accelerator.device} processes={n} passed")
accelerator.end_training()

The process-count assertion prevents a launcher silently running a single process from appearing to pass. The reduction exercises device communication; the optimizer step checks PyTorch’s distributed wrapper and synchronized parameters. Accelerate supplies the launcher and wrapper.

Run on one GPU

uv run python check_distributed.py --expected-processes 1

Expect rank=0 device=cuda processes=1 passed (the device may include :0). An assertion or NCCL error is a failed environment check.

Launch two GPU processes

Pass explicit launch options so the command does not depend on a previously saved Accelerate configuration:

uv run accelerate launch --multi_gpu --num_processes 2 --num_machines 1 --mixed_precision no check_distributed.py --expected-processes 2

Expect one passed line from each rank, with cuda:0 and cuda:1 and processes=2. Output order varies. Both lines must appear, and the command must exit successfully.

The equivalent PyTorch launcher is:

uv run torchrun --standalone --nproc_per_node=2 check_distributed.py --expected-processes 2

It must produce the same two passing ranks. --standalone configures rendezvous for this single host. More than two GPUs or multiple hosts need their own communication checks; passing this probe does not establish multi-node networking.

Reproduce the environment

Commit pyproject.toml, .python-version, and the lockfileA file that records the exact version of every installed package, so everyone working on the project gets identical installs. uv.lock. A collaborator installs the recorded versions with:

uv sync --locked

uv reports the synchronized environment. Rerun the two-process check on each deployment host: the lockfile records Python packages, not host drivers, device access, or GPU topology.