Integração com PEFT

O TRL suporta métodos de PEFT (Parameter-Efficient Fine-Tuning) para treinamento de modelos com eficiência de memória. O PEFT permite o ajuste fino de grandes modelos de linguagem treinando apenas um pequeno número de parâmetros adicionais, mantendo o modelo base congelado, o que reduz significativamente os custos computacionais e os requisitos de memória.

Este guia aborda como usar o PEFT com diferentes treinadores do TRL, incluindo as técnicas LoRA, QLoRA e prompt tuning.

Para um exemplo completo e funcional, consulte o notebook de SFT com LoRA/QLoRA.

Instalação

Para usar PEFT com TRL, instale as dependências necessárias:

pip install trl[peft]

Para suporte a QLoRA (quantização de 4 bits e 8 bits), instale também:

pip install bitsandbytes

Início rápido

Todos os treinadores do TRL suportam PEFT por meio do argumento peft_config. A maneira mais simples de habilitar o PEFT é usar a interface de linha de comando com a flag --use_peft:

python trl/scripts/sft.py \
    --model_name_or_path Qwen/Qwen2-0.5B \
    --dataset_name trl-lib/Capybara \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16 \
    --output_dir Qwen2-0.5B-SFT-LoRA

Alternativamente, você pode passar uma configuração PEFT diretamente no seu código Python:

from peft import LoraConfig
from trl import SFTTrainer

# Configure LoRA
peft_config = LoraConfig(
    r=32,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

# Configure training - note the higher learning rate for LoRA (10x base rate)
training_args = SFTConfig(
    learning_rate=2.0e-4,  # 10x the base rate (2.0e-5) for LoRA
    ...
)

# Create trainer with PEFT
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
)

Três formas de configurar o PEFT

O TRL oferece três métodos diferentes para configurar o PEFT, cada um adequado a diferentes casos de uso:

1. Usando flags de CLI (mais simples)

A maneira mais fácil de habilitar o PEFT é usar a flag --use_peft com a interface de linha de comando. Esse método é ideal para experimentos rápidos e configurações padrão:

python trl/scripts/sft.py \
    --model_name_or_path Qwen/Qwen2-0.5B \
    --dataset_name trl-lib/Capybara \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16 \
    --lora_dropout 0.05 \
    --output_dir Qwen2-0.5B-SFT-LoRA

Vantagens: Configuração rápida, sem necessidade de código

Desvantagens: Limitado ao LoRA, com menos opções de personalização

2. Passando peft_config para o Trainer (Recomendado)

Para mais controle, passe uma configuração PEFT diretamente para o trainer. Esta é a abordagem recomendada para a maioria dos casos de uso:

from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

peft_config = LoraConfig(
    r=32,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    peft_config=peft_config,  # Pass config here
)

Vantagens: Controle total, suporta todos os métodos PEFT (LoRA, Prompt Tuning, etc.)

Desvantagens: Exige código Python

3. Aplicando PEFT diretamente ao modelo (Avançado)

Para máxima flexibilidade, você pode aplicar PEFT ao seu modelo antes de passá-lo ao trainer:

from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM
from trl import SFTConfig, SFTTrainer

# Load base model
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B")

# Apply PEFT configuration
peft_config = LoraConfig(
    r=32,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, peft_config)

# Pass PEFT-wrapped model to trainer
trainer = SFTTrainer(
    model=model,  # Already has PEFT applied
    args=training_args,
    train_dataset=dataset,
    # Note: no peft_config needed here
)

Vantagens: Controle máximo, útil para arquiteturas de modelos personalizadas ou configurações complexas

Desvantagens: Mais verboso, requer compreensão dos detalhes internos do PEFT

Considerações sobre a taxa de aprendizado

Ao usar LoRA ou outros métodos PEFT, você geralmente precisa usar uma taxa de aprendizado mais alta (aproximadamente 10x) em comparação com o fine-tuning completo. Isso ocorre porque os métodos PEFT treinam apenas uma pequena fração dos parâmetros, exigindo uma taxa de aprendizado maior para alcançar atualizações de parâmetros semelhantes.

Taxas de aprendizado recomendadas:

Por que 10x? Os adaptadores LoRA possuem significativamente menos parâmetros treináveis do que o modelo completo. Uma taxa de aprendizado mais alta compensa essa redução no número de parâmetros, garantindo um treinamento eficaz. Para uma explicação detalhada, veja este post no blog.

Para mais boas práticas sobre o uso eficaz do LoRA, consulte a documentação LoRA Without Regret.

PEFT com diferentes treinadores

Os treinadores do TRL suportam configurações PEFT para vários paradigmas de treinamento. Abaixo estão exemplos detalhados para cada treinador principal.

python trl/scripts/sft.py \
    --model_name_or_path Qwen/Qwen2-0.5B \
    --dataset_name trl-lib/Capybara \
    --learning_rate 2.0e-4 \
    --num_train_epochs 1 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16 \
    --output_dir Qwen2-0.5B-SFT-LoRA

from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

# Configure LoRA
peft_config = LoraConfig(
    r=32,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "v_proj"],  # optional: specify target modules
)

# Configure training with higher learning rate for LoRA
training_args = SFTConfig(
    learning_rate=2.0e-4,  # 10x the base rate for LoRA
    ...
)

# Create trainer with PEFT config
trainer = SFTTrainer(
    model="Qwen/Qwen2-0.5B",  # can pass model name or loaded model
    args=training_args,
    train_dataset=dataset,
    peft_config=peft_config,  # pass PEFT config here
)
trainer.train()

python trl/scripts/dpo.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback_binarized \
    --learning_rate 5.0e-6 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16 \
    --output_dir Qwen2-0.5B-DPO-LoRA

from peft import LoraConfig
from trl import DPOConfig, DPOTrainer

# Configure LoRA
peft_config = LoraConfig(
    r=32,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

# Configure training with higher learning rate for LoRA
training_args = DPOConfig(
    learning_rate=5.0e-6,  # 10x the base rate for DPO with LoRA
    ...
)

# Create trainer with PEFT config
trainer = DPOTrainer(
    model="Qwen/Qwen2-0.5B",  # can pass model name or loaded model
    args=training_args,
    train_dataset=dataset,
    peft_config=peft_config,  # pass PEFT config here
)
trainer.train()

python trl/scripts/grpo.py \
    --model_name_or_path Qwen/Qwen2-0.5B \
    --dataset_name trl-lib/math-reasoning \
    --learning_rate 1.0e-5 \
    --per_device_train_batch_size 2 \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16 \
    --output_dir Qwen2-0.5B-GRPO-LoRA

from peft import LoraConfig
from trl import GRPOConfig, GRPOTrainer

# Configure LoRA
peft_config = LoraConfig(
    r=32,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

# Configure training with higher learning rate for LoRA
training_args = GRPOConfig(
    learning_rate=1.0e-5,  # 10x the base rate for GRPO with LoRA
    ...
)

# Create trainer with PEFT config
trainer = GRPOTrainer(
    model="Qwen/Qwen2-0.5B",  # can pass model name or loaded model
    args=training_args,
    train_dataset=dataset,
    peft_config=peft_config,  # pass PEFT config here
)
trainer.train()

QLoRA: Adaptação de Baixo Rango Quantizada

QLoRA combina quantização de 4 bits com LoRA para permitir o ajuste fino de modelos muito grandes em hardware de consumo. Essa técnica pode reduzir os requisitos de memória em até 4x em comparação com o LoRA padrão.

Como o QLoRA funciona

1. Quantização de 4 bits: O modelo base é carregado em precisão de 4 bits usando bitsandbytes
2. Pesos congelados: Os pesos do modelo quantizado permanecem congelados durante o treinamento
3. Adaptadores LoRA: Somente os parâmetros do adaptador LoRA são treinados em precisão mais alta
4. Eficiência de memória: Permite o ajuste fino de modelos como Llama-70B em uma única GPU de consumo

Usando QLoRA com TRL

Basta combinar load_in_4bit=True com a configuração do PEFT:

Linha de Comando

python trl/scripts/sft.py \
    --model_name_or_path meta-llama/Llama-2-7b-hf \
    --dataset_name trl-lib/Capybara \
    --load_in_4bit \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 16 \
    --output_dir Llama-2-7b-QLoRA

Exemplo em Python

Passe o quantization_config diretamente para o trainer junto com o peft_config — o trainer carrega e quantiza o modelo para você. O mesmo argumento quantization_config está disponível em SFTTrainer, DPOTrainer, GRPOTrainer e RLOOTrainer.

import torch

from peft import LoraConfig
from transformers import BitsAndBytesConfig
from trl import SFTConfig, SFTTrainer

# Configure 4-bit quantization
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# Configure LoRA
peft_config = LoraConfig(
    r=32,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

# Configure training with higher learning rate for LoRA
training_args = SFTConfig(
    learning_rate=2.0e-4,  # 10x the base rate for QLoRA
    ...
)

# Create trainer with quantization and PEFT config
trainer = SFTTrainer(
    model="meta-llama/Llama-2-7b-hf",
    args=training_args,
    train_dataset=dataset,
    quantization_config=bnb_config,
    peft_config=peft_config,
)

trainer.train()

Opções de Configuração do QLoRA

O BitsAndBytesConfig oferece várias opções para otimizar a memória e o desempenho:

import torch

from transformers import BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",  # or "fp4"
    bnb_4bit_compute_dtype=torch.bfloat16,  # Compute dtype for 4-bit base models
    bnb_4bit_use_double_quant=True,  # Nested quantization for additional memory savings
)

Parâmetros de configuração:

• bnb_4bit_quant_type: Tipo de dados de quantização ("nf4" ou "fp4"). NF4 é recomendado.
• bnb_4bit_compute_dtype: O dtype usado para o cálculo. Use bfloat16 para uma estabilidade de treinamento melhor.
• bnb_4bit_use_double_quant: Ativa a quantização aninhada para economizar cerca de 0,4 bits adicionais por parâmetro.

Quantização de 8 bits

Para uma precisão ligeiramente maior com uma redução menor de memória, você pode usar quantização de 8 bits:

from transformers import BitsAndBytesConfig, AutoModelForCausalLM

bnb_config = BitsAndBytesConfig(load_in_8bit=True)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,
    device_map="auto",
)

Ou via linha de comando:

python trl/scripts/sft.py \
    --model_name_or_path meta-llama/Llama-2-7b-hf \
    --load_in_8bit \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16

Checkpoints pré-quantizados

Alguns checkpoints são liberados quantizados, como gpt-oss em MXFP4 ou Mistral Medium 3.5 em FP8. Esses pesos não podem ser treinados in place, portanto, quando você passa um checkpoint assim pelo nome, o trainer o carrega dequantizado, com ou sem um peft_config. O modelo então ocupa a memória dos seus pesos dequantizados, por exemplo, cerca de 42 GB em bfloat16 para gpt-oss-20b, contra 13 GB para seu checkpoint MXFP4. Para carregá-lo em bfloat16, como para qualquer outro modelo:

from trl import SFTConfig

training_args = SFTConfig(model_init_kwargs={"dtype": "bfloat16"})

Os checkpoints FP8 exigem o transformers 5.0.0 ou posterior.

Prompt Tuning

O Prompt tuning é outra técnica de PEFT que aprende prompts suaves (embeddings contínuos) adicionados ao início da entrada, mantendo o modelo inteiro congelado. Isso é particularmente eficaz para modelos grandes.

Como o Prompt Tuning Funciona

1. Tokens virtuais: Adiciona incorporações contínuas aprendíveis (tokens virtuais) à entrada
2. Modelo Congelado: O modelo base inteiro permanece congelado
3. Prompts específicos de tarefa: Cada tarefa aprende seus próprios embeddings de prompt
4. Eficiência extrema: apenas os embeddings do prompt são treinados (tipicamente 8-20 tokens)

Usando Prompt Tuning com TRL

from peft import PromptTuningConfig, PromptTuningInit, TaskType
from trl import SFTConfig, SFTTrainer

# Configure Prompt Tuning
peft_config = PromptTuningConfig(
    task_type=TaskType.CAUSAL_LM,
    prompt_tuning_init=PromptTuningInit.TEXT,
    num_virtual_tokens=8,
    prompt_tuning_init_text="Classify if the tweet is a complaint or not:",
    tokenizer_name_or_path="Qwen/Qwen2-0.5B",
)

# Configure training with higher learning rate for Prompt Tuning
training_args = SFTConfig(
    learning_rate=2.0e-2,  # Prompt Tuning typically uses 1e-2 to 3e-2
    ...
)

# Create trainer with PEFT config
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    peft_config=peft_config,  # pass PEFT config here
)

trainer.train()

Configuração do Prompt Tuning

from peft import PromptTuningConfig, PromptTuningInit, TaskType

peft_config = PromptTuningConfig(
    task_type=TaskType.CAUSAL_LM,  # Task type
    prompt_tuning_init=PromptTuningInit.TEXT,  # Initialize from text
    num_virtual_tokens=8,  # Number of virtual tokens
    prompt_tuning_init_text="Your initialization text here",
    tokenizer_name_or_path="model_name",
)

Parâmetros de configuração:

• task_type: O tipo de tarefa (TaskType.CAUSAL_LM para modelagem de linguagem)
• prompt_tuning_init: Método de inicialização (TEXT, RANDOM)
• num_virtual_tokens: Número de tokens virtuais a serem adicionados no início (geralmente 8-20)
• prompt_tuning_init_text: Texto para inicializar os tokens virtuais (ao usar a inicialização TEXT)
• tokenizer_name_or_path: Tokenizador para inicialização a partir de texto

Prompt Tuning vs LoRA

Configurações avançadas de PEFT

Parâmetros de Configuração do LoRA

from peft import LoraConfig

peft_config = LoraConfig(
    r=16,  # LoRA rank
    lora_alpha=32,  # LoRA scaling factor
    lora_dropout=0.05,  # Dropout probability
    bias="none",  # Bias training strategy
    task_type="CAUSAL_LM",  # Task type
    target_modules=["q_proj", "v_proj"],  # Modules to apply LoRA
    modules_to_save=None,  # Additional modules to train
)

Parâmetros principais:

• r: rank do LoRA (valores típicos: 8, 16, 32, 64). Um rank maior significa mais parâmetros, mas potencialmente um melhor desempenho.
• lora_alpha: Fator de escala (geralmente 2x o rank). Controla a magnitude das atualizações do LoRA.
• lora_dropout: Probabilidade de dropout para as camadas LoRA (típico: 0.05-0.1).
• target_modules: Quais módulos aplicar o LoRA. Opções comuns:
  • ["q_proj", "v_proj"]: Consulta e valor de atenção (eficiente em memória)
  • ["q_proj", "k_proj", "v_proj", "o_proj"]: Todas as projeções de atenção
  • ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]: Todas as camadas lineares
• modules_to_save: Módulos adicionais para treinar completamente (por exemplo, ["embed_tokens", "lm_head"])

Seleção do Módulo Alvo

Você pode especificar quais módulos aplicar o LoRA. Padrões comuns:

# Minimal (most memory efficient)
target_modules=["q_proj", "v_proj"]

# Attention only
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]

# All linear layers (best performance, more memory)
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]

Usando Argumentos de Linha de Comando

Os scripts do TRL aceitam parâmetros PEFT via linha de comando:

python trl/scripts/sft.py \
    --model_name_or_path Qwen/Qwen2-0.5B \
    --dataset_name trl-lib/Capybara \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16 \
    --lora_dropout 0.05 \
    --lora_target_modules q_proj v_proj \
    --output_dir output

Flags disponíveis:

• --use_peft: Ativa o PEFT
• --lora_r: rank do LoRA (padrão: 16)
• --lora_alpha: alpha do LoRA (padrão: 32)
• --lora_dropout: dropout do LoRA (padrão: 0.05)
• --lora_target_modules: Módulos-alvo (separados por espaço)
• --lora_modules_to_save: Módulos adicionais para treinar
• --use_rslora: Ativa o LoRA com Estabilização de Rank
• --use_dora: Ativa Weight-Decomposed LoRA (DoRA)
• --load_in_4bit: Ativa a quantização de 4 bits (QLoRA)
• --load_in_8bit: Ativa a quantização de 8 bits

Salvando e carregando modelos PEFT

Salvando

Após o treinamento, salve seus adaptadores PEFT:

# Save the adapters
trainer.save_model("path/to/adapters")

# Or manually
model.save_pretrained("path/to/adapters")

Isso salva apenas os pesos do adaptador (~alguns MB) em vez do modelo completo (~vários GB).

Carregando

Carregue um modelo PEFT para inferência:

from transformers import AutoModelForCausalLM
from peft import PeftModel

# Load base model
base_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B")

# Load PEFT adapters
model = PeftModel.from_pretrained(base_model, "path/to/adapters")

# Optionally merge adapters into base model for faster inference
model = model.merge_and_unload()

Enviando para o Hub

Você pode compartilhar facilmente seus adaptadores PEFT no Hugging Face Hub:

# Push adapters to Hub
model.push_to_hub("username/model-name-lora")

# Load from Hub
from peft import PeftModel
model = PeftModel.from_pretrained(base_model, "username/model-name-lora")

Treinamento Multi-GPU

O PEFT funciona perfeitamente com o suporte multi-GPU do TRL:

torchrun --nproc_per_node 8 trl/scripts/sft.py \
    --model_name_or_path Qwen/Qwen2-0.5B \
    --dataset_name trl-lib/Capybara \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16

Para QLoRA com múltiplas GPUs, o modelo base é automaticamente fragmentado:

torchrun --nproc_per_node 8 trl/scripts/sft.py \
    --model_name_or_path meta-llama/Llama-2-70b-hf \
    --load_in_4bit \
    --use_peft \
    --lora_r 32

Paralelismo de Pipeline Ingênuo (NPP) para Modelos Grandes

Para modelos muito grandes (>60B parâmetros), o TRL suporta Naive Pipeline Parallelism (NPP), que distribui o modelo e os adaptadores entre várias GPUs. As ativações e os gradientes são comunicados entre as GPUs, suportando tanto int8 quanto outros tipos de dados.

Como usar o NPP

Carregue seu modelo com um device_map personalizado para dividi-lo entre vários dispositivos:

from transformers import AutoModelForCausalLM
from peft import LoraConfig

# Create custom device map (see accelerate documentation)
device_map = {
    "model.embed_tokens": 0,
    "model.layers.0": 0,
    # ... distribute layers across GPUs
    "lm_head": 0,  # Must be on GPU 0
}

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-70b-hf",
    device_map=device_map,
    peft_config=lora_config,
)

• Mantenha o módulo lm_head na primeira GPU (dispositivo 0) para evitar erros
• Veja este tutorial sobre mapas de dispositivos para a configuração adequada
• Execute os scripts de treinamento diretamente (não com torchrun): python script.py
• O paralelismo de dados ainda não é suportado com NPP

Recursos

Exemplos e cadernos do TRL

• Nota de demonstração de SFT com LoRA/QLoRA - Exemplo funcional completo mostrando as implementações de LoRA e QLoRA
• Directory de Exemplos do TRL - Coleção de scripts de treinamento demonstrando PEFT com diferentes treinadores
• Receitas do TRL Cookbook - Guias passo a passo para cenários comuns de treinamento PEFT

Documentação

• Documentação do PEFT - Documentação oficial da biblioteca PEFT
• Documentação do TRL - Documentação completa do TRL com guias de treinadores
• LoRA sem arrependimentos - Boas práticas para usar o LoRA de forma eficaz

Artigos de Pesquisa

• Artigo do LoRA - Metodologia e resultados originais do LoRA
• Artigo QLoRA - Ajuste fino eficiente com quantização de 4 bits
• Artigo sobre Prompt Tuning - O poder da escala para o ajuste de prompts eficiente em parâmetros