神经网络Transformer架构中的迭代式参数更新 概念引入：传统优化与Transformer的挑战 神经网络训练的本质是通过调整模型参数（权重和偏置）来最小化损失函数，这个过程称为 参数更新 。最经典的方法是 梯度下降 ：计算损失函数关于每个参数的梯度（导数），然后沿梯度反方向（即减小损失的方向）以一定步长（学习率）更新参数。 在训练标准深度神经网络时，通常使用 反向传播 算法计算梯度，然后应用如SGD、Adam等优化器进行一次前向传播和一次反向传播后，即更新一次参数。我们可称此为“单步”或“直接”参数更新。 然而，当模型规模变得极其庞大，如拥有数千亿参数的Transformer大语言模型时，直接应用此范式面临巨大挑战： 内存限制 。在单个训练步骤中，同时存储模型所有参数的优化器状态（如动量）、梯度和参数本身，所需内存远超单个加速器（如GPU）的容量。 核心思想：分而治之的参数更新 迭代式参数更新 是一种专门为训练超大规模模型设计的、解决内存瓶颈的技术策略。其核心思想是： 将完整的模型参数集划分为多个组或“子集”，在单个训练步骤中，只将其中一个子集加载到设备内存中，进行计算图和梯度计算，并立即更新这个子集的参数，然后循环处理其他参数子集。 这与传统方法（所有参数同时参与一次前向/反向传播，然后同时更新）形成对比。迭代式更新通过 分片 和 串行化 处理，用 更长的计算时间（多次迭代）来换取对单设备内存需求的极大降低 ，使得在有限硬件上训练巨大模型成为可能。 关键技术实现：ZeRO优化器 迭代式参数更新最著名和系统的实现是微软DeepSpeed库中的 ZeRO 系列优化器（特别是ZeRO-Offload和ZeRO-3）。 ZeRO-3阶段 的工作原理是深度示例： 参数分区 ： 在训练开始前，将整个模型的数十亿参数在数据并行进程组的所有GPU（或CPU内存）之间进行 均匀分区 。每个GPU只持久存储并“负责”更新分配给它的那一部分参数。 前向传播 ： 当进行前向传播时，对于当前层所需的参数，如果它们不在当前GPU上，则通过 点对点通信 从持有该参数分区的GPU实时获取。计算完成后，非本地参数可以被丢弃，仅保留计算结果。 反向传播 ： 反向传播过程类似。计算梯度时，需要哪部分参数，就从持有它的GPU获取。计算得到的梯度同样只与对应的参数分区相关联。 迭代式更新 ： 关键步骤在此。当反向传播为某个参数分区计算出梯度后， 持有该分区的GPU会立即使用优化器（如Adam）更新这部分参数 。由于每个GPU只负责更新自己分区内的参数，它只需要存储这部分参数的优化器状态（如动量、方差），内存开销锐减。 循环往复 ： 模型训练的一个“全局步骤”由多个这样的“子步骤”构成。系统会遍历所有参数分区，依次完成它们所需的前向/反向计算和更新。从全局看，所有参数都基于同一批数据（在数据并行下）被更新了一次，但更新是分批（迭代式）完成的。 优势与权衡 核心优势 ： 极大降低单设备内存消耗 。使得用远小于模型参数量总数的设备内存来训练模型成为可能，例如在40GB显存的GPU上训练千亿参数模型。 主要代价 ： 增加通信开销和延长单步训练时间 。由于需要频繁地在设备间获取参数和同步梯度，网络通信带宽可能成为瓶颈。同时，由于串行处理参数分区，完成一个全局训练步骤所需的计算时间比传统并行方式更长。 适用场景 ： 是训练 参数规模远超单个设备内存容量 的巨型Transformer模型的必备技术之一。对于小模型，其通信开销可能得不偿失。 与其他技术的关联 与 模型并行 ： 两者都用于解决大模型内存问题，但思路不同。模型并行是将模型的不同层或算子拆分到不同设备上，是计算图的拆分。迭代式参数更新（如ZeRO）通常是在 数据并行 的基础上，对参数状态进行分区和按需调度，两者常结合使用形成更高效的混合并行策略。 与 梯度累积 ： 梯度累积是通过多个小批次累加梯度来模拟大批次训练，主要解决显存不足以容纳大批次数据的问题。迭代式参数更新主要解决的是参数和优化器状态的内存问题。两者可同时使用，分别应对数据和参数两方面的内存挑战。 与 混合精度训练 ： 迭代式参数更新常与混合精度训练（使用FP16/BF16计算，用FP32存储主参数副本）紧密结合，因为后者本身也能大幅节省内存，两者协同进一步压榨硬件潜力。