Distributed training: DP, DDP, FSDP, pipeline parallel

Trong distributed systems, có một câu nói nổi tiếng: “There are only two hard things in computer science: cache invalidation and naming things”. Trong distributed training, có ba thứ khó: chia data, chia model, và đồng bộ gradient. Tất cả còn lại là wrapper xung quanh ba thứ này.

Bài này không train model thật, vì điều đó cần cluster 8+ GPU. Thay vào đó, bài giải thích mental model của 4 paradigm chia trong training: Data Parallel (DP), Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP / ZeRO), và Pipeline Parallel. Mỗi cái giải quyết một bottleneck khác nhau.

Sau bài này bạn đọc paper training của LLM lớn (Megatron, DeepSpeed, Llama-3 paper) không bị lạc. Khi cần setup multi-GPU thực tế (ngay cả 2 GPU), biết chọn đúng paradigm.

Mental model: phân loại theo cái gì bị chia

Có 3 thứ trong training có thể chia qua nhiều GPU:

1. Data: mỗi GPU thấy batch khác nhau.
2. Model: mỗi GPU giữ một phần của model (params, gradients, optimizer state).
3. Layer / Sequence: mỗi GPU xử lý một range layer hoặc một range của sequence.

Bốn paradigm map vào đây:

Paradigm	Data	Model weight	Optimizer state	Layer
DP (Data Parallel)	chia	replicate	replicate	replicate
DDP	chia	replicate	replicate	replicate
FSDP / ZeRO-3	chia	chia	chia	replicate
Tensor Parallel	replicate (per group)	chia matrix	chia	replicate
Pipeline Parallel	chia (micro-batch)	replicate (per stage)	replicate	chia

Quy tắc chung:

• Model nhỏ (fits 1 GPU): dùng DDP.
• Model lớn (không fit 1 GPU): dùng FSDP hoặc combo.
• Model cực lớn (>100B): combo FSDP + Tensor Parallel + Pipeline Parallel.

Phần 1: Data Parallel (DP), pattern cũ

DP là pattern đơn giản nhất, có từ thời CNN.

GPU 0:  full model + batch[0:32]
GPU 1:  full model + batch[32:64]
GPU 2:  full model + batch[64:96]
GPU 3:  full model + batch[96:128]

Sau forward + backward:
GPU 0: gradient_0
GPU 1: gradient_1
GPU 2: gradient_2
GPU 3: gradient_3

Avg gradient = mean(grad_0, grad_1, grad_2, grad_3)
Cập nhật weights bằng avg gradient.

Mỗi GPU giữ full model. Chia batch qua các GPU. Sau backward, average gradient và update.

PyTorch:

model = nn.DataParallel(model, device_ids=[0, 1, 2, 3])
output = model(batch)

Vấn đề của DP:

1. Single-process multi-thread trong Python. Python GIL chặn parallelism. Slow.
2. Master GPU (device 0) làm việc nhiều hơn: gom output, broadcast lại. Memory imbalance.
3. Scaling kém vượt 2-4 GPU.

DP đã obsolete. PyTorch khuyến nghị dùng DDP thay thế.

Phần 2: Distributed Data Parallel (DDP), pattern hiện tại

DDP fix mọi vấn đề của DP bằng cách chạy một process per GPU, không phải một process multi-thread.

Process 0 (GPU 0): full model + batch[0:32]
Process 1 (GPU 1): full model + batch[32:64]
Process 2 (GPU 2): full model + batch[64:96]
Process 3 (GPU 3): full model + batch[96:128]

Forward + backward (parallel hoàn toàn):
  P0: grad_0
  P1: grad_1
  P2: grad_2
  P3: grad_3

All-reduce (NCCL) đồng bộ gradient:
  Mọi process: avg_grad = mean(grad_0, grad_1, grad_2, grad_3)

Update: mỗi process tự update weights (vì avg_grad đều giống nhau).

Key insight: mỗi process độc lập update, nhưng sau all-reduce, gradient giống nhau nên weights sync.

PyTorch:

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)
    model = MyModel().to(rank)
    model = DDP(model, device_ids=[rank])

    for batch in dataloader:
        output = model(batch)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()

    cleanup()

Chạy với torchrun:

torchrun --nproc_per_node=4 train.py

Spawn 4 process, mỗi process 1 GPU. NCCL backend xử lý all-reduce hiệu quả qua NVLink/InfiniBand.

DDP scaling tốt tới ~256 GPU trong một node. Vượt qua đó cần tinh chỉnh communication.

Hạn chế DDP: mỗi GPU phải fit full model + full optimizer state. Không train được model lớn hơn VRAM 1 GPU.

Phần 3: FSDP / ZeRO, chia model qua các GPU

DeepSpeed ZeRO (2020) và PyTorch FSDP (2022) là cùng một ý tưởng: thay vì replicate model trên mọi GPU, chia model qua các GPU.

Ba “stage” của ZeRO:

ZeRO-1: chia optimizer state.

• Weights, gradients: replicate.
• Optimizer state (m, v của Adam): chia qua N GPU.
• Cắt memory ~2x cho optimizer.

ZeRO-2: chia optimizer state + gradients.

• Weights: replicate.
• Gradients: chia. Sau backward, gradient được reduce-scatter xuống GPU sở hữu shard đó.
• Optimizer: chia.
• Cắt memory ~3x total.

ZeRO-3 (FSDP): chia tất cả.

• Weights: chia. Khi cần forward một layer, gather shard từ các GPU khác.
• Gradients: chia (như ZeRO-2).
• Optimizer: chia.
• Cắt memory ~N x với N GPU.

FSDP flow chi tiết:

Mỗi GPU giữ 1/N của mỗi param.

Forward layer L:
  1. All-gather: gom shard từ tất cả GPU -> full weight của L
  2. Compute forward với full weight
  3. Discard full weight, giữ shard của mình

Backward layer L:
  1. All-gather: gom shard
  2. Compute backward, sinh gradient full
  3. Reduce-scatter: cộng gradient và chia về shard
  4. Discard full weight, full gradient

Optimizer step:
  Mỗi GPU update shard của mình (vì optimizer state cũng chia)

Trade-off: thêm communication (all-gather + reduce-scatter) cho mỗi layer mỗi step. Với GPU NVLink trong 1 node, overhead nhỏ. Với GPU qua InfiniBand cross-node, overhead lớn hơn.

PyTorch FSDP:

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

model = MyTransformer().to(rank)
model = FSDP(
    model,
    auto_wrap_policy=transformer_auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
)

Sau khi wrap, model có thể train với weight tổng >> VRAM 1 GPU.

Llama-3 70B train trên 1024 H100 với FSDP. 70B weight FP32 = 280GB, chia 1024 GPU thì mỗi GPU chỉ giữ ~280MB. Còn lại VRAM 80GB dùng cho activation + compute.

Phần 4: Tensor Parallel, chia matrix

DDP và FSDP đều chia theo data hoặc theo shard params. Tensor Parallel chia matrix multiplication qua nhiều GPU.

Trong một linear layer Y = X @ W, ma trận W có shape [in, out]. Tensor Parallel chia W theo cột:

W = [W_0 | W_1 | W_2 | W_3]   (chia 4 cột)

GPU 0 giữ W_0, compute Y_0 = X @ W_0
GPU 1 giữ W_1, compute Y_1 = X @ W_1
...
Output Y = concat(Y_0, Y_1, Y_2, Y_3)

Hoặc chia theo hàng:

W = [W_0; W_1; W_2; W_3]   (chia 4 hàng)
X = [X_0 | X_1 | X_2 | X_3]

GPU i compute Y_i = X_i @ W_i
Output Y = sum(Y_0, Y_1, Y_2, Y_3)   (all-reduce sum)

Megatron-LM của NVIDIA là implement điển hình. Trong một attention block:

• Q, K, V projection: chia theo cột (mỗi GPU giữ vài attention head)
• Output projection: chia theo hàng (all-reduce sum cuối)

Tensor Parallel cần GPU rất gần nhau vì all-reduce mỗi layer mỗi step. Thường chỉ dùng trong một node (8 GPU NVLink). Cross-node bandwidth không đủ.

GPT-4 (rumored): TP=8, PP=16, DP=128. Tổng 16384 GPU.

Phần 5: Pipeline Parallel, chia layer

Pipeline Parallel chia model theo layer:

GPU 0: layers 0-23
GPU 1: layers 24-47
GPU 2: layers 48-71
GPU 3: layers 72-95

Một micro-batch đi qua GPU 0 -> GPU 1 -> GPU 2 -> GPU 3. Vấn đề: GPU 1, 2, 3 idle khi GPU 0 đang forward.

Fix: GPipe chia batch thành nhiều micro-batch, pipeline chúng:

time -->
GPU 0: m0 m1 m2 m3 m0(bwd) m1(bwd) m2(bwd) m3(bwd)
GPU 1:    m0 m1 m2 m3 m0(bwd) m1(bwd) m2(bwd) m3(bwd)
GPU 2:       m0 m1 m2 m3 m0(bwd) m1(bwd) m2(bwd) m3(bwd)
GPU 3:          m0 m1 m2 m3 m0(bwd) m1(bwd) m2(bwd) m3(bwd)

GPU 1, 2, 3 chỉ idle ở “bubble” đầu và cuối. Với nhiều micro-batch, bubble nhỏ so với tổng time.

Pipeline Parallel có overhead truyền activation giữa GPU. Phù hợp cross-node (vì chỉ truyền giữa 2 GPU adjacent, không all-reduce).

Phần 6: Combo thực tế cho LLM lớn

Llama-3-405B paper mô tả setup:

• TP = 8 (trong 1 node, NVLink)
• PP = 16 (cross-node, InfiniBand)
• FSDP = 16 (cross-node)

Tổng: 8 * 16 * 16 = 2048 GPU.

Mỗi node 8 GPU NVLink: TP nội bộ. 16 node thành 1 pipeline group: PP. 16 pipeline group: FSDP (replicate model theo data).

Setup này gọi là 3D Parallelism. Megatron-DeepSpeed framework support.

Quy tắc thực tế:

Số GPU	Setup khuyến nghị
1	Pure (no parallelism)
2-8	DDP nếu model fit, FSDP nếu không
8-32	FSDP
32-128	FSDP + checkpoint, hoặc TP=8 + FSDP
128-1024	TP=8 + FSDP, hoặc TP=8 + PP + FSDP
1024+	3D Parallelism

Pitfall: scaling không tuyến tính

Một dev có lần train 1B model trên 8x A100, throughput 30K tokens/s. Scale lên 16x A100, mong đợi 60K tokens/s. Thực tế: 42K.

Lý do: communication overhead. Với DDP, all-reduce gradient mỗi step tốn ~5% time trên 8 GPU. Lên 16 GPU, overhead lên 12%. Lên 64 GPU, 25%. Scaling efficiency giảm dần.

Fix:

1. Overlap communication và compute. DDP có cờ gradient_as_bucket_view=True để async all-reduce trong khi backward tiếp tục.
2. Tăng batch size. Communication tỷ lệ với param size, không phải batch size. Batch lớn = computation per step lớn = overhead ratio nhỏ.
3. Dùng InfiniBand thay Ethernet. NCCL on InfiniBand ~5x nhanh hơn Ethernet.

Bài học: scaling không free. Khi double GPU, kỳ vọng 1.7-1.8x throughput, không phải 2x. Linear scaling chỉ đạt được trong setup được tune kỹ.

Cheatsheet

Paradigm	Mỗi GPU giữ	Khi nào dùng
Pure (no parallel)	Full model	Model fit 1 GPU
DDP	Full model, batch chia	Model fit, scale data
FSDP / ZeRO-3	1/N model	Model > VRAM 1 GPU
Tensor Parallel	1/N của matrix W	Trong 1 node NVLink
Pipeline Parallel	1 stage layer	Cross-node, large model

Communication primitive	Mục đích
All-reduce	Gom + broadcast (DDP gradient avg)
All-gather	Gom toàn bộ shard về mọi GPU (FSDP forward)
Reduce-scatter	Gom + chia (FSDP backward gradient)
Send/recv	Truyền 1-1 (Pipeline Parallel)

Tool	Layer	Khi nào dùng
torch.nn.DataParallel	DP	Không bao giờ (obsolete)
torch.nn.parallel.DistributedDataParallel	DDP	1-8 GPU, model fit
torch.distributed.fsdp.FullyShardedDataParallel	FSDP	Model không fit 1 GPU
deepspeed	ZeRO 1/2/3	Alternative FSDP, nhiều feature hơn
megatron-deepspeed	TP + PP + ZeRO	Model > 100B

Backend	Hardware
nccl	GPU NVIDIA (default)
gloo	CPU, hoặc GPU không có NCCL
mpi	HPC cluster

Lời kết

Distributed training là phần khó nhất của LLM, nhưng cũng là kiến thức gating: nếu hiểu nó, bạn vào được nhóm dev có thể train model lớn. Trên thế giới có lẽ chỉ vài chục nghìn người làm việc này hàng ngày, và họ là những ML engineer giá trị nhất hiện tại.

Hands-on song song:

1. Setup multi-GPU thử nghiệm. Nếu không có 2 GPU, dùng RunPod / Vast.ai thuê instance 2x RTX 4090 ~$0.5/h. Chạy DDP example của PyTorch (https://pytorch.org/tutorials/intermediate/ddp_tutorial.html). Verify 2 process spawn đúng, all-reduce chạy.
2. Convert một training script đơn lẻ thành DDP. Step: thêm init_process_group, wrap model với DDP(), thay torch.utils.data.DataLoader bằng DistributedSampler. Đo throughput single GPU vs 2 GPU DDP.
3. Nếu có 4+ GPU, thử FSDP. Train một model 1B (ví dụ pythia-1b từ HuggingFace) với batch size lớn. Verify memory mỗi GPU < 1/4 model weight.
4. Đọc Llama-3 paper, section “Training Infrastructure”. Megatron-Llama notation TP=8, PP=16, DP=16. Vẽ ra ASCII diagram để hiểu mapping 16384 GPU thành 3D grid.

Bài 18 sẽ chuyển sang Part 5: LoRA và QLoRA. Sau khi đã hiểu pretraining đắt như thế nào (Llama-3-8B mất 100K GPU-hour = $1.5M trên cloud), fine-tuning là cách rẻ để adapt model. LoRA cắt từ 100% params xuống 0.1%, vẫn giữ 95% performance. Đây là kỹ thuật mọi dev nên biết, kể cả khi không train từ zero.