🧩 Distributed Training & Parallelism

Why distributed training?

Modern models don't fit on one GPU. A 7B-parameter model trained in mixed precision needs roughly 16 bytes per parameter just for its training state — around 112 GB — before you count activations. A single 64 GB accelerator can't hold that, and a 70B model is ten times worse. Distributed training is the set of techniques for splitting one training job across many GPUs so the model fits and trains faster.

There are two distinct problems, and it helps to keep them separate:

• "It doesn't fit." The model's parameters, gradients, and optimizer state are too big for one GPU's memory. You solve this by partitioning the model across GPUs.
• "It's too slow." Even if it fits, one GPU would take months. You solve this by replicating the work across GPUs and processing more data in parallel.

The art of distributed training is combining the right techniques for your model, your hardware, and your interconnect.

Where the memory goes

Before the parallelism strategies make sense, you need to know what's eating memory. For mixed-precision (fp16/bf16) training with the Adam optimizer, each parameter carries:

• fp16 weights — 2 bytes
• fp16 gradients — 2 bytes
• fp32 master weights (a high-precision copy for stable updates) — 4 bytes
• Adam momentum + variance (two fp32 values) — 8 bytes

That's 16 bytes per parameter, and the optimizer state (the last 12 bytes) dominates. On top of this sit activations — the intermediate tensors saved for the backward pass — which scale with batch size and sequence length. Every parallelism technique below is really a strategy for cutting one of these four buckets down per GPU.

Data parallelism (DP)

The simplest and most common. Every GPU holds a complete copy of the model and processes a different slice of the batch. After each backward pass, GPUs all-reduce their gradients (average them together) so every replica applies the same update and stays in sync.

Data parallelism scales throughput — 8 replicas process ~8× the data per step. But it does not solve "it doesn't fit": every GPU still holds the whole model. PyTorch's classic DistributedDataParallel (DDP) is pure data parallelism.

Tensor parallelism (TP)

Now we partition the model itself. Tensor parallelism (also called intra-layer or Megatron-style parallelism) splits the big matrix multiplications inside each layer across GPUs — for example, cutting a weight matrix into column-shards, each GPU computing part of the output, then combining. Each GPU holds only 1/TP of that layer's weights.

TP is the most powerful memory lever, but it is the chattiest: because the layer's math is split, GPUs must communicate (an all-reduce) inside every forward and backward pass. That only stays fast over a very high-bandwidth link, so tensor parallelism is almost always kept within a single node (NVLink / Infinity Fabric), rarely across the slower network between nodes.

Pipeline parallelism (PP)

Instead of splitting each layer, pipeline parallelism splits the stack of layers into sequential stages, one group of layers per GPU. Data flows through like an assembly line: GPU 1 runs layers 1–8, hands its activations to GPU 2 for layers 9–16, and so on. Communication is cheap — only the activations at stage boundaries cross the wire.

The catch is the pipeline bubble. While the first microbatch is still working through the early stages, the later stages sit idle; the same happens in reverse as the pipeline drains. With P stages and m microbatches, the idle fraction is (P−1)/(m+P−1). The fix is to push more microbatches through so every stage stays busy — which is why pipeline parallelism always comes with microbatching.

ZeRO / FSDP — sharding the redundancy

Data parallelism wastes memory: every replica keeps an identical copy of the weights, gradients, and optimizer state. ZeRO (Zero Redundancy Optimizer, from DeepSpeed) and its PyTorch-native equivalent FSDP (Fully Sharded Data Parallel) fix this by sharding those states across the data-parallel replicas instead of duplicating them:

• Stage 1: shard the optimizer state (the biggest chunk — 12 of the 16 bytes). Huge saving, almost no extra communication.
• Stage 2: also shard the gradients.
• Stage 3 (= FSDP): also shard the weights themselves. Nothing is fully replicated. Each GPU stores only its slice and all-gathers the full weights for a layer just in time to compute it, then discards them. Maximum memory savings, at the cost of that extra communication each step.

ZeRO is the technique that lets you train a model far larger than one GPU's memory using "just" data parallelism — you keep the simple DP programming model but drop the redundancy.

Sequence / context parallelism

For very long sequences, the activations (not the weights) become the bottleneck — attention memory grows with sequence length. Sequence parallelism splits the activations along the sequence dimension across GPUs, and context parallelism extends this to make million-token contexts trainable. These complement tensor parallelism, which shards along the hidden dimension.

Expert parallelism (for MoE)

A Mixture-of-Experts model has many expert sub-networks but activates only a few per token. Expert parallelism places different experts on different GPUs; each token is routed to the GPUs holding its chosen experts (an all-to-all communication). This is how models with trillions of total parameters stay trainable — only a fraction is active per token.

Putting it together: 3D (and ND) parallelism

Real large-scale runs combine these into a mesh. A common recipe: tensor parallelism within a node (fast link, chatty comms), pipeline parallelism across a few nodes (cheap comms), and data parallelism (with ZeRO) across the rest for throughput. The GPU count factors as total = TP × PP × DP. Add expert and sequence parallelism and you get 4D or 5D parallelism. Choosing the split is a balancing act between memory, interconnect speed, and the pipeline bubble.

The explorer above is a device-mesh planner: pick a model size and GPU count, then dial tensor, pipeline, and data parallelism plus ZeRO and watch the per-GPU memory bar cross (or clear) the 64 GB budget. Try loading a 70B model on one GPU (it won't fit), then split it with tensor and pipeline parallelism, then turn on ZeRO-3 — and watch the memory come down.

PyTorch vs JAX — two philosophies

You'll implement all of this in one of two ecosystems, and they take genuinely different approaches. The short version: PyTorch is imperative and you place the shards; JAX is compiled and the compiler places them.

PyTorch: eager, explicit, incremental. Code runs line by line (eager mode), which makes it easy to debug and inspect. Distribution is something you add to your model: wrap it in DDP for data parallelism, or FSDP / DTensor for sharding, and you reason fairly directly about which collective (all-reduce, all-gather) happens where. torch.compile can then trace and optimize hot paths, but the default mental model is "Python runs, and I opt into distribution." This makes PyTorch flexible and beginner-friendly, and it's why most open models ship PyTorch code.

JAX: functional, compiled, declarative. JAX code is pure functions that you jit-compile; the XLA compiler fuses and optimizes the whole computation graph ahead of time, often yielding excellent throughput especially on TPUs. For distribution you don't hand-place collectives — you describe a device mesh and annotate how arrays are sharded across it (shard_map, or jit with sharding constraints), and XLA's GSPMD partitioner figures out the communication for you. You say "this tensor is split this way across these devices," and the compiler inserts the all-reduces.

How to choose, in one line: PyTorch trades some peak performance for flexibility, debuggability, and ecosystem size; JAX trades some ease-of-debugging for compiler-driven performance and clean, scalable sharding — historically the default for the largest TPU runs. Both express the same parallelism concepts above; they just differ in whether you or the compiler wires up the communication.