2026 年 1 月 26 日，2025 冬季 InfiniTensor 大模型与人工智能系统训练营训练系统系列第三课《大模型训练中的显存优化》开讲。

本节课将系统讲解大模型训练中通过 ZeRO 分片、激活值重计算、模型并行等核心技术，以通信和计算为代价，系统性优化显存占用的方法。

数据并行（DDP）与显存冗余

1. DDP 流程：

• 每张卡上模型相同，数据分片。
• 反向计算后，梯度通过 AllReduce 汇总。
• DDP 主要解决算力扩展问题，对显存占用的优化有限。

2. ZeRO 冗余优化器：

• 目标： 在保持数据并行语义不变的前提下，系统性地消除参数、梯度和优化器状态的冗余复制。
• 三级优化方案：

ZeRO Stage 1：

• 优化器状态按 DP size 切分。
• 前向、反向流程不变，梯度更新时使用各自优化器状态切片。
• 各卡仅更新其负责的参数分片，随后通过 AllGather 同步更新后的参数。

ZeRO Stage 2：

• 梯度状态也按 DP size 切分。
• 反向计算中，梯度 AllReduce 替换为 ReduceScatter。
• 优化器更新时使用各自梯度切片。

ZeRO Stage 3：

• 参数状态也按 DP size 切分。
• 前向、反向计算中涉及完整参数时，临时进行 AllGather。
• 优化器更新后无需再进行 AllGather。

激活值检查点

1. 定义：

• 前向过程中只保留少数激活值检查点，反向过程中重新计算未保存的激活值。

2. 实现原理：

• 将前向图分为若干子图，子图边缘保存检查点。
• 反向过程中，未找到的激活值从最近检查点重新计算。

3. 子图切分策略：

• Selective Checkpointing：基于先验知识或人为调优选择性地保存激活值。
• Structured Checkpointing：根据模型结构进行结构化划分。

模型并行策略

张量并行（TP）：参数按维度切分

• 对模型内单个参数进行切分，下发到不同 rank。
• 减少参数、梯度和优化器状态的显存占用。

流水并行（PP）：模型按层切分

• 对模型不同阶段进行切分，下发到不同 rank。
• 减少单个 rank 上的模型参数量，进而减少显存占用。

序列并行（SP）：沿序列维度切分

• 沿序列维度或上下文维度进行切分，减少激活值显存占用。

Offloading 技术

• 将优化器状态移至 CPU 或硬盘
• 针对消费级显卡和小显存场景
• 性能影响较大，需异步拷贝优化

其他优化技术

• 混合精度训练：FP16/BF16+FP32 组合使用
• 梯度累积：减小 Micro-batch size
• 算子融合：将多个算子合并执行，减少中间激活占用
• 内存布局优化减少碎片

总结

本课程系统讲解了大模型训练中通过 ZeRO 分片、激活值重计算、模型并行等核心技术，以通信和计算为代价，系统性优化显存占用的方法。

💡 学习贴士

更多细节及完整课程内容，请观看直播或查看课程回放。

链接：https://www.infinitensor.com/camp/winter2025

📚 训练营报名持续进行中

扫码报名与加入群聊，或点击 「阅读原文」 直达官网报名