Cuando introdujimos los aceleradores Gaudí en las instancias EC2 DL1 de Amazon, nos enfrentamos a un desafío que amenazaba toda la implementación. Las cifras de rendimiento no sólo fueron decepcionantes; fueron desastrosos. Los modelos que requerían capacitación efectiva experimentaban hasta un 50% de degradación de su rendimiento al escalar a través de múltiples nodos. ¿El problema? Una topología de red que enrutaba todos los bytes de datos a través de la memoria del host, provocando un cuello de botella que socavaba todo lo que Gaudí fue diseñado para hacer.
Dirigí el esfuerzo de ingeniería para abordar este problema, que finalmente resultó en el desarrollo de lo que ahora llamamos Peer Direct. Es una característica que transformó la forma en que los aceleradores de Gaudí se comunican en entornos de nube, y su historia tiene algunas lecciones útiles sobre el entrenamiento de IA distribuida a escala.
El problema con las NIC de host
Gaudí fue diseñado con la NIC (tarjeta de interfaz de red) integrada directamente en el silicio. Cada chip tiene diez interfaces de red que pueden manejar 100 Gbps y admiten RDMA con RoCE v2, lo que permite que los dispositivos accedan a la memoria de los demás directamente sin necesidad de la CPU o Esta arquitectura es altamente eficiente para cargas de trabajo de entrenamiento de IA, donde operaciones colectivas como AllReduce necesitan acumular gradientes de docenas o cientos de dispositivos por iteración de entrenamiento.
Pero las implementaciones en la nube no siempre cumplen con arquitecturas perfectas. Cuando Amazon probó Gaudi para instancias DL1, tuvo que utilizar NIC de host ordinarias en lugar de la red integrada de Gaudi. Las razones eran pragmáticas: ahorro de costos y la logística de trabajar alrededor de la infraestructura del centro de datos existente para acomodar una nueva topología de red. Desde su perspectiva empresarial, aprovechar la infraestructura de red establecida tenía mucho sentido.
Desde el punto de vista del rendimiento, fue un desastre. En lugar de transferencias RDMA de igual a igual entre tarjetas Gaudí, toda la comunicación fue por el largo camino. Los datos tenían que ser duplicados desde la memoria de gran ancho de banda de Gaudí a la DRAM del host, procesados por la CPU del host, enviados a la NIC del host a través de TCP/IP, recibidos por el host lejano y duplicados nuevamente en la memoria del Gaudí lejano. Todos los saltos agregados causaron latencia, robaron ciclos de CPU y agregaron restricciones de ancho de banda que arruinaron por completo la escalabilidad del entrenamiento distribuido.
El déficit de rendimiento era tan grave que uno se preguntaba si algún día valdría la pena implementarlo. No se trataba de una optimización trivial; Era una amenaza existencial para todo el acuerdo con AWS.
Por qué el rendimiento es tan importante
Vale la pena saber por qué una pérdida de rendimiento del 50% es tan desastrosa en la vida de los modelos de entrenamiento, y especialmente en los modelos grandes como el GPT-5. Ahora se necesitan semanas o meses para entrenar modelos de lenguaje enormes, incluso en grupos descomunales. Si juega con modelos que tienen miles de millones o billones de parámetros, cada punto porcentual de rendimiento se traduce directamente en tiempo y dólares.
Consideremos la economía. Si se necesitan 30 días para entrenar un modelo en lugar de 15, no sólo estás esperando más; estás pagando por el doble de tiempo de cálculo. A escala de la nube, con cientos o miles de aceleradores en uso continuo, esto suma millones de dólares. Peor aún, reduce a la mitad la velocidad de iteración. En un mundo competitivo de IA donde las empresas compiten por desarrollar modelos mejorados, duplicar el número de pruebas en el mismo período de tiempo puede ser la distinción entre estar delante y detrás.
El costo ambiental también es crucial. Los modelos grandes requieren mucha electricidad para enseñar. Un mejor rendimiento significa menos tiempo de cálculo, lo que reduce a la mitad el consumo de energía y las emisiones de carbono. A medida que aumenta la presión sobre la industria de la IA para que reduzca su huella de carbono, las ganancias en eficiencia ya no son un lujo sino más bien una necesidad.
La solución que diseñamos, Peer Direct, ofrecía un rendimiento similar al de RDMA cuando el diseño de la red física no era adecuado para RDMA normal. Necesitábamos acceso directo a la memoria entre dispositivos Gaudí en diferentes sistemas sin atravesar la memoria del host, pero en NIC del host que no fueron diseñadas para esto en primer lugar.
El habilitador fue AWS Elastic Fabric Adapter, una interfaz de red de alto rendimiento para cargas de trabajo de HPC e IA en EC2. EFA proporciona comunicaciones de derivación del sistema operativo de baja latencia, generalmente con una latencia inferior a 10 microsegundos. EFA proporciona una semántica similar a RDMA utilizando libfabric, una biblioteca de comunicación en el espacio del usuario que proporciona una interfaz común entre varias tecnologías de red.
La tarea consistía en combinar libfabric con la Biblioteca de Comunicación Colectiva de La Habana, HCCL, que maneja todas las cargas de trabajo de capacitación distribuidas. HCCL se construyó partiendo del supuesto de RDMA nativo utilizando las NIC en chip de Gaudí. Necesitábamos crear un puente que permitiera a HCCL aprovechar libfabric de forma transparente para las comunicaciones sin comprometer sus garantías de rendimiento y su semántica de comunicación.
La solución necesitaba varios avances técnicos. Primero, introdujimos un sistema de registro de memoria que permitió a libfabric acceder directamente a la memoria de gran ancho de banda de Gaudí. Utilizamos el marco DMA-BUF del kernel de Linux, que proporciona un mecanismo compartido para compartir buffers de controladores de dispositivos. Cuando HCCL necesita transferir datos, el controlador Gaudi proporciona un descriptor de archivo DMA-BUF para la región de memoria, que libfabric puede utilizar para crear transferencias RDMA directamente desde la memoria del dispositivo.
En segundo lugar, incluimos un caché LRU para registros de memoria. El registro de la memoria es caro; Implica llamadas al kernel y operaciones de configuración que pueden causar una sobrecarga significativa. Al almacenar en caché la asignación de direcciones de memoria a sus identificadores libfabric, podríamos reutilizar registros en regiones de acceso activo, eliminando la mayor parte de la sobrecarga de registro del entrenamiento real.
El resultado fue un canal de comunicación que se parecía a esto: HCCL llama al contenedor OFI, que llama al identificador libfabric en caché para realizar una transferencia RDMA directamente desde la memoria Gaudi de origen a la memoria Gaudi de destino, sin que ninguna de las CPU sea llamada. El contenedor OFI se introdujo para mantener limpia la base del código y evitar inclusiones directas de encabezados: es una biblioteca liviana que se vincula dinámicamente a HCCL y permite el uso de libfabric sin requerir integración directa.
Una vez completada la transferencia, libfabric informa a través de una cola de finalización y HCCL continúa el cálculo con los datos recibidos recientemente.
La experiencia de desarrollo
Construir Peer Direct implicó aventurarse en nuevos territorios con cronogramas ajustados. Libfabric aún no era popular en el campo de los aceleradores de IA. No había mucha documentación pública disponible y la discusión fue escasa. Se hizo más énfasis en sumergirse en el código fuente de libfabric y en la ingeniería inversa basada en la experimentación.
La comunicación con los ingenieros de AWS fue primordial, pero la zona horaria estaba limitada. Trabajar con un equipo con doce horas de anticipación significó que las iteraciones de depuración tuvieran plazos de entrega de 24 horas. Cada problema necesitaba una documentación cuidadosa y una comunicación adecuada, ya que la colaboración en tiempo real no era posible.
Había mucho en juego ya que toda la implementación de DL1 dependía del funcionamiento de esta funcionalidad. Los retrasos habrían frustrado el lanzamiento de un producto importante. Nadie en nuestro equipo tenía un conocimiento profundo de los aspectos internos de libfabric, por lo que estábamos aprendiendo una base de código compleja mientras diseñábamos una integración crítica simultáneamente.
Los resultados
Cuando realmente implementamos Peer Direct, las mejoras en la velocidad fueron todo el esfuerzo que valió la pena. Vimos un aumento de rendimiento de 1,5 a 2 veces para operaciones colectivas en un tamaño de mensaje de 32 MB. En mensajes más grandes, el rendimiento fue aún más sorprendente, con un rendimiento hasta 1,76 veces mejor con un tamaño de mensaje de 256 MB. La sobrecarga de la CPU creó un cuello de botella que desapareció por completo.
Lo más significativo es que estas mejoras de microbenchmark se tradujeron directamente en el rendimiento del entrenamiento del modelo real. Al entrenar el modelo DeepSpeed BERT de Habana con 5 mil millones de parámetros en 128 dispositivos Gaudí, vimos una ganancia sustancial en el rendimiento. Los modelos que utilizan métodos de optimización de memoria más agresivos, como ZeRO-2, que dependen más de la operación colectiva, se beneficiaron desproporcionadamente de Peer Direct.
PeerDirect fue uno de los principales facilitadores del rendimiento de Gaudi en instancias AWS DL1, permitiendo que la capacitación distribuida a gran escala se ejecutara sin esfuerzo el día del lanzamiento. Más allá de este impacto inicial, el esfuerzo sentó las bases para futuras funciones de comunicación de alto rendimiento y demostró que los aceleradores de IA nativos de la nube podrían seguir siendo competitivos a pesar de las limitaciones de la infraestructura de la nube.
La experiencia me recordó una lección importante en ingeniería de sistemas: a menudo las mejoras de rendimiento más importantes no resultan de optimizar el camino rápido, sino de evitar por completo los desvíos injustificados. Durante el entrenamiento de IA distribuida, hacer que los datos viajen directamente a través de los aceleradores sin copias innecesarias y sin intervención de la CPU es lo que distingue a un sistema que funciona de uno que escala.
¿Conclusiones clave? Una “aprendizaje” importante de este proyecto es que los supuestos sobre la topología de la red deben probarse en el punto más temprano del proceso de capacitación distribuida. Como muchas de las pilas de aceleradores se crearon en base a un entorno idealizado, no tienen en cuenta los saltos adicionales, las capas de traducción y/o los factores basados en costos que existen en los entornos de nube. Por lo tanto, antes de centrarse en optimizar el nivel de modelo o el nivel de kernel, los ingenieros deben realizar microbenchmarking colectivos simples en la topología deseada. Si la eficiencia de escalado disminuye drásticamente al aumentar el número de nodos o el tamaño de los mensajes, la razón probable es la ruta de datos, no el kernel. Al identificar el desvío entre el host y la memoria desde el principio, los ingenieros pueden centrar sus esfuerzos donde tendrán el mayor impacto.
Otra importante lección aprendida fue la necesidad de tratar tanto el registro de memoria como la transferencia de datos como cuestiones de rendimiento de primera clase. La sobrecarga de registro de memoria puede exceder en gran medida el tiempo dedicado a la comunicación si cada transferencia de datos requiere un nuevo registro. El caché LRU para recuerdos registrados fue una adición poco glamorosa a HCCL; sin embargo, eliminó efectivamente una fuente sistémica de latencia e hizo que la ruta RDMA fuera viable para cargas de trabajo del mundo real. Al desarrollar sistemas distribuidos, los ingenieros deben perfilar no sólo el ancho de banda de red disponible sino también los costos del ciclo de vida asociados con la asignación de buffers, su registro y la eliminación de esos registros. Pequeños cambios en estas rutas de control pueden dar como resultado grandes aumentos en el rendimiento de un extremo a otro.
Finalmente, el método de integración utilizado en este proyecto proporciona un patrón de integración. En lugar de reescribir HCCL para usar libfabric directamente, creamos una capa de abstracción delgada que mantuvo la semántica existente mientras reemplazaba la capa de transporte subyacente. Esto proporcionó varios beneficios, incluida la minimización del riesgo, la reducción de la rotación de código y la posibilidad de realizar pruebas incrementales. Los equipos que enfrentan un desafío similar (es decir, adaptar bibliotecas de comunicación nativas del acelerador a estructuras nativas de la nube) deberían intentar aislar la capa de transporte, mantener la semántica colectiva y crear interfaces pequeñas y comprobables entre las dos. Esto no solo permite un desarrollo más rápido sino que también permite un soporte más sencillo para futuros backends de transporte.
Divulgación: trabajo como gerente de equipo de ejecución de IA en Intel. Las perspectivas compartidas en este artículo son mías.