Как работает FP8 quantization на H100 (Transformer Engine)?

Краткий тезис

FP8 quantization на H100 с помощью Transformer Engine (TE) позволяет автоматически конвертировать тензоры из FP16/BF16 в FP8 на лету, используя scaling factors для минимизации потери точности. Это даёт до 2× ускорения вычислений и снижение потребления памяти в 2 раза за счёт использования 8-битных тензоров, при минимальном влиянии на точность, особенно с fine-tuning и обработкой outliers. TE интегрирован в библиотеки NeMo и Megatron‑LM, упрощая внедрение.

1. Термин: FP8 quantization

FP8 quantization — это формат чисел с плавающей точкой, использующий 8 бит вместо 16 (FP16) или 32 (FP32). Существует два основных варианта: E4M3 (4 бита экспонента, 3 мантисса) для весов и активаций, и E5M2 (5 бит экспонента, 2 мантисса) для градиентов. Это позволяет снизить объём памяти и ускорить матричные умножения на H100, у которого есть специальные блоки для обработки FP8.

2. Архитектура: H100 и Tensor Cores

H100 (Hopper) от NVIDIA содержит четвёртое поколение Tensor Cores, поддерживающее FP8 вычисления. Тактовые частоты и пропускная способность позволяют выполнять операции FP8 в 2 раза быстрее, чем FP16. Transformer Engine — это программная библиотека NVIDIA, которая автоматически определяет, где можно использовать FP8, управляет динамическим масштабированием (scaling) и форматированием тензоров без вмешательства пользователя.

3. Принцип работы: автоматическая конвертация и scaling factors

Процесс состоит из шагов:

1. На каждом шаге обучения/инференса Transformer Engine анализирует распределение значений тензоров (веса, активации, градиенты).
2. Вычисляет scaling factor (коэффициент масштабирования) как максимальное абсолютное значение, делённое на максимальное представимое число в FP8 (например, для E4M3 это 448). Это позволяет уместить весь диапазон значений в FP8 без переполнения.
3. Тензор умножается на scaling factor (обычно 1/макс) и округляется до FP8.
4. После операции матричного умножения обратное масштабирование (деление на scaling factor) возвращает результат в FP16/BF16.
5. Scaling factor хранится как 8-битное число (или пара чисел для разных тензоров) и передаётся между операциями.

4. Форматы E4M3 и E5M2

E5M2 имеет больший диапазон, но меньшую точность — подходит для градиентов, где важнее избежать переполнения.

5. Роль Transformer Engine

Transformer Engine (TE) — это программная надстройка над CUDA и Tensor Cores, которая:

• Автоматически выбирает между FP8 и FP16/BF16 для каждого слоя на основе статистик (outliers, variance).
• Управляет динамическим масштабированием: подгоняет scaling factor на каждом шаге.
• Поддерживает ретроспективное масштабирование (delayed scaling) для компенсации сильных выбросов.
• Интегрирован в популярные фреймворки: Megatron‑LM, NeMo, PyTorch через модуль transformer_engine.pytorch.

6. Ускорение и экономия памяти

• Ускорение: Матричные умножения в FP8 выполняются примерно в 2 раза быстрее на H100 за счёт вдвое меньшего размера данных и более эффективного использования Tensor Cores.
• Память: Веса и активации занимают в 2 раза меньше места (8 бит вместо 16). Это позволяет увеличить размер модели или batch size на той же карте.
• Градиенты: могут оставаться в FP8 или FP16 в зависимости от реализации; для точности часто градиенты хранят в FP16/32.

7. Точность и fine-tuning

Изначальный loss of accuracy при FP8 quantization может составлять 1–2% на сложных задачах (например, LLM fine-tuning). Однако с помощью:

• Fine-tuning в смешанной точности (FP32 master weights, FP8 forward/backward) — accuracy восстанавливается до уровня FP16.
• Outlier‑aware scaling — обработка выбросов (например, через блоки LayerNorm в FP16).
• Gradient scaling — динамическое масштабирование градиентов для предотвращения underflow. Минимизирует потери. Для inference (без обучения) loss of accuracy часто ниже 0.5% при правильной калибровке.

8. Сравнение с другими методами квантизации

FP8 выигрывает в балансе скорости и точности благодаря динамическому масштабированию и поддержке аппаратуры.

9. Практические аспекты: когда использовать FP8?

• Обучение больших языковых моделей (LLM): FP8 позволяет ускорить обучение и снизить VRAM, что критично для моделей >10B параметров.
• Инференс: FP8 inference даёт меньшую задержку при сохранении качества.
• Не рекомендуется: для задач, требующих экстремальной точности (например, научные расчёты) или на старых GPU без поддержки FP8.

10. Интеграция с PyTorch

Пример использования Transformer Engine в PyTorch:

import transformer_engine.pytorch as te
from transformer_engine.common.recipe import Format, DelayedScaling

# Рецепт для FP8 с динамическим масштабированием
fp8_recipe = DelayedScaling(
    fp8_format=Format.HYBRID,  # E4M3 для весов, E5M2 для градиентов
    amax_history_len=1024,
    amax_compute_algo="max",
)

# Замена Linear слоя на TE Linear
model = te.Linear(in_features, out_features, bias=True)

# В цикле обучения с FP8 autocast
from transformer_engine.pytorch import fp8_autocast
with fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    output = model(input)

Пет-проект для закрепления

Задача: Сравнить производительность и точность FP8 quantization на H100 (или эмулировать на CPU с помощью симуляции FP8).

Инструменты:

• PyTorch
• Transformer Engine (если есть доступ к H100) или torch.fake_quantize для эмуляции FP8.
• Модель: небольшой transformer (например, GPT‑2 small).

Шаги:

1. Обучить модель в FP16 на небольшом датасете (например, wikitext‑2).
2. Загрузить веса и применить FP8 квантизацию с помощью Transformer Engine (или вручную с помощью scaling).
3. Запустить инференс на валидации и измерить perplexity и скорость.
4. Сравнить с FP16 baseline: скорость (токенов/сек) и точность (perplexity).
5. Провести fine‑tuning в FP8 на нескольких эпохах и оценить восстановление точности.

Ожидаемый результат: FP8 ускорит инференс в ~1.8×, а после fine‑tuning perplexity будет близка к FP16 (разница <0.5).

Связь с другими вопросами

Навигация

• Предыдущий: 311
• Следующий: 313
• Индекс: 00. Индекс разборов