ZenFlow: Новое решение для обучения больших языковых моделей
Команда DeepSpeed представила ZenFlow — новый движок для оффлоадинга, который решает одну из главных проблем в обучении больших языковых моделей (LLM): задержки GPU, вызванные работой CPU. Традиционные фреймворки, такие как ZeRO-Offload и ZeRO-Infinity, часто оставляют дорогие GPU без дела из-за медленных обновлений CPU и передачи данных по PCIe. Например, дообучение модели Llama 2-7B на 4-х A100 GPU с полным оффлоадингом может увеличить время выполнения шага с 0,5 секунд до более чем 7 секунд, что соответствует замедлению в 14 раз. ZenFlow устраняет эти задержки, разделяя вычисления GPU и CPU с помощью важностно-осведомленного пайплайнинга, достигая до 5 раз ускорения по сравнению с ZeRO-Offload и снижая задержки GPU более чем на 85%.
Как работает ZenFlow
ZenFlow включает несколько инновационных функций:
- Важностно-осведомленные обновления градиентов: ZenFlow приоритизирует наиболее значимые градиенты для немедленных обновлений на GPU, откладывая менее критичные градиенты для асинхронного накопления на CPU. Это сокращает трафик градиентов на каждом шаге почти на 50% и снижает нагрузку на PCIe в 2 раза по сравнению с ZeRO-Offload.
- Асинхронное накопление градиентов: Неприоритетные градиенты группируются и обновляются асинхронно на CPU, скрывая работу CPU за вычислениями GPU. Это гарантирует активность GPU, максимизируя использование аппаратного обеспечения.
- Легковесный выбор градиентов: ZenFlow заменяет полный сбор градиентов на легковесный прокси-норму градиента по столбцам, сокращая объем коммуникации более чем в 4000 раз с минимальным влиянием на точность. Это позволяет эффективно масштабироваться на многогрупповых кластерах.
- Никаких изменений в коде, минимальная конфигурация: ZenFlow интегрирован в DeepSpeed и требует лишь небольших изменений в конфигурации JSON. Пользователи могут задать параметры, такие как topk_ratio (например, 0.05 для топ-5% градиентов), и включить адаптивные стратегии.
- Автоматическая настройка производительности: Движок адаптирует интервалы обновлений в реальном времени, исключая необходимость в ручной настройке и обеспечивая максимальную эффективность по мере изменения динамики обучения.
Показатели производительности
ZenFlow демонстрирует впечатляющие показатели:
- До 5 раз ускорение по всему процессу.
- Снижение задержек GPU более чем на 85%.
- Приблизительно в 2 раза меньший трафик PCIe.
- Отсутствие потери точности на тестах GLUE.
- Легковесный выбор градиентов для эффективного масштабирования.
- Автонастройка без необходимости ручной настройки параметров.
Практическое применение
ZenFlow служит как расширение для ZeRO-Offload от DeepSpeed, не требуя изменений в коде — только обновления конфигурации в файле JSON DeepSpeed. Пример использования доступен в репозитории DeepSpeedExamples, который включает пример дообучения ZenFlow на тестах GLUE. Пользователи могут выполнить это с помощью простого скрипта, следуя инструкциям по настройке, представленным в README репозитория.
Пример конфигурации
"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
}
}
Как начать
Для получения подробного руководства обратитесь к примеру дообучения DeepSpeed-ZenFlow и официальному учебнику для пошаговых инструкций.
Итог
ZenFlow представляет собой значительный шаг вперед для тех, кто обучает или дообучает большие языковые модели на ограниченных ресурсах GPU. Эффективно устраняя задержки GPU, вызванные CPU, он обеспечивает более высокую производительность и снижает затраты на обучение без ущерба для точности модели. Этот подход особенно полезен для организаций, масштабирующих рабочие нагрузки LLM на разнородном оборудовании или стремящихся максимизировать использование GPU в облачных или локальных кластерах.
Для технических команд сочетание автоматической настройки, минимальной конфигурации и бесшовной интеграции с DeepSpeed делает ZenFlow доступным и мощным инструментом. Предоставленные примеры и документация упрощают быстрые эксперименты и развертывание.
ZenFlow переопределяет оффлоадинг для обучения LLM, обеспечивая обучение без задержек и высокую пропускную способность с минимальными накладными расходами на конфигурацию — это ценный инструмент для всех, кто стремится расширить границы масштабного ИИ.
Посмотрите техническую статью, страницу на GitHub и блог. Также не стесняйтесь подписаться на нас в Twitter и присоединиться к нашему сообществу из более чем 100 000 участников на ML SubReddit, а также подписаться на нашу рассылку.
