Esta página foi traduzida pela API Cloud Translation.

Inferência do PyTorch do JetStream em VMs TPU v6e

Este tutorial mostra como usar o JetStream para exibir 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:

Faça login na sua Conta do Google. Se ainda não tiver uma, crie uma nova conta.
No console do Google Cloud, selecione ou crie um projeto do Google Cloud na página do seletor de projetos.
Ative o faturamento para o projeto Google Cloud . O faturamento é obrigatório para todo o uso do Google Cloud .
Instale os componentes da gcloud alfa.
Execute o comando a seguir para instalar a versão mais recente dos componentes gcloud.
```
gcloud components update
```
Ative a API TPU usando o comando gcloud a seguir no Cloud Shell. Também é possível ativá-la no console do Google Cloud.
```
gcloud services enable tpu.googleapis.com
```

Crie uma identidade de serviço para a VM da TPU.

gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE

Crie uma conta de serviço de TPU e conceda acesso aos serviços do Google Cloud .

As contas de serviço permitem que o serviço de TPU do Google Cloud acesse outros serviços do Google Cloud. Recomendamos usar uma conta de serviço gerenciado pelo usuário. Siga estes guias para criar e conceder funções. Os seguintes papéis são necessários:
- Administrador de TPU: necessário para criar uma TPU
- Administrador de armazenamento: necessário para acessar o Cloud Storage
- Logs Writer: necessário para gravar registros com a API Logging.
- Gravador de métricas do Monitoring: necessário para gravar métricas no Cloud Monitoring
Faça a autenticação com o 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

Entre em contato com a equipe de vendas ou de contas do Cloud TPU para solicitar a cota de TPU e fazer perguntas sobre a capacidade.

Provisionar o ambiente do Cloud TPU

É possível provisionar TPUs v6e com o GKE, com o GKE e o XPK ou como recursos em fila.

Pré-requisitos

Verifique se o projeto tem cota TPUS_PER_TPU_FAMILY suficiente, que especifica o número máximo de ícones que você pode acessar no projeto Google Cloud .
Este tutorial foi testado com a seguinte configuração:
- Python 3.10 or later
- Versões noturnas do software:
  - 0.4.32.dev20240912 noturno JAX
  - LibTPU noturno 0.1.dev20240912+nightly
- Versões estáveis do software:
  - JAX + JAX Lib de v0.4.35
Verifique se o projeto tem cota suficiente de TPU para:
- Cota da VM de TPU
- Quota de endereço IP
- Cota do Hyperdisk equilibrado
Permissões do projeto do usuário
- Se você estiver usando o GKE com XPK, consulte Permissões do console do Google Cloud na conta de usuário ou de serviço para conferir as permissões necessárias para executar o XPK.

Criar variáveis de ambiente

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

export NODE_ID=TPU_NODE_ID # TPU name
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=v6e-4
export ZONE=us-central2-b
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=YOUR_SERVICE_ACCOUNT
export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID
export VALID_DURATION=VALID_DURATION

# Additional environment variable needed for Multislice:
export NUM_SLICES=NUM_SLICES

# Use a custom network for better performance as well as to avoid having the
# default network becoming overloaded.
export NETWORK_NAME=${PROJECT_ID}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

Descrições de sinalizações de comando

Variável	Descrição
NODE_ID	O ID atribuído pelo usuário do TPU que é criado quando a solicitação de recurso em fila é alocada.
PROJECT_ID	Google Cloud nome do projeto. Use um projeto existente ou crie um novo.
ZONA	Consulte o documento Regiões e zonas de TPU para conferir as zonas compatíveis.
ACCELERATOR_TYPE	Consulte a documentação Tipos de acelerador para saber quais são os tipos de aceleradores compatíveis.
RUNTIME_VERSION	`v2-alpha-tpuv6e`
SERVICE_ACCOUNT	Esse é o endereço de e-mail da sua conta de serviço, que pode ser encontrado em console do Google Cloud -> IAM -> Contas de serviço Por exemplo: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com
NUM_SLICES	O número de fatias a serem criadas (necessário apenas para fatias múltiplas)
QUEUED_RESOURCE_ID	O ID de texto atribuído pelo usuário da solicitação de recurso em fila.
VALID_DURATION	O período em que a solicitação de recurso em fila é válida.
NETWORK_NAME	O nome de uma rede secundária a ser usada.
NETWORK_FW_NAME	O nome de um firewall de rede secundário a ser usado.

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 v6e-4 \
        --runtime-version v2-alpha-tpuv6e \
        --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