Integración de RapidFire AI

RapidFire AI es un framework de ejecución de experimentos de código abierto que se integra con TRL para transformar el enfoque de «entrenar una configuración a la vez» en una comparación en tiempo real y lado a lado de múltiples configuraciones en las mismas GPU(s), lo que permite iterar sobre hiperparámetros, configuraciones de LoRA, esquemas de prompts y abstracciones 16–24 veces más rápido, sin hardware adicional.

Enlaces: GitHub · Documentación · Probar en Colab

¿Por qué usar RapidFire AI con TRL?

Al realizar el ajuste fino o el entrenamiento posterior con TRL, normalmente es necesario:

• Pruebe diferentes configuraciones de hiperparámetros
• Comparar diferentes configuraciones de LoRA
• Pruebe diferentes esquemas de indicaciones
• Ejecutar estudios de ablación

Cómo funciona

RapidFire AI emplea programación adaptativa basada en fragmentos:

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]

Esto permite:

• Comparación temprana de configuraciones en los mismos subconjuntos de datos de forma incremental
• Aprovechamiento eficiente de la GPU y minimización de los tiempos de inactividad
• Seguimiento en tiempo real y automatizado de métricas de experimentos
• Control dinámico sobre las ejecuciones en curso para incentivar más experimentación

Características principales

• 16-24 veces mayor rendimiento de experimentación en comparación con el entrenamiento secuencial.
• Casi sin cambios de código - simples APIs de configuración de reemplazo directo que simplemente envuelven las APIs de configuración existentes de TRL y PEFT.
• Operaciones de control interactivas - control en tiempo real para detener, reanudar y clonar-modificar (con o sin inicio en caliente) ejecuciones de entrenamiento en curso.
• Integración con Fully Sharded Data Parallel (FSDP) para entrenar modelos grandes que no caben en una sola GPU, fragmentando parámetros, gradientes y estados del optimizador entre varias GPUs.
• Compatibilidad total con transformers, PEFT, SFTTrainer, DPOTrainer y GRPOTrainer.
• Seguimiento de experimentos intercambiable: MLflow (predeterminado), TensorBoard y Trackio, habilitados de forma individual o en combinación.
• Soporte para Google Colab sin configuración: cuadernos de tutorial con un solo clic para SFT, DPO y GRPO en GPUs T4 gratuitas.
• Listo para producción: Ya se utiliza en entornos de producción con ejemplos de funcionamiento completos.

Instalación

Prerrequisitos

• Python 3.12.x
• GPU de NVIDIA con capacidad de cálculo 7.x o 8.x (se requieren varias GPUs para FSDP)
• CUDA Toolkit 11.8+
• PyTorch 2.8+

pip install

pip install rapidfireai

Una vez instalado, autentícate con Hugging Face e inicializa 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

El panel de control estará disponible en http://localhost:8853, donde puedes monitorear y controlar experimentos en tiempo real.

Inicio rápido: Entrenamiento SFT con múltiples configuraciones

Aquí hay un ejemplo completo que muestra cómo entrenar varias configuraciones de SFT de forma concurrente:

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()

Qué sucede durante la ejecución

Cuando ejecutas este ejemplo:

1. Expansión de la configuración: 2 configuraciones base × 2 configuraciones PEFT = 4 ejecuciones de entrenamiento en total
2. Programación por bloques: Los datos de entrenamiento se dividen en bloques y los 4 configuraciones se entrenan de forma concurrente
3. Intercambio de GPU: Los modelos se intercambian dentro y fuera de la memoria de GPU según los límites de los fragmentos
4. Seguimiento en tiempo real: Todas las métricas visibles en el panel de control en http://localhost:8853
5. Control interactivo: detenga, reanude o clone y modifique cualquier configuración desde el panel de control

Esto ofrece un rendimiento 16-24 veces superior en comparación con el entrenamiento de cada configuración de forma secuencial.

Entrenadores de TRL admitidos

SFTTrainer

Use RFSFTConfig como reemplazo directo de 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
)

Cuaderno de ejemplo: SFT para atención al cliente

DPOTrainer

Use RFDPOConfig como reemplazo directo de 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
)

Cuaderno de ejemplo: DPO para alineación de preferencias

GRPOTrainer

Use RFGRPOConfig como reemplazo directo de 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
)

Cuaderno de ejemplo: GRPO para razonamiento matemático

Conceptos principales

Entrenamiento concurrente basado en fragmentos

RapidFire AI divide los datos de entrenamiento en fragmentos y alterna entre configuraciones:

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]
...

Este enfoque maximiza la utilización de la GPU y permite la comparación temprana de configuraciones, manteniendo la estabilidad del entrenamiento mediante el checkpointing automático.

Operaciones de control interactivo (IC Ops)

A través del panel de RapidFire AI, puede controlar dinámicamente los experimentos en ejecución:

• Detener: Pausar una configuración (se realiza un punto de control automáticamente)
• Reanudar: Continuar desde el último punto de control
• Clonar y modificar: Duplica una configuración con modificaciones (la nueva ejecución comienza desde cero)
• Clonar y modificar con inicio cálido: Clonar, modificar e inicializar a partir de los pesos del modelo padre
• Eliminar: Elimina ejecuciones fallidas o no deseadas

Esto permite la experimentación adaptativa, donde puedes detener configuraciones con bajo rendimiento de forma temprana y clonar las prometedoras con hiperparámetros ajustados.

Experimentación con múltiples configuraciones

Use RFGridSearch o RFRandomSearch para generar automáticamente combinaciones de configuración:

# 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)

Funciones avanzadas

Integración de PEFT/LoRA

Soporte completo para el 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"
)

Funciones de recompensa personalizadas (GRPO)

Defina varias funciones de recompensa para el entrenamiento de 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
)

Soporte para múltiples GPU

RapidFire AI detecta y utiliza automáticamente todas las GPUs disponibles. De forma predeterminada, el programador distribuye las configuraciones independientes entre las GPUs (paralelismo de datos entre configuraciones), por lo que no se requiere ninguna configuración especial para ejecutar N configuraciones en N GPUs de forma concurrente.

Para los modelos que no caben en una sola GPU, RapidFire AI también admite Fully Sharded Data Parallel (FSDP) para dividir una única configuración entre varias GPUs; consulte la siguiente sección.

Entrenamiento multi-GPU con FSDP

Cuando un modelo es demasiado grande para una sola GPU, active FSDP directamente a través de los argumentos de entrenamiento de RFSFTConfig o RFDPOConfig: los mismos campos fsdp y fsdp_config que expone TrainingArguments de 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"},
)

Puntos clave:

• FSDP funciona de forma transparente con la programación basada en fragmentos de RapidFire AI, las operaciones de IC (detener / reanudar / clonar y modificar, con o sin inicio en caliente) y todos los backends de seguimiento de métricas admitidos.
• FSDP es totalmente compatible con PEFT / LoRA: los pesos de los adaptadores LoRA se recopilan de manera eficiente entre fragmentos al guardar puntos de control.
• FSDP se combina con la búsqueda en cuadrícula y la búsqueda aleatoria: cada configuración expandida obtiene su propia ejecución de entrenamiento fragmentada.

Cuadernos de ejemplo:

• SFT con FSDP (lite, modelo pequeño)
• SFT con FSDP (modelo grande)
• DPO con FSDP

Backends de seguimiento de experimentos

RapidFire AI admite tres backends de registro de métricas que pueden usarse de forma individual o conjunta: MLflow (el predeterminado para instalaciones locales), TensorBoard (el predeterminado en Google Colab) y Trackio.

Seleccione uno o más backends al iniciar el servidor con la bandera --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

También están disponibles variables de entorno equivalentes:

• RF_MLFLOW_ENABLED (por defecto true, o false en Colab)
• RF_TENSORBOARD_ENABLED (por defecto false, o true en Colab)
• RF_TRACKIO_ENABLED (por defecto false)

Los tres backends reciben las mismas métricas (pérdida, puntuaciones de evaluación, tasa de aprendizaje, etc.) y respetan los eventos del ciclo de vida de ejecución de IC Ops, por lo que puedes usar, por ejemplo, Trackio para compartir de forma ligera junto con MLflow para un panel local completo.

Ejecución en Google Colab

RapidFire AI se ejecuta en las GPU T4 gratuitas de Google Colab, con cuadernos de tutorial para flujos de trabajo de SFT, DPO, GRPO y RAG / ingeniería de contexto. En Colab, TensorBoard es el backend de seguimiento predeterminado (MLflow está deshabilitado por simplicidad) y los comandos habituales rapidfireai init / rapidfireai start se ejecutan directamente desde las celdas del cuaderno, sin necesidad de acceso a la terminal.

Comience: RapidFire AI en Google Colab.

Buenas prácticas

Ajuste de la granularidad de los fragmentos

El parámetro num_chunks controla la frecuencia de intercambio:

# 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)

Regla general: Comience con num_chunks=4 y ajústelo según el tamaño del conjunto de datos y el número de configuraciones.

Gestión de memoria

Para modelos grandes, usa cuantización:

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 rendimiento

Según los benchmarks internos que comparan el entrenamiento secuencial frente al entrenamiento concurrente de RapidFire AI:

Benchmarks realizados en NVIDIA A100 40GB con los modelos TinyLlama-1.1B y Llama-3.2-1B

Solución de problemas

Para obtener orientación sobre la solución de problemas, consulte la RapidFire AI Troubleshooting Guide.

Recursos adicionales

• Colab Notebook: RapidFire AI en Google Colab
• Documentación: oss-docs.rapidfire.ai
• GitHub: RapidFireAI/rapidfireai
• PyPI: pypi.org/project/rapidfireai
• Discord: Únete a nuestro Discord
• Cuadernos de tutorial: Repositorio de GitHub

Obtén más información sobre RapidFire AI en su repositorio oficial y documentación.