Команда DeepSpeed представила ZenFlow — новый механизм разгрузки, призванный преодолеть серьёзное узкое место в обучении больших языковых моделей (LLM): задержки GPU, вызванные работой CPU.
Традиционные фреймворки, такие как ZeRO-Offload и ZeRO-Infinity, хоть и снижают нагрузку на память GPU за счёт переноса оптимизаторов и градиентов в память CPU, но всё равно оставляют дорогостоящие GPU большую часть времени каждого этапа обучения без дела — в ожидании медленных обновлений CPU и передачи данных через PCIe.
Например, при тонкой настройке Llama 2-7B на 4× A100 GPUs с полной разгрузкой время шага может увеличиться с 0,5 с до более чем 7 с — замедление в 14 раз.
Как работает ZenFlow
- Обновление градиентов с учётом важности. ZenFlow определяет наиболее значимые градиенты для немедленного обновления GPU, а менее важные градиенты накапливаются асинхронно на CPU. Это снижает трафик градиентов на каждом шаге почти на 50% и нагрузку на шину PCIe примерно в 2 раза по сравнению с ZeRO-Offload.
- Ограниченное асинхронное накопление на CPU. Некритичные градиенты объединяются и обновляются асинхронно на CPU, скрывая работу CPU за вычислениями GPU. Это обеспечивает постоянную занятость GPU, предотвращая задержки и максимизируя использование оборудования.
- Лёгкий выбор градиентов. ZenFlow заменяет полный градиент AllGather на лёгкий прокси-нормы градиентов по столбцам, сокращая объём передачи данных более чем в 4000 раз с минимальным влиянием на точность. Это позволяет эффективно масштабировать работу на кластерах с несколькими GPU.
Практическое использование
- Интеграция. ZenFlow — это расширение для DeepSpeed ZeRO-Offload. Не требуется изменений в коде; достаточно обновить конфигурацию в файле DeepSpeed JSON.
- Пример использования. В репозитории DeepSpeedExamples есть пример настройки ZenFlow на бенчмарке GLUE. Пользователи могут запустить его с помощью простого скрипта (bash finetunegptglue.sh), следуя инструкциям по настройке в README репозитория. Пример демонстрирует разгрузку оптимизатора CPU с асинхронными обновлениями ZenFlow, предоставляя практическую отправную точку для экспериментов.
Пример конфигурации:
“`json
{
“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,
“fullwarmup_rounds”: 0,
“overlap_step”: true
}
}
}
“`
Основные моменты производительности
- До 5× ускорения сквозной передачи данных. Более быстрая сходимость, снижение затрат.
- Более чем на 85% сокращение задержек GPU. Более высокая загрузка GPU.
- ≈2× снижение трафика PCIe. Снижение нагрузки на полосу пропускания кластера.
- Отсутствие потери точности на тестах GLUE. Сохранение качества модели.
- Лёгкий выбор градиентов. Эффективное масштабирование до кластеров с несколькими GPU.
- Автоматическая настройка. Не требуется ручная настройка параметров.
Резюме
ZenFlow — это значительный шаг вперёд для всех, кто обучает или настраивает большие языковые модели на ограниченных ресурсах GPU. Эффективно устраняя задержки GPU, вызванные работой CPU, он обеспечивает более высокую пропускную способность и снижает общую стоимость обучения без ущерба для точности модели.
Подход особенно ценен для организаций, масштабирующих рабочие нагрузки LLM на гетерогенном оборудовании, или стремящихся максимально использовать возможности GPU в облаке или локальных кластерах.
Для технических команд сочетание автоматической настройки, минимальной конфигурации и бесшовной интеграции с DeepSpeed делает ZenFlow одновременно доступным и мощным. Предоставленные примеры и документация снижают барьер для внедрения, обеспечивая быстрое экспериментирование и развёртывание.
ZenFlow переопределяет разгрузку для обучения LLM, обеспечивая бесперебойную настройку с высокой пропускной способностью при минимальных затратах на конфигурацию — обязательно попробуйте всем, кто расширяет границы крупномасштабного ИИ.
1. Какие проблемы решает ZenFlow при обучении больших языковых моделей (LLM)?
ZenFlow решает проблему задержек GPU, вызванных работой CPU. Традиционные фреймворки снижают нагрузку на память GPU, но оставляют GPU большую часть времени каждого этапа обучения без дела — в ожидании медленных обновлений CPU и передачи данных через PCIe.
2. Какие основные принципы работы ZenFlow?
Основные принципы работы ZenFlow включают:
* обновление градиентов с учётом важности, где наиболее значимые градиенты немедленно обновляются на GPU, а менее важные накапливаются асинхронно на CPU;
* ограниченное асинхронное накопление на CPU, где некритичные градиенты объединяются и обновляются асинхронно на CPU, скрывая работу CPU за вычислениями GPU;
* лёгкий выбор градиентов, где полный градиент AllGather заменяется на лёгкий прокси-нормы градиентов по столбцам, сокращая объём передачи данных более чем в 4000 раз с минимальным влиянием на точность.
3. Как интегрировать ZenFlow в DeepSpeed?
ZenFlow — это расширение для DeepSpeed ZeRO-Offload. Для интеграции не требуется изменений в коде; достаточно обновить конфигурацию в файле DeepSpeed JSON.
4. Какие основные моменты производительности демонстрирует ZenFlow?
Основные моменты производительности ZenFlow включают:
* до 5× ускорения сквозной передачи данных;
* более чем на 85% сокращение задержек GPU;
* ≈2× снижение трафика PCIe;
* отсутствие потери точности на тестах GLUE;
* лёгкий выбор градиентов;
* автоматическая настройка.
5. Для кого особенно ценен подход, предложенный в ZenFlow?
Подход, предложенный в ZenFlow, особенно ценен для организаций, масштабирующих рабочие нагрузки LLM на гетерогенном оборудовании, или стремящихся максимально использовать возможности GPU в облаке или локальных кластерах. Для технических команд сочетание автоматической настройки, минимальной конфигурации и бесшовной интеграции с DeepSpeed делает ZenFlow одновременно доступным и мощным.