Mettre à l’échelle les LLMs avec le parallélisme multi-GPU et multi-nœud

Traduction automatique Cet article a été traduit automatiquement depuis la version originale en anglais.

Les charges de travail impliquant de grands modèles dépassent les capacités d’un seul GPU pour différentes raisons. Un job d’entraînement peut manquer de mémoire à cause de l’état de l’optimiseur. Un autre peut être limité par les activations de séquences longues. Même lorsqu’un modèle tient en mémoire, il peut ne pas atteindre son objectif de débit. Chaque problème appelle une partition et un schéma de communication différents.

Ce guide pratique présente les principales stratégies de parallélisme et les contraintes qui les sous-tendent, en s’appuyant sur l’Ultra-Scale Playbook de Hugging Face. L’objectif est de montrer ce que chaque partitionnement apporte, quelles données il communique et quand des combinaisons deviennent nécessaires.

TL;DR. Le parallélisme de données répliqué augmente le débit d’entraînement lorsqu’une réplique tient sur un appareil. Le parallélisme de données entièrement shardé partitionne l’état du modèle, mais ajoute des all-gathers de paramètres et des reduce-scatters de gradients. Le parallélisme tensoriel, pipeline, contextuel et par experts partitionne respectivement les calculs des couches, la profondeur, la séquence et les couches mixture-of-experts (MoE). Ne les combinez qu’après avoir identifié la contrainte principale de mémoire ou de communication.

Ce guide suppose que vous maîtrisez la rétropropagation, les couches Transformer et une boucle d’entraînement PyTorch standard.

Commencer par deux budgets mémoire

L’entraînement et l’inférence n’ont pas la même empreinte mémoire.

training peak ≈ parameters
              + gradients
              + optimizer state
              + saved activations
              + temporary buffers
              + communication buffers
              + allocator headroom

inference peak ≈ resident weights
               + key/value (KV) cache
               + runtime workspace
               + communication buffers
               + allocator headroom

Un modèle de 70 milliards de paramètres a une limite inférieure de 140 Go décimaux pour les seuls poids en BF16. Ce chiffre renseigne peu sur l’entraînement, où les gradients, l’état de l’optimiseur, les poids maîtres et les activations peuvent dominer. Il ne permet pas non plus de dimensionner le serving, où la politique de cache, la longueur des séquences, la concurrence des batchs et la quantification jouent un rôle.

Profilez l’architecture exacte, la précision, la longueur des séquences, le micro-batch, l’optimiseur, la politique de checkpointing et le runtime. Relevez la mémoire maximale allouée et réservée, le nombre de tokens par seconde, le temps passé dans les kernels et le temps exposé dans les opérations collectives.

Le passage à l’échelle distribué commence par un goulot d’étranglement mesuréLe passage à l’échelle distribué commence par un goulot d’étranglement mesuré

Chaque dimension de parallélisme implique un compromis

Pour chaque stratégie, demandez-vous quelle dimension du tenseur est partitionnée, quel état est répliqué et quelle opération collective entre dans le chemin critique.

StratégiePartitionneGain principalCommunication introduite
Parallélisme de données répliquébatchdébit d’entraînementall-reduce des gradients
Parallélisme de données entièrement shardéparamètres, gradients et état de l’optimiseur entre les membres d’un groupe de parallélisme de données (DP)mémoire de l’état du modèleall-gather des paramètres, reduce-scatter des gradients
Parallélisme tensorieldimensions des matrices ou de l’attention à l’intérieur des couchespoids et activations des couchesopérations collectives au sein des blocs Transformer
Parallélisme pipelinegroupes de couchesprofondeur du modèle et état par étapeactivations point à point et bulles de scheduling
Parallélisme contextueldimension de la séquencemémoire des activations de séquences longueséchange des clés/valeurs ou de l’attention entre les membres du groupe de séquence
Parallélisme par expertsexperts MoE et tokens routéscapacité d’experts par rankdispatch et combinaison des tokens, généralement via all-to-all

Le gain mémoire n’est pas un multiplicateur fixe. Il dépend du degré de sharding, de ce qui reste répliqué, de l’état temporairement non shardé, de la politique appliquée aux activations, du padding, du déséquilibre et des buffers.

Parallélisme de données répliqué : du débit sans gain de capacité

Entraînement avec parallélisme de données répliquéEntraînement avec parallélisme de données répliqué

Le parallélisme de données répliqué, généralement appelé distributed data parallel, conserve une réplique complète de l’entraînement sur chaque rank. La documentation PyTorch de DistributedDataParallel décrit ce modèle fondé sur la réplication et la synchronisation des gradients. Chaque rank traite un micro-batch différent, puis les gradients sont synchronisés avant l’étape de l’optimiseur.

Utilisez-le lorsque l’état complet de l’entraînement tient en mémoire avec une marge de sécurité suffisante, et lorsque le batch global peut augmenter ou que l’accumulation des gradients peut être ajustée. Ses principaux avantages sont une sémantique simple et une implémentation mature qui recouvre le calcul du backward avec la réduction des gradients par buckets.

Ajouter des ranks peut être contre-productif lorsque le batch local devient trop petit. Cela peut également arriver lorsque le réseau ne parvient pas à masquer l’all-reduce ou lorsque la livraison des entrées provoque des blocages. Le batch d’optimisation souhaité peut ne pas être extensible.

Parallélisme de données entièrement shardé : remplacer la mémoire d’état par des opérations collectives

Exécution avec parallélisme de données entièrement shardéExécution avec parallélisme de données entièrement shardé

Le parallélisme de données entièrement shardé stocke les shards des paramètres, des gradients et de l’optimiseur sur un groupe de ranks. Le papier ZeRO décrit ce partitionnement de l’état, et l’API FSDP2 de PyTorch l’implémente. Les paramètres d’une couche sont réunis par all-gather pour le calcul, puis peuvent être re-shardés. Les gradients sont renvoyés à leurs propriétaires par reduce-scatter.

La documentation actuelle de PyTorch distingue l’API fully_shard utilisée avec le parallélisme de données entièrement shardé version 2 (FSDP2) de l’ancien wrapper FullyShardedDataParallel. FSDP2 regroupe les communications selon les modules auxquels fully_shard est appliqué et recommande une application bottom-up afin que les groupes de couches puissent recouvrir communication et calcul.

from torch.distributed.fsdp import fully_shard
from torch.optim import AdamW

# Apply bottom-up: each block becomes a communication group.
for block in model.transformer.blocks:
    fully_shard(block)

# Shard remaining root parameters such as embeddings and output projection.
fully_shard(model)

# Construct the optimizer after parameters have become sharded distributed tensors (DTensors).
optimizer = AdamW(model.parameters(), lr=learning_rate)

Il s’agit d’un schéma structurel, pas d’un launcher complet. Les device meshes, la mixed precision, le checkpointing, l’initialisation, l’état de l’optimiseur et les distributed checkpoints doivent être cohérents avec la stack d’entraînement.

Le sharding est intéressant lorsque l’état du modèle constitue la contrainte principale et que le calcul des couches peut masquer suffisamment le trafic des opérations collectives. Il peut représenter un mauvais compromis pour les petits modèles, les liens lents, les petites couches ou les configurations dont le groupe de shards traverse une mauvaise frontière topologique.

Parallélisme tensoriel : partitionner les calculs des couches

Partitionnement tensoriel d’une couchePartitionnement tensoriel d’une couche

Le parallélisme tensoriel partitionne l’algèbre linéaire à l’intérieur d’une couche. On peut notamment citer les projections column-parallel et row-parallel. Le guide de parallélisme de NVIDIA documente ce partitionnement au niveau des couches. Les résultats partiels nécessitent des opérations collectives au sein des blocs Transformer ; la latence et la bande passante ont donc un impact répété pendant les passes forward et backward.

Utilisez-le lorsqu’une couche ou ses activations ne tiennent pas en mémoire, ou lorsque les matmuls sont suffisamment grands pour que des kernels partitionnés soient plus performants qu’un calcul sur un seul rank. Mappez le groupe tensor-parallel vers le domaine de communication le plus rapide disponible, puis mesurez. Un degré élevé de parallélisme tensoriel peut réduire chaque matrice locale jusqu’à faire chuter l’efficacité des kernels, tout en augmentant le coût des opérations collectives.

Le sequence parallelism est souvent associé au parallélisme tensoriel pour éviter de répliquer certains calculs sur les activations. Il est distinct du context parallelism appliqué à l’ensemble de la séquence d’entrée du modèle.

Parallélisme pipeline : partitionner la profondeur et le temps de scheduling

Étapes du pipeline et micro-batchesÉtapes du pipeline et micro-batches

Le parallélisme pipeline place différents groupes de couches sur différentes étapes et transmet les activations entre elles. Les micro-batches permettent aux étapes de travailler simultanément. Le papier GPipe utilise ce scheduling pour les réseaux de neurones géants.

Il réduit l’état du modèle par étape et peut diminuer le volume de communication franchissant une frontière plus lente par rapport aux opérations collectives tensor-parallel effectuées couche par couche. Ses coûts comprennent les bulles, les transferts d’activations, le déséquilibre entre étapes, un scheduling plus complexe, ainsi qu’une récupération et un checkpointing plus difficiles.

Pour un scheduling simple de type GPipe, avec p étapes et m micro-batches, la fraction idéalisée de la bulle forward est approximativement :

(p - 1) / (m + p - 1)

Les schedulings réels peuvent utiliser des variantes one-forward/one-backward, interleaved ou zero-bubble, et un coût inégal des couches peut dominer la formule. Choisissez les frontières entre étapes à partir des temps et de la mémoire mesurés, et non du nombre de couches identique.

Parallélisme contextuel : partitionner les activations de séquences longues

Échange d’attention avec parallélisme contextuelÉchange d’attention avec parallélisme contextuel

Le parallélisme contextuel distribue la dimension de la séquence. La documentation NVIDIA sur le parallélisme contextuel décrit ce partitionnement ainsi que l’échange des clés/valeurs nécessaire à l’attention. Chaque rank possède un shard de la séquence, tandis que l’attention échange les informations nécessaires pour préserver la sémantique du contexte complet. Les implémentations peuvent utiliser des anneaux point à point, des all-gathers, des all-to-all ou des combinaisons hiérarchiques.

Cette stratégie réduit la mémoire des activations lors de l’entraînement avec un contexte long, mais réplique les poids au sein du groupe contextuel et introduit des communications liées à l’attention. Le gain dépend du type d’attention, du masquage causal, de la longueur de la séquence, de la recomputation et de la manière dont les groupes contextuels se composent avec les groupes tensoriels et de données.

Ne la sélectionnez pas à partir d’un seuil universel de 8K, 32K ou 100K. Profilez la mémoire des activations et les communications de l’attention pour l’architecture réelle.

Parallélisme par experts : uniquement pour une architecture MoE

Routage des tokens avec parallélisme par expertsRoutage des tokens avec parallélisme par experts

Le parallélisme par experts distribue les experts des couches mixture-of-experts. Le guide de parallélisme de NVIDIA documente ce placement des experts et sa combinaison avec d’autres dimensions de parallélisme. Le routeur envoie les représentations des tokens vers les experts sélectionnés, puis combine leurs résultats. Seuls les experts sélectionnés calculent pour chaque token, mais l’ensemble des poids des experts doit tout de même être stocké et réparti pour le serving.

Le parallélisme par experts n’est pas un bouton d’optimisation pour un modèle dense. Il fait partie d’une architecture MoE. Ses enjeux comprennent l’équilibrage de charge, les limites de capacité, l’all-to-all des tokens, les tokens supprimés ou paddés, les pertes auxiliaires et le déséquilibre en cas de panne. Suivez le nombre de tokens par expert, l’entropie du routage, les dépassements de capacité, le temps de communication et la qualité par route.

Composer une configuration à partir de la topologie

Les systèmes d’entraînement de grands modèles denses sans groupe de parallélisme par experts utilisent couramment le produit des tailles de groupes de parallélisme de données (DP), tensoriel (TP), pipeline (PP) et contextuel (CP) :

world size = DP × TP × PP × CP

Lorsque le parallélisme par experts (EP) constitue un groupe indépendant, le guide de parallélisme de NVIDIA calcule le total comme suit :

total GPUs = TP × PP × CP × EP × DP

Utilisez le mesh pris en charge par le framework plutôt que de multiplier des configurations non supportées.

Construisez la configuration dans cet ordre :

  1. Dessinez les domaines de communication : liens GPU-à-GPU, switches, frontières de non-uniform memory access (NUMA), fabric des nœuds, sursouscription et chemin de stockage.
  2. Placez les opérations collectives fréquentes et sensibles à la latence, généralement TP, dans le domaine adapté le plus rapide.
  3. Choisissez les groupes FSDP ou de DP répliqué à partir de la capacité et de la bande passante restantes.
  4. Ajoutez PP lorsque le placement en profondeur ou le trafic inter-domaines y gagne, en équilibrant le temps et la mémoire mesurés de chaque étape.
  5. Ajoutez CP uniquement pour la contrainte liée à la séquence, et le parallélisme par experts (EP) uniquement pour la topologie d’experts du modèle.
  6. Vérifiez la divisibilité du nombre de têtes, des dimensions cachées, des couches, des experts, du batch et de la séquence pour le mesh candidat.
  7. Évaluez plusieurs meshes valides. Les heuristiques tenant compte de la topologie sélectionnent des candidats, pas des gagnants.

Deux clusters disposant du même nombre de GPUs peuvent préférer des configurations différentes, car la bande passante des liens, la hiérarchie des switches, le rattachement des CPUs et la contention réseau diffèrent.

L’entraînement et le serving nécessitent des décisions distinctes

L’inférence ne transporte généralement ni gradients ni état de l’optimiseur ; les configurations d’entraînement de type FSDP ne se transposent donc pas automatiquement.

Pour le serving, posez-vous les questions suivantes :

  • Une réplique tient-elle en mémoire avec les poids, le KV cache, l’espace de travail et la concurrence cible ?
  • Le débit est-il mieux servi par un plus grand nombre de répliques indépendantes ou par le sharding d’une seule réplique ?
  • Le TP réduit-il suffisamment la pression sur les poids et le cache par rank pour compenser les communications par couche ?
  • Le PP est-il efficacement pris en charge pour le modèle et le scheduler de requêtes ?
  • Comment le prefill et le decode sollicitent-ils différemment le calcul, la bande passante mémoire et l’interconnexion ?
  • Que devient la latence de queue lorsque les requêtes ont des longueurs de prompt et de sortie différentes ?

Évaluez le serveur complet avec son scheduler, sa quantification, sa distribution du contexte, sa politique de batching et la forme du trafic. Le nombre de tokens par seconde en entraînement ne permet pas de prédire le time to first token ni la latence inter-tokens en serving.

Mesurer honnêtement une configuration de mise à l’échelle

Pour chaque configuration candidate, relevez :

  • le modèle, le code, le runtime, les kernels et l’identité de la topologie ;
  • le batch global et local, la distribution des séquences et le nombre de tokens ;
  • la mémoire maximale par catégorie, lorsque cette information est disponible ;
  • le nombre utile de tokens par seconde et l’utilisation des opérations flottantes du modèle (FLOP), lorsque celle-ci est calculée de manière cohérente ;
  • le temps exposé dans les opérations all-reduce, all-gather, reduce-scatter, all-to-all et point à point ;
  • les blocages d’entrée, le temps de checkpointing, le comportement après redémarrage et la distribution des stragglers ;
  • la loss d’entraînement ou la parité des sorties du serving par rapport à la baseline.

Comparez délibérément le strong scaling et le weak scaling. Le strong scaling conserve une quantité totale de travail constante lorsque le nombre de ranks augmente. Le weak scaling conserve une quantité de travail constante par rank ; le travail total augmente donc avec le nombre de ranks. Un pourcentage étiqueté « efficacité du scaling » n’a aucun sens sans connaître ce dénominateur et cette baseline.

Conclusion

Le parallélisme établit une correspondance entre un goulot d’étranglement mesuré, une dimension de tenseur et un schéma de communication. La réplication, le sharding, le partitionnement des couches, le staging, le partitionnement de la séquence et le routage par experts lèvent chacun une contrainte différente et créent un mode de défaillance différent.

Inventoriez la charge de travail, dessinez la topologie, générez des meshes valides et profilez-les. La configuration gagnante est celle qui tient en mémoire avec une marge de sécurité et qui minimise les communications exposées pour le job réellement exécuté.

Références