Étapes de DeepSpeed ZeRO : quel état de l’entraînement est réparti ?
Traduction automatique
Cet article a été traduit automatiquement depuis la version originale en anglais.
ZeRO-1 répartit l’état de l’optimiseur entre les GPUs en parallélisme de données ; ZeRO-2 répartit aussi les gradients, et ZeRO-3 également les paramètres du modèle. Testez l’étape la plus basse qui permet à la charge de travail de tenir en mémoire avec un débit acceptable. Une étape supérieure économise davantage de mémoire pour l’état du modèle, mais peut nécessiter plus de communication.
Ces étapes changent l’emplacement de stockage de l’état de l’entraînement. Elles ne réduisent pas automatiquement la mémoire des activations.
Un exemple de mémoire avec des hypothèses explicites
Supposons que l’entraînement utilise des poids et des gradients en FP16, ainsi que les poids maîtres d’Adam en FP32 et deux moments de l’optimiseur. Cela alloue 2 octets aux poids, 2 aux gradients et 12 à l’état de l’optimiseur par paramètre : 16 octets au total.
Pour 7.5 milliards de paramètres sur huit GPUs, le calcul de l’état du modèle donne :
| Étape | État réparti entre les GPUs | GB d’état du modèle par GPU |
|---|---|---|
| Parallélisme de données classique | Aucun | 120 |
| ZeRO-1 | État de l’optimiseur | 41.25 |
| ZeRO-2 | État de l’optimiseur et gradients | 28.125 |
| ZeRO-3 | Les trois composants | 15 |
Ces valeurs sont en GB décimaux. Elles excluent les activations, les paramètres temporairement rassemblés, le surcoût de l’allocateur et les tampons du runtime. Une pile d’entraînement qui stocke les données avec d’autres précisions nécessite un calcul différent.
L’article sur ZeRO explique les partitions de l’état et analyse leur communication. ZeRO-3 reconstitue les paramètres nécessaires par des opérations all-gather ; l’analyse de la répartition complète dans l’article atteint environ 1.5 fois le volume de communication du parallélisme de données classique. Ce rapport résulte d’une analyse fondée sur ses propres hypothèses, et ne constitue pas une garantie pour tous les réseaux et toutes les configurations.
Choisir selon le profil mémoire
Si les paramètres tiennent en mémoire et que les états de l’optimiseur provoquent l’échec, testez l’étape 1 ou 2. Si les paramètres entièrement répliqués ne tiennent pas en mémoire, testez l’étape 3. Si les activations dominent, combinez l’étape choisie avec le checkpointing des activations ou un lot d’entraînement plus petit.
DeepSpeed permet aussi de déporter l’état vers le CPU et le stockage NVMe. Consultez la configuration actuelle de ZeRO et mesurez le coût des transferts sur la machine et le stockage réellement utilisés. La capacité supplémentaire obtenue peut aider, mais des transferts lents peuvent réduire le débit.
La section ZeRO du LLM Engineering Guide fournit un calcul exécutable pour d’autres nombres de paramètres et groupes de GPU.