El ejemplo de AWS ejecuta entrenamiento distribuido con Ray Train en Amazon EKS usando el contenedor publicado por AWS, sin construir una imagen propia. Configura Qwen2.5-1.5B con PyTorch FSDP y bfloat16 en ocho GPU NVIDIA L4 repartidas entre dos instancias g6.12xlarge. Para que NCCL use EFA, los workers deben estar en la misma zona de disponibilidad y la misma VPC.

Arquitectura del ejemplo

AWS asigna funciones distintas a un nodo de sistema y a los dos nodos GPU. El primero aloja el operador KubeRay y el Ray head, que no reserva GPU; cada nodo GPU ejecuta un pod worker con cuatro GPU y un dispositivo EFA.

ComponenteConfiguración del ejemploFunción
Nodo de sistemaUna instancia m7i.xlarge, sin GPU asignadas al Ray headAloja el operador KubeRay y coordina el clúster Ray
Nodos de trabajoDos instancias g6.12xlarge, cada una con cuatro NVIDIA L4 y una interfaz EFAEjecutan los ocho workers de Ray Train, con los rangos 0–3 en una instancia y 4–7 en la otra

El contenedor publicado se utiliza sin cambios tanto en el Ray head como en los workers: public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1. AWS lista en esa imagen Ray 2.58.0, PyTorch 2.13.0, torchvision 0.28.0, CUDA 13.0.2 y NCCL 2.29.7.

Requisitos y ubicación de EFA

El ejemplo requiere capacidad bajo demanda de 96 vCPU de la familia G para las dos instancias g6.12xlarge, o una reserva de capacidad bajo demanda existente. AWS incluye g6.12xlarge entre los tipos compatibles con EFA; su disponibilidad depende de la región, así que compruébala antes de crear el clúster.

EFA, o Elastic Fabric Adapter, transporta la comunicación entre instancias que necesita el ejemplo. AWS indica que el tráfico EFA no puede cruzar zonas de disponibilidad ni VPC. Por eso, los dos workers deben compartir ambos límites. El tráfico IP normal de ENA es distinto y sigue siendo enrutable.

Configurar Ray Train y EFA

En cada pod worker, el manifiesto de Kubernetes solicita cuatro GPU y un dispositivo EFA. También fija requests y limits en cuatro GPU, una EFA, 40 CPU y 160 GiB de memoria. Mantén sincronizado el valor de num-gpus de Ray con la solicitud de GPU de Kubernetes: el ejemplo deriva ambos de la misma variable.

La configuración de red usa estas variables de entorno:

  • FI_PROVIDER=efa
  • NCCL_NET_PLUGIN=ofi
  • NCCL_SOCKET_IFNAME=eth0
  • NCCL_DEBUG=INFO

Las instancias usan subredes privadas en una VPC, con una NAT Gateway para descargar el contenedor y los pesos del modelo desde Hugging Face. El trabajo se envía desde el propio pod Ray head mediante kubectl exec; el ejemplo no expone el puerto 8265 fuera del clúster.

Desplegar el clúster y enviar el trabajo

Diagrama conceptual de una región con varias zonas de disponibilidad y una VPC que abarca subredes; dos workers comparten zona y VPC para comunicarse por EFA.

Sigue este orden para que cada componente esté disponible cuando lo necesite el siguiente:

  1. Configura scripts/env.sh. Define las variables del clúster, la región, los nodos, la imagen y el entrenamiento. Así, los scripts posteriores usan la misma configuración.
  2. Ejecuta deploy_all.sh. El script comprueba la capacidad y crea el clúster y el grupo de nodos. La opción --plan previsualiza las acciones de infraestructura; --train encadena el entrenamiento después de configurar la infraestructura.
  3. Instala o verifica los complementos de GPU y EFA. El complemento de dispositivo debe hacer que esos recursos estén disponibles para Kubernetes antes de iniciar los pods workers.
  4. Instala KubeRay. El operador gestiona el recurso RayCluster que usarán los workers.
  5. Ejecuta deploy_ray_train_job.sh. El script genera y aplica el manifiesto del trabajo y pone en marcha el RayCluster.
  6. Comprueba los recursos asignables en los nodos GPU. Antes de esperar a los pods, verifica que Kubernetes anuncie GPU y EFA disponibles y que los nodos tengan la etiqueta nvidia.com/gpu.present=true.
  7. Envía el trabajo desde el Ray head pod. El ejemplo usa kubectl exec para hacerlo dentro del clúster, sin publicar externamente el puerto 8265.
  8. Inspecciona el transporte elegido por NCCL. Confirma que la ruta pasa por EFA; el script del ejemplo da error si EFA no forma parte de ella.

Comprobar recursos y evitar el fallback a sockets

Una solicitud de EFA en el pod permite al complemento montar el dispositivo. Si falta o el recurso no está disponible, NCCL puede recurrir a sockets TCP y el trabajo aún podría completarse con ese transporte. Comprueba tanto los recursos asignables de cada nodo como el transporte seleccionado por NCCL; que los pods arranquen, por sí solo, no confirma el uso de EFA.

Las requests y limits iguales del ejemplo también configuran los pods con calidad de servicio Guaranteed de Kubernetes. Revisa que los valores de CPU, memoria, GPU y EFA del manifiesto correspondan a los recursos que ofrece cada worker.

Qué demuestra el entrenamiento

El código usa identificadores de tokens aleatorios, no un conjunto de datos real. AWS plantea sustituirlos por una canalización de Ray Data o DataLoader para realizar un ajuste fino con datos reales. Por tanto, el ejemplo ejecuta el flujo distribuido y sus colectivas, pero no presenta un ajuste fino sobre un corpus.

AWS afirma que, con el mismo código del ejemplo, la ruta EFA duplicó el throughput frente a TCP. La afirmación se limita a ese ejemplo: AWS no aporta cifras de rendimiento ni condiciones del benchmark para trasladar el resultado a otras cargas.