DeepSpeed - HackMD

# DeepSpeed ## 目录 * 以下给出[DeepSpeed官方](https://www.deepspeed.ai/training/)所总结的training feature * Distributed Training with Mixed Precision * 16-bit mixed precision * Single-GPU/Multi-GPU/Multi-Node * Model Parallelism * Support for Custom Model Parallelism * Integration with Megatron-LM * Pipeline Parallelism * 3D Parallelism * The Zero Redundancy Optimizer * Optimizer State and Gradient Partitioning * Activation Partitioning * Constant Buffer Optimization * Contiguous Memory Optimization * ZeRO-Offload * Leverage both CPU/GPU memory for model training * Support 10B model training on a single GPU * Ultra-fast dense transformer kernels * Sparse attention * Memory- and compute-efficient sparse kernels * Support 10x long sequences than dense * Flexible support to different sparse structures * 1-bit Adam, 0/1 Adam and 1-bit LAMB * Custom communication collective * Up to 26x communication volume saving * Additional Memory and Bandwidth Optimizations * Smart Gradient Accumulation * Communication/Computation Overlap * Training Features * Simplified training API * Gradient Clipping * Automatic loss scaling with mixed precision * Training Optimizers * Fused Adam optimizer and arbitrary torch.optim.Optimizer * Memory bandwidth optimized FP16 Optimizer * Large Batch Training with LAMB Optimizer * Memory efficient Training with ZeRO Optimizer * CPU-Adam * Training Agnostic Checkpointing * Advanced Parameter Search * Learning Rate Range Test * 1Cycle Learning Rate Schedule * Simplified Data Loader * Data Efficiency * Efficient data sampling via curriculum learning and efficient data routing via random layerwise token dropping * Up to 2x data and 2x time saving during GPT-3/BERT pretraining and GPT/ViT finetuning * Or further improve model quality under the same data/time * Curriculum Learning * A curriculum learning-based data pipeline that presents easier or simpler examples earlier during training * Stable and 3.3x faster GPT-2 pre-training with 8x/4x larger batch size/learning rate while maintaining token-wise convergence speed * Complementary to many other DeepSpeed features * Note that the Data Efficiency Library above provides more general curriculum learning support. This legacy curriculum learning feature is still supported but we recommend to use the Data Efficiency Library. * Progressive Layer Dropping * Efficient and robust compressed training * Up to 2.5x convergence speedup for pre-training * Performance Analysis and Debugging * Mixture of Experts (MoE) * 以下给出[DeepSpeed官方](https://www.microsoft.com/en-us/research/blog/deepspeed-accelerating-large-scale-model-inference-and-training-via-system-optimizations-and-compression/)所总结的inference feature * Inference-adapted parallelism * Inference-optimized CUDA kernels * Effective quantize-aware training * 此外，DeepSpeed还支持模型压缩功能 ## Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism * 概述: 针对transformer结构设计模型切分方法，将切分后的模型分发到多张卡上并行加速 * 方法: * 切分MLP * $Y = GeLU(XA)$, $Z = YB$ * $A = \begin{bmatrix}A_1 & A_2\end{bmatrix} B=\begin{bmatrix}B_1 \\ B_2\end{bmatrix}$ * $\begin{bmatrix}Y_1 & Y_2\end{bmatrix} = [GeLU(XA_1), GeLU(XA_2)]$ * $Z = Y_1B_1 + Y_2B_2$ > 下标代表所在的GPU编号，最后计算Z时用all-reduce求和 * 切分self-attention * 利用self-attention中多头的特性，将每个头平均分配给多个卡即可 * | MLP切分方式 | self-attension切分方式 | | -------- | -------- | | ![](https://hackmd.io/_uploads/r1IZPI7Vh.png) | ![](https://hackmd.io/_uploads/S1b7D8XVh.png) | * 切分输出层 * 到输出层时，向量大小为$b$x$l$x$v$，其中$b$是batch-size, $l$是输入语句长度，$v$是单字表大小 * 因为$v$非常大，如果直接通信会造成通信量过大，先算softmax再传$b$x$l$结果即可(如果要entropy呢?) * 参考资料 * paper: https://arxiv.org/pdf/1909.08053.pdf * https://www.youtube.com/watch?v=mk8FuiDmf0I ## ZeRO: Memory Optimizations Toward Training Trillion Parameter Models * 背景一: 如何提速? 当拥有多张卡的时候，直觉上有以下几种方法提速 * 数据并行(data parallelism): 将整个batch的数据切分为多个min-batch，每张卡上跑一个mini-batch，计算结束后将各个mini-batch的梯度求平均即可得到整体的梯度 * 优点: 实现简单、计算与通信量低 * 缺点: 每个节点都要维护一份一模一样的模型，浪费显存 * 模型并行(model parallelism): 将模型的每一层都切分为多个tensor，每张卡上计算一部份，每计算完一层都需要进行整合 * 优点: 节省显存 * 缺点: 每计算完一层都需要跟其他卡沟通、通信量大。切分后的模型粒度过小，将影响计算速度。 * 流水并行(pipeline parallelism): 将模型按层切割，每张卡上计算其中n层，可以通过流水线的方式实现以加速 * 优点: 节省显存 * 缺点: * 为了使用流水线，通常会拆分为多个mini-batch，使得batch norm等操作困难(mini-batch间有时序关系) * 尽管使用流水线，仍然会使效率降低(pipeline bubble)，通常mini batch数量越多，效率降低的越不显著 > ZeRO 希望同时具有DP通信量低与MP/PP节省内存的优点 * 背景二: 如何降显存? 通常显存因为以下几种原因被占用 * 模型参数、梯度与优化器参数: * 若使用混合精度训练，并使用Adam优化器，且模型参数数量为$\Phi$，则该模型在训练时总共需要占用16$\Phi$ bytes的显存，包含: * 参数(fp16, $2\Phi$)、 * 梯度(fp16, $2\Phi$) * 优化器状态: 参数(fp32, $4\Phi$), Adam's momentum(fp32, $4\Phi$), Adam's variance(fp32, $4\Phi$) > 可以看到优化器状态是占用显存的主要因素，ZeRO-DP通过将优化器状态切割成多份后，分别记录在多张卡上达到降低显存占用的效果。 * 激活值(Activations): * ![](https://i.imgur.com/2JD0SKt.gif) > source gif: https://siboehm.com/articles/22/data-parallel-training * Activations即上图中的cache，在前向传播时被保留，用以计算反向传播 * Activation checkpointing: 可以适当丢弃部份activations，需要时重新前向计算即可，以增加33%计算量为代价能够节省8x的内存。 > ZeRO-R 通过将Activation checkpointing切割成多份后，分别记录在多张卡上达到降低显存占用的效果 * 缓存(Temporary buffers): * 模型在执行gradient all-reduce或gradient norm等操作时，可能会开启过大的缓存 > ZeRO-R 通过固定临时缓存大小，避免此问题 * 不连续的显存(Memory Fragmentation): * 尽管显存足够，若找不到连续且足够大的显存空间，也会有out of memory的问题 * 造成显存不连续的主要原因是不同向量生命周期不同，如activation checkpoint与梯度会被长期保留，而其余activation将在使用后被销毁 > ZeRO-R 通过开辟固定的显存空间来存放会被长期保留的向量，避免此问题 * 算法细节 ZeRO-DP(data parallelism) > ZeRO-DP是以DP为基础改进而得的算法，以下先介绍all-reduce一种集群通信的方法，接着介绍传统DP方法，最后介绍ZeRO-DP算法的两个变体 * all-reduce: ![](https://i.imgur.com/KU8914j.png) > source image: https://zhuanlan.zhihu.com/p/504957661 * 概述: 通过环形通信，将不同节点上的数相加，并将结果更新到所有结点 * 分为两个部份reduce-scatter与all-gather (节点数=n, 单笔数据$a_1,b_1...$大小=$\Phi$) * reduce-scatter: 每个节点每一步都加上自己右手边节点的数，重复n次即可 * all-gather: 将蓝色部份内容依次向左传递，重复n次即可 * reduce-scatter与all-gather每步传递与接收的总数据量都是$\Phi$ * 称all-reduce通信量为2$\Phi$ (实际上要传n-1次，可以使用流水线加速) * DP(data parallelism) * 每个节点都要维护一个一模一样的模型，将数据分为多个mini-batch * 每个节点计算出该mini-batch的梯度，$w_i, i=1,..,n$, 其显存占用大小为$\Phi$ * 用all-reduce将所有mini-batch的梯度求和，并在各自节点上更新模型 > 注意: 理论上reduce-scatter后，每个节点就可以得到求和后的总梯度，只是求和的顺序不同，但由于计算机上浮点数加法不满足交换率，为了确保所有结点模型更新后参数完全相同，因此还是需要进行all-gather步骤 [[1]](https://siboehm.com/articles/22/data-parallel-training) * ZeRO-DP ($P_{os+g}$) * 将模型梯度与优化器状态分成$n$份，分别记录在$n$个节点上 * 每个节点计算出自己mini-batch的梯度，用reduce-scatter把所有结点的梯度相加 * 更新模型参数，用all-gather将更新后参数传给其他节点 > 此处求和顺序不同没有关系，只要确保更新后所有结点模型参数相同即可 * 通信量与传统DP相同、显存占用变为: $2+\frac{2+4+4+4}{n}$ (n很大时为8x) * ZeRO-DP ($P_{os+g+p}$) * 将模型参数也分成$n$份，分别记录在$n$个节点上 * 前向传播:通过广播从其他节点取得其他部份的结果，计算自己负责的mini-batch * 反向传播:通过广播从其他节点取得其他部份的结果，计算自己负责的mini-batch * 得到梯度后用reduce-scatter把所有结点的梯度相加 * 直接更新自己那部份参数就可以 * 通信量是传统DP 1.5倍(前向、反向、reduce-scatter)且需要广播，显存占用变为$\frac{16}{n}$ * 算法细节 ZeRO-R: Optimizing Residual State Memory * 使用Activation checkpointing * 将Activation checkpointing切割成多份后，分别记录在多张卡(类似ZeRO-DP) * 固定临时缓存大小 * 开辟固定的显存空间来存放会被长期保留的向量(包含: activation checkpoints与梯度) > 问题: > 1. 并且不懂通信的细节，感觉通信过程中也会占用到很多的缓存，不知道为啥并没有讨论。 > 2. $P_{os+g+p}$要求节点能广播，理论上广播是环状传递效率的一半，这点应该纳入考虑。 > 3. $P_{os+g+p}$与MP的区别? > 4. Activation checkpoint是需要人工选择哪些参数被保留，才能够最大化效率吗，还是有标准的选择方法($P_p$中模型的切分方法同理)。 > 5. 流水线加速: 为了流水线切割为多个mini-batch会影响性能很多吗，要切割成多少个mini-batch对性能才不会有影响? * 参考资料 * paper: https://arxiv.org/pdf/1910.02054.pdf * https://basicv8vc.github.io/posts/zero/ * https://siboehm.com/articles/22/data-parallel-training * https://zhuanlan.zhihu.com/p/504957661 ## ZeRO-Offload Democratizing Billion-Scale Model Training * 概述: 将优化器状态与梯度传到CPU以降低显存占用 * 方法 * 模型流程设计 ![](https://hackmd.io/_uploads/Hkd_qq7Eh.png) * 设计思路 * 上左图是混合精度训练流程图，箭头上代表通信数据量 * $M$为模型参数量，$B$为batch-size * 首先，计算量来说前向与反向传播都是$O(MB)$，参数更新则是$O(M)$ * 因此，只要$B$不是太小，将参数更新放在CPU上运行不太影响效率 * 并且，前向与反向传播都必须在GPU上运行 * 需要在这两组操作中间切一刀，使得通信量最小，得到应该切过两个2M的箭头 * 考虑到CPU上放尽可能多的操作/数据，最终切分方式如上右图所示 > ![](https://hackmd.io/_uploads/Skt7acXNn.png) > 事实上，GPU与CPU之间通信所花的时间比想象来说少的多，因为计算反向传播需要时间，后面的层算完后该层梯度就可以开始先往CPU传了，参数更新时同理，如上图所示 * 多卡场景 * 当有n张卡时，与ZeRO-DP的$P_{os+g}$同理，可以将梯度与模型状态分为n等份，分别存在n张卡中，更进一步，offload到CPU上 * 优化CPU优化器 * 通过SIMD vector instruction, Loop unrolling, OMP multithreading优化，没有深入研究 * One-Step Delayed Parameter Update * 当$B$不是太大时，CPU上做参数更新可能影响性能 * 前N步正常训练，N+1步以后延迟一步更新 * 对性能存在负面影响，看起来不适合直接加入RL训练中 * 参考资料 * paper: https://arxiv.org/pdf/2101.06840.pdf * https://zhuanlan.zhihu.com/p/513571706 ## Ultra-fast dense transformer kernels * 参考资料: https://www.deepspeed.ai/2020/05/27/fastest-bert-training.html#a-advanced-fused-kernels-to-reduce-data-movement ## Sparse attention * 参考资料: https://www.deepspeed.ai/2020/09/08/sparse-attention.html ## 1-bit Adam * 概述: 通过压缩梯度的方式，降低通信开销 * 方法 * 以SGD为例: * $x_t$为$t$时刻模型参数，$g_t$为$t$时刻梯度，$e_t$为截至$t$时刻为止的累积误差，$C(\cdot)$为压缩函数 * 更新方式: $x_{t+1} = x_t - \gamma C[g_{t+1} + e_t]$ * 这样方法能够保证误差不会累积，不影响收敛性  * 参考资料 * paper: https://arxiv.org/pdf/2108.02102.pdf * https://www.deepspeed.ai/2020/09/08/onebit-adam-blog-post.html ## Inference * 概述 * 多卡并行将造成通信增加与计算粒度下降等问题 * 多卡并行效率通过模型并行解决 * 计算粒度下降通过kernel优化解决 * 需要用户自行选用合适并行程度 * 多卡并行推断 * DeepSpeed支持transformer类模型的自动切分(model parallelism) * 针对推断设计的kernel * Deep fusion: * 合并多个操作到同个kernel以加速 * 包含: 算子优化、针对Transformer结构的优化 * Inference-customized GeMM(General Matrix Multiplication): * batch-size过小，导致模型加载的时间远大于计算的时间 * DeepSpeed优化kernel，使其模型加载速率增大 * 当batch_size介于1到10时，效率高20%  ## ZeroQuant Efficient and Affordable Post-Training Quantization for Large-Scale Transformers * 方法 * 分组量化: * 大模型参数，同一行参数对应同一个输入特征维度，因此数值上接近 * 反之，不同行间数值差距极大，因此不能用一个固定范围来做量化 * 采用分组量化，将一层参数(共$n$行)分为$g$组，$n\geq g$ * Token-wise量化: * 对于不同的token，activation会很不一样 * 采用Token-wise量化，对于每组token都动态计算其范围 * kernel融合: * 量化与反量化操作耗时 * 通过kernel融合加速，如: BN与量化融合 * 参考资料 * paper: https://arxiv.org/pdf/2206.01861.pdf * https://mp.weixin.qq.com/s/iWEGUYOe32bqxxGohjqIhw