FSDP ou DeepSpeed ZeRO : lequel convient à ma configuration d’entraînement ?
Traduction automatique
Cet article a été traduit automatiquement depuis la version originale en anglais.
Commencez par FSDP lorsqu’une infrastructure d’entraînement PyTorch nécessite un partitionnement distribué de l’état. Testez DeepSpeed si vous avez besoin de ses niveaux ZeRO sélectionnables ou de ses fonctions de déchargement vers le CPU et le NVMe. Aucun framework n’offre systématiquement les meilleures performances ; comparez la configuration requise sur votre cluster.
Les deux peuvent répartir les paramètres, les gradients et l’état de l’optimiseur entre les GPUs. Leurs APIs, leurs politiques et leur gestion des checkpoints diffèrent.
Le partitionnement complet suit le même processus de base
Avec un partitionnement complet des paramètres, chaque GPU stocke une partie du modèle. Avant de calculer un groupe de couches, le framework rassemble les paramètres nécessaires à ce groupe. Il effectue ensuite le calcul, libère les paramètres rassemblés selon sa politique de repartitionnement, puis combine et répartit les gradients avec reduce-scatter.
Les choix de regroupement et de préchargement déterminent la mémoire temporaire et le chevauchement de la communication avec le calcul. Partitionner chaque petite opération séparément peut rendre la communication prépondérante ; rassembler de très grands groupes peut créer un pic temporaire de mémoire.
| Besoin | Candidat à tester |
|---|---|
| APIs natives de PyTorch et état de tenseurs distribués | FSDP2 |
| Partitionnement de l’état de l’optimiseur avec paramètres toujours répliqués | DeepSpeed ZeRO-1 ou ZeRO-2 |
| Partitionnement complet de l’état | FSDP ou ZeRO-3 |
| Déchargement de l’état vers le NVMe | DeepSpeed ZeRO-Infinity |
Le tutoriel FSDP2 utilise fully_shard et des paramètres DTensor. Il propose aussi des politiques de repartitionnement : conserver les paramètres après la passe avant consomme davantage de mémoire, mais réduit les rassemblements ultérieurs. Créez l’optimiseur après avoir appliqué les transformations de partitionnement nécessaires.
Comparez une étape d’entraînement complète
Gardez constants le modèle, les données, la précision, la longueur de séquence, la taille de lot effective et la topologie des GPUs. Mesurez le pic de mémoire allouée, la durée d’une étape et la communication. Vérifiez qu’un checkpoint sauvegardé peut restaurer l’état du modèle, l’état de l’optimiseur et la progression d’entraînement prévue.
Hugging Face Accelerate peut sélectionner FSDP ou DeepSpeed via des configurations de lancement pour du code préparé à ces intégrations. Cela ne signifie pas qu’un script d’entraînement quelconque peut basculer sans modification, notamment pour les checkpoints, le déchargement et la gestion d’optimiseurs personnalisés.
La comparaison de FSDP dans le LLM Engineering Guide relie ces choix à la mémoire utilisée par l’état d’entraînement.