Mejorar el rendimiento de un modelo con bfloat16

De forma predeterminada, las TPUs realizan operaciones de multiplicación de matrices con valores bfloat16 y acumulaciones con valores IEEE float32. Si se usan números de punto flotante de precisión reducida, se reduce el tiempo de convergencia sin perder precisión.

El intervalo dinámico de bfloat16 y float32 son equivalentes. Sin embargo, bfloat16 usa la mitad del espacio de memoria. Para obtener más información sobre el bfloat16 rendimiento, consulta A Study of BFLOAT16 for Deep Learning Training (Estudio de BFLOAT16 para el entrenamiento de aprendizaje profundo).

Usar bfloat16 de forma explícita

Aunque la conversión automática de formatos en las TPUs te permite no tener que pensar en la precisión numérica, puedes mejorar el rendimiento si conviertes los valores explícitamente a bfloat16. Hay dos motivos para convertir valores explícitamente a bfloat16:

1. Almacenar valores en formato bfloat16 se ahorra memoria en el chip, lo que permite que las TPUs de Cloud entrenen modelos más grandes o usen tamaños de lote mayores.

2. Algunas operaciones están limitadas por el ancho de banda de la memoria, lo que significa que el tiempo que se tarda en cargar datos de la memoria puede ralentizar el tiempo total dedicado a realizar el cálculo. Almacenar los operandos y los resultados de esas operaciones en formato bfloat16 se reduce la cantidad de datos que se deben transferir, lo que mejora la velocidad general.

Para empezar, te recomendamos que adquieras experiencia con uno de los modelos de referencia de TPU de Cloud. Después, la guía de herramientas de creación de perfiles y la guía de solución de problemas proporcionan información técnica detallada para ayudarte a crear y optimizar modelos de aprendizaje automático por tu cuenta.

Detalles de la conversión de formato

El compilador de XLA inserta automáticamente la conversión de formato de float32 a bfloat16. En las TPU, el esquema de redondeo de la conversión es redondeo al entero par más cercano y desbordamiento a inf. Además, la bfloat16 en las TPUs de Cloud no admite números subnormales, por lo que todos los números subnormales se convierten en cero durante la conversión. Los valores especiales, como NaN y inf, se conservan en la conversión.

El compilador de XLA también inserta automáticamente la conversión de formato de bfloat16 a float32. Como float32 puede representar todos los valores exactos de bfloat16, la conversión añade 16 ceros en los bits de la mantisa. Los valores especiales se conservan en la conversión.

Los puntos de control obtenidos de un modelo entrenado en TPUs de Cloud se pueden implementar en otras plataformas de hardware (por ejemplo, inferencia o ajuste fino en CPUs o GPUs) sin necesidad de realizar conversiones manuales exhaustivas.