AWS lanza Ray Train DLC para simplificar el entrenamiento de IA distribuido en EKS

Amazon Web Services (AWS) ha lanzado el Deep Learning Container (DLC) de Ray Train, una imagen preconfigurada diseñada para agilizar el entrenamiento distribuido de machine learning en Amazon Elastic Kubernetes Service (EKS). Publicado en octubre de 2026, este contenedor integra Ray, PyTorch y las bibliotecas de red necesarias para la computación de alto rendimiento, permitiendo a los ingenieros desplegar trabajos de entrenamiento multinodo sin gestionar complejas cadenas de dependencias.

Qué ocurrió

El principal cuello de botella al escalar el entrenamiento de PyTorch en clústeres suele ser la infraestructura de contenedores, más que el código del modelo en sí. Los ingenieros dedican mucho tiempo a alinear los entornos de ejecución de CUDA, las compilaciones de PyTorch, la NVIDIA Collective Communications Library (NCCL) y las bibliotecas de espacio de usuario de Elastic Fabric Adapter (EFA). Las incompatibilidades entre estos componentes pueden provocar degradaciones silenciosas del rendimiento, donde NCCL recurre a sockets TCP más lentos en lugar de utilizar conexiones EFA de alta velocidad. Para abordar esto, AWS ha ensamblado, probado y publicado una imagen unificada en Amazon Elastic Container Registry (ECR) Public.

Este nuevo DLC es construido y mantenido por AWS, que parchea la pila tecnológica y la valida mediante pruebas de entrenamiento con KubeRay y EFA multinodo. La versión incluye Ray 2.58.0, PyTorch 2.13.0, CUDA 13.0.2 y NCCL 2.29.7, junto con el plugin AWS OFI NCCL que conecta NCCL con EFA. Al ofrecer esta imagen pública, AWS elimina la necesidad de que los equipos construyan y mantengan sus propias imágenes de contenedor cada vez que se actualizan estas bibliotecas subyacentes.

La guía técnica adjunta demuestra cómo usar este DLC para ajustar fino (fine-tuning) el modelo Qwen2.5-1.5B utilizando PyTorch Fully Sharded Data Parallel (FSDP) en ocho GPUs NVIDIA L4. La configuración emplea dos instancias g6.12xlarge en una sola zona de disponibilidad, gestionadas por el operador KubeRay. Esta configuración permite que las operaciones colectivas entre nodos se ejecuten inmediatamente sobre EFA, ya que el DLC ya incorpora la pila de red requerida.

Cómo funciona

El Ray Train DLC funciona como una única imagen agnóstica al rol, que sirve tanto a los pods head de Ray como a los worker. En Kubernetes, el operador KubeRay utiliza un manifiesto RayCluster para crear los pods necesarios, apuntando a los workers hacia el nodo head para la coordinación. Dado que la imagen contiene todos los controladores y bibliotecas necesarios, los usuarios no tienen que construir imágenes personalizadas. En su lugar, configuran el clúster mediante variables de entorno y solicitudes de recursos en el manifiesto de Kubernetes.

La gestión de la red se realiza automáticamente a través de variables de entorno preconfiguradas. El DLC utiliza por defecto EFA para la comunicación mediante la configuración FI_PROVIDER=efa y excluye las interfaces de loopback de la configuración de NCCL. Si un trabajo requiere paquetes Python adicionales no incluidos en la imagen base, la función runtime_env de Ray los instala en el momento del envío, manteniendo la imagen base intacta. Esta separación garantiza que el contenedor permanezca estable mientras ofrece flexibilidad para necesidades experimentales específicas.

Dentro del contenedor, Ray Train gestiona la formación del grupo de procesos distribuido. Los ingenieros escriben su bucle de entrenamiento como una función estándar de Python, que Ray ejecuta en todas las GPUs disponibles. Para cargas de trabajo FSDP, Ray divide los pesos del modelo, los gradientes y los estados del optimizador entre las GPUs, reduciendo los requisitos de memoria por dispositivo. El DLC asegura que el intenso tráfico de comunicación generado por las operaciones gather y reduce-scatter de FSDP viaje a través de la interfaz EFA de alto ancho de banda en lugar de los sockets de red estándar.

Detalles clave

• Etiqueta de imagen: El ejemplo utiliza public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1, que incluye Ray 2.58.0 y PyTorch 2.13.0.
• Soporte de hardware: La arquitectura de referencia utiliza instancias g6.12xlarge con GPUs NVIDIA L4 e interfaces EFA, desplegadas en una sola zona de disponibilidad.
• Pila de red: Incluye CUDA 13.0.2, NCCL 2.29.7 y el plugin AWS OFI NCCL, probados para garantizar el uso de EFA en lugar de la alternativa TCP.
• Método de configuración: La forma del clúster y los ajustes de transporte se definen mediante variables de entorno y solicitudes de recursos de Kubernetes, no reconstruyendo la imagen.
• Estimación de costos: La configuración de ejemplo cuesta aproximadamente $9.50 por hora, cubriendo dos instancias GPU, el plano de control de EKS y una NAT Gateway.
• Compatibilidad de controladores: El punto de entrada del contenedor habilita automáticamente las bibliotecas de compatibilidad hacia adelante de CUDA si el controlador NVIDIA del host es más antiguo que la versión de CUDA de la imagen.

Por qué importa

Para los ingenieros de software y los equipos de ML, este lanzamiento reduce significativamente la sobrecarga operativa del entrenamiento distribuido. Gestionar versiones compatibles de los controladores CUDA, NCCL y EFA es un proceso frágil que a menudo conduce a problemas sutiles de rendimiento. Al proporcionar una imagen pública probada, AWS permite a los equipos centrarse en el desarrollo del modelo en lugar de depurar la infraestructura. La capacidad de fijar versiones específicas o usar etiquetas móviles para actualizaciones de seguridad ofrece flexibilidad tanto para la investigación reproducible como para la estabilidad en producción.

La integración con KubeRay y EKS también simplifica la gestión del clúster. Los ingenieros pueden definir toda su infraestructura de entrenamiento en un único manifiesto, especificando el número de GPUs y las interfaces EFA como recursos estándar de Kubernetes. Este enfoque aprovecha las herramientas existentes de Kubernetes y evita la necesidad de scripts personalizados para gestionar los eventos del ciclo de vida del clúster de Ray. El resultado es una ruta más fiable y escalable para ejecutar trabajos de ajuste fino de grandes modelos de lenguaje en la infraestructura de AWS.

Qué puedes hacer

• Revisar el repositorio de ejemplo: Accede al código y a los manifiestos proporcionados para entender la estructura de un trabajo de Ray Train en EKS.
• Comprobar la disponibilidad de cuotas: Asegúrate de que tu cuenta de AWS tenga suficiente capacidad On-Demand para instancias de la familia G, específicamente 96 vCPUs para los dos nodos g6.12xlarge.
• Instalar las herramientas requeridas: Configura AWS CLI v2, eksctl, kubectl, Helm y envsubst para gestionar el clúster y enviar trabajos.
• Probar el DLC localmente: Descarga la imagen pública desde ECR para verificar la compatibilidad con tu entorno de desarrollo local antes de desplegarla en el clúster.
• Configurar variables de entorno: Utiliza el script scripts/env.sh proporcionado para establecer parámetros del clúster como región, tipo de instancia y etiqueta de imagen.
• Monitorizar el rendimiento de la red: Habilita NCCL_DEBUG=INFO en tu manifiesto para confirmar que se está utilizando EFA para la comunicación entre nodos durante el entrenamiento.