Hugging Face Transformers 通过 DeepSpeed 和 FairScale 实现 ZeRO 集成
Hugging Face Transformers v4.2.0 引入了对 DeepSpeed 和 FairScale 的实验性支持，通过集成零冗余优化器 (ZeRO) 来允许用户训练更大的模型并增加批次大小 (batch size)，从而优化 GPU 显存使用情况。这种集成允许用户通过 Trainer API，使用 --sharded_ddp (FairScale) 和 --deepspeed (DeepSpeed) 命令行参数直接使用这些优化功能。
用于多 GPU 训练的 ZeRO 显存优化
通过 DeepSpeed 或 FairScale 集成 ZeRO，与标准的分布式数据并行 (DDP) 基准相比，可以显著减少训练和评估时间，同时增加最大可能的批次大小 (BS)。在两块 24GB Titan RTX GPU 上使用 t5-large 模型进行基准测试时，观察到了以下性能提升：
| Method | Max Batch Size | Train Time | Eval Time |
|---|---|---|---|
| Baseline (DDP) | 16 | 30.9458 | 56.3310 |
| fp16 | 20 | 21.4943 | 53.4675 |
| sharded_ddp (FairScale) | 30 | 25.9085 | 47.5589 |
| sharded_ddp + fp16 | 30 | 17.3838 | 45.6593 |
| DeepSpeed (no CPU offload) | 40 | 10.4007 | 34.9289 |
| DeepSpeed (with CPU offload) | 50 | 20.9706 | 32.1409 |
DeepSpeed 在批次大小和训练速度方面表现出最高的增益，而 FairScale 则因其仅需单个命令行参数而无需配置文件，被认为更容易部署。
用于大模型的单 GPU 训练
DeepSpeed 通过 CPU offloading 实现了对原本会超出单个 GPU 显存容量的模型进行训练。在单块 24GB RTX-3090 显卡上使用 t5-3b 模型进行测试时，标准的单 GPU 设置即使在批次大小为 1 时也会因显存溢出 (OOM) 错误而失败。通过使用 DeepSpeed，该模型在批次大小为 20 时成功完成了训练，且系统仅在批次大小为 30 时才触发 OOM。
ZeRO 的技术原理
ZeRO (Zero Redundancy Optimizer) 通过在数据并行训练中引入分布式数据存储来优化显存。ZeRO 不会在所有 GPU 上复制整个模型状态，而是将参数、梯度和优化器状态划分到可用的 GPU 之间。
分布式分区
每个 GPU 仅存储参数、梯度和优化器状态的一个分片 (shard)。在运行时，每个 GPU 会实时从其他参与的 GPU 获取特定层所需的必要数据，从而确保数据存储零重叠。
ZeRO-Offload
ZeRO-Offload 将特定的处理和显存需求从 GPU 转移到宿主 CPU，这对于将海量模型（如 t5-3b）适配到有限的 GPU 硬件上至关重要。
显存碎片管理
DeepSpeed 通过独立管理 GPU 显存来分离长期和短期分配，从而解决 GPU 显存碎片化问题——即尽管总显存可用，但由于没有足够大的连续块，导致出现 OOM 错误。
部署与集成
ZeRO 优化不需要对模型架构进行修改；只需更改训练代码即可。Hugging Face Trainer 的用户可以使用以下标志来实施这些功能：
--sharded_ddp: 激活 FairScale 集成。--deepspeed: 激活 DeepSpeed 集成（需要一个 JSON 配置文件）。
提到的未来增强功能包括 DeepSpeed Sparse Attention、1-bit Adam，以及预计将在 FairScale 和 DeepSpeed 中都支持的模型参数分片功能。