Distributed LLM Training: Từ Data Parallelism đến 3D Parallelism

1. Giới Thiệu 🚀

Năm 2026, các Large Language Models đã đạt tới hàng trăm tỷ tham số. Llamà 3 70B cần 140GB VRAM chỉ cho weights ở FP16. DeepSeek-V3 với 671B tham số là điều không tưởng trên một GPU đơn. Bài toán đặt ra: làm sao phân tán việc training qua nhiều GPU một cách hiệu quả?

Các chiến lược parallelism cần biết:

• Data Parallelism (DDP) sao chép model lên mỗi GPU, chia nhỏ batch
• Fully Sharded Data Parallelism (FSDP) shard model parameters ra tất cả GPU
• Tensor Parallelism (TP) chia nhỏ từng phép toán tensor
• Pipeline Parallelism (PP) chia model theo chiều sâu layers

2. Thách Thức VRAM & Memory 💾

Trước khi đi vào distributed training, bạn cần hiểu chính xác một model LLM cần bao nhiêu VRAM khi training. Có 4 thành phần chính:

• Model Parameters (W): 2 bytes/tham số (FP16) x N tham số. Llamà 3 8B = 16GB.
• Gradients (dW): Bằng với parameters. 8B = 16GB.
• Optimizer States: AdamW lưu 2 states (momentum + variance) x 4 bytes (FP32) x N = 8 bytes/tham số.
• Activations: Phụ thuộc batch size, seq_len, hidden_dim. Thường ~5-10GB cho 8B model.

Thực tế, các framework dùng mixed precision training: weights FP16 cho forward/backward, optimizer states FP32 cho update. Tổng: 16GB (W FP16) + 16GB (gradients FP16) + 16GB (optimizer FP32) + activations = ~56GB cho 8B model.

So sánh VRAM GPU phổ biến: RTX 4090 (24GB), A100 (40/80GB), H100 (80GB). Rõ ràng: không GPU đơn nào training được Llamà 3 70B.

3. Data Parallelism & DDP 📊

📌 3.1 Cách Hoạt Động

Data Parallelism là chiến lược đơn giản nhất: mỗi GPU giữ một bản sao đầy đủ của model, nhưng chỉ xử lý một phần của batch. Gradient từ tất cả GPU được đồng bộ (all-reduce), model parameters cập nhật đồng nhất.

PyTorch triển khai qua DistributedDataParallel (DDP) multi-process, mỗi GPU một process riêng, tránh GIL bottleneck. Dùng backend NCCL cho GPU.

🎯 3.2 Code Example: DDP

# train_ddp.py PyTorch DDP với 4 GPU
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler

def setup(rank, world_size):
    dist.init_process_group(backend="nccl",
        init_method="env://",
        world_size=world_size, rank=rank)
    torch.cuda.set_device(rank)

def train(rank, world_size):
    setup(rank, world_size)
    model = MyLLM().to(rank)
    ddp_model = DistributedDataParallel(model, device_ids=[rank])
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = DataLoader(dataset, batch_size=8, sampler=sampler)
    optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=3e-4)
    for epoch in range(10):
        sampler.set_epoch(epoch)
        for batch in loader:
            x, y = batch[0].to(rank), batch[1].to(rank)
            loss = ddp_model(x, y)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    cleanup()

# torchrun --nproc_per_node=4 train_ddp.py
if __name__ == "__main__":
    world_size = int(os.environ["WORLD_SIZE"])
    rank = int(os.environ["LOCAL_RANK"])
    train(rank, world_size)

📊 3.3 DDP Scaling Benchmark

Hạn chế: Mỗi GPU giữ full model copy waste memory. Model 70B cần 140GB/GPU bất khả thi. Communication cost O(n^2) mỗi GPU phải gửi/nhận gradient từ tất cả GPU khác.

4. Fully Sharded Data Parallelism (FSDP) 📦

📌 4.1 Tại Sao Cần FSDP?

DDP giới hạn bởi VRAM vì mỗi GPU chứa full model. FSDP giải quyết bằng cách shard (phân mảnh) model parameters, gradients, và optimizer states ra tất cả GPU chỉ tập hợp (all-gather) khi cần cho forward/backward. Đây là ý tưởng của DeepSpeed ZeRO (Zero Redundancy Optimizer).

FSDP lấy cảm hứng từ ZeRO Stage 3. PyTorch đã tích hợp FSDP native từ 2023, với API đơn giản hơn ZeRO.

📌 4.2 Cách Hoạt Động

FSDP wrap model với các FSDP units (thường là từng transformer layer). Mỗi unit chỉ giữ 1/N tham số trên GPU hiện tại. Khi forward qua unit đó, nó thực hiện all-gather để tập hợp đủ tham số, tính toán, rồi free tham số không cần thiết.

🎯 4.3 Code Example: FSDP

# fsdp_train.py PyTorch FSDP
from torch.distributed.fsdp import (
    FullyShardedDataParallel, MixedPrecision,
    ShardingStrategy, BackwardPrefetch
)
from torch.distributed.fsdp.wrap import (
    transformer_auto_wrap_policy
)
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper, apply_activation_checkpointing
)
import functools

def setup_fsdp(model):
    # 1. Auto wrap tung Transformer layer
    auto_wrap_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={TransformerBlock}
    )
    # 2. Mixed precision BF16
    mp_policy = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    )
    # 3. FSDP model
    fsdp_model = FullyShardedDataParallel(
        model,
        auto_wrap_policy=auto_wrap_policy,
        mixed_precision=mp_policy,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        device_id=torch.cuda.current_device(),
        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
        limit_all_gathers=True,
    )
    # 4. Activation checkpointing giam memory 50-70%
    apply_activation_checkpointing(
        fsdp_model,
        checkpoint_wrapper_fn=checkpoint_wrapper,
        check_fn=lambda m: isinstance(m, TransformerBlock)
    )
    return fsdp_model

# So sanh memory footprint (model 70B FP16 = 140GB):
# DDP với 8 GPU: 140GB / GPU (impossible)
# FSDP ZeRO-2: 70GB / GPU (gradients sharded)
# FSDP ZeRO-3: 17.5GB / GPU (params+grads+states sharded)
# FSDP ZeRO-3 + CP: ~9GB / GPU (them activation checkpointing)

5. Tensor Parallelism (TP) ✂️

📌 5.1 Khi Nào Cần Tensor Parallelism?

FSDP giải quyết được memory problem bạn có thể training model 70B trên 8 GPU thay vì 2 H100. Nhưng vẫn còn một vấn đề: single operation không fit vào GPU. Ví dụ: linear layer với shape [hidden_dim, 4 hidden_dim] = [8192, 32768] = 256M tham số chỉ trong 1 layer. Khi forward, activation tensor có kích thước [batch_size, seq_len, 32768] hàng GB chỉ cho một phép nhân mà trận.

Tensor Parallelism giải quyết bằng cách chia nhỏ các phép toán tensor ra nhiều GPU. Thay vì tính toán nguyên mà trận trên 1 GPU, mỗi GPU giữ 1 phần của weight mà trận, tính toán trên phần dữ liệu của nó, rồi đồng bộ kết quả.

⚔️ 5.2 Column-wise vs Row-wise Parallel

Có 2 cách chính để shard một linear layer:

• Column-wise (shard theo cột): Chia weight W thành [W1 W2] theo chiều cột. Mỗi GPU nhận input X với phần weight của nó output [XW1 XW2]. Cần all-reduce để gộp kết quả.
• Row-wise (shard theo hàng): Chia W thành [W1; W2] theo hàng. Mỗi GPU nhận input riêng X1/X2 với phần weight output cộng dồn qua all-reduce.

# Megatron-LM style Tensor Parallelism
class ColumnParallelLinear(nn.Module):
    def __init__(self, in_dim, out_dim, world_size, rank):
        super().__init__()
        # Mới GPU chi giu 1/world_size cot
        self.weight = nn.Parameter(
            torch.randn(in_dim, out_dim // world_size)
        )
        self.world_size = world_size

    def forward(self, x):
        local_out = torch.matmul(x, self.weight)
        dist.all_reduce(local_out) # Gop từ tat ca GPU
        return local_out

# Megatron-LM launch:
# --tensor-model-parallel-size 8
# --num-layers 40 --hidden-size 8192
# --ffn-hidden-size 32768
# --num-attention-heads 64
# --micro-batch-size 1

📌 5.3 Communication trong TP

TP yêu cầu rất nhiều communication mỗi layer có 1 all-reduce cho forward và 1 cho backward. Do đó, TP chỉ hiệu quả khi các GPU kết nối với nhau qua NVLink (băng thông 900 GB/s intra-node). Dùng TP cross-node (InfiniBand 400 Gb/s = 50 GB/s) sẽ bị bottleneck communication.

Thực tế: TP giới hạn trong 1 node (8 GPU với H100 DGX). Nếu cần >8 GPU cho TP, dùng TP=8 kết hợp với PP (pipeline) thay vì tăng TP.

Ví dụ cụ thể: Llamà 70B với TP=8 trên 1 node H100 (8x80GB). Mỗi GPU giữ 1/8 weight = 17.5GB weights + grads + optim. Tổng ~35GB/GPU còn dư cho activations và micro-batch size nhỏ.

6. Pipeline Parallelism (PP) 🔗

📌 6.1 Ý Tưởng Cốt Lõi

Pipeline Parallelism chia model theo chiều sâu mỗi GPU (hoặc nhóm GPU) chịu trách nhiệm cho một nhóm layers liên tiếp (stage). Ví dụ: model 40 layers trên 4 GPU GPU 0: layers 1-10, GPU 1: layers 11-20, GPU 2: layers 21-30, GPU 3: layers 31-40.

Mỗi micro-batch chạy lần lượt qua các stages GPU 0 forward xong, truyền activation sang GPU 1, GPU 1 forward, v.v. Nhưng nếu chỉ chạy tuần tự, GPU sẽ có rất nhiều idle time đây gọi là "pipeline bubble".

Thời gian biến: F=Forward micro-batch, B=Backward micro-batch. 1F1B giảm bubble bằng cách bắt đầu backward sớm, không chờ forward hết tất cả micro-batches.

🎯 6.3 Code Example: PP

# PiPPy (PyTorch native PP):
from torch.distributed.pipeline.sync import Pipe

# Chia model thành 4 stages
model = nn.Sequential(*layers)
pipe_model = Pipe(model, chunks=8) # 8 micro-batches

# DeepSpeed PP config
# {
#     "pipeline": {
#         "stages": 4,
#         "partition_method": "parameters",
#         "pipe_parallel_size": 4,
#         "gradient_accumulation_steps": 8,
#         "train_micro_batch_size_per_gpu": 2
#     }
# }

7. 3D Parallelism với DeepSpeed 🧬

📌 7.1 Ket Hop Ca Ba

3D Parallelism chính là cuộc chơi 1x1: kết hợp tất cả các hướng di chuyển (data, tensor, pipeline), người chiến thắng là dog model lớn nhất. Đây là cách Meta training Llamà 3 405B (chỉ 128 DP group, 32 node), NVIDIA training các model hàng 100B-500B, và các Open-Source labs training có quy mô tầm lại.

Cach ba chiệu parallelism ket hop trong thực tế:

• TP=8 (1 node) chia tung layer cho 8 GPU trong node qua NVLink 900 GB/s.
• PP=4 (4 nodes) chia model thanh 4 stages, mới stage gom 1 node TP=8. Tong 4x8=32 GPU.
• DP=4 (16 nodes) nhan ban ca pipeline len 4 ban sao, mới ban 32 GPU. Tong 4x32=128 GPU tren 16 nodes.

🏭 7.2 Real-World Config: DeepSpeed + Megatron

# ds_config.json DeepSpeed 3D Parallelism
{
    "train_batch_size": 256,
    "train_micro_batch_size_per_gpu": 4,
    "gradient_accumulation_steps": 8,

    "zero_optimization": {
        "stage": 3,
        "allgather_partitions": true,
        "allgather_bucket_size": 5e8,
        "reduce_scatter": true,
        "reduce_bucket_size": 5e8,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "stage3_prefetch_bucket_size": 5e8,
        "stage3_param_persistence_threshold": 1e6,
        "stage3_max_live_parameters": 5e8,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "fp16": { "enabled": true, "auto_cast": true },
    "bf16": { "enabled": true },
    "tensor_parallel": { "enabled": true, "tp_size": 8 },
    "pipeline": {
        "enabled": true, "stages": 4,
        "partition_method": "type:[TransformerLayer]",
        "activation_checkpointing": {
            "enabled": true, "partition_activations": true
        }
    }
}

# Train script
import deepspeed
model = MyLLM()
model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args, model=model,
    model_parameters=model.parameters(),
    config_params=ds_config
)
for batch in dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

# deepspeed train.py --deepspeed_config ds_config.json

📌 7.3 Case Study: Meta Llamà 3 405B

Meta trained Llamà 3 405B tren 16,384 GPU H100 (2048 nodes DGX H100, 8 GPU/node) su dung 3D Parallelism:

• TP=8 trong 1 node (NVLink 900 GB/s)
• PP=8 qua 8 nodes (InfiniBand 400 Gb/s)
• DP=128 256 nodes / 8 PP = 32 DP groups x 2 = 64 (thực tế phuc tap hon)
• Dung Megatron-LM + DeepSpeed ZeRO-3
• Thoi gian training: ~54 ngay cho 405B tham so

8. Use Cases Thực Te 🏢

📌 8.1 Meta Llamà 3 Series

• Llamà 3 8B: DDP tren 64 GPU H100 (8 nodes x 8 GPU). Scale efficiency ~85%.
• Llamà 3 70B: FSDP ZeRO-3 + TP=8 tren 128 GPU (16 nodes). Memory: ~25GB/GPU.
• Llamà 3 405B: 3D Parallelism TP=8 x PP=8 x DP=128 tren 16K GPU. Day là một trong nhung training run lon nhat cong khai.

📌 8.2 NVIDIA Megatron-LM

Megatron-LM là framework tiên phong về Tensor + Pipeline Parallelism. NVIDIA dùng nó để train các model như Megatron-Turing NLG 530B và Nemotron. Megatron-LM hiện nay là một phần của NeMo Megatron framework và được tích hợp với DeepSpeed qua deepspeed.megatron.

⚔️ 8.3 Startup vs Enterprise Decision Guide

9. Pitfalls & Troubleshooting 🚨

📌 9.1 Debug Checklist

• Kiem tra NCCL connectivity: python -m torch.distributed.run --nproc_per_node=8 --nnodes=2 test_nccl.py
• Kiem tra memory leak: torch.cuda.memory_summary() sau mới 100 steps
• Kiem tra communication overlap: Profiler torch.profiler với activities=ProfilerActivity.CUDA
• Kiem tra FSDP wrapping: fsdp_model.summarize() hoac print(fsdp_model) de xem policy co dung khong
• Tao cluster: torchrun --nproc_per_node=8 --nnodes=2 --rdzv_endpoint=main_node:29500 train.py

💻 Code Example: DeepSpeed Training

import deepspeed
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

ds_config = {
    "train_batch_size": 16,
    "fp16": {"enabled": True},
    "zero_optimization": {"stage": 3}
}

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model, config=ds_config
)

for batch in dataloader:
    loss = model_engine(batch, labels=batch).loss
    model_engine.backward(loss)
    model_engine.step()

📌 10. Kết Luận

Bai viet nay da di qua 4 chien luoc distributed training từ co ban đến năng cao. Diem chinh can nho:

• DDP: Don gian, scaling tot nho hon 7B. Mới GPU giu full model.
• FSDP: Shard model parameters, tiet kiem memory 8x so với DDP. Can cho 7B-70B.
• TP: Chia tung layer, can NVLink toc do cao. Dung trong 1 node.
• PP: Chia model theo chiệu sau layers, phu hop multi-node. Can micro-batches lon de giam bubble.
• 3D Parallelism: Ket hop DP+TP+PP cho model 70B+. Su dung boi Meta, NVIDIA, va các lab lon.

Khong co "silver bullet" cho distributed training. Mới model size, cluster size, va budget co config parallel khac nhau. Nguyen tac chung: toi uu communication bottleneck truoc, memory bottleneck sau.