目录
主要参考:
分布式训练基础 – 数据并行,模型并行,流水线并行
DeepSpeed Zero 的各个 stage 介绍
针对Zero 的各个stage,这里有三个点需要额外再说一下:
- 各个stage,要实现将某一部分参数分配到不同GPU,需要一个 分区映射表,用于指明哪个GPU储存哪个变量的哪一部分,这个映射在整个训练过程是固定的,方便后续查找和更新。例如:stage 1 是把优化器参数分区储存了,那么在进行优化器更新步骤时,就需要先根据分区映射表在不同GPU上找到当前步骤需要的数据,再来完成相应计算。且在计算完成后需要做参数广播(必要时)。这也可以解释越是后面的stage,其速度越慢,因为需要更多额外的时间在参数分区查找和数据通信上面。
- 不难发现实际当中 activates 的内存占用往往比模型参数和梯度都要大得多,但 Zero 当中没用对这部分变量做分区存储,原因主要有两点:
- 模型参数,梯度等变量是与input无关的,即无论input是什么,batch size多大,需要用到的变量都是这些,因此可以认为这些变量在各个GPU上都保留一份是纯复制,是冗余的;但activates相关的变量于input强相关,输入长,batch size大,activates内存就大,各个GPU上由于输入的数据不同,其activates变量也不一样,不涉及重复存储的问题
\quad- activates内存优化的方法是梯度检查点(checkpointing)和重计算策略,而不是简单的分区存储。检查点技术通过在前向传播时只存储部分关键激活值(称为“检查点”),在反向传播时重新计算丢弃的部分,以此降低内存占用,并不会引入跨设备的数据传输延迟。
- 同样是将模型参数分区放在不同GPU上,Zero stage3 和 模型并行 还是有较大区别:
- 模型并行 是将不同layer的参数放在不同GPU上,一个mini-batch的数据进来,要依次在不同GPU上进行计算。每个 GPU 直接存储自己负责的完整layer,前向和反向传播过程中不需要对同一层参数做额外的跨设备聚合,因为每个层的数据都在本地。
\quad- Zero stage3 的模型参数分区,是对每个参数张量进行细粒度分片,也就是说,一个张量(无论它属于网络的哪个层)会被划分成多个部分,每个 GPU 只保存其中的一部分。由于每个 GPU 只持有部分参数(例如,一个权重矩阵被分成多份),在前向传播或反向传播时可能需要从其他 GPU 聚合(如通过 All-Gather 操作)各个分片来获得完整的参数信息,然后进行计算。更新后,同样需要将分片的数据更新到各自所在的 GPU。这里使用的映射表记录了每个参数分片的归属,从而协调数据传输和更新过程。
\quad- 模型并行 一般是当模型本身太大,单个 GPU 无法装下整个模型或中间激活,必须跨 GPU 分布计算。且每个 GPU 不仅存储自己负责的 layer 的模型参数,同时也负责该部分计算产生的中间激活值和对应的梯度,这是与 Zero stage3 不同的地方。而 Zero stage3 大类上还是数据并行,即模型前向和反向计算单卡可行,只不过由于数据并行复制导致内存浪费,通过分片消除冗余存储。