AWS lance le DLC Ray Train pour simplifier l'entraînement IA distribué sur EKS

Amazon Web Services (AWS) a lancé le Deep Learning Container (DLC) Ray Train, une image préconfigurée conçue pour rationaliser l'entraînement distribué en apprentissage automatique sur Amazon Elastic Kubernetes Service (EKS). Publié en octobre 2026, ce conteneur regroupe Ray, PyTorch et les bibliothèques réseau nécessaires au calcul haute performance, permettant aux ingénieurs de déployer des tâches d'entraînement multi-nœuds sans gérer de chaînes de dépendances complexes.

Ce qui s'est passé

Le principal goulot d'étranglement lors de la mise à l'échelle de l'entraînement PyTorch sur des clusters est souvent l'infrastructure des conteneurs plutôt que le code du modèle lui-même. Les ingénieurs consacrent généralement beaucoup de temps à aligner les environnements d'exécution CUDA, les builds PyTorch, la bibliothèque NVIDIA Collective Communications Library (NCCL) et les bibliothèques utilisateurs Elastic Fabric Adapter (EFA). Les incompatibilités entre ces composants peuvent entraîner une dégradation silencieuse des performances, où NCCL se rabat sur des sockets TCP plus lents au lieu d'utiliser les connexions EFA à grande vitesse. Pour remédier à cela, AWS a assemblé, testé et publié une image unifiée dans Amazon Elastic Container Registry (ECR) Public.

Ce nouveau DLC est construit et maintenu par AWS, qui applique des correctifs à la pile technologique et la valide via des tests d'entraînement KubeRay multi-nœuds avec EFA. La version inclut Ray 2.58.0, PyTorch 2.13.0, CUDA 13.0.2 et NCCL 2.29.7, ainsi que le plugin AWS OFI NCCL qui connecte NCCL à EFA. En fournissant cette image publique, AWS supprime la nécessité pour les équipes de construire et maintenir leurs propres images de conteneurs à chaque mise à jour de ces bibliothèques sous-jacentes.

Le guide technique accompagnant cette publication démontre comment utiliser ce DLC pour affiner le modèle Qwen2.5-1.5B en utilisant PyTorch Fully Sharded Data Parallel (FSDP) sur huit GPU NVIDIA L4. La configuration utilise deux instances g6.12xlarge dans une seule zone de disponibilité, gérées par l'opérateur KubeRay. Cette configuration permet aux collectives inter-nœuds de fonctionner immédiatement sur EFA, car le DLC intègre déjà la pile réseau requise.

Comment cela fonctionne

Le DLC Ray Train fonctionne comme une image unique, agnostique vis-à-vis des rôles, servant à la fois les pods head et worker de Ray. Sur Kubernetes, l'opérateur KubeRay utilise un manifeste RayCluster pour créer les pods nécessaires, indiquant aux workers le nœud head pour la coordination. Comme l'image contient tous les pilotes et bibliothèques nécessaires, les utilisateurs n'ont pas besoin de construire des images personnalisées. Ils configurent plutôt le cluster via des variables d'environnement et des demandes de ressources dans le manifeste Kubernetes.

La gestion du réseau est automatisée grâce à des variables d'environnement préconfigurées. Le DLC utilise par défaut EFA pour la communication via le paramètre FI_PROVIDER=efa et exclut les interfaces de bouclage de la configuration NCCL. Si une tâche nécessite des packages Python supplémentaires non inclus dans l'image de base, la fonctionnalité runtime_env de Ray les installe au moment de la soumission, laissant l'image de base inchangée. Cette séparation garantit que le conteneur reste stable tout en offrant de la flexibilité pour des besoins expérimentaux spécifiques.

À l'intérieur du conteneur, Ray Train gère la formation des groupes de processus distribués. Les ingénieurs écrivent leur boucle d'entraînement comme une fonction Python standard, que Ray exécute sur tous les GPU disponibles. Pour les charges de travail FSDP, Ray divise les poids du modèle, les gradients et les états de l'optimiseur entre les GPU, réduisant les exigences mémoire par périphérique. Le DLC garantit que le trafic de communication intense généré par les opérations gather et reduce-scatter de FSDP circule sur l'interface EFA à large bande passante plutôt que sur les sockets réseau standards.

Détails clés

• Tag de l'image : L'exemple utilise public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1, qui inclut Ray 2.58.0 et PyTorch 2.13.0.
• Prise en charge matérielle : L'architecture de référence utilise des instances g6.12xlarge avec des GPU NVIDIA L4 et des interfaces EFA, déployées dans une seule zone de disponibilité.
• Pile réseau : Inclut CUDA 13.0.2, NCCL 2.29.7 et le plugin AWS OFI NCCL, testés pour garantir l'utilisation d'EFA plutôt qu'un repli sur TCP.
• Méthode de configuration : La forme du cluster et les paramètres de transport sont définis via des variables d'environnement et des demandes de ressources Kubernetes, et non en reconstruisant l'image.
• Estimation des coûts : La configuration d'exemple coûte environ 9,50 $ par heure, couvrant deux instances GPU, le plan de contrôle EKS et une passerelle NAT.
• Compatibilité des pilotes : Le point d'entrée du conteneur active automatiquement les bibliothèques de compatibilité ascendante CUDA si le pilote NVIDIA hôte est plus ancien que la version CUDA de l'image.

Pourquoi c'est important

Pour les ingénieurs logiciels et les équipes ML, cette version réduit considérablement la charge opérationnelle de l'entraînement distribué. Gérer des versions compatibles de CUDA, NCCL et des pilotes EFA est un processus fragile qui conduit souvent à des problèmes de performance subtils. En fournissant une image publique testée, AWS permet aux équipes de se concentrer sur le développement de modèles plutôt que sur le débogage de l'infrastructure. La possibilité d'épingler des versions spécifiques ou d'utiliser des tags mobiles pour les mises à jour de sécurité offre une flexibilité tant pour la recherche reproductible que pour la stabilité en production.

L'intégration avec KubeRay et EKS simplifie également la gestion des clusters. Les ingénieurs peuvent définir toute leur infrastructure d'entraînement dans un seul manifeste, spécifiant le nombre de GPU et les interfaces EFA comme des ressources Kubernetes standard. Cette approche tire parti des outils Kubernetes existants et évite le besoin de scripts personnalisés pour gérer les événements du cycle de vie du cluster Ray. Le résultat est une voie plus fiable et évolutive pour exécuter des tâches d'affinage de grands modèles de langage sur l'infrastructure AWS.

Ce que vous pouvez faire

• Examiner le dépôt d'exemple : Accédez au code et aux manifestes fournis pour comprendre la structure d'une tâche Ray Train sur EKS.
• Vérifier la disponibilité des quotas : Assurez-vous que votre compte AWS dispose d'une capacité On-Demand suffisante pour les instances de la famille G, spécifiquement 96 vCPUs pour les deux nœuds g6.12xlarge.
• Installer les outils requis : Configurez AWS CLI v2, eksctl, kubectl, Helm et envsubst pour gérer le cluster et soumettre les tâches.
• Tester le DLC localement : Téléchargez l'image publique depuis ECR pour vérifier la compatibilité avec votre environnement de développement local avant le déploiement sur le cluster.
• Configurer les variables d'environnement : Utilisez le script fourni scripts/env.sh pour définir les paramètres du cluster tels que la région, le type d'instance et le tag de l'image.
• Surveiller les performances réseau : Activez NCCL_DEBUG=INFO dans votre manifeste pour confirmer qu'EFA est utilisé pour la communication inter-nœuds pendant l'entraînement.