Transferência de memória de GPU de alto desempenho no AWS Sagemaker Hyperpod

A infraestrutura moderna de aprendizado profundo geralmente exige a transferência de grandes quantidades de dados entre GPUs em diferentes máquinas. Na Perplexity, encontramos um desafio técnico único: transferir com eficiência regiões de memória de GPU não contíguas entre máquinas à velocidade máxima possível. Nossa plataforma de destino, instâncias AWS p5, oferece uma impressionante largura de banda de rede de 3200 Gbps por meio de 32 placas de rede. Este artigo compartilha nossa jornada de construção de uma solução de rede personalizada de alto desempenho que atinge 97,1% dessa largura de banda teórica.

O Desafio Técnico

Nosso caso de uso apresentou vários requisitos técnicos principais:

• Transferência de alta largura de banda entre GPUs remotas de blocos de memória não contíguos
• Capacidade de adicionar ou remover nós dinamicamente de implantações do Kubernetes sem interromper as operações em andamento
• Suporte para padrões de comunicação ponto a ponto

Embora a biblioteca NCCL da NVIDIA seja o padrão de fato para aprendizado profundo distribuído, ela não era ideal para o nosso caso de uso:

• O NCCL se destaca na comunicação coletiva, mas requer o estabelecimento de um "mundo" estático, o que exige a reinicialização de todo o cluster ao ajustar os nós participantes.
• O modelo de comunicação síncrona do NCCL adiciona complexidade à nossa carga de trabalho assíncrona
• Queríamos controle direto sobre nossos padrões de transferência de memória para otimização
• Construir nossa própria solução proporcionou oportunidades valiosas de aprendizado

Redes Modernas de Alto Desempenho

Para entender nossa solução, vamos primeiro explorar como as redes modernas de alto desempenho diferem das redes tradicionais.

A maioria das redes que usamos diariamente depende de protocolos TCP/IP, onde os aplicativos se comunicam com a placa de rede por meio do kernel do sistema operacional usando sockets. No entanto, as redes de alto desempenho usam RDMA (Remote Direct Memory Access) - uma pilha de hardware e software completamente diferente que permite acesso direto à memória entre máquinas sem envolver a CPU.

A AWS fornece o Elastic Fabric Adapter (EFA), uma interface de rede personalizada que implementa o protocolo personalizado da Amazon chamado Scalable Reliable Datagram (SRD). Ao contrário das redes TCP/IP tradicionais, onde os dados devem ser copiados várias vezes entre o espaço do usuário, o espaço do kernel e os buffers de rede, o EFA com RDMA permite a transferência direta de dados entre a memória da GPU e a placa de rede, ignorando a CPU completamente.

Filosofia de Design de Rede de Alto Desempenho

A construção de sistemas de rede de alto desempenho exige repensar várias premissas fundamentais:

1. Propriedade de Buffer: Ao contrário dos sockets tradicionais, onde o kernel gerencia os buffers de rede e exige cópias entre o espaço do usuário e o espaço do kernel, o RDMA exige que os aplicativos gerenciem seus próprios buffers. Quando um aplicativo inicia uma operação de rede, ele transfere a propriedade do buffer para a placa de rede até que a operação seja concluída, eliminando a necessidade de cópia de dados.
2. Registro de Memória: Os aplicativos devem registrar regiões de memória com o kernel do sistema operacional. O kernel configura mapeamentos de endereços virtuais que permitem que a CPU, as GPUs e as placas de rede entendam os mesmos endereços virtuais. Este registro é uma operação única que permite transferências subsequentes de dados com cópia zero.
3. Plano de Controle vs Plano de Dados: Redes de alto desempenho separam as operações em duas categorias:
   1. As operações do plano de controle (como configuração de conexão e registro de memória) passam pelo kernel para garantir a segurança
   2. As operações do plano de dados (transferência real de dados) ignoram o kernel para obter o desempenho máximo
4. Recepção Antes da Transmissão: Sem buffers gerenciados pelo kernel, os aplicativos devem pré-publicar operações de recebimento, especificando onde os dados recebidos devem ser colocados. Esta é uma mudança fundamental em relação ao modelo de socket, onde os aplicativos podem receber dados a qualquer momento.
5. Conclusão Baseada em Polling: Em vez de aguardar notificações do kernel por meio de mecanismos como epoll, os aplicativos sondam diretamente as filas de conclusão de hardware. Isso elimina a sobrecarga de chamadas de sistema e permite a reação imediata às operações concluídas.
6. Consciência de Topologia de Hardware: Compreender e otimizar para a topologia de hardware é crucial para alcançar o desempenho máximo.

Entendendo a Topologia de Hardware

As instâncias AWS p5 possuem uma arquitetura interna sofisticada. Como mostrado abaixo, cada instância contém dois sockets de CPU formando dois nós NUMA, com cada nó NUMA se conectando a quatro switches PCIe:

Expandir

Sob cada switch PCIe, encontramos:

• Quatro placas de rede EFA de 100 Gbps
• Uma GPU NVIDIA H100
• Um SSD NVMe

Os caminhos de dados para transferências TCP/IP e RDMA demonstram diferenças fundamentais em suas abordagens:

Expandir

Com TCP/IP (lado esquerdo), os dados devem ser copiados várias vezes:

1. Da memória da GPU para o buffer do aplicativo no espaço do usuário
2. Do buffer do espaço do usuário para o buffer do kernel
3. Do buffer do kernel para a placa de rede
4. O processo inverso ocorre no lado receptor

Cada operação de cópia consome ciclos de CPU e largura de banda de memória. O aplicativo também deve alternar o contexto entre o espaço do usuário e o espaço do kernel para cada operação de rede.

Em contraste, o RDMA (lado direito) fornece transferência de dados com verdadeira cópia zero:

1. A placa de rede lê diretamente da memória da GPU
2. Os dados viajam diretamente para a placa de rede remota
3. A placa de rede remota grava diretamente na memória da GPU de destino
4. O aplicativo só precisa verificar uma fila de conclusão (CQ) no espaço do usuário para saber quando a transferência foi concluída

Com o RDMA e o emparelhamento de hardware adequado, a transferência de dados entre duas GPUs requer apenas a travessia do switch PCIe local e da rede:

Expandir

Em contraste, as transferências TCP/IP devem copiar dados várias vezes através da memória principal, causando congestionamento significativo no barramento PCIe:

Expandir

Construindo com o libfabric

Usamos o libfabric - um framework que fornece uma interface genérica para serviços de fabric. Nossa implementação usa dois tipos de operações RDMA:

• RDMA bidirecional (SEND/RECV) para mensagens de controle que transportam metadados sobre regiões de memória (por exemplo, deslocamentos e tamanhos).
• RDMA WRITE unidirecional para transferência real de dados, onde cada operação WRITE lida com um bloco de memória contíguo.

Jornada para o Desempenho Máximo

Nosso desenvolvimento progrediu através de vários estágios:

1. Implementação de transferência básica de mensagens unidirecionais usando SEND/RECV
2. Extensão para comunicação bidirecional
3. Adição de GPUDirect RDMA WRITE para transferência direta entre GPUs
4. Expansão para lidar com múltiplas transferências concorrentes
5. Introdução de enfileiramento de operações para maior robustez
6. Alcançada 97,4% de utilização de largura de banda em uma única placa de rede

Ao escalar para 32 placas de rede, implementamos várias otimizações cruciais:

1. Enfileiramento de Operações: Em vez de enviar operações diretamente para as placas de rede, mantemos uma fila em nível de aplicativo. Isso oferece robustez contra congestionamento de rede e simplifica o modelo de programação.
2. Aquecimento de Rede: O estabelecimento prévio de conexões melhorou o desempenho de inicialização.
3. Multithreading: Tópicos dedicados para as operações de rede de cada GPU.
4. Fixação de Núcleo de CPU: Vinculação de tópicos a núcleos de CPU específicos para evitar efeitos NUMA e falhas de cache.
5. Sharding de Estado: Redução de contenção entre tópicos.
6. Agrupamento de Operações: Envio de várias operações em conjunto.
7. Postagem Preguiçosa de Operações: As operações são enfileiradas primeiro no aplicativo. Após pesquisar a fila de conclusão, tentamos enviar as operações pendentes para a placa de rede, garantindo o uso eficiente dos recursos de rede.
8. Alocação de Recursos Sensível a NUMA: Aloque recursos do libfabric, como filas de conclusão, no nó NUMA correto para minimizar a latência de acesso à memória.

Por meio dessas otimizações, alcançamos um desempenho final de 3.108 Gbps - 97,1% da largura de banda máxima teórica.

O vídeo abaixo mostra nosso programa de linha de comando em ação. Ele transfere blocos não contíguos de memória de GPU para um nó remoto, alcançando uma velocidade de transmissão de 3108,283 Gbps - demonstrando uma utilização de largura de banda quase teórica da infraestrutura de rede:

Conclusão

Construir um sistema de rede de alto desempenho requer a compreensão tanto da arquitetura de hardware quanto dos princípios de design de sistemas. Embora bibliotecas como o NCCL ofereçam excelentes soluções para padrões comuns, às vezes soluções personalizadas são necessárias para requisitos específicos.

Nossa jornada demonstra que alcançar um desempenho de rede próximo ao teórico é possível com atenção cuidadosa à arquitetura do sistema, topologia de hardware e várias técnicas de otimização. O segredo não é apenas entender os componentes individuais, mas como eles interagem para formar um sistema completo.

A análise técnica completa desta jornada, incluindo detalhes de implementação e exemplos de código, está disponível em nosso repositório de código aberto.

Citação

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