Building a distributed training framework from first principles

In this video, I'll build a distributed training framework from first principles using PyTorch. We'll derive everything from first principles, explaining the mathematics behind distributed training as well as every auxiliary component we build along the way, including Multi-head Latent Attention (MLA), Rotary Position Embeddings (RoPE) and YaRN. I'll explain process groups, ranks, collective communication operations, distributed autograd, device meshes and distributed tensors, and then use these foundations to implement the major forms of parallelism used to train modern language models. We'll derive and code pipeline parallelism, data parallelism, FSDP, HSDP, tensor parallelism and context parallelism, including pipeline schedules, block matrix multiplication, model sharding and ring attention. We'll then combine these techniques and examine how tensor layouts and gradients must be tracked as data moves between devices. Most importantly, we won't study each parallelism technique in isolation. We'll combine pipeline, data, tensor, context and expert parallelism into a single working framework and follow the tensors, communication operations and gradients as they move across GPUs. By the end, you'll understand not only how these techniques work individually, but how modern distributed training systems compose them to train large dense and mixture-of-experts models. We'll also build the important components of a modern transformer, including rotary embeddings, long-context extensions, multi-head latent attention and a mixture-of-experts architecture. This leads into expert parallelism, expert tensor parallelism, token dispatching, all-to-all communication and combining different parallel dimensions through dense and sparse device meshes. We'll also build the supporting training infrastructure, including the training loop, optimizer, learning-rate schedule, checkpointing and metrics. The goal is to give you the foundations required to understand how distributed training frameworks work internally and how to design and implement new parallelization strategies yourself. No prior knowledge of distributed training is required. The full source code can be found at: github.com/hkproj/torchfeather I highly recommend having watched the below videos or having had exposure to the topics described in the videos: - Flash Attention (youtube.com/watch?v=zy8ChVd_oTM) - Coding a Transformer model: youtube.com/watch?v=ISNdQcPhsts Non