Entrenar Llama 2 con Megatron-LM en máquinas virtuales A3 Mega

Información general

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

Antes de empezar

Sigue estos pasos 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.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  3. Verify that billing is enabled for your Google Cloud project.

  4. Enable the GKE API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

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

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  6. Verify that billing is enabled for your Google Cloud project.

  7. Enable the GKE API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    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 column 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 Conceder acceso.
    4. En el campo Nuevos principales, introduce tu identificador de usuario. Normalmente, se trata de la dirección de correo de una cuenta de Google.

    5. En la lista Selecciona un rol, elige un rol.
    6. Para conceder más roles, haz clic en Añadir otro rol y añade cada rol adicional.
    7. Haz clic en Guardar.
    8. Crear un clúster A3 Mega

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

      Configurar un entorno

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

        export CLUSTER_NAME=CLUSTER_NAME
        export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
        export PROJECT_ID=PROJECT_ID
        

        Haz los cambios siguientes:

        • CLUSTER_NAME: el nombre de tu clúster de GKE A3 Mega que tiene habilitadas las funciones GPUDirect-TCPXO y de varias redes.
        • CONTROL_PLANE_LOCATION: la ubicación de Compute Engine del plano de control de tu clúster. Proporciona una región para los clústeres regionales o una zona para los clústeres zonales.
        • PROJECT_ID: tu ID de proyecto Google Cloud .
      2. Configura Google Cloud CLI para que use tus Google Cloud credenciales para la autenticación:

        gcloud auth login
        

        Para obtener más información, consulta Autenticarse para usar la CLI de Google Cloud.

      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. Obtén las credenciales de tu clúster de GKE:

        gcloud container clusters get-credentials ${CLUSTER_NAME} \
          --location=${CONTROL_PLANE_LOCATION} \
          --project=${PROJECT_ID}
        
      5. Si aún no lo has hecho, 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
        

      Usar el programador basado en la topología para desplegar tus pods

      Puedes usar el programador compatible con la topología para desplegar tus pods de GKE en nodos que tengan una topología de GPU específica.

      En los siguientes comandos de kubectl, usarás los archivos directamente desde un repositorio. También puedes clonar el repositorio de forma local y los kubectl comandos pueden hacer referencia a los archivos locales.

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

      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 etiquetas 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
        

      Ejecutar la carga de trabajo

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

      1. Crea un segmento de Cloud Storage y un repositorio de Docker. En scripts/setup-and-configure-resources.sh script, sustituye los nombres del bucket y del repositorio por los que has creado y, a continuación, ejecuta la secuencia de comandos:

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

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

      Lanzar la prueba comparativa de Megatron-LM Llama2

      1. Edita el archivo helm/values.yaml para especificar tu segmento de Cloud Storage y la imagen de Docker que has creado en las secciones anteriores. Para ver algunas configuraciones de ejemplo, consulta sample-configurations.

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

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

        Sustituye HELM_EXPERIMENT_NAME por un nombre arbitrario para tu experimento.

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

      Limpieza

      Para evitar que se apliquen cargos en tu cuenta de Google Cloud por los recursos utilizados en esta página, sigue estos pasos.

      Elimina el clúster de GKE:

      Ve a la página Clústeres:

      Ir a Clústeres

      1. Marca la casilla de CLUSTER_NAME.
      2. Haz clic en Eliminar.
      3. Para confirmar la eliminación, escribe CLUSTER_NAME y haz clic en Eliminar.

      Eliminar el segmento de Cloud Storage

      Ve a la página Segmentos:

      Ir a Contenedores

      1. Marca la casilla del segmento de Cloud Storage que has creado para esta guía de inicio rápido.

      2. Haz clic en Eliminar.

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

      Siguientes pasos