Resumo Técnico: Compreensão e Melhoria do Desempenho de Comunicação em Inferência de LLMs Multi-nó

Declaração do Problema

À medida que os Modelos de Linguagem de Grande Escala (LLMs) aumentam de tamanho, seus requisitos de memória frequentemente excedem a capacidade de uma única GPU, tornando necessária a inferência distribuída em múltiplos nós. Embora estratégias de paralelismo de modelo, como Paralelismo de Tensores (TP) e Paralelismo de Pipeline (PP), sejam bem estudadas para inferência em nó único, seu desempenho em configurações multi-nó permanece pouco explorado. A inferência multi-nó introduz desafios significativos devido a latências inter-nó mais elevadas em comparação com conexões intra-nó (por exemplo, NVLink).

Os autores identificam que, embora o TP seja frequentemente superior para cargas de trabalho limitadas por memória e intensivas em decodificação, ele sofre com sobrecargas de comunicação substanciais devido a operações frequentes de all-reduce. Por outro lado, o Paralelismo Híbrido (HP), que combina TP e PP, reduz a sobrecarga de comunicação, mas falha em reduzir efetivamente o tempo de computação em regimes intensivos em decodificação devido à natureza de multiplicações de matrizes pequenas. Um gargalo crítico identificado é o desempenho subótimo das implementações padrão de all-reduce NCCL para tamanhos de mensagem pequenos (128 KB a 2 MB) típicos das fases de decodificação de LLMs, particularmente ao escalar entre nós.

Metodologia

Os autores conduziram um estudo sistemático de desempenho utilizando uma combinação de motores de inferência existentes de última geração e um protótipo de pesquisa personalizado:

1. Configuração Experimental:

   • Hardware: Os experimentos foram realizados no sistema Perlmutter (GPUs NVIDIA A100, interconexão HPE Slingshot-11) e no sistema Vista (GPUs NVIDIA GH200, interconexão InfiniBand).
   • Modelos: Llama 3.1 70B e 405B (Instruct) em precisão bf16.
   • Cargas de Trabalho: Dois regimes principais foram testados: "Intensivo em Pré-preenchimento" (comprimentos de prompt grandes) e "Intensivo em Decodificação" (sequências de geração longas).
   • Motores: O estudo utilizou vLLM e SGLang juntamente com YALIS, um motor de inferência de código aberto personalizado desenvolvido pelos autores para facilitar experimentos controlados em ambientes HPC multi-nó. O YALIS apresenta uma camada unificada de definição de modelo, execução baseada em Torch Compile e paralelismo de tensores baseado em AxoNN.

2. Abordagem de Análise:

   • Os autores realizaram experimentos de escalonamento forte, medindo a latência de lote ponta a ponta à medida que a contagem de GPUs aumentava de configurações de nó único para 32 nós.
   • Eles decomporam o tempo de execução em computação (Matmul, Outra Comp.), comunicação e tempo ocioso para identificar gargalos.
   • Benchmarks sintéticos de GEMM foram utilizados para isolar os efeitos do micro-loteamento (PP) versus divisão de tensores (TP) em kernels de multiplicação de matrizes.

3. Solução Proposta (NVRAR):

   • Para abordar o gargalo de comunicação, os autores desenvolveram o NVRAR, um algoritmo hierárquico de all-reduce baseado em duplicação recursiva usando NVSHMEM.
   • Projeto: O NVRAR opera em três fases:
     1. Reduce-Scatter Intra-nó: Utiliza NCCL para reduzir dados dentro de um nó.
     2. Duplicação Recursiva Inter-nó: Utiliza primitivas RMA não bloqueantes do NVSHMEM (put_nbi) para realizar all-reduce entre nós usando uma topologia de duplicação recursiva.
     3. All-Gather Intra-nó: Utiliza NCCL para reunir os dados reduzidos de volta a todas as GPUs dentro de um nó.
   • Otimizações:
     • Comunicação Não Bloqueante em Blocos: Sobreposição de computação e comunicação processando dados em blocos.
     • Cargas Úteis Fundidas: Combina dados e flags de sincronização em uma única carga útil de 8 bytes para evitar primitivas de sinalização explícita caras no Slingshot/InfiniBand.
     • Sincronização por Número de Sequência: Utiliza números de sequência únicos para sincronização global para evitar a sobrecarga de latência de barreiras globais (fence/quiet) enquanto garante a correção.

Contribuições Principais

1. Estudo Sistemático de Desempenho: O artigo fornece uma análise detalhada dos comportamentos de escalonamento de TP versus HP em configurações multi-nó, identificando que o TP é geralmente superior para cargas de trabalho intensivas em decodificação, mas é limitado por sobrecargas de comunicação, enquanto o HP luta para reduzir o tempo de computação nesses regimes.
2. Motor YALIS: O desenvolvimento do YALIS, um motor de inferência orientado à pesquisa projetado para instrumentabilidade e experimentação controlada em ambientes HPC multi-nó.
3. Algoritmo NVRAR: A proposta e implementação do NVRAR, um algoritmo de all-reduce personalizado otimizado para o regime de mensagens pequenas característico das fases de decodificação de LLMs. Ele aproveita o NVSHMEM para alcançar latência inferior às implementações padrão do NCCL em redes Slingshot e InfiniBand.
4. Integração e Avaliação: A integração do NVRAR no YALIS e no vLLM, demonstrando melhorias de desempenho ponta a ponta para modelos grandes (Llama 3.1 405B) e modelos de Mistura de Especialistas (MoE) (Qwen3-235B).

Resultados

• Comportamento de Escalonamento: Tanto TP quanto HP exibem escalonamento forte pobre em configurações multi-nó. O HP supera o TP em regimes limitados por computação (intensivo em pré-preenchimento), enquanto o TP supera o HP em regimes limitados por memória (intensivo em decodificação).
• Gargalos de Comunicação: O all-reduce NCCL escala mal entre nós para mensagens pequenas (128 KB–2 MB), frequentemente performando pior do que soluções baseadas em MPI nesses regimes específicos.
• Desempenho do NVRAR:
  • Microbenchmarks: O NVRAR alcança 1,9×–3,6× menor latência do que o NCCL para tamanhos de mensagem entre 128 KB e 2 MB nas interconexões HPE Slingshot e InfiniBand.
  • Inferência Ponta a Ponta: Integrado ao YALIS e vLLM, o NVRAR reduz a latência de lote ponta a ponta em até 1,72× para o modelo Llama 3.1 405B em cargas de trabalho intensivas em decodificação usando paralelismo de tensores.
  • Modelos MoE: O NVRAR também melhora a taxa de transferência para modelos MoE (por exemplo, Qwen3-235B) acelerando o componente TP da estratégia de paralelismo, alcançando até 1,14× maior taxa de transferência em comparação com configurações baseadas em NCCL.

Significância

O artigo afirma que, à medida que os LLMs crescem e a inferência se move para clusters multi-nó, a eficiência da comunicação torna-se o determinante primário do desempenho. Os autores demonstram que bibliotecas de comunicação padrão (NCCL) não são otimizadas para os tamanhos de mensagem específicos e restrições de latência das fases de decodificação de LLMs. Ao introduzir o NVRAR, eles fornecem uma solução direcionada que reduz significativamente a sobrecarga de comunicação, permitindo um escalonamento mais eficiente do paralelismo de tensores para modelos grandes. Este trabalho destaca a necessidade de primitivas de comunicação personalizadas para sistemas de inferência distribuída de próxima geração e valida a utilidade de abordagens hierárquicas baseadas em NVSHMEM sobre bibliotecas de comunicação coletiva padrão em ambientes específicos de interconexão HPC.