FSDP или DeepSpeed ZeRO: что подходит для моей конфигурации обучения?
Автоматический перевод
Эта статья была автоматически переведена с оригинальной английской версии.
Начните с FSDP, если вашему стеку обучения на PyTorch нужен распределённый шардинг состояния. Протестируйте DeepSpeed, если вам нужны его выбираемые стадии ZeRO или выгрузка состояния на CPU и NVMe. Ни один фреймворк не выигрывает по производительности во всех случаях; сравните нужную конфигурацию на своём кластере.
Оба могут распределять параметры, градиенты и состояние оптимизатора между GPUs. Их APIs, политики и работа с чекпоинтами различаются.
Полный шардинг следует одному и тому же базовому процессу
При полном шардинге параметров каждая GPU хранит часть модели. Перед вычислениями для группы слоёв фреймворк собирает параметры, которые нужны этой группе. Затем он выполняет вычисления, освобождает собранные параметры согласно своей политике повторного шардинга и объединяет и распределяет градиенты через reduce-scatter.
Выбор группировки и предварительной загрузки определяет временное потребление памяти и совмещение обмена данными с вычислениями. Если шардировать каждую мелкую операцию отдельно, обмен данными может занимать большую часть времени; сбор параметров для очень больших групп может создавать временный пик потребления памяти.
| Требование | Кандидат для проверки |
|---|---|
| Нативные APIs PyTorch и распределённое состояние тензоров | FSDP2 |
| Шардинг состояния оптимизатора при сохранении репликации параметров | DeepSpeed ZeRO-1 или ZeRO-2 |
| Полный шардинг состояния | FSDP или ZeRO-3 |
| Выгрузка состояния на NVMe | DeepSpeed ZeRO-Infinity |
Руководство по FSDP2 использует fully_shard и параметры DTensor. Оно также описывает политики повторного шардинга: сохранение параметров после прямого прохода требует больше памяти, но сокращает число последующих операций сбора. Создавайте оптимизатор после применения необходимых преобразований шардинга.
Сравнивайте полный шаг обучения
Зафиксируйте модель, данные, точность, длину последовательности, эффективный размер батча и топологию GPUs. Измерьте пиковый объём выделенной памяти, время шага и обмен данными. Проверьте, что сохранённый чекпоинт позволяет восстановить состояние модели, состояние оптимизатора и нужный прогресс обучения.
Hugging Face Accelerate позволяет выбирать FSDP или DeepSpeed через конфигурации запуска для кода, подготовленного к этим интеграциям. Это не означает, что любой скрипт обучения можно переключить без изменений, особенно в части чекпоинтов, выгрузки состояния и работы с нестандартными оптимизаторами.
Сравнение FSDP в LLM Engineering Guide связывает эти решения с памятью, которую занимает состояние обучения.