大模型训练与推理的显存优化实战：从激活重计算到混合精度

一、为什么显存优化是大模型生命线的核心

2026 年，万亿参数模型已从实验室走向生产部署。训练一个 100B 参数的模型，仅参数本身就需要 200GB 显存（FP16），加上优化器状态、梯度和激活值，总显存需求轻松突破 1TB。推理场景虽稍好，但 KV Cache 的线性增长依然让长上下文服务成为显存黑洞。

显存优化技术已不再是"锦上添花"，而是决定模型能否落地的工程关键。

二、激活重计算（Activation Checkpointing）

2.1 核心思想

训练时只保存部分层的激活值，其余层在反向传播时重新计算。这是一种用计算换显存的典型策略。

显存占用 = 保存的激活 + 参数 + 优化器状态
         ↓
选择性丢弃大部分激活，反向传播时重算

2.2 PyTorch 实战

import torch
from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def forward(self, x):
        # 使用 checkpoint 包装，中间激活不保存
        x = checkpoint(self._forward_block, x, use_reentrant=False)
        return x
    
    def _forward_block(self, x):
        x = x + self.attention(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

2.3 重计算策略选择

策略	显存节省	计算开销	适用场景
Full Checkpoint	~60%	+33%	训练显存极致紧张
Selective	~30%	+10%	FFN 层不重算，只重算 Attention
No Checkpoint	0%	0%	推理或显存充足

选择性重算的关键洞察：Attention 层的激活值显存占用大但计算量相对小，FFN 层反之。因此对 Attention 层做重算是性价比最高的选择。

三、混合精度训练（Mixed Precision Training）

3.1 FP16/BF16/FP8 对比

格式	位数	指数位	精度	硬件支持	显存
FP32	32	8	高	通用	4字节/参数
FP16	16	5	中	Volta+	2字节/参数
BF16	16	8	中	Ampere+	2字节/参数
FP8	8	4/5	低	Hopper+	1字节/参数

BF16 凭借与 FP32 相同的动态范围，成为 2026 年训练的事实标准。FP8 则在推理侧快速普及。

3.2 PyTorch AMP 实战

from torch.amp import autocast, GradScaler

scaler = GradScaler()

for batch in dataloader:
    with autocast(device_type='cuda', dtype=torch.bfloat16):
        output = model(batch)
        loss = criterion(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

3.3 Loss Scaling 的数学本质

FP16 的最小正正规数约 6×10⁻⁸，梯度下溢是混合精度训练的致命问题。Loss Scaling 将损失乘以常数 S 后再反向传播，使梯度回到 FP16 可表示范围：

前向: loss_scaled = loss × S
反向: grad_scaled = ∂(loss_scaled)/∂w = S × ∂loss/∂w
更新: w = w - lr × (grad_scaled / S)

GradScaler 自动检测梯度 Inf/NaN，动态调整 S：连续 N 次成功加倍 S，出现溢出则减半 S。

四、优化器状态分片：ZeRO 三阶段

DeepSpeed ZeRO（Zero Redundancy Optimizer）通过消除数据并行中的显存冗余实现线性扩展。

ZeRO 三阶段演进

Stage 1: 分片优化器状态
├── 每 GPU 只存 1/N 的优化器状态（Adam 的 m 和 v）
└── 显存降至: params + grads + optimizer_states/N

Stage 2: + 分片梯度
├── 每 GPU 只存 1/N 的梯度
└── 显存降至: params + (grads + optimizer_states)/N

Stage 3: + 分片参数
├── 每 GPU 只存 1/N 的参数
├── 前向/反向时通过 AllGather 临时收集完整参数
└── 显存降至: total_params/N + 激活值

4.2 ZeRO 实战配置

{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e7,
    "stage3_prefetch_bucket_size": 5e7
  }
}

CPU Offload 将优化器状态和参数更新卸载到主机内存，用 PCIe 带宽换 GPU 显存。在 2026 年 PCIe 5.0 x16（64GB/s）的场景下，配合 overlap_comm 可将通信与计算重叠，几乎不增加训练时间。

五、PagedAttention：推理 KV Cache 的革命

5.1 操作系统内存分页的思想迁移

传统 KV Cache 为每个序列预留最大长度的连续显存空间，导致：

• 60%+ 显存浪费（序列长度差异大）
• 无法动态共享相同 prefix 的缓存

vLLM 提出的 PagedAttention 借鉴 OS 虚拟内存分页机制：

• KV Cache 被切成固定大小的 block（如 16 tokens）
• 逻辑块地址 → 物理块地址的映射表（类似页表）
• 不同序列可共享相同前缀的物理块（Copy-on-Write）

5.2 显存利用率对比

场景	传统 KV Cache	PagedAttention
短序列平均	~40% 利用率	~96% 利用率
Beam Search	× beam_size 内存	CoW 共享内存
Prefix Caching	不支持	自动复用

六、量化技术体系

6.1 训练后量化（PTQ）

方法	精度	显存节省	延迟影响
W8A8 (INT8)	高	50%	-10%
W4A16 (GPTQ)	中高	75%	+5%
W4A4 (AWQ)	中	87.5%	+15%
W3A16 (GGUF)	中低	81%	+30%

6.2 量化感知训练（QAT）

在训练过程中模拟量化噪声，让模型适应低精度表示。核心公式：

前向: x_quant = round(clamp(x/scale, -128, 127)) × scale
直通估计器: ∂L/∂x ≈ ∂L/∂x_quant

2026 年，NVIDIA Hopper 架构原生支持 FP8 Tensor Core，使得 QAT 到 FP8 的迁移成本极低。

七、统一内存与异构计算

7.1 CUDA Unified Memory

# 使用统一内存自动在 CPU/GPU 间迁移数据
import torch

# PyTorch 的 Unified Memory 抽象
with torch.cuda.stream(torch.cuda.Stream()):
    # 预取数据到 GPU
    data = data.to('cuda', non_blocking=True)
    torch.cuda.current_stream().synchronize()

CUDA Unified Memory + HBM3e 的分层存储架构使得单卡可运行更大模型，但需要精心设计访问模式以避免性能抖动。

7.2 异构显存池

┌─────────────────────────────────────┐
│           异构显存池                  │
├──────────┬──────────┬───────────────┤
│ 本地 HBM │ 远程 GPU │ 主机 DDR5     │
│ 80GB     │ NVLink   │ 512GB         │
│ 快       │ 中       │ 慢            │
└──────────┴──────────┴───────────────┘
         LRU 自动迁移策略

八、2026 年显存优化的工程实践总结

训练场景

• BF16 + GradScaler：默认选择，几乎零精度损失
• ZeRO Stage 2 + CPU Offload：中大规模（10B+）训练标配
• Selective Activation Checkpointing：对 Attention 层做选择性重算
• FlashAttention-3：利用 IO Aware Attention 减少 HBM 访问

推理场景

• PagedAttention：vLLM 已成事实标准
• W4A16 AWQ/GPTQ：4-bit 量化几乎无感知精度损失
• Prefix Caching：多轮对话场景节省 60%+ 重复计算
• Chunked Prefill：将长 prefill 与 decode 交错，避免延迟抖动

显存优化的决策树

模型训练?
├── 是 → 参数量 < 1B? → 单卡 BF16
│        │            └── 否 → ZeRO Stage 2
│        └── 显存仍不足? → + Activation Checkpointing
│                       └── + CPU Offload
└── 否 (推理) → 延迟敏感? → FP8 量化
                         └── 吞吐优先? → W4A16 + PagedAttention

九、未来展望

2026 年及以后，显存优化技术将沿着三个方向演进：

• 硬件-软件协同设计：GPU 直接支持稀疏性结构化（如 2:4 稀疏），框架自动利用
• 编译期优化：MLIR/Tracer 编译器自动插入重计算和量化决策
• 分布式共享显存：跨节点 NVLink + CXL 构建全局内存池，打破单卡显存天花板

显存优化的终极目标是让算法工程师不再被显存束缚，专注于模型架构和训练策略本身。随着硬件迭代和软件栈成熟，这一目标正在逐步接近现实。