Inferência do PyTorch do JetStream em VMs TPU v6e

Este tutorial mostra como usar o JetStream para disponibilizar modelos do PyTorch na TPU v6e. O JetStream é um mecanismo otimizado para capacidade de processamento e memória para inferência de modelos de linguagem grandes (LLMs) em dispositivos XLA (TPUs). Neste tutorial, você vai executar o benchmark de inferência para o modelo Llama2-7B.

Antes de começar

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

  1. Siga o guia Configurar o ambiente do Cloud TPU para configurar um projeto Google Cloud , configurar o CLI do Google Cloud, ativar a API Cloud TPU e garantir que você tenha acesso para usar o Cloud TPU.

  2. Faça a autenticação com Google Cloud e configure o projeto e a zona padrão para a Google Cloud CLI.

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

Capacidade segura

Quando estiver tudo pronto para garantir a capacidade da TPU, consulte Cotas da Cloud TPU para mais informações. Se você tiver outras dúvidas sobre como garantir a capacidade, entre em contato com a equipe de vendas ou de conta do Cloud TPU.

Provisionar o ambiente do Cloud TPU

É possível provisionar VMs do TPU com o GKE, com o GKE e o XPK ou como recursos em fila.

Pré-requisitos

  • Verifique se o projeto tem cota de TPUS_PER_TPU_FAMILY suficiente, que especifica o número máximo de chips que você pode acessar no projetoGoogle Cloud .
  • Verifique se o projeto tem cota suficiente de TPU para:
    • Cota da VM de TPU
    • Cota de endereços IP
    • Quota do Hyperdisk equilibrado
  • Permissões do projeto do usuário

Criar variáveis de ambiente

No 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 de sinalizações de comando

Variável Descrição
PROJECT_ID Google Cloud nome do projeto. Use um projeto existente ou crie um novo.
TPU_NAME O nome da TPU.
ZONE Consulte o documento Regiões e zonas de TPU para saber quais zonas são compatíveis.
ACCELERATOR_TYPE O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.
RUNTIME_VERSION A versão do software do Cloud TPU.
SERVICE_ACCOUNT O endereço de e-mail da sua conta de serviço . Para encontrá-lo, acesse a página Contas de serviço no console do Google Cloud.

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

QUEUED_RESOURCE_ID O ID de texto atribuído pelo usuário da solicitação de recurso em fila.

Provisionar um 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 status do recurso em fila.

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

Para uma lista completa de status de solicitações de recursos em fila, consulte a documentação de Recursos em fila.

Conectar-se à TPU usando SSH

  gcloud compute tpus tpu-vm ssh ${TPU_NAME}

Executar a comparação PyTorch Llama2-7B do JetStream

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

Quando o comparativo de mercado de inferência for concluído, limpe os recursos da TPU.

Limpar

Exclua a TPU:

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