Module 14. Learning from Feedback: RLHF, DPO, RLOO, and GRPO

# Install only missing dependencies; ordinary Colab already has them.
import importlib.util
import subprocess
import sys

packages = {"numpy": "numpy", "matplotlib": "matplotlib"}
missing = [pip_name for module, pip_name in packages.items() if importlib.util.find_spec(module) is None]
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])
print('Environment ready.')

import hashlib
import json
import math
from dataclasses import dataclass
from typing import Literal

import matplotlib.pyplot as plt
import numpy as np

_REFERENCE_JSON = '{\r\n  "baseline_estimator": {\r\n    "bernoulli_p": 0.3,\r\n    "full_group_mean_expected": 0.18375,\r\n    "group_size": 8,\r\n    "true_gradient": 0.21\r\n  },\r\n  "bt": {\r\n    "deltas": [\r\n      0.5,\r\n      1.0,\r\n      2.0,\r\n      4.0\r\n    ],\r\n    "inverse_margin": 1.0986122886681098,\r\n    "inverse_win_rate": 0.75,\r\n    "probabilities": [\r\n      0.6224593312018546,\r\n      0.7310585786300049,\r\n      0.8807970779778824,\r\n      0.9820137900379085\r\n    ]\r\n  },\r\n  "dpo": {\r\n    "batch": {\r\n      "accuracy": [\r\n        true,\r\n        true,\r\n        false\r\n      ],\r\n      "beta": 0.1,\r\n      "completion_mask": [\r\n        [\r\n          1,\r\n          1,\r\n          1,\r\n          1,\r\n          0\r\n        ],\r\n        [\r\n          1,\r\n          1,\r\n          1,\r\n          1,\r\n          0\r\n        ],\r\n        [\r\n          1,\r\n          1,\r\n          1,\r\n          1,\r\n          0\r\n        ]\r\n      ],\r\n      "loss": [\r\n        0.6117553388706229,\r\n        0.6611752126141188,\r\n        0.7709570477895321\r\n      ],\r\n      "margin": [\r\n        0.16999999999999993,\r\n        0.06500000000000004,\r\n        -0.15000000000000002\r\n      ],\r\n      "mean_loss": 0.681295866424758,\r\n      "sequence_logps": {\r\n        "policy_chosen": [\r\n          -12.0,\r\n          -12.0,\r\n          -10.5\r\n        ],\r\n        "policy_rejected": [\r\n          -11.5,\r\n          -11.5,\r\n          -11.0\r\n        ],\r\n        "reference_chosen": [\r\n          -13.0,\r\n          -12.8,\r\n          -10.0\r\n        ],\r\n        "reference_rejected": [\r\n          -10.8,\r\n          -11.65,\r\n          -12.0\r\n        ]\r\n      },\r\n      "token_logps": {\r\n        "policy_chosen": [\r\n          [\r\n            -2.0,\r\n            -3.0,\r\n            -3.5,\r\n            -3.5,\r\n            99.0\r\n          ],\r\n          [\r\n            -3.0,\r\n            -3.0,\r\n            -3.0,\r\n            -3.0,\r\n            99.0\r\n          ],\r\n          [\r\n            -2.0,\r\n            -2.5,\r\n            -3.0,\r\n            -3.0,\r\n            99.0\r\n          ]\r\n        ],\r\n        "policy_rejected": [\r\n          [\r\n            -2.0,\r\n            -2.5,\r\n            -3.0,\r\n            -4.0,\r\n            99.0\r\n          ],\r\n          [\r\n            -2.5,\r\n            -2.5,\r\n            -3.0,\r\n            -3.5,\r\n            99.0\r\n          ],\r\n          [\r\n            -2.0,\r\n            -2.5,\r\n            -3.0,\r\n            -3.5,\r\n            99.0\r\n          ]\r\n        ],\r\n        "reference_chosen": [\r\n          [\r\n            -2.5,\r\n            -3.0,\r\n            -3.5,\r\n            -4.0,\r\n            99.0\r\n          ],\r\n          [\r\n            -3.2,\r\n            -3.2,\r\n            -3.2,\r\n            -3.2,\r\n            99.0\r\n          ],\r\n          [\r\n            -2.0,\r\n            -2.5,\r\n            -2.5,\r\n            -3.0,\r\n            99.0\r\n          ]\r\n        ],\r\n        "reference_rejected": [\r\n          [\r\n            -1.8,\r\n            -2.2,\r\n            -3.0,\r\n            -3.8,\r\n            99.0\r\n          ],\r\n          [\r\n            -2.5,\r\n            -2.65,\r\n            -3.0,\r\n            -3.5,\r\n            99.0\r\n          ],\r\n          [\r\n            -2.5,\r\n            -3.0,\r\n            -3.0,\r\n            -3.5,\r\n            99.0\r\n          ]\r\n        ]\r\n      },\r\n      "weight": [\r\n        0.45760205922564906,\r\n        0.4837557189379275,\r\n        0.5374298453437495\r\n      ]\r\n    },\r\n    "scalar_cases": {\r\n      "pair_II_beta_01": {\r\n        "beta": 0.1,\r\n        "loss": 0.6611752126141188,\r\n        "margin": 0.06500000000000004,\r\n        "policy_chosen": -12.0,\r\n        "policy_rejected": -11.5,\r\n        "reference_chosen": -12.8,\r\n        "reference_rejected": -11.65,\r\n        "weight": 0.4837557189379275\r\n      },\r\n      "pair_I_beta_01": {\r\n        "beta": 0.1,\r\n        "loss": 0.6117553388706229,\r\n        "margin": 0.16999999999999993,\r\n        "policy_chosen": -12.0,\r\n        "policy_rejected": -11.5,\r\n        "reference_chosen": -13.0,\r\n        "reference_rejected": -10.8,\r\n        "weight": 0.45760205922564906\r\n      },\r\n      "pair_I_beta_05": {\r\n        "beta": 0.5,\r\n        "loss": 0.35586506844219595,\r\n        "margin": 0.8499999999999996,\r\n        "policy_chosen": -12.0,\r\n        "policy_rejected": -11.5,\r\n        "reference_chosen": -13.0,\r\n        "reference_rejected": -10.8,\r\n        "weight": 0.2994328575260271\r\n      }\r\n    }\r\n  },\r\n  "grpo": {\r\n    "advantages": [\r\n      1.2909944487358056,\r\n      -0.7745966692414834,\r\n      -0.7745966692414834,\r\n      1.2909944487358056,\r\n      -0.7745966692414834,\r\n      -0.7745966692414834,\r\n      1.2909944487358056,\r\n      -0.7745966692414834\r\n    ],\r\n    "degenerate_probability": 0.05771361999999997,\r\n    "group_size": 8,\r\n    "mean": 0.375,\r\n    "pass_probability": 0.3,\r\n    "rewards": [\r\n      1.0,\r\n      0.0,\r\n      0.0,\r\n      1.0,\r\n      0.0,\r\n      0.0,\r\n      1.0,\r\n      0.0\r\n    ],\r\n    "std": 0.4841229182759271\r\n  },\r\n  "memory": {\r\n    "bytes_per_parameter": {\r\n      "frozen_bf16": 2,\r\n      "train": 16\r\n    },\r\n    "dpo_grpo_gib": 134.61336493492126,\r\n    "frozen_gib": 14.957040548324585,\r\n    "n_params": 8030000000.0,\r\n    "ppo_gib": 269.22672986984253,\r\n    "ratio": 2.0,\r\n    "train_gib": 119.65632438659668\r\n  },\r\n  "rollout": {\r\n    "accelerator_flops": 990000000000000.0,\r\n    "accelerators": 8,\r\n    "completion_tokens": 4000,\r\n    "decode_tokens_per_second_per_accelerator": 209.4,\r\n    "generated_tokens": 8192000,\r\n    "generation_minutes": 81.50270614453994,\r\n    "generations_per_prompt": 8,\r\n    "mfu": 0.4,\r\n    "prompts": 256,\r\n    "ratio": 39.25108921772578,\r\n    "training_minutes": 2.0764444444444448\r\n  },\r\n  "schema_version": 2\r\n}\r\n'
_REFERENCE_SHA256 = '8d1a0ad86aeeff2b0bc870f319d1bd107a12d611999df62fb8e0fb119ff74799'
assert hashlib.sha256(_REFERENCE_JSON.encode("utf-8")).hexdigest() == _REFERENCE_SHA256
REF = json.loads(_REFERENCE_JSON)
assert REF["schema_version"] == 2

np.set_printoptions(precision=6, suppress=True)
print('Module-14 reference payload unpacked and verified', _REFERENCE_SHA256[:12])

def stable_sigmoid(x: np.ndarray | float) -> np.ndarray:
    """Compute sigmoid(x) without exponential overflow."""

    array = np.asarray(x, dtype=np.float64)
    return np.exp(-np.logaddexp(0.0, -array))


def logit(probability: np.ndarray | float) -> np.ndarray:
    """Invert the sigmoid for probabilities strictly between zero and one."""

    p = np.asarray(probability, dtype=np.float64)
    if np.any(~np.isfinite(p)) or np.any((p <= 0.0) | (p >= 1.0)):
        raise ValueError("probabilities must be finite and strictly between 0 and 1")
    return np.log(p) - np.log1p(-p)


bt = REF["bt"]
deltas = np.asarray(bt["deltas"], dtype=np.float64)
probabilities = stable_sigmoid(deltas)
np.testing.assert_allclose(probabilities, bt["probabilities"], rtol=0, atol=1e-15)
inverse = float(logit(bt["inverse_win_rate"]))
assert abs(inverse - bt["inverse_margin"]) < 1e-15
assert stable_sigmoid(-1000.0) == 0.0 and stable_sigmoid(1000.0) == 1.0
print('BT table and the 75% → ln 3 inverse read reproduced')
for delta, probability in zip(deltas, probabilities):
    print("  Δr = {:.1f} -> P = {:.3f}".format(delta, probability))

@dataclass(frozen=True)
class DPOResult:
    """Vectorized DPO diagnostics for one scalar or a broadcastable batch."""

    margin: np.ndarray
    loss: np.ndarray
    weight: np.ndarray
    accuracy: np.ndarray


def dpo_loss(
    policy_chosen: np.ndarray | float,
    policy_rejected: np.ndarray | float,
    reference_chosen: np.ndarray | float,
    reference_rejected: np.ndarray | float,
    beta: float = 0.1,
) -> DPOResult:
    """Compute stable DPO from completion-level log-probability sums.

    All four inputs must broadcast to one shape. `weight` is sigmoid(-margin), the
    positive magnitude of d loss / d margin. The complete parameter gradient also
    contains beta and derivatives of the sequence log-probabilities.
    """

    if not math.isfinite(beta) or beta <= 0:
        raise ValueError("beta must be finite and positive")
    arrays = np.broadcast_arrays(
        np.asarray(policy_chosen, dtype=np.float64),
        np.asarray(policy_rejected, dtype=np.float64),
        np.asarray(reference_chosen, dtype=np.float64),
        np.asarray(reference_rejected, dtype=np.float64),
    )
    if not all(np.all(np.isfinite(array)) for array in arrays):
        raise ValueError("sequence log-probabilities must be finite")
    pc, pr, rc, rr = arrays
    margin = beta * ((pc - rc) - (pr - rr))
    loss = np.logaddexp(0.0, -margin)       # softplus(-margin)
    weight = stable_sigmoid(-margin)
    return DPOResult(margin=margin, loss=loss, weight=weight, accuracy=margin > 0)


scalar_cases = REF["dpo"]["scalar_cases"]
for name, case in scalar_cases.items():
    result = dpo_loss(
        case["policy_chosen"], case["policy_rejected"],
        case["reference_chosen"], case["reference_rejected"], case["beta"],
    )
    assert abs(float(result.margin) - case["margin"]) < 1e-14
    assert abs(float(result.loss) - case["loss"]) < 1e-14
    assert abs(float(result.weight) - case["weight"]) < 1e-14

extreme = dpo_loss(0.0, 0.0, -10_000.0, 0.0, beta=0.1)
assert np.isfinite(extreme.loss) and float(extreme.loss) < 1e-300
hard_wrong = dpo_loss(-10_000.0, 0.0, 0.0, 0.0, beta=0.1)
assert np.isfinite(hard_wrong.loss) and abs(float(hard_wrong.loss) - 1000.0) < 1e-12
print('Scalar DPO trace and extreme margins verified')

def masked_sequence_logps(token_logps: np.ndarray, completion_mask: np.ndarray) -> np.ndarray:
    """Sum `[batch, tokens]` logps only where the completion mask is true."""

    values = np.asarray(token_logps, dtype=np.float64)
    mask = np.asarray(completion_mask, dtype=bool)
    if values.ndim != 2 or values.shape != mask.shape:
        raise ValueError("token_logps and completion_mask must share [batch, tokens]")
    if np.any(mask.sum(axis=1) == 0):
        raise ValueError("every sequence needs at least one completion token")
    if not np.all(np.isfinite(values[mask])):
        raise ValueError("active token log-probabilities must be finite")
    return np.where(mask, values, 0.0).sum(axis=1, dtype=np.float64)


batch = REF["dpo"]["batch"]
mask = np.asarray(batch["completion_mask"], dtype=bool)
token_logps = {name: np.asarray(values, dtype=np.float64) for name, values in batch["token_logps"].items()}
sequence_logps = {name: masked_sequence_logps(values, mask) for name, values in token_logps.items()}
for name, expected in batch["sequence_logps"].items():
    np.testing.assert_allclose(sequence_logps[name], expected, rtol=0, atol=1e-14)

result = dpo_loss(
    sequence_logps["policy_chosen"], sequence_logps["policy_rejected"],
    sequence_logps["reference_chosen"], sequence_logps["reference_rejected"],
    beta=batch["beta"],
)
np.testing.assert_allclose(result.margin, batch["margin"], rtol=0, atol=1e-14)
np.testing.assert_allclose(result.loss, batch["loss"], rtol=0, atol=1e-14)
np.testing.assert_allclose(result.weight, batch["weight"], rtol=0, atol=1e-14)
assert result.accuracy.tolist() == batch["accuracy"]
assert abs(float(result.loss.mean()) - batch["mean_loss"]) < 1e-14

# Changing an ignored sentinel cannot affect a sequence score.
mutated = token_logps["policy_chosen"].copy()
mutated[~mask] = -1e9
np.testing.assert_array_equal(masked_sequence_logps(mutated, mask), sequence_logps["policy_chosen"])
print('Three pairs passed mask → sequence logp → DPO; one is deliberately misordered')
print("margin:", result.margin, "loss:", result.loss, "accuracy:", result.accuracy)

margin_grid = np.linspace(-8.0, 8.0, 801)
loss_grid = np.logaddexp(0.0, -margin_grid)
weight_grid = stable_sigmoid(-margin_grid)
assert np.all(np.diff(loss_grid) < 0) and np.all(np.diff(weight_grid) < 0)

plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 12,
                     "axes.labelsize": 13, "xtick.labelsize": 11,
                     "ytick.labelsize": 11, "legend.fontsize": 11,
                     "text.color": "#2E3440", "axes.labelcolor": "#2E3440",
                     "axes.edgecolor": "#D8DEE9", "grid.color": "#D8DEE9"})
fig, ax = plt.subplots(figsize=(11, 5.2), dpi=150, layout="constrained")
ax.plot(margin_grid, loss_grid, label='DPO loss = softplus(−z)', color="#5E81AC", linewidth=2)
ax.plot(margin_grid, weight_grid, label='Local weight = sigmoid(−z)', color="#D08770", linewidth=2)
ax.scatter(result.margin, result.loss, marker="o", label='Pair losses', color="#5E81AC", zorder=4)
ax.scatter(result.margin, result.weight, marker="s", label='Pair weights', color="#D08770", zorder=4)
ax.axvline(0.0, color="#D8DEE9", linewidth=0.8)
ax.set_xlabel('DPO margin, z')
ax.set_ylabel('Value')
ax.set_ylim(-0.05, 2.5)
ax.grid(alpha=0.25)
ax.legend(ncol=2)
plt.show()

@dataclass(frozen=True)
class GroupAdvantages:
    """Standardized reward signal and the statistics used to create it."""

    values: np.ndarray
    mean: float
    std: float
    valid: bool


def group_standardized_advantages(
    rewards: np.ndarray,
    *,
    min_std: float = 1e-12,
    on_degenerate: Literal["raise", "zeros"] = "raise",
) -> GroupAdvantages:
    """Standardize one reward group and make zero-variance handling explicit."""

    r = np.asarray(rewards, dtype=np.float64)
    if r.ndim != 1 or r.size < 2 or not np.all(np.isfinite(r)):
        raise ValueError("rewards must be a finite one-dimensional group of size >= 2")
    if not math.isfinite(min_std) or min_std < 0:
        raise ValueError("min_std must be finite and non-negative")
    if on_degenerate not in ("raise", "zeros"):
        raise ValueError("unknown zero-variance policy")
    mean = float(r.mean())
    std = float(r.std(ddof=0))
    if std <= min_std:
        if on_degenerate == "zeros":
            return GroupAdvantages(np.zeros_like(r), mean, std, False)
        if on_degenerate != "raise":
            raise ValueError("on_degenerate must be 'raise' or 'zeros'")
        raise ValueError("degenerate group: reward standard deviation is zero")
    return GroupAdvantages((r - mean) / std, mean, std, True)


grpo = REF["grpo"]
group = group_standardized_advantages(np.asarray(grpo["rewards"], dtype=np.float64))
np.testing.assert_allclose(group.values, grpo["advantages"], rtol=0, atol=1e-14)
assert abs(group.mean - grpo["mean"]) < 1e-14
assert abs(group.std - grpo["std"]) < 1e-14
assert abs(float(group.values.sum())) < 1e-14
expected_degenerate = (1 - grpo["pass_probability"]) ** grpo["group_size"] + grpo["pass_probability"] ** grpo["group_size"]
assert abs(expected_degenerate - grpo["degenerate_probability"]) < 1e-15

try:
    group_standardized_advantages(np.zeros(8))
except ValueError as exc:
    assert "degenerate" in str(exc)
else:
    raise AssertionError("degenerate rewards must not silently produce NaN")
zero_policy = group_standardized_advantages(np.ones(8), on_degenerate="zeros")
assert not zero_policy.valid and np.all(zero_policy.values == 0)
print('The 3/8 GRPO trace, zero-sum advantages, and degenerate policy verified')

@dataclass(frozen=True)
class ClippedSurrogate:
    """Diagnostics for the token-level clipped policy-gradient term."""

    loss: float
    ratio: np.ndarray
    clipped_ratio: np.ndarray
    token_objective: np.ndarray


def grpo_clipped_surrogate(
    new_token_logps: np.ndarray,
    old_token_logps: np.ndarray,
    advantages: np.ndarray,
    completion_mask: np.ndarray,
    *,
    clip_low: float = 0.2,
    clip_high: float = 0.2,
) -> ClippedSurrogate:
    """Compute the masked token-level PPO/GRPO clipped surrogate.

    Inputs have shape `[group, tokens]`; advantages have shape `[group]`.
    The returned loss is the negative mean over active completion tokens.
    """

    new = np.asarray(new_token_logps, dtype=np.float64)
    old = np.asarray(old_token_logps, dtype=np.float64)
    mask = np.asarray(completion_mask, dtype=bool)
    adv = np.asarray(advantages, dtype=np.float64)
    if new.ndim != 2 or new.shape != old.shape or new.shape != mask.shape:
        raise ValueError("new, old, and mask must share [group, tokens]")
    if adv.shape != (new.shape[0],):
        raise ValueError("advantages must have shape [group]")
    if not np.all(np.isfinite(new[mask])) or not np.all(np.isfinite(old[mask])) or not np.all(np.isfinite(adv)):
        raise ValueError("active logps and advantages must be finite")
    if mask.sum() == 0:
        raise ValueError("the batch must contain active completion tokens")
    if not (math.isfinite(clip_low) and math.isfinite(clip_high) and 0 <= clip_low < 1 and clip_high >= 0):
        raise ValueError("clip_low must be in [0,1); clip_high must be non-negative")

    log_ratio = np.where(mask, new, 0.0) - np.where(mask, old, 0.0)
    if np.any(np.abs(log_ratio[mask]) > 700):
        raise FloatingPointError("token log-ratio is too large for finite exponentiation")
    ratio = np.exp(log_ratio)
    clipped_ratio = np.clip(ratio, 1.0 - clip_low, 1.0 + clip_high)
    expanded_adv = adv[:, None]
    token_objective = np.minimum(ratio * expanded_adv, clipped_ratio * expanded_adv)
    loss = -float(np.where(mask, token_objective, 0.0).sum() / mask.sum())
    return ClippedSurrogate(loss, ratio, clipped_ratio, token_objective)


old = np.full((8, 3), -2.0, dtype=np.float64)
desired_ratio = np.asarray([1.5, 1.1, 0.9, 1.0, 0.5, 1.0, 1.3, 0.7])[:, None]
new = old + np.log(np.broadcast_to(desired_ratio, old.shape))
token_mask = np.asarray([[1, 1, 0], [1, 1, 1], [1, 0, 0], [1, 1, 1],
                         [1, 1, 0], [1, 1, 1], [1, 1, 1], [1, 0, 0]], dtype=bool)
clipped = grpo_clipped_surrogate(new, old, group.values, token_mask)
assert math.isfinite(clipped.loss)
assert np.isclose(clipped.clipped_ratio[0, 0], 1.2)   # positive A: high ratio is capped
assert np.isclose(clipped.clipped_ratio[4, 0], 0.8)   # negative A: low ratio is capped
assert not np.allclose(clipped.ratio[token_mask], clipped.clipped_ratio[token_mask])
print('Clipping activated for both positive and negative advantages', "loss =", round(clipped.loss, 6))

baseline_ref = REF["baseline_estimator"]
p = baseline_ref["bernoulli_p"]
G = baseline_ref["group_size"]
true_gradient = baseline_ref["true_gradient"]

rng = np.random.default_rng(7)
trials = 120_000
actions = (rng.random((trials, G)) < p).astype(np.float64)
rewards = actions.copy()
scores = actions - p                         # d log Bernoulli(a; sigmoid(theta)) / d theta

raw = np.mean(rewards * scores, axis=1)
reward_sum = rewards.sum(axis=1, keepdims=True)
leave_one_out = (reward_sum - rewards) / (G - 1)
loo = np.mean((rewards - leave_one_out) * scores, axis=1)
full_group_mean = rewards.mean(axis=1, keepdims=True)
full = np.mean((rewards - full_group_mean) * scores, axis=1)

expected_full = baseline_ref["full_group_mean_expected"]
assert abs(raw.mean() - true_gradient) < 0.0015
assert abs(loo.mean() - true_gradient) < 0.0015
assert abs(full.mean() - expected_full) < 0.0015
assert raw.var() / loo.var() > 1.8
print('Raw and leave-one-out matched the exact gradient; the full-group mean showed the expected scale bias')
print("true = {:.5f} · raw = {:.5f} · LOO = {:.5f} · full-group = {:.5f}".format(
    true_gradient, raw.mean(), loo.mean(), full.mean()
))
print("Var(raw)/Var(LOO) = {:.2f}".format(raw.var() / loo.var()))

def stage_memory_gib(n_params: float, stack: list[str]) -> float:
    """Estimate parameter-state memory in GiB under the module's floor model."""

    if not math.isfinite(n_params) or n_params <= 0:
        raise ValueError("n_params must be finite and positive")
    bytes_per_parameter = {"train": 16, "frozen_bf16": 2}
    try:
        total_bytes = sum(bytes_per_parameter[kind] * n_params for kind in stack)
    except KeyError as exc:
        raise ValueError("unknown stack component: " + str(exc.args[0])) from exc
    return total_bytes / 2**30


def rollout_scenario_minutes(
    *,
    prompts: int,
    generations_per_prompt: int,
    completion_tokens: int,
    n_params: float,
    accelerators: int,
    decode_tokens_per_second_per_accelerator: float,
    accelerator_flops: float,
    mfu: float,
) -> dict[str, float]:
    """Estimate generated tokens, generation time, and one 6ND training pass."""

    integer_args = (prompts, generations_per_prompt, completion_tokens, accelerators)
    if any(not isinstance(value, int) or value <= 0 for value in integer_args):
        raise ValueError("count arguments must be positive integers")
    positive = (n_params, decode_tokens_per_second_per_accelerator, accelerator_flops)
    if any(not math.isfinite(value) or value <= 0 for value in positive):
        raise ValueError("rates and n_params must be finite and positive")
    if not math.isfinite(mfu) or not (0 < mfu <= 1):
        raise ValueError("mfu must be in (0, 1]")
    tokens = prompts * generations_per_prompt * completion_tokens
    generation_minutes = tokens / (accelerators * decode_tokens_per_second_per_accelerator) / 60
    training_minutes = 6 * n_params * tokens / (accelerators * accelerator_flops * mfu) / 60
    return {
        "generated_tokens": float(tokens),
        "generation_minutes": generation_minutes,
        "training_minutes": training_minutes,
        "ratio": generation_minutes / training_minutes,
    }


memory = REF["memory"]
n_params = memory["n_params"]
train_gib = stage_memory_gib(n_params, ["train"])
frozen_gib = stage_memory_gib(n_params, ["frozen_bf16"])
ppo_gib = stage_memory_gib(n_params, ["train", "train", "frozen_bf16", "frozen_bf16"])
dpo_gib = stage_memory_gib(n_params, ["train", "frozen_bf16"])
np.testing.assert_allclose(
    [train_gib, frozen_gib, ppo_gib, dpo_gib],
    [memory["train_gib"], memory["frozen_gib"], memory["ppo_gib"], memory["dpo_grpo_gib"]],
    rtol=0, atol=1e-12,
)
assert abs(ppo_gib / dpo_gib - 2.0) < 1e-14

rollout = REF["rollout"]
estimate = rollout_scenario_minutes(
    prompts=rollout["prompts"],
    generations_per_prompt=rollout["generations_per_prompt"],
    completion_tokens=rollout["completion_tokens"],
    n_params=n_params,
    accelerators=rollout["accelerators"],
    decode_tokens_per_second_per_accelerator=rollout["decode_tokens_per_second_per_accelerator"],
    accelerator_flops=rollout["accelerator_flops"],
    mfu=rollout["mfu"],
)
for key in ("generated_tokens", "generation_minutes", "training_minutes", "ratio"):
    assert abs(estimate[key] - rollout[key]) < 1e-10
print('GiB footprint and rollout scenario reproduced')
print("PPO = {:.1f} GiB · DPO + ref = {:.1f} GiB · precached ref = {:.1f} GiB".format(
    ppo_gib, dpo_gib, train_gib
))
resource_line = (
    "траектории = {:.2f} млн токенов · генерация = {:.1f} мин · обучение = {:.1f} мин · отношение = {:.1f}:1"
    if 'en' == "ru"
    else "rollouts = {:.2f}M tokens · generation = {:.1f} min · training = {:.1f} min · ratio = {:.1f}:1"
)
print(resource_line.format(
    estimate["generated_tokens"] / 1e6, estimate["generation_minutes"],
    estimate["training_minutes"], estimate["ratio"]
))

import itertools
actions_exact = np.array(list(itertools.product((0., 1.), repeat=G)))
probabilities = (p**actions_exact * (1-p)**(1-actions_exact)).prod(1)
exact_scores = actions_exact-p
centered = actions_exact-actions_exact.mean(1, keepdims=True)
exact_full = float(probabilities @ (centered*exact_scores).mean(1))
exact_loo = float(probabilities @ (centered*G/(G-1)*exact_scores).mean(1))
np.testing.assert_allclose([exact_full, exact_loo], [.18375, .21], atol=1e-14, rtol=0)
assert np.isclose(probabilities.sum(), 1)
# Masked padding must not leak nonfinite placeholders into exponentiation.
padded_new, padded_old = new.copy(), old.copy()
padded_new[~token_mask], padded_old[~token_mask] = np.nan, np.inf
padded = grpo_clipped_surrogate(padded_new, padded_old, group.values, token_mask)
assert np.isclose(padded.loss, clipped.loss) and np.all(np.isfinite(padded.ratio))
try:
    group_standardized_advantages(np.array([0., 1.]), on_degenerate="typo")
except ValueError:
    pass
else:
    raise AssertionError("invalid policy must be rejected for every group")
print("Exact expectations:", exact_full, exact_loo)

core_namespace = {}
exec('import math\nimport numpy as np\n\n\ndef stable_sigmoid(x):\n    """Numerically stable sigmoid."""\n    x = np.asarray(x, dtype=np.float64)\n    return np.exp(-np.logaddexp(0.0, -x))\n\n\ndef dpo_loss(policy_lp_w, policy_lp_l, ref_lp_w, ref_lp_l, beta=0.1):\n    """DPO from completion-token log-probability sums."""\n    if not math.isfinite(beta) or beta <= 0:\n        raise ValueError("beta must be finite and positive")\n\n    pw, pl, rw, rl = np.broadcast_arrays(\n        np.asarray(policy_lp_w, dtype=float),\n        np.asarray(policy_lp_l, dtype=float),\n        np.asarray(ref_lp_w, dtype=float),\n        np.asarray(ref_lp_l, dtype=float),\n    )\n    if not all(np.all(np.isfinite(v)) for v in (pw, pl, rw, rl)):\n        raise ValueError("log probabilities must be finite")\n\n    z = beta * ((pw - rw) - (pl - rl))\n    loss = np.logaddexp(0.0, -z)\n    local_slope = stable_sigmoid(-z)\n    return loss, local_slope, z > 0\n\n\ndef grpo_advantages(rewards, *, on_degenerate="raise"):\n    """Population normalization with explicit zero-variance handling."""\n    r = np.asarray(rewards, dtype=np.float64)\n    if r.ndim != 1 or r.size < 2 or not np.all(np.isfinite(r)):\n        raise ValueError("rewards must be a finite one-dimensional group of size >= 2")\n\n    if on_degenerate not in ("raise", "zeros"):\n        raise ValueError("on_degenerate must be raise or zeros")\n    std = float(r.std(ddof=0))\n    if std == 0.0:\n        if on_degenerate == "zeros":\n            return np.zeros_like(r)\n        raise ValueError("the group contains no relative signal")\n    return (r - r.mean()) / std\n\n\ndef stage_memory_gib(n_params, stack):\n    """Teaching estimate of persistent model states, excluding activations."""\n    if not math.isfinite(n_params) or n_params <= 0:\n        raise ValueError("n_params must be finite and positive")\n    bytes_per_parameter = {"train": 16, "frozen_bf16": 2}\n    return sum(bytes_per_parameter[kind] * n_params for kind in stack) / 2**30\n\n\nppo = stage_memory_gib(\n    8.03e9,\n    ["train", "train", "frozen_bf16", "frozen_bf16"],\n)\ndpo = stage_memory_gib(8.03e9, ["train", "frozen_bf16"])\n', core_namespace)
np.testing.assert_allclose(core_namespace['grpo_advantages'](grpo['rewards']), group.values)
assert abs(core_namespace['ppo']-ppo_gib) < 1e-12

"""Complete CPU feedback loops on a restricted one-token answer space.

Prompt: is a > b? Answer: NO or YES, followed by a deterministic stop.
This is a contextual bandit, not a language model or a reasoning benchmark.
"""
import copy
import io
import math
import importlib.util
import subprocess
import sys
if importlib.util.find_spec('torch') is None:
    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', 'torch'])
import torch
from torch import nn
from torch.nn import functional as F

torch.set_num_threads(1)
torch.manual_seed(140)
pairs = [(a, b) for a in range(8) for b in range(8)]
train_pairs = [p for p in pairs if abs(p[0]-p[1]) != 4]
test_pairs = [p for p in pairs if abs(p[0]-p[1]) == 4]
assert len(train_pairs) == 56 and len(test_pairs) == 8
assert not set(train_pairs) & set(test_pairs)

def features(prompts):
    # Deliberately useful, declared features; no claim of learned arithmetic.
    return torch.tensor([[1., (a-b)/7.] for a, b in prompts])

def verifier(prompts, answer_ids):
    correct = torch.tensor([int(a > b) for a, b in prompts])
    return answer_ids.eq(correct).float()

class AnswerPolicy(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.zeros(2, 2))
    def forward(self, prompts):
        return F.log_softmax(features(prompts) @ self.weights, dim=-1)

def chosen_logps(model, prompts, ids):
    # One stochastic completion token; the fixed stop has log probability zero.
    return model(prompts).gather(1, ids[:, None]).squeeze(1)

@torch.no_grad()
def metrics(model, prompts, reference):
    lp = model(prompts)
    correct = torch.tensor([int(a > b) for a, b in prompts])
    return dict(correct_probability=float(lp.exp().gather(1, correct[:, None]).mean()),
                greedy_accuracy=float(verifier(prompts, lp.argmax(-1)).mean()),
                kl=float((lp.exp() * (lp-reference(prompts))).sum(-1).mean()))

def reload_check(model):
    buffer = io.BytesIO()
    torch.save(model.state_dict(), buffer)
    buffer.seek(0)
    restored = AnswerPolicy()
    restored.load_state_dict(torch.load(buffer, weights_only=True))
    assert torch.equal(model(test_pairs), restored(test_pairs))

def run_dpo(beta=.5, *, steps=140, seed=140, init_std=0., verbose=True):
    """Train on fixed pairs; seed controls only optional random initialization.

    The reference copies the initial policy. Match seed and init_std across beta
    values to hold both starting weights and cached reference scores constant.
    With init_std=0 the full-batch experiment is deterministic for every seed.
    """
    if not math.isfinite(beta) or beta <= 0:
        raise ValueError('beta must be finite and positive')
    if isinstance(steps, bool) or not isinstance(steps, int) or steps <= 0:
        raise ValueError('steps must be a positive integer')
    if not math.isfinite(init_std) or init_std < 0:
        raise ValueError('init_std must be finite and nonnegative')
    model = AnswerPolicy()
    if init_std > 0:
        generator = torch.Generator().manual_seed(seed)
        with torch.no_grad():
            model.weights.normal_(0., init_std, generator=generator)
    reference = copy.deepcopy(model).requires_grad_(False)
    frozen = reference.weights.detach().clone()
    chosen = torch.tensor([int(a > b) for a, b in train_pairs])
    rejected = 1-chosen
    with torch.no_grad():
        cached_gap = chosen_logps(reference, train_pairs, chosen)-chosen_logps(reference, train_pairs, rejected)
        initial_gap = chosen_logps(model, train_pairs, chosen)-chosen_logps(model, train_pairs, rejected)
        initial_loss = float(F.softplus(-beta*(initial_gap-cached_gap)).mean())
    optimizer = torch.optim.Adam(model.parameters(), lr=.12)
    before = metrics(model, test_pairs, reference)
    for _ in range(steps):
        gap = chosen_logps(model, train_pairs, chosen)-chosen_logps(model, train_pairs, rejected)
        loss = F.softplus(-beta*(gap-cached_gap)).mean()
        optimizer.zero_grad(); loss.backward(); optimizer.step()
    with torch.no_grad():
        gap = chosen_logps(model, train_pairs, chosen)-chosen_logps(model, train_pairs, rejected)
        final_loss = float(F.softplus(-beta*(gap-cached_gap)).mean())
        pair_accuracy = float((gap-cached_gap > 0).float().mean())
    after = metrics(model, test_pairs, reference)
    # Check the computation, not a quality threshold tied to one beta and budget.
    assert math.isfinite(final_loss) and torch.isfinite(model.weights).all()
    assert torch.equal(frozen, reference.weights) and reference.weights.grad is None
    reload_check(model)
    result = dict(beta=beta, steps=steps, seed=seed, init_std=init_std,
                  initial_weights=frozen.tolist(), before=before, after=after,
                  initial_loss=initial_loss, loss=final_loss,
                  pair_accuracy=pair_accuracy, pairs=len(train_pairs))
    if verbose:
        print('DPO:', result)
    return result

def run_grpo():
    torch.manual_seed(141)
    model = AnswerPolicy()
    reference = copy.deepcopy(model).requires_grad_(False)
    frozen = reference.weights.detach().clone()
    optimizer = torch.optim.Adam(model.parameters(), lr=.08)
    before = metrics(model, test_pairs, reference)
    group_size, steps, updates_per_rollout = 8, 100, 2
    zero_groups = 0
    clips = []
    expanded_prompts = [p for p in train_pairs for _ in range(group_size)]
    for _ in range(steps):
        # This snapshot describes the sampler. The reference stays fixed throughout.
        with torch.no_grad():
            sampler_lp = model(train_pairs)
            answers = torch.multinomial(sampler_lp.exp(), group_size, replacement=True)
            old_lp = sampler_lp.gather(1, answers)
            rewards = verifier(expanded_prompts, answers.flatten()).reshape(-1, group_size)
            centered = rewards-rewards.mean(1, keepdim=True)
            std = rewards.std(1, correction=0, keepdim=True)
            valid = std > 0
            advantages = torch.where(valid, centered/std.clamp_min(1e-8), torch.zeros_like(centered))
            zero_groups += int((~valid).sum())
        for _ in range(updates_per_rollout):
            lp = model(train_pairs)
            new_lp = lp.gather(1, answers)
            ratio = (new_lp-old_lp).exp()
            clipped = ratio.clamp(.8, 1.2)
            surrogate = torch.minimum(ratio*advantages, clipped*advantages).mean()
            # Exact categorical KL is available for this two-action teaching model.
            kl = (lp.exp()*(lp-reference(train_pairs))).sum(-1).mean()
            loss = -surrogate + .02*kl
            optimizer.zero_grad(); loss.backward(); optimizer.step()
            clips.append(float(((ratio < .8) | (ratio > 1.2)).float().mean()))
    after = metrics(model, test_pairs, reference)
    assert after['correct_probability'] > before['correct_probability'] + .3
    assert not advantages.requires_grad and not old_lp.requires_grad
    assert torch.equal(frozen, reference.weights) and reference.weights.grad is None
    reload_check(model)
    result = dict(before=before, after=after, generated_answers=steps*len(train_pairs)*group_size,
                  optimizer_steps=steps*updates_per_rollout, zero_groups=zero_groups,
                  mean_clip_fraction=sum(clips)/len(clips), zero_group_policy='zero advantages; keep KL')
    print('GRPO:', result)
    return result

if __name__ == '__main__':
    dpo_training_result = run_dpo()
    grpo_training_result = run_grpo()

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

"""Complete CPU feedback loops on a restricted one-token answer space.

Prompt: is a > b? Answer: NO or YES, followed by a deterministic stop.
This is a contextual bandit, not a language model or a reasoning benchmark.
"""
import copy
import io
import math
import importlib.util
import subprocess
import sys
if importlib.util.find_spec('torch') is None:
    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', 'torch'])
import torch
from torch import nn
from torch.nn import functional as F

torch.set_num_threads(1)
torch.manual_seed(140)
pairs = [(a, b) for a in range(8) for b in range(8)]
train_pairs = [p for p in pairs if abs(p[0]-p[1]) != 4]
test_pairs = [p for p in pairs if abs(p[0]-p[1]) == 4]
assert len(train_pairs) == 56 and len(test_pairs) == 8
assert not set(train_pairs) & set(test_pairs)

def features(prompts):
    # Deliberately useful, declared features; no claim of learned arithmetic.
    return torch.tensor([[1., (a-b)/7.] for a, b in prompts])

def verifier(prompts, answer_ids):
    correct = torch.tensor([int(a > b) for a, b in prompts])
    return answer_ids.eq(correct).float()

class AnswerPolicy(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.zeros(2, 2))
    def forward(self, prompts):
        return F.log_softmax(features(prompts) @ self.weights, dim=-1)

def chosen_logps(model, prompts, ids):
    # One stochastic completion token; the fixed stop has log probability zero.
    return model(prompts).gather(1, ids[:, None]).squeeze(1)

@torch.no_grad()
def metrics(model, prompts, reference):
    lp = model(prompts)
    correct = torch.tensor([int(a > b) for a, b in prompts])
    return dict(correct_probability=float(lp.exp().gather(1, correct[:, None]).mean()),
                greedy_accuracy=float(verifier(prompts, lp.argmax(-1)).mean()),
                kl=float((lp.exp() * (lp-reference(prompts))).sum(-1).mean()))

def reload_check(model):
    buffer = io.BytesIO()
    torch.save(model.state_dict(), buffer)
    buffer.seek(0)
    restored = AnswerPolicy()
    restored.load_state_dict(torch.load(buffer, weights_only=True))
    assert torch.equal(model(test_pairs), restored(test_pairs))

def run_dpo(beta=.5, *, steps=140, seed=140, init_std=0., verbose=True):
    """Train on fixed pairs; seed controls only optional random initialization.

    The reference copies the initial policy. Match seed and init_std across beta
    values to hold both starting weights and cached reference scores constant.
    With init_std=0 the full-batch experiment is deterministic for every seed.
    """
    if not math.isfinite(beta) or beta <= 0:
        raise ValueError('beta must be finite and positive')
    if isinstance(steps, bool) or not isinstance(steps, int) or steps <= 0:
        raise ValueError('steps must be a positive integer')
    if not math.isfinite(init_std) or init_std < 0:
        raise ValueError('init_std must be finite and nonnegative')
    model = AnswerPolicy()
    if init_std > 0:
        generator = torch.Generator().manual_seed(seed)
        with torch.no_grad():
            model.weights.normal_(0., init_std, generator=generator)
    reference = copy.deepcopy(model).requires_grad_(False)
    frozen = reference.weights.detach().clone()
    chosen = torch.tensor([int(a > b) for a, b in train_pairs])
    rejected = 1-chosen
    with torch.no_grad():
        cached_gap = chosen_logps(reference, train_pairs, chosen)-chosen_logps(reference, train_pairs, rejected)
        initial_gap = chosen_logps(model, train_pairs, chosen)-chosen_logps(model, train_pairs, rejected)
        initial_loss = float(F.softplus(-beta*(initial_gap-cached_gap)).mean())
    optimizer = torch.optim.Adam(model.parameters(), lr=.12)
    before = metrics(model, test_pairs, reference)
    for _ in range(steps):
        gap = chosen_logps(model, train_pairs, chosen)-chosen_logps(model, train_pairs, rejected)
        loss = F.softplus(-beta*(gap-cached_gap)).mean()
        optimizer.zero_grad(); loss.backward(); optimizer.step()
    with torch.no_grad():
        gap = chosen_logps(model, train_pairs, chosen)-chosen_logps(model, train_pairs, rejected)
        final_loss = float(F.softplus(-beta*(gap-cached_gap)).mean())
        pair_accuracy = float((gap-cached_gap > 0).float().mean())
    after = metrics(model, test_pairs, reference)
    # Check the computation, not a quality threshold tied to one beta and budget.
    assert math.isfinite(final_loss) and torch.isfinite(model.weights).all()
    assert torch.equal(frozen, reference.weights) and reference.weights.grad is None
    reload_check(model)
    result = dict(beta=beta, steps=steps, seed=seed, init_std=init_std,
                  initial_weights=frozen.tolist(), before=before, after=after,
                  initial_loss=initial_loss, loss=final_loss,
                  pair_accuracy=pair_accuracy, pairs=len(train_pairs))
    if verbose:
        print('DPO:', result)
    return result

def run_grpo():
    torch.manual_seed(141)
    model = AnswerPolicy()
    reference = copy.deepcopy(model).requires_grad_(False)
    frozen = reference.weights.detach().clone()
    optimizer = torch.optim.Adam(model.parameters(), lr=.08)
    before = metrics(model, test_pairs, reference)
    group_size, steps, updates_per_rollout = 8, 100, 2
    zero_groups = 0
    clips = []
    expanded_prompts = [p for p in train_pairs for _ in range(group_size)]
    for _ in range(steps):
        # This snapshot describes the sampler. The reference stays fixed throughout.
        with torch.no_grad():
            sampler_lp = model(train_pairs)
            answers = torch.multinomial(sampler_lp.exp(), group_size, replacement=True)
            old_lp = sampler_lp.gather(1, answers)
            rewards = verifier(expanded_prompts, answers.flatten()).reshape(-1, group_size)
            centered = rewards-rewards.mean(1, keepdim=True)
            std = rewards.std(1, correction=0, keepdim=True)
            valid = std > 0
            advantages = torch.where(valid, centered/std.clamp_min(1e-8), torch.zeros_like(centered))
            zero_groups += int((~valid).sum())
        for _ in range(updates_per_rollout):
            lp = model(train_pairs)
            new_lp = lp.gather(1, answers)
            ratio = (new_lp-old_lp).exp()
            clipped = ratio.clamp(.8, 1.2)
            surrogate = torch.minimum(ratio*advantages, clipped*advantages).mean()
            # Exact categorical KL is available for this two-action teaching model.
            kl = (lp.exp()*(lp-reference(train_pairs))).sum(-1).mean()
            loss = -surrogate + .02*kl
            optimizer.zero_grad(); loss.backward(); optimizer.step()
            clips.append(float(((ratio < .8) | (ratio > 1.2)).float().mean()))
    after = metrics(model, test_pairs, reference)
    assert after['correct_probability'] > before['correct_probability'] + .3
    assert not advantages.requires_grad and not old_lp.requires_grad
    assert torch.equal(frozen, reference.weights) and reference.weights.grad is None
    reload_check(model)
    result = dict(before=before, after=after, generated_answers=steps*len(train_pairs)*group_size,
                  optimizer_steps=steps*updates_per_rollout, zero_groups=zero_groups,
                  mean_clip_fraction=sum(clips)/len(clips), zero_group_policy='zero advantages; keep KL')
    print('GRPO:', result)
    return result


result_dpo=run_dpo()

result_grpo=run_grpo()