Guia de desempenho da Cloud TPU
A primeira etapa para 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 perfis 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 está vinculado à entrada
As TPUs executam 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 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.O tamanho do lote é muito pequeno devido à fragmentação (dividindo lotes entre núcleos)
O ambiente de execução da TPU divide um lote entre 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 precisa ser 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 a XLA faça parte da base de código padrão do TensorFlow, ela também pode ser usada nos 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 mais informações sobre o 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 divididos em blocos de 128 x 8. Quando os dados de uma computação de matriz não preenchem um bloco de 128 x 8, o compilador XLA preenche os tensores. O preenchimento apresenta duas desvantagens:
- Os tensores preenchidos utilizam pouco 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.
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 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 múltiplos de 128.
o tamanho total do lote seja um múltiplo de 1024 (128 por núcleo de TPU) e os tamanhos das dimensões do recurso sejam múltiplos de 8.
Usando um tamanho de lote de 1024 e dimensões de recurso 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. 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 = 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 à 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)`