ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Реализовать FP8 инференс на H100

1. Цель задачи

Освоить использование библиотеки Transformer Engine для выполнения инференса LLM в формате FP8 на GPU H100. Научиться настраивать FP8-квантование, проводить сравнение производительности и качества с базовым FP16 инференсом и добиться ускорения не менее чем в 2 раза при сохранении точности предсказаний.

Ключевой результат Рабочий скрипт инференса модели в FP8 с измеримым ускорением ≥2x относительно FP16 и падением метрик качества (perplexity, accuracy на калибровочном датасете) не более 1%.

2. Исходные данные

Если нет реального H100 — симулируем:
FP8 инференс возможен только на H100 (или B100/B200). Если доступ к H100 отсутствует:

1. Используйте эмуляцию FP8 через библиотеку quantlib или ручное квантование в int8 (не даёт аппаратного ускорения).
2. Запустите FP16 baseline и имитируйте ускорение, искусственно добавляя задержки к FP16 (как учебное упражнение).
   Внимание: реальное ускорение 2x без H100 достигнуть не удастся — задача переходит в теоретическое исследование.

3. Технологический стек

4. Этапы выполнения

Этап 1: Настройка окружения (30–60 минут)

Действия

1. Проверить доступность H100

   nvidia-smi --query-gpu=name,compute_cap --format=csv
   # Должно быть "NVIDIA H100" и compute capability >= 9.0
   

2. Установить зависимости

   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
   pip install transformer-engine pyte[nccl] huggingface_hub datasets evaluate
   pip install nvidia-nsight-systems   # nsys
   

3. Скачать модель (LLaMA-2 7B или меньшую, например, TinyLlama 1.1B для теста)

   from transformers import AutoModelForCausalLM, AutoTokenizer
   model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
   model_fp16 = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
   tokenizer = AutoTokenizer.from_pretrained(model_name)
   model_fp16.to("cuda")
   

4. Подготовить калибровочный датасет
   • Загрузить 1024 примера из datasets.load_dataset("c4", split="train") или другого источника.
   • Каждый пример обрезать до 512 токенов.
   • Сохранить как список тензоров input_ids.

Ожидаемый результат этапа
Среда готова, модель загружена в FP16, калибровочный датасет загружен и токенизирован.

Этап 2: Базовый FP16 инференс и профилирование (1–2 часа)

Действия

1. Реализовать функцию инференса в FP16

   @torch.inference_mode()
   def infer_fp16(model, input_ids, max_new_tokens=128):
       outputs = model.generate(input_ids, max_new_tokens=max_new_tokens,
                                 use_cache=True, num_beams=1)
       return outputs
   

2. Измерить время инференса на 50 батчах (batch_size=1, длина генерации 128 токенов)

   import time
   times = []
   for batch in calibration_data[:50]:
       start = torch.cuda.Event(enable_timing=True)
       end = torch.cuda.Event(enable_timing=True)
       start.record()
       _ = infer_fp16(model, batch.unsqueeze(0))
       end.record()
       torch.cuda.synchronize()
       times.append(start.elapsed_time(end))
   avg_fp16 = sum(times) / len(times)
   print(f"FP16 avg time: {avg_fp16:.1f} ms")
   

3. Записать профиль с помощью nsys profile -o fp16_profile python script.py — сохранить для анализа.
4. Оценить качество FP16: вычислить perplexity на 200 примерах калибровочного датасета (базовый уровень).

Ожидаемый результат этапа

• Среднее время инференса FP16 (T_fp16).
• Perplexity FP16 (PPL_fp16).
• Профиль в формате .nsys-rep.

Этап 3: Реализация FP8 инференса с Transformer Engine (2–3 часа)

Действия

1. Конвертировать модель в формат TE
   Используйте автоматическую замену слоёв:

   import transformer_engine.pytorch as te
   from transformer_engine.common import recipe
   
   # Определяем FP8 рецепт
   fp8_recipe = recipe.DelayedScaling(margin=0, interval=4, fp8_format=recipe.Format.E4M3)
   
   # Оборачиваем модель в TE (работает для линейных слоёв и layer norm)
   model_fp8 = te.utils.convert_to_te_model(model_fp16)
   model_fp8.to("cuda")
   model_fp8.eval()
   

2. Калибровка FP8 масштабов
   Прогнать калибровочный датасет (64-128 батчей) с включённым режимом калибровки:

   with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe, calibrating=True):
       for batch in calibration_data[:128]:
           _ = model_fp8(batch.unsqueeze(0))
   

3. Инференс в FP8

   @torch.inference_mode()
   def infer_fp8(model, input_ids, max_new_tokens=128):
       with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe, calibrating=False):
           outputs = model.generate(input_ids, max_new_tokens=max_new_tokens,
                                     use_cache=True, num_beams=1)
       return outputs
   

4. Измерить время аналогично этапу 2 (50 батчей). Получить T_fp8.

5. Верификация качества:

   • Вычислить perplexity на тех же 200 примерах (PPL_fp8).
   • Убедиться, что |PPL_fp8 - PPL_fp16| / PPL_fp16 < 0.01.

Ожидаемый результат этапа

• Среднее время инференса FP8 (T_fp8).
• Perplexity FP8 (PPL_fp8).
• Коэффициент ускорения S = T_fp16 / T_fp8.

Этап 4: Оптимизация для достижения 2x ускорения (2–3 часа)

Если S < 2.0, выполнить оптимизации:

Конкретные шаги

1. Проверить выравнивание — добавить паддинг до hidden_size % 16 == 0 (обычно уже выполняется).
2. Увеличить batch_size до 2, 4, 8 — измерить throughput (токенов/сек). FP8 выигрывает при больших размерах.
3. Использовать te.Sequential вместо стандартного forward — оптимизация памяти.
4. Профилирование — запустить nsys для FP8 и найти узкие места (например, кастомизированные операции, не покрытые TE).
5. Включить torch.compile (экспериментально, может быть несовместимо с TE).

После каждой оптимизации пересчитывать T_fp8 и качество.

Ожидаемый результат этапа

• Ускорение S >= 2.0 (или обоснование, почему не достигнуто).
• Обновлённый профиль и конфигурация.

Этап 5: Оформление результатов (30–60 минут)

Действия

1. Создать отчёт в виде Jupyter Notebook или Markdown, содержащий:
   • Таблицу сравнения времени и perplexity.
   • График ускорения для разных batch_size.
   • Лучший набор гиперпараметров FP8.
2. Добавить инструкцию по воспроизведению (как запустить).
3. Зафиксировать код в Git (один скрипт fp8_inference.py, поддерживающий оба режима).

Ожидаемый результат этапа
Готовый артефакт (notebook/отчёт), скрипт и профилировочные файлы.

5. Критерии приемки (Definition of Done)

• 1. Код инференса работает как в FP16, так и в FP8 (выбор через аргумент --dtype).
• 2. Время инференса в FP8 меньше как минимум в 2 раза по сравнению с FP16 (batch_size=1).
• 3. Perplexity на калибровочном датасете ухудшился не более чем на 1% относительно FP16.
• 4. Для калибровки использовано не более 128 примеров.
• 5. Профиль Nsight Systems сохранён и может быть загружен для визуализации.
• 6. Отчёт содержит явное сравнение (таблица, график) и обоснование выбора параметров FP8.
• 7. Все зависимости зафиксированы в requirements.txt или environment.yml.
• 8. Код проходит flake8 (стиль) и не содержит лишнего вывода.

6. Ожидаемый результат

Основной артефакт

• Директория проекта со структурой:

  fp8_inference/
  ├── fp8_inference.py        # Главный скрипт (аргументы: --dtype fp16/fp8, --model, --batch_size)
  ├── requirements.txt
  ├── calibration_data.json   # Токенизированные примеры (опционально)
  ├── profile_fp16.nsys-rep   # Профиль FP16
  ├── profile_fp8.nsys-rep    # Профиль FP8
  └── report.md               # Отчёт с результатами
  

Содержание отчёта

• Цель и методология.
• Таблица времени (ms) и throughput (токенов/с) для FP16 и FP8 (batch_size=1,2,4,8).
• Perplexity до и после.
• Ускорение (speedup).
• График (гистограмма) сравнения.
• Выводы: достигнуто ли 2x и почему (или почему нет).

Дополнительные результаты (опционально):

• Сравнение с TensorRT-LLM FP8 (если есть желание расширить).
• Анализ числовой ошибки на уровне выводов (максимальное относительное отклонение).

7. Возможные сложности и их решение

8. Бюджет времени (оценка)

Примечание Для первого выполнения задачи с учётом возможных ошибок и отладки закладывайте 2 полных рабочих дня.

9. Связанные вопросы из базы знаний

10. Чек-лист самопроверки

• Я проверил совместимость версий PyTorch, Transformer Engine и CUDA.
• Я получил ускорение не менее 2x (или документировал причины его отсутствия).
• Я сравнил perplexity на одинаковых входных данных и убедился в падении ≤1%.
• Я сохранил профили nsys и добавил их в репозиторий (или указал, как воспроизвести).
• Мой код поддерживает запуск как в FP16, так и в FP8 через один аргумент командной строки.
• Я зафиксировал гиперпараметры рецепта (margin, interval, формат) и обосновал их выбор.
• Я убедился, что инференс выполняется без ошибок и не вылетает по памяти.