Méthodes d'entraînement à basse précision

La sortie de nouveaux types de matériel a conduit à l’émergence de nouveaux paradigmes d’entraînement qui les exploitent mieux. Actuellement, cela se traduit par un entraînement en précision 8 bits à l’aide de paquets tels que TransformersEngine (TE) ou torchao (FP8 natif de PyTorch).

Pour une introduction aux sujets abordés aujourd'hui, nous vous recommandons de consulter le guide d'utilisation de la basse précision, car cette documentation y fera régulièrement référence.

Un graphique rapide

Voici un graphique rapide montrant les différentes précisions en bits pour chaque solution pendant l'entraînement :

TransformersEngine

TransformersEngine est la première solution à tenter l’entraînement en virgule flottante 8 bits. Elle fonctionne en utilisant des couches de remplacement directes pour certaines d’entre elles dans un modèle, qui exploitent leur moteur FP8 pour réduire le nombre de bits (par exemple de 32 à 8) sans dégrader la précision finale du modèle.

Plus précisément, Accelerate trouvera et remplacera les couches suivantes par des versions TransformersEngine :

• nn.LayerNorm pour te.LayerNorm
• nn.Linear pour te.Linear

En conséquence, nous obtenons un modèle dont la plupart des couches sont en BF16, tandis que certaines couches sont en FP8, ce qui réduit partiellement la mémoire utilisée.

D’après nos observations, les gains de performance obtenus avec TransformerEngine ne deviennent réellement visibles que lorsque la grande majorité des couches du modèle est constituée de ces deux types de couches à remplacer. En conséquence, seuls les modèles de grande taille ont montré des améliorations de performance, lorsque le nombre de paramètres atteint et dépasse plusieurs milliards.

Le TransformerEngine peut recevoir de nombreux arguments différents qui personnalisent la manière dont il effectue les calculs FP8 et ce qu’ils font. La liste complète des arguments est disponible ci-dessous :

• margin : la marge à utiliser pour le mise à l'échelle du gradient.
• interval : L'intervalle à utiliser pour la fréquence de recalcul du facteur d'échelle.
• fp8_format``: The format to use for the FP8 recipe. Must be one ofHYBRIDorE4M3. (GenerallyHYBRIDfor training,E4M3` for evaluation)
• amax_history_len : La longueur de l’historique à utiliser pour le calcul du facteur d’échelle
• amax_compute_algo : l’algorithme à utiliser pour le calcul du facteur d’échelle. Doit être max ou most_recent.
• override_linear_precision : Indique s'il faut exécuter les GEMM fprop, dgrad et wgrad en plus haute précision.

Vous pouvez personnaliser chacun de ces éléments dans le cadre de utils.TERecipeKwargs afin d'aider à optimiser les performances de vos modèles.

Si l’on observe le graphique mentionné précédemment, TE convertit simplement les couches de calcul en FP8, tandis que le reste est en FP32. En conséquence, cette approche utilise la plus grande quantité de mémoire, mais elle présente l’avantage de garantir la perte minimale en précision finale lors de l’entraînement.