mlxsmith 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

mlxsmith/train/online_dpo.py

@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+from typing import Iterable, Optional
+
+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, clip_grad_norm
+from ..llm.registry import get_llm_backend
+from ..llm.backend import BackendNotAvailable
+from ..sdk.losses import preference_loss
+from ..verifiers.llm_judge import verify as judge_verify
+from .lora import LoRAConfig
+
+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 row.get("question") 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_online_dpo(
+    project_root: Path,
+    cfg: ProjectConfig,
+    data_path: Path,
+    model_id_or_path: str,
+    accel: str,
+    *,
+    judge_model: Optional[str] = None,
+    judge_backend: str = "mlx-lm",
+    rubric: Optional[str] = None,
+    group_size: Optional[int] = None,
+    max_new_tokens: Optional[int] = None,
+    temperature: Optional[float] = None,
+    judge_mock_response: Optional[str | list[str]] = None,
+) -> RunPaths:
+    run = new_run(project_root, "online_dpo")
+    snapshot_config(cfg.model_dump(), run.config_snapshot_path)
+
+    prompts = list(_iter_prompts(data_path))
+    if not prompts:
+        raise RuntimeError("No prompts found in online DPO dataset")
+
+    backend = get_backend(accel)
+    backend.patch()
+    console.print(f"[bold]ONLINE-DPO[/bold] run: {run.run_dir.name} accel={backend.name}")
+
+    llm = get_llm_backend(cfg.model.backend)
+    base_model, adapter_path, _meta = resolve_model_spec(project_root, model_id_or_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.\nmodel={model_id_or_path}\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,
+        optimizer=cfg.train.optimizer,
+        optimizer_kwargs=cfg.train.optimizer_kwargs,
+    )
+
+    loss_type = str(cfg.pref.loss_type or cfg.pref.algo)
+    beta = float(cfg.pref.beta)
+    kl_coeff = float(cfg.pref.kl_coeff)
+    delta = float(cfg.pref.delta)
+
+    total = int(cfg.train.iters)
+    grad_accum = max(1, int(cfg.train.grad_accum))
+    max_grad_norm = float(getattr(cfg.train, "max_grad_norm", 0.0))
+    group = int(group_size or cfg.rft.rollouts or 4)
+    max_new = int(max_new_tokens or cfg.rft.max_new_tokens)
+    temp = float(temperature if temperature is not None else cfg.rft.temperature)
+
+    rng = random.Random(cfg.train.seed)
+    accum_grads = None
+    accum_loss = 0.0
+    accum_count = 0
+
+    def _next_mock(idx: int) -> Optional[str]:
+        if judge_mock_response is None:
+            return None
+        if isinstance(judge_mock_response, list):
+            if not judge_mock_response:
+                return None
+            return judge_mock_response[min(idx, len(judge_mock_response) - 1)]
+        return judge_mock_response
+
+    for step in range(1, total + 1):
+        prompt = rng.choice(prompts)
+        candidates: list[tuple[str, float]] = []
+        for k in range(group):
+            gen = llm.generate(
+                prompt,
+                max_new_tokens=max_new,
+                temperature=temp,
+                seed=rng.randint(0, 2**31 - 1),
+            )
+            completion = gen.text[len(prompt) :] if gen.text.startswith(prompt) else gen.text
+            res = judge_verify(
+                prompt,
+                completion,
+                str(run.artifacts_dir),
+                model=judge_model,
+                backend=judge_backend,
+                rubric=rubric,
+                reward_mode="score",
+                mock_response=_next_mock(k),
+            )
+            reward = float(getattr(res, "reward", 0.0))
+            candidates.append((completion, reward))
+
+        if len(candidates) < 2:
+            continue
+        chosen, chosen_r = max(candidates, key=lambda x: x[1])
+        rejected, rejected_r = min(candidates, key=lambda x: x[1])
+        if chosen == 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):
+            return preference_loss(
+                llm,
+                chosen_ids,
+                rejected_ids,
+                prompt_len_chosen=p_len_c,
+                prompt_len_rejected=p_len_r,
+                algo=loss_type,
+                beta=beta,
+                reference_backend=ref_llm,
+                kl_coeff=kl_coeff,
+                train_on_prompt=bool(cfg.train.train_on_prompt),
+                delta=delta,
+            )
+
+        lval, grads = llm.value_and_grad(loss_fn)
+        accum_loss += float(lval.item()) if hasattr(lval, "item") else float(lval)
+        accum_count += 1
+        if grads is not None:
+            accum_grads = tree_add(accum_grads, grads)
+
+        if step % grad_accum == 0:
+            if accum_grads is not None:
+                scaled = tree_scale(accum_grads, 1.0 / grad_accum)
+                if max_grad_norm > 0:
+                    scaled = clip_grad_norm(scaled, max_grad_norm)
+                llm.apply_grads(opt, scaled)
+            accum_grads = None
+            accum_loss = 0.0
+            accum_count = 0
+
+        if step % cfg.train.log_every == 0 or step == 1 or step == total:
+            avg_loss = accum_loss / max(1, accum_count) if accum_count else float(lval)
+            write_jsonl(
+                run.metrics_path,
+                [
+                    {
+                        "ts": now_ts(),
+                        "step": step,
+                        "kind": "online_dpo",
+                        "algo": loss_type,
+                        "loss": avg_loss,
+                        "reward_best": chosen_r,
+                        "reward_worst": rejected_r,
+                        "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": "online_dpo",
+                },
+            )
+
+    console.print(f"[green]Saved adapter[/green] {run.adapter_dir}")
+    return run

mlxsmith/train/pref.py

@@ -10,7 +10,8 @@ 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 ..util import write_jsonl, now_ts, tree_add, tree_scale, clip_grad_norm
+from ..sdk.losses import preference_loss
 from ..llm.registry import get_llm_backend
 from ..llm.backend import BackendNotAvailable
 from .lora import LoRAConfig
@@ -28,7 +29,8 @@ def run_pref(project_root: Path, cfg: ProjectConfig, data_dir: Path, base_model_
 
     backend = get_backend(accel)
     backend.patch()
-    console.print(f"[bold]PREF[/bold] run: {run.run_dir.name} algo={cfg.pref.algo} accel={backend.name}")
+    loss_type = str(cfg.pref.loss_type or cfg.pref.algo)
+    console.print(f"[bold]PREF[/bold] run: {run.run_dir.name} algo={loss_type} accel={backend.name}")
 
     prefs_path = data_dir / "train.jsonl"
     if not prefs_path.exists():
@@ -79,10 +81,16 @@ def run_pref(project_root: Path, cfg: ProjectConfig, data_dir: Path, base_model_
         except BackendNotAvailable:
             ref_llm = None
 
-    opt, _params = llm.optimizer_and_params(lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
+    opt, _params = llm.optimizer_and_params(
+        lr=cfg.train.lr,
+        weight_decay=cfg.train.weight_decay,
+        optimizer=cfg.train.optimizer,
+        optimizer_kwargs=cfg.train.optimizer_kwargs,
+    )
 
     beta = float(cfg.pref.beta)
     kl_coeff = float(cfg.pref.kl_coeff)
+    delta = float(cfg.pref.delta)
     rng = random.Random(cfg.train.seed)
     total = int(cfg.train.iters)
     grad_accum = max(1, int(cfg.train.grad_accum))
@@ -114,30 +122,19 @@ def run_pref(project_root: Path, cfg: ProjectConfig, data_dir: Path, base_model_
                 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
+            return preference_loss(
+                llm,
+                chosen_ids,
+                rejected_ids,
+                prompt_len_chosen=p_len_c,
+                prompt_len_rejected=p_len_r,
+                algo=loss_type,
+                beta=beta,
+                reference_backend=ref_llm,
+                kl_coeff=kl_coeff,
+                train_on_prompt=train_on_prompt,
+                delta=delta,
+            )
 
         lval, grads = llm.value_and_grad(loss_fn)
         if grads is not None:
@@ -145,7 +142,11 @@ def run_pref(project_root: Path, cfg: ProjectConfig, data_dir: Path, base_model_
 
         if step % grad_accum == 0:
             if accum_grads is not None:
-                llm.apply_grads(opt, tree_scale(accum_grads, 1.0 / grad_accum))
+                scaled = tree_scale(accum_grads, 1.0 / grad_accum)
+                max_grad_norm = float(getattr(cfg.train, "max_grad_norm", 0.0))
+                if max_grad_norm > 0:
+                    scaled = clip_grad_norm(scaled, max_grad_norm)
+                llm.apply_grads(opt, scaled)
             accum_grads = None
 
         if step % cfg.train.log_every == 0 or step == 1 or step == total:
@@ -156,9 +157,10 @@ def run_pref(project_root: Path, cfg: ProjectConfig, data_dir: Path, base_model_
                         "ts": now_ts(),
                         "step": step,
                         "kind": "pref",
-                        "algo": cfg.pref.algo,
+                        "algo": loss_type,
                         "beta": beta,
                         "kl_coeff": kl_coeff,
+                        "delta": delta,
                         "loss": float(lval.item()) if hasattr(lval, "item") else float(lval),
                         "accel": backend.name,
                     }

mlxsmith/train/rft.py

@@ -46,12 +46,13 @@ def _rollout_token_env(
     max_steps: int,
     temperature: float,
     seed: int,
-) -> tuple[list[int], int, str, float, dict, int]:
+) -> tuple[list[int], int, str, float, dict, int, list[float]]:
     obs = env.initial_observation()
     obs_tokens, reward, done, info = _normalize_observation(obs)
     prompt_len = len(obs_tokens)
     full_tokens = list(obs_tokens)
     gen_tokens = 0
+    behavior_logprobs: list[float] = []
 
     for idx in range(max_steps):
         if done:
@@ -67,6 +68,8 @@ def _rollout_token_env(
             logprobs=0,
         )
         new_token = int(gen.token_ids[-1])
+        if gen.logprobs:
+            behavior_logprobs.append(float(gen.logprobs[-1]))
         full_tokens.append(new_token)
         gen_tokens += 1
 
@@ -77,7 +80,62 @@ def _rollout_token_env(
         obs_tokens = list(step.observation) if step.observation else list(full_tokens)
 
     completion = llm.decode(full_tokens[prompt_len:])
-    return full_tokens, prompt_len, completion, reward, info, gen_tokens
+    return full_tokens, prompt_len, completion, reward, info, gen_tokens, behavior_logprobs
+
+
+def _pg_loss(
+    llm,
+    token_ids: list[int],
+    *,
+    prompt_len: int,
+    advantage: float,
+    behavior_logprobs: list[float] | None,
+    loss_type: str,
+    epsilon_low: float,
+    epsilon_high: float,
+    token_level: bool,
+    ref_llm=None,
+    kl_coeff: float = 0.0,
+):
+    mx = llm.mx  # type: ignore
+    logp = None
+    if token_level:
+        token_logps, _ = llm.token_logprobs(
+            token_ids, prompt_len=prompt_len, top_k=0, include_prompt=False
+        )
+        if not token_logps:
+            return mx.array(0.0)
+        if loss_type == "dapo" and behavior_logprobs:
+            n = min(len(token_logps), len(behavior_logprobs))
+            total = mx.array(0.0)
+            for lp, bp in zip(token_logps[:n], behavior_logprobs[:n]):
+                ratio = mx.exp(mx.array(lp) - mx.array(bp))
+                clipped = mx.minimum(
+                    mx.maximum(ratio, mx.array(1.0 - epsilon_low)),
+                    mx.array(1.0 + epsilon_high),
+                )
+                total = total + clipped
+            loss = -mx.array(float(advantage)) * total / mx.array(float(n))
+        else:
+            avg_logp = sum(token_logps) / float(len(token_logps))
+            loss = -mx.array(float(advantage)) * mx.array(avg_logp)
+    else:
+        logp = llm.sequence_logprob(token_ids, prompt_len=prompt_len)
+        if loss_type == "dapo" and behavior_logprobs:
+            behavior = sum(behavior_logprobs)
+            ratio = mx.exp(logp - mx.array(float(behavior)))
+            clipped = mx.minimum(mx.maximum(ratio, mx.array(1.0 - epsilon_low)), mx.array(1.0 + epsilon_high))
+            loss = -mx.array(float(advantage)) * clipped
+        else:
+            loss = -mx.array(float(advantage)) * logp
+
+    if ref_llm is not None and kl_coeff > 0:
+        if logp is None:
+            logp = llm.sequence_logprob(token_ids, prompt_len=prompt_len)
+        ref_logp = ref_llm.sequence_logprob(token_ids, prompt_len=prompt_len)
+        loss = loss + mx.array(float(kl_coeff)) * (logp - ref_logp)
+
+    return loss
 
 
 def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_path: Path, base_model_path: Path, accel: str) -> RunPaths:
@@ -146,7 +204,12 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
         except BackendNotAvailable:
             ref_llm = None
 
-    opt, _params = llm.optimizer_and_params(lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
+    opt, _params = llm.optimizer_and_params(
+        lr=cfg.train.lr,
+        weight_decay=cfg.train.weight_decay,
+        optimizer=cfg.train.optimizer,
+        optimizer_kwargs=cfg.train.optimizer_kwargs,
+    )
 
     rng = random.Random(cfg.train.seed)
     total_iters = int(cfg.train.iters)
@@ -155,6 +218,10 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
     max_new = int(getattr(cfg.rft, "max_new_tokens", 256))
     kl_coeff = float(cfg.rft.kl_coeff)
     normalize_adv = bool(cfg.rft.normalize_advantage)
+    loss_type = str(cfg.rft.loss_type or cfg.rft.algo)
+    epsilon_low = float(getattr(cfg.rft, "epsilon_low", 0.2))
+    epsilon_high = float(getattr(cfg.rft, "epsilon_high", epsilon_low))
+    token_level = bool(getattr(cfg.rft, "token_level_loss", False))
 
     if token_env_spec is not None:
         base_name = env.get("name") or "token_env"
@@ -204,7 +271,7 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
                         seed=rng.randint(0, 2**31 - 1),
                     )
 
-                token_ids, prompt_len, completion, reward, info, gen_count = _rollout_token_env(
+                token_ids, prompt_len, completion, reward, info, gen_count, behavior_logprobs = _rollout_token_env(
                     llm,
                     env_instance,
                     max_steps=max_new,
@@ -225,27 +292,35 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
                             continue
 
                 passed = bool(info.get("passed", reward > 0.0))
-                gens.append((token_ids, prompt_len, completion, passed, reward, info))
+                gens.append((token_ids, prompt_len, completion, passed, reward, info, behavior_logprobs))
 
             gen_elapsed = max(time.time() - gen_start, 1e-6)
             tps = gen_tokens / gen_elapsed
 
-            mean_r = sum(r for *_rest, r, _info in gens) / max(1, len(gens))
+            mean_r = sum(r for *_rest, r, _info, _bp in gens) / max(1, len(gens))
             std_r = (
-                sum((r - mean_r) ** 2 for *_rest, r, _info in gens) / max(1, len(gens))
+                sum((r - mean_r) ** 2 for *_rest, r, _info, _bp in gens) / max(1, len(gens))
             ) ** 0.5
-            advs = [r - mean_r for *_rest, r, _info in gens]
-            if normalize_adv and std_r > 1e-6:
+            advs = [r - mean_r for *_rest, r, _info, _bp in gens]
+            if loss_type != "dr_grpo" and normalize_adv and std_r > 1e-6:
                 advs = [a / std_r for a in advs]
 
             def loss_fn(_model):
                 loss = llm.mx.array(0.0)  # type: ignore
-                for (token_ids, prompt_len, _comp, _passed, _reward, _info), adv in zip(gens, advs):
-                    logp = llm.sequence_logprob(token_ids, prompt_len=prompt_len)
-                    pg = -llm.mx.array(float(adv)) * logp  # type: ignore
-                    if ref_llm is not None and kl_coeff > 0:
-                        ref_logp = ref_llm.sequence_logprob(token_ids, prompt_len=prompt_len)
-                        pg = pg + llm.mx.array(kl_coeff) * (logp - ref_logp)  # type: ignore
+                for (token_ids, prompt_len, _comp, _passed, _reward, _info, bps), adv in zip(gens, advs):
+                    pg = _pg_loss(
+                        llm,
+                        token_ids,
+                        prompt_len=prompt_len,
+                        advantage=float(adv),
+                        behavior_logprobs=bps,
+                        loss_type=loss_type,
+                        epsilon_low=epsilon_low,
+                        epsilon_high=epsilon_high,
+                        token_level=token_level,
+                        ref_llm=ref_llm,
+                        kl_coeff=kl_coeff,
+                    )
                     loss = loss + pg
                 return loss / llm.mx.array(float(len(gens)))  # type: ignore
 
@@ -256,8 +331,8 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
             best_idx = max(range(len(gens)), key=lambda i: gens[i][4])
             best = gens[best_idx]
             pass_at_1 = 1.0 if gens[0][3] else 0.0
-            pass_at_k = 1.0 if any(passed for *_g, passed, _r, _i in gens) else 0.0
-            acceptance = sum(1 for *_g, passed, _r, _i in gens if passed) / max(1, len(gens))
+            pass_at_k = 1.0 if any(g[3] for g in gens) else 0.0
+            acceptance = sum(1 for g in gens if g[3]) / max(1, len(gens))
 
             latency_summary = latency_summary_ms(verifier_latencies_ms)
             per_verifier_summary = {
@@ -269,7 +344,7 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
                     "ts": now_ts(),
                     "step": step,
                     "kind": "rft",
-                    "algo": cfg.rft.algo,
+                    "algo": loss_type,
                     "task_id": task_id,
                     "mean_reward": mean_r,
                     "std_reward": std_r,
@@ -290,7 +365,7 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
                     metrics["verifier_latency_ms_by_path"] = per_verifier_summary
                 write_jsonl(run.metrics_path, [metrics])
 
-            for (token_ids, prompt_len, completion, passed, reward, _info) in gens:
+            for (token_ids, prompt_len, completion, passed, reward, _info, _bps) in gens:
                 if passed:
                     prompt_text = llm.decode(token_ids[:prompt_len]) if prompt_len > 0 else ""
                     write_jsonl(
@@ -331,14 +406,16 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
         per_verifier_latencies: dict[str, list[float]] = {}
 
         for k in range(rollouts):
-            gen = llm.generate(
+            gen = llm.generate_with_logprobs(
                 prompt,
                 max_new_tokens=max_new,
                 temperature=temperature,
                 seed=rng.randint(0, 2**31 - 1),
+                logprobs=0,
             )
             completion = gen.text[len(prompt) :] if gen.text.startswith(prompt) else gen.text
             gen_tokens += max(0, len(gen.token_ids) - gen.prompt_len)
+            behavior_logprobs = list(gen.logprobs) if gen.logprobs is not None else []
 
             wdir = ensure_dir(run.artifacts_dir / task_id / f"step_{step:06d}" / f"rollout_{k:02d}")
             if "tests" in task:
@@ -359,27 +436,35 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
 
             passed = bool(getattr(res, "passed", False))
             reward = float(getattr(res, "reward", 0.0))
-            gens.append((gen, completion, passed, reward))
+            gens.append((gen, completion, passed, reward, behavior_logprobs))
 
         gen_elapsed = max(time.time() - gen_start, 1e-6)
         tps = gen_tokens / gen_elapsed
 
-        mean_r = sum(r for *_rest, r in gens) / max(1, len(gens))
+        mean_r = sum(r for *_rest, r, _bp in gens) / max(1, len(gens))
         std_r = (
-            sum((r - mean_r) ** 2 for *_rest, r in gens) / max(1, len(gens))
+            sum((r - mean_r) ** 2 for *_rest, r, _bp in gens) / max(1, len(gens))
         ) ** 0.5
-        advs = [r - mean_r for *_rest, r in gens]
-        if normalize_adv and std_r > 1e-6:
+        advs = [r - mean_r for *_rest, r, _bp in gens]
+        if loss_type != "dr_grpo" and normalize_adv and std_r > 1e-6:
             advs = [a / std_r for a in advs]
 
         def loss_fn(_model):
             loss = llm.mx.array(0.0)  # type: ignore
-            for (gen, _comp, _passed, _reward), adv in zip(gens, advs):
-                logp = llm.sequence_logprob(gen.token_ids, prompt_len=gen.prompt_len)
-                pg = -llm.mx.array(float(adv)) * logp  # type: ignore
-                if ref_llm is not None and kl_coeff > 0:
-                    ref_logp = ref_llm.sequence_logprob(gen.token_ids, prompt_len=gen.prompt_len)
-                    pg = pg + llm.mx.array(kl_coeff) * (logp - ref_logp)  # type: ignore
+            for (gen, _comp, _passed, _reward, bps), adv in zip(gens, advs):
+                pg = _pg_loss(
+                    llm,
+                    list(gen.token_ids),
+                    prompt_len=gen.prompt_len,
+                    advantage=float(adv),
+                    behavior_logprobs=bps,
+                    loss_type=loss_type,
+                    epsilon_low=epsilon_low,
+                    epsilon_high=epsilon_high,
+                    token_level=token_level,
+                    ref_llm=ref_llm,
+                    kl_coeff=kl_coeff,
+                )
                 loss = loss + pg
             return loss / llm.mx.array(float(len(gens)))  # type: ignore
 
@@ -390,8 +475,8 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
         best_idx = max(range(len(gens)), key=lambda i: gens[i][3])
         best = gens[best_idx]
         pass_at_1 = 1.0 if gens[0][2] else 0.0
-        pass_at_k = 1.0 if any(passed for _g, _c, passed, _r in gens) else 0.0
-        acceptance = sum(1 for *_rest, passed, _reward in gens if passed) / max(1, len(gens))
+        pass_at_k = 1.0 if any(g[2] for g in gens) else 0.0
+        acceptance = sum(1 for g in gens if g[2]) / max(1, len(gens))
 
         latency_summary = latency_summary_ms([t * 1000.0 for t in verifier_times])
         per_verifier_summary = {
@@ -406,10 +491,10 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
                         "ts": now_ts(),
                         "step": step,
                         "kind": "rft",
-                        "algo": cfg.rft.algo,
-                        "task_id": task_id,
-                        "mean_reward": mean_r,
-                        "std_reward": std_r,
+                    "algo": loss_type,
+                    "task_id": task_id,
+                    "mean_reward": mean_r,
+                    "std_reward": std_r,
                         "best_reward": best[3],
                         "best_passed": best[2],
                         "pass@1": pass_at_1,
@@ -429,7 +514,7 @@ def run_rft(project_root: Path, cfg: ProjectConfig, env_path: Path, verifier_pat
                 ],
             )
 
-        for (gen, completion, passed, reward) in gens:
+        for (gen, completion, passed, reward, _bp) in gens:
             if passed:
                 write_jsonl(
                     accepted_path,