Study guide
6.2 Distributed Training Strategies & Architectures
AWS Certified Machine Learning Engineer — Associate MLA-C01
Track your progress as you read
Finished sections get a check mark.
1 Chapter 1: Introduction & Exam Strategy
1.1 AWS MLA-C01 Blueprint, Format & Strategy 1.2 The Machine Learning Engineering Lifecycle on AWS
2 Domain 1: Data Ingestion & Storage Architecture for ML
2.1 Ingesting Batch and Streaming Data for ML 2.2 High-Performance Storage for Training and Inference 2.3 Data Cataloging, Governance, and Querying
3 Domain 1: Feature Engineering & Data Transformation
3.1 Distributed Feature Engineering with AWS Glue & DataBrew 3.2 Interactive Data Preparation with SageMaker Data Wrangler 3.3 Feature Management with SageMaker Feature Store
4 Domain 1: Data Quality, Integrity & Ground Truth Labeling
4.1 Data Validation, Schema Drift, and Quality Rules 4.2 Handling Missing Values, Outliers & Class Imbalance 4.3 Data Labeling with SageMaker Ground Truth & Ground Truth Plus
5 Domain 2: Model Selection & Framework Architecture
5.1 SageMaker Built-in Algorithms & Selection Criteria 5.2 Custom Models with Script Mode and BYOC 5.3 Foundation Models, JumpStart, and Amazon Bedrock
6 Domain 2: SageMaker Model Training & Hyperparameter Optimization
6.1 SageMaker Training Jobs & Cost Optimization 6.2 Distributed Training Strategies & Architectures 6.3 Hyperparameter Optimization with Automatic Model Tuning (AMT)
7 Domain 2: Model Evaluation, Tracking & Model Registry
7.1 ML Evaluation Metrics & Validation Strategies 7.2 Experiment Tracking with SageMaker Experiments 7.3 Model Governance with SageMaker Model Registry
8 Domain 3: SageMaker Inference Options & Endpoint Architecture
8.1 Real-Time, Serverless, and Asynchronous Inference 8.2 High-Throughput Offline Inference with Batch Transform 8.3 Multi-Model Endpoints (MME) & Multi-Container Endpoints (MCE)
9 Domain 3: Deployment Strategies & Infrastructure Scaling
9.1 Production Deployment Patterns & Traffic Shifting 9.2 Endpoint Auto-Scaling & Hardware Acceleration
10 Domain 3: MLOps Pipelines & Workflow Orchestration
10.1 SageMaker Pipelines Core Architecture 10.2 Advanced Pipeline Patterns & Integration 10.3 End-to-End MLOps CI/CD & Automated Retraining
11 Domain 4: Production Model Monitoring with SageMaker Model Monitor
11.1 Data Quality & Model Quality Monitoring 11.2 Model Bias & Model Explainability Drift 11.3 Operational Observability with Amazon CloudWatch
12 Domain 4: Responsible AI, Bias Detection & Governance
12.1 Pre-Training and Post-Training Bias with SageMaker Clarify 12.2 Model Documentation & Auditing with Model Cards & Dashboard 12.3 Safety, Guardrails & Generative AI Governance
13 Domain 4: ML Security, Network Isolation & Cost Optimization
13.1 IAM Roles and Least Privilege for ML 13.2 Network Isolation, VPC Endpoints & Data Encryption 13.3 ML Cost Optimization & Resource Governance
6.2 Distributed Training Strategies & Architectures
Key Takeaways Explain This Topic AI
Distributed training divides into Data Parallelism (replicates model across all GPUs and splits mini-batches) and Model Parallelism (shards model layers/weights across GPUs when model parameters exceed a single GPU's VRAM).
SageMaker Distributed Data Parallel (SMDDP) optimizes communication collectives using AWS Nitro hardware offload and Elastic Fabric Adapter (EFA), delivering near-linear multi-node scaling compared to standard PyTorch DDP.
Model Parallelism encompasses Tensor Parallelism (intra-layer matrix splitting across GPUs via high-speed NVLink) and Pipeline Parallelism (inter-layer sequential layer partitioning across nodes with micro-batching).
PyTorch Fully Sharded Data Parallel (FSDP) and DeepSpeed ZeRO (Stages 1, 2, 3) eliminate memory redundancy by sharding optimizer states, gradients, and model parameters across data-parallel workers.
High-performance distributed training requires Cluster Placement Groups to minimize physical network latency within an Availability Zone, and Elastic Fabric Adapter (EFA) OS-bypass networking providing up to 3,200 Gbps bandwidth on P5 instances.
Last updated: August 2026
Distributed Training Strategies & Architectures
When training deep neural networks, computer vision models, or large language models (LLMs), single-GPU compute and memory limits quickly become insurmountable bottlenecks. Scaling these workloads requires Distributed Training across clusters composed of multiple nodes and multiple GPUs.
For the AWS Certified Machine Learning Engineer — Associate (MLA-C01) exam, you must understand the architectural distinction between Data Parallelism and Model Parallelism , know how to implement AWS-native distributed libraries ( SageMaker Distributed Data Parallel [SMDDP] and SageMaker Model Parallelism [SMP] ), utilize open-source standards ( PyTorch FSDP and DeepSpeed ZeRO ), and configure critical underlying infrastructure such as Elastic Fabric Adapter (EFA) and Cluster Placement Groups .
1. When to Distribute: The Core Scaling Decision Framework
Before selecting a distributed training strategy, you must diagnose the primary physical constraint: is the bottleneck dataset volume or model parameter size ? +-----------------------------------------------------------------------------------------+
- | DISTRIBUTED TRAINING DECISION MATRIX |
- | |
- | Does the complete model (weights + gradients + optimizer states) |
- | fit comfortably into the VRAM of a SINGLE GPU? |
- | | |
- | +---> YES: Bottleneck is dataset size / training epoch speed. |
- | | ---> DATA PARALLELISM |
- | | - Replicate identical model onto every GPU |
- | | - Partition training dataset (mini-batches) across GPUs |
- | | - Synchronize gradients via AllReduce (SMDDP / PyTorch DDP) |
- | | |
- | +---> NO: Model parameters exceed single GPU memory (e.g. >10B parameters). |
- | ---> MODEL PARALLELISM & SHARDED DATA PARALLEL |
- | - Model parameters sharded across GPUs |
- | - Approaches: Tensor Parallelism, Pipeline Parallelism, |
- | PyTorch FSDP, or DeepSpeed ZeRO-3 |
- +-----------------------------------------------------------------------------------------+
- 2. Data Parallelism & SageMaker Distributed Data Parallel (SMDDP)
In Data Parallelism (DP) , the entire neural network architecture and its weights are replicated identically across every GPU worker. The overall training batch (global batch size) is divided into smaller micro-batches (local batch size) processed concurrently by each GPU. +-----------------------------------------------------------------------------------------+
- | DATA PARALLELISM GRADIENT SYNC |
- | |
- | [Worker GPU 0] ---> Forward Pass (Batch 0) ---> Backward Pass ---> Local Gradients 0 |
- | [Worker GPU 1] ---> Forward Pass (Batch 1) ---> Backward Pass ---> Local Gradients 1 |
- | [Worker GPU 2] ---> Forward Pass (Batch 2) ---> Backward Pass ---> Local Gradients 2 |
- | [Worker GPU 3] ---> Forward Pass (Batch 3) ---> Backward Pass ---> Local Gradients 3 |
- | |
- | | |
- | v |
- | [ALLREDUCE COLLECTIVE COMMUNICATION (SMDDP / NCCL / EFA)] |
- | | |
- | v |
- | [All GPUs receive identical averaged Global Gradients] |
- | [All GPUs execute optimizer step ---> Weights Synchronized] |
- +-----------------------------------------------------------------------------------------+
- Parameter Server vs. Ring-AllReduce / SMDDP
Parameter Server Architecture (Legacy):
Dedicated parameter server nodes collect gradients from worker nodes, calculate parameter updates, and broadcast new weights back.
Bottleneck: Parameter server network bandwidth saturates as worker count increases, creating an I/O bottleneck.
Ring-AllReduce (NCCL Standard):
GPUs are organized in a logical ring. Each GPU sends and receives data only from its immediate neighbor in chunks.
Communication cost is independent of the number of GPUs, depending only on total model parameter size.
SageMaker Distributed Data Parallel (SMDDP):
AWS-optimized collective communication library designed specifically for the AWS cloud infrastructure.
Hardware Offload: Offloads AllReduce collective communication algorithms directly onto AWS Nitro System hardware and utilizes Elastic Fabric Adapter (EFA) .
Performance: Achieves near-linear scaling efficiency (up to 95% scaling efficiency across hundreds of GPUs), significantly outperforming standard NCCL over multi-node clusters.
Configuring SMDDP in SageMaker Python SDK from sagemaker.pytorch import PyTorch
- smddp_estimator = PyTorch(
- entry_point='train_resnet.py',
- source_dir='src',
- role=role,
- instance_count=8, # 8 multi-GPU nodes
- instance_type='ml.p4d.24xlarge', # 8 x NVIDIA A100 per node = 64 GPUs total
- framework_version='2.1.0',
- py_version='py310',
- distribution={
- 'smdistributed': {
- 'dataparallel': {
- 'enabled': True # Enables SMDDP communication collective
- }
- }
- }
- )
- 3. Model Parallelism: Tensor vs. Pipeline Parallelism
When a model exceeds the memory capacity of a single GPU (for example, a 13B, 70B, or 405B parameter LLM requiring hundreds of gigabytes just to store parameters, activations, and optimizer states), Model Parallelism is required to partition the model architecture across multiple GPUs. +-----------------------------------------------------------------------------------------+
- | TENSOR PARALLELISM VS. PIPELINE PARALLELISM |
- | |
- | 1. TENSOR PARALLELISM (Intra-Layer Sharding): |
- | - Splits individual weight matrices (e.g., Attention QKV, MLP) within a single layer. |
- | - High communication frequency; requires ultra-high bandwidth NVLink inside a node. |
- | |
- | Layer 1 Matrix W: [ GPU 0: W[:, :k] ] <-- NVLink Sync --> [ GPU 1: W[:, k:] ] |
- | |
- | 2. PIPELINE PARALLELISM (Inter-Layer Sharding): |
- | - Partitions sequential layers across different nodes or GPUs. |
- | - Layer 1-8 on Node 0 ---> Layer 9-16 on Node 1 ---> Layer 17-24 on Node 2. |
- | - Requires micro-batching (1F1B schedule) to minimize pipeline idle bubbles. |
- +-----------------------------------------------------------------------------------------+
- Tensor Parallelism (Intra-Layer Sharding)
Mechanism: Matrix operations within an individual layer (such as the linear projections in multi-head self-attention or feed-forward networks in transformers) are sliced across GPUs (Column-Parallel and Row-Parallel linear layers).
Communication Requirement: Requires frequent AllGather and ReduceScatter operations after every transformer block.
Placement: Must run within the same physical node connected via high-bandwidth NVLink (up to 900 GB/s on NVIDIA A100/H100) rather than across network cables.
Pipeline Parallelism (Inter-Layer Sharding)
Mechanism: The model's sequential layers are divided into stages, with each stage assigned to a different GPU or node.
The "Pipeline Bubble" Problem: Early stages sit idle waiting for backward passes from later stages.
Micro-Batching (1F1B Schedule): The global batch is split into micro-batches. Workers interleave one forward pass (1F) with one backward pass (1B) across micro-batches, keeping all pipeline stages active and reducing the bubble fraction.
4. Sharded Data Parallelism: DeepSpeed ZeRO & PyTorch FSDP
While traditional Model Parallelism requires intrusive modifications to model code, Sharded Data Parallelism provides an elegant alternative that allows training massive models using familiar data-parallel paradigms.
In standard deep learning training with Adam optimizer, GPU memory is consumed by three primary components:
Model Parameters (Weights): 1 × 1\times 1 × (4 bytes per parameter in FP32, 2 bytes in FP16/BF16).
Gradients: 1 × 1\times 1 × (2 or 4 bytes per parameter).
Optimizer States (Adam): 3 × 3\times 3 × (FP32 master weights [4B] + Momentum [4B] + Variance [4B] = 12 bytes per parameter). +-----------------------------------------------------------------------------------------+
- | DEEPSPEED ZERO & PYTORCH FSDP MEMORY STAGES |
- | |
- | Standard DDP: [ Parameters ] [ Gradients ] [ Optimizer States ] (Duplicated) |
- | |
- | ZeRO-Stage 1: [ Parameters ] [ Gradients ] [ Optimizer Sharded across GPUs ] |
- | ---> 4x memory reduction |
- | |
- | ZeRO-Stage 2: [ Parameters ] [ Gradients Sharded ] [ Optimizer Sharded ] |
- | ---> 8x memory reduction |
- | |
- | ZeRO-Stage 3 / [ Parameters Sharded ] [ Gradients Sharded ] [ Optimizer Sharded] |
- | PyTorch FSDP: ---> Linear memory reduction proportional to total GPU count! |
- | ---> Parameters gathered on-the-fly during forward/backward pass |
- +-----------------------------------------------------------------------------------------+
Comparison of Sharding Strategies:
ZeRO Stage 1: Optimizer states are partitioned across data parallel workers. Each GPU updates only its shard of optimizer states (reducing total memory footprint by ~4x).
ZeRO Stage 2: Both optimizer states and gradients are partitioned across workers (~8x reduction).
ZeRO Stage 3 / PyTorch FSDP (Fully Sharded Data Parallel): Optimizer states, gradients, and model parameters are all sharded across GPUs. During the forward pass, each layer's full parameters are dynamically gathered via an AllGather operation, computed, and immediately freed from memory. PyTorch FSDP is natively supported in SageMaker Script Mode.
5. Networking Infrastructure: Placement Groups & EFA
Distributed training performance is heavily bound by inter-node network throughput and packet latency. When running multi-node clusters, standard cloud networking introduces jitter that degrades AllReduce synchronization. +-----------------------------------------------------------------------------------------+
- | DISTRIBUTED TRAINING NETWORKING INFRASTRUCTURE |
- | |
- | [Cluster Placement Group] (Single Availability Zone) |
- | +---------------------------------------------------------------------------------+ |
- | | Node 1 (ml.p4d.24xlarge) Node 2 (ml.p4d.24xlarge) | |
- | | +-----------------------------+ +-----------------------------+ | |
- | | | 8x A100 GPUs (NVLink Mesh) | | 8x A100 GPUs (NVLink Mesh) | | |
- | | +-----------------------------+ +-----------------------------+ | |
- | | | | | |
- | | v v | |
- | | [EFA OS-Bypass] [EFA OS-Bypass] | |
- | | (Libfabric / SRD) (Libfabric / SRD) | |
- | | | | | |
- | | +============== 400 Gbps Fabric ===============+ | |
- | +---------------------------------------------------------------------------------+ |
- +-----------------------------------------------------------------------------------------+
- 1. Cluster Placement Groups
Mechanism: Enforces that all EC2 instances in the training fleet are physically located within the same rack or adjacent racks inside a single Availability Zone (AZ).
Benefit: Delivers the lowest possible inter-node latency (single-digit microseconds) and the highest bi-directional bandwidth, which is mandatory for multi-node AllReduce synchronization.
2. Elastic Fabric Adapter (EFA)
Mechanism: A custom network device for EC2 instances that provides Operating System (OS) bypass .
Protocol: Uses AWS's proprietary Scalable Reliable Datagram (SRD) protocol, which routes network packets across multi-path network fabrics dynamically to avoid congestion hotspots and packet out-of-order latency.
Interface: Deep learning frameworks communicate directly with hardware via Libfabric and NCCL/SMDDP without involving the Linux kernel network stack, achieving sub-10 microsecond latency and up to 3,200 Gbps network bandwidth on ml.p5 instances.