Core ML

Distributed Training

You cannot fit a 70-billion parameter model into a single GPU. It physically doesn't fit in the memory. Distributed training is the art of breaking the model, the data, or the optimizer across hundreds of GPUs so they can train together as one massive brain.

When a model outgrows a single GPU, you must distribute the workload. You can copy the model and split the data (Data Parallelism), split the layers (Pipeline Parallelism), or split the tensors themselves (Tensor Parallelism).

Why Does This Exist?

In the early days of Deep Learning (circa 2014), a state-of-the-art model like VGG-16 had 138 million parameters. You could easily fit the entire model and a batch of images into the 12GB of VRAM on a single Nvidia Titan X.

Today, Llama 3 has 70 billion parameters. Just to store those parameters in standard 16-bit precision requires 140GB of VRAM. The largest GPU in the world (the Nvidia H100) only has 80GB of VRAM. You literally cannot fit the model onto one chip. Furthermore, even if you could, training it on one chip would take 400 years.

Distributed Training is the infrastructure layer of modern AI. It provides the mathematical and networking frameworks required to stitch dozens, hundreds, or tens of thousands of GPUs together over high-speed networks (like NVLink and InfiniBand) so they act as a single, massive computer.

Think of It Like This

Imagine you need to translate a 10,000-page book from English to Spanish in one day. One translator cannot do it. You need to hire 100 translators.

Data Parallelism: You print 100 copies of the dictionary (the Model). You give each translator their own dictionary, and hand each of them 100 pages of the book (the Data). Pipeline Parallelism: The book is too complex for one person. Translator A reads the English and writes a rough outline. They hand the paper to Translator B, who writes the Spanish prose. Translator B hands it to Translator C, who checks the grammar. Tensor Parallelism: Translating a single, massive sentence is too hard. Translator A translates the nouns, while Translator B translates the verbs simultaneously, and they merge the sentence together before moving to the next.

How It Actually Works (The 3D Parallelism Grid)

Modern distributed training for Foundation Models usually combines three distinct paradigms simultaneously, often called 3D Parallelism.

1. Data Parallelism (DP)

The simplest approach. Every GPU holds an exact, complete copy of the model. You take your massive dataset and chop it into smaller batches. GPU 1 looks at Batch 1, GPU 2 looks at Batch 2. They each calculate their gradients (how the model should change). Before they update their weights, they synchronize their gradients across the network (an All-Reduce operation) so every GPU updates in exactly the same way. Constraint: The entire model must fit on a single GPU.

2. Pipeline Parallelism (PP)

When the model is too big for one GPU, you slice it vertically by depth. If you have a 40-layer Transformer and 4 GPUs, you put Layers 1-10 on GPU 1, Layers 11-20 on GPU 2, etc. GPU 1 processes the data, sends the intermediate activations to GPU 2, and so on. Constraint: GPUs spend a lot of time sitting idle waiting for the previous GPU to finish (the "Pipeline Bubble").

3. Tensor Parallelism (TP)

When a single layer (like a massive Attention matrix) is too big for one GPU, you slice it horizontally. You literally cut the mathematical matrix in half. GPU 1 calculates the left half of the matrix multiplication, GPU 2 calculates the right half, and they combine the result instantly. Constraint: Requires massive, ultra-fast networking (like NVLink) because the GPUs have to communicate constantly during a single calculation.

Show Me the Code

You rarely write distributed communication primitives yourself. You use frameworks like PyTorch's DistributedDataParallel (DDP) or HuggingFace Accelerate. Here is how simple Data Parallelism is in PyTorch.

1import torch2import torch.nn as nn3from torch.nn.parallel import DistributedDataParallel as DDP4import torch.distributed as dist5
6# 1. Initialize the process group (connect the GPUs together)7# This script is launched on multiple GPUs simultaneously by the OS8dist.init_process_group(backend="nccl")9local_rank = dist.get_rank()10
11# 2. Assign this specific script to a specific GPU12torch.cuda.set_device(local_rank)13
14# 3. Create the model and move it to this specific GPU15model = MyMassiveNeuralNetwork().to(local_rank)16
17# 4. Wrap the model in DDP. 18# DDP will automatically handle synchronizing the gradients across the network!19model = DDP(model, device_ids=[local_rank])20
21# 5. Normal training loop!22optimizer = torch.optim.Adam(model.parameters())23
24for batch in dataloader:25    # DDP magically synchronizes gradients during the backward pass26    loss = model(batch.to(local_rank)).sum()27    loss.backward()28    29    optimizer.step()30    optimizer.zero_grad()

Watch Out For

The Communication Bottleneck. In distributed training, GPUs are incredibly fast at doing math, but the network cables connecting them are relatively slow. If you use too many GPUs, the training process actually slows down because the GPUs spend 80% of their time sitting idle, waiting for data to travel over the Ethernet cable. Scaling efficiency is the hardest problem in distributed ML. You must balance compute time vs. communication time.

The Quick Version

Distributed Training is required when a model's parameters or dataset exceed the physical limits of a single GPU.
Data Parallelism (DP): Split the data, copy the model.
Pipeline Parallelism (PP): Split the model's layers sequentially across GPUs.
Tensor Parallelism (TP): Split individual mathematical matrices across GPUs.
Modern LLM training uses all three simultaneously (3D Parallelism) orchestrated by frameworks like PyTorch DDP or Megatron-LM.