Optimum + ONNX Runtime: 更轻松、更快速地训练您的 Hugging Face 模型

简介

语言、视觉和语音领域基于 Transformer 的模型正变得越来越大，以支持面向最终用户的复杂多模态用例。模型大小的增加直接影响了训练这些模型以及随着模型规模增大而扩展它们所需的资源。Hugging Face 和微软的 ONNX Runtime 团队正在合作，致力于在微调大型语言、语音和视觉模型方面取得进展。Hugging Face 的 Optimum 库通过与 ONNX Runtime 训练的集成，为许多流行的 Hugging Face 模型提供了开放的解决方案，能够将训练时间缩短 35% 或更多。我们在此详细介绍 Hugging Face Optimum 和 ONNX Runtime 训练生态系统，并通过性能数据突显使用 Optimum 库的好处。

性能结果

下图显示，在训练时使用 ONNX Runtime 和 DeepSpeed ZeRO Stage 1 的情况下，Hugging Face 模型通过 Optimum 实现了 39% 到 130% 的惊人加速。性能测量是在选定的 Hugging Face 模型上进行的，以 PyTorch 作为基准运行，仅使用 ONNX Runtime 进行训练作为第二次运行，而 ONNX Runtime + DeepSpeed ZeRO Stage 1 作为最终运行，显示出最大的增益。基准 PyTorch 运行使用的优化器是 AdamW 优化器，而 ORT 训练运行则使用 Fused Adam 优化器。这些运行是在一个配备 8 个 GPU 的单个 Nvidia A100 节点上执行的。

有关开启 Optimum 以进行训练加速的配置设置的更多详细信息，请参见此处。这些运行使用的版本信息如下：

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 库

Hugging Face 是一个快速发展的开放社区和平台，旨在推动优秀机器学习的普及化。继 Transformers 库成功之后，我们为了满足社区的需求，将模态从 NLP 扩展到音频和视觉，现在已覆盖机器学习的各种用例。目前在 Hugging Face Hub 上，有超过 12 万个免费且可访问的各种机器学习任务的模型检查点，1.8 万个数据集，以及 2 万个机器学习演示应用。然而，将 Transformer 模型规模化地应用于生产环境仍然是业界的一大挑战。尽管精度很高，但基于 Transformer 的模型的训练和推理可能耗时且昂贵。

为了满足这些需求，Hugging Face 构建了两个开源库：Accelerate 和 Optimum。其中，🤗 Accelerate 专注于开箱即用的分布式训练，而 🤗 Optimum 作为 Transformers 的扩展，通过最大化利用用户目标硬件的效率来加速模型训练和推理。Optimum 集成了 ONNX Runtime 等机器学习加速器以及 Intel 的 Habana Gaudi 等专用硬件，因此用户可以在训练和推理中获得显著的速度提升。此外，Optimum 在继承与 Transformers 相同易用性的同时，无缝集成了 Hugging Face 的其他工具。开发者可以轻松调整他们的工作，以更少的计算能力实现更低的延迟。

ONNX Runtime 训练

ONNX Runtime 加速大型模型训练，对于流行的 HuggingFace Transformer 模型，单独使用可将吞吐量提高达 40%，与 DeepSpeed 结合使用时可提高达 130%。ONNX Runtime 已经集成到 Optimum 中，通过 Hugging Face 的 Optimum 训练框架实现更快的训练。

ONNX Runtime 训练通过多种内存和计算优化实现了如此高的吞吐量提升。内存优化使 ONNX Runtime 能够最大化批处理大小并高效利用可用内存，而计算优化则加快了训练时间。这些优化包括但不限于高效的内存规划、内核优化、用于 Adam 优化器的多张量应用 (multi tensor apply) (它将应用于模型所有参数的逐元素更新批量化为一次或几次内核启动)、FP16 优化器 (它消除了大量设备到主机的内存复制)、混合精度训练以及图优化，如图节点融合和节点消除。ONNX Runtime 训练支持 NVIDIA 和 AMD GPU，并提供自定义算子的可扩展性。

简而言之，它使 AI 开发者能够充分利用他们熟悉的生态系统，如 PyTorch 和 Hugging Face，并在他们选择的目标设备上使用 ONNX Runtime 的加速功能，从而节省时间和资源。

Optimum 中的 ONNX Runtime 训练

Optimum 提供了一个 ORTTrainer API，它扩展了 Transformers 中的 Trainer，使用 ONNX Runtime 作为加速后端。ORTTrainer 是一个易于使用的 API，包含功能完备的训练循环和评估循环。它支持超参数搜索、混合精度训练和多 GPU 分布式训练等功能。ORTTrainer 使 AI 开发者能够在训练 Transformers 模型时组合 ONNX Runtime 和其他第三方加速技术，这有助于进一步加速训练并充分利用硬件性能。例如，开发者可以将 ONNX Runtime 训练与 Transformers 的 Trainer 中集成的分布式数据并行和混合精度训练结合起来。此外，ORTTrainer 使得将 ONNX Runtime 训练与 DeepSpeed ZeRO-1 结合变得容易，后者通过分区优化器状态来节省内存。预训练或微调完成后，开发者既可以保存训练好的 PyTorch 模型，也可以使用 Optimum 为 ONNX Runtime 实现的 API 将其转换为 ONNX 格式，以便于推理部署。就像 Trainer 一样，ORTTrainer 与 Hugging Face Hub 完全集成：训练结束后，用户可以将其模型检查点上传到他们的 Hugging Face Hub 帐户。

那么具体来说，用户应该如何使用 Optimum 来利用 ONNX Runtime 的训练加速功能呢？如果您已经在使用 Trainer，您只需修改几行代码即可享受到上述所有改进。主要需要进行两处替换。首先，将 Trainer 替换为 ORTTrainer，然后将 TrainingArguments 替换为 ORTTrainingArguments，后者包含了训练器用于训练和评估的所有超参数。ORTTrainingArguments 扩展了 TrainingArguments，以应用一些由 ONNX Runtime 赋能的额外参数。例如，用户可以应用 Fused Adam 优化器以获得额外的性能提升。以下是一个示例：

-from transformers import Trainer, TrainingArguments
+from optimum.onnxruntime import ORTTrainer, ORTTrainingArguments

# Step 1: Define training arguments
-training_args = TrainingArguments(
+training_args = ORTTrainingArguments(
    output_dir="path/to/save/folder/",
-   optim = "adamw_hf",
+   optim = "adamw_ort_fused",
    ...
)

# Step 2: Create your ONNX Runtime Trainer
-trainer = Trainer(
+trainer = ORTTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
+   feature="sequence-classification",
    ...
)

# Step 3: Use ONNX Runtime for training!🤗
trainer.train()

展望未来

Hugging Face 团队正在努力开源更多大型模型，并通过训练和推理的加速工具降低用户从中受益的门槛。我们正与 ONNX Runtime 训练团队合作，为更新、更大的模型架构 (包括 Whisper 和 Stable Diffusion) 带来更多的训练优化。微软还将其最先进的训练加速技术打包到了 Azure Container for PyTorch 中。这是一个轻量级的精选环境，包括 DeepSpeed 和 ONNX Runtime，旨在提高 AI 开发者使用 PyTorch 进行训练的生产力。除了大型模型训练，ONNX Runtime 训练团队还在为边缘学习构建新的解决方案——在内存和功耗受限的设备上进行训练。

开始使用

我们邀请您查看以下链接，以了解更多关于 Optimum ONNX Runtime 训练的信息，并开始为您的 Hugging Face 模型使用它。

• Optimum ONNX Runtime 训练文档
• Optimum ONNX Runtime 训练示例
• Optimum Github 仓库
• ONNX Runtime 训练示例
• ONNX Runtime 训练 Github 仓库
• ONNX Runtime
• DeepSpeed 和 ZeRO 教程
• Azure Container for PyTorch

---

🏎感谢阅读！如果您有任何问题，欢迎通过 Github 或论坛与我们联系。您也可以在 Twitter 或 LinkedIn 上与我联系。