FSDP2

Fully Sharded Data Parallel (FSDP2) распределяет модель, градиенты и состояния оптимизатора по нескольким GPU. Перед вычислениями каждый GPU собирает полный набор параметров из всех фрагментов, а затем освобождает их. Фрагментация позволяет обучать модели, которые не помещаются в память одного GPU, но требует больше коммуникации, чем DDP. Используйте FSDP, если ваша модель или состояния оптимизатора не помещаются на одном GPU.

                      ┌─────────────────┐
                      │  training data  │
                      └────────┬────────┘
            ┌──────────────────┼──────────────────┐
            │ shard 0          │ shard 1          │ shard 2
            ▼                  ▼                  ▼
     ┌─────────────┐    ┌─────────────┐    ┌─────────────┐
     │  param      │    │  param      │    │  param      │
     │  shard 0    │    │  shard 1    │    │  shard 2    │
     │  GPU 0      │    │  GPU 1      │    │  GPU 2      │
     └──────┬──────┘    └──────┬──────┘    └──────┬──────┘
            │                  │                  │
            └──────── all-gather (params) ────────┘
                               │
                    full params on each GPU
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
         forward             forward             forward
            │                  │                  │
            └───── reduce-scatter (grads) ────────┘
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
     grad shard 0       grad shard 1       grad shard 2
     optim shard 0      optim shard 1      optim shard 2
        step               step               step

Стратегии шардирования

FSDP2 управляет шардированием с помощью ~TrainingArguments.fsdp_config. Установите fsdp=True, чтобы включить FSDP, и задайте reshard_after_forward в конфигурации FSDP, чтобы выбрать баланс между потреблением памяти и пропускной способностью.

auto_wrap_policy управляет тем, как модули оборачиваются в единицы FSDP. По умолчанию установлено значение "TRANSFORMER_BASED_WRAP", которое оборачивает трансформерные слои модели. Без оборачивания ("NO_WRAP") вся модель становится одной единицей FSDP, и вы теряете преимущество в памяти от шардирования.

Настройка FSDP

Эти поля управляют тем, как FSDP2 оборачивает, разделяет и загружает модель. reshard_after_forward и auto_wrap_policy описаны в разделе Стратегии шардирования.

• cpu_offload переносит параметры и градиенты на CPU, когда они не используются, чтобы экономить память GPU.

• transformer_layer_cls_to_wrap определяет слой трансформера, который нужно обернуть в единицу FSDP, когда auto_wrap_policy равно "TRANSFORMER_BASED_WRAP". Каждая единица управляет собственными операциями сбора и распределения. Во время прямого прохода собираются только параметры текущей единицы. Параметры предыдущих единиц освобождаются для экономии памяти.

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

• min_num_params задаёт минимальное количество параметров в модуле для обёртки на основе размера. Используется только тогда, когда auto_wrap_policy равно "SIZE_BASED_WRAP".

• state_dict_type управляет форматом контрольной точки. По умолчанию установлено значение "FULL_STATE_DICT" для создания одной контрольной точки, совместимой с Transformers. Используйте "SHARDED_STATE_DICT", чтобы получить по одному файлу контрольной точки на каждый ранг; это быстрее для больших моделей. Разбитые контрольные точки можно загрузить только обратно в FSDP, поэтому для финальной контрольной точки, которую вы хотите распространять или загружать вне FSDP, сохраняйте её в формате "FULL_STATE_DICT".

• cpu_ram_efficient_loading загружает контрольную точку с диска только на ранге 0. Остальные GPU инициализируют пустую модель и получают веса через broadcast, что исключает загрузку большой модели в оперативную память CPU несколькими процессами.

• activation_checkpointing пересчитывает активации во время обратного прохода вместо их хранения. Используйте это вместо gradient checkpointing в TrainingArguments. Установка обоих параметров приводит к ошибке.

Настройте обучение с FSDP с помощью файла конфигурации Accelerate или файла конфигурации FSDP, переданного в fsdp_config.

accelerate launch train.py

{
  "version": 2,
  "reshard_after_forward": true,
  "cpu_offload": false,
  "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
  "transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"],
  "state_dict_type": "FULL_STATE_DICT",
  "cpu_ram_efficient_loading": true,
  "activation_checkpointing": true
}

from transformers import TrainingArguments

TrainingArguments(
    ...,
    fsdp=True,
    fsdp_config="path/to/fsdp.json",
)

Для моделей mixture-of-experts параметр fsdp_size можно использовать совместно с tp_size и ep_size, чтобы распределить эксперты по одному измерению сетки, а всё остальное — по другому. См. экспертный параллелизм.

Оценка с параллелизмом контекста

Контекстный параллелизм (CP) с FSDP2 настраивается через cp_size в ParallelismConfig Accelerate. При запуске передавайте оба параметра: --use_parallelism_config и --parallelism_config_cp_size. Непустой parallelism_config в YAML-файле Accelerate также включает use_parallelism_config. Если указан только --parallelism_config_cp_size, он игнорируется, и cp_size=1 остаётся без изменений.

В FSDP2 CP распределяет последовательность по рангам, так что каждый GPU хранит срез контекста (allgather или alltoall). Это отличается от Ulysses sequence parallelism, который работает через DeepSpeed ALST и использует sp_size в том же конфиге. Выберите CP или Ulysses SP. Не устанавливайте оба параметра cp_size и sp_size больше 1.

evaluate() использует тот же контекст параллелизма контекста, что и train(), через prediction_step(). Каждый ранг обрабатывает свой фрагмент последовательности, пока этот контекст активен. eval_loss — это корректное значение функции потерь (не умноженное на cp_size). Ранняя остановка и load_best_model_at_end используют эту метрику.

Logits остаются разбитыми по последовательности, и ничего не собирает их в предсказания полной длины. Предсказания и compute_metrics по logits не поддерживаются при CP.

Повторно используйте существующий запуск FSDP2 + cp_size, когда вам нужен eval_loss.

Если вы переопределяете prediction_step() и считаете токены самостоятельно, вычислите num_items_in_batch до входа в контекст контекстно-параллельного режима. Пока этот контекст активен, буферы последовательностей разделяются на месте между рангами. Подсчёт после этого даст заниженное значение для батча. По умолчанию Trainer уже выполняет подсчёт заранее, поэтому вам это нужно только в том случае, если вы переопределяете prediction_step.

Дальнейшие шаги

• См. DDP для обучения с параллелизмом данных, если ваша модель помещается на одном GPU.
• См. DeepSpeed для оптимизации ZeRO и выгрузки в NVMe.
• Для FSDP на TPU с PyTorch/XLA задайте xla, xla_fsdp_settings и xla_fsdp_grad_ckpt в ~TrainingArguments.fsdp_config.
• Read the FSDP chapter from The Ultra-Scale Playbook for more information about how FSDP works.