Something went wrong. Try again.
End-to-end Tinker model training and serving with Eve.
Something went wrong. Try again.
20 kB · 531 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532from __future__ import annotations
import argparseimport asyncioimport contextlibimport hashlibimport jsonimport osimport randomfrom datetime import UTC, datetimefrom importlib.metadata import versionfrom pathlib import Pathfrom typing import Any
import numpy as np
from training.data import dataset_sha256, load_split, validate_casesfrom training.evaluator import evaluate_native, serializable_rowsfrom training.recipes import all_recipes, get_recipefrom training.recipes.base import ROOT, Case, Recipe
BASE_MODEL = "Qwen/Qwen3.5-4B"RENDERER_NAME = "qwen3_5_disable_thinking"COOKBOOK_COMMIT = "0b5c01eaee49bdb0d476f4f383e1c0fb9aced590"ARTIFACTS_DIR = ROOT / "artifacts"
def build_batches(item_count: int, steps: int, batch_size: int, seed: int) -> list[list[int]]: if item_count < 1: raise ValueError("item_count must be positive") if batch_size < 1 or batch_size > item_count: raise ValueError("batch_size must be between 1 and the training-set size") if steps < 1: raise ValueError("steps must be positive")
batches: list[list[int]] = [] epoch = 0 while len(batches) < steps: indices = list(range(item_count)) random.Random(seed + epoch).shuffle(indices) for start in range(0, item_count, batch_size): batch = indices[start : start + batch_size] if len(batch) == batch_size: batches.append(batch) if len(batches) == steps: break epoch += 1 return batches
def oversample_cases( cases: tuple[Case, ...], kinds: tuple[str, ...], factor: int) -> tuple[Case, ...]: if factor < 1: raise ValueError("oversample factor must be positive") unknown = set(kinds) - {case.kind for case in cases} if unknown: raise ValueError(f"unknown oversample kinds: {sorted(unknown)}") expanded = list(cases) if factor > 1: selected = [case for case in cases if case.kind in kinds] for _ in range(factor - 1): expanded.extend(selected) return tuple(expanded)
def promotion_result( baseline_accuracy: float, candidate_accuracy: float, minimum_accuracy: float, minimum_gain: float,) -> dict[str, Any]: gain = candidate_accuracy - baseline_accuracy return { "minimum_accuracy": minimum_accuracy, "minimum_gain": minimum_gain, "baseline_accuracy": baseline_accuracy, "candidate_accuracy": candidate_accuracy, "observed_gain": gain, "passed": candidate_accuracy >= minimum_accuracy and gain >= minimum_gain, }
def source_fingerprint(recipe: Recipe) -> str: digest = hashlib.sha256() paths = ( recipe.instructions_path, ROOT / "training" / "data.py", ROOT / "training" / "evaluator.py", ROOT / "training" / "run.py", ROOT / "training" / "recipes" / "base.py", ROOT / "training" / "recipes" / "registry.py", ROOT / "training" / "recipes" / f"{recipe.id.replace('-', '_')}.py", recipe.data_dir / "train.jsonl", recipe.data_dir / "validation.jsonl", recipe.data_dir / "test.jsonl", ) for path in paths: digest.update(str(path.relative_to(ROOT)).encode()) digest.update(b"\0") digest.update(path.read_bytes()) return digest.hexdigest()
def _write_private_json(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC descriptor = os.open(path, flags, 0o600) try: with os.fdopen(descriptor, "w") as handle: json.dump(payload, handle, indent=2, ensure_ascii=False) handle.write("\n") except BaseException: with contextlib.suppress(OSError): os.close(descriptor) raise path.chmod(0o600)
def _artifact_path(recipe: Recipe, kind: str, stamp: str) -> Path: return ARTIFACTS_DIR / f"{recipe.id}-{kind}-{stamp}" / "result.json"
def _base_receipt(recipe: Recipe, kind: str, renderer_name: str) -> dict[str, Any]: cases = {split: load_split(recipe, split) for split in ("train", "validation", "test")} validate_cases(recipe, cases) return { "schema": 1, "kind": kind, "created_at": datetime.now(UTC).isoformat(), "recipe": recipe.id, "base_model": BASE_MODEL, "renderer": renderer_name, "dataset_sha256": dataset_sha256(recipe), "source_sha256": source_fingerprint(recipe), "counts": {split: len(rows) for split, rows in cases.items()}, "packages": { "tinker": version("tinker"), "tinker-cookbook-commit": COOKBOOK_COMMIT, }, }
async def _runtime(): import tinker from tinker_cookbook.renderers import get_renderer
service = tinker.ServiceClient() sampler = await service.create_sampling_client_async(base_model=BASE_MODEL) renderer = get_renderer(RENDERER_NAME, sampler.get_tokenizer()) return service, sampler, renderer, RENDERER_NAME
async def run_baseline(recipe: Recipe, args: argparse.Namespace) -> tuple[dict[str, Any], Path]: _service, sampler, renderer, renderer_name = await _runtime() rows, score = await evaluate_native( recipe, sampler, renderer, load_split(recipe, "validation"), recipe.instructions_path.read_text(), max_tokens=args.max_tokens, ) stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") receipt = { **_base_receipt(recipe, "baseline", renderer_name), "split": "validation", "score": score, "rows": serializable_rows(rows), } path = _artifact_path(recipe, "baseline", stamp) _write_private_json(path, receipt) return receipt, path
async def verify_checkpoints(service, paths: tuple[str, ...]) -> dict[str, Any]: import tinker
rest = service.create_rest_client() parsed = [tinker.types.ParsedCheckpointTinkerPath.from_tinker_path(path) for path in paths] run_ids = {item.training_run_id for item in parsed} if len(run_ids) != 1: raise ValueError("saved checkpoints do not belong to one training run") run_id = next(iter(run_ids)) training_run, checkpoints = await asyncio.gather( rest.get_training_run_async(run_id), rest.list_checkpoints_async(run_id), ) by_path = {checkpoint.tinker_path: checkpoint for checkpoint in checkpoints.checkpoints} missing = [path for path in paths if path not in by_path] if missing: raise RuntimeError(f"Tinker REST readback did not contain saved checkpoints: {missing}") return { "training_run_id": training_run.training_run_id, "base_model": training_run.base_model, "is_lora": training_run.is_lora, "lora_rank": training_run.lora_rank, "corrupted": training_run.corrupted, "checkpoints": [ { "path": path, "type": by_path[path].checkpoint_type, "public": by_path[path].public, "size_bytes": by_path[path].size_bytes, } for path in paths ], }
async def run_training(recipe: Recipe, args: argparse.Namespace) -> tuple[dict[str, Any], Path]: import tinker from tinker_cookbook.hyperparam_utils import get_lr from tinker_cookbook.renderers import TrainOnWhat from tinker_cookbook.supervised.data import conversation_to_datum
service, base_sampler, renderer, renderer_name = await _runtime() instructions = recipe.instructions_path.read_text() train_cases = load_split(recipe, "train") oversample_kinds = tuple(args.oversample_kind or recipe.default_oversample_kinds) oversample_factor = ( args.oversample_factor if args.oversample_factor is not None else recipe.default_oversample_factor ) minimum_accuracy = ( args.minimum_accuracy if args.minimum_accuracy is not None else recipe.default_minimum_accuracy ) minimum_gain = ( args.minimum_gain if args.minimum_gain is not None else recipe.default_minimum_gain ) effective_train_cases = oversample_cases( train_cases, oversample_kinds, oversample_factor, ) validation_cases = load_split(recipe, "validation") test_cases = load_split(recipe, "test")
baseline_rows, baseline_score = await evaluate_native( recipe, base_sampler, renderer, validation_cases, instructions, max_tokens=args.max_tokens, )
learning_rate = args.learning_rate if args.learning_rate is not None else get_lr(BASE_MODEL) training_client = await service.create_lora_training_client_async( base_model=BASE_MODEL, rank=args.rank, ) training_info = await training_client.get_info_async() datums = [ conversation_to_datum( [ {"role": "system", "content": instructions}, {"role": "user", "content": case.prompt}, { "role": "assistant", "content": json.dumps(case.target, separators=(",", ":")), }, ], renderer, max_length=args.max_length, train_on_what=TrainOnWhat.LAST_ASSISTANT_MESSAGE, ) for case in effective_train_cases ] batches = build_batches(len(datums), args.steps, args.batch_size, args.seed) losses: list[float] = [] train_tokens = 0 for step, indices in enumerate(batches, start=1): batch = [datums[index] for index in indices] train_tokens += sum(datum.model_input.length for datum in batch) forward_future = await training_client.forward_backward_async(batch, "cross_entropy") optimizer_future = await training_client.optim_step_async( tinker.AdamParams(learning_rate=learning_rate) ) forward = await forward_future.result_async() await optimizer_future.result_async() logprobs = np.concatenate( [output["logprobs"].tolist() for output in forward.loss_fn_outputs] ) weights = np.concatenate([datum.loss_fn_inputs["weights"].tolist() for datum in batch]) loss = float(-np.dot(logprobs, weights) / weights.sum()) losses.append(loss) print(f"step {step:02d}/{args.steps}: loss={loss:.4f}", flush=True)
stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") checkpoint_name = f"eve-{recipe.id}-{stamp}" sampler_future = training_client.save_weights_for_sampler(checkpoint_name) state_future = training_client.save_state(checkpoint_name) sampler_receipt, state_receipt = await asyncio.gather( sampler_future.result_async(), state_future.result_async(), ) checkpoint_readback = await verify_checkpoints( service, (sampler_receipt.path, state_receipt.path), )
candidate_sampler = await service.create_sampling_client_async(model_path=sampler_receipt.path) candidate_rows, candidate_score = await evaluate_native( recipe, candidate_sampler, renderer, validation_cases, instructions, max_tokens=args.max_tokens, ) promotion = promotion_result( baseline_score["exact_accuracy"], candidate_score["exact_accuracy"], minimum_accuracy, minimum_gain, )
test_receipt: dict[str, Any] | None = None if promotion["passed"]: baseline_test_rows, baseline_test_score = await evaluate_native( recipe, base_sampler, renderer, test_cases, instructions, max_tokens=args.max_tokens, ) candidate_test_rows, candidate_test_score = await evaluate_native( recipe, candidate_sampler, renderer, test_cases, instructions, max_tokens=args.max_tokens, ) test_receipt = { "opened": True, "baseline": { "score": baseline_test_score, "rows": serializable_rows(baseline_test_rows), }, "candidate": { "score": candidate_test_score, "rows": serializable_rows(candidate_test_rows), }, }
receipt = { **_base_receipt(recipe, "training", renderer_name), "run": { "model_id": training_info.model_id, "sampler_path": sampler_receipt.path, "state_path": state_receipt.path, }, "hyperparameters": { "steps": args.steps, "batch_size": args.batch_size, "learning_rate": learning_rate, "rank": args.rank, "seed": args.seed, "max_length": args.max_length, "max_tokens": args.max_tokens, "oversample_kinds": list(oversample_kinds), "oversample_factor": oversample_factor, "effective_train_examples": len(effective_train_cases), }, "training": { "losses": losses, "train_tokens": train_tokens, }, "validation": { "baseline": { "score": baseline_score, "rows": serializable_rows(baseline_rows), }, "candidate": { "score": candidate_score, "rows": serializable_rows(candidate_rows), }, }, "promotion": promotion, "test": test_receipt, "checkpoint_readback": checkpoint_readback, } artifact_path = _artifact_path(recipe, "training", stamp) _write_private_json(artifact_path, receipt)
if promotion["passed"]: active = { "schema": 2, "recipe": recipe.id, "model_path": sampler_receipt.path, "promoted_at": datetime.now(UTC).isoformat(), "validation_accuracy": candidate_score["exact_accuracy"], "dataset_sha256": dataset_sha256(recipe), "source_sha256": source_fingerprint(recipe), "artifact": str(artifact_path.relative_to(ROOT)), } _write_private_json(recipe.active_checkpoint_path, active)
return receipt, artifact_path
def dry_run(recipe: Recipe, args: argparse.Namespace) -> dict[str, Any]: from tinker_cookbook.hyperparam_utils import get_lr
train = load_split(recipe, "train") validation = load_split(recipe, "validation") test = load_split(recipe, "test") learning_rate = args.learning_rate if args.learning_rate is not None else get_lr(BASE_MODEL) oversample_kinds = tuple(args.oversample_kind or recipe.default_oversample_kinds) oversample_factor = ( args.oversample_factor if args.oversample_factor is not None else recipe.default_oversample_factor ) minimum_accuracy = ( args.minimum_accuracy if args.minimum_accuracy is not None else recipe.default_minimum_accuracy ) minimum_gain = ( args.minimum_gain if args.minimum_gain is not None else recipe.default_minimum_gain ) effective_train = oversample_cases(train, oversample_kinds, oversample_factor) return { **_base_receipt(recipe, "dry-run", RENDERER_NAME), "hyperparameters": { "steps": args.steps, "batch_size": args.batch_size, "learning_rate": learning_rate, "rank": args.rank, "seed": args.seed, "max_length": args.max_length, "max_tokens": args.max_tokens, "oversample_kinds": list(oversample_kinds), "oversample_factor": oversample_factor, "effective_train_examples": len(effective_train), }, "planned_calls": { "baseline_validation_samples": len(validation), "candidate_validation_samples": len(validation), "train_examples_per_step": args.batch_size, "train_steps": args.steps, "held_out_test_samples_after_promotion": 2 * len(test), "train_examples_available": len(train), "effective_train_examples": len(effective_train), }, "promotion_gate": { "minimum_accuracy": minimum_accuracy, "minimum_gain": minimum_gain, }, }
def parser() -> argparse.ArgumentParser: result = argparse.ArgumentParser(description="Train and evaluate an Eve + Tinker recipe") result.add_argument( "--recipe", choices=tuple(recipe.id for recipe in all_recipes()), default="capability-placement", ) action = result.add_mutually_exclusive_group(required=True) action.add_argument("--dry-run", action="store_true", help="print the run contract only") action.add_argument( "--baseline", action="store_true", help="evaluate the raw base on validation" ) action.add_argument( "--train", action="store_true", help="train, gate, and conditionally promote" ) result.add_argument("--steps", type=int, default=32) result.add_argument("--batch-size", type=int, default=8) result.add_argument("--learning-rate", type=float) result.add_argument("--rank", type=int, default=16) result.add_argument("--seed", type=int, default=20260803) result.add_argument("--max-length", type=int, default=768) result.add_argument("--max-tokens", type=int, default=384) result.add_argument( "--oversample-kind", action="append", help="recipe kind to repeat in training; defaults come from the selected recipe", ) result.add_argument("--oversample-factor", type=int) result.add_argument("--minimum-accuracy", type=float) result.add_argument("--minimum-gain", type=float) return result
async def async_main(args: argparse.Namespace) -> None: recipe = get_recipe(args.recipe) if args.dry_run: print(json.dumps(dry_run(recipe, args), indent=2)) return if "TINKER_API_KEY" not in os.environ or not os.environ["TINKER_API_KEY"].strip(): raise SystemExit("TINKER_API_KEY is required for live Tinker calls") if args.baseline: receipt, path = await run_baseline(recipe, args) print( f"baseline exact={receipt['score']['exact']}/{receipt['score']['total']} " f"artifact={path.relative_to(ROOT)}" ) return receipt, path = await run_training(recipe, args) candidate = receipt["validation"]["candidate"]["score"] print( f"candidate exact={candidate['exact']}/{candidate['total']} " f"promoted={receipt['promotion']['passed']} artifact={path.relative_to(ROOT)}" )
def main() -> None: args = parser().parse_args() if args.minimum_accuracy is not None and not 0 <= args.minimum_accuracy <= 1: raise SystemExit("--minimum-accuracy must be in [0, 1]") if args.minimum_gain is not None and not 0 <= args.minimum_gain <= 1: raise SystemExit("--minimum-gain must be in [0, 1]") if args.oversample_factor is not None and args.oversample_factor < 1: raise SystemExit("--oversample-factor must be positive") recipe = get_recipe(args.recipe) unknown_kinds = set(args.oversample_kind or ()) - set(recipe.kinds) if unknown_kinds: raise SystemExit(f"unknown --oversample-kind for {recipe.id}: {sorted(unknown_kinds)}") asyncio.run(async_main(args))
if __name__ == "__main__": main()