Las simulaciones de dinámica molecular (MD) están entre las cargas de trabajo más exigentes de la ciencia computacional. Con ellas, los investigadores observan el comportamiento de los átomos con un detalle extraordinario, desde el plegamiento de proteínas hasta el descubrimiento de fármacos y materiales.
Pero tienen un costo altísimo. Estas simulaciones modelan desde cientos de miles hasta decenas de millones de átomos, avanzando femtosegundos por paso a lo largo de miles de millones de pasos. Como el tamaño del problema suele ser fijo, la paralelización debe repartir el trabajo entre todos los recursos disponibles a la vez, un régimen conocido como strong scaling.
El cuello de botella no era la GPU, era la CPU
GROMACS es uno de los paquetes de MD más usados del mundo. El hardware moderno y la paralelización heterogénea CPU-GPU han llevado el rendimiento al régimen de submilisegundos, alcanzando 100 a 200 microsegundos por paso de tiempo repartidos entre varias GPUs. Lograr strong scaling a ese nivel exige una latencia de kernel bajísima, pero la comunicación entre GPUs sigue siendo el freno fundamental a medida que la simulación crece.
La mayoría de las aplicaciones de HPC a gran escala, GROMACS incluido, usan la interfaz de paso de mensajes (MPI) para comunicarse entre procesos. El problema es que MPI fue diseñado para una ejecución centrada en la CPU. Cuando GROMACS corre sobre GPUs, el flujo de comunicación obliga a la GPU a detenerse mientras la CPU orquesta las transferencias de datos y luego le avisa a la GPU que puede seguir. En el halo exchange de GROMACS, el algoritmo que comparte los datos de átomos de frontera entre dominios vecinos, ese traspaso se repite en las tres dimensiones espaciales y consume más del 50% del tiempo total de CPU en los picos de iteración.
¿Cómo se saca a la CPU del camino crítico?
La clave está en eliminar los traspasos CPU-GPU con comunicación nativa entre GPUs. NVIDIA NVSHMEM es una librería de acceso remoto a memoria basada en el modelo de espacio de direcciones global particionado de OpenSHMEM. NVSHMEM permite que los propios kernels de la GPU inicien las transferencias de datos, retirando a la CPU de la ruta crítica y habilitando un mejor solapamiento entre cómputo y comunicación.
El halo exchange de GROMACS usa un mecanismo de reenvío por etapas: los datos se retransmiten a través de rangos intermedios dentro de cada dimensión, en uno o más pasos llamados pulses. La reformulación reemplaza las barreras gruesas por señales finas por cada pulse, expresadas directamente en un kernel fusionado.
Del MPI tradicional al intercambio nativo en GPU
El intercambio basado en MPI exige dos sincronizaciones CPU-GPU bloqueantes por cada pulse: una antes de invocar MPI, para garantizar que el kernel de empaquetado terminó, y otra antes de que el siguiente pulse pueda consumir los datos reenviados. Con una descomposición en 3D y un pulse por dimensión, eso equivale a seis esperas bloqueantes por paso para las coordenadas y seis más para las fuerzas: 12 en total. Cada una suma latencia a la ruta crítica y golpea la tasa de iteración de GROMACS.
El nuevo diseño funde el empaquetado, el envío remoto y la espera de finalización en un solo kernel por cada pulse. La CPU sigue emitiendo los lanzamientos en orden, pero desaparecen las barreras CPU-GPU intermedias: el orden correcto lo garantiza el propio stream de CUDA. Gracias a esta fusión de kernels consciente de las dependencias, los lanzamientos de kernel por paso bajan de seis a uno.
Resultados: hasta 2x de rendimiento
La lógica de transporte también es consciente de la interconexión: enruta las escrituras directas por NVLink cuando los pares están conectados por ese enlace, y usa operaciones put de NVSHMEM sobre fabricas RDMA (como InfiniBand) para el resto. El enfoque además aprovecha el motor TMA de la arquitectura NVIDIA Hopper para las transferencias remotas masivas.
Las pruebas en el superordenador NVIDIA Eos y en clústeres NVIDIA GB200 NVL72 mostraron hasta 2x de mejora en el escalado, tanto dentro de un nodo como entre nodos, frente a MPI con soporte para GPU, sobre todo en sistemas limitados por latencia. El método se generaliza a cualquier aplicación de HPC con patrones de halo exchange, aunque estandarizar las primitivas de comunicación nativas de GPU sigue siendo un desafío abierto.
