Настройка FSDP (Fully Sharded Data Parallel) для обучения

FSDP PyTorch: почему обучение большой модели упирается в память

Обучение большой модели упирается не в число параметров само по себе, а в то, что каждый ускоритель обязан держать у себя. Полное состояние модели — веса, градиенты и состояние оптимизатора — при распределённом обучении дублируется на каждом устройстве. Состояние оптимизатора при этом в несколько раз тяжелее самих весов: на каждый параметр хранится несколько накопленных величин.

Как только суммарный объём перестаёт помещаться в память ускорителя, обучение либо падает с ошибкой выделения, либо начинается вынужденная экономия: режут размер пакета, включают накопление градиентов, дробят слои вручную. Каждый такой шаг — компромисс: свободной памяти больше, но время шага растёт, а часть устройств простаивает в ожидании.

FSDP PyTorch закрывает это шардированием состояния модели: веса, градиенты и состояние оптимизатора разбиваются между устройствами, а не хранятся целиком на каждом. Перед вычислением слоя его параметры собираются на нужных ускорителях, после — снова распускаются, а состояние оптимизатора остаётся разбитым постоянно. Именно там освобождается основная часть памяти, и обучение больших моделей перестаёт зависеть от того, влезает ли модель целиком в один ускоритель.

Разбираем, как это закрывается на практике: какие режимы шардирования выбирать, как настраивать группы устройств и где чаще всего теряется производительность.

Что даёт шардирование состояния модели и оптимизатора

Шардирование состояния модели и оптимизатора распределяет параметры, градиенты и моменты по устройствам вместо дублирования на каждом. Упирается обучение обычно не в саму модель, а в её оптимизатор: веса ещё помещаются, а состояния — уже нет. В PyTorch FSDP это включается режимом шардирования, и дальше результат зависит от того, как настроены обмен и контрольные точки.

Проверяем по пунктам, что меняется измеримо.

• Размер пакета. Освобождённая память уходит под активации, поэтому глобальный размер пакета растёт без смены оборудования. Больший пакет даёт более ровный градиент и позволяет поднять скорость обучения без потери сходимости.
• Время шага обучения. Обмен градиентами перекрывается вычислениями, поэтому шаг удлиняется заметно медленнее, чем растёт объём данных. Итоговое время обучения сокращается кратно, а не на проценты.
• Устойчивость к сбоям. Шардированное состояние собирается в контрольную точку без дублирования, поэтому падение отдельного процесса не означает потерю прогресса. Уходит класс отказов, связанных с переполнением памяти на одном устройстве.
• Время восстановления. Запуск из последней контрольной точки занимает минуты: состояние поднимается в том же разбиении, и обучение продолжается с того же шага.

Проверять это стоит на своём профиле нагрузки — у разных архитектур узкое место смещается между памятью оптимизатора, активациями и сетью обмена. Мы снимаем замеры до и после, фиксируем размер пакета, время шага и время восстановления, и только затем считаем режим настроенным.

Сравнение режимов шардирования и их ограничений

Выбор режима определяется двумя величинами: сколько памяти освобождается на каждом ускорителе и во что обходится обмен градиентами и параметрами. Полное шардирование режет и параметры, и градиенты, и состояния оптимизатора — выигрыш по памяти максимальный, но и пересылок больше всего.

Группы параметров позволяют точечно выводить из шардирования отдельные слои — эмбеддинги, нормировки, выходную голову — когда их дробление даёт больше обмена, чем экономии. Ниже — три режима и границы, за которыми каждый проигрывает.

Режим закрепляют по замерам, а не по описанию: смотрят долю времени в обмене и утилизацию памяти на рабочем размере пакета. В PyTorch переключение между режимами — это смена конфигурации обёртки, саму модель переписывать не нужно.

Как проходит внедрение: от аудита до передачи команде

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

1. Аудит конфигурации обучения. Разбираем текущий рецепт: размер батча и его разбиение, режим точности, состав слоёв, планировщик скорости обучения. Фиксируем базовую линию — время шага, потребление памяти, кривую сходимости на коротком прогоне.

2. План шардирования. Решаем, что разносить по устройствам: параметры, градиенты, состояние оптимизатора, а что оставить реплицированным. Определяем, какие блоки оборачивать целиком, чтобы не платить за лишний обмен.

3. Настройка FSDP. Собираем конфигурацию под выбранный план: политика обёртки, порядок перестроения параметров, накопление градиентов. Проверяем, что запуск воспроизводится на том же числе устройств.

4. Проверка сходимости на реальной задаче. Сравниваем с базовой линией по шагам, а не только по финальной метрике: расхождение в начале указывает на точность или порядок обмена. Правим до выхода на полный прогон.

5. Контрольные точки. Настраиваем сохранение и загрузку распределённого состояния, восстановление после сбоя узла и ротацию: прогон продолжается с последней точки, а не с нуля.

6. Передача команде. Отдаём конфигурации, схему шардирования, регламент запуска и разбора сбоев, проводим совместный прогон на вашем стенде. Дальше инженеры меняют рецепт сами.

Конфигурации, схемы и регламенты, которые получает заказчик

Всё, что остаётся после внедрения, лежит в вашем репозитории и в документации, а не в переписке и не в голове отдельного инженера. Человек, который придёт на проект после нас, разворачивает обучение по тем же файлам и получает тот же результат. Ниже — что именно фиксируется и передаётся.

• Конфигурации запуска — параметры обучения и разбивки модели в системе контроля версий: изменения проходят согласование, а не правятся прямо на сервере, поэтому любой запуск воспроизводим.
• Схемы шардирования — как FSDP в PyTorch раскладывает параметры и состояния оптимизатора по устройствам и сколько памяти резервирует, плюс расчёт под вашу модель и число узлов.
• Политики доступа — кто вправе запускать обучение, менять конфигурации и выгружать веса; роли разведены, доступ к данным и артефактам минимально необходимый.
• Регламент восстановления — периодичность контрольных точек, срок хранения и порядок перезапуска с последней сохранённой после отказа узла.
• Документация — развёртывание с нуля, описание конфигураций и порогов, при которых обучение деградирует, разбор типовых ошибок.
• Наблюдаемость — кривые потерь и пропускной способности по шагам: сразу видно, упирается ли обучение в данные, сеть или ускоритель.

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

Кейс: обучение большой модели перестало упираться в память

Задача была типовой для больших моделей: обучение не влезало в память ускорителей при осмысленном размере пакета. Команда шла на компромиссы — уменьшала пакет, урезала длину контекста, дробила слои вручную, — но каждый второй запуск всё равно падал с ошибкой нехватки памяти. Инженерные часы уходили не на качество модели, а на перезапуски и подбор конфигурации.

Что было до

Базовый замер на реальной задаче: шаг обучения занимал около 3,8 секунды, успешно доходили до контрольной точки 6 запусков из 10, остальные падали по памяти. За рабочую неделю набиралось 12–14 таких отказов, часть ускорителей простаивала, а возобновление после сбоя требовало ручной сборки состояния. Увеличение пакета, которое улучшило бы сходимость, было невозможно в принципе.

Что изменили

Обучение перевели на FSDP в PyTorch — параметры, градиенты и состояния оптимизатора шардируются по устройствам, поэтому полная копия модели больше не обязана жить в памяти одного ускорителя. Дополнительно включили пересчёт активаций на прямом проходе, выровняли размеры пакетов между узлами и перевели сохранение контрольных точек в шардированный вид — состояние собирается только при возобновлении.

Что стало

Шаг обучения сократился до 2,1 секунды, доля успешных запусков выросла до 10 из 10, падений по памяти нет уже месяц. Размер пакета удалось поднять вдвое — сходимость стала ровнее, а откат на предыдущую контрольную точку занимает минуты вместо ручной пересборки. Мы замеряем такие метрики на каждом проекте: они показывают, окупилась ли настройка распределения.

Чем отличается FSDP PyTorch от обычного распределённого обучения?

Классический параллелизм данных держит полную копию модели на каждом ускорителе: процессы берут разные части батча, считают градиенты и синхронизируют их перед шагом оптимизатора. Пока веса, состояния оптимизатора и активации помещаются в память одного устройства, схема масштабируется почти линейно. Как только модель вырастает, реплики перестают влезать целиком — и добавление устройств больше не помогает.

Параллелизм модели решает задачу иначе: слои или блоки разрезаются между устройствами, и каждое держит только свою часть весов. Плата — ручное разбиение, балансировка нагрузки и постоянные пересылки активаций между ступенями. Такой подход масштабирует объём модели, но плохо переносится между архитектурами и требует отдельной отладки под каждую.

FSDP PyTorch занимает промежуточное положение. Параметры, градиенты и состояния оптимизатора шардируются по всем процессам группы: перед вычислением слой собирается целиком, после — освобождается. На каждом ускорителе постоянно живёт лишь доля состояния модели, а не полная копия, при этом цикл обучения остаётся привычным — модули оборачиваются, дальше работа идёт как с обычной моделью.

Отсюда и критерий, когда применять: малая и средняя модель — параллелизм данных, меньше обменов и проще отладка; модель не влезает в память одного ускорителя даже с накоплением градиентов — шардирование состояния. Для сверхглубоких сетей имеет смысл комбинация: FSDP внутри узла, параллелизм данных между узлами. Разница подходов сводится к одному вопросу — масштабируем пропускную способность реплик, объём состояния или и то и другое сразу.

Технические сомнения перед запуском шардированного обучения

Сомнения инженерного уровня обычно одни и те же: встанет ли FSDP на существующий код модели, выдержат ли связки между узлами обмен градиентами, что делать с накопленными контрольными точками и не просядет ли качество при смешанной точности. Разбираем по порядку.

FSDP в PyTorch не требует переписывать архитектуру: он подключается политикой обёртки поверх модулей, а оптимизатор и цикл обучения остаются на месте. Отдельно проверяем требования к оборудованию: пределом масштабирования чаще становится межузловой обмен, а не объём памяти на устройстве.

Что делать с накопленными контрольными точками?

Старые файлы остаются валидными: разбираем их по слоям и собираем в шардированный формат отдельным скриптом, обратная конвертация даёт откат на неразделённый запуск. Проверка — сверка выхода модели до и после конвертации на одном шаге.

Хватит ли связок между узлами и как быть со смешанной точностью?

Замеряем долю времени на обмен между узлами: если она растёт, увеличиваем размер шарда и сокращаем число синхронизаций. Смешанная точность включается на уровне обёрток; риск в пересчёте нормализации снимается приведением типов на границах блоков и коротким прогоном.

Обсудим задачу: что уточнить перед стартом работ

Чем конкретнее вводные по задаче, тем меньше допущений в оценке. Нам не нужен доступ к продакшену на этом этапе — достаточно ответов на вопросы ниже, чтобы понять объём работ и риски.

• Текущий цикл обучения: на чём запускается сейчас, какие узкие места уже видны — нехватка памяти на устройстве, простой ускорителей, долгие паузы на сохранении состояния.
• Размер модели и состав памяти: сколько параметров, какие слои самые тяжёлые, что именно не помещается — веса, градиенты, состояния оптимизатора или промежуточные активации.
• Конфигурация оборудования: число ускорителей, как они связаны между собой, пропускная способность канала связи — от этого зависит выбор схемы шардирования и то, окупится ли она.
• Схема размещения: что раскладываем по устройствам, какие блоки оставляем целиком на каждом, нужна ли перестройка при смене числа узлов.
• Работа с чекпоинтами: в каком виде сохраняем состояние, требуется ли возобновление обучения с той же точки после сбоя и как быстро оно должно происходить.
• Сценарий отказа: что считаем допустимым — потеря шага, перезапуск с последней контрольной точки, восстановление на меньшем числе устройств.
• Критерии приёмки: по каким метрикам судим результат — время шага, ускорение при росте числа устройств, сходимость и итоговое качество модели.

Дальше мы возвращаемся с оценкой и планом: объём работ, сроки, состав команды и точки контроля. Для FSDP в PyTorch отдельным пунктом предлагаем пилотный запуск на вашей модели и реальном участке данных — он показывает фактическое ускорение и снимает вопросы по совместимости до начала основных работ.

Опишите задачу в том виде, в каком она есть — черновиком, без подготовки. Уточняющие вопросы зададим сами, а сроки работ и границы ответственности зафиксируем до старта.