FSDP2

Fully Sharded Data Parallel (FSDP2) divide el modelo, los gradientes y los estados del optimizador en fragmentos entre las GPU. Antes del cálculo, cada GPU reúne un conjunto completo de parámetros de todos los fragmentos y luego los libera. La fragmentación permite entrenar modelos más grandes que la memoria de una sola GPU, a costa de más comunicación que DDP. Usa FSDP cuando tu modelo o los estados del optimizador no quepan en una sola 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

Estrategias de fragmentación

FSDP2 controla el particionamiento con ~TrainingArguments.fsdp_config. Establece fsdp=True para habilitar FSDP y configura reshard_after_forward en la configuración de FSDP para elegir el equilibrio entre memoria y rendimiento.

auto_wrap_policy controla cómo se envuelven los módulos en unidades de FSDP. Su valor predeterminado es "TRANSFORMER_BASED_WRAP", que envuelve las capas transformer del modelo. Sin envoltura ("NO_WRAP"), el modelo completo es una sola unidad de FSDP y se pierde el beneficio de memoria del particionamiento.

Configurar FSDP

Estos campos controlan cómo FSDP2 envuelve, fragmenta y carga el modelo. reshard_after_forward y auto_wrap_policy se tratan en Estrategias de fragmentación.

• cpu_offload descarga los parámetros y los gradientes a la CPU cuando no están en uso para ahorrar memoria de GPU.

• transformer_layer_cls_to_wrap define la capa de transformer que se envolverá en una unidad de FSDP cuando auto_wrap_policy sea "TRANSFORMER_BASED_WRAP". Cada unidad gestiona sus propias operaciones de gather y scatter. Solo los parámetros de la unidad actual se recogen durante la pasada hacia adelante. Los parámetros de las unidades anteriores se liberan para ahorrar memoria.

  Envolver solo el modelo de nivel superior no genera ahorros de memoria de GPU. Envolver cada capa Linear individual hace que la comunicación entre unidades sea muy costosa. Deje este campo vacío y FSDP leerá el valor de la definición del modelo.

• min_num_params establece el número mínimo de parámetros por módulo para el envoltado basado en tamaño. Solo se utiliza cuando auto_wrap_policy es "SIZE_BASED_WRAP".

• state_dict_type controla el formato del checkpoint. El valor predeterminado es "FULL_STATE_DICT" para un único checkpoint compatible con Transformers. Utilice "SHARDED_STATE_DICT" para un archivo de checkpoint por rank, lo cual es más rápido para modelos grandes. Los checkpoints fragmentados solo se pueden cargar de nuevo en FSDP, por lo que guarde un "FULL_STATE_DICT" para el checkpoint final que desee compartir o cargar fuera de FSDP.

• cpu_ram_efficient_loading carga el checkpoint desde el disco solo en el rank 0. Las otras GPUs inicializan un modelo vacío y reciben los pesos mediante broadcast, evitando que múltiples procesos carguen un modelo grande en la RAM de la CPU.

• activation_checkpointing recalcula las activaciones durante la pasada inversa en lugar de almacenarlas. Utilice esto en lugar de gradient checkpointing en TrainingArguments. Establecer ambos provoca un error.

Configure el entrenamiento con FSDP mediante un archivo de configuración de Accelerate o un archivo de configuración de FSDP pasado a 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",
)

Para los modelos de mezcla de expertos, fsdp_size puede combinarse con tp_size y ep_size para fragmentar los expertos en una dimensión de la malla y todo lo demás en la otra. Consulte paralelismo de expertos.

Evaluar con paralelismo de contexto

El paralelismo de contexto (CP) con FSDP2 se configura a través de cp_size en ParallelismConfig de Accelerate. Al iniciar, pase tanto --use_parallelism_config como --parallelism_config_cp_size. Un parallelism_config no vacío en el YAML de Accelerate también activa use_parallelism_config. Un --parallelism_config_cp_size aislado se ignora y mantiene cp_size=1.

Bajo FSDP2, CP divide la secuencia entre los rangos para que cada GPU conserve una porción del contexto (allgather o alltoall). Esto difiere de Ulysses sequence parallelism, que se ejecuta a través de DeepSpeed ALST y utiliza sp_size en la misma configuración. Elige entre CP o Ulysses SP. No establezcas tanto cp_size como sp_size por encima de 1.

evaluate() utiliza el mismo contexto de paralelismo de contexto que train() a través de prediction_step(). Cada rank ejecuta su fragmento de secuencia mientras ese contexto está activo. eval_loss es la pérdida correcta (no multiplicada por cp_size). El early stopping y load_best_model_at_end utilizan esa métrica.

Los logits permanecen fragmentados a lo largo de la secuencia y nada los reensambla en predicciones de longitud completa. Las predicciones y compute_metrics sobre logits no son compatibles con CP.

Reutiliza tu lanzamiento existente de FSDP2 + cp_size cuando necesites eval_loss.

Si creas una subclase de prediction_step() y cuentas los tokens tú mismo, cuenta num_items_in_batch antes de entrar en el contexto de paralelismo de contexto. Mientras ese contexto está activo, los búferes de secuencia se fragmentan in situ entre los rangos. Contar después subestima el lote. El Trainer predeterminado ya cuenta primero, y solo necesitas esto si sobrescribes prediction_step.

Próximos pasos

• Consulte DDP para el entrenamiento paralelo de datos cuando su modelo quepa en una sola GPU.
• Consulte DeepSpeed para la optimización ZeRO y el descargado a NVMe.
• Para FSDP en TPUs con PyTorch/XLA, configure xla, xla_fsdp_settings y xla_fsdp_grad_ckpt en ~TrainingArguments.fsdp_config.
• Lea el capítulo de FSDP de The Ultra-Scale Playbook para obtener más información sobre cómo funciona FSDP.