Hugging Face Optimum 集成 ONNX Runtime 训练加速，提速 35% 以上

引言

基于 Transformer 的语言、视觉和语音模型正变得越来越大，以支持面向终端客户的复杂多模态用例。模型规模的增加直接影响训练这些模型所需的资源，并随着规模的增大而扩展。Hugging Face 和 Microsoft 的 ONNX Runtime 团队正在合作，推动大型语言、语音和视觉模型微调方面的进步。Hugging Face 的 Optimum 库通过与 ONNX Runtime 的集成进行训练，提供了一个开放解决方案，可将许多流行的 Hugging Face 模型的训练时间缩短 35% 或更多。我们介绍 Hugging Face Optimum 和 ONNX Runtime Training 生态系统的详细信息，并用性能数据突出使用 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 Training 运行使用 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 是一个快速发展的开放社区和平台，旨在让优秀的机器学习民主化。我们将模态从 NLP 扩展到音频和视觉，如今在 Transformers 库成功之后，已覆盖机器学习领域的各种用例，以满足社区的需求。现在在 Hugging Face Hub 上，有超过 12 万个免费且可访问的模型检查点，适用于各种机器学习任务，1.8 万个数据集，以及 2 万个机器学习演示应用。然而，将 Transformer 模型扩展到生产环境仍然是行业面临的挑战。尽管准确率高，但基于 Transformer 的模型的训练和推理可能耗时且昂贵。

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

ONNX Runtime Training

ONNX Runtime 可将大模型训练加速，独立运行时吞吐量最高提升 40%，与 DeepSpeed 结合使用时，对于流行的基于 HuggingFace transformer 的模型，吞吐量最高提升 130%。ONNX Runtime 已作为 Optimum 的一部分集成，并通过 Hugging Face 的 Optimum 训练框架实现更快的训练。

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

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

Optimum 中的 ONNX Runtime Training

Optimum 提供了一个 ORTTrainer API，扩展了 Transformers 中的 Trainer，以使用 ONNX Runtime 作为加速后端。ORTTrainer 是一个易于使用的 API，包含功能完整的训练循环和评估循环。它支持超参数搜索、混合精度训练和多 GPU 分布式训练等功能。ORTTrainer 使 AI 开发者能够在训练 Transformers 模型时组合 ONNX Runtime 和其他第三方加速技术，这有助于进一步加速训练并充分发挥硬件性能。例如，开发者可以将 ONNX Runtime Training 与 Transformers 的 Trainer 中集成的分布式数据并行和混合精度训练相结合。此外，ORTTrainer 使得将 ONNX Runtime Training 与 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 Optimizer 以获得额外的性能提升。以下是一个示例：

-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，旨在提高使用 PyTorch 进行训练的 AI 开发者的生产力。除了大型模型训练，ONNX Runtime 训练团队还在构建用于边缘学习的新解决方案——在内存和功耗受限的设备上进行训练。

开始使用

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

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