File size: 7,220 Bytes
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 | """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)
|