mlxsmith 0.1.0__py3-none-any.whl

mlxsmith/train/lora.py

@@ -0,0 +1,280 @@
+"""LoRA adapter utilities.
+
+This module prefers MLX-LM's LoRA utilities and adapter format when available.
+Fallback implementations are provided for environments without MLX/MLX-LM.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+def _require_mlx():
+    import mlx.core as mx  # type: ignore
+    import mlx.nn as nn  # type: ignore
+    from mlx.utils import tree_flatten  # type: ignore
+
+    return mx, nn, tree_flatten
+
+
+def _try_mlx_lm_utils():
+    try:
+        from mlx_lm.tuner import utils as tuner_utils  # type: ignore
+        from mlx_lm.utils import save_config as mlx_save_config  # type: ignore
+    except Exception:
+        return None, None
+    return tuner_utils, mlx_save_config
+
+
+@dataclass
+class LoRAConfig:
+    r: int = 16
+    alpha: int = 32
+    dropout: float = 0.0
+    target_modules: list[str] | None = None
+    num_layers: int = 0  # 0 => all layers
+    scale: float | None = None
+    fine_tune_type: str = "lora"  # lora | dora | full
+
+
+class LoRALinear:
+    """Minimal fallback LoRA wrapper for mlx.nn.Linear."""
+
+    def __init__(self, base_linear: Any, *, r: int, alpha: int, dropout: float = 0.0):
+        mx, nn, _ = _require_mlx()
+        self.nn = nn
+        self.mx = mx
+
+        self.base = base_linear
+        self.r = int(r)
+        self.alpha = float(alpha)
+        self.scale = float(alpha) / float(r) if r > 0 else 0.0
+        self.dropout = nn.Dropout(p=float(dropout)) if dropout and dropout > 0 else None
+
+        w = getattr(base_linear, "weight")
+        out_dim, in_dim = int(w.shape[0]), int(w.shape[1])
+
+        self.A = mx.random.normal((self.r, in_dim)) * 0.01
+        self.B = mx.zeros((out_dim, self.r))
+
+    def __call__(self, x):
+        mx = self.mx
+        base_w = mx.stop_gradient(self.base.weight)
+        y = x @ base_w.T
+        if getattr(self.base, "bias", None) is not None:
+            y = y + mx.stop_gradient(self.base.bias)
+        z = x
+        if self.dropout is not None:
+            z = self.dropout(z)
+        z = (z @ self.A.T) @ self.B.T
+        return y + z * self.scale
+
+
+def _iter_named_modules(root: Any, prefix: str = ""):
+    for name in dir(root):
+        if name.startswith("_"):
+            continue
+        try:
+            obj = getattr(root, name)
+        except Exception:
+            continue
+        if callable(getattr(obj, "__call__", None)) and hasattr(obj, "__dict__"):
+            full = f"{prefix}{name}" if not prefix else f"{prefix}.{name}"
+            yield full, root, name, obj
+
+
+def inject_lora(
+    model: Any,
+    *,
+    r: int = 16,
+    alpha: int = 32,
+    dropout: float = 0.0,
+    target_modules: list[str] | None = None,
+) -> int:
+    _mx, _nn, _ = _require_mlx()
+    targets = target_modules or ["q_proj", "k_proj", "v_proj", "o_proj"]
+    wrapped = 0
+
+    for full, parent, attr, obj in list(_iter_named_modules(model)):
+        if not hasattr(obj, "weight"):
+            continue
+        if not any(attr.endswith(t) or full.endswith(t) for t in targets):
+            continue
+        cls_name = obj.__class__.__name__.lower()
+        if "linear" not in cls_name:
+            continue
+        try:
+            setattr(parent, attr, LoRALinear(obj, r=r, alpha=alpha, dropout=dropout))
+            wrapped += 1
+        except Exception:
+            continue
+
+    return wrapped
+
+
+def lora_parameters(model: Any) -> dict[str, Any]:
+    params: dict[str, Any] = {}
+    for full, _parent, _attr, obj in list(_iter_named_modules(model)):
+        if obj.__class__.__name__ == "LoRALinear" or (hasattr(obj, "A") and hasattr(obj, "B")):
+            try:
+                params[f"{full}.A"] = obj.A
+                params[f"{full}.B"] = obj.B
+            except Exception:
+                pass
+    return params
+
+
+def _scale_for_config(cfg: LoRAConfig) -> float:
+    if cfg.scale is not None:
+        return float(cfg.scale)
+    if cfg.r == 0:
+        return float(cfg.alpha)
+    return float(cfg.alpha) / float(cfg.r)
+
+
+def _keys_for_target_modules(model: Any, target_modules: list[str]) -> set[str]:
+    keys: set[str] = set()
+    # Prefer model.named_modules if present (MLX-LM models support this)
+    if hasattr(model, "named_modules"):
+        for name, _mod in model.named_modules():
+            if any(name.endswith(t) for t in target_modules):
+                keys.add(name)
+        return keys
+
+    # Fallback: walk attributes
+    for full, _parent, _attr, obj in list(_iter_named_modules(model)):
+        if any(full.endswith(t) for t in target_modules):
+            keys.add(full)
+    return keys
+
+
+def apply_lora(model: Any, cfg: LoRAConfig) -> dict:
+    """Apply LoRA layers (prefer MLX-LM utilities) and return adapter config."""
+    tuner_utils, _save_cfg = _try_mlx_lm_utils()
+    scale = _scale_for_config(cfg)
+    keys = None
+    if cfg.target_modules:
+        keys = sorted(_keys_for_target_modules(model, cfg.target_modules))
+
+    if tuner_utils is not None and hasattr(tuner_utils, "linear_to_lora_layers"):
+        # MLX-LM format
+        config = {
+            "rank": int(cfg.r),
+            "scale": float(scale),
+            "dropout": float(cfg.dropout),
+        }
+        if keys:
+            config["keys"] = keys
+        num_layers = int(cfg.num_layers)
+        tuner_utils.linear_to_lora_layers(
+            model,
+            num_layers,
+            config,
+            use_dora=(cfg.fine_tune_type == "dora"),
+        )
+        return {
+            "fine_tune_type": cfg.fine_tune_type,
+            "num_layers": num_layers,
+            "lora_parameters": config,
+        }
+
+    # Fallback to local LoRA injection
+    inject_lora(
+        model,
+        r=cfg.r,
+        alpha=cfg.alpha,
+        dropout=cfg.dropout,
+        target_modules=cfg.target_modules,
+    )
+    return {
+        "fine_tune_type": "lora",
+        "num_layers": 0,
+        "lora_parameters": {
+            "rank": int(cfg.r),
+            "scale": float(scale),
+            "dropout": float(cfg.dropout),
+        },
+    }
+
+
+def save_adapter(model: Any, out_dir: str | Path, *, adapter_config: dict, metadata: dict | None = None) -> None:
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    tuner_utils, mlx_save_config = _try_mlx_lm_utils()
+    try:
+        mx, _nn, tree_flatten = _require_mlx()
+    except Exception:
+        mx = None
+        tree_flatten = None
+
+    if mx is not None and hasattr(model, "trainable_parameters") and tree_flatten is not None:
+        adapter_weights = dict(tree_flatten(model.trainable_parameters()))
+        try:
+            mx.save_safetensors(str(out / "adapters.safetensors"), adapter_weights)
+        except Exception:
+            # fallback to numpy
+            import numpy as np
+
+            arrays = {k: mx.array(v).to_numpy() for k, v in adapter_weights.items()}
+            np.savez(out / "lora.npz", **arrays)
+    else:
+        # fallback to local LoRA params
+        params = lora_parameters(model)
+        if params:
+            import numpy as np
+
+            mx_local, _nn_local, _ = _require_mlx()
+            arrays = {k: mx_local.array(v).to_numpy() for k, v in params.items()}
+            np.savez(out / "lora.npz", **arrays)
+
+    config_path = out / "adapter_config.json"
+    if mlx_save_config is not None:
+        cfg = dict(adapter_config)
+        mlx_save_config(cfg, config_path)
+    else:
+        config_path.write_text(json.dumps(adapter_config, indent=2), encoding="utf-8")
+
+    if metadata is not None:
+        (out / "adapter_metadata.json").write_text(
+            json.dumps(metadata, indent=2), encoding="utf-8"
+        )
+
+
+def load_adapter_config(adapter_dir: str | Path) -> dict | None:
+    path = Path(adapter_dir) / "adapter_config.json"
+    if not path.exists():
+        return None
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def apply_adapter(model: Any, adapter_dir: str | Path) -> dict | None:
+    adapter_dir = Path(adapter_dir)
+    adapter_cfg = load_adapter_config(adapter_dir)
+    tuner_utils, _ = _try_mlx_lm_utils()
+    if adapter_cfg is not None and tuner_utils is not None and hasattr(tuner_utils, "load_adapters"):
+        tuner_utils.load_adapters(model, str(adapter_dir))
+        return adapter_cfg
+
+    # Fallback: load local lora.npz into LoRALinear wrappers
+    lora_file = adapter_dir / "lora.npz"
+    if not lora_file.exists():
+        return adapter_cfg
+
+    import numpy as np
+
+    weights = dict(np.load(lora_file))
+    mx, _nn, _ = _require_mlx()
+    # apply
+    for full, _parent, _attr, obj in list(_iter_named_modules(model)):
+        if hasattr(obj, "A") and hasattr(obj, "B"):
+            key_a = f"{full}.A"
+            key_b = f"{full}.B"
+            if key_a in weights:
+                obj.A = mx.array(weights[key_a])
+            if key_b in weights:
+                obj.B = mx.array(weights[key_b])
+    return adapter_cfg

mlxsmith/train/pref.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+
+from rich.console import Console
+
+from ..accel import get_backend
+from ..config import ProjectConfig
+from ..models import resolve_model_spec
+from ..runs import RunPaths, new_run, snapshot_config
+from ..util import write_jsonl, now_ts, tree_add, tree_scale
+from ..llm.registry import get_llm_backend
+from ..llm.backend import BackendNotAvailable
+from .lora import LoRAConfig
+
+console = Console()
+
+
+def _load_pref_rows(path: Path) -> list[dict]:
+    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
+
+
+def run_pref(project_root: Path, cfg: ProjectConfig, data_dir: Path, base_model_path: Path, accel: str) -> RunPaths:
+    run = new_run(project_root, "pref")
+    snapshot_config(cfg.model_dump(), run.config_snapshot_path)
+
+    backend = get_backend(accel)
+    backend.patch()
+    console.print(f"[bold]PREF[/bold] run: {run.run_dir.name} algo={cfg.pref.algo} accel={backend.name}")
+
+    prefs_path = data_dir / "train.jsonl"
+    if not prefs_path.exists():
+        raise RuntimeError("Preference data missing. Expect data/prefs/train.jsonl with {prompt, chosen, rejected}")
+    rows = _load_pref_rows(prefs_path)
+
+    llm = get_llm_backend(cfg.model.backend)
+    base_model, adapter_path, adapter_meta = resolve_model_spec(project_root, str(base_model_path), cfg)
+
+    try:
+        llm.load(
+            base_model,
+            max_seq_len=cfg.model.max_seq_len,
+            dtype=cfg.model.dtype,
+            trust_remote_code=cfg.model.trust_remote_code,
+        )
+        if adapter_path:
+            llm.apply_adapter(str(adapter_path))
+        else:
+            lora_cfg = LoRAConfig(
+                r=cfg.lora.r,
+                alpha=cfg.lora.alpha,
+                dropout=cfg.lora.dropout,
+                target_modules=list(cfg.lora.target_modules or []),
+                num_layers=cfg.lora.num_layers,
+                scale=cfg.lora.scale,
+                fine_tune_type=cfg.lora.fine_tune_type,
+            )
+            llm.apply_lora_from_config(lora_cfg)
+    except BackendNotAvailable as e:
+        console.print(f"[yellow]MLX backend unavailable[/yellow]: {e}")
+        (run.adapter_dir / "ADAPTER.txt").write_text(
+            f"Backend unavailable in this environment.\nbase={base_model}\naccel={backend.name}\n",
+            encoding="utf-8",
+        )
+        return run
+
+    ref_llm = None
+    if cfg.pref.reference_model:
+        ref_llm = get_llm_backend(cfg.model.backend)
+        try:
+            ref_llm.load(
+                cfg.pref.reference_model,
+                max_seq_len=cfg.model.max_seq_len,
+                dtype=cfg.model.dtype,
+                trust_remote_code=cfg.model.trust_remote_code,
+            )
+        except BackendNotAvailable:
+            ref_llm = None
+
+    opt, _params = llm.optimizer_and_params(lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
+
+    beta = float(cfg.pref.beta)
+    kl_coeff = float(cfg.pref.kl_coeff)
+    rng = random.Random(cfg.train.seed)
+    total = int(cfg.train.iters)
+    grad_accum = max(1, int(cfg.train.grad_accum))
+    train_on_prompt = bool(getattr(cfg.train, "train_on_prompt", False))
+
+    accum_grads = None
+    for step in range(1, total + 1):
+        row = rng.choice(rows)
+        prompt = row.get("prompt") or ""
+        chosen = row.get("chosen") or row.get("accepted") or ""
+        rejected = row.get("rejected") or row.get("rejected_response") or ""
+        if not (prompt and chosen and rejected):
+            continue
+
+        prompt_ids = llm.encode(prompt)
+        chosen_ids = llm.encode(prompt + chosen)
+        rejected_ids = llm.encode(prompt + rejected)
+        p_len_c = len(prompt_ids)
+        p_len_r = len(prompt_ids)
+        max_len = int(cfg.model.max_seq_len)
+        if max_len:
+            if len(chosen_ids) > max_len:
+                overflow = len(chosen_ids) - max_len
+                chosen_ids = chosen_ids[overflow:]
+                p_len_c = max(0, p_len_c - overflow)
+            if len(rejected_ids) > max_len:
+                overflow = len(rejected_ids) - max_len
+                rejected_ids = rejected_ids[overflow:]
+                p_len_r = max(0, p_len_r - overflow)
+
+        def loss_fn(_model):
+            logp_c = llm.sequence_logprob(chosen_ids, prompt_len=p_len_c)
+            logp_r = llm.sequence_logprob(rejected_ids, prompt_len=p_len_r)
+            ref_diff = 0.0
+            if ref_llm is not None:
+                ref_logp_c = ref_llm.sequence_logprob(chosen_ids, prompt_len=p_len_c)
+                ref_logp_r = ref_llm.sequence_logprob(rejected_ids, prompt_len=p_len_r)
+                ref_diff = ref_logp_c - ref_logp_r
+            diff = (logp_c - logp_r) - ref_diff
+
+            if cfg.pref.algo == "orpo":
+                # ORPO loss = NLL(chosen) - beta * log(sigmoid(diff))
+                nll = llm.sft_loss(chosen_ids, train_on_prompt=train_on_prompt, prompt_len=p_len_c)
+                or_loss = -beta * llm.mx.log(llm.mx.sigmoid(diff))  # type: ignore
+                loss = nll + or_loss
+            else:
+                # DPO loss
+                scaled = llm.mx.array(beta) * diff  # type: ignore
+                loss = llm.mx.log1p(llm.mx.exp(-scaled))  # type: ignore
+
+            if ref_llm is not None and kl_coeff > 0:
+                # Simple KL penalty on chosen responses
+                kl = (logp_c - ref_logp_c) if ref_llm is not None else 0.0
+                loss = loss + llm.mx.array(kl_coeff) * kl  # type: ignore
+            return loss
+
+        lval, grads = llm.value_and_grad(loss_fn)
+        if grads is not None:
+            accum_grads = tree_add(accum_grads, grads)
+
+        if step % grad_accum == 0:
+            if accum_grads is not None:
+                llm.apply_grads(opt, tree_scale(accum_grads, 1.0 / grad_accum))
+            accum_grads = None
+
+        if step % cfg.train.log_every == 0 or step == 1 or step == total:
+            write_jsonl(
+                run.metrics_path,
+                [
+                    {
+                        "ts": now_ts(),
+                        "step": step,
+                        "kind": "pref",
+                        "algo": cfg.pref.algo,
+                        "beta": beta,
+                        "kl_coeff": kl_coeff,
+                        "loss": float(lval.item()) if hasattr(lval, "item") else float(lval),
+                        "accel": backend.name,
+                    }
+                ],
+            )
+
+        if step % cfg.train.save_every == 0 or step == total:
+            llm.save_adapter(
+                str(run.adapter_dir),
+                metadata={
+                    "base_model": base_model,
+                    "source_adapter": str(adapter_path) if adapter_path else None,
+                    "run": run.run_dir.name,
+                    "kind": "pref",
+                },
+            )
+
+    console.print(f"[green]Saved adapter[/green] {run.adapter_dir}")
+    return run