| """Train a Jev-style decision model on Qwen3.8-27B with Unsloth's FastDecisionModel (Clef head). |
| |
| One LoRA run, sized to a wall-clock budget measured on the GPU it runs on: |
| * LocalLLaMA/typed-decisions train (soft labels), always all of it |
| * n4ze3m/typed-decisions-synth (DeepSeek soft labels, 100 other workflows) |
| * Unsloth's public decision mixture (12 human-labelled sources, augmented) |
| Temperatures are fitted on a typed-decisions train holdout. Evaluated on typed-decisions test and |
| on held-out ANLI / WinoGrande rows that no training source contains. |
| |
| No TypeSafe / Jev output is used anywhere in training, calibration or model selection. |
| |
| python train_dm.py --model /content/models/Qwen3.8-27B --train-minutes 40 |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import random |
| import sys |
| import time |
|
|
| T0 = time.time() |
|
|
|
|
| def log(*parts): |
| print(f"[{time.strftime('%H:%M:%S')} +{(time.time() - T0) / 60:.1f}m]", *parts, flush = True) |
|
|
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model", default = "Qwen/Qwen3.8-27B") |
| parser.add_argument("--base-name", default = "Qwen/Qwen3.8-27B", help = "recorded as the adapters' base") |
| parser.add_argument("--load-in-4bit", action = "store_true") |
| parser.add_argument("--max-seq-length", type = int, default = 2048) |
| parser.add_argument("--lora-r", type = int, default = 32) |
| parser.add_argument("--lr", type = float, default = 1e-4) |
| parser.add_argument("--head-lr", type = float, default = 3e-4) |
| parser.add_argument("--batch-size", type = int, default = 8) |
| parser.add_argument("--grad-accum", type = int, default = 2) |
| parser.add_argument("--synth-rows", type = int, default = 3000) |
| parser.add_argument("--public-rows", type = int, default = 2400) |
| parser.add_argument("--train-minutes", type = float, default = 40.0) |
| parser.add_argument("--max-items", type = int, default = None, help = "cap on training items (smoke)") |
| parser.add_argument("--probe-steps", type = int, default = 3) |
| parser.add_argument("--eval-limit", type = int, default = None) |
| parser.add_argument("--heldout", default = "/content/out/heldout_rows.jsonl") |
| parser.add_argument("--seed", type = int, default = 3407) |
| parser.add_argument("--out", default = "/content/out") |
| parser.add_argument("--name", default = "final") |
| parser.add_argument("--ckpt-repo", default = None, help = "private HF repo for the checkpoint") |
| args = parser.parse_args() |
| os.makedirs(args.out, exist_ok = True) |
|
|
| from unsloth import DecisionTrainer, FastDecisionModel |
| from unsloth.models import decision_datasets as dd |
|
|
| import torch |
| from datasets import load_dataset |
| from transformers import TrainerCallback, TrainingArguments |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from td_metrics import parsed, score_cases |
|
|
| PUBLIC_SOURCES = ( |
| "banking77", |
| "clinc150", |
| "mnli", |
| "snli", |
| "wanli", |
| "boolq", |
| "ag_news", |
| "sst5", |
| "mmlu", |
| "commonsense_qa", |
| "arc", |
| "prompt_injections", |
| ) |
|
|
|
|
| class Progress(TrainerCallback): |
| def __init__(self, name, limit_s = None): |
| self.name, self.limit_s, self.start, self.steps = name, limit_s, None, [] |
|
|
| def on_train_begin(self, args, state, control, **kwargs): |
| self.start = time.time() |
| log(f"{self.name}: train begin, max_steps={state.max_steps}") |
|
|
| def on_step_end(self, args, state, control, **kwargs): |
| self.steps.append(time.time()) |
| if self.limit_s and time.time() - self.start > self.limit_s: |
| log(f"{self.name}: time limit {self.limit_s / 60:.1f}m reached at step {state.global_step}") |
| control.should_training_stop = True |
|
|
| def on_log(self, args, state, control, logs = None, **kwargs): |
| if not logs: |
| return |
| recent = self.steps[-6:] |
| per_step = (recent[-1] - recent[0]) / (len(recent) - 1) if len(recent) > 1 else None |
| left = (state.max_steps - state.global_step) * per_step if per_step else None |
| log( |
| f"{self.name}: step {state.global_step}/{state.max_steps}", |
| json.dumps({k: round(v, 5) if isinstance(v, float) else v for k, v in logs.items()}), |
| f"s/step={per_step:.2f}" if per_step else "", |
| f"eta={left / 60:.1f}m" if left else "", |
| ) |
|
|
|
|
| def gpu_mem(): |
| if not torch.cuda.is_available(): |
| return {} |
| return { |
| "allocated_gb": round(torch.cuda.memory_allocated() / 1e9, 1), |
| "peak_gb": round(torch.cuda.max_memory_allocated() / 1e9, 1), |
| } |
|
|
|
|
| def synth_rows(limit, seed, cleaner): |
| |
| rows = [] |
| for row in load_dataset("n4ze3m/typed-decisions-synth", split = "train"): |
| state = parsed(row["state"]) if row.get("state_is_json") else row["state"] |
| questions, teacher = parsed(row["questions"]), parsed(row["teacher"]) |
| if not isinstance(questions, dict) or not isinstance(teacher, dict): |
| continue |
| gold = {} |
| for qid, question in questions.items(): |
| answer = teacher.get(qid) |
| if not isinstance(answer, dict): |
| continue |
| if question.get("type") == "noul" and answer.get("noul") is not None: |
| p = float(answer["noul"]) |
| gold[qid] = {"probabilities": {"false": 1.0 - p, "true": p}} |
| elif answer.get("probabilities"): |
| gold[qid] = {"probabilities": answer["probabilities"]} |
| questions = {qid: q for qid, q in questions.items() if qid in gold} |
| if not questions: |
| continue |
| candidate = {"source": "synth", "state": state, "questions": questions, "gold": gold} |
| if not cleaner.contaminated(candidate): |
| rows.append(candidate) |
| random.Random(seed).shuffle(rows) |
| rng = random.Random(seed + 1) |
| return [dd.augment_row(row, rng) for row in rows[:limit]] |
|
|
|
|
| def public_rows(n_rows, seed, eval_texts): |
| per_source = max(1, -(-n_rows // len(PUBLIC_SOURCES))) |
| pools = {} |
| for name in PUBLIC_SOURCES: |
| try: |
| pools[name] = dd.load_source(name, "train", limit = 2 * per_source, seed = seed) |
| except Exception as exc: |
| log(f"skipping {name}: {type(exc).__name__}: {str(exc)[:200]}") |
| log("public sources", json.dumps({k: len(v) for k, v in pools.items()})) |
| return dd.build_decision_mixture(pools, n_rows = n_rows, seed = seed, decontaminate_against = eval_texts) |
|
|
|
|
| def predict_cases(model, tokenizer, rows, name): |
| FastDecisionModel.for_inference(model) |
| cases, start = [], time.time() |
| for i, row in enumerate(rows, 1): |
| answers = FastDecisionModel.predict(model, tokenizer, parsed(row["state"]), parsed(row["questions"])) |
| cases.append( |
| { |
| "id": row["id"], |
| "workflow": row["workflow"], |
| "questions": row["questions"], |
| "gold": row["gold"], |
| "pred": {qid: {k: v for k, v in a.items() if k != "legend"} for qid, a in answers.items()}, |
| } |
| ) |
| if i % 100 == 0: |
| log(f"{name}: predicted {i}/{len(rows)} ({(time.time() - start) / i:.2f}s/case)") |
| metrics = score_cases(cases) |
| with open(f"{args.out}/{name}_preds.jsonl", "w", encoding = "utf-8") as f: |
| for case in cases: |
| f.write(json.dumps(case) + "\n") |
| with open(f"{args.out}/{name}_metrics.json", "w", encoding = "utf-8") as f: |
| json.dump(metrics, f, indent = 1) |
| log(f"{name.upper()}_METRICS", json.dumps(metrics)) |
| return metrics |
|
|
|
|
| def training_args(max_steps, lr, head_lr_unused = None): |
| return TrainingArguments( |
| output_dir = f"{args.out}/_trainer", |
| per_device_train_batch_size = args.batch_size, |
| gradient_accumulation_steps = args.grad_accum, |
| learning_rate = lr, |
| lr_scheduler_type = "cosine", |
| warmup_steps = max(1, round(0.05 * max_steps)), |
| weight_decay = 0.01, |
| max_steps = max_steps, |
| bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported(), |
| logging_steps = 1 if max_steps <= 10 else 5, |
| report_to = "none", |
| save_strategy = "no", |
| seed = args.seed, |
| dataloader_num_workers = 0, |
| ) |
|
|
|
|
| def main(): |
| log("args", json.dumps(vars(args))) |
| model, tokenizer = FastDecisionModel.from_pretrained( |
| args.model, |
| decision_head = "clef", |
| max_seq_length = args.max_seq_length, |
| load_in_4bit = args.load_in_4bit, |
| random_state = args.seed, |
| ) |
| model.decision_config["base_model"] = args.base_name |
| log("loaded", type(model.encoder).__name__, json.dumps(model.head.config), json.dumps(gpu_mem())) |
|
|
| td_test = list(load_dataset("LocalLLaMA/typed-decisions", "all", split = "test")) |
| if args.eval_limit: |
| td_test = td_test[: args.eval_limit] |
| heldout = [] |
| if args.heldout and os.path.exists(args.heldout): |
| heldout = [json.loads(line) for line in open(args.heldout, encoding = "utf-8") if line.strip()] |
| if args.eval_limit: |
| heldout = heldout[: args.eval_limit] |
| eval_texts = [parsed(row["state"]) for row in td_test + heldout] |
| cleaner = dd.Decontaminator(eval_texts) |
|
|
| |
| |
| |
| test_states = {json.dumps(parsed(row["state"]), sort_keys = True) for row in td_test} |
| rng = random.Random(args.seed + 2) |
| td_rows = [ |
| dd.augment_row(r, rng) |
| for r in dd.load_source("typed_decisions", "train") |
| if json.dumps(r["state"], sort_keys = True) not in test_states |
| ] |
| td_items, td_report = FastDecisionModel.build_dataset(td_rows, tokenizer, model) |
| td_train, td_holdout = FastDecisionModel.split_holdout(td_items, args.seed) |
| synth_items, synth_report = FastDecisionModel.build_dataset(synth_rows(args.synth_rows, args.seed, cleaner), tokenizer, model) |
| pub_items, pub_report = FastDecisionModel.build_dataset(public_rows(args.public_rows, args.seed, eval_texts), tokenizer, model) |
| for name, items, report in (("typed", td_items, td_report), ("synth", synth_items, synth_report), ("public", pub_items, pub_report)): |
| lengths = sorted(len(item["input_ids"]) for item in items) or [0] |
| log(f"data {name}: items={len(items)} skipped={report['skipped']} truncated={report['truncated']} " |
| f"tokens p50={lengths[len(lengths) // 2]} p90={lengths[int(len(lengths) * 0.9)]} max={lengths[-1]}") |
| log(f"typed train={len(td_train)} holdout={len(td_holdout)}") |
|
|
| model = FastDecisionModel.get_peft_model(model, r = args.lora_r, lora_alpha = args.lora_r, random_state = args.seed) |
| log(f"trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.1f}M", json.dumps(gpu_mem())) |
|
|
| |
| per_step = args.batch_size * args.grad_accum |
| pool = td_train + synth_items + pub_items |
| probe = random.Random(args.seed + 3).sample(pool, min(len(pool), per_step * args.probe_steps)) |
| progress = Progress("probe") |
| DecisionTrainer( |
| model = model, tokenizer = tokenizer, train_dataset = probe, head_learning_rate = 0.0, |
| args = training_args(args.probe_steps, 0.0), callbacks = [progress], |
| ).train() |
| stamps = progress.steps |
| probe_tokens = sum(len(item["input_ids"]) for item in probe) / args.probe_steps |
| s_step = (stamps[-1] - stamps[0]) / (len(stamps) - 1) if len(stamps) > 1 else stamps[0] - progress.start |
| s_token = s_step / probe_tokens |
| log(f"probe: {s_step:.2f}s/step for {probe_tokens:.0f} tokens/step -> {s_token * 1e3:.3f} ms/token", json.dumps(gpu_mem())) |
|
|
| |
| budget_tokens = args.train_minutes * 60 / s_token / 1.03 |
| rng = random.Random(args.seed + 4) |
| rng.shuffle(synth_items) |
| rng.shuffle(pub_items) |
| selected = list(td_train) |
| used = sum(len(item["input_ids"]) for item in selected) |
| extra = [x for pair in zip(synth_items, pub_items) for x in pair] + synth_items[len(pub_items):] + pub_items[len(synth_items):] |
| for item in extra: |
| if used >= budget_tokens or (args.max_items and len(selected) >= args.max_items): |
| break |
| selected.append(item) |
| used += len(item["input_ids"]) |
| if args.max_items: |
| selected = selected[: args.max_items] |
| rng.shuffle(selected) |
| steps = max(1, len(selected) // per_step) |
| typed_ids = {id(x) for x in td_train} |
| counts = {"typed": sum(id(x) in typed_ids for x in selected), "total": len(selected)} |
| log(f"selected {len(selected)} items ({used / 1e6:.2f}M tokens) -> {steps} steps, est {used * s_token / 60:.1f}m", json.dumps(counts)) |
|
|
| progress = Progress("train", limit_s = args.train_minutes * 60 * 1.15) |
| trainer = DecisionTrainer( |
| model = model, tokenizer = tokenizer, train_dataset = selected, head_learning_rate = args.head_lr, |
| args = training_args(steps, args.lr), callbacks = [progress], |
| ) |
| out = trainer.train() |
| log("train done", json.dumps(out.metrics), json.dumps(gpu_mem())) |
|
|
| calibration = FastDecisionModel.calibrate(model, tokenizer, td_holdout) if td_holdout else {} |
| calibration = {k: v for k, v in calibration.items() if not isinstance(v, (list, dict))} |
| log("calibration", json.dumps(calibration), "temperatures", model.decision_config.get("temperature")) |
|
|
| metrics = {"typed_decisions_test": predict_cases(model, tokenizer, td_test, f"{args.name}_td_test")} |
| if heldout: |
| metrics["heldout"] = predict_cases(model, tokenizer, heldout, f"{args.name}_heldout") |
|
|
| report = { |
| "base_model": args.base_name, |
| "train_items": len(selected), |
| "train_tokens": used, |
| "steps": trainer.state.global_step, |
| "planned_steps": steps, |
| "probe_s_per_step": s_step, |
| "train_metrics": out.metrics, |
| "data": {"typed_train": len(td_train), "typed_holdout": len(td_holdout), "synth_pool": len(synth_items), "public_pool": len(pub_items)}, |
| "calibration_on_typed_holdout": calibration, |
| "temperatures": model.decision_config.get("temperature"), |
| "eval": metrics, |
| "gpu": torch.cuda.get_device_name() if torch.cuda.is_available() else "cpu", |
| "args": vars(args), |
| } |
| folder = f"{args.out}/{args.name}" |
| model.decision_config["base_model"] = args.base_name |
| model.save_pretrained(folder, tokenizer) |
| with open(f"{folder}/train_report.json", "w", encoding = "utf-8") as f: |
| json.dump(report, f, indent = 1) |
| log(f"saved {folder}: {sorted(os.listdir(folder))}") |
| if args.ckpt_repo: |
| from huggingface_hub import HfApi |
|
|
| api = HfApi() |
| api.create_repo(args.ckpt_repo, private = True, exist_ok = True) |
| api.upload_folder(folder_path = folder, repo_id = args.ckpt_repo, path_in_repo = args.name) |
| for path in os.listdir(args.out): |
| if path.startswith(args.name) and path.endswith((".json", ".jsonl")): |
| api.upload_file(path_or_fileobj = f"{args.out}/{path}", path_in_repo = f"eval/{path}", repo_id = args.ckpt_repo) |
| log(f"pushed {args.name} to {args.ckpt_repo} (private)") |
| log("ALL_DONE", json.dumps({k: v.get("accuracy") for k, v in metrics.items()})) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
