Hugging Face Optimum and ONNX Runtime Training Integration

Hugging Face and Microsoft have integrated ONNX Runtime into the Optimum library, providing an open solution that improves training times by 35% or more for many popular Hugging Face models. This collaboration aims to reduce the computational resources and time required to fine-tune large-scale language, speech, and vision models.

Training Performance Gains

Integrating ONNX Runtime with the Optimum library delivers significant acceleration for Hugging Face models. Benchmarks conducted on a single Nvidia A100 node with 8 GPUs show throughput improvements ranging from 39% to 130% when combining ONNX Runtime with DeepSpeed ZeRO Stage 1.

Key configuration details for these benchmarks include:

  • Baseline: PyTorch using the AdamW optimizer.
  • Accelerated: ONNX Runtime using the Fused Adam Optimizer.
  • Maximum Gain: ONNX Runtime + DeepSpeed ZeRO Stage 1.

Environment Specifications:

  • 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

The Optimum Library Ecosystem

Optimum is an extension of the Transformers library designed to accelerate model training and inference by maximizing the efficiency of targeted hardware. While the Accelerate library focuses on distributed training, Optimum integrates machine learning accelerators like ONNX Runtime and specialized hardware such as Intel's Habana Gaudi to lower latency and reduce computing power requirements.

Technical Optimizations in ONNX Runtime Training

ONNX Runtime (ORT) achieves throughput improvements of up to 40% standalone and up to 130% when used with DeepSpeed. These gains are driven by several memory and compute optimizations:

  • Memory Optimizations: Efficient memory planning allows for maximized batch sizes and better utilization of available memory.
  • Compute Optimizations:
    • Kernel Optimizations: General improvements to execution speed.
    • Multi Tensor Apply for Adam Optimizer: Batches elementwise updates for all model parameters into a few kernel launches.
    • FP16 Optimizer: Reduces device-to-host memory copies.
    • Mixed Precision Training: Improves speed and reduces memory footprint.
    • Graph Optimizations: Includes node fusions and node eliminations.

ONNX Runtime Training supports both NVIDIA and AMD GPUs and allows for the use of custom operators.

Implementing ORTTrainer in Optimum

Optimum introduces the ORTTrainer API, which extends the Transformers Trainer to use ONNX Runtime as the backend. This API provides a feature-complete training and evaluation loop supporting hyperparameter search, mixed-precision training, and multi-GPU distributed training.

ORTTrainer allows developers to compose ONNX Runtime with other acceleration techniques, such as Distributed Data Parallel (DDP) and DeepSpeed ZeRO-1 (which partitions optimizer states to save memory). After training, models can be saved as PyTorch models or converted to ONNX format for optimized inference deployment.

To migrate from the standard Trainer to ORTTrainer, developers need to make two primary changes:

  1. Replace Trainer with ORTTrainer.
  2. Replace TrainingArguments with ORTTrainingArguments, enabling access to ORT-specific features like the adamw_ort_fused optimizer.

Future Roadmap

Hugging Face and Microsoft are collaborating to bring these training optimizations to larger model architectures, including Stable Diffusion and Whisper. Additionally, Microsoft has introduced the Azure Container for PyTorch, a curated environment containing DeepSpeed and ONNX Runtime to streamline productivity for PyTorch developers. Efforts are also underway to develop solutions for "learning on the edge," focusing on training on devices with constrained memory and power.

Sources