Модуль 12. Инфраструктура RL для LLM: память, роллауты и асинхронность

from __future__ import annotations

import os
import random
import sys
import time
from pathlib import Path

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


def find_repo_root(start: Path) -> Path:
    """Find the course root regardless of the notebook working directory."""
    start = start.resolve()
    for candidate in (start, *start.parents):
        if (candidate / "modules/module-12/module_12_reference.py").is_file():
            return candidate
    raise FileNotFoundError("Could not locate the RL for LLM repository root")


ROOT = find_repo_root(Path.cwd())
MODULE_DIR = ROOT / "modules/module-12"
if str(MODULE_DIR) not in sys.path:
    sys.path.insert(0, str(MODULE_DIR))

from module_12_reference import (
    BYTES_PER_GIB,
    checkpoint_cadence,
    completion_token_mask,
    effective_sample_size,
    human_bytes,
    ideal_sharded_memory,
    kv_cache_bytes,
    model_state_bytes,
    optimal_checkpoint_interval,
    peak_memory,
    pipeline_report,
    rollout_contract_errors,
    straggler_report,
    synchronization_schedule,
    synchronization_transfer_time,
)
from module_12_pipeline_sim import (
    PipelineConfig,
    checkpoint_as_dict,
    make_checkpoint,
    recovery_drill,
    result_row,
    simulate_pipeline,
)

PROFILE = os.getenv("RL_LLM_PROFILE", "smoke").strip().lower()
RUN_REAL_MODEL = False
if PROFILE not in {"smoke", "full", "real-model"}:
    raise ValueError("RL_LLM_PROFILE must be smoke, full, or real-model")

DURATION = 240.0 if PROFILE == "full" else 120.0
SEED = 12016
np.random.seed(SEED)
random.seed(SEED)


print(f"Корень курса: {ROOT}")
print(f"Профиль: {PROFILE}; виртуальная длительность={DURATION:.0f} c; seed={SEED}")

profiles = {
    "7B / 12-byte full-training profile": model_state_bytes(
        7_000_000_000,
        parameter_bytes=2,
        gradient_bytes=2,
        optimizer_state_bytes=8,
        master_parameter_bytes=0,
    ),
    "7B / 16-byte full-training profile": model_state_bytes(
        7_000_000_000,
        parameter_bytes=2,
        gradient_bytes=2,
        optimizer_state_bytes=8,
        master_parameter_bytes=4,
    ),
    "7B base + 1% trainable adapters": model_state_bytes(
        7_070_000_000,
        trainable_parameters=70_000_000,
        parameter_bytes=2,
        gradient_bytes=2,
        optimizer_state_bytes=8,
        master_parameter_bytes=4,
    ),
}

memory_rows = []
for name, profile in profiles.items():
    memory_rows.append({
        "profile": name,
        "logical bytes": profile.total_bytes,
        "GiB": profile.total_bytes / BYTES_PER_GIB,
        "human": human_bytes(profile.total_bytes),
    })

display(pd.DataFrame(memory_rows))

kv_16_gib = kv_cache_bytes(
    num_sequences=64,
    tokens_per_sequence=4096,
    num_layers=32,
    num_kv_heads=4,
    head_dim=128,
    element_bytes=2,
)
print("KV-cache teaching configuration:", human_bytes(kv_16_gib))
assert kv_16_gib == 16 * BYTES_PER_GIB

phase_report = peak_memory({
    "rollout": [14 * BYTES_PER_GIB, kv_16_gib, 2 * BYTES_PER_GIB],
    "training": [14 * BYTES_PER_GIB, 8 * BYTES_PER_GIB, 28 * BYTES_PER_GIB],
    "weight publication": [14 * BYTES_PER_GIB, 14 * BYTES_PER_GIB, 2 * BYTES_PER_GIB],
})
print("Peak phase:", phase_report.phase, human_bytes(phase_report.peak_bytes))
display(pd.DataFrame([
    {"phase": name, "GiB": value / BYTES_PER_GIB}
    for name, value in phase_report.phase_totals.items()
]))

pipe = pipeline_report(rollout_seconds=60.0, train_seconds=20.0, batches=10)
stragglers = straggler_report([3, 3, 4, 4, 5, 5, 6, 18])
print("Two-stage pipeline:", pipe)
print("Synchronous straggler group:", stragglers)
assert np.isclose(pipe.sequential_total, 800.0)
assert np.isclose(pipe.pipeline_total, 620.0)
assert np.isclose(stragglers.utilization, 1 / 3)

profile_configs = {
    "sequential": PipelineConfig(
        mode="sequential",
        duration_seconds=DURATION,
        seed=SEED,
    ),
    "bounded async / sync every update": PipelineConfig(
        mode="bounded_async",
        duration_seconds=DURATION,
        max_queue_tokens=7000,
        max_policy_lag=2,
        sync_every_updates=1,
        seed=SEED,
    ),
    "deeper async / sync every 3 updates": PipelineConfig(
        mode="bounded_async",
        duration_seconds=DURATION,
        max_queue_tokens=7000,
        max_policy_lag=6,
        sync_every_updates=3,
        seed=SEED,
    ),
}
results = {name: simulate_pipeline(cfg) for name, cfg in profile_configs.items()}
summary = pd.DataFrame([
    {"profile": name, **result_row(result)}
    for name, result in results.items()
])
display(summary)




# Profile rankings are observations, not invariants under parameter changes.
assert all(np.isfinite(r.trained_tokens_per_second) and r.trained_tokens_per_second >= 0 for r in results.values())

fig, ax = plt.subplots(figsize=(9.5, 4.8))
for name, result in results.items():
    if result.config.mode != "bounded_async":
        continue
    x = [row["time"] for row in result.timeline]
    y = [row["queue_tokens"] for row in result.timeline]
    ax.plot(x, y, label=name)
ax.axhline(7000, linestyle="--", linewidth=1.5, label="new-rollout admission threshold")
ax.set_xlabel("time, s")
ax.set_ylabel("tokens in queue")
ax.set_title("A bounded queue converts overload into observable backpressure")
ax.grid(alpha=0.25)
ax.legend()
plt.show()

queue_table = pd.DataFrame([
    {
        "profile": name,
        "maximum observed queue, tokens": result.max_queue_tokens_observed,
        "backpressure, s": result.backpressure_seconds,
        "discarded rollouts": result.discarded_rollouts,
    }
    for name, result in results.items()
])
display(queue_table)

batch_rows = []
for name, result in results.items():
    for batch in result.batches:
        batch_rows.append({
            "profile": name,
            "update": batch.update,
            "mean lag": batch.mean_lag,
            "max lag": batch.max_lag,
            "normalized ESS": batch.normalized_ess,
        })
batch_df = pd.DataFrame(batch_rows)
display(batch_df.groupby("profile").agg(
    updates=("update", "count"),
    mean_lag=("mean lag", "mean"),
    p95_lag=("mean lag", lambda x: x.quantile(0.95)),
    mean_nESS=("normalized ESS", "mean"),
    min_nESS=("normalized ESS", "min"),
))

fig, ax = plt.subplots(figsize=(8.6, 4.6))
for name, part in batch_df.groupby("profile"):
    if name == "sequential":
        continue
    ax.scatter(part["mean lag"], part["normalized ESS"], s=20, alpha=0.65, label=name)
ax.set_xlabel("mean policy lag in a batch, versions")
ax.set_ylabel("normalized ESS")
ax.set_ylim(0, 1.03)
ax.set_title("The simulator exposes a speed–freshness trade-off, not a theorem")
ax.grid(alpha=0.25)
ax.legend()
plt.show()

good_record = {
    "sample_id": "run12-p003-r01",
    "prompt_id": "p003",
    "policy_version": 17,
    "prompt_tokens": [101, 42, 17],
    "response_tokens": [314, 159, 2],
    "behavior_logprobs": [-0.31, -1.07, -0.22],
    "decoding": {"temperature": 0.8, "top_p": 1.0},
    "reward": {"correctness": 1.0, "format": 0.0},
    "stop_reason": "eos",
    "verifier_version": "arith-v4",
}
broken_record = dict(good_record)
broken_record.pop("behavior_logprobs")
broken_record["policy_version"] = -1

print("Valid record errors:", rollout_contract_errors(good_record))
print("Broken record errors:", rollout_contract_errors(broken_record))
assert rollout_contract_errors(good_record) == []
assert rollout_contract_errors(broken_record)

full_bf16_payload = 7_000_000_000 * 2
adapter_payload = 70_000_000 * 2
rows = []
for name, payload in (("full 7B bf16 weights", full_bf16_payload), ("1% bf16 adapter", adapter_payload)):
    rows.append({
        "payload": name,
        "GiB": payload / BYTES_PER_GIB,
        "transfer over 25 GB/s at 80% efficiency, s": synchronization_transfer_time(
            payload,
            25.0,
            setup_seconds=0.15,
            efficiency=0.80,
        ),
    })
display(pd.DataFrame(rows))

cadences = []
for every in (1, 2, 4, 8):
    schedule = synchronization_schedule(sync_every=every, updates=16)
    cadences.append({
        "sync every K updates": every,
        "synchronizations": schedule.synchronization_count,
        "mean lag": schedule.mean_lag,
        "max lag": schedule.max_lag,
    })
display(pd.DataFrame(cadences))

checkpoint_seconds = 180.0
mtbf_seconds = 10 * 3600.0
interval = optimal_checkpoint_interval(checkpoint_seconds, mtbf_seconds)
cadence = checkpoint_cadence(
    interval,
    checkpoint_seconds=checkpoint_seconds,
    mtbf_seconds=mtbf_seconds,
)
print("First-order Young interval, seconds:", interval)
print("Checkpoint overhead model:", cadence)

rng = np.random.default_rng(SEED)
_ = rng.integers(0, 10_000, size=8)
manifest = make_checkpoint(
    optimizer_step=3,
    current_version=3,
    published_version=3,
    dataset_cursor=8,
    rng=rng,
)
display(pd.DataFrame([checkpoint_as_dict(manifest)]).drop(columns=["rng_state"]))

drill = recovery_drill(seed=SEED, before=8, after=8)
print("Uninterrupted:", drill.uninterrupted)
print("Correct resume:", drill.resumed, "match=", drill.exact_match)
print("Weights-only resume:", drill.broken_resume, "match=", drill.broken_match)
assert drill.exact_match and not drill.broken_match

if not RUN_REAL_MODEL:
    print("Optional real-model branch disabled; review the environment and set RUN_REAL_MODEL=True.")
else:
    import hashlib
    import importlib.metadata
    import re
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    model_id = "Qwen/Qwen3-0.6B"
    revision = os.environ.get("RL_LLM_MODEL_REVISION", "")
    if not re.fullmatch(r"[0-9a-f]{40}", revision):
        raise ValueError("Set RL_LLM_MODEL_REVISION to the reviewed model's immutable 40-character commit SHA")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dtype = (torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16) if device.type == "cuda" else torch.float32
    tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision, padding_side="left")
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(model_id, revision=revision, torch_dtype=dtype).to(device)
    model.eval()
    prompts = [
        "Explain why an unbounded rollout queue is dangerous.",
        "Give one invariant for atomic rollout weight publication.",
        "Distinguish policy lag from training-inference mismatch.",
    ]
    formatted = [tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}], tokenize=False,
        add_generation_prompt=True, enable_thinking=False,
    ) for prompt in prompts]
    encoded = tokenizer(formatted, add_special_tokens=False, return_tensors="pt", padding=True).to(device)
    generation_config = dict(max_new_tokens=64, do_sample=True, temperature=0.8,
                             top_p=1.0, top_k=0, pad_token_id=tokenizer.pad_token_id)

    def synchronize():
        if device.type == "cuda":
            torch.cuda.synchronize()

    # Download/loading and this one warm-up forward are outside both timings.
    with torch.inference_mode():
        model(**encoded, use_cache=False)
    torch.manual_seed(SEED)
    synchronize()
    started = time.perf_counter()
    with torch.inference_mode():
        generated = model.generate(**encoded, **generation_config)
    synchronize()
    generation_seconds = time.perf_counter() - started
    prompt_width = encoded["input_ids"].shape[1]
    suffix = generated[:, prompt_width:]
    eos_ids = model.generation_config.eos_token_id
    eos_ids = [] if eos_ids is None else ([eos_ids] if isinstance(eos_ids, int) else eos_ids)
    suffix_mask = completion_token_mask(suffix.cpu().numpy(), eos_ids)
    lengths = suffix_mask.sum(axis=1).tolist()
    attention_mask = torch.cat((encoded["attention_mask"], torch.as_tensor(
        suffix_mask, dtype=encoded["attention_mask"].dtype, device=device)), dim=1)
    position_ids = attention_mask.long().cumsum(-1) - 1
    position_ids.masked_fill_(attention_mask == 0, 0)
    # Teacher forcing scores the fixed generated sequences, with prompt PAD and
    # suffix PAD masked. Explicit positions preserve left-padding semantics.
    synchronize()
    started = time.perf_counter()
    with torch.inference_mode():
        scored = model(input_ids=generated, attention_mask=attention_mask,
                       position_ids=position_ids, use_cache=False)
    synchronize()
    forward_seconds = time.perf_counter() - started
    assert scored.logits.shape[:2] == generated.shape
    print({
        "model": model_id, "revision": revision, "dtype": str(dtype), "seed": SEED,
        "device": str(device), "padding_side": tokenizer.padding_side,
        "template_sha256": hashlib.sha256(tokenizer.chat_template.encode()).hexdigest(),
        "enable_thinking": False, "generation_config": generation_config,
        "versions": {name: importlib.metadata.version(name) for name in ("torch", "transformers")},
        "generation_seconds": generation_seconds, "teacher_forced_forward_seconds": forward_seconds,
        "new_tokens_including_eos": lengths, "active_generated_tokens": sum(lengths),
        "generated_tokens_per_second": sum(lengths) / generation_seconds,
        "active_forward_tokens": int(attention_mask.sum()),
        "forward_tokens_per_second": int(attention_mask.sum()) / forward_seconds,
    })
    print("One local batch; no backward pass or optimizer. This is not a cluster benchmark.")