Стадии DeepSpeed ZeRO: какие данные обучения распределяются между GPU?
Автоматический перевод
Эта статья была автоматически переведена с оригинальной английской версии.
ZeRO-1 распределяет состояние оптимизатора между GPUs при параллелизме по данным, ZeRO-2 дополнительно распределяет градиенты, а ZeRO-3 — ещё и параметры модели. Проверяйте самую низкую стадию, при которой задача помещается в память и обеспечивает приемлемый throughput. Более высокая стадия экономит больше памяти состояния модели, но может требовать больше обмена данными.
Эти стадии меняют место хранения данных обучения. Они не уменьшают память активаций автоматически.
Пример расчёта памяти с явными допущениями
Предположим, обучение использует веса и градиенты в FP16, а также основные копии весов Adam в FP32 и два момента оптимизатора. На каждый параметр тогда приходится 2 байта для весов, 2 для градиентов и 12 для состояния оптимизатора: всего 16 байт.
Для 7.5 миллиарда параметров на восьми GPUs расчёт памяти состояния модели даёт:
| Стадия | Данные, распределённые между GPUs | GB состояния модели на GPU |
|---|---|---|
| Обычный параллелизм по данным | Никакие | 120 |
| ZeRO-1 | Состояние оптимизатора | 41.25 |
| ZeRO-2 | Состояние оптимизатора и градиенты | 28.125 |
| ZeRO-3 | Все три компонента | 15 |
Это десятичные GB. Расчёт не включает активации, временно собранные параметры, накладные расходы аллокатора и буферы рантайма. Если стек обучения хранит данные с другой точностью, расчёт нужно изменить.
Статья о ZeRO объясняет разделение состояния и анализирует обмен данными. ZeRO-3 восстанавливает нужные параметры через операции all-gather; анализ полного распределения в статье даёт объём обмена примерно в 1.5 раза больше, чем при стандартном параллелизме по данным. Это соотношение получено при допущениях анализа и не гарантируется для любой сети и конфигурации.
Выбирайте по профилю памяти
Если параметры помещаются в память, а сбой вызывают состояния оптимизатора, проверьте стадию 1 или 2. Если полностью реплицированные параметры не помещаются, проверьте стадию 3. Если основную часть памяти занимают активации, сочетайте выбранную стадию с чекпоинтингом активаций или меньшим батчем обучения.
DeepSpeed также поддерживает перенос состояния в CPU и NVMe. Проверьте текущую конфигурацию ZeRO и измерьте затраты на передачу данных на реальном хосте и накопителе. Дополнительная ёмкость может помочь, но медленная передача данных способна снизить throughput.
Раздел о ZeRO в LLM Engineering Guide содержит исполняемый расчёт для другого числа параметров и других групп GPU.