Параллелизм по данным с полным разделением состояний
Определение
FSDP2 — механизм параллельного обучения по данным в PyTorch, который разделяет параметры, градиенты и состояния оптимизатора между участниками. Для вычислений параметры собираются; здесь описана версия PyTorch 2.9.
Также ищут: FSDP · FSDP2 · fully sharded data parallel · fully_shard
Зачем это знать
Память и обмены FSDP2 зависят от групп параметров и групп GPU. Проверьте, какие параметры собираются вместе, как долго хранятся полные копии и между сколькими устройствами делятся состояния.
Карта понятия
Что подтверждено
В PyTorch 2.9 FSDP2 разделяет параметры, градиенты и состояния оптимизатора между участниками группы разделения. Для прямого прохода параметры собираются операцией all-gather; после вычисления градиентов reduce-scatter объединяет градиенты участников и распределяет результат.
Граница применимости: API torch.distributed.fsdp.fully_shard в документации PyTorch 2.9. Описаны основные операции группы разделения, а не полный перечень обменов для любой сетки и настройки.
В PyTorch 2.9 reshard_after_forward=True освобождает полные копии параметров после прямого прохода и повторно собирает их перед обратным; локальные части сохраняются. False оставляет полные копии и избегает повторного сбора. Значение по умолчанию None означает True для некорневого модуля и False для корневого.
Граница применимости: Здесь разобраны True, False и None для fully_shard в PyTorch 2.9; API также допускает целочисленную настройку. Это поведение параметра не гарантирует меньший пик памяти всего запуска или лучшее время шага.
В PyTorch FSDP2 каждый вызов fully_shard формирует группу параметров для совместного обмена, исключая параметры, уже включённые в другие группы. Вызовы делают от вложенных модулей к корню; вызов для корня объединяет оставшиеся параметры. Группы по слоям позволяют уменьшать пик памяти и совмещать обмены с вычислениями. Если вызвать fully_shard только для корня, все параметры окажутся в одной группе; документация обычно не рекомендует такую схему.
Граница применимости: Рекомендация документации PyTorch 2.9 для fully_shard. Речь о группах параметров, а не о составе группы GPU. Возможность снизить пик и совместить обмены с вычислениями не гарантирует ускорение любого запуска; одной группе не приписывается отсутствие экономии всех состояний.
В двумерной DeviceMesh PyTorch FSDP2 делит параметры по второй оси сетки и повторяет их по первой. При форме 2×4 и размещении (Replicate(), Shard(0)) получаются две группы разделения по четыре GPU. Ноль в Shard(0) обозначает ось тензора параметров, а не ось сетки GPU.
Граница применимости: Размещение по умолчанию для двумерной сетки fully_shard в PyTorch 2.9. Размеры заданы в порядке осей сетки. Равные размеры локальных частей в числовом упражнении — отдельное условие упражнения, а не гарантия для любого тензора.
Что проверить
- Зафиксируйте размеры и порядок осей DeviceMesh. В двумерной сетке FSDP2 первая ось отвечает за репликацию, вторая — за разделение.
- Отдельно укажите группы параметров, созданные вызовами fully_shard, и число GPU, между которыми делится каждая группа. Запишите reshard_after_forward.
- Измерьте пик памяти и время полного шага в целевой конфигурации. Наличие FSDP2 само по себе не подтверждает достаточность памяти или ускорение.
Первоисточники
- torch.distributed.fsdp.fully_shard
Документация FSDP2: сбор параметров и объединение с распределением градиентов, группы параметров, одномерная и двумерная DeviceMesh, reshard_after_forward. Ссылка [source] проверена: ведёт на тег v2.9.1, а не на неопределённую текущую ветвь.
Связанные понятия
Курсы по теме
В каталоге
- GPU для обучения с FSDP — Сопоставьте объём памяти GPU и соединения между ними с выбранными группами устройств FSDP и измеренным временем шага.