Integração com vLLM

Este documento o guiará pelo processo de uso do vLLM com o TRL para uma geração mais rápida em métodos online, como GRPO e Online DPO. Primeiro, resumiremos um resumo (tl;dr) de como usar o vLLM com o TRL e, em seguida, entraremos nos detalhes de como isso funciona por baixo dos panos.

O TRL atualmente suporta apenas versões do vLLM de 0.23.0 a 0.33.x. Certifique-se de ter uma versão dentro dessa faixa instalada para evitar problemas de compatibilidade.

Os seguintes treinadores atualmente suportam geração com vLLM:

• GRPOTrainer
• RLOOTrainer
• experimental.online_dpo.OnlineDPOTrainer

🚀 Como posso usar o vLLM com o TRL para acelerar o treinamento?

💡 Observação: Recursos necessários para este exemplo específico: um único nó com 8 GPUs.

Ao usar o vLLM com o TRL, o servidor vLLM e o trainer devem ser executados em dispositivos CUDA separados para evitar conflitos. Para orientações sobre como configurar isso corretamente, consulte Modes of using vLLM during training.

Primeiro, instale o vLLM usando o seguinte comando:

pip install "trl[vllm]"

Em seguida, execute o servidor em GPUs específicas (por exemplo, GPUs 0-3):

CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-7B --tensor-parallel-size 4 \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

Quando o servidor estiver em execução, você pode usá-lo para gerar completions para o treinamento. No exemplo abaixo, estamos usando os diferentes trainers suportados, utilizando o servidor vLLM para a geração. Os argumentos --tensor-parallel-size e --data-parallel-size controlam como o modelo e os dados são divididos entre as GPUs.

Neste exemplo, dividimos um modelo em 4 GPUs com paralelismo de tensor. Em seguida, execute o script de treinamento em GPUs diferentes (por exemplo, GPUs 4-7) passando use_vllm=True nos argumentos de treinamento da seguinte forma:

Exemplo de um script train.py simples:

from datasets import load_dataset
from trl import GRPOTrainer, GRPOConfig
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-7B",
    args=GRPOConfig(use_vllm=True, vllm_mode="server"),
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)

trainer.train()

from datasets import load_dataset
from trl.experimental.online_dpo import OnlineDPOConfig, OnlineDPOTrainer
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = OnlineDPOTrainer(
    model="Qwen/Qwen2.5-7B",
    args=OnlineDPOConfig(use_vllm=True, vllm_mode="server"),
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)

trainer.train()

from datasets import load_dataset
from trl import RLOOTrainer, RLOOConfig
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = RLOOTrainer(
    model="Qwen/Qwen2.5-7B",
    args=RLOOConfig(use_vllm=True, vllm_mode="server"),
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)

trainer.train()

E o comando de treinamento em GPUs separadas do servidor:

CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch train.py

Por que usar o vLLM?

Métodos online geram complementações durante o treinamento, e gerá-las com o próprio generate do modelo é o gargalo. O vLLM serve essas complementações muito mais rápido, graças a técnicas como PagedAttention.

Como o TRL usa o servidor 🔍

O trainer solicita completions no endpoint compatível com OpenAI /v1/completions, enviando os IDs dos tokens do prompt. Prompts multimodais seguem um caminho diferente: o servidor processa as imagens por conta própria, e os recursos resultantes são pareados com os mesmos IDs de tokens em /inference/v1/generate, já que nenhum endpoint compatível com OpenAI aceita IDs de tokens e imagens ao mesmo tempo.

O servidor apenas gera. Após cada passo do otimizador, o trainer transmite os pesos atualizados para ele via NCCL, anunciando-os com /start_weight_update e /update_weights e confirmando-os com /finish_weight_update.

Uso avançado

🍷 Mais opções de personalização com o vLLM?

Você pode personalizar a configuração do servidor passando qualquer argumento do vllm serve, por exemplo --tensor-parallel-size, --data-parallel-size, --max-model-len, --enable-prefix-caching, --enforce-eager, --kv-cache-dtype ou --trust-remote-code. Execute vllm serve --help para ver a lista completa.

Apenas os seguintes são exigidos pelo TRL:

trl vllm-serve está obsoleto: agora ele apenas monta este comando e executa o servidor do vLLM. Ele imprime o comando exato de vllm serve que executa, para que você possa copiá-lo e dispensar o wrapper.

💆🏻‍♀️ Qual é a melhor configuração distribuída?

Escale a geração com --tensor-parallel-size. O paralelismo de dados não ajuda mais modelos densos: desde o vLLM PR #30739 (lançado na 0.14.0), o escalonamento de paralelismo de dados offline para modelos não-MoE não é suportado.

vLLM com Backend Transformers

O vLLM pode usar o backend Transformers para implementações de modelos, o que funciona tanto para LLMs quanto para VLMs. Para habilitar isso, defina vllm_model_impl="transformers" na sua configuração ou passe-o por meio do argumento de linha de comando.

Para mais detalhes, confira vLLM Transformers Backend.

Exemplo:

CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-VL-3B-Instruct \
    --tensor-parallel-size 1 --port 8000 --enforce-eager --model-impl transformers \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1 \
    --enable-scale-out  # vLLM 0.30.0 and later, drop on earlier versions

Modos de uso do vLLM durante o treinamento

O TRL suporta dois modos para integrar o vLLM durante o treinamento: modo colocalizado (padrão) e modo servidor.

Modo Colocate

No modo colocate, o vLLM é executado dentro do processo do trainer e compartilha a memória da GPU com o modelo de treinamento. Isso evita a inicialização de um servidor separado e pode melhorar a utilização da GPU, mas pode levar a contenda de memória nas GPUs de treinamento. Este é o modo padrão.

Exemplo de configuração:

from trl import GRPOConfig

training_args = GRPOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

from trl.experimental.online_dpo import OnlineDPOConfig

training_args = OnlineDPOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

from trl import RLOOConfig

training_args = RLOOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

Modo Servidor

No modo servidor, o vLLM é executado como um processo separado em GPUs dedicadas e se comunica com o treinador via HTTP. Essa configuração é ideal se você tiver GPUs dedicadas à inferência.

Exemplo de configuração:

from trl import GRPOConfig

training_args = GRPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

from trl.experimental.online_dpo import OnlineDPOConfig

training_args = OnlineDPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

from trl import RLOOConfig

training_args = RLOOConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

Consulte a documentação do trainer que você está usando para obter detalhes específicos sobre o uso e os parâmetros do vLLM.

Para reduzir o uso de memória da GPU ao executar o vLLM, considere ativar o modo de suspensão do vLLM.