Интеграция RapidFire AI

RapidFire AI — это open-source-фреймворк для выполнения экспериментов, который интегрируется с TRL и превращает подход «обучение одной конфигурации за раз» в сравнение множества конфигураций в реальном времени на одних и тех же GPU — так вы можете итеративно подбирать гиперпараметры, настройки LoRA, схемы промптов и проводить абляции в 16–24 раза быстрее без дополнительного оборудования.

Ссылки: GitHub · Документация · Попробовать в Colab

Зачем использовать RapidFire AI вместе с TRL?

При дообучении или пост-обучении с помощью TRL обычно необходимо:

• Попробуйте различные конфигурации гиперпараметров
• Сравнение различных настроек LoRA
• Тестирование различных схем промптов
• Проведение исследований с абляцией

Как это работает

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, которые просто оборачивают существующие конфигурационные API TRL и PEFT.
• Интерактивные операции управления — управление в реальном времени для остановки, возобновления и клонирования с модификацией (с предварительным прогревом или без него) тренировочных прогонов в процессе выполнения.
• Интеграция с Fully Sharded Data Parallel (FSDP) для обучения больших моделей, которые не помещаются на одном GPU, за счёт распределения параметров, градиентов и состояний оптимизатора по нескольким GPU.
• Полная совместимость с transformers, PEFT, SFTTrainer, DPOTrainer и GRPOTrainer.
• Подключаемый трекинг экспериментов: MLflow (по умолчанию), TensorBoard и Trackio, которые можно включать по отдельности или в комбинации.
• Поддержка Google Colab без настройки: учебные ноутбуки в один клик для SFT, DPO и GRPO на бесплатных GPU T4.
• Готовность к продакшену: Уже используется в производственных средах с полностью рабочими примерами.

Установка

Предварительные требования

• Python 3.12.x
• NVIDIA GPU с вычислительной способностью 7.x или 8.x (для 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()

Что происходит во время выполнения

При запуске этого примера:

1. Расширение конфигураций: 2 базовые конфигурации × 2 конфигурации PEFT = 4 всего запусков обучения
2. Пакетное планирование: обучающие данные разделяются на пакеты, и все 4 конфигурации обучаются параллельно
3. Обмен данных с GPU: модели загружаются и выгружаются из памяти GPU на основе границ блоков
4. Отслеживание в реальном времени: все метрики отображаются на панели управления по адресу http://localhost:8853
5. Интерактивное управление: остановка, возобновление или клонирование с модификацией любой конфигурации из панели управления

Это обеспечивает в 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
)

Пример блокнота: 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
)

Пример блокнота: 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
)

Пример блокнота: 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 вы можете динамически управлять запускаемыми экспериментами:

• Остановить: приостановить конфигурацию (чекпоинт сохраняется автоматически)
• Возобновление: Продолжить с последней контрольной точки
• Clone-Modify: Дублирование конфигурации с изменениями (новый запуск начинается с нуля)
• Clone-Modify с Warm Start: клонирование и модификация с инициализацией из весов родительской модели
• Удалить: удалить неудачные или ненужные запуски

Это позволяет проводить адаптивные эксперименты, при которых вы можете останавливать плохо работающие конфигурации на ранней стадии и клонировать перспективные, внося корректировки в гиперпараметры.

Эксперименты с несколькими конфигурациями

Используйте 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 (data-parallel между конфигурациями), поэтому для одновременного запуска N конфигураций на N GPU не требуется специальной настройки.

Для моделей, которые не помещаются на одном GPU, RapidFire AI также поддерживает Fully Sharded Data Parallel (FSDP) для распределения одной конфигурации по нескольким GPU — см. следующий раздел.

Обучение на нескольких GPU с FSDP

Если модель слишком велика для одного GPU, включите FSDP непосредственно через аргументы обучения RFSFTConfig или RFDPOConfig — те же поля fsdp и fsdp_config, которые предоставляет TrainingArguments 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"},
)

Ключевые моменты:

• FSDP прозрачно работает с чанк-ориентированным планированием RapidFire AI, IC Ops (остановка / возобновление / клонирование с модификацией с warm-starting или без него) и всеми поддерживаемыми бэкендами отслеживания метрик.
• FSDP полностью совместим с PEFT / LoRA — веса адаптеров LoRA эффективно собираются по всем шардам при сохранении контрольных точек.
• FSDP сочетается с поиском по сетке и случайным поиском: каждая развернутая конфигурация получает собственный запуск обучения с шардированием.

Примеры блокнотов:

• SFT с FSDP (lite, небольшая модель)
• SFT с FSDP (большая модель)
• DPO с 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, или false в Colab)
• RF_TENSORBOARD_ENABLED (по умолчанию false, или true в Colab)
• RF_TRACKIO_ENABLED (по умолчанию false)

Все три бэкенда получают одни и те же метрики (loss, оценки при валидации, learning rate и т. д.) и учитывают события жизненного цикла запуска IC Ops, поэтому, например, можно использовать Trackio для лёгкого обмена данными наряду с MLflow для полноценной локальной панели управления.

Запуск в Google Colab

RapidFire AI работает на бесплатных GPU T4 в Google Colab и предоставляет учебные ноутбуки для рабочих процессов SFT, DPO, GRPO и RAG / контекстной инженерии. В Colab по умолчанию используется TensorBoard в качестве бэкенда отслеживания (MLflow отключён для простоты), а стандартные команды rapidfireai init / rapidfireai start выполняются непосредственно из ячеек ноутбука — доступ к терминалу не требуется.

Начните работу: RapidFire AI в Google Colab.

Лучшие практики

Настройка гранулярности чанков

Параметр 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:

Бенчмарки, выполненные на NVIDIA A100 40GB с моделями TinyLlama-1.1B и Llama-3.2-1B

Устранение неполадок

По вопросам устранения неполадок см. RapidFire AI Troubleshooting Guide.

Дополнительные ресурсы

• Ноутбук Colab: RapidFire AI в Google Colab
• Документация: oss-docs.rapidfire.ai
• GitHub: RapidFireAI/rapidfireai
• PyPI: pypi.org/project/rapidfireai
• Discord: Присоединяйтесь к нашему Discord
• Ноутбуки с учебными материалами: Репозиторий GitHub

Подробнее о RapidFire AI можно узнать в их официальном репозитории и документации.