Treinar em uma TPU de host único usando Pax


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. Foco da Pax sobre como simplificar o ML em grande escala compartilhando componentes de infraestrutura com frameworks de ML atuais e utilizar a Biblioteca de modelos Praxis para modularidade.

Objetivos

  • Configurar recursos da 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 CLI do Google Cloud contém ferramentas e bibliotecas para interação com os produtos e serviços da Google Cloud CLI. Se você não instalou instale-a agora seguindo as instruções em Como instalar a Google Cloud CLI.

Configurar o comando gcloud

Execute gcloud auth list para ver as 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 seguinte comando gcloud 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 da TPU. Conecte-se via SSH diretamente à VM da TPU. É possível executar código arbitrário, instalar pacotes, conferir registros e depurar código diretamente na VM da TPU.

Crie sua VM da TPU executando o seguinte comando em um o Cloud Shell ou o terminal do computador 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 versões de TPU v2 e v3 ou tpu-vm-v4-base para TPUs v4.

$ gcloud compute tpus tpu-vm create tpu-name \
--zone zone \
--accelerator-type accelerator-type \
--version version

Conecte-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 De username@projectname a 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 reconhece 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 da TPU

Agora você pode executar qualquer código Pax que quiser. O comando lm_cloud exemplos são um ótimo lugar para começar a executar modelos no Pax. Por exemplo, o seguintes comandos treinam Modelo de linguagem SPMD com base em transformador de parâmetros 2B em dados sintéticos.

Os comandos a seguir mostram a saída de treinamento para um modelo de linguagem SPMD. Ele treina 300 passos 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 passos

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: