Entraînement en précision mixte : autocasting, loss scaling et limites numériques

Ce que c'est

L'entraînement en précision mixte combine plusieurs formats à virgule flottante au sein d'un même entraînement. Les opérations matricielles coûteuses peuvent utiliser FP16 ou BF16, tandis que les opérations sensibles restent en FP32. La précision mixte automatique (AMP) choisit les formats selon des règles propres à chaque opération et à chaque périphérique ; cela ne revient pas à convertir tous les tenseurs en demi-précision. Les bénéfices dépendent du modèle, du matériel et de la charge de travail. [1][2]

Fonctionnement

L'autocasting choisit la précision de chaque opération. Le loss scaling traite un problème distinct : de très petits gradients en FP16 peuvent tomber à zéro. Multiplier la perte par un facteur d'échelle amplifie les gradients de la passe arrière ; retirer ce facteur avant la mise à jour de l'optimiseur rétablit leur amplitude voulue. Un scaler dynamique détecte les gradients non finis, ignore les mises à jour de l'optimiseur concernées et ajuste son facteur d'échelle. Il ne permet pas pour autant de rendre valides tous les modèles numériquement instables. [1][2]

Démarche pratique

Commencez par établir une référence en FP32, avec des données et des règles d'évaluation fixes. Pour un entraînement PyTorch classique en mode eager, encadrez la passe avant et le calcul de la perte avec torch.autocast, puis effectuez la passe arrière en dehors de ce contexte. En FP16, utilisez torch.amp.GradScaler lorsque c'est pertinent. Ne convertissez pas manuellement tout le modèle simplement pour activer l'autocasting. [1]

Si vous appliquez un écrêtage des gradients (clipping), remettez-les à l'échelle d'origine (unscale) avant de les écrêter. Appelez les méthodes step et update du scaler dans l'ordre indiqué par la documentation, puis réinitialisez les gradients pour la mise à jour suivante. Sauvegardez l'état du scaler avec les checkpoints du modèle et de l'optimiseur. [2] En cas d'accumulation de gradients, conservez un seul facteur d'échelle pour l'ensemble du lot effectif. L'exécution compilée et les opérations personnalisées demandent leurs propres vérifications de précision, spécifiques à chaque version.

Exemple illustratif

Prenons un classifieur dont le temps d'exécution est dominé par les convolutions, mais dont la perte finale nécessite une plus grande plage numérique. L'autocasting peut exécuter les convolutions éligibles en précision réduite tout en évaluant en FP32 une perte sensible prise en charge. Si une étape produit des gradients infinis, le scaler peut ignorer cette mise à jour ; il s'agit d'un événement de diagnostic, non d'une preuve que l'entraînement a convergé. Cet exemple décrit un comportement, pas une accélération mesurée.

Limites et vérification

FP16 et BF16 n'ont pas la même plage numérique. Un modèle entraîné pour BF16 peut déborder en FP16, dont la plus grande valeur finie est 65 504. Augmenter le facteur d'échelle de la perte ne corrige pas un calcul de la passe avant qui déborde. BF16 évite généralement le recours à une mise à l'échelle des gradients comme en FP16, mais introduit tout de même des erreurs d'arrondi et doit être validé. [1]

Comparez les métriques de validation et les courbes de perte avec la référence FP32, inspectez les valeurs non finies et les mises à jour ignorées, puis mesurez le pic de mémoire et le débit après la phase de warm-up. Les modèles petits ou limités par les entrées peuvent en tirer peu de bénéfice. [2] Indiquez la précision utilisée, le matériel, les versions logicielles et les écarts observés. Cet article ne rapporte aucune expérience locale.

Sources

[1] PyTorch: Automatic Mixed Precision package [2] PyTorch: Automatic Mixed Precision recipe