LLM Systems 15: Distributed Data Parallel Training

0. Recap

• Overall idea: partition the data, distribute the forward/backward
• Parameter Server
  • server to update and distribute parameters, worker to get local grad
• NCCL Multi-GPU communication
  • using ring and batching to reduce the latency for Broadcast
• Data Parallel Training via All Reduce
  • Efficient Ring AllReduce (ScatterReduce+AllGather

1. Distributed Data Parallel Training

Distributed Data Parallel

• Same as Data Parallel
• multiple nodes, each with multiple GPUs
  • Create replicas of a model on multiple nodes
  • Each model performs the forward pass and the backward pass independently
  • Gather average gradients across nodes
  • Optimizers run locally (identical)

2. Design and implementation of Distributed Data Parallel

Design Goal of DDP

• Non-intrusive: Developers should be able to reuse the local training script with minimal modifications.
• Interceptive: The API needs to allow the implementation to intercept various signals and trigger appropriate algorithms promptly. The API must expose as many optimization opportunities as possible to the internal implementation.

Setting up the Distributed Process

• World size
  • total number of processes W
• Global rank
  • global process id
• Local rank
  • local process id

示例图：

Launch Distributed Processes

The launch.py (torch/distributed/launch.py) will pass world size, global rank, master address, master port via env vars, and local rank as a commandline parameter to every instance.

• 关键环境变量：
  • MASTER_ADDR：主节点（Rank 0 所在机器）的 IP 地址，用于各个进程建立网络连接和握手（Rendezvous）。
  • MASTER_PORT：主节点监听的网络端口号。
  • RANK (Global Rank)：当前进程在所有节点中的全局进程 ID（从 0 到 WORLD_SIZE - 1）。
  • WORLD_SIZE：本次分布式训练的总进程数。

Implementation

• Gradient Bucketing
  • 具体介绍略，见 paper。

实现代码：

// The function `autograd_hook` is called after the gradient for a
// model parameter has been accumulated into its gradient tensor.
// This function is only to be called from the autograd thread.
void Reducer::autograd_hook(size_t index) {
	mark_variable_ready(index);
}

void Reducer::mark_variable_ready(size_t variable_index) {
	const auto& bucket_index = variable_locators_[variable_index];
	auto& bucket = buckets_[bucket_index.bucket_index];
	if (--bucket.pending == 0) {
		mark_bucket_ready(bucket_index.bucket_index);
	}
}

void Reducer::mark_bucket_ready(size_t bucket_index) {
	for (; next_bucket_ < buckets_.size() && buckets_[next_bucket_].pending == 0; next_bucket_++) {
		num_buckets_ready_++;
		auto& bucket = buckets_[next_bucket_];
		all_reduce_bucket(bucket);
	}
}

void Reducer::all_reduce_bucket(Bucket& bucket) {
	auto variables_for_bucket = get_variables_for_bucket(next_bucket_, bucket);
	const auto& tensor = bucket.gradients;
	GradBucket grad_bucket(next_bucket_, buckets_.size(), tensor, bucket.offsets,
							bucket.lengths, bucket.sizes_vec, variables_for_bucket);
	bucket.future_work = run_comm_hook(grad_bucket);
}

3. Code walkthrough: Using DDP in PyTorch

代码链接：link

核心代码即注释：

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.data.distributed import DistributedSampler

# =====================================================================
# 步骤 1: DDP 环境初始化与进程组建立
# =====================================================================
def setup_ddp():
    # 使用 torchrun 启动时，会自动注入 LOCAL_RANK, RANK, WORLD_SIZE 等环境变量
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    # 强制将当前进程绑定到节点内的物理 GPU 上
    torch.cuda.set_device(local_rank)
    return local_rank

def cleanup_ddp():
    # 训练结束后销毁通信组，释放底层资源与 CUDA Stream
    dist.destroy_process_group()

# =====================================================================
# 步骤 2: 核心训练流程
# =====================================================================
def main():
    local_rank = setup_ddp()
    rank = dist.get_rank()            # 当前进程的全局 Rank ID (0 ~ world_size-1)
    world_size = dist.get_world_size() # 总进程数 / GPU 总数

    # --- A. 数据集与 DistributedSampler 绑定 ---
    dataset = TensorDataset(torch.randn(1000, 10), torch.randint(0, 2, (1000,)))
    # 使用 DistributedSampler 将全局数据均匀切分给各个 GPU，避免数据重复
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    train_loader = DataLoader(dataset, batch_size=32, sampler=sampler)

    # --- B. 模型初始化与 DDP 包装 ---
    # 1. 实例化本地模型并移至对应的 GPU
    model = nn.Sequential(nn.Linear(10, 50), nn.ReLU(), nn.Linear(50, 2)).to(local_rank)
    # 2. 用 DDP 包装模型（自动注册 Autograd Hook，实现梯度分桶与异步 AllReduce）
    model = DDP(model, device_ids=[local_rank])

    criterion = nn.CrossEntropyLoss()
    # 学习率可根据卡数等比例缩放（常用经验法则：lr = lr_base * world_size）
    optimizer = optim.SGD(model.parameters(), lr=0.01 * world_size)

    # --- C. 训练循环 ---
    for epoch in range(1, 4):
        # 【关键点】每轮 Epoch 必须调用 set_epoch，保证各卡上 shuffle 的随机种子同步更新
        sampler.set_epoch(epoch)
        model.train()
        epoch_loss = 0.0

        for x, y in train_loader:
            x, y = x.to(local_rank), y.to(local_rank)

            optimizer.zero_grad()
            outputs = model(x)
            loss = criterion(outputs, y)

            # 执行 backward 时，DDP 会按桶自动触发异步 AllReduce 跨卡同步梯度
            loss.backward()
            optimizer.step()

            epoch_loss += loss.item()

        # --- D. 指标跨卡汇总与 Rank 0 状态保存 ---
        # 使用 dist.all_reduce 归约所有卡上的 Loss，算得真实的全局 Loss
        loss_tensor = torch.tensor(epoch_loss, device=local_rank)
        dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM)
        global_avg_loss = (loss_tensor / world_size / len(train_loader)).item()

        # 仅在 Rank 0 主进程中执行日志打印与模型保存，避免多卡重复写盘
        if rank == 0:
            print(f"| Epoch {epoch:02d} | Global Avg Loss: {global_avg_loss:.4f} |")
            # 保存时使用 model.module 剥离 DDP 包装，只留原始权重的 state_dict
            torch.save(model.module.state_dict(), "model_checkpoint.pt")

    cleanup_ddp()

if __name__ == "__main__":
    main()

• 启动命令：使用 torchrun 启动（例如单机 2 卡训练）：

torchrun --nproc_per_node=2 main.py

• 核心组件要点：
  1. setup_ddp：依赖 torchrun 注入的环境变量 LOCAL_RANK 完成设备绑定。
  2. DistributedSampler：按卡划分数据集，每个 Epoch 必须执行 sampler.set_epoch(epoch) 才能确保随机洗牌正常。
  3. DDP(model, device_ids=[local_rank])：接管模型梯度通信，在反向传播计算时流水线化发起分桶 AllReduce。
  4. if rank == 0:：控制打印与保存，防止多卡同时写日志导致死锁或覆盖。
  5. model.module：保存模型权重时加上 .module，便于后续在单卡/CPU 上载入标准 .state_dict()。

4. Summary

• Data Parallel via All Reduce
• Distributed Data Parallel Training
  • gradient bucketing
  • overlay backward and AllReduce communication