TRL documentation
Tích hợp RapidFire AI
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 GPU | Chỉ TRL | TRL + RapidFire AI |
|---|---|---|
| Chiến lược huấn luyện | Chạy tuần tự N cấu hình | Chạ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ất | Trự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ện | Không | Có (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 startBả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:
- 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
- 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
- Đổ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
- 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 - Đ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 trackioCác biến môi trường tương đương cũng có sẵn:
RF_MLFLOW_ENABLED(mặc định làtrue, hoặcfalsetrong Colab)RF_TENSORBOARD_ENABLED(mặc định làfalse, hoặctruetrong Colab)RF_TRACKIO_ENABLED(mặc địnhfalse)
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ản | Thời gian tuần tự | Thời gian RapidFire AI | Tốc độ tăng |
|---|---|---|---|
| 4 cấu hình, 1 GPU | 120 phút | 7,5 phút | 16× |
| 8 cấu hình, 1 GPU | 240 phút | 12 phút | 20× |
| 4 cấu hình, 2 GPU | 60 phút | 4 phút | 15× |
| 8 cấu hình, 4 GPU | 60 phút | 3 phút | 20× |
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
- Colab Notebook: RapidFire AI trong Google Colab
- Tài liệu: oss-docs.rapidfire.ai
- GitHub: RapidFireAI/rapidfireai
- PyPI: pypi.org/project/rapidfireai
- Discord: Tham gia Discord của chúng tôi
- Sổ tay hướng dẫn: Kho GitHub
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