Treinar o Llama2 com o Megatron-LM nas máquinas virtuais A3 Mega

Visão geral

Neste guia de início rápido, você vai aprender a executar uma carga de trabalho PyTorch Megatron-LM baseado em contêiner no A3 Mega. O código está disponível neste repositório do GitHub: megatron-gke.

Antes de começar

Siga estas etapas para ativar a API 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.

  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.

      Acessar o IAM
    2. Selecione o projeto.
    3. Clique em Conceder acesso.

    4. No campo Novos principais, digite seu identificador de usuário. Normalmente, é o endereço de e-mail de uma Conta do Google.

    5. Na lista Selecionar papel, escolha um.
    6. Para conceder outros papéis, adicione-os clicando em Adicionar outro papel.
    7. Clique em Salvar.

      8. Criar um cluster do A3 Mega

      Criar um cluster A3 Mega GKE com GPUDirect-TCPXO e várias redes. Para mais informações, consulte Maximizar a largura de banda da rede GPU com GPUDirect e várias redes.

      Configurar o ambiente

      1. Criar variáveis de ambiente para alguns parâmetros comuns

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

        Substitua:

        • CLUSTER_NAME: o nome do cluster do A3 Mega GKE com GPUDirect-TCPXO e várias redes ativadas.
        • CONTROL_PLANE_LOCATION: o local do Compute Engine do plano de controle do cluster. Forneça uma região para clusters regionais ou uma zona para clusters zonais.
        • PROJECT_ID: o ID do projeto Google Cloud .

      2. Configure a Google Cloud CLI para usar suas credenciais do Google Cloud para autenticação:

        gcloud auth login

        Para mais informações, consulte Autenticar para usar a CLI do Google Cloud.

      3. Instale kubectl e o plug-in gcloud CLI do GKE:

        sudo apt-get install kubectl
sudo apt-get install google-cloud-sdk-gke-gcloud-auth-plugin

      4. Busque credenciais para seu cluster do GKE:

        gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --location=${CONTROL_PLANE_LOCATION} \
  --project=${PROJECT_ID}

      5. Instale o Helm se ainda não tiver feito isso:

        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 o programador com reconhecimento de topologia para implantar pods

      É possível usar o programador com reconhecimento de topologia para implantar pods do GKE em nós com uma topologia de GPU especificada.

      Nos comandos kubectl a seguir, você usará os arquivos diretamente de um repositório de dados. Como alternativa, é possível clonar o repositório localmente e os comandos kubectl podem fazer referência aos arquivos locais.

      Para mais informações, consulte Programador de topologia.

      1. Configurar a conta de serviço:

        kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/service-account.yaml

      2. Instale os scripts do programador de topologia em um 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. Instale o daemonset do rótulo de topologia e o pod do programador de topologia:

        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. Observe as ações do programador de topologia:

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

      Executar a carga de trabalho

      Criar o Dockerfile e enviar para o Google Cloud Artifact Registry

      1. Crie um bucket do Cloud Storage e um repositório do Docker. No scripts/setup-and-configure-resources.sh script, substitua os nomes do bucket e do repositório criados por você e, em seguida, execute o script:

        bash scripts/setup-and-configure-resources.sh

      2. Crie e envie a imagem pytorch-megatron:23.11-py3 para seu repositório. Verifique se o nome do repositório do Docker no arquivo scripts/build-and-push-docker-image.sh corresponde ao nome do repositório usado no script scripts/setup-and-configure-resources.sh. Também é possível editar o nome da tag de imagem Docker antes do envio.

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

      Iniciar o comparativo de mercado Megatron-LM Llama2

      1. Edite o arquivo helm/values.yaml para especificar o bucket do Cloud Storage e a imagem Docker criada nas seções anteriores. Para alguns exemplos de configurações, consulte sample-configurations.

      2. Opcional: também é possível editar o arquivo selected-configuration.sh para especificar quaisquer alterações feitas na configuração padrão do Helm.

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

        Substitua HELM_EXPERIMENT_NAME por um nome arbitrário para o experimento.

      O experimento grava métricas da ferramenta de criação de perfil da Nsight Systems no bucket do Cloud Storage especificado no diretório megatron-experiments.

      Limpar

      Para evitar cobranças na conta do Google Cloud pelos recursos usados nesta página, siga as etapas abaixo.

      Exclua o cluster do GKE:

      Acessar a página de clusters:

      Acessar Clusters

      1. Marque a caixa de seleção de CLUSTER_NAME.
      2. Clique em Excluir.
      3. Para confirmar a exclusão, digite CLUSTER_NAME e clique em Excluir.

      Excluir o bucket do Cloud Storage

      Acessar a página Buckets:

      Acessar buckets

      1. Marque a caixa de seleção do bucket do Cloud Storage que você criou para este guia de início rápido.

      2. Clique em Excluir.

      3. Para confirmar a exclusão, digite DELETE e clique em Excluir.

      A seguir