File size: 7,880 Bytes
5ba1bb4 ec9ed0b 5ba1bb4 ec9ed0b 5ba1bb4 ec9ed0b 5ba1bb4 ec9ed0b 5ba1bb4 ec9ed0b 5ba1bb4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """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.")
if "article_examples_sha256" in manifest:
current = hashlib.sha256(Path(__file__).with_name("article_examples.json").read_bytes()).hexdigest()
if manifest["article_examples_sha256"] != current:
raise ValueError("Article behavior examples changed after preparation.")
records, seen, groups = {}, 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)
group = row.get("group_id", key)
if group in groups and groups[group] != split:
raise ValueError("A source packet appears in multiple dataset splits.")
groups[group] = split
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_chat_code_and_article_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 & Writing candidate\n\nText-only QLoRA adapter. "
"This candidate has validation-loss measurements only; no claim of better chat, coding or factual article quality. "
"See training_report.json and data_manifest.json. "
"Training data: Smol-Magpie-Ultra, Smol-rewrite, Smol-summarize (Apache-2.0), authored fictional article behavior examples (GPL-3.0-only), 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)
|