DeepSpeed ZeRO

Définition

DeepSpeed ZeRO (Zero Redundancy Optimizer) est un framework d'optimisation pour l'entraînement distribué de très grands modèles de deep learning, développé par Microsoft Research. Il permet d'entraîner des modèles avec des centaines de milliards de paramètres sur des clusters GPU en éliminant la redondance de mémoire inhérente à la parallélisation de données standard. Dans la parallélisation de données classique (Data Parallelism), chaque GPU maintient une copie complète du modèle, des gradients et des états d'optimiseur (Adam conserve le momentum et la variance = 3× la mémoire des paramètres). Pour un modèle de 7B paramètres en FP16 : paramètres ≈ 14 GB, états Adam FP32 ≈ 48 GB → ~62 GB par GPU, hors activations. ZeRO résout ce problème en trois niveaux de partitionnement progressifs : ZeRO-1 partitionne les états d'optimiseur sur les GPUs (économie 4×), ZeRO-2 partitionne également les gradients (économie supplémentaire), ZeRO-3 partitionne les paramètres du modèle eux-mêmes (économie totale = N_GPU×). ZeRO-3 permet ainsi d'entraîner un modèle dont les paramètres dépassent la mémoire d'un seul GPU, en ne maintenant sur chaque device que sa portion de paramètres. ZeRO-Infinity étend le concept en utilisant le NVMe SSD comme mémoire supplémentaire : les paramètres, gradients et activations peuvent être déchargés sur des SSDs NVMe, permettant théoriquement d'entraîner des modèles de taille arbitraire sur un cluster de taille raisonnable. Cette approche est utilisée pour l'entraînement de modèles expérimentaux dépassant 1T paramètres. Pour les équipes MLOps, DeepSpeed est intégré nativement dans HuggingFace Accelerate et Transformers via des fichiers de configuration JSON. L'alternative PyTorch native est FSDP (Fully Sharded Data Parallel), souvent préféré pour sa meilleure intégration avec l'écosystème PyTorch et ses performances comparables.