FSDP2

Fully Sharded Data Parallel (FSDP2) répartit le modèle, les gradients et les états de l’optimiseur sur plusieurs GPU. Avant le calcul, chaque GPU rassemble un jeu complet de paramètres à partir de tous les shards, puis les libère ensuite. Cette répartition permet d’entraîner des modèles plus volumineux que la mémoire d’un seul GPU, au prix d’une communication plus importante que DDP. Utilisez FSDP lorsque votre modèle ou les états de votre optimiseur ne tiennent pas sur un seul 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

Stratégies de sharding

FSDP2 contrôle le fractionnement via ~TrainingArguments.fsdp_config. Définissez fsdp=True pour activer FSDP, et définissez reshard_after_forward dans la configuration FSDP pour choisir le compromis entre mémoire et débit.

auto_wrap_policy contrôle la manière dont les modules sont regroupés en unités FSDP. Par défaut, il vaut "TRANSFORMER_BASED_WRAP", ce qui enveloppe les couches transformer du modèle. Sans regroupement ("NO_WRAP"), le modèle entier constitue une seule unité FSDP et vous perdez l’avantage mémoire apporté par le fractionnement.

Configurer FSDP

Ces champs contrôlent la manière dont FSDP2 enveloppe, fragmente et charge le modèle. reshard_after_forward et auto_wrap_policy sont abordés dans Stratégies de fragmentation.

• cpu_offload transfère les paramètres et les gradients vers le CPU lorsqu'ils ne sont pas utilisés, afin d'économiser la mémoire GPU.

• transformer_layer_cls_to_wrap définit la couche de transformer à encapsuler dans une unité FSDP lorsque auto_wrap_policy est "TRANSFORMER_BASED_WRAP". Chaque unité gère ses propres opérations de gather et de scatter. Seuls les paramètres de l’unité courante sont rassemblés pendant le passage avant. Les paramètres des unités précédentes sont libérés pour économiser de la mémoire.

  Envelopper uniquement le modèle de plus haut niveau ne permet aucune économie de mémoire GPU. Envelopper chaque couche Linear individuelle rend la communication entre unités très coûteuse. Laissez ce champ vide et FSDP lira la valeur depuis la définition du modèle.

• min_num_params définit le nombre minimal de paramètres par module pour l’emballage basé sur la taille. Il n’est utilisé que lorsque auto_wrap_policy est "SIZE_BASED_WRAP".

• state_dict_type contrôle le format du point de contrôle. La valeur par défaut est "FULL_STATE_DICT" pour un point de contrôle unique compatible avec Transformers. Utilisez "SHARDED_STATE_DICT" pour obtenir un fichier de point de contrôle par rang, ce qui est plus rapide pour les grands modèles. Les points de contrôle fragmentés ne peuvent être rechargés que dans FSDP ; conservez donc un "FULL_STATE_DICT" pour le point de contrôle final que vous souhaitez partager ou charger en dehors de FSDP.

• cpu_ram_efficient_loading charge le checkpoint depuis le disque uniquement sur le rang 0. Les autres GPU initialisent un modèle vide et reçoivent les poids par diffusion, évitant ainsi que plusieurs processus chargent un grand modèle dans la RAM CPU.

• activation_checkpointing recalcule les activations lors de la passe arrière au lieu de les stocker. Utilisez ceci à la place de gradient checkpointing dans TrainingArguments. Définir les deux provoque une erreur.

Configurez l’entraînement FSDP avec soit un fichier de configuration Accelerate, soit un fichier de configuration FSDP passé à 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",
)

Pour les modèles mixture-of-experts, fsdp_size peut être combiné avec tp_size et ep_size pour sharder les experts sur une dimension du mesh et tout le reste sur l'autre. Voir expert parallelism.

Évaluer avec le parallélisme de contexte

Le parallélisme de contexte (CP) avec FSDP2 est configuré via cp_size de ParallelismConfig d’Accelerate. Au lancement, passez à la fois --use_parallelism_config et --parallelism_config_cp_size. Un parallelism_config non vide dans le fichier YAML d’Accelerate active également use_parallelism_config. Un --parallelism_config_cp_size isolé est ignoré et maintient cp_size=1.

Sous FSDP2, CP répartit la séquence entre les rangs de sorte que chaque GPU détient une portion du contexte (allgather ou alltoall). Cela diffère du parallélisme de séquence Ulysses, qui passe par DeepSpeed ALST et utilise sp_size dans la même configuration. Choisissez CP ou SP Ulysses. Ne définissez pas à la fois cp_size et sp_size à une valeur supérieure à 1.

evaluate() utilise le même contexte de parallélisme de contexte que train() via prediction_step(). Chaque rang exécute son fragment de séquence pendant que ce contexte est actif. eval_loss est la perte correcte (non multipliée par cp_size). L’arrêt anticipé et load_best_model_at_end utilisent cette métrique.

Les logits restent fragmentés le long de la séquence, et rien ne les réassemble en prédictions de longueur complète. Les prédictions et compute_metrics sur les logits ne sont pas pris en charge sous CP.

Réutilisez votre lancement existant FSDP2 + cp_size lorsque vous avez besoin de eval_loss.

Si vous sous-classez prediction_step() et comptez les jetons vous-même, comptez num_items_in_batch avant d’entrer dans le contexte du parallélisme de contexte. Tant que ce contexte est actif, les tampons de séquences sont partagés sur place entre les rangs. Compter ensuite sous-estime le lot. Le Trainer par défaut compte d’abord, et vous n’avez besoin de cela que si vous surchargez prediction_step.

Étapes suivantes

• Voir DDP pour l’entraînement parallèle sur les données lorsque votre modèle tient sur une seule GPU.
• Voir DeepSpeed pour l'optimisation ZeRO et le déchargement NVMe.
• Pour FSDP sur TPU avec PyTorch/XLA, définissez xla, xla_fsdp_settings et xla_fsdp_grad_ckpt dans ~TrainingArguments.fsdp_config.
• Lisez le chapitre FSDP du The Ultra-Scale Playbook pour plus d'informations sur le fonctionnement de FSDP.