AWS et NVIDIA sécurisent les longs entraînements PyTorch sur EKS

AWS et NVIDIA présentent une intégration de NVIDIA Resiliency Extension, ou NVRx, avec PyTorch Fully Sharded Data Parallel sur Amazon Elastic Kubernetes Service. La solution vise les entraînements distribués exécutés pendant plusieurs heures ou plusieurs jours sur plusieurs nœuds. Une panne de GPU peut provoquer des délais d’attente NCCL, des redémarrages désynchronisés et une perte de progression. La sauvegarde synchrone ajoute aussi des périodes d’attente, puisqu’elle bloque tous les rangs pendant les écritures. Selon l’article, cette étape a représenté jusqu’à 40 % du temps total dans les configurations étudiées. Les résultats présentés reposent sur des GPU H100, de deux à huit nœuds.

NVRx reste modulaire et s’installe comme une couche Python avec pip install nvidia-resiliency-ext. L’outil ne demande ni noyaux personnalisés, ni dérivé de PyTorch, ni recompilation. Ses primitives s’intègrent à un script FSDP existant par des imports ordinaires. Trois fonctions sont utilisées dans l’exemple. TorchAsyncCheckpoint déporte l’écriture vers un processus d’arrière-plan. inprocess.Wrapper relance la fonction d’entraînement après certaines erreurs transitoires. Enfin, ft_launcher redémarre les travailleurs après des arrêts plus sévères. Ces fonctions restent indépendantes, afin d’adapter le mécanisme à la classe de panne rencontrée. Le modèle et le code d’entraînement restent inchangés dans l’intégration décrite.

NVRx découple sauvegarde et reprise

La sauvegarde asynchrone remplace l’appel bloquant par async_save(), qui transmet l’état du modèle à un processus d’arrière-plan et rend immédiatement la main à l’entraînement. Un appel à finalize_async_save() finalise ensuite l’écriture précédente avant la sauvegarde suivante ou la fin du travail. Avec FSDP LOCAL_STATE_DICT, chaque rang écrit directement son fragment. Cette approche évite un rassemblement complet des données et un goulot d’étranglement associé au rang zéro. Pour les erreurs transitoires, NVRx interrompt le groupe de processus actif, vérifie les GPU, les liens NVLink et les interfaces réseau, puis réintègre les survivants depuis le dernier point de contrôle.

Le redémarrage dur repose sur ft_launcher, qui couvre les cas qu’un redémarrage dans le processus ne peut pas intercepter. Le mécanisme traite notamment les signaux SIGKILL, les arrêts pour dépassement de mémoire et certains blocages du système d’exploitation. Chaque rang exécute un client de surveillance. Le lanceur compare les battements de cœur à des délais définis par la ligne de commande. En cas d’arrêt ou de défaillance, il termine les survivants, récupère la mémoire GPU et crée de nouveaux travailleurs dans le même travail. Les travailleurs restaurés relisent le dernier point de contrôle. L’orchestrateur du cluster couvre séparément la perte d’un nœud.

EKS relie calcul et stockage

L’architecture AWS utilise Amazon EKS, des groupes de nœuds autogérés et des instances EC2 p5.48xlarge. Chaque instance fournit huit GPU NVIDIA H100 de 80 Go et 32 interfaces Elastic Fabric Adapter. EFA annonce une bande passante de 3 200 Gbit/s pour les opérations NCCL all-reduce. Les travaux s’exécutent comme des Jobs Kubernetes, avec des Services sans tête pour la découverte des pairs par DNS. Les plugins NVIDIA et EFA exposent les ressources au planificateur. L’affinité de nœud et les tolérances facilitent l’allocation complète des huit GPU par nœud. Amazon ECR conserve l’image contenant PyTorch, NVRx et le code du modèle.

Le stockage partagé repose sur Amazon FSx for Lustre en configuration SCRATCH_2, avec une capacité de 1,2 To. Le système de fichiers est monté dans chaque pod via le pilote CSI de FSx. Les sauvegardes synchrones et asynchrones y écrivent leurs états. Après une panne, les travailleurs récupérés y lisent également leur point de contrôle. L’article place FSx dans la même zone de disponibilité que les nœuds GPU afin de réduire la latence de lecture. Cette lecture domine le temps de récupération aux échelles étudiées, plutôt que le mécanisme de redémarrage lui-même. La reproduction demande notamment EKS 1.28 ou une version ultérieure, PyTorch 2.9+ et NVRx 0.4.1 pour le banc de test.

Le parcours reproductible part d’un script FSDP lancé avec torchrun et une sauvegarde distribuée synchrone. Une panne de travailleur y arrête l’ensemble du travail, qui doit reprendre depuis le dernier point de contrôle. L’exemple ajoute ensuite séparément la sauvegarde asynchrone, le redémarrage dans le processus et le lanceur de redémarrage dans le travail. Le dépôt awsome-distributed-ai fournit le cas de test NVRx, tandis que les architectures EKS sont documentées dans ce dépôt AWS. La documentation de NVIDIA Resiliency Extension complète les composants logiciels utilisés.