Integração do RapidFire AI

RapidFire AI é um framework de execução de experimentos de código aberto que se integra ao TRL para transformar o “treinar uma configuração por vez” em comparação em tempo real, lado a lado, de muitas configurações nas mesmas GPU(s) — permitindo iterar em hiperparâmetros, configurações de LoRA, esquemas de prompt e ablações 16–24× mais rápido, sem hardware adicional.

Links: GitHub · Documentação · Experimente no Colab

Por que usar RapidFire AI com TRL?

Ao fazer fine-tuning ou pós-treino com TRL, você geralmente precisa:

• Experimente diferentes configurações de hiperparâmetros
• Compare diferentes configurações de LoRA
• Teste diferentes esquemas de prompt
• Execute estudos de ablação

Como Funciona

A RapidFire AI utiliza agendamento adaptativo baseado em blocos:

GPU Timeline (Single GPU):
Chunk 1: [Config A] → [Config B] → [Config C] → [Config D]
Chunk 2: [Config A] → [Config B] → [Config C] → [Config D]
Chunk 3: [Config A] → [Config B] → [Config C] → [Config D]

Isso permite:

• Comparação inicial de configurações em subconjuntos de dados iguais de forma incremental
• Aproveitamento eficiente da GPU e minimização de tempos ociosos
• Rastreamento de métricas de experimentos em tempo real e automatizado
• Controle dinâmico sobre execuções em andamento para incentivar mais experimentação

Principais recursos

• 16-24× maior throughput de experimentação em comparação ao treinamento sequencial.
• Quase nenhuma alteração de código - APIs de configuração simples que apenas envolvem as APIs de configuração existentes do TRL e do PEFT.
• Operações de controle interativo - controle em tempo real para interromper, retomar e clonar/modificar (com ou sem inicialização quente) execuções de treinamento em andamento.
• Integração com Fully Sharded Data Parallel (FSDP) para treinar modelos grandes que não cabem em uma única GPU, distribuindo parâmetros, gradientes e estados do otimizador entre várias GPUs.
• Compatibilidade total com transformers, PEFT, SFTTrainer, DPOTrainer e GRPOTrainer.
• Rastreamento de experimentos plugável: MLflow (padrão), TensorBoard e Trackio, habilitados individualmente ou em combinação.
• Suporte ao Google Colab sem configuração: notebooks de tutorial com um clique para SFT, DPO e GRPO em GPUs T4 gratuitas.
• Pronto para produção: Já utilizado em ambientes de produção com exemplos funcionais completos.

Instalação

Pré-requisitos

• Python 3.12.x
• GPU da NVIDIA com Compute Capability 7.x ou 8.x (múltiplas GPUs são necessárias para FSDP)
• CUDA Toolkit 11.8+
• PyTorch 2.8+

pip install

pip install rapidfireai

Depois de instalado, autentique-se no Hugging Face e inicialize o RapidFire AI:

# Authenticate with Hugging Face
hf auth login --token YOUR_TOKEN

# Workaround for current issue: https://github.com/huggingface/xet-core/issues/527
pip uninstall -y hf-xet

# Initialize RapidFire AI
rapidfireai init

# Start the RapidFire AI server
rapidfireai start

O painel estará disponível em http://localhost:8853, onde você pode monitorar e controlar experimentos em tempo real.

Início rápido: Treinamento SFT com múltiplas configurações

Aqui está um exemplo completo mostrando como treinar várias configurações de SFT simultaneamente:

from rapidfireai import Experiment
from rapidfireai.automl import List, RFGridSearch, RFModelConfig, RFLoraConfig, RFSFTConfig
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load dataset
dataset = load_dataset("bitext/Bitext-customer-support-llm-chatbot-training-dataset")
train_dataset = dataset["train"].select(range(128)).shuffle(seed=42)
eval_dataset = dataset["train"].select(range(100, 124)).shuffle(seed=42)

# Define data formatting function
def formatting_function(row):
    return {
        "prompt": [
            {"role": "system", "content": "You are a helpful customer support assistant."},
            {"role": "user", "content": row["instruction"]},
        ],
        "completion": [
            {"role": "assistant", "content": row["response"]}
        ]
    }

# Initialize experiment
experiment = Experiment(experiment_name="sft-customer-support")

# Define multiple LoRA configurations to compare
peft_configs = List([
    RFLoraConfig(r=8, lora_alpha=16, lora_dropout=0.1,
                 target_modules=["q_proj", "v_proj"], bias="none"),
    RFLoraConfig(r=32, lora_alpha=64, lora_dropout=0.1,
                 target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], bias="none")
])

# Define multiple training configurations
# 2 base configs × 2 PEFT configs = 4 total training runs
config_set = List([
    RFModelConfig(
        model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        peft_config=peft_configs,
        training_args=RFSFTConfig(  # Wraps TRL's SFTConfig
            learning_rate=1e-3,
            per_device_train_batch_size=4,
            max_steps=128,
            fp16=True,
        ),
        model_type="causal_lm",
        model_kwargs={"device_map": "auto", "dtype": "auto", "use_cache": False},
        formatting_func=formatting_function,
    ),
    RFModelConfig(
        model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        peft_config=peft_configs,
        training_args=RFSFTConfig(
            learning_rate=1e-4,  # Different learning rate
            per_device_train_batch_size=4,
            max_steps=128,
            fp16=True,
        ),
        model_type="causal_lm",
        model_kwargs={"device_map": "auto", "dtype": "auto", "use_cache": False},
        formatting_func=formatting_function,
    )
])

# Define model creation function
def create_model(model_config):
    model = AutoModelForCausalLM.from_pretrained(
        model_config["model_name"],
        **model_config["model_kwargs"]
    )
    tokenizer = AutoTokenizer.from_pretrained(model_config["model_name"])
    return (model, tokenizer)

# Create grid search over all configurations
config_group = RFGridSearch(configs=config_set, trainer_type="SFT")

# Run all 4 configurations concurrently with chunk-based scheduling
experiment.run_fit(config_group, create_model, train_dataset, eval_dataset,
                   num_chunks=4, seed=42)

# End experiment
experiment.end()

O Que Acontece Durante a Execução

Quando você executa este exemplo:

1. Expansão de configuração: 2 configurações base × 2 configurações PEFT = 4 execuções de treinamento no total
2. Agendamento baseado em blocos: os dados de treinamento são divididos em blocos e todos os 4 configs são treinados simultaneamente
3. Alternância de GPU: Os modelos são movidos para dentro e para fora da memória da GPU com base nos limites dos blocos
4. Rastreamento em tempo real: Todas as métricas visíveis no painel em http://localhost:8853
5. Controle interativo: interrompa, retome ou clone e modifique qualquer configuração a partir do painel

Isso entrega 16-24× mais throughput em comparação ao treinamento de cada configuração sequencialmente!

Treinadores TRL suportados

SFTTrainer

Use RFSFTConfig como substituto direto para SFTConfig:

from rapidfireai.automl import RFSFTConfig

training_args = RFSFTConfig(
    learning_rate=5e-5,
    per_device_train_batch_size=4,
    num_train_epochs=3,
    max_length = 512,
    # ... all other SFTConfig parameters supported
)

Exemplo de Notebook: SFT para Suporte ao Cliente

DPOTrainer

Use RFDPOConfig como substituto direto para DPOConfig:

from rapidfireai.automl import RFDPOConfig

training_args = RFDPOConfig(
    beta=0.1,
    loss_type="sigmoid",
    max_length=1024,
    learning_rate=5e-4,
    # ... all other DPOConfig parameters supported
)

Exemplo de Notebook: DPO para alinhamento de preferências

GRPOTrainer

Use RFGRPOConfig como substituto direto para GRPOConfig:

from rapidfireai.automl import RFGRPOConfig

training_args = RFGRPOConfig(
    learning_rate=5e-6,
    num_generations=8,
    max_completion_length=256,
    # ... all other GRPOConfig parameters supported
)

Exemplo de Notebook: GRPO para Raciocínio Matemático

Conceitos Fundamentais

Treinamento Concorrente Baseado em Blocos

RapidFire AI divide os dados de treinamento em blocos e alterna entre configurações:

GPU Timeline (Single GPU):
Chunk 1: [Config A] → [Config B] → [Config C] → [Config D]
Chunk 2: [Config A] → [Config B] → [Config C] → [Config D]
Chunk 3: [Config A] → [Config B] → [Config C] → [Config D]
...

Essa abordagem maximiza a utilização da GPU e permite a comparação precoce de configurações, mantendo a estabilidade do treinamento por meio de checkpointing automático.

Operações de Controle Interativo (IC Ops)

Por meio do painel do RapidFire AI, você pode controlar dinamicamente os experimentos em execução:

• Parar: Pausa uma configuração (checkpoint automático)
• Resumo: Continuar a partir do último checkpoint
• Clone-Modify: Duplica uma configuração com modificações (a nova execução começa do zero)
• Clone-Modify com Warm Start: Clone-modify e inicialize a partir dos pesos do pai
• Excluir: Remove execuções com falha ou indesejadas

Isso permite experimentação adaptativa, na qual você pode interromper configurações com baixo desempenho precocemente e clonar as promissoras com hiperparâmetros ajustados.

Experimentação Multi-Config

Use RFGridSearch ou RFRandomSearch para gerar automaticamente combinações de configuração:

# Grid search: tests all combinations
config_group = RFGridSearch(configs=config_list, trainer_type="SFT")

# Random search: samples N configurations
config_group = RFRandomSearch(configs=config_list, trainer_type="DPO", num_samples=10)

Recursos Avançados

Integração PEFT/LoRA

Suporte completo para ajuste fino eficiente de parâmetros:

from rapidfireai.automl import RFLoraConfig
from peft import TaskType

lora_config = RFLoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=64,
    lora_alpha=64,
    lora_dropout=0.1,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    bias="none"
)

Funções de Recompensa Personalizadas (GRPO)

Defina múltiplas funções de recompensa para o treinamento do GRPO:

def correctness_reward(prompts, completions, answer, **kwargs):
    """Reward for correct answers"""
    responses = [completion[0]['content'] for completion in completions]
    extracted = [extract_answer(r) for r in responses]
    return [2.0 if r == a else 0.0 for r, a in zip(extracted, answer)]

def format_reward(completions, **kwargs):
    """Reward for proper formatting"""
    import re
    pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, r) for r in responses]
    return [0.5 if match else 0.0 for match in matches]

# Use in model config
config = RFModelConfig(
    reward_funcs=[correctness_reward, format_reward],
    # ... other parameters
)

Suporte a múltiplas GPUs

O RapidFire AI detecta e utiliza automaticamente todas as GPUs disponíveis. Por padrão, o agendador distribui configurações independentes entre as GPUs (paralelismo de dados entre configurações), portanto, nenhum ajuste especial é necessário para executar N configurações em N GPUs simultaneamente.

Para modelos que não cabem em uma única GPU, o RapidFire AI também suporta Fully Sharded Data Parallel (FSDP) para dividir uma única configuração entre várias GPUs — consulte a próxima seção.

Treinamento Multi-GPU com FSDP

Quando um modelo é grande demais para uma única GPU, ative o FSDP diretamente pelos argumentos de treinamento de RFSFTConfig ou RFDPOConfig — os mesmos campos fsdp e fsdp_config expostos pelo TrainingArguments do Hugging Face:

from rapidfireai.automl import RFModelConfig, RFSFTConfig, RFLoraConfig

model_config = RFModelConfig(
    model_name="meta-llama/Llama-3.1-8B-Instruct",
    peft_config=RFLoraConfig(
        r=16, lora_alpha=32, lora_dropout=0.05,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        bias="none",
    ),
    training_args=RFSFTConfig(
        learning_rate=2e-4,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        fsdp="full_shard auto_wrap",
        fsdp_config={
            "sharding_strategy": "FULL_SHARD",
            "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
            "backward_prefetch": "backward_pre",
            "forward_prefetch": True,
            "use_orig_params": False,
            "cpu_ram_efficient_loading": True,
            "offload_params": True,
            "sync_module_states": True,
            "limit_all_gathers": True,
        },
    ),
    model_type="causal_lm",
    model_kwargs={"dtype": "auto"},
)

Principais pontos:

• O FSDP funciona de forma transparente com o agendamento baseado em blocos da RapidFire AI, as operações IC (parar / retomar / clonar e modificar com ou sem inicialização quente) e todos os backends de rastreamento de métricas suportados.
• FSDP é totalmente compatível com PEFT / LoRA — os pesos dos adaptadores LoRA são coletados de forma eficiente entre os shards ao salvar checkpoints.
• FSDP se combina com busca em grade e busca aleatória: cada configuração expandida recebe sua própria execução de treinamento particionada.

Cadernos de exemplo:

• SFT com FSDP (lite, modelo pequeno)
• SFT com FSDP (modelo grande)
• DPO com FSDP

Backends de Rastreamento de Experimentos

O RapidFire AI suporta três backends de registro de métricas que podem ser usados individualmente ou em conjunto: MLflow (o padrão para instalações locais), TensorBoard (o padrão no Google Colab) e Trackio.

Selecione um ou mais backends na inicialização do servidor com a flag --tracking-backends:

# MLflow only (default on local installs)
rapidfireai start --tracking-backends mlflow

# TensorBoard only
rapidfireai start --tracking-backends tensorboard

# Any combination
rapidfireai start --tracking-backends mlflow tensorboard trackio

Variáveis de ambiente equivalentes também estão disponíveis:

• RF_MLFLOW_ENABLED (padrão true, ou false no Colab)
• RF_TENSORBOARD_ENABLED (padrão false, ou true no Colab)
• RF_TRACKIO_ENABLED (padrão false)

Todos os três backends recebem as mesmas métricas (loss, pontuações de avaliação, taxa de aprendizado, etc.) e respeitam os eventos de ciclo de vida de execução do IC Ops, permitindo que você use, por exemplo, o Trackio para compartilhamento leve em conjunto com o MLflow para um dashboard local completo.

Executando no Google Colab

O RapidFire AI roda em GPUs T4 gratuitas do Google Colab, com notebooks de tutorial para fluxos de trabalho de SFT, DPO, GRPO e RAG / engenharia de contexto. No Colab, o TensorBoard é o backend de rastreamento padrão (o MLflow é desativado para simplificar), e os comandos habituais rapidfireai init / rapidfireai start são executados diretamente das células do notebook — sem necessidade de acesso ao terminal.

Comece: RapidFire AI no Google Colab.

Boas Práticas

Ajustando a granularidade dos blocos

O parâmetro num_chunks controla a frequência de troca:

# Fewer chunks = less overhead, less frequent comparison
experiment.run_fit(..., num_chunks=2)

# More chunks = more overhead, more frequent comparison
experiment.run_fit(..., num_chunks=16)

Regra geral: Comece com num_chunks=4 e ajuste com base no tamanho do conjunto de dados e no número de configurações.

Gerenciamento de Memória

Para modelos grandes, use quantização:

from transformers import BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
)

model_kwargs = {
    "quantization_config": bnb_config,
    "device_map": "auto",
}

Benchmarks de Desempenho

Com base em benchmarks internos comparando o treinamento sequencial ao treinamento concorrente do RapidFire AI:

Benchmarks realizados em NVIDIA A100 40GB com os modelos TinyLlama-1.1B e Llama-3.2-1B

Solução de problemas

Para orientações de solução de problemas, consulte o RapidFire AI Troubleshooting Guide.

Recursos adicionais

• Colab Notebook: RapidFire AI no Google Colab
• Documentação: oss-docs.rapidfire.ai
• GitHub: RapidFireAI/rapidfireai
• PyPI: pypi.org/project/rapidfireai
• Discord: Participe do nosso Discord
• Cadernos de tutorial: Repositório do GitHub

Saiba mais sobre o RapidFire AI no seu repositório oficial e na documentação.