UC Berkeley crea mKernel: una biblioteca única para sincronizar GPU en clústeres

El equipo UCCL de UC Berkeley ha lanzado mKernel, una biblioteca CUDA que pretende revolucionar la sincronización de GPU en clústeres grandes. La biblioteca reúne en un único persistent kernel la comunicación local entre GPU, la sincronización entre servidores y el cálculo, y evita así los cambios de contexto constantes.

El cuello de botella en los clústeres a gran escala

En los centros de datos a gran escala, la sincronización entre GPU es uno de los principales cuellos de botella de ancho de banda. Cuando los ingenieros entrenan modelos grandes de forma distribuida, se topan con un problema que tradicionalmente se resolvía de manera muy ineficiente.

El enfoque anterior era por capas: los ingenieros usaban tres herramientas y bibliotecas distintas. La primera servía para la comunicación rápida dentro de un mismo servidor (NVLink, que funciona mediante enlaces de alta velocidad entre GPU). La segunda, para sincronizar servidores distintos a través de la red (RDMA, que usa adaptadores de red especializados). La tercera, para el cálculo en sí.

Cada paso entre estos tres sistemas congela todo el pipeline de la GPU: tiene que detenerse, cambiar de contexto, descargar parte de la memoria, cargar datos nuevos y solo entonces continuar. En clústeres con miles de GPU, estos retrasos de microsegundos se acumulan hasta convertirse en minutos de rendimiento perdido.

Cómo resuelve mKernel el problema

mKernel cambia la filosofía de raíz. En lugar de tres sistemas separados, todas las operaciones (comunicación local, sincronización de red y cálculo) se ejecutan en un único persistent kernel. Es un programa que reside de forma permanente en la memoria de la GPU y nunca se descarga a la CPU.

La arquitectura combina tres componentes:

• NVLink para la comunicación entre GPU de un mismo servidor — ofrece una velocidad 10-20x superior a PCIe gracias a canales directos de alta velocidad entre GPU
• RDMA mediante adaptadores de red para sincronizar servidores — evita la sobrecarga de la CPU y la latencia del sistema operativo al enviar datos entre máquinas
• Cálculo denso integrado directamente en el kernel — la GPU trabaja con datos locales sin cambios de contexto, se sincroniza con sus vecinas y pasa de inmediato a la siguiente tarea

Esto significa que una GPU puede pasar sin obstáculos de la comunicación local por NVLink a la sincronización global por RDMA y a su propio cálculo, todo en un único fragmento de código y sin paradas.

Un ejemplo concreto: cómo funciona en la práctica

En el entrenamiento distribuido, una GPU no puede avanzar hasta que las demás GPU de otros servidores terminan sus cálculos y sincronizan los gradientes. Con el enfoque antiguo, la GPU simplemente espera con el pipeline totalmente vacío.

Con mKernel el proceso es distinto: la GPU sigue calculando con los datos ya cargados, se sincroniza al mismo tiempo con sus vecinas mediante NVLink y RDMA, y pasa de inmediato a la siguiente oleada de entrenamiento sin pausas. Es como la cadena de montaje de una fábrica de automóviles, que no se detiene mientras la pieza viaja a la siguiente estación.

Por qué es crítico para los centros de datos

El entrenamiento distribuido de modelos grandes es una de las tareas computacionales más complejas del desarrollo actual de IA. Cuando se usan 1000 GPU a la vez (y las grandes empresas trabajan con más), incluso una pequeña ralentización de la sincronización puede desperdiciar entre el 20 % y el 30 % de todos los recursos.

mKernel promete eliminar por completo esa sobrecarga.

En las primeras pruebas, los investigadores ya ven resultados: en clústeres multinodo, la biblioteca acelera la sincronización entre 2 y 3 veces en operaciones típicas. Esto es especialmente importante para el mecanismo de atención (attention) de los transformadores, donde la sincronización de gradientes entre GPU es la parte más costosa de todo el trabajo.

Qué significa

mKernel es una señal de que la programación de GPU entra en la era de los sistemas integrados. Antes, los ingenieros escribían el código por capas: primero el cálculo, luego la sincronización y después la transferencia. Ahora la frontera entre ellas se difumina. Eso implica un cálculo más rápido en los centros de datos y un entrenamiento de modelos grandes más accesible y veloz. Y, sobre todo, la próxima generación de sistemas distribuidos se diseñará de una forma completamente nueva.