Hugging Face와 Microsoft가 Optimum + ONNX Runtime 통합으로 트랜스포머 모델 학습 속도를 39~130% 가속화

Context

트랜스포머 기반 모델의 크기가 증가하면서 학습에 필요한 컴퓨팅 리소스가 기하급수적으로 증가하고 있다. 기존 PyTorch 기반 학습 파이프라인은 메모리 효율성이 낮고 학습 시간이 길어서 대규모 언어, 음성, 비전 모델의 파인튜닝이 비용이 많이 드는 문제가 있었다.

Technical Solution

• Optimum 라이브러리 도입: Transformers의 확장 라이브러리로서 ONNX Runtime, Intel Habana Gaudi 등의 가속기를 통합하여 하드웨어 효율성 극대화
• ORTTrainer API 제공: 기존 Transformers의 Trainer를 ONNX Runtime 백엔드로 확장하여 하이퍼파라미터 서치, 혼합 정밀도 학습, 다중 GPU 분산 학습 지원
• 메모리 최적화 기법 적용: 효율적 메모리 계획(Memory Planning), Fused Adam Optimizer의 다중 텐서 일괄 처리, FP16 Optimizer를 통한 디바이스-호스트 메모리 복사 제거
• 연산 최적화 기법 적용: 커널 최적화, 혼합 정밀도 학습(Mixed Precision Training), 그래프 최적화(노드 융합 및 제거) 구현
• DeepSpeed ZeRO Stage 1과 조합: 옵티마이저 상태 분할을 통한 메모리 절감으로 추가 성능 향상 달성
• ORTTrainingArguments 추가: adamw_ort_fused 옵티마이저 등 ONNX Runtime 고유 인자를 TrainingArguments에 추가하여 기존 코드 호환성 유지

Impact

• 단일 ONNX Runtime 사용 시 PyTorch 대비 최대 40% 학습 처리량 향상
• ONNX Runtime + DeepSpeed ZeRO Stage 1 조합 사용 시 39~130% 성능 가속화(NVIDIA A100 8GPU 노드 기준)
• 인기 있는 Hugging Face 모델들에서 35% 이상의 학습 시간 단축

Key Takeaway

대규모 모델 학습에서 메모리 최적화와 연산 최적화, 분산 학습 기법을 계층적으로 조합하면 단순 기본값 대비 수십 배의 성능 향상을 얻을 수 있으며, Hugging Face 에코시스템 내에서 API 호환성을 유지하면서도 하드웨어 특화 최적화를 투명하게 적용할 수 있다.