6.2 Distributed Training Strategies & Architectures

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

1. 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.

2. 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.

3. 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:

1. Model Parameters (Weights): $1\times$ (4 bytes per parameter in FP32, 2 bytes in FP16/BF16).
2. Gradients: $1\times$ (2 or 4 bytes per parameter).
3. Optimizer States (Adam): $3\times$ (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.