Оглавление

Как сообщает PyTorch, летом 2025 года в DeepSpeed появилось новое расширение ZenFlow — движок оффлоадинга без простоев для обучения больших языковых моделей. Технология решает ключевую проблему традиционных подходов: когда мощные GPU простаивают в ожидании медленных CPU-вычислений и передач данных по PCIe.

Проблема традиционного оффлоадинга

Оффлоадинг стал стандартным подходом для масштабирования тонкой настройки больших языковых моделей за пределы ограничений памяти GPU. Фреймворки вроде ZeRO-Offload уменьшают использование памяти GPU, перенося градиенты и состояния оптимизатора на CPU. Однако они создают новый бутылочное горлышко: дорогие GPU часто простаивают в ожидании медленных CPU-обновлений и передач данных по PCIe.

Анализ времени простоя графических процессоров из-за обновлений центрального процессора

Источник: pytorch.org

Рисунок 2: ZeRO-Offload вызывает повторяющиеся простои GPU из-за блокирующих CPU-обновлений и передач по PCIe, приводя к >60% времени прохода на шаг при обучении Llama 2-7B на 4× A100.

На практике включение оффлоадинга при обучении Llama 2-7B на 4× A100 GPU может увеличить каждый шаг с 0.5s до более чем 7s — замедление в 14 раз.

Распределение важности градиентов: вклад 1% наиболее значимых параметров

Источник: pytorch.org

Рисунок 3: В ZeRO-Offload обновления оптимизатора на стороне CPU и передачи по PCIe доминируют во времени итерации, оставляя GPU бездействующим более 5 секунд.

Как работает ZenFlow

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

Диаграмма архитектуры ZenFlow с разделением GPU и CPU

Источник: pytorch.org

Рисунок 1: ZenFlow — это движок оффлоадинга без простоев DeepSpeed для обучения LLM. Он разъединяет обновления GPU и CPU, prioritizating важные градиенты для немедленного обновления на GPU и откладывая остальные для асинхронного накопления на стороне CPU.

Ключевые особенности ZenFlow

  • Нулевые простои GPU: Топ-k важных градиентов обновляются немедленно на GPU; низкоприоритетные градиенты обрабатываются асинхронно на CPU
  • Асинхронность с ограниченной устареваемостью: Стратегия сохраняет сходимость при разъединении выполнения CPU и GPU
  • Автоматическая настройка: ZenFlow адаптирует интервалы обновления во время выполнения на основе динамики градиентов

Производительность

ZenFlow демонстрирует впечатляющие результаты:

  • До 5× ускорения от конца к концу по сравнению с ZeRO-Offload и 6.3× по сравнению с ZeRO-Infinity
  • >85% сокращение простоев GPU на узлах A100/H100
  • ≈2× меньше трафика PCIe (1.13× размера модели на шаг против 2× в ZeRO)
  • Сохраняет или улучшает точность на GLUE (OPT-350M → Llama-13B)

Технически это прорывное решение, которое наконец-то позволяет использовать оффлоадинг без катастрофических потерь производительности. Важность-ориентированный подход к градиентам — это то, что индустрия ждала давно, и реализация в ZenFlow выглядит чрезвычайно грамотной с инженерной точки зрения.

Мотивация проектирования

Обучение больших моделей с оффлоадингом может экономить память GPU, но часто ценой производительности. Синхронные фреймворки оффлоадинга оставляют GPU бездействующим, пока CPU выполняет полный шаг оптимизатора и передает обновленные параметры обратно на GPU.

Визуализация времени простоя графического процессора при использовании технологии ZeRO-Offload

Источник: pytorch.org

Рисунок 4: Обновления CPU доминируют во времени шага, вызывая >60% простоя GPU из-за плохого перекрытия с вычислениями.

Для Llama-2-7B с 4× A100, путь CPU может занимать более 4s, в то время как обратный проход занимает примерно 2s, поэтому более 60% каждой итерации — это чистое время ожидания GPU.

Детализация времени итерации ZeRO-Offload, показывающая доминирование процессора

Источник: pytorch.org

Рисунок 5: Топ 1% градиентов может составлять более 85% норм градиентов.

Архитектурные решения ZenFlow

ZenFlow построен вокруг трех ключевых идей, которые разделяют критические и некритические обновления градиентов при минимизации коммуникационных бутылочных горлышек.

Идея 1: Importance-Aware Top-k Gradient Update

Не все градиенты одинаково влияют на обучение. ZenFlow вводит важность-ориентированный дизайн, который приоритизирует обновления для топ-k наиболее значимых градиентов. Эти градиенты обновляются непосредственно на GPU, используя его высокую пропускную способность вычислений.

Идея 2: Bounded-Asynchronous CPU Accumulation

Асинхронное накопление позволяет CPU оставаться занятым, пока GPU выполняет другие вычисления. Применяется окно накопления для некритических градиентов, позволяя им накапливаться в течение нескольких итераций перед обновлением.

Больше деталей в статье на arXiv.