#!/usr/bin/env python3 """Unsloth QLoRA with token-safe projects and complete Drive resume checkpoints.""" from __future__ import annotations import argparse import hashlib import importlib.metadata import json import os import re import signal import sys import time from pathlib import Path ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT / "src")) from vibegame import BASE_MODEL, SYSTEM_PROMPT BASE_REVISION = "4858886f896bd05db823557fc5cbc4ac28342af7" ORIGINAL_BASE = "Qwen/Qwen2.5-Coder-7B-Instruct" ORIGINAL_REVISION = "c03e6d358207e414f1eca0bb1891e29f1db0e242" TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] def atomic_json(path: Path, value): path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_name(path.name + ".tmp") with tmp.open("w") as f: json.dump(value, f, indent=2) f.flush() os.fsync(f.fileno()) tmp.replace(path) def digest(path): h = hashlib.sha256() with Path(path).open("rb") as f: for b in iter(lambda: f.read(1024 * 1024), b""): h.update(b) return h.hexdigest() def prompt_key(text): return re.sub(r"\W+", " ", text.casefold()).strip() def read_jsonl(path): with Path(path).open() as f: return [json.loads(line) for line in f if line.strip()] def check_splits(train, dev, blind): """Blind prompts are read solely to enforce exclusion; never tokenized for loss.""" if len(blind) != 500: raise ValueError(f"Expected exactly 500 blind evaluation prompts; got {len(blind)}") sets = [] for name, rows in [("train", train), ("dev", dev), ("blind", blind)]: keys = [prompt_key(row["instruction"]) for row in rows] if not rows or len(keys) != len(set(keys)): raise ValueError(f"Empty or duplicate prompts in {name}") sets.append(set(keys)) if sets[0] & sets[1] or sets[0] & sets[2] or sets[1] & sets[2]: raise ValueError("Prompt contamination between training, dev, or blind evaluation") output_keys = [{hashlib.sha256(json.dumps(json.loads(r["output"]), sort_keys=True).encode()).hexdigest() for r in rows} for rows in (train, dev)] if output_keys[0] & output_keys[1]: raise ValueError("Identical project outputs appear in train and dev") def tokenize_rows(rows, tokenizer, maximum, split): kept, rejected = [], [] for i, row in enumerate(rows): output = json.loads(row["output"]) files = output.get("files") if not isinstance(files, dict) or "project.godot" not in files or not any(p.endswith(".gd") for p in files): raise ValueError(f"{split} row {i}: expected complete {{files: {{path: text}}}} project") messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": row["instruction"]}] prefix = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True) full = tokenizer.apply_chat_template(messages + [{"role": "assistant", "content": row["output"]}], tokenize=True, add_generation_prompt=False) if full[:len(prefix)] != prefix: raise ValueError("Tokenizer chat template prefix changed; refusing wrong assistant loss mask") if len(full) > maximum: rejected.append({"row": i, "tokens": len(full), "instruction_sha256": hashlib.sha256(row["instruction"].encode()).hexdigest()}) continue labels = [-100] * len(prefix) + full[len(prefix):] if all(label == -100 for label in labels): raise ValueError("Empty assistant response") kept.append({"input_ids": full, "attention_mask": [1] * len(full), "labels": labels}) if not kept: raise ValueError(f"No complete {split} projects fit {maximum} tokens") return kept, rejected def preflight(output): import torch if not os.environ.get("COLAB_RELEASE_TAG") and not Path("/content").is_dir(): raise RuntimeError("This training entry point requires a Google Colab T4 runtime") if not torch.cuda.is_available() or torch.cuda.device_count() != 1: raise RuntimeError("Select the free Colab T4 runtime; no paid GPU is provisioned by this code") props = torch.cuda.get_device_properties(0) if "T4" not in props.name or props.total_memory > 17 * 1024**3: raise RuntimeError(f"Expected one 16GB-class Tesla T4, got {props.name}") drive = Path("/content/drive/MyDrive").resolve() output = output.resolve() if not drive.is_dir() or drive not in output.parents: raise RuntimeError("All training outputs must be in mounted /content/drive/MyDrive") output.mkdir(parents=True, exist_ok=True) probe = output / ".write-probe" probe.write_text("Drive write check") probe.unlink() return {"gpu": props.name, "vram_bytes": props.total_memory, "payment_check": "Code creates no paid resources. Subscription tier is not exposed to Python; select the free T4 in Colab UI."} def latest_complete_checkpoint(output): candidates = sorted(output.glob("checkpoint-*"), key=lambda p: int(p.name.split("-")[-1]), reverse=True) for path in candidates: marker = path / "COMPLETE.json" if not marker.exists(): continue meta = json.loads(marker.read_text()) if all((path / n).exists() and digest(path / n) == value for n, value in meta["sha256"].items()): return path return None def main(): p = argparse.ArgumentParser(description=__doc__) p.add_argument("--train", type=Path, required=True) p.add_argument("--dev", type=Path, required=True) p.add_argument("--blind-prompts", type=Path, required=True) p.add_argument("--output", type=Path, required=True) p.add_argument("--max-seq-length", type=int, default=4096) p.add_argument("--epochs", type=int, choices=[2, 3], default=2) p.add_argument("--save-steps", type=int, default=100) p.add_argument("--session-hours", type=float, default=10.5) p.add_argument("--resume", choices=["auto", "never"], default="auto") args = p.parse_args() if not 512 <= args.max_seq_length <= 4096 or args.save_steps < 1 or not 0 < args.session_hours <= 11: p.error("Use sequence length 512..4096, save_steps >=1, session_hours <=11") os.environ.update({"WANDB_DISABLED": "true", "HF_HUB_DISABLE_TELEMETRY": "1", "TOKENIZERS_PARALLELISM": "false"}) # Import Unsloth first, as its patches must precede transformers imports. # isort: off from unsloth import FastLanguageModel import torch from bitsandbytes.nn import Linear4bit from datasets import Dataset from transformers import BitsAndBytesConfig, DataCollatorForSeq2Seq, Trainer, TrainerCallback, TrainingArguments # isort: on hardware = preflight(args.output) train, dev, blind = map(read_jsonl, [args.train, args.dev, args.blind_prompts]) check_splits(train, dev, blind) validation_manifest = args.train.parent / "validation_manifest.jsonl" if not validation_manifest.exists(): raise RuntimeError("Missing validation_manifest.jsonl. Run build_dataset.py before training.") admitted = {row["output_sha256"] for row in read_jsonl(validation_manifest) if row.get("accepted")} if any(hashlib.sha256(row["output"].encode()).hexdigest() not in admitted for row in train + dev): raise RuntimeError("Training/development data includes a project not admitted by Godot validation") checkpoint = latest_complete_checkpoint(args.output) if args.resume == "auto" else None signature = {"base_model": BASE_MODEL, "base_revision": BASE_REVISION, "original_base": ORIGINAL_BASE, "original_revision": ORIGINAL_REVISION, "max_seq_length": args.max_seq_length, "epochs": args.epochs, "rank": 32, "alpha": 64, "seed": 3407, "target_modules": TARGETS, "train_sha256": digest(args.train), "dev_sha256": digest(args.dev), "validation_manifest_sha256": digest(validation_manifest), "blind_sha256": digest(args.blind_prompts), "system_prompt_sha256": hashlib.sha256(SYSTEM_PROMPT.encode()).hexdigest()} manifest_path = args.output / "run_manifest.json" if manifest_path.exists(): previous = json.loads(manifest_path.read_text()) if previous["signature"] != signature: raise RuntimeError("Dataset or training configuration changed. Choose a new run directory.") if checkpoint is None and (args.output / "adapter").exists(): raise RuntimeError("This run already has a final adapter; choose a new run directory to retrain") versions = {x: importlib.metadata.version(x) for x in ["unsloth", "transformers", "torch", "peft", "datasets", "bitsandbytes", "accelerate"]} if checkpoint and not manifest_path.exists(): embedded = checkpoint / "run_manifest.json" if not embedded.exists(): raise RuntimeError("Restored checkpoint lacks run_manifest.json; copy its original manifest before resuming") previous = json.loads(embedded.read_text()) if previous.get("signature") != signature: raise RuntimeError("Restored checkpoint dataset or configuration differs from current run") if checkpoint and previous.get("versions") != versions: raise RuntimeError("Installed training package versions changed; restore versions from run_manifest.json before resuming") atomic_json(manifest_path, {"signature": signature, "versions": versions, "hardware": hardware, "checkpoint": str(checkpoint) if checkpoint else None, "started_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}) quant = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.float16) model, tokenizer = FastLanguageModel.from_pretrained( model_name=BASE_MODEL, revision=BASE_REVISION, max_seq_length=args.max_seq_length, dtype=torch.float16, load_in_4bit=True, quantization_config=quant, device_map={"": 0}, trust_remote_code=False, use_exact_model_name=True) # Stored bnb checkpoints may carry bf16 defaults; enforce T4-supported fp16 compute. for module in model.modules(): if isinstance(module, Linear4bit): module.compute_dtype = torch.float16 module.compute_type_is_set = True if module.weight.quant_state is not None and module.weight.quant_state.quant_type != "nf4": raise RuntimeError("Base checkpoint is not NF4") model = FastLanguageModel.get_peft_model( model, r=32, target_modules=TARGETS, lora_alpha=64, lora_dropout=0, bias="none", use_gradient_checkpointing="unsloth", random_state=3407, use_rslora=False, loftq_config=None) model.config.use_cache = False tokenizer.padding_side = "right" if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token train_tokens, rejected_train = tokenize_rows(train, tokenizer, args.max_seq_length, "train") dev_tokens, rejected_dev = tokenize_rows(dev, tokenizer, args.max_seq_length, "dev") atomic_json(args.output / "token_filter.json", {"train_kept": len(train_tokens), "dev_kept": len(dev_tokens), "train_rejected": rejected_train, "dev_rejected": rejected_dev, "rule": "Reject whole projects; never truncate."}) start = time.monotonic() stop_requested = {"value": False} def request_stop(*_): stop_requested["value"] = True signal.signal(signal.SIGTERM, request_stop) signal.signal(signal.SIGINT, request_stop) class DurableCheckpoint(TrainerCallback): def on_step_end(self, args, state, control, **kwargs): if stop_requested["value"] or time.monotonic() - start >= args_session_seconds: control.should_save = True control.should_training_stop = True return control def on_epoch_end(self, args, state, control, **kwargs): control.should_save = True return control def on_save(self, args, state, control, **kwargs): target = Path(args.output_dir) / f"checkpoint-{state.global_step}" # Trainer finished saving optimizer, scheduler, RNG and adapter before this callback. required = ["adapter_config.json", "adapter_model.safetensors", "optimizer.pt", "scheduler.pt", "rng_state.pth", "trainer_state.json"] missing = [n for n in required if not (target / n).exists()] if missing: raise RuntimeError(f"Incomplete checkpoint (will not resume): {missing}") tokenizer.save_pretrained(target) atomic_json(target / "run_manifest.json", json.loads(manifest_path.read_text())) atomic_json(target / "COMPLETE.json", {"global_step": state.global_step, "sha256": {f.name: digest(f) for f in target.iterdir() if f.is_file() and f.name != "COMPLETE.json"}}) return control args_session_seconds = args.session_hours * 3600 training_args = TrainingArguments( output_dir=str(args.output), num_train_epochs=args.epochs, per_device_train_batch_size=1, per_device_eval_batch_size=1, gradient_accumulation_steps=16, learning_rate=2e-4, lr_scheduler_type="cosine", warmup_ratio=0.03, optim="adamw_8bit", weight_decay=0.01, fp16=True, bf16=False, gradient_checkpointing=True, max_grad_norm=1.0, save_strategy="steps", save_steps=args.save_steps, save_total_limit=None, save_only_model=False, eval_strategy="epoch", prediction_loss_only=True, eval_accumulation_steps=1, logging_steps=5, report_to="none", seed=3407, data_seed=3407, dataloader_num_workers=0, remove_unused_columns=True, save_safetensors=True) trainer = Trainer(model=model, args=training_args, train_dataset=Dataset.from_list(train_tokens), eval_dataset=Dataset.from_list(dev_tokens), data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True, label_pad_token_id=-100), processing_class=tokenizer, callbacks=[DurableCheckpoint()]) result = trainer.train(resume_from_checkpoint=str(checkpoint) if checkpoint else None) trainer.save_state() complete = trainer.state.global_step >= trainer.state.max_steps metrics = dict(result.metrics) metrics.update({"status": "complete" if complete else "paused_for_session_limit", "global_step": trainer.state.global_step, "target_steps": trainer.state.max_steps, "peak_vram_bytes": torch.cuda.max_memory_allocated(), "blind_eval_used_for_optimization": False}) if complete: model.save_pretrained(args.output / "adapter", safe_serialization=True) tokenizer.save_pretrained(args.output / "adapter") atomic_json(args.output / "adapter" / "run_manifest.json", json.loads(manifest_path.read_text())) metrics.update(trainer.evaluate()) atomic_json(args.output / "adapter" / "TRAINING_COMPLETE.json", metrics) atomic_json(args.output / "training_metrics.json", metrics) print(json.dumps(metrics, indent=2)) if not complete: print("Session budget reached. Reconnect to a free T4 and rerun this notebook; it resumes the last complete Drive checkpoint.") sys.exit(75) if __name__ == "__main__": main()