Como usar o bfloat16 com modelos do TensorFlow

As pesquisas em machine learning (ML) mostram que muitos modelos de ML podem tolerar uma aritmética de precisão mais baixa sem a degradação da acurácia convergente. Ao usar bfloat16, muitos modelos alcançam resultados com a mesma acurácia convergente oriunda de números de ponto flutuante de 32 bits, e alguns deles apresentam uma acurácia até melhor.

Este documento refere-se ao treinamento de precisão mista no sentido de armazenar ativações e gradientes na memória usando o formato bfloat16. Para ver mais informações, consulte Treinamento de precisão mista. Este documento não aborda o uso de bfloat16 na unidade de matriz (MXU, na sigla em inglês) no Cloud TPU.

Os tópicos seguintes aplicam-se aos modelos de ML que usam o TensorFlow:

  • descrição do ponto flutuante cerebral de 16 bits personalizado do Google, bfloat16
  • vantagens de desempenho ao usar o bfloat16 na memória em modelos de ML em hardware compatível, como o Cloud TPU
  • como armazenar ativações e gradientes na memória usando bfloat16 em um modelo de TPU no TensorFlow

Por padrão, o TensorFlow armazena todas as variáveis em ponto flutuante de 32 bits (fp32). Usar o bfloat16 para ativações e gradientes acelera o tempo da etapa do dispositivo e diminui o uso de memória. Veja Como alterar seu modelo para determinar os benefícios de usar o bfloat16 para ativações e gradientes no seu modelo.

O formato de ponto flutuante bfloat16

O formato bfloat16 é [1:8:7], com um bit de sinal, oito bits exponenciais e sete bits de mantissa, além de um bit de mantissa implícito. Por comparação, o formato padrão de ponto flutuante de 16 bits (fp16) é [1: 5: 10]. Observe que o formato fp16 tem apenas 5 bits exponenciais. Por causa dessas características, o bfloat16 tem um intervalo dinâmico maior que o fp16. O intervalo do bfloat16 é útil para elementos como gradientes, que podem estar fora do intervalo dinâmico do fp16 e, assim, demandariam uma perda de dimensionamento. O formato bfloat16 pode representar esses gradientes diretamente. Além disso, é possível usar o formato bfloat16 para representar com acurácia todos os números inteiros [-256, 256]. Isso permite codificar um int8 em bfloat16 sem perda de acurácia.

A figura a seguir mostra três formatos de pontos flutuantes:

  • fp32: ponto flutuante IEEE de precisão simples
  • fp16: ponto flutuante IEEE de meia precisão
  • bfloat16: ponto flutuante cerebral de 16 bits

image

O intervalo dinâmico do bfloat16 é maior que o intervalo do fp16.

Vantagens de desempenho e uso de memória

A Cloud TPU é compatível com o armazenamento de valores, como ativações e gradientes, no formato bfloat16. O uso do bfloat16 reduz o tamanho dos dados na memória e permite que modelos maiores se encaixem na mesma quantidade de memória. O uso do bfloat16 também pode reduzir a rematerialização, o que melhora o tempo de duração do passo.

Algumas operações são ligadas à largura de banda da memória, o que significa que ela determina o tempo gasto nessas operações. Armazenar entradas e saídas de operações com limite de largura de banda de memória no formato bfloat16 reduz a quantidade de dados que precisam ser transferidos, melhorando a velocidade das operações.

A tabela a seguir mostra melhorias observadas nos experimentos internos.

image

Como alterar o modelo para usar bfloat16

Por padrão, ativações, gradientes e pesos são armazenados em fp32 na memória. Você pode usar bfloat16 para ativações e gradientes, deixando pesos em fp32, e então comparar o desempenho do seu modelo usando bfloat16 ou fp32 para determinar os benefícios.

  1. Execute o modelo em fp32 usando capture_tpu_profile.

  2. Para visualizar o tempo de duração do passo e a acurácia convergente do modelo, use o visualizador de perfil no TensorBoard. Para mais detalhes, consulte Como usar as ferramentas da Cloud TPU.

  3. Transmita a entrada para bfloat16 no pipeline de entrada dentro do analisador de registros para que a conversão possa ser feita em paralelo, e não no final de input_fn. Isso faz com que todos gradientes e ativações sejam convertidos no modelo para bfloat16.

    Por exemplo:

    image = tf.cast(image, tf.bfloat16)
    
  4. Crie sua rede sob o escopo bfloat16 e, em seguida, converta as saídas do modelo em float32.

Depois de configurar seu modelo para usar tf.bfloat16 para ativações, verifique o seguinte para ver o impacto de bfloat16 no seu modelo:

  1. Execute o modelo com bfloat16 usando capture_tpu_profile.
  2. Para visualizar o tempo de duração do passo e a acurácia convergente do modelo, use o visualizador de perfil no TensorBoard. Para mais detalhes, consulte Como usar as ferramentas da Cloud TPU.
  3. Compare o tempo da etapa de bfloat16 e fp32. O tempo da etapa normalmente melhora com o uso do bfloat16.
  4. Compare a acurácia convergente de bfloat16 e fp32. Normalmente, eles são idênticos, mas os valores podem ser melhores ou piores do que o esperado.

Se você ainda não sabe que intervalo de variação esperar do modelo, faça várias execuções para determinar a variação de execução para execução na acurácia convergente.

Se o perfil mostrar que o tempo de processamento é mais rápido, mas o pipeline de entrada se tornou um gargalo, otimize o pipeline de entrada para conseguir uma vantagem de velocidade ainda maior. Para orientações gerais sobre como melhorar o desempenho do pipeline do TensorFlow, consulte Desempenho do pipeline de entrada de dados (em inglês).

Uma prática recomendada para treinamento e inferência é usar a mesma precisão para ambos. É possível treinar usando fp32 para ativações e, em seguida, executar a inferência com bfloat16, ou vice-versa. Se você optar por tipos de precisão diferentes, verifique a acurácia convergente usando a precisão que foi usada para a inferência.