DeepSpeed Zero 解读

针对Zero 的各个stage，这里有三个点需要额外再说一下：

1. 各个stage，要实现将某一部分参数分配到不同GPU，需要一个 分区映射表，用于指明哪个GPU储存哪个变量的哪一部分，这个映射在整个训练过程是固定的，方便后续查找和更新。例如：stage 1 是把优化器参数分区储存了，那么在进行优化器更新步骤时，就需要先根据分区映射表在不同GPU上找到当前步骤需要的数据，再来完成相应计算。且在计算完成后需要做参数广播（必要时）。这也可以解释越是后面的stage，其速度越慢，因为需要更多额外的时间在参数分区查找和数据通信上面。
2. 不难发现实际当中 activates 的内存占用往往比模型参数和梯度都要大得多，但 Zero 当中没用对这部分变量做分区存储，原因主要有两点:

1. 模型参数，梯度等变量是与input无关的，即无论input是什么，batch size多大，需要用到的变量都是这些，因此可以认为这些变量在各个GPU上都保留一份是纯复制，是冗余的；但activates相关的变量于input强相关，输入长，batch size大，activates内存就大，各个GPU上由于输入的数据不同，其activates变量也不一样，不涉及重复存储的问题
   \quad
2. activates内存优化的方法是梯度检查点（checkpointing）和重计算策略，而不是简单的分区存储。检查点技术通过在前向传播时只存储部分关键激活值（称为“检查点”），在反向传播时重新计算丢弃的部分，以此降低内存占用，并不会引入跨设备的数据传输延迟。

3. 同样是将模型参数分区放在不同GPU上，Zero stage3 和 模型并行 还是有较大区别：

1. 模型并行 是将不同layer的参数放在不同GPU上，一个mini-batch的数据进来，要依次在不同GPU上进行计算。每个 GPU 直接存储自己负责的完整layer，前向和反向传播过程中不需要对同一层参数做额外的跨设备聚合，因为每个层的数据都在本地。
   \quad
2. Zero stage3 的模型参数分区，是对每个参数张量进行细粒度分片，也就是说，一个张量（无论它属于网络的哪个层）会被划分成多个部分，每个 GPU 只保存其中的一部分。由于每个 GPU 只持有部分参数（例如，一个权重矩阵被分成多份），在前向传播或反向传播时可能需要从其他 GPU 聚合（如通过 All-Gather 操作）各个分片来获得完整的参数信息，然后进行计算。更新后，同样需要将分片的数据更新到各自所在的 GPU。这里使用的映射表记录了每个参数分片的归属，从而协调数据传输和更新过程。
   \quad
3. 模型并行 一般是当模型本身太大，单个 GPU 无法装下整个模型或中间激活，必须跨 GPU 分布计算。且每个 GPU 不仅存储自己负责的 layer 的模型参数，同时也负责该部分计算产生的中间激活值和对应的梯度，这是与 Zero stage3 不同的地方。而 Zero stage3 大类上还是数据并行，即模型前向和反向计算单卡可行，只不过由于数据并行复制导致内存浪费，通过分片消除冗余存储。