补全llm知识体系的地基：并行计算、ZeRO

并行计算/分布式

• 充分运用多机、多卡的算力，来组织大模型权重的加载，模型的训练、数据的分配、梯度的更新等
• 核心思想：抓准资源瓶颈，进行合适的调度

NCCL与通信原语

• NCCL是nvidia同机多卡间通信的协议，包含几个基本的操作，称为原语

• Broadcast：将一份数据复制到所有显卡上

• reduce：对所有显卡上的数据求和，存到某一个显卡上

• all gather：把分布在不同显卡上的数据concate，保存在所有显卡上

• reduce scatter：对所有显卡上的数据求和，结果分块存放在每个显卡上

• all reduce：对所有显卡上的数据求和，每个显卡上都有求和结果

  • （等价于reduce scatter+all gather）
  • Pytorch DDP、DeepSpeed和Megatron-LM都采用ring-All-reduce，每个节点只和相邻节点通信，能够防止主节点瓶颈
  • 实现思路：假设只有最基本的操作：a = a + b和a = b，共n张显卡参与
    • reduce scatter阶段：每个数据块向相邻显卡对应数据块规约，重复n-1次，每张显卡将拥有一个完全规约的数据块
    • all gather阶段：每个数据块向相邻显卡复制，重复n-1次，完成全部广播

详细参考：https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html

主流方式：

• 数据并行：
  • 在不同的显卡上加载同样的模型权重（broadcast）
  • 将一个batch数据分块(minibatch)加载到不同显卡上
  • 计算出梯度之后all reduce，此时每张卡上都有全量的梯度更新方向
  • 每张卡独立进行参数更新
• 模型并行：一个模型/tensor占据了过大的显存，故把模型垂直切割成多份，每张卡都有一层的一部分
• 流水线并行：将一个模型的不同layer放到不同的卡上，同一批数据在不同计算阶段在不同的卡之间以激活值的形式流转。这一方法通信量较少，但是对于空间占用、模型分割和任务调度则需要很高的适应性

ZeRO：Zero Redundancy Optimizer

• 在数据并行的基础上，对显存占用进行深度优化，减少冗余
• 显存占用主要包括：
  • model state：
    • 模型参数：
    • 梯度
    • 优化器
  • activation：
    • 传入参数
    • hidden state

ZeRO 1：设法节省Optimizer state

• 模型训练主要包含子阶段：前向传播、反向传播和模型更新。其中，前向传播和反向传播阶段并不需要优化器参与，只有模型更新阶段需要利用参数和优化器计算出新的参数
• 每个进程独立保留一部分优化器状态（这依赖于这个进程维护着哪部分参数）
• 在完成反向传播时，将梯度/动量/二阶动量等更新到对应部分的优化器状态中
• 完成对应部分的参数更新
• 所有进程之间进行all gather，将所有部分更新后的参数连接形成一个完整的参数更新

ZeRO 2：切分gradient

• 在ZeRO 1的基础上，要利用一部分优化器更新一部分参数，只需要这一部分参数对应的梯度就可以了。其他梯度完全可以被释放
• 在数据并行架构中，每个进程计算了自己mini batch数据对应的全部梯度
• 进行all reduce，梯度聚合后分发到所有进程中，各进程保留自己所需部分即可

ZeRO 3：切分model Parameter

• 每个layer都被切分成多块（切割方式：intra-layer）
• 每次计算之前，依靠广播获得需要的参数，依靠all-gather来聚合结果
• 主要思路是依靠更多的通信来减少同一份参数的冗余，即communcation和space的trade-off