Self-distill Nemotron-3-Nano proofs with Tinker

"""Reinforcement.tech compiled Tinker loop

H1: Reinforcement: Build Your Own Reward Model

INPUT
  Signal : Demonstrations (Self-distillation (SDFT))
  In     : Expert demonstrations you want to keep.
  Task   : Formal Reasoning Engines — Checkable traces
  Model  : nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 (MOE, 30B-A3B)

OUTPUT
  Loop         : Runnable distillation loop.
  Loss         : cross_entropy on on-policy SDFT rewrites
  LoRA rank    : 32
  Steps        : 50

Tinker primitives used in this file:
  sample            generate on-policy rollouts
  forward_backward  accumulate LoRA gradients
  optim_step        Adam update on the adapter
  save_state        checkpoint weights + optimizer

Requires:
  uv pip install tinker
  export TINKER_API_KEY=...
  Docs: https://tinker-docs.thinkingmachines.ai/tinker/quickstart/

Dataset JSONL (one record per line). Pass --data PATH.
  environment : {"prompt": str, "metadata": {...}}
  dpo         : {"prompt": str, "chosen": str, "rejected": str}
  sdft        : {"prompt": str, "completion": str}
  metadata is passed through to the verifier (expected answer, test file, Lean goal).
  --eval-data uses the same schema on a held-out set (never optim_step).
  Without --data this file runs one built-in example (a smoke test, not training).

CLI (overrides the module constants; --seed is only an args field):
  --data --eval-data --eval-every --log --steps --rank --lr --out --seed
  --resume --group-size --prompts-per-step --max-tokens --temperature
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import random
from pathlib import Path

try:
    import tinker
    from tinker import types
except ImportError:  # --help and load_dataset work without Tinker installed
    tinker = None
    types = None

BASE_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
LORA_RANK = 32
LEARNING_RATE = 5e-4
STEPS = 50
MAX_TOKENS = 256
TEMPERATURE = 0.5
GROUP_SIZE = 8
PROMPTS_PER_STEP = 1
EVAL_EVERY = 10
SAVE_EVERY = 10
RUN_NAME = Path(__file__).stem
DATASET_KEYS = ("prompt", "completion",)

LARGE_MODEL_NOTE = ""


def parse_args():
    parser = argparse.ArgumentParser(description="Reinforcement.tech compiled Tinker loop")
    parser.add_argument("--data", help="JSONL dataset path")
    parser.add_argument("--eval-data", help="Held-out JSONL. Same schema as --data. Never used for optim_step.")
    parser.add_argument("--eval-every", type=int, default=EVAL_EVERY)
    parser.add_argument("--log", help="Append per-step JSONL: {step, mean_reward, n_datums, loss}")
    parser.add_argument("--steps", type=int, default=STEPS)
    parser.add_argument("--rank", type=int, default=LORA_RANK)
    parser.add_argument("--lr", type=float, default=float(LEARNING_RATE))
    parser.add_argument("--out", default=RUN_NAME, help="Run name / checkpoint prefix")
    parser.add_argument(
        "--seed",
        type=int,
        default=0,
        help="Shuffle seed for --data / --eval-data. Not copied in apply_args; rows_for_run reads args.seed.",
    )
    parser.add_argument("--resume", help="tinker:// path printed by save_state (weights + optimizer)")
    parser.add_argument("--group-size", type=int, default=GROUP_SIZE)
    parser.add_argument("--prompts-per-step", type=int, default=PROMPTS_PER_STEP)
    parser.add_argument("--max-tokens", type=int, default=MAX_TOKENS)
    parser.add_argument("--temperature", type=float, default=TEMPERATURE)
    return parser.parse_args()


def apply_args(args) -> None:
    global LORA_RANK, STEPS, LEARNING_RATE, RUN_NAME
    global GROUP_SIZE, PROMPTS_PER_STEP, MAX_TOKENS, TEMPERATURE, EVAL_EVERY
    # --seed is intentionally not a module constant. rows_for_run(args) and the
    # eval shuffle read args.seed so two shuffles stay independent of apply_args.
    LORA_RANK = args.rank
    STEPS = args.steps
    LEARNING_RATE = args.lr
    RUN_NAME = args.out
    GROUP_SIZE = args.group_size
    PROMPTS_PER_STEP = args.prompts_per_step
    MAX_TOKENS = args.max_tokens
    TEMPERATURE = args.temperature
    EVAL_EVERY = args.eval_every


def require_key() -> None:
    if not os.environ.get("TINKER_API_KEY"):
        raise SystemExit("Set TINKER_API_KEY before running this loop.")


def load_dataset(path: str) -> list[dict]:
    """Read JSONL. One record per line. Schema: see the module docstring."""
    rows: list[dict] = []
    with open(path, encoding="utf-8") as handle:
        for line_no, line in enumerate(handle, 1):
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            missing = [key for key in DATASET_KEYS if key not in row]
            if missing:
                raise SystemExit(f"{path}:{line_no} missing {missing}")
            rows.append(row)
    if not rows:
        raise SystemExit(f"{path} had no JSONL rows.")
    return rows


def rows_for_run(args) -> list[dict]:
    if args.data:
        rows = load_dataset(args.data)
    else:
        print("No --data given; running on 1 built-in example. This is a smoke test, not training.")
        rows = list(EXAMPLE_ROWS)
    rng = random.Random(args.seed)
    rng.shuffle(rows)
    needed = max(args.steps, 1) * max(args.prompts_per_step, 1)
    if needed > len(rows):
        print(f"warning: --steps/--prompts-per-step need {needed} rows; dataset has {len(rows)}; cycling.")
    return rows


def eval_rows_for_run(args) -> list[dict]:
    if not args.eval_data:
        return []
    rows = load_dataset(args.eval_data)
    rng = random.Random(args.seed)
    rng.shuffle(rows)
    return rows


def print_config(args, rows: list[dict], eval_rows: list[dict]) -> None:
    print("=== run config ===")
    print(f"model={BASE_MODEL}")
    print(f"rank={LORA_RANK}  steps={STEPS}  lr={LEARNING_RATE}  seed={args.seed}")
    print(f"group_size={GROUP_SIZE}  prompts_per_step={PROMPTS_PER_STEP}")
    print(f"max_tokens={MAX_TOKENS}  temperature={TEMPERATURE}")
    print(f"data={args.data or '(EXAMPLE_ROWS)'}  n_train={len(rows)}")
    print(f"eval_data={args.eval_data or '(none)'}  n_eval={len(eval_rows)}  eval_every={EVAL_EVERY}")
    print(f"log={args.log or '(none)'}  resume={args.resume or '(none)'}  out={RUN_NAME}")
    print("==================")
    if LARGE_MODEL_NOTE:
        print(LARGE_MODEL_NOTE)


def append_metrics(path: str | None, record: dict) -> None:
    if not path:
        return
    out = Path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    with out.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record) + "\n")


def metric_loss(result) -> float | None:
    metrics = getattr(result, "metrics", None) or {}
    if hasattr(metrics, "get"):
        for key in ("loss", "cross_entropy"):
            if metrics.get(key) is not None:
                return metrics[key]
    loss = getattr(result, "loss", None)
    return float(loss) if loss is not None else None


async def connect(args):
    if tinker is None or types is None:
        raise SystemExit("uv pip install tinker  — then rerun.")
    service = tinker.ServiceClient()
    if args.resume:
        training = await service.create_training_client_from_state_with_optimizer_async(args.resume)
        print(f"resumed weights+optimizer from {args.resume}")
    else:
        training = await service.create_lora_training_client_async(
            base_model=BASE_MODEL,
            rank=LORA_RANK,
            user_metadata={"product": "reinforcement.tech", "run": RUN_NAME},
        )
    tokenizer = training.get_tokenizer()
    return service, training, tokenizer


async def sampling_client(training):
    """Ephemeral on-policy sampler. Do not pass name= — it is deprecated and ignored."""
    return await training.save_weights_and_get_sampling_client_async()


async def checkpoint(training, step: int) -> None:
    if STEPS == 0:
        return
    if step % SAVE_EVERY != 0 and step != STEPS - 1:
        return
    saved = await training.save_state_async(name=f"{RUN_NAME}-step-{step}")
    result = await saved.result_async()
    path = getattr(result, "path", None)
    print(f"checkpoint name={RUN_NAME}-step-{step} path={path}")
    print("resume later with --resume PATH")

# Built-in smoke-test demo. Used only when --data is absent.
DEMO = "theorem add_comm (a b : Nat) : a + b = b + a := by omega"
STUDENT_PROMPT = "Solve the task. Do not look at a demonstration."
EXAMPLE_ROWS = [{"prompt": STUDENT_PROMPT, "completion": DEMO}]


def as_sft_datum(tokenizer, prompt: str, completion: str) -> types.Datum:
    prompt_tokens = tokenizer.encode(prompt)
    completion_tokens = tokenizer.encode(completion)
    full = prompt_tokens + completion_tokens
    n_prefix = max(len(prompt_tokens) - 1, 0)
    return types.Datum(
        model_input=types.ModelInput.from_ints(tokens=full[:-1]),
        loss_fn_inputs=dict(
            target_tokens=full[1:],
            weights=[0.0] * n_prefix + [1.0] * len(completion_tokens),
        ),
    )


async def forward_ce(training, datums) -> float:
    future = await training.forward_async(datums, loss_fn="cross_entropy")
    result = await future.result_async()
    loss = metric_loss(result)
    return float(loss) if loss is not None else 0.0


async def evaluate(eval_rows, training, tokenizer) -> float:
    losses: list[float] = []
    for row in eval_rows:
        datum = as_sft_datum(tokenizer, str(row["prompt"]), str(row["completion"]))
        losses.append(await forward_ce(training, [datum]))
    return sum(losses) / max(len(losses), 1)


async def train(args) -> None:
    require_key()
    rows = rows_for_run(args)
    eval_rows = eval_rows_for_run(args)
    print_config(args, rows, eval_rows)
    _service, training, tokenizer = await connect(args)
    preview_prompt_text = str(rows[0]["prompt"]) if rows else STUDENT_PROMPT

    for step in range(STEPS):
        sampling = await sampling_client(training)
        datums = []
        for offset in range(PROMPTS_PER_STEP):
            row = rows[(step * PROMPTS_PER_STEP + offset) % len(rows)]
            teacher = (
                "You are the teacher. The expert demonstration is below. "
                "Rewrite it in your own on-policy style. Keep the same facts and tools.\n\n"
                f"DEMO:\n{row['completion']}"
            )
            teacher_prompt = types.ModelInput.from_ints(tokenizer.encode(teacher))
            rewrite = await sampling.sample_async(
                prompt=teacher_prompt,
                num_samples=1,
                sampling_params=types.SamplingParams(max_tokens=MAX_TOKENS, temperature=TEMPERATURE),
            )
            on_policy = tokenizer.decode(rewrite.sequences[0].tokens)
            datums.append(as_sft_datum(tokenizer, str(row["prompt"]), on_policy))

        fwdbwd = await training.forward_backward_async(datums, loss_fn="cross_entropy")
        optim = await training.optim_step_async(types.AdamParams(learning_rate=LEARNING_RATE))
        result = await fwdbwd.result_async()
        await optim.result_async()
        await checkpoint(training, step)
        loss = metric_loss(result)
        append_metrics(args.log, {"step": step, "mean_reward": None, "n_datums": len(datums), "loss": loss})
        print(f"step {step:03d}  {result.metrics}")
        if eval_rows and step % EVAL_EVERY == 0:
            eval_loss = await evaluate(eval_rows, training, tokenizer)
            append_metrics(
                args.log,
                {"step": step, "mean_reward": None, "n_datums": len(eval_rows), "loss": eval_loss, "split": "eval"},
            )
            print(f"step {step:03d}  eval_loss={eval_loss:.3f}  n={len(eval_rows)}")

    sampling = await sampling_client(training)
    preview_prompt = types.ModelInput.from_ints(tokenizer.encode(preview_prompt_text))
    preview = await sampling.sample_async(
        prompt=preview_prompt,
        num_samples=1,
        sampling_params=types.SamplingParams(max_tokens=MAX_TOKENS, temperature=TEMPERATURE),
    )
    print("sample:", tokenizer.decode(preview.sequences[0].tokens))


if __name__ == "__main__":
    args = parse_args()
    apply_args(args)
    asyncio.run(train(args))