Inferência do JetStream PyTorch em VMs da TPU v6e

Este tutorial mostra como usar o JetStream para publicar modelos PyTorch em TPUs v6e. O JetStream é um motor otimizado para débito e memória para inferência de grandes modelos de linguagem (GMLs) em dispositivos XLA (TPUs). Neste tutorial, vai executar o teste de referência de inferência para o modelo Llama2-7B.

Antes de começar

Prepare-se para aprovisionar uma TPU v6e com 4 chips:

  1. Siga o guia Configure o ambiente do Cloud TPU para configurar um Google Cloud projeto, configurar a CLI do Google Cloud, ativar a API Cloud TPU e garantir que tem acesso à utilização de Cloud TPUs.

  2. Autentique-se com o Google Cloud e configure o projeto e a zona predefinidos para a CLI Google Cloud.

    gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE

Capacidade segura

Quando estiver pronto para proteger a capacidade da TPU, consulte as cotas da Cloud TPU para mais informações sobre as cotas da Cloud TPU. Se tiver mais perguntas sobre a proteção da capacidade, contacte a equipa de vendas ou da conta do Cloud TPU.

Aprovisione o ambiente do Cloud TPU

Pode aprovisionar VMs de TPU com o GKE, com o GKE e o XPK ou como recursos em fila.

Pré-requisitos

  • Verifique se o seu projeto tem quota suficiente, que especifica o número máximo de chips aos quais pode aceder no seu projeto.TPUS_PER_TPU_FAMILYGoogle Cloud
  • Verifique se o seu projeto tem quota de TPUs suficiente para:
    • Quota de VMs da TPU
    • Quota de endereços IP
    • Quota de Hyperdisk Balanced
  • Autorizações do projeto do utilizador

Crie variáveis de ambiente

Num Cloud Shell, crie as seguintes variáveis de ambiente:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-central2-b
export ACCELERATOR_TYPE=v6e-4
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

Descrições das variáveis de ambiente

Variável Descrição
PROJECT_ID O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
TPU_NAME O nome da TPU.
ZONE A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
ACCELERATOR_TYPE O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
RUNTIME_VERSION A versão do software do Cloud TPU.
SERVICE_ACCOUNT O endereço de email da sua conta de serviço. Pode encontrá-la acedendo à página Contas de serviço na Google Cloud consola.

Por exemplo: tpu-service-account@PROJECT_ID.iam.gserviceaccount.com

QUEUED_RESOURCE_ID O ID de texto atribuído pelo utilizador do pedido de recurso em fila.

Aprovisione uma TPU v6e

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
        --node-id ${TPU_NAME} \
        --project ${PROJECT_ID} \
        --zone ${ZONE} \
        --accelerator-type ${ACCELERATOR_TYPE} \
        --runtime-version ${RUNTIME_VERSION} \
        --service-account ${SERVICE_ACCOUNT}

Use os comandos list ou describe para consultar o estado do recurso em fila.

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project ${PROJECT_ID} --zone ${ZONE}

Para ver uma lista completa dos estados dos pedidos de recursos em fila, consulte a documentação Recursos em fila.

Estabeleça ligação à TPU através do SSH

  gcloud compute tpus tpu-vm ssh ${TPU_NAME}

Execute o benchmark JetStream PyTorch Llama2-7B

Para configurar o JetStream-PyTorch, converter os pontos de verificação do modelo e executar o teste de referência de inferência, siga as instruções no repositório do GitHub.

Quando o teste de referência de inferência estiver concluído, certifique-se de que limpa os recursos da TPU.

Limpar

Elimine a TPU:

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --force \
      --async