Melhorar o desempenho do modelo com bfloat16

Por padrão, as TPUs realizam operações de multiplicação de matriz com valores bfloat16 e acumulações com valores float32 do IEEE. O uso de números de ponto flutuante com precisão reduzida diminui o tempo de convergência sem perda de acurácia.

Os intervalos dinâmicos de bfloat16 e float32 são equivalentes. No entanto, bfloat16 usa metade do espaço de memória. Para mais informações sobre o desempenho de bfloat16, consulte Um estudo do BFLOAT16 para treinamento de aprendizado profundo.

Usar explicitamente bfloat16

Embora a conversão automática de formato em TPUs ajude a evitar considerações de precisão numérica, é possível melhorar o desempenho transmitindo valores explicitamente para bfloat16. Há dois motivos para converter valores explicitamente em bfloat16:

1. Como armazenar valores no formato bfloat16 economiza memória no chip, os Cloud TPUs podem treinar modelos maiores ou usar tamanhos de lote maiores.

2. Algumas operações são vinculadas à largura de banda de memória, o que significa que o tempo necessário para carregar dados da memória pode diminuir o tempo total gasto na execução do cálculo. Armazenar operandos e saídas dessas operações no formato bfloat16 reduz a quantidade de dados que precisam ser transferidos, o que melhora a velocidade geral.

Para começar, recomendamos que você tenha alguma experiência com um dos modelos de referência do Cloud TPU. Além disso, o guia de ferramentas de criação de perfil e o guia de solução de problemas vão fornecer informações técnicas detalhadas para ajudar você a criar e otimizar modelos de machine learning por conta própria.

Detalhes sobre a conversão de formato

A conversão do formato float32 para bfloat16 é inserida automaticamente pelo compilador XLA. Na TPU, o esquema de arredondamento da conversão consiste em arredondar para o número par mais próximo, resultando em inf em caso de estouro. Além disso, como o bfloat16 no Cloud TPU não aceita subnormais, todos eles são zerados durante a conversão. Valores especiais, como NaN e inf, são preservados durante a conversão.

A conversão de formato de bfloat16 para float32 também é inserida automaticamente pelo compilador XLA. Como float32 pode representar todos os valores exatos em bfloat16, a conversão adiciona 16 zeros nos bits da parte fracionária. Os valores especiais são preservados na conversão.

Os checkpoints recebidos de um modelo treinado em Cloud TPUs podem ser implantados em outras plataformas de hardware, como inferência ou ajuste em CPUs ou GPUs, sem conversões manuais extensas.