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 à divisã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 pode dar a impressão de 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 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 para 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ária para um tensor e pode levar a 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 a 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 matrizes de transmissão.

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

No exemplo seguinte, a transmissão implícita na adição de um vetor e de 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)`

Recomendações de desempenho para a arquitetura de dois chiplets Ironwood

O modelo de programação Ironwood permite-lhe aceder a dois dispositivos TPU em vez da arquitetura de núcleo lógico único (também conhecido como MegaCore) usada nas gerações anteriores (TPU v4 e v5p). Esta alteração melhora a rentabilidade e a eficiência do fabrico do chip. Embora isto represente uma mudança arquitetónica, o novo design garante que pode reutilizar os modelos de software existentes com alterações mínimas.

Para alcançar o melhor desempenho com a arquitetura de dois chiplets, recomendamos as seguintes abordagens:

• Use o paralelismo de tensores em vários chiplets: a interface D2D de largura de banda elevada foi concebida para um paralelismo de tensores eficiente. Recomendamos que divida os tensores pelos dois dispositivos no chip.

• Use coletivos hierárquicos: para maximizar a eficiência da comunicação, tire partido da hierarquia de rede de dois níveis: a ligação D2D ultrarrápida entre os chiplets no chip e as ligações ICI rápidas numa fatia. Quando usa o paralelismo automático com SPMD (single program, multiple data), o compilador XLA processa isto automaticamente gerando operações coletivas hierárquicas. Quando particiona manualmente o modelo, também deve estruturar os padrões de comunicação em torno desta hierarquia. Priorizar a comunicação entre os dois dispositivos no mesmo chip antes de comunicar com dispositivos noutros chips.

• Sobreponha a comunicação com a computação: para maximizar a utilização do hardware, transfira as operações de comunicação coletiva, como a redução total, para os SparseCores. Estas operações, que não estão associadas à unidade de multiplicação de matrizes (MXU), podem ser executadas nos SparseCores em simultâneo, enquanto os TensorCores continuam os respetivos cálculos. Esta técnica pode recuperar algumas das vantagens de desempenho inerentes às operações fundidas na arquitetura MegaCore anterior.

• Transferência para o SparseCore para incorporações: no design de dois chiplets, as tabelas de incorporação podem ser divididas nas HBMs de ambos os chiplets. Para evitar a degradação do desempenho devido à falta de memória partilhada, transfira as operações de recolha de incorporações para o SparseCore. Esta estratégia usa a interligação D2D de alta velocidade para transferir eficientemente vetores de incorporação entre os chiplets. Para mais informações sobre o SparseCore e os modelos de incorporação, consulte o artigo Uma análise detalhada do SparseCore para modelos de incorporação (conteúdo extenso).

Para mais informações sobre a arquitetura Ironwood na TPU7x, consulte o artigo TPU7x (Ironwood).