Guia de desempenho da Cloud TPU

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

Desempenho 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 que haja um fluxo constante de dados sendo carregados nela. A maneira como isso é feito depende de como você carrega e pré-processa seu 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. O tamanho do lote é muito pequeno por causa da fragmentação (dividir lotes entre núcleos)

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

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

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 PyTorch e JAX. Os modelos do Cloud TPU são convertidos em um gráfico do XLA, que é compilado pelo XLA em um executável de TPU. Para saber mais sobre XLA, consulte XLA: como otimizar o compilador para machine learning (em inglês).

Preenchimento

Para usar a memória da TPU de maneira eficiente, estruture seus dados para que possam ser divididos em blocos de 128 x 8. Quando os dados para o cálculo de uma matriz não preenchem um bloco inteiro de 128 x 8, o compilador XLA preenche os tensores. O padding tem duas desvantagens:

1. O uso de tensores com padding subutiliza o núcleo da TPU.
2. O padding aumenta a quantidade de armazenamento de memória on-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, você pode determinar a quantidade de preenchimento realizado usando a ferramenta Visualizador de memória. Para evitar o preenchimento, escolha as dimensões do tensor adequadas para a TPU.

Dimensões do tensor

O compilador XLA arredonda os tamanhos dos tensores armazenados na memória HBM de TPU para realizar cálculos com mais eficiência. Esse padding acontece de maneira transparente 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 ambiente de execução da TPU estabelece 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 precisa ser verdadeira:

1. O tamanho total do lote precisa ser um múltiplo de 64 (8 por núcleo de TPU) e os tamanhos da dimensão do recurso precisam ser um múltiplo de 128.

2. O tamanho total do lote precisa ser um múltiplo de 1024 (128 por núcleo de TPU) e os tamanhos da dimensão do recurso precisam ser um múltiplo de 8.

Usar um tamanho de lote de 1.024 e dimensões de recurso que são um múltiplo de 128 resulta em melhor eficiência, embora isso possa não ser possível para todos os modelos.

Observação :a dimensão de atributo refere-se 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 podem seguir essa regra, especialmente a primeira e a última camada da rede. Isso não é um problema, a maioria dos modelos exige uma certa quantidade de padding.

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 constituintes que precisam ser executadas em conjunto.

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 = tmp[i] * z[i];
    }

Com a fusão, os acessos à matriz acontecem ao mesmo tempo:

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

Nesse exemplo, o número de idas e voltas de 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.

Embora as transmissões possam muitas vezes ser fundidas com os consumidores, forçar uma transmissão pode resultar em mau 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)`