在 PyTorch 分布式训练中，混合精度训练（AMP，Automatic Mixed Precision）和梯度累积（Gradient Accumulation）是两种常用的优化技术，可以显著提升训练效率和模型性能。以下是它们的实现方法和关键注意事项，结合之前的 DDP 流程进行扩展。

------

1. 混合精度训练（AMP）

混合精度训练通过使用 FP16（半精度浮点）和 FP32（单精度浮点）的组合，减少显存占用并加速计算。PyTorch 提供了torch.cuda.amp或torch.npu.amp（昇腾 NPU）模块来简化实现。

关键步骤

  1. 初始化 AMP：使用GradScaler管理梯度缩放。

  2. 包装前向和反向传播：在autocast上下文中执行计算。

  3. 更新权重：使用scaler.step()和scaler.update()。

代码示例

【python】
 from torch.cuda.amp import GradScaler, autocast  # GPU用cuda.amp，NPU用npu.amp

def train(rank, world_size):
    model = ...  # 初始化模型和DDP（见前文）
    train_loader = prepare_dataloader(rank, world_size)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    scaler = GradScaler()  # 初始化梯度缩放器

    for epoch in range(10):
        model.train()
        train_loader.sampler.set_epoch(epoch)

        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            data = data.view(data.size(0), -1)

            optimizer.zero_grad()
            
            # 使用autocast混合精度计算
            with autocast():
                output = model(data)
                loss = criterion(output, target)
            
            # 反向传播（梯度缩放）
            scaler.scale(loss).backward()
            scaler.step(optimizer)  # 更新权重
            scaler.update()        # 更新缩放因子

            if batch_idx % 100 == 0 and rank == 0:
                print(f"Epoch: {epoch} | Batch: {batch_idx} | Loss: {loss.item()}")

关键点

  • autocast()：自动选择 FP16/FP32 计算，减少显存占用。

  • GradScaler：防止 FP16 梯度下溢（underflow）。

  • 昇腾 NPU：将torch.cuda.amp替换为torch.npu.amp，并确保后端为hccl。

------

2. 梯度累积

梯度累积通过在多个批次上累积梯度，模拟更大批次（batch size）的训练效果，适用于显存不足的场景。

关键步骤

  1. 设置累积步数：例如accum_steps=4表示每 4 个批次更新一次权重。

  2. 跳过优化器更新：仅在累积足够步数后执行optimizer.step()。

  3. 同步梯度：DDP 会自动处理跨进程的梯度同步。

代码示例

【python】
 def train(rank, world_size):
    model = ...  # 初始化模型和DDP
    train_loader = prepare_dataloader(rank, world_size)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    accum_steps = 4  # 累积4个批次的梯度

    for epoch in range(10):
        model.train()
        train_loader.sampler.set_epoch(epoch)
        optimizer.zero_grad()  # 初始化梯度

        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            data = data.view(data.size(0), -1)

            output = model(data)
            loss = criterion(output, target) / accum_steps  # 损失归一化
            loss.backward()  # 反向传播（梯度累积）

            # 每accum_steps更新一次权重
            if (batch_idx + 1) % accum_steps == 0:
                optimizer.step()
                optimizer.zero_grad()

            if batch_idx % 100 == 0 and rank == 0:
                print(f"Epoch: {epoch} | Batch: {batch_idx} | Loss: {loss.item() * accum_steps}")

关键点

  • 损失归一化：loss = loss / accum_steps避免梯度放大。

  • 梯度清零：在optimizer.step()后调用zero_grad()。

  • 与 AMP 结合：在autocast上下文中使用梯度累积：

【python】
 with autocast():
      output = model(data)
      loss = criterion(output, target) / accum_steps
  scaler.scale(loss).backward()
  if (batch_idx + 1) % accum_steps == 0:
      scaler.step(optimizer)
      scaler.update()
      optimizer.zero_grad()

------

3. 完整流程整合

将 DDP、AMP 和梯度累积结合的完整训练流程：

【python】
 def train(rank, world_size):
    model = ...  # 初始化模型和DDP
    train_loader = prepare_dataloader(rank, world_size)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    scaler = GradScaler()
    accum_steps = 4

    for epoch in range(10):
        model.train()
        train_loader.sampler.set_epoch(epoch)
        optimizer.zero_grad()

        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            data = data.view(data.size(0), -1)

            with autocast():
                output = model(data)
                loss = criterion(output, target) / accum_steps
            
            scaler.scale(loss).backward()

            if (batch_idx + 1) % accum_steps == 0:
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()

            if batch_idx % 100 == 0 and rank == 0:
                print(f"Epoch: {epoch} | Batch: {batch_idx} | Loss: {loss.item() * accum_steps}")

------

常见问题与解决

Q1: 混合精度训练损失爆炸

  • 原因：梯度缩放因子（GradScaler）设置不当。

  • 解决：调整init_scale（默认2**16）或使用scaler.update()动态调整。

Q2: 梯度累积效果不明显

  • 原因：累积步数（accum_steps）设置过小，或学习率未调整。

  • 解决：增大accum_steps或按比例调整学习率（如lr *= accum_steps）。

Q3: 多进程打印冲突

  • 原因：所有进程同时打印日志。

  • 解决：仅rank == 0的进程打印日志。

------

总结

  1. 混合精度训练：通过autocast和GradScaler减少显存占用并加速计算。

  2. 梯度累积：通过多批次累积梯度，模拟更大批次训练。

  3. 结合使用：在autocast上下文中实现梯度累积，并注意损失归一化和梯度清零。