Guia de desempenho da Cloud TPU

A primeira etapa para resolver problemas de desempenho da TPU é criar um perfil do modelo. Para mais informações sobre como capturar um perfil de desempenho, consulte Como criar um perfil do modelo no Cloud TPU.

Performance do modelo de TPU

Esta seção descreve problemas gerais que podem reduzir o desempenho do modelo e como você pode resolvê-los.

1. O modelo é limitado pela entrada

   As TPUs realizam cálculos muito rapidamente. Para garantir que a TPU não fique ociosa, é importante verificar se há um fluxo constante de dados sendo carregados nela. A maneira de fazer isso depende de como você carrega e pré-processa o conjunto de dados. Por exemplo, é possível ler arquivos de dados em paralelo usando tf.data.TFRecordset() e o parâmetro num_parallel_reads.

2. Tamanho do lote muito pequeno devido à fragmentação (dividindo lotes entre núcleos)

   O tempo de execução da TPU divide um lote em todos os oito núcleos de um dispositivo de TPU (por exemplo, v2-8 ou v3-8). Se você especificar um tamanho de lote global de 128, cada núcleo vai receber um tamanho de lote de 16 (128 / 8).

   Para otimizar o uso de memória, use o maior tamanho de lote que cabe na memória da TPU. Cada núcleo de TPU usa registros vetoriais bidimensionais de 8 X 128 para processar multiplicações de matrizes. Em geral, o tamanho do lote deve ser divisível uniformemente por 8 ou 128.

3. Ajuste do gerenciamento de memória

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

• Descrição:TPU_PREMAPPED_BUFFER_SIZE define o tamanho do buffer de memória do host (em bytes) que é pré-mapeado e fixado para uso pelo ambiente de execução da TPU em transferências de dados (por exemplo, DMA). O valor padrão é 4294967296 bytes. O valor precisa ser um múltiplo de 2^12 (4 KB = 4 * 1024 bytes = 4096 = 2^12).

  Os exemplos a seguir são valores válidos de TPU_PRE_MAPPED_BUFFER_SIZE.

      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 esse tamanho pode melhorar o desempenho da transferência de dados entre o host e o dispositivo de TPU, especialmente para cargas de trabalho com tensores grandes ou comunicação frequente entre host e dispositivo. No entanto, isso também aumenta a quantidade de memória de host fixada, reduzindo a memória disponível para outros processos.

  Tamanho do buffer

  Se a região de buffer pré-mapeada não for grande o suficiente para alocar memória durante a execução do programa, a carga de trabalho vai falhar e retornar um erro RESOURCE_EXHAUSTED semelhante a:

  "Falha ao alocar buffer da região pré-mapeada com: RESOURCE_EXHAUSTED: Tentando alocar allocation_size. Isso não foi possível. Há available_size gratuitos."

  Se o buffer for muito grande, a inicialização da TPU poderá levar muito mais tempo (potencialmente mais de 15 segundos), fazendo parecer que a TPU está travada.

  Para diagnosticar isso, inspecione os registros de tempo de execução da TPU. Esses registros detalham as operações realizadas, incluindo o pré-mapeamento de buffers. Você pode encontrar os registros em /tmp/tpu_logs/tpu_driver.INFO ou imprimi-los diretamente no console definindo a variável de ambiente TPU_STDERR_LOG_LEVEL=0. Essa configuração vai gerar uma saída semelhante a esta:

   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.
  

• Uso:se o buffer pré-mapeado for muito pequeno ou muito grande, defina manualmente o tamanho do buffer usando as 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, você pode:

   export TPU_PREMAPPED_BUFFER_SIZE=4294967296
  

  para definir o tamanho do buffer 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 você suspeitar que a transferência de dados do host para o dispositivo é um gargalo. Monitore o uso da memória do host e o desempenho do modelo para encontrar um equilíbrio ideal. O valor padrão geralmente é suficiente para a maioria dos casos de uso.

Otimizações do compilador XLA

O XLA é um compilador para machine learning que pode produzir binários para TPUs, CPUs, GPUs e outras plataformas. Embora o XLA faça parte da base de código padrão do TensorFlow, ele também pode ser usado em modelos do PyTorch e do JAX. Os modelos para o Cloud TPU são traduzidos em um gráfico do XLA, que é compilado pelo mesmo em um executável de TPU. Para mais informações sobre o XLA, consulte XLA: como otimizar o compilador para machine learning.

Preenchimento

Para usar a memória da TPU com eficiência, estruture seus dados para que eles possam ser divididos em blocos de 128 x 8. Quando os dados de uma computação de matriz não preenchem um bloco inteiro de 128 x 8, o compilador XLA preenche os tensores. O preenchimento apresenta duas desvantagens:

1. Os tensores preenchidos utilizam pouco o núcleo da TPU.
2. O preenchimento aumenta a quantidade de armazenamento em memória on-chip necessária para um tensor, o que pode resultar em erro de falta de memória.

O preenchimento é realizado automaticamente pelo compilador XLA quando necessário. No entanto, é possível determinar a quantidade de preenchimento usando a ferramenta de visualização de memória. Para evitar o preenchimento, escolha dimensões de tensor que sejam adequadas para a TPU.

Dimensões do tensor

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

Tamanho do lote

O compilador XLA arredonda os tamanhos dos tensores armazenados na memória HBM da TPU para realizar 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.

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

1. O tamanho total do lote precisa ser um múltiplo de 64 (8 por núcleo de TPU), e as dimensões do recurso precisam ser um múltiplo de 128.

2. 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.

Usar um tamanho de lote de 1024 e dimensões de recurso que sejam um múltiplo de 128 resulta na melhor eficiência, embora isso não seja possível em todos os modelos.

Observação: 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.

Fusão

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

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

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

Esse código é aproximadamente equivalente ao seguinte código falso:

    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 fusão, os acessos à matriz acontecem ao mesmo tempo:

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

Nesse exemplo, o número de idas e voltas à memória é reduzido e o XLA não precisa alocar espaço para "tmp".

A fusão é um tipo de otimização essencial e traz diversos benefícios para o Cloud TPU:

• Reduz as transferências de memória ao eliminar a necessidade de armazenar resultados intermediários na memória principal, que é um processo lento.
• Proporciona maior utilização das unidades de hardware que, caso contrário, ficariam ociosas.
• Pode reduzir a quantidade de memória usada pelo modelo porque menos buffers precisam estar ativos ao mesmo tempo.

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 a mesma forma que a matriz. Para saber mais, consulte o guia de transmissão de matrizes.

Muitas vezes, as transmissões podem ser fundidas com os elementos que as consomem, mas obrigar uma transmissão pode resultar em baixo desempenho e aumento do uso da memória.

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

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