| |
| """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"}) |
| |
| |
| from unsloth import FastLanguageModel |
| import torch |
| from bitsandbytes.nn import Linear4bit |
| from datasets import Dataset |
| from transformers import BitsAndBytesConfig, DataCollatorForSeq2Seq, Trainer, TrainerCallback, TrainingArguments |
| |
| 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) |
| |
| 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}" |
| |
| 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() |
|
|
