Solução de problemas do JAX: TPU

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

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

Problemas gerais do JAX

Se você tiver problemas durante o desenvolvimento do seu modelo de treinamento ou do treinamento com o JAX, consulte as Perguntas frequentes sobre o JAX (em inglês).

Para erros de programação mais gerais que podem ser encontrados ao escrever um aplicativo de treinamento com o JAX, consulte Erros do JAX (em inglês).

Criar perfis de desempenho do JAX

É possível entender como os recursos de TPU estão sendo usados com as ferramentas descritas na seção Como criar perfis de desempenho do JAX (em inglês).

Resolver problemas de memória

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

O criador de perfil de memória do dispositivo JAX pode ser usado para:

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 ocorrer ao executar um programa JAX em uma TPU.

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

Tudo será executado na TPU, desde que o JAX não exiba "No GPU/TPU found, falling back to CPU" (Nenhuma GPU/TPU encontrada, recorrendo à CPU).

Para verificar se a TPU está ativa, consulte jax.devices(), em que serão exibidos vários dispositivos TPU. Também é possível verificar programaticamente com assert jax.devices()[0].platform == 'tpu'.

RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No TPU Platform available.

Essa mensagem de erro de tempo 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 tempo de execução do JAX e tente de novo.

Resolver problemas de TPU e 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 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 é compatível).