4.2 PyTorch data parallelism từ nguyên lý đến thực chiến

Data parallelism là chiến lược song song cơ bản và dùng nhiều nhất: mỗi GPU giữ một bản sao mô hình đầy đủ, xử lý shard dữ liệu khác nhau, rồi đồng bộ gradient. Logic tiến hóa và nguyên lý bộ nhớ/giao tiếp nằm ở 4.1 Data parallelism chi tiết. Bài này nối tiếp nguyên lý, tập trung thực chiến kỹ thuật: nền giao tiếp PyTorch, code DDP/FSDP chạy được đầy đủ, khởi động đa máy đa GPU và các lỗi thường gặp, để đưa data parallelism vào script train thật.

📑 Mục lục

• 1. Nền tảng giao tiếp phân tán trong PyTorch
• 2. DDP: phương án chuẩn của data parallelism
• 3. FSDP: tối ưu bộ nhớ bằng parameter sharding
• 4. DDP vs FSDP: cách chọn
• 5. torchrun và thực chiến train đa máy
• 6. Hướng dẫn các lỗi thường gặp khi train phân tán
• Tóm tắt
• Checklist tự kiểm
• Tài liệu tham khảo

---

1. Nền tảng giao tiếp phân tán trong PyTorch

Trước khi hiểu DDP và FSDP, cần nắm lớp giao tiếp mà chúng dựa vào.

1.1 Khái niệm cốt lõi

1.2 Chọn backend giao tiếp

PyTorch hỗ trợ ba backend:

📌 Điểm then chốt: train phân tán trên GPU gần như luôn dùng backend NCCL. NCCL (NVIDIA Collective Communications Library) được tối ưu sâu cho NVLink, NVSwitch, InfiniBand và các kết nối NVIDIA khác.

1.3 Collective communication primitives

Đồng bộ dữ liệu trong distributed training dựa trên collective communication. Các primitive dùng nhiều nhất:

💡 Gợi ý: hiểu AllReduce = ReduceScatter + AllGather là chìa khóa để thấy khác biệt mẫu giao tiếp giữa DDP và FSDP.

1.4 Khởi tạo giao tiếp

Mọi job phân tán đều bắt đầu bằng khởi tạo process group:

import torch.distributed as dist

# Cách 1: init từ biến môi trường (hay dùng nhất, kèm torchrun)
dist.init_process_group(backend="nccl")

# Cách 2: truyền tham số tường minh
dist.init_process_group(
    backend="nccl",
    init_method="tcp://192.168.1.1:29500",
    world_size=8,
    rank=0,
)

Khi dùng torchrun, các biến MASTER_ADDR, MASTER_PORT, RANK, WORLD_SIZE, LOCAL_RANK, v.v. được set sẵn, nên trong code chỉ cần gọi dist.init_process_group(backend="nccl").

---

2. DDP: phương án chuẩn của data parallelism

DDP (DistributedDataParallel) là đường data parallel chín muồi và dùng rộng nhất trong PyTorch.

2.1 Nguyên lý làm việc của DDP

Ba câu tóm cơ chế cốt lõi:

1. Init: Broadcast tham số của rank 0 sang mọi process để điểm xuất phát giống nhau
2. Forward + backward: mỗi process độc lập tính gradient trên shard dữ liệu của mình
3. Đồng bộ gradient: trong quá trình backward, DDP dùng AllReduce lấy trung bình gradient mọi process

⚠️ Lưu ý: đồng bộ gradient của DDP không đợi cả backward xong mới bắt đầu. Nó dùng cơ chế Bucket: đóng gói gradient theo thứ tự tính backward thành các Bucket; mỗi Bucket sẵn thì AllReduce ngay, để tính toán overlap với giao tiếp. Đó là tối ưu hiệu năng then chốt của DDP.

2.2 Code train DDP đầy đủ

Ví dụ DDP chạy được ngay:

import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms


def setup():
    """Khởi tạo môi trường phân tán"""
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank


def cleanup():
    """Dọn môi trường phân tán"""
    dist.destroy_process_group()


def main():
    local_rank = setup()
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # 1. Dựng mô hình và bọc bằng DDP
    model = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10),
    ).to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    # 2. Dựng data loader (bắt buộc dùng DistributedSampler)
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.view(-1)),
    ])
    dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=2)

    # 3. Optimizer và hàm loss
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    # 4. Vòng lặp train
    model.train()
    for epoch in range(5):
        sampler.set_epoch(epoch)  # Quan trọng: mỗi epoch seed xáo khác nhau
        total_loss = 0.0
        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(local_rank), target.to(local_rank)

            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()  # DDP tự AllReduce tại đây
            optimizer.step()

            total_loss += loss.item()

        if rank == 0:
            avg_loss = total_loss / len(dataloader)
            print(f"Epoch {epoch}: avg_loss = {avg_loss:.4f}")

    cleanup()


if __name__ == "__main__":
    main()

Lệnh khởi động:

# Một máy 4 GPU
torchrun --nproc_per_node=4 train_ddp.py

2.3 Các điểm cần nhớ khi dùng DDP

Vài chi tiết dễ bỏ sót nhưng rất quan trọng:

1. Bắt buộc dùng DistributedSampler

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)

Không dùng DistributedSampler thì mỗi process đọc cùng dữ liệu, data parallelism mất ý nghĩa.

2. Mỗi epoch phải gọi sampler.set_epoch(epoch)

for epoch in range(num_epochs):
    sampler.set_epoch(epoch)  # Xáo cách chia dữ liệu giữa các epoch
    for batch in dataloader:
        ...

Không gọi set_epoch sẽ khiến mọi epoch các process nhận cùng một cách chia dữ liệu.

3. Chỉ log và lưu trên rank 0

if dist.get_rank() == 0:
    print(f"Loss: {loss.item()}")
    torch.save(model.module.state_dict(), "checkpoint.pt")  # Lưu ý: model.module

⚠️ Lưu ý: lưu mô hình phải dùng model.module.state_dict() chứ không phải model.state_dict(). DDP thêm tiền tố module. vào tên tham số; lưu trực tiếp sẽ lệch key lúc load.

2.4 Giới hạn của DDP

DDP đơn giản và hiệu quả, nhưng có một giới hạn căn bản: mỗi GPU giữ bản sao mô hình đầy đủ. Nghĩa là:

• Tham số + gradient + optimizer state phải nhét vừa bộ nhớ một GPU
• Với mô hình 7B trở lên, một GPU 80 GB thường không đủ

Đó là vấn đề FSDP giải quyết.

---

3. FSDP: tối ưu bộ nhớ bằng parameter sharding

FSDP (FullyShardedDataParallel) là hiện thực hóa ý tưởng DeepSpeed ZeRO ngay trong PyTorch: không để mỗi card lưu cả mô hình đầy đủ, mà chia shard tham số, gradient và optimizer state.

3.1 Hiểu FSDP qua ZeRO

Nhận xét cốt lõi của ZeRO (Zero Redundancy Optimizer): trong DDP mỗi GPU dư thừa lưu đầy đủ optimizer state, gradient và tham số; đó là lãng phí lớn.

ZeRO gỡ dư thừa dần qua ba giai đoạn:

Ví dụ mô hình 7B, 8 GPU (ước lượng 16 byte/tham số theo paper ZeRO: tham số FP16 2B + gradient FP16 2B + optimizer state FP32 12B):

• DDP: mỗi GPU ~112 GB (không nhét nổi H100 80 GB)
• FSDP (FULL_SHARD): mỗi GPU ~$\frac{112}{8}$ = ~14 GB (nhét dễ)

3.2 Luồng làm việc của FSDP

Luồng FSDP phức tạp hơn DDP: “ghép tham số khi cần” và “dùng xong bỏ”.

1. Trạng thái ban đầu: mỗi GPU chỉ lưu shard tham số mình phụ trách
2. Forward: khi tính một tầng, AllGather thu mọi shard, ghép tham số đầy đủ → tính → giải phóng shard không cục bộ
3. Backward: AllGather lại tham số → tính gradient → ReduceScatter chia shard gradient về từng GPU → giải phóng tham số đầy đủ
4. Cập nhật: mỗi GPU chỉ cập nhật tham số và optimizer state của shard mình

📌 Điểm then chốt: FSDP đổi giao tiếp lấy bộ nhớ. Forward và backward mỗi tầng cần một AllGather + một ReduceScatter, lưu lượng giao tiếp lớn hơn DDP, nhưng bộ nhớ giảm mạnh.

3.3 Chiến lược sharding của FSDP

PyTorch FSDP cung cấp bốn chiến lược, khá linh hoạt:

💡 Gợi ý: HYBRID_SHARD là chiến lược khuyến nghị khi train đa máy. Trong máy đi FSDP (NVLink băng thông cao), giữa máy đi data parallel (InfiniBand chậm hơn): cân bằng tốt giữa hiệu quả giao tiếp và tiết kiệm bộ nhớ.

3.4 Code train FSDP đầy đủ

import os
import functools
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms


def setup():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank


def cleanup():
    dist.destroy_process_group()


class SimpleMLP(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=2048, num_layers=8, output_dim=10):
        super().__init__()
        layers = []
        for i in range(num_layers):
            in_d = input_dim if i == 0 else hidden_dim
            out_d = output_dim if i == num_layers - 1 else hidden_dim
            layers.append(nn.Linear(in_d, out_d))
            if i < num_layers - 1:
                layers.append(nn.ReLU())
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)


def main():
    local_rank = setup()
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # 1. Tạo mô hình
    model = SimpleMLP().to(local_rank)

    # 2. Cấu hình mixed precision
    bf16_policy = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    )

    # 3. Chính sách auto-wrap (submodule vượt ngưỡng tham số được shard riêng)
    wrap_policy = functools.partial(
        size_based_auto_wrap_policy,
        min_num_params=1_000_000,
    )

    # 4. Bọc mô hình bằng FSDP
    model = FSDP(
        model,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        mixed_precision=bf16_policy,
        auto_wrap_policy=wrap_policy,
        device_id=local_rank,
    )

    # 5. Nạp dữ liệu
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.view(-1)),
    ])
    dataset = datasets.MNIST("./data", train=True, download=(rank == 0), transform=transform)
    dist.barrier()  # Đợi rank 0 tải xong
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=2)

    # 6. Optimizer
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    # 7. Vòng lặp train
    model.train()
    for epoch in range(5):
        sampler.set_epoch(epoch)
        total_loss = 0.0
        for data, target in dataloader:
            data, target = data.to(local_rank), target.to(local_rank)

            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        if rank == 0:
            print(f"Epoch {epoch}: avg_loss = {total_loss / len(dataloader):.4f}")

    cleanup()


if __name__ == "__main__":
    main()

3.5 Lưu và load mô hình FSDP

Lưu mô hình FSDP phức tạp hơn DDP vì tham số đang ở dạng shard. Nên dùng Distributed Checkpoint (DCP) của PyTorch, kèm helper trong torch.distributed.checkpoint.state_dict để lấy/ghi lại state shard đúng cách:

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import (
    get_state_dict,
    set_state_dict,
)

# Lưu (mọi rank tham gia): helper trích đúng state model / optimizer đã shard của FSDP
model_sd, optim_sd = get_state_dict(model, optimizer)
dcp.save(
    {"model": model_sd, "optim": optim_sd},
    checkpoint_id="./checkpoints/epoch_5",
)

# Load (mọi rank tham gia): lấy placeholder shard hiện tại, load đổ tại chỗ, rồi ghi lại
model_sd, optim_sd = get_state_dict(model, optimizer)
dcp.load(
    {"model": model_sd, "optim": optim_sd},
    checkpoint_id="./checkpoints/epoch_5",
)
set_state_dict(model, optimizer, model_state_dict=model_sd, optim_state_dict=optim_sd)

⚠️ Lưu ý: đừng lưu mô hình FSDP bằng torch.save(model.state_dict()) một mình. Hành vi state_dict của FSDP phụ thuộc state_dict_type; lưu trực tiếp có thể chỉ ghi shard cục bộ. DCP + get_state_dict / set_state_dict mới xử lý đúng lưu/khôi phục shard, và hỗ trợ resharding (lưu 8 GPU, khôi phục 4 GPU vẫn được).

3.6 FSDP2: API thế hệ sau

PyTorch đang đẩy API FSDP mới: torch.distributed.fsdp.fully_shard (thường gọi FSDP2). Khác với kiểu module-wrapper của FSDP1, FSDP2 dùng per-parameter sharding có thể kết hợp:

from torch.distributed.fsdp import fully_shard

# Phong cách FSDP2: gọi fully_shard trực tiếp trên module
for layer in model.layers:
    fully_shard(layer)
fully_shard(model)

Ưu thế của FSDP2:

• Kết hợp tốt hơn với tensor parallelism (TP) và các chiến lược song song khác
• Dựa trên trừu tượng DTensor (Distributed Tensor); API thống nhất hơn
• Kiểm soát độ hạt shard linh hoạt hơn

💡 Gợi ý: tính đến PyTorch 2.x, FSDP2 đã dùng được nhưng vẫn iterate nhanh. Dự án mới có thể thử FSDP2; môi trường production FSDP1 vẫn thường ổn định hơn.

---

4. DDP vs FSDP: cách chọn

Gợi ý chọn:

• ✅ Tham số + optimizer state nhét vừa một GPU → dùng DDP
• ✅ Một GPU không nhét nổi mô hình đầy đủ → dùng FSDP
• ✅ Train đa máy + mô hình lớn → dùng FSDP + HYBRID_SHARD
• ✅ Mô hình rất lớn (hàng chục tỷ trở lên) → cân nhắc FSDP + TP + PP (Megatron-LM hoặc DeepSpeed)

---

5. torchrun và thực chiến train đa máy

5.1 torchrun là gì

torchrun là công cụ khởi động distributed training được PyTorch khuyến nghị (thay torch.distributed.launch đã deprecated). Khả năng cốt lõi:

• Tự set biến môi trường: RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT, v.v.
• Hỗ trợ elastic training: node có thể join/leave động
• Cơ chế chịu lỗi: worker chết có thể tự restart

5.2 Khởi động một máy nhiều GPU

# Một máy 4 GPU
torchrun --nproc_per_node=4 train.py --batch_size 64

# Chỉ định master port (tránh xung đột cổng)
torchrun --nproc_per_node=4 --master_port=29501 train.py

5.3 Khởi động đa máy đa GPU

Trên mỗi máy lần lượt chạy torchrun, cùng thông tin rendezvous (hội tụ):

# Máy 0 (master)
torchrun \
    --nnodes=2 \
    --nproc_per_node=4 \
    --node_rank=0 \
    --rdzv_backend=c10d \
    --rdzv_endpoint=192.168.1.100:29500 \
    train.py

# Máy 1
torchrun \
    --nnodes=2 \
    --nproc_per_node=4 \
    --node_rank=1 \
    --rdzv_backend=c10d \
    --rdzv_endpoint=192.168.1.100:29500 \
    train.py

5.4 Elastic training

torchrun hỗ trợ elastic training: số node có thể đổi trong một khoảng:

# Ít nhất 1 máy, nhiều nhất 4 máy
torchrun \
    --nnodes=1:4 \
    --nproc_per_node=4 \
    --rdzv_id=my_training_job \
    --rdzv_backend=c10d \
    --rdzv_endpoint=192.168.1.100:29500 \
    train.py

⚠️ Lưu ý: elastic training đòi hỏi code train lưu và khôi phục checkpoint đúng. Khi node join hoặc leave, torchrun restart mọi Worker; script phải resume từ checkpoint gần nhất.

5.5 Lấy thông tin phân tán trong code train

Khi khởi động bằng torchrun, đọc thông tin phân tán từ biến môi trường:

import os

local_rank = int(os.environ["LOCAL_RANK"])     # Số GPU trong máy
rank = int(os.environ["RANK"])                  # Id process toàn cục
world_size = int(os.environ["WORLD_SIZE"])      # Tổng số process
local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])  # Số process trên máy này

---

6. Hướng dẫn các lỗi thường gặp khi train phân tán

6.1 Deadlock (hang)

Lỗi phổ biến nhất là “treo không chạy”. Nguyên nhân thường gặp:

6.2 Biến môi trường debug hay dùng

# Bật log debug NCCL
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL

# Tăng timeout collective (backend NCCL mặc định 10 phút); set trong code:
# dist.init_process_group(..., timeout=datetime.timedelta(seconds=3600))

# Timeout heartbeat của watchdog: thread giám sát phát hiện collective treo bao lâu
# thì buộc dừng process (PyTorch 2.x)
# Lưu ý: nó quản "phát hiện treo / abort", không phải kéo dài timeout một lần giao tiếp
export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=600

# Log phân tán của PyTorch
export TORCH_DISTRIBUTED_DEBUG=DETAIL

# Chỉ định GPU nhìn thấy (tránh lệch số GPU)
export CUDA_VISIBLE_DEVICES=0,1,2,3

6.3 Tính nhất quán của random seed

Quản lý tính ngẫu nhiên trong distributed training cần chú ý:

import torch
import random
import numpy as np

def set_seed(seed, rank):
    # Init mô hình cần seed giống nhau (để tham số các rank DDP khớp)
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

    # Augmentation v.v. cần seed khác nhau (tránh dữ liệu các rank giống hệt)
    torch.manual_seed(seed + rank)

📌 Điểm then chốt: DDP đòi hỏi tham số khởi tạo các rank giống nhau (lúc dựng DDP sẽ Broadcast từ rank 0), nhưng data augmentation, Dropout, v.v. nên khác nhau giữa các rank; nếu không “data parallel” mất tính đa dạng.

6.4 Phóng learning rate

Effective batch size của data parallelism = batch size một GPU × world_size. Khi world_size tăng, thường phóng learning rate cùng tỉ lệ:

$$lr_{\text{effective}} = lr_{\text{base}} \times \text{world\_size}$$

Nhưng phóng tuyến tính khi world_size rất lớn có thể làm train mất ổn định. Thực tế thường kết hợp learning-rate warmup:

# Warmup tuyến tính + phóng LR tuyến tính
base_lr = 1e-4
effective_lr = base_lr * world_size
warmup_steps = 1000

for step in range(total_steps):
    if step < warmup_steps:
        lr = effective_lr * (step / warmup_steps)
    else:
        lr = effective_lr  # hoặc thêm cosine decay sau
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

💡 Gợi ý: phóng learning rate tuyến tính là quy tắc kinh nghiệm, không phải luật tuyệt đối. Ở quy mô lớn (hàng trăm GPU trở lên), có thể cần phóng theo $\sqrt{\text{world\_size}}$ (ý tưởng của LAMB/LARS) để giữ ổn định.

---

📝 Tóm tắt

Bài này nối nguyên lý 4.1, nói thấu thực chiến data parallelism trong PyTorch:

• Nền giao tiếp: World Size / Rank / Local Rank, backend NCCL, collective AllReduce / AllGather / ReduceScatter
• DDP: phương án data parallel chín muồi nhất, phù hợp mô hình một GPU nhét được, kèm code chạy được. Cốt lõi là AllReduce gradient theo Bucket, kèm DistributedSampler + set_epoch và lưu mô hình đúng cách
• FSDP: tối ưu bộ nhớ bằng parameter sharding, phù hợp mô hình một GPU không nhét nổi, kèm code đầy đủ và bốn chiến lược sharding. Cốt lõi là AllGather + ReduceScatter ghép khi cần; lưu mô hình nên dùng DCP
• torchrun: công cụ khởi động khuyến nghị, hỗ trợ một máy, đa máy, elastic training
• Vấn đề thường gặp: soi deadlock, quản lý random seed, phóng learning rate

Nắm các kỹ năng thực chiến này rồi, bước tiếp theo có thể đi sâu 3D parallelism (TP + PP + DP), DeepSpeed ZeRO, Megatron-LM và các kỹ thuật phân tán cao hơn.

---

🎯 Checklist tự kiểm

• Giải thích được khác biệt cốt lõi DDP vs FSDP (bản sao đầy đủ vs parameter sharding)
• Vẽ được luồng dữ liệu AllReduce đồng bộ gradient trong DDP
• Nói được ba giai đoạn ZeRO mỗi giai đoạn shard gì, tương ứng chiến lược FSDP nào
• Dùng PyTorch DDP đổi một script train một GPU thành nhiều GPU
• Cấu hình đúng sharding strategy, mixed precision và auto-wrap của FSDP
• Dùng torchrun khởi động một máy nhiều GPU và đa máy nhiều GPU
• Giải thích vai trò DistributedSampler và vì sao cần set_epoch
• Liệt kê được 3 nguyên nhân deadlock phân tán thường gặp và cách soi
• Tính được effective batch size dưới data parallelism và phóng learning rate đúng
• Dùng Distributed Checkpoint (DCP) lưu và khôi phục mô hình FSDP

---

📚 Tài liệu tham khảo

• PyTorch Distributed Communication Package — torch.distributed
• PyTorch DistributedDataParallel — Getting Started
• PyTorch FSDP — FullyShardedDataParallel
• PyTorch torchrun (Elastic Launch)
• PyTorch Distributed Checkpoint
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel
• NCCL Documentation