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 precisão convergente. Ao usar bfloat16, muitos modelos alcançam resultados com a mesma precisão convergente oriunda de números de ponto flutuante de 32 bits, e alguns deles apresentam uma precisão 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. Isso significa que o bfloat16 tem um intervalo dinâmico maior que o fp16. Isso é ú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, você pode usar o formato bfloat16 para representar com precisão todos os números inteiros [-256, 256]. Isso permite codificar um int8 em bfloat16 sem perda de precisão.

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

imagem

Observe que o intervalo dinâmico de bfloat16 é maior que o de fp16.

Vantagens de desempenho e uso de memória

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

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 em nossos experimentos internos.

imagem

Como alterar seu 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. Use o visualizador de perfil no TensorBoard para ver o tempo da etapa e a precisão convergente do modelo. Consulte Como usar as ferramentas do Cloud TPU para ver mais detalhes.

  3. Transmita a entrada para bfloat16 no seu pipeline de entrada dentro do analisador de registros para que a conversão possa ser feita em paralelo, não no final de input_fn. Isso converte todas as ativações e gradientes no modelo para bfloat16.

    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. Use o visualizador de perfil no TensorBoard para ver o tempo da etapa e a precisão convergente do modelo. Consulte Como usar as ferramentas do Cloud TPU para ver mais detalhes.
  3. Compare o tempo da etapa de bfloat16 e fp32. O tempo da etapa normalmente melhora com o uso do bfloat16.
  4. Compare a precisão convergente de bfloat16 e fp32. Normalmente ela é idêntica, mas pode ser melhor ou pior.

Poderão ser necessárias várias execuções para medir cada variação na precisão convergente, caso você ainda não saiba qual intervalo de variação esperar do seu modelo.

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

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.

Esta página foi útil? Conte sua opinião sobre:

Enviar comentários sobre…