Pular para o conteúdo principal

Treinamento GRPO

Treine modelos de linguagem usando Group Relative Policy Optimization (GRPO) com ambientes de recompensa personalizados. Em vez de um modelo de recompensa, você fornece um módulo Python com uma classe de ambiente que executa episódios multi-turno e retorna pontuações.

Visão Geral

GRPO se diferencia do PPO de uma forma importante:
  • PPO requer um modelo de recompensa pré-treinado para pontuar respostas
  • GRPO usa um ambiente personalizado que você escreve — ele gera múltiplas completações por prompt, pontua-as através do seu ambiente, e otimiza a política relativa ao grupo
Isso torna o GRPO ideal para treinamento agêntico onde as recompensas vêm da execução de tarefas (uso de ferramentas, execução de código, interações multi-turno) em vez de um modelo de recompensa estático.

Início Rápido

aitraining llm --train \
  --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
  --trainer grpo \
  --rl-env-module my_envs.hotel_env \
  --rl-env-class HotelEnv \
  --rl-num-generations 4 \
  --rl-max-new-tokens 256

Python API

from autotrain.trainers.clm.params import LLMTrainingParams
from autotrain.project import AutoTrainProject

params = LLMTrainingParams(
    model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
    project_name="grpo-agent",

    trainer="grpo",
    rl_env_module="my_envs.hotel_env",
    rl_env_class="HotelEnv",
    rl_num_generations=4,

    # Parâmetros RL compartilhados
    rl_max_new_tokens=256,
    rl_temperature=1.0,
    rl_kl_coef=0.1,
    rl_clip_range=0.2,

    epochs=1,
    batch_size=4,
    lr=1e-5,
    peft=True,
    lora_r=16,
    lora_alpha=32,
)

project = AutoTrainProject(params=params, backend="local", process=True)
project.create()

Interface do Ambiente

Você implementa uma classe Python com 3 métodos:
from datasets import Dataset

class MyEnv:
    def build_dataset(self, tokenizer) -> Dataset:
        """Retorna um Dataset do HuggingFace com uma coluna 'prompt'.

        Pode incluir colunas extras (ex: case_idx) que são passadas
        como kwargs para score_episode().
        """
        return Dataset.from_dict({
            "prompt": ["Você é um agente de reservas de hotel...", ...],
            "case_idx": [0, 1, ...],
        })

    def score_episode(self, model, tokenizer, completion, case_idx) -> float:
        """Executa um episódio multi-turno a partir da completação e retorna um score.

        Args:
            model: O modelo atual sendo treinado
            tokenizer: O tokenizador
            completion: O texto gerado pelo modelo
            case_idx: Índice do dataset (ou qualquer coluna extra)

        Returns:
            Float entre 0.0 e 1.0
        """
        return score

    def get_tools(self) -> list[dict]:
        """Retorna schemas de ferramentas para geração (opcional).

        Returns:
            Lista de dicts de definição de ferramentas (formato OpenAI function calling)
        """
        return []

Ambiente com Configuração

Passe configuração JSON para seu ambiente via --rl-env-config:
aitraining llm --train \
  --trainer grpo \
  --rl-env-module my_envs.hotel_env \
  --rl-env-class HotelEnv \
  --rl-env-config '{"max_turns": 5, "difficulty": "hard"}'
O JSON é parseado e passado como **kwargs para o construtor do seu ambiente:
class HotelEnv:
    def __init__(self, max_turns=3, difficulty="normal"):
        self.max_turns = max_turns
        self.difficulty = difficulty

Requisitos

O treinamento GRPO requer que --rl-env-module e --rl-env-class sejam especificados. Eles são validados na inicialização — se qualquer um estiver faltando, o treinamento falhará com uma mensagem de erro clara.
GRPO usa o GRPOTrainer do TRL (requer TRL >= 0.28.0). O lado de padding do tokenizador é automaticamente definido como left, conforme exigido pelo GRPO.

Parâmetros

Parâmetros Específicos do GRPO

ParâmetroFlag CLIPadrãoDescrição
rl_env_module--rl-env-moduleNoneCaminho do módulo Python para o ambiente (obrigatório)
rl_env_class--rl-env-classNoneNome da classe no módulo do ambiente (obrigatório)
rl_num_generations--rl-num-generations4Número de completações por prompt
use_vllm--use-vllmFalseUsar vLLM para geração mais rápida
vllm_mode--vllm-modecolocateModo vLLM: colocate ou server
vllm_gpu_memory_utilization--vllm-gpu-memory-utilization0.3Fração da memória GPU para vLLM (modo colocate)
vllm_server_url--vllm-server-urlNoneURL do servidor vLLM externo (modo server)
vllm_tensor_parallel_size--vllm-tensor-parallel-size1GPUs para paralelismo tensorial do vLLM
vllm_server_gpus--vllm-server-gpus1GPUs dedicadas ao servidor vLLM (subtraídas do treinamento)

Parâmetros RL Compartilhados (PPO + GRPO)

ParâmetroFlag CLIPadrãoDescrição
rl_kl_coef--rl-kl-coef0.1Penalidade de divergência KL (beta no GRPOConfig)
rl_clip_range--rl-clip-range0.2Faixa de clipping (epsilon no GRPOConfig)
rl_env_config--rl-env-configNoneConfiguração JSON passada ao construtor do ambiente
rl_max_new_tokens--rl-max-new-tokens128Máximo de tokens por completação
rl_top_k--rl-top-k50Amostragem top-k
rl_top_p--rl-top-p1.0Amostragem top-p (núcleo)
rl_temperature--rl-temperature1.0Temperatura de geração

Aceleração com vLLM

Use vLLM para geração de completações significativamente mais rápida durante o treinamento GRPO:
aitraining llm --train --trainer grpo \
  --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
  --rl-env-module my_envs.hotel_env \
  --rl-env-class HotelEnv \
  --use-vllm \
  --vllm-gpu-memory-utilization 0.3
Dois modos disponíveis:
  • colocate (padrão) — vLLM compartilha a GPU com o treinamento. Ajuste --vllm-gpu-memory-utilization (padrão 0.3) para controlar a divisão de memória.
  • server — vLLM roda em GPUs dedicadas. Os processos de treinamento são automaticamente reduzidos por --vllm-server-gpus.
# Modo server: 8 GPUs total, 2 para vLLM, 6 para treinamento
aitraining llm --train --trainer grpo \
  --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
  --rl-env-module my_envs.hotel_env \
  --rl-env-class HotelEnv \
  --use-vllm \
  --vllm-mode server \
  --vllm-server-gpus 2 \
  --vllm-tensor-parallel-size 2
vLLM requer instalação separada: pip install aitraining[vllm] (requer vllm>=0.14.0).
GRPO não requer --data-path — o dataset é construído pelo método build_dataset() do seu ambiente.

Como Funciona

  1. Ambiente carrega — Seu módulo é importado via importlib.import_module(), classe instanciada com config opcional
  2. Dataset construídoenv.build_dataset(tokenizer) retorna prompts
  3. Modelo gera — GRPO gera rl_num_generations completações por prompt
  4. Ambiente pontuaenv.score_episode() é chamado para cada completação, retornando 0.0-1.0
  5. GRPO otimiza — A política é atualizada relativa às pontuações do grupo (melhores completações recebem maior peso)

GRPO vs PPO

CaracterísticaPPOGRPO
Fonte de recompensaModelo de recompensa pré-treinadoAmbiente personalizado (seu código)
Sinal de treinamentoPontuações do modelo de recompensaPontuações de episódios do ambiente (0-1)
Completações por prompt1Múltiplas (rl_num_generations)
Ideal paraRLHF geralTreinamento agêntico, uso de ferramentas, multi-turno
RequerCaminho do modelo de recompensaMódulo Python env + classe
Versão TRL>= 0.26.0>= 0.28.0

Melhores Práticas

  1. Comece com ambientes simples — Valide que a pontuação funciona antes de lógica multi-turno complexa
  2. Use rl_num_generations pequeno — Comece com 4, aumente se precisar de mais diversidade
  3. Pontue entre 0 e 1 — Use o intervalo completo; evite sempre retornar 0 ou 1
  4. Teste seu ambiente independentemente — Certifique-se de que build_dataset() e score_episode() funcionam antes de treinar
  5. Use LoRA — GRPO com fine-tuning completo requer memória significativa; LoRA torna prático
  6. Taxas de aprendizado pequenas — Comece com 1e-5, mesma orientação do PPO

Próximos Passos