DeepSpeed ZeRO

DeepSpeed ZeRO (Zero Redundancy Optimizer) élimine la redondance mémoire dans l’entraînement distribué en fragmentant les états de l’optimiseur, les gradients et les paramètres sur plusieurs GPU. ZeRO comporte trois stades, chacun fragmentant davantage d’états que le précédent. DeepSpeed prend également en charge le déchargement vers la mémoire CPU ou NVMe pour des économies supplémentaires. Chaque stade et niveau de déchargement supplémentaire réduit la mémoire de pointe, au prix d’une communication inter-GPU accrue.

                  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 sharde les gradients et les états de l’optimiseur avec une surcharge de communication inférieure à celle de ZeRO-3. N’utilisez ZeRO-3 que si votre modèle ne tient pas sur les GPU avec ZeRO-2.

Installation

Installez DeepSpeed depuis PyPI, ou installez Transformers avec l'extra deepspeed.

pip install deepspeed
pip install transformers[deepspeed]

En cas d’erreurs d’installation liées à CUDA, consultez la documentation DeepSpeed CUDA. Installer à partir du code source est l’option la plus fiable, car elle correspond exactement à votre matériel et inclut des fonctionnalités non encore disponibles dans la version publiée sur PyPI.

Configurer

Trainer intègre DeepSpeed via l’argument deepspeed, qui accepte un fichier de configuration JSON. En alternative, utilisez un fichier de configuration Accelerate à la place de TrainingArguments.

TrainingArguments

Fichier de configuration d'Accelerate

Utilisez "auto" dans votre configuration pour les valeurs que vous souhaitez que DeepSpeed remplisse à partir de TrainingArguments. Si vous souhaitez spécifier explicitement une valeur, assurez-vous d'utiliser la même valeur pour l'argument DeepSpeed et TrainingArguments.

Consultez la référence DeepSpeed Configuration JSON pour obtenir la liste complète des options de configuration DeepSpeed.

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

Configuration 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`

Passez la configuration à l'argument 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

Étapes de ZeRO

Sélectionnez une configuration de stage ZeRO à utiliser comme point de départ.

ZeRO-1

ZeRO-2

ZeRO-3

Les champs suivants sont importants pour personnaliser l'entraînement.

zero_optimization définit le stade ZeRO.
```
{ "zero_optimization": { "stage": 3 } }
```
Définissez les arguments de taille de lot et d’accumulation des gradients sur "auto". Si vous les définissez manuellement avec des valeurs qui ne correspondent pas à TrainingArguments, l’entraînement se poursuit silencieusement avec les valeurs incorrectes.
```
{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto"
}
```
bf16 définit la précision de l’entraînement. Définissez-le sur "auto" afin qu’il reflète le drapeau bf16 dans TrainingArguments.
```
{ "bf16": { "enabled": "auto" } }
```
stage3_gather_16bit_weights_on_model_save effectue un all-gather sur tous les GPU avant l'enregistrement, reconstruisant les tenseurs complets à partir de leurs fragments. Il s'agit d'un argument ZeRO-3.
```
{
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
```
Définissez overlap_comm sur true pour masquer la latence de l’opération all-reduce derrière la passe arrière. allgather_bucket_size et reduce_bucket_size arbitrent entre la vitesse de communication et la mémoire GPU. Des valeurs plus faibles utilisent moins de mémoire mais ralentissent la communication.
```
{
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}
```
offload_optimizer déplace l’optimiseur en mémoire CPU. Pour économiser encore plus de mémoire, déplacez également les paramètres du modèle avec offload_param (ZeRO-3 uniquement). Définissez pin_memory sur true pour accélérer les transferts CPU-GPU, mais cela verrouille la RAM, la rendant indisponible pour les autres processus.
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    }
}
```

optimizer et scheduler sont par défaut l’optimiseur et le planificateur configurés dans TrainingArguments. Définissez-les sur "auto" afin que DeepSpeed lise les valeurs depuis TrainingArguments, sauf si vous avez besoin d’un optimiseur natif DeepSpeed comme 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 vous déchargez l’optimiseur, définissez zero_force_ds_cpu_optimizer sur false pour utiliser l’optimiseur CPU Adam de DeepSpeed.

{
    "zero_force_ds_cpu_optimizer": false
}

Points de contrôle

DeepSpeed enregistre les points de contrôle dans un format fragmenté qui ne peut pas être chargé directement avec from_pretrained(). Définissez load_best_model_at_end() sur True pour que Trainer suive et recharge le meilleur point de contrôle à la fin de l’entraînement.

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")

Définir save_only_model=True évite l’enregistrement de l’état complet de l’optimiseur, ce qui signifie qu’il est impossible de recharger le meilleur modèle à la fin de l’entraînement. Définissez également stage3_gather_16bit_weights_on_model_save: true pour reconstruire les poids complets à partir de leurs fragments. Cette étape est requise pour enregistrer un artefact de modèle 16 bits consolidé ou un état dictionnaire 16 bits avec ZeRO-3. Transformers lève une erreur lorsque save_only_model=True est combiné à load_best_model_at_end=True.

Pour reprendre l’entraînement avec des configurations de parallélisme différentes, consultez le guide Universal Checkpointing de DeepSpeed.

Étapes suivantes

Lisez le chapitre Zero Redundancy Optimizer du The Ultra-Scale Playbook pour en savoir plus sur le fonctionnement de ZeRO.
Lisez les articles ZeRO : Memory Optimizations Toward Training Trillion Parameter Models, Democratizing Billion-Scale Model Training et Breaking the GPU Memory Wall for Extreme Scale Deep Learning.

Update on GitHub