Gerenciar eventos de manutenção com pods do Cloud TPU

Visão geral

Os nós e VMs de TPU são instâncias de VMs do Compute Engine com hardware de TPU anexado. As VMs do Compute Engine estão sujeitas a eventos de manutenção de VMs do Compute Engine. Cada TPU é conectada a uma VM do Compute Engine. Portanto, o uso de mais TPUs (por exemplo, em um pod de TPU) aumenta a probabilidade de uma das VMs encontrar um evento de manutenção.

Observação: as diferenças entre VMs e nós de TPU estão descritas em Arquitetura do sistema.

Neste documento, discutimos várias abordagens para processar eventos de manutenção de jobs de treinamento de longa duração nos Cloud TPUs.

Como usar checkpoints para recuperação rápida de eventos de manutenção

Checkpoints são fundamentais para recuperações curtas de eventos de manutenção e precisam ser salvos com frequência: uma boa regra é salvar checkpoints aproximadamente a cada hora. Não criar checkpoints com frequência suficiente corre o risco de perder muito progresso do treinamento devido a eventos de manutenção ou outras interrupções no treinamento.

Os checkpoints geralmente se referem a todos os parâmetros salvos usados no treinamento, como pesos do modelo. O tempo necessário para salvar um checkpoint pode variar da ordem de segundos à ordem de minutos.

Embora a maioria dos eventos de manutenção seja recuperada automaticamente e os jobs de treinamento continuem sem intervenção manual, pode haver casos extremos em que o job não é reiniciado e continua automaticamente. Quando isso acontece, é necessário excluir e recriar os recursos da TPU e reiniciar o job de treinamento a partir de um checkpoint salvo. Para mais informações sobre como detectar e recuperar a partir de falhas de recuperação automática, consulte Detectar e recuperar de falhas de TPU abaixo.

Os mecanismos usados para salvar e carregar checkpoints são diferentes para cada framework de ML. Os modelos compatíveis do Cloud TPU geralmente têm checkpoints integrados. Para mais informações sobre checkpoints, consulte: TensorFlow 2.x, PyTorch ou JAX/flax.

Como detectar eventos de manutenção

É possível detectar se e quando um evento de manutenção ocorre na sua TPU usando o seguinte comando gcloud describe:

VMs de TPU

$ gcloud compute tpus tpu-vm describe tpu-name --zone=zone  | grep 'health'

Nós da TPU

$ gcloud compute tpus describe tpu-name --zone=zone | grep 'health'

A saída desse comando exibe o estado atual da TPU e uma descrição do evento de manutenção mais recente. A saída será semelhante a esta:

health: HEALTHY
healthDescription: The TPU had a maintenance event at 2022-01-26T03:44:36.265703305Z

Gerenciar eventos de manutenção

Há várias maneiras de reduzir interrupções em eventos de manutenção.

1. Salvar os checkpoints periodicamente

   No cenário ideal, quando um "evento de interrupção" acontece, o treinamento é retomado a partir do checkpoint mais recente.

2. Novas tentativas de script de treinamento

   O script de treinamento pode parar como resultado de um "evento de interrupção". É possível usar um script bash para repetir continuamente o script de treinamento até que ele seja concluído. Cada nova tentativa precisa continuar a partir do checkpoint mais recente. Portanto, os scripts de nova tentativa precisam ser sempre usados em conjunto com checkpoints.

   Os pipelines de treinamento prontos para produção precisam usar um sistema de gerenciamento de recursos, como o Google Kubernetes Engine (GKE). Atualmente, o GKE é compatível com nós de TPU, mas não com VMs de TPU. Caso contrário, será possível implementar um script bash para repetir continuamente o script de treinamento até a conclusão. Exemplo:

   Com o nó da TPU:

   (a partir da VM) while ! python3 [training command]; do sleep 1; done

   Com a VM da TPU:

   while ! gcloud compute tpus tpu-vm ssh ${TPU_NAME} --command "python3 [training command]"; do sleep 1; done
   

   Observe que você precisa executar o comando da VM da TPU em um shell do Google Cloud ou em um terminal, não na VM da TPU.

3. Detectar e recuperar de falhas da TPU

   Quando uma TPU não se recupera de um evento de manutenção, é possível usar um script de recuperação para detectar o estado da TPU, além de excluir e recriar a TPU. Um exemplo desse script pode ser encontrado aqui. Consulte Como gerenciar TPUs para mais detalhes sobre como excluir e recriar manualmente as TPUs.

   Ao criar ou recriar uma VM de TPU, é possível especificar um script de inicialização com o parâmetro --metadata startup-script. Um script de inicialização é executado sempre que uma VM de TPU é criada. Consulte Executar scripts de instalação padrão para mais informações.