diff --git a/fusion_mlx/admin/fine_tune_route.py b/fusion_mlx/admin/fine_tune_route.py index 8f0feac..73bed09 100644 --- a/fusion_mlx/admin/fine_tune_route.py +++ b/fusion_mlx/admin/fine_tune_route.py @@ -856,4 +856,100 @@ async def delete_reward_job( return {"status": "deleted"} +@_router.post("/api/fine-tune/reward/score") +async def score_reward_endpoint( + request: Request, + is_admin: bool = Depends(require_admin), +): + # Score (prompt, completions) under a trained reward adapter (#431). + # Loads standalone (separate from inference pool), attaches the trained + # value head, scores each completion, evicts. Closes the Phase1 (#424) + # reward-model -> Phase2 (#363 GRPO) loop: this URL is passed as GRPO's + # config.reward_endpoint, which POSTs {prompt, completions} -> {rewards}. + # GRPO's callback protocol is fixed at {prompt, completions} (no + # model_id/adapter_name), so those MUST be supplied as query params in + # the reward_endpoint URL, e.g. .../reward/score?key=T&model_id=X&adapter_name=Y. + body = await request.json() + + model_id = body.get("model_id", "") or request.query_params.get("model_id", "") + adapter_name = body.get("adapter_name", "") or request.query_params.get( + "adapter_name", "" + ) + prompt = body.get("prompt", "") + completions = body.get("completions", []) + + if not model_id: + raise HTTPException(status_code=400, detail="model_id is required") + if not prompt: + raise HTTPException(status_code=400, detail="prompt is required") + if not completions or not isinstance(completions, list): + raise HTTPException( + status_code=400, detail="completions (non-empty list) required" + ) + if not adapter_name: + raise HTTPException( + status_code=400, + detail="adapter_name is required (a trained reward adapter)", + ) + + svc = _get_service() + model_path = svc._resolve_model_path(model_id) + if model_path is None: + raise HTTPException(status_code=404, detail=f"Model not found: {model_id}") + + from fusion_mlx.training.service import ADAPTER_BASE_DIR + + adapter_path = str(ADAPTER_BASE_DIR / model_id / adapter_name) + import os + + if not os.path.isdir(adapter_path): + raise HTTPException( + status_code=404, + detail=f"Adapter not found: {model_id}/{adapter_name}", + ) + + from fusion_mlx.training.reward import RewardScoreResult + from fusion_mlx.training.reward import score_text as reward_score_text + + logger.info( + "reward score endpoint: model=%s adapter=%s n_completions=%d", + model_path, + adapter_path, + len(completions), + ) + + def _run(): + import mlx_lm.utils as mlx_utils + + model, tokenizer = mlx_utils.load(model_path, adapter_path=adapter_path) + try: + rewards = reward_score_text( + model, tokenizer, model_path, prompt, completions, adapter_path + ) + return RewardScoreResult( + rewards=rewards, model_id=model_id, adapter_name=adapter_name + ) + finally: + del model + del tokenizer + import gc + + import mlx.core as mx + + gc.collect() + mx.clear_cache() + logger.info("reward score endpoint: model evicted") + + try: + result = await asyncio.to_thread(_run) + except ValueError as e: + logger.exception("reward score endpoint: bad adapter") + raise HTTPException(status_code=400, detail=str(e)) + except Exception as e: + logger.exception("reward score endpoint failed") + raise HTTPException(status_code=500, detail=f"Scoring failed: {e}") + + return result.to_dict() + + router = _router diff --git a/fusion_mlx/training/reward.py b/fusion_mlx/training/reward.py index 6feb8b3..f9ba6a3 100644 --- a/fusion_mlx/training/reward.py +++ b/fusion_mlx/training/reward.py @@ -220,3 +220,101 @@ def save_adapter(self, adapter_path): weights = dict(tree_flatten(self.model.trainable_parameters())) mx.save_safetensors(adapter_path, weights) logger.info("REWARD: saved reward adapter to %s", adapter_path) + + +@dataclass +class RewardScoreResult: + rewards: list + model_id: str + adapter_name: str + + def to_dict(self): + return { + "rewards": self.rewards, + "model_id": self.model_id, + "adapter_name": self.adapter_name, + } + + +def score_text(model, tokenizer, model_path, prompt, completions, adapter_path=None): + # Inference-time scalar reward scoring for (prompt, completion) pairs, + # using a trained reward adapter (LoRA backbone + value head, #424). + # Mirror of RewardTrainer._score but non-differentiable: forward the + # concatenated sequence, take the last-token hidden, project via the + # value head to a scalar. Loads value_head weights from the adapter's + # safetensors (the standard mlx_utils.load adapter_path path applies + # LoRA but does not restore the custom value_head submodule). + import os + + from safetensors import safe_open + + logger.info( + "reward score_text: model=%s adapter=%s prompt_len=%d n_completions=%d", + model_path, + adapter_path, + len(prompt), + len(completions), + ) + + head = None + if adapter_path: + weights_file = os.path.join(adapter_path, "adapters.safetensors") + vh_keys = {} + if os.path.isfile(weights_file): + with safe_open(weights_file, "mlx") as f: + for k in list(f.keys()): + if k.startswith("value_head."): + vh_keys[k] = f.get_tensor(k) + if "value_head.proj.weight" in vh_keys: + hidden = int(vh_keys["value_head.proj.weight"].shape[1]) + head = _ValueHead(hidden) + head.load_weights( + [ + ("proj.weight", vh_keys["value_head.proj.weight"]), + ("proj.bias", vh_keys["value_head.proj.bias"]), + ] + ) + logger.info("reward score_text: value head loaded hidden=%d", hidden) + else: + logger.warning( + "reward score_text: adapter has no value_head weights (%s), " + "scoring with untrained head", + weights_file, + ) + + if getattr(model, "value_head", None) is None: + if head is not None: + model.value_head = head + else: + raise ValueError( + "reward score_text: model has no value_head and adapter " + "provided no value_head weights; not a reward model" + ) + + def _score_one(prompt_ids, completion_ids): + full = mx.concatenate([prompt_ids, completion_ids]) + trunk = getattr(model, "transformer", None) or getattr(model, "model", None) + if trunk is not None: + hidden = trunk(full[None, :]) + if isinstance(hidden, tuple): + hidden = hidden[0] + hidden = hidden[0] + else: + logits = _model_forward_logits(model, full) + if isinstance(logits, tuple): + logits = logits[0] + n_comp = int(completion_ids.shape[0]) + hidden = mx.mean(logits[0, -n_comp:, :].astype(mx.float32), axis=0) + hidden = mx.expand_dims(hidden, 0) + s = model.value_head(hidden) + mx.eval(s) + return float(s) + + rewards = [] + prompt_ids = mx.array(tokenizer.encode(prompt)) + for completion in completions: + completion_ids = mx.array(tokenizer.encode(completion)) + rewards.append(_score_one(prompt_ids, completion_ids)) + + logger.info("reward score_text: rewards=%s", rewards) + return rewards diff --git a/tests/unit/test_reward_score_route.py b/tests/unit/test_reward_score_route.py new file mode 100644 index 0000000..190a607 --- /dev/null +++ b/tests/unit/test_reward_score_route.py @@ -0,0 +1,217 @@ +# Route tests for /admin/api/fine-tune/reward/score (#431). +# Minimal FastAPI app + dependency override for require_admin. The model +# load (mlx_utils.load) and reward scoring (reward.score_text) are both +# monkey-patched so no real model is loaded in CI. + +from __future__ import annotations + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fusion_mlx.admin.auth import require_admin +from fusion_mlx.admin.fine_tune_route import set_fine_tune_context +from fusion_mlx.admin.routes import router as admin_router + + +class _FakeEntry: + def __init__(self, model_path="/tmp/fake-model", model_type="llm"): + self.model_path = model_path + self.model_type = model_type + self.engine = None + + +class _FakePool: + def __init__(self, model_path="/tmp/fake-model"): + self._entries = {"m1": _FakeEntry(model_path=model_path)} + + def get_entry(self, model_id): + return self._entries.get(model_id) + + def unload_if_idle_unpinned(self, model_id): + return False + + +class _FakeService: + def __init__(self, model_path="/tmp/fake-model"): + self._pool = _FakePool(model_path=model_path) + + def set_engine_pool(self, pool): + self._pool = pool + + def _resolve_model_path(self, model_id): + entry = self._pool.get_entry(model_id) + if entry is not None and hasattr(entry, "model_path"): + return entry.model_path + return model_id + + +def _build_app(service=None): + app = FastAPI() + set_fine_tune_context(_FakePool(), service or _FakeService()) + app.include_router(admin_router) + app.dependency_overrides[require_admin] = lambda: True + return app + + +def _patch_load_and_score(monkeypatch, rewards_out): + # Stub mlx_utils.load (called inside the handler's _run) and + # reward.score_text so no real model is touched. + import fusion_mlx.training.reward as _reward + + captured = {} + + def fake_score_text( + model, tokenizer, model_path, prompt, completions, adapter_path=None + ): + captured.update( + prompt=prompt, + completions=completions, + adapter_path=adapter_path, + ) + return list(rewards_out) + + monkeypatch.setattr(_reward, "score_text", fake_score_text) + + import mlx_lm.utils as _mlx_utils + + def fake_load(model_path, adapter_path=None): + captured["load_model_path"] = model_path + captured["load_adapter_path"] = adapter_path + return ("fake-model", "fake-tokenizer") + + monkeypatch.setattr(_mlx_utils, "load", fake_load) + return captured + + +def test_reward_score_happy_path(monkeypatch, tmp_path): + fake_adapter = tmp_path / "m1" / "rm_adapter" + fake_adapter.mkdir(parents=True) + monkeypatch.setattr("fusion_mlx.training.service.ADAPTER_BASE_DIR", tmp_path) + + captured = _patch_load_and_score(monkeypatch, [6.7, -7.2]) + + app = _build_app() + client = TestClient(app) + + resp = client.post( + "/admin/api/fine-tune/reward/score", + json={ + "model_id": "m1", + "adapter_name": "rm_adapter", + "prompt": "What is 2+2?", + "completions": ["4", "five"], + }, + ) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["rewards"] == [6.7, -7.2] + assert body["model_id"] == "m1" + assert body["adapter_name"] == "rm_adapter" + assert captured["prompt"] == "What is 2+2?" + assert captured["completions"] == ["4", "five"] + assert captured["load_adapter_path"] is not None + + +def test_reward_score_query_params(monkeypatch, tmp_path): + # GRPO's callback protocol sends only {prompt, completions} in the body; + # model_id/adapter_name must come from query params in the reward_endpoint URL. + fake_adapter = tmp_path / "m1" / "rm_adapter" + fake_adapter.mkdir(parents=True) + monkeypatch.setattr("fusion_mlx.training.service.ADAPTER_BASE_DIR", tmp_path) + + captured = _patch_load_and_score(monkeypatch, [1.0]) + + app = _build_app() + client = TestClient(app) + + resp = client.post( + "/admin/api/fine-tune/reward/score" "?model_id=m1&adapter_name=rm_adapter", + json={"prompt": "Hi", "completions": ["x"]}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["rewards"] == [1.0] + assert captured["prompt"] == "Hi" + + +def test_reward_score_missing_model_id(): + app = _build_app() + client = TestClient(app) + resp = client.post( + "/admin/api/fine-tune/reward/score", + json={"adapter_name": "a", "prompt": "Hi", "completions": ["x"]}, + ) + assert resp.status_code == 400 + assert "model_id" in resp.json()["detail"] + + +def test_reward_score_missing_adapter_name(): + app = _build_app() + client = TestClient(app) + resp = client.post( + "/admin/api/fine-tune/reward/score", + json={"model_id": "m1", "prompt": "Hi", "completions": ["x"]}, + ) + assert resp.status_code == 400 + assert "adapter_name" in resp.json()["detail"] + + +def test_reward_score_missing_completions(): + app = _build_app() + client = TestClient(app) + resp = client.post( + "/admin/api/fine-tune/reward/score", + json={"model_id": "m1", "adapter_name": "a", "prompt": "Hi"}, + ) + assert resp.status_code == 400 + assert "completions" in resp.json()["detail"] + + +def test_reward_score_adapter_not_found(monkeypatch, tmp_path): + monkeypatch.setattr("fusion_mlx.training.service.ADAPTER_BASE_DIR", tmp_path) + app = _build_app() + client = TestClient(app) + resp = client.post( + "/admin/api/fine-tune/reward/score", + json={ + "model_id": "m1", + "adapter_name": "nope", + "prompt": "Hi", + "completions": ["x"], + }, + ) + assert resp.status_code == 404 + assert "Adapter not found" in resp.json()["detail"] + + +def test_reward_score_not_a_reward_model(monkeypatch, tmp_path): + # Adapter dir exists but score_text raises ValueError -> endpoint 400. + fake_adapter = tmp_path / "m1" / "sft_only" + fake_adapter.mkdir(parents=True) + monkeypatch.setattr("fusion_mlx.training.service.ADAPTER_BASE_DIR", tmp_path) + + import fusion_mlx.training.reward as _reward + + def fake_score_text( + model, tokenizer, model_path, prompt, completions, adapter_path=None + ): + raise ValueError("not a reward model") + + monkeypatch.setattr(_reward, "score_text", fake_score_text) + + import mlx_lm.utils as _mlx_utils + + monkeypatch.setattr(_mlx_utils, "load", lambda p, adapter_path=None: ("m", "t")) + + app = _build_app() + client = TestClient(app) + resp = client.post( + "/admin/api/fine-tune/reward/score", + json={ + "model_id": "m1", + "adapter_name": "sft_only", + "prompt": "Hi", + "completions": ["x"], + }, + ) + assert resp.status_code == 400 + assert "not a reward model" in resp.json()["detail"]