Solução de problemas do JAX – TPU

Neste guia, você recebe orientações para a solução de problemas do JAX de modo a identificar e resolver problemas que podem ocorrer ao treinar modelos do JAX no Cloud TPU.

Para um guia mais geral sobre como usar o Cloud TPU, consulte o Guia de início rápido do JAX.

Observação: se você não conseguir resolver o problema usando esse guia, consulte Como receber suporte.

Problemas gerais do JAX

Se você tiver problemas durante o desenvolvimento do modelo de treinamento ou durante o treinamento com o JAX, consulte as Perguntas frequentes sobre o JAX.

Para erros de programação mais gerais que podem acontecer ao escrever um aplicativo de treinamento com o JAX, consulte Erros do JAX.

Criar perfis de desempenho do JAX

Entenda como os recursos de TPU estão sendo usados com as ferramentas descritas em Como criar perfis de desempenho do JAX.

Resolver problemas de memória

Observação: a memória da TPU (memória de alta largura de banda) é compartilhada entre os TensorCores em cada chip de TPU. Isso permite uma coordenação eficiente entre os TensorCores no chip.

É possível monitorar como a memória é usada com o criador de perfil de memória de dispositivos JAX, mas não é possível gerenciar diretamente como ela é usada.

O criador de perfil de memória de dispositivos JAX pode ser usado para o seguinte:

• descobrir quais matrizes e executáveis estão na memória da TPU em um determinado momento; ou
• monitorar vazamentos de memória.

Não é possível especificar como a memória da TPU é alocada para operações específicas. Para mais informações sobre problemas de desempenho da TPU específicos do JAX, consulte Observações sobre o desempenho ao usar TPUs com o JAX.

Resolver problemas de TPU

As seções a seguir descrevem como resolver alguns problemas comuns que podem acontecer ao executar um programa JAX em uma TPU.

Como verificar se a TPU está em execução?

Tudo será executado na TPU, a menos que o JAX exiba: Nenhuma GPU/TPU encontrada; recorrendo à CPU.

Para verificar se a TPU está ativa, consulte jax.devices(), onde serão exibidos vários dispositivos de TPU. Também é possível fazer a verificação programática com assert jax.devices()[0].platform == 'tpu'.

RuntimeError: Não foi possível inicializar a TPU de back-end: INDISPONÍVEL: Não há uma plataforma de TPU disponível.

Essa mensagem de erro de ambiente de execução ou a descoberta do seguinte em /tmp/tpu_logs/tpu_driver.WARNING na VM de TPU: W1118 17:40:20.985243 23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:xxxxxx pode indicar que você está executando a versão errada da VM de TPU.

Verifique se você está executando a versão atual do ambiente de execução do JAX e tente de novo.

Resolver problemas de TPUs e do GKE

Para ajudar na solução de problemas, ative o registro detalhado no manifesto de carga de trabalho do GKE e forneça os registros ao suporte do GKE.

TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0

As seções a seguir descrevem mensagens de erro relacionadas a configurações e resoluções de TPU e do GKE.

Nenhum endpoint disponível para o serviço jobset-webhook-service

Esse erro significa que o JobSet não foi instalado corretamente. Verifique se os pods do Kubernetes de implantação do jobset-controller-manager estão em execução. Para mais informações, consulte a documentação de solução de problemas do JobSet.

Falha na inicialização da TPU: falha na conexão

Verifique se a versão do nó do GKE é 1.30.4-gke.1348000 ou mais recente (o GKE 1.31 não é aceito).