深度学习中的并行策略概述：1 单GPU优化

深度学习中的并行策略概述：1 单GPU优化

1 Training Larger Models on a Single GPU

在讨论模型的“扩展”时，往往会想到在多个GPU或多台机器上进行模型训练。不过，即便是在单个GPU上，也存在多种方法来训练更大规模的模型并提升其效率。本文将探讨包括混合精度训练 Mixed Precision Training、激活函数检查点 Activation Recomputation 、梯度累积 Gradient Accumulation 在内的一些技术。这些技术主要致力于降低训练过程中的内存占用，因为在单设备训练中，内存常常是受限的资源。此外，当在多个GPU或TPU上进行训练时，这些技术同样适用。

混合精度训练（Mixed Precision Training）

混合精度训练是一种结合16位和32位浮点数以加速模型训练的技术。16位浮点数（如float16和bfloat16）用于加速计算并减少内存使用，而32位浮点数（float32）用于关键计算以保持数值稳定。bfloat16因其较宽的表示范围，可以在无需损失缩放的情况下使用，适合作为float32的替代品，以节省内存并接近其性能。

通过将模型的特征和激活值转换为bfloat16，而保持权重和优化器状态为float32，来实现混合精度训练。这种方法在保持高精度更新的同时，减少了内存占用并提高了训练速度。如果模型过大无法完全载入内存，还可以考虑对模型参数和优化器应用更低精度的处理。通过这种方式，可以在不牺牲模型性能的前提下，优化内存使用和训练效率。

PyTorch从1.6版本开始内置了torch.cuda.amp，用于自动混合精度训练。使用自动混合精度（AMP）可以减少显存占用并加速训练。

from torch.cuda.amp import autocast, GradScaler

# 初始化模型、优化器
model = TransformerModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 初始化梯度缩放器
scaler = GradScaler()

# 训练循环
for data, target in dataloader:
    optimizer.zero_grad()
    
    # 使用autocast上下文管理器实现混合精度
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    
    # 缩放损失，以避免在反向传播时出现梯度下溢
    scaler.scale(loss).backward()
    
    # 梯度缩放器会根据需要调整梯度
    scaler.step(optimizer)
    
    # 更新缩放器
    scaler.update()

激活检查点（Activation Checkpointing / Activation Recomputation）

梯度检查点技术通过在反向传播中重新计算部分激活值，以计算换取内存。这种方法在前向传播时仅保存部分激活值，其余在反向传播时重新计算。对于像Transformer这样的大内存占用模型，当激活值的存储成为瓶颈时，此技术尤其有效，因为重新计算激活值的成本通常低于存储它们。

以一个简化的仅包含MLP块的Transformer为例。每个MLP块由两个全连接层组成，中间是GELU激活函数，使用bfloat16格式的激活值（每个激活值2字节）。设批量大小为B，序列长度为S，隐藏层大小为H。前向传播中激活值的内存消耗总计为（2BSH+8BSH+8BSH+BSH）19BSH字节。

采用梯度检查点，可以选择仅保留大小为2BSH的输入张量，并在反向传播中重新计算其他激活值。这几乎可以减少90%的激活值内存消耗，代价是需要在反向传播中重新计算这些激活值。

激活检查点是一种节省显存的技术，它通过在反向传播时重新计算前向传播中的中间激活值，而不是保存它们。

from torch.utils.checkpoint import checkpoint

# 定义一个使用激活检查点的Transformer模型部分
def partial_forward(model, inputs):
    activation = model(inputs)
    return activation

# 在训练循环中使用激活检查点
for data, target in dataloader:
    optimizer.zero_grad()
    
    # 使用checkpoint函数实现激活检查点
    with autocast():
        activation = checkpoint(partial_forward, model, data)
    
    # 继续计算损失
    output = model(activation)
    loss = criterion(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

梯度累积（Gradient Accumulation）

训练大型模型时，常在批量大小和内存需求间权衡。增加批量大小能够提高梯度估计的准确性，但也会增加内存消耗。面对硬件内存限制批量大小的情况，梯度累积技术通过累积多个小批量的梯度来模拟大批量训练的效果。每个小批量独立处理，所有小批量处理完毕后再统一更新模型。这种方法在内存受限但需要大批量训练时特别有用。然而，梯度累积的缺点在于其串行处理小批量数据，未能实现并行化，因此需要确保即使在小批量训练时也能最大化硬件的利用率。

在示意图中，展示了一个总批量大小为8的梯度累积过程，该过程通过4个大小为2的小批量（称为迷你批次）来实现。每处理完一个迷你批次，即可释放其前向传播和反向传播过程中产生的所有中间数组，随后继续处理下一个迷你批次。待所有迷你批次处理完毕，便执行优化器的更新步骤。这种方法使得我们能够以仅需相当于批量大小2的内存需求，模拟出批量大小为8的训练效果。

梯度累积允许模拟更大的批量大小，通过累积多个小批量的梯度。

# 设置梯度累积参数
accum_iter = 4

# 训练循环
for batch_idx, (data, target) in enumerate(dataloader):
    optimizer.zero_grad()
    
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    
    # 累积梯度
    loss = loss / accum_iter
    loss.backward()
    
    # 每累积一定步数后更新参数
    if ((batch_idx + 1) % accum_iter == 0) or (batch_idx + 1 == len(dataloader)):
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

2 Profiling and Scaling Single-GPU Transformer Models

下面是一个完整的示例代码，展示了如何在PyTorch中使用混合精度训练（AMP）、激活检查点（Activation Checkpointing）和梯度累积（Gradient Accumulation）来优化Transformer模型的训练：

import torch
from torch.cuda.amp import autocast, GradScaler
from torch.utils.checkpoint import checkpoint
from torch import nn, optim

# 假设TransformerModel是你要训练的模型
class TransformerModel(nn.Module):
    # 这里只是一个示例结构，你需要根据实际情况定义你的Transformer模型
    def __init__(self):
        super(TransformerModel, self).__init__()
        self.encoder = nn.TransformerEncoderLayer(d_model=512, nhead=8)
        self.decoder = nn.TransformerDecoderLayer(d_model=512, nhead=8)

    def forward(self, src, tgt):
        src = self.encoder(src)
        tgt = self.decoder(tgt, src)
        return tgt

# 初始化模型、优化器和损失函数
model = TransformerModel().cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
scaler = GradScaler()

# 设置梯度累积参数
accum_num = 4
accum_iter = 0

# 定义一个使用激活检查点的Transformer模型部分
def partial_forward(model, inputs):
    activation = model.encoder(inputs)
    return activation

# 训练循环
for data, target in dataloader:
    optimizer.zero_grad()
    
    # 使用checkpoint函数实现激活检查点
    with autocast():
        # 假设data和target已经是适当的tensor并且已经移到了GPU上
        src = checkpoint(partial_forward, model, data)
        
        # 继续计算损失
        output = model.decoder(target, src)
        loss = criterion(output, target)
    
    # 累积梯度
    scaler.scale(loss).backward()
    accum_iter += 1
    
    # 每累积一定步数后更新参数
    if accum_iter % accum_num == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()
        accum_iter = 0