Scaling: Distributed Training, FSDP, DeepSpeed
Your 124M model trained on one GPU. Now try 7 billion parameters. The model doesn't fit in memory. The data takes weeks on a single machine. Distributed training isn't optional at scale. It's the only path forward. Explain the three types of parallelism (data, tensor, pipeline) and when each is necessary based on model and cluster size. Implement data-parallel training using PyTorch DDP with gradient synchronization across multiple GPUs. Calculate the memory budget for a given model size (weights + optimizer states + gradients + activations) to determine the minimum hardware. Configure FSDP or DeepSpeed ZeRO stages to shard model states across GPUs and fit models that exceed single-GPU memory. A 7B parameter model in FP16 needs 14GB just for the weights. Adam optimizer stores two additional copies of every parameter (first and second moment estimates). That is another 28GB. Gradients during backpropagation add 14GB more. You are at 56GB before a single activation is stored. An NVIDIA A100 has 80GB of memory. 56GB out of 80GB consumed. That leaves 24GB for activations -- the intermediate values computed during the forward pass that must be kept alive for backpropagation. For a 2048-token sequence with a 4096-dimensional model, a single layer's activations use about 64MB. With 32 layers, you need 2GB per sample. A batch size of 8 requires 16GB. You…
Scaling: Distributed Training, FSDP, DeepSpeed: Your 124M model trained on one GPU. Now try 7 billion parameters. The model doesn't fit in memory. The data…
This free lesson is part of the AI Engineering from Scratch curriculum. Read the full explanation, run the lesson code, and verify the result in the interactive reader or from the repository source.
Browse the complete course catalog or open this lesson on GitHub.