TRL documentation

Tích hợp RapidFire AI

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.18.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Tích hợp RapidFire AI

RapidFire AI là một framework thực hiện thí nghiệm mã nguồn mở, tích hợp với TRL để biến quy trình “huấn luyện từng cấu hình một” thành so sánh song song theo thời gian thực nhiều cấu hình trên cùng một GPU — giúp bạn lặp lại các tham số siêu, cài đặt LoRA, phương án prompt và các phép loại trừ nhanh hơn 16–24 lần mà không cần thêm phần cứng.

Liên kết: GitHub · Tài liệu · Thử trong Colab

Tại sao nên sử dụng RapidFire AI với TRL?

Khi tinh chỉnh hoặc huấn luyện sau với TRL, bạn thường cần:

  • Thử các cấu hình siêu tham số khác nhau
  • So sánh các cài đặt LoRA khác nhau
  • Kiểm tra các phương án prompt khác nhau
  • Chạy các nghiên cứu ablation
Kịch bản: so sánh N cấu hình huấn luyện trên cùng một GPUChỉ TRLTRL + RapidFire AI
Chiến lược huấn luyệnChạy tuần tự N cấu hìnhChạy N cấu hình đồng thời
Khi nào bạn có thể so sánh các cấu hình?Sau khi tất cả các lần chạy hoàn tấtTrực tiếp, từ khối đầu tiên
Dừng các thí nghiệm kém hiệu quả / nhân bản các thí nghiệm hiệu quả trong khi đang huấn luyệnKhôngCó (Các thao tác điều khiển tương tác)

Cách hoạt động

RapidFire AI sử dụng lịch phân bổ theo chunk thích ứng:

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]

Điều này cho phép:

  • So sánh sớm các cấu hình trên các tập dữ liệu con giống nhau theo cách tăng dần
  • Tận dụng hiệu quả GPU và giảm thiểu thời gian chờ
  • Theo dõi chỉ số thử nghiệm theo thời gian thực và tự động
  • Kiểm soát động đối với các lần chạy đang diễn ra để khuyến khích thử nghiệm nhiều hơn

Tính năng chính

  • Thông lượng thử nghiệm cao hơn 16-24 lần so với huấn luyện tuần tự.
  • Gần như không cần thay đổi mã - các API cấu hình đơn giản, có thể thay thế trực tiếp, chỉ bao quanh các API cấu hình TRL và PEFT hiện có.
  • Thao tác điều khiển tương tác - điều khiển thời gian thực để dừng, tiếp tục và nhân bản-sửa đổi (có hoặc không có khởi động ấm) các lượt huấn luyện đang chạy.
  • Tích hợp với Fully Sharded Data Parallel (FSDP) để huấn luyện các mô hình lớn không vừa trên một GPU duy nhất bằng cách chia nhỏ tham số, gradient và trạng thái optimizer trên nhiều GPU.
  • Tương thích hoàn toàn với transformers, PEFT, SFTTrainer, DPOTrainer và GRPOTrainer.
  • Theo dõi thí nghiệm có thể thay thế: MLflow (mặc định), TensorBoard và Trackio, được bật riêng lẻ hoặc kết hợp.
  • Hỗ trợ Google Colab không cần cấu hình: các notebook hướng dẫn một cú nhấp cho SFT, DPO và GRPO trên GPU T4 miễn phí.
  • Sẵn sàng cho môi trường sản xuất: Đã được sử dụng trong các môi trường sản xuất với các ví dụ hoạt động hoàn chỉnh.

Cài đặt

Yêu cầu tiên quyết

  • Python 3.12.x
  • GPU NVIDIA với Compute Capability 7.x hoặc 8.x (cần nhiều GPU cho FSDP)
  • Bộ công cụ CUDA 11.8+
  • PyTorch 2.8+

pip install

pip install rapidfireai

Sau khi cài đặt, xác thực với Hugging Face và khởi tạ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

Bảng điều khiển sẽ có sẵn tại http://localhost:8853, nơi bạn có thể theo dõi và kiểm soát các thí nghiệm theo thời gian thực.

Bắt đầu nhanh: Huấn luyện SFT với nhiều cấu hình

Dưới đây là một ví dụ hoàn chỉnh cho thấy cách huấn luyện đồng thời nhiều cấu hình 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()

Những gì xảy ra trong quá trình thực thi

Khi bạn chạy ví dụ này:

  1. Mở rộng cấu hình: 2 cấu hình cơ sở × 2 cấu hình PEFT = 4 lần chạy huấn luyện
  2. Lịch trình theo khối: Dữ liệu huấn luyện được chia thành các khối, và cả 4 cấu hình đều được huấn luyện đồng thời
  3. Đổi GPU: Các mô hình được đưa vào/ra khỏi bộ nhớ GPU dựa trên các ranh giới chunk
  4. Theo dõi thời gian thực: Tất cả các chỉ số hiển thị trên bảng điều khiển tại http://localhost:8853
  5. Điều khiển tương tác: Dừng, tiếp tục hoặc nhân bản-sửa đổi bất kỳ cấu hình nào từ bảng điều khiển

Điều này mang lại hiệu suất cao hơn 16-24 lần so với việc huấn luyện từng cấu hình một cách tuần tự!

Các Trainer TRL được hỗ trợ

SFTTrainer

Sử dụng RFSFTConfig như một thay thế trực tiếp cho 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
)

Ví dụ Notebook: SFT cho Hỗ trợ Khách hàng

DPOTrainer

Sử dụng RFDPOConfig như một thay thế trực tiếp cho 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
)

Ví dụ Notebook: DPO cho việc căn chỉnh sở thích

GRPOTrainer

Sử dụng RFGRPOConfig như một thay thế trực tiếp cho 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
)

Ví dụ Notebook: GRPO cho suy luận Toán học

Khái niệm cốt lõi

Huấn luyện đồng thời dựa trên Chunk

RapidFire AI chia dữ liệu huấn luyện thành các khối và luân phiên giữa các cấu hình:

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

Cách tiếp cận này tối đa hóa mức sử dụng GPU và cho phép so sánh các cấu hình sớm trong khi vẫn duy trì sự ổn định của quá trình huấn luyện thông qua việc lưu checkpoint tự động.

Các thao tác điều khiển tương tác (IC Ops)

Thông qua bảng điều khiển RapidFire AI, bạn có thể kiểm soát động các thí nghiệm đang chạy:

  • Dừng: Tạm dừng một cấu hình (được lưu checkpoint tự động)
  • Tiếp tục: Tiếp tục từ checkpoint gần nhất
  • Clone-Modify: Sao chép một cấu hình với các chỉnh sửa (chạy mới bắt đầu từ đầu)
  • Clone-Modify với Khởi động Ấm: Clone-modify và khởi tạo từ trọng số của mô hình cha
  • Xóa: Loại bỏ các lần chạy thất bại hoặc không mong muốn

Điều này cho phép thử nghiệm thích ứng, nơi bạn có thể dừng sớm các cấu hình hoạt động kém và nhân bản các cấu hình triển vọng với các siêu tham số đã được điều chỉnh.

Thử nghiệm Cấu hình Đa

Sử dụng RFGridSearch hoặc RFRandomSearch để tự động tạo các tổ hợp cấu hình:

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

Tính năng nâng cao

Tích hợp PEFT/LoRA

Hỗ trợ đầy đủ cho tinh chỉnh hiệu quả tham số:

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

Hàm phần thưởng tùy chỉnh (GRPO)

Định nghĩa nhiều hàm phần thưởng cho huấn luyện 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
)

Hỗ trợ đa GPU

RapidFire AI tự động phát hiện và sử dụng tất cả các GPU khả dụng. Mặc định, trình điều phối phân phối các cấu hình độc lập giữa các GPU (song song dữ liệu giữa các cấu hình), do đó không cần thiết lập đặc biệt để chạy N cấu hình trên N GPU đồng thời.

Đối với các mô hình không vừa trên một GPU duy nhất, RapidFire AI cũng hỗ trợ Fully Sharded Data Parallel (FSDP) để chia nhỏ một cấu hình duy nhất trên nhiều GPU — xem phần tiếp theo.

Đào tạo đa GPU với FSDP

Khi một model quá lớn cho một GPU đơn, hãy bật FSDP trực tiếp thông qua các training args của RFSFTConfig hoặc RFDPOConfig — cùng các trường fsdp và fsdp_config được cung cấp bởi TrainingArguments của 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"},
)

Các điểm chính:

  • FSDP hoạt động trong suốt với lịch trình dựa trên chunk của RapidFire AI, IC Ops (dừng / tiếp tục / nhân bản-sửa đổi có hoặc không có khởi động ấm) và tất cả các backend theo dõi chỉ số được hỗ trợ.
  • FSDP tương thích hoàn toàn với PEFT / LoRA — trọng số của bộ thích ứng LoRA được thu thập hiệu quả trên các shard khi lưu checkpoint.
  • FSDP kết hợp với tìm kiếm lưới và tìm kiếm ngẫu nhiên: mỗi cấu hình được mở rộng sẽ có một lượt huấn luyện phân mảnh riêng.

Ví dụ Notebook:

Các nền tảng theo dõi thí nghiệm

RapidFire AI hỗ trợ ba backend ghi log chỉ số có thể được sử dụng riêng lẻ hoặc kết hợp với nhau: MLflow (mặc định cho cài đặt cục bộ), TensorBoard (mặc định trong Google Colab) và Trackio.

Chọn một hoặc nhiều backend khi khởi động máy chủ bằng cờ --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

Các biến môi trường tương đương cũng có sẵn:

  • RF_MLFLOW_ENABLED (mặc định là true, hoặc false trong Colab)
  • RF_TENSORBOARD_ENABLED (mặc định là false, hoặc true trong Colab)
  • RF_TRACKIO_ENABLED (mặc định false)

Cả ba backend đều nhận cùng các chỉ số (loss, điểm đánh giá, learning rate, v.v.) và tuân thủ các sự kiện vòng đời chạy của IC Ops, vì vậy bạn có thể, ví dụ, sử dụng Trackio để chia sẻ nhẹ nhàng cùng với MLflow cho một dashboard cục bộ đầy đủ.

Chạy trong Google Colab

RapidFire AI chạy trên GPU T4 miễn phí của Google Colab, kèm theo các notebook hướng dẫn cho các quy trình SFT, DPO, GRPO và RAG / kỹ thuật ngữ cảnh. Trong Colab, TensorBoard là backend theo dõi mặc định (MLflow bị vô hiệu hóa để đơn giản hóa), và các lệnh thông thường rapidfireai init / rapidfireai start được chạy trực tiếp từ các ô notebook — không cần truy cập terminal.

Bắt đầu: RapidFire AI trong Google Colab.

Thực hành tốt nhất

Điều chỉnh độ hạt của Chunk

Tham số num_chunks kiểm soát tần suất hoán đổi:

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

Quy tắc chung: Bắt đầu với num_chunks=4 và điều chỉnh dựa trên kích thước bộ dữ liệu và số lượng cấu hình.

Quản lý bộ nhớ

Đối với các mô hình lớn, hãy sử dụng lượng tử hóa:

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 Hiệu năng

Dựa trên các bài kiểm tra hiệu năng nội bộ so sánh huấn luyện tuần tự và huấn luyện đồng thời của RapidFire AI:

Kịch bảnThời gian tuần tựThời gian RapidFire AITốc độ tăng
4 cấu hình, 1 GPU120 phút7,5 phút16×
8 cấu hình, 1 GPU240 phút12 phút20×
4 cấu hình, 2 GPU60 phút4 phút15×
8 cấu hình, 4 GPU60 phút3 phút20×

Các bài kiểm tra hiệu năng được thực hiện trên NVIDIA A100 40GB với các mô hình TinyLlama-1.1B và Llama-3.2-1B

Xử lý sự cố

Để có hướng dẫn khắc phục sự cố, hãy xem RapidFire AI Troubleshooting Guide.

Tài nguyên bổ sung

Tìm hiểu thêm về RapidFire AI trong kho mã nguồn chính thức và tài liệu của họ.

Update on GitHub