Distributed LLM Training 2026: FSDP, YaFSDP, DeepSpeed & Zero-3

2.1 Data Parallelism (DP) — Baseline 📊

Replicate model trên mọi GPU, split batch theo data dimension. Mỗi GPU tính forward/backward độc lập, all-reduce gradients sau backward.

• Pros: Đơn giản, scaling gần linear với batch size nhỏ
• Cons: Model + optimizer + gradients phải fit VRAM từng GPU → giới hạn model size
• Variants: DDP (DistributedDataParallel) — PyTorch standard, bucket-based gradient sync

2.2 Tensor Parallelism (TP) — Intra-Layer Sharding 🔧

Split weight matrices theo row/column dimension. Mỗi GPU cầm 1 shard của weight, all-reduce/all-gather trong forward/backward.

• Megatron-LM style: Column parallel (linear1) + Row parallel (linear2) trong MLP/Attention
• Pros: Giảm memory per GPU tuyến tính với TP degree
• Cons: Cần high-bandwidth interconnect (NVLink/NVSwitch), communication overhead mỗi layer
• Typical TP degree: 2-8 (trong 1 node)

2.3 Pipeline Parallelism (PP) — Inter-Layer Sharding 📦

Split model theo layer dimension. GPU 1 cầm layers 1-10, GPU 2 cầm layers 11-20... Micro-batches flow qua pipeline.

• Pros: Memory per GPU = model_layers / PP_degree, communication chỉ ở boundary layers
• Cons: Pipeline bubbles (idle time), load imbalance nếu layers không đều
• Schedules: GPipe (sync), 1F1B (async), ZeroBubble (hybrid)
• Typical PP degree: 2-8 (cross-node)

2.4 Sequence Parallelism (SP) — Activations Sharding 📝

Split activations theo sequence dimension. Quan trọng cho long-context training (32k-1M tokens).

• Ring Attention / Ring Attention v2: Split K/V theo sequence, ring-all-reduce attention
• Ulysses / Megatron-CP: All-to-all communication cho attention heads
• Pros: Cho phép context length > VRAM limit, linear scaling với sequence length

2.5 Context Parallelism (CP) — Long Context Specialist 📝

Variant của SP tối ưu cho attention. DeepSpeed-Ulysses, Megatron-CP, xFormers-ring.

• CP degree thường = TP degree hoặc SP degree
• Quan trọng cho LLM 128k-1M context (Llama-3.1, Nemotron-3, Qwen2.5)

3.1 Kiến Trúc FSDP & Sharding Strategies 🏗️

FSDP Memory Model (per GPU):

┌─────────────────────────────────────────────────────────────────┐
│                      FULL MODEL (7B params)                     │
│  Param: 14GB  |  Optim: 56GB (AdamW: 4x)  |  Grad: 14GB        │
│  Total: ~84GB per GPU (IMPOSSIBLE on single GPU)               │
└─────────────────────────────────────────────────────────────────┘
                              ↓ FSDP Sharding (DP=8)
┌─────────────────────────────────────────────────────────────────┐
│                      SHARDED PER GPU (DP rank)                  │
│  Param Shard: 1.75GB  |  Optim Shard: 7GB  |  Grad Shard: 1.75GB │
│  + Activation: ~6GB  |  Peak: ~16.5GB  ✅ FITS A100 40GB       │
└─────────────────────────────────────────────────────────────────┘

3 Sharding Strategies chính:

# FSDP Setup (PyTorch 2.3+)
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy, CPUOffload
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import functools

# Auto-wrap policy cho transformer layers
auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={LlamaDecoderLayer}  # hoặc GPTNeoXLayer, etc.
)

# FSDP config
fsdp_config = {
    "sharding_strategy": ShardingStrategy.FULL_SHARD,
    "cpu_offload": CPUOffload(offload_params=False),  # True để offload params to CPU
    "auto_wrap_policy": auto_wrap_policy,
    "backward_prefetch": BackwardPrefetch.BACKWARD_PRE,  # Prefetch next layer params
    "forward_prefetch": True,
    "use_orig_params": True,  # Cho phép optimizer hoạt động trên original params
    "sync_module_states": True,
    "param_init_fn": None,
}

model = FSDP(model, **fsdp_config)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5, foreach=True)  # foreach=True quan trọng cho perf

3.2 FSDP vs DDP: Memory & Communication 📊

4.1 Innovations: Buckets, Overlap, All-Gather 🔬

FSDP vs YaFSDP Communication Pattern:

FSDP (Sequential):
┌─────────────────────────────────────────────────────────────┐
│ Layer 1: [All-Gather Params] → [Forward] → [All-Gather]... │
│ Layer 2: [All-Gather Params] → [Forward] → [All-Gather]... │
│ Layer 3: [All-Gather Params] → [Forward] → [All-Gather]... │
│ ...                                                         │
│ Idle GPU time giữa all-gather và compute                    │
└─────────────────────────────────────────────────────────────┘

YaFSDP (Pipelined + Buckets):
┌─────────────────────────────────────────────────────────────┐
│ Bucket 1 (Layers 1-4):  [All-Gather] → [Forward]            │
│ Bucket 2 (Layers 5-8):           [All-Gather] → [Forward]   │
│ Bucket 3 (Layers 9-12):                 [All-Gather] → ...  │
│                                                         │
│ Overlap: All-Gather(Bucket N+1) || Forward(Bucket N)       │
│ Pre-fetch: Next bucket params trước khi cần                │
└─────────────────────────────────────────────────────────────┘

Key Innovations:

• Bucket-based sharding: Group layers vào buckets, all-gather per bucket thay vì per-layer → giảm kernel launches, better NIC utilization
• Communication-computation overlap: Double buffering — prefetch bucket N+1 trong khi compute bucket N
• Hierarchical all-gather: Intra-node NVLink all-gather → inter-node NCCL all-gather, tận dụng bandwidth hierarchy
• Optimizer state sharding: Tương tự FSDP nhưng overlap reduce-scatter với backward
• FP8 support: Native FP8 training (H100/H200), quantization-aware communication

# YaFSDP Setup (pip install yafsdp)
from yafsdp import YaFSDP

# Config tương tự FSDP nhưng có thêm overlap params
yafsdp_config = {
    "sharding_strategy": "full_shard",
    "bucket_size": 100_000_000,  # ~100M params per bucket
    "overlap_comm_compute": True,
    "prefetch_buckets": 2,  # Prefetch N buckets ahead
    "hierarchical_allgather": True,
    "fp8": True,  # H100+ only
    "use_orig_params": True,
}

model = YaFSDP(model, **yafsdp_config)

4.2 Benchmarks: YaFSDP vs FSDP 📈

5.1 Zero Stages: 1 → 2 → 3 → Offload 📊

# DeepSpeed Zero-3 Config (ds_config.json)
{
  "train_batch_size": 512,
  "train_micro_batch_size_per_gpu": 2,
  "gradient_accumulation_steps": 32,
  
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e7,
    "stage3_prefetch_bucket_size": 5e7,
    "stage3_param_persistence_threshold": 1e5,
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true,
    
    # Offload to CPU (optional)
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    }
  },
  
  "fp16": {
    "enabled": false
  },
  "bf16": {
    "enabled": true
  },
  
  "gradient_clipping": 1.0,
  "prescale_gradients": false,
  "wall_clock_breakdown": false
}

5.2 3D Parallelism & Hybrid Engine 🌐

DeepSpeed mạnh nhất ở hybrid 3D parallelism — kết hợp DP + TP + PP seamless. Megatron-LM style TP + PP với ZeRO DP.

# DeepSpeed 3D Parallelism Config
{
  "zero_optimization": {
    "stage": 3
  },
  
  "tensor_parallel": {
    "enabled": true,
    "tp_size": 4,
    "tp_mode": "megatron"
  },
  
  "pipeline": {
    "enabled": true,
    "pp_size": 2,
    "schedule": "1f1b",
    "micro_batches": 4
  },
  
  # Total GPUs = dp_size × tp_size × pp_size
  # Example: 128 GPUs = 16 × 4 × 2
}

6.1 Decision Matrix 🎯

6.2 Khuyến Nghị Theo Workload 🎯

7.1 Communication Optimization 📡

• NCCL Tuning: Set NCCL_ALGO=RING cho small messages, TREE cho large. NCCL_PROTO=LL (low latency) cho small, LL128 cho large.
• Topology Awareness: Sử dụng NCCL_TOPO_FILE hoặc torch.distributed.init_device_mesh để map ranks đúng NVLink domain.
• Bucket Sizing: FSDP/YaFSDP bucket ~50-200M params. DeepSpeed reduce_bucket_size ~50M. Tune theo model size.
• Overlap Strategy: Enable backward_prefetch=BACKWARD_PRE (FSDP), overlap_comm=true (DeepSpeed), overlap_comm_compute=true (YaFSDP).
• Gradient Accumulation: Large effective batch = small micro-batch × grad_accum. Reduces comm frequency. Typical: 32-128 steps.

# NCCL Env Vars cho optimal perf
export NCCL_ALGO=RING,TREE
export NCCL_PROTO=LL,LL128
export NCCL_MIN_NCHANNELS=4
export NCCL_MAX_NCHANNELS=16
export NCCL_BUFFSIZE=8388608  # 8MB buffer
export NCCL_NTHREADS=4
export NCCL_NSOCKS_PERTHREAD=4
export CUDA_DEVICE_MAX_CONNECTIONS=1  # Quan trọng cho overlap

# Cho multi-node Ethernet (non-NVLink)
export NCCL_IB_DISABLE=0
export NCCL_IB_GID_INDEX=3
export NCCL_IB_TC=106
export NCCL_IB_SL=5
export NCCL_IB_TIMEOUT=22

7.2 Memory Optimization Checklist 💾

# Activation Checkpointing Pattern (FSDP/YaFSDP/DS)
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper, CheckpointImpl, apply_activation_checkpointing
)

# Checkpoint every N transformer layers
non_reentrant_wrapper = lambda module: checkpoint_wrapper(
    module, checkpoint_impl=CheckpointImpl.NO_REENTRANT
)

apply_activation_checkpointing(
    model,
    checkpoint_wrapper_fn=non_reentrant_wrapper,
    check_fn=lambda m: isinstance(m, LlamaDecoderLayer)  # hoặc layer class tương ứng
)

7.3 Monitoring & Debugging 🔍

• PyTorch Profiler: torch.profiler.profile với schedule=torch.profiler.schedule(wait=1, warmup=1, active=3), trace NCCL kernels.
• NCCL Debug: NCCL_DEBUG=INFO hoặc VERSION để xem ring/tree selection, bus bandwidth.
• DeepSpeed Report: deepspeed --report generates HTML report với memory/comm breakdown.
• Memory Snapshot: torch.cuda.memory._dump_snapshot() + torch.cuda.memory._load_snapshot() visualize bằng memory_viz.
• TensorBoard: Log MFU (Model FLOPs Utilization), throughput, memory, loss curves.
• Key Metrics: MFU > 40% (good), > 50% (excellent). Comm/Compute ratio < 0.3. Memory fragmentation < 10%.

# MFU Calculation
def calculate_mfu(model_params, batch_size, seq_len, gpus, step_time, dtype="bf16"):
    # 6 × params × batch × seq (forward + backward)
    flops_per_token = 6 * model_params
    tokens_per_step = batch_size * seq_len * gpus
    total_flops = flops_per_token * tokens_per_step
    
    # GPU peak FLOPS (BF16 Tensor Core)
    gpu_peak_flops = {"a100": 312e12, "h100": 989e12, "h200": 989e12}[gpu_type.lower()]
    
    theoretical = gpu_peak_flops * gpus * step_time
    mfu = total_flops / theoretical
    return mfu

# Target: MFU > 0.4 (40%)