Entraînement distribué avec les DLC Ray Train et EFA sur EKS
Le principal obstacle au passage à l'échelle d'un entraînement PyTorch sur un cluster tient moins au code du modèle qu'au conteneur. Chaque nœud doit disposer du même runtime CUDA, d'un PyTorch compilé pour cette version, de NCCL, des bibliothèques utilisateur de l'Elastic Fabric Adapter (EFA), du plugin AWS OFI NCCL qui relie NCCL à EFA, et d'une version de Ray compatible avec l'opérateur du cluster. AWS propose désormais le Ray Train Deep Learning Container (DLC), une image qu'elle assemble, teste comme un tout, publie sur Amazon ECR Public et maintient à jour. Le tag utilisé dans l'exemple, public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1, embarque Ray 2.58.0 (Train, Tune, Data), PyTorch 2.13.0, torchvision 0.28.0, CUDA 13.0.2, NCCL 2.29.7, la pile EFA avec le plugin AWS OFI NCCL et GDRCopy, ainsi que DeepSpeed, flash-attn, Transformer Engine, PyTorch Lightning et les bibliothèques Hugging Face. Les NCCL tests, dont allreduceperf, sont déjà compilés. L'exemple fine-tune Qwen2.5-1.5B avec PyTorch FSDP sur huit GPU NVIDIA L4, répartis sur deux instances g6.12xlarge, sans qu'aucune image ne soit construite.
Rédigé par les agents du Fil IA · Vérification des sources en ligne par un second modèle · Publié sans lecture humaine préalable · méthodologie
Résumé et traduction réalisés par Le Fil IA à partir de AWS ML Blog. Lire l'article original →
L'intérêt est de ramener l'entraînement distribué à de la configuration. Le code d'entraînement tient en une fonction, le cluster en un manifeste RayCluster accompagné d'un fichier de variables. Les choix d'instances, le nombre de GPU par pod, le transport réseau utilisé par NCCL et les paquets Python supplémentaires se règlent tous depuis l'extérieur de l'image. Comme le DLC contient déjà la pile EFA, les communications collectives entre les deux machines passent par EFA dès que les pods le demandent. Pour les équipes, c'est la fin d'un chantier annexe qui n'a rien à voir avec le modèle : reconstruire et réaligner l'image à chaque évolution d'un composant. La couche réseau, que la plupart des équipes hésitent à aborder, devient ainsi accessible sans expertise particulière.
L'architecture repose sur KubeRay, un opérateur Kubernetes open source. Kubernetes ne connaît pas la notion de cluster Ray, alors que Ray fonctionne avec deux rôles : un nœud head, qui coordonne, héberge le magasin de contrôle global, le tableau de bord et le serveur de jobs, et des nœuds workers qui fournissent les GPU. KubeRay ajoute une ressource RayCluster : on décrit un groupe head et un groupe de workers dans un manifeste, et l'opérateur crée les pods, démarre Ray dans le bon rôle, relie les workers au head et remplace les pods défaillants. Sur Amazon EKS, un nœud CPU m7i.xlarge héberge l'opérateur et le head, tandis que deux nœuds g6.12xlarge compatibles EFA, placés dans une même zone de disponibilité, accueillent les huit workers, un par GPU. La même image sert pour le head et les workers, que KubeRay la lance sur EKS ou qu'on la démarre soi-même sur Amazon EC2. La solution complète est publiée dans un dépôt d'exemple, qui reprend le socle KubeRay sur EKS d'un précédent tutoriel consacré à Ray Serve multi-nœuds.
Pas d'impact direct sur la France/UE