DeepSpeed ZeRo Stage的介绍

"时间是金"

Posted by 王川 on July 24, 2023

deepspeed中,zero 1/2/3 的区别是什么?

DeepSpeed是由微软开发的优化库,用于训练深度学习模型,旨在提高大规模分布式训练的性能和可扩展性。DeepSpeed的一个关键特性是Zero优化,专注于在分布式训练期间进行内存优化和减少通信开销。

DeepSpeed的Zero优化分为三个阶段:Zero1、Zero2和Zero3。每个阶段引入了额外的内存和通信优化。我们来详细介绍每个阶段:

  1. Zero1(ZeRO-Offload): Zero1是DeepSpeed中的第一阶段优化。在这个阶段,优化器状态和梯度被分割到数据并行(模型并行保持不变)。优化器状态保存在CPU内存中,而梯度放置在GPU内存中。这种方法有助于减少GPU内存消耗,与标准分布式训练相比。但是,对于具有非常大内存占用的模型来说,可能还不足够。

  2. Zero2(ZeRO-Offload + Offload Activations): Zero2在Zero1优化的基础上进一步发展,通过将一些模型激活值也从GPU内存转移到CPU内存来进行优化。在前向传播过程中,某些激活值被转移到CPU内存,以减少GPU内存使用。这允许训练更大的模型,否则这些模型的尺寸将超出GPU内存。然而,这也带来了GPU和CPU之间额外的通信开销。

  3. Zero3(ZeRO-Infinity): Zero3是DeepSpeed中Zero优化的最高阶段。它是Zero2的进化版本,旨在消除CPU内存对模型大小的限制。Zero3采用了”activation checkpointing”的方法,在反向传播过程中动态地重新计算某些中间激活值,避免在内存中存储它们。通过避免存储所有中间激活值,Zero3可以处理具有极大内存需求的模型,并最小化GPU内存消耗。这种技术特别适用于具有海量内存要求的模型,否则无法在单个GPU上训练。

总之,Zero1通过将优化器状态划分到CPU内存和梯度划分到GPU内存来减少GPU内存消耗。Zero2进一步优化,通过将一些模型激活值转移到CPU内存。Zero3则更进一步,引入激活值检查点技术,最小化GPU内存需求,实现训练极大模型的可能性。

关于优化器状态和梯度

在深度学习中,优化器(Optimizer)是用于调整模型参数以最小化损失函数的算法。在分布式训练中,由于涉及多个计算设备(通常是多个GPU),每个设备都需要访问模型的参数和梯度信息。这就引出了两个关键概念:优化器状态(Optimizer State)和梯度(Gradients)。

  1. 优化器状态(Optimizer State): 优化器状态是指在训练过程中,优化器内部维护的一些额外的参数或统计信息。这些信息通常用于跟踪模型的训练状态和帮助更新模型参数。例如,在随机梯度下降(SGD)优化器中,优化器状态可能包含学习率(learning rate)、动量(momentum)等超参数的当前值,以及每个参数的历史梯度信息。这些状态参数在每个训练步骤中都会更新,并在后续步骤中使用。

  2. 梯度(Gradients): 在深度学习中,梯度表示损失函数相对于模型参数的变化率。在训练过程中,我们通过将输入数据传递到模型中,计算出损失函数值,然后通过反向传播算法计算模型参数相对于损失函数的梯度。这些梯度用于更新模型参数,使得模型能够更好地拟合训练数据。

在Zero1优化中,优化器状态和梯度被划分到不同的内存区域:

  • 优化器状态(Optimizer State):它是指优化器内部的参数或统计信息,比如学习率和动量等。在Zero1中,这些状态信息被存储在CPU内存中。因为优化器状态通常较小,将它们放在CPU内存上不会占用太多的内存空间。

  • 梯度(Gradients):它是指模型参数相对于损失函数的变化率。在Zero1中,梯度被存储在GPU内存中。由于梯度的大小通常比较大,将它们放在GPU内存上有助于加快计算和更新模型参数的速度。

通过将优化器状态和梯度划分到不同的内存区域,Zero1可以显著减少GPU内存的占用,从而允许在相同的GPU内存容量下训练更大的模型或使用更大的批次大小,提高训练性能和效率。

关于ZeRo2中的激活值

在深度学习中,模型激活值是指神经网络中每个神经元在前向传播过程中产生的输出值。当输入数据通过神经网络的各层传递时,每个神经元都会进行加权和激活操作,然后将结果作为输出传递到下一层。这些输出值就是模型的激活值。

更具体地说,对于神经网络中的每个神经元,输入数据(来自前一层的输出或输入层的输入)将与该神经元的权重进行加权求和,然后通过激活函数进行处理,得到该神经元的输出值。这个输出值将成为下一层的输入,并继续传递下去。

常见的激活函数包括ReLU(Rectified Linear Unit)、Sigmoid、Tanh等。它们的作用是引入非线性性,使得神经网络能够学习复杂的非线性关系,从而提高模型的表达能力。

在Zero2优化中,模型激活值被引入了CPU内存的优化。在传统的分布式训练中,每个GPU都需要存储模型的完整激活值,这可能导致显存不足的问题,特别是在处理大型深度学习模型时。

为了解决这个问题,Zero2将部分模型激活值从GPU内存转移到CPU内存。具体来说,Zero2采用了一种技术称为”Offload Activations”,在前向传播过程中,某些激活值不再保留在GPU内存中,而是被临时存储在CPU内存中。当进行反向传播和参数更新时,需要这些激活值时,再从CPU内存中恢复。这样就有效地减少了GPU内存的使用量。

值得注意的是,模型激活值的转移并不会影响模型的训练过程,因为在反向传播时,所需的梯度信息仍然保留在GPU内存中,而模型激活值只是临时存储在CPU内存中,不会影响梯度计算和参数更新的正确性。通过这种方式,Zero2使得可以在有限的GPU内存下训练更大的模型,从而提高模型的性能和训练效率。