DeepSpeed ZeRO

DeepSpeed ZeRO (Zero Redundancy Optimizer) elimina la redundancia de memoria en el entrenamiento distribuido al fragmentar los estados del optimizador, los gradientes y los parámetros entre las GPU. ZeRO tiene tres etapas, cada una fragmenta más estado que la anterior. DeepSpeed también admite la descarga a la memoria CPU o NVMe para lograr un mayor ahorro. Cada etapa adicional y nivel de descarga reduce la memoria pico, a costa de una mayor comunicación entre GPU.

                  params        grads       opt states
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-1          │██████████│ │██████████│ │███░░░░░░░│  GPU 0
                │██████████│ │██████████│ │░░░███░░░░│  GPU 1
                │██████████│ │██████████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-2          │██████████│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │██████████│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │██████████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-3          │███░░░░░░░│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │░░░███░░░░│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │░░░░░░████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
  █ resident    ░ held on another GPU

ZeRO-2 fragmenta los gradientes y los estados del optimizador con una sobrecarga de comunicación menor que ZeRO-3. Utilice ZeRO-3 solo cuando su modelo no quepa en las GPUs con ZeRO-2.

Instalación

Instala DeepSpeed desde PyPI, o instala Transformers con el extra deepspeed.

pip install deepspeed
pip install transformers[deepspeed]

Si te encuentras con errores de instalación relacionados con CUDA, consulta la documentación de DeepSpeed CUDA. Instalar desde el código fuente es la opción más fiable, ya que se adapta a tu hardware exacto e incluye funciones que aún no están disponibles en la versión de PyPI.

Configurar

Trainer integra DeepSpeed a través del argumento deepspeed, que acepta un archivo de configuración JSON. Alternativamente, use un archivo de configuración de Accelerate en lugar de TrainingArguments.

TrainingArguments

Archivo de configuración de Accelerate

Use "auto" en su configuración para los valores que quiera que DeepSpeed complete desde TrainingArguments. Si desea especificar un valor explícitamente, asegúrese de usar el mismo valor para el argumento de DeepSpeed y TrainingArguments.

Consulte la referencia de DeepSpeed Configuration JSON para obtener una lista completa de las opciones de configuración de DeepSpeed.

{
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto" }
    },
    "fp16": { "enabled": "auto" }
}

Configuración de DeepSpeed	TrainingArguments
`train_micro_batch_size_per_gpu`	`per_device_train_batch_size`
`gradient_accumulation_steps`	`gradient_accumulation_steps`
`optimizer.params.lr`	`learning_rate`
`fp16.enabled`	`fp16`

Pase la configuración al argumento deepspeed.

from transformers import TrainingArguments

args = TrainingArguments(
    deepspeed="path/to/deepspeed_config.json",
    ...
)

# DeepSpeed launcher
deepspeed --num_gpus 4 train.py

# torchrun
torchrun --nproc_per_node 4 train.py

# Accelerate
accelerate launch --num_processes 4 train.py

Etapas de ZeRO

Seleccione una configuración de etapa de ZeRO para usar como punto de partida.

ZeRO-1

ZeRO-2

ZeRO-3

Los siguientes campos son importantes para personalizar el entrenamiento.

zero_optimization establece la etapa de ZeRO.
```
{ "zero_optimization": { "stage": 3 } }
```
Establezca los argumentos de tamaño de lote y acumulación de gradientes en "auto". Si los establece manualmente en valores que no coinciden con TrainingArguments, el entrenamiento continúa en silencio con los valores incorrectos.
```
{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto"
}
```
bf16 establece la precisión del entrenamiento. Establécelo en "auto" para que refleje la bandera bf16 en TrainingArguments.
```
{ "bf16": { "enabled": "auto" } }
```
stage3_gather_16bit_weights_on_model_save realiza un all-gather en todas las GPUs antes de guardar, reconstruyendo los tensores completos a partir de sus fragmentos. Este es un argumento de ZeRO-3.
```
{
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
```
Establezca overlap_comm en true para ocultar la latencia de all-reduce detrás de la pasada inversa. allgather_bucket_size y reduce_bucket_size intercambian la velocidad de comunicación por memoria de GPU. Los valores más bajos usan menos memoria pero ralentizan la comunicación.
```
{
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}
```
offload_optimizer descarga el optimizador a la memoria CPU. Para ahorrar aún más memoria, también descargue los parámetros del modelo con offload_param (solo ZeRO-3). Establezca pin_memory en true para acelerar las transferencias CPU-GPU, pero esto bloquea la RAM que no estará disponible para otros procesos.
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    }
}
```

optimizer y scheduler se establecen por defecto como el optimizador y el programador configurados en TrainingArguments. Establézcalos en "auto" para que DeepSpeed lea los valores de TrainingArguments, a menos que necesite un optimizador nativo de DeepSpeed como LAMB.

{
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto", "betas": "auto", "eps": "auto", "weight_decay": "auto" }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": { "total_num_steps": "auto", "warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto" }
    }
}

Si estás descargando el optimizador, establece zero_force_ds_cpu_optimizer en false para usar el optimizador CPU Adam de DeepSpeed.

{
    "zero_force_ds_cpu_optimizer": false
}

Puntos de control

DeepSpeed guarda los puntos de control en un formato fragmentado que no se puede cargar directamente con from_pretrained(). Establezca load_best_model_at_end() en True para que Trainer rastree y recargue el mejor punto de control al final del entrenamiento.

from transformers import TrainingArguments, Trainer

args = TrainingArguments(
    deepspeed="ds_config_zero3.json",
    load_best_model_at_end=True,
    ...
)
# after training, save a normal transformers checkpoint
trainer.save_model("./best-model")

Establecer save_only_model=True omite el guardado del estado completo del optimizador, lo que significa que no se puede recargar el mejor modelo al final del entrenamiento. También establezca stage3_gather_16bit_weights_on_model_save: true para reconstruir los pesos completos a partir de sus fragmentos. Esto es necesario para guardar un artefacto de modelo consolidado de 16 bits o un state dict de 16 bits con ZeRO-3. Transformers genera un error cuando save_only_model=True se combina con load_best_model_at_end=True.

Para reanudar con diferentes configuraciones de paralelismo, consulte la guía de Universal Checkpointing de DeepSpeed.

Próximos pasos

Lea el capítulo Zero Redundancy Optimizer de The Ultra-Scale Playbook para obtener más información sobre cómo funciona ZeRO.
Lea los artículos de ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, Democratizing Billion-Scale Model Training y Breaking the GPU Memory Wall for Extreme Scale Deep Learning.

Update on GitHub