diff --git a/CLAUDE.md b/CLAUDE.md index eb2e274..444d76a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -5,7 +5,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co ## What this is ShadowLM Trainer is a fine-tuning SDK: load any open model, train it with any of -13 methods, on any hardware, then own the weights. The headline use case is +15 methods, on any hardware, then own the weights. The headline use case is "shadowing" — moving one task off a rented frontier model onto a small model you own, by capturing real agent traffic (`slm.capture()`), judging episodes, and training on them — without modifying the agent (the model API is the only diff --git a/README.md b/README.md index 3af37a9..ef25959 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@

License: MIT Python 3.10+ - Methods + Methods Batteries included

@@ -31,8 +31,9 @@ print(model.generate("What is the capital of France?")) # inference model.save("out/", fmt="adapter") # ship it ``` -Change `method="lora"` to `qlora`, `dora`, `full`, `dpo`, `grpo`, `more`, `bitfit`, -`prompt`, `ptuning`, `adapter`, `cpt`, `more_plus` — and nothing else changes. That's the idea. +Change `method="lora"` to `qlora`, `dora`, `full`, `dpo`, `grpo`, `sdft`, `sdpo`, +`more`, `bitfit`, `prompt`, `ptuning`, `adapter`, `cpt`, `more_plus` — and nothing else +changes. That's the idea. ## What ShadowLM is for @@ -72,7 +73,7 @@ The whole **capture → judge → train → own a shadowLM** loop runs on these: | Block | What it does | API | |-------|--------------|-----| | **Capture proxy** | drop-in OpenAI endpoint that records your agent's traffic into trajectories — agent unchanged | `slm.capture()` | -| **13 methods** | LoRA · QLoRA · DoRA · full · CPT · DPO · GRPO · MoRE · MoRE+ · BitFit · prompt · p-tuning · adapter | `method=` | +| **15 methods** | LoRA · QLoRA · DoRA · full · CPT · DPO · GRPO · SDFT · SDPO · MoRE · MoRE+ · BitFit · prompt · p-tuning · adapter | `method=` | | **Judge → train** | score episodes with an LLM judge, train with trajectory-GRPO or DPO | `judge_group` | | **APO** | optimize the *prompt* instead of weights — same capture/judge front end, no GPU | `slm.optimize_prompt()` | | **VERL RL** | production multi-GPU GRPO (vLLM rollouts + FSDP) for cluster-scale RL | `backend="verl"` | @@ -99,6 +100,8 @@ spec (adapter kind, base requirements, data rendering), never the method name. | `cpt` | continued pretraining on raw domain text | either | 5e-5 | | `dpo` | preference optimization on `{prompt, chosen, rejected}` | either | 5e-6 | | `grpo` | RL from reward functions or scored `TrajectoryGroup`s | either | 5e-6 | +| `sdft` | **on-policy self-distillation** — the demo-conditioned model teaches itself; learns without forgetting | either | 1e-5 | +| `sdpo` | **RL via self-distillation** — the feedback-conditioned self-teacher densely rescores each rollout | either | 1e-5 | | `more` | **mixture of retrieval experts** — facts fused into attention | either | 1e-4 | | `more_plus` | **decoupled MoE** — per-fact final-FFN LoRA experts, BM25+semantic routed, cache-safe merge | **unquantized** | 1e-4 | | `bitfit`| train only the bias terms (~0.1% of params) | **unquantized** | 5e-4 | @@ -210,7 +213,7 @@ API — nothing reimplemented — to turn the blocks into a one-click migration: ``` [x] SDK — datasets → finetune → inference on mlx / torch / remote -[x] 13 methods incl. MoRE, MoRE+ (decoupled MoE), trajectory GRPO, judge rewards +[x] 15 methods incl. SDFT (self-distillation), SDPO (RL via self-distillation), MoRE, MoRE+ (decoupled MoE), trajectory GRPO, judge rewards [x] Capture proxy · shadow accelerator · any-hardware [x] Remote backend + reference server + the studio dashboard + CLI [ ] Studio orchestration — decision inbox · eval gates · shadow router · switch diff --git a/examples/README.md b/examples/README.md index a665ff9..c3c451f 100644 --- a/examples/README.md +++ b/examples/README.md @@ -29,6 +29,8 @@ python examples/remote/grpo.py | `cpt` | ✅ | ✅ | ✅ | raw text | — | | `dpo` | ✅ | ✅ | ✅ | preference pairs | — | | `grpo` | ✅ | ✅ | ✅ | prompts + reward fn | — | +| `sdft` | ✅ | ✅ | ✅ | chat | — | +| `sdpo` | ✅ | ✅ | ✅ | prompts + reward fn | — | | `more` | ✅ | ✅ | ✅ | facts | — | | `more_plus` | ✅ | ✅ | ✅ | facts | unquantized | | `bitfit` | ✅ | ✅ | ✅ | chat | unquantized + bias params | @@ -52,12 +54,14 @@ The `data/` folder holds tiny sample datasets so the examples are self-contained | file | format | used by | |------|--------|---------| -| `data/chat.jsonl` | chat (`messages`) | lora, qlora, dora, full, bitfit, prompt, ptuning, adapter | +| `data/chat.jsonl` | chat (`messages`) | lora, qlora, dora, full, sdft, bitfit, prompt, ptuning, adapter | | `data/preference.jsonl` | preference (`prompt/chosen/rejected`) | dpo | | `data/domain.jsonl` | raw text (`text`) | cpt | | `data/facts.jsonl` | instruction (`instruction/output`) | more, more_plus | -`grpo` defines its prompts and reward function inline in each script. +`grpo` and `sdpo` define their prompts and reward function inline in each +script (an `sdpo` reward fn may return `(score, feedback)` pairs — the feedback +becomes the self-teacher's in-context signal). There's also `shadowlm_qa.jsonl` — a chat dataset *about ShadowLM itself*, handy for a quick end-to-end finetune that teaches a small model to answer questions diff --git a/examples/mlx/sdft.py b/examples/mlx/sdft.py new file mode 100644 index 0000000..868fc38 --- /dev/null +++ b/examples/mlx/sdft.py @@ -0,0 +1,23 @@ +"""sdft · mlx backend + +SDFT — on-policy self-distillation: the model samples its own answers and is +pulled toward itself reading the golden response in-context, so it learns the +task with far less forgetting than SFT. Steps are slower than lora (each one +rolls out completions). Holds a second frozen copy of the base as the teacher. +Run from the repo root: + python examples/mlx/sdft.py +""" +import shadowlm as slm + + +def main(): + ds = slm.Dataset.from_jsonl("examples/data/chat.jsonl") + model = slm.load("mlx-community/Qwen2.5-0.5B-Instruct-bf16", backend="mlx") + run = model.finetune(ds, method="sdft", max_steps=30, + sdft_max_completion_length=64) + print("final loss:", run.loss, run.sparkline()) + model.save("out/mlx_sdft", fmt="adapter") + + +if __name__ == "__main__": + main() diff --git a/examples/mlx/sdpo.py b/examples/mlx/sdpo.py new file mode 100644 index 0000000..f421e8f --- /dev/null +++ b/examples/mlx/sdpo.py @@ -0,0 +1,36 @@ +"""sdpo · mlx backend + +SDPO — RL via self-distillation: the feedback-conditioned model teaches itself. +Reward fns may return (score, feedback) pairs; the feedback (and any successful +sibling rollout) becomes the self-teacher's in-context signal. +Run from the repo root: + python examples/mlx/sdpo.py +""" +import shadowlm as slm + + +prompts = [ + {"prompt": "What port does the ShadowLM studio serve on? Answer with just the number.", "answer": "8329"}, + {"prompt": "Which backend is ShadowLM's production training path? Answer with one word.", "answer": "torch"}, + {"prompt": "What does the M in MoRE stand for? Answer with one word.", "answer": "mixture"}, +] + +# 1.0 on a hit; on a miss, (0.0, hint) — the hint reaches the self-teacher. +def reward(prompts, completions, answer=None, types=None, **kwargs): + golds = answer if isinstance(answer, list) else [answer] * len(completions) + return [1.0 if (g or "").lower() in c.lower() + else (0.0, f"A correct answer contains {g!r}.") + for c, g in zip(completions, golds)] + + +def main(): + model = slm.load("mlx-community/Qwen2.5-0.5B-Instruct-bf16", backend="mlx") + run = model.finetune(prompts, method="sdpo", reward_fns=[reward], + max_steps=30, sdpo_group_size=4, + sdpo_max_completion_length=64) + print("final loss:", run.loss, run.sparkline()) + model.save("out/mlx_sdpo", fmt="adapter") + + +if __name__ == "__main__": + main() diff --git a/examples/remote/sdft.py b/examples/remote/sdft.py new file mode 100644 index 0000000..512f55d --- /dev/null +++ b/examples/remote/sdft.py @@ -0,0 +1,26 @@ +"""sdft · remote backend + +SDFT — on-policy self-distillation from demonstrations: learns the task from +plain chat rows with far less forgetting than SFT. + +Point SHADOWLM_API_URL at your server (defaults to http://127.0.0.1:8329). +Run from the repo root: + python examples/remote/sdft.py +""" +import os +import shadowlm as slm + +os.environ.setdefault("SHADOWLM_API_URL", "http://127.0.0.1:8329") + + +def main(): + ds = slm.Dataset.from_jsonl("examples/data/chat.jsonl") + model = slm.load("Qwen/Qwen3-8B", backend="remote") + run = model.finetune(ds, method="sdft", max_steps=60, + sdft_max_completion_length=128) + print("final loss:", run.loss, run.sparkline()) + model.save("out/remote_sdft", fmt="adapter") + + +if __name__ == "__main__": + main() diff --git a/examples/remote/sdpo.py b/examples/remote/sdpo.py new file mode 100644 index 0000000..1b69607 --- /dev/null +++ b/examples/remote/sdpo.py @@ -0,0 +1,39 @@ +"""sdpo · remote backend + +SDPO — RL via self-distillation: the feedback-conditioned model teaches itself. +Reward fns may return (score, feedback) pairs; the feedback (and any successful +sibling rollout) becomes the self-teacher's in-context signal. + +Point SHADOWLM_API_URL at your server (defaults to http://127.0.0.1:8329). +Run from the repo root: + python examples/remote/sdpo.py +""" +import os +import shadowlm as slm + +os.environ.setdefault("SHADOWLM_API_URL", "http://127.0.0.1:8329") + + +prompts = [ + {"prompt": "What port does the ShadowLM studio serve on? Answer with just the number.", "answer": "8329"}, + {"prompt": "Which backend is ShadowLM's production training path? Answer with one word.", "answer": "torch"}, + {"prompt": "What does the M in MoRE stand for? Answer with one word.", "answer": "mixture"}, +] + +# 1.0 on a hit; on a miss, (0.0, hint) — the hint reaches the self-teacher. +def reward(prompts, completions, answer=None, types=None, **kwargs): + golds = answer if isinstance(answer, list) else [answer] * len(completions) + return [1.0 if (g or "").lower() in c.lower() + else (0.0, f"A correct answer contains {g!r}.") + for c, g in zip(completions, golds)] + + +def main(): + model = slm.load("Qwen/Qwen3-8B", backend="remote") + run = model.finetune(prompts, method="sdpo", reward_fns=[reward], max_steps=30) + print("final loss:", run.loss, run.sparkline()) + model.save("out/remote_sdpo", fmt="adapter") + + +if __name__ == "__main__": + main() diff --git a/examples/torch/sdft.py b/examples/torch/sdft.py new file mode 100644 index 0000000..2a2493f --- /dev/null +++ b/examples/torch/sdft.py @@ -0,0 +1,23 @@ +"""sdft · torch backend + +SDFT — on-policy self-distillation: the model samples its own answers and is +pulled toward the same model reading the golden response in-context (adapters +disabled), so it learns the task with far less forgetting than SFT. Steps are +slower than lora — each one rolls out completions. +Run from the repo root: + python examples/torch/sdft.py +""" +import shadowlm as slm + + +def main(): + ds = slm.Dataset.from_jsonl("examples/data/chat.jsonl") + model = slm.load("Qwen/Qwen3-8B", backend="torch", device="cuda") + run = model.finetune(ds, method="sdft", max_steps=60, + sdft_max_completion_length=128) + print("final loss:", run.loss, run.sparkline()) + model.save("out/torch_sdft", fmt="adapter") + + +if __name__ == "__main__": + main() diff --git a/examples/torch/sdpo.py b/examples/torch/sdpo.py new file mode 100644 index 0000000..cd1c18a --- /dev/null +++ b/examples/torch/sdpo.py @@ -0,0 +1,35 @@ +"""sdpo · torch backend + +SDPO — RL via self-distillation: the feedback-conditioned model teaches itself. +Reward fns may return (score, feedback) pairs; the feedback (and any successful +sibling rollout) becomes the self-teacher's in-context signal. +Run from the repo root: + python examples/torch/sdpo.py +""" +import shadowlm as slm + + +prompts = [ + {"prompt": "What port does the ShadowLM studio serve on? Answer with just the number.", "answer": "8329"}, + {"prompt": "Which backend is ShadowLM's production training path? Answer with one word.", "answer": "torch"}, + {"prompt": "What does the M in MoRE stand for? Answer with one word.", "answer": "mixture"}, +] + +# 1.0 on a hit; on a miss, (0.0, hint) — the hint reaches the self-teacher. +def reward(prompts, completions, answer=None, types=None, **kwargs): + golds = answer if isinstance(answer, list) else [answer] * len(completions) + return [1.0 if (g or "").lower() in c.lower() + else (0.0, f"A correct answer contains {g!r}.") + for c, g in zip(completions, golds)] + + +def main(): + model = slm.load("Qwen/Qwen3-8B", backend="torch", device="cuda") + run = model.finetune(prompts, method="sdpo", reward_fns=[reward], + max_steps=60) + print("final loss:", run.loss, run.sparkline()) + model.save("out/torch_sdpo", fmt="adapter") + + +if __name__ == "__main__": + main() diff --git a/frontend/src/pages/Train.tsx b/frontend/src/pages/Train.tsx index b6d8f46..f88b327 100644 --- a/frontend/src/pages/Train.tsx +++ b/frontend/src/pages/Train.tsx @@ -81,6 +81,27 @@ const EXTRA: Record = { { key: "grpo_group_size", label: "Group size", kind: "int", def: "4", hint: "completions per prompt" }, { key: "grpo_max_completion_length", label: "Max completion", kind: "int", def: "256" }, ], + sdft: [ + { key: "sdft_alpha", label: "Alpha (divergence)", kind: "float", def: "0", + hint: "0 = forward KL, 1 = reverse KL, between = JSD" }, + { key: "sdft_max_completion_length", label: "Max completion", kind: "int", def: "512", + hint: "on-policy tokens sampled per prompt" }, + { key: "sdft_temperature", label: "Rollout temperature", kind: "float", def: "1", + hint: "0 = greedy" }, + ], + sdpo: [ + { key: "sdpo_alpha", label: "Alpha (divergence)", kind: "float", def: "0.5", + hint: "0 = forward KL, 1 = reverse KL, 0.5 = JSD" }, + { key: "sdpo_group_size", label: "Group size", kind: "int", def: "4", + hint: "rollouts per prompt" }, + { key: "sdpo_max_completion_length", label: "Max completion", kind: "int", def: "256" }, + { key: "sdpo_temperature", label: "Rollout temperature", kind: "float", def: "1", + hint: "0 = greedy" }, + { key: "sdpo_success_threshold", label: "Success threshold", kind: "float", def: "1", + hint: "reward that makes a rollout a reusable solution" }, + { key: "sdpo_teacher_ema", label: "Teacher EMA", kind: "float", def: "0.05", + hint: "0 = frozen teacher, 1 = live student" }, + ], prompt: [{ key: "num_virtual_tokens", label: "Virtual tokens", kind: "int", def: "16" }], ptuning: [{ key: "num_virtual_tokens", label: "Virtual tokens", kind: "int", def: "16" }], }; @@ -115,14 +136,14 @@ const FAMILY: Record = { lora: "peft", qlora: "peft", dora: "peft", adapter: "peft", bitfit: "peft", prompt: "peft", ptuning: "peft", full: "sft", cpt: "sft", - dpo: "rl", grpo: "rl", + dpo: "rl", grpo: "rl", sdft: "rl", sdpo: "rl", more: "memory", more_plus: "memory", }; const FAMILY_LABEL: Record = { peft: "PEFT · parameter-efficient", sft: "SFT · full & continued pretraining", - rl: "Preference & RL", + rl: "Preference · RL · distillation", memory: "Memory · retrieval", other: "Other", }; @@ -183,7 +204,9 @@ export default function Train({ methods }: { methods: MethodInfo[] }) { const advParams = advancedParams(methodInfo); const configParams = allParams(methodInfo); // held-out eval only applies when it's meaningful and not already provided - const useHoldout = evalSplit && methodInfo?.trainer !== "grpo" && !meta?.eval_split; + const noHoldoutTrainers = ["grpo", "sdft", "sdpo"]; + const useHoldout = evalSplit && !meta?.eval_split + && !noHoldoutTrainers.includes(methodInfo?.trainer ?? ""); const ready = Boolean(ds && model && method); const canNext = [Boolean(ds), Boolean(model), Boolean(method), true][step]; @@ -457,7 +480,7 @@ export default function Train({ methods }: { methods: MethodInfo[] }) {
eval uses the dataset's own {meta.eval_split} split
- ) : methodInfo?.trainer === "grpo" ? null : ( + ) : noHoldoutTrainers.includes(methodInfo?.trainer ?? "") ? null : (