Evaluate every checkpoint

git clone https://github.com/unionai/unionai-examples
cd unionai-examples/v2/tutorials/checkpoint_evals
python train.py

if __name__ == "__main__":
    import evals

    flyte.init_from_config()
    if "--leaderboard" in sys.argv:
        run = flyte.run(refresh_leaderboard, train_run=sys.argv[sys.argv.index("--leaderboard") + 1])
    else:
        # The evaluator must be deployed, with its trigger active, before the first checkpoint lands.
        flyte.deploy(evals.eval_env)
        if "--smoke" in sys.argv:
            run = flyte.run(train_and_watch, profile="smoke", train_examples=1_600, n_checkpoints=2, wait_minutes=20)
        else:
            run = flyte.run(train_and_watch)
    print(run.url)

BASE_MODEL = "HuggingFaceTB/SmolLM2-360M"

# The trainer publishes every checkpoint as a new version of this artifact, partitioned by
# training run, step and profile ("full" or "smoke").
CHECKPOINT_ARTIFACT = "math-sft-checkpoint"
# Every eval run publishes its scores as a new version of this artifact, same partitions.
RESULTS_ARTIFACT = "math-sft-checkpoint-evals"
# The checkpoint that wins the ship rule is promoted to this artifact.
RELEASE_ARTIFACT = "math-sft-release-candidate"

# One entry per benchmark slice. Bump a benchmark's version to re-score only that benchmark on
# every checkpoint; the other benchmarks stay cached. `limit` is the full-profile sample size.
BENCHMARKS = {
    "gsm8k": {"kind": "generate", "version": "1", "limit": 0, "label": "GSM8K (exact match)", "role": "target"},
    "arc_easy": {"kind": "choice", "version": "1", "limit": 0, "label": "ARC-Easy (MC)", "role": "reasoning"},
    "mmlu": {"kind": "choice", "version": "1", "limit": 2000, "label": "MMLU (MC)", "role": "guard"},
}
DEFAULT_SUITE = ["gsm8k", "arc_easy", "mmlu"]
SMOKE_LIMIT = 40

# The ship rule: the best GSM8K checkpoint whose MMLU score is at most this many points below
# the base model's.
GUARD_BENCHMARK = "mmlu"
GUARD_MAX_DROP = 2.0

train_env = flyte.TaskEnvironment(
    name="ckpt-evals-train",
    image=gpu_image,
    resources=flyte.Resources(cpu=6, memory="48Gi", gpu="L40s:1", disk="80Gi"),
    env_vars={"TOKENIZERS_PARALLELISM": "false"},
)

driver_env = flyte.TaskEnvironment(
    name="ckpt-evals-driver",
    image=driver_image,
    resources=flyte.Resources(cpu=1, memory="2Gi"),
    depends_on=[train_env],
)

@train_env.task(report=True)
async def finetune(
    train_run: str,
    profile: str = "full",
    train_examples: int = 50_000,
    n_checkpoints: int = 8,
    learning_rate: float = 3e-5,
    micro_batch: int = 8,
    grad_accum: int = 4,
    max_len: int = 512,
    seed: int = 0,
) -> str:
    """Full fine-tune on one GPU; save and publish checkpoints without pausing for uploads."""
    import math
    import os
    import random

    import torch
    from safetensors.torch import save_file
    from transformers import AutoModelForCausalLM, AutoTokenizer

    torch.manual_seed(seed)
    tok = AutoTokenizer.from_pretrained(BASE_MODEL)
    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, dtype=torch.float32).cuda()
    data = load_sft_data(tok, train_examples, max_len, seed)
    total = len(data) // (micro_batch * grad_accum)
    save_at = checkpoint_steps(total, n_checkpoints)
    opt = torch.optim.AdamW(model.parameters(), lr=learning_rate, betas=(0.9, 0.95), weight_decay=0.0)
    warmup = max(1, total // 30)

    def lr_at(s: int) -> float:  # linear warmup, cosine decay to 10%
        if s < warmup:
            return learning_rate * (s + 1) / warmup
        return learning_rate * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * (s - warmup) / max(1, total - warmup))))

    loop = asyncio.get_running_loop()
    uploads: list = []
    history: list[dict] = []
    published: list[dict] = []
    started = time.time()

    def save(step: int, loss: float | None) -> None:
        # Saving takes seconds and runs on the training thread. Uploading and publishing run on
        # the event loop, so the next optimizer step starts right away.
        path = f"/tmp/ckpt/step-{step:05d}"
        os.makedirs(path, exist_ok=True)
        sd = {k: v.detach().to(torch.bfloat16).cpu().contiguous() for k, v in model.state_dict().items()
              if not (model.config.tie_word_embeddings and k == "lm_head.weight")}
        save_file(sd, f"{path}/model.safetensors", metadata={"format": "pt"})
        model.config.save_pretrained(path)
        model.generation_config.save_pretrained(path)
        tok.save_pretrained(path)
        published.append({"step": step, "loss": loss, "saved_s": round(time.time() - started, 1)})
        uploads.append(asyncio.run_coroutine_threadsafe(
            publish_checkpoint(path, step, train_run, profile, loss), loop))

    def train() -> None:
        model.train()
        order = list(range(len(data)))
        random.Random(seed).shuffle(order)
        save(0, None)
        for step in range(1, total + 1):
            for g in opt.param_groups:
                g["lr"] = lr_at(step - 1)
            step_loss = 0.0
            for m in range(grad_accum):
                lo = ((step - 1) * grad_accum + m) * micro_batch
                batch = [data[i] for i in order[lo:lo + micro_batch]]
                width = max(len(ids) for ids, _ in batch)
                x = torch.full((len(batch), width), tok.pad_token_id or tok.eos_token_id)
                y = torch.full((len(batch), width), -100)
                mask = torch.zeros((len(batch), width), dtype=torch.long)
                for r, (ids, plen) in enumerate(batch):
                    x[r, :len(ids)] = torch.tensor(ids)
                    mask[r, :len(ids)] = 1
                    y[r, plen:len(ids)] = torch.tensor(ids[plen:])  # loss on the answer only
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    loss = model(input_ids=x.cuda(), attention_mask=mask.cuda(), labels=y.cuda()).loss / grad_accum
                loss.backward()
                step_loss += loss.item()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            opt.zero_grad(set_to_none=True)
            if step % 10 == 0 or step == total:
                history.append({"step": step, "loss": round(step_loss, 4), "lr": lr_at(step - 1),
                                "elapsed_s": round(time.time() - started, 1)})
            if step in save_at:
                save(step, step_loss)
            if step % 25 == 0 or step in save_at:
                flyte.report.replace(training_html(history, published, total, train_examples, BASE_MODEL))
                flyte.report.flush()

    await asyncio.to_thread(train)
    versions = await asyncio.gather(*[asyncio.wrap_future(u) for u in uploads])
    flyte.report.replace(training_html(history, published, total, train_examples, BASE_MODEL))
    flyte.report.flush()
    return json.dumps({"steps": save_at, "versions": versions, "history": history, "published": published,
                       "total_steps": total, "train_s": round(time.time() - started, 1)})

async def publish_checkpoint(path: str, step: int, train_run: str, profile: str, loss: float | None) -> str:
    """Upload a saved checkpoint and register it as a new version of the checkpoint artifact.

    The evaluator's trigger fires on this call; training does not wait for it.
    """
    from flyte.remote import Artifact

    ckpt = await Dir.from_local(path)
    art = await Artifact.create.aio(
        ckpt,
        name=CHECKPOINT_ARTIFACT,
        version=f"{train_run}-step{step:05d}",
        kind="model",
        description=f"{BASE_MODEL} math SFT, step {step} of {train_run}",
        partitions={"train_run": train_run, "step": str(step), "profile": profile},
        attrs={"base_model": BASE_MODEL, "train_loss": f"{loss:.4f}" if loss is not None else "n/a"},
    )
    return art.version

# Generative evals: one L4 per (checkpoint, benchmark). vLLM runs in the task's process, and
# uses PyTorch's sampler instead of FlashInfer's, which would JIT-compile with nvcc at startup.
generate_env = flyte.TaskEnvironment(
    name="ckpt-evals-generate",
    image=gpu_image,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu="L4:1", disk="40Gi"),
    env_vars={"VLLM_ENABLE_V1_MULTIPROCESSING": "0", "VLLM_USE_FLASHINFER_SAMPLER": "0",
              "TOKENIZERS_PARALLELISM": "false"},
)

# Multiple-choice evals score every answer option's log-likelihood, one L4 per benchmark.
choice_env = flyte.TaskEnvironment(
    name="ckpt-evals-choice",
    image=gpu_image,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu="L4:1", disk="40Gi"),
    env_vars={"TOKENIZERS_PARALLELISM": "false"},
)

eval_env = flyte.TaskEnvironment(
    name="ckpt-evals",
    image=driver_image,
    resources=flyte.Resources(cpu=1, memory="2Gi"),
    depends_on=[generate_env, choice_env],
)

# The cache key is the task inputs (checkpoint, benchmark, benchmark version, limit) under one
# explicit cache version, so editing a report does not throw away finished evals.
EVAL_CACHE = flyte.Cache("override", version_override="checkpoint-evals-1")

on_checkpoint = flyte.Trigger(
    name="eval_new_checkpoint",
    description="Run the eval suite on every new version of the math SFT checkpoint",
    automation=flyte.OnArtifact(name=CHECKPOINT_ARTIFACT),
    inputs={
        "checkpoint": flyte.TriggeredArtifact,
        "step": flyte.TriggeredPartition("step"),
        "train_run": flyte.TriggeredPartition("train_run"),
        "profile": flyte.TriggeredPartition("profile"),
    },
)

@eval_env.task(triggers=[on_checkpoint], report=True, produces_artifacts=True)
async def evaluate_checkpoint(
    checkpoint: Dir, step: str, train_run: str, profile: str = "full", suite: list[str] = DEFAULT_SUITE
) -> File:
    """Fan the suite out over one checkpoint and publish the scores as an artifact version."""
    started = time.time()
    calls = []
    for name in suite:
        spec = BENCHMARKS[name]
        limit = SMOKE_LIMIT if profile == "smoke" else spec["limit"]
        task = run_generate if spec["kind"] == "generate" else run_choice
        calls.append(task.aio(checkpoint=checkpoint, benchmark=name, version=spec["version"], limit=limit))
    scores = {s["benchmark"]: s for s in map(json.loads, await asyncio.gather(*calls))}

    result = {"train_run": train_run, "step": int(step), "profile": profile, "checkpoint": checkpoint.path,
              "eval_run": flyte.ctx().action.run_name, "started_at": started, "finished_at": time.time(),
              "scores": scores}
    await flyte.report.replace.aio(scorecard_html(result), do_flush=True)

    path = "/tmp/scores.json"
    with open(path, "w") as f:
        json.dump(result, f)
    summary = ", ".join(f"{k} {v['accuracy']:.1f}" for k, v in scores.items())
    return artifacts.new(await File.from_local(path), artifacts.Metadata(
        name=RESULTS_ARTIFACT,
        description=f"step {int(step)} of {train_run}: {summary}",
        partitions={"train_run": train_run, "step": step, "profile": profile},
        attrs={k: f"{v['accuracy']:.2f}" for k, v in scores.items()},
    ))

@generate_env.task(cache=EVAL_CACHE, retries=1)
def run_generate(checkpoint: Dir, benchmark: str, version: str, limit: int) -> str:
    """Score one checkpoint on one generative benchmark."""
    model_dir = checkpoint.download_sync()
    return json.dumps(score_generate(model_dir, load_items(benchmark, limit, version), version))

@choice_env.task(cache=EVAL_CACHE, retries=1)
def run_choice(checkpoint: Dir, benchmark: str, version: str, limit: int) -> str:
    """Score one checkpoint on one multiple-choice benchmark."""
    model_dir = checkpoint.download_sync()
    return json.dumps(score_choice(model_dir, load_items(benchmark, limit, version), benchmark, version))

def load_items(benchmark: str, limit: int, version: str = "1") -> list[dict]:
    """A fixed, seeded slice of one benchmark's test split as plain dicts."""
    from datasets import load_dataset

    if benchmark == "gsm8k":
        ds = load_dataset("openai/gsm8k", "main", split="test")
        rows = [ds[i] for i in _sample(len(ds), limit)]
        return [{"question": r["question"], "gold": r["answer"].split("####")[-1].strip().replace(",", "")}
                for r in rows]
    if benchmark in ("arc_easy", "arc_challenge"):
        name = "ARC-Easy" if benchmark == "arc_easy" else "ARC-Challenge"
        ds = load_dataset("allenai/ai2_arc", name, split="test")
        rows = [ds[i] for i in _sample(len(ds), limit)]
        prompt = ARC_PROMPTS[version]
        return [{"context": prompt.format(question=r["question"].strip()), "choices": list(r["choices"]["text"]),
                 "gold": list(r["choices"]["label"]).index(r["answerKey"])} for r in rows]
    if benchmark == "mmlu":
        ds = load_dataset("cais/mmlu", "all", split="test")
        keep = [i for i, s in enumerate(ds["subject"]) if not any(m in s for m in MATH_SUBJECTS)]
        rows = [ds[keep[j]] for j in _sample(len(keep), limit)]
        return [{"context": f"The following is a question about {r['subject'].replace('_', ' ')}.\n\n"
                            f"Question: {r['question'].strip()}\nAnswer:",
                 "choices": list(r["choices"]), "gold": int(r["answer"])} for r in rows]
    raise ValueError(f"unknown benchmark {benchmark!r}")

_NUM = r"-?\d[\d,]*(?:\.\d+)?"

def extract_answer(text: str) -> str | None:
    """The final number: after "The answer is" if the model says it, else the last number."""
    # This takes the first "answer is", not the last. Fine-tuned outputs say it once, at the end.
    # The published GSM8K v1 scores were computed this way; switching to the last match changes
    # results, so do it together with a GSM8K version bump.
    m = re.search(rf"answer is:?\s*\$?\s*({_NUM})", text, flags=re.IGNORECASE)
    if m is None:
        nums = re.findall(_NUM, text)
        if not nums:
            return None
        found = nums[-1]
    else:
        found = m.group(1)
    found = found.replace(",", "").rstrip(".")
    try:
        v = float(found)
    except ValueError:
        return None
    return str(int(v)) if v.is_integer() else str(v)

def score_generate(model_dir: str, items: list[dict], version: str, n_examples: int = 24) -> dict:
    """GSM8K: greedy generation with vLLM, exact match on the final number."""
    from vllm import LLM, SamplingParams

    started = time.time()
    llm = LLM(model=model_dir, dtype="bfloat16", max_model_len=2048, gpu_memory_utilization=0.85,
              enforce_eager=True, seed=0)
    params = SamplingParams(temperature=0.0, max_tokens=400, stop=["\nQuestion:", "\n\n\n"])
    outs = llm.generate([SFT_PROMPT.format(question=it["question"]) for it in items], params)
    texts = [o.outputs[0].text.strip() for o in outs]
    preds = [extract_answer(t) for t in texts]
    correct = [p is not None and float(p) == float(it["gold"]) for p, it in zip(preds, items)]
    examples = [{"question": it["question"], "gold": it["gold"], "pred": p, "ok": ok, "output": t[:900]}
                for it, p, ok, t in list(zip(items, preds, correct, texts))[:n_examples]]
    return _summary("gsm8k", version, correct, time.time() - started, examples)

def score_choice(model_dir: str, items: list[dict], benchmark: str, version: str, batch_tokens: int = 24_576) -> dict:
    """Multiple choice in the cloze form: pick the answer text with the highest log-likelihood.

    Small models score near chance when they have to answer with a letter, so this scores each
    option's own text after the question, normalized by its length in characters (as OLMES does
    for small models).
    """
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    started = time.time()
    tok = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForCausalLM.from_pretrained(model_dir, dtype=torch.bfloat16).cuda().eval()
    pairs = []  # (item index, token ids, number of answer tokens, answer length in characters)
    for i, it in enumerate(items):
        ctx = tok.encode(it["context"], add_special_tokens=False)
        for choice in it["choices"]:
            ans = tok.encode(" " + choice.strip(), add_special_tokens=False)
            pairs.append((i, (ctx + ans)[-1024:], len(ans), max(1, len(choice.strip()))))
    score = [0.0] * len(pairs)
    pad = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
    order = sorted(range(len(pairs)), key=lambda k: len(pairs[k][1]))  # batch similar lengths
    with torch.inference_mode():
        b = 0
        while b < len(order):
            # Sorted by length, so the last row is the widest: fill the batch up to a token budget.
            n = max(1, batch_tokens // len(pairs[order[min(b + 255, len(order) - 1)]][1]))
            ks = order[b:b + min(n, 256)]
            b += len(ks)
            ids = [pairs[k][1] for k in ks]
            width = max(len(t) for t in ids)
            # Right padding keeps position ids exact for every row.
            x = torch.tensor([t + [pad] * (width - len(t)) for t in ids]).cuda()
            mask = torch.tensor([[1] * len(t) + [0] * (width - len(t)) for t in ids]).cuda()
            logp = torch.log_softmax(model(input_ids=x, attention_mask=mask).logits.float(), dim=-1)
            for r, k in enumerate(ks):
                _, t, n_ans, n_chars = pairs[k]
                pos = torch.arange(len(t) - n_ans, len(t), device=x.device)
                score[k] = logp[r, pos - 1, x[r, pos]].sum().item() / n_chars  # token t is predicted at t-1
    by_item: list[list[float]] = [[] for _ in items]
    for k, (i, _, _, _) in enumerate(pairs):  # pairs are in item, then option, order
        by_item[i].append(score[k])
    correct = [max(range(len(sc)), key=sc.__getitem__) == it["gold"] for sc, it in zip(by_item, items)]
    return _summary(benchmark, version, correct, time.time() - started, [])

@driver_env.task(report=True)
async def train_and_watch(
    profile: str = "full",
    train_examples: int = 50_000,
    n_checkpoints: int = 8,
    learning_rate: float = 3e-5,
    wait_minutes: int = 40,
) -> str:
    """Train in the background, redraw the leaderboard as scores arrive, then promote the winner."""
    train_run = flyte.ctx().action.run_name
    started = time.time()
    trainer = asyncio.create_task(finetune(train_run=train_run, profile=profile, train_examples=train_examples,
                                           n_checkpoints=n_checkpoints, learning_rate=learning_rate))
    info, rows, deadline = None, {}, None
    while True:
        rows = await collect_results(train_run)
        times = await checkpoint_times(train_run)
        if trainer.done() and info is None:
            info = json.loads(trainer.result())  # re-raises if training failed
            deadline = time.time() + wait_minutes * 60
        expected = len(info["steps"]) if info else n_checkpoints + 1
        await flyte.report.replace.aio(
            leaderboard_html(train_run, rows, times, started, info, expected, final=False), do_flush=True)
        if info and (len(rows) >= expected or time.time() > deadline):
            break
        await asyncio.sleep(30)

    best = pick_best(rows)
    release = await promote(train_run, rows, best["step"]) if best else None
    await flyte.report.replace.aio(
        leaderboard_html(train_run, rows, times, started, info, expected, final=True, release=release), do_flush=True)
    return json.dumps({"best_step": best["step"] if best else None, "release_version": release,
                       "evaluated": sorted(rows), "expected": info["steps"]})

async def collect_results(train_run: str) -> dict[int, dict]:
    """The newest published scores for each checkpoint of one training run, keyed by step."""
    from flyte.remote import Artifact

    rows: dict[int, dict] = {}
    async for art in Artifact.listall.aio(name=RESULTS_ARTIFACT, train_run=train_run):
        step = int(art.partitions["step"])
        if step in rows:  # listall is newest first; a backfill publishes a newer version
            continue
        f = await art.to_python.aio(File)
        async with f.open("rb") as fh:
            rows[step] = json.loads(bytes(await fh.read()))
        rows[step]["published_at"] = art.pb2.created_at.ToDatetime().timestamp()
    return rows

async def checkpoint_times(train_run: str) -> dict[int, float]:
    from flyte.remote import Artifact

    return {int(a.partitions["step"]): a.pb2.created_at.ToDatetime().timestamp()
            async for a in Artifact.listall.aio(name=CHECKPOINT_ARTIFACT, train_run=train_run)}

def pick_best(rows: dict[int, dict]) -> dict | None:
    """The ship rule: highest GSM8K among checkpoints that keep MMLU within GUARD_MAX_DROP of step 0."""
    if 0 not in rows:
        return None
    floor = (_acc(rows[0], GUARD_BENCHMARK) or 0.0) - GUARD_MAX_DROP
    ok = [r for s, r in rows.items() if s > 0 and (_acc(r, GUARD_BENCHMARK) or 0.0) >= floor]
    return max(ok, key=lambda r: (_acc(r, "gsm8k") or 0.0, -r["step"])) if ok else None

async def promote(train_run: str, rows: dict[int, dict], best: int) -> str:
    """Publish the winning checkpoint again, as the release-candidate artifact with a model card."""
    from flyte.remote import Artifact

    ckpt = await Artifact.get.aio(CHECKPOINT_ARTIFACT, train_run=train_run, step=str(best))
    card = await artifacts.Card.create_from.aio(content=release_card(train_run, rows, best), format="html",
                                                card_type="model")
    scores = rows[best]["scores"]
    art = await Artifact.create.aio(
        await ckpt.to_python.aio(Dir),
        name=RELEASE_ARTIFACT,
        version=f"{train_run}-step{best:05d}",
        kind="model",
        description=f"step {best} of {train_run}: " + ", ".join(f"{k} {v['accuracy']:.1f}" for k, v in scores.items()),
        attrs={"train_run": train_run, "step": str(best), **{k: f"{v['accuracy']:.2f}" for k, v in scores.items()}},
        card=card,
        parents=[ckpt.artifact_version_id],
    )
    return art.version

    "arc_easy": {"kind": "choice", "version": "2", "limit": 0, "label": "ARC-Easy (MC)", "role": "reasoning"},
    "arc_challenge": {"kind": "choice", "version": "2", "limit": 0, "label": "ARC-Challenge (MC)",
                      "role": "reasoning"},

def backfill(train_run: str, extra: list[str]) -> list[str]:
    """Re-run the suite plus ``extra`` benchmarks on every checkpoint of one training run.

    Benchmarks that already ran on a checkpoint come back from the cache; only ``extra`` runs.
    """
    from flyte.remote import Artifact

    unknown = [b for b in extra if b not in BENCHMARKS]
    if unknown:
        raise ValueError(f"add {unknown} to BENCHMARKS in config.py before backfilling")
    suite = DEFAULT_SUITE + [b for b in extra if b not in DEFAULT_SUITE]
    urls = []
    for art in Artifact.listall(name=CHECKPOINT_ARTIFACT, train_run=train_run, latest_per_partition=True):
        p = art.partitions
        run = flyte.run(evaluate_checkpoint, checkpoint=art, step=p["step"], train_run=train_run,
                        profile=p["profile"], suite=suite)
        urls.append(run.url)
    return urls

python evals.py --backfill <train_run> arc_challenge
python train.py --leaderboard <train_run>