FastTRL 🚀

TRL'nin tamamen yeniden yazılmış, daha hızlı ve güvenilir versiyonu.

Özellikler

Kurulum

pip install -e .
# PEFT + bitsandbytes ile:
pip install -e ".[all]"

Hızlı Başlangıç

1. SFT (Supervised Fine-Tuning)

from fasttrl import SFTTrainer, SFTConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-8B")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3-8B")
dataset = load_dataset("trl-lib/ultrachat_200k", split="train_sft")

config = SFTConfig(
    output_dir="./sft_output",
    model_name="meta-llama/Llama-3-8B",
    max_seq_length=2048,
    num_train_epochs=3,
    per_device_train_batch_size=4,
    learning_rate=2e-5,
    packing=True,           # Sequence packing → daha hızlı eğitim
    bf16=True,              # bfloat16 mixed precision
    neftune_noise_alpha=5,  # NEFTune gürültü artırımı
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    args=config,
    train_dataset=dataset,
)
trainer.train()

2. DPO (Direct Preference Optimization)

from fasttrl import DPOTrainer, DPOConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

model = AutoModelForCausalLM.from_pretrained("./sft_output/final_model")
tokenizer = AutoTokenizer.from_pretrained("./sft_output/final_model")
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train_prefs")

config = DPOConfig(
    output_dir="./dpo_output",
    beta=0.1,               # KL katsayısı (küçük = daha az kısıtlama)
    loss_type="sigmoid",    # Orijinal DPO
    max_length=1024,
    max_prompt_length=512,
    learning_rate=5e-7,
    bf16=True,
    precompute_ref_log_probs=True,  # Hız için önhesaplama
)

trainer = DPOTrainer(
    model=model,
    tokenizer=tokenizer,
    args=config,
    train_dataset=dataset,
)
trainer.train()

3. PPO (RLHF)

from fasttrl import PPOTrainer, PPOConfig
from fasttrl.models import AutoModelForCausalLMWithValueHead
from transformers import AutoTokenizer, pipeline

# Policy model (value head ile)
model = AutoModelForCausalLMWithValueHead.from_pretrained("./sft_output/final_model")
tokenizer = AutoTokenizer.from_pretrained("./sft_output/final_model")

# Ödül modeli
reward_pipeline = pipeline(
    "text-classification",
    model="./reward_output/final_model",
    device=0,
)

def reward_fn(texts):
    results = reward_pipeline(texts)
    return [r["score"] for r in results]

config = PPOConfig(
    output_dir="./ppo_output",
    batch_size=64,
    mini_batch_size=16,
    ppo_epochs=4,
    learning_rate=1e-5,
    init_kl_coef=0.2,
    target_kl=6.0,
    adap_kl_ctrl=True,
)

trainer = PPOTrainer(
    model=model,
    tokenizer=tokenizer,
    args=config,
)

# Eğitim döngüsü
for batch in your_dataloader:
    # 1) Yanıt üret
    responses = trainer.generate(batch["input_ids"])

    # 2) Ödülleri hesapla
    texts = [tokenizer.decode(r) for r in responses]
    scores = [torch.tensor(s) for s in reward_fn(texts)]

    # 3) PPO adımı
    stats = trainer.step(batch["input_ids"], responses, scores)
    print(stats)

4. GRPO (Group Relative Policy Optimization)

from fasttrl import GRPOTrainer, GRPOConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Kural tabanlı ödül fonksiyonu (örn: matematiksel doğruluk)
def math_reward(completions, **kwargs):
    rewards = []
    for c in completions:
        # Basit ödül: sayı içeriyorsa +1
        rewards.append(1.0 if any(ch.isdigit() for ch in c) else 0.0)
    return rewards

config = GRPOConfig(
    output_dir="./grpo_output",
    num_generations=8,      # Her prompt için 8 yanıt
    max_new_tokens=256,
    beta=0.04,
    epsilon=0.2,
    learning_rate=1e-6,
)

trainer = GRPOTrainer(
    model=model,
    tokenizer=tokenizer,
    args=config,
    reward_funcs=[math_reward],
    train_dataset=math_dataset,
)
trainer.train()

5. Reward Model Eğitimi

from fasttrl import RewardTrainer, RewardConfig
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from datasets import load_dataset

model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=1)
tokenizer = AutoTokenizer.from_pretrained("gpt2")
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train_prefs")

config = RewardConfig(
    output_dir="./reward_output",
    max_length=512,
    learning_rate=1e-5,
    margin=0.5,  # Chosen-rejected arasında minimum fark
)

trainer = RewardTrainer(
    model=model,
    tokenizer=tokenizer,
    args=config,
    train_dataset=dataset,
)
trainer.train()

Desteklenen DPO Kayıp Fonksiyonları

Hız Karşılaştırması

Konfigürasyon Kaydetme/Yükleme

# Kaydet
config.save("./my_config.json")

# Yükle
config = SFTConfig.load("./my_config.json")

Lisans

MIT License