09-分布式训练¶

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

def setup(rank, world_size):
    """初始化分布式环境"""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # 初始化进程组
    # nccl: NVIDIA GPU通信后端（最快）
    # gloo: CPU通信后端（兼容性好）
    dist.init_process_group(
        backend='nccl',  # GPU用nccl, CPU用gloo
        rank=rank,
        world_size=world_size
    )
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def train_ddp(rank, world_size, epochs=10):
    """DDP训练主函数"""
    setup(rank, world_size)

    # 1. 模型 → 指定GPU → DDP包装
    model = nn.Sequential(
        nn.Linear(784, 256), nn.ReLU(),
        nn.Linear(256, 128), nn.ReLU(),
        nn.Linear(128, 10)
    ).to(rank)

    ddp_model = DDP(model, device_ids=[rank])

    # 2. DistributedSampler确保每张卡看到不同的数据
    train_dataset = ...  # your dataset
    sampler = DistributedSampler(
        train_dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )
    dataloader = DataLoader(  # DataLoader批量加载数据，支持shuffle和多进程
        train_dataset,
        batch_size=64,
        sampler=sampler,
        num_workers=4,
        pin_memory=True
    )

    optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    # 3. 训练循环
    for epoch in range(epochs):
        sampler.set_epoch(epoch)  # 重要！确保每epoch数据打乱方式不同
        ddp_model.train()  # train()开启训练模式

        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自动在backward时同步梯度
            optimizer.step()  # 根据梯度更新模型参数

            if rank == 0 and batch_idx % 100 == 0:
                print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}")  # .item()将单元素张量转为Python数值

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

    cleanup()

# 启动: torchrun --nproc_per_node=4 train.py
# 注意: torch.distributed.launch 已废弃，请统一使用 torchrun
if __name__ == "__main__":
    import torch.multiprocessing as mp
    world_size = torch.cuda.device_count()
    mp.spawn(train_ddp, args=(world_size,), nprocs=world_size, join=True)

# DDP关键技术点:

# 1. 梯度同步: AllReduce操作
#    所有GPU在backward()结束时自动同步梯度
#    使用Ring-AllReduce算法（带宽最优）

# 2. 梯度桶(Gradient Buckets):
#    不是每个参数单独同步，而是打包成桶
#    桶大小默认25MB，可调:
ddp_model = DDP(model, device_ids=[rank], bucket_cap_mb=25)

# 3. 计算与通信重叠:
#    DDP在backward()期间就开始同步已计算好的梯度
#    → 通信和计算并行进行

# 4. 广播初始参数:
#    DDP初始化时将rank 0的参数广播到所有其他rank
#    → 确保所有卡从相同的参数开始训练

import torch
import torch.nn as nn
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
    BackwardPrefetch,
    CPUOffload,
)
from torch.distributed.fsdp.wrap import (
    size_based_auto_wrap_policy,
    transformer_auto_wrap_policy,
)
from functools import partial

def setup_fsdp_model(model, rank):
    """配置FSDP模型"""

    # 混合精度策略
    mixed_precision_policy = MixedPrecision(
        param_dtype=torch.bfloat16,      # 参数存储精度
        reduce_dtype=torch.bfloat16,     # 梯度通信精度
        buffer_dtype=torch.bfloat16,     # Buffer精度
    )

    # 自动包装策略 (按参数量)
    auto_wrap_policy = partial(
        size_based_auto_wrap_policy,
        min_num_params=1e6  # 参数>1M的子模块单独分片
    )

    # 对Transformer模型，更推荐按层包装:
    # from transformers.models.llama.modeling_llama import LlamaDecoderLayer
    # auto_wrap_policy = partial(
    #     transformer_auto_wrap_policy,
    #     transformer_layer_cls={LlamaDecoderLayer}
    # )

    # FSDP包装
    fsdp_model = FSDP(
        model,

        # 分片策略
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        # FULL_SHARD: 完全分片（= ZeRO-3，显存最省）
        # SHARD_GRAD_OP: 分片梯度和优化器（= ZeRO-2）
        # NO_SHARD: 不分片（= DDP）

        # 混合精度
        mixed_precision=mixed_precision_policy,

        # 自动包装
        auto_wrap_policy=auto_wrap_policy,

        # 预取优化（提前收集下一层参数，计算通信重叠）
        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,

        # CPU Offload (极端显存不足时使用，会降低速度)
        # cpu_offload=CPUOffload(offload_params=True),

        device_id=rank,
        use_orig_params=True,  # 兼容非FSDP优化器
    )

    return fsdp_model

def train_fsdp(rank, world_size):
    """FSDP训练示例"""
    setup(rank, world_size)

    # 创建模型（先在CPU上，FSDP会自动管理设备）
    model = create_large_model()
    fsdp_model = setup_fsdp_model(model, rank)

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

    for epoch in range(num_epochs):
        for batch in dataloader:
            batch = {k: v.to(rank) for k, v in batch.items()}

            loss = fsdp_model(**batch).loss
            loss.backward()

            # 梯度裁剪（FSDP需要特殊处理）
            fsdp_model.clip_grad_norm_(max_norm=1.0)

            optimizer.step()
            optimizer.zero_grad()

    # FSDP模型保存
    # 方法1: 全量保存（收集到rank 0）
    from torch.distributed.fsdp import FullStateDictConfig, StateDictType

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

    cleanup()

import torch
import torch.nn as nn
import torch.distributed as dist

class ColumnParallelLinear(nn.Module):  # 继承nn.Module定义神经网络层
    """列并行线性层 (简化实现)"""

    def __init__(self, in_features, out_features, world_size, rank):  # __init__构造方法，创建对象时自动调用
        super().__init__()  # super()调用父类方法
        assert out_features % world_size == 0  # assert断言，条件为False时抛出异常
        self.local_out = out_features // world_size
        self.rank = rank
        self.world_size = world_size

        # 每张卡只存储 1/world_size 的权重列
        self.weight = nn.Parameter(
            torch.randn(in_features, self.local_out) * 0.01
        )
        self.bias = nn.Parameter(torch.zeros(self.local_out))

    def forward(self, x):
        # 本地计算
        local_output = x @ self.weight + self.bias

        # AllGather收集所有卡的输出
        output_list = [torch.zeros_like(local_output) for _ in range(self.world_size)]
        dist.all_gather(output_list, local_output)

        return torch.cat(output_list, dim=-1)

class RowParallelLinear(nn.Module):
    """行并行线性层 (简化实现)"""

    def __init__(self, in_features, out_features, world_size, rank):
        super().__init__()
        assert in_features % world_size == 0
        self.local_in = in_features // world_size
        self.rank = rank
        self.world_size = world_size

        self.weight = nn.Parameter(
            torch.randn(self.local_in, out_features) * 0.01
        )
        self.bias = nn.Parameter(torch.zeros(out_features))

    def forward(self, x):
        # x已经是切分后的 (每张卡只有部分特征)
        local_output = x @ self.weight

        # AllReduce求和
        dist.all_reduce(local_output, op=dist.ReduceOp.SUM)

        return local_output + self.bias

# 核心通信操作

# 1. Broadcast: 一个发送，所有接收
#    rank 0: [A, B, C, D] → 所有rank: [A, B, C, D]
dist.broadcast(tensor, src=0)

# 2. AllReduce: 所有归约，所有接收结果
#    rank 0: [1, 2]  rank 1: [3, 4]
#    → 所有rank: [4, 6]  (sum)
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

# 3. AllGather: 所有收集
#    rank 0: [A]  rank 1: [B]  rank 2: [C]
#    → 所有rank: [A, B, C]
output_list = [torch.zeros_like(tensor) for _ in range(world_size)]
dist.all_gather(output_list, tensor)

# 4. ReduceScatter: 归约后分发
#    rank 0: [A0, A1]  rank 1: [B0, B1]
#    → rank 0: [A0+B0]  rank 1: [A1+B1]
dist.reduce_scatter(output, input_list, op=dist.ReduceOp.SUM)

# 5. AllToAll: 全交换
#    rank 0: [A0, A1]  rank 1: [B0, B1]
#    → rank 0: [A0, B0]  rank 1: [A1, B1]
dist.all_to_all(output_list, input_list)

import torch.profiler

# 使用PyTorch Profiler分析分布式训练性能
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA
    ],
    schedule=torch.profiler.schedule(
        wait=1, warmup=1, active=3, repeat=1
    ),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'),
    record_shapes=True,
    with_stack=True
) as prof:
    for step, (data, target) in enumerate(dataloader):
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        prof.step()

# 查看结果: tensorboard --logdir=./logs
# 重点关注:
# 1. 计算 vs 通信时间比（理想 > 80% 计算）
# 2. GPU利用率（理想 > 90%）
# 3. 通信等待时间（有无通信瓶颈）

# Accelerate: 最简单的分布式训练方式（支持DDP/FSDP/DeepSpeed）
# pip install accelerate

from accelerate import Accelerator
from torch.utils.data import DataLoader

def train_with_accelerate():
    # 一行初始化
    accelerator = Accelerator(
        mixed_precision='bf16',  # 混合精度
        gradient_accumulation_steps=4,
    )

    model = create_model()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    dataloader = DataLoader(dataset, batch_size=32)
    scheduler = get_scheduler("cosine", optimizer, ...)

    # Accelerate自动处理:
    # - 模型放到正确的GPU
    # - 数据自动分片
    # - 梯度自动同步
    model, optimizer, dataloader, scheduler = accelerator.prepare(
        model, optimizer, dataloader, scheduler
    )

    for epoch in range(num_epochs):
        model.train()
        for batch in dataloader:
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                accelerator.backward(loss)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_main_process:
                print(f"Loss: {loss.item():.4f}")

    # 保存
    accelerator.wait_for_everyone()
    unwrapped = accelerator.unwrap_model(model)
    accelerator.save(unwrapped.state_dict(), "model.pth")

# 配置: accelerate config  (交互式配置DDP/FSDP/DeepSpeed)
# 启动: accelerate launch train.py

# 分布式训练最常见的bug: 进程死锁

# 原因1: 不同rank执行了不同的代码路径
# 错误示例:
if rank == 0:
    output = model(data)  # 只在rank 0执行 → 其他rank等待AllReduce → 死锁！

# 正确: 所有rank必须执行相同的集合操作

# 原因2: 数据不一致导致不同rank的iteration数不同
# 解决: 使用DistributedSampler + drop_last=True

# 调试工具:
# NCCL_DEBUG=INFO torchrun ... → 打印详细的NCCL通信日志
# TORCH_DISTRIBUTED_DEBUG=DETAIL → PyTorch分布式调试信息

# 梯度检查点: 不保存中间激活值，反向传播时重新计算
from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def forward(self, x):
        # 用checkpoint包装 → 不保存此层的激活值
        return checkpoint(self._forward_impl, x, use_reentrant=False)

    def _forward_impl(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x