ZenFlow ускоряет до 5×: безстолковые offload-обновления для обучения LLM

ZenFlow — расширение DeepSpeed, созданное для устранения простоев GPU, вызванных обновлениями на CPU при offloading'е оптимизаторов и градиентов. За счёт разделения вычислений GPU и CPU с помощью importance-aware конвейеризации и ограниченно-асинхронных обновлений ZenFlow поддерживает GPU загруженными и даёт значительный прирост производительности при минимальной настройке.

Почему простои GPU критичны

Offloading оптимизаторов и градиентов в память CPU снижает потребление GPU-памяти, но традиционные решения, такие как ZeRO-Offload и ZeRO-Infinity, часто приводят к простаиванию GPU, пока CPU выполняет обновления и идут передачи по PCIe. В некоторых сценариях шаг обучения может замедлиться с 0.5s до более 7s, как в случае с Llama 2-7B на 4× A100 при полном offload'е — это 14× замедление. Такие простои дорого обходятся по времени и ресурсам.

Основные методы ZenFlow

• Importance-Aware обновления градиентов: ZenFlow сразу обновляет наиболее важные топ-k градиенты на GPU, а менее значимые накапливает асинхронно на CPU. Это снижает трафик градиентов почти на 50% и уменьшает нагрузку на PCIe примерно в 2× по сравнению с ZeRO-Offload.
• Ограниченно-асинхронная агрегация на CPU: некритичные градиенты батчатся и обновляются на CPU асинхронно, скрывая работу CPU за вычислениями GPU, что предотвращает простои.
• Лёгкий выбор градиентов: вместо полного AllGather градиентов используется прокси по нормам столбцов градиентов, что уменьшает объём коммуникации более чем в 4,000× и практически не влияет на точность.
• Автоматическая подстройка: интервалы обновления адаптируются в рантайме (параметры select_interval и update_interval можно ставить в "auto"), что убирает необходимость ручной тонкой настройки.

Результаты производительности

ZenFlow обеспечивает до 5× ускорения end-to-end по сравнению с ZeRO-Offload и сокращает простои GPU более чем на 85%. Также уменьшается трафик по PCIe примерно в 2×, при этом качество моделей сохраняется на GLUE и других бенчмарках.

Интеграция и использование

ZenFlow интегрируется как расширение ZeRO-Offload в DeepSpeed и не требует изменения кода — достаточно обновить JSON-конфигурацию DeepSpeed. Доступны параметры вроде topk_ratio для выбора доли градиентов, select_strategy, select_interval и update_interval. Поддерживается автоматическое определение параметров для упрощённой развёртки.

Пример конфигурации из документации DeepSpeed:

"zero_optimization": {
  "stage": 2,
  "offload_optimizer": {
    "device": "cpu",
    "pin_memory": true
  },
  "zenflow": {
    "topk_ratio": 0.05,
    "select_strategy": "auto",
    "select_interval": "auto",
    "update_interval": 4,
    "full_warm_up_rounds": 0,
    "overlap_step": true
  }
}

В репозитории DeepSpeedExamples есть пример дообучения с ZenFlow на GLUE и скрипт bash finetune_gpt_glue.sh для старта. Следуйте README репозитория для установки и запуска.

Кому это полезно

Наибольшую выгоду увидят команды, которые дообучают большие модели на ограниченных или гетерогенных кластерах GPU. ZenFlow особенно актуален, когда offloading необходим для размещения модели, но пропускной способности CPU и задержки PCIe грозят снизить производительность GPU.

Итоговая идея

ZenFlow меняет подход к offloading'у: он приоритизирует важные градиенты, скрывает CPU-работу за GPU-вычислениями и использует лёгкий отбор градиентов для минимизации коммуникации. Это даёт беспростойное, высокопроизводительное дообучение при минимальной конфигурации и без правки кода.