Модуль 6. LLM как политика: токены, логарифмы вероятностей и KL

from __future__ import annotations

from pathlib import Path
import os
import sys
import math
import numpy as np
import matplotlib.pyplot as plt


def find_repo_root(start: Path) -> Path:
    for candidate in (start, *start.parents):
        if (candidate / "modules/module-06/module_6_reference.py").is_file():
            return candidate
    raise FileNotFoundError("Could not find the RL for LLM repository root")


REPO_ROOT = find_repo_root(Path.cwd().resolve())
MODULE_DIR = REPO_ROOT / "modules/module-06"
sys.path.insert(0, str(MODULE_DIR))

from module_6_reference import (
    behavior_logprob,
    causal_response_logprobs,
    classify_stop_event,
    decoding_distribution_with_order,
    exact_kl_from_logits,
    kl_estimator_moments,
    masked_sequence_metrics,
    sequence_logprob_ratio,
    stable_log_softmax,
    validate_rollout_batch,
)

PROFILE = os.getenv("RL_LLM_M6_PROFILE", "smoke").strip().lower()
if PROFILE not in {"smoke", "full", "real-model"}:
    raise ValueError("RL_LLM_M6_PROFILE must be smoke, full, or real-model")
RUN_REAL_MODEL = os.getenv("RL_LLM_M6_RUN_REAL_MODEL", "0") == "1" or PROFILE == "real-model"
SEED = 20260805
RNG = np.random.default_rng(SEED)
N_MONTE_CARLO = 20_000 if PROFILE == "smoke" else 250_000
print({"profile": PROFILE, "run_real_model": RUN_REAL_MODEL, "repo": str(REPO_ROOT)})

large_logits = np.array([[1000.0, 999.0, 997.0], [-1000.0, -1001.0, -1002.0]])
log_probs = stable_log_softmax(large_logits, axis=-1)
probs = np.exp(log_probs)
np.testing.assert_allclose(probs.sum(axis=-1), 1.0, atol=1e-12)
assert np.all(np.isfinite(log_probs))
selected_ids = np.array([0, 2])
selected = log_probs[np.arange(2), selected_ids]
print("selected log-probabilities:", selected)
print("row sums:", probs.sum(axis=-1))

B, L, V = 2, 7, 9
input_ids = np.array([
    [0, 1, 2, 4, 5, 6, 8],
    [0, 3, 2, 7, 1, 8, 8],
], dtype=np.int64)
response_mask = np.array([
    [0, 0, 0, 1, 1, 1, 1],
    [0, 0, 0, 1, 1, 1, 0],
], dtype=bool)
model_logits = RNG.normal(0.0, 1.2, size=(B, L, V))
# Make the realized response tokens plausible but not deterministic.
for b in range(B):
    for t in range(L - 1):
        model_logits[b, t, input_ids[b, t + 1]] += 1.5

selected_lp, shifted_mask = causal_response_logprobs(model_logits, input_ids, response_mask)
metrics = masked_sequence_metrics(selected_lp, shifted_mask)
print("selected shape:", selected_lp.shape)
print("shifted masks:\n", shifted_mask.astype(int))
print("token counts:", metrics.token_count)
print("sum log-probabilities:", np.round(metrics.sum_logprob, 4))
print("perplexities:", np.round(metrics.perplexity, 4))

fig, ax = plt.subplots(figsize=(9.0, 3.6))
for b in range(B):
    ax.plot(np.arange(L - 1), selected_lp[b], marker="o", label=f"sequence {b}")
    inactive = ~shifted_mask[b]
    ax.scatter(np.arange(L - 1)[inactive], selected_lp[b][inactive], marker="x", s=65)
ax.set_xlabel("shifted token position")
ax.set_ylabel("selected log-probability")
ax.set_title("Causal alignment: circles are scored positions; crosses are masked")
ax.grid(True)
ax.legend()
plt.show()

near_q = np.array([0.40, 0.30, 0.20, 0.10])
near_p = np.array([0.35, 0.30, 0.22, 0.13])
tail_q = np.array([0.999, 0.001])
tail_p = np.array([0.9, 0.1])

near = kl_estimator_moments(near_p, near_q)
tail = kl_estimator_moments(tail_p, tail_q)
print("near-policy moments:", {k: round(v, 6) for k, v in near.items()})
print("heavy-ratio-tail moments:", {k: round(v, 6) for k, v in tail.items()})

# Cross-check the near-policy exact value through logits.
near_exact_from_logits = float(exact_kl_from_logits(np.log(near_q), np.log(near_p)))
assert abs(near_exact_from_logits - near["exact_kl"]) < 1e-12

labels = ["k1", "k2", "k3"]
x = np.arange(3)
width = 0.34
fig, ax = plt.subplots(figsize=(8.5, 4.6))
ax.bar(x - width / 2, [near[f"{k}_std"] for k in labels], width, label="near policies")
ax.bar(x + width / 2, [tail[f"{k}_std"] for k in labels], width, label="rare large p/q")
ax.set_yscale("log")
ax.set_xticks(x, labels)
ax.set_ylabel("exact standard deviation")
ax.set_title("Estimator variance depends on the ratio tail")
ax.grid(True, axis="y", which="both")
ax.legend()
plt.show()

def estimator_samples(p: np.ndarray, q: np.ndarray, n: int, rng: np.random.Generator):
    ids = rng.choice(len(q), size=n, p=q)
    log_r = np.log(p[ids]) - np.log(q[ids])
    return {
        "k1": -log_r,
        "k2": 0.5 * log_r**2,
        "k3": (p[ids] / q[ids] - 1.0) - log_r,
    }

samples = estimator_samples(tail_p, tail_q, N_MONTE_CARLO, RNG)
checkpoints = np.unique(np.geomspace(20, N_MONTE_CARLO, 90).astype(int))
fig, ax = plt.subplots(figsize=(9.0, 4.6))
for name, values in samples.items():
    running = np.cumsum(values)[checkpoints - 1] / checkpoints
    ax.plot(checkpoints, running, label=name)
ax.axhline(tail["exact_kl"], linestyle="--", linewidth=1.4, label="exact KL")
ax.set_xscale("log")
ax.set_xlabel("number of behavior samples")
ax.set_ylabel("running sample mean")
ax.set_title("Finite-sample behavior under a rare ratio tail")
ax.grid(True)
ax.legend()
plt.show()
print("final means:", {name: float(np.mean(values)) for name, values in samples.items()})

order_logits = np.array([1.720, 3.523, 1.987])
order_a = decoding_distribution_with_order(
    order_logits, temperature=1.7, top_k=2, top_p=0.7,
    order=("temperature", "top_k", "top_p"),
)
order_b = decoding_distribution_with_order(
    order_logits, temperature=1.7, top_k=2, top_p=0.7,
    order=("temperature", "top_p", "top_k"),
)
print("temperature → top-k → top-p:", np.round(order_a, 6))
print("temperature → top-p → top-k:", np.round(order_b, 6))
assert np.count_nonzero(order_a) == 1 and np.count_nonzero(order_b) == 2

chosen = 2
print("chosen token logprob under A:", behavior_logprob(order_a, chosen))
print("chosen token logprob under B:", behavior_logprob(order_b, chosen))
ratio = sequence_logprob_ratio(
    np.log(np.array([0.55, 0.42, 0.61])),
    np.log(np.array([0.50, 0.45, 0.30])),
)
print("sequence probability ratio:", ratio)

fig, ax = plt.subplots(figsize=(7.8, 4.0))
positions = np.arange(len(order_logits))
ax.bar(positions - 0.18, order_a, 0.36, label="T→k→p")
ax.bar(positions + 0.18, order_b, 0.36, label="T→p→k")
ax.set_xticks(positions, [f"token {i}" for i in positions])
ax.set_ylabel("behavior probability")
ax.set_title("Processor order can change support")
ax.legend(); ax.grid(True, axis="y")
plt.show()

stop_examples = [
    classify_stop_event("eos", terminal_reasons=("eos",), truncation_reasons=("length", "timeout")),
    classify_stop_event("length", terminal_reasons=("eos",), truncation_reasons=("length", "timeout")),
    classify_stop_event("success_and_timeout", terminal_reasons=("success_and_timeout",), truncation_reasons=("success_and_timeout",)),
    classify_stop_event("server_cancelled", terminal_reasons=("eos",), truncation_reasons=("length", "timeout")),
]
for item in stop_examples:
    print(item)

batch = {
    "response_ids": np.array([[5, 6, 2], [7, 8, 0], [4, 2, 0]]),
    "response_mask": np.array([[1, 1, 1], [1, 1, 0], [1, 1, 0]]),
    "old_logprobs": np.array([[-0.2, -0.4, -0.1], [-0.3, -0.5, np.nan], [-0.6, -0.2, np.nan]]),
    "ref_logprobs": np.array([[-0.25, -0.45, -0.12], [-0.35, -0.55, np.nan], [-0.5, -0.3, np.nan]]),
    "sequence_rewards": np.array([1.0, 0.2, 1.0]),
    "terminated": np.array([True, False, True]),
    "truncated": np.array([False, True, True]),
    "stop_reason": np.array(["eos", "length", "success_and_timeout"], dtype=object),
}
summary = validate_rollout_batch(batch)
print(summary)
assert summary.batch_size == 3 and summary.active_tokens == 7
assert summary.terminated_sequences == 2 and summary.truncated_sequences == 2

def teacher_forced_generated_logprobs(model, sequences, prompt_length: int):
    import torch
    with torch.no_grad():
        logits = model(sequences).logits[:, :-1, :]
        log_probs = torch.log_softmax(logits.float(), dim=-1)
        targets = sequences[:, 1:]
        selected = log_probs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
    # generated token at absolute position prompt_length is predicted by logits prompt_length-1
    return selected[:, prompt_length - 1:]


def inferred_transformers_stop_reason(generated_ids, eos_ids, max_new_tokens: int):
    eos_set = {int(x) for x in eos_ids}
    if generated_ids and int(generated_ids[-1]) in eos_set:
        return {"stop_reason": "eos_inferred", "terminated": True, "truncated": False, "inferred": True}
    if len(generated_ids) >= max_new_tokens:
        return {"stop_reason": "length_inferred", "terminated": False, "truncated": True, "inferred": True}
    return {"stop_reason": "unknown_inferred", "terminated": False, "truncated": False, "inferred": True}

print("Real-model helpers are defined; no model has been downloaded or executed yet.")

REAL_MODEL_RESULT = None
if RUN_REAL_MODEL:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    model_id = "Qwen/Qwen3-0.6B"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype="auto",
        device_map="auto",
    )
    prompt = "Explain in one sentence why a sampled token log-ratio may be negative while KL is non-negative."
    encoded = tokenizer(prompt, return_tensors="pt").to(model.device)
    max_new_tokens = 32
    with torch.no_grad():
        generated = model.generate(
            **encoded,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.8,
            top_p=0.9,
            return_dict_in_generate=True,
            output_scores=True,
        )
    transition = model.compute_transition_scores(
        generated.sequences,
        generated.scores,
        getattr(generated, "beam_indices", None),
        normalize_logits=True,
    )
    prompt_length = encoded["input_ids"].shape[1]
    teacher = teacher_forced_generated_logprobs(model, generated.sequences, prompt_length)
    generated_count = transition.shape[1]
    teacher = teacher[:, :generated_count]
    max_delta = float(torch.max(torch.abs(transition.float() - teacher.float())).cpu())
    new_ids = generated.sequences[0, prompt_length:].tolist()
    eos_ids = tokenizer.eos_token_id
    eos_ids = [eos_ids] if isinstance(eos_ids, int) else list(eos_ids or [])
    stop = inferred_transformers_stop_reason(new_ids, eos_ids, max_new_tokens)
    REAL_MODEL_RESULT = {
        "model_id": model_id,
        "generated_tokens": len(new_ids),
        "max_transition_teacher_delta": max_delta,
        **stop,
        "text": tokenizer.decode(generated.sequences[0], skip_special_tokens=True),
    }
    print(REAL_MODEL_RESULT)
else:
    print("Optional real-model branch skipped. Set RL_LLM_M6_RUN_REAL_MODEL=1 to run it explicitly.")

assert np.allclose(np.exp(log_probs).sum(axis=-1), 1.0)
assert np.isclose(near["k1_mean"], near["exact_kl"], atol=1e-14)
assert np.isclose(near["k3_mean"], near["exact_kl"], atol=1e-14)
assert tail["k3_std"] > 10 * tail["k1_std"]
assert np.count_nonzero(order_a) != np.count_nonzero(order_b)
assert stop_examples[-1].terminated is False and stop_examples[-1].truncated is False

COURSE_RESULT = {
    "synthetic_only": not RUN_REAL_MODEL,
    "causal_batch_perplexity": np.asarray(metrics.perplexity).round(6).tolist(),
    "near_exact_kl": near["exact_kl"],
    "heavy_tail_exact_kl": tail["exact_kl"],
    "heavy_tail_k3_std": tail["k3_std"],
    "rollout_active_tokens": summary.active_tokens,
    "real_model_result": REAL_MODEL_RESULT,
}
print(COURSE_RESULT)