noao-chat-4-rl阶段训练

# RL 超参数
rl_config = {
    # 数据
    "source": "sft",              # 从 SFT 模型开始
    "task": GSM8K,                # 数学推理任务
    
    # 采样
    "device_batch_size": 8,       # 推理批次
    "examples_per_step": 16,      # 每步处理的问题数
    "num_samples": 16,            # 每个问题采样 16 个答案
    "max_new_tokens": 256,        # 最长生成 256 tokens
    "temperature": 1.0,           # 采样温度（探索）
    "top_k": 50,                  # Top-K 采样
    
    # 训练
    "unembedding_lr": 0.004,
    "embedding_lr": 0.2,
    "matrix_lr": 0.02,
    "init_lr_frac": 0.05,         # 从 5% 开始（比 SFT 高）
    
    # 评估
    "eval_every": 60,
    "eval_examples": 400,
}

# 标准 PPO 目标函数：
L = E[min(ratio * A, clip(ratio, 1-ε, 1+ε) * A)] - β * KL(π||π_ref)
    ↑                                              ↑
    PPO Clip                                      KL 约束

# 简化后（本项目）：
L = E[log π(a|s) * A]
    ↑
    简单的策略梯度（REINFORCE）

# LLM 的本质
def policy(state):
    """
    输入: state = 当前上下文 (问题 + 已生成的文本)
    输出: action = 下一个 token 的概率分布
    """
    logits = model(state)
    probs = softmax(logits)
    return probs

# 伪代码
for question in dataset:
    # 1. 采样多个答案
    answers = model.sample(question, num_samples=16)
    
    # 2. 计算每个答案的奖励
    rewards = [evaluate(question, answer) for answer in answers]
    
    # 3. 计算梯度
    for answer, reward in zip(answers, rewards):
        # log_prob: 生成这个答案的对数概率
        # reward: 这个答案的得分
        gradient += log_prob(answer) * reward
    
    # 4. 更新参数
    model.update(gradient)

rewards = [0.8, 0.9, 0.7, 0.85]  # 都是正数
# 所有答案的概率都会增加！
# 但我们只想增加"更好"的答案

rewards = [0.8, 0.9, 0.7, 0.85]
mean_reward = 0.8125

advantages = [
    0.8 - 0.8125 = -0.0125,   # 比平均差 → 降低概率
    0.9 - 0.8125 = +0.0875,   # 比平均好 → 增加概率
    0.7 - 0.8125 = -0.1125,   # 比平均差 → 降低概率
    0.85 - 0.8125 = +0.0375,  # 比平均好 → 增加概率
]

# 序列级别
answer = "The answer is 4"
reward = 1.0  # 整个答案的奖励

# Token 级别
tokens = ["The", "answer", "is", "4"]
# 每个 token 应该得到多少奖励？

# 所有 token 获得相同的奖励
for token in tokens:
    token_reward = reward  # 1.0

# 实现中：
advantages = advantages.unsqueeze(-1)  # (B,) → (B, 1)
# 广播到 (B, T)，每个 token 位置都是相同值

@torch.no_grad()
def get_batch():
    for example_idx in range(ddp_rank, len(train_task), ddp_world_size):
        # 获取一个数学题
        conversation = train_task[example_idx]
        # {
        #     "messages": [
        #         {"role": "user", "content": "What is 2+2?"},
        #         {"role": "assistant", "content": "4"}
        #     ]
        # }

        # 去掉 assistant 的回答，只保留问题
        tokens = tokenizer.render_for_completion(conversation)
        
        # tokens:
        # [<|bos|>, <|user_start|>, "What", "is", "2", "+", "2", "?", 
        #  <|user_end|>, <|assistant_start|>]
        #                           ↑
        #                    到这里停止，让模型生成
        
        prefix_length = len(tokens)  # 记录前缀长度

        model.eval()  # 评估模式
        generated_token_sequences = []
        masks = []
        
        # 分批生成，避免 OOM
        num_sampling_steps = num_samples // device_batch_size
        # 16 // 8 = 2 步
        
        for sampling_step in range(num_sampling_steps):
            seed = hash((step, example_idx, sampling_step)) & 0x7FFFFFFF
            
            with autocast_ctx:
                sequences_batch, masks_batch = engine.generate_batch(
                    tokens,
                    num_samples=device_batch_size,  # 8
                    max_tokens=max_new_tokens,      # 256
                    temperature=1.0,                # 随机采样！
                    top_k=50,
                    seed=seed,
                )
            
            generated_token_sequences.extend(sequences_batch)
            masks.extend(masks_batch)

# temperature=0.0 (贪婪解码)
# 每次都选概率最高的 token
# 16 个样本会完全一样 ❌

# temperature=1.0 (采样)
# 按照概率分布随机选择
# 16 个样本各不相同 ✅

        # 对每个生成的答案计算奖励
        rewards = []
        for sample_tokens in generated_token_sequences:
            # 提取生成的部分（去掉 prompt）
            generated_tokens = sample_tokens[prefix_length:]
            
            # 解码为文本
            generated_text = tokenizer.decode(generated_tokens)
            # "The answer is 4"
            
            # 计算奖励
            reward = train_task.reward(conversation, generated_text)
            # reward: 0.0 或 1.0（对于 GSM8K）
            
            rewards.append(reward)

def reward(self, conversation, assistant_response):
    # 从标准答案中提取数字
    ref_num = extract_answer(conversation['messages'][-1])
    # "#### 10" → "10"
    
    # 从生成的回答中提取数字
    pred_num = extract_answer(assistant_response)
    # "The answer is 10" → "10"
    
    # 比较
    is_correct = float(pred_num == ref_num)
    return is_correct  # 0.0 或 1.0

        # 填充序列到相同长度
        max_length = max(len(seq) for seq in generated_token_sequences)
        
        padded_sequences = [
            seq + [pad_token] * (max_length - len(seq))
            for seq in generated_token_sequences
        ]
        
        padded_masks = [
            mask + [0] * (max_length - len(mask))
            for mask in masks
        ]
        
        # 转为张量
        ids = torch.tensor(padded_sequences, device=device)  # (16, T)
        mask_ids = torch.tensor(padded_masks, device=device)  # (16, T)
        
        # 构造 inputs 和 targets
        inputs = ids[:, :-1]    # (16, T-1)
        targets = ids[:, 1:]    # (16, T-1)
        
        # 应用 mask：只训练生成的部分
        targets[mask_ids[:, 1:] == 0] = -1
        
        rewards = torch.tensor(rewards, device=device)  # (16,)

        # 计算优势函数
        mu = rewards.mean()  # 平均奖励
        advantages = rewards - mu  # 相对优势
        
        # 例如：
        # rewards = [0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0]
        # mu = 0.625
        # advantages = [-0.625, 0.375, -0.625, 0.375, ...]

# 输入
question_tokens: [10, 20, 30, 40, 50]  # 问题 (5 tokens)

# 采样 16 次，每次生成不同的答案
sample_1: [10, 20, 30, 40, 50, 60, 70, 80]        # +3 tokens
sample_2: [10, 20, 30, 40, 50, 61, 72, 83, 94]    # +4 tokens
...
sample_16: [10, 20, 30, 40, 50, 65, 75]           # +2 tokens

# 填充到最长
max_length = 9  # 最长的序列
padded_samples: (16, 9)

# inputs & targets
inputs:  (16, 8)
targets: (16, 8)

# rewards & advantages
rewards: (16,)      # [0, 1, 0, 1, ...]
advantages: (16,)   # [-0.625, 0.375, ...]

def reward(conversation, assistant_response):
    # 提取答案
    ref = extract_answer(conversation)  # "#### 42"
    pred = extract_answer(assistant_response)  # "#### 42"
    
    # 比较
    return 1.0 if ref == pred else 0.0

def advanced_reward(conversation, assistant_response):
    reward = 0.0
    
    # 1. 最终答案正确 +1.0
    if is_correct_answer(conversation, assistant_response):
        reward += 1.0
    
    # 2. 使用工具 +0.2
    if has_calculator_call(assistant_response):
        reward += 0.2
    
    # 3. 步骤清晰 +0.1
    if has_clear_steps(assistant_response):
        reward += 0.1
    
    # 4. 格式规范 +0.1
    if has_proper_format(assistant_response):
        reward += 0.1
    
    return reward

# 标准方法
advantages = (rewards - rewards.mean()) / rewards.std()
# 问题：rewards 是 0/1，std 可能很小或为 0

# 简化方法
advantages = rewards - rewards.mean()
# 优点：简单，稳定，效果相当

# 假设 16 个样本的奖励
rewards = torch.tensor([
    0, 1, 0, 1, 1, 0, 1, 1,  # 前 8 个
    0, 1, 0, 1, 1, 0, 1, 0   # 后 8 个
])  # 10 个正确，6 个错误

mean = 0.625
std = 0.484

# 标准方法
advantages_standard = (rewards - 0.625) / 0.484
# [-1.29, 0.77, -1.29, 0.77, 0.77, -1.29, 0.77, 0.77,
#  -1.29, 0.77, -1.29, 0.77, 0.77, -1.29, 0.77, -1.29]

# 简化方法
advantages_simple = rewards - 0.625
# [-0.625, 0.375, -0.625, 0.375, 0.375, -0.625, ...]

# 效果类似：正确的>0，错误的<0

reward = 1.0  # 标量

logits: (B, T, vocab_size)
targets: (B, T)
# 需要每个 token 位置都有奖励信号

# advantages: (B,) = (16,)
advantages = advantages.unsqueeze(-1)  # (16, 1)

# logp: (B, T) = (16, 256)
logp = -model(inputs, targets, loss_reduction='none')

# 计算策略梯度目标
pg_obj = (logp * advantages).sum()
#         ↑       ↑
#     (16, 256) * (16, 1)  → 广播为 (16, 256)
#     每个 token 都乘以相同的 advantage

# 假设一个样本，advantage = 0.375

logp = [-0.5, -0.3, -0.8, -0.2, -0.4]  # 5 个 token
advantage = 0.375

# 广播后
pg_obj_per_token = [
    -0.5 * 0.375 = -0.1875,
    -0.3 * 0.375 = -0.1125,
    -0.8 * 0.375 = -0.3000,
    -0.2 * 0.375 = -0.0750,
    -0.4 * 0.375 = -0.1500,
]

# 总和
pg_obj = -0.825

# 准备
batch_iterator = get_batch()  # 生成器

for step in range(num_steps):
    
    # ===== 评估阶段 =====
    if step % eval_every == 0:
        model.eval()
        # 评估 Pass@k
        passk = evaluate_pass_at_k(model, val_task)
        print(f"Pass@1: {passk[0]:.4f}, Pass@8: {passk[7]:.4f}")
        model.train()
    
    # ===== 训练阶段 =====
    rewards_list = []
    sequence_lengths = []
    
    # 处理多个问题
    for example_step in range(examples_per_rank):
        # 1. 获取一个问题的 16 个采样
        sequences, inputs, targets, rewards, advantages = next(batch_iterator)
        
        model.train()
        
        # 2. 分批处理（避免 OOM）
        num_passes = inputs.size(0) // device_batch_size
        # 16 // 8 = 2 passes
        
        for pass_idx in range(num_passes):
            # 取出一个 mini-batch
            b0, b1 = pass_idx * device_batch_size, (pass_idx + 1) * device_batch_size
            inputs_batch = inputs[b0:b1]     # (8, T)
            targets_batch = targets[b0:b1]   # (8, T)
            advantages_batch = advantages[b0:b1]  # (8,)
            
            # 3. 计算 log 概率
            with autocast_ctx:
                logp = -model(inputs_batch, targets_batch, loss_reduction='none')
                # logp: (8, T)，每个 token 的 log P(token | context)
            
            # 4. 计算策略梯度目标
            pg_obj = (logp * advantages_batch.unsqueeze(-1)).sum()
            
            # 5. 归一化
            num_valid = (targets_batch >= 0).sum().clamp(min=1)
            pg_obj = pg_obj / (num_valid * num_passes * examples_per_rank)
            
            # 6. 转为损失（最小化）
            loss = -pg_obj
            loss.backward()
        
        # 记录统计
        rewards_list.append(rewards.mean().item())
        sequence_lengths.extend(len(seq) for seq in sequences)
    
    # ===== 更新参数 =====
    lrm = get_lr_multiplier(step)
    for opt in optimizers:
        for group in opt.param_groups:
            group["lr"] = group["initial_lr"] * lrm
    
    for opt in optimizers:
        opt.step()
    model.zero_grad()
    
    # ===== 日志 =====
    mean_reward = sum(rewards_list) / len(rewards_list)
    print(f"Step {step} | Reward: {mean_reward:.4f}")

# inputs:  (8, 256) - token IDs
# targets: (8, 256) - 目标 IDs（包含 -1）
# advantages: (8,) - 每个序列的优势

# 步骤 1: 前向传播
with autocast_ctx:
    # 模型计算每个 token 的 NLL
    nll = model(inputs, targets, loss_reduction='none')
    # nll: (8, 256)
    
    # 转为 log 概率
    logp = -nll  # (8, 256)

# 步骤 2: 应用优势函数
advantages = advantages.unsqueeze(-1)  # (8, 1)
weighted_logp = logp * advantages      # (8, 256)

# 步骤 3: 求和
pg_obj = weighted_logp.sum()  # 标量

# 步骤 4: 归一化
num_valid = (targets >= 0).sum()  # 有效 token 数量
pg_obj = pg_obj / num_valid

# 步骤 5: 转为损失
loss = -pg_obj  # 要最小化

# 正 advantage (reward 高)
advantage = +0.375
pg_obj = log P(tokens) * 0.375  # 正值
loss = -pg_obj  # 负值
# 最小化负值 = 最大化 log P(tokens)
# → 增加这些 token 的概率 ✅

# 负 advantage (reward 低)
advantage = -0.625
pg_obj = log P(tokens) * (-0.625)  # 负值
loss = -pg_obj  # 正值
# 最小化正值 = 减小 log P(tokens)
# → 降低这些 token 的概率 ✅

def get_lr_multiplier(it):
    # 简单的线性衰减
    lrm = 1.0 - it / num_steps
    return lrm

# SFT 初始学习率
sft_initial_lr = matrix_lr * 0.02  # 0.02 * 0.02 = 0.0004

# RL 初始学习率
rl_initial_lr = matrix_lr * 0.05  # 0.02 * 0.05 = 0.001

# RL 的学习率是 SFT 的 2.5 倍！

def evaluate_pass_at_k(model, task, k=8):
    """
    对每个问题生成 k 个答案
    计算至少有一个正确的比例
    """
    correct_count = 0
    total_count = 0
    
    for question in task:
        # 生成 k 个答案
        answers = model.generate(question, num_samples=k)
        
        # 检查是否至少有一个正确
        is_any_correct = any(
            task.evaluate(question, answer)
            for answer in answers
        )
        
        if is_any_correct:
            correct_count += 1
        total_count += 1
    
    pass_at_k = correct_count / total_count
    return pass_at_k

# Pass@1: 首次尝试的准确率
# 例如：0.45 = 45% 的题第一次就做对

# Pass@8: 给 8 次机会的准确率
# 例如：0.72 = 72% 的题在 8 次尝试中至少对一次

# Pass@k 越高，说明模型：
# - 探索能力越强
# - 找到正确答案的概率越高
# - 即使不确定，也有机会试出来

# 输入：完整对话
conversation = {
    "messages": [
        {"role": "user", "content": "What is 2+2?"},
        {"role": "assistant", "content": "4"}  # ← 有标准答案
    ]
}

# Tokenize
ids, mask = tokenizer.render_conversation(conversation)

# 训练
inputs = ids[:-1]
targets = ids[1:]  # 已知的正确答案
loss = model(inputs, targets)  # 直接计算损失
loss.backward()

# 输入：只有问题
conversation = {
    "messages": [
        {"role": "user", "content": "What is 2+2?"}
        # 没有 assistant 答案！
    ]
}

# 生成 16 个答案
tokens = tokenizer.render_for_completion(conversation)
answers = model.generate(tokens, num_samples=16, temperature=1.0)
# ["4", "four", "2+2", "The answer is 4", ...]

# 评估每个答案
rewards = [evaluate(conv, ans) for ans in answers]
# [1.0, 0.0, 0.0, 1.0, ...]

# 计算优势
advantages = rewards - rewards.mean()

# 训练（策略梯度）
for answer, advantage in zip(answers, advantages):
    logp = log_prob(answer)
    loss = -logp * advantage
    loss.backward()

# 标准交叉熵
loss = F.cross_entropy(
    logits.view(-1, vocab_size),  # (B*T, V)
    targets.view(-1),              # (B*T,)
    ignore_index=-1
)

# 每个 token 都有明确的目标
# 只要预测对了，loss 就低

# 策略梯度
logp = -model(inputs, targets, loss_reduction='none')  # (B, T)
pg_obj = (logp * advantages.unsqueeze(-1)).sum()
loss = -pg_obj / num_valid_tokens

# token 本身没有"对错"
# 只有整个序列的"好坏"（advantage）
# 好的序列 → 增加概率
# 坏的序列 → 降低概率

rewards = [1.0] * 16  # 全对
mean = 1.0
advantages = [0.0] * 16  # 都一样好

# 这种情况下，不会更新太多
# 因为没有"相对好坏"的区分

# RL 核心超参数
rl_config = {
    # 采样
    "num_samples": 16,         # 每个问题采样数
    "temperature": 1.0,        # 采样温度
    "top_k": 50,               # Top-K 采样
    
    # 训练
    "examples_per_step": 16,   # 每步处理的问题数
    "device_batch_size": 8,    # 推理批次大小
    "init_lr_frac": 0.05,      # 初始学习率系数
    
    # 评估
    "eval_every": 60,          # 评估频率
    "eval_examples": 400,      # 评估样本数
}