Как ускорить обучение моделей Hugging Face на 130% с помощью Optimum и ONNX Runtime

Почему обучение моделей такое медленное и как это исправить

Современные языковые модели растут в размерах быстрее, чем мощность доступного железа. Дообучение крупной LLM на пользовательских данных занимает дни и недели, съедает бюджеты на облачные GPU и тормозит цикл разработки. Разработчики из Hugging Face и Microsoft предложили прямой ответ на эту проблему. Их решение - библиотека Optimum с интеграцией ONNX Runtime для обучения. Прирост скорости составляет от 35% до 130% на популярных архитектурах, а изменения в коде минимальны.

Главный барьер для многих команд - убеждение, что ускорение обучения требует глубокой переработки кодовой базы или покупки специализированного оборудования. Optimum ломает этот стереотип. Вы остаётесь в привычном API Hugging Face, просто заменяете класс Trainer на ORTTrainer. Ниже разберём, как это работает и как применить в своём проекте.

Ранее мы рассказывали, как Hugging Face добился 100-кратного ускорения вывода нейросетей через компиляцию под конкретное железо. Теперь те же принципы оптимизации пришли в обучение.

Что такое Optimum и ONNX Runtime: краткий ликбез

Новые инструменты часто вызывают сопротивление: «ещё одна библиотека, которую придётся изучать». Разберём по сути - что делают Optimum и ONNX Runtime и зачем они нужны именно вам.

Optimum: мост между Hugging Face и ускорением

Optimum - это библиотека от Hugging Face, которая даёт единый интерфейс для разных методов оптимизации моделей. Она скрывает технические детали за знакомым API. Вам не нужно разбираться в тонкостях компиляции графов или квантования - Optimum берёт эту работу на себя.

Библиотека поддерживает несколько бэкендов: ONNX Runtime, Intel Neural Compressor, Habana Gaudi. Для обучения ключевой компонент - ORTTrainer, который заменяет стандартный Trainer из библиотеки Transformers. Optimum изначально создавался для упрощения запуска моделей в продакшене, но теперь его возможности распространяются на полный цикл обучения.

ONNX Runtime: не только для инференса

ONNX Runtime часто воспринимают как движок для быстрого выполнения готовых моделей. Это верно, но неполно. Microsoft развивает в нём модуль ORTModule, который интегрируется с PyTorch и ускоряет прямое и обратное распространение сигнала. ORTModule перехватывает вычислительный граф PyTorch, оптимизирует его и выполняет на том же оборудовании быстрее.

Фактически вы получаете прозрачное ускорение без смены фреймворка. Код остаётся на PyTorch, модель остаётся моделью Hugging Face, а вычисления идут через оптимизированный рантайм ONNX.

Как это работает: три кита ускорения обучения

Прирост скорости на 35–130% достигается комбинацией трёх механизмов. Каждый из них по отдельности даёт небольшой выигрыш, вместе они создают кумулятивный эффект.

Графовые оптимизации: меньше операций - быстрее результат

PyTorch строит вычислительный граф динамически, шаг за шагом. ONNX Runtime перехватывает этот граф и применяет статические оптимизации до начала вычислений. Слияние слоёв нормализации и активации в одну операцию устраняет лишние проходы по памяти. Удаление избыточных тензорных преобразований сокращает число операций. Переупорядочивание узлов графа улучшает использование кэша процессора и GPU.

Результат: та же математика выполняется за меньшее число шагов. На небольших моделях вроде BERT-base прирост от графовых оптимизаций составляет 15–25%.

Память и вычисления: баланс для больших моделей

Обучение крупных моделей упирается в объём видеопамяти. Увеличить размер батча нельзя - не хватает гигабайт. Уменьшить батч - падает utilisation GPU и растёт время. ONNX Runtime применяет gradient checkpointing - технику, при которой промежуточные активации не хранятся в памяти, а пересчитываются заново при обратном проходе. Это компромисс: больше вычислений, но радикально меньше памяти.

ORTModule автоматически выбирает, для каких слоёв применить checkpointing, а какие оставить без изменений. Ручной подбор таких точек занял бы часы, автоматический занимает секунды. Освободившуюся память можно направить на увеличение батча, что напрямую ускоряет обучение.

Mixed precision: используем 16-битные вычисления без потери точности

Современные GPU эффективнее работают с 16-битными числами с плавающей точкой. Tensor Cores на архитектурах Volta и новее выдают кратно больше операций в секунду на FP16 и BF16 по сравнению с FP32. ONNX Runtime автоматически приводит вычисления к половинной точности там, где это безопасно для сходимости модели.

BF16 особенно интересен: он сохраняет тот же диапазон значений, что и FP32, просто с меньшей точностью. Риск расхождения обучения минимален. Автоматическое смешанное обучение в ONNX Runtime избавляет от ручной расстановки приведения типов в коде - вы просто включаете флаг, и рантайм делает всё сам.

Пошаговое руководство: добавляем ускорение в ваш проект

Переход на Optimum и ONNX Runtime занимает минуты. Покажем на примере дообучения модели для классификации текста.

Установка и первые шаги

Установите пакет с зависимостями для ONNX Runtime:

pip install optimum[onnxruntime]

Эта команда подтянет саму Optimum, ONNX Runtime и все необходимые привязки к PyTorch. После установки доступны два ключевых класса: ORTTrainer и ORTTrainingArguments.

Адаптация существующего кода обучения

Стандартный скрипт обучения на Hugging Face выглядит так:

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(output_dir="./results"),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()

С Optimum изменения минимальны - меняете импорты и класс тренера:

from optimum.onnxruntime import ORTTrainer, ORTTrainingArguments

trainer = ORTTrainer(
    model=model,
    args=ORTTrainingArguments(output_dir="./results"),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    feature="sequence-classification",
)
trainer.train()

Параметр feature указывает тип задачи: sequence-classification, token-classification, question-answering. Optimum использует его для выбора оптимальных оптимизаций под конкретную архитектуру.

Запуск и замер производительности

Для измерения реального прироста добавьте замер времени:

import time
start = time.time()
trainer.train()
print(f"Обучение заняло: {time.time() - start:.2f} секунд")

На GPU класса V100 для BERT-base дообучение на задачах классификации текста ускоряется на 40–60%. На A100 с поддержкой BF16 прирост достигает 90–130% за счёт эффективного использования Tensor Cores.

Если вы экспериментируете с гиперпараметрами моделей, обратите внимание на руководство по настройке гиперпараметров NLP-моделей с Ray Tune - комбинация с Optimum даёт двойной выигрыш: быстрее итерации и быстрее каждая итерация.

Какие модели и задачи выигрывают больше всего

Optimum с ONNX Runtime поддерживает широкий спектр архитектур, но прирост неравномерен. Лучшие результаты показывают:

• BERT и RoBERTa - классификация текста, NER, вопросно-ответные системы
• GPT-2 - генерация текста
• T5 - задачи seq2seq: суммаризация, перевод
• DistilBERT - облегчённые модели, где важна каждая миллисекунда

Ограничения связаны с нестандартными операциями в кастомных архитектурах. Если модель использует операции, не поддерживаемые ONNX Runtime, ORTTrainer сообщит об ошибке конвертации. Перед внедрением в продакшен-пайплайн проверьте модель на тестовом прогоне.

Сравнение с другими подходами: когда выбирать Optimum

Optimum не единственный способ ускорить обучение. DeepSpeed и FSDP оптимизируют распределённое обучение на множестве GPU. PyTorch Lightning даёт гибкость в организации цикла обучения. У каждого подхода своя ниша.

Optimum выигрывает, когда нужен быстрый результат без перестройки кодовой базы. Вы жертвуете тонкой настройкой оптимизаций под конкретный случай, но получаете работающее ускорение за несколько строк кода. Для команд, которые хотят сократить время экспериментов и не готовы погружаться в детали распределённых вычислений, это оптимальный выбор.

DeepSpeed и FSDP оправданы, когда модель не помещается на один GPU даже с gradient checkpointing. Optimum решает обратную задачу: ускорить обучение на одном или двух GPU без усложнения инфраструктуры. Hugging Face и Google уже показали аналогичный подход для TPU - интеграция PyTorch/XLA дала ускорение без переписывания кода. Optimum делает то же самое для GPU через ONNX Runtime.

Заключение: стоит ли переходить на Optimum и ONNX Runtime

Ускорение обучения на 35–130% - не маркетинговое обещание, а воспроизводимый результат на стандартных бенчмарках. Порог входа низкий: установка одного пакета и замена двух классов в коде. Решение подходит для большинства типовых задач с моделями Hugging Face.

Попробуйте на своём проекте: возьмите текущий скрипт обучения, замените Trainer на ORTTrainer и сравните время дообучения. Скорее всего, вы увидите прирост уже на первой итерации. Документация и примеры доступны в репозитории Optimum на GitHub.

Есть вопрос или заметили неточность? Напишите нам - мы проверяем обратную связь и обновляем материалы.