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

A infraestrutura moderna de aprendizaxe profunda adoita requirir a transferencia de grandes cantidades de datos entre GPUs en diferentes máquinas. En Perplexity, atopamos un reto técnico único: transferir de maneira eficiente rexións de memoria de GPU non contiguas entre máquinas á máxima velocidade posible. A nosa plataforma de destino, as instancias p5 de AWS, ofrece un impresionante ancho de banda de rede de 3200 Gbps a través de 32 tarxetas de rede. Este artigo comparte a nosa traxectoria construíndo unha solución de rede personalizada de alto rendemento que acada o 97,1 % deste ancho de banda teórico.

O reto técnico

O noso caso de uso presentaba varios requisitos técnicos clave:

• Transferencia de alto ancho de banda entre GPUs remotas de fragmentos de memoria non contiguos
• Capacidade para engadir ou eliminar nós dinamicamente de desprazamentos de Kubernetes sen interromper as operacións en curso
• Compatibilidade con patróns de comunicación entre iguais (peer-to-peer)

Aínda que a biblioteca NCCL de NVIDIA é o estándar de facto para a aprendizaxe profunda distribuída, non era a ideal para o noso caso de uso:

• NCCL destaca na comunicación colectiva, mais require establecer un "mundo" estático, o cal esixe reiniciar todo o clúster ao axustar os nós participantes.
• O modelo de comunicación síncrona de NCCL engade complexidade para a nosa carga de traballo asíncrona
• Queriamos un control directo sobre os nosos patróns de transferencia de memoria para a optimización
• Construír a nosa propia solución achegou valiosas oportunidades de aprendizaxe

Redes modernas de alto rendemento

Para comprender a nosa solución, imos explorar primeiro en que se diferencian as redes modernas de alto rendemento das redes tradicionais.

A maioría das redes que usamos a diario dependen de protocolos TCP/IP, nos que as aplicacións se comunican coa tarxeta de rede a través do núcleo do sistema operativo empregando sockets. Con todo, as redes de alto rendemento usan RDMA (Remote Direct Memory Access), unha pila de hardware e software completamente diferente que permite o acceso directo á memoria entre máquinas sen involucrar a CPU.

AWS ofrece Elastic Fabric Adapter (EFA), unha interface de rede personalizada que implementa o protocolo personalizado de Amazon chamado Scalable Reliable Datagram (SRD). A diferenza das redes TCP/IP tradicionais, nas que os datos deben copiarse varias veces entre o espazo de usuario, o espazo do núcleo e os búferes de rede, EFA con RDMA permite a transferencia directa de datos entre a memoria da GPU e a tarxeta de rede, omitindo por completo a CPU.

Filosofía do deseño de redes de alto rendemento

Construír sistemas de rede de alto rendemento require repensar varias suposicións fundamentais:

1. Propiedade do búfer: A diferenza dos sockets tradicionais, nos que o núcleo xestiona os búferes de rede e require copiar entre o espazo de usuario e o espazo do núcleo, RDMA require que as aplicacións xestionen os seus propios búferes. Cando unha aplicación inicia unha operación de rede, transfire a propiedade do búfer á tarxeta de rede ata que a operación se completa, eliminando a necesidade de copiar datos.
2. Rexistro de memoria: as aplicacións deben rexistrar rexións de memoria co núcleo do sistema operativo. O núcleo configura as asignacións de enderezos virtuais que permiten que a CPU, as GPUs e as tarxetas de rede comprendan os mesmos enderezos virtuais. Este rexistro é unha operación dunha soa vez que permite transferencias de datos posteriores sen copias.
3. Plano de control vs. plano de datos: as redes de alto rendemento separan as operacións en dúas categorías:
   1. As operacións do plano de control (como a configuración da conexión e o rexistro de memoria) pasan polo núcleo para garantir a seguridade
   2. As operacións do plano de datos (a transferencia real de datos) omiten o núcleo para acadar o máximo rendemento
4. Recepción antes da transmisión: sen búferes xestionados polo núcleo, as aplicacións deben publicar previamente as operacións de recepción, especificando onde se deben colocar os datos entrantes. Este é un cambio fundamental respecto ao modelo de sockets, no que as aplicacións poden recibir datos en calquera momento.
5. Finalización baseada en sondaxe: en vez de esperar polas notificacións do núcleo a través de mecanismos como epoll, as aplicacións interrogan directamente as colas de finalización do hardware. Isto elimina a sobrecarga das chamadas ao sistema e permite unha reacción inmediata ás operacións completadas.
6. Conciencia da topoloxía de hardware: comprender e optimizar a topoloxía de hardware é crucial para acadar o máximo rendemento.

Comprender a topoloxía do hardware

As instancias p5 de AWS teñen unha arquitectura interna sofisticada. Como se mostra a continuación, cada instancia contén dous zócalos de CPU que forman dous nós NUMA, e cada nodo NUMA conéctase a catro conmutadores PCIe:

Ampliar

Baixo cada conmutador PCIe, atopamos:

• Catro tarxetas de rede EFA de 100 Gbps
• Unha GPU NVIDIA H100
• Unha unidade SSD NVMe

As rutas de datos para as transferencias TCP/IP e RDMA demostran diferenzas fundamentais nos seus enfoques:

Ampliar

Con TCP/IP (lado esquerdo), os datos deben copiarse varias veces:

1. Desde a memoria da GPU ata o búfer da aplicación no espazo de usuario
2. Desde o búfer do espazo de usuario ata o búfer do núcleo
3. Desde o búfer do núcleo ata a tarxeta de rede
4. O proceso inverso ocorre no lado receptor

Cada operación de copia consume ciclos de CPU e ancho de banda de memoria. A aplicación tamén debe cambiar de contexto entre o espazo de usuario e o espazo do núcleo para cada operación de rede.

Pola contra, RDMA (lado dereito) ofrece unha verdadeira transferencia de datos sen copias (zero-copy):

1. A tarxeta de rede le directamente desde a memoria da GPU
2. Os datos viaxan directamente á tarxeta de rede remota
3. A tarxeta de rede remota escribe directamente na memoria da GPU de destino
4. A aplicación só precisa comprobar unha cola de finalización (CQ) no espazo de usuario para saber cando rematou a transferencia

Con RDMA e unha emparexaxe de hardware adecuada, a transferencia de datos entre dúas GPUs só require atravesar o conmutador PCIe local e a rede:

Ampliar

Pola contra, as transferencias TCP/IP deben copiar os datos varias veces a través da memoria principal, o que causa unha conxestión significativa no bus PCIe:

Ampliar

Construción con libfabric

Empregamos libfabric, un marco que ofrece unha interface xenérica para servizos de tecido (fabric). A nosa implementación usa dous tipos de operacións RDMA:

• RDMA bilateral (SEND/RECV) para mensaxes de control que conteñen metadados sobre rexións de memoria (por exemplo, desprazamentos e tamaños).
• RDMA WRITE unilateral para a transferencia real de datos, onde cada operación WRITE xestiona un fragmento de memoria contiguo.

Traxectoria cara ao rendemento máximo

O noso desenvolvemento progresou a través de varias etapas:

1. Implementouse a transferencia básica de mensaxes unidireccionais mediante SEND/RECV
2. Estendeuse á comunicación bidireccional
3. Engadiuse GPUDirect RDMA WRITE para a transferencia directa entre GPUs
4. Amplíouse para xestionar múltiples transferencias concorrentes
5. Introduciuse a cola de operacións para maior robustez
6. Acadouse o 97,4 % de utilización do ancho de banda nunha única tarxeta de rede

Ao escalar a 32 tarxetas de rede, implementamos varias optimizacións cruciais:

1. Cola de operacións: en vez de enviar operacións directamente ás tarxetas de rede, mantemos unha cola a nivel de aplicación. Isto proporciona robustez fronte á conxestión da rede e simplifica o modelo de programación.
2. Esquentamento da rede: preestablecer conexións mellorou o rendemento de inicio.
3. Multifío: fíos dedicados para as operacións de rede de cada GPU.
4. Fixación de núcleos de CPU: vincular fíos a núcleos de CPU específicos para evitar efectos NUMA e fallos na caché.
5. Fragmentación do estado: reducir a contención entre fíos.
6. Lotes de operacións: enviar múltiples operacións conxuntamente.
7. Envío de operacións diferido: as operacións quéixanse primeiro na aplicación. Despois de sondar a cola de finalización, intentamos enviar as operacións pendentes á tarxeta de rede, garantindo un uso eficiente dos recursos de rede.
8. Asignación de recursos con conciencia NUMA: asignar recursos de libfabric, como colas de finalización, no nodo NUMA correcto para minimizar a latencia de acceso á memoria.

A través destas optimizacións, acadamos un rendemento final de 3108 Gbps, o 97,1 % do máximo teórico de ancho de banda.

O vídeo seguinte mostra o noso programa de liña de comandos en acción. Transfire fragmentos non contiguos de memoria de GPU a un nodo remoto, acadando unha velocidade de transmisión de 3108,283 Gbps, o que demostra unha utilización do ancho de banda case teórica da infraestrutura de rede:

Conclusión

Construír un sistema de rede de alto rendemento require comprender tanto a arquitectura de hardware coma os principios de deseño de sistemas. Aínda que bibliotecas como NCCL ofrecen excelentes solucións para patróns comúns, ás veces son necesarias solucións personalizadas para requisitos específicos.

A nosa traxectoria demostra que é posible acadar un rendemento de rede case teórico prestándolle unha coidadosa atención á arquitectura do sistema, á topoloxía do hardware e a varias técnicas de optimización. A clave non é só comprender os compoñentes individuais, senón como interactúan para formar un sistema completo.

A análise técnica en profundidade desta traxectoria, incluídos os detalles de implementación e os exemplos de código, está dispoñible no noso repositorio de código aberto.

Cita

@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}
}