HBM内存墙下的训练与推理优化：梯度检查点与ZeRO的工程权衡

内存墙的现实：当 HBM 成本占据芯片三分之二

AI 芯片的成本结构正在发生根本性变化。根据 Epoch AI 的追踪数据，HBM（高带宽内存）在 AI 芯片组件成本中的占比已攀升至近三分之二。这一比例的背后是模型规模的指数级增长 —— 从 Bert-Large 的 3 亿参数到 Megatron-Turing-530B 的 5300 亿参数，仅用了不到四年时间。当内存容量和带宽成为制约模型扩展的首要瓶颈时，单纯依靠增加 HBM 容量已不再是经济可行的解决方案。

内存墙（Memory Wall）的本质在于：模型参数、优化器状态、激活值三者的内存占用与 HBM 容量之间存在结构性矛盾。以 175B 参数的 GPT-3 为例，半精度推理需要约 350GB 内存来存储模型权重，这远超单张 A100-40GB 的容量。训练阶段的内存需求更为严峻 ——Adam 优化器需要存储动量和二阶矩，使每参数的内存占用膨胀至 12 字节以上。

训练阶段的内存优化策略

梯度检查点：以计算换内存

梯度检查点（Gradient Checkpointing）是解决激活值内存爆炸的经典策略。传统训练需要存储每一层的激活值以计算梯度，而检查点技术仅保留关键层的激活，在反向传播时重新计算被丢弃的中间结果。

工程权衡的核心在于检查点密度的选择。在 Transformer 架构中，一个经验法则是对注意力层和前馈层进行选择性检查点。对于参数量在 10B-100B 范围的模型，检查点策略通常可减少 40%-60% 的峰值内存占用，代价是增加约 20%-30% 的前向计算时间。这种交换在内存受限场景下往往是划算的 —— 当 GPU 内存成为瓶颈时，释放出的内存可用于增大 batch size，从而提升整体吞吐。

可落地的配置参数：

• 检查点间隔：每 2-4 个 Transformer 块设置一个检查点
• 优先级：优先对激活值占用大的层（如大维度 MLP）启用检查点
• 精度混合：结合 FP16/BF16 训练，将激活值以半精度存储

ZeRO 优化器：分阶段内存分区

DeepSpeed 的 ZeRO（Zero Redundancy Optimizer）通过数据并行中的内存分区策略，将优化器状态、梯度和模型参数分散到多个设备。ZeRO 分为三个阶段，每个阶段对应不同的内存节省程度和通信开销：

ZeRO Stage 1：仅分区优化器状态，内存节省约 4 倍（Adam 优化器状态占 8 字节 / 参数，分区后降至 2 字节 / 参数）

ZeRO Stage 2：额外分区梯度，内存节省约 8 倍

ZeRO Stage 3：进一步分区模型参数，实现与数据并行度成比例的内存节省。在 64 路数据并行下，单卡内存占用可降低至原始需求的 1/64

关键权衡：Stage 3 虽然内存效率最高，但需要在每次前向和后向传播时进行参数收集（all-gather）和梯度规约（reduce-scatter），通信量显著增加。在实际部署中，对于参数量在 100B 以上的模型，通常采用 Stage 2 作为甜点配置，在内存节省和通信开销之间取得平衡。

推理阶段的内存优化策略

ZeRO-Inference：异构内存流式加载

训练优化的成果延伸至推理领域。ZeRO-Inference 针对单卡 GPU 内存无法容纳大模型的场景，将模型权重常驻于 CPU DRAM 或 NVMe SSD，仅在计算时逐层流式加载到 GPU HBM。

核心设计决策：完全卸载（Full Offload）vs 部分卸载（Partial Offload）。DeepSpeed 的评估表明，当模型无法完全装入 GPU 时，将有限的 HBM 全部用于增大 batch size，而非用于缓存部分模型权重，能带来更高的 token 生成吞吐。以 OPT-30B 在 V100-32GB 上的测试为例，完全卸载到 CPU 内存的吞吐为 43 tokens / 秒，而部分卸载 12B 参数到 GPU 的吞吐仅为 18 tokens / 秒。

** 层预取（Layer Prefetching）** 是关键的性能优化。通过 PCIe 链路在计算当前层时预取下一层权重，可以隐藏传输延迟。对于 NVMe 卸载场景，预取可将性能提升 13%-21%，因为 CPU 内存作为缓冲层允许 NVMe 到 CPU 和 CPU 到 GPU 的传输重叠执行。

适用场景判定：ZeRO-Inference 专为吞吐优先的批处理场景设计。对于延迟敏感的在线服务，仍需采用模型并行将完整模型装入多张 GPU 的 HBM。

KV Cache 的内存管理

在自回归生成任务中，KV Cache 的内存占用随序列长度线性增长。对于长文本生成，KV Cache 可能成为比模型权重更严峻的内存瓶颈。

工程实践：

• 动态分配：根据最大生成长度预分配 KV Cache，避免频繁的内存分配和拷贝
• 分页管理：借鉴 vLLM 的 PagedAttention，将 KV Cache 分页存储，提高内存利用率和共享效率
• 量化压缩：对 KV Cache 进行 INT8 或 FP8 量化，在可接受的精度损失下减半内存占用

可落地的参数配置清单

基于上述分析，以下是可直接应用的配置参数：

训练场景（DeepSpeed 配置）：

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true
  },
  "gradient_checkpointing": true,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000
  }
}

推理场景（ZeRO-Inference 配置）：

{
  "zero_optimization": {
    "stage": 3,
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true
  }
}

关键阈值参考：

• 当模型参数量 > GPU 显存容量（GB）× 2.5 时，启用 ZeRO-Inference
• 当序列长度 > 2048 时，启用 KV Cache 分页管理
• 当 batch size 可增大至原来的 1.5 倍以上时，梯度检查点的计算开销可被吞吐收益抵消

结论

HBM 内存墙并非不可逾越的障碍，而是推动算法 - 系统协同设计的催化剂。梯度检查点、ZeRO 系列优化器和异构内存管理共同构成了一套分层递进的内存优化工具箱。工程团队需要根据模型规模、硬件配置和任务特性（训练 vs 推理、延迟 vs 吞吐）进行策略组合。随着模型向万亿参数规模演进，这些内存优化技术将从可选方案变为必备基础设施。

参考来源：

• DeepSpeed ZeRO-Inference 技术文档（https://www.deepspeed.ai/2022/09/09/zero-inference.html）
• Epoch AI 关于 AI 芯片内存成本的分析
• MLSys 课程关于内存优化的技术资料

ai-systems