Guia de desempenho da Cloud TPU
A primeira etapa ao solucionar problemas de desempenho da TPU é criar o perfil do modelo. Para mais informações sobre como capturar um perfil de desempenho, consulte Como criar um perfil de desempenho 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.
O modelo está vinculado à entrada
As TPUs realizam cálculos muito rapidamente. Para garantir que o TPU não fique ocioso, é importante ter um fluxo constante de dados sendo carregados no TPU. A maneira como isso é feito 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.O tamanho do lote é muito pequeno devido à fragmentação (dividir lotes entre núcleos)
O ambiente de execução da TPU divide um lote em todos os 8 núcleos de um dispositivo 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 matriz. 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 para o Cloud TPU são traduzidos em um gráfico XLA, que é compilado pelo XLA em um executável de TPU. Para saber mais sobre o XLA, consulte XLA: como otimizar o compilador para machine learning.
Preenchimento
Para usar a memória de TPU de forma eficiente, estruture seus dados de forma que eles possam ser agrupados em 128 x 8 blocos. Quando os dados para o cálculo de uma matriz não preenchem um 128 x 8, o compilador XLA preenche os tensores. O preenchimento apresenta duas desvantagens:
- Os tensores preenchidos subutilizam o núcleo da TPU.
- O preenchimento aumenta a quantidade de armazenamento em memória on-chip necessária para um tensor e pode resultar em um erro de falta de memória.
Embora o preenchimento seja realizado automaticamente pelo compilador XLA quando necessário, pode determinar a quantidade de padding realizada usando a ferramenta de visualização de memória. Para evitar o preenchimento, escolha dimensões de tensor adequadas para TPU.
Dimensões do tensor
O compilador XLA arredonda os tamanhos dos tensores armazenados na memória HBM da TPU para realizar cálculos com mais eficiência. Esse preenchimento acontece de modo 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 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:
o tamanho total do lote seja um múltiplo de 64 (8 por núcleo de TPU) e os tamanhos das dimensões do recurso sejam um múltiplo de 128.
o tamanho total do lote deve ser um múltiplo de 1024 (128 por núcleo de TPU); e os tamanhos das dimensões de recursos devem ser múltiplos de 8.
Usando um tamanho de lote de 1024 e dimensões de recursos que sejam um múltiplo de 128, você garante a melhor eficiência, mesmo que isso não seja possível em todos os modelos.
Fusão
A Fusion é uma técnica geral que o compilador XLA usa para otimizar programas. Um A operação fundida é a combinação de várias operações constituintes sejam executados 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 viagens de ida e volta da 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 consome, mas forçar 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)`