Mise à l'échelle de l'apprentissage par renforcement MoE sur Amazon EKS avec EFA et DeepEP pour un débit supérieur de 40 %
aws ·
Cet article présente une architecture optimisée pour mettre à l'échelle l'apprentissage par renforcement (RL) basé sur des Mélanges d'Experts (MoE) sur Amazon EKS, en répondant aux défis de la coordination du calcul hétérogène et des communications à haut débit. La solution combine le service Elastic Kubernetes Service (Amazon EKS), l'Elastic Fabric Adapter (EFA) et DeepEP pour équilibrer la génération des rollout et l'entraînement de la politique. Une innovation majeure réside dans l'utilisation de DeepEP sur EFA, qui remplace les collectifs standard de type all-to-all par des noyaux spécialisés afin de réduire la surcharge de communication dans le parallélisme des experts. Des tests réalisés sur 48 instances P5en ont démontré une augmentation de 40 % du débit agrégé des rollout RL par rapport aux configurations de référence. L'architecture exploite également les instances Spot Amazon EC2 pour la génération des rollout, réduisant ainsi les coûts sans perturber les charges de travail d'entraînement de la politique étroitement couplées. Des groupes de nœuds distincts sont utilisés pour les tâches accélérées par GPU, les environnements basés sur CPU et les tampons d'expérience optimisés en mémoire, garantissant ainsi une mise à l'échelle indépendante. Cette approche permet une co-conception efficace à travers toute la pile, en équilibrant les contraintes de calcul, de mémoire et de bande passante. Le résultat est une voie flexible et rentable pour les charges de travail de RL distribué à grande échelle impliquant des centaines ou des milliers d'accélérateurs.
aws