Dimensionando a inferência LLM: inovações em paralelismo tensorial, paralelismo de contexto e paralelismo especializado

• Na Meta, estamos constantemente ampliando os limites dos sistemas de inferência LLM para potencializar aplicativos como o Meta AI App.
• Estamos compartilhando como desenvolveu e implementou técnicas avançadas de paralelismo para otimize as principais métricas de desempenho relacionadas à eficiência de recursos, rendimento e latência.

A rápida evolução dos grandes modelos de linguagem (LLMs) inaugurou uma nova era de aplicações baseadas em IA, desde agentes de conversação até geração avançada de conteúdo. No entanto, a implantação destes modelos massivos em escala para inferência em tempo real apresenta desafios significativos, particularmente no que diz respeito à obtenção de alto rendimento, baixa latência e melhor eficiência de recursos.

Nosso principal objetivo é otimizar as principais métricas de desempenho:

• Eficiência de recursos: Maximizando a utilização da GPU para melhorar a eficiência operacional.
• Taxa de transferência (consultas/s): Atendendo mais usuários processando um maior volume de solicitações.
• Latência: Minimizando os tempos de resposta para uma experiência de usuário perfeita. Isso inclui:
  • Time-to-first-token (TTFT) para pré-preenchimento: O tempo que leva para a primeira parte da resposta aparecer, idealmente abaixo de 350 ms.
  • Tempo para token incremental (TTIT) para decodificação: A latência entre palavras subsequentes, visando menos de 25 ms.

Essas métricas destacam as demandas computacionais distintas da inferência LLM: o pré-preenchimento exige muita computação, enquanto a decodificação exige muita largura de banda da memória. Para enfrentar esses desafios e permitir a implantação de grandes modelos, desenvolvemos e implementamos técnicas avançadas de paralelismo.

Os dois estágios da inferência LLM

Uma tarefa típica de inferência generativa do LLM se desenvolve em dois estágios:

1. Estágio de pré-preenchimento: Este estágio processa o prompt de entrada (que pode ter milhares de tokens) para gerar um cache de valor-chave (KV) para cada camada do transformador do LLM. O pré-preenchimento é vinculado à computação, porque o mecanismo de atenção é dimensionado quadraticamente com o comprimento da sequência.
2. Estágio de decodificação: Este estágio utiliza e atualiza incrementalmente o cache KV para gerar tokens (palavras) um por um. A decodificação é vinculado à memóriajá que o tempo de E/S da memória de leitura domina o tempo de atenção, com os pesos do modelo e o cache KV ocupando a maior parte da memória.

Resolvendo gargalos com paralelismo

Para dimensionar a inferência LLM de forma eficaz, especialmente para lidar com contextos longos e modelos massivos, empregamos três tipos principais de paralelismo de inferência:

1. Paralelismo tensorial (TP), que melhora o ajuste de modelos grandes em várias GPUs e alcança alto rendimento que um único dispositivo não pode fornecer. Envolve a fragmentação de camadas individuais do modelo, como blocos de atenção e camadas perceptron multicamadas (MLP), em blocos menores e independentes que podem ser executados em diferentes dispositivos.

Um desafio no paralelismo de tensores é a operação de comunicação “allreduce”, que pode contribuir com até 30% da latência ponta a ponta. Para mitigar isso, desenvolvemos acesso direto a dados (DDA) algoritmos:

• Algoritmo plano DDA: Melhora a latência de redução de tamanho de mensagem pequena, permitindo que cada classificação carregue diretamente a memória de outras classificações e execute operações de redução local. Isso reduz a latência de O(N) para O(1), aumentando a quantidade de troca de dados de O(n) para O(n^2).
• Algoritmo de árvore DDA: Divide o allreduce em duas fases (reduce-scatter e all-gather) e usa acesso direto aos dados em cada etapa. Isso move a mesma quantidade de dados que o algoritmo de anel, mas reduz a latência a um fator constante, tornando-o adequado para tamanhos de mensagens um pouco maiores.

Nossas soluções DDA demonstram acelerações significativas em relação a linhas de base como NCCL (NVIDIA Collective Communications Library) e RCCL (ROCm Communication Collectives Library for AMD GPUs). Por exemplo, com AMD MI300X, alcançamos paridade de desempenho geral com Nvidia H100, com DDA superando a linha de base RCCL em 10-50% para decodificação (tamanhos de mensagens pequenos) e produzindo aceleração de 10-30% para pré-preenchimento, resultando em aproximadamente 10% de redução no TTIT.

2. Paralelismo de contexto (CP), o que facilita o gerenciamento e o processamento de contextos extremamente longos, como o Capacidades de token de 1M/10M introduzidas com Llama 4. A inferência de contexto longo apresenta desafios únicos:

• Calcular: Os FLOPs de atenção densa são dimensionados quadraticamente com o comprimento do contexto, levando à dominação da computação de atenção.
• Memória: O cache KV cresce linearmente com o contexto.
• Comunicação: A latência de comunicação aumenta ao paralelizar vários hosts.

Nós implementamos duas variantes de paralelismo de contexto no módulo de atenção, muitas vezes referido como “atenção do anel”:

• Passa-KV: Nesta abordagem, os tokens de entrada são divididos em várias classificações de CP. Cada classificação calcula sua porção de tensores de consulta, chave e valor. Em seguida, os tensores de chave e valor são trocados entre as classificações para permitir interações de atenção em todo o contexto.
• Passar-Q: Semelhante ao Pass-KV, mas os tensores de consulta são trocados entre as classificações.

Nossas otimizações de paralelismo de contexto, combinadas com um kernel de atenção rápida, permitiram um desempenho notável para recursos de contexto longo. Conseguimos menos de um minuto para um milhão de tokens em um único host H100 e menos de um minuto para 10 milhões de tokens usando inferência distribuída em vários hosts H100 (por exemplo, 32 hosts H100). Com o Llama 3 405B, demonstramos escalonamento quase linear, alcançando o pré-preenchimento de 128 mil tokens em 3,8 segundos com CP em 16 nós e o pré-preenchimento de 1 milhão de tokens em 77 segundos.

3. Paralelismo especializado (PE)o que ajuda a dimensionar modelos de mistura de especialistas (MoE), onde um grande número de “especialistas” (módulos de rede neural) torna impossível ajustar o modelo inteiro em um único host. Na inferência baseada em EP, utilizamos um padrão de comunicação de dois disparos, todos para todos, para trocar tokens entre o paralelismo de dados e as classificações de paralelismo especializado com base no roteamento.

A comunicação todos-para-todos pode contribuir de 10 a 30% para a latência ponta a ponta, especialmente para mensagens decodificadas (100 KB a 2 MB). Para otimizar isso, estamos explorando soluções que incluem:

• Dinâmico de todos para todos: Envio de subblocos de dados para vizinhos remotos.
• Persistente de todos para todos: Resolvendo lentidão causada principalmente por troca de identificador de memória, balanceamento de carga de rede e sobrecarga de CPU.

Olhando para o Futuro: Inferência Desagregada e Desafios Futuros

Para otimizar ainda mais a inferência LLM, estamos avançando em direção Paralelismo ND (CP, PP, EP, TP entre nós, com DP separado) e desagregando camadas de pré-preenchimento e decodificação. Isso permite um melhor equilíbrio de recursos e o potencial de uso de hardware heterogêneo, onde hardware com muita computação é usado para pré-preenchimento e hardware com muita largura de banda de memória para decodificação. Este paralelismo multidimensional pode ajudar a desbloquear o serviço e a avaliação de modelos colossais.

Os desafios futuros neste espaço incluem:

• Design de tecido de nuvem: Otimizando a infraestrutura de nuvem subjacente para cargas de trabalho LLM.
• Comunicação indo para o kernel (kernel fundido): Integração de operações de comunicação diretamente em kernels computacionais para maior eficiência.
• Kernel iniciado pelo dispositivo: Permitir que os dispositivos iniciem operações diretamente, reduzindo a sobrecarga da CPU.

Esses avanços na paralelização e melhorias no nível do sistema ajudaram a viabilizar a próxima geração de aplicativos de IA e ampliar os limites do que os LLMs podem alcançar. Estamos comprometidos com a inovação contínua para garantir inferência LLM eficiente e escalonável para milhões de usuários em todo o mundo.