ZeRO et FSDP : Révolutionner l'entraînement IA sur plusieurs GPU

Introduction aux défis de l'entraînement IA sur GPU

Dans le domaine de l'intelligence artificielle, l'entraînement de modèles massifs comme GPT-3, qui compte 175 milliards de paramètres, nécessite une gestion efficace des ressources. Le Distributed Data Parallelism (DDP) est une méthode couramment utilisée pour accélérer l'entraînement en répartissant les données sur plusieurs GPU. Cependant, cette technique introduit un problème majeur : la redondance mémoire. Chaque GPU conserve une copie complète des paramètres du modèle, des gradients et des états de l'optimiseur, ce qui entraîne un gaspillage considérable de VRAM précieuse.

La redondance mémoire dans le DDP

Pour comprendre l'impact de la redondance mémoire, examinons ce qui consomme réellement de la mémoire lors de l'entraînement d'un modèle. Un modèle avec N paramètres doit gérer plusieurs éléments :

• Les paramètres du modèle, qui sont les poids des neurones.
• Les gradients, qui sont calculés pour chaque paramètre.
• Les états de l'optimiseur, tels que ceux utilisés par Adam, qui nécessitent le stockage du premier et du second moment pour chaque paramètre.
• Les activations, qui sont les sorties intermédiaires stockées pendant le passage avant pour être utilisées dans le passage arrière.

Les trois premiers éléments augmentent proportionnellement à la taille du modèle et sont dupliqués sur chaque GPU dans le DDP. Les activations, en revanche, dépendent de la taille du lot, de la longueur de la séquence et du nombre de neurones, et sont uniques à chaque GPU car chaque GPU traite des données différentes. ZeRO ne modifie pas la gestion des activations.

Pour un modèle de 7 milliards de paramètres utilisant l'optimiseur Adam et le format de précision FP32, la mémoire utilisée se répartit comme suit :

• Paramètres : 7 milliards multipliés par 4 octets, soit 28 Go.
• Gradients : 7 milliards multipliés par 4 octets, soit 28 Go.
• États de l'optimiseur : 7 milliards multipliés par 2 fois 4 octets, soit 56 Go.

Ainsi, chaque GPU dans un système DDP utilise 112 Go de mémoire, sans compter les activations qui ajoutent une charge supplémentaire. Des techniques comme le checkpointing des activations peuvent être utilisées pour réduire cette charge en supprimant certaines activations et en les recomputant au besoin, mais cela sort du cadre de cet article.

ZeRO : Une solution innovante

ZeRO, ou Zero Redundancy Optimizer, propose une approche pour réduire cette redondance mémoire. Il se décline en plusieurs versions, de ZeRO-1 à ZeRO-3, chacune apportant des améliorations progressives.

ZeRO-1 : Réduction des états de l'optimiseur

La première version, ZeRO-1, se concentre sur le partitionnement des états de l'optimiseur. Dans ce modèle, chaque GPU conserve une copie complète des paramètres du modèle et des gradients, mais ne stocke qu'une fraction des états de l'optimiseur, soit 1/N où N est le nombre de GPU. Cela signifie que chaque GPU ne met à jour qu'une portion correspondante des paramètres.

Le processus d'entraînement avec ZeRO-1 suit ces étapes :

• Passage avant : chaque GPU traite son propre micro-lot de données.
• Passage arrière : les gradients sont calculés.
• All-reduce des gradients : chaque GPU reçoit tous les gradients.
• Étape de l'optimiseur : chaque GPU met à jour sa partition de paramètres.
• All-gather des paramètres : synchronisation des paramètres mis à jour entre les GPU.

Implémentation de ZeRO-1

Pour illustrer le fonctionnement de ZeRO-1, voici une implémentation simplifiée en Python :

import torch.distributed as dist

def __init__(self, model, optimizer_cls):
    self.model = model
    self.rank = dist.get_rank()
    self.world_size = dist.get_world_size()
    self.param_shards = list()  # chaque rang ne détient que sa part des états de l'optimiseur
    self.param_metadata = list()  # métadonnées pour reconstruire les parts
    
    for param in self.model.parameters():
        original_shape = param.data.shape
        flat = param.data.view(-1)
        numel = flat.numel()
        remainder = numel % self.world_size
        pad_size = (self.world_size - remainder) % self.world_size
        padded_numel = numel + pad_size
        shard_size = padded_numel // self.world_size
        shard_start = self.rank * shard_size
        shard_end = shard_start + shard_size
        
        self.param_metadata.append({
            "original_shape": original_shape,
            "padded_numel": padded_numel,
            "shard_size": shard_size,
            "shard_start": shard_start,
            "shard_end": shard_end,
        })
        
        if pad_size > 0:
            flat_padded = torch.cat([flat, flat.new_zeros(pad_size)])
        else:
            flat_padded = flat
            
        shard = flat_padded[shard_start:shard_end].clone()
        shard.requires_grad_(True)
        self.param_shards.append(shard)
        
    self.optimizer = optimizer_cls(self.param_shards)

def training_step(self, inputs, targets, loss_fn):
    output = self.model(inputs)  # passage avant
    loss = loss_fn(output, targets)  # calcul de la perte
    loss.backward()  # passage arrière
    self._sync_gradients()  # all-reduce des gradients entre les GPU
    self.optimizer.step()  # mise à jour de la part locale des paramètres
    self._sync_params()  # synchronisation des paramètres du modèle
    # réinitialiser les gradients pour l'étape suivante
    for param in self.model.parameters():
        param.grad = None

def _sync_gradients(self):
    for idx, param in enumerate(self.model.parameters()):
        # Code pour synchroniser les gradients

Cette implémentation démontre comment ZeRO-1 parvient à partitionner les états de l'optimiseur tout en conservant les paramètres et les gradients complets sur chaque GPU. Cela permet une utilisation plus efficace de la mémoire, facilitant ainsi l'entraînement de modèles IA de grande envergure.