Silvecer los flujos de trabajo de aprendizaje automático con Skypilot en Amazon Sagemaker Hyperpod

Esta publicación está coescrita con Zhanghao Wu, cocreador de Skypilot.

El rápido avance de los modelos generativos de IA y Foundation (FMS) ha aumentado significativamente los requisitos de trabajo de recursos computacionales para las cargas de trabajo de aprendizaje automático (ML). Las tuberías ML modernas requieren sistemas eficientes para distribuir cargas de trabajo a través de recursos de cómputo acelerados, al tiempo que se asegura de que la productividad del desarrollador siga siendo alta. Las organizaciones necesitan soluciones de infraestructura que no solo sean poderosas sino también flexibles, resistentes y directas de administrar.

Skypilot es un marco de código abierto que simplifica la ejecución de cargas de trabajo ML al proporcionar una capa de abstracción unificada que ayuda a los ingenieros de ML a ejecutar sus cargas de trabajo en diferentes recursos de calculación sin administrar complejidades de infraestructura subyacentes. Ofrece una interfaz simple de alto nivel para aprovisionar recursos, programar trabajos y administrar la capacitación distribuida en múltiples nodos.

Amazon Sagemaker Hyperpod es una infraestructura especialmente diseñada para desarrollar e implementar FMS a gran escala. Sagemaker HyperPod no solo proporciona la flexibilidad para crear y usar su propia pila de software, sino que también proporciona un rendimiento óptimo a través de la misma colocación de instancias de la columna, así como la resistencia incorporada. La combinación de la resistencia del Hyperpod de Sagemaker y la eficiencia de Skypilot proporciona un marco poderoso para ampliar sus cargas de trabajo generativas de IA.

En esta publicación, compartimos cómo Sagemaker HyperPod, en colaboración con Skypilot, está racionalización de los flujos de trabajo de desarrollo de IA. Esta integración hace que nuestra infraestructura de GPU avanzada sea más accesible para los ingenieros de ML, mejorando la productividad y la utilización de recursos.

Desafíos de organizar cargas de trabajo de aprendizaje automático

Kubernetes se ha vuelto popular para las cargas de trabajo de ML debido a su escalabilidad y sus ricas herramientas de código abierto. Sagemaker Hyperpod orquestado en Servicio de Kubernetes de Amazon Elastic (Amazon EKS) combina la potencia de Kubernetes con el entorno resistente de Sagemaker Hyperpod diseñado para capacitar a grandes modelos. El apoyo de Amazon EKS en Sagemaker HyperPod fortalece la resiliencia a través de controles de salud profunda, recuperación automatizada de nodos y capacidades de trabajo automático de trabajo, proporcionando capacitación ininterrumpida para trabajos a gran escala y de larga duración.

Los ingenieros de ML hacen la transición de los entornos tradicionales de VM o locales a menudo enfrentan una curva de aprendizaje empinada. La complejidad de Kubernetes manifiesta y la gestión del clúster puede plantear desafíos significativos, potencialmente ralentizando los ciclos de desarrollo y la utilización de recursos.

Además, los equipos de infraestructura de IA enfrentaron el desafío de equilibrar la necesidad de herramientas de gestión avanzadas con el deseo de proporcionar una experiencia fácil de usar para sus ingenieros de ML. Requerían una solución que pudiera ofrecer control de alto nivel y facilidad de uso para las operaciones diarias.

Sagemaker Hyperpod con Skypilot

Para abordar estos desafíos, nos asociamos con Skypilot para mostrar una solución que utiliza las fortalezas de ambas plataformas. Sagemaker Hyperpod sobresale en la gestión de los recursos e instancias de cómputo subyacentes, proporcionando la infraestructura robusta necesaria para exigir cargas de trabajo de IA. Skypilot complementa esto ofreciendo una capa intuitiva para la gestión de empleo, el desarrollo interactivo y la coordinación del equipo.

A través de esta asociación, podemos ofrecer a nuestros clientes lo mejor de ambos mundos: la infraestructura potente y escalable de Sagemaker Hyperpod, combinada con una interfaz fácil de usar que reduce significativamente la curva de aprendizaje para los ingenieros de ML. Para los equipos de infraestructura de IA, esta integración proporciona capacidades de gestión avanzadas al tiempo que simplifica la experiencia para sus ingenieros de ML, creando una situación de ganar-ganar para todos los interesados.

Skypilot ayuda a los equipos de IA a ejecutar sus cargas de trabajo en diferentes infraestructuras con una interfaz unificada de alto nivel y una gestión poderosa de recursos y trabajos. Un ingeniero de IA puede traer su marco de IA y especificar los requisitos de recursos para el trabajo; Skypilot programará de manera inteligente las cargas de trabajo en la mejor infraestructura: encuentre las GPU disponibles, aprovisione la GPU, ejecute el trabajo y administre su ciclo de vida.

Descripción general de la solución

La implementación de esta solución es sencilla, ya sea que esté trabajando con los clústeres existentes de Sagemaker HyperPod o configurando una nueva implementación. Para los grupos existentes, puede conectarse usando Interfaz de línea de comandos de AWS (AWS CLI) comandos para actualizar su kubeconfig y verificar la configuración. Para nuevas implementaciones, lo guiamos a través de la configuración del servidor API, la creación de grupos y la configuración de opciones de red de alto rendimiento como Elastic Fabric Adapter (EFA).

El siguiente diagrama ilustra la arquitectura de la solución.

En las siguientes secciones, mostramos cómo ejecutar trabajos de Skypilot para capacitación distribuida de múltiples nodos en Sagemaker HyperPod. Revisamos el proceso de crear un clúster de Sagemaker HyperPod, instalar Skypilot, crear un clúster de Skypilot e implementar un trabajo de capacitación Skypilot.

Requisitos previos

Debes tener los siguientes requisitos previos:

• Un clúster HyperPod de Sagemaker existente con Amazon EKS (para crear uno, consulte Implemente su clúster HyperPod). Debe aprovisionar una sola instancia de ml.p5.48xLarge para las muestras de código en las siguientes secciones.
• Acceso a la CLI de AWS y kubectl Herramientas de línea de comando.
• Un entorno de Python para instalar Skypilot.

Crea un clúster Hyperpod de Sagemaker

Puedes crear un clúster de EKS con un solo AWS CloudFormation pila siguiendo las instrucciones en Uso de CloudFormationconfigurado con una nube privada virtual (VPC) y recursos de almacenamiento.

Para crear y administrar clústeres de Sagemaker HyperPod, puede usar el Consola de gestión de AWS o AWS Cli. Si usa la CLI de AWS, especifique la configuración del clúster en un archivo JSON y elija el clúster EKS creado desde la pila CloudFormation como el orquestador del clúster Sagemaker HyperPod. Luego crea los nodos de los trabajadores de clúster con NodeRecovery empezar a Automatic para habilitar la recuperación de nodo automático y para OnStartDeepHealthChecksagregar InstanceStress y InstanceConnectivity para habilitar los controles de salud profundas. Vea el siguiente código:

cat > cluster-config.json << EOL
{
    "ClusterName": "hp-cluster",
    "Orchestrator": {
        "Eks": {
            "ClusterArn": "${EKS_CLUSTER_ARN}"
        }
    },
    "InstanceGroups": [
        {
            "InstanceGroupName": "worker-group-1",
            "InstanceType": "ml.p5.48xlarge",
            "InstanceCount": 2,
            "LifeCycleConfig": {
                "SourceS3Uri": "s3://${BUCKET_NAME}",
                "OnCreate": "on_create.sh"
            },
            "ExecutionRole": "${EXECUTION_ROLE}",
            "ThreadsPerCore": 1,
            "OnStartDeepHealthChecks": [
                "InstanceStress",
                "InstanceConnectivity"
            ],
        },
  ....
    ],
    "VpcConfig": {
        "SecurityGroupIds": [
            "$SECURITY_GROUP"
        ],
        "Subnets": [
            "$SUBNET_ID"
        ]
    },
    "ResilienceConfig": {
        "NodeRecovery": "Automatic"
    }
}
EOL

Puedes agregar Instancestorageconfigs para aprovisionar y montar adicionales Tienda de bloques elástica de Amazon (Amazon EBS) volúmenes en Sagemaker Hyperpod Nodos.

Para crear el clúster usando el API SAGEMAKER HYPERPODejecute el siguiente comando AWS CLI:

aws sagemaker create-cluster \ 
--cli-input-json file://cluster-config.json

Ahora está listo para configurar Skypilot en su clúster HyperPod de Sagemaker.

Conéctese a su clúster EKS HyperPod EKS de Sagemaker

Desde su entorno AWS CLI, ejecute el AWS EKS Update-KubeConfig comandar para actualizar su archivo de configuración de Kube local (ubicado en ~/.kube/config) con las credenciales y la configuración necesarias para conectarse a su clúster EKS utilizando el kubectl Comando (proporcione su nombre específico de clúster EKS):

aws eks update-kubeconfig --name $EKS_CLUSTER_NAME

Puede verificar que está conectado al clúster EKS ejecutando el siguiente comando:

kubectl config current-context

Instale Skypilot con soporte de Kubernetes

Use el siguiente código para instalar Skypilot con Kubernetes Soporte usando PIP:

pip install skypilot[kubernetes]

Esto instala la última compilación de Skypilot, que incluye las integraciones necesarias de Kubernetes.

Verificar la conexión de Skypilot con el clúster EKS

Compruebe si Skypilot puede conectarse a su clúster Kubernetes:

sky check k8s

La salida debe parecer similar al siguiente código:

Checking credentials to enable clouds for SkyPilot.
Kubernetes: enabled [compute]

To enable a cloud, follow the hints above and rerun: sky check
If any problems remain, refer to detailed docs at: https://docs.skypilot.co/en/latest/getting-started/installation.html

🎉 Enabled clouds 🎉
Kubernetes [compute]
Active context: arn:aws:eks:us-east-2:XXXXXXXXXXXXX:cluster/sagemaker-hyperpod-eks-cluster

Using SkyPilot API server: http://127.0.0.1:46580

Si esta es la primera vez que usa Skypilot con este clúster Kubernetes, puede ver un aviso para crear etiquetas GPU para sus nodos. Siga las instrucciones ejecutando el siguiente código:

python -m sky.utils.kubernetes.gpu_labeler --context <your-eks-context>

Este script ayuda a Skypilot a identificar qué recursos de GPU están disponibles en cada nodo en su clúster. El Trabajo de etiquetado de GPU Puede tomar unos minutos dependiendo de la cantidad de recursos de GPU en su clúster.

Descubra las GPU disponibles en el clúster

Para ver qué recursos de GPU están disponibles en su clúster HyperPod de Sagemaker, use el siguiente código:

sky show-gpus --cloud k8s

Esto enumerará los tipos de GPU disponibles y sus recuentos. Tenemos dos instancias P5.48xLarge, cada una equipada con 8 GPU H100 NVIDIA:

 Kubernetes GPUs
GPU REQUESTABLE_QTY_PER_NODE TOTAL_GPUS TOTAL_FREE_GPUS
H100 1, 2, 4, 8 16 16

Kubernetes per node accelerator availability
NODE_NAME GPU_NAME TOTAL_GPUS FREE_GPUS
hyperpod-i-00baa178bc31afde3 H100 8 8
hyperpod-i-038beefa954efab84 H100 8 8

Lanzar un entorno de desarrollo interactivo

Con Skypilot, puede lanzar un clúster Skypilot para el desarrollo interactivo:

sky launch -c dev --gpus H100

Este comando crea un entorno de desarrollo interactivo (IDE) con una única GPU H100 y sincronizará el directorio de trabajo local con el clúster. Skypilot maneja la creación de POD, la asignación de recursos y la configuración del IDE.

Considered resources (1 node):
-------------------------------------------------------------------------------------------------------------------------------------------------------------------
 CLOUD        INSTANCE            vCPUs   Mem(GB)   ACCELERATORS   REGION/ZONE                                                                 COST ($)   CHOSEN   
-------------------------------------------------------------------------------------------------------------------------------------------------------------------
 Kubernetes   2CPU--8GB--H100:1   2       8         H100:1         arn:aws:eks:us-east-2:XXXXXXXXXX:cluster/sagemaker-hyperpod-eks-cluster   0.00          ✔     
------------------------------------------------------------------------------------------------------------------------------------------------------------------
Launching a new cluster 'dev'. Proceed? [Y/n]: Y
• Launching on Kubernetes.
Pod is up.
✔ Cluster launched: dev. View logs: sky api logs -1 sky-2025-05-05-15-28-47-523797/provision. log
• Syncing files.
Run commands not specified or empty.
Useful Commands
Cluster name: dey
To log into the head VM:   ssh dev
To submit a job:           sky exec dev yaml_file
To stop the cluster:       sky stop dev
To teardown the cluster:   sky down dev

Después de que se inicie, puede conectarse a su IDE:

ssh dev

Esto le brinda una carcasa interactiva en su IDE, donde puede ejecutar su código, instalar paquetes y realizar experimentos ML.

Trabajos de capacitación para ejecutar

Con Skypilot, puede ejecutar trabajos de capacitación distribuidos en su clúster Sagemaker HyperPod. El siguiente es un ejemplo de lanzar un trabajo de capacitación distribuido utilizando un archivo de configuración YAML.

Primero, cree un archivo llamado train.yaml Con la configuración de su trabajo de capacitación:

resources:
    accelerators: H100

num_nodes: 1

setup: |
    git clone --depth 1 https://github.com/pytorch/examples || true
    cd examples
    git filter-branch --prune-empty --subdirectory-filter distributed/minGPT-ddp
    # SkyPilot's default image on AWS/GCP has CUDA 11.6 (Azure 11.5).
    uv venv --python 3.10
    source .venv/bin/activate
    uv pip install -r requirements.txt "numpy<2" "torch"

run: |
    cd examples
    source .venv/bin/activate
    cd mingpt
    export LOGLEVEL=INFO

    MASTER_ADDR=$(echo "$SKYPILOT_NODE_IPS" | head -n1)
    echo "Starting distributed training, head node: $MASTER_ADDR"

    torchrun \
    --nnodes=$SKYPILOT_NUM_NODES \
    --nproc_per_node=$SKYPILOT_NUM_GPUS_PER_NODE \
    --master_addr=$MASTER_ADDR \
    --master_port=8008 \
    --node_rank=${SKYPILOT_NODE_RANK} \
    main.py

Luego, lance su trabajo de capacitación:

sky launch -c train train.yaml

Esto crea un trabajo de capacitación en un solo nodos P5.48xLarge, equipado con 8 GPU H100 NVIDIA. Puede monitorear la salida con el siguiente comando:

sky logs train

Ejecución de trabajos de capacitación de múltiples nodos con EFA

El adaptador de tela elástico (EFA) es una interfaz de red para Nube de cómputo elástica de Amazon (Amazon EC2) instancias que le permiten ejecutar aplicaciones que requieren altos niveles de comunicaciones entre nodos a escala en AWS a través de su interfaz de hardware de derivación del sistema operativo personalizado. Esto permite que las aplicaciones se comuniquen directamente con el hardware de la red al tiempo que evita el núcleo del sistema operativo, reduciendo significativamente la latencia y la sobrecarga de la CPU. Este acceso directo al hardware es particularmente beneficioso para las cargas de trabajo ML distribuidas donde la comunicación frecuente entre nodos durante la sincronización de gradiente puede convertirse en un cuello de botella. Mediante el uso de instancias habilitadas para la EFA, como P5.48xLarge o P6-B200.48xLarge, los científicos de datos pueden escalar sus trabajos de capacitación en múltiples nodos mientras mantienen la comunicación de baja latencia y alto nivel de ancho de banda para la capacitación distribuida eficiente, reduciendo el tiempo de capacitación y la mejora de la utilización de recursos para los trabajos de IA a gran escala.

El siguiente fragmento de código muestra cómo incorporar esto en su trabajo de Skypilot:

name: nccl-test-efa

resources:
  cloud: kubernetes
  accelerators: H100:8
  image_id: docker:public.ecr.aws/hpc-cloud/nccl-tests:latest

num_nodes: 2

envs:
  USE_EFA: "true"

run: |
  if [ "${SKYPILOT_NODE_RANK}" == "0" ]; then
    echo "Head node"

    # Total number of processes, NP should be the total number of GPUs in the cluster
    NP=$(($SKYPILOT_NUM_GPUS_PER_NODE * $SKYPILOT_NUM_NODES))

    # Append :${SKYPILOT_NUM_GPUS_PER_NODE} to each IP as slots
    nodes=""
    for ip in $SKYPILOT_NODE_IPS; do
      nodes="${nodes}${ip}:${SKYPILOT_NUM_GPUS_PER_NODE},"
    done
    nodes=${nodes::-1}
    echo "All nodes: ${nodes}"

    # Set environment variables
    export PATH=$PATH:/usr/local/cuda-12.2/bin:/opt/amazon/efa/bin:/usr/bin
    export LD_LIBRARY_PATH=/usr/local/cuda-12.2/lib64:/opt/amazon/openmpi/lib:/opt/nccl/build/lib:/opt/amazon/efa/lib:/opt/aws-ofi-nccl/install/lib:/usr/local/nvidia/lib:$LD_LIBRARY_PATH
    export NCCL_HOME=/opt/nccl
    export CUDA_HOME=/usr/local/cuda-12.2
    export NCCL_DEBUG=INFO
    export NCCL_BUFFSIZE=8388608
    export NCCL_P2P_NET_CHUNKSIZE=524288
    export NCCL_TUNER_PLUGIN=/opt/aws-ofi-nccl/install/lib/libnccl-ofi-tuner.so

    if [ "${USE_EFA}" == "true" ]; then
      export FI_PROVIDER="efa"
    else
      export FI_PROVIDER=""
    fi

    /opt/amazon/openmpi/bin/mpirun \
      --allow-run-as-root \
      --tag-output \
      -H $nodes \
      -np $NP \
      -N $SKYPILOT_NUM_GPUS_PER_NODE \
      --bind-to none \
      -x FI_PROVIDER \
      -x PATH \
      -x LD_LIBRARY_PATH \
      -x NCCL_DEBUG=INFO \
      -x NCCL_BUFFSIZE \
      -x NCCL_P2P_NET_CHUNKSIZE \
      -x NCCL_TUNER_PLUGIN \
      --mca pml ^cm,ucx \
      --mca btl tcp,self \
      --mca btl_tcp_if_exclude lo,docker0,veth_def_agent \
      /opt/nccl-tests/build/all_reduce_perf \
      -b 8 \
      -e 2G \
      -f 2 \
      -g 1 \
      -c 5 \
      -w 5 \
      -n 100
  else
    echo "Worker nodes"
  fi

config:
  kubernetes:
    pod_config:
      spec:
        containers:
        - resources:
            limits:
              
              vpc.amazonaws.com/efa: 32
            requests:
              
              vpc.amazonaws.com/efa: 32

Limpiar

Para eliminar su clúster Skypilot, ejecute el siguiente comando:

sky down <cluster_name>

Para eliminar el clúster HyperPod creado en esta publicación, puede usar la consola AI Sagemaker o el siguiente comando AWS CLI:

aws sagemaker delete-cluster --cluster-name <cluster_name>

La eliminación del clúster tomará unos minutos. Puede confirmar la eliminación exitosa después de no ver grupos en la consola AI Sagemaker.

Si usó la pila CloudFormation para crear recursos, puede eliminarlo utilizando el siguiente comando:

aws cloudformation delete-stack --stack-name <stack_name>

Conclusión

Al combinar las sólidas capacidades de infraestructura de Sagemaker HyperPod con la interfaz fácil de usar Skypilot, hemos mostrado una solución que ayuda a los equipos a centrarse en la innovación en lugar de la complejidad de la infraestructura. Este enfoque no solo simplifica las operaciones, sino que también mejora la productividad y la utilización de recursos en todas las organizaciones de todos los tamaños. Para comenzar, consulte Skypilot en el Soporte de Amazon EKS en Amazon Sagemaker HyperPod Workshop.

Sobre los autores

Roy alela es un arquitecto senior de soluciones especializadas de IA/ML en AWS. Ayuda a los clientes de AWS, desde pequeñas nuevas empresas hasta grandes empresas, a arrastrar y implementar modelos de base de manera eficiente en AWS. Le apasiona los problemas de optimización computacional y mejora el rendimiento de las cargas de trabajo de IA.

Zhanghao wu es co-creador del Proyecto de Código Abierto de Skypilot y posee un doctorado en Ciencias de la Computación de UC Berkeley. Trabaja en Skypilot Core, Arquitectura de cliente cliente, trabajos administrados y mejorando la experiencia de IA en diversas infraestructura en la nube en general.

Ankit anand es especialista en el mercado de la Fundación Senior Modelos (GTM) en AWS. Se asocia con los principales constructores de modelos de IA generativos, clientes estratégicos y equipos de servicio de AWS para habilitar la próxima generación de cargas de trabajo AI/ML en AWS. La experiencia de ANKIT incluye experiencia en gestión de productos dentro de la industria de servicios financieros para el comercio de alta frecuencia y de baja latencia y desarrollo de negocios para Amazon Alexa.