Escale los flujos de trabajo de aprendizaje automático con Amazon SageMaker Studio y Amazon SageMaker HyperPod

Escalar los flujos de trabajo de aprendizaje automático (ML) desde los prototipos iniciales hasta la implementación de producción a gran escala puede ser una tarea desalentadora, pero la integración de Amazon SageMaker Estudio y HiperPod de Amazon SageMaker ofrece una solución simplificada a este desafío. A medida que los equipos avanzan desde la prueba de concepto hasta los modelos listos para producción, a menudo tienen dificultades para gestionar de manera eficiente las crecientes necesidades de infraestructura y almacenamiento. Esta integración aborda estos obstáculos al proporcionar a los científicos de datos y a los ingenieros de ML un entorno integral que respalda todo el ciclo de vida de ML, desde el desarrollo hasta la implementación a escala.

En esta publicación, lo guiamos a través del proceso de escalar sus cargas de trabajo de ML usando SageMaker Studio y SageMaker HyperPod.

Descripción general de la solución

La implementación de la solución consta de los siguientes pasos de alto nivel:

1. Configure su entorno y los permisos para acceder a los clústeres de Amazon HyperPod en SageMaker Studio.
2. Cree un espacio JupyterLab y monte un Amazon FSx para brillo sistema de archivos a su espacio. Esto elimina la necesidad de migrar datos o cambiar el código a medida que escala. Esto también mitiga los posibles problemas de reproducibilidad que a menudo surgen de las discrepancias de datos en las diferentes etapas del desarrollo del modelo.
3. Ahora puede utilizar SageMaker Studio para descubrir los clústeres de SageMaker HyperPod y ver los detalles y las métricas del clúster. Cuando tiene acceso a varios clústeres, esta información puede ayudarle a comparar las especificaciones de cada clúster, la utilización actual y el estado de la cola de los clústeres para identificar cuál cumple con los requisitos de su tarea de aprendizaje automático específica.
4. Usamos un cuaderno de muestra para mostrar cómo conectarse al clúster y ejecutar un trabajo de entrenamiento de Meta Llama 2 con PyTorch FSDP en su clúster Slurm.
5. Después de enviar el trabajo de larga duración al clúster, puede monitorear las tareas directamente a través de la interfaz de usuario de SageMaker Studio. Esto puede ayudarle a obtener información en tiempo real sobre sus flujos de trabajo distribuidos y permitirle identificar rápidamente cuellos de botella, optimizar la utilización de recursos y mejorar la eficiencia general del flujo de trabajo.

Este enfoque integrado no solo agiliza la transición del prototipo a la capacitación a gran escala, sino que también mejora la productividad general al mantener una experiencia de desarrollo familiar incluso cuando se escala a cargas de trabajo de nivel de producción.

Requisitos previos

Complete los siguientes pasos previos:

1. Cree un clúster de SageMaker HyperPod Slurm. Para obtener instrucciones, consulte la Taller de Amazon SageMaker HyperPod o Tutorial para empezar a utilizar SageMaker HyperPod.
2. Asegúrate de tener el última versión del Interfaz de línea de comandos de AWS (AWS CLI).
3. Cree un usuario en el nodo principal de Slurm o en el nodo de inicio de sesión con un UID mayor que 10000. Consulte Multiusuario para obtener instrucciones para crear un usuario.
4. Etiquete el clúster de SageMaker HyperPod con la clave hyperpod-cluster-filesystem. Este es el ID del sistema de archivos FSx for Lustre asociado con el clúster SageMaker HyperPod. Esto es necesario para que Amazon SageMaker Studio monte FSx for Lustre en los espacios de Jupyter Lab y Code Editor. Utilice el siguiente fragmento de código para agregar una etiqueta a un clúster de SageMaker HyperPod existente:

   aws sagemaker add-tags --resource-arn <cluster_ARN> \
   --tags Key=hyperpod-cluster-filesystem,Value=<fsx_id>

Configura tus permisos

En las siguientes secciones, describimos los pasos para crear un Amazon SageMaker dominio, cree un usuario, configure un espacio de SageMaker Studio y conéctese al clúster de SageMaker HyperPod. Al final de estos pasos, debería poder conectarse a un clúster de SageMaker HyperPod Slurm y ejecutar una carga de trabajo de entrenamiento de muestra. Para seguir las instrucciones de configuración, debe tener privilegios de administrador. Complete los siguientes pasos:

1. Crear un nuevo Gestión de acceso e identidad de AWS (IAM) función de ejecución con AmazonSageMakerAcceso completo adscrito al rol. Adjunte también la siguiente política JSON a la función, que permite a SageMaker Studio acceder al clúster de SageMaker HyperPod. Asegúrese de que la relación de confianza en el rol permita al sagemaker.amazonaws.com servicio para asumir este rol.

{
    "Version": "2012-10-17",            
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "ssm:StartSession",
                "ssm:TerminateSession"
            ],
            "Resource": "*"    
        }
{
            "Effect": "Allow",
            "Action": [
                "sagemaker:CreateCluster",
                "sagemaker:ListClusters"
            ],
            "Resource": "*"    
        },
        {
            "Effect": "Allow",
            "Action": [
                "sagemaker:DescribeCluster",
                "sagemaker:DescribeClusterNode",
                "sagemaker:ListClusterNodes",
                "sagemaker:UpdateCluster",
                "sagemaker:UpdateClusterSoftware"
            ],
            "Resource": "arn:aws:sagemaker:region:account-id:cluster/*"    
        }
    ]
}

2. Para utilizar el rol que creó para acceder al encabezado del clúster SageMaker HyperPod o al nodo de inicio de sesión usando Gerente de sistemas AWSnecesitas agregar una etiqueta a este rol de IAM, donde Tag Key = “SSMSessionRunAs” y Tag Value = “<posix user>”. El usuario POSIX es el usuario que está configurado en el nodo principal de Slurm. Systems Manager utiliza este usuario para ejecutar en el nodo principal.
3. Cuando activa la compatibilidad con Ejecutar como, impide que Session Manager inicie sesiones utilizando el ssm-user cuenta en un nodo administrado. Para habilitar Ejecutar como en el Administrador de sesiones, complete los siguientes pasos:
   1. En la consola del Administrador de sesión, elija Preferenciasluego elige Editar.
   2. No especifique ningún nombre de usuario. El nombre de usuario se seleccionará de la etiqueta de rol. SSMSessionRunAs que creaste anteriormente.
   3. En el Perfil de shell de Linux sección, ingrese /bin/bash.
   4. Elegir Ahorrar.

4. Cree un nuevo dominio de SageMaker Studio con la función de ejecución creada anteriormente junto con otros parámetros necesarios para acceder al clúster de SageMaker HyperPod. Utilice el siguiente script para crear el dominio y reemplazar las variables de exportación en consecuencia. Aquí, VPC_ID y Subnet_ID son los mismos que la VPC y la subred del clúster SageMaker HyperPod. El EXECUTION_ROLE_ARN es el rol que creaste anteriormente.

export DOMAIN_NAME=<domain name>
export VPC_ID=vpc_id-for_hp_cluster
export SUBNET_ID=private_subnet_id
export EXECUTION_ROLE_ARN=execution_role_arn
export FILE_SYSTEM_ID=fsx id
export USER_UID=10000
export USER_GID=1001
export REGION=us-east-2

cat > user_settings.json << EOL
{
    "ExecutionRole": "$EXECUTION_ROLE_ARN",
    "CustomPosixUserConfig":
    {
        "Uid": $USER_UID,
        "Gid": $USER_GID
    },
    "CustomFileSystemConfigs":
    [
        {
            "FSxLustreFileSystemConfig":
            {
                "FileSystemId": "$FILE_SYSTEM_ID",
                "FileSystemPath": "$FILE_SYSTEM_PATH"
            }
        }
    ]
}
EOL

aws sagemaker create-domain \
--domain-name $DOMAIN_NAME \
--vpc-id $VPC_ID \
--subnet-ids $SUBNET_ID \
--auth-mode IAM \
--default-user-settings file://user_settings.json \
--region $REGION 

El UID y el GID en la configuración anterior están configurados en 10000 y 1001 por defecto; esto se puede anular según el usuario creado en Slurm, y este UID/GID se utiliza para otorgar permisos al sistema de archivos FSx para Lustre. Además, configurar esto a nivel de dominio le da a cada usuario el mismo UID. Para tener un UID separado para cada usuario, considere configurar CustomPosixUserConfig al crear el perfil de usuario.

5. Después de crear el dominio, debe adjuntar SecurityGroupIdForInboundNfs creado como parte de la creación de dominio para todos los ENI del volumen FSx Lustre:
   1. Localice el Sistema de archivos elástico de Amazon (Amazon EFS) sistema de archivos asociado con el dominio y el grupo de seguridad correspondiente adjunto al mismo. Puede encontrar el sistema de archivos EFS en la consola de Amazon EFS; está etiquetado con el ID del dominio, como se muestra en la siguiente captura de pantalla.
   2. Recoja el grupo de seguridad correspondiente, que se denominará inbound-nfs-<domain-id> y se puede encontrar en el Red pestaña.
      
   3. En la consola FSx for Lustre, elija Para ver todos los ENI, consulte la consola de Amazon EC2. Esto mostrará todos los ENI adjuntos a FSx para Lustre. Alternativamente, puede encontrar ENI utilizando la CLI de AWS o llamando al fsx:describeFileSystems
   4. Para cada ENI, adjunte el SecurityGroupIdForInboundNfs del dominio a éste.

Como alternativa, puede utilizar el siguiente script para buscar y adjuntar automáticamente grupos de seguridad a los ENI asociados con el volumen FSx for Lustre. Reemplace el REGION, DOMAIN_IDy FSX_ID atributos en consecuencia.

#!/bin/bash

export REGION=us-east-2
export DOMAIN_ID=d-xxxxx
export FSX_ID=fs-xxx

export EFS_ID=$(aws sagemaker describe-domain --domain-id $DOMAIN_ID --region $REGION --query 'HomeEfsFileSystemId' --output text)
export MOUNT_TARGET_ID=$(aws efs describe-mount-targets --file-system-id $EFS_ID --region $REGION --query 'MountTargets[0].MountTargetId' --output text)
export EFS_SG=$(aws efs describe-mount-target-security-groups --mount-target-id $MOUNT_TARGET_ID --query 'SecurityGroups[0]' --output text)
echo "security group associated with the Domain $EFS_SG"

echo "Adding security group to FSxL file system ENI's"
# Get the network interface IDs associated with the FSx file system
NETWORK_INTERFACE_IDS=$(aws fsx describe-file-systems --file-system-ids $FILE_SYSTEM_ID --query "FileSystems[0].NetworkInterfaceIds" --output text)
# Iterate through each network interface and attach the security group
for ENI_ID in $NETWORK_INTERFACE_IDS; do
aws ec2 modify-network-interface-attribute --network-interface-id $ENI_ID --groups $EFS_SG
echo "Attached security group $EFS_SG to network interface $ENI_ID"
done

Sin este paso, la creación de la aplicación fallará y generará un error.

6. Después de crear el dominio, puede utilizarlo para crear un perfil de usuario. Reemplace el valor DOMAIN_ID por el creado en el paso anterior.

export DOMAIN_ID=d-xxx
export USER_PROFILE_NAME=test
export REGION=us-east-2

aws sagemaker create-user-profile \
--domain-id $DOMAIN_ID \
--user-profile-name$USER_PROFILE_NAME \
--region $REGION

Cree un espacio JupyterLab y monte el sistema de archivos FSx para Lustre

Cree un espacio usando el sistema de archivos FSx para Lustre con el siguiente código:

export SPACE_NAME=hyperpod-space
export DOMAIN_ID=d-xxx
export USER_PROFILE_NAME=test
export FILE_SYSTEM_ID=fs-xxx
export REGION=us-east-2

aws sagemaker create-space --domain-id $DOMAIN_ID \
--space-name $SPACE_NAME \
--space-settings "AppType=JupyterLab,CustomFileSystems=[{FSxLustreFileSystem={FileSystemId=$FILE_SYSTEM_ID}}]" \
--ownership-settings OwnerUserProfileName=$USER_PROFILE_NAME  --space-sharing-settings SharingType=Private  \
--region $REGION

Crea una aplicación usando el espacio con el siguiente código:

export SPACE_NAME=hyperpod-space
export DOMAIN_ID=d-xxx
export APP_NAME=test-app
export INSTANCE_TYPE=ml.t3.medium
export REGION=us-east-2
export IMAGE_ARN=arn:aws:sagemaker:us-east-2:081975978581:image/sagemaker-distribution-cpu

aws sagemaker create-app --space-name $SPACE_NAME \
--resource-spec '{"InstanceType":"$INSTANCE_TYPE","SageMakerImageArn":"$IMAGE_ARN"}' \
--domain-id  $DOMAIN_ID --app-type JupyterLab --app-name $APP_NAME --region $REGION

Descubra clústeres en SageMaker Studio

Ahora debería tener todo listo para acceder al clúster SageMaker HyperPod usando SageMaker Studio. Complete los siguientes pasos:

1. En la consola de SageMaker, elija Configuraciones de administrador, Dominios.
2. Localice el perfil de usuario que creó e inicie SageMaker Studio.
3. Bajo Calcular en el panel de navegación, elija Clústeres de HyperPod.

Aquí puede ver los clústeres de SageMaker HyperPod disponibles en la cuenta.

4. Identifique el clúster adecuado para su carga de trabajo de capacitación observando los detalles del clúster y las métricas de hardware del clúster.

También puede obtener una vista previa del clúster eligiendo el icono de flecha.

También puedes ir al Ajustes y Detalles pestañas para encontrar más información sobre el clúster.

Trabaje en SageMaker Studio y conéctese al clúster

También puede iniciar JupyterLab o Code Editor, que monta el volumen del clúster FSx para Lustre para desarrollo y depuración.

1. En SageMaker Studio, elija Empezar en y elige JupyterLab.
2. Elija un espacio que tenga montado el sistema de archivos FSx for Lustre para obtener un entorno consistente y reproducible.

El Sistema de archivos de clúster La columna identifica qué espacio tiene montado el sistema de archivos del clúster.

Esto debería iniciar JupyterLab con el volumen FSx for Lustre montado. De forma predeterminada, debería ver el cuaderno de introducción en su carpeta de inicio, que tiene instrucciones paso a paso para ejecutar un trabajo de entrenamiento de Meta Llama 2 con PyTorch FSDP en el clúster Slurm. Este cuaderno de ejemplo demuestra cómo puede utilizar los cuadernos de SageMaker Studio para realizar la transición desde la creación de prototipos de su script de entrenamiento hasta la ampliación de sus cargas de trabajo en varias instancias en el entorno del clúster. Además, debería ver el sistema de archivos FSx para Lustre que montó en su espacio JupyterLab en /home/sagemaker-user/custom-file-systems/fsx_lustre.

Supervisar las tareas en SageMaker Studio

Puede ir a SageMaker Studio y elegir el clúster para ver una lista de tareas actualmente en la cola de Slurm.

Puede elegir una tarea para obtener detalles adicionales de la tarea, como la programación y el estado del trabajo, los detalles del uso de recursos y el envío y los límites del trabajo.

También puede realizar acciones como liberar, volver a poner en cola, suspender y mantener estas tareas de Slurm utilizando la interfaz de usuario.

Limpiar

Complete los siguientes pasos para limpiar sus recursos:

1. Eliminar el espacio:

aws —region <REGION> sagemaker delete-space \
--domain-id <DomainId> \
--space-name <SpaceName>

2. Eliminar el perfil de usuario:

aws —region <REGION> sagemaker delete-user-profile \
--domain-id <DomainId> \
--user-profile-name <UserProfileName>

3. Eliminar el dominio. Para conservar el volumen EFS, especifique HomeEfsFileSystem=Retain.

aws —region <REGION> sagemaker delete-domain \
--domain-id <DomainId> \
--retention-policy HomeEfsFileSystem=Delete

4. Eliminar el clúster SageMaker HyperPod.
5. Eliminar el rol de IAM tu creaste.

Conclusión

En esta publicación, exploramos un enfoque para optimizar sus flujos de trabajo de ML utilizando SageMaker Studio. Demostramos cómo puede pasar sin problemas de crear un prototipo de su script de capacitación dentro de SageMaker Studio a ampliar su carga de trabajo en múltiples instancias en un entorno de clúster. También explicamos cómo montar el volumen del clúster FSx para Lustre en sus espacios de SageMaker Studio para obtener un entorno reproducible consistente.

Este enfoque no solo agiliza su proceso de desarrollo sino que también le permite iniciar trabajos de larga duración en los clústeres y monitorear cómodamente su progreso directamente desde SageMaker Studio.

Le recomendamos que pruebe esto y comparta sus comentarios en la sección de comentarios.

Un agradecimiento especial a Durga Sury (Sr. ML SA), Monidipa Chakraborty (Sr. SDE) y Sumedha Swamy (Sr. Manager PMT) por su apoyo al lanzamiento de esta publicación.

Acerca de los autores

Arun Kumar Lokanatha es arquitecto sénior de soluciones de aprendizaje automático en el equipo de Amazon SageMaker. Se especializa en grandes cargas de trabajo de capacitación de modelos de lenguaje, ayudando a los clientes a crear cargas de trabajo de LLM utilizando SageMaker HyperPod, trabajos de capacitación de SageMaker y capacitación distribuida de SageMaker. Fuera del trabajo, le gusta correr, hacer senderismo y cocinar.

Pooja Karadgi es gerente técnico senior de productos en Amazon Web Services. En AWS, forma parte del equipo de Amazon SageMaker Studio y ayuda a crear productos que satisfagan las necesidades de administradores y científicos de datos. Comenzó su carrera como ingeniera de software antes de hacer la transición a la gestión de productos. Fuera del trabajo, le gusta elaborar planificadores de viajes en hojas de cálculo, al más puro estilo MBA. Dado el tiempo que invierte en la creación de estos planificadores, está claro que siente un profundo amor por viajar, además de una gran pasión por el senderismo.