Guia de desempenho da Cloud TPU
A primeira etapa para resolver problemas de desempenho da TPU é criar o perfil do seu modelo. Para mais informações sobre como capturar um perfil de desempenho, consulte Como criar o perfil do seu 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.
O modelo tem limite de entrada.
As TPUs realizam cálculos muito rapidamente. Para garantir que a TPU não seja inativa, é importante garantir que haja um fluxo estável de dados sendo carregado na TPU. A maneira como isso é feito depende de como você carrega e pré-processa o conjunto de dados. Por exemplo, você pode ler arquivos de dados em paralelo usando tf.data.TFRecordset() e o parâmetro
num_parallel_reads
.O tamanho do lote é muito pequeno devido à fragmentação (divisão de lotes entre núcleos)
O ambiente 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 receberá um tamanho de lote de 16 (128 / 8).
Para otimizar o uso da memória, use o maior tamanho de lote que caiba 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 precisa ser igualmente divisível por 8 ou 128.
Otimizações do compilador XLA
O XLA é um compilador de 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 XLA que o XLA compila para um executável da TPU. Para mais informações sobre XLA, consulte XLA: como otimizar o compilador para machine learning.
Preenchimento
Para usar a memória da TPU de maneira eficiente, estruture seus dados para que eles possam ser colocados em blocos de 128 x 8. Quando os dados de um cálculo de matriz não preenchem um bloco 128 x 8 inteiro, o compilador XLA preenche os tensores. Existem duas desvantagens no preenchimento:
- Os tensores acolchoados subutilizam o núcleo da TPU.
- O padding aumenta a quantidade de armazenamento de memória no chip necessária para um tensor e pode levar a um erro de memória insuficiente.
Embora o preenchimento seja executado automaticamente pelo compilador XLA quando necessário, é possível determinar a quantidade de preenchimento usando a ferramenta visualizador de memória. É possível evitar o preenchimento escolhendo dimensões de tensores que sejam adequadas para a TPU.
Dimensões do tensor
O compilador XLA arredonda os tamanhos dos tensores armazenados na memória HBM da TPU para executar computações com mais eficiência. Esse preenchimento acontece de forma 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 exibe 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 opções precisa ser verdadeira:
O tamanho total do lote precisa ser um múltiplo de 64 (oito por núcleo de TPU), e os tamanhos de dimensão do recurso precisam ser um múltiplo de 128.
O tamanho total do lote precisa ser um múltiplo de 1024 (128 por núcleo de TPU), e os tamanhos de dimensão do recurso precisam ser um múltiplo de 8.
O uso de um tamanho de lote de 1.024 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.
Fusão
A fusão é uma técnica geral que o compilador XLA usa para otimizar programas. Uma operação combinada é a combinação de várias operações constituintes que serão executadas em combinação.
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 na 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 ser fundidas com os consumidores, forçar uma transmissão pode resultar em desempenho ruim 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)`