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

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

O Desafio Técnico

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

• 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 Kubernetes sem interromper as operações em curso
• Suporte a padrões de comunicação ponto a ponto

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

• O NCCL se destaca na comunicação coletiva, mas exige o estabelecimento de um "mundo" estático, o que requer a reinicialização de todo o cluster ao ajustar os nós participantes.
• O modelo de comunicação síncrona do NCCL adiciona complexidade para 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 valiosas oportunidades de aprendizado

Redes Modernas de Alto Desempenho

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

A maioria das redes que usamos diariamente depende de protocolos TCP/IP, onde as aplicações se comunicam com a placa de rede através do kernel do sistema operacional usando sockets. No entanto, 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 Elastic Fabric Adapter (EFA), uma interface de rede personalizada que implementa o protocolo personalizado da Amazon chamado Scalable Reliable Datagram (SRD). Ao contrário do networking TCP/IP tradicional, 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 Redes de Alto Desempenho

Construir 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 buffers de rede e exige cópia entre o espaço do usuário e o espaço do kernel, o RDMA exige que as aplicações gerenciem seus próprios buffers. Quando uma aplicação inicia uma operação de rede, ela 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: As aplicações 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 de cópia zero.
3. Plano de Controle vs Plano de Dados: Redes de alto desempenho separam as operações em duas categorias:
   1. 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 desempenho máximo
4. Recepção Antes da Transmissão: Sem buffers gerenciados pelo kernel, as aplicações devem pré-postar operações de recepção, especificando onde os dados recebidos devem ser colocados. Esta é uma mudança fundamental em relação ao modelo de socket, onde as aplicações podem receber dados a qualquer momento.
5. Conclusão baseada em Polling: Em vez de esperar por notificações do kernel por meio de mecanismos como epoll, as aplicações pesquisam diretamente as filas de conclusão do hardware. Isso elimina a sobrecarga de chamadas de sistema e permite reação imediata a operações concluídas.
6. Consciência da Topologia de Hardware: Compreender e otimizar para a topologia de hardware é crucial para atingir 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 conectando-se 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 da aplicação 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 reverso ocorre no lado receptor

Cada operação de cópia consome ciclos de CPU e largura de banda de memória. A aplicação 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) oferece transferência de dados de cópia zero real:

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 escreve diretamente na memória da GPU de destino
4. A aplicação 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 RDMA e pareamento de hardware adequado, transferir dados entre duas GPUs requer apenas percorrer o switch PCIe local e a rede:

Expandir

Em contraste, 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 libfabric

Usamos libfabric - uma estrutura 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 carregam 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 de Pico

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

1. Implementamos transferência básica de mensagens unidirecionais usando SEND/RECV
2. Estendido para comunicação bidirecional
3. Adicionado GPUDirect RDMA WRITE para transferência direta GPU-GPU
4. Expandido para lidar com múltiplas transferências simultâneas
5. Introduzida fila de operações para robustez
6. Atingimos 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 aplicação. Isso oferece robustez contra congestionamento de rede e simplifica o modelo de programação.
2. Aquecimento de Rede: O pré-estabelecimento de conexões melhorou o desempenho na inicialização.
3. Multithreading: Tópicos dedicados para as operações de rede de cada GPU.
4. Pinagem de Núcleos de CPU: Vincular tópicos a núcleos de CPU específicos para evitar efeitos NUMA e perda de cache.
5. Sharding de Estado: Reduzindo a contenção entre tópicos.
6. Batch de Operações: Enviando múltiplas operações juntas.
7. Postagem de Operações Preguiçosas: As operações são primeiro enfileiradas na aplicação. 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 ciente de NUMA: Alocar 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, atingindo uma velocidade de transmissão de 3108,283 Gbps - demonstrando utilização de largura de banda quase teórica da infraestrutura de rede:

Conclusão

Construir um sistema de rede de alto desempenho exige compreender tanto a arquitetura de hardware quanto os princípios de design de sistemas. Embora bibliotecas como NCCL forneç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 atingir desempenho de rede quase teórico é possível com atenção cuidadosa à arquitetura do sistema, topologia de hardware e várias técnicas de otimização. A chave não é apenas entender os componentes individuais, mas como eles interagem para formar um sistema completo.

O mergulho técnico completo 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}
}