code/scripts/train.py

#!/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()