第7章 NCCL与多卡通信¶

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(rank, world_size):
    """初始化分布式环境"""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()  # super()调用父类方法
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 256),
            nn.ReLU(),
            nn.Linear(256, 10),
        )

    def forward(self, x):
        return self.classifier(self.features(x))

def train(rank, world_size, epochs=5):
    """DDP训练函数"""
    setup(rank, world_size)

    # 创建模型并移到GPU
    model = SimpleModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    # 数据加载
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)

    # 分布式采样器：确保每个GPU看到不同的数据
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=2)

    # 优化器
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        sampler.set_epoch(epoch)  # 确保每个epoch数据打乱不同
        ddp_model.train()

        total_loss = 0
        for batch_idx, (data, target) in enumerate(dataloader):  # enumerate同时获取索引和值
            data, target = data.to(rank), target.to(rank)

            optimizer.zero_grad()
            output = ddp_model(data)
            loss = criterion(output, target)
            loss.backward()   # 反向传播时DDP自动AllReduce梯度
            optimizer.step()

            total_loss += loss.item()

        # 只在rank 0打印
        if rank == 0:
            avg_loss = total_loss / len(dataloader)
            print(f"Epoch {epoch+1}, Avg Loss: {avg_loss:.4f}")

    # 保存模型（只在rank 0保存）
    if rank == 0:
        torch.save(ddp_model.module.state_dict(), "model.pt")

    cleanup()

# 启动方式
if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    torch.multiprocessing.spawn(train, args=(world_size,), nprocs=world_size)

import torch
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy

def train_fsdp(rank, world_size):
    setup(rank, world_size)

    # 创建大模型
    model = LargeTransformerModel().to(rank)

    # 混合精度策略
    mixed_precision = MixedPrecision(
        param_dtype=torch.bfloat16,     # 参数用bf16
        reduce_dtype=torch.float32,      # 梯度规约用fp32
        buffer_dtype=torch.bfloat16,     # buffer用bf16
    )

    # 自动分片策略: 参数量>1M的层独立分片
    auto_wrap_policy = size_based_auto_wrap_policy(min_num_params=1_000_000)

    # 创建FSDP模型
    fsdp_model = FSDP(
        model,
        sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO-3
        mixed_precision=mixed_precision,
        auto_wrap_policy=auto_wrap_policy,
        device_id=rank,
    )

    optimizer = torch.optim.AdamW(fsdp_model.parameters(), lr=1e-4)

    for epoch in range(num_epochs):
        for batch in dataloader:
            inputs = batch["input_ids"].to(rank)
            labels = batch["labels"].to(rank)

            outputs = fsdp_model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()

            # 梯度裁剪
            fsdp_model.clip_grad_norm_(1.0)

            optimizer.step()
            optimizer.zero_grad()

    # 保存FSDP模型
    from torch.distributed.fsdp import FullStateDictConfig, StateDictType

    save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
    with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, save_policy):
        state_dict = fsdp_model.state_dict()
        if rank == 0:
            torch.save(state_dict, "model_fsdp.pt")

    cleanup()

# deepspeed_config.json (ZeRO Stage 2)
"""
{
    "train_batch_size": 32,
    "gradient_accumulation_steps": 4,
    "gradient_clipping": 1.0,

    "fp16": {
        "enabled": true,
        "initial_scale_power": 16
    },

    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "overlap_comm": true,
        "contiguous_gradients": true
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 1e-4,
            "betas": [0.9, 0.999],
            "weight_decay": 0.01
        }
    }
}
"""

# 训练代码
import deepspeed

def train_deepspeed():
    model = MyModel()

    # DeepSpeed初始化
    model_engine, optimizer, _, scheduler = deepspeed.initialize(
        model=model,
        config="deepspeed_config.json",
    )

    for epoch in range(num_epochs):
        for batch in dataloader:
            inputs = batch["input_ids"].to(model_engine.device)
            labels = batch["labels"].to(model_engine.device)

            outputs = model_engine(inputs)
            loss = criterion(outputs, labels)

            model_engine.backward(loss)
            model_engine.step()

def gradient_compression(gradients, compression_ratio=0.01):
    """Top-K梯度压缩"""
    flat = gradients.flatten()
    k = max(1, int(flat.numel() * compression_ratio))

    # 选择绝对值最大的k个梯度
    values, indices = torch.topk(flat.abs(), k)

    # 只传输Top-K梯度
    compressed_values = flat[indices]

    return compressed_values, indices, flat.shape

def gradient_decompress(values, indices, shape):
    """解压梯度"""
    flat = torch.zeros(shape.numel(), device=values.device)
    flat[indices] = values
    return flat.reshape(shape)

# DDP中的通信计算重叠示意
"""
正常执行:
  反向计算Layer N    → AllReduce Layer N    → 反向计算Layer N-1 → ...
  |---- compute ----|---- communicate ----|---- compute ----|

重叠执行:
  反向计算Layer N    → 反向计算Layer N-1    → 反向计算Layer N-2
  |---- compute ----|---- compute ------  |
                    |---- AllReduce N ----||---- AllReduce N-1 --|

  总时间 ≈ max(compute, communicate) 而非 compute + communicate
"""

import torch.distributed as dist
import time

def benchmark_nccl(rank, world_size, data_sizes_mb=[1, 10, 100, 1000]):
    """NCCL通信性能测试"""
    setup(rank, world_size)

    results = []
    for size_mb in data_sizes_mb:
        numel = size_mb * 1024 * 1024 // 4  # float32
        tensor = torch.randn(numel, device=f"cuda:{rank}")

        # 预热
        for _ in range(5):
            dist.all_reduce(tensor)
        torch.cuda.synchronize()

        # 计时
        start = time.perf_counter()
        n_iters = 50
        for _ in range(n_iters):
            dist.all_reduce(tensor)
        torch.cuda.synchronize()
        elapsed = (time.perf_counter() - start) / n_iters

        bandwidth = (2 * size_mb * (world_size - 1) / world_size) / elapsed / 1024

        if rank == 0:
            results.append({
                "size_mb": size_mb,
                "latency_ms": elapsed * 1000,
                "bandwidth_gbps": bandwidth,
            })
            print(f"Size: {size_mb}MB, Latency: {elapsed*1000:.2f}ms, "
                  f"BW: {bandwidth:.1f} GB/s")

    cleanup()
    return results