Intel Sapphire Rapids: ускорение обучения PyTorch-моделей без замены кода

Проблема: почему обучение на CPU всё ещё медленное, а GPU - не всегда выход

Дообучение языковых моделей среднего размера на обычных процессорах часто превращается в многочасовое ожидание. Один эпох для DistilBERT на сервере с Intel Xeon предыдущего поколения занимает 3,5 часа. Это неудобно для экспериментов и замедляет выход продукта в продакшен.

Графические ускорители решают проблему скорости, но создают новые: высокую стоимость аренды, дефицит чипов и привязку к специфическому стеку. Многие команды используют Hugging Face Transformers как стандарт и не готовы переписывать код под GPU-библиотеки.

Есть альтернативный путь. Процессоры Intel Sapphire Rapids с инструкциями AMX и библиотека Intel Extension for PyTorch (IPEX) дают кратное ускорение без изменения кода. На одном узле время обучения падает с 3,5 часов до 26 минут, а на четырёх - до 7,5 минут. Это почти идеальное линейное масштабирование на стандартных облачных инстансах Amazon EC2 R7iz.

Что такое Intel Sapphire Rapids и инструкции AMX простыми словами

Sapphire Rapids - четвёртое поколение серверных процессоров Intel Xeon. Главное новшество для задач глубокого обучения - набор инструкций AMX (Advanced Matrix Extensions). Это аппаратный блок, заточенный под одну операцию: перемножение матриц.

Матричные умножения - сердце глубокого обучения

Обучение нейросетей на 80-90% состоит из перемножений матриц. Каждый прямой проход сигнала через слои модели и каждый шаг обратного распространения ошибки - это умножение матриц весов на матрицы активаций. Процессор выполняет миллиарды таких операций за одну эпоху.

Предыдущие инструкции AVX-512 ускоряли векторные вычисления, но матричные операции выполняли через последовательность векторных. AMX делает то же самое за один такт. Результат - рост пропускной способности в несколько раз на тех же частотах. Sapphire Rapids с AMX обрабатывает BF16-умножения быстрее, чем Ice Lake с AVX-512, без увеличения энергопотребления.

Кейс: дообучение DistilBERT на одном узле - 26 минут вместо 3,5 часов

Проведён эксперимент по дообучению модели DistilBERT на задаче классификации текстов. Сравнивались два инстанса Amazon EC2: R7iz на Sapphire Rapids и эквивалентный инстанс предыдущего поколения на Ice Lake. Код Hugging Face Transformers не менялся - использовалась стандартная реализация Trainer.

На Ice Lake один эпох занял 210 минут. На Sapphire Rapids с активированным IPEX и BF16-режимом - 26 минут. Ускорение в 8 раз на идентичном коде и сопоставимом уровне энергопотребления. Это разница между «запустил и ушёл пить кофе» и «запустил и ждёшь полдня».

BF16: половинная точность без потери качества

BF16 (Brain Floating Point) - формат чисел с плавающей запятой, который использует 16 бит вместо 32. В отличие от старого FP16, он сохраняет тот же диапазон значений, что и FP32, но жертвует точностью младших разрядов. Для обучения нейросетей это допустимо: градиенты устойчивы к небольшим погрешностям, а скорость вычислений удваивается.

AMX в Sapphire Rapids аппаратно поддерживает BF16-умножение с накоплением в FP32. Это значит, что матрицы перемножаются в быстром половинном формате, а результат суммируется в полной точности. Качество обучения не страдает, а пропускная способность растёт радикально.

Масштабирование: как четыре узла сокращают время до 7,5 минут

Распределённое обучение на четырёх узлах R7iz сократило время эпоха до 7,5 минут. Это в 3,5 раза быстрее одного узла и в 28 раз быстрее исходного Ice Lake. Ускорение близко к линейному - редкий случай для распределённых систем, где накладные расходы на синхронизацию обычно съедают часть выигрыша.

Кластер построен на базе четырёх инстансов Amazon EC2 R7iz, соединённых высокоскоростной сетью. Каждый узел обрабатывает свою порцию данных, затем градиенты усредняются между всеми узлами. Ключевой компонент, обеспечивший линейное масштабирование, - библиотека Intel oneAPI Collective Communications Library (oneCCL).

Intel oneAPI CCL: эффективная коммуникация между узлами

oneAPI CCL реализует операцию all-reduce - коллективный обмен градиентами между узлами. Стандартные реализации на основе MPI или Gloo добавляют задержки, пропорциональные размеру модели и числу узлов. oneCCL использует кольцевую топологию и асинхронную передачу данных: пока один узел считает градиенты, другой уже передаёт предыдущий пакет. Накладные расходы на синхронизацию минимизированы, поэтому добавление новых узлов даёт почти пропорциональный прирост скорости.

Ни строчки кода менять не нужно: магия IPEX

Intel Extension for PyTorch (IPEX) - библиотека, которая автоматически оптимизирует модель под Sapphire Rapids. Она перехватывает вызовы PyTorch и заменяет стандартные операции на версии, использующие AMX и BF16. Разработчику достаточно добавить две строки в скрипт обучения:

import intel_extension_for_pytorch as ipex
model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=torch.bfloat16)

Весь остальной код - загрузка данных через Hugging Face Datasets, токенизация, Trainer с его коллбэками - остаётся без изменений. IPEX работает с моделями из Hugging Face Transformers, включая BERT, DistilBERT, RoBERTa и другие архитектуры на основе Transformer-энкодеров.

Ранее для ускорения на CPU требовалось вручную настраивать планировщик задач, выбирать реализации операций и экспериментировать с параметрами. Инженеры Intel уже демонстрировали подобные оптимизации для инференса небольших моделей на ноутбуках. IPEX делает то же самое для обучения, но автоматически и для серверных процессоров.

Sapphire Rapids как альтернатива GPU: когда это выгодно

Для моделей с числом параметров до 100-150 миллионов CPU-кластер на Sapphire Rapids часто экономичнее GPU. Инстансы R7iz доступны в большинстве регионов AWS, не требуют бронирования и стоят дешевле сопоставимых по памяти GPU-инстансов.

Масштабирование на CPU проще: добавление узлов линейно ускоряет обучение без переписывания кода под распределённые стратегии. Hugging Face Optimum и IPEX берут на себя оптимизацию под конкретное железо. Для команд, которые уже используют Hugging Face Transformers и не хотят осваивать CUDA, это путь с минимальным порогом входа.

Ограничения тоже есть. Для моделей уровня GPT-3 или больших диффузионных сетей GPU остаются предпочтительнее из-за объёма видеопамяти HBM и зрелости экосистемы. Но для дообучения BERT, DistilBERT, RoBERTa и схожих архитектур Sapphire Rapids даёт сопоставимую с GPU скорость при меньшей стоимости и большей доступности. Оптимизация вывода Transformer-моделей идёт параллельно по многим направлениям, и CPU-направление - одно из самых динамичных.

Заключение: что это значит для практиков машинного обучения

Sapphire Rapids с AMX и IPEX сокращает время дообучения моделей среднего размера в 8 раз на одном узле и в 28 раз на четырёх - без изменения кода Hugging Face Transformers. Распределённое обучение через oneAPI CCL масштабируется почти линейно на стандартных облачных инстансах R7iz.

Если вы дообучаете BERT-подобные модели и упираетесь в скорость CPU, попробуйте перенести обучение на инстанс с Sapphire Rapids и активировать IPEX с BF16. Две строки кода могут сэкономить часы ожидания. Альтернативные аппаратные решения появляются регулярно, но CPU-путь остаётся самым доступным для быстрого старта.

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