Mixed precision solvers with half-precision floating point numbers for Lattice QCD on A64FX processor

Este estudio demuestra que el uso de números de coma flotante de media precisión (FP16) en solucionadores de precisión mixta para QCD de red en el procesador A64FX es práctico y estable, gracias a la implementación de pasos de reescalado que evitan la inestabilidad numérica y mantienen el aumento en el número de iteraciones por debajo del 20% en comparación con la versión de doble precisión.

Lo esencial

¡Hola! Imagina que estás intentando resolver un rompecabezas gigante y extremadamente complejo. Este rompecabezas representa el universo de las partículas subatómicas (como los quarks y los gluones) y se llama Lattice QCD (Cromodinámica Cuántica en Red).

El problema es que este rompecabezas es tan enorme que, para resolverlo, necesitas una computadora súper potente. Pero aquí está el truco: las computadoras actuales a veces se vuelven lentas porque intentan hacer los cálculos con una precisión de "microscopio" (números de 64 dígitos), lo cual es como usar un telescopio para leer una etiqueta de precio en una tienda. Es demasiado detalle y gasta mucha energía y tiempo.

Los científicos de este documento (del RIKEN y otras instituciones japonesas) se preguntaron: ¿Qué pasa si usamos una precisión más simple, como "lupa" (números de 16 dígitos), para la mayoría del trabajo?

Aquí te explico lo que hicieron, usando analogías sencillas:

1. El Problema: El "Hielo" en la Computadora

Imagina que tienes un equipo de trabajadores (la computadora) intentando mover cajas muy pequeñas.

• Precisión Alta (FP64): Usan guantes de invierno muy gruesos. Pueden manejar cualquier cosa, pero son lentos y pesados.
• Precisión Baja (FP16): Usan guantes finos. Son muy rápidos y pueden mover muchas cajas a la vez. ¡Pero hay un problema! Si las cajas son demasiado pequeñas (números muy cercanos a cero), los guantes finos no las sienten. Las cajas se "pierden" o se convierten en cero. A esto los científicos lo llaman desbordamiento inferior (underflow).

En sus pruebas iniciales, cuando intentaron usar los "guantes finos" (FP16) directamente, el sistema se rompió porque perdieron información crítica de las cajas pequeñas. El cálculo se volvía inestable y fallaba.

2. La Solución: El "Ajuste de Volumen" (Rescaling)

Para solucionar esto, los autores (Kanamori y su equipo) idearon una estrategia inteligente llamada rescalado (o ajuste de escala).

Imagina que tienes un grifo que gotea muy poco (un número pequeño). Si intentas medirlo con un balde gigante, no verás nada. Pero, ¿qué pasa si pones un embudo pequeño debajo del grifo para concentrar el agua? ¡Ahora sí puedes medirla!

• El truco: Antes de que los números se vuelvan demasiado pequeños para que la computadora de "guantes finos" los vea, el algoritmo amplifica (multiplica) temporalmente todos los números.
• El proceso:
  1. Amplificar: Toman el problema y lo hacen "más grande" para que los números pequeños no se pierdan.
  2. Resolver rápido: Usan la computadora rápida (FP16) para hacer los cálculos pesados.
  3. Corregir: Al final, vuelven a ajustar la escala para obtener la respuesta exacta.

Además, usaron un sistema de doble equipo:

• Un equipo rápido (FP16) hace el trabajo sucio y rápido.
• Un equipo de supervisión (FP64, el de "guantes gruesos") revisa el trabajo y corrige los pequeños errores que el equipo rápido pudo haber cometido.

3. El Resultado: ¡Velocidad de Rayo!

Gracias a este método de "amplificar y corregir", lograron algo increíble en la supercomputadora Fugaku (una de las más potentes del mundo, ubicada en Japón):

• Velocidad: Su nuevo método fue dos veces más rápido que el método anterior que usaba precisión media (FP32) y tres veces más rápido que el método tradicional de alta precisión (FP64).
• Estabilidad: El sistema ya no se rompía por los números pequeños.
• Eficiencia: Aunque tuvieron que hacer un par de pasos extra para corregir, el ahorro de tiempo fue enorme.

En Resumen

Imagina que quieres pintar un mural gigante.

• Antes: Usabas pinceles finos y lentos para asegurar que cada detalle fuera perfecto. Tardabas días.
• Ahora: Usas rodillos grandes y rápidos (FP16) para pintar el 90% del mural. Si ves que falta un poco de detalle, usas un pincel fino (FP64) solo para retocar esos puntos.

¿Por qué importa esto?
Esto es vital para el futuro de la ciencia. Las nuevas computadoras (como las que usarán en el próximo superordenador japonés "Fugaku NEXT") están diseñadas para ser muy rápidas con estos "guantes finos" (FP16), especialmente para Inteligencia Artificial. Este trabajo demuestra que podemos usar esa velocidad para resolver los problemas más difíciles de la física, como entender de qué está hecho el universo, sin sacrificar la precisión final.

¡Es como aprender a correr más rápido sin tropezar!

Resumen técnico

Título: Solutores de precisión mixta con números de punto flotante de media precisión (FP16) para QCD de red en el procesador A64FX

1. Planteamiento del Problema

Las simulaciones de Cromodinámica Cuántica de Red (QCD de red) requieren una gran cantidad de operaciones de punto flotante, tradicionalmente realizadas con precisión doble (FP64) para garantizar la estabilidad numérica. Sin embargo, la demanda de aplicaciones de Inteligencia Artificial ha impulsado el desarrollo de hardware optimizado para operaciones de baja precisión, específicamente media precisión (FP16).

El procesador A64FX (utilizado en el superordenador Fugaku) soporta la extensión vectorial escalable (SVE) de ARM, que permite realizar operaciones FP16 con un rendimiento cuatro veces superior al FP64. El desafío principal es que los algoritmos de solucionadores lineales iterativos estándar, diseñados para FP32 o FP64, fallan o se vuelven inestables cuando se aplican directamente con FP16 debido a:

• Desbordamiento (Underflow): Los vectores de residuo disminuyen rápidamente durante la iteración, cayendo por debajo del rango mínimo representable en FP16 (aprox. $6.10 \times 10^{-5}$), lo que provoca que los valores se conviertan en cero.
• Inestabilidad numérica: La falta de precisión y el rango dinámico limitado de FP16 degradan la convergencia de algoritmos como BiCGStab.

2. Metodología

Los autores investigan la viabilidad de utilizar FP16 en solucionadores de precisión mixta para la matriz de fermiones de Wilson (la formulación más simple de la acción de quarks discretizada). Su enfoque se basa en dos pilares:

• Algoritmo de Refinamiento Iterativo con Reescalado:

  • Se propone un esquema donde el sistema lineal se resuelve en el exterior con FP64 (precisión completa), pero el precondicionador interno (el solucionador más costoso) se ejecuta en FP16.
  • Para evitar el underflow, se introduce un factor de reescalado ($s$) en el paso de refinamiento iterativo (algoritmo de Richardson). Esto normaliza el vector de residuo antes de convertirlo a FP16, asegurando que los valores se mantengan dentro del rango dinámico útil.
  • Se recalcula este factor de escala en cada iteración utilizando aritmética FP32/FP64 para mantener la precisión global.

• Estabilización del Solucionador BiCGStab:

  • Se adapta el algoritmo BiCGStab (estándar para matrices no hermitianas complejas) para incluir reescalado dinámico.
  • Se introducen factores de reescalado ($\gamma$ y $\lambda$) para los vectores de residuo y de solución, respectivamente.
    • $\gamma$: Evita el underflow en los vectores de residuo.
    • $\lambda$: Evita el desbordamiento (overflow) en el vector de solución si el sistema tiene autovalores muy pequeños.
  • Se implementa una recálculo de los factores de escala basándose en la relación de las normas de los vectores antes y después del reescalado.

• Implementación en Hardware (A64FX):

  • Se utiliza el tipo de datos _Float16 (estándar ISO/IEC) en lugar de __fp16 para aprovechar las instrucciones aritméticas nativas de FP16 del procesador A64FX.
  • Se emplea la extensión vectorial escalable (SVE) con un tamaño de "tile" (baldosa) de $8 \times 4$ para FP16, optimizando el uso de los 512 bits del vector SIMD.
  • Las sumas parciales se realizan en FP32 dentro de los hilos OpenMP y la reducción global en FP64 para mantener la precisión en la acumulación de normas.

3. Contribuciones Clave

• Estabilidad Numérica en FP16: Demostración de que el uso directo de FP16 en QCD de red es inestable, pero que la introducción de pasos de reescalado en el refinamiento iterativo y en el solucionador interno (BiCGStab) restaura la estabilidad y la convergencia.
• Algoritmo Híbrido Propuesto: Desarrollo de una variante de BiCGStab que combina reescalado de residuo y solución, junto con la recálculo dinámico de factores de escala, superando las limitaciones de las bibliotecas existentes que no soportan FP16 de manera robusta en este contexto.
• Optimización para A64FX: Implementación eficiente utilizando las instrucciones SVE de ARM, logrando un rendimiento máximo en operaciones de punto flotante de media precisión.

4. Resultados

Los experimentos se realizaron en el superordenador Fugaku (RIKEN) utilizando una red de $32^3 \times 64$ y un parámetro de salto $\kappa = 0.13$.

• Convergencia:
  • El solucionador FP16 sin reescalado falló, mostrando una convergencia extremadamente lenta y estancamiento (stagflation) debido a errores de redondeo y underflow.
  • Con la metodología propuesta (reescalado), el solucionador FP16 convergió exitosamente. El número de iteraciones adicionales en comparación con la versión FP64 se mantuvo dentro de un 20% (ej. ~850-920 multiplicaciones matriz-vector en FP16 frente a 775 en FP64).
• Rendimiento y Tiempo de Ejecución:
  • Velocidad de Punto Flotante: Se alcanzaron 8249 GFlops en FP16, comparado con 3895 GFlops en FP32 y 2045 GFlops en FP64.
  • Tiempo de Ejecución: El solucionador FP16 fue aproximadamente 2 veces más rápido que el de FP32 y 3 veces más rápido que el de FP64 para resolver el sistema completo.
  • El tiempo de ejecución total se redujo de 1.39 segundos (FP64) a 0.46 segundos (FP16 óptimo).
• Análisis de Underflow: Se observó que sin reescalado, la mayoría de los elementos del vector de entrada se convertían en cero rápidamente. Con los factores de escala adecuados (ej. $s=128$ o $4096$), se eliminó casi por completo el underflow, permitiendo que la información del vector de entrada se propague eficientemente.

5. Significado y Conclusiones

El trabajo demuestra que es práctico y altamente eficiente utilizar aritmética de media precisión (FP16) en simulaciones de QCD de red de gran escala en arquitecturas modernas como A64FX, siempre que se implementen técnicas de estabilización adecuadas.

• Impacto: Esto permite acelerar significativamente las simulaciones de física de partículas en el superordenador Fugaku y prepara el terreno para futuros sistemas (como Fugaku NEXT) que integrarán aceleradores GPU con núcleos Tensor Core optimizados para FP16.
• Generalización: Aunque el estudio se centró en la matriz de fermiones de Wilson, los autores sugieren que los métodos de reescalado propuestos son aplicables a formulaciones más complejas (como fermiones clover o de pared de dominio) y a otros solucionadores iterativos en diversas aplicaciones científicas.
• Futuro: Se identifica la necesidad de extender estos métodos a matrices más complejas y evaluar el rendimiento utilizando el formato BF16 (Bfloat16), que ofrece un rango dinámico mayor que FP16, aunque con menos bits de fracción.