Автоматическое определение архитектуры и установка
Инструмент NVIDIA Transformer Engine (TE) интегрируется в PyTorch через пакет transformer_engine[pytorch]. При инициализации система автоматически определяет архитектуру GPU, вычислительную способность (Compute Capability) и объем памяти. Это позволяет выбрать оптимальный режим выполнения:
- TE Capable: GPU поддерживает fused-ядра TE (требует Compute Capability ≥ 8.0, например, Ampere и новее).
- FP8 Capable: GPU поддерживает аппаратные тензорные ядра FP8 (требует Compute Capability ≥ 8.9, например, Hopper H100).
- Fallback: На старых архитектурах (например, T4) или при отсутствии поддержки TE используется чистый PyTorch.
Ключевые компоненты и Fused Kernels
Transformer Engine предлагает набор оптимизированных модулей, которые заменяют стандартные слои PyTorch для снижения накладных расходов на передачу данных между ядрами. Основные компоненты включают:
te.Linear: Оптимизированное линейное преобразование.te.LayerNorm: Нормализация слоя с fused-вычислениями.te.LayerNormLinearиte.LayerNormMLP: Комбинированные операции, объединяющие нормализацию и полносвязные слои.te.TransformerLayer: Полностью слитый блок трансформера, включающий внимание и MLP.
Рецепт FP8: Delayed Scaling
Для достижения максимальной производительности на поддерживаемом оборудовании используется формат FP8 (E4M3 и E5M2). NVIDIA рекомендует использовать рецепт DelayedScaling, который управляет масштабированием тензоров и историей амплитуд (amax) для предотвращения переполнения. В примере используется гибридный формат HYBRID с длиной истории amax_history_len=16 и алгоритмом вычисления максимума "max".
Бенчмаркинг и сравнение производительности
В рамках тестирования была создана компактная модель MiniGPT (96 токенов, 768 скрытых состояний, 12 слоев) и обучена на синтетических последовательностях. Сравнение проводилось между чистой реализацией PyTorch и версией с использованием TE. Ниже приведены ключевые метрики, которые измеряются в процессе тестирования:
| Параметр | Описание | Значение/Примечание |
|---|---|---|
| Архитектура GPU | Определяется через torch.cuda | Зависит от железа (CC >= 8.0 для TE) |
| Точность вычислений | BF16 или FP8 | FP8 доступен только на Hopper/Ada |
| Размер батча (bench) | batch_size | 32 |
| Итерации бенчмарка | iters | 30 (после 10 warmup) |
| Метрика времени | ms/step | Среднее время одного шага обучения |
| Метрика памяти | Peak Memory (GB) | torch.cuda.max_memory_allocated |
Практическое применение
Использование Transformer Engine позволяет не только ускорить обучение за счет fused-ядер, но и снизить потребление памяти благодаря FP8. Код демонстрирует полную цепочку: от установки пакета и проверки совместимости GPU до создания модели, настройки оптимизатора AdamW и запуска цикла обучения с измерением метрик. Это критически важно для разработчиков, работающих с большими языковыми моделями (LLM) на инфраструктуре NVIDIA.
Источник: MarkTechPost ↗