Colossal Ai 调研

[thinksoso]

colossal ai 支持特性

• 数据并行
• 流水线并行
• 1维, 2维, 2.5维, 3维张量并行
• 序列并行
• 混合精度
• 零冗余优化器 (ZeRO)
• 友好的trainer和engine
• 可扩展新的并行方式

[thinksoso]

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采用了递增阶段的优化方式，在早期进行的优化可以被用在后续的优化阶段。

[thinksoso]

并行框架

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的组成：

1. Optimizer States Optimizer在进行梯度更新时所需要用到的数据，例如SGD中的Momentum以及使用混合精度训练时的Float32 Master Parameter。
2. Gradient：在反向传播后所产生的梯度信息，决定其更新方向
3. 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