Модуль 9. RLVR и GRPO: проверяемые награды, групповые оценки и устойчивость

from __future__ import annotations

INFO_X = 'вероятность успеха одного ответа p'
INFO_Y = 'вероятность информативной группы'
FINITE_REWARD = 'ожидаемая награда'
FINITE_ENTROPY = 'энтропия'
FINITE_X = 'шаг точного обновления'
BEFORE_TRAIN = 'до RL, train'
AFTER_TRAIN = 'после RL, train'
BEFORE_HELDOUT = 'до RL, held-out'
AFTER_HELDOUT = 'после RL, held-out'
HIST_ACCURACY = 'точность свежей генерации'
HIST_ENTROPY = 'энтропия'
HIST_INFO = 'доля информативных групп'
HIST_X = 'обновление'
RATIO_MIN = 'минимальное отношение'
RATIO_MEAN = 'среднее отношение'
RATIO_MAX = 'максимальное отношение'
COVERAGE_BEFORE = 'до RL'
COVERAGE_AFTER = 'после RL'
COVERAGE_Y = 'доля правильных ответов'
REAL_MODEL_SKIPPED = 'Реальная ветвь выключена. Установите RUN_REAL_MODEL=True только в подготовленной GPU-среде.'


from pathlib import Path
import sys
import json
from dataclasses import asdict

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
from IPython.display import display

_candidates = [Path.cwd(), Path.cwd().parent, Path.cwd().parent.parent]
ROOT = next((p for p in _candidates if (p / "modules" / "module-09").exists()), None)
if ROOT is None:
    raise FileNotFoundError("Could not locate the RL for LLM repository root")
MODULE_DIR = ROOT / "modules" / "module-09"
sys.path.insert(0, str(MODULE_DIR))

from module_9_reference import (
    categorical_entropy,
    clip_activity,
    clipped_surrogate_terms,
    dapo_token_weights,
    expected_draws_for_informative_groups,
    flatten_token_weights,
    group_relative_advantages,
    informative_group_probability,
    original_grpo_token_weights,
    dr_grpo_token_weights,
    pass_at_k_independent,
    response_contributions,
    synthetic_mode_training,
    verify_numeric_answer,
)
from module_9_training import (
    GRPOConfig,
    arithmetic_prompts,
    collect_rollouts,
    run_toy_grpo_experiment,
)

PROFILE = "smoke"  # "full" uses a larger CPU budget
RUN_REAL_MODEL = False
SMOKE_MODE = True
SEED = 9
np.random.seed(SEED)
torch.manual_seed(SEED)
print({"root": str(ROOT), "profile": PROFILE, "real_model": RUN_REAL_MODEL})

cases = [
    (r"Final: \\boxed{3/7}", "0.428571", 1),
    ("The reference number 42 is mentioned, but the answer is 41", "42", 0),
    (r"First \\boxed{2}, then corrected: \\boxed{3}", "3", 1),
    ("No numerical answer", "5", 0),
    ("1/0", "0", 0),
]
rows = []
for text, reference, expected in cases:
    observed = verify_numeric_answer(text, reference)
    rows.append({"text": text, "reference": reference, "observed": observed, "expected": expected})
verifier_table = pd.DataFrame(rows)
display(verifier_table)
assert verifier_table["observed"].tolist() == verifier_table["expected"].tolist()

p_grid = np.linspace(0.001, 0.999, 400)
fig, ax = plt.subplots(figsize=(8.5, 4.8))
for group_size in (2, 4, 8, 16):
    values = [informative_group_probability(float(p), group_size) for p in p_grid]
    ax.plot(p_grid, values, label=f"G={group_size}")
ax.set_xlabel(INFO_X)
ax.set_ylabel(INFO_Y)
ax.set_ylim(0, 1.02)
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()

examples = pd.DataFrame({
    "p": [0.01, 0.10, 0.50, 0.90, 0.99],
})
examples["P_info, G=8"] = examples["p"].map(lambda p: informative_group_probability(p, 8))
examples["draws for 100 informative prompts"] = examples["p"].map(
    lambda p: expected_draws_for_informative_groups(100, p, 8)
)
display(examples.round(4))

rewards = np.array([1.0, 1.0, 0.0, 0.0])
advantages = group_relative_advantages(rewards, scale=True)
raw_advantages = group_relative_advantages(rewards, scale=False)
ratios = np.array([1.35, 0.85, 1.10, 0.70])
terms = clipped_surrogate_terms(ratios, advantages, eps_low=0.2, eps_high=0.2)
activity = clip_activity(ratios, advantages, eps_low=0.2, eps_high=0.2)
display(pd.DataFrame({
    "reward": rewards,
    "centered": raw_advantages,
    "standardized": advantages,
    "ratio": ratios,
    "surrogate": terms,
    "clip_active": activity,
}))
assert abs(float(advantages.sum())) < 1e-10

reward_vector = np.array([1.0, 0.0, 1.0])
lengths = [2, 5, 8]
weights = {
    "original_grpo": original_grpo_token_weights(reward_vector, lengths),
    "dr_grpo": dr_grpo_token_weights(reward_vector, lengths, fixed_length=8),
    "dapo": dapo_token_weights(reward_vector, lengths),
}
summary = []
for name, token_weights in weights.items():
    contributions = response_contributions(token_weights)
    summary.append({
        "reduction": name,
        "response_1": contributions[0],
        "response_2": contributions[1],
        "response_3": contributions[2],
        "absolute_token_weight_sum": np.abs(flatten_token_weights(token_weights)).sum(),
    })
display(pd.DataFrame(summary).round(4))

finite_rows = []
for step in range(61):
    state = synthetic_mode_training(beta_entropy=0.0, learning_rate=0.18, steps=step)
    finite_rows.append({
        "step": step,
        "expected_reward": state["expected_reward"],
        "entropy": categorical_entropy(state["probabilities"]),
        "pass_at_4": pass_at_k_independent(state["successful_mass"], 4),
    })
finite_table = pd.DataFrame(finite_rows)
display(finite_table.iloc[[0, 9, 19, 39, 60]].round(4))
fig, ax = plt.subplots(figsize=(8.5, 4.8))
ax.plot(finite_table["step"], finite_table["expected_reward"], label=FINITE_REWARD)
ax.plot(finite_table["step"], finite_table["entropy"], label=FINITE_ENTROPY)
ax.plot(finite_table["step"], finite_table["pass_at_4"], label="pass@4")
ax.set_xlabel(FINITE_X)
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()

experiment = run_toy_grpo_experiment(profile=PROFILE, seed=SEED, loss_type="dapo")
metrics = pd.DataFrame({
    BEFORE_TRAIN: experiment["before_train"],
    AFTER_TRAIN: experiment["after_train"],
    BEFORE_HELDOUT: experiment["before_heldout"],
    AFTER_HELDOUT: experiment["after_heldout"],
}).T
selected_columns = ["accuracy", "mean_reward", "mean_length", "entropy", "teacher_forced_kl", "pass_at_4"]
display(metrics[selected_columns].round(4))
print(asdict(experiment["config"]))

history = pd.DataFrame(experiment["history"])
display(history.round(4))
fig, ax = plt.subplots(figsize=(9.2, 5.0))
ax.plot(history["update"], history["eval_accuracy"], marker="o", label=HIST_ACCURACY)
ax.plot(history["update"], history["eval_entropy"], marker="o", label=HIST_ENTROPY)
ax.plot(history["update"], history["informative_fraction"], marker="o", label=HIST_INFO)
ax.set_xlabel(HIST_X)
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()

fig, ax = plt.subplots(figsize=(9.2, 4.8))
ax.plot(history["update"], history["ratio_min"], label=RATIO_MIN)
ax.plot(history["update"], history["ratio_mean"], label=RATIO_MEAN)
ax.plot(history["update"], history["ratio_max"], label=RATIO_MAX)
ax.axhline(1.0, linewidth=1.0, linestyle="--")
ax.set_xlabel(HIST_X)
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()

coverage = pd.DataFrame({
    "before": experiment["before_train_by_prompt"],
    "after": experiment["after_train_by_prompt"],
}).sort_index()
display(coverage)
fig, ax = plt.subplots(figsize=(10.5, 4.8))
x = np.arange(len(coverage))
width = 0.38
ax.bar(x - width / 2, coverage["before"], width, label=COVERAGE_BEFORE)
ax.bar(x + width / 2, coverage["after"], width, label=COVERAGE_AFTER)
ax.set_xticks(x, coverage.index, rotation=45, ha="right")
ax.set_ylabel(COVERAGE_Y)
ax.set_ylim(0, 1.05)
ax.legend()
fig.tight_layout()
plt.show()

inspect_config = GRPOConfig(group_size=4, prompt_batch=2, updates=1, update_epochs=1, loss_type="dapo")
generator = torch.Generator().manual_seed(SEED + 777)
groups, discarded = collect_rollouts(
    experiment["policy"], experiment["vocab"], experiment["prompts"], inspect_config,
    generator=generator, generator_version="notebook-inspection",
)
first_group = groups[0]
display(pd.DataFrame([{
    "prompt_id": row.prompt_id,
    "response": " ".join(row.response_tokens),
    "reward": row.reward,
    "accuracy": row.accuracy,
    "status": row.verifier_status,
    "length": len(row.response_tokens),
    "terminated": row.terminated,
} for row in first_group]))
print({"discarded_fraction": discarded, "generator_version": first_group[0].generator_version})

SOFTWARE_SNAPSHOT = {
    "trl": "1.9.2",
    "transformers": "5.14.1",
    "datasets": "5.0.1",
    "peft": "0.20.0",
    "accelerate": "1.14.0",
    "bitsandbytes": "0.50.0",
    "model": "Qwen/Qwen3-0.6B",
}
print(SOFTWARE_SNAPSHOT)

if RUN_REAL_MODEL:
    from datasets import load_dataset
    from peft import LoraConfig
    from trl import GRPOConfig as TRLGRPOConfig, GRPOTrainer
    from trl.rewards import accuracy_reward

    train_dataset = load_dataset("trl-lib/DeepMath-103K", split="train[:128]")
    training_args = TRLGRPOConfig(
        output_dir="module-09-qwen3-grpo",
        learning_rate=5e-6,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        max_completion_length=192,
        num_generations=4,
        num_iterations=1,
        loss_type="dapo",
        scale_rewards="group",
        importance_sampling_level="token",
        epsilon=0.2,
        epsilon_high=0.28,
        beta=0.0,
        mask_truncated_completions=True,
        max_steps=2 if SMOKE_MODE else 80,
        logging_steps=1,
        report_to="none",
    )
    peft_config = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.05, task_type="CAUSAL_LM")
    trainer = GRPOTrainer(
        model=SOFTWARE_SNAPSHOT["model"],
        reward_funcs=accuracy_reward,
        args=training_args,
        train_dataset=train_dataset,
        peft_config=peft_config,
    )
    trainer.train()
else:
    print(REAL_MODEL_SKIPPED)