Solução de problemas do TensorFlow – TPU

Este guia, além das perguntas frequentes, ajuda usuários que estão treinando modelos do TensorFlow no Cloud TPU a resolver problemas. Para resolver problemas de treinamento com o PyTorch ou o JAX, consulte os documentos de solução de problemas desses frameworks:

Confira abaixo alguns guias mais gerais sobre como usar o Cloud TPU:

Visão geral

Os problemas comuns que podem acontecer nos Cloud TPUs se enquadram nas seguintes categorias:

  1. Problemas de conexão com a TPU

  2. Como depurar erros comuns

  3. Como reduzir o uso da memória

  4. Como melhorar a velocidade do treinamento

  5. Como depurar quedas na acurácia do modelo

Problema ao se conectar ao servidor da TPU

Esta seção descreve como resolver problemas quando o TensorFlow para de responder ou exibe um erro ao se conectar à TPU. A etapa de compilação do gráfico da TPU pode levar muito tempo nos modelos grandes. Portanto, execute o script por pelo menos cinco minutos antes de confirmar que ele parou de responder.

A primeira etapa é verificar se o problema está no próprio servidor ou no pipeline de treinamento do TensorFlow. Para isso, execute o programa do TensorFlow e verifique se ele funciona corretamente. Se ainda houver problemas de conexão, isso vai confirmar que o problema é no servidor da TPU. Nesse caso:

  1. Execute este comando para listar as TPUs disponíveis: Substitua zone e project-id pela zona e pelo ID do projeto.

    (vm)$ gcloud compute tpus tpu-vm list --zone zone --project project-id

    Isso exibe uma saída como esta: 

    NAME       ZONE           ACCELERATOR_TYPE  NETWORK_ENDPOINT   NETWORK  RANGE          STATUS
TPU_NAME   us-central1-b  v2-8              10.240.1.2:8470    default  10.240.1.0  READY

  2. Verifique se a TPU está listada como READY.

  3. Se a TPU não constar como READY ou você ainda estiver com problemas de conexão, use o seguinte para reiniciar manualmente o servidor:

    (vm)$ gcloud compute tpus tpu-vm stop TPU_NAME && gcloud compute tpus tpu-vm start TPU_NAME

    Esse comando pode levar alguns minutos para ser concluído.

  4. Execute novamente o comando gcloud compute tpus tpu-vm list e espere a TPU entrar no estado READY. Isso pode levar alguns minutos.

  5. Tente executar o programa novamente.

  6. Se você ainda estiver com problemas, peça ajuda usando um dos mecanismos descritos em Como receber suporte.

Se o código for executado corretamente, mas o modelo ainda parar de responder, isso vai indicar que o problema provavelmente está no pipeline de treinamento. Para a depuração, comece substituindo a TPUStrategy no código pela estratégia padrão. Quando você usa a estratégia padrão, sempre que strategy.scope() ou strategy.run() são usados, o modelo é executado na CPU (ou na GPU, se presente) em vez de na TPU. Se o modelo for executado na CPU e não na TPU, isso vai indicar um problema específico da TPU. Se ele ainda não funcionar, a prática recomendada será depurar o problema na CPU.

Perda da conexão ssh durante o treinamento

A conexão ssh com o Cloud TPU pode expirar durante um treinamento de longa duração, principalmente ao usar o Cloud Shell. Aqui, não há saída no console da TPU, e pode parecer que ela parou o treinamento. Para evitar isso, execute a sessão de treinamento com um multiplexador de terminal ou uma ferramenta de gerenciamento de sessão, como tmux ou screen. Isso mantém a conexão ssh ativa, seja qual for a duração do treinamento.

Como depurar erros comuns

Nesta seção, descrevemos como resolver erros comuns que podem acontecer ao treinar modelos no Cloud TPU.

Não é possível criar uma TPU

Ao criar um Cloud TPU, você pode receber o seguinte erro:

googleapiclient.errors.HttpError: < HttpError 403 when requesting https://content-tpu.googleapis.com/v1/projects/{PROJECT}/locations/{ZONE}/nodes/{TPU_NAME}?alt=json returned "Request had insufficient authentication scopes."

Esse é um problema de permissões que pode ser resolvido com o seguinte comando:

gcloud auth login --update-adc

Esse comando atualiza as Application Default Credentials (ADC) e resolve o problema. Para mais informações, consulte gcloud auth login.

Formatos dinâmicos não aceitos

Mensagem de erro

ValueError: shape [Shape] must have a fixed size for dimension
d that is known at graph construction time.

Frameworks e configurações afetados

Essa mensagem só aparece durante a compilação do XLA com o TensorFlow.

Detalhes

Para executar um modelo na TPU, o Cloud TPU o compila usando o compilador XLA. Essa etapa de compilação aumenta bastante a velocidade de treinamento e o uso da memória. No entanto, os formatos ou tamanhos de dimensão de todos os tensores no grafo precisam ser conhecidos no momento da compilação do grafo. Se não for possível determinar formatos nesse momento, a compilação da TPU vai falhar com um erro como o acima.

Uma operação comum que retorna um formato dinâmico é dataset.batch(batch_size), já que o número de amostras restantes em um fluxo pode ser menor que o tamanho do lote. Portanto, ao realizar treinamentos na TPU, defina drop remainder=True como dataset.batch. Isso pode eliminar as últimas amostras de um arquivo para garantir que cada lote tenha o formato estático de batch_size. Por exemplo:

dataset = tf.data.Dataset.range(8)
dataset = dataset.batch(3, drop_remainder=True)

Operação do TensorFlow indisponível

Mensagem de erro

NotFoundError: No registered 'OpName' OpKernel for XLA_TPU_JIT
devices compatible with node

Frameworks e configurações afetados

Essa mensagem pode aparecer durante treinamentos com o TensorFlow.

Detalhes

O modelo usa uma operação do TensorFlow que não está disponível na TPU.

Para conferir uma lista de operações disponíveis na TPU, planos de suporte futuro e sugestões de soluções alternativas, consulte o guia sobre operações disponíveis do TensorFlow.

Mensagem de erro de falta de memória

Mensagem de erro

ResourceExhaustedError: Ran out of memory in memory space hbm; used:
YYY; limit: 7.48G.

Frameworks e configurações afetados

Essa mensagem pode aparecer durante treinamentos com o TensorFlow, o PyTorch ou o JAX.

Detalhes

Cada Cloud TPU é formado por oito núcleos de TPU. As TPUs v2 têm 8 GB e as v3 têm 16 GB de RAM ou de memória de alta largura de banda (HBM). Essa memória é usada para armazenar os tensores de peso (variáveis), bem como os de resultados intermediários necessários para o cálculo do gradiente. Se o modelo for muito grande para caber na RAM da TPU, a inicialização vai falhar e a mensagem de erro acima será exibida. Consulte a seção sobre como reduzir o uso da memória para encontrar mais ajuda.

Dicas para reduzir o uso da memória:

Problemas que interrompem a execução

Quando há um erro no TensorFlow durante a execução da TPU, o script às vezes para de responder em vez de ir para o shell. Se isso acontecer, pressione CTRL+C no teclado para acionar SIGQUIT e encerrar o Python imediatamente.

Da mesma forma, pressionar CTRL+C durante a execução da TPU não encerra o TensorFlow imediatamente. Em vez disso, ele aguarda o fim do loop de iteração atual para ser encerrado corretamente.

Se você receber um novo erro ao se reconectar à TPU depois de sair dessa maneira, redefina manualmente o servidor da TPU com estes comandos: 

gcloud compute tpus tpu-vm stop tpu-name --zone=zone
gcloud compute tpus tpu-vm start tpu-name --zone=zone

em que tpu-name é extraído da primeira coluna exibida pelo comando gcloud compute tpus tpu-vm list e zone é a zona mostrada na segunda coluna.

Preenchimento excessivo de tensor

Possível causa do problema de memória

Os tensores na memória da TPU são preenchidos. Isso significa que a TPU arredonda os tamanhos dos tensores armazenados na memória para executar cálculos com mais eficiência. O preenchimento acontece de modo evidente no nível do hardware e não afeta os resultados. No entanto, em alguns casos, o preenchimento aumenta bastante o uso de memória e o tempo de execução.

Como reduzir o uso de memória

O software da TPU exibe os tensores na memória para aumentar a eficiência computacional e minimizar o preenchimento. Esse processo de exibição da memória é complexo. No entanto, para conseguir os melhores resultados, o modelo precisa cumprir a regra geral a seguir. Para reduzir a sobrecarga de memória e aumentar a eficiência computacional, é necessário que:

  • o tamanho total do lote seja um múltiplo de 64 (8 por núcleo de TPU) e as dimensões do recurso sejam um múltiplo de 128;

    ou

  • o tamanho total do lote seja um múltiplo de 1024 (128 por núcleo de TPU) e as dimensões do recurso sejam um múltiplo de 8.

Usando um tamanho de lote de 1024 e dimensões de recurso que sejam um múltiplo de 128, você garante a melhor eficiência, mesmo que isso não seja possível em todos os modelos. Para maior clareza, a dimensão de recurso se refere ao tamanho oculto de uma camada completamente conectada ou ao número de canais de saída em uma convolução. Nem todas as camadas seguem essa regra, especialmente a primeira e a última da rede. Isso não é um problema, e é esperado que a maioria dos modelos exija uma certa quantidade de preenchimento.

Como reduzir o uso da memória

Se ocorrer um erro de falta de memória ao executar o modelo na TPU, reduza o uso da memória.

As maneiras mais eficazes de reduzir o uso da memória são:

  • Reduzir o preenchimento excessivo de tensores.
  • Reduzir o tamanho dos lotes.

O lote ou o modelo é muito grande

Possível causa do problema de memória

Ao treinar uma rede neural em uma CPU, GPU ou TPU, o uso da memória é proveniente do seguinte:

  1. O uso da memória é proporcional ao número de pesos no modelo.
  2. Armazenamento de ativações intermediárias da passagem direta necessárias para calcular a passagem reversa. O uso da memória é diretamente proporcional aos tamanhos do lote e da camada e ao número de camadas.

Portanto, a memória exigida por um modelo é muito dependente do tamanho do lote.

A memória exigida por um modelo depende do número de camadas na rede.

O ambiente de execução da TPU tenta otimizar os operadores para ajustar o modelo na memória. Isso é chamado de rematerialização e é parecido com a criação de checkpoints de gradiente. No entanto, nem sempre ele é capaz de fazer isso.

Como reduzir o uso de memória

Reduza lentamente o tamanho do lote até que ele caiba na memória, verificando se ele é múltiplo de 64. O tamanho de lote por núcleo precisa ser múltiplo de 8. Lembre-se de que tamanhos de lote maiores são mais eficientes na TPU. O tamanho de lote total de 1024 (128 por núcleo) costuma ser um bom ponto de partida.

Se não for possível executar o modelo na TPU mesmo com um tamanho de lote pequeno, como 64, reduza o número de camadas ou os tamanhos delas.

Como melhorar a velocidade do treinamento

Nesta seção, descrevemos várias maneiras de melhorar a velocidade do treinamento ao executar o modelo na TPU, mas com uma velocidade mais lenta do que o esperado. Consulte o guia de desempenho para conferir outras sugestões sobre como melhorar o desempenho do treinamento.

Poucas etapas por execução por loop de treinamento

Descrição do problema de desempenho

Transmitir o argumento steps_per_execution para Model.compile controla quantas etapas de treinamento são executadas entre os callbacks do host. Cada callback do host exige uma comunicação significativa entre a CPU do host do servidor da TPU e o dispositivo de TPU. Portanto, se steps_per_execution for muito pequeno, o treinamento poderá ficar mais lento.

Como saber se o modelo foi afetado

Se um perfil de TPU revelar callbacks frequentes da CPU do host entre as etapas do dispositivo de TPU, o treinamento poderá se beneficiar de um valor steps_per_execution maior.

Como mitigar

Defina steps_per_execution como um valor maior. Embora steps_per_execution possa ser definido como um valor grande, é importante ter em mente que as mensagens de registro e o salvamento de um checkpoint só podem ocorrer depois que o número especificado de etapas for executado.

Gargalo de processamento de entrada

Descrição do problema de desempenho

Enquanto a TPU realiza o treinamento em um determinado bloco de dados, a função de processamento de entrada prepara o próximo bloco na CPU. Se a função de entrada levar mais tempo que a função de modelo, a TPU vai ficar inativa enquanto a função de entrada recupera os dados.

Como saber se o modelo foi afetado

Siga as instruções em Ferramentas do Cloud TPU: analisador do pipeline de entrada para conferir a respectiva análise no TensorBoard:

imagem

A página de análise do pipeline de entrada exibe um resumo claro que mostra se o modelo tem um gargalo do processamento de entrada. Ela também mostra o tempo de execução por operação para ajudar na identificação de operações problemáticas.

Como mitigar

Há várias mitigações possíveis no carregamento de dados com a API Dataset:

  1. Armazenar os dados como uma coleção de estruturas tf.train.Example em arquivos TFRecord e carregá-los com TFRecordDataset. Consulte o tutorial da API Dataset ou o tutorial do ResNet para conferir exemplos.
  2. Usar dataset.cache() ou dataset.prefetch() para armazenar em buffer os dados de entrada. Isso evita que lentidões esporádicas no acesso aos arquivos criem um gargalo.
  3. Especificar o parâmetro num_parallel_calls da função dataset.map() para ativar operações map() com várias linhas de execução. Uma heurística para o valor de num_parallel_calls é usar o número de núcleos de CPU disponíveis.
  4. O pré-processamento de dados off-line é caro, mas pode ser pago uma única vez para que não seja necessário gerá-lo em todos os períodos de cada treinamento.

Todo o processamento de entrada é executado em CPUs no servidor da TPU, e não na máquina local. Portanto, a velocidade da máquina local não é um fator.

Tempos de etapa lentos e baixa utilização de MXU

Descrição do problema de desempenho

A Cloud TPU realiza convoluções e multiplicações de matrizes a velocidades incrivelmente altas. A maioria das outras operações do TensorFlow conta com implementações eficientes na TPU. No entanto, essa não é a principal qualidade da TPU em relação a outros hardwares. Portanto, é necessário que as convoluções ou multiplicações de matrizes sejam predominantes no modelo para aproveitar ao máximo a TPU.

Como saber se o modelo foi afetado

Os sintomas que você vai notar nesse caso são tempos de etapa lentos e baixa utilização de MXU, que aparecem quando você cria um perfil do desempenho.

Como mitigar

Tente reduzir o número de operações que não são multiplicações de matrizes. Depois de reduzir o número de multiplicações de matriz, faça um novo comparativo para verificar se o desempenho está aceitável nas TPUs.

Preenchimento excessivo de tensor

Descrição do problema de desempenho

A TPU preenche os tensores na memória para poder usar as unidades computacionais com eficiência. O preenchimento pode aumentar o uso da memória e da respectiva largura de banda. Consulte a seção sobre preenchimento de tensor para entender e solucionar os problemas relacionados.

Capacidade de processamento lenta e baixo uso da memória

Descrição do problema de desempenho

Como regra geral, o uso de lotes maiores aumenta a velocidade de treinamento na TPU em termos de amostras/segundo.

Como saber se o modelo foi afetado

O tamanho do lote de qualquer modelo precisa sempre ser pelo menos 64 (8 por núcleo de TPU), já que a TPU preenche os tensores nesse tamanho. O tamanho ideal ao treinar na TPU é de 1024 (128 por núcleo de TPU). Isso elimina as ineficiências relacionadas à transferência de memória e preenchimento.

Como mitigar

A prática recomendada é usar o maior tamanho de lote que caiba na memória e seja múltiplo de 64. Para fazer isso da maneira mais fácil, comece com 1024. Se ocorrer um erro de falta de memória, reduza o tamanho do lote até que o modelo seja executado adequadamente. Alterar o tamanho do lote de um modelo requer o ajuste de outros hiperparâmetros para conseguir a mesma acurácia, como a taxa de aprendizado. No entanto, isso precisa ser avaliado de acordo com o caso.

Tamanhos de camada muito pequenos

Descrição do problema de desempenho

Mesmo quando as convoluções ou multiplicações de matrizes são predominantes no modelo, a TPU pode não funcionar com eficiência total se os tensores de entrada forem pequenos. Em comparação a outros hardwares, a TPU é executada com mais eficiência quando o tamanho do lote e os tamanhos da camada são grandes. Por exemplo, dimensão ≥ 512.

Como saber se o modelo foi afetado

Como regra geral, os tamanhos de camada menores que 128 geram baixa eficiência na TPU. Isso acontece porque 128 é a dimensão integrada da unidade de multiplicação de matrizes da TPU. Nas camadas completamente conectadas, o tamanho mínimo oculto de 512 é recomendado para gerar alta eficiência. Em geral, as camadas convolucionais não precisam ser tão grandes quanto as completamente conectadas para atingir o mesmo nível de eficiência.

Como mitigar

Se os tamanhos pequenos de camada no modelo forem motivados principalmente pela velocidade do treinamento, faça novamente a comparação dos modelos com camadas maiores na TPU. Por exemplo, dobrar o tamanho da saída de uma camada de 256 para 512 só aumenta o tempo de treinamento em 20%, mesmo que o modelo execute o dobro de cálculos.

Perfil de modelo no nível operacional

Em geral, é bom medir o uso de memória e o tempo de execução no nível operacional para identificar gargalos de desempenho. Saiba como fazer isso
no guia Ferramentas do Cloud TPU: visualizador de trace.

Como depurar quedas na acurácia do modelo

Um dos objetivos do ecossistema do Cloud TPU é que qualquer modelo treinado em uma CPU ou GPU atinja uma acurácia muito semelhante ao ser treinado na TPU, com pequenos ajustes em hiperparâmetros como o tamanho do lote e a taxa de aprendizado. No entanto, os usuários às vezes observam uma queda na acurácia ao treinar modelos na TPU. A depuração desses problemas pode ser muito frustrante devido à natureza aleatória do treinamento da rede neural. Nesta seção, aprenda a identificar a causa raiz de qualquer queda na acurácia ao fazer a portabilidade de um modelo para a TPU.

Noções básicas sobre fragmentação de dados (paralelismo)

Um dos principais objetivos do TensorFlow é fazer as operações produzirem resultados quase idênticos, sejam elas executadas em CPU, GPU ou TPU. Há certas exceções, como as operações aleatórias. Em geral, se você notar qualquer diferença significativa entre a saída de operações não aleatórias na TPU e na CPU, registre um bug.

No entanto, para o pipeline de treinamento como um todo, há uma grande diferença entre o treinamento em CPUs/GPUs e em TPUs. Nos treinamentos em uma TPU, o TensorFlow executa a fragmentação de dados. Cada Cloud TPU contém oito núcleos de TPU, que funcionam como unidades de processamento independentes. Em cada etapa do treinamento, todos os núcleos de TPU recebem um lote de dados, calculam os gradientes de peso, trocam esses gradientes com os outros núcleos de TPU e calculam a atualização de peso. Por padrão, a média de perda é calculada com base nos núcleos, mas é possível somá-la alterando o parâmetro de CrossShardOptimizer.

Se for possível calcular a perda total do modelo como a média ou soma de perdas independentes por amostra, esse procedimento será matematicamente similar ao treinamento em um único lote grande.

A operação mais comum que não é independente por amostra é a normalização de lote, que é executada em cada lote por núcleo separadamente. Por exemplo, se o tamanho total do lote for 128, o tamanho por núcleo será 16 e cada um dos oito núcleos vai executar a normalização de lote nas 16 amostras. Em alguns casos, realizar a normalização em lotes pequenos, como aqueles com menos de 32 amostras, reduziu a acurácia. No cenário ideal, o tamanho total do lote deve ser grande. Por exemplo, de 256 a 1.024. No entanto, se esse tamanho for muito grande para caber na memória, será necessário avaliar o efeito da fragmentação de acordo com o caso.

Como depurar o treinamento da TPU de vários núcleos

Se o modelo tem a mesma perda na CPU e na TPU de núcleo único, o problema é provavelmente um dos seguintes:

(a) A degradação é por conta da variância aleatória natural ao treinar modelos neurais com diferentes inicializações.

(b) A degradação é por conta de um problema relacionado à fragmentação de dados na TPU.

Para determinar se (a) é o problema, treine novamente o modelo completo na CPU/GPU e na TPU de vários núcleos usando a mesma inicialização de peso.

Se você tiver certeza de que a queda na acurácia tem significância estatística, estes serão os problemas mais prováveis relacionados à fragmentação de dados:

  1. Se o modelo usa a normalização de lote, um tamanho total de lote menor que 256, como menos de 32 por núcleo, pode reduzir a acurácia.
  2. As funções de perda em lote são afetadas pela fragmentação. Essas funções costumam ser muito especializadas. Por exemplo, Karras e outros autores (2017) usam um discriminador de lotes ao treinar uma rede adversária generativa.

Solução de problemas de configuração da gcloud

Problema
gcloud components update exibe a seguinte mensagem de erro:
ERROR: (gcloud.components.update)
You cannot perform this action because the Cloud SDK component manager is
disabled for this installation.
Solução

Para usar a CLI do Google Cloud, você precisa usar uma instalação que não seja gerenciada por um gerenciador de pacotes.

  1. Execute este comando para remover a instalação atual da gcloud CLI:

    sudo apt-get remove google-cloud-sdk

  2. Siga as instruções para instalar a CLI do Google Cloud.

Problema

O comando gcloud compute tpus tpu-vm ssh TPU_NAME --zone ZONE exibe esta mensagem de erro:

Waiting for SSH key to propagate.
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ERROR: (gcloud.compute.tpus.tpu-vm.ssh) Could not SSH into the instance.  It is possible that your SSH key has not propagated to the instance yet. Try running this command again.  If you still cannot connect, verify that the firewall and instance are set to accept ssh traffic.
Solução

Ocorreu um erro na propagação da chave SSH. Tente mover as chaves geradas automaticamente para um local de backup para forçar gcloud a recriá-las:

mv ~/.ssh/google_compute_engine ~/.ssh/old-google_compute_engine
mv ~/.ssh/google_compute_engine.pub ~/.ssh/old-google_compute_engine.pub

Registros de depuração

Os frameworks JAX, PyTorch e TensorFlow, disponíveis com o Cloud TPU, acessam TPUs usando uma biblioteca compartilhada chamada libtpu, que está presente em todas as VMs de TPU. Essa biblioteca inclui o compilador XLA que é usado para compilar programas de TPU, o ambiente de execução de TPU usado para executar programas compilados e o driver de TPU usado pelo ambiente de execução para acesso de baixo nível à TPU.

A biblioteca libtpu registra informações que podem ser úteis para depuração. Por padrão, esses registros são gravados em /tmp/tpu_logs em cada VM do Cloud TPU. As variáveis de ambiente abaixo podem ser definidas antes de iniciar o treinamento para modificar o comportamento da geração de registros:

TPU_LOG_DIR: o diretório em que os registros são gravados.
O local padrão do diretório é /tmp/tpu_logs. O diretório será criado se ainda não existir, mas nenhum diretório pai será criado. Se ocorrer um erro ao encontrar ou criar o diretório especificado, uma mensagem será exibida em stderr, mas ela não vai interromper o programa e a geração de registros será desativada. Defina o nome do diretório como disabled para desativar completamente a geração de registros em disco.
TPU_MIN_LOG_LEVEL: a gravidade mínima que será registrada em disco.
As opções são 0 (INFO), 1 (WARNING), 2 (ERROR) e 3 (FATAL). O padrão é 0.
TPU_STDERR_LOG_LEVEL: a gravidade mínima que será registrada em stderr, além de no disco, se aplicável.
As opções são as mesmas de TPU_MIN_LOG_LEVEL. O padrão é 3.
TPU_MAX_LOG_SIZE_MB: o tamanho máximo em megabytes de cada arquivo de registro.
Um novo arquivo de registro será iniciado automaticamente quando o anterior atingir aproximadamente esse tamanho. O padrão é 1024.