Hugging Face Optimum 및 ONNX Runtime 훈련 통합

Hugging Face와 Microsoft는 ONNX Runtime을 Optimum 라이브러리에 통합하여 많은 인기 있는 Hugging Face 모델의 훈련 시간을 35% 이상 단축시키는 오픈 솔루션을 제공했습니다. 이번 협업은 대규모 언어, 음성 및 비전 모델을 미세 조정하는 데 필요한 계산 자원과 시간을 줄이는 것을 목표로 합니다.

훈련 성능 향상

ONNX Runtime을 Optimum 라이브러리와 통합하면 Hugging Face 모델에 대해 상당한 가속을 제공합니다. 8개의 GPU를 갖춘 단일 Nvidia A100 노드에서 수행된 벤치마크는 ONNX Runtime과 DeepSpeed ZeRO Stage 1을 결합했을 때 처리량이 39%에서 130%까지 향상됨을 보여줍니다.

이 벤치마크의 주요 구성 세부 사항은 다음과 같습니다:

  • Baseline: AdamW 옵티마이저를 사용하는 PyTorch.
  • Accelerated: Fused Adam Optimizer를 사용하는 ONNX Runtime.
  • Maximum Gain: 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%까지 처리량을 향상시킵니다. 이러한 향상은 여러 메모리 및 연산 최적화에 의해 이루어집니다:

  • Memory Optimizations: 효율적인 메모리 계획을 통해 배치 크기를 최대화하고 사용 가능한 메모리를 더 잘 활용할 수 있습니다.
  • Compute Optimizations:
    • Kernel Optimizations: 실행 속도의 전반적인 개선.
    • Multi Tensor Apply for Adam Optimizer: 모든 모델 파라미터에 대한 요소별 업데이트를 몇 번의 커널 실행으로 배치합니다.
    • FP16 Optimizer: 디바이스와 호스트 간 메모리 복사를 감소시킵니다.
    • Mixed Precision Training: 속도를 향상시키고 메모리 사용량을 줄입니다.
    • Graph Optimizations: 노드 융합 및 노드 제거를 포함합니다.

ONNX Runtime Training은 NVIDIA와 AMD GPU 모두를 지원하며 커스텀 연산자의 사용을 허용합니다.

Optimum에서 ORTTrainer 구현

Optimum은 ORTTrainer API를 도입하여 Transformers Trainer를 확장하고 ONNX Runtime을 백엔드로 사용합니다. 이 API는 하이퍼파라미터 탐색, 혼합 정밀도 훈련, 다중 GPU 분산 훈련을 지원하는 완전한 기능의 훈련 및 평가 루프를 제공합니다.

ORTTrainer는 개발자가 Distributed Data Parallel(DDP) 및 DeepSpeed ZeRO-1(옵티마이저 상태를 분할하여 메모리를 절약)과 같은 다른 가속 기술과 ONNX Runtime을 결합할 수 있게 합니다. 훈련이 끝난 후 모델은 PyTorch 모델로 저장하거나 ONNX 형식으로 변환하여 최적화된 추론 배포에 사용할 수 있습니다.

표준 Trainer에서 ORTTrainer로 마이그레이션하려면 개발자는 두 가지 주요 변경을 해야 합니다:

  1. Trainer를 ORTTrainer로 교체합니다.
  2. TrainingArguments를 ORTTrainingArguments로 교체하여 adamw_ort_fused 옵티마이저와 같은 ORT 전용 기능에 접근할 수 있게 합니다.

향후 로드맵

Hugging Face와 Microsoft는 Stable Diffusion 및 Whisper와 같은 대형 모델 아키텍처에 이러한 훈련 최적화를 적용하기 위해 협업하고 있습니다. 또한 Microsoft는 DeepSpeed와 ONNX Runtime을 포함한 맞춤형 환경인 Azure Container for PyTorch를 도입하여 PyTorch 개발자의 생산성을 높였습니다. 메모리와 전력이 제한된 디바이스에서 훈련하는 '엣지 학습' 솔루션을 개발하기 위한 노력도 진행 중입니다.

Sources