shadowlm 0.1.0__py3-none-any.whl

shadowlm/__init__.py

@@ -0,0 +1,43 @@
+"""ShadowLM Trainer — a fine-tuning SDK.
+
+Any open model. Any harness. Any method.
+
+    import shadowlm as slm
+
+    ds    = slm.Dataset.from_jsonl("data.jsonl").as_chat()
+    model = slm.load("mlx-community/Qwen2.5-0.5B-Instruct-4bit")
+    run   = model.finetune(ds, method="lora", max_steps=60)
+    print(run.loss, run.sparkline())
+    print(model.generate("Hello!"))
+    model.save("out/", fmt="adapter")
+
+datasets → finetune → inference. mlx on Apple Silicon, torch on CUDA (or CPU)
+— accelerated by the shadow layer.
+"""
+
+from . import methods, runs
+from .capture import CaptureProxy, capture
+from .data import Dataset
+from .models import Model, Reply, load
+from .rl import Trajectory, TrajectoryGroup, judge_group
+from .training import Metric, TrainConfig, TrainingRun
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "CaptureProxy",
+    "capture",
+    "Dataset",
+    "Model",
+    "Reply",
+    "load",
+    "methods",
+    "runs",
+    "Metric",
+    "TrainConfig",
+    "TrainingRun",
+    "Trajectory",
+    "TrajectoryGroup",
+    "judge_group",
+    "__version__",
+]

shadowlm/_quiet.py

@@ -0,0 +1,23 @@
+"""Swallow a backend's own stdout/stderr so shadowLM owns the console.
+
+Backends (mlx-lm, transformers, huggingface_hub) print their own progress lines
+and tqdm bars. shadowLM prints its own clean output to the real terminal
+(`sys.__stdout__`, captured in models.py), which these redirects don't touch — so
+we silence the backend chatter without losing our own. `SHADOWLM_DEBUG=1` shows
+the raw output.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+
+
+@contextlib.contextmanager
+def quiet_backend():
+    if os.environ.get("SHADOWLM_DEBUG"):
+        yield
+        return
+    with open(os.devnull, "w") as devnull, \
+            contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(devnull):
+        yield

shadowlm/accel.py

@@ -0,0 +1,70 @@
+"""The shadow accelerator — shadowLM's in-house training optimization layer.
+
+It sits on top of whichever backend is active (mlx / torch) and turns on the
+memory and throughput optimizations that are *safe for the current model and
+hardware*: gradient checkpointing, fused attention kernels, fused optimizers.
+
+Modes:
+    "none"    off — plain training
+    "shadow"  force every applicable optimization on
+    "auto"    enable what actually helps at the current size (the default)
+
+`plan()` is pure and side-effect free; each backend reads the returned `ShadowPlan`
+and applies the flags it understands.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+# Below this layer count, the optimizations cost more overhead than they save, so
+# "auto" leaves them off (forcing "shadow" still turns them on).
+_BIG_MODEL_LAYERS = 24
+
+
+@dataclass
+class ShadowPlan:
+    grad_checkpoint: bool = False
+    flash_attention: bool = False
+    fused_optimizer: bool = False
+    enabled: list[str] = field(default_factory=list)
+
+    @property
+    def active(self) -> bool:
+        return bool(self.enabled)
+
+    @property
+    def note(self) -> str:
+        if not self.enabled:
+            return "[shadow] no extra optimizations needed at this size"
+        return "[shadow] enabled: " + ", ".join(self.enabled)
+
+
+def plan(mode: str, *, backend: str, n_layers: int = 0, has_flash: bool = False) -> ShadowPlan:
+    """Decide which shadow optimizations to apply.
+
+    backend: "mlx" | "torch". n_layers: the model's transformer-block count.
+    has_flash: whether a flash-attention kernel is importable (torch only).
+    """
+    if mode not in ("none", "shadow", "auto"):
+        raise ValueError(f"unknown accelerator {mode!r} (expected auto|shadow|none)")
+    if mode == "none":
+        return ShadowPlan()
+    force = mode == "shadow"
+    big = n_layers >= _BIG_MODEL_LAYERS
+    p = ShadowPlan()
+
+    if backend == "mlx":
+        if force or big:
+            p.grad_checkpoint = True
+            p.enabled.append("gradient checkpointing")
+    else:  # torch (cuda / cpu)
+        if has_flash:
+            p.flash_attention = True
+            p.enabled.append("flash-attention-2")
+        if force or big:
+            p.grad_checkpoint = True
+            p.enabled.append("gradient checkpointing")
+        p.fused_optimizer = True
+        p.enabled.append("fused optimizer")
+    return p

shadowlm/ascii.py

@@ -0,0 +1,79 @@
+"""Startup banner for shadow training."""
+
+from __future__ import annotations
+
+import functools
+import os
+
+_HEART = r"""
+╔══════════════════════════════════════════════════════╗
+║             ♥                         ♥♥             ║
+║            ♥♥♥                       ♥♥♥♥            ║
+║           ♥♥♥♥♥                      ♥♥♥♥♥           ║
+║          ♥♥♥♥♥♥♥                    ♥♥♥♥♥♥♥          ║
+║         ♥♥♥♥♥♥♥♥                   ♥♥♥♥♥♥♥♥          ║
+║         ♥♥♥♥♥♥♥♥♥                 ♥♥♥♥♥♥♥♥♥♥         ║
+║         ♥♥♥♥♥♥♥♥♥♥               ♥♥♥♥♥♥♥♥♥♥♥         ║
+║         ♥♥♥♥♥♥♥♥♥♥♥             ♥♥♥♥♥♥♥♥♥♥♥          ║
+║          ♥♥♥♥♥♥♥♥♥♥♥           ♥♥♥♥♥♥♥♥♥♥♥♥          ║
+║           ♥♥♥♥♥♥♥♥♥♥♥         ♥♥♥♥♥♥♥♥♥♥♥♥           ║
+║            ♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥            ║
+║             ♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥              ║
+║               ♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥                ║
+║                ♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥                 ║
+║                  ♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥♥                   ║
+║                    ♥♥♥♥♥♥♥♥♥♥♥♥♥                     ║
+║                      ♥♥♥♥♥♥♥                         ║
+║                        ♥♥                            ║
+║                                                      ║
+╚══════════════════════════════════════════════════════╝"""
+
+_NAME = r"""
+███████╗ ██╗                    ██╗                     ██╗      ███╗   ███╗
+██╔════╝ ██║                    ██║                     ██║      ████╗ ████║
+███████╗ ███████╗  █████╗   ██████║  █████╗  ██╗ █╗ ██╗ ██║      ██╔████╔██║
+╚════██║ ██╔══██╗ ██╔══██║ ██╔══██║ ██╔══██╗ ██║███╗██║ ██║      ██║╚██╔╝██║
+███████║ ██║  ██║ ╚██████║ ╚██████║ ╚█████╔╝ ╚███╔███╔╝ ███████╗ ██║ ╚═╝ ██║
+╚══════╝ ╚═╝  ╚═╝  ╚═════╝  ╚═════╝  ╚════╝   ╚══╝╚══╝  ╚══════╝ ╚═╝     ╚═╝
+                            T  R  A  I  N  E  R
+               any open model  ·  any harness  ·  any method
+                    from  Lyzr  Research  Labs  ·  slm♥
+"""
+
+# The full banner prints on the first training session of the process; later
+# sessions get a compact one-liner so sweeps don't drown in hearts.
+_sessions = 0
+
+
+def run_on_main_rank(fn):
+    """Only run on the main process — rank 0 in a distributed (multi-GPU) launch.
+
+    Reads the usual launcher env vars; a plain single-process run is rank 0, so
+    the banner prints normally. Keeps worker ranks quiet once the studio fans
+    training out across GPUs.
+    """
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        try:
+            rank = int(os.environ.get("RANK") or os.environ.get("LOCAL_RANK") or 0)
+        except ValueError:  # launcher set RANK to something non-numeric
+            rank = 0
+        if rank == 0:
+            return fn(*args, **kwargs)
+        return None
+
+    return wrapper
+
+
+@run_on_main_rank
+def print_ascii_art(*, once: bool = True) -> None:
+    """Print the ShadowLM Trainer banner — full art first, a compact line after."""
+    global _sessions
+    _sessions += 1
+    if once and _sessions > 1:
+        print(f"\nslm♥ ShadowLM Trainer · training session #{_sessions}\n")
+        return
+    print(_HEART)
+    print(_NAME)
+    print("Starting training session...\n")

shadowlm/backends/__init__.py

@@ -0,0 +1,64 @@
+"""Backend selection.
+
+`select_backend("auto")` picks the backend for the current hardware:
+CUDA → torch, else Apple Silicon → mlx, else torch on CPU. If no backend is
+installed, `load()` says what to install. Force one with backend="mlx" / "torch"
+(and device="cpu" to pin torch to the CPU).
+"""
+
+from __future__ import annotations
+
+from .base import Backend, Callbacks, FinetuneResult
+
+__all__ = ["Backend", "Callbacks", "FinetuneResult", "select_backend"]
+
+
+def _mlx_available() -> bool:
+    from .mlx import MLXBackend
+    return MLXBackend.is_available()
+
+
+def _torch_available() -> bool:
+    from .torch import TorchBackend
+    return TorchBackend.is_available()
+
+
+def _no_backend() -> RuntimeError:
+    return RuntimeError(
+        "No training backend available. Install one:\n"
+        "  • Apple Silicon:  pip install shadowlm[mlx]\n"
+        "  • CUDA / CPU:     pip install shadowlm[torch]"
+    )
+
+
+def select_backend(name: str = "auto", *, accelerator: str = "auto",
+                   device: str = "auto") -> Backend:
+    name = (name or "auto").lower()
+
+    if name == "auto":
+        from .torch import TorchBackend
+        if TorchBackend.has_cuda():
+            name = "torch"
+        elif _mlx_available():
+            name = "mlx"
+        elif _torch_available():
+            name, device = "torch", "cpu"
+        else:
+            raise _no_backend()
+
+    if name == "mlx":
+        from .mlx import MLXBackend
+        if not MLXBackend.is_available():
+            raise RuntimeError("mlx backend needs Apple Silicon + shadowlm[mlx].")
+        return MLXBackend(accelerator=accelerator)
+
+    if name == "torch":
+        # CPU is just this backend with device="cpu" (no separate "cpu" backend).
+        from .torch import TorchBackend
+        if not TorchBackend.is_available():
+            raise RuntimeError(
+                "torch backend needs shadowlm[torch] (torch, transformers, trl, peft, datasets)."
+            )
+        return TorchBackend(device=device, accelerator=accelerator)
+
+    raise ValueError(f"unknown backend {name!r} (expected auto|mlx|torch)")

shadowlm/backends/base.py

@@ -0,0 +1,99 @@
+"""The backend contract.
+
+A `Backend` is a stateful object that holds a loaded model and knows how to
+finetune it and generate from it. Everything user-facing (`Model`, `load`,
+`finetune`, `generate`) is backend-agnostic — swap mlx ↔ torch without touching
+the SDK surface.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Callable
+
+from ..data import Dataset
+from ..training import Metric, TrainConfig
+
+
+class Callbacks:
+    """Bridges a backend's training loop back to the caller.
+
+    A backend calls `.step(metric)` for each logged step and `.log(msg)` for
+    status lines, and checks `.stopped()` to honour cancellation.
+    """
+
+    def __init__(
+        self,
+        on_step: Callable[[Metric], None] | None = None,
+        on_log: Callable[[str], None] | None = None,
+        on_eval: Callable[[Metric], None] | None = None,
+        should_stop: Callable[[], bool] | None = None,
+    ) -> None:
+        self._on_step = on_step
+        self._on_log = on_log
+        self._on_eval = on_eval
+        self._should_stop = should_stop
+
+    def step(self, metric: Metric) -> None:
+        if self._on_step:
+            self._on_step(metric)
+
+    def eval(self, metric: Metric) -> None:
+        if self._on_eval:
+            self._on_eval(metric)
+
+    def log(self, message: str) -> None:
+        if self._on_log:
+            self._on_log(message)
+
+    def stopped(self) -> bool:
+        return bool(self._should_stop and self._should_stop())
+
+
+class FinetuneResult:
+    """What a backend returns from `finetune`: where the adapter/model landed."""
+
+    def __init__(self, checkpoint: str, final_loss: float | None) -> None:
+        self.checkpoint = checkpoint
+        self.final_loss = final_loss
+
+
+class Backend(ABC):
+    name: str = "base"
+
+    def __init__(self, *, device: str = "auto", accelerator: str = "auto") -> None:
+        # device interpretation is backend-specific: mlx → gpu, torch → cuda|cpu.
+        # accelerator selects the shadow optimization layer: none | shadow | auto.
+        self.device = device
+        self.accelerator = accelerator
+        self.model_name: str | None = None
+
+    @classmethod
+    def is_available(cls) -> bool:
+        """Whether this backend can actually run in the current environment."""
+        return True
+
+    @abstractmethod
+    def load(self, name: str, *, load_in_4bit: bool, max_seq_length: int,
+             adapter: str | None = None, **kwargs) -> None:
+        ...
+
+    @abstractmethod
+    def finetune(self, dataset: Dataset, config: TrainConfig, callbacks: Callbacks,
+                 output_dir: str, eval_dataset: Dataset | None = None,
+                 reward_fns: list | None = None) -> FinetuneResult:
+        ...
+
+    @abstractmethod
+    def generate(self, prompt: str, *, max_new_tokens: int, temperature: float,
+                 top_p: float, **kwargs) -> str:
+        ...
+
+    def chat(self, messages: list[dict], *, tools: list[dict] | None = None,
+             max_new_tokens: int, temperature: float, top_p: float, **kwargs) -> str:
+        """Multi-turn chat through the tokenizer's chat template, with optional
+        tool schemas. Returns the raw assistant text (tool-call markup included)."""
+        raise NotImplementedError(f"{self.name} backend does not implement chat()")
+
+    def save(self, path: str, *, fmt: str = "adapter") -> str:
+        raise NotImplementedError(f"{self.name} backend does not implement save()")