| """Cloud GPU QLoRA training. Never falls back to training on the user's CPU.""" |
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| import shutil |
|
|
| from .data import encode_example, normalize_messages, prompt_key |
| from .runtime import CONFIG |
|
|
|
|
| def read_data(folder): |
| manifest = json.loads((folder / "manifest.json").read_text(encoding="utf-8")) |
| if manifest["config"] != CONFIG: |
| raise ValueError("Data manifest does not match the pinned experiment configuration.") |
| records, seen = {}, set() |
| for split in ["train", "validation", "test"]: |
| raw = (folder / f"{split}.jsonl").read_bytes() |
| if hashlib.sha256(raw).hexdigest() != manifest["sha256"][split]: |
| raise ValueError(f"Data changed after preparation: {split}") |
| records[split] = [json.loads(line) for line in raw.decode().splitlines() if line] |
| if not records[split] or len(records[split]) != manifest["counts"][split]: |
| raise ValueError(f"Unexpected split size: {split}") |
| for row in records[split]: |
| key = prompt_key(normalize_messages(row)) |
| if key != row["prompt_hash"] or key in seen: |
| raise ValueError("Duplicate or inconsistent prompts across dataset splits.") |
| seen.add(key) |
| return records |
|
|
|
|
| def collate(examples, pad_token_id): |
| import torch |
| length = max(len(item["input_ids"]) for item in examples) |
| return {key: torch.tensor([item[key] + [pad] * (length - len(item[key])) for item in examples]) |
| for key, pad in [("input_ids", pad_token_id), ("attention_mask", 0), ("labels", -100)]} |
|
|
|
|
| def train(args): |
| import torch |
| if not torch.cuda.is_available(): |
| raise RuntimeError("Cloud CUDA GPU required. No training or model download was started.") |
| from datasets import Dataset |
| from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training |
| from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig |
| from transformers import Trainer, TrainingArguments, set_seed |
|
|
| output = Path(args.output) |
| if output.exists() and any(output.iterdir()) and not args.resume_from_checkpoint: |
| raise FileExistsError("Use an empty output directory or explicitly resume a checkpoint.") |
| records = read_data(Path(args.data)) |
| set_seed(CONFIG["seed"]) |
| tokenizer = AutoTokenizer.from_pretrained( |
| CONFIG["base_model"], revision=CONFIG["base_revision"], trust_remote_code=False) |
| tokenizer.padding_side = "right" |
| dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 |
| base = AutoModelForCausalLM.from_pretrained( |
| CONFIG["base_model"], revision=CONFIG["base_revision"], trust_remote_code=False, |
| dtype=dtype, device_map={"": torch.cuda.current_device()}, attn_implementation="sdpa", |
| quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", |
| bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=dtype)) |
| base.config.use_cache = False |
| base = prepare_model_for_kbit_training(base, use_gradient_checkpointing=True, |
| gradient_checkpointing_kwargs={"use_reentrant": False}) |
| model = get_peft_model(base, LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, |
| target_modules="all-linear", bias="none", task_type=TaskType.CAUSAL_LM)) |
| model.peft_config["default"].base_model_name_or_path = CONFIG["base_model"] |
| model.peft_config["default"].revision = CONFIG["base_revision"] |
| datasets = {split: Dataset.from_list([ |
| encode_example(tokenizer, row["messages"], CONFIG["max_sequence_length"]) for row in records[split] |
| ]) for split in ["train", "validation"]} |
| settings = TrainingArguments( |
| output_dir=str(output), num_train_epochs=1, max_steps=args.max_steps, |
| per_device_train_batch_size=1, per_device_eval_batch_size=1, |
| gradient_accumulation_steps=16, learning_rate=5e-5, warmup_ratio=0.05, |
| lr_scheduler_type="cosine", logging_steps=5, eval_strategy="steps", eval_steps=25, |
| save_strategy="steps", save_steps=25, save_total_limit=2, |
| load_best_model_at_end=True, metric_for_best_model="eval_loss", |
| greater_is_better=False, bf16=dtype == torch.bfloat16, fp16=dtype == torch.float16, |
| gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False}, |
| optim="adamw_torch", report_to="none", seed=CONFIG["seed"], |
| dataloader_num_workers=0, remove_unused_columns=False) |
| trainer = Trainer(model=model, args=settings, processing_class=tokenizer, |
| train_dataset=datasets["train"], eval_dataset=datasets["validation"], |
| data_collator=lambda batch: collate(batch, tokenizer.pad_token_id)) |
| baseline = trainer.evaluate() |
| trained = trainer.train(resume_from_checkpoint=args.resume_from_checkpoint) |
| final = trainer.evaluate() |
| destination = output / "adapter" |
| trainer.save_model(str(destination)) |
| tokenizer.save_pretrained(destination) |
| shutil.copy2(Path(args.data) / "manifest.json", destination / "data_manifest.json") |
| report = {"config": CONFIG, "baseline_validation": baseline, "final_validation": final, |
| "training": trained.metrics, "counts": {k: len(v) for k, v in records.items()}, |
| "status": "candidate_requires_answer_and_code_evaluation", "max_steps": args.max_steps} |
| (destination / "training_report.json").write_text(json.dumps(report, indent=2), encoding="utf-8") |
| (destination / "README.md").write_text( |
| "---\nlicense: gpl-3.0\nlibrary_name: peft\npipeline_tag: text-generation\n" |
| f"base_model: {CONFIG['base_model']}\nbase_model_relation: adapter\n" |
| "datasets:\n- HuggingFaceTB/smoltalk\ntags:\n- lora\n- code\n- experimental\n---\n\n" |
| "# PocketVision Chat & Code candidate\n\nText-only QLoRA adapter. " |
| "This candidate has validation-loss measurements only; no claim of better chat or coding accuracy. " |
| "See training_report.json and data_manifest.json. " |
| "Training data: Smol-Magpie-Ultra (Apache-2.0) and the Self-OSS-Instruct subset " |
| "(original: bigcode/self-oss-instruct-sc2-exec-filter-50k, ODC-BY). " |
| "Original base weights remain Apache-2.0; project code and adapter use GPL-3.0-only.\n", |
| encoding="utf-8") |
| for source in ["LICENSE", "LICENSE-BASE-MODEL.txt"]: |
| shutil.copy2(Path(__file__).resolve().parent.parent / source, destination / source) |
| print(f"Candidate saved to {destination}. Evaluate answers before publishing a release.") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="Train text-only Chat & Code on a cloud CUDA GPU") |
| parser.add_argument("--data", default="chatcode-data") |
| parser.add_argument("--output", required=True, help="Prefer mounted persistent cloud storage") |
| parser.add_argument("--max-steps", type=int, default=-1, help="25 for a smoke run; -1 for one epoch") |
| parser.add_argument("--resume-from-checkpoint") |
| args = parser.parse_args() |
| if args.max_steps != -1 and args.max_steps < 25: |
| parser.error("Use at least 25 steps so validation and checkpoint saving occur.") |
| train(args) |
|
|
