PyPI - mlxsmith - Versions diffs - 0.1.0__py3-none-any.whl - Mend

mlxsmith 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (69) hide show

mlxsmith/__init__.py +2 -0
mlxsmith/accel/__init__.py +10 -0
mlxsmith/accel/base.py +17 -0
mlxsmith/accel/none.py +13 -0
mlxsmith/accel/zmlx_backend.py +42 -0
mlxsmith/adapters.py +46 -0
mlxsmith/api/__init__.py +48 -0
mlxsmith/api/handlers.py +1217 -0
mlxsmith/api/schemas.py +436 -0
mlxsmith/auth.py +88 -0
mlxsmith/bench.py +102 -0
mlxsmith/cli.py +950 -0
mlxsmith/config.py +543 -0
mlxsmith/config_models.py +261 -0
mlxsmith/data.py +493 -0
mlxsmith/envs/__init__.py +33 -0
mlxsmith/envs/system.py +388 -0
mlxsmith/envs/token_env.py +191 -0
mlxsmith/eval.py +112 -0
mlxsmith/infer.py +140 -0
mlxsmith/llm/__init__.py +16 -0
mlxsmith/llm/backend.py +126 -0
mlxsmith/llm/interface.py +212 -0
mlxsmith/llm/mlx_lm_backend.py +509 -0
mlxsmith/llm/mock_backend.py +228 -0
mlxsmith/llm/registry.py +12 -0
mlxsmith/models.py +257 -0
mlxsmith/orchestrator/__init__.py +25 -0
mlxsmith/orchestrator/daemon.py +454 -0
mlxsmith/orchestrator/inference_worker.py +496 -0
mlxsmith/orchestrator/queue.py +355 -0
mlxsmith/orchestrator/trainer_worker.py +437 -0
mlxsmith/rlm/__init__.py +8 -0
mlxsmith/rlm/corpus.py +74 -0
mlxsmith/rlm/gating.py +90 -0
mlxsmith/rlm/generate.py +249 -0
mlxsmith/rlm/history.py +12 -0
mlxsmith/rlm/inference.py +150 -0
mlxsmith/rlm/loop.py +1297 -0
mlxsmith/rlm/mutate.py +82 -0
mlxsmith/rlm/trainer.py +73 -0
mlxsmith/rlm/weights.py +263 -0
mlxsmith/runs.py +44 -0
mlxsmith/sdk/__init__.py +392 -0
mlxsmith/sdk/future.py +486 -0
mlxsmith/sdk/losses.py +262 -0
mlxsmith/sdk/sampling_client.py +729 -0
mlxsmith/sdk/training_client.py +676 -0
mlxsmith/server.py +376 -0
mlxsmith/train/__init__.py +0 -0
mlxsmith/train/distill.py +279 -0
mlxsmith/train/lora.py +280 -0
mlxsmith/train/pref.py +180 -0
mlxsmith/train/rft.py +458 -0
mlxsmith/train/sft.py +151 -0
mlxsmith/util.py +174 -0
mlxsmith/verifiers/__init__.py +3 -0
mlxsmith/verifiers/compose.py +109 -0
mlxsmith/verifiers/docker_verifier.py +111 -0
mlxsmith/verifiers/jsonschema.py +54 -0
mlxsmith/verifiers/pytest_verifier.py +82 -0
mlxsmith/verifiers/regex.py +15 -0
mlxsmith/verifiers/types.py +10 -0
mlxsmith-0.1.0.dist-info/METADATA +163 -0
mlxsmith-0.1.0.dist-info/RECORD +69 -0
mlxsmith-0.1.0.dist-info/WHEEL +5 -0
mlxsmith-0.1.0.dist-info/entry_points.txt +2 -0
mlxsmith-0.1.0.dist-info/licenses/LICENSE +21 -0
mlxsmith-0.1.0.dist-info/top_level.txt +1 -0

mlxsmith/server.py ADDED Viewed

@@ -0,0 +1,376 @@
+"""MLXSmith FastAPI server with OpenAI-compatible chat completions.
+This module provides:
+- OpenAI-compatible chat completions with streaming support
+- Internal rollout endpoint (tokens + logprobs)
+- Adapter hot-reload
+- RLM state and history endpoints
+- Web UI and RLM monitor (when enabled)
+"""
+from __future__ import annotations
+import json
+import time
+import uuid
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse, HTMLResponse
+from pydantic import BaseModel
+from .config import ProjectConfig
+from .models import resolve_model_spec
+from .llm.registry import get_llm_backend
+# Import new API handlers
+from .api.handlers import create_router, InternalAuthMiddleware
+from .api.schemas import (
+    ChatMessage,
+    ChatRequest,
+    ChatResponse,
+    RolloutRequest,
+    RolloutResponse,
+    AdapterReloadRequest,
+    AdapterReloadResponse,
+)
+# Re-export schemas for backward compatibility
+__all__ = [
+    "create_app",
+    "ChatMessage",
+    "ChatRequest",
+    "ChatResponse",
+    "RolloutRequest",
+    "RolloutResponse",
+    "AdapterReloadRequest",
+    "AdapterReloadResponse",
+]
+def _ui_html() -> str:
+    return """<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1" />
+  <title>mlxsmith serve</title>
+  <style>
+    :root {
+      --bg: #f7f2ea;
+      --ink: #231f20;
+      --accent: #0b6e4f;
+      --accent-2: #b44a1d;
+      --card: #fff8ee;
+    }
+    * { box-sizing: border-box; }
+    body {
+      margin: 0;
+      font-family: "Space Grotesk", "Avenir Next", "Segoe UI", sans-serif;
+      color: var(--ink);
+      background:
+        radial-gradient(circle at 10% 10%, rgba(11,110,79,0.12), transparent 35%),
+        radial-gradient(circle at 90% 20%, rgba(180,74,29,0.12), transparent 40%),
+        var(--bg);
+    }
+    header {
+      padding: 28px 32px;
+      border-bottom: 1px solid rgba(35,31,32,0.1);
+      display: flex;
+      align-items: center;
+      justify-content: space-between;
+      gap: 16px;
+    }
+    .title {
+      font-size: 20px;
+      letter-spacing: 0.02em;
+      text-transform: uppercase;
+    }
+    .wrap {
+      display: grid;
+      grid-template-columns: 1.2fr 0.8fr;
+      gap: 24px;
+      padding: 24px 32px 40px;
+    }
+    .card {
+      background: var(--card);
+      border: 1px solid rgba(35,31,32,0.08);
+      border-radius: 16px;
+      padding: 18px;
+      box-shadow: 0 6px 20px rgba(35,31,32,0.06);
+    }
+    h2 { margin: 0 0 12px; font-size: 18px; }
+    textarea {
+      width: 100%;
+      min-height: 120px;
+      border: 1px solid rgba(35,31,32,0.2);
+      border-radius: 12px;
+      padding: 12px;
+      font: inherit;
+      background: #fffdf9;
+    }
+    .row { display: flex; gap: 12px; margin-top: 10px; }
+    button {
+      border: 0;
+      border-radius: 10px;
+      padding: 10px 14px;
+      font: inherit;
+      cursor: pointer;
+    }
+    .btn-primary { background: var(--accent); color: #fff; }
+    .btn-secondary { background: var(--accent-2); color: #fff; }
+    .log {
+      background: #14110f;
+      color: #fdf7f0;
+      border-radius: 12px;
+      padding: 12px;
+      height: 300px;
+      overflow: auto;
+      font-family: "Iosevka", "Menlo", monospace;
+      font-size: 13px;
+    }
+    .muted { color: rgba(35,31,32,0.6); font-size: 12px; }
+    @media (max-width: 900px) {
+      .wrap { grid-template-columns: 1fr; }
+    }
+  </style>
+</head>
+<body>
+  <header>
+    <div class="title">mlxsmith serve</div>
+    <div class="muted">OpenAI-compatible API + RLM monitor</div>
+  </header>
+  <div class="wrap">
+    <section class="card">
+      <h2>Chat</h2>
+      <textarea id="prompt" placeholder="Ask the model anything..."></textarea>
+      <div class="row">
+        <button class="btn-primary" id="send">Send</button>
+        <button class="btn-secondary" id="stream">Stream</button>
+        <a class="btn-secondary" href="/rlm/monitor" style="text-decoration:none;">RLM Monitor</a>
+      </div>
+      <div class="log" id="output"></div>
+    </section>
+    <section class="card">
+      <h2>Status</h2>
+      <div id="status" class="muted">Loading…</div>
+      <div class="muted" style="margin-top:12px;">Tip: enable <code>serve.ui</code> in mlxsmith.yaml to keep this page on by default.</div>
+    </section>
+  </div>
+  <script>
+    const output = document.getElementById('output');
+    const statusEl = document.getElementById('status');
+    const promptEl = document.getElementById('prompt');
+    const baseBody = () => ({
+      messages: [{role: 'user', content: promptEl.value}],
+      max_tokens: 256
+    });
+    async function refreshStatus() {
+      try {
+        const res = await fetch('/internal/rlm/state');
+        if (!res.ok) return;
+        const data = await res.json();
+        statusEl.textContent = JSON.stringify(data, null, 2);
+      } catch (err) {
+        statusEl.textContent = 'Status unavailable';
+      }
+    }
+    refreshStatus();
+    setInterval(refreshStatus, 5000);
+    document.getElementById('send').onclick = async () => {
+      output.textContent = '';
+      const res = await fetch('/v1/chat/completions', {
+        method: 'POST',
+        headers: {'Content-Type': 'application/json'},
+        body: JSON.stringify(baseBody())
+      });
+      const data = await res.json();
+      output.textContent = data.choices?.[0]?.message?.content || '';
+    };
+    document.getElementById('stream').onclick = async () => {
+      output.textContent = '';
+      const body = baseBody();
+      body.stream = true;
+      const res = await fetch('/v1/chat/completions', {
+        method: 'POST',
+        headers: {'Content-Type': 'application/json'},
+        body: JSON.stringify(body)
+      });
+      const reader = res.body.getReader();
+      const decoder = new TextDecoder('utf-8');
+      let buf = '';
+      while (true) {
+        const { value, done } = await reader.read();
+        if (done) break;
+        buf += decoder.decode(value, {stream: true});
+        const parts = buf.split('\\n\\n');
+        buf = parts.pop() || '';
+        for (const part of parts) {
+          const line = part.trim();
+          if (!line.startsWith('data:')) continue;
+          const payload = line.replace('data:', '').trim();
+          if (payload === '[DONE]') return;
+          try {
+            const obj = JSON.parse(payload);
+            const delta = obj.choices?.[0]?.delta?.content || '';
+            output.textContent += delta;
+          } catch (e) {}
+        }
+      }
+    };
+  </script>
+</body>
+</html>
+"""
+def _monitor_html() -> str:
+    return """<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1" />
+  <title>RLM Monitor</title>
+  <style>
+    body {
+      margin: 0;
+      font-family: "Space Grotesk", "Avenir Next", "Segoe UI", sans-serif;
+      background: #f4f6f1;
+      color: #222;
+    }
+    header {
+      padding: 24px 32px;
+      background: #1f3b2c;
+      color: #f4f6f1;
+      display: flex;
+      justify-content: space-between;
+      align-items: center;
+    }
+    .wrap { padding: 24px 32px; display: grid; gap: 16px; }
+    .card {
+      background: #fff;
+      border-radius: 14px;
+      padding: 16px;
+      box-shadow: 0 8px 20px rgba(0,0,0,0.08);
+    }
+    canvas { width: 100%; height: 220px; }
+    pre { margin: 0; font-size: 12px; }
+  </style>
+</head>
+<body>
+  <header>
+    <div>RLM Monitor</div>
+    <a href="/" style="color:#f4f6f1;text-decoration:none;">Back to Serve</a>
+  </header>
+  <div class="wrap">
+    <section class="card">
+      <canvas id="chart" width="900" height="220"></canvas>
+    </section>
+    <section class="card">
+      <pre id="state">Loading…</pre>
+    </section>
+  </div>
+  <script>
+    const canvas = document.getElementById('chart');
+    const ctx = canvas.getContext('2d');
+    const stateEl = document.getElementById('state');
+    function draw(history) {
+      ctx.clearRect(0,0,canvas.width,canvas.height);
+      if (!history.length) return;
+      const scores = history.map(h => h.adapter_score || 0);
+      const max = Math.max(...scores, 1);
+      const min = Math.min(...scores, 0);
+      const pad = 20;
+      const w = canvas.width - pad * 2;
+      const h = canvas.height - pad * 2;
+      ctx.strokeStyle = '#1f3b2c';
+      ctx.lineWidth = 2;
+      ctx.beginPath();
+      scores.forEach((s, i) => {
+        const x = pad + (i / Math.max(1, scores.length - 1)) * w;
+        const y = pad + (1 - (s - min) / (max - min || 1)) * h;
+        if (i === 0) ctx.moveTo(x, y); else ctx.lineTo(x, y);
+      });
+      ctx.stroke();
+    }
+    async function refresh() {
+      try {
+        const h = await fetch('/internal/rlm/history').then(r => r.json());
+        draw(h);
+        const s = await fetch('/internal/rlm/state').then(r => r.json());
+        stateEl.textContent = JSON.stringify(s, null, 2);
+      } catch (e) {
+        stateEl.textContent = 'History unavailable';
+      }
+    }
+    refresh();
+    setInterval(refresh, 5000);
+  </script>
+</body>
+</html>
+"""
+def create_app(model_spec: str, cfg: ProjectConfig) -> FastAPI:
+    """Create and configure the FastAPI application.
+    Args:
+        model_spec: Model specification (path or HF repo ID)
+        cfg: Project configuration
+    Returns:
+        Configured FastAPI application
+    """
+    app = FastAPI(
+        title="mlxsmith",
+        description="MLXSmith API server for local LLM inference and RLM training",
+        version="0.1.0",
+    )
+    # Load LLM backend
+    llm = get_llm_backend(cfg.model.backend)
+    base_model, adapter_path, _meta = resolve_model_spec(Path.cwd(), model_spec, cfg)
+    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))
+    current_adapter = str(adapter_path) if adapter_path else None
+    # Add authentication middleware for internal endpoints
+    app.add_middleware(
+        InternalAuthMiddleware,
+        api_token=None,  # Set via MLXSMITH_API_TOKEN env var
+        internal_prefix="/internal",
+        public_paths=["/health", "/v1/chat/completions", "/docs", "/openapi.json"],
+    )
+    # Create and include the API router
+    router = create_router(
+        llm_backend=llm,
+        base_model=base_model,
+        current_adapter=current_adapter,
+        cfg=cfg,
+    )
+    app.include_router(router)
+    # Add UI routes if enabled
+    if cfg.serve.ui:
+        @app.get("/")
+        def ui_root():
+            return HTMLResponse(_ui_html())
+        @app.get("/rlm/monitor")
+        def ui_monitor():
+            return HTMLResponse(_monitor_html())
+    return app

mlxsmith/train/__init__.py ADDED Viewed

File without changes

mlxsmith/train/distill.py ADDED Viewed

@@ -0,0 +1,279 @@
+from __future__ import annotations
+import json
+import random
+from pathlib import Path
+from typing import Iterable
+from rich.console import Console
+from ..accel import get_backend
+from ..config import ProjectConfig
+from ..llm.interface import compute_logprobs
+from ..models import resolve_model_spec
+from ..runs import RunPaths, new_run, snapshot_config
+from ..llm.registry import get_llm_backend
+from ..llm.backend import BackendNotAvailable
+from ..sdk.losses import importance_sampling_loss
+from ..train.lora import LoRAConfig
+from ..util import ensure_dir, write_jsonl, now_ts, tree_add, tree_scale
+from .sft import run_sft
+console = Console()
+def _iter_prompts(path: Path) -> Iterable[str]:
+    for line in path.read_text(encoding="utf-8").splitlines():
+        if not line.strip():
+            continue
+        row = json.loads(line)
+        prompt = row.get("prompt") or row.get("instruction") or row.get("input") or ""
+        if not prompt and "messages" in row:
+            msgs = row.get("messages") or []
+            if msgs:
+                prompt = "\n".join([m.get("content", "") for m in msgs])
+        if prompt:
+            yield str(prompt)
+def run_distill(
+    project_root: Path,
+    cfg: ProjectConfig,
+    data_path: Path,
+    *,
+    teacher_model: str,
+    student_model: str,
+    accel: str,
+    mode: str = "offline",
+    max_new_tokens: int = 256,
+    temperature: float = 0.7,
+) -> RunPaths:
+    run = new_run(project_root, "distill")
+    snapshot_config(cfg.model_dump(), run.config_snapshot_path)
+    mode = mode.lower()
+    prompts = list(_iter_prompts(data_path))
+    if not prompts:
+        raise RuntimeError("No prompts found in distillation dataset")
+    backend = get_backend(accel)
+    backend.patch()
+    console.print(f"[bold]DISTILL[/bold] run: {run.run_dir.name} mode={mode} accel={backend.name}")
+    teacher = get_llm_backend(cfg.model.backend)
+    student = get_llm_backend(cfg.model.backend)
+    base_teacher, _, _meta_t = resolve_model_spec(project_root, teacher_model, cfg)
+    base_student, adapter_path, _meta_s = resolve_model_spec(project_root, student_model, cfg)
+    try:
+        teacher.load(
+            base_teacher,
+            max_seq_len=cfg.model.max_seq_len,
+            dtype=cfg.model.dtype,
+            trust_remote_code=cfg.model.trust_remote_code,
+        )
+        student.load(
+            base_student,
+            max_seq_len=cfg.model.max_seq_len,
+            dtype=cfg.model.dtype,
+            trust_remote_code=cfg.model.trust_remote_code,
+        )
+    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.\nteacher={teacher_model}\nstudent={student_model}\n",
+            encoding="utf-8",
+        )
+        return run
+    distill_dir = ensure_dir(run.artifacts_dir / "distill_data")
+    train_path = distill_dir / "train.jsonl"
+    if mode == "opd":
+        rows = [{"prompt": prompt} for prompt in prompts]
+        write_jsonl(train_path, rows)
+        console.print(f"[green]Distill dataset (OPD)[/green] {train_path}")
+        if adapter_path:
+            student.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,
+            )
+            student.apply_lora_from_config(lora_cfg)
+        opt, _params = student.optimizer_and_params(lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
+        rng = random.Random(cfg.train.seed)
+        total = int(cfg.train.iters)
+        grad_accum = max(1, int(cfg.train.grad_accum))
+        max_len = int(cfg.model.max_seq_len)
+        accum_grads = None
+        def _to_float(val):
+            if hasattr(val, "item"):
+                try:
+                    return float(val.item())
+                except Exception:
+                    pass
+            return float(val)
+        for step in range(1, total + 1):
+            prompt = rng.choice(prompts)
+            gen = student.generate_with_logprobs(
+                prompt,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                logprobs=0,
+            )
+            completion = gen.text[len(prompt) :] if gen.text.startswith(prompt) else gen.text
+            behavior_logprobs = list(gen.logprobs) if gen.logprobs else []
+            if not behavior_logprobs:
+                continue
+            teacher_logprobs: list[float] = []
+            teacher_avg = None
+            teacher_res = compute_logprobs(
+                teacher,
+                prompt,
+                completion,
+                top_k=0,
+                max_seq_len=max_len,
+            )
+            if teacher_res.token_logprobs:
+                teacher_logprobs = [float(lp) for lp in teacher_res.token_logprobs]
+            else:
+                prompt_ids = teacher.encode(prompt)
+                ids = teacher.encode(prompt + completion)
+                if max_len and len(ids) > max_len:
+                    overflow = len(ids) - max_len
+                    ids = ids[overflow:]
+                    prompt_len = max(0, len(prompt_ids) - overflow)
+                else:
+                    prompt_len = len(prompt_ids)
+                teacher_len = max(0, len(ids) - prompt_len)
+                if teacher_len > 0:
+                    teacher_total = teacher.sequence_logprob(ids, prompt_len=prompt_len)
+                    teacher_avg = _to_float(teacher_total) / float(teacher_len)
+            ids = list(gen.token_ids)
+            prompt_len = gen.prompt_len
+            if max_len and len(ids) > max_len:
+                overflow = len(ids) - max_len
+                ids = ids[overflow:]
+                if overflow >= prompt_len:
+                    removed_completion = overflow - prompt_len
+                    prompt_len = 0
+                    if removed_completion > 0:
+                        behavior_logprobs = behavior_logprobs[removed_completion:]
+                        if teacher_logprobs:
+                            teacher_logprobs = teacher_logprobs[removed_completion:]
+                else:
+                    prompt_len = max(0, prompt_len - overflow)
+            completion_len = max(0, len(ids) - prompt_len)
+            if completion_len == 0 or not behavior_logprobs:
+                continue
+            if teacher_logprobs:
+                n = min(len(teacher_logprobs), len(behavior_logprobs), completion_len)
+                if n <= 0:
+                    continue
+                if completion_len != n:
+                    ids = ids[: prompt_len + n]
+                    completion_len = n
+                behavior_logprobs = behavior_logprobs[:n]
+                teacher_logprobs = teacher_logprobs[:n]
+                behavior_logprob = sum(behavior_logprobs)
+                advantage = sum(t - b for t, b in zip(teacher_logprobs, behavior_logprobs)) / float(n)
+            else:
+                if teacher_avg is None:
+                    continue
+                if len(behavior_logprobs) > completion_len:
+                    behavior_logprobs = behavior_logprobs[:completion_len]
+                behavior_logprob = sum(behavior_logprobs)
+                if not behavior_logprobs:
+                    continue
+                behavior_avg = behavior_logprob / float(len(behavior_logprobs))
+                advantage = teacher_avg - behavior_avg
+            def loss_fn(_model):
+                return importance_sampling_loss(
+                    student,
+                    ids,
+                    prompt_len=prompt_len,
+                    advantage=advantage,
+                    behavior_logprob=behavior_logprob,
+                )
+            lval, grads = student.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:
+                    student.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": "distill_opd",
+                            "loss": _to_float(lval),
+                            "advantage": float(advantage),
+                            "accel": backend.name,
+                        }
+                    ],
+                )
+            if step % cfg.train.save_every == 0 or step == total:
+                student.save_adapter(
+                    str(run.adapter_dir),
+                    metadata={
+                        "base_model": base_student,
+                        "source_adapter": str(adapter_path) if adapter_path else None,
+                        "run": run.run_dir.name,
+                        "kind": "distill_opd",
+                    },
+                )
+        child = run
+    else:
+        rows = []
+        for prompt in prompts:
+            gen = teacher.generate(
+                prompt,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+            )
+            completion = gen.text[len(prompt) :] if gen.text.startswith(prompt) else gen.text
+            rows.append({"prompt": prompt, "response": completion})
+        write_jsonl(train_path, rows)
+        console.print(f"[green]Distill dataset (offline)[/green] {train_path}")
+        child = run_sft(project_root, cfg, distill_dir, student_model, accel)
+    write_jsonl(
+        run.metrics_path,
+        [
+            {
+                "ts": now_ts(),
+                "kind": "distill",
+                "mode": mode,
+                "teacher": teacher_model,
+                "student": student_model,
+                "child_run": str(child.run_dir),
+                "samples": len(prompts),
+            }
+        ],
+    )
+    return run