ИИ и технологии
Масштабирование FlashAttention, обучение на нескольких GPU и квантизация весов
Современная ИИ-инфраструктура опирается на аппаратно-ориентированные ядра attention, распределённый параллелизм на нескольких GPU и низкобитную квантизацию, чтобы эффективно обучать и обслуживать огромные модели. Эти три техники устраняют разрыв между теоретическим масштабом модели и физическими ограничениями GPU.
1. FlashAttention-2 и FlashAttention-3 на GPU Hopper (H100)
Хотя оригинальный FlashAttention решил проблему узкого места пропускной способности памяти, он всё же оставлял значительную часть сырой вычислительной производительности современных GPU неиспользованной. FlashAttention-2 (2023) и FlashAttention-3 (2024) переработали координацию потоков GPU и аппаратных блоков.
- FlashAttention-2 (Улучшенное разбиение работы).
- Сокращение операций вне matmul. Тензорные ядра GPU выполняют матричные умножения (matmul) примерно в 16× быстрее, чем стандартные ядра CUDA выполняют скалярные операции вроде перемасштабирования Softmax. FlashAttention-2 переписал математику онлайн-Softmax, отложив перемасштабирование до самого конца внутреннего цикла и сократив число медленных инструкций вне matmul.
- Параллелизация по длине последовательности. Исходная версия распараллеливала работу в основном по размеру батча и головам attention, оставляя ядра GPU простаивать при обучении с длинными последовательностями и малыми батчами. FlashAttention-2 добавляет параллелизм по измерению длины последовательности и меняет порядок циклов (внешний цикл по блокам Q, внутренний — по K и V), минимизируя медленные записи в разделяемую память.
- FlashAttention-3 (Использование аппаратных возможностей Hopper H100).
- Асинхронное перемещение данных (TMA). Архитектура NVIDIA Hopper представила Tensor Memory Accelerator (TMA) — выделенный аппаратный блок, который копирует данные между глобальной HBM и внутрикристальной SRAM, не занимая вычислительные потоки.
- Специализация warp'ов (модель «производитель-потребитель»). FlashAttention-3 делит warp'ы GPU (группы по 32 потока) на специализированные роли. Warp'ы-производители используют TMA для предварительной загрузки будущих тайлов K и V в SRAM, а warp'ы-потребители подают текущие тайлы в тензорные ядра (инструкции WGMMA). Это одновременно накладывает передачи памяти на матричные вычисления и расчёты Softmax.
- Поддержка низкой точности FP8. FlashAttention-3 задействует FP8-тензорные ядра Hopper, удваивая сырую пропускную способность по сравнению с FP16. Чтобы предотвратить потерю точности от выбросов активаций в FP8, используется блочная квантизация (независимое масштабирование малых тайлов) и инкогерентная обработка (умножение Q и K на случайные ортогональные матрицы для сглаживания экстремальных значений перед квантизацией).
2. Стратегии масштабирования на несколько GPU (3D-параллелизм)
Когда модель имеет десятки или сотни миллиардов параметров, её веса, градиенты и состояния оптимизатора превышают VRAM одного GPU. Крупномасштабное обучение объединяет три различные стратегии параллелизации — так называемый 3D-параллелизм.
| Стратегия | Как разделяет работу | Узкое место коммуникации | Лучший сценарий применения |
|---|---|---|---|
| Параллелизм данных (DP / FSDP) | Реплицирует модель на все GPU; каждый GPU обрабатывает свой батч данных. | All-Reduce / All-Gather градиентов и шардированных весов между GPU. | Масштабирование пропускной способности обучения на сотни GPU, когда слои помещаются на одном устройстве. |
| Тензорный параллелизм (TP) | Разделяет отдельные весовые матрицы (слои) по вертикали или горизонтали между GPU. | All-Reduce после каждого слоя Transformer (требует сверхбыстрого NVLink). | Огромные слои, не помещающиеся на одном GPU, строго в пределах одного серверного узла с 8 GPU. |
| Конвейерный параллелизм (PP) | Разделяет модель по глубине, размещая последовательные группы слоёв на разных GPU. | Передачи активаций «точка-точка» (P2P) между стадиями конвейера. | Масштабирование на несколько серверных узлов, соединённых более медленными сетями InfiniBand/Ethernet. |
- Параллелизм данных & ZeRO / FSDP. В стандартном параллелизме данных каждый GPU хранит полную копию модели, расточительно расходуя память. Современные фреймворки используют ZeRO (Zero Redundancy Optimizer) или PyTorch FSDP (Fully Sharded Data Parallel), которые шардируют параметры, градиенты и состояния оптимизатора по всем GPU, собирая веса каждого слоя точно к моменту вычисления.
- Тензорный параллелизм (внутри слоя). Вместо дублирования огромного матричного умножения прямого распространения Y = XW тензорный параллелизм разделяет W по столбцам между GPU (W = [W_1, W_2]). Каждый GPU одновременно вычисляет свою половину выхода и обменивается частичными результатами по высокоскоростному NVLink (900 ГБ/с на H100).
- Конвейерный параллелизм (между слоями). Слои 1–8 находятся на GPU 1, слои 9–16 — на GPU 2 и так далее. Чтобы нижестоящие GPU не простаивали в ожидании завершения ранних слоёв («пузырь конвейера»), фреймворки режут каждый обучающий батч на меньшие микробатчи, так что все GPU остаются занятыми в режиме сборочного конвейера.
3. Квантизация (GPTQ, AWQ) и ядра деквантизации на лету
Во время инференса (генерации токенов по одному) большие языковые модели ограничены пропускной способностью памяти, а не вычислениями. Загрузка 70 миллиардов 16-битных параметров (140 ГБ) из VRAM лишь для генерации одного токена оставляет тензорные ядра в ожидании передач памяти. Квантизация сжимает веса из 16-битных чисел с плавающей точкой в 4-битные целые (INT4), сокращая объём памяти и время передачи почти в 4×.
- GPTQ (Постобучающая квантизация через калибровку по Гессиану).
- Наивное округление 16-битных весов до 4-битных целых разрушает точность модели. GPTQ пропускает небольшой калибровочный набор данных через модель и использует математику второго порядка (обратную матрицу Гессиана), чтобы измерить чувствительность каждого веса.
- Он квантизирует веса столбец за столбцом; всякий раз, когда округление одного веса вносит ошибку, GPTQ математически корректирует оставшиеся неквантизованные веса в этой строке, чтобы погасить ошибку.
- AWQ (Квантизация весов с учётом активаций).
- Исследователи обнаружили, что примерно от 0,1% до 1% весовых каналов значительно важнее остальных, поскольку они умножаются на большие значения активаций.
- Вместо того чтобы хранить эти критические веса в 16-битном формате (что создаёт неудобные, медленные структуры памяти со смешанной точностью), AWQ масштабирует вверх эти значимые весовые каналы перед квантизацией и масштабирует вниз соответствующие входные активации во время выполнения, сохраняя точность и удерживая 100% весов в едином 4-битном формате.
- Как специальные ядра CUDA выполняют деквантизацию на лету (W4A16).
- GPU не могут напрямую умножать 16-битную активацию (FP16) на 4-битный целочисленный вес (INT4). Вместо этого специализированные ядра инференса (такие как Marlin или ExLlama) выполняют вычисление W4A16 (4-битные веса, 16-битные активации):
- 1. Сжатое чтение из HBM: Ядро загружает упакованные 4-битные веса (восемь значений INT4, упакованных в одно 32-битное целое) из медленной глобальной VRAM, сокращая использование пропускной способности памяти в 4×.
- 2. Распаковка на уровне регистров: Внутри сверхбыстрых регистров GPU ядро CUDA использует быстрые битовые сдвиги (>>) и битовые маски (&), чтобы распаковать целые INT4, и применяет сохранённые коэффициенты масштаба/нулевой точки, преобразуя их обратно в значения FP16 на лету.
- 3. Немедленные вычисления на тензорных ядрах: Свежераспакованные веса FP16 подаются прямо в тензорные ядра вместе с активациями FP16, и расширенные 16-битные веса ни разу не записываются обратно в VRAM.