llm_lora.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""LLM LoRA / QLoRA adapter on the same Recipe / Run spine."""
 
from __future__ import annotations
 
import json
from pathlib import Path
from typing import Any
 
from aquin.recipe.schema import Checkpoint, EvalGate, Recipe, RecipeError
from aquin.run import Run
 
 
def _read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open(encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            if not isinstance(row, dict):
                raise RecipeError(f"jsonl rows must be objects, got {type(row).__name__}")
            rows.append(row)
    if not rows:
        raise RecipeError(f"No rows in {path}")
    return rows
 
 
def _row_text(row: dict[str, Any], text_field: str) -> str:
    if text_field in row and row[text_field] is not None:
        return str(row[text_field])
    if "prompt" in row and "completion" in row:
        return f"{row['prompt']}\n{row['completion']}"
    if "messages" in row and isinstance(row["messages"], list):
        parts = []
        for m in row["messages"]:
            if isinstance(m, dict):
                parts.append(f"{m.get('role', 'user')}: {m.get('content', '')}")
        return "\n".join(parts)
    raise RecipeError(
        f"Each jsonl row needs {text_field!r}, or prompt+completion, or messages."
    )
 
 
def _overlap(pred: str, reference: str) -> float:
    ref = [t for t in reference.lower().split() if t]
    if not ref:
        return 1.0
    got = set(pred.lower().split())
    return sum(1 for t in ref if t in got) / len(ref)
 
 
def train_llm(recipe: Recipe, recorder: Run, out_dir: Path) -> dict[str, Any]:
    try:
        import torch
        from datasets import Dataset
        from peft import LoraConfig, TaskType, get_peft_model
        from transformers import (
            AutoModelForCausalLM,
            AutoTokenizer,
            DataCollatorForLanguageModeling,
            Trainer,
            TrainingArguments,
        )
    except ImportError as exc:
        raise RecipeError(
            "LLM LoRA train needs transformers, peft, datasets, and torch."
        ) from exc
 
    data_path = Path(str(recipe.data["path"]))
    rows = _read_jsonl(data_path)
    text_field = str(recipe.data.get("text_field") or "text")
    texts = [_row_text(r, text_field) for r in rows]
 
    base = str(recipe.train["base"])
    rank = int(recipe.train.get("rank") or 8)
    lora_alpha = int(recipe.train.get("lora_alpha") or 16)
    lr = float(recipe.train.get("lr") or 2e-4)
    epochs = float(recipe.train.get("epochs") or 1)
    max_steps = recipe.train.get("max_steps")
    max_seq = int(recipe.train.get("max_seq_len") or 512)
    batch = int(recipe.train.get("batch_size") or 1)
    accum = int(recipe.train.get("grad_accum") or 4)
 
    tokenizer = AutoTokenizer.from_pretrained(base, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
 
    model = AutoModelForCausalLM.from_pretrained(base)
    model.config.pad_token_id = tokenizer.pad_token_id
    lora_kwargs: dict[str, Any] = dict(
        task_type=TaskType.CAUSAL_LM,
        r=rank,
        lora_alpha=lora_alpha,
        lora_dropout=float(recipe.train.get("dropout") or 0.05),
    )
    modules = recipe.train.get("target_modules")
    if modules:
        lora_kwargs["target_modules"] = modules
    model = get_peft_model(model, LoraConfig(**lora_kwargs))
 
    ds = Dataset.from_dict({"text": texts})
 
    def tokenize(batch: dict[str, list[str]]) -> dict[str, Any]:
        return tokenizer(
            batch["text"],
            truncation=True,
            max_length=max_seq,
            padding=False,
        )
 
    tokenized = ds.map(tokenize, batched=True, remove_columns=["text"])
    ckpt_dir = out_dir / "checkpoints" / "lora"
    ckpt_dir.mkdir(parents=True, exist_ok=True)
 
    args = TrainingArguments(
        output_dir=str(out_dir / "hf_trainer"),
        per_device_train_batch_size=batch,
        gradient_accumulation_steps=accum,
        learning_rate=lr,
        num_train_epochs=epochs,
        max_steps=int(max_steps) if max_steps is not None else -1,
        logging_steps=1,
        save_strategy="no",
        report_to=[],
        fp16=bool(torch.cuda.is_available()),
        bf16=False,
        remove_unused_columns=False,
    )
    collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
    losses: list[float] = []
 
    class _Log(Trainer):
        def log(self, logs: dict[str, float], *rest: Any, **kwargs: Any) -> None:  # type: ignore[override]
            super().log(logs, *rest, **kwargs)
            if "loss" in logs:
                losses.append(float(logs["loss"]))
                recorder.log(int(self.state.global_step or 0), loss=float(logs["loss"]))
 
    trainer = _Log(
        model=model,
        args=args,
        train_dataset=tokenized,
        data_collator=collator,
    )
    trainer.train()
    model.save_pretrained(str(ckpt_dir))
    tokenizer.save_pretrained(str(ckpt_dir))
    recorder.signal(int(trainer.state.global_step or 1), "lora checkpoint saved")
 
    mean_loss = float(sum(losses) / len(losses)) if losses else None
    metrics: dict[str, Any] = {"train_loss": mean_loss, "n_rows": len(texts)}
    inspect = {"n_rows": len(texts), "base": base, "rank": rank}
 
    probes_path = recipe.eval.get("probes")
    min_score = recipe.eval.get("min_score")
    details: dict[str, Any] = {"train_loss": mean_loss}
    if probes_path:
        probes = _read_jsonl(Path(str(probes_path)))
        model.eval()
        scores: list[float] = []
        rows_out: list[dict[str, Any]] = []
        max_new = int(recipe.eval.get("max_tokens") or 32)
        for p in probes:
            prompt = str(p.get("prompt") or p.get("text") or "")
            ref = str(p.get("reference") or p.get("completion") or "")
            inputs = tokenizer(prompt, return_tensors="pt")
            inputs = {k: v.to(model.device) for k, v in inputs.items()}
            with torch.no_grad():
                out = model.generate(**inputs, max_new_tokens=max_new, do_sample=False)
            pred = tokenizer.decode(out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True)
            sc = _overlap(pred, ref) if ref else 0.0
            scores.append(sc)
            rows_out.append({"prompt": prompt, "response": pred, "score": sc})
        score = float(sum(scores) / len(scores)) if scores else 0.0
        metric = str(recipe.eval.get("metric") or "overlap")
        skipped = min_score is None
        passed = True if skipped else score >= float(min_score)
        details["probes"] = rows_out
        gate = EvalGate(
            passed=passed,
            metric=metric,
            score=round(score, 6),
            min_score=None if min_score is None else float(min_score),
            skipped=skipped,
            details=details,
        )
        metrics["eval_overlap"] = score
    else:
        metric = "train_loss"
        skipped = min_score is None
        if skipped or mean_loss is None:
            passed = True
            score = mean_loss
        else:
            # lower is better when gating on train loss
            score = mean_loss
            passed = mean_loss <= float(min_score)
        gate = EvalGate(
            passed=passed,
            metric=metric,
            score=None if score is None else round(float(score), 6),
            min_score=None if min_score is None else float(min_score),
            skipped=skipped,
            details=details,
        )
 
    return {
        "checkpoint": Checkpoint(
            path=str(ckpt_dir),
            kind="peft",
            step=int(trainer.state.global_step or 1),
            extra={"base": base, "rank": rank},
        ),
        "gate": gate,
        "metrics": metrics,
        "inspect": inspect,
    }