La technologie derrière l'entraînement de BLOOM : 384 GPU A100 et Megatron-DeepSpeed

L'équipe BigScience de Hugging Face détaille l'infrastructure matérielle et logicielle utilisée pour entraîner le modèle multilingue BLOOM de 176 milliards de paramètres. L'entraînement a mobilisé 384 GPU NVIDIA A100 80GB répartis sur 48 nœuds, avec une architecture GPT-3 enrichie, un jeu de données de 350 milliards de tokens couvrant 46 langues, et une durée d'environ 3,5 mois sur le supercalculateur Jean Zay. Le framework Megatron-DeepSpeed combine le parallélisme de données ZeRO, le parallélisme tensoriel de Megatron-LM et le parallélisme de pipeline pour répartir l'entraînement, tandis que le format BF16 et des kernels CUDA fusionnés ont permis de stabiliser et d'accélérer le processus.