Entraînement avec ZeRO (Zero Redundancy Optimizer)

Définition

L'entraînement avec ZeRO, pour Zero Redundancy Optimizer, est une technique d'optimisation mémoire développée par Microsoft et intégrée à la bibliothèque DeepSpeed, qui élimine la redondance des états d'entraînement, poids, gradients et états de l'optimiseur, entre les GPU participant à un entraînement distribué. ZeRO se décline en plusieurs étages : le stage 1 partitionne les états de l'optimiseur, le stage 2 ajoute le partitionnement des gradients, et le stage 3 partitionne également les poids du modèle eux-mêmes entre tous les GPU. Chaque dispositif ne conserve alors qu'une fraction de ces données, reconstituée à la volée par communication lors des calculs nécessaires, ce qui réduit drastiquement l'empreinte mémoire par GPU comparé au parallélisme de données classique. ZeRO permet ainsi d'entraîner des modèles de dizaines à centaines de milliards de paramètres sur des clusters GPU standards, là où une réplication complète du modèle serait impossible. Une variante, ZeRO-Offload, déporte une partie des calculs et des états vers la mémoire CPU ou le stockage NVMe pour entraîner des modèles encore plus grands sur du matériel limité. Le compromis principal est un surcoût de communication réseau, croissant avec le stage utilisé. Pour un DSI, ZeRO est un facteur déterminant du dimensionnement matériel nécessaire à un projet de pré-entraînement ou de fine-tuning de grands modèles.