6.1 Distributed Training Architectures Overview — How to Glue a Giant Model to Your Cluster Without Crying

You've already learned how to make models smaller and faster with quantization, pruning, and compression. Great! Now we ask: how do you scale training of the models that remain — or the compressed ones that still don't fit a single GPU? Welcome to distributed training architecture land: the part where we split the model, the work, and sometimes our sanity.

Why this matters (and why it builds on pruning/quantization)

You reduced parameter precision and cut redundant weights — tiny win for inference latency. But when fine-tuning billion-parameter beasts, memory for optimizer states, activations, and gradients still explodes. Distributed architectures aren’t just performance boosters; they change how memory and communication are allocated. Think of quantization/pruning as slimming the dragon; distributed architectures are the scaffolding you strap on to ride it.

The big categories — quick tour

• Data Parallelism (DP) — Everyone has a full copy of the model; each GPU processes a slice of data and gradients are synchronized. Simple, robust, but limited by device memory.
• Tensor (or Operator) Parallelism (TP) — Split the layers' linear algebra across devices (e.g., splitting large matrix multiplications). Great for extremely wide layers.
• Pipeline Parallelism (PP) — Split layers into stage-chunks and stream micro-batches through a pipeline. Reduces peak memory per device for activations, but needs careful scheduling.
• Parameter Sharding (ZeRO family / FSDP) — Shard model parameters, optimizer states, and/or gradients across devices to reduce per-GPU memory. Comes in stages and flavors.
• Hybrid approaches — Most real-world setups mix DP + TP + PP + ZeRO/FSDP. It’s like making a training lasagna: layers matter.

Data Parallelism — the beginner’s cheat code

• When to use: model fits per-GPU memory comfortably. Great for small-to-medium models and straightforward fine-tuning.
• Pros: Simple, easy scaling across nodes, well-supported in frameworks.
• Cons: Memory duplicates model on each GPU, so inefficient for huge models.

Imagine copying a 7B model onto each GPU and shouting, "Synchronize!" after every batch. Loud and simple.

Tensor (Model) Parallelism — slice the math

• What: Split large linear algebra operations (e.g., a 32k x 32k matmul) across GPUs.
• Pros: Enables training of extremely wide layers that exceed single-GPU memory.
• Cons: Requires tight inter-GPU comms; increases implementation complexity.

Used by Megatron-LM and similar systems. It’s surgical: you break a matrix into shards and do the multiply across GPUs.

Pipeline Parallelism — assembly line for layers

• What: Partition model layers into stages on different GPUs and stream micro-batches.
• Pros: Reduces activation memory and enables training of deep models across many devices.
• Cons: Pipeline bubbles, increased latency, and tricky checkpointing; needs micro-batching and gradient accumulation.

Tip: combine PP with activation recomputation to save memory at the cost of extra compute.

ZeRO (DeepSpeed) and FSDP (PyTorch): the sharding showdown

Short version: both aim to eliminate memory redundancy, but they have different design choices and trade-offs.

ZeRO (Zero Redundancy Optimizer)

ZeRO breaks down memory into three major kinds of state and shards them across data-parallel ranks:

• Stage 1: shard optimizer states (huge for Adam).
• Stage 2: shard optimizer + gradients.
• Stage 3: shard optimizer + gradients + parameters (fully sharded).

Table — ZeRO stages at-a-glance:

ZeRO integrates with DeepSpeed and provides elastic checkpointing, memory/comms optimizations, and robust scaling.

FSDP (Fully Sharded Data Parallel)

FSDP is PyTorch-native and shards parameters and optimizer state per layer, exposing fine-grained control. It performs local forward/backward with on-the-fly parameter gathering.

• Pros: Tight integration with PyTorch, flexible, often better for mixed workloads and checkpointing.
• Cons: Implementation details matter: sharding granularity, flattening vs per-parameter, and communication scheduling.

Practical contrast

• Checkpointing & recovery: DeepSpeed/ZeRO often have more batteries-included features; FSDP gives more PyTorch-native control.
• Communication pattern: ZeRO Stage 3 tends to use reduce-scatter/all-gather across the world; FSDP commonly uses point-to-point and per-layer gather/scatter.
• When to choose: Use ZeRO for massive-scale, production-focused training with DeepSpeed. Use FSDP if you want deep PyTorch integration and layer-level sharding control.

Code sketch (very conceptual):

# Pseudocode: ZeRO-like training step
for microbatch in microbatches:
  forward(microbatch)  # params may be local or remote depending on stage
  backward()            # grads computed locally
  reduce_scatter(gradients)  # shards grads across ranks
  optimizer_step_on_shards() # update optimizer states on shards
  all_gather(params_if_needed)

How pipeline, tensor, and sharding play together

Real-life setups: tensor parallelism handles gigantic linear layers, pipeline parallelism stitches model stages across nodes, and ZeRO/FSDP shards parameter/optimizer memory so each GPU holds just its slice. This lets teams train models that would otherwise require hundreds of GPUs.

Considerations:

• Bandwidth vs latency: TP & ZeRO require high interconnect bandwidth (InfiniBand / NVLink). On ethernet, you’ll be throttled.
• Memory vs compute trade-offs: Activation recomputation + sharding reduces memory but raises compute cost.
• Interaction with compression: Quantized or pruned weights reduce memory and comm volume — this can reduce the intensity of sharding needed. But some compressed formats complicate gradient updates and optimizer state handling.

Practical checklist — pick your weapons

Ask yourself:

1. Does the model fit a single GPU? If yes → DP + mixed precision + quantization might suffice.
2. Is the model wide (huge matrices)? Consider tensor parallelism.
3. Is it deep and you need to reduce activation memory? Consider pipeline parallelism.
4. Are optimizer states dominating memory? Consider ZeRO Stage 1/2.
5. Do you want full minimal per-GPU memory footprint? ZeRO Stage 3 or FSDP.
6. What's your network? If it's not high-bandwidth, hybrid strategies minimizing all-gather barriers are preferable.

Closing — TL;DR (but with flair)

Distributed training is less a single tool and more an orchestra: data-parallel drums, tensor-parallel brass, pipeline strings, and ZeRO/FSDP conducting the memory section. Combine them thoughtfully with the compression tricks you already mastered. The goal: minimize per-GPU memory, keep communication sane, and not spend your budget on a van full of GPUs.

Final note: start with the simplest setup that works, measure memory/comm hotspots, then incrementally add sharding, tensor splits, or pipeline stages. It’s way easier to debug one type of parallelism at a time than to wake up to a production job that deadlocks across 128 GPUs.

Next up: we’ll dig into concrete orchestration with DeepSpeed + ZeRO configs and FSDP best practices — how to tune those stage knobs without setting your cluster on fire.