From ae1779a9a83fc8f9f36019875efed522bf488b9c Mon Sep 17 00:00:00 2001 From: Giorgio Gilestro Date: Sun, 5 Jul 2026 20:27:53 +0100 Subject: [PATCH] =?UTF-8?q?specialise:=20extract=20train=5Flora=5Fon=5Ftas?= =?UTF-8?q?ks=20=E2=80=94=20reusable=20answer-only=20SFT=20primitive?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Factor the LoRA SFT loop out of train_specialist into train_lora_on_tasks(base, tasks, out_dir, ...), which fine-tunes a fresh adapter on an arbitrary list of (prompt, answer) pairs. train_specialist now generates its family's tasks and delegates. Backward-compatible (same behaviour for the merge/moe/directed experiments); the primitive is available for future loops that train on arbitrary or model-generated pairs. Co-Authored-By: Claude Opus 4.8 --- src/llm/specialise.py | 38 +++++++++++++++++++++++++++++++------- 1 file changed, 31 insertions(+), 7 deletions(-) diff --git a/src/llm/specialise.py b/src/llm/specialise.py index 13a26c8..aade829 100644 --- a/src/llm/specialise.py +++ b/src/llm/specialise.py @@ -27,15 +27,25 @@ def _encode(tok, task, device): return ids, labels -def train_specialist(base_name: str, family: str, out_dir: str, *, n_train: int = 600, - epochs: int = 3, lr: float = 2e-4, batch_size: int = 8, r: int = 16, - alpha: int = 32, seed: int = 0, device: str = "cuda", hard: bool = False) -> str: - """Fine-tune a LoRA specialist on ``family`` and save the adapter to ``out_dir``. +def train_lora_on_tasks(base_name: str, tasks: list, out_dir: str, *, epochs: int = 3, + lr: float = 2e-4, batch_size: int = 8, r: int = 16, alpha: int = 32, + seed: int = 0, device: str = "cuda") -> str: + """Fine-tune a fresh LoRA adapter on an arbitrary list of ``tasks`` and save it to ``out_dir``. + + The reusable answer-only SFT primitive: each task supplies a ``.prompt`` and a ``.answer`` (which + may be a *ground-truth* answer, for specialisation, or a *model-generated* one, for the + self-consumption loop). The prompt tokens are masked out of the loss. A new adapter is trained from + the frozen base each call (replace dynamics), so a lineage of generations is a pure function of the + data each generation supplies. Args: - hard (bool): train on the harder task variant (must match the eval difficulty). + base_name (str): HF id of the frozen base model. + tasks (list): objects with ``.prompt`` and ``.answer`` string attributes. + out_dir (str): where to save the adapter. + epochs, lr, batch_size, r, alpha, seed, device: standard SFT knobs. - Returns the adapter directory path. + Returns: + str: ``out_dir``. """ import torch from peft import LoraConfig, get_peft_model @@ -52,7 +62,6 @@ def train_specialist(base_name: str, family: str, out_dir: str, *, n_train: int model = get_peft_model(model, lora) model.train() - tasks = make_tasks(family, n_train, seed=seed, hard=hard) encoded = [_encode(tok, t, device) for t in tasks] opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr) rng = np.random.default_rng(seed) @@ -77,3 +86,18 @@ def train_specialist(base_name: str, family: str, out_dir: str, *, n_train: int Path(out_dir).mkdir(parents=True, exist_ok=True) model.save_pretrained(out_dir) return out_dir + + +def train_specialist(base_name: str, family: str, out_dir: str, *, n_train: int = 600, + epochs: int = 3, lr: float = 2e-4, batch_size: int = 8, r: int = 16, + alpha: int = 32, seed: int = 0, device: str = "cuda", hard: bool = False) -> str: + """Fine-tune a LoRA specialist on ``family`` and save the adapter to ``out_dir``. + + Args: + hard (bool): train on the harder task variant (must match the eval difficulty). + + Returns the adapter directory path. + """ + tasks = make_tasks(family, n_train, seed=seed, hard=hard) + return train_lora_on_tasks(base_name, tasks, out_dir, epochs=epochs, lr=lr, + batch_size=batch_size, r=r, alpha=alpha, seed=seed, device=device)