Guia de desempenho do Cloud TPU

O primeiro passo na resolução de problemas de desempenho da TPU é criar um perfil do seu modelo. Para mais informações sobre como capturar um perfil de desempenho, consulte o artigo Criar perfis do seu modelo na Cloud TPU.

Desempenho do modelo de TPU

Esta secção descreve problemas gerais que podem reduzir o desempenho do modelo e como os pode resolver.

1. O modelo está associado à entrada

   As TPUs fazem cálculos muito rapidamente. Para garantir que a TPU não está inativa, é importante certificar-se de que existe um fluxo constante de dados a ser carregado na TPU. A forma como isto é feito depende da forma como carrega e pré-processa o seu conjunto de dados. Por exemplo, pode ler ficheiros de dados em paralelo usando tf.data.TFRecordset() e o parâmetro num_parallel_reads.

2. O tamanho do lote é demasiado pequeno devido à fragmentação (divisão de lotes entre núcleos)

   O tempo de execução da TPU divide um lote nos 8 núcleos de um dispositivo TPU (por exemplo, v2-8 ou v3-8). Se especificar um tamanho do lote global de 128, cada núcleo recebe um tamanho do lote de 16 (128 / 8).

   Para uma utilização ideal da memória, use o maior tamanho do lote que se ajuste à memória da TPU. Cada núcleo da TPU usa registos vetoriais bidimensionais de 8 x 128 para processar multiplicações de matrizes. Em geral, o tamanho do lote deve ser divisível por 8 ou 128.

3. Sintonização da gestão de memória

   Pode usar as variáveis de ambiente TPU_PREMAPPED_BUFFER_SIZE para ajustar os comportamentos de tempo de execução de baixo nível.

• Descrição: TPU_PREMAPPED_BUFFER_SIZE define o tamanho da memória intermédia da memória do anfitrião (em bytes) que é pré-mapeada e fixada para utilização pelo tempo de execução da TPU para transferências de dados (por exemplo, DMA). O valor predefinido é de 4294967296 bytes. O valor tem de ser um múltiplo de 2^12 (4 KB = 4 * 1024 bytes = 4096 = 2^12).

  Os seguintes exemplos são valores TPU_PRE_MAPPED_BUFFER_SIZE válidos.

      17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped).
      40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).
  

• Impacto: aumentar este tamanho pode melhorar potencialmente o desempenho da transferência de dados entre o anfitrião e o dispositivo TPU, especialmente para cargas de trabalho com tensores grandes ou comunicação frequente entre o anfitrião e o dispositivo. No entanto, também aumenta a quantidade de memória do anfitrião fixada, reduzindo a memória disponível para outros processos.

  Tamanho da memória intermédia

  Se a região de buffer pré-mapeada não for suficientemente grande para alocar memória durante a execução do programa, a carga de trabalho falha e devolve um erro RESOURCE_EXHAUSTED semelhante ao seguinte:

  "A atribuição do buffer da região pré-mapeada falhou com: RESOURCE_EXHAUSTED: A tentar atribuir allocation_size. Isso não foi possível. Existem available_size gratuitos."

  Se o buffer for excessivamente grande, a inicialização da TPU pode demorar muito mais tempo (potencialmente mais de 15 segundos), o que faz parecer que a TPU está bloqueada.

  Para diagnosticar este problema, inspecione os registos de tempo de execução da TPU. Estes registos detalham as operações que estão a ser realizadas, incluindo o pré-mapeamento de buffers. Pode encontrar os registos em /tmp/tpu_logs/tpu_driver.INFO ou imprimi-los diretamente na consola definindo a variável de ambiente TPU_STDERR_LOG_LEVEL=0. Esta definição gera um resultado semelhante ao seguinte:

   I0604 12:45:24.926233   62136 tpu_hal.cc:214] Starting premapped memory manager initialization...
   I0604 12:45:29.411218   62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0
   I0604 12:45:29.411244   61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s
   I0604 12:45:29.411267   61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296
   ```
  
  This output will tell you how long it took to initialize the TPU and
  the size of the premapped buffer.
  

• Utilização: se o buffer pré-mapeado for demasiado pequeno ou demasiado grande, pode definir manualmente o tamanho do buffer através das seguintes variáveis de ambiente.

  TPU_PREMAPPED_BUFFER_SIZE: Sets the total size (in bytes) of the
  pre-mapped buffer region.
  TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: Sets the maximum size of
  a single buffer that can be allocated from the pre-mapped region.
  

  Por exemplo, pode:

   export TPU_PREMAPPED_BUFFER_SIZE=4294967296
  

  para definir o tamanho da memória intermédia e:

   export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES
   ```
   to enable it.
  
   This export sets the size to the default.
  

• Orientação: ajuste o valor de TPU_PREMAPPED_BUFFER_SIZE se suspeitar que a transferência de dados do dispositivo anfitrião é um gargalo. Monitorize a utilização da memória do anfitrião e o desempenho do modelo para encontrar um equilíbrio ideal. Normalmente, o valor predefinido é suficiente para a maioria dos exemplos de utilização.

Otimizações do compilador XLA

O XLA é um compilador para aprendizagem automática que pode produzir ficheiros binários para TPUs, CPUs, GPUs e outras plataformas. Embora o XLA faça parte da base de código padrão do TensorFlow, também pode ser usado em modelos do PyTorch e do JAX. Os modelos para a Cloud TPU são traduzidos num gráfico XLA, que o XLA compila num ficheiro executável da TPU. Para mais informações sobre a XLA, consulte o artigo XLA: compilador de otimização para aprendizagem automática.

Preenchimento

Para usar a memória da TPU de forma eficiente, estruture os seus dados de modo que possam ser divididos em blocos de 128 x 8. Quando os dados de um cálculo de matriz não preenchem um bloco de 128 x 8 completo, o compilador XLA preenche os tensores. Existem duas desvantagens no preenchimento:

1. Os tensores com preenchimento não usam totalmente o núcleo da TPU.
2. O preenchimento aumenta a quantidade de armazenamento de memória no chip necessário para um tensor e pode originar um erro de falta de memória.

Embora o preenchimento seja realizado automaticamente pelo compilador XLA quando necessário, pode determinar a quantidade de preenchimento realizado através da ferramenta de visualização de memória. Pode evitar o preenchimento escolhendo dimensões de tensores adequadas para a TPU.

Dimensões do tensor

Para alcançar o pico de FLOPs, as dimensões da multiplicação de matrizes devem ser maiores do que o tamanho da MXU para a versão da TPU que está a usar. O tamanho da MXU é de 256 x 256 para a versão v6e e de 128 x 128 para versões anteriores à v6e. Para mais informações, consulte o artigo Arquitetura do sistema Cloud TPU.

Tamanho do lote

O compilador XLA arredonda os tamanhos dos tensores armazenados na memória HBM da TPU para realizar cálculos de forma mais eficiente. Este preenchimento ocorre de forma transparente ao nível do hardware e não afeta os resultados. No entanto, em determinados casos, o preenchimento pode resultar num aumento significativo da utilização de memória e do tempo de execução.

O tempo de execução da TPU organiza os tensores na memória para maximizar a eficiência computacional e minimizar o preenchimento. Para minimizar a sobrecarga de memória e maximizar a eficiência computacional, uma das seguintes condições tem de ser verdadeira:

1. O tamanho total do lote deve ser um múltiplo de 64 (8 por núcleo da TPU) e os tamanhos das dimensões das caraterísticas devem ser um múltiplo de 128.

2. O tamanho total do lote deve ser um múltiplo de 1024 (128 por núcleo da TPU) e os tamanhos das dimensões das caraterísticas devem ser um múltiplo de 8.

A utilização de um tamanho do lote de 1024 e dimensões de caraterísticas que sejam um múltiplo de 128 resulta na melhor eficiência, embora isto possa não ser possível para todos os modelos.

Nota: a dimensão da funcionalidade refere-se ao tamanho oculto de uma camada totalmente ligada ou ao número de canais de saída numa convolução. Nem todas as camadas podem estar em conformidade com esta regra, especialmente a primeira e a última camadas da rede. Não há problema. A maioria dos modelos requer alguma quantidade de preenchimento.

Fusion

A fusão é uma técnica geral que o compilador XLA usa para otimizar programas. Uma operação unida é a combinação de várias operações constituintes que vão ser executadas em combinação.

Por exemplo, considere a seguinte série de operações:

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

Este código é aproximadamente equivalente ao seguinte pseudocódigo:

    for (i = 0; i < element_count; i++) {
      tmp[i] = x[i] + y[i];
    }

    for (i = 0; i < element_count; i++) {
      result[i] = tmp[i] * z[i];
    }

Com a união, os acessos à matriz ocorrem em simultâneo:

    for (i = 0; i < element_count; i++) {
      result[i] = (x[i] + y[i]) * z[i];
    }

Neste exemplo, o número de viagens de ida e volta à memória é reduzido e o XLA não precisa de alocar espaço para "tmp".

A união é uma otimização crítica e beneficia a Cloud TPU de várias formas:

• Reduz as transferências de memória, eliminando a necessidade de armazenar resultados intermédios na memória principal, que é lenta.
• Permite uma maior utilização das unidades de hardware que, de outra forma, não seriam utilizadas.
• Pode reduzir a utilização de memória de um modelo, uma vez que são necessários menos buffers para estarem ativos em simultâneo.

Transmissão

A transmissão ocorre implicitamente quando dois tensores com formas diferentes, mas compatíveis, são combinados.

Por exemplo, tf.add(vector, matrix) requer que o vetor seja transmitido para a forma da matriz. O resultado da operação tem o mesmo formato que a matriz. Para mais detalhes, consulte o guia sobre como transmitir matrizes.

Embora as transmissões possam ser frequentemente unidas aos respetivos consumidores, forçar uma transmissão pode resultar num desempenho fraco e num aumento da utilização de memória.

No exemplo seguinte, a transmissão implícita na adição de um vetor e uma matriz não pode ser fundida com o argmax, o que resulta numa transmissão materializada:

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`