Transferencia de memoria de GPU de alto rendimiento en AWS Sagemaker Hyperpod

La infraestructura moderna de aprendizaje profundo a menudo requiere transferir grandes cantidades de datos entre GPU en distintas máquinas. En Perplexity, nos enfrentamos a un desafío técnico único: transferir de manera eficiente regiones de memoria de GPU no contiguas entre máquinas a la máxima velocidad posible. Nuestra plataforma de destino, las instancias p5 de AWS, ofrece un impresionante ancho de banda de red de 3200 Gbps a través de 32 tarjetas de red. Este artículo comparte nuestra travesía para construir una solución de red personalizada de alto rendimiento que alcanza el 97.1 % de este ancho de banda teórico.

El desafío técnico

Nuestro caso de uso presentó varios requisitos técnicos clave:

• Transferencia de alto ancho de banda entre GPU remotas de fragmentos de memoria no contiguos
• Capacidad para agregar o eliminar nodos de forma dinámica de implementaciones de Kubernetes sin interrumpir las operaciones en curso
• Compatibilidad con patrones de comunicación entre pares

Si bien la biblioteca NCCL de NVIDIA es el estándar de facto para el aprendizaje profundo distribuido, no era ideal para nuestro caso de uso:

• NCCL destaca en la comunicación colectiva, pero requiere establecer un «mundo» estático, lo que exige reiniciar todo el clúster al ajustar los nodos participantes.
• El modelo de comunicación sincrónica de NCCL añade complejidad para nuestra carga de trabajo asincrónica
• Queríamos un control directo sobre nuestros patrones de transferencia de memoria para la optimización
• Construir nuestra propia solución brindó valiosas oportunidades de aprendizaje

Redes modernas de alto rendimiento

Para entender nuestra solución, primero exploremos en qué se diferencian las redes modernas de alto rendimiento de las redes tradicionales.

La mayoría de las redes que usamos a diario se basan en protocolos TCP/IP, donde las aplicaciones se comunican con la tarjeta de red a través del núcleo del sistema operativo utilizando sockets. Sin embargo, las redes de alto rendimiento utilizan RDMA (Acceso Directo a Memoria Remota), una pila de hardware y software completamente diferente que permite el acceso directo a la memoria entre máquinas sin la intervención de la CPU.

AWS proporciona Elastic Fabric Adapter (EFA), una interfaz de red personalizada que implementa el protocolo personalizado de Amazon llamado Scalable Reliable Datagram (SRD). A diferencia de las redes TCP/IP tradicionales, donde los datos deben copiarse varias veces entre el espacio de usuario, el espacio del núcleo y los búferes de red, EFA con RDMA permite la transferencia directa de datos entre la memoria de la GPU y la tarjeta de red, omitiendo por completo la CPU.

Filosofía del diseño de redes de alto rendimiento

La construcción de sistemas de red de alto rendimiento requiere replantear varios supuestos fundamentales:

1. Propiedad de los búferes: A diferencia de los sockets tradicionales, donde el núcleo administra los búferes de red y requiere copiar entre el espacio de usuario y el espacio del núcleo, RDMA requiere que las aplicaciones administren sus propios búferes. Cuando una aplicación inicia una operación de red, transfiere la propiedad del búfer a la tarjeta de red hasta que se completa la operación, lo que elimina la necesidad de copiar datos.
2. Registro de memoria: Las aplicaciones deben registrar regiones de memoria con el núcleo del sistema operativo. El núcleo configura asignaciones de direcciones virtuales que permiten que la CPU, las GPU y las tarjetas de red comprendan las mismas direcciones virtuales. Este registro es una operación única que habilita transferencias de datos posteriores sin copia.
3. Plano de control frente a plano de datos: Las redes de alto rendimiento separan las operaciones en dos categorías:
   1. Las operaciones del plano de control (como la configuración de conexiones y el registro de memoria) pasan por el núcleo para garantizar la seguridad
   2. Las operaciones del plano de datos (la transferencia de datos real) omiten el núcleo para lograr el máximo rendimiento
4. Recepción antes de la transmisión: Sin búferes administrados por el núcleo, las aplicaciones deben publicar previamente las operaciones de recepción, especificando dónde deben colocarse los datos entrantes. Este es un cambio fundamental respecto al modelo de sockets, donde las aplicaciones pueden recibir datos en cualquier momento.
5. Finalización basada en sondeo: En lugar de esperar notificaciones del núcleo a través de mecanismos como epoll, las aplicaciones sondean directamente las colas de finalización del hardware. Esto elimina la sobrecarga de las llamadas al sistema y permite una reacción inmediata ante las operaciones completadas.
6. Conciencia de la topología de hardware: Comprender y optimizar la topología del hardware es fundamental para alcanzar el rendimiento máximo.

Comprender la topología de hardware

Las instancias p5 de AWS tienen una arquitectura interna sofisticada. Como se muestra a continuación, cada instancia contiene dos sockets de CPU que forman dos nodos NUMA, y cada nodo NUMA se conecta a cuatro conmutadores PCIe:

Ampliar

Debajo de cada conmutador PCIe, encontramos:

• Cuatro tarjetas de red EFA de 100 Gbps
• Una GPU NVIDIA H100
• Una unidad SSD NVMe

Las rutas de datos para las transferencias TCP/IP y RDMA demuestran diferencias fundamentales en sus enfoques:

Ampliar

Con TCP/IP (lado izquierdo), los datos deben copiarse varias veces:

1. De la memoria de la GPU al búfer de la aplicación en el espacio de usuario
2. Del búfer del espacio de usuario al búfer del núcleo
3. Del búfer del núcleo a la tarjeta de red
4. El proceso inverso ocurre en el lado receptor

Cada operación de copia consume ciclos de CPU y ancho de banda de memoria. La aplicación también debe realizar un cambio de contexto entre el espacio de usuario y el espacio del núcleo para cada operación de red.

En contraste, RDMA (lado derecho) proporciona una transferencia de datos de auténtica copia cero:

1. La tarjeta de red lee directamente de la memoria de la GPU
2. Los datos viajan directamente a la tarjeta de red remota
3. La tarjeta de red remota escribe directamente en la memoria de la GPU de destino
4. La aplicación solo necesita verificar una cola de finalización (CQ) en el espacio de usuario para saber cuándo ha finalizado la transferencia

Con RDMA y el emparejamiento de hardware adecuado, la transferencia de datos entre dos GPU solo requiere atravesar el conmutador PCIe local y la red:

Ampliar

En contraste, las transferencias TCP/IP deben copiar los datos varias veces a través de la memoria principal, lo que provoca una congestión significativa en el bus PCIe:

Ampliar

Construcción con libfabric

Utilizamos libfabric, un marco que proporciona una interfaz genérica para servicios de entramado. nuestra implementación utiliza dos tipos de operaciones RDMA:

• RDMA bidireccional (SEND/RECV) para mensajes de control que transportan metadatos sobre regiones de memoria (por ejemplo, desplazamientos y tamaños).
• RDMA WRITE unidireccional para la transferencia de datos real, donde cada operación WRITE maneja un fragmento de memoria contiguo.

Camino hacia el rendimiento máximo

Nuestro desarrollo avanzó a través de varias etapas:

1. Implementación de transferencia básica de mensajes unidireccionales utilizando SEND/RECV
2. Ampliación a la comunicación bidireccional
3. Incorporación de GPUDirect RDMA WRITE para transferencia directa entre GPU
4. Ampliación para gestionar múltiples transferencias concurrentes
5. Introducción de colas de operaciones para mayor solidez
6. Se alcanzó una utilización del ancho de banda del 97.4 % en una sola tarjeta de red

Al escalar a 32 tarjetas de red, implementamos varias optimizaciones cruciales:

1. Cola de operaciones: En lugar de enviar operaciones directamente a las tarjetas de red, mantenemos una cola a nivel de aplicación. Esto proporciona solidez frente a la congestión de la red y simplifica el modelo de programación.
2. Calentamiento de la red: El establecimiento previo de conexiones mejoró el rendimiento de inicio.
3. Multihilo: Hilos dedicados para las operaciones de red de cada GPU.
4. Asignación de núcleos de CPU: Vinculación de hilos a núcleos de CPU específicos para evitar efectos NUMA y fallos de caché.
5. Fragmentación de estados: Reducción de la contención entre hilos.
6. Lote de operaciones: Envío conjunto de múltiples operaciones.
7. Publicación perezosa de operaciones: Las operaciones se ponen primero en cola en la aplicación. Tras sondear la cola de finalización, intentamos enviar las operaciones pendientes a la tarjeta de red, garantizando un uso eficiente de los recursos de red.
8. Asignación de recursos compatible con NUMA: Asignación de recursos de libfabric, como las colas de finalización, en el nodo NUMA correcto para minimizar la latencia de acceso a la memoria.

A través de estas optimizaciones, alcanzamos un rendimiento final de 3108 Gbps, el 97.1 % del ancho de banda máximo teórico.

El siguiente vídeo muestra nuestro programa de línea de comandos en acción. Transfiere fragmentos de memoria de GPU no contiguos a un nodo remoto, alcanzando una velocidad de transmisión de 3108.283 Gbps y demostrando una utilización del ancho de banda casi teórica de la infraestructura de red:

Conclusión

La construcción de un sistema de red de alto rendimiento requiere comprender tanto la arquitectura de hardware como los principios de diseño de sistemas. Si bien bibliotecas como NCCL ofrecen excelentes soluciones para patrones comunes, a veces son necesarias soluciones personalizadas para requisitos específicos.

Nuestra travesía demuestra que es posible alcanzar un rendimiento de red casi teórico prestando especial atención a la arquitectura del sistema, la topología de hardware y diversas técnicas de optimización. La clave no es solo comprender los componentes individuales, sino cómo interactúan para formar un sistema completo.

El análisis técnico detallado de esta travesía, que incluye detalles de implementación y ejemplos de código, está disponible en nuestro repositorio de código abierto.

Citación

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}