Este documento apresenta uma breve introdução sobre como trabalhar com o Pax em uma TPU de host único (v2-8, v3-8, v4-8).
O Pax é um framework para configurar e executar experimentos de aprendizado de máquina no JAX. O foco do Pax é simplificar o ML em escala, compartilhando componentes de infraestrutura com frameworks de ML existentes e utilizando a biblioteca de modelagem Praxis para modularidade.
Objetivos
- Configurar recursos de TPU para treinamento
- Instalar o Pax em uma TPU de host único
- Treinar um modelo SPMD baseado em transformador usando o Pax
Antes de começar
Execute os comandos a seguir para configurar gcloud para usar
seu projeto da Cloud TPU e instalar os componentes necessários para treinar
um modelo que executa o Pax em uma TPU de host único.
Instalar a CLI do Google Cloud
A Google Cloud CLI contém ferramentas e bibliotecas para interagir com os produtos e serviços da Google Cloud CLI. Se você ainda não instalou a CLI, faça isso agora usando as instruções em Como instalar a Google Cloud CLI.
Configurar o comando gcloud
Execute gcloud auth list para conferir suas contas disponíveis.
$ gcloud config set account account
$ gcloud config set project project-id
Ativar a API Cloud TPU
Ative a API Cloud TPU usando o comando gcloud abaixo
no Cloud Shell.
Também é possível ativá-la no
Console do Google Cloud.
$ gcloud services enable tpu.googleapis.com
Execute o comando a seguir para criar uma identidade de serviço (uma conta de serviço).
$ gcloud beta services identity create --service tpu.googleapis.com
Criar uma VM de TPU
Com as VMs do Cloud TPU, seu modelo e código são executados diretamente na VM do TPU. Você faz o SSH diretamente na VM da TPU. É possível executar código arbitrário, instalar pacotes, conferir registros e depurar código diretamente na VM da TPU.
Crie a VM de TPU executando o comando a seguir em um Cloud Shell ou no terminal do computador em que a Google Cloud CLI está instalada.
Defina o zone com base na disponibilidade no seu contrato. Consulte Regiões e zonas do Cloud TPU, se necessário.
Defina a variável accelerator-type como v2-8, v3-8 ou v4-8.
Defina a variável version como tpu-vm-base
para as versões v2 e v3 da TPU ou tpu-vm-v4-base para as TPUs v4.
$ gcloud compute tpus tpu-vm create tpu-name \ --zone zone \ --accelerator-type accelerator-type \ --version version
Conectar-se à VM do Google Cloud TPU
Use o comando a seguir para se conectar à VM da TPU:
$ gcloud compute tpus tpu-vm ssh tpu-name --zone zone
Quando você faz login na VM, o prompt do shell muda de
username@projectname para username@vm-name:
Instalar o Pax na VM do Google Cloud TPU
Instale o Pax, o JAX e o libtpu na VM da TPU usando os seguintes comandos:
(vm)$ python3 -m pip install -U pip \ python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Verificação do sistema
Para verificar se tudo está instalado corretamente, verifique se o JAX encontra os núcleos do TPU:
(vm)$ python3 -c "import jax; print(jax.device_count())"
O número de núcleos de TPU será exibido. Ele deve ser 8 se você estiver usando uma v2-8 ou v3-8 ou 4 se estiver usando uma v4-8.
Como executar o código Pax em uma VM de TPU
Agora é possível executar qualquer código do Pax que você quiser. Os exemplos de lm_cloud são um ótimo lugar para começar a executar modelos no Pax. Por exemplo, os comandos a seguir treinam um modelo de linguagem SPMD com base no transformador parâmetro 2B com dados sintéticos.
Os comandos a seguir mostram a saída de treinamento para um modelo de linguagem SPMD. Ele é treinado por 300 etapas em aproximadamente 20 minutos.
(vm)$ python3 .local/lib/python3.10/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=job_log_dir
Na fatia v4-8, a saída precisa incluir:
Perdas e tempos de etapa
Tensor de resumo em step=step_# loss = loss
Tensor de resumo em step=step_# Steps per second x
Limpar
Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.
Quando terminar de usar a VM de TPU, siga estas etapas para limpar os recursos.
Desconecte-se da instância do Compute Engine, caso ainda não tenha feito isso:
(vm)$ exit
Exclua o Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name --zone zone
A seguir
Para mais informações sobre o Cloud TPU, consulte: