Training: 4-bit-aware LoRA (--qat4), 4-bit scoring (--round4), pipelines v0.5/v0.6, eval results
b6e964c verified | #!/usr/bin/env python3 | |
| """Fine-tune the auditor model with LoRA, fully offline, on CPU (or GPU if present). | |
| The base model must already be in a local folder (downloaded once, see training/README.md). | |
| No network access is used during training (HF_HUB_OFFLINE is forced on). | |
| Usage (from the endpoint-audit folder, with the training environment): | |
| training/.venv/bin/python training/train_lora.py --benchmark # time 10 steps, estimate the full run | |
| training/.venv/bin/python training/train_lora.py --max-train 3000 # train on a balanced subset | |
| training/.venv/bin/python training/train_lora.py --max-train 3000 --resume # continue after a stop | |
| ... --qat4 --init-adapter training/models/auditor-lora-v0.5/final # make a trained adapter survive 4-bit WebLLM builds | |
| --qat4 trains through the same 4-bit rounding the browser build uses (q4f16_1: groups of 32 along the input, scale | |
| max|w|/7 in float16), so the merged model still works once mlc_quantize.py rounds it. Without it, model v0.5 (0.5B) | |
| answered well in full precision and wrote junk in 4 bits. The rounding is applied to base + LoRA weights together on | |
| every forward pass, with the gradient passed straight through it. | |
| Loss is computed on the assistant answer only. Examples longer than --max-len tokens are skipped | |
| (never truncated), so the model always sees complete prompts and answers. | |
| """ | |
| import argparse | |
| import json | |
| import math | |
| import os | |
| import random | |
| import time | |
| os.environ["HF_HUB_OFFLINE"] = "1" | |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" | |
| import torch # noqa: E402 | |
| from peft import LoraConfig, PeftModel, get_peft_model # noqa: E402 | |
| from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402 | |
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| DEFAULT_BASE = os.path.join(ROOT, "training", "models", "Qwen2.5-0.5B-Instruct") | |
| LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] | |
| # PyTorch 2.2 (the last build for Intel Macs) returns wrong attention results with the "sdpa" kernel when | |
| # gradients are on and more than one thread is used, which drives training to NaN. "eager" is correct and faster here. | |
| ATTENTION = "eager" | |
| def round4(w): | |
| """w as the browser will see it: float16, then 4-bit groups of 32 (same as mlc_quantize.quantize, q4f16_1).""" | |
| out = w.shape[0] | |
| g = w.detach().half().float().reshape(out, -1, 32) | |
| scale = (g.abs().amax(-1, keepdim=True).half() * torch.tensor(1 / 7, dtype=torch.float16)).float() | |
| x = torch.where(scale == 0, torch.zeros_like(g), g / scale) | |
| q = (torch.sign(x) * torch.floor(x.abs() + 0.5)).clamp(-7, 7) | |
| return (q * scale).reshape(w.shape) | |
| def lora_forward_q4(layer): | |
| """Replace a PEFT LoRA linear's forward: merge, round to 4 bits, pass the gradient straight through.""" | |
| base, name = layer.base_layer, layer.active_adapters[0] | |
| lora_a, lora_b, scaling = layer.lora_A[name], layer.lora_B[name], layer.scaling[name] | |
| def forward(x, *args, **kwargs): | |
| w = base.weight + (lora_b.weight @ lora_a.weight) * scaling | |
| w = w + (round4(w) - w).detach() | |
| return torch.nn.functional.linear(x, w, base.bias) | |
| layer.forward = forward | |
| def use_4bit(model): | |
| """Train against 4-bit weights: embeddings (also the output layer: tied) once, LoRA layers on every pass.""" | |
| from peft.tuners.lora import Linear as LoraLinear | |
| emb = model.get_input_embeddings() | |
| with torch.no_grad(): | |
| emb.weight.copy_(round4(emb.weight)) | |
| layers = [m for m in model.modules() if isinstance(m, LoraLinear)] | |
| for layer in layers: | |
| lora_forward_q4(layer) | |
| print(f"4-bit-aware training: {len(layers)} LoRA layers + embeddings rounded as in q4f16_1") | |
| def load_examples(path): | |
| with open(path, encoding="utf-8") as handle: | |
| return [json.loads(line) for line in handle] | |
| def balanced_subset(examples, limit, seed): | |
| """Pick up to `limit` examples spread evenly over task, OS and status.""" | |
| if not limit or limit >= len(examples): | |
| return examples | |
| groups = {} | |
| for item in examples: | |
| meta = item["meta"] | |
| groups.setdefault((meta["task"], meta["os"], meta.get("status", "")), []).append(item) | |
| rng = random.Random(seed) | |
| for items in groups.values(): | |
| rng.shuffle(items) | |
| chosen, keys = [], sorted(groups) | |
| while len(chosen) < limit and keys: | |
| for key in list(keys): | |
| if groups[key]: | |
| chosen.append(groups[key].pop()) | |
| if len(chosen) == limit: | |
| break | |
| else: | |
| keys.remove(key) | |
| rng.shuffle(chosen) | |
| return chosen | |
| def encode(tokenizer, example, max_len): | |
| """Token ids with labels masked (-100) everywhere except the assistant answer.""" | |
| messages = example["messages"] | |
| prompt_ids = tokenizer.apply_chat_template(messages[:-1], tokenize=True, add_generation_prompt=True) | |
| full_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False) | |
| if full_ids[:len(prompt_ids)] != prompt_ids or len(full_ids) > max_len: | |
| return None | |
| labels = [-100] * len(prompt_ids) + full_ids[len(prompt_ids):] | |
| return {"input_ids": full_ids, "labels": labels} | |
| def batches(encoded, batch_size, pad_id): | |
| for start in range(0, len(encoded), batch_size): | |
| chunk = encoded[start:start + batch_size] | |
| width = max(len(e["input_ids"]) for e in chunk) | |
| ids = torch.full((len(chunk), width), pad_id, dtype=torch.long) | |
| labels = torch.full((len(chunk), width), -100, dtype=torch.long) | |
| mask = torch.zeros((len(chunk), width), dtype=torch.long) | |
| for row, item in enumerate(chunk): | |
| n = len(item["input_ids"]) | |
| ids[row, :n] = torch.tensor(item["input_ids"]) | |
| labels[row, :n] = torch.tensor(item["labels"]) | |
| mask[row, :n] = 1 | |
| yield ids, labels, mask | |
| def evaluate(model, encoded, pad_id, batch_size): | |
| model.eval() | |
| total, count = 0.0, 0 | |
| for ids, labels, mask in batches(encoded, batch_size, pad_id): | |
| total += model(input_ids=ids, attention_mask=mask, labels=labels).loss.item() | |
| count += 1 | |
| model.train() | |
| return total / max(count, 1) | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| parser.add_argument("--base", default=DEFAULT_BASE, help="local folder with the base model") | |
| parser.add_argument("--data", default=os.path.join(ROOT, "training", "out")) | |
| parser.add_argument("--output", default=os.path.join(ROOT, "training", "models", "auditor-lora")) | |
| parser.add_argument("--max-train", type=int, default=3000, help="balanced subset size (0 = all)") | |
| parser.add_argument("--max-val", type=int, default=200) | |
| parser.add_argument("--max-len", type=int, default=2048, help="skip examples longer than this many tokens") | |
| parser.add_argument("--epochs", type=float, default=1.0) | |
| parser.add_argument("--batch-size", type=int, default=1) | |
| parser.add_argument("--grad-accum", type=int, default=8) | |
| parser.add_argument("--lr", type=float, default=2e-4) | |
| parser.add_argument("--lora-r", type=int, default=16) | |
| parser.add_argument("--lora-alpha", type=int, default=32) | |
| parser.add_argument("--save-every", type=int, default=50, help="optimizer steps between checkpoints") | |
| parser.add_argument("--eval-every", type=int, default=100) | |
| parser.add_argument("--threads", type=int, default=os.cpu_count()) | |
| parser.add_argument("--seed", type=int, default=13) | |
| parser.add_argument("--benchmark", action="store_true", help="run 10 optimizer steps, report speed and estimated total time") | |
| parser.add_argument("--resume", action="store_true", help="continue from the last checkpoint in --output") | |
| parser.add_argument("--qat4", action="store_true", help="train through 4-bit rounding (for WebLLM q4 builds)") | |
| parser.add_argument("--init-adapter", help="start from this trained adapter instead of a fresh one") | |
| args = parser.parse_args() | |
| torch.manual_seed(args.seed) | |
| torch.set_num_threads(args.threads) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Device: {device}, threads: {args.threads}, base: {args.base}") | |
| tokenizer = AutoTokenizer.from_pretrained(args.base, local_files_only=True) | |
| pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id | |
| train_raw = balanced_subset(load_examples(os.path.join(args.data, "train.jsonl")), args.max_train, args.seed) | |
| val_raw = balanced_subset(load_examples(os.path.join(args.data, "validation.jsonl")), args.max_val, args.seed) | |
| train = [e for e in (encode(tokenizer, x, args.max_len) for x in train_raw) if e] | |
| val = [e for e in (encode(tokenizer, x, args.max_len) for x in val_raw) if e] | |
| # Group similar lengths to reduce padding, but keep task order mixed across groups. | |
| train.sort(key=lambda e: len(e["input_ids"])) | |
| groups = [train[i:i + args.batch_size] for i in range(0, len(train), args.batch_size)] | |
| random.Random(args.seed).shuffle(groups) | |
| train = [e for g in groups for e in g] | |
| tokens = sum(len(e["input_ids"]) for e in train) | |
| print(f"Train examples: {len(train)} ({len(train_raw) - len(train)} skipped as longer than {args.max_len} tokens), " | |
| f"{tokens:,} tokens; validation: {len(val)}") | |
| model = AutoModelForCausalLM.from_pretrained(args.base, local_files_only=True, torch_dtype=torch.float32, | |
| attn_implementation=ATTENTION) | |
| model.config.use_cache = False | |
| model.gradient_checkpointing_enable() | |
| model.enable_input_require_grads() | |
| last_checkpoint = os.path.join(args.output, "last") | |
| state_path = os.path.join(last_checkpoint, "trainer_state.json") | |
| if args.resume and os.path.exists(state_path): | |
| model = PeftModel.from_pretrained(model, last_checkpoint, is_trainable=True) | |
| with open(state_path) as handle: | |
| state = json.load(handle) | |
| print(f"Resuming from step {state['step']}") | |
| elif args.init_adapter: | |
| model = PeftModel.from_pretrained(model, args.init_adapter, is_trainable=True) | |
| state = {"step": 0, "micro": 0, "history": [], "init_adapter": args.init_adapter} | |
| print(f"Starting from adapter {args.init_adapter}") | |
| else: | |
| model = get_peft_model(model, LoraConfig(r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=0.05, | |
| target_modules=LORA_TARGETS, task_type="CAUSAL_LM")) | |
| state = {"step": 0, "micro": 0, "history": []} | |
| model.print_trainable_parameters() | |
| if args.qat4: | |
| use_4bit(model) | |
| state["qat4"] = True | |
| model.to(device) | |
| model.train() | |
| optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.lr, weight_decay=0.0) | |
| micro_per_epoch = math.ceil(len(train) / args.batch_size) | |
| total_steps = math.ceil(micro_per_epoch * args.epochs / args.grad_accum) | |
| if args.benchmark: | |
| total_steps = min(total_steps, state["step"] + 10) | |
| warmup = max(1, int(0.03 * total_steps)) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR( | |
| optimizer, lambda s: min(1.0, (s + 1) / warmup) * max(0.0, 1 - s / max(total_steps, 1)) if s >= warmup | |
| else (s + 1) / warmup) | |
| for _ in range(state["step"]): | |
| scheduler.step() | |
| def save(tag): | |
| path = os.path.join(args.output, tag) | |
| model.save_pretrained(path) | |
| tokenizer.save_pretrained(path) | |
| with open(os.path.join(path, "trainer_state.json"), "w") as handle: | |
| json.dump(state, handle, indent=1) | |
| batch_list = list(batches(train, args.batch_size, pad_id)) | |
| started = time.time() | |
| step_start = time.time() | |
| running = 0.0 | |
| bad_steps = 0 | |
| micro = state["micro"] | |
| while state["step"] < total_steps: | |
| ids, labels, mask = batch_list[micro % len(batch_list)] | |
| loss = model(input_ids=ids.to(device), attention_mask=mask.to(device), labels=labels.to(device)).loss | |
| (loss / args.grad_accum).backward() | |
| running += loss.item() | |
| micro += 1 | |
| if micro % args.grad_accum: | |
| continue | |
| grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) | |
| if not torch.isfinite(grad_norm): | |
| # Never apply a NaN/inf update: it would destroy every weight. | |
| optimizer.zero_grad() | |
| running = 0.0 | |
| bad_steps += 1 | |
| print(f"skipped a step with a non-finite gradient ({bad_steps} in a row)", flush=True) | |
| if bad_steps >= 3: | |
| raise SystemExit("Stopping: 3 non-finite gradients in a row. The last good checkpoint is kept.") | |
| continue | |
| bad_steps = 0 | |
| optimizer.step() | |
| scheduler.step() | |
| optimizer.zero_grad() | |
| state["step"] += 1 | |
| state["micro"] = micro | |
| seconds = time.time() - step_start | |
| step_start = time.time() | |
| avg = running / args.grad_accum | |
| running = 0.0 | |
| state["history"].append({"step": state["step"], "loss": round(avg, 4), "seconds": round(seconds, 1)}) | |
| # Estimate from recent steps only, so a pause (e.g. the Mac sleeping) does not distort it. | |
| recent = [h["seconds"] for h in state["history"][-10:] if h["seconds"] < 1800] or [seconds] | |
| remaining = (total_steps - state["step"]) * sum(recent) / len(recent) | |
| print(f"step {state['step']}/{total_steps} loss {avg:.4f} {seconds:.1f}s/step ~{remaining / 3600:.1f} h left", flush=True) | |
| if not args.benchmark and state["step"] % args.save_every == 0: | |
| save("last") | |
| save(f"step-{state['step']:04d}") # numbered copy: a later bad save can never overwrite it | |
| if not args.benchmark and val and state["step"] % args.eval_every == 0: | |
| val_loss = evaluate(model, val, pad_id, args.batch_size) | |
| state["history"][-1]["val_loss"] = round(val_loss, 4) | |
| print(f" validation loss {val_loss:.4f}", flush=True) | |
| if args.benchmark: | |
| per_step = sum(h["seconds"] for h in state["history"][-10:]) / max(len(state["history"][-10:]), 1) | |
| full = math.ceil(micro_per_epoch * args.epochs / args.grad_accum) | |
| print(f"\nBENCHMARK: {per_step:.1f} s per optimizer step ({args.grad_accum} examples). " | |
| f"Full run of {full} steps on {len(train)} examples: about {full * per_step / 3600:.1f} hours.") | |
| return | |
| save("last") | |
| save("final") | |
| if val: | |
| print(f"Final validation loss: {evaluate(model, val, pad_id, args.batch_size):.4f}") | |
| print(f"Adapter saved to {os.path.join(args.output, 'final')}") | |
| if __name__ == "__main__": | |
| main() | |