6 篇博文 含有标签「Training」

def dpo_loss(
    policy_chosen_logp,   # π_θ(y_w|x)
    policy_rejected_logp, # π_θ(y_l|x)
    ref_chosen_logp,      # π_ref(y_w|x)
    ref_rejected_logp,    # π_ref(y_l|x)
    beta=0.1,             # temperature
):
    """
    DPO loss: maximize the gap between chosen and rejected responses.
    β controls how far the policy can deviate from the reference.
    """
    import torch

    policy_log_ratio = policy_chosen_logp - policy_rejected_logp
    ref_log_ratio = ref_chosen_logp - ref_rejected_logp
    logits = policy_log_ratio - ref_log_ratio

    loss = -torch.nn.functional.logsigmoid(beta * logits)
    return loss.mean()

def grpo_loss(
    policy_model,
    reference_model,
    prompts,           # batch of prompts
    reward_func,       # callable: (prompt, response) -> float
    group_size=4,      # G: responses per prompt
    beta=0.04,         # KL penalty coefficient
    epsilon=0.2,       # clipping parameter
):
    """
    GRPO loss as used in DeepSeek-R1.

    Key difference from PPO:
    - No critic/value model
    - Advantage computed via group normalization
    - Clipping on both sides like PPO
    """
    import torch
    import torch.nn.functional as F

    total_loss = []
    total_kl = []

    for prompt in prompts:
        # Step 1: Sample G responses from current policy
        responses = []
        rewards = []
        with torch.no_grad():
            for _ in range(group_size):
                resp_ids, logps = policy_model.generate(
                    prompt, return_logprobs=True, do_sample=True
                )
                responses.append((resp_ids, logps))

                # Get reward (rule-based or model-based)
                resp_text = policy_model.tokenizer.decode(resp_ids[0])
                r = reward_func(prompt, resp_text)
                rewards.append(r)

            rewards = torch.tensor(rewards)

        # Step 2: Group normalization for advantage
        advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-8)

        # Step 3: Compute loss for each positive-advantage response
        for i, (resp_ids, old_logps) in enumerate(responses):
            if advantages[i] <= 0:
                continue

            # Current policy logprobs
            new_logps = policy_model.get_logprobs(prompt, resp_ids)

            # PPO-style ratio
            ratio = torch.exp(new_logps - old_logps)

            # Clipped objective
            clipped_ratio = torch.clamp(
                ratio, 1 - epsilon, 1 + epsilon
            )
            policy_loss = -torch.min(
                ratio * advantages[i],
                clipped_ratio * advantages[i]
            ).mean()

            # KL penalty to reference model
            ref_logps = reference_model.get_logprobs(prompt, resp_ids)
            kl = (new_logps - ref_logps).mean()

            loss = policy_loss + beta * kl
            total_loss.append(loss)
            total_kl.append(kl.item())

    if total_loss:
        return torch.stack(total_loss).mean(), sum(total_kl) / len(total_kl)
    return torch.tensor(0.0), 0.0

from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. 加载模型
model_name = "Qwen/Qwen2.5-7B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    attn_implementation="flash_attention_2",
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# 2. 准备数据集（只需要 prompts）
dataset = load_dataset("openai/gsm8k", "main", split="train[:2000]")

def format_prompt(example):
    return {
        "prompt": f"Question: {example['question']}\nAnswer: "
    }

dataset = dataset.map(format_prompt)

# 3. 定义 reward function
def math_reward(prompt, completion):
    """
    Rule-based reward for math reasoning.
    简单粗暴：提取 \boxed{} 中的答案，跟标准答案比对。
    """
    import re

    # Extract answer from completion
    match = re.search(r'\\boxed\{([^}]+)\}', completion)
    if not match:
        return 0.0

    predicted = match.group(1).strip()

    # Extract ground truth (GSM8K format: "#### 42")
    # We need the ground truth from somewhere — in practice
    # you'd pass it alongside the prompt
    # Simplified here for illustration
    return 1.0 if predicted  # placeholder

# 实际使用时，用 GRPOTrainer 自带的 reward 机制：
# 可以传 reward_funcs 列表，支持 rule-based + model-based 混合

# 4. 配置 GRPO
grpo_config = GRPOConfig(
    output_dir="./grpo-qwen-math",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,

    # GRPO 特定参数
    num_generation_per_prompt=4,   # G (group size)
    max_prompt_length=512,
    max_completion_length=256,
    temperature=0.9,

    # KL 控制
    beta=0.04,                     # KL penalty coefficient

    # PPO-style clipping
    epsilon=0.2,

    # 学习率
    learning_rate=1e-6,
    lr_scheduler_type="cosine",
    warmup_ratio=0.1,

    # 日志
    logging_steps=10,
    save_steps=500,

    bf16=True,
    gradient_checkpointing=True,
)

# 5. 启动训练
trainer = GRPOTrainer(
    model=model,
    tokenizer=tokenizer,
    args=grpo_config,
    train_dataset=dataset,
    reward_funcs=[math_reward],
)

trainer.train()
trainer.save_model("./grpo-qwen-math-final")

from unsloth import FastLanguageModel, is_bfloat16_supported

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="Qwen/Qwen2.5-7B-Instruct",
    max_seq_length=2048,
    load_in_4bit=True,           # QLoRA
    fast_inference=True,
)

model = FastLanguageModel.get_peft_model(
    model,
    r=64,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_alpha=128,
    use_gradient_checkpointing="unsloth",
)

# 后续跟标准 TRL GRPOTrainer 一样使用

def adaptive_beta(current_kl, target_kl=0.01, beta_min=0.01, beta_max=0.5):
    """
    Adjust β to keep KL divergence near target.
    当 KL 过大时增大 β，过小时减小 β。
    """
    if current_kl > target_kl * 2:
        return min(beta_max, beta * 2)
    elif current_kl < target_kl / 2:
        return max(beta_min, beta / 2)
    return beta