邝炜瑞, Weirui Kuang's Homepage.

Multi-machine multi-card Training

Cover Image for Multi-machine multi-card Training
邝炜瑞, Weirui Kuang
邝炜瑞, Weirui Kuang

多机多卡训练原理学习笔记

流水线并行 - Pipeline Parallelism

流水线并行的核心思想是:在朴素模型并行的基础上,进一步引入数据并行的办法,即把原先的数据再划分成若干个batch,送入GPU进行训练。micro-batch

数据并行 - Data Parallelism

  • DP
    • 一般单机多卡
    • GPU上的模型参数保持一致
    • 聚合再下发梯度的操作,称为AllReduce
    • 梯度异步更新 - 解决通讯开销大
  • DDP
    • 多机场景 - DP中Server上的通讯压力均衡转到各个Worker
    • Ring AllReduce
  • ZeRO 优化器状态/梯度/参数 三个层面的用完即抛,需要即补 (模型并行的形式,数据并行的实质。)
    • ZeRO的思想就是用通讯换显存
    • ZeRO-3 数据并行+张量并行
    • ZeRO-1:将优化器状态(例如,对于Adam优化器,32位权重,第一和第二矩估计)分割在多个进程中,使每个进程只更新自己的部分。
    • ZeRO-2:将32位的梯度也分割,使每个进程只保留与自己的优化器状态对应的梯度。
    • ZeRO-3:将参数也分割,并使用CPU或NVMe来卸载部分参数,以进一步减少GPU内存占用。
  • FSDP
    • Pytorch 的FSDP是一种数据并行的训练方法,它实际上就是ZeRO-3

模型张量并行 - Tensor Parallesim

对GPU的利用度不高

当K越大,即GPU的数量越多时,空置的比例接近1,即GPU的资源都被浪费掉了

  • Megatron-LM
    • 按行/列切分W
    • 行:
    • 列:

数据并行+张量并行

  • TP + DP (Megatron + ZeRO)

参考资料

猛猿:图解大模型训练之:流水线并行(Pipeline Parallelism),以Gpipe为例

猛猿:图解大模型训练之:数据并行上篇(DP, DDP与ZeRO)

猛猿:图解大模型训练之:数据并行下篇(ZeRO,零冗余优化)

猛猿:图解大模型系列之:张量模型并行,Megatron-LM

猛猿:图解大模型系列之:Megatron源码解读1,分布式环境初始化