多GPU并行训练中的显存分配优化技巧

在大模型训练中，多GPU并行训练是提升训练效率的关键技术。然而，显存分配不当会导致训练中断或效率低下。本文将分享几个实用的显存优化技巧。

1. 使用torch.cuda.set_per_process_memory_fraction()控制显存占用

import torch
# 限制每个进程使用80%的GPU显存
torch.cuda.set_per_process_memory_fraction(0.8)

这在多任务并行时特别有效，避免显存溢出。

2. 合理设置batch size与gradient accumulation 当单个batch无法容纳更多数据时，可使用梯度累积技术：

accumulation_steps = 4
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps  # 梯度缩放
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

3. 使用torch.utils.checkpoint优化内存使用 对于深层网络，可以开启梯度检查点：

from torch.utils.checkpoint import checkpoint
# 在前向传播中启用检查点
output = checkpoint(model, input_tensor)

这些技巧可显著减少显存占用，提升训练稳定性。建议根据实际硬件配置调整参数。