PyTorch : un checkpoint FSDP de 522 s à 9 s

Ce qui se passe

Sur un entraînement distribué, la panne n'est pas un incident : c'est le régime normal. Avec un taux de défaillance annualisé d'environ 1 % par GPU, un job de 256 cartes tournant 30 jours a 19 % de chances d'être interrompu, et 57 % à 1 024 cartes. Mesuré plutôt que modélisé, c'est pire : le supercalculateur Delta et ses 608 H100 connaissent une panne toutes les 1,9 heure — soit, ramené à 32 GPU, une moyenne de 36 heures entre deux interruptions.

La métrique qui compte est le goodput : la part du temps que les accélérateurs passent à calculer, plutôt qu'à attendre ou à refaire du travail perdu. Deux sous-systèmes la décident, et tous deux sont traités comme des détails d'implémentation : le checkpointing et le chargement des données.

Le mécanisme

Le premier checkpoint qu'écrit une équipe est un torch.save sur le rang 0. Il rassemble l'état de tous les rangs sur un seul processus, puis écrit un fichier unique de façon synchrone — pendant que les GPU attendent.

Le Distributed Checkpoint de PyTorch renverse la conception : chaque rang écrit son propre fragment en parallèle, accompagné d'un petit fichier .metadata qui décrit comment les fragments se recomposent. Le temps de sauvegarde décroît en 1/N, et surtout la reprise peut se faire sur un nombre différent de GPU : DCP replanifie quels octets vont à quel rang, ce qui rend possible le redémarrage sur un cluster amputé.

Contrairement à une idée reçue, ce n'est pas réservé aux modèles shardés : en DDP aussi, où chaque rang détient une réplique identique, DCP découpe l'état et écrit en parallèle. C'est de plus la même API le jour où l'on passe à FSDP.

La sauvegarde asynchrone (async_save) fait le reste : une copie rapide vers un tampon d'attente, puis un envoi en tâche de fond recouvert par la suite de l'entraînement. La boucle ne paie que la copie.

Ce que ça change

Un checkpoint bon marché autorise un checkpoint fréquent, et c'est là que le calcul bascule. Le travail perdu par panne vaut en moyenne la moitié de l'intervalle : diviser l'intervalle par dix divise par dix le temps de reprise. Au rythme d'interruptions relevé sur Llama 3 — environ 8,6 par jour — sauvegarder toutes les deux heures revient à jeter 8,6 heures de calcul quotidien, soit un goodput de 64 %. Toutes les trente minutes, la perte tombe à 2,15 heures, et le goodput monte à 91 %.

Encore faut-il que la reprise soit automatique : au redémarrage, le job doit retrouver le dernier checkpoint complet et ignorer celui que le crash a laissé à moitié écrit. Le .metadata de DCP, écrit seulement une fois tous les fragments posés, sert exactement de marqueur de complétude.

Le guide insiste enfin sur une panne qui ne produit aucun message d'erreur : checkpointer le modèle sans checkpointer la position dans le jeu de données. Au redémarrage, le dataloader repart du début de l'époque, ré-entraîne sur des exemples déjà vus et en saute d'autres. Le job se termine normalement, les métriques sont simplement moins bonnes. Corriger demande de sérialiser un offset d'échantillon ou de shard — et donc de rendre l'ordre déterministe, ce qui implique de sauvegarder aussi les graines et l'état des générateurs aléatoires.

Source : Databricks, Fast, fault-tolerant PyTorch training on AI Runtime