3 篇博文 含有标签「Reasoning」

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

import torch
import torch.nn.functional as F
from typing import List, Optional

class EarlyStopReasoner:
    def __init__(
        self,
        model,
        tokenizer,
        entropy_threshold: float = 0.3,
        min_steps: int = 10,
        max_steps: int = 1000,
        window_size: int = 5,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.entropy_threshold = entropy_threshold
        self.min_steps = min_steps
        self.max_steps = max_steps
        self.window_size = window_size
        
    def _estimate_answer_entropy(
        self, 
        logits: torch.Tensor, 
        generated_ids: List[int]
    ) -> float:
        """估计当前推理状态下最终答案的条件熵"""
        
        # 获取候选答案 token 的概率分布
        probs = F.softmax(logits[0, -1, :], dim=-1)
        
        # 只考虑词汇表顶部的 K 个候选
        top_k = 50
        top_probs, top_indices = torch.topk(probs, top_k)
        
        # 计算部分熵 (partial entropy)
        entropy = -torch.sum(top_probs * torch.log(top_probs + 1e-10))
        
        # 归一化到 [0, 1]
        max_entropy = torch.log(torch.tensor(top_k, dtype=torch.float))
        normalized_entropy = entropy / max_entropy
        
        return normalized_entropy.item()
    
    @torch.no_grad()
    def generate(self, prompt: str) -> tuple[str, int, float]:
        inputs = self.tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"]
        attention_mask = inputs.get("attention_mask")
        
        # 将输入移到模型所在设备
        device = next(self.model.parameters()).device
        input_ids = input_ids.to(device)
        if attention_mask is not None:
            attention_mask = attention_mask.to(device)
        
        past_key_values = None
        generated = input_ids[0].tolist()
        entropy_history = []
        early_stopped = False
        
        for step in range(self.max_steps):
            with torch.no_grad():
                outputs = self.model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    past_key_values=past_key_values,
                    use_cache=True,
                )
                
                logits = outputs.logits
                past_key_values = outputs.past_key_values
                
                # 计算不确定性
                if step >= self.min_steps:
                    entropy = self._estimate_answer_entropy(logits, generated)
                    entropy_history.append(entropy)
                    
                    # 使用滑动窗口平均平滑不确定性估计
                    if len(entropy_history) >= self.window_size:
                        avg_entropy = sum(
                            entropy_history[-self.window_size:]
                        ) / self.window_size
                        
                        if avg_entropy < self.entropy_threshold:
                            early_stopped = True
                            break
                
                # 贪婪解码（简化，实际可以使用采样）
                next_token = logits[0, -1].argmax(dim=-1).unsqueeze(0)
                input_ids = next_token.unsqueeze(0)
                generated.append(next_token.item())
                
                if next_token.item() == self.tokenizer.eos_token_id:
                    break
        
        decoded = self.tokenizer.decode(generated)
        total_steps = len(entropy_history) if early_stopped else step + 1
        
        return decoded, total_steps, entropy_history[-1] if entropy_history else 1.0

import torch
import torch.nn as nn

class RewardPredictor(nn.Module):
    """轻量级奖励预测器，基于当前的 hidden state 预测最终答案正确概率"""
    
    def __init__(self, hidden_size: int = 4096):
        super().__init__()
        self.predictor = nn.Sequential(
            nn.Linear(hidden_size, 512),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(512, 128),
            nn.GELU(),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )
        
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # hidden_states: [batch, seq_len, hidden_size]
        # 使用最后一个位置的 hidden state
        last_hidden = hidden_states[:, -1, :]  # [batch, hidden_size]
        reward = self.predictor(last_hidden)   # [batch, 1]
        return reward.squeeze(-1)              # [batch]

class AdaptiveReasoningController:
    """
    多级自适应推理控制器
    策略 1: 简单问题直接输出（< 200 tokens）
    策略 2: 中等复杂度问题使用 Early Stopping
    策略 3: 复杂问题使用完整推理 + Re-FORC
    """
    
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.es_reasoner = EarlyStopReasoner(model, tokenizer)
        
    def estimate_complexity(self, prompt: str) -> str:
        """快速估算问题复杂度"""
        # 使用一个小型分类器或启发式规则
        prompt_len = len(prompt.split())
        
        if prompt_len < 20 and "?" in prompt:
            return "simple"
        elif "prove" in prompt.lower() or "explain" in prompt.lower():
            return "complex"
        else:
            return "medium"
    
    def generate(self, prompt: str) -> tuple[str, int]:
        complexity = self.estimate_complexity(prompt)
        
        if complexity == "simple":
            # 直接贪婪解码，限制最大 token 数
            output = self._direct_generate(prompt, max_tokens=256)
            return output, output["usage"]["completion_tokens"]
        
        elif complexity == "medium":
            # Early Stopping
            text, steps, final_entropy = self.es_reasoner.generate(prompt)
            return text, steps
        
        else:
            # 完整推理
            return self._full_reasoning(prompt)

# vLLM 配置示例：启用 prefix caching + 动态批处理
from vllm import LLM, SamplingParams

llm = LLM(
    model="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
    tensor_parallel_size=2,
    max_model_len=32768,
    enable_prefix_caching=True,    # 缓存公共前缀
    max_num_batched_tokens=65536,   # 支持长序列批处理
    gpu_memory_utilization=0.90,   # 留出 KV cache 空间
)

sampling_params = SamplingParams(
    temperature=0.0,
    max_tokens=16384,
    early_stopping=True,           # 使用 vLLM 的 early stopping
    stop_token_ids=[151643],       # 自定义停止 token
)

def prepare_reasoning_training_data(trajectories):
    """
    准备训练数据：在推理链中插入 <STOP_REASONING> token
    当模型已经得出正确答案时，插入停止标记
    """
    processed = []
    for prompt, reasoning_chain, answer, is_correct in trajectories:
        if not is_correct:
            continue  # 只使用正确轨迹
        
        # 找到第一次出现正确答案的位置
        first_correct_pos = find_first_correct_position(
            reasoning_chain, answer
        )
        
        # 在该位置后插入 <STOP_REASONING> 和最终答案
        augmented_chain = (
            reasoning_chain[:first_correct_pos]
            + " <STOP_REASONING> "
            + answer
        )
        
        processed.append({
            "prompt": prompt,
            "completion": augmented_chain,
        })
    
    return processed