Qu'est-ce que l'entraînement en précision mixte et qu'est-ce qui peut mal tourner ?

La précision mixte exécute la plupart des opérations en 16 bits (fp16 ou bf16) tout en gardant une copie maîtresse des poids et certaines réductions en fp32. Cela divise à peu près la mémoire par deux et exploite les tensor cores, donc c'est nettement plus rapide. Le risque avec fp16 est le dépassement par le bas : les petits gradients s'arrondissent à zéro, d'où l'existence de la mise à l'échelle de la perte. bf16 conserve la plage d'exposants du fp32, ce qui l'évite largement.

Vous gardez une copie maîtresse des poids en fp32 mais vous exécutez les passes avant et arrière en 16 bits, donc les activations et les gradients pèsent moitié moins et les produits matriciels tapent dans les tensor cores. En pratique c'est à peu près la moitié de la mémoire et souvent près du double du débit, ce qui peut faire la différence entre un entraînement qui tient sur votre matériel et un qui ne tient pas. L'échec classique, c'est le dépassement par le bas en fp16. Le fp16 n'a que cinq bits d'exposant, donc les très petits gradients tombent à zéro, et les petits gradients sont exactement ceux qui comptent en fin d'entraînement. C'est ce que corrige la mise à l'échelle dynamique de la perte : multipliez la perte par un grand facteur avant la passe arrière pour que les gradients atterrissent dans la plage représentable, annulez la mise à l'échelle avant le pas d'optimiseur, et réduisez le facteur quand vous voyez des infinis. Le bf16 garde huit bits d'exposant, la même plage que le fp32, et le paie avec moins de bits de mantisse, donc il déborde par le bas bien moins souvent et n'a en général besoin d'aucune mise à l'échelle de la perte. Mon choix par défaut aujourd'hui, c'est bf16 sur tout ce qui est Ampere ou plus récent. Et je garde les réductions comme le softmax, la layer norm et la perte elle-même en fp32, parce que sommer des milliers de termes en 16 bits perd vite en exactitude.