# benchmark_chunked_nll_chunk_size.py
"""
Micro-benchmark `_chunked_cross_entropy_loss` across chunk sizes.

Runs forward + backward on synthetic `(hidden, weight, labels)` tensors and reports peak GPU
memory + wall time. The un-chunked baseline (full `[N_valid, V]` logits) is also timed as an
upper bound.
"""

import argparse
import json
import time

import torch
import torch.nn.functional as F

from trl.trainer.sft_trainer import _chunked_cross_entropy_loss


VOCAB_SIZE = 151936  # Qwen2.5 / Qwen3 family
CHUNK_SIZES = [64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
ITERS = 10
DTYPE = torch.bfloat16


def make_inputs(n_valid, hidden_size):
    # Shift means `n_valid + 1` tokens yield exactly `n_valid` valid (non-ignored) positions.
    hidden = torch.randn(1, n_valid + 1, hidden_size, device="cuda", dtype=DTYPE, requires_grad=True)
    weight = torch.randn(VOCAB_SIZE, hidden_size, device="cuda", dtype=DTYPE, requires_grad=True)
    labels = torch.randint(0, VOCAB_SIZE, (1, n_valid + 1), device="cuda")
    return hidden, weight, labels


def bench(fn, n_valid, hidden_size):
    hidden, weight, labels = make_inputs(n_valid, hidden_size)
    try:
        fn(hidden, weight, labels).backward()  # warmup
        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()
        start = time.perf_counter()
        for _ in range(ITERS):
            hidden.grad = weight.grad = None
            fn(hidden, weight, labels).backward()
        torch.cuda.synchronize()
        return torch.cuda.max_memory_allocated() / 1e6, (time.perf_counter() - start) * 1000 / ITERS
    except torch.cuda.OutOfMemoryError:
        return None, None
    finally:
        del hidden, weight, labels
        torch.cuda.empty_cache()


def reference_loss(h, w, lbl):
    """What the standard nll path does: materialize full logits."""
    shift_h = h[..., :-1, :].reshape(-1, h.size(-1))
    shift_l = lbl[..., 1:].reshape(-1)
    return F.cross_entropy(shift_h.float() @ w.float().t(), shift_l, ignore_index=-100)


def chunked_loss(chunk_size):
    return lambda h, w, lbl: _chunked_cross_entropy_loss(h, w, lbl, chunk_size=chunk_size)[0]


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--hidden_size", type=int, required=True)
    parser.add_argument("--n_valid", type=int, required=True)
    parser.add_argument("--output", type=str, default=None)
    args = parser.parse_args()

    torch.manual_seed(0)
    print(f"N_valid={args.n_valid}, vocab={VOCAB_SIZE}, hidden={args.hidden_size}, dtype=bfloat16, iters={ITERS}")
    print(f"{'chunk_size':>12} {'peak_MB':>10} {'time_ms':>10}")

    ref = None
    peak, ms = bench(reference_loss, args.n_valid, args.hidden_size)
    if peak is None:
        print(f"{'reference':>12} {'OOM':>10} {'—':>10}")
    else:
        print(f"{'reference':>12} {peak:>10.1f} {ms:>10.2f}")
        ref = {"peak_mb": peak, "time_ms": ms}

    results = []
    for cs in CHUNK_SIZES:
        if cs > args.n_valid:
            continue
        peak, ms = bench(chunked_loss(cs), args.n_valid, args.hidden_size)
        if peak is None:
            print(f"{cs:>12} {'OOM':>10} {'—':>10}")
            results.append({"chunk_size": cs, "oom": True})
        else:
            print(f"{cs:>12} {peak:>10.1f} {ms:>10.2f}")
            results.append({"chunk_size": cs, "peak_mb": peak, "time_ms": ms})

    if args.output:
        shape = {"vocab_size": VOCAB_SIZE, "hidden_size": args.hidden_size, "n_valid": args.n_valid, "dtype": "bfloat16"}
        payload = {"shape": shape, "iters": ITERS, "reference": ref, "results": results}
        with open(args.output, "w") as f:
            json.dump(payload, f, indent=2)
        print(f"\nwrote {args.output}")


if __name__ == "__main__":
    main()

# /// script
# dependencies = ["matplotlib", "numpy"]
# ///
import argparse
import json
import pathlib

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.lines import Line2D


def find_knee(chunk_sizes, time_ms, tol=0.50):
    """Smallest chunk_size whose time is within `tol` of the minimum (memory-priority)."""
    threshold = time_ms.min() * (1 + tol)
    for cs, t in zip(chunk_sizes, time_ms):
        if t <= threshold:
            return int(cs), float(t)
    return int(chunk_sizes[-1]), float(time_ms[-1])


def describe(shape, all_shapes):
    """Human-readable label with 'small/medium/big model' and 'long sequence' tags."""
    h, n = shape["hidden_size"], shape["n_valid"]
    hs = {s["hidden_size"] for s in all_shapes}
    ns = {s["n_valid"] for s in all_shapes}
    tags = []
    if len(hs) > 1:
        tags.append({max(hs): "big model", min(hs): "small model"}.get(h, "medium model"))
    if len(ns) > 1 and n == max(ns):
        tags.append("long sequence")
    return f"H={h}, N_valid={n}" + (f"  ({', '.join(tags)})" if tags else "")


def load_runs(dir_path):
    datas = []
    for p in sorted(pathlib.Path(dir_path).glob("*.json")):
        data = json.loads(p.read_text())
        rows = [r for r in data["results"] if not r.get("oom")]
        if rows:
            datas.append((data, rows))
    datas.sort(key=lambda d: (d[0]["shape"]["hidden_size"], d[0]["shape"]["n_valid"]))
    shapes = [d[0]["shape"] for d in datas]
    return [
        {
            "label": describe(data["shape"], shapes),
            "reference": data.get("reference"),
            "chunk_sizes": np.array([r["chunk_size"] for r in rows]),
            "peak_gb": np.array([r["peak_mb"] for r in rows]) / 1024,
            "time_ms": np.array([r["time_ms"] for r in rows]),
        }
        for data, rows in datas
    ]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dir", default="/tmp/chunk_size_bench")
    parser.add_argument("--out", default="chunk_size_pareto.png")
    args = parser.parse_args()

    runs = load_runs(args.dir)
    if not runs:
        raise SystemExit(f"No JSON files found in {args.dir}")

    colors = [plt.get_cmap("tab10")(i) for i in range(len(runs))]
    fig, ax = plt.subplots(figsize=(9.5, 7))
    fig.suptitle("Chunked CE: memory ↔ time trade-off", fontsize=14)

    for r, c in zip(runs, colors):
        ax.plot(r["time_ms"], r["peak_gb"], "-o", color=c, markersize=6, linewidth=2, label=r["label"])
        for cs, mem, t in zip(r["chunk_sizes"], r["peak_gb"], r["time_ms"]):
            ax.annotate(str(int(cs)), (t, mem), textcoords="offset points", xytext=(6, 4), fontsize=8, color=c)

        knee_cs, knee_t = find_knee(r["chunk_sizes"], r["time_ms"])
        knee_mem = r["peak_gb"][r["chunk_sizes"] == knee_cs][0]
        ax.scatter([knee_t], [knee_mem], color=c, marker="*", s=320, edgecolor="black", linewidth=1.2, zorder=6)

        if r["reference"] is not None:
            ref_t, ref_mem = r["reference"]["time_ms"], r["reference"]["peak_mb"] / 1024
            ax.scatter([ref_t], [ref_mem], color=c, marker="s", s=100, edgecolor="black", linewidth=1, zorder=5)
            ax.annotate("ref", (ref_t, ref_mem), textcoords="offset points", xytext=(8, -3),
                        fontsize=9, fontweight="bold", color=c)

    ax.set(xscale="log", yscale="log", xlabel="time per step (ms)", ylabel="peak memory (GB)")
    ax.set_yscale("log", base=2)
    ax.grid(True, which="both", alpha=0.3)

    shape_handles = [Line2D([0], [0], marker="o", linestyle="-", color=c, label=r["label"])
                     for r, c in zip(runs, colors)]
    glyph_handles = [
        Line2D([0], [0], marker="*", color="gray", linestyle="", markeredgecolor="black", markersize=14,
               label="knee (within 50% of min time)"),
        Line2D([0], [0], marker="s", color="gray", linestyle="", markeredgecolor="black", markersize=9,
               label="un-chunked reference"),
    ]
    ax.add_artist(ax.legend(handles=shape_handles, loc="upper left", fontsize=10, title="shape"))
    ax.legend(handles=glyph_handles, loc="upper right", fontsize=10, title="markers")

    print("\nknee per shape:")
    for r in runs:
        knee_cs, knee_t = find_knee(r["chunk_sizes"], r["time_ms"])
        knee_mem = r["peak_gb"][r["chunk_sizes"] == knee_cs][0]
        print(f"  {r['label']}  →  chunk_size = {knee_cs}  ({knee_mem:.2f} GB, {knee_t:.1f} ms)")

    plt.tight_layout(rect=[0, 0, 1, 0.96])
    plt.savefig(args.out, dpi=140, bbox_inches="tight")
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()