Solução de problemas do TensorFlow: TPU

Este guia, junto com as Perguntas frequentes, oferece ajuda na solução de problemas para usuários que estão treinando modelos do TensorFlow no Cloud TPU. Se você estiver solucionando problemas de treinamento do Pytorch ou do JAX, consulte os documentos de solução de problemas para estes frameworks:

Para ver guias mais gerais sobre como usar a Cloud TPU, consulte:

Visão geral

Os problemas comuns encontrados com as Cloud TPUs se enquadram nas seguintes categorias:

  1. Problemas na conexão com a TPU

  2. Como depurar erros comuns

  3. Como reduzir o uso de memória

  4. Como aumentar a velocidade do treinamento

  5. Como depurar quedas na acurácia do modelo

Problema ao se conectar ao servidor da TPU

Nesta seção, descrevemos como solucionar situações em que 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 canal de treinamento do TensorFlow. Para isso, execute o tutorial do MNIST usando o URL de servidor da TPU e verifique se funciona corretamente. Caso você ainda encontre problemas de conexão no tutorial do MNIST, isso confirma que o problema é no servidor da TPU. Nesse caso:

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

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

    Isso exibe um resultado como este:

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

  2. Verifique se o valor passado para --tpu (demo-tpu no exemplo acima) está correto e se a TPU consta como READY.

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

    (vm)$ gcloud compute tpus stop $TPU_SERVER_NAME && gcloud compute tpus start $TPU_SERVER_NAME

    No exemplo acima, $TPU_SERVER_NAME é demo-tpu. Pode ser que isso leve alguns minutos.

  4. Execute novamente o comando ... tpus list acima e espere a TPU entrar no estado READY. Isso pode levar algum tempo.

  5. Tente executar o tutorial do MNIST novamente.

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

Se o exemplo do MNIST for executado corretamente, mas o modelo ainda parar de responder, é provável que o problema esteja no pipeline de treinamento. Para depurar isso, substitua a TPUStrategy no código pela estratégia padrão. Quando você usa a estratégia padrão, sempre que usar strategy.scope() ou strategy.run(), o modelo é executado na CPU (ou GPU, se presente) em vez da TPU. Se o modelo for executado na CPU e não na TPU, será necessário haver um problema específico da TPU. Se ele ainda não for executado, a prática recomendada é depurar o problema na CPU.

Perda de conexão com ssh durante o treinamento

A conexão ssh com a Cloud TPU pode expirar durante um treinamento de longa duração, principalmente se você estiver usando o Cloud Shell. Neste momento, não há saída para o console da TPU e pode parecer que a TPU 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 manterá a conexão ssh ativa, independentemente da duração do treinamento.

Como depurar erros comuns

Não é possível criar uma TPU

Ao criar uma Cloud TPU, talvez você veja 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."

Este é um problema de permissões e pode ser resolvido executando o seguinte comando:

gcloud auth login --update-adc

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

Não é possível usar o sistema de arquivos local

Mensagem de erro

InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

Estruturas e configurações afetadas

Essa mensagem pode ocorrer durante o treinamento com o TensorFlow usando a arquitetura de nó da TPU.

Detalhes

Todos os arquivos de entrada e o diretório do modelo precisam usar um caminho do bucket do Cloud Storage (gs://bucket-name/...). Esse bucket precisa ser acessível do servidor da TPU. Todo o processamento de dados e a criação de checkpoints do modelo são executados no servidor da TPU, e não na máquina local. Para mais informações sobre como configurar corretamente o Cloud Storage para usá-lo com a TPU, consulte o guia Como conectar-se a buckets do Cloud Storage.

Tipo de dados incompatível

Mensagem de erro

TypeError: DataType is not a supported TPU infeed type.

Estruturas e configurações afetadas

Essa mensagem pode ocorrer durante o treinamento com o TensorFlow usando a arquitetura de nó da TPU.

Detalhes

No momento, somente os tipos de dados tf.float32, tf.int32, tf.bfloat16 e tf.bool são compatíveis com a TPU. Outros tipos de dados comuns, como tf.uint8, tf.string e tf.int64, precisam ser convertidos em um dos tipos de dados compatíveis durante o pré-processamento de dados (ou seja, no pipeline tf.data.Dataset).

Veja um exemplo de conversão na função decode_image usada no treinamento MNIST.

Formas dinâmicas incompatíveis

Mensagem de erro

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

Estruturas e configurações afetadas

Essa mensagem é exibida somente durante a compilação do XLA com o TensorFlow.

Detalhes

Para executar um modelo na TPU, o TensorFlow o compila usando o compilador XLA. Embora essa etapa de compilação melhore significativamente a velocidade de treinamento e o uso de memória, as formas (tamanhos de dimensão) de todos os tensores no gráfico precisam ser conhecidas no momento da compilação. Se não for possível determinar formas nesse momento, a compilação da TPU falhará com um erro como o acima.

Uma operação comum que retorna uma forma dinâmica é 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 treinar na TPU, defina drop remainder=True para dataset.batch. Isso possivelmente elimina as últimas amostras de um arquivo para garantir que cada lote tenha uma forma estática de batch_size. 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

Estruturas e configurações afetadas

Essa mensagem pode aparecer durante o treinamento com o TensorFlow.

Detalhes

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

Para ver uma lista de operações disponíveis na TPU, juntamente com planos para compatibilidade futura 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.

Estruturas e configurações afetadas

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

Detalhes

Cada Cloud TPU é composta de oito núcleos de TPU, as TPUs v2 têm 8 GB e as TPUs v3 têm 16 GB de RAM (ou HBM, memória de alta largura de banda). 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 falhará, e a mensagem de erro acima será exibida. Consulte a seção sobre como reduzir o uso de memória para encontrar mais ajuda.

Dicas para reduzir o uso de memória:

Problemas que interrompem a execução

Quando há erro no TensorFlow durante a execução da TPU, o script às vezes deixa de responder em vez de ser encerrado e voltar para o shell. Se isso acontecer, pressione CTRL+\ 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 ciclo de iteração atual para se encerrar corretamente.

Se você encontrar novos erros ao se reconectar à TPU depois de sair dessa maneira, redefina manualmente o servidor da TPU com os comandos:

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

em que tpu-name é retirado da primeira coluna mostrada pelo comando gcloud compute tpus list e zone é a zona mostrada na segunda coluna.

Preenchimento excessivo do 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 totalmente 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, você precisará reduzir o uso da memória.

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

  • Reduza o preenchimento excessivo de tensores
  • Reduzir o tamanho do lote

Modelo ou lote 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 de um destes:

  1. O uso da memória é proporcional ao número de pesos no modelo.
  2. armazenamento de ativações intermediárias da transmissão para frente necessário para calcular a transmissão para trás. 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 (chamado de rematerialização, semelhante ao checkpoint de gradiente). No entanto, nem sempre é possível fazer isso.

Como reduzir o uso de memória

Reduza lentamente o tamanho do lote até que ele caiba na memória, garantindo que o tamanho total do lote seja múltiplo de 64 (o tamanho do 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 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 (por exemplo, 64), reduza o número de camadas ou os tamanhos delas.

Como melhorar a velocidade do treinamento

Nesta seção, descrevemos várias maneiras possíveis de melhorar a velocidade do treinamento caso você consiga executar o modelo na TPU, mas a velocidade é mais lenta que o esperado. Consulte o Guia de desempenho para ver 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 callbacks de host. Cada callback de host requer uma comunicação significativa entre a CPU do host e o dispositivo da TPU e, por isso, se steps_per_execution for muito pequeno, ele poderá atrasar o treinamento.

Como saber se o modelo foi afetado

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

Como mitigar

Defina steps_per_execution como um valor maior. Observe que steps_per_execution pode ser definido como um valor alto, mas lembre-se de que registrar mensagens e salvar um checkpoint só pode ocorrer depois que o número especificado de etapas tiver sido 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 do modelo, a TPU ficará ociosa enquanto a função de entrada estiver recuperando dados.

Como saber se o modelo foi afetado

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

image

A página de análise do canal de entrada exibe um resumo claro que mostra se o modelo tem gargalo do processamento de entrada. Ela também exibe o tempo de execução por operação para identificar as 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 (em inglês) ou do ResNet para ver exemplos.
  2. Usar dataset.cache() e/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 simples 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 localizadas no servidor da TPU, e não na máquina local. Portanto, a velocidade da máquina local não é um fator.

Tempos 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ê verá nesse caso são tempos lentos de passo acoplados a uma baixa utilização de MXU mostrada ao criar um perfil de 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 matrizes, faça uma nova comparação para ver se o desempenho é aceitável em TPUs.

Preenchimento excessivo do 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 lenta e baixo uso de 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 se encaixe na memória e seja um múltiplo de 64. Para isso, o jeito mais fácil é começar com 1024. Se ocorrer um erro de falta de memória, reduza o tamanho do lote até que o modelo seja executado com êxito. 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 nativa da unidade de multiplicação de matriz da TPU. Nas camadas completamente conectadas, o tamanho mínimo oculto de 512 é recomendado para gerar alta eficiência. As camadas convolucionais normalmente não precisam ser tão grandes quanto as completamente conectadas para atingir o mesmo nível de eficiência.

Como mitigar

Se a motivação principal para tamanhos de camada pequenos no modelo for a velocidade do treinamento, faça uma nova comparação dos modelos com camadas maiores na TPU. Por exemplo, aumentar o tamanho da saída de uma camada de 256 para 512 só pode aumentar o tempo de treinamento em 20%, mesmo que o modelo execute duas vezes os cálculos.

Criação de perfil de modelo no nível operacional

Geralmente, é bom medir o uso de memória e tempo de execução no nível operacional para identificar gargalos de desempenho. Para instruções sobre como fazer isso,
consulte o guia Ferramentas da Cloud TPU: visualizador de traces.

Como depurar quedas na acurácia do modelo

Um dos objetivos do ecossistema da Cloud TPU é que qualquer modelo treinado atualmente em uma CPU ou GPU atinja uma acurácia muito semelhante no treinamento 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. Veja nesta seção orientações sobre como identificar a principal causa de qualquer queda na acurácia ao fazer a portabilidade 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ê encontrar qualquer diferença significativa entre a saída de operações não aleatórias na TPU e na CPU, denuncie como um bug.

No entanto, para o pipeline de treinamento como um todo, há uma diferença significativa entre o treinamento na CPU/GPU e na TPU. Ao treinar 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. Para cada etapa no treinamento, cada núcleo de TPU recebe um lote de dados, calcula os gradientes de peso, troca os gradientes com os outros núcleos de TPU e, em seguida, calcula a atualização de peso. Por padrão, a média de perda é calculada 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, este 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. Cada um dos 8 núcleos executará a normalização de lote nas 16 amostras. Em alguns casos, foi realizada uma normalização em lotes pequenos (por exemplo, menos de 32) para reduzir a precisão. No cenário ideal, o tamanho total do lote deve ser grande. Por exemplo, de 256 a 1.024. 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.

Treinamento determinista

Um motivo por que é difícil depurar diferenças na acurácia do modelo é que, em diferentes frameworks (TensorFlow, PyTorch, JAX), o software de treinamento usa a inicialização do peso diferente e o embaralhamento de dados sempre que um modelo é treinado. É recomendável modificar o procedimento de treinamento para que ele seja determinístico, de modo que várias execuções produzam modelos quase idênticos. Nesta seção, mostramos como executar o tutorial do MNIST de maneira determinista:

  1. Gere um arquivo de checkpoint inicial ao executar uma única etapa na CPU. Essa etapa é usada para conseguir a inicialização de peso determinista. Além disso, use uma sugestão aleatória fixa para qualquer função aleatória no modelo.
# Run training for 1 step to create an initial checkpoint.
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/init_output \
  --random_seed=12345 \
  --iterations=1
  --train_steps=1
  1. Modifique as funções de embaralhamento de dados na função de entrada para usar uma sugestão aleatória. Isso já foi feito no tutorial do MNIST. É algo que funciona para as operações de processamento de dados de entrada porque elas são sempre executadas na CPU. As operações aleatórias na função do modelo podem não ser deterministas entre a TPU e a CPU. Sempre que você chamar uma operação aleatória, transmita uma semente fixa para garantir os mesmos resultados entre as execuções. Exemplo:
# In the flag definitions
tf.flags.DEFINE_integer("batch_size", None, "Random seed for training")

# In the input_fn
if FLAGS.random_seed is not None:
dataset = dataset.shuffle(seed=FLAGS.random_seed)
  1. Execute o mesmo modelo duas vezes na CPU para verificar se o treinamento é determinístico. O treinamento precisa ser realizado para um número razoável de etapas (por exemplo, 1.000), mas não precisa ser executado para convergência.

    Como o treinamento na CPU é parecido com um treinamento de núcleo único na TPU, use um tamanho de lote que caiba em um único núcleo de TPU. Normalmente, é o tamanho total do lote dividido por 8. O TensorFlow não garante o determinismo bit a bit entre as execuções, mas a perda será muito parecida:

Copiar as ponderações iniciais

gsutil mkdir ${STORAGE_BUCKET}/cpu_output_1
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_1
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_2
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_2

Corrida 1

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

Saída 1

accuracy = 0.9910644, global_step = 1000, loss = 0.025323588

Corrida 2

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

Saída 2

accuracy = 0.9910644, global_step = 1000, loss = 0.025323414

Treinamento de núcleo único na TPU

Depois que for possível executar o tutorial do MNIST com determinismo, a próxima etapa é replicar os resultados treinados pela CPU na TPU. Você usa um único núcleo de TPU para identificar se o problema está relacionado à fragmentação de dados ou ao próprio mecanismo de execução da TPU.

Veja como executar uma avaliação e treinamento de núcleo único no tutorial do MNIST:

Usar a mesma inicialização de peso que a CPU

gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/tpu_output

Faça um treinamento de 1.000 passos

python mnist.py \
    --use_tpu=True \
    --master=$GRPC_SERVER \
    --train_file=${STORAGE_BUCKET}/data/train.tfrecords \
    --model_dir=${STORAGE_BUCKET}/tpu_output \
    --random_seed=12345 \
    --num_shards=1 \
    --batch_size=128 \
    --train_steps=1000 \
    --eval_steps=10

Saída

  accuracy = 0.9910644, global_step = 1000, loss = 0.02514153

A perda não corresponderá exatamente ao modelo treinado pela CPU, mas será parecida. Caso contrário, isso indicará que você encontrou um bug no mecanismo de execução da TPU. Antes de enviar um relatório de bug, verifique o seguinte:

  1. Se você está passando num_shards=1 para TPUConfig.

  2. Se você não tem operações aleatórias na função do modelo e se qualquer operação aleatória na função de entrada está sendo sugerida corretamente.

  3. se você está usando o mesmo arquivo de ponto de verificação inicial no treinamento da CPU e da TPU.

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

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, como mostrado acima.

Se você tiver certeza de que a queda na acurácia tem significância estatística, estes sã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 (por exemplo, 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) (em inglês) usam um discriminador de lotes ao treinar uma rede adversária generativa.

Solução de problemas de VM da TPU

Os problemas e soluções a seguir são aplicáveis apenas às configurações de VM da TPU.

Solução de problemas de configuração do 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 o gcloud com a VM da TPU, é necessário usar uma instalação do gcloud que não seja gerenciada por um gerenciador de pacotes. Siga estas etapas para instalar gcloud do código-fonte:
  sudo apt-get remove google-cloud-sdk
  curl -O https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  tar -xzf google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  ./google-cloud-sdk/install.sh
  source ~/.bashrc
Problema

A execução de qualquer comando começando com gcloud alpha compute tpus tpu-vm exibe as seguintes informações:

ERROR: (gcloud.alpha.compute.tpus) Invalid choice: 'tpu-vm'.
Solução

Isso acontece quando o repositório do componente não foi atualizado corretamente. Para verificar isso, execute gcloud --version. A primeira linha da saída precisa ser "Google Cloud SDK HEAD". Se a saída for diferente, a atualização não aconteceu. Se isso acontecer, tente atualizar os componentes do gcloud com o comando a seguir.

gcloud components update

Se o erro persistir, tente reinstalar o gcloud com o comando:

gcloud components reinstall
Problema

O comando gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} exibe a seguinte 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.alpha.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

Depurar registros

Os frameworks de Cloud TPU compatíveis, JAX, PyTorch e TensorFlow acessam as TPUs por meio de uma biblioteca compartilhada chamada libtpu, presente em todas as VMs de TPU. Essa biblioteca inclui o compilador XLA usado para compilar programas da TPU, o ambiente de execução da TPU usado para executar programas compilados e o driver da 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 no /tmp/tpu_logs em cada VM da Cloud TPU. As variáveis de ambiente a seguir podem ser definidas antes do início do treinamento para modificar o comportamento de geração de registros:

TPU_LOG_DIR: o diretório em que os registros são gravados.
O local do diretório é /tmp/tpu_logs por padrão. O diretório será criado se ainda não existir, mas nenhum diretório pai será criado. Se um erro for encontrado ou criado no diretório especificado, uma mensagem será impressa em stderr, mas não 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 no disco.
TPU_MIN_LOG_LEVEL: a gravidade mínima que será registrada no disco
As opções são 0 (INFORMAÇÕES), 1 (AVISO), 2 (ERRO) e 3 (FATAL). O padrão é 0.
TPU_STDERR_LOG_LEVEL: além da gravidade, que será registrada como stderr, além do disco
As opções são as mesmas da TPU_MIN_LOG_LEVEL. O padrão é 3.
TPU_MAX_LOG_SIZE_MB: o tamanho máximo em megabytes de cada arquivo de registros
Um novo arquivo de registro será iniciado automaticamente quando o anterior atingir aproximadamente esse tamanho. O padrão é 1024.