Esta página foi traduzida pela API Cloud Translation.

Preservar o progresso do treinamento usando o Autocheckpoint

Historicamente, quando uma VM TPU requer manutenção, o procedimento é iniciado imediatamente, sem deixar tempo para que os usuários realizem ações de preservação do progresso, como salvar um ponto de verificação. Isso é mostrado na Figura 1(a).

Diagrama mostrando o impacto da manutenção do host com e sem verificação automática

Figura 1. Ilustração do recurso de checkpoint automático: (a) Sem o checkpoint automático, o progresso do treinamento do último checkpoint é perdido quando há um evento de manutenção. (b) Com o ponto de verificação automático, o progresso do treinamento desde o último ponto de verificação pode ser preservado quando há um evento de manutenção.

Você pode usar o Autocheckpoint (Figura 1(b)) para preservar o progresso do treinamento, configurando o código para salvar um ponto de controle não programado quando um evento de manutenção ocorrer. Quando um evento de manutenção ocorre, o progresso desde o último ponto de verificação é salvo automaticamente. O recurso funciona em fatias únicas e em fatias múltiplas.

O recurso de checkpoint automático funciona com frameworks que podem capturar sinais SIGTERM e, em seguida, salvar um checkpoint. Os frameworks compatíveis incluem:

MaxText,
Pax,
JAX com Orbax.

Como usar o checkpoint automático

O recurso de verificação automática fica desativado por padrão. Ao criar um TPU ou solicitar um recurso enfileirado, ative o ponto de verificação automático adicionando a flag --autocheckpoint-enabled ao provisionar o TPU. Com o recurso ativado, o Cloud TPU executa as etapas abaixo quando recebe a notificação de um evento de manutenção:

Capturar o sinal SIGTERM enviado para o processo usando o dispositivo TPU
Aguarde até que o processo seja encerrado ou 5 minutos tenham se passado, o que ocorrer primeiro.
Realizar a manutenção das fatias afetadas

A infraestrutura usada pelo Autocheckpoint é independente do framework de ML. Qualquer framework de ML pode oferecer suporte ao Autocheckpoint se puder capturar o sinal SIGTERM e iniciar um processo de checkpoint.

No código do aplicativo, é necessário ativar os recursos de verificação automática fornecidos pelo framework de ML. No Pax, por exemplo, isso significa ativar as flags de linha de comando ao iniciar o treinamento. Para mais informações, consulte o Guia de início rápido do Checkpoint automático com Pax. Nos bastidores, os frameworks salvam um ponto de verificação não programado quando um sinal SIGTERM é recebido, e a VM de TPU afetada passa por manutenção quando a TPU não está mais em uso.

Guia de início rápido: checkpoint automático com MaxText

O MaxText é um LLM de alto desempenho, escalonável de forma arbitrária, de código aberto e bem testado, escrito em Python/JAX puro para Cloud TPUs. O MaxText contém toda a configuração necessária para usar o recurso de verificação automática.

O arquivo README MaxText descreve duas maneiras de executar o MaxText em escala:

Como usar multihost_runner.py, recomendado para experimentos
Como usar multihost_job.py, recomendado para produção

Ao usar multihost_runner.py, ative o ponto de controle automático definindo a flag autocheckpoint-enabled ao provisionar o recurso na fila.

Ao usar multihost_job.py, ative o ponto de verificação automático especificando a flag de linha de comando ENABLE_AUTOCHECKPOINT=true ao iniciar o job.

Guia de início rápido: verificação automática com Pax em uma única fatia

Esta seção mostra um exemplo de como configurar e usar o Autocheckpoint com Pax em uma única fatia. Com a configuração adequada:

Um ponto de controle será salvo quando um evento de manutenção ocorrer.
A Cloud TPU vai realizar a manutenção nas VMs afetadas depois que o ponto de verificação for salvo.
Quando a manutenção da Cloud TPU for concluída, você poderá usar a VM TPU normalmente.

Use a flag autocheckpoint-enabled ao criar a VM TPU ou solicitar um recurso na fila.

Exemplo:

Defina as variáveis de ambiente:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=zone-you-want-to-use
export ACCELERATOR_TYPE=your-accelerator-type
export RUNTIME_VERSION=tpu-ubuntu2204-base

Descrições das variáveis de ambiente

Variável	Descrição
`PROJECT_ID`	O ID do seu Google Cloud projeto. Use um projeto existente ou crie um novo.
`TPU_NAME`	O nome da TPU.
`ZONE`	A zona em que a VM TPU será criada. Para mais informações sobre as zonas com suporte, consulte Regiões e zonas de TPU.
`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.

Defina o ID e a zona do projeto na configuração ativa:

gcloud config set project $PROJECT_ID
gcloud config set compute/zone $ZONE

Crie uma TPU:

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --accelerator-type $ACCELERATOR_TYPE \
    --version $RUNTIME_VERSION \
    --autocheckpoint-enabled

Conecte-se à TPU usando SSH:

gcloud compute tpus tpu-vm ssh $TPU_NAME

Instalar o Pax em uma única fatia

O recurso de verificação automática funciona nas versões 1.1.0 e mais recentes do Pax. Na VM da TPU, instale o jax[tpu] e o paxml mais recente:
```
pip install paxml && pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```

Configure o modelo LmCloudSpmd2B. Antes de executar o script de treinamento, mude ICI_MESH_SHAPE para [1, 8, 1]:

@experiment_registry.register
class LmCloudSpmd2B(LmCloudSpmd):

    """SPMD model with 2B params.

    Global batch size = 2 * 2 * 1 * 32 = 128
    """
    PERCORE_BATCH_SIZE = 8

    NUM_LAYERS = 18
    MODEL_DIMS = 3072
    HIDDEN_DIMS = MODEL_DIMS * 4

    CHECKPOINT_POLICY = layers.AutodiffCheckpointType.SAVE_NOTHING
    ICI_MESH_SHAPE = [1, 8, 1]

Inicie o treinamento com a configuração adequada.

O exemplo a seguir mostra como configurar o modelo LmCloudSpmd2B para salvar pontos de verificação acionados pelo Autocheckpoint em um bucket do Cloud Storage. Substitua your-storage-bucket pelo nome de um bucket existente ou crie um novo bucket.
```
export JOB_LOG_DIR=gs://your-storage-bucket

{ python3 .local/lib/python3.10/site-packages/paxml/main.py \
    --jax_fully_async_checkpoint=1 \
    --exit_after_ondemand_checkpoint=1 \
    --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2B \
    --job_log_dir=$JOB_LOG_DIR; } 2>&1 | tee pax_logs.txt
```
Observe as duas flags transmitidas para o comando:
- jax_fully_async_checkpoint: com essa flag ativada, orbax.checkpoint.AsyncCheckpointer será usado. A classe AsyncCheckpointer salva automaticamente um ponto de controle quando o script de treinamento recebe um sinal SIGTERM.
- exit_after_ondemand_checkpoint: com essa flag ativada, o processo da TPU é encerrado depois que o checkpoint automático é salvo, o que aciona a manutenção para ser realizada imediatamente. Se você não usar essa flag, o treinamento vai continuar depois que o checkpoint for salvo e o Cloud TPU vai esperar que um tempo limite ocorra (5 minutos) antes de realizar a manutenção necessária.

Ponto de verificação automático com Orbax

O recurso de verificação automática não é limitado a MaxText ou Pax. Qualquer framework que possa capturar o sinal SIGTERM e iniciar um processo de verificação funciona com a infraestrutura fornecida pelo Autocheckpoint. O Orbax, um namespace que oferece bibliotecas de utilitários comuns para usuários do JAX, oferece esses recursos.

Conforme explicado na documentação do Orbex, esses recursos são ativados por padrão para os usuários do orbax.checkpoint.CheckpointManager. O método save que é chamado após cada etapa verifica automaticamente se um evento de manutenção está iminente e, se for o caso, salva um ponto de controle, mesmo que o número da etapa não seja um múltiplo de save_interval_steps. A documentação do GitHub (link em inglês) também ilustra como fazer com que o treinamento saia após salvar um Autocheckpoint, com uma modificação no código do usuário.