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.

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

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

Pontos de verificação 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.

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 ponto de verificação 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.

Gerenciar eventos de manutenção

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

  1. Salvar pontos de verificação 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 deve continuar a partir do ponto de verificação mais recente, portanto, os scripts de nova tentativa devem ser sempre usados em conjunto com checkpoints.

    Os pipelines de treinamento de produção precisam usar o GKE ou outros sistemas de gerenciamento de recursos. 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 VM de TPU: while ! gcloud compute tpus tpu-vm ssh ${TPU_NAME} --command "python3 [training command]"; do sleep 1; done

    Você precisa executar o comando de VM da TPU no GCP ou em um terminal, e 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.