Hugging Face Optimum 与 ONNX Runtime 训练集成

Hugging Face 与 Microsoft 将 ONNX Runtime 集成到 Optimum 库中，提供了一种开放的解决方案，使许多流行的 Hugging Face 模型的训练时间提升 35% 以上。此合作旨在降低大规模语言、语音和视觉模型微调所需的计算资源和时间。

训练性能提升

将 ONNX Runtime 与 Optimum 库结合，可为 Hugging Face 模型带来显著加速。在单节点 Nvidia A100（8 GPU）上进行的基准测试显示，结合 DeepSpeed ZeRO Stage 1 使用 ONNX Runtime 时，吞吐量提升范围为 39% 到 130%。

这些基准测试的关键配置包括：

  • 基线：使用 AdamW 优化器的 PyTorch。
  • 加速：使用融合 Adam 优化器的 ONNX Runtime。
  • 最大增益：ONNX Runtime + DeepSpeed ZeRO Stage 1。

环境规格：

  • PyTorch: 1.14.0.dev20221103+cu116
  • ORT: 1.14.0.dev20221103001+cu116
  • DeepSpeed: 0.6.6
  • HuggingFace: 4.24.0.dev0
  • Optimum: 1.4.1.dev0
  • Cuda: 11.6.2

Optimum 库生态系统

Optimum 是 Transformers 库的扩展，旨在通过最大化目标硬件的效率来加速模型训练和推理。Accelerate 库侧重于分布式训练，而 Optimum 则集成了 ONNX Runtime 等机器学习加速器以及 Intel 的 Habana Gaudi 等专用硬件，以降低延迟并减少计算功耗需求。

ONNX Runtime 训练中的技术优化

ONNX Runtime（ORT）单独可实现最高约 40% 的吞吐量提升，结合 DeepSpeed 使用时可达 130%。这些增益来源于多项内存和计算优化：

  • 内存优化：高效的内存规划能够最大化批量大小并更好地利用可用内存。
  • 计算优化：
    • 内核优化：通用的执行速度提升。
    • Adam 优化器的多张量应用：将所有模型参数的逐元素更新批处理为少量内核调用。
    • FP16 优化器：减少设备到主机的内存拷贝。
    • 混合精度训练：提升速度并降低内存占用。
    • 图优化：包括节点融合和节点消除。

ONNX Runtime Training 支持 NVIDIA 与 AMD GPU，并允许使用自定义算子。

在 Optimum 中实现 ORTTrainer

Optimum 引入了 ORTTrainer API，扩展了 Transformers 的 Trainer 以使用 ONNX Runtime 作为后端。该 API 提供了功能完整的训练与评估循环，支持超参数搜索、混合精度训练以及多 GPU 分布式训练。

ORTTrainer 使开发者能够将 ONNX Runtime 与其他加速技术（如 Distributed Data Parallel（DDP）和 DeepSpeed ZeRO‑1——后者通过划分优化器状态来节省内存）组合使用。训练完成后，模型可以保存为 PyTorch 模型，或转换为 ONNX 格式以进行优化的推理部署。

从标准 Trainer 迁移到 ORTTrainer，开发者需要进行两项主要更改：

  1. 将 Trainer 替换为 ORTTrainer。
  2. 将 TrainingArguments 替换为 ORTTrainingArguments，以启用 ORT 特有的功能，例如 adamw_ort_fused 优化器。

未来路线图

Hugging Face 与 Microsoft 正在合作，将这些训练优化推广至更大的模型架构，包括 Stable Diffusion 与 Whisper。此外，Microsoft 已推出 Azure Container for PyTorch，这是一个包含 DeepSpeed 与 ONNX Runtime 的精选环境，旨在提升 PyTorch 开发者的生产力。与此同时，团队也在研发“边缘学习”解决方案，聚焦于在内存和功耗受限的设备上进行训练。

Sources