DeepSpeed超大规模模型训练工具
2021年 2 月份发布了 DeepSpeed。这是一个开源深度学习训练优化库，包含的一个新的显存优化技术—— ZeRO（零冗余优化器），通过扩大规模，提升速度，控制成本，提升可用性，极大地推进了大模型训练能力。DeepSpeed 已经帮助研究人员，开发了图灵自然语言生成模型( Turing-NLG)，在发表时，为世界上最大的语言模型（拥有 170 亿参数），有着最佳的精度。在2021年 5 月份发布了 ZeRO-2——支持有着 2000 亿参数的模型训练，与最新技术相比，训练速度可达 10 倍——系列计算、IO 和收敛优化功能，从而助力最快速的 BERT 训练。自那时起，持续高速进行创新，不断突破深度学习模型训练的速度和规模的边界。
这些进展不仅会推动深度学习训练走向极致，同时也让这份技术的使用范围更加广泛——上至数据科学家们在超算上训练，下至在低端集群甚至仅仅一张 GPU 上训练。具体来说，DeepSpeed 加入了 4 项系统性新技术来进一步拓展 AI at Scale 倡议。推动了微软的AI产品与平台的创新。这些技术提供了极为高效的计算、显存和通信的利用效率，并助力训练有着十亿至万亿量级参数的模型。支持超长输入序列，无论在单卡GPU、千卡GPU的高端集群上，还是在慢速以太网的低端集群上，均可以使用。
- 用 3D 并行化实现万亿参数模型训练： DeepSpeed 实现了三种并行方法的灵活组合：ZeRO 支持的数据并行，流水线并行和张量切片模型并行。3D 并行性适应了不同工作负载的需求，支持具有万亿参数的超大型模型，实现了近乎完美的显存扩展性和吞吐量扩展效率。提高的通信效率，使用户可以在网络带宽有限的常规群集上，以 2-7 倍的速度训练有数十亿参数的模型。
- ZeRO-Offload 使 GPU 单卡，能够训练 10 倍大的模型： 为了同时利用 CPU 和 GPU 内存来训练大型模型，扩展了 ZeRO-2。用户在使用带有单张英伟达 V100 GPU 的机器时，可以在不耗尽显存的情况下，运行多达 130 亿个参数的模型，模型规模扩展至现有方法的10倍，保持有竞争力的吞吐量。此功能使数十亿参数的模型训练，更加大众化，，并为许多深度学习从业人员，打开了一扇探索更大更好的模型的窗户。
- 通过 DeepSpeed Sparse Attention，用6倍速度执行10倍长的序列： DeepSpeed提供了稀疏 attention kernel ——一种工具性技术，可支持长序列的模型输入，包括文本输入，图像输入和语音输入。与经典的稠密 Transformer 相比，支持的输入序列长一个数量级，在保持相当的精度下，获得最高 6 倍的执行速度提升。还比最新的稀疏实现快 1.5–3 倍。稀疏 kernel 灵活支持稀疏格式，使用户能够通过自定义稀疏结构进行创新。
- 1 比特 Adam 减少 5 倍通信量： Adam 是一个在大规模深度学习模型训练场景下的，有效的（也许是最广为应用的）优化器。与通信效率优化算法往往不兼容。在跨设备进行分布式扩展时，通信开销可能成为瓶颈。推出了一种 1 比特 Adam 新算法，高效实现。该算法最多可减少 5 倍通信量，实现了与Adam相似的收敛率。在通信受限的场景下，观察到分布式训练速度，提升了 3.5 倍，这使得该算法，可以扩展到不同类型的 GPU 群集和网络环境。
3D 并行：扩展至万亿参数模型
随着现代 GPU 群集上计算量的快速增长，训练具有惊人的功能的、强大的万亿参数模型，不再是遥不可及的，可能在不久的将来就能实现。DeepSpeed 结合了三项强大的技术，可以训练数万亿规模的模型，扩展到数千个 GPU：数据并行训练，模型并行训练和流水线并行训练。这三者的共生让深度学习训练的规模，远远超出了单独使用每种策略。3D 并行解决了训练万亿参数模型的两个基本挑战：显存效率和计算效率。DeepSpeed 可以扩展至在显存中，放下最巨大的模型，不会牺牲速度。
了解训练巨大模型的显存和计算效率的挑战
显存效率：训练万亿参数模型所需的显存，远远超出了单张 GPU 的显存大小。在使用 Adam 优化器，进行混合精度训练时，存储模型状态量（参数、梯度和优化器状态量）需要约 16TB 的显存。最先进的英伟达 A100 GPU 只有 40 GB 的显存。仅仅为了存储模型状态，需要 400 张这样的 GPU。
激活函数额外消耗的显存，随 batch 大小而增加。batch 设置为1的情况下，训练万亿参数模型，产生超过 1 TB 的激活函数用的显存（后文称为激活显存）。用 checkpoint 处理激活显存，用计算换显存，可以将显存减少到大约20 GB，对于训练而言，仍然过高了。
必须在多个 GPU 设备之间，有效地划分模型状态量和激活显存，才能让这种大模型，在不耗尽显存的情况下，开始训练。
计算效率：经估算端到端训练一个万亿参数的模型，大约需要 5000 Zflops（即 5 后面带有 24 个零；这个估算结果基于 OpenAI 的研究 law of scaling）。训练一个模型需要 4000 张 A100，以 50％ 的计算效率运行大约 100 天。
尽管大型超级计算 GPU 集群，可以拥有超过 4000 个 GPU，但是由于 batch 大小的限制，要在这种规模上实现高计算效率，仍然是一项挑战。计算效率随着计算时间，对通信时间的比例的增加而增加。该比例与 batch 大小成正比。训练模型的 batch 大小，有一个上限——超过这个上限，收敛情况会明显变差。
实际上最大的模型之一，GPT-3 的训练 batch 大小约 1500。如果使用大约 4000 张 GPU, 即使可以自由设置 batch 大小为 4000，每张卡上的 batch 大小，也只有 1，这将影响扩展性。
理解数据并行、模型并行和流水线并行之间的权衡
数据并行是深度学习中的一种普遍使用的技术。每批输入的训练数据，都在数据并行的 worker 之间平分。反向传播后，需要通信规约梯度，保证优化器在各个 worker 上进行相同的更新。数据并行性，具有几个明显的优势，包括计算效率高和实现起来工作量小。数据并行的 batch 大小，随 worker 数量提高，往往无法在不影响收敛性的情况下，一直增加 batch 大小。