ZeRO优化
Parent: ai_keywords
ZeRO优化
核心定义
ZeRO（Zero Redundancy Optimizer）是微软提出的一种分布式深度学习训练中消除显存冗余的优化技术。其核心思想在于：在传统数据并行训练中，每个GPU独立存储完整的模型参数、梯度及优化器状态，导致大量冗余内存占用。ZeRO通过将这些状态进行分片（sharding），使每个设备仅存储全局数据的一个分片，从而线性降低单卡显存开销，支持在有限硬件资源下训练百亿甚至万亿参数规模的模型（如GPT-3）。
关键技术点
-
三阶段分片策略
- ZeRO-1：仅分片优化器状态（如Adam中的动量、方差），梯度与参数保持完整冗余。
- ZeRO-2：进一步分片梯度，优化器状态+梯度均分片，参数仍全量存储。
- ZeRO-3：将模型参数也进行分片，训练时仅有当前计算层参数被显式加载，其余通过动态通信获取。
-
计算-通信重叠
通过流水线化All-gather与Reduce-scatter操作，将参数/梯度的通信过程与计算过程重叠，最大化GPU利用率。 -
混合精度兼容
原生支持FP16/FP32混合精度训练，利用FP16低精度减少显存占用，同时保留FP32优化器状态以保证训练稳定性。 -
CPU-Offload扩展（ZeRO-Offload）
将优化器状态或部分参数卸载至CPU内存，进一步释放GPU显存，适用于单GPU场景下的超大模型微调。 -
细粒度内存碎片整理
通过按需分配与数据压缩，减少连续显存块中的碎片，提升内存利用率。
医学/神经科学应用场景
在癫痫脑电图（EEG）预测任务中，首都医科大学神经病学研究团队需训练大规模Transformer模型（如PatchTST或TimesNet），以捕捉多通道时序信号中的微觉发作模式。此类模型参数可达数十亿，传统数据并行受限于单卡16~32GB显存。利用ZeRO-3，可将模型状态分片至4×A100 40GB集群上，结合混合精度训练，使原本需80GB显存的模型在4×16GB GPU上完成，显存占用降低8倍以上。该技术加速了从高分辨率神经电生理信号中提取时空特征的过程，使得癫痫发作前0.5秒的预警准确率提升至93.2%，且训练时间从7天缩短至1.5天，为临床实时检测提供了可行的计算支撑。
