Simular el comportamiento de moléculas y materiales a escala atómica es una de las tareas computacionalmente más exigentes en ciencia de materiales y química computacional. Los métodos de mecánica cuántica, como la teoría del funcional de la densidad (DFT), ofrecen alta precisión, pero su costo computacional los hace inviables para sistemas con miles de átomos. Los potenciales interatómicos clásicos resuelven el problema de escala, pero sacrifican exactitud al depender de formas funcionales fijas que no capturan la complejidad del enlace químico.

En ese contexto, los potenciales interatómicos de aprendizaje automático (MLIP, por sus siglas en inglés) han ganado terreno como alternativa que combina la velocidad de los métodos clásicos con la precisión de los cuánticos, entrenando redes neuronales sobre datos de referencia ab initio. Sin embargo, su implementación enfrenta un desafío técnico crítico: pequeños errores numéricos en el cálculo de fuerzas pueden acumularse y comprometer la estabilidad de las trayectorias moleculares a lo largo de simulaciones prolongadas.

Tensor Cores como palanca de rendimiento

Las GPU modernas de NVIDIA incorporan unidades especializadas —los Tensor Cores— diseñadas para ejecutar multiplicaciones matriciales en formatos de menor precisión numérica con un rendimiento significativamente superior al de las operaciones FP32 estándar. Dos estrategias distintas fueron evaluadas sobre el modelo Point Edge Transformer (PET), un MLIP basado en arquitectura transformer desarrollado por el Lab COSMO de la EPFL.

Sobre la arquitectura Blackwell, el método BF16x9 descompone cada operando FP32 en tres componentes BF16 y combina nueve productos parciales para reconstruir el resultado con precisión equivalente a FP32. Su activación a través de cuBLAS no requiere modificaciones en el modelo. Sobre la arquitectura Ampere, el método TF32x3 utiliza un componente TF32 principal y uno residual, calculando tres productos parciales mediante la biblioteca CUTLASS.

Resultados cuantitativos en sistemas reales

Las pruebas se realizaron con el modelo PET-MAD pequeño (26 millones de parámetros) en sistemas de agua de 192 y 17,280 átomos, y con un modelo grande de 255 millones de parámetros en 192 átomos. Los resultados son concretos: en la GPU Blackwell, BF16x9 aceleró las multiplicaciones matriciales del modelo entre 2.2 y 2.5 veces en el sistema de mayor escala. A nivel de dinámica molecular completa, la estrategia cuBLAS logró aceleraciones de 1.2 veces en el sistema pequeño y 1.7 veces en el de 17,280 átomos, manteniendo el error por debajo del umbral de no determinismo propio del modelo.

En la GPU A100 con TF32x3, las multiplicaciones del modelo se ejecutaron entre 1.7 y 1.9 veces más rápido. El modelo pequeño alcanzó 1.5 veces de aceleración en el sistema grande, con un error de fuerza equivalente a tres veces la diferencia entre ejecuciones FP32 consecutivas —un margen considerado aceptable para la mayoría de aplicaciones de simulación. El modelo grande obtuvo 1.4 veces de aceleración en 192 átomos, con un error de 3.9 × 10⁻⁶ eV/Å.

El método TF32 de una sola pasada mostró mayor velocidad, pero introdujo degradación en precisión y problemas de estabilidad en trayectorias largas, lo que lo descalifica para simulaciones de producción donde la conservación de energía es un requisito.

Implicaciones para la escala y la arquitectura de cómputo científico

Los resultados establecen que es posible obtener aceleraciones sustanciales en MLIP sin rediseñar los modelos ni reentrenar desde cero, simplemente aprovechando las capacidades nativas de las GPU actuales mediante estrategias de precisión mixta controlada. Esto tiene consecuencias directas sobre el dimensionamiento de infraestructura: simulaciones que hoy requieren clústeres de múltiples nodos podrían ejecutarse en configuraciones más compactas, reduciendo costos operativos y tiempos de iteración en proyectos de descubrimiento de materiales o diseño molecular.

La dependencia de los resultados respecto al tamaño del sistema y la proporción de tiempo dedicado a multiplicaciones matriciales indica que la ganancia no es uniforme: sistemas más grandes y modelos con mayor densidad de operaciones lineales son los que extraen mayor beneficio. Esa variable debe considerarse al evaluar si la adopción de estos métodos justifica el esfuerzo de integración en pipelines de simulación existentes.