RapidFire AI 集成

RapidFire AI 是一个开源的实验执行框架，它与 TRL 集成，将“一次训练一个配置”转变为在同一 GPU 上对多个配置进行实时并排比较——因此，您可以在不增加额外硬件的情况下，以快 16–24 倍的速度迭代超参数、LoRA 设置、提示词方案和消融实验。

链接：GitHub · 文档 · 在 Colab 中试用

为什么将 RapidFire AI 与 TRL 配合使用？

在使用 TRL 进行微调或后训练时，你通常需要：

尝试不同的超参数配置
比较不同的 LoRA 设置
测试不同的提示方案
运行消融实验

场景：在同一 GPU 上比较 N 个训练配置	仅使用 TRL	TRL + RapidFire AI
训练策略	按顺序运行 N 个配置	并发运行 N 个配置
何时可以比较配置？	所有运行完成后	实时，从第一个块开始
在训练中途停止表现不佳的模型 / 克隆表现优异的模型	No	是（交互式控制操作）

工作原理

RapidFire AI 采用自适应分块调度：

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]

这可以实现：

在相同数据子集上对配置进行早期增量比较
高效的 GPU 利用率并最小化空闲时间
实时且自动化的实验指标跟踪
动态控制正在运行的任务，以鼓励更多实验

主要功能

与顺序训练相比，实验吞吐量提高了 16-24 倍。
几乎无需修改代码 - 简单的即插即用配置 API，直接封装现有的 TRL 和 PEFT 配置 API。
交互式控制操作 - 实时控制正在运行的训练任务，支持停止、恢复以及克隆修改（可选择是否使用热启动）。
通过与 Fully Sharded Data Parallel (FSDP) 集成，将参数、梯度和优化器状态分片到多个 GPU 上，以训练无法在单个 GPU 上容纳的大模型。
完全兼容 Transformers、PEFT、SFTTrainer、DPOTrainer 和 GRPOTrainer。
可插拔的实验跟踪：MLflow（默认）、TensorBoard 和 Trackio，可单独启用或组合使用。
零配置 Google Colab 支持：在免费 T4 GPU 上提供 SFT、DPO 和 GRPO 的一键式教程笔记本。
生产就绪：已在生产环境中使用，并附有完整的可运行示例。

安装

先决条件

Python 3.12.x
具有 7.x 或 8.x 计算能力的 NVIDIA GPU（FSDP 需要多 GPU）
CUDA Toolkit 11.8+
PyTorch 2.8+

pip install

pip install rapidfireai

安装完成后，使用 Hugging Face 进行身份验证并初始化 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

仪表板将在 http://localhost:8853 上可用，您可以在此实时监控和控制实验。

快速入门：使用多个配置进行 SFT 训练

下面是一个完整的示例，展示了如何同时训练多个 SFT 配置：

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

执行期间发生了什么

运行此示例时：

配置扩展：2 个基础配置 × 2 个 PEFT 配置 = 共 4 次训练运行
基于分块的调度：训练数据被划分为多个分块，所有 4 个配置并发训练
GPU 交换：模型根据块边界在 GPU 内存中换入/换出
实时跟踪：在 http://localhost:8853 的仪表板中可见的所有指标
交互式控制：从仪表板停止、恢复或克隆并修改任何配置

与依次训练每个配置相比，这带来了 16-24 倍更高的吞吐量！

支持的 TRL 训练器

SFTTrainer

将 RFSFTConfig 作为 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
)

示例 Notebook：用于客户支持的 SFT

DPOTrainer

将 RFDPOConfig 作为 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
)

示例 Notebook：用于偏好对齐的 DPO

GRPOTrainer

将 RFGRPOConfig 作为 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
)

示例 Notebook：用于数学推理的 GRPO

核心概念

基于分块的并发训练

RapidFire AI 将训练数据划分为多个块，并在不同配置之间交替使用：

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

这种方法能够最大化 GPU 利用率，并支持在训练过程中提前比较不同配置，同时通过自动检查点机制保持训练的稳定性。

交互式控制操作（IC Ops）

通过 RapidFire AI 仪表板，您可以动态控制正在运行的实验：

停止：暂停配置（会自动创建检查点）
恢复：从上一个检查点继续
克隆并修改：复制并修改配置（新运行从零开始）
使用热启动进行克隆-修改：克隆-修改并从父模型的权重进行初始化
删除：移除失败或不再需要的运行

这使得自适应实验成为可能，您可以提前停止表现不佳的配置，并克隆那些有前景的配置，同时对其超参数进行微调。

多配置实验

使用 RFGridSearch 或 RFRandomSearch 自动生成配置组合：

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

高级功能

PEFT/LoRA 集成

全面支持参数高效微调：

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

自定义奖励函数（GRPO）

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

多 GPU 支持

RapidFire AI 会自动检测并利用所有可用的 GPU。默认情况下，调度器会将独立的配置分配到各个 GPU 上（在配置之间进行数据并行），因此无需特殊设置即可在 N 个 GPU 上并发运行 N 个配置。

对于无法在单个 GPU 上运行的模型，RapidFire AI 还支持 Fully Sharded Data Parallel (FSDP)，可将单个配置分片到多个 GPU 上——请参阅下一节。

使用 FSDP 进行多 GPU 训练

当模型对于单个 GPU 而言过大时，可以直接通过 RFSFTConfig 或 RFDPOConfig 的训练参数启用 FSDP —— 即 Hugging Face TrainingArguments 所暴露的相同的 fsdp 和 fsdp_config 字段：

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

要点：

FSDP 与 RapidFire AI 的基于块的调度、IC Ops（带或不带热启动的停止 / 恢复 / 克隆-修改）以及所有受支持的指标跟踪后端透明地协同工作。
FSDP 与 PEFT / LoRA 完全兼容——在保存检查点时，LoRA 适配器权重会在各个分片之间高效收集。
FSDP 可与网格搜索和随机搜索组合使用：每个展开的配置都会获得一个独立的分片训练运行。

示例笔记本：

实验跟踪后端

RapidFire AI 支持三种指标日志记录后端，可以单独使用或组合使用：MLflow（本地安装的默认选项）、TensorBoard（Google Colab 中的默认选项）以及 Trackio。

在服务器启动时，使用 --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

还提供了等效的环境变量：

RF_MLFLOW_ENABLED（默认值为 true，在 Colab 中为 false）
RF_TENSORBOARD_ENABLED（默认值为 false，在 Colab 中为 true）
RF_TRACKIO_ENABLED（默认值为 false）

所有三种后端都会接收相同的指标（损失、评估分数、学习率等），并遵循 IC Ops 运行生命周期事件，因此您可以例如使用 Trackio 进行轻量级共享，同时使用 MLflow 提供完整的本地仪表板。

在 Google Colab 中运行

RapidFire AI 运行在免费的 Google Colab T4 GPU 上，并提供用于 SFT、DPO、GRPO 以及 RAG / 上下文工程工作流的教程笔记本。在 Colab 中，TensorBoard 是默认的跟踪后端（为简化起见，MLflow 已被禁用），并且通常的 rapidfireai init / rapidfireai start 命令可以直接从笔记本单元格中运行——无需终端访问权限。

入门指南：在 Google Colab 中使用 RapidFire AI。

最佳实践

调整块粒度

num_chunks 参数控制交换频率：

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

经验法则：从 num_chunks=4 开始，并根据数据集大小和配置数量进行调整。

内存管理

对于大型模型，请使用量化：

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",
}

性能基准

基于内部基准测试，比较了顺序训练与 RapidFire AI 并发训练：

场景	顺序时间	RapidFire AI 时间	加速
4 个配置，1 个 GPU	120 分钟	7.5 分钟	16×
8 个配置，1 个 GPU	240 分钟	12 分钟	20×
4 个配置，2 个 GPU	60 分钟	4 分钟	15×
8 个配置，4 个 GPU	60 分钟	3 分钟	20×

在 NVIDIA A100 40GB 上使用 TinyLlama-1.1B 和 Llama-3.2-1B 模型进行的基准测试

故障排除

有关故障排除指南，请参阅 RapidFire AI Troubleshooting Guide。

其他资源

Colab Notebook：在 Google Colab 中使用 RapidFire AI
文档：oss-docs.rapidfire.ai
GitHub：RapidFireAI/rapidfireai
PyPI：pypi.org/project/rapidfireai
Discord：加入我们的 Discord
教程笔记本：GitHub 仓库

在官方仓库和文档中了解有关 RapidFire AI 的更多信息。

Update on GitHub