Resolução de problemas do JAX – TPU

Este guia fornece indicações para informações de resolução de problemas do JAX para ajudar a identificar e resolver problemas que possa encontrar durante a preparação de modelos JAX na Cloud TPU.

Para aceder a um guia mais geral de introdução ao Cloud TPU, consulte o início rápido do JAX.

Problemas gerais do JAX

Se tiver problemas ao desenvolver o seu modelo de preparação ou preparar com o JAX, consulte as Perguntas frequentes sobre o JAX.

Para ver mais erros de programação gerais que pode encontrar ao escrever uma aplicação de preparação com o JAX, consulte Erros do JAX.

Crie perfis do desempenho do JAX

Pode compreender como os seus recursos de TPU estão a ser usados através das ferramentas descritas em Criar perfis do desempenho do JAX.

Resolva problemas de memória

Pode monitorizar a forma como a memória é usada com o profiler de memória do dispositivo JAX, mas não pode gerir diretamente a forma como é usada.

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

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

Resolva problemas com a TPU

As secções seguintes descrevem como resolver alguns problemas comuns que pode encontrar quando executa um programa JAX numa TPU.

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

Tudo é executado na TPU, desde que o JAX não imprima "No GPU/TPU found, falling back to CPU." (Nenhuma GPU/TPU encontrada. A usar a CPU).

Pode verificar se a TPU está ativa consultando jax.devices(), onde deve ver vários dispositivos TPU apresentados, ou verificar programaticamente com: assert jax.devices()[0].platform == 'tpu'.

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

Esta mensagem de erro de tempo de execução ou a deteção 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 está a executar a versão errada da VM de TPU.

Verifique se está a executar a versão atual do tempo de execução do JAX e tente novamente.

Resolva problemas do TPU e do GKE

Para ajudar na resolução de problemas, ative o registo detalhado no manifesto da carga de trabalho do GKE e, em seguida, faculte os registos ao apoio técnico do GKE.

TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0

As secções seguintes descrevem mensagens de erro relacionadas com configurações e resoluções de TPU e GKE.

no endpoints available for service 'jobset-webhook-service'

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

Falha na inicialização da TPU: não foi possível estabelecer ligação

Certifique-se de que a versão do nó do GKE é 1.30.4-gke.1348000 ou posterior (o GKE 1.31 não é suportado).