Escalar el aprendizaje por refuerzo MoE en Amazon EKS con EFA y DeepEP con un 40 % más de rendimiento
aws ·
Este artículo detalla una arquitectura optimizada para escalar el aprendizaje por refuerzo de Mezcla de Expertos (MoE) en Amazon EKS, abordando desafíos en la coordinación de cómputo heterogéneo y la comunicación de alto rendimiento. La solución combina Amazon Elastic Kubernetes Service (Amazon EKS), Elastic Fabric Adapter (EFA) y DeepEP para equilibrar la generación de rollout con el entrenamiento de políticas. Una innovación clave es el uso de DeepEP sobre EFA, que reemplaza las colectivas all-to-all estándar con kernels especializados para reducir la sobrecarga de comunicación en la paralelización de expertos. Las pruebas en 48 instancias P5en demostraron un aumento del 40 % en el throughput agregado de rollout de RL en comparación con configuraciones base. La arquitectura también aprovecha las instancias EC2 Spot de Amazon para la generación de rollout, reduciendo costos sin interrumpir las cargas de trabajo de entrenamiento de políticas estrechamente acopladas. Se utilizan grupos de nodos separados para tareas aceleradas por GPU, entornos basados en CPU y buffers de experiencia optimizados en memoria, lo que garantiza un escalado independiente. Este enfoque permite una co-diseño eficiente a través de toda la pila, equilibrando las limitaciones de cómputo, memoria y ancho de banda. El resultado es una ruta flexible y rentable para cargas de trabajo de RL distribuido a gran escala que involucran cientos o miles de aceleradores.
aws