Page
1
(This page has no text content)
Page
2
Distributed AI Systems A practical guide to building scalable training, inference, and serving systems for production AI Fuheng Wu
Page
3
(This page has no text content)
Page
4
(This page has no text content)
Page
5
Distributed AI Systems Copyright © 2026 Packt Publishing All rights reserved. No part of this book may be reproduced, stored in a retrieval system, or transmitted in any form or by any means, without the prior written permission of the publisher, except in the case of brief quotations embedded in critical articles or reviews. Every effort has been made in the preparation of this book to ensure the accuracy of the information presented. However, the information contained in this book is sold without warranty, either express or implied. Neither the author, nor Packt Publishing or its dealers and distributors, will be held liable for any damages caused or alleged to have been caused directly or indirectly by this book. Packt Publishing has endeavored to provide trademark information about all of the companies and products mentioned in this book by the appropriate use of capitals. However, Packt Publishing cannot guarantee the accuracy of this information. Portfolio Director: Kunal Chaudhari Relationship Lead: Rithika Shetty Project Manager: Ankit Maroli Content Engineer: Deepayan Bhattacharjee Technical Editor: Aditya Bharadwaj Copy Editor: Deepayan Bhattacharjee Indexer: Tejal Soni Production Designer: Ganesh Bhadwalkar Growth Lead: Vinishka Kalra First published: June 2026 Production reference: 1290626 Published by Packt Publishing Ltd. Grosvenor House 11 St Paul's Square Birmingham B3 1RB, UK. ISBN 978-1-80730-171-2 www.packtpub.com
Page
6
To my parents, Chuntao He and Zongyuan Wu. – Fuheng Wu
Page
7
Foreword I have known Fuheng Wu for many years. We have both worked at Bloomberg and Uber, and we have stayed in the same field ever since—as I went on to Meta, Anyscale, and NVIDIA, we kept trading notes on distributed AI. At Uber, we exchanged technical ideas constantly: I was working on OpenSearch while he was building distributed AI systems with Petastorm and training models with Horovod on Michelangelo. I have seen that production-first instinct up close—the debugging sessions, the cluster failures, the pressure to ship under real constraints. When he told me he was writing this book, I had no doubt it would be grounded in that experience, not textbook abstraction. That distinction has never mattered more, because AI has entered an era where progress increasingly depends not only on better models, but on the systems capable of training and serving them. That shift is happening in one of the most transformative epochs in the history of technology. The rapid evolution of artificial intelligence—moving from niche statistical models to gargantuan large language models and multimodal systems spanning billions or even trillions of parameters—has fundamentally shifted our relationship with computing. Yet, as the mathematical formulation of these models remains elegantly bound to deep learning fundamentals, the engineering reality of executing them has exploded in complexity. Today, the bottleneck to AI progress is no longer just algorithmic creativity; it is fundamentally a system engineering challenge. Building a neural network on a single GPU has become a commoditized task, abstracted away by modern software libraries. However, taking that same network and distributing its training across thousands of specialized chips, or orchestrating low-latency, high-throughput inference for millions of concurrent users, is an entirely different beast. It requires a profound, multidisciplinary mastery over physical hardware, high-speed networking topologies, memory management tricks, complex software orchestration, and algorithmic nuances. It is a domain where a single misplaced communication bottleneck can degrade hardware utilization from pristine efficiency to a crawl, rendering millions of dollars of compute infrastructure idle. For too long, knowledge in the distributed AI space has been siloed. Engineers looking to master this domain have had to piece together disparate academic papers, fragmented blog posts, tribal knowledge from elite tech companies, and chaotic open-source repository documentation. There has been a glaring gap between theoretical high-level explanations of parallelism and the actual code required to run stable, fault-tolerant workloads—from launching multi-node training on HPC clusters to serving models reliably in production.
Page
8
Distributed AI Systems: A practical guide to building scalable training, inference, and serving systems for production AI is the comprehensive blueprint the industry has desperately needed—and a book I wish Fuheng had written years ago. Fuheng Wu has done something remarkable with this text. He has demystified the black box of infrastructure engineering for AI. Rather than treating training, inference, and cluster management as isolated domains, this book treats them as an integrated lifecycle. It systematically walks the reader through the complete operational pipeline—from understanding the silicon and high-speed interconnects that form the physical bedrock, to mastering PyTorch DDP, FSDP, DeepSpeed ZeRO optimization, and Megatron-style tensor, pipeline, and expert parallelism for multi-node scaling. Crucially, this book recognizes that training is only half the battle. In real-world AI deployments, serving these models efficiently is where the economic and engineering viability of AI is won or lost. By diving deep into distributed inference fundamentals, vLLM, and SGLang, Fuheng equips engineers with the tools to tame memory fragmentation and unlock massive throughput. The book also covers running multi-node training on HPC clusters with SLURM and assembling a production serving stack with Kubernetes, GPU scheduling, and observability—so your systems don't just work in theory, they remain resilient and maintainable under real workloads. Whether you are a machine learning engineer seeking to scale your models beyond the limits of a single accelerator, a platform engineer tasked with building a modern AI cluster, or an architect designing the next generation of enterprise AI infrastructure, this book will serve as your definitive guide. Fuheng's hands-on, code-first approach translates intimidating infrastructure challenges into predictable, repeatable engineering methodologies. As you read these chapters, I encourage you to not just absorb the code, but to deeply study the underlying patterns of resource estimation, memory management, and network communication. The frameworks of today will inevitably evolve, but the distributed systems principles meticulously laid out in this book will remain the foundation of AI infrastructure for years to come. Gang Zhao Staff Engineer at NVIDIA, Anyscale, Meta, Uber June 2026
Page
9
Contributors About the author Fuheng Wu is a Principal ML Tech Lead at Oracle Generative AI, specializing in distributed training, inference, and GPU systems for enterprise AI workloads. He has delivered core components of Oracle's large-scale vision and document-AI models and Oracle Code Assist (Oracle's AI coding copilot), and co-authored a Microsoft–Oracle blog on high-performance deep learning with ONNX Runtime. A graduate of the Singapore–MIT Alliance, where he studied linear algebra with Gilbert Strang of MIT, he began his career in applied machine learning at NetEase, then worked on Wall Street in quantitative finance—pricing convertible bonds at Bloomberg and developing algorithmic trading strategies at WorldQuant—before joining Uber, where he built distributed training systems on Michelangelo with Petastorm and Horovod. Fuheng has hands-on experience across PyTorch Distributed, DeepSpeed, Megatron-style parallelism, Kubernetes GPU clusters, and production LLM serving. He has contributed to open-source projects including SGLang, vLLM, genai-bench, pyLLaMA, chatLLaMA, and HiQ, and is the author of Mathematics for AI and Machine Learning. He teaches Generative AI foundations at the McCombs School of Business, University of Texas at Austin, and mentors students preparing for AI Olympiads. Writing as Xuan Xin, he is also the author of the poetic memoir Above the Clouds: From Mountain Springs to the Edge of AI in Silicon Valley and the music album Let's Train the Model (2024), which renders deep learning and AI concepts as music. Beyond technology, he is a Zen calligrapher and volunteer math tutor. I am grateful to receive detailed feedback on early drafts of this book from: Liu Yong — Ph.D, NUS, Huaizheng Zhang — Staff Research Scientist, ByteDance, Junwei Zhang — GenAI | Engineer | Moderator | 18K+ followers | x{Uber, DoorDash}, Junnan Li — Research Director @Salesforce; Time-Series, CUA, and Multimodal AI (BLIP), Jerry Liu — Co-founder/CEO @ LlamaIndex, Henry Mao — Co-founder @ smithery.ai | Prev. co-founder of Jenni.ai (exited), Doris Xin — CEO & Co-Founder @ Disarray | UC Berkeley CS PhD | Ex-LinkedIn, Google, Databricks, You Yang — NUS Professor, Juncong Jack Wu — USA AI Olympiad Nationals Bronze, 2026; Canada AI Olympiad National Qualifier Distinguished Honor Roll, National Team Selection Top 18; USACO Silver
Page
10
About the reviewer Xuanyu Zhang is a machine learning engineer specializing in distributed AI infrastructure, large language model serving, and inference performance optimization. His work focuses on the systems layer behind modern generative AI, including model-serving runtimes, speculative decoding, quantization, continuous batching, GPU performance engineering, and multi-GPU and multi-node deployment of large-scale models. Xuanyu has worked on production machine learning and AI infrastructure at leading technology and AI companies, including Apple, Together AI, and Mistral AI. Across these roles, he has contributed to the design, optimization, and deployment of high-performance systems for large-scale model inference, bridging cutting-edge AI research with reliable production engineering. His experience spans both practical system implementation and deep technical evaluation of distributed training and inference architectures. As a reviewer of Distributed AI Systems, Xuanyu brought an industry practitioner's perspective to the book, with particular attention to the accuracy, clarity, and practical usefulness of its explanations and exercises on large language model training and inference systems. His review focused on ensuring that the material would be technically rigorous, accessible to engineers and researchers, and grounded in real-world AI infrastructure challenges.
Page
11
(This page has no text content)
Page
12
Table of Contents Preface xxxi Free benefits with your book ............................................................................ xxxv Chapter 1: Introduction to Modern Distributed AI 1 Overview ............................................................................................................. 1 Why modern AI requires distribution .................................................................... 2 The scale challenge • 4 Estimating model resource requirements • 5 Training memory requirements • 6 Inference memory requirements • 16 GPU requirements estimation • 17 Real-World considerations • 17 The evolution from classic ML to foundation models • 18 The modern AI model lifecycle ............................................................................. 19 Training: the learning phase • 20 Inference: the prediction/generation phase • 20 Serving: the production system • 21 Decision framework: when do you need distributed systems? ............................... 21 Decision tree: quick reference • 22 Understanding the decision tree • 22 Hands-On: running distributed training ............................................................. 23 Environment setup • 24 Single-GPU baseline • 25 Multi-GPU distributed training • 26 Extended training: CIFAR-10 with 20 epochs • 28 Distributed inference: throughput scaling • 30 Distributed AI fundamentals with PyTorch .......................................................... 31
Page
13
The distributed AI stack • 31 Process Groups and Ranks • 34 Initializing the process group • 34 Collective operations • 36 AllReduce • 36 AllGather • 37 Broadcast • 39 Reduce • 40 Gather • 42 Scatter • 43 ReduceScatter • 45 AlltoAll • 46 Choosing the right operation • 48 DistributedDataParallel (DDP) • 48 DistributedSampler • 49 Launching distributed jobs • 50 Code summary ................................................................................................... 50 References ........................................................................................................... 51 Exercises ............................................................................................................. 51 Process group initialization • 51 Manual AllReduce implementation • 52 Broadcast with verification • 53 DistributedSampler wrapper • 54 Gradient synchronization timing • 55 Expected learning outcomes • 56 Chapter 2: GPU Hardware, Networking, and Parallelism Strategies 57 Computational power: AI clusters and metrics .................................................... 57 What is computational power? • 57 Why clusters? • 59 AI clusters: built for training and inference • 60 Key metrics for AI clusters • 61 Table of Contents xii
Page
14
Central Processing Unit (CPU) ............................................................................ 64 CPU Architecture Basics • 64 CPU-GPU Interaction • 65 NUMA and CPU affinity • 65 CPU requirements for distributed training • 66 Graphics Processing Unit (GPU) .......................................................................... 66 Understanding GPU memory and compute architecture • 66 Key architecture milestones • 69 What these numbers mean for training • 69 Graphics Processing Unit (GPU) product families • 70 Choosing the right GPU • 71 Architecture features that matter • 71 Tensor Processing Unit (TPU) ............................................................................. 72 Why TPUs exist • 72 TPU architecture: systolic arrays • 73 TPU generations • 73 TPU pod architecture • 73 TPU vs GPU: when to use which • 74 TPU programming model • 75 Sparse core (since v4) • 75 Neural Processing Unit (NPU) ............................................................................. 76 What makes NPUs different • 76 NPU architecture: AI cores and memory • 77 Training vs inference NPUs • 77 NPU software stack • 78 NPU vs. GPU: when to choose which • 78 NPU interconnects and scaling • 78 The NPU landscape • 79 High-Speed interconnects: the network backbone ............................................... 79 Chip programming systems: SPMD and CUDA ..................................................... 81 Programming models vs. execution models • 81 SPMD: single program, multiple data • 81 xiii Table of Contents
Page
15
CUDA's execution model: SIMT • 82 Why SIMT over SIMD? • 83 CUDA thread indexing • 83 Memory hierarchy in CUDA • 84 How frameworks use CUDA • 84 SPMD in distributed training • 85 AMD and other alternatives • 85 Distributed communication: patterns and primitives .......................................... 85 Parallelism: core strategies ................................................................................. 86 Parallelism taxonomy table • 87 The three questions for parallelism determination • 88 Data parallelism: replicated and sharded • 89 Model parallelism: compute sharding • 89 Combining Strategies • 90 Composition patterns (hybrid parallelism) • 90 Strategy selection: choosing the right approach ................................................... 91 Training strategy decision tree • 91 Key considerations for training • 94 Inference strategy decision tree • 94 Key considerations for inference • 96 Practical tips • 97 Hands-On: hardware inspection and bandwidth testing ...................................... 97 Environment Setup • 97 Step 1: inspect GPU hardware • 97 Step 2: inspect hardware topology • 98 Step 3: measure Single-GPU memory bandwidth • 99 Step 4: benchmark Inter-GPU communication • 100 Step 5: analyze results • 102 Code summary .................................................................................................. 103 Notes and references ......................................................................................... 104 Exercises .......................................................................................................... 104 GPU hardware inspector • 104 Table of Contents xiv
Page
16
Memory requirement calculator • 105 Topology detection and analysis • 106 Parallelism strategy selector • 107 Expected learning outcomes • 109 Chapter 3: Distributed Training with PyTorch DDP 111 How DDP works internally ................................................................................. 111 The basic flow • 112 Data Parallel (DP) vs. Distributed Data Parallel (DDP) • 113 Gradient bucketing: why it matters • 114 Communication–Computation overlap • 114 The AllReduce operation • 115 Mixed precision and gradient scaling • 116 Buffer Synchronization • 117 DDP forward pass implementation details • 117 DDP Computation–Communication overlap implementation • 118 Setting up Single-Node DDP .............................................................................. 118 Using torchrun (recommended) • 118 Understanding the environment variables • 121 Using DistributedSampler • 121 DataLoader internals for distributed training • 124 Device selection best practices • 125 A complete single node example • 126 Setting up Multi-Node DDP ............................................................................... 128 Multi-Node rchitecture • 129 Launching Multi-Node training • 129 Using SLURM for Multi-Node launch • 130 Network configuration • 131 Environment variables for Multi-Node • 132 A complete Multi-Node example • 132 Debugging and troubleshooting ......................................................................... 133 Hangs: the most common problem • 134 xv Table of Contents
Page
17
Wrong results or inconsistent gradients • 136 CUDA out of memory • 137 Performance issues • 139 Network debugging • 141 Profiling DDP performance ................................................................................ 141 Using torch.profiler.profile for DDP analysis • 141 Analyzing Computation–Communication overlap • 145 Profiling Multi-Node DDP • 146 Optimizing DDP performance ........................................................................... 148 Mixed precision training • 148 Gradient accumulation • 149 Communication overlap optimization • 150 Tuning bucket sizes • 151 Finding unused parameters • 151 Static graph optimization • 152 Checkpointing and resuming distributed jobs .................................................... 152 Saving checkpoints • 153 Loading checkpoints • 153 Saving RNG state for reproducibility • 154 Atomic checkpointing • 155 Checkpointing best practices • 155 Advanced DDP features ..................................................................................... 156 Gradient hooks • 156 Communication hooks • 156 Handling uneven inputs with join() • 157 Best practices and common patterns .................................................................. 158 Elastic data parallelism .................................................................................... 160 How elastic training works • 160 Launching elastic training • 161 When to use elastic training • 162 Complete example: DDP with a transformer ....................................................... 162 Summary .......................................................................................................... 162 Table of Contents xvi
Page
18
Code summary .................................................................................................. 163 References ......................................................................................................... 164 Exercises ........................................................................................................... 164 Implement manual gradient synchronization • 164 Compare DP vs DDP performance • 165 Implement gradient bucketing • 166 Implement Fault-Tolerant DDP training • 167 Analyze communication patterns with NCCL • 168 Expected learning outcomes • 170 Chapter 4: Scaling with Fully Sharded Data Parallel (FSDP) 171 From DDP to FSDP ............................................................................................. 171 Why FSDP enables Larger-Than-Memory models ............................................... 172 How FSDP works • 173 FSDP1 vs FSDP2: the evolution • 174 Original FSDP (FSDP1) • 175 FSDP2: the per-parameter sharding API ............................................................. 177 Device mesh: the foundation • 178 Key parameters • 179 Hierarchical sharding • 181 A complete working example: T5 summarization with FSDP .............................. 182 Checkpointing with FSDP2 ................................................................................ 187 Using the DCP API (recommended) • 187 Manual sharded checkpointing • 188 Full state dict for evaluation • 190 Prefetching: optimizing communication ............................................................ 191 Forward prefetching • 191 Backward prefetching • 192 When to use prefetching • 192 Activation checkpointing and offloading ............................................................ 192 Why activation checkpointing matters with FSDP • 193 Using activation checkpointing • 193 xvii Table of Contents
Page
19
CPU offloading • 194 When to use offloading • 195 Performance optimization ................................................................................. 195 Profiling FSDP training • 195 Optimizing communication • 196 Optimizing activation memory • 197 Memory profiling • 197 Multi-node FSDP training .................................................................................. 197 Setting up multi-node FSDP • 198 Network configuration for multi-node • 198 Checkpointing on multi-node • 199 SLURM integration • 199 Scaling considerations • 200 Debugging FSDP issues ..................................................................................... 201 Out of Memory (OOM) • 201 Hanging or deadlocked • 201 Slow training • 201 Incorrect results • 202 Debugging tools • 202 Comparing FSDP2 with ZeRO and DDP ............................................................. 202 Memory and performance • 203 When to use what • 203 Practical tips .................................................................................................... 203 Initialization best practices • 204 Data loading and gradient clipping • 204 Mixed precision • 204 Progressive Optimization • 204 Advanced topics ............................................................................................... 205 Hybrid sharding (HSDP) • 205 Compiler integration • 205 Other integrations • 205 FSDP via SPMD for TPU/XLA • 205 Table of Contents xviii
Page
20
Summary ......................................................................................................... 207 Code Summary ................................................................................................. 208 Useful links ...................................................................................................... 208 Notes and references ......................................................................................... 209 Exercises .......................................................................................................... 210 Implement manual parameter sharding • 210 Compare FSDP sharding strategies • 211 Implement custom FSDP wrapping policy • 212 Implement mixed precision with FSDP • 213 Implement FSDP checkpointing • 214 Expected learning outcomes • 215 Chapter 5: Beyond State Sharding with DeepSpeed and Megatron 217 Beyond state sharding ....................................................................................... 217 ZeRO stage 1: optimizer state partitioning .......................................................... 218 ZeRO stage 2: optimizer state and gradient partitioning .................................... 220 ZeRO stage 3: full sharding (like FSDP) ............................................................... 221 ZeRO-Offload: CPU memory extension .............................................................. 222 ZeRO-Infinity: NVMe offload for massive models ............................................... 223 ZeRO++: Communication-Optimized ZeRO ....................................................... 226 Megatron: computation parallelism as a second axis ......................................... 228 State sharding vs. computation sharding • 228 Tensor Parallelism: sharding the layer • 229 Pipeline Parallelism: sharding the depth • 231 Sequence parallelism and context parallelism: long contexts • 233 DeepSpeed-Ulysses: an alternative to ring attention • 236 Expert Parallelism: scaling MoE models • 238 Why FSDP2 cannot replace Megatron • 240 Hybrid Parallelism: combining state and computation sharding • 241 Megatron core: Production-Ready library • 242 Megatron-FSDP: optimized state sharding • 242 Distributed optimizer: Memory-Efficient optimization • 243 xix Table of Contents