DeepSpeed框架：大纲和资料梳理

最新推荐文章于 2024-08-26 20:33:29 发布

zenRRan

最新推荐文章于 2024-08-26 20:33:29 发布

阅读量332

点赞数

原文链接：https://mp.weixin.qq.com/s?__biz=MzI3ODgwODA2MA==&mid=2247522978&idx=5&sn=28a2784d483bf63448e0288a5283f489&chksm=eb538e31dc240727305dcdd7babfa1b23fee6ee8753b0110289ea0a75fa45e9c26fbc16a1a94&scene=126&sessionid=0

版权

来自：NLP工程化

进NLP群—>加入NLP交流群

DeepSpeed是一个深度学习优化软件套件，使分布式训练和推理变得简单、高效和有效。它可以做些什么呢？训练/推理具有数十亿或数万亿参数的密集或稀疏模型；实现出色的系统吞吐量并有效扩展到数千个GPU；在资源受限的GPU系统上进行训练/推理；实现前所未有的低延迟和高吞吐量的推理；以低成本实现极限压缩，实现无与伦比的推理延迟和模型尺寸减小。特别说明，DeepSpeed在Windows上仅支持推理，不支持训练。

一.DeepSpeed四大创新支柱
1.DeepSpeed-Training
DeepSpeed提供了系统创新的融合，使大规模深度学习训练变得有效、高效，大大提高了易用性，并在可能的规模方面重新定义了深度学习训练版图。ZeRO、3D-Parallelism、DeepSpeed-MoE、ZeRO-Infinity等创新属于培训支柱[2]。
2.DeepSpeed-Inference
DeepSpeed汇集了tensor、pipeline、expert和ZeRO-parallelism等并行技术的创新，并将它们与高性能定制推理内核、通信优化和异构内存技术相结合，以前所未有的规模实现推理，同时实现无与伦比的延迟、吞吐量和性能。降低成本。这种推理系统技术的系统组合属于推理支柱[3]。
3.DeepSpeed-Compression
为了进一步提高推理效率，DeepSpeed为研究人员和从业人员提供易于使用且组合灵活的压缩技术来压缩他们的模型，同时提供更快的速度、更小的模型大小并显着降低的压缩成本。此外，ZeroQuant和XTC等SoTA在压缩方面的创新也包含在压缩支柱下[4]。
4.DeepSpeed4Science
觉得离自己好遥远，就不介绍这个了[5]。
5.DeepSpeed软件架构
（1）APIs：配置参数都在ds_config.json文件中，上层通过简单的API接口就可以训练模型和推断。
（2）RunTime：DeepSpeed的核心运行时组件，使用Python语言实现，负责管理、执行和优化性能。它承担了将训练任务部署到分布式设备的功能，包括数据分区、模型分区、系统优化、微调、故障检测以及检查点的保存和加载等任务。
（3）DeepSpeed的底层内核组件，使用C++和CUDA实现。它优化计算和通信过程，提供了一系列底层操作等。

二.分布式模型训练
分布式模型训练包括：数据并行，模型并行（张量并行和流水线并行），混合并行。如下所示：
1.数据并行
在数据并行系统中，每个计算设备都有整个神经网络模型的完整副本（Model Replica），进行迭代时，每个计算设备只分配了一个批次数据样本的子集，并根据该批次样本子集的数据进行网络模型的前向计算。如下所示：
2.模型并行
模型并行(Model Parallelism)往往用于解决单节点内存不足的问题。模型并行可以从计算图角度，以下两种形式进行切分：按模型的层切分到不同设备，即层间并行或算子间并行(Inter-operator Parallelism)，也称之为流水线并行(Pipeline Parallelism，PP)；将计算图层内的参数切分到不同设备，即层内并行或算子内并行(Intra-operator Parallelism)，也称之为张量并行(Tensor Parallelism，TP)。如下所示：
3.混合并行
混合并行(Hybrid Parallelism，HP)是将多种并行策略如数据并行、流水线并行和张量并行等进行混合使用。通过结合不同的并行策略，混合并行可以充分发挥各种并行策略的优点，以最大程度地提高计算性能和效率。如下所示：

4.DeepSpeed 3D并行策略示意图
图中给出了包含32个计算设备进行3D并行的例子。神经网络的各层分为4个流水线阶段。每个流水线阶段中的层在4个张量并行计算设备之间进一步划分。最后，每个流水线阶段有两个数据并行实例，使用ZeRO内存优化在这2个副本之间划分优化器状态量。如下所示：

5.典型HPC集群硬件组成
整个计算集群包含大量带有计算加速设备的服务器。每个服务器中往往有多个计算加速设备(通常2-16个)。多个服务器会被放置在一个机柜(Rack)中，服务器通过架顶交换机(Top of Rack Switch，ToR)连接网络。在架顶交换机满载的情况下，可以通过在架顶交换机间增加骨干交换机(Spine Switch)进一步接入新的机柜。这种连接服务器的拓扑结构往往是一个多层树(Multi-Level Tree)。如下所示：

三.如何存储参数
1.Model States
解析：模型本身相关且必须存储的内容，如下所示：
（1）Parameters：模型参数
（2）Gradients：模型梯度
（3）Optimizer States：Adam优化算法中的momentum和variance
2.Residual States
解析：非模型本身必须，但在训练过程中产生的内容，如下所示：
（1）Activation：激活值
（2）Temporary Buffers：临时存储
（3）Unusable Fragmented Memory：碎片化存储空间
3.混合精度训练
解析：混合精度训练就是一部分参数使用FP16（4B）存储，另一部分参数使用FP16或BF16（2B）存储，以此来减轻存储压力。FP32中第31位为符号位，第30到第23位用于表示指数，第22到第0位用于表示尾数。FP16中第15位为符号位，第14到第10位用于表示指数，第9到第用于表示尾数。BF16中第15位为符号位，第14到第7位用于表示指数，第6到第0位用于表示尾数。如下所示：
混合精度优化过程，如下所示：

四.如何优化储存
零冗余优化器（ZeRO）是一种用于大规模分布式深度学习的新颖内存优化技术，它可以大大减少模型和数据并行性所需的资源，同时可以大量增加可训练的参数数量。本质就是用完即弃，需要再补。如下所示：
ZeRO是一系列显存优化方法的统称，它分为ZeRO-DP（Zero Redundancy Optimizer-Data Parallel）和ZeRO-R（Zero Redundancy Optimizer-Reduce）两部分。如下所示：
1.ZeRO-DP
（1）ZeRO-1
ZeRO-1对优化器状态都进行分片，占用内存为原始的1/4，通信容量与数据并行性相同。
（2）ZeRO-2
ZeRO-2对优化器状态和梯度都进行分片，占用内存为原始的1/8，通信容量与数据并行性相同。
（3）ZeRO-3
ZeRO-3对优化器状态、梯度以及模型参数都进行分片，内存减少与数据并行度和复杂度成线性关系，同时通信容量是数据并行性的1.5倍。
论文：ZeRO：Memory Optimizations Toward Training Trillion Parameter Models：https://arxiv.org/pdf/1910.02054.pdf

2.ZeRO-R
（1）激活检查点(Partitioned Activation Checkpointing)
activation的存储非常灵活，更需要设计好实验。
（2）临时缓冲区(Constant Size Buffer)
固定大小的内存Buffer。
（3）空间管理(Memory Defragmentation)
主要是将碎片化内存空间重新整合为连续存储空间。

3.ZeRO-Offload
本质就是显存不够，内存来凑。ZeRO-Offload分为Offload Strategy和Offload Schedule两部分，前者解决如何在GPU和CPU间划分模型的问题，后者解决如何调度计算和通信的问题。
论文：ZeRO-Offload：Democratizing Billion-Scale Model Training：https://arxiv.org/pdf/2101.06840.pdf

4.ZeRO-Infinity
ZeRO-Infinity是ZeRO-3的拓展，允许通过使用NVMe固态硬盘扩展GPU和CPU内存来训练大型模型。
论文：ZeRO-Infinity：Breaking the GPU Memory Wall for Extreme Scale Deep Learning：https://arxiv.org/pdf/2104.07857.pdf

本文只是一个大纲框架，后续打算针对DeepSpeed和Megatron-LM框架写一个源码解读教程，以前总认为写软件没啥意思，现在发现分布式机器学习系统还是很有意思的，如果算法给力，直观可见的节省存储，提高训练效率，降低推理的延时和增加高吞吐量。

参考文献：
[1]DeepSpeed项目地址：https://github.com/microsoft/DeepSpeed
[2]DeepSpeed训练：https://www.deepspeed.ai/training/
[3]DeepSpeed推理：https://www.deepspeed.ai/inference
[4]DeepSpeed压缩：https://www.deepspeed.ai/compression
[5]DeepSpeed4Science网站和教程：https://deepspeed4science.ai/和https://www.deepspeed.ai/deepspeed4science/
[6]DeepSpeed入门：https://www.deepspeed.ai/getting-started/
[7]DeepSpeed JSON配置：https://www.deepspeed.ai/docs/config-json/
[8]DeepSpeed API文档：https://deepspeed.readthedocs.io/en/latest/
[9]DeepSpeed教程：https://www.deepspeed.ai/tutorials/
[10]DeepSpeed博客：https://www.deepspeed.ai/posts/
[11]deepspeed-mii：https://github.com/microsoft/deepspeed-mii
[12]Transformer DeepSpeed集成：https://huggingface.co/docs/transformers/main/main_classes/deepspeed
[13]DeepSpeed官方文档：https://deepspeed.readthedocs.io/en/latest/
[14]ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters：https://www.microsoft.com/en-us/research/blog/ZeRO-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/
[15]《大规模语言模型：从理论到实践》