Notes du FSDP

Nuances de la prélecture FSDP

Pour les regroupements de données en aval chevauchants avec calcul en aval, deux mécanismes sont possibles :

• Prélecture en aval implicite (toujours activée)
• Prélecture en aval explicite (forward_prefetch=True)

Le préchargement implicite vers l'avant consiste à émettre les requêtes All-Gather à partir d'un flux CUDA distinct pour permettre le chevauchement d'une requête All-Gather avec le calcul vers l'avant émis avant elle (du point de vue du processeur). Par exemple, si nous avons : All-Gather de couche 0 -> Calcul vers l'avant de couche 0 -> All-Gather de couche 1 -> ..., alors la requête All-Gather de couche 1 peut chevaucher la requête All-Gather de couche 0, même si le processus léger du microprocesseur l'a émise après. (Le premier calcul All-Gather ne pourra chevaucher aucun autre calcul.)

Le préchargement explicite vers l'avant consiste à modifier l'ordre d'émission des processus légers du microprocesseur : par exemple : All-Gather de couche 0 -> All-Gather de couche 1 -> Calcul vers l'avant de couche 0 -> .... En mode impatient, il est impossible de savoir quelle couche est la couche suivante (par exemple, la couche 1 dans l'exemple) lorsque l'exécution se poursuit sur la couche 0. Par conséquent, la prélecture explicite vers l'avant ne doit être utilisée que pour les modèles dont l'ordre d'exécution est fixe d'itération en itération (que l'on appelle parfois « graphe statique »). FLAVA est un exemple de modèle qui ne satisfait pas à cette contrainte.

La prélecture explicite vers l'avant permet uniquement de gagner du temps lors de l'émission des noyaux de calcul vers l'avant d'une couche, mais au prix de l'allocation du tenseur de sortie du prochain regroupement complet alors que le tenseur actuel est encore utilisé. En émettant le prochain regroupement complet avant les noyaux de calcul vers l'avant actuels, le prochain regroupement complet peut démarrer plus tôt sur le GPU. Pour la plupart des charges de travail LLM, ce n'est pas le cas ; il n'y a donc aucune raison d'activer forward_prefetch=True.

En revanche, pour le mode rétrograde, nous devons utiliser la prélecture explicite vers l'arrière, sinon il n'y aura aucun chevauchement entre la communication et le calcul. Cela s'explique par le fait que nous utilisons un seul groupe de processus NCCL pour les opérations de regroupement et de réduction de dispersion (en partie parce que, dans les versions antérieures de NCCL, il n'était pas sûr d'en utiliser plusieurs simultanément sur le même périphérique et sur les mêmes rangs). Un seul groupe de processus NCCL signifie un seul flux NCCL interne sur lequel les opérations de réduction de dispersion et de regroupement s'exécutent en série. Ainsi, à moins de réorganiser explicitement l'ordre des émissions CPU pour qu'elles soient du regroupement suivant au regroupement actuel, la réduction de dispersion actuelle bloquerait le regroupement suivant et donc le calcul inverse suivant, empêchant ainsi le chevauchement de la réduction de dispersion actuelle.

Taille de la charge utile de communication

Dans FSDP, les communications sont :

• Toutes les communications sont regroupées sur les paramètres en aval
• Toutes les communications sont regroupées sur les paramètres en aval
• Réduction de la dispersion sur les gradients en aval

Si le point de contrôle d'activation (checkpoint()) est utilisé, aucune communication supplémentaire n'est générée, car les paramètres sont de toute façon préchargés lors du retour en arrière.

Dans la conception FSDP, la charge utile de communication par rang est déterminée comme suit : chaque appel à FullyShardedDataParallel crée un groupe de communication composé des paramètres de module.parameters(), à l'exception de ceux déjà affectés à une instance FullyShardedDataParallel imbriquée. Par exemple, pour Llama, si vous appliquez FullyShardedDataParallel à chaque bloc transformateur et au module racine, il existe alors un groupe de communication pour chaque bloc transformateur, puis un groupe de communication avec l'intégration initiale et la linéarité finale. Chaque groupe de communication correspond à un seul appel de regroupement global et à un seul appel de réduction de la dispersion. Ainsi, la manière dont vous appliquez FullyShardedDataParallel détermine la taille de la communication. En général, appliquer FSDP à chaque bloc transformateur est une bonne heuristique pour les LLM, et il est difficile de faire mieux compte tenu de la conception actuelle.

Prenons l'exemple d'un modèle basé sur Transformer, fragmenté sur 8 GPU. Le fragmentation s'effectue uniquement au niveau du bloc Transformer. Chaque bloc Transformer contient 1,6 milliard de paramètres, au format fp32 (4 octets chacun). Cela signifie qu'une fois fragmenté, chaque bloc Transformer contiendra 0,2 milliard de paramètres sur chaque rang.

• La passe avant communiquera par blocs de 0,2 x 4 = 0,8 Go en mode tout-collecte.
• La passe arrière communiquera deux fois 0,8 Go chacun (1 x tout-collecte et 1 x réduction-dispersion).

Autrement dit, il y aura 3 communications avec une charge utile de 0,8 Go chacune. Si le modèle était composé de 10 blocs transformateurs, il y aurait un total de 30 communications, soit 30 x 0,8 = 24 Go.

Pour formaliser la taille de la charge utile par communication et par rang, on obtient : total_transformer_block_params_in_B*dtype_bytes/num_gpus (Go).

Notez que dans cet exemple, nous n'avons pas inclus les communications supplémentaires nécessaires à l'intégration, devant également être prises en compte. Le calcul dépendrait de la corrélation entre les intégrations d'entrée et de sortie. Si elles ne sont pas liées, le nombre de communications sera multiplié par deux.

Tailles des tampons FSDP

Commençons par les tampons alloués aux communications :

Le transfert forward nécessite actuellement deux fois la taille du tampon de collecte totale. Voici pourquoi :

Comme expliqué dans la section « Nuances du préchargement FSDP », dans le cas d'un préchargement explicite (forward_prefetch=True`), couche 0 : collecte totale -> couche 0 : calcul de transfert -> couche 1 : collecte totale, deux tampons de taille « all-gather » sont nécessaires, car l'un est utilisé pour le transfert actuel, tandis que l'autre effectue le préchargement.

Alors que le préchargement implicite (forward_prefetch=False, par défaut) de la même séquence ne devrait théoriquement nécessiter qu'un seul tampon, il s'agit en réalité de deux fois la taille du tampon de collecte totale. En effet, dans la conception FSDP à paramètres plats, nous ne copions pas le tampon de collecte totale. Les paramètres utilisés pour le calcul sont directement visualisés dans le tampon de collecte globale (en fait, le principal avantage du « paramètre plat » réside précisément dans cette raison). Dans ce cas, alors que la collecte globale de la couche 1 chevauche la collecte globale de la couche 0, cette dernière utilise les paramètres visualisés dans le tampon de collecte globale de la couche 0.

Une question naturelle se pose alors : quand utiliser forward_prefetch=False ? Pour les modèles à graphes statiques (comme la plupart des LLM), il existe une raison technique majeure. En pratique, nous avons ajouté cette option rapidement pour certains modèles internes gourmands en ressources CPU et n'avons pas testé chaque chemin de code avec elle lors des tests unitaires ; nous sommes donc moins confiants. La valeur forward_prefetching=False peut être légèrement plus simple à utiliser, car nous n'avons pas besoin de vérifier l'ordre de collecte globale enregistré comme un éventuel « mode d'échec » ; la collecte globale d'un module se trouve toujours sous sa propre étiquette record_function dans sa trace de profileur.

La fonction « backward » nécessite actuellement au moins deux fois la taille du tampon de regroupement global, voire un peu plus. Voici pourquoi :

La conception actuelle de FSDP utilise recordStream pour gérer les allocations produites dans un flux et consommées dans un autre, ce qui peut entraîner une utilisation mémoire plus importante que prévu. Cette quantité supplémentaire peut être « non déterministe », car elle dépend du timing du noyau GPU par rapport au CPU. Le paramètre limit_all_gathers=True permet d'atténuer ce problème ; pour plus de détails, consultez la discussion sur FSDP et CUDACachingAllocator.

Fonctionnement de FSDP existant avec autograd :

• Le FSDP existant regroupe le flat_param, qui est la feuille autograd.
• Il appelle torch.split pour obtenir des vues 1D dans le flat_param correspondant à ses paramètres d'origine.
• Il appelle torch.view à chaque division 1D pour afficher le retour vers ND.
• Cela signifie qu'en mode rétrograde, nous obtenons ViewBackward (ND -> 1D) et SplitWithSizesBackward (qui est une concaténation). Plus précisément, chaque gradient individuel est calculé comme une allocation distincte, et une concaténation explicite est utilisée pour construire le tampon d'entrée de réduction de la dispersion. Cela implique en réalité une taille de tampon doublée pour la réduction de la dispersion à ce point mémoire maximal.

En résumé, pour le mode rétrograde, il faut environ deux fois la taille du tampon pour réduire la dispersion, plus les effets de RecordStream.

Deuxièmement, parlons des tampons supplémentaires :

Une fois les paramètres fragmentés collectés à partir de tous les rangs, ils nécessitent un tampon supplémentaire de total_transformer_block_params_in_B*dtype_bytes pour l'ensemble des paramètres. Ainsi, pour reprendre l'exemple précédent, si chaque bloc transformateur contient 1,6 B de paramètres et que les paramètres sont au format fp32, le tampon serait de 1,6 * 4 = 6,4 Go.

Deux de ces tampons sont nécessaires, car l'un est en cours d'utilisation et l'autre en prélecture.

En résumé, nous avons :

• 2 tampons de communication de « total_transformer_block_params_in_B*dtype_bytes/num_gpus »
• 2 tampons de paramètres de bloc de transformateur non fragmentés de « total_transformer_block_params_in_B*dtype_bytes »

ou si vous avez suivi l'exemple :

• 2*1,6*4/8=1,6 Go
• 2**1,6*4=12,8 Go
• et un total de 14,4 Go.

Voyons maintenant brièvement ce qui arrive aux représentations continues, car nous les avons exclues des calculs :

• Compte tenu de la règle mentionnée dans la note commençant par « la taille du tampon de communication est déterminée comme suit », nous pouvons analyser la situation comme suit :
• Supposons que nous appliquions FSDP au module racine (par exemple, la classe Transformer). Supposons que nous appliquions ensuite FSDP à chaque bloc Transformer (par exemple, la classe TransformerBlock).
• Le plus souvent, la représentation continue et la projection linéaire finale sont des enfants directs de la classe Transformer racine.
• D'après notre règle, cela signifie que la représentation continue et la projection linéaire finale sont affectées au paramètre plat de la classe Transformer racine.
• Une autre règle spéciale s'applique : la racine ne libère pas ses paramètres après le transfert, car ils seront de toute façon immédiatement regroupés dans le transfert arrière.
• En résumé, le paramètre plat de la racine, y compris la représentation continue et la projection finale, est regroupé pour commencer le transfert avant et conservé dans la mémoire GPU jusqu'à la fin du transfert arrière. Si l'intégration et le linéaire final ne sont pas liés par leur pondération, nous pourrions appliquer FSDP à l'intégration et au linéaire final. Pour les paramètres liés par leur pondération, nous exigeons qu'ils fassent partie du même paramètre plat (sinon, ils seraient comptés deux fois). Cela permettrait de libérer l'intégration après son utilisation en forward et de ne la rassembler que vers la fin de backward.
• Nous espérons que cela donne une meilleure idée : chaque module FSDP reçoit des paramètres assignés dans son module.parameters, à l'exception de ceux déjà assignés à un autre module FSDP imbriqué, et le forward du module FSDP définit l'intervalle de temps de fonctionnement de ses paramètres. Par conséquent, la structure imbriquée nn.Module peut affecter la planification all-gather/free et, par conséquent, les performances mémoire/débit.