Distributed & Multi-GPU

How to scale your work across multiple GPUs — and when you actually need to.

No deep technical background needed: this page explains the concepts step by step. The code examples can be copy-pasted and adapted.

1. Why parallelise?#

Before reaching for multi-GPU, ask yourself: do I actually need it?

Think of a GPU like a workbench: the model and the data have to fit on the bench (the VRAM) for you to work. You need parallelisation in three situations:

• Your model doesn't fit on one GPU — e.g. a 70B-parameter model in fp16 needs ~140 GB of VRAM, but one H100 card has 94 GB. The model literally cannot be loaded.
• Training is too slow — splitting the work across several GPUs makes each training epoch faster.
• You need bigger batches than one GPU can hold in memory.

You do not need it when your model fits comfortably on a single GPU. Multi-GPU adds complexity, and the GPUs spend time communicating with each other — sometimes a single GPU is actually faster. A 7B model (~14 GB in fp16) runs perfectly on a single H100.

How does a GPU work? (in plain words)#

A regular processor (CPU) is like one very smart worker doing tasks one after another. A GPU is like thousands of simple workers all doing the same small calculation at the same time. Neural networks are mostly enormous multiplication tables — exactly the kind of repetitive work GPUs excel at.

Two numbers matter for us:

• VRAM — the GPU's memory: how big a model (and data) can fit on it
• Memory bandwidth — how fast data moves inside the card: how quickly it works

Curious how GPU programming actually works under the hood? See GPU programming basics from Marc Lelarge's deep learning course (dataflowr / MVA).

What fits on our machines?#

Rule of thumb: a model needs ~2 GB of VRAM per billion parameters in fp16 (standard precision).

If your model doesn't fit, you have two options:

1. Quantize — compress the model to int8 or int4, dividing its memory footprint by 2 or 4 (see AI Tools)
2. Parallelise across GPUs — what this page explains

2. The three parallelisation techniques#

There are three main strategies. Don't worry about mastering all of them — in practice the tools (Accelerate, vLLM) handle the details for you. What matters is understanding the idea behind each, so you pick the right one.

The reference on this topic is the Ultrascale Playbook by HuggingFace — excellent reading if you want to go deeper.

Data parallelism — same model, different data#

The idea: photocopy the model onto every GPU. Each GPU studies a different slice of the data, then they compare notes and average their findings.

            Training batch
           /      |       \
      slice 1  slice 2  slice 3
         |        |        |
    ┌────────┐┌────────┐┌────────┐
    │ GPU 0  ││ GPU 1  ││ GPU 2  │
    │ full   ││ full   ││ full   │
    │ model  ││ model  ││ model  │
    └────────┘└────────┘└────────┘
         \        |        /
        gradients averaged

• Use when: the model fits on one GPU but training takes too long
• Effect: with 4 GPUs, training is roughly 4× faster
• Limit: the model must fit entirely on each GPU

Tensor parallelism — each layer split across GPUs#

The idea: the model is too big for one GPU, so each layer is cut in pieces: GPU 0 holds the left half of every layer, GPU 1 the right half. They constantly exchange intermediate results.

              Input
                |
    ┌───────────┴───────────┐
    │                       │
┌────────────┐       ┌────────────┐
│   GPU 0    │ <───> │   GPU 1    │
│ layer 1: A │       │ layer 1: B │
│ layer 2: A │       │ layer 2: B │
│ layer 3: A │       │ layer 3: B │
└────────────┘       └────────────┘
  GPUs talk to each other at every layer

• Use when: the model is too large for one GPU — mostly for inference
• Cost: lots of communication between GPUs. Our H100 NVL cards have NVLink (a fast direct connection) which makes this efficient — but only between NVLink pairs: on the 4×H100, GPUs (0,1) and (2,3) are linked, other combinations have no NVLink and are much slower (see the NVLink rule)

Pipeline parallelism — layers split into stages#

The idea: like an assembly line. GPU 0 handles the first layers, passes the result to GPU 1 which handles the middle layers, and so on.

        ┌──────────┐   ┌──────────┐   ┌──────────┐
Input ─>│  GPU 0   │──>│  GPU 1   │──>│  GPU 2   │─> Output
        │ layers   │   │ layers   │   │ layers   │
        │  1-16    │   │  17-32   │   │  33-48   │
        └──────────┘   └──────────┘   └──────────┘

• Use when: very large models, usually combined with the other strategies
• Cost: like any assembly line, workers can sit idle waiting for the previous station

Which one do I need?#

In practice, you rarely choose by hand: the tools choose for you. Here is the cheat sheet:

Multi-GPU mainly applies to the 2×H100 and 4×H100. Only the 4×H100 has NVLink, by pairs: a 2-GPU job there should always run on an NVLink pair. The 2×H100 has no NVLink — data parallelism works well, tensor parallelism is slower than on an NVLink pair.

bash

CUDA_VISIBLE_DEVICES=0,1 accelerate launch train.py   # or 2,3

3. The simple way — Accelerate (recommended)#

HuggingFace Accelerate is the easiest path to multi-GPU. You barely change your code — Accelerate figures out the rest.

Install:

pip install accelerate

Adapt your training script — only these lines change:

from accelerate import Accelerator

accelerator = Accelerator()

model, optimizer, dataloader = accelerator.prepare(
    model, optimizer, dataloader
)

for batch in dataloader:
    outputs = model(**batch)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

Configure once (it asks simple questions: how many GPUs, etc.), then launch:

accelerate config
accelerate launch train.py

That's it: your training now runs on all the GPUs you made visible.

4. The PyTorch way — DDP#

If you prefer raw PyTorch, use DistributedDataParallel:

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
rank = dist.get_rank()
torch.cuda.set_device(rank)

model = MyModel().to(rank)
model = DDP(model, device_ids=[rank])

Launch with torchrun:

CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 train.py

--nproc_per_node must match the number of GPUs you made visible. Use all four GPUs of the 4×H100 (CUDA_VISIBLE_DEVICES=0,1,2,3 … --nproc_per_node=4) only when they are all free and your job really needs them.

5. Serving large models — vLLM#

For inference on models too big for one GPU, vLLM splits the model automatically:

# Llama 3.3 70B in bf16 split across 2 GPUs (best on an NVLink pair of the 4×H100)
CUDA_VISIBLE_DEVICES=0,1 vllm serve meta-llama/Llama-3.3-70B-Instruct \
  --host 127.0.0.1 --port 8000 --api-key "$VLLM_API_KEY" \
  --tensor-parallel-size 2

See Good practices with vLLM — in particular how to stop it without leaving GPU memory allocated.

In plain HuggingFace code, one argument does the same job: device_map="auto" spreads the model across all the GPUs it can see — so choose them first with CUDA_VISIBLE_DEVICES:

# run with: CUDA_VISIBLE_DEVICES=0,1 python script.py
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.3-70B-Instruct",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

6. Training large models — FSDP & DeepSpeed#

To train (not just run) a model too big for one GPU, the trick is called sharding: instead of copying everything everywhere, each GPU stores only a fraction of the model, the gradients, and the optimizer.

Both tools below integrate with Accelerate — same workflow as before:

FSDP (built into PyTorch)#

accelerate config
# choose "FSDP" when asked
accelerate launch train.py

DeepSpeed#

pip install deepspeed
accelerate config
# choose "DeepSpeed" stage 2 or 3
accelerate launch train.py

Rule of thumb: start with FSDP, which is simpler. Move to DeepSpeed only if you need its advanced features.

7. Make it fast — easy optimisations#

Before parallelising, try these: they often give bigger gains for much less effort. Each is one or two lines of code:

bf16 — half the memory, for free#

Train in bf16 instead of fp32: half the memory, faster, and stable on H100s:

training_args = TrainingArguments(bf16=True, ...)

Flash Attention 2 — faster attention#

A smarter implementation of the attention mechanism, for a significant speedup:

model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-Instruct-v0.3",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2"
)

Gradient checkpointing — trade speed for memory#

If you run out of memory, this frees a lot of it at the cost of ~20% slower training:

model.gradient_checkpointing_enable()

Gradient accumulation — big batches without the memory#

Simulate a large batch by accumulating several small ones:

training_args = TrainingArguments(
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,   # effective batch = 32
)

torch.compile — free speedup#

One line, PyTorch optimises your model automatically:

model = torch.compile(model)

8. Feed your GPUs — efficient data loading#

A very common and invisible problem: the GPU computes faster than the data arrives, so it sits idle waiting. Like a chef waiting for ingredients.

            ┌──────────┐
            │ Worker 1 │──┐
┌────────┐  ├──────────┤  │  ┌─────────┐   ┌─────┐
│Dataset │─>│ Worker 2 │──┼─>│  Queue  │──>│ GPU │
│on disk │  ├──────────┤  │  │prefetch │   └─────┘
└────────┘  │ Worker N │──┘  └─────────┘
            └──────────┘
  several CPU processes prepare the data in
  parallel so the GPU never waits

The fix: give your DataLoader several workers:

dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=8,        # parallel data loading
    pin_memory=True,      # faster CPU to GPU transfer
    prefetch_factor=2     # preload batches in advance
)

Guidelines:

• num_workers: start with 4–8; increase if the GPU still waits
• pin_memory=True: always, when training on GPU
• Don't go too high: each worker consumes RAM

How to know if your GPU is waiting:

watch -n 1 nvidia-smi

Look at GPU-Util: it should stay high (90-100%) during training. If it keeps dropping to 0%, your data pipeline is too slow: increase num_workers.

9. Profiling — measure before optimising#

Golden rule: never optimise blindly. Measure first where the time actually goes: the answer is often surprising.

Level 1 — nvidia-smi (10 seconds)#

watch -n 1 nvidia-smi

Level 2 — PyTorch Profiler (5 minutes)#

Shows exactly which operations take the most time:

import torch
from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    profile_memory=True
) as prof:
    output = model(batch)
    loss = output.loss
    loss.backward()

print(prof.key_averages().table(
    sort_by="cuda_time_total", row_limit=10
))

If the top entries are matrix multiplications, that's normal: your GPU does useful work. If they are data copies or waiting, fix your data pipeline first.

Level 3 — visualise in TensorBoard#

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./logs")
) as prof:
    for step, batch in enumerate(dataloader):
        train_step(batch)
        prof.step()
        if step >= 10:
            break

pip install torch-tb-profiler
tensorboard --logdir=./logs

Forward the port to see it in your browser — run on your laptop:

ssh -L 6006:localhost:6006 4h100

Then open http://localhost:6006.

Full guide: PyTorch Profiler tutorial

10. Custom GPU kernels — for the curious#

You will most likely never need this: PyTorch already uses highly optimised low-level code. This section exists so you know it exists.

A "kernel" is a small program that runs directly on the GPU. Writing your own only makes sense when:

• One unusual operation dominates your profiling results
• You invented a new architecture component no library covers

And even then, try in this order, and stop at the first level that solves your problem:

1. torch.compile — one line, PyTorch generates optimised kernels for you:

   model = torch.compile(model)

2. Existing libraries — Flash Attention, xFormers… someone probably already wrote your kernel
3. Triton — write GPU kernels in Python syntax, no C++ needed (tutorials)
4. CUDA C++ — maximum control, maximum effort; only for very specialised work

To understand the fundamentals of GPU programming: GPU programming basics (Marc Lelarge, dataflowr / MVA).

Useful resources#

• Ultrascale Playbook — the reference on large-scale training
• HuggingFace — Efficient Training on Multiple GPUs
• Accelerate documentation
• PyTorch DDP tutorial
• DeepSpeed tutorials
• PyTorch Profiler tutorial
• GPU programming basics — Marc Lelarge (dataflowr / MVA)

What's next?#