TRL 对齐训练库（结合源码）

# 拼接 prompt + chosen 和 prompt + rejected
prompt_chosen_ids = [example["prompt_ids"] + example["chosen_ids"] for example in examples]
prompt_rejected_ids = [example["prompt_ids"] + example["rejected_ids"] for example in examples]
# completion_mask: 0 = prompt tokens, 1 = completion tokens
chosen_mask = [[0] * len(example["prompt_ids"]) + [1] * len(example["chosen_ids"]) for example in examples]
rejected_mask = [[0] * len(example["prompt_ids"]) + [1] * len(example["rejected_ids"]) for example in examples]
# 合并为一个批次
input_ids = prompt_chosen_ids + prompt_rejected_ids  # 前半 chosen，后半 rejected

from trl import SFTTrainer, SFTConfig

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",  # 模型名称或路径或 PreTrainedModel 实例
    args=SFTConfig(...),                   # SFT 训练配置
    train_dataset=dataset,                 # 训练数据集
    eval_dataset=eval_dataset,             # 评估数据集（可选）
    processing_class=tokenizer,            # tokenizer 或 processor
    data_collator=None,                    # 数据整理器（默认自动选择）
    formatting_func=None,                  # 数据格式化函数（可选）
    callbacks=None,                        # 训练回调
    peft_config=None,                      # PEFT 配置（如 LoRA）
)

from trl import SFTConfig

config = SFTConfig(
    output_dir="./results",
    # === 训练参数 ===
    learning_rate=2e-5,              # 学习率（默认 2e-5，比 transformers 默认的 5e-5 更小）
    per_device_train_batch_size=4,   # 每设备训练批次大小
    gradient_accumulation_steps=8,   # 梯度累积步数
    num_train_epochs=3,              # 训练轮数

    # === 序列长度与打包 ===
    max_length=1024,                 # 最大序列长度（默认 1024）
    packing=False,                   # 是否启用序列打包（默认 False）
    packing_strategy="bfd",          # 打包策略："bfd"、"bfd_split"、"wrapped"
    padding_free=False,              # 无 padding 模式（需 FlashAttention）

    # === 损失控制 ===
    completion_only_loss=None,       # 仅在 completion 部分计算损失
    assistant_only_loss=False,       # 仅在 assistant 部分计算损失
    loss_type="nll",                 # 损失类型："nll"、"chunked_nll"、"dft"

    # === 显存优化 ===
    gradient_checkpointing=True,     # 梯度检查点（默认 True）
    activation_offloading=False,     # 激活卸载到 CPU
    bf16=True,                       # BF16 混合精度（默认 True）

    # === 数据处理 ===
    dataset_text_field="text",       # 文本列名
    dataset_num_proc=None,           # 数据处理进程数
    truncation_mode="keep_start",    # 截断模式
)

from trl import DPOTrainer, DPOConfig

trainer = DPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",      # 待训练模型
    ref_model=None,                            # 参考模型（None 时使用初始策略）
    args=DPOConfig(...),                       # DPO 训练配置
    train_dataset=dataset,                     # 偏好数据集
    processing_class=tokenizer,                # tokenizer
    data_collator=None,                        # 数据整理器（默认 DataCollatorForPreference）
    peft_config=None,                          # PEFT 配置
    loss_type=["sigmoid"],                     # 损失类型列表
)

from trl import DPOConfig

config = DPOConfig(
    output_dir="./results",
    # === 核心训练参数 ===
    learning_rate=1e-6,               # 学习率（默认 1e-6，比 SFT 更小）
    beta=0.1,                         # KL 约束系数（β 越大，偏离参考模型越少）

    # === 损失类型 ===
    loss_type=["sigmoid"],            # 损失类型列表，支持 15+ 种
    # "sigmoid": 标准 DPO 损失
    # "hinge": 铰链损失
    # "ipo": IPO 损失（β 为正则化参数 τ）
    # "robust": 鲁棒 DPO（配合 label_smoothing）
    # "discopop": DiscoPOP 损失
    # "sft": 纯 SFT 损失
    # "sigmoid_norm": 标准化 sigmoid DPO
    loss_weights=None,                # 多损失组合的权重

    # === 参考模型控制 ===
    precompute_ref_log_probs=False,   # 预计算参考模型 log probs（节省训练时显存）
    sync_ref_model=False,             # 同步参考模型（TR-DPO）
    ref_model_mixup_alpha=0.6,        # TR-DPO 混合系数 α
    ref_model_sync_steps=512,         # TR-DPO 同步频率

    # === f-散度控制 ===
    f_divergence_type="reverse_kl",   # f-散度类型："reverse_kl"、"forward_kl"、"js_divergence"、"alpha_divergence"

    # === 序列长度 ===
    max_length=1024,                  # 最大序列长度
    padding_free=False,               # 无 padding 模式

    # === 其他 ===
    disable_dropout=True,             # 禁用 dropout（默认 True）
    label_smoothing=0.0,              # 标签平滑（Robust DPO/EXO 使用）
)

from trl import GRPOTrainer, GRPOConfig

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",       # 待训练模型
    reward_funcs=accuracy_reward,               # 奖励函数（支持多种类型）
    args=GRPOConfig(...),                       # GRPO 训练配置
    train_dataset=dataset,                      # 训练数据集（需含 "prompt" 列）
    processing_class=tokenizer,                 # tokenizer
    reward_processing_classes=None,             # 奖励模型对应的 tokenizer
    peft_config=None,                           # PEFT 配置
    tools=None,                                 # 工具函数列表（Agent 训练）
    rollout_func=None,                          # 自定义生成函数
    environment_factory=None,                   # 环境工厂（Agent 训练）
)

from trl import GRPOConfig

config = GRPOConfig(
    output_dir="./results",
    # === 核心训练参数 ===
    learning_rate=1e-6,                # 学习率（默认 1e-6）
    beta=0.0,                          # KL 系数（0.0 时不加载参考模型；DeepSeek-R1 用 0.001）
    num_generations=8,                 # 每个 prompt 的生成数 G（默认 8，最小 2）
    epsilon=0.2,                       # PPO 裁剪参数 ε（默认 0.2）
    num_iterations=1,                  # 每批次的迭代次数 μ

    # === 生成参数 ===
    max_completion_length=256,         # 生成最大长度（默认 256）
    temperature=1.0,                   # 采样温度
    top_p=1.0,                         # Top-p 采样
    generation_batch_size=None,        # 生成批次大小

    # === 损失类型 ===
    loss_type="dapo",                  # 损失类型（默认 dapo）
    # "grpo": 序列长度归一化（有长度偏差）
    # "dapo": 全局活跃 token 归一化（消除长度偏差，默认）
    # "dr_grpo": 全局常数归一化
    # "bnpo": 本地批次归一化
    # "cispo": 裁剪重要性采样权重

    # === 奖励处理 ===
    scale_rewards="group",             # 奖励缩放："group"、"batch"、"none"
    reward_weights=None,               # 各奖励函数权重
    mask_truncated_completions=False,  # 屏蔽被截断的 completion

    # === vLLM 加速 ===
    use_vllm=False,                    # 是否使用 vLLM 生成
    vllm_mode="colocate",              # vLLM 模式："colocate"、"server"

    # === 参考模型同步 ===
    sync_ref_model=False,              # TR-DPO 式参考模型同步
    ref_model_mixup_alpha=0.6,
    ref_model_sync_steps=512,
)

from trl import KTOTrainer, KTOConfig

# 注意：KTOTrainer 当前位于 trl.experimental
trainer = KTOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    ref_model=None,                    # 参考模型
    args=KTOConfig(...),
    train_dataset=dataset,             # 需含 prompt + completion + label（desirable/undesirable）
    processing_class=tokenizer,
)

from trl import RewardTrainer, RewardConfig
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "Qwen/Qwen2.5-0.5B-Instruct", num_labels=1
)

trainer = RewardTrainer(
    model=model,
    args=RewardConfig(
        output_dir="./results",
        learning_rate=1e-4,                  # 默认 1e-4
        max_length=1024,
        center_rewards_coefficient=0.01,     # 奖励中心化系数（推荐 0.01）
        disable_dropout=True,
    ),
    train_dataset=dataset,              # 偏好数据集（chosen + rejected）
    processing_class=tokenizer,
)

# 计算 log ratios
chosen_logratios = chosen_logps - ref_chosen_logps      # log(π_θ/π_ref) for chosen
rejected_logratios = rejected_logps - ref_rejected_logps  # log(π_θ/π_ref) for rejected
delta_score = chosen_scores - rejected_scores
# sigmoid DPO 损失
per_sequence_loss = -F.logsigmoid(self.beta * delta_score)

# 重要性采样比率
log_ratio = per_token_logps - old_per_token_logps
coef_1 = torch.exp(log_ratio)  # ρ_t

# 裁剪目标
coef_2 = torch.clamp(coef_1, 1 - epsilon, 1 + epsilon)
per_token_loss1 = coef_1 * advantages
per_token_loss2 = coef_2 * advantages
per_token_loss = -torch.min(per_token_loss1, per_token_loss2)

# KL 散度
per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1

from trl import SFTTrainer, SFTConfig
from datasets import load_dataset
from transformers import AutoTokenizer

# 加载模型和 tokenizer
model_name = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 加载对话数据集（messages 格式）
dataset = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft[:5000]")

# 配置训练参数
sft_config = SFTConfig(
    output_dir="./sft-output",
    learning_rate=2e-5,                      # SFT 学习率
    per_device_train_batch_size=4,            # 每设备批次大小
    gradient_accumulation_steps=8,            # 梯度累积，等效批次大小=4*8=32
    num_train_epochs=1,                       # 训练轮数
    max_length=2048,                          # 最大序列长度
    packing=True,                             # 启用序列打包
    assistant_only_loss=True,                 # 仅在 assistant 回复上计算损失
    gradient_checkpointing=True,              # 梯度检查点
    bf16=True,                                # BF16 混合精度
    logging_steps=10,
    save_steps=500,
)

# 创建 Trainer 并训练
trainer = SFTTrainer(
    model=model_name,
    args=sft_config,
    train_dataset=dataset,
    processing_class=tokenizer,
)

trainer.train()
trainer.save_model("./sft-output/final")

from trl import DPOTrainer, DPOConfig
from datasets import load_dataset

# 加载偏好数据集
# 数据集需包含 "prompt"、"chosen"、"rejected" 列
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train[:10000]")

# 配置 DPO 训练参数
dpo_config = DPOConfig(
    output_dir="./dpo-output",
    learning_rate=1e-6,                       # DPO 学习率（比 SFT 更小）
    per_device_train_batch_size=2,             # 每设备批次大小
    gradient_accumulation_steps=16,            # 梯度累积
    num_train_epochs=1,
    beta=0.1,                                 # KL 约束系数
    loss_type=["sigmoid"],                    # 标准 DPO 损失
    max_length=1024,                          # 最大序列长度
    gradient_checkpointing=True,
    bf16=True,
)

# 创建 Trainer 并训练
# ref_model=None 时自动使用初始策略作为参考
trainer = DPOTrainer(
    model="./sft-output/final",               # 使用 SFT 后的模型
    ref_model=None,                           # 参考模型 = 初始策略
    args=dpo_config,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model("./dpo-output/final")

from trl import GRPOTrainer, GRPOConfig
from trl.rewards import accuracy_reward
from datasets import load_dataset

# 加载数学推理数据集
dataset = load_dataset("trl-lib/DeepMath-103K", split="train[:5000]")

# 自定义奖励函数
def format_reward(completions, **kwargs):
    """检查回答是否包含正确的 LaTeX 格式"""
    rewards = []
    for completion in completions:
        if "\\boxed{" in completion:
            rewards.append(1.0)
        else:
            rewards.append(-0.5)
    return rewards

# 配置 GRPO 训练参数
grpo_config = GRPOConfig(
    output_dir="./grpo-output",
    learning_rate=1e-6,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    num_train_epochs=1,
    beta=0.001,                               # KL 系数（DeepSeek-R1 推荐 0.001）
    num_generations=8,                        # 每个 prompt 生成 8 个响应
    max_completion_length=512,                # 最大生成长度
    epsilon=0.2,                              # PPO 裁剪参数
    loss_type="dapo",                         # DAPO 损失（消除长度偏差）
    temperature=1.0,                          # 采样温度
    scale_rewards="group",                    # 组内奖励标准化
    mask_truncated_completions=True,          # 屏蔽被截断的 completion
    gradient_checkpointing=True,
    bf16=True,
    log_completions=True,                     # 日志记录生成样本
)

# 创建 Trainer 并训练
trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    reward_funcs=[accuracy_reward, format_reward],  # 多奖励函数
    reward_weights=[0.8, 0.2],                     # 奖励权重
    args=grpo_config,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model("./grpo-output/final")

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

# 加载偏好数据集
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train[:10000]")

# 创建序列分类模型（num_labels=1 输出标量奖励）
model = AutoModelForSequenceClassification.from_pretrained(
    "Qwen/Qwen2.5-0.5B-Instruct",
    num_labels=1,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
tokenizer.pad_token = tokenizer.eos_token

# 配置奖励模型训练
reward_config = RewardConfig(
    output_dir="./reward-output",
    learning_rate=1e-4,                       # 奖励模型学习率
    per_device_train_batch_size=4,
    num_train_epochs=1,
    max_length=1024,
    center_rewards_coefficient=0.01,          # 奖励中心化系数
    disable_dropout=True,
    gradient_checkpointing=True,
    bf16=True,
)

# 创建 Trainer 并训练
trainer = RewardTrainer(
    model=model,
    args=reward_config,
    train_dataset=dataset,
    processing_class=tokenizer,
)

trainer.train()

from trl import DPOTrainer, DPOConfig
from datasets import load_dataset

# 使用 prompt 数据集（无需预收集偏好对）
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train[:5000]")

# 在线 DPO 配置
dpo_config = DPOConfig(
    output_dir="./online-dpo-output",
    learning_rate=5e-7,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=16,
    beta=0.1,
    loss_type=["sigmoid"],
    max_length=1024,
    gradient_checkpointing=True,
    bf16=True,
)

# 在线 DPO：使用同一模型同时作为策略模型和生成模型
# 每个训练步骤中：
# 1. 对 prompt 生成两个响应
# 2. 使用奖励模型对两个响应打分
# 3. 将高分响应作为 chosen，低分响应作为 rejected
# 4. 计算 DPO 损失并更新策略
trainer = DPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    args=dpo_config,
    train_dataset=dataset,
)

trainer.train()

from trl import SFTTrainer, SFTConfig
from peft import LoraConfig, TaskType

# 配置 LoRA
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,                                     # LoRA 秩
    lora_alpha=32,                            # LoRA alpha
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # 目标层
    bias="none",
)

# 在 SFTTrainer 中使用 LoRA
trainer = SFTTrainer(
    model="Qwen/Qwen2.5-7B-Instruct",        # 大模型也可用 LoRA 训练
    args=SFTConfig(
        output_dir="./lora-sft-output",
        learning_rate=2e-5,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=8,
        max_length=2048,
        bf16=True,
    ),
    train_dataset=dataset,
    peft_config=lora_config,                  # 传入 LoRA 配置
)

trainer.train()

# 保存 LoRA adapter（仅保存增量参数）
trainer.save_model("./lora-sft-output/adapter")