Oneflow-Inc
Colossal Ai 调研
colossal ai 支持特性
- 数据并行
- 流水线并行
- 1维, 2维, 2.5维, 3维张量并行
- 序列并行
- 混合精度
- 零冗余优化器 (ZeRO)
- 友好的trainer和engine
- 可扩展新的并行方式
Colossal-AI 架构
多重并行
目前主流的AI并行解决方案，如微软的DeepSpeed和英伟达的威震天，都使用3D并行，即数据并行、流水并行和1D张量并行。在数据和流水线并行的兼容下，Colossal-AI 进一步发展了 2D 张量并行、3D 张量并行和 2.5D 张量并行。此外，对于大图、视频、长文本等数据，采用了序列并行，将长序列分成子序列分布在一系列设备上。因此，Colossal-AI 可以将并行度从目前最高的 3 维提升到 5 维甚至 6 维。
大规模优化器
Colossal-AI使用大规模优化器，如LAMB和LARS，将batch size从512扩大到65536，在保证精度的同时大大缩短了模型训练时间。
自适应任务调度
不光根据gpu数量来调度任务，也考虑了batch_size等因素，实现弹性调度。
零冗余优化器
加入了DeepSpeed的ZeRO。ZeRO通过聚合数据并行的计算和内存资源，减少模型训练中的每个设备(GPU)的内存和计算需求。ZeRO通过在分布式训练的可用设备(GPU和CPU)上划分各种模型训练状态(权重、梯度和优化器状态)，减少了每个GPU的内存消耗。 ZeRO采用了递增阶段的优化方式，在早期进行的优化可以被用在后续的优化阶段。
并行框架
Megatron
特点：
- 层内模型并行 1D tensor parallelism 目的： 减少GPU占用和通信
Model Parallel Transformers
transformer layer = self-attention + mlp
MLP的并行
MLP 矩阵乘法(GEMM)加上非线性层 $Y=GeLU(XA)$
并行的一种方式：把权重矩阵A横向划分，把输入张量纵向划分
$X=[X_1,X_2], A=[A_1,A_2]^T$
因为GeLU是非线性变换，$GeLU(a+b)!=GeLU(a)+GeLU(b)$ 所以需要同步
并行的另一种方式: 把权重矩阵A纵向划分
$A=[A_1, A_2]$
好处是不需要同步
所以最优的方式是，对一个两层的MLP，把第一层的权重矩阵按列划分，把第二层的权重矩阵按行划分(只有第一层需要激活函数)，这样整个过程不需要同步
第一层 $[Y1,Y2]=[GeLU(XA_1),GeLU(XA_2)]$
第二层 $[Z_1,Z_2] = [Y_1,Y_2] *[A_1, A_2]^T = [Drop(Y_1A_1),Drop(Y_2A_2)]$
只有两个地方需要reduce，一个是输入dropout层之前，一个是反向传播时的f的函数
class f(torch.autograd.Function):
def forward(ctx, x):
return x
def backward(ctx, gradient):
all_reduce(gradient)
return gradient
相对应的，g在forward时候做all_reduce
self attention 并行
把Q，K，V按列划分，使得一个head的操作都在一个GPU，
这样一个transformers的层只要四次all reduce操作，前向两次，后向两次
DeepSpeed
- 3D 并行
- ZeRO
- 在每个GPU上保存一部分优化器状态量和梯度 对optimizer 阶段的重构，从原来上来就allreduce，各个节点都保持一份weight, gradient以及adam 需要mean, var等，变成只在一个节点保存，通信变成了reduce和broadcast；2，各个层次的参数生命周期不重叠，每个节点仅保存计算时需要的参数。
- ZeRO-offload 把优化器状态量和梯度移到内存
- sparse transformers 稀疏注意力
3D 并行
https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/
数据并行+模型并行+pipeline并行
pipeline并行
Gpipe https://arxiv.org/pdf/1811.06965.pdf
We first split a mini-batch of training examples into smaller micro-batches, then pipeline the execution of each set of micro-batches over cells. We apply synchronous mini-batch gradient descent for training, where gradients are accumulated across all micro-batches in a mini-batch and applied at the end of a mini-batch
ZeRO
https://zhuanlan.zhihu.com/p/394064174 https://www.youtube.com/watch?v=czgA-MbAdvA&t=2550s 大规模训练中的显存占用可以分为 Model States 与 Activation 两部分，而 ZeRO 就是为了解决 Model States 而诞生的一项技术。
Model States的组成：
- Optimizer States Optimizer在进行梯度更新时所需要用到的数据，例如SGD中的Momentum以及使用混合精度训练时的Float32 Master Parameter。
- Gradient：在反向传播后所产生的梯度信息，决定其更新方向
- Model Parameter 模型参数，通过数据“学习”的东西
数据并行，所有的model state dict都需要保存一份
计算参数量对应的内存
a*16，a是aB参数，B是十亿
ZeRO 则在数据并行的基础上，引入了对冗余Model States的优化。使用 ZeRO 后，各个进程之后只保存完整状态的1/GPUs，互不重叠，不再存在冗余。在本文中，我们就以这个 7.5B 参数量的模型为例，量化各个级别的 ZeRO 对于内存的优化表现。
ZeRO 有三个不同级别，分别对应对 Model States 不同程度的分割 (Paritition)： - ZeRO-1：分割Optimizer States； - ZeRO-2：分割Optimizer States与Gradients； - ZeRO-3：分割Optimizer States、Gradients与Parameters；
stage one 在不同数据上计算梯度，平均梯度
稀疏注意力
https://spaces.ac.cn/archives/6853 注意力矩阵是$O(n^2)$的，需要把注意力矩阵变成稀疏的来减少内存
https://zhuanlan.zhihu.com/p/34534763 稀疏矩阵的乘法 https://zhuanlan.zhihu.com/p/342711915
Atrous Self Attention 膨胀注意力，要求每个元素只跟它距离为k,2k,3k,...的元素关联 Local Self Attention 局部注意力，要求每个元素只与前后k个元素以及自身有关 Sparse Attention = Atrous + Local
colossal ai
增加了高维度并行，并综合了deepspeed