Intégration de RapidFire AI

RapidFire AI est un framework open source d’exécution d’expériences qui s’intègre à TRL pour transformer l’approche « entraîner une configuration à la fois » en comparaison en temps réel et côte à côte de nombreuses configurations sur les mêmes GPU(s) — vous permettant ainsi d’itérer sur les hyperparamètres, les réglages LoRA, les schémas de prompts et les ablations 16 à 24 fois plus rapidement, sans matériel supplémentaire.

Liens : GitHub · Docs · Essayer dans Colab

Pourquoi utiliser RapidFire AI avec TRL ?

Lors de l'ajustement fin ou de l'entraînement postérieur avec TRL, vous avez généralement besoin de :

• Essayer différentes configurations d'hyperparamètres
• Comparer différents paramètres LoRA
• Tester différents schémas de prompts
• Réaliser des études d'ablation

Fonctionnement

RapidFire AI utilise un ordonnancement adaptatif par blocs :

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]

Cela permet :

• Comparaison préliminaire des configurations sur des sous-ensembles de données identiques, de manière incrémentale
• Optimisation de l'utilisation du GPU et minimisation des temps d'inactivité
• Suivi en temps réel et automatisé des métriques d'expérimentation
• Contrôle dynamique des exécutions en cours pour encourager davantage l'expérimentation

Fonctionnalités clés

• 16 à 24 fois plus de débit d'expérimentation par rapport à l'entraînement séquentiel.
• Presque aucune modification de code - des API de configuration simples à intégrer qui encapsulent simplement les API de configuration existantes de TRL et PEFT.
• Opérations de contrôle interactif - contrôle en temps réel pour arrêter, reprendre et cloner-modifier (avec ou sans amorçage) des exécutions d'entraînement en cours.
• Intégration avec Fully Sharded Data Parallel (FSDP) pour l’entraînement de grands modèles qui ne tiennent pas sur une seule GPU, en répartissant les paramètres, les gradients et les états de l’optimiseur sur plusieurs GPU.
• Compatibilité totale avec Transformers, PEFT, SFTTrainer, DPOTrainer et GRPOTrainer.
• Suivi d'expériences modulaire : MLflow (par défaut), TensorBoard et Trackio, activables individuellement ou en combinaison.
• Prise en charge de Google Colab sans configuration : des notebooks de tutoriels en un clic pour SFT, DPO et GRPO sur des GPU T4 gratuits.
• Prêt pour la production : Déjà utilisé dans des environnements de production avec des exemples fonctionnels complets.

Installation

Prérequis

• Python 3.12.x
• GPU NVIDIA avec une capacité de calcul 7.x ou 8.x (plusieurs GPU requis pour FSDP)
• CUDA Toolkit 11.8+
• PyTorch 2.8+

pip install

pip install rapidfireai

Une fois installé, authentifiez-vous auprès de Hugging Face et initialisez 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

Le tableau de bord sera disponible à l’adresse http://localhost:8853, où vous pourrez surveiller et contrôler les expériences en temps réel.

Démarrage rapide : entraînement SFT avec plusieurs configurations

Voici un exemple complet montrant comment entraîner plusieurs configurations SFT en parallèle :

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

Ce qui se passe pendant l'exécution

Lorsque vous exécutez cet exemple :

1. Extension de la configuration : 2 configurations de base × 2 configurations PEFT = 4 exécutions d'entraînement au total
2. Planification par blocs : les données d'entraînement sont divisées en blocs, et les 4 configurations s'entraînent simultanément
3. Échange sur GPU : Les modèles sont chargés/déchargés de la mémoire GPU en fonction des limites de blocs
4. Suivi en temps réel : toutes les métriques sont visibles dans le tableau de bord à l'adresse http://localhost:8853
5. Contrôle interactif : arrêter, reprendre ou cloner-modifier n'importe quelle configuration depuis le tableau de bord

Cela offre un débit 16 à 24 fois supérieur par rapport à l'entraînement de chaque configuration de manière séquentielle !

Formateurs TRL pris en charge

SFTTrainer

Utilisez RFSFTConfig comme remplacement direct 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
)

Carnet d'exemple : SFT pour le support client

DPOTrainer

Utilisez RFDPOConfig comme remplacement direct 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
)

Carnet d'exemples : DPO pour l'alignement des préférences

GRPOTrainer

Utilisez RFGRPOConfig comme remplacement direct 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
)

Carnet d'exemple : GRPO pour le raisonnement mathématique

Concepts fondamentaux

Entraînement concurrentiel par blocs

RapidFire AI divise les données d'entraînement en segments et alterne entre les configurations :

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

Cette approche maximise l'utilisation des GPU et permet une comparaison précoce des configurations tout en maintenant la stabilité de l'entraînement grâce à la sauvegarde automatique des points de contrôle.

Opérations de contrôle interactif (IC Ops)

Grâce au tableau de bord RapidFire AI, vous pouvez contrôler dynamiquement les expériences en cours :

• Arrêter : Mettre en pause une configuration (point de contrôle enregistré automatiquement)
• Reprendre : Continuer depuis le dernier point de contrôle
• Clone-Modify : Dupliquer une configuration avec des modifications (la nouvelle exécution repart de zéro)
• Clone-Modify avec démarrage chaud : Clone-modify et initialisation à partir des poids du parent
• Supprimer : Supprimer les exécutions échouées ou non souhaitées

Cela permet une expérimentation adaptative où vous pouvez arrêter tôt les configurations sous-performantes et cloner les prometteuses avec des hyperparamètres ajustés.

Expérimentation multi-configuration

Utilisez RFGridSearch ou RFRandomSearch pour générer automatiquement des combinaisons de configurations :

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

Fonctions avancées

Intégration PEFT/LoRA

Prise en charge complète du fine-tuning efficace en termes de paramètres :

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

Fonctions de récompense personnalisées (GRPO)

Définir plusieurs fonctions de récompense pour l'entraînement 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
)

Prise en charge multi-GPU

RapidFire AI détecte et exploite automatiquement tous les GPU disponibles. Par défaut, le planificateur répartit les configurations indépendantes sur les GPU (parallélisme de données entre les configurations), de sorte qu’aucune configuration particulière n’est nécessaire pour exécuter N configurations sur N GPU simultanément.

Pour les modèles qui ne tiennent pas sur un seul GPU, RapidFire AI prend également en charge Fully Sharded Data Parallel (FSDP) pour répartir une seule configuration sur plusieurs GPU — voir la section suivante.

Entraînement multi-GPU avec FSDP

Lorsqu’un modèle est trop volumineux pour une seule GPU, activez FSDP directement via les arguments d’entraînement de RFSFTConfig ou RFDPOConfig — les mêmes champs fsdp et fsdp_config exposés par 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"},
)

Points clés :

• FSDP fonctionne de manière transparente avec la planification par lots de RapidFire AI, les opérations IC (arrêt / reprise / clonage-modification avec ou sans amorçage) et tous les backends de suivi des métriques pris en charge.
• FSDP est entièrement compatible avec PEFT / LoRA — les poids des adaptateurs LoRA sont collectés efficacement entre les shards lors de l'enregistrement des points de contrôle.
• FSDP se combine avec la recherche sur grille et la recherche aléatoire : chaque configuration dépliée reçoit sa propre exécution d'entraînement fractionné.

Carnets d'exemples :

• SFT avec FSDP (lite, petit modèle)
• SFT avec FSDP (grand modèle)
• DPO avec FSDP

Backends de suivi des expériences

RapidFire AI prend en charge trois backends de journalisation des métriques, utilisables individuellement ou ensemble : MLflow (par défaut pour les installations locales), TensorBoard (par défaut dans Google Colab) et Trackio.

Sélectionnez un ou plusieurs backends au démarrage du serveur avec le drapeau --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

Des variables d'environnement équivalentes sont également disponibles :

• RF_MLFLOW_ENABLED (par défaut true, ou false dans Colab)
• RF_TENSORBOARD_ENABLED (par défaut false, ou true dans Colab)
• RF_TRACKIO_ENABLED (par défaut false)

Les trois backends reçoivent les mêmes métriques (perte, scores d'évaluation, taux d'apprentissage, etc.) et respectent les événements du cycle de vie d'une exécution IC Ops, ce qui vous permet, par exemple, d'utiliser Trackio pour un partage léger en parallèle de MLflow pour un tableau de bord local complet.

Exécution dans Google Colab

RapidFire AI fonctionne sur des GPU T4 Google Colab gratuits, avec des notebooks de tutoriel pour les flux de travail SFT, DPO, GRPO et RAG / ingénierie du contexte. Dans Colab, TensorBoard est le backend de suivi par défaut (MLflow est désactivé pour simplifier), et les commandes habituelles rapidfireai init / rapidfireai start s’exécutent directement depuis les cellules du notebook — aucun accès au terminal n’est requis.

Commencer : RapidFire AI dans Google Colab.

Bonnes pratiques

Réglage de la granularité des blocs

Le paramètre num_chunks contrôle la fréquence d'échange :

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

Règle générale : commencez avec num_chunks=4 et ajustez en fonction de la taille du jeu de données et du nombre de configurations.

Gestion de la mémoire

Pour les grands modèles, utilisez la quantification :

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 performance

Sur la base de benchmarks internes comparant l’entraînement séquentiel et l’entraînement concurrent de RapidFire AI :

Benchmarks effectués sur NVIDIA A100 40GB avec les modèles TinyLlama-1.1B et Llama-3.2-1B

Dépannage

Pour des conseils de dépannage, consultez le RapidFire AI Troubleshooting Guide.

Ressources supplémentaires

• Carnet Colab : RapidFire AI dans Google Colab
• Documentation : oss-docs.rapidfire.ai
• GitHub : RapidFireAI/rapidfireai
• PyPI : pypi.org/project/rapidfireai
• Discord : Rejoignez notre Discord
• Carnets de tutoriel : Dépôt GitHub

En savoir plus sur RapidFire AI dans leur dépôt officiel et leur documentation.