Transferencia de memoria GPU de alto rendimiento en AWS Sagemaker Hyperpod

La infraestructura moderna de aprendizaje profundo a menudo requiere la transferencia de 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 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 nuestro viaje en la construcción de una solución de red personalizada de alto rendimiento que logra 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 añadir o eliminar nodos de forma dinámica de las implementaciones de Kubernetes sin interrumpir las operaciones en curso
• Soporte para patrones de comunicación entre pares (peer-to-peer)

Aunque 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 sobresale 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 síncrona de NCCL añade complejidad a nuestra carga de trabajo asíncrona
• Queríamos un control directo sobre nuestros patrones de transferencia de memoria para optimizar
• Construir nuestra propia solución proporcionó valiosas oportunidades de aprendizaje

Redes modernas de alto rendimiento

Para entender nuestra solución, primero exploremos cómo difieren las redes modernas de alto rendimiento de las redes tradicionales.

La mayoría de las redes que usamos a diario dependen de protocolos TCP/IP, donde las aplicaciones se comunican con la tarjeta de red a través del núcleo del sistema operativo usando 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 involucrar a 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 la red TCP/IP tradicional, 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 completamente la CPU.

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

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

1. Propiedad del búfer: A diferencia de los sockets tradicionales, donde el núcleo gestiona los búferes de red y requiere copias entre el espacio de usuario y el espacio del núcleo, RDMA requiere que las aplicaciones gestionen 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 la operación se completa, eliminando la necesidad de copiar datos.
2. Registro de memoria: Las aplicaciones deben registrar las regiones de memoria con el núcleo del sistema operativo. El núcleo establece asignaciones de direcciones virtuales que permiten que la CPU, las GPU y las tarjetas de red entiendan las mismas direcciones virtuales. Este registro es una operación única que permite transferencias de datos posteriores de cero copias.
3. Plano de control vs 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 la conexión y el registro de memoria) pasan por el núcleo para garantizar la seguridad
   2. Las operaciones del plano de datos (transferencia real de datos) omiten el núcleo para obtener el máximo rendimiento
4. Recepción antes de la transmisión: Sin búferes gestionados 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 socket 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 llamadas al sistema y permite una reacción inmediata a las operaciones completadas.
6. Conciencia de la topología de hardware: Comprender y optimizar para la topología de hardware es crucial para lograr el máximo rendimiento.

Comprensión de la topología del hardware

Las instancias p5 de AWS tienen una arquitectura interna sofisticada. Como se muestra a continuación, cada instancia contiene dos zócalos de CPU que forman dos nodos NUMA, donde 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
• Un 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. Desde la memoria de la GPU al búfer de la aplicación en el espacio de usuario
2. Desde el búfer del espacio de usuario al búfer del núcleo
3. Desde el 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 cambiar de contexto entre el espacio de usuario y el espacio del núcleo para cada operación de red.

Por el contrario, RDMA (lado derecho) proporciona una transferencia de datos de cero copias real:

1. La tarjeta de red lee directamente desde 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 comprobar una cola de finalización (CQ) en el espacio de usuario para saber cuándo se ha completado la transferencia

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

Ampliar

Por el contrario, las transferencias TCP/IP deben copiar datos varias veces a través de la memoria principal, causando una congestión significativa del bus PCIe:

Ampliar

Construyendo con libfabric

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

• RDMA de dos lados (SEND/RECV) para mensajes de control que transportan metadatos sobre regiones de memoria (p. ej., desplazamientos y tamaños).
• RDMA WRITE de un lado para la transferencia real de datos, donde cada operación WRITE gestiona un fragmento de memoria contiguo.

Viaje hacia el máximo rendimiento

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

1. Implementamos una transferencia de mensajes unidireccional básica usando SEND/RECV
2. Extendimos a la comunicación bidireccional
3. Añadimos GPUDirect RDMA WRITE para la transferencia directa entre GPU
4. Expandimos para manejar múltiples transferencias concurrentes
5. Introdujimos la puesta en cola de operaciones para la robustez
6. Logramos 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. Puesta en cola de operaciones: En lugar de enviar operaciones directamente a las tarjetas de red, mantenemos una cola a nivel de aplicación. Esto proporciona robustez contra la congestión de la red y simplifica el modelo de programación.
2. Calentamiento de red: El establecimiento previo de conexiones mejoró el rendimiento al inicio.
3. Multihilo: Hilos dedicados para las operaciones de red de cada GPU.
4. Anclaje de núcleos de CPU: Vincular hilos a núcleos de CPU específicos para evitar efectos NUMA y errores de caché.
5. Fragmentación de estado: Reducción de la contención entre hilos.
6. Agrupación de operaciones: Envío de múltiples operaciones juntas.
7. Publicación diferida de operaciones: Las operaciones se ponen primero en cola en la aplicación. Después de sondear la cola de finalización, intentamos enviar las operaciones pendientes a la tarjeta de red, asegurando un uso eficiente de los recursos de la red.
8. Asignación de recursos con reconocimiento NUMA: Asignar recursos de libfabric como colas de finalización en el nodo NUMA correcto para minimizar la latencia de acceso a la memoria.

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

El video a continuación muestra nuestro programa de línea de comandos en acción. Transfiere fragmentos no contiguos de memoria de GPU a un nodo remoto, logrando una velocidad de transmisión de 3108.283 Gbps, lo que demuestra una utilización del ancho de banda casi teórica de la infraestructura de red:

Conclusión

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

Nuestro viaje demuestra que lograr un rendimiento de red casi teórico es posible con una atención cuidadosa a la arquitectura del sistema, la topología del hardware y varias 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 completo de este viaje, incluidos los 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}
}