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

Visão geral

As VMs de TPU são instâncias de VMs do Compute Engine com hardware de TPU conectado. As VMs do Compute Engine estão sujeitas a eventos de manutenção de VMs do Compute Engine. Cada TPU está 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.

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 informações sobre como detectar e se recuperar de falhas de recuperação automática, consulte Detectar e recuperar de falhas de TPU.

Os mecanismos usados para salvar e carregar os checkpoints são diferentes para cada framework de ML. Os modelos do Cloud TPU com suporte geralmente têm checkpoint integrado. 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 ocorre um evento de manutenção na TPU usando o seguinte comando gcloud describe:

$ gcloud compute tpus tpu-vm 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 precisará ser semelhante a esta:

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

Logs de eventos de manutenção

Acesse os registros de auditoria de eventos do sistema para ver os registros históricos de eventos de manutenção na TPU.

No menu de navegação do console do Google Cloud, clique em Compute Engine > Instâncias de VM e pesquise, por exemplo:

"tpu.nodes.terminate" OR "tpu.nodes.restart"

No período de pesquisa, todas as interrupções e reparos dos workers da TPU são exibidos. Os registros mostram a data e a hora do evento, o tipo de evento e, para eventos de "encerramento", o motivo do encerramento em protoPayload.metadata.terminateReason.

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 ocorre um "evento de interrupção", o treinamento é retomado do último checkpoint.

  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 Production precisam usar um sistema de gerenciamento de recursos, como o Google Kubernetes Engine (GKE). Para mais informações sobre como usar o Google Kubernetes Engine com VMs de TPU, consulte Implantar cargas de trabalho de TPU. Você pode implementar um script bash para repetir continuamente o script de treinamento até a conclusão. Exemplo:

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

  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 e excluir e recriar a TPU. Veja um exemplo desse script neste link. Consulte Como gerenciar TPUs para ver detalhes sobre como excluir e recriar TPUs manualmente.

    Ao criar ou recriar uma VM da 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 da TPU é criada. Consulte Executar scripts de instalação padrão para mais informações.