共计 1521 个字符，预计需要花费 4 分钟才能阅读完成。

理论计算

要准确预估显存需求，我们需要了解模型参数、激活值（activations）和优化器状态（optimizer states）这三部分对显存的占用。这里给出一个简单的计算公式：

总显存 = 参数显存 + 激活值显存 + 优化器状态显存

对于 2B（20 亿）参数的模型：

1. 参数显存：
2. FP32 精度：每个参数占 4 字节，2B 参数需要 2×10⁹ × 4 = 8 GB

3. FP16/BF16 精度：每个参数占 2 字节，显存减半为 4 GB

4. 优化器状态显存：

5. 使用 Adam 优化器时，每个参数需要存储动量（momentum）和方差（variance）：

   • FP32 优化器：2×10⁹ × 4 × 2 = 16 GB（每个参数 8 字节）
   • FP16 优化器：2×10⁹ × 2 × 2 = 8 GB（每个参数 4 字节）

6. 激活值显存：

7. 取决于 batch size 和序列长度，通常占总显存的 30%-50%
8. 示例：batch size=32 时约需 6 -10 GB

优化策略

梯度检查点（Gradient Checkpointing）

通过牺牲部分计算时间换取显存节省，核心思想是只保存部分层的激活值，其余层在反向传播时重新计算。

import torch
from torch.utils.checkpoint import checkpoint

class BigModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = torch.nn.Linear(1024, 1024)
        self.layer2 = torch.nn.Linear(1024, 1024)

    def forward(self, x):
        # 只在反向传播时保留 layer2 的激活
        x = checkpoint(self.layer1, x)  # 注意：传递函数对象而非调用结果
        x = self.layer2(x)
        torch.cuda.synchronize()  # CUDA 同步点确保计时准确
        return x

混合精度训练（AMP）

自动混合精度训练可以大幅减少显存占用并提升计算速度：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

torch.cuda.empty_cache()  # 建议在每个 epoch 结束后调用

模型并行策略选择

1. Tensor 并行：适合单个层参数极大的情况（如大型 FFN 层）
2. 优势：通信仅在层内进行

3. 劣势：需要修改模型架构

4. Pipeline 并行：适合层数多的模型

5. 优势：实现相对简单
6. 劣势：存在流水线气泡（bubble）开销

性能实测

生产环境陷阱

1. 分布式训练通信开销：
2. 小规模集群（≤8 节点）建议用 All-Reduce

3. 大规模集群考虑使用 Ring-Allreduce

4. 激活值内存碎片化：

5. 定期调用torch.cuda.empty_cache()
6. 使用 torch.cuda.memory_summary() 监控：

print(torch.cuda.memory_summary(device=None, abbreviated=False))

开放性问题

当显存不足时，CPU offloading 虽然能扩展可用内存，但会带来：
1. 数据传输延迟增加 5 -10 倍
2. 需要精确计算 PCIe 带宽瓶颈
3. 如何平衡计算强度和通信开销？