| """LoRA / QLoRA fine-tuning on Sinhala MCQs. |
| |
| Training data = jsonl in the same format as the dev set (question, options, |
| answer, level). The prompt is built by the SAME code as inference, and the |
| loss is only on the answer token - exactly what predict.py scores. |
| |
| python train.py --config configs/train_qlora_qwen25_7b.yaml |
| |
| After training, predict with the adapter: |
| model.adapter_path: outputs/<name>/adapter (in a predict config) |
| |
| Rules to remember: |
| * never train on test questions (rule 04) |
| * LoRA parameters count towards the 8B limit (base + adapter) |
| * data generated with closed APIs (GPT/Claude/Gemini) must be disclosed |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import shutil |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| from sinmmlu.config import load_config |
| from sinmmlu.data import has_gold, load_questions |
| from sinmmlu.metrics import accuracy_by_level, format_report |
| from sinmmlu.modeling import describe_budget, load_model, load_tokenizer, pick_dtype, prepare_for_kbit_training, set_determinism |
| from sinmmlu.prompts import build_messages, render_prompt |
| from sinmmlu.runinfo import append_result, write_run_info |
| from sinmmlu.scoring import Scorer, label_token_id |
|
|
|
|
| def build_examples(records: list[dict], tok, cfg: dict) -> list[dict]: |
| """One example = prompt tokens + the correct label token. |
| |
| Options are rotated (deterministically) so the right answer appears in |
| every position equally often - otherwise the model can learn "the answer |
| is usually 2" instead of learning the content. |
| """ |
| p = cfg["prompt"] |
| max_len = cfg["inference"]["max_length"] |
| rotate = cfg["train"].get("rotate_options", True) |
| examples, skipped = [], 0 |
| for i, r in enumerate(records): |
| n = len(r["options"]) |
| shift = (i % n) if rotate else 0 |
| options = r["options"][shift:] + r["options"][:shift] |
| gold_pos = (r["gold"] - shift) % n |
| msgs = build_messages(r, p["template"], None, options) |
| text, prefix = render_prompt(tok, msgs, p["use_chat_template"], p["chat_template_kwargs"]) |
| ids = tok(text, add_special_tokens=False)["input_ids"] |
| if len(ids) + 1 > max_len: |
| skipped += 1 |
| continue |
| target = label_token_id(tok, text, prefix + str(gold_pos + 1)) |
| examples.append({"input_ids": ids + [target], "labels": [-100] * len(ids) + [target]}) |
| if skipped: |
| print(f"[train] skipped {skipped} examples longer than {max_len} tokens") |
| return examples |
|
|
|
|
| class Collator: |
| def __init__(self, pad_id: int): |
| self.pad_id = pad_id |
|
|
| def __call__(self, batch: list[dict]) -> dict: |
| width = max(len(b["input_ids"]) for b in batch) |
| ids, labels, mask = [], [], [] |
| for b in batch: |
| pad = width - len(b["input_ids"]) |
| ids.append([self.pad_id] * pad + b["input_ids"]) |
| labels.append([-100] * pad + b["labels"]) |
| mask.append([0] * pad + [1] * len(b["input_ids"])) |
| return { |
| "input_ids": torch.tensor(ids), |
| "labels": torch.tensor(labels), |
| "attention_mask": torch.tensor(mask), |
| } |
|
|
|
|
| def main(argv=None) -> int: |
| from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training |
| from transformers import Trainer, TrainingArguments |
|
|
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--config", required=True) |
| ap.add_argument("--who", default="") |
| ap.add_argument("--max-steps", type=int, default=-1, help="stop early (debugging)") |
| args = ap.parse_args(argv) |
|
|
| cfg = load_config(args.config) |
| tc = cfg["train"] |
| set_determinism(cfg["seed"]) |
| if not tc["train_file"]: |
| sys.exit("train.train_file is required") |
|
|
| train_records = load_questions(tc["train_file"], cfg["data"]["answer_base"]) |
| if not has_gold(train_records): |
| sys.exit("training file has no gold answers") |
| print(f"[train] {len(train_records)} training questions from {tc['train_file']}") |
|
|
| tok = load_tokenizer(cfg["model"]) |
| model = load_model(cfg["model"], for_training=True) |
| if cfg["model"]["load_in_4bit"] and cfg["model"].get("text_only"): |
| |
| model = prepare_for_kbit_training(model, use_gradient_checkpointing=tc["gradient_checkpointing"]) |
| elif cfg["model"]["load_in_4bit"]: |
| model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=tc["gradient_checkpointing"]) |
| elif tc["gradient_checkpointing"]: |
| model.gradient_checkpointing_enable() |
| model.enable_input_require_grads() |
|
|
| lora = LoraConfig( |
| r=tc["lora_r"], |
| lora_alpha=tc["lora_alpha"], |
| lora_dropout=tc["lora_dropout"], |
| target_modules=tc["target_modules"], |
| task_type="CAUSAL_LM", |
| ) |
| model = get_peft_model(model, lora) |
| trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| print(f"[lora] adapter parameters: {trainable:,}") |
|
|
| examples = build_examples(train_records, tok, cfg) |
| steps_per_epoch = math.ceil(len(examples) / (tc["batch_size"] * tc["grad_accum"])) |
| total_steps = args.max_steps if args.max_steps > 0 else steps_per_epoch * tc["epochs"] |
| out_dir = Path(tc["output_dir"]) |
| dtype = pick_dtype(cfg["model"]["dtype"]) |
| on_gpu = torch.cuda.is_available() |
| targs = TrainingArguments( |
| output_dir=str(out_dir / "checkpoints"), |
| num_train_epochs=tc["epochs"], |
| max_steps=args.max_steps, |
| learning_rate=tc["learning_rate"], |
| per_device_train_batch_size=tc["batch_size"], |
| gradient_accumulation_steps=tc["grad_accum"], |
| lr_scheduler_type="cosine", |
| warmup_steps=math.ceil(0.03 * total_steps), |
| logging_steps=tc["logging_steps"], |
| save_strategy="no", |
| report_to=[], |
| seed=cfg["seed"], |
| data_seed=cfg["seed"], |
| bf16=on_gpu and dtype == torch.bfloat16, |
| fp16=on_gpu and dtype == torch.float16, |
| remove_unused_columns=False, |
| dataloader_num_workers=0, |
| ) |
| trainer = Trainer(model=model, args=targs, train_dataset=examples, data_collator=Collator(tok.pad_token_id)) |
| trainer.train() |
|
|
| adapter_dir = out_dir / "adapter" |
| model.save_pretrained(adapter_dir) |
| tok.save_pretrained(adapter_dir) |
| shutil.copy(args.config, out_dir / "train_config.yaml") |
| total = sum(p.numel() for p in model.parameters()) |
| print(f"[save] adapter -> {adapter_dir}") |
| print(f"[budget] base + adapter as loaded: {describe_budget(total)}") |
|
|
| extra = {"train_file": tc["train_file"], "n_train": len(examples), "adapter_params": trainable} |
| if tc["eval_file"]: |
| model.eval() |
| eval_records = load_questions(tc["eval_file"], cfg["data"]["answer_base"]) |
| preds = Scorer(model, tok, cfg).predict(eval_records) |
| result = accuracy_by_level([r["level"] for r in eval_records], preds, [r["gold"] for r in eval_records]) |
| print(f"[eval] {tc['eval_file']}\n" + format_report(result)) |
| cfg_logged = json.loads(json.dumps(cfg)) |
| cfg_logged["model"]["adapter_path"] = str(adapter_dir) |
| append_result(cfg_logged, tc["eval_file"], result, args.who, "after training") |
| extra["eval_result"] = result |
| write_run_info(out_dir / "train.run.json", cfg, extra) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
