Multi-GPU Setups: Splitting Models Across GPUs

Llama 3 405B in FP16 is 810 GB. The biggest single GPU has 192 GB of VRAM. You need a minimum of 5 GPUs just to hold the weights — and that's without KV cache, activations, or any breathing room.

When your model doesn't fit on one GPU, you split it across multiple GPUs. But how you split it matters enormously. The wrong strategy can make 8 GPUs slower than 4. The right strategy scales nearly linearly.

There are three fundamental approaches — data parallelism, tensor parallelism, and pipeline parallelism — and most real systems use all three simultaneously. This episode is your guide to when and why.

Data Parallelism: Same Model, Different Data

The simplest approach. Every GPU holds a complete copy of the model. Each GPU processes a different batch of data.

GPU 0: Full model copy → processes batch 0 → gradients_0

GPU 1: Full model copy → processes batch 1 → gradients_1

GPU 2: Full model copy → processes batch 2 → gradients_2

GPU 3: Full model copy → processes batch 3 → gradients_3

↓
          All-Reduce: average gradients across GPUs
                    ↓
          All GPUs update their model copy

When Data Parallelism Works

Situation Works? Why

The Communication Cost

After each training step, GPUs must synchronize their gradients. This is the All-Reduce operation:

Data to communicate = 2 × model_size (send + receive gradients)

For a 7B model in FP32:

All-Reduce data = 2 × 28 GB = 56 GB per step
With NVLink (900 GB/s):   56 / 900 = 0.062 seconds
With PCIe (64 GB/s):      56 / 64 = 0.875 seconds  ← 14x slower!

This is why interconnect bandwidth matters. PCIe-connected multi-GPU training can spend more time communicating than computing.

DDP and FSDP

DDP (Distributed Data Parallel): Standard PyTorch data parallelism. Each GPU has a full model copy. Simple and efficient.
FSDP (Fully Sharded Data Parallel): Instead of each GPU holding a full model copy, each GPU holds a shard (fraction) of the model. Parameters are gathered on-demand when needed, then discarded. This means a model that doesn't fit on one GPU CAN be trained with data parallelism.

DDP with 4 GPUs:

GPU 0: full 7B copy (28 GB) + optimizer (56 GB) = 84 GB  ← Needs A100 80GB

GPU 1: full 7B copy (28 GB) + optimizer (56 GB) = 84 GB

FSDP with 4 GPUs:

GPU 0: 1/4 of 7B (7 GB) + 1/4 optimizer (14 GB) = 21 GB  ← Fits on RTX 4090!

GPU 1: 1/4 of 7B (7 GB) + 1/4 optimizer (14 GB) = 21 GB

FSDP trades communication for memory. More communication (gathering parameters before each layer), but far less VRAM per GPU.

Tensor Parallelism: Splitting Individual Layers

Tensor parallelism splits each layer across GPUs. Instead of each GPU running a complete layer, each GPU runs a slice of every layer.

Without Tensor Parallelism (1 GPU):

Layer 1: Weight matrix [4096 × 4096] → full computation on GPU 0

With Tensor Parallelism (4 GPUs):

Layer 1: Split weight [4096 × 4096] into 4 pieces

GPU 0: Weight slice [4096 × 1024] → partial result
    GPU 1: Weight slice [4096 × 1024] → partial result
    GPU 2: Weight slice [4096 × 1024] → partial result
    GPU 3: Weight slice [4096 × 1024] → partial result
    → All-Reduce: combine partial results → full output

How Matrix Multiply Gets Split

For a weight matrix W of shape [N × M], with tensor parallelism degree P:

Column parallel (split output):

W = [W₁W₂W₃W₄] (split columns across 4 GPUs)

GPU 0: x × W₁ = y₁ (each GPU gets partial output)

GPU 1: x × W₂ = y₂

GPU 2: x × W₃ = y₃

GPU 3: x × W₄ = y₄

All-Gather: y = [y₁, y₂, y₃, y₄] (combine outputs)

Row parallel (split input):

W = [W₁; W₂; W₃; W₄]  (split rows across 4 GPUs)

GPU 0: x₁ × W₁ = partial_0

GPU 1: x₂ × W₂ = partial_1

GPU 2: x₃ × W₃ = partial_2

GPU 3: x₄ × W₄ = partial_3

All-Reduce: y = partial_0 + partial_1 + partial_2 + partial_3

In practice, transformers alternate between column-parallel and row-parallel splits across the two linear layers in FFN and the QKV projections in attention. This cancels out communication — each layer only needs one All-Reduce.

When Tensor Parallelism Works

Situation Works? Why

Critical requirement: Tensor parallelism needs high-bandwidth interconnect (NVLink). Every token requires All-Reduce across GPUs. With PCIe, the communication overhead destroys performance.

Communication per layer per token (All-Reduce):

Data = 2 × hidden_dim × bytes_per_value / TP_degree

For Llama 3 70B, FP16, TP=8:

Per layer: 2 × 8192 × 2 / 8 = 4 KB per token per layer

× 80 layers = 320 KB per token

At NVLink (900 GB/s): 0.36 μs per token ← Fine

At PCIe (64 GB/s): 5 μs per token ← Significant overhead

🔧

Pipeline Parallelism: Splitting by Layers

Pipeline parallelism puts different layers on different GPUs:

GPU 0: Layers 0-19 (Embedding + first 20 layers)

GPU 1: Layers 20-39

GPU 2: Layers 40-59

GPU 3: Layers 60-79 (last 20 layers + output head)

Data flows: GPU 0 → GPU 1 → GPU 2 → GPU 3

The Pipeline Bubble

The problem: GPUs are idle while waiting for data from the previous stage.

Time →

GPU 0: [Batch 0] [idle] [idle] [idle] [Batch 1] ...

GPU 1: [idle] [Batch 0] [idle] [idle] [idle] ...

GPU 2: [idle] [idle] [Batch 0] [idle] [idle] ...

GPU 3: [idle] [idle] [idle] [Batch 0] [idle] ...

Each GPU is idle 75% of the time! This is the "pipeline bubble."

Micro-batching: The Fix

Split each batch into micro-batches and pipeline them:

Time →

GPU 0: [μ0] [μ1] [μ2] [μ3] [idle] [idle]

GPU 1: [μ0] [μ1] [μ2] [μ3] [idle]

GPU 2: [μ0] [μ1] [μ2] [μ3]

GPU 3: [μ0] [μ1] [μ2] [μ3]

Much better! GPUs are busy most of the time.

Bubble fraction ≈ (num_stages - 1) / num_microbatches

With 4 pipeline stages and 32 micro-batches, the bubble is only (4-1)/32 ≈ 9% — much more efficient.

When Pipeline Parallelism Works

Situation Works? Why

Key advantage: Pipeline parallelism only transfers activations between GPUs — a single vector per layer per token. Much less data than tensor parallelism's All-Reduce.

Inter-GPU transfer per token:

Tensor Parallelism: 320 KB (All-Reduce every layer)

Pipeline Parallelism: 16 KB (one activation vector at stage boundary)

📊

Combining All Three: The 3D Approach

Real production systems combine all three methods:

Example: Training 405B model on 128 GPUs

Data Parallelism: DP=4 (4 model replicas, different data)

Tensor Parallelism: TP=8 (each layer split across 8 GPUs)

Pipeline Parallelism: PP=4 (layers split into 4 stages)

Total GPUs: 4 × 8 × 4 = 128

DP Group 0              DP Group 1
                ┌─────────────────┐    ┌─────────────────┐

Pipeline Stage 0│ TP: 8 GPUs │ │ TP: 8 GPUs │

Pipeline Stage 1│ TP: 8 GPUs │ │ TP: 8 GPUs │

Pipeline Stage 2│ TP: 8 GPUs │ │ TP: 8 GPUs │

Pipeline Stage 3│ TP: 8 GPUs │ │ TP: 8 GPUs │

└─────────────────┘    └─────────────────┘

Which to Use When

Strategy Use When Interconnect Needed Scales

The Decision Tree

Does the model fit on 1 GPU?

├── YES → Use 1 GPU (or Data Parallelism for throughput)

├── NO → Does it fit on GPUs within one server (with NVLink)?

│   ├── YES → Tensor Parallelism (TP=2,4,8)

│ └── NO → Pipeline Parallelism across servers

│ + Tensor Parallelism within server

└── Training?

└── Add Data Parallelism on top for throughput

NVLink: The Enabler

NVLink makes multi-GPU work viable. Without it, tensor parallelism is impractical:

8 GPU server with NVLink (DGX H100):

GPU-to-GPU: 900 GB/s

Any GPU can talk to any GPU at full speed

TP=8 works beautifully

8 GPU server with PCIe:

GPU-to-GPU: 64 GB/s (through CPU)

14x slower than NVLink

TP=8 is a disaster — communication bottleneck

Rule of thumb:

NVLink → TP up to 8

PCIe → TP=2 max, prefer PP instead

Inference-Specific Parallelism

For serving models, the strategies differ from training:

Tensor Parallelism for Latency

When you need fast single-request response time, TP reduces latency because all GPUs work on every token:

70B on 1 GPU:  reads 140 GB weights → limited by bandwidth
70B on 8 GPUs: each reads 17.5 GB weights → 8x faster (minus communication)

Practical speedup with TP=8: ~5-6x (not full 8x due to communication)

Expert Parallelism for MoE

MoE models have a natural parallelism strategy — put different experts on different GPUs:

Mixtral 8x7B with 8 GPUs:

GPU 0: Expert 0 + shared attention

GPU 1: Expert 1 + shared attention

...

GPU 7: Expert 7 + shared attention

Each token → router decides experts → data goes to those GPUs → results combine

This is Expert Parallelism (EP) — and it maps perfectly to MoE's sparse activation. Since only 2 experts are active per token, only 2 GPUs do real work for each token.

🔬

Practical Multi-GPU Configurations

Model Recommended Setup GPUs Notes

🛡️

Practical Takeaways

Data parallelism for throughput — same model on each GPU, different data, scale training linearly

Tensor parallelism for latency — split layers across GPUs, all work on every token, requires NVLink

Pipeline parallelism for scale — different layers on different GPUs, works across servers

NVLink is mandatory for tensor parallelism — PCIe is 14x slower and kills performance

FSDP lets you train models that don't fit on one GPU — shard weights and optimizer across GPUs

Real systems combine all three — TP within a node, PP across nodes, DP for throughput

📦

What's Next?

Episode 67: Apple Silicon for AI — Think Macs can't do AI? Think again. Unified memory, Metal, MLX — Apple Silicon has a unique architecture that makes it surprisingly competitive for inference. We'll benchmark it against NVIDIA.

← Previous

Ep 65: VRAM Math

Next →

Ep 67: Apple Silicon for AI

Next: Episode 67 — Apple Silicon for AI

A Mac Studio with M4 Ultra can run Llama 3 70B unquantized in FP16. No multi-GPU setup. No $30,000 H100. Here's why.

This is part of a 98-episode series covering AI engineering from tokens to production deployment.