Проблема: HBM как узкое горлышко масштабирования
Обучение больших языковых моделей (LLM) всё чаще упирается не в вычислительную мощность, а в объём высокоскоростной памяти GPU (HBM). Весы, градиенты, состояния оптимизатора и промежуточные активации конкурируют за ресурсы. NVIDIA предлагает альтернативу активному пересчёту (activation rematerialization): Host Offloading в фреймворке JAX. Метод переносит выбранные активации в зафиксированную память хоста (pinned host memory) во время прямого прохода и подгружает их обратно для обратного прохода.
Почему это работает на NVIDIA Grace Blackwell?
Ключ к успеху — не просто наличие памяти, а скорость соединения. В системах NVIDIA Grace Blackwell CPU и GPU связаны через NVLink-C2C с двусторонней пропускной способностью 900 ГБ/с. В будущих платформах Vera Rubin этот показатель удвоится до 1.8 ТБ/с. Это делает память хоста практически равноценной staging-областью для данных, если правильно организовать асинхронный перенос.
Результаты на MaxText: DeepSeek-V3 671B
Эксперименты проводились на системе NVIDIA GB200 NVL72 (128 GPU) с моделью DeepSeek-V3 671B (sparse MoE). Оптимизация включала три компонента:
- Host Offloading (выгрузка активаций MLA, MoE и MLP).
- Latency Hiding Scheduler (LHS) — скрытие задержек.
- Pipelined transfers — конвейерная передача данных.
Комбинация этих техник позволила достичь пропускной способности 908.2 TFLOPs/s/device. Это на 57% быстрее, чем при использовании активного пересчёта, и на 67.7% быстрее, чем простой offloading без планировщика.
Таблица производительности и памяти
Сравнение политик размещения активаций для DeepSeek-V3 671B. Обратите внимание на возможность использования большего микро-батча и отсутствие ошибок OOM (Out Of Memory) при оптимизированном подходе.
| Конфигурация | Микро-батч | Глобальный батч | Пропускная способность (TFLOPs/s/устр.) | Пик GPU памяти (GiB) | Память хоста (GiB) |
|---|---|---|---|---|---|
| Offload + LHS + Pipelined | 1024 | 1024 | 908.2 | 165.2 | 145.1 |
| No offload + LHS + Rematerialization | 1024 | 1024 | 578.3 | 151.3 | 0.0 |
| Offload (без LHS и pipelining) | 1024 | 1024 | 541.6 | 145.6 | 145.1 |
| No offload + LHS + Save on device | 256 | 256 | 425.3 | 113.3 | 0.0 |
| No offload + LHS + Save on device | 1024 | 1024 | OOM | 0.0 | 0.0 |
Результаты на Llama 3.1 405B
Для плотной модели Llama 3.1 405B (batch size 2, seq len 8192, NVFP4 квантование) выгрузка QKV-активаций с LHS увеличила пропускную способность с 2,669 до 2,746 TFLOPs/s/device. Хотя прирост менее выражен, чем у MoE-моделей, он подтверждает универсальность метода для оптимизации памяти.
Вывод: Синергия софта и железа
NVIDIA подчёркивает, что на обычных кластерах стриминг таких объёмов активаций в память хоста парализовал бы тренировку. Только благодаря тесной интеграции XLA-компилятора, кастомных флагов планирования и высокоскоростных каналов NVLink-C2C, перенос данных происходит асинхронно, не блокируя вычисления. Это открывает доступ к конфигурациям батчей, которые ранее были невозможны из-за нехватки HBM.
Источник: NVIDIA dev blog ↗