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.

    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.

      Acessar o IAM
    2. Selecionar um projeto.
    3. Clique em CONCEDER ACESSO.
    4. No campo Novos principais, insira seu identificador de usuário. Normalmente, é o endereço de e-mail de uma Conta do Google.

    5. Na lista Selecionar um papel, escolha um.
    6. Para conceder outros papéis, clique em Adicionar outro papel e adicione cada papel adicional.
    7. Clique em Salvar.

    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 REGION=REGION
    export ZONE=ZONE
    export PROJECT_ID=PROJECT_ID
    

    Substitua:

    • CLUSTER_NAME: o nome do cluster do A3 Mega GKE com GPUDirect-TCPXO e várias redes ativadas.
    • REGION: a região em que você criou o cluster.
    • ZONE: a zona em que você criou o cluster.
    • PROJECT_ID: é seu ID do projeto no Google Cloud.
  2. Configurar a CLI do Google Cloud 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} \
      --zone=${ZONE} \
      --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 Artifact Registry do Google Cloud

  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 estas etapas.

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