Fortalecendo a competitividade das GPUs AMD para inferência de LLM (2023)

• MLC-LLM permite compilar e implantar LLMs em GPUs AMD com ROCm, fazendo com que a Radeon RX 7900 XTX alcance 80% da velocidade da RTX 4090 e 94% da RTX 3090 Ti na inferência com Llama 2 7B/13B
• As três GPUs comparadas têm 24 GB de memória e largura de banda semelhante, então, na inferência de LLM sensível à latência, o gargalo de memória pesa mais do que o desempenho FP16
• A diferença da AMD vinha mais da falta de suporte e otimização de software do que do hardware em si, e o investimento em ROCm junto com compilação de machine learning ajuda a reduzir o custo de suporte a vários backends
• O ROCm foi aplicado à Radeon 7900 XTX, e o Vulkan ao APU AMD do SteamDeck, mostrando que é possível rodar Llama-7B quantizado em 4 bits também em dispositivos AMD de consumo
• Os resultados são baseados em benchmark de lote único, e otimizações de attention podem alterar os números tanto da AMD quanto da NVIDIA, então a interpretação deve considerar uma margem de erro de 10%

---

Inferência de LLM em GPU AMD com ROCm

• O MLC-LLM permite compilar e implantar LLMs em GPUs AMD usando ROCm
• O desempenho da Radeon RX 7900 XTX com Llama 2 7B/13B fica no seguinte nível
  • 80% da velocidade da NVIDIA GeForce RTX 4090
  • 94% da velocidade da NVIDIA GeForce RTX 3090 Ti
• Além do ROCm, também há suporte a Vulkan, ampliando o alcance de implantação de LLMs para dispositivos como o SteamDeck com APU AMD

Comparação de hardware e gargalo real

• A AMD RX 7900 XTX ocupa, em especificações, uma posição comparável à RTX 4090 e à RTX 3090 Ti
  • As três GPUs têm 24 GB de memória, podendo carregar modelos do mesmo tamanho
  • A largura de banda de memória também é semelhante entre elas
  • A RTX 4090 tem desempenho FP16 2x maior que a RX 7900 XTX, e a RTX 3090 Ti tem 1,3x mais
• A inferência de LLM sensível à latência é, em grande parte, limitada por gargalo de memória, então aqui a diferença de desempenho FP16 não é o gargalo principal
• A RX 7900 XTX é 40% mais barata que a RTX 4090
• A RTX 3090 Ti é um produto da geração anterior, então a comparação de preço é mais difícil e serve como ponto de referência

A causa da defasagem da AMD e a abordagem da MLC

• O motivo de a AMD ter ficado atrás estava mais próximo da falta de suporte e otimização de software para esses modelos do que do hardware
• Duas mudanças ajudam a reduzir essa diferença
  • A AMD vem investindo na stack ROCm para recuperar terreno
  • A compilação de machine learning reduz o custo de suporte geral de software para vários backends
• Em vez de escrever manualmente kernels separados para ROCm ou CUDA, a compilação de machine learning compila as cargas de trabalho e automatiza as otimizações
• O MLC-LLM é uma solução de implantação de LLM baseada em compilação de machine learning, construída sobre o Apache TVM Unity
  • Oferece um fluxo de desenvolvimento centrado em Python
  • Inclui transformação de grafo computacional, otimização de layout e escalonamento de kernels de GPU e APIs nativas de implantação
  • Suporta CUDA, Metal, ROCm, Vulkan e OpenCL
  • Vai de GPUs de servidor até dispositivos móveis como iPhone e Android

Formas de suporte a GPU e APU AMD

• Os caminhos de suporte para hardware AMD incluem ROCm, OpenCL, Vulkan e WebGPU
  • ROCm é a stack que a AMD tem promovido recentemente e possui componentes semelhantes ao CUDA
  • Vulkan é um padrão gráfico moderno, com amplo suporte em vários dispositivos GPU
  • WebGPU é um padrão web moderno que permite executar computação no navegador
• Há pouco software de machine learning construído fora do CUDA, e o custo de engenharia para duplicar a stack para cada novo hardware ou modelo de programação de GPU é alto
• A MLC dá suporte a vários caminhos com geração automática de código, sem reescrever kernels de GPU para cada backend
• O desempenho real depende da qualidade do runtime de GPU de baixo nível e da disponibilidade em cada plataforma

Implementação das otimizações em ROCm

• Foi escolhido ROCm para a Radeon 7900 XTX e Vulkan para o APU do SteamDeck
• A stack ROCm funcionou imediatamente, e graças ao pipeline de desenvolvimento baseado em Python do TVM Unity foi possível adicionar uma versão otimizada em poucas horas
• O suporte a ROCm reaproveitou componentes existentes
  • Todo o pipeline da MLC para alvos já existentes, como CUDA e Metal
    • Planejamento de memória
    • Fusão de operadores
  • O espaço genérico de otimização de kernels de GPU escrito em TVM TensorIR
  • O fluxo de geração de código ROCm do TVM, que gera kernels ROCm de baixo nível por meio do LLVM
• O código gerado pode ser exportado como biblioteca compartilhada ou estática e chamado via CLI, Python ou API REST

Condições e interpretação do benchmark com Llama 2

• O benchmark mediu Llama 2 7B e 13B com quantização em 4 bits
• O desempenho de decodificação foi medido inserindo um único token de prompt e gerando 512 tokens
• Todos os resultados são baseados em inferência de lote único
• Com base no lançamento do ROCm 5.6, o desempenho de inferência em lote único chegou a 80% da velocidade da NVIDIA 4090
• A linha de base em CUDA era considerada de ponta para essa carga de trabalho naquele momento
• Ainda há espaço para melhorias, como otimizações melhores de attention, e se essas otimizações entrarem no MLC, os números tanto da AMD quanto da NVIDIA podem melhorar
• Se a otimização for implementada apenas no lado da NVIDIA, a diferença pode crescer de 20% para 30%, então é recomendável interpretar os números com uma margem de erro de 10%

Condições de execução direta e exemplo

• Foram fornecidos wheel pré-compilados e instruções de execução para reproduzir o benchmark
• A condição de execução é uma GPU AMD em Linux com ROCm 5.6 ou superior funcionando
• A instalação do pacote MLC pré-compilado com ROCm habilitado segue a documentação de teste do MLC-LLM
• O exemplo em Python carrega o modelo Llama-2-7b-chat-hf-q4f16_1 com mlc_chat.ChatModule e mede o desempenho com benchmark_generate("Hi", generate_length=512)
• O MLC-LLM também oferece uma CLI interativa, mas no ROCm a CLI precisa ser compilada a partir do código-fonte, seguindo a documentação de build da CLI

Uso de Vulkan e memória unificada no SteamDeck

• O SteamDeck é usado como exemplo de uma família mais ampla de dispositivos AMD com APU
• A VRAM de GPU disponível no ROCm é limitada a 4 GB pela BIOS
• O driver Mesa Vulkan permite que buffers usem até 16 GB por meio de memória unificada, ultrapassando esse limite
• Essa capacidade de memória é suficiente para executar Llama-7B quantizado em 4 bits
• Como resultado, confirmou-se a possibilidade de suporte a LLM também em vários dispositivos AMD de consumo

Próximos trabalhos e links do projeto

• Na era da IA generativa, a disponibilidade de hardware se tornou uma questão importante
• A compilação de machine learning pode aliviar esse problema ao viabilizar implantação genérica de alto desempenho em vários backends de hardware
• A pesquisa atual está focada em GPUs de consumo
• Pela experiência anterior, as otimizações da MLC para GPUs de consumo às vezes se generalizam também para GPUs de nuvem, como da RTX 4090 para A100 e A10g
• As próximas áreas de trabalho são as seguintes
  • Suporte a batching e multi-GPU
  • Integração com o ecossistema PyTorch
  • Suporte a mais formas de quantização e arquiteturas de modelo
  • Expansão da otimização automática para mais backends de hardware
• A NVIDIA ainda mantém posição de liderança com inovação contínua, e o cenário pode mudar com novo hardware como a H100 e a evolução do software
• As informações sobre o MLC-LLM estão disponíveis na página do projeto, e o código-fonte está no repositório do GitHub