book

Distributed AI Systems

Name: Distributed AI Systems
Author: Fuheng Wu
ISBN: 9781807301712

by Fuheng Wu

June 2026

Intermediate

558 pages

13h 48m

English

Packt Publishing

Read now

Unlock full access

Preface
Free benefits with your book
Chapter 1: Introduction to Modern Distributed AI
OverviewWhy modern AI requires distributionThe scale challengeEstimating model resource requirementsTraining memory requirementsInference memory requirementsGPU requirements estimationReal-World considerationsThe evolution from classic ML to foundation modelsThe modern AI model lifecycleTraining: the learning phaseInference: the prediction/generation phaseServing: the production systemDecision framework: when do you need distributed systems?Decision tree: quick referenceUnderstanding the decision treeHands-On: running distributed trainingEnvironment setupSingle-GPU baselineMulti-GPU distributed trainingExtended training: CIFAR-10 with 20 epochsDistributed inference: throughput scalingDistributed AI fundamentals with PyTorchThe distributed AI stackProcess Groups and RanksInitializing the process groupCollective operationsAllReduceAllGatherBroadcastReduceGatherScatterReduceScatterAlltoAllChoosing the right operationDistributedDataParallel (DDP)DistributedSamplerLaunching distributed jobsCode summaryReferencesExercisesProcess group initializationManual AllReduce implementationBroadcast with verificationDistributedSampler wrapperGradient synchronization timingExpected learning outcomes
Chapter 2: GPU Hardware, Networking, and Parallelism Strategies
Computational power: AI clusters and metricsWhat is computational power?Why clusters?AI clusters: built for training and inferenceKey metrics for AI clustersCentral Processing Unit (CPU)CPU Architecture BasicsCPU-GPU InteractionNUMA and CPU affinityCPU requirements for distributed trainingGraphics Processing Unit (GPU)Understanding GPU memory and compute architectureKey architecture milestonesWhat these numbers mean for trainingGraphics Processing Unit (GPU) product familiesChoosing the right GPUArchitecture features that matterTensor Processing Unit (TPU)Why TPUs existTPU architecture: systolic arraysTPU generationsTPU pod architectureTPU vs GPU: when to use whichTPU programming modelSparse core (since v4)Neural Processing Unit (NPU)What makes NPUs differentNPU architecture: AI cores and memoryTraining vs inference NPUsNPU software stackNPU vs.GPU: when to choose whichNPU interconnects and scalingThe NPU landscapeHigh-Speed interconnects: the network backboneChip programming systems: SPMD and CUDAProgramming models vs.execution modelsSPMD: single program, multiple dataCUDA's execution model: SIMTWhy SIMT over SIMD?CUDA thread indexingMemory hierarchy in CUDAHow frameworks use CUDASPMD in distributed trainingAMD and other alternativesDistributed communication: patterns and primitivesParallelism: core strategiesParallelism taxonomy tableThe three questions for parallelism determinationData parallelism: replicated and shardedModel parallelism: compute shardingCombining StrategiesComposition patterns (hybrid parallelism)Strategy selection: choosing the right approachTraining strategy decision treeKey considerations for trainingInference strategy decision treeKey considerations for inferencePractical tipsHands-On: hardware inspection and bandwidth testingEnvironment SetupStep 1: inspect GPU hardwareStep 2: inspect hardware topologyStep 3: measure Single-GPU memory bandwidthStep 4: benchmark Inter-GPU communicationStep 5: analyze resultsCode summaryNotes and referencesExercisesGPU hardware inspectorMemory requirement calculatorTopology detection and analysisParallelism strategy selectorExpected learning outcomes
Chapter 3: Distributed Training with PyTorch DDP
How DDP works internallyThe basic flowData Parallel (DP) vs.Distributed Data Parallel (DDP)Gradient bucketing: why it mattersCommunication–Computation overlapThe AllReduce operationMixed precision and gradient scalingBuffer SynchronizationDDP forward pass implementation detailsDDP Computation–Communication overlap implementationSetting up Single-Node DDPUsing torchrun (recommended)Understanding the environment variablesUsing DistributedSamplerDataLoader internals for distributed trainingDevice selection best practicesA complete single node exampleSetting up Multi-Node DDPMulti-Node rchitectureLaunching Multi-Node trainingUsing SLURM for Multi-Node launchNetwork configurationEnvironment variables for Multi-NodeA complete Multi-Node exampleDebugging and troubleshootingHangs: the most common problemWrong results or inconsistent gradientsCUDA out of memoryPerformance issuesNetwork debuggingProfiling DDP performanceUsing torch.profiler.profile for DDP analysisAnalyzing Computation–Communication overlapProfiling Multi-Node DDPOptimizing DDP performanceMixed precision trainingGradient accumulationCommunication overlap optimizationTuning bucket sizesFinding unused parametersStatic graph optimizationCheckpointing and resuming distributed jobsSaving checkpointsLoading checkpointsSaving RNG state for reproducibilityAtomic checkpointingCheckpointing best practicesAdvanced DDP featuresGradient hooksCommunication hooksHandling uneven inputs with join()Best practices and common patternsElastic data parallelismHow elastic training worksLaunching elastic trainingWhen to use elastic trainingComplete example: DDP with a transformerSummaryCode summaryReferencesExercisesImplement manual gradient synchronizationCompare DP vs DDP performanceImplement gradient bucketingImplement Fault-Tolerant DDP trainingAnalyze communication patterns with NCCLExpected learning outcomes
Chapter 4: Scaling with Fully Sharded Data Parallel (FSDP)
From DDP to FSDPWhy FSDP enables Larger-Than-Memory modelsHow FSDP worksFSDP1 vs FSDP2: the evolutionOriginal FSDP (FSDP1)FSDP2: the per-parameter sharding APIDevice mesh: the foundationKey parametersHierarchical shardingA complete working example: T5 summarization with FSDPCheckpointing with FSDP2Using the DCP API (recommended)Manual sharded checkpointingFull state dict for evaluationPrefetching: optimizing communicationForward prefetchingBackward prefetchingWhen to use prefetchingActivation checkpointing and offloadingWhy activation checkpointing matters with FSDPUsing activation checkpointingCPU offloadingWhen to use offloadingPerformance optimizationProfiling FSDP trainingOptimizing communicationOptimizing activation memoryMemory profilingMulti-node FSDP trainingSetting up multi-node FSDPNetwork configuration for multi-nodeCheckpointing on multi-nodeSLURM integrationScaling considerationsDebugging FSDP issuesOut of Memory (OOM)Hanging or deadlockedSlow trainingIncorrect resultsDebugging toolsComparing FSDP2 with ZeRO and DDPMemory and performanceWhen to use whatPractical tipsInitialization best practicesData loading and gradient clippingMixed precisionProgressive OptimizationAdvanced topicsHybrid sharding (HSDP)Compiler integrationOther integrationsFSDP via SPMD for TPU/XLASummaryCode SummaryUseful linksNotes and referencesExercisesImplement manual parameter shardingCompare FSDP sharding strategiesImplement custom FSDP wrapping policyImplement mixed precision with FSDPImplement FSDP checkpointingExpected learning outcomes
Chapter 5: Beyond State Sharding with DeepSpeed and Megatron
Beyond state shardingZeRO stage 1: optimizer state partitioningZeRO stage 2: optimizer state and gradient partitioningZeRO stage 3: full sharding (like FSDP)ZeRO-Offload: CPU memory extensionZeRO-Infinity: NVMe offload for massive modelsZeRO++: Communication-Optimized ZeROMegatron: computation parallelism as a second axisState sharding vs.computation shardingTensor Parallelism: sharding the layerPipeline Parallelism: sharding the depthSequence parallelism and context parallelism: long contextsDeepSpeed-Ulysses: an alternative to ring attentionExpert Parallelism: scaling MoE modelsWhy FSDP2 cannot replace MegatronHybrid Parallelism: combining state and computation shardingMegatron core: Production-Ready libraryMegatron-FSDP: optimized state shardingDistributed optimizer: Memory-Efficient optimizationFP8 training: Next-Generation precisionWhen do you need Megatron?Real-World training configurationsComplete training example with MegatronHybrid Parallelism in practiceThe Two-Axis viewBuilding a hybrid configurationWhy this worksChoosing your configurationOperational realitiesPerformance optimization best practicesConfiguration patternsChoosing the right strategyA decision frameworkUnderstanding the Trade-OffsConcrete memory exampleSummaryCode summaryUseful linksReferencesExercisesCompare ZeRO stagesImplement CPU offloadingImplement Tensor ParallelismImplement Pipeline ParallelismConfigure Hybrid ParallelismExpected learning outcomes
Chapter 6: Distributed Inference and vLLM
From training to inferenceIntroduction to vLLMPrerequisitesInstallationDocker SetupInstall and run from package manager (uv, conda, pip)Compile, install, and run from local sourceOffline inferenceOnline inferenceKV cacheDecoder-Only transformer architectureText generation process: prefill and decodeThe decode phase and its inefficienciesKV cache solutionPagedAttention: solving KV cache fragmentationThe KV cache fragmentation problemHow PagedAttention worksEliminating padding FLOPsConnection to distributed inferenceMotivation: the "out of memory" problemOverview of the vLLM architectureOverview of parallelism strategies in vLLMTensor Parallelism (TP)The linear algebra foundationsBenefits of Tensor ParallelismWhen to use Tensor ParallelismData Parallelism (DP) for throughput scalingCombining Data Parallelism with other strategiesDeployment modesBenefits and Trade-OffsPipeline Parallelism (PP)The pipeline bubble problemChunked prefill for smoother pipelinesExpert Parallelism (EP): A modifier flag for MoE modelsUnderstanding MoE architectureDecoder layer structureExpert structureRouting mechanismHow Expert Parallelism modifies behaviorCombining parallelism strategiesTP + PP: the standard Multi-Node configurationExpert Parallelism combinationsTP + DP: scaling both waysThe EP activation constraintHands-On examplesBasic vLLM setup with Tensor ParallelismMulti-Node with TP and PPCustom chunked prefill configurationProfiling with Nsight SystemsBest practicesSummaryCode summaryUseful linksReferencesExercisesImplement KV cache managementImplement continuous batchingBenchmark vLLM throughputImplement speculative decodingDeploy vLLM with OpenAI-Compatible APIExpected learning outcomesSubscribe to Deep Engineering
Chapter 7: Cross-Request Optimization with SGLang
A different philosophy for distributed inferencePrerequisitesInstallationDocker SetupOverview of the SGLang architectureSGLang core theoryRadixAttention: prefix cache reuseStructured output decoding with XGrammarOperator fusionZero-overhead schedulerRouter-based distributed architectureSGLang model gateway architectureWhen Router-based architecture excelsPrefill/Decode disaggregationPD disaggregation architectureTransfer enginesPD disaggregation setupDistributed inference architecture and parallelism strategiesTensor Parallelism: weight shardingPipeline Parallelism: layer shardingData Parallelism: Request-level replicationHybrid Parallelism: combining strategiesCommunication patternsMulti-node SGLang deploymentMulti-node with Tensor ParallelismRouter-based multi-node deploymentRouter policies and load balancingSession affinity and cache localityFault toleranceExpert for MoE modelsSpeculative DecodingData-Parallel attentionProduction deployment patternsHands-on examplesBasic multi-node deploymentPD disaggregation with MooncakeTesting session affinityMoE model with EP+TPSummaryCode summaryUseful linksReferenceExercisesImplement RadixAttention cacheBenchmark prefix cachingImplementing constrained decodingImplement Multi-turn conversation with stateImplement parallel generation with forkExpected learning outcomes
Chapter 8: Running Distributed Training with SLURM
Introduction to clusters for HPC and AI trainingWhy SLURM for AI training?SLURM architecture: A brief overviewSetting up SLURM for Multi-GPU trainingSimulating a Multi-Node clusterQuick setup and verificationSubmitting distributed training jobsInteractive execution with srunBatch submission with sbatchUnderstanding SLURM environment variablesLaunching distributed training frameworks with SLURMDDP with SLURMFSDP with SLURM {#sec:slurm-fsdp-example}DeepSpeed with SLURMMegatron-LM with SLURMAdvanced SLURM features for trainingJob arrays for hyperparameter tuningInteractive jobs with sallocJob dependenciesCheckpointing and job resumptionMonitoring and debuggingJob monitoringLogging and outputProfiling distributed trainingBest practicesResource allocationCheckpointing strategyHandling failuresTroubleshooting common issuesNodes not availableGPU allocation issuesCommunication ErrorsJob hangingCode summaryUseful linksReferencesExercisesWrite a basic SLURM job scriptImplement automatic checkpointing with SLURMConfigure Multi-Node NCCL communicationImplement job array for hyperparameter searchMonitor and profile SLURM jobsExpected learning outcomes
Chapter 9: Production LLM Serving Stack
Anatomy of a production LLM serving systemRequest routing and traffic managementRouting strategiesLoad balancingCanary deployments and A/B testingOperations: observability, reliability, and costObservabilityReliability and fault toleranceCost optimizationDeploying LLM serving on kubernetesLocal development with k3dDeploying vLLM on k3dCleanupMulti-Model and Multi-Engine servingMulti-Model routingMulti-Engine routingKubernetes deployment with llm-dWhat is llm-d?Key featuresDeployment architectureGetting started with llm-dWell-Lit pathsMonitoring and observabilityChoosing Your PathMulti-Model servingComparing k3d and llm-d approachesSummaryCode summaryUseful linksReferencesExercisesBuild an API gateway with model routingImplement rate limiting middlewareImplement canary deployment with automated rollbackAdd distributed tracing with OpenTelemetryDeploy multi-model serving on k3dExpected learning outcomes

Chapter 10: Distributed Benchmarking and Performance Optimization
The performance gap in distributed systemsWhy benchmarking mattersBenchmarking fundamentalsUnderstanding percentile latenciesEfficiency metricsBenchmarking methodologyTraining benchmarkingKey training metricsPyTorch profilerNVIDIA Nsight SystemsCustom training benchmarkMeasuring scaling efficiencyValidating distributed training accuracyNetwork and communication profilingScaling bottleneck analysisInference benchmarkingKey inference metricsBenchmarking tools: genai-benchCustom inference benchmarksCold start vs warm performanceReasoning and multi-step modelsValidating inference accuracySummaryCode summaryUseful linksReferenceExercisesMultiple choice questionsShort answer questionsAnswer KeyExpected learning outcomes
Chapter 11: The Evolving Landscape of Distributed AI
The evolution of distributed AIWhere we stand todayThe new bottlenecksEmerging trends at a glanceMixture of experts: the dominant architectureThe Edge-cloud continuumWhy edge AI matters nowPractical On-device modelsEdge-cloud coordinationParallelism at scaleThe communication revolutionHierarchical parallelismSequence parallelism for long contextsRing attentionWhen things go wrong: fault toleranceThe checkpointing challengeDetecting and handling failuresElastic trainingBeyond text: multimodal distributed trainingThe Cross-modal challengeHandling multiple modalitiesFederated learning: an alternative paradigmThe federated paradigmWhere federated learning thrivesThe challenges of heterogeneityModern federated techniquesFederated learning for LLMsThe flower frameworkWhen to consider federated learningAgentic AI: the next frontierWhy agents need distributionMulti-agent patternsDistributed tool executionReasoning agentsThe economics of scaleWorkload distributionResource managementMulti-tenant GPU sharingGradient compressionPreparing for the futureSkills to developStaying currentSummaryCode summaryUseful linksReferencesExercisesMultiple choice questionsShort answer questionsHands-on exerciseImplement basic MoE routing
Chapter 12: Unlock Your Exclusive Benefits
Unlock this Book's Free Benefits in 3 Easy StepsAppendixOpenAI-Compatible API Endpoints
Other Books You May Enjoy
Subscribe to Deep Engineering
Index

Content preview from Distributed AI Systems

5 Beyond State Sharding with DeepSpeed and Megatron

Extending memory capacity and sharding computation for very large models

The future is already here. It's just unevenly distributed.

William Gibson, Writer

Beyond state sharding

In the previous chapter, we explored FSDP—PyTorch's approach to sharding parameters, gradients, and optimizer states across GPUs. FSDP2's full sharding is functionally equivalent to DeepSpeed's ZeRO Stage 3[1]: both eliminate memory redundancy by ensuring each GPU holds only 1/N of the training state.

State sharding solves the memory problem, but it doesn't change how computation happens. Every GPU still executes the same operations on the same model architecture—just with different data batches. For the largest models ...

Become an O’Reilly member and get unlimited access to this title plus top books and audiobooks from O’Reilly and nearly 200 top publishers, thousands of courses curated by job role, 150+ live events each month,
and much more.

Read now

Unlock full access

More than 5,000 organizations count on O’Reilly

O’Reilly covers everything we've got, with content to help us build a world-class technology community, upgrade the capabilities and competencies of our teams, and improve overall team performance as well as their engagement.

Julian F.

Head of Cybersecurity

I wanted to learn C and C++, but it didn't click for me until I picked up an O'Reilly book. When I went on the O’Reilly platform, I was astonished to find all the books there, plus live events and sandboxes so you could play around with the technology.

Addison B.

Field Engineer

I’ve been on the O’Reilly platform for more than eight years. I use a couple of learning platforms, but I'm on O'Reilly more than anybody else. When you're there, you start learning. I'm never disappointed.

Amir M.

Data Platform Tech Lead

I'm always learning. So when I got on to O'Reilly, I was like a kid in a candy store. There are playlists. There are answers. There's on-demand training. It's worth its weight in gold, in terms of what it allows me to do.

Mark W.

Embedded Software Engineer

Publisher Resources

ISBN: 9781807301712

Cloud Computing

Data Engineering

Data Science

AI & ML

Programming Languages

Software Architecture

IT/Ops

Security

Design

Business

Soft Skills

Distributed AI Systems

by Fuheng Wu

5

Beyond State Sharding with DeepSpeed and Megatron

Beyond state sharding

Become an O’Reilly member and get unlimited access to this title plus top books and audiobooks from O’Reilly and nearly 200 top publishers, thousands of courses curated by job role, 150+ live events each month,
and much more.