Inferência com o Gemma usando Dataflow e vLLM

NOV 13, 2024
Danny McCormick Senior Software Engineer Dataflow

Os modelos de linguagem grandes (LLMs), como o Gemma, são potentes e versáteis. Eles podem traduzir idiomas, escrever diferentes tipos de conteúdo de texto e responder às suas perguntas de maneira informativa. No entanto, implantá-los em produção, especialmente para casos de uso de streaming, pode ser um desafio significativo.

Esta postagem do blog explorará como usar duas ferramentas de última geração, o vLLM e o Dataflow, para implantar LLMs em escala com eficiência e com o mínimo de código. Primeiro, mostraremos como o vLLM usa lotes contínuos para disponibilizar LLMs com mais eficiência. Em segundo lugar, descreveremos como o gerenciador de modelos do Dataflow simplifica a implantação do vLLM e de outros frameworks de modelos grandes.


O que é o vLLM?

O vLLM é uma biblioteca de código aberto projetada especificamente para inferência de LLMs de alta capacidade de processamento e baixa latência. Ele otimiza a disponibilização de LLMs empregando várias técnicas especializadas, incluindo os lotes contínuos.

Para entender como funcionam os lotes contínuos, vejamos primeiro como os modelos, tradicionalmente, criam lotes de entradas. As GPUs se destacam no processamento paralelo, no qual vários cálculos são feitos simultaneamente. A criação de lotes permite que a GPU utilize todos os seus núcleos disponíveis para trabalhar em um lote inteiro de dados de uma só vez, em vez de processar cada entrada individualmente. Isso acelera significativamente o processo de inferência; muitas vezes, a realização da inferência em oito registros de entrada de uma só vez usa recursos semelhantes à realização da inferência em um único registro.

Você pode pensar nos lotes como sendo semelhantes à cozinha de um restaurante: em vez de preparar cada prato individualmente, o chef pode agrupar pedidos semelhantes e prepará-los juntos, economizando tempo e recursos. Se você tiver oito fogões, o tempo e o esforço para preparar uma única omelete ou oito será similar.

without batching, inference is serial

Há algumas desvantagens, no entanto, nos lotes tradicionais. O mais importante, neste contexto, é que os lotes não funcionam tão bem quando os tempos das inferências diferem. Como a maioria dos frameworks não tem acesso ou conhecimento dos mecanismos subjacentes para a realização da inferência, eles normalmente apenas esperam até que todas as solicitações sejam concluídas antes de pegar um novo lote. Isso significa que um único registro lento pode consumir toda a capacidade de uma GPU, mesmo que os outros registros do lote tenham sido concluídos, resultando em jobs mais lentos e caros.

batches of varying lengths lead to wasted compute

Ao executar inferências com um modelo de linguagem grande (LLM), esperar que um lote inteiro seja concluído pode ser proibitivamente lento e caro. Isso porque a quantidade de tempo necessária para gerar uma inferência para um registro tem uma correlação de 1 para 1 com o tamanho do registro. Por exemplo, imagine a criação de um lote com as duas solicitações a seguir para um LLM:

  1. Qual é a capital do México?

2. Quais são algumas das diferenças e semelhanças culturais entre o México e os Estados Unidos?


É esperada uma resposta curta para a pergunta (1) e uma resposta longa para a pergunta (2). Como leva muito mais tempo para responder à pergunta (2), no entanto, temos que esperar que essa pergunta seja concluída, enquanto ela monopoliza nossa GPU, antes de podermos retornar qualquer um dos resultados do lote.

prompts of varying length lead to wasted compute

Os lotes contínuos permitem que o vLLM atualize os lotes enquanto a solicitação ainda está em execução. Ele consegue isso ao se beneficiar da maneira pela qual os LLMs realizam inferências: por um processo de repetição da geração do próximo token em sua resposta. Então, na verdade, na geração da frase "A capital do México é a Cidade do México", estamos executando a inferência sete vezes (uma vez por palavra de saída). Em vez de criar lotes de entradas uma vez, a técnica de lotes contínuos do vLLM permite que ele recalcule um lote toda vez que as execuções do LLM gerarem um conjunto de tokens para um lote. Ele pode adicionar solicitações ao lote em tempo real e retornar resultados antecipados quando um registro de um lote estiver completamente concluído.

continuous batching fully utilizes the GPU

Os lotes dinâmicos do vLLM, e outras otimizações, demonstraram melhorar a capacidade de processamento de inferências em duas a quatro vezes para LLMs populares em alguns casos, o que faz dele uma ferramenta muito útil para a disponibilização de modelos. Para obter mais informações sobre o vLLM, confira este artigo.


Uso do vLLM no Dataflow

A implantação de uma instância do vLLM em um pipeline de streaming pode ser complexa. Normalmente, você teria que:

  • Ativar um único servidor do vLLM no meio de seu pipeline. Isso geralmente não é simples, já que a maioria dos sistemas de streaming ativa vários processos worker, então é preciso escolher um líder para ativar um processo do vLLM dedicado.

  • Assegurar que todos os processos worker possam se comunicar com esse servidor Singleton.

  • Monitorar o serviço e colocá-lo em produção para que ele seja tolerante a falhas no servidor do vLLM.


Isso envolve muito multiprocessamento, o que pode ser demorado e propenso a erros, e requer conhecimento especializado. Isso também requer uma compreensão profunda da topologia subjacente. Essa topologia pode mudar se você também quiser experimentar diferentes configurações de máquina (por exemplo, para comparar o desempenho entre uma máquina com 8 núcleos e uma com 16 núcleos).

Felizmente, o Dataflow simplifica esse processo com seu gerenciador de modelos. Esse recurso abstrai as complexidades do gerenciamento e da implantação de modelos dentro de um pipeline. Por padrão, o Dataflow provisiona um processo worker por núcleo disponível em suas máquinas worker. Esses processos são responsáveis por lidar com a E/S do worker e com quaisquer transformações realizadas nos dados, e operam de forma totalmente independente. Para a maioria dos pipelines, incluindo pipelines de preparação de dados para casos de uso de ML, essa é a topologia ideal.

Beam worker processes operate independently in parallel

O problema está nos pipelines que precisam disponibilizar modelos grandes, como um dos modelos Gemma. Não é barato nem rápido carregar uma cópia de modelos grandes em todos os processos porque os pipelines provavelmente terão problemas de memória. A topologia ideal para a maioria desses pipelines é carregar uma única cópia do modelo grande.

O gerenciador de modelos do Dataflow foi criado para permitir que os usuários controlem o número exato de cópias de um modelo que são implantadas no pipeline, independentemente da topologia de rede. Quando você aplica a transformação RunInference, o Dataflow é capaz de entender a sua intenção para que possa criar a topologia ideal para o pipeline e implantar o número ideal de modelos. Você só precisa fornecer alguns parâmetros de configuração.

the model manager allows Beam workers to share a single model

Ao usar o vLLM, em vez de carregar um modelo, o gerenciador de modelos do Dataflow ativa uma única instância do vLLM em um processo de inferência dedicado. Os processos worker podem, então, enviar registros para essa instância de maneira eficiente para a realização da inferência.

the model manager allows Beam to spin up a dedicated vLLM process

Esse gerenciador de modelos permite que o Dataflow se beneficie totalmente dos lotes contínuos do vLLM; quando um worker recebe uma solicitação de entrada, ele a anexa de forma assíncrona à fila de solicitações do vLLM e aguarda a resposta, permitindo que o vLLM crie lotes dinamicamente com tantas solicitações quantas for possível.

O gerenciador de modelos do Dataflow e a transformação RunInference tornam incrivelmente fácil incorporar o vLLM a um pipeline. Basta especificar alguns detalhes de configuração e algumas linhas de código. Como o Dataflow é capaz de entender a sua intenção subjacente, ele configura todo o restante da topologia do pipeline para você. Em cinco linhas de código, você pode escrever um pipeline completo para ler os dados, executá-los via vLLM e enviá-los para um coletor.

model_handler = VLLMCompletionsModelHandler('google/gemma-2b')
with beam.Pipeline() as p:
  _ = (p | beam.ReadFromSource(<config>)
         | RunInference(model_handler) # Send the prompts to vLLM and get responses.
         | beam.WriteToSink(<config>))

Você pode encontrar um pipeline completo e executá-lo aqui: https://cloud.google.com/dataflow/docs/notebooks/run_inference_vllm


Desempenho

O vLLM aumenta significativamente o desempenho da inferência de LLMs em pipelines do Dataflow. Para comparar o desempenho do vLLM com um pipeline simples usando lotes de tamanho fixo, executamos dois pipelines com um único worker com GPUs T4. Cada pipeline leu prompts do conjunto de dados P3, os executou no modelo google/gemma-2b e registrou o resultado.

Com o uso de uma estratégia de lotes simples (padrão), foram necessárias 59.137 horas de vCPU para processar 10.000 prompts. Com o vLLM com lotes contínuos, foram necessárias apenas 2.481 horas de vCPU para processar os mesmos 10.000 prompts. Isso é uma melhoria de mais de 23 vezes!

Há algumas ressalvas aqui: especificamente, nenhum ajuste foi feito em nenhum dos pipelines, e o pipeline simples provavelmente teria um desempenho muito melhor se fosse ajustado para usar lotes maiores ou mais uniformes. Dito isso, essa é a magia do vLLM; com menos de 20 linhas de código e nenhum trabalho de ajuste, conseguimos produzir um pipeline de disponibilização de LLMs de alto desempenho! Se quiséssemos comparar outro modelo, poderíamos fazê-lo de maneira eficiente, alterando uma única string em nosso gerenciador de modelos!


Próximas etapas

Ao combinar o poder do vLLM e do Dataflow, você pode implantar e escalonar LLMs com eficiência para seus aplicativos de streaming com facilidade. Para saber mais sobre como fazer isso, experimente executar este notebook de exemplo: https://cloud.google.com/dataflow/docs/notebooks/run_inference_vllm.

Para saber mais sobre os modelos Gemma e algumas das coisas que você pode fazer com eles, confira os documentos do Gemma: https://ai.google.dev/gemma/docs

Para saber mais sobre o vLLM e alguns dos outros mecanismos que ele utiliza para otimizar a disponibilização, acesse os documentos do vLLM: https://docs.vllm.ai/en/latest/