Entrena Llama2 con Megatron-LM en máquinas virtuales A3 Mega


Descripción general

En esta guía de inicio rápido, aprenderás a ejecutar una carga de trabajo de PyTorch basada en contenedores Megatron-LM en A3 Mega. El código está disponible en este repositorio de GitHub: megatron-gke.

Antes de comenzar

Siga los pasos que se indican a continuación para habilitar la API de Google Kubernetes Engine (GKE):

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. Enable the GKE API.

    Enable the API

  5. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  6. Make sure that billing is enabled for your Google Cloud project.

  7. Enable the GKE API.

    Enable the API

  8. Make sure that you have the following role or roles on the project: roles/container.admin, roles/compute.networkAdmin, roles/iam.serviceAccountUser

    Check for the roles

    1. In the Google Cloud console, go to the IAM page.

      Go to IAM
    2. Select the project.
    3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

    4. For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.

    Grant the roles

    1. In the Google Cloud console, go to the IAM page.

      Ir a IAM
    2. Selecciona el proyecto.
    3. Haz clic en Grant access.
    4. En el campo Principales nuevas, ingresa tu identificador de usuario. Esta suele ser la dirección de correo electrónico de una Cuenta de Google.

    5. En la lista Seleccionar un rol, elige un rol.
    6. Para otorgar funciones adicionales, haz clic en Agregar otro rol y agrega cada rol adicional.
    7. Haz clic en Guardar.

    Crea un clúster A3 Mega

    Crea un clúster de GKE A3 Mega con GPUDirect-TCPXO y varias redes. Para obtener más información, consulta Maximiza el ancho de banda de red de la GPU con GPUDirect y varias redes.

Configure su entorno

  1. Crea variables de entorno para algunos parámetros comunes

    export CLUSTER_NAME=CLUSTER_NAME
    export REGION=REGION
    export ZONE=ZONE
    export PROJECT_ID=PROJECT_ID
    

    Reemplaza lo siguiente:

    • CLUSTER_NAME: el nombre del clúster de GKE Mega A3 que tiene habilitadas GPUDirect-TCPXO y redes múltiples.
    • REGION: Es la región en la que creaste el clúster.
    • ZONE: Es la zona en la que creaste el clúster.
    • PROJECT_ID: Es el ID de tu proyecto de Google Cloud.
  2. Configura Google Cloud CLI para usar tus credenciales de Google Cloud para la autenticación:

    gcloud auth login
    

    Si deseas obtener más información, consulta Autentica para usar Google Cloud CLI.

  3. Instala kubectl y el complemento de gcloud CLI de GKE:

    sudo apt-get install kubectl
    sudo apt-get install google-cloud-sdk-gke-gcloud-auth-plugin
    
  4. Recupera las credenciales para tu clúster de GKE:

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
      --zone=${ZONE} \
      --project=${PROJECT_ID}
    
  5. Si aún no lo está, instala Helm:

    curl -fsSL -o get_helm.sh https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3
    chmod 700 get_helm.sh
    ./get_helm.sh && rm get_helm.sh
    sudo chmod +x /usr/local/bin/helm
    

Usa un programador adaptado a la topología para implementar tus Pods

Puedes usar el programador adaptado a la topología para implementar tus Pods de GKE en nodos que tienen una topología de GPU especificada.

En los siguientes comandos de kubectl, usarás los archivos directamente desde un repositorio. De manera alternativa, puedes clonar el repositorio de forma local y los comandos kubectl pueden hacer referencia a los archivos locales en su lugar.

Para obtener más información, consulta Programador de topologías.

  1. Configura la cuenta de servicio:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/service-account.yaml
    
  2. Instala las secuencias de comandos del programador de topología en un ConfigMap:

    curl -OL  https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/schedule-daemon.py
    curl -OL  https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/label-nodes-daemon.py
    
    kubectl -n kube-system create configmap topology-scheduler-scripts \
        --from-file=schedule-daemon.py=schedule-daemon.py \
        --from-file=label-nodes-daemon.py=label-nodes-daemon.py
    
  3. Instala el daemonset de la etiqueta de topología y el Pod del programador de topología:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/label-nodes-daemon.yaml
    $ kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/schedule-daemon.yaml
    
  4. Observa las acciones del programador de topología:

    kubectl -n kube-system logs topology-scheduler-pod
    

Ejecuta la carga de trabajo

Compila el Dockerfile y envíalo a Artifact Registry de Google Cloud

  1. Crea un bucket de Cloud Storage y un repositorio de Docker. En scripts/setup-and-configure-resources.sh script, reemplaza los nombres del bucket y del repositorio por los que creaste y, luego, ejecuta la secuencia de comandos:

    bash scripts/setup-and-configure-resources.sh
    
  2. Compila y envía la imagen pytorch-megatron:23.11-py3 a tu repositorio. Asegúrate de que el nombre del repositorio de Docker en el archivo scripts/build-and-push-docker-image.sh coincida con el nombre del repositorio que usaste en la secuencia de comandos scripts/setup-and-configure-resources.sh. También puedes editar el nombre de la etiqueta de imagen de Docker antes de enviarla.

    bash scripts/build-and-push-docker-image.sh
    

Lanza la comparativa de Megatron-LM Llama2

  1. Edita el archivo helm/values.yaml para especificar el bucket de Cloud Storage y la imagen de Docker creada en las secciones anteriores. Para ver algunas configuraciones de ejemplo, consulta configuración de muestra.

  2. Opcional: También puedes editar el archivo selected-configuration.sh para especificar los cambios que realizaste en la configuración predeterminada de Helm.

    helm install HELM_EXPERIMENT_NAME helm/ --values helm/values.yaml
    

    Reemplaza HELM_EXPERIMENT_NAME por un nombre arbitrario para tu experimento.

El experimento escribe métricas de la herramienta de creación de perfiles de Nsight Systems en el bucket de Cloud Storage especificado en el directorio megatron-experiments.

Limpia

Sigue estos pasos para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos que usaste en esta página.

Borra el clúster de GKE:

Ve a la página Clústeres:

Ir a los clústeres

  1. Selecciona la casilla de verificación correspondiente a CLUSTER_NAME.
  2. Haz clic en Borrar.
  3. Para confirmar la eliminación, escribe CLUSTER_NAME y haz clic en Borrar.

Borra el bucket de Cloud Storage

Ve a la página Buckets:

Ir a Buckets

  1. Selecciona la casilla de verificación del bucket de Cloud Storage que creaste para esta guía de inicio rápido.

  2. Haz clic en Borrar.

  3. Para confirmar la eliminación, escribe DELETE y haz clic en Borrar.

¿Qué sigue?