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

mlxsmith/rlm/loop.py

@@ -49,6 +49,28 @@ from .weights import (
 console = Console()
 
 
+def _run_task_verifier(cfg: ProjectConfig, task_prompt: str, completion: str, workdir: Path) -> tuple[bool, float, float]:
+    """Execute verifier for a task and return (passed, reward, latency_ms)."""
+    t0 = time.time()
+    if cfg.rlm.verifier_backend == "docker":
+        res = docker_verify(
+            task_prompt,
+            completion,
+            str(workdir),
+            timeout_s=int(cfg.rlm.verifier_timeout_s),
+            image=cfg.rlm.docker_image,
+            memory_mb=int(cfg.rlm.docker_memory_mb),
+            cpus=float(cfg.rlm.docker_cpus),
+            pids=int(cfg.rlm.docker_pids),
+        )
+    else:
+        res = pytest_verify(task_prompt, completion, str(workdir), timeout_s=int(cfg.rlm.verifier_timeout_s))
+    latency_ms = (time.time() - t0) * 1000.0
+    passed = bool(getattr(res, "passed", False))
+    reward = float(getattr(res, "reward", 0.0))
+    return passed, reward, latency_ms
+
+
 def _score_from_eval(result_path: Path) -> float:
     try:
         data = json.loads(result_path.read_text(encoding="utf-8"))
@@ -199,7 +221,12 @@ def run_rlm(
                 trust_remote_code=cfg.model.trust_remote_code,
             )
 
-        opt, _params = train_llm.optimizer_and_params(lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
+        opt, _params = train_llm.optimizer_and_params(
+            lr=cfg.train.lr,
+            weight_decay=cfg.train.weight_decay,
+            optimizer=cfg.train.optimizer,
+            optimizer_kwargs=cfg.train.optimizer_kwargs,
+        )
 
         corpus_rows = load_corpus(corpus_path, max_size=int(rlm_cfg.corpus_max))
         existing_prompts = [row.get("prompt", "") for row in corpus_rows if row.get("prompt")]
@@ -371,9 +398,14 @@ def run_rlm(
 # Multi-Process Orchestrated RLM
 # =============================================================================
 
-from ..orchestrator.queue import MessageQueue, MessageType, Message  # noqa: E402
-from ..orchestrator.inference_worker import InferenceConfig, run_inference_worker  # noqa: E402
-from ..orchestrator.trainer_worker import TrainerConfig, run_trainer_worker  # noqa: E402
+def _lazy_import_orchestrator():
+    """Lazy import to break circular dependency with orchestrator module."""
+    global MessageQueue, MessageType, Message
+    global InferenceConfig, run_inference_worker
+    global TrainerConfig, run_trainer_worker
+    from ..orchestrator.queue import MessageQueue, MessageType, Message  # noqa: E402
+    from ..orchestrator.inference_worker import InferenceConfig, run_inference_worker  # noqa: E402
+    from ..orchestrator.trainer_worker import TrainerConfig, run_trainer_worker  # noqa: E402
 
 
 @dataclass
@@ -402,6 +434,7 @@ class RLMOrchestrator:
         iterations: int = 50,
         resume: bool = False,
     ):
+        _lazy_import_orchestrator()
         self.project_root = project_root
         self.cfg = cfg
         self.model_spec = model_spec
@@ -492,6 +525,8 @@ class RLMOrchestrator:
             trust_remote_code=self.cfg.model.trust_remote_code,
             lr=self.cfg.train.lr,
             weight_decay=self.cfg.train.weight_decay,
+            optimizer=self.cfg.train.optimizer,
+            optimizer_kwargs=self.cfg.train.optimizer_kwargs,
             kl_coeff=self.cfg.rft.kl_coeff,
             normalize_advantage=self.cfg.rft.normalize_advantage,
             lora_r=self.cfg.lora.r,
@@ -602,30 +637,12 @@ class RLMOrchestrator:
                 tests_dir = ensure_dir(wdir / "tests")
                 (tests_dir / "test_task.py").write_text(task.tests, encoding="utf-8")
                 
-                t0 = time.time()
-                if self.cfg.rlm.verifier_backend == "docker":
-                    res = docker_verify(
-                        task.prompt,
-                        completion,
-                        str(wdir),
-                        timeout_s=int(self.cfg.rlm.verifier_timeout_s),
-                        image=self.cfg.rlm.docker_image,
-                        memory_mb=int(self.cfg.rlm.docker_memory_mb),
-                        cpus=float(self.cfg.rlm.docker_cpus),
-                        pids=int(self.cfg.rlm.docker_pids),
-                    )
-                else:
-                    from ..verifiers.pytest_verifier import verify as pytest_verify
-                    res = pytest_verify(
-                        task.prompt,
-                        completion,
-                        str(wdir),
-                        timeout_s=int(self.cfg.rlm.verifier_timeout_s),
-                    )
-                latency_ms = (time.time() - t0) * 1000.0
-                
-                passed = bool(getattr(res, "passed", False))
-                reward = float(getattr(res, "reward", 0.0))
+                passed, reward, latency_ms = _run_task_verifier(
+                    self.cfg,
+                    task.prompt,
+                    completion,
+                    wdir,
+                )
                 
                 rollouts.append(Rollout(
                     task_id=task.id,
@@ -706,30 +723,12 @@ class RLMOrchestrator:
                 tests_dir = ensure_dir(wdir / "tests")
                 (tests_dir / "test_task.py").write_text(task.tests, encoding="utf-8")
 
-                t0 = time.time()
-                if self.cfg.rlm.verifier_backend == "docker":
-                    res = docker_verify(
-                        task.prompt,
-                        completion,
-                        str(wdir),
-                        timeout_s=int(self.cfg.rlm.verifier_timeout_s),
-                        image=self.cfg.rlm.docker_image,
-                        memory_mb=int(self.cfg.rlm.docker_memory_mb),
-                        cpus=float(self.cfg.rlm.docker_cpus),
-                        pids=int(self.cfg.rlm.docker_pids),
-                    )
-                else:
-                    from ..verifiers.pytest_verifier import verify as pytest_verify
-                    res = pytest_verify(
-                        task.prompt,
-                        completion,
-                        str(wdir),
-                        timeout_s=int(self.cfg.rlm.verifier_timeout_s),
-                    )
-                latency_ms = (time.time() - t0) * 1000.0
-
-                passed = bool(getattr(res, "passed", False))
-                reward = float(getattr(res, "reward", 0.0))
+                passed, reward, latency_ms = _run_task_verifier(
+                    self.cfg,
+                    task.prompt,
+                    completion,
+                    wdir,
+                )
 
                 rollouts.append(
                     Rollout(
@@ -1171,28 +1170,7 @@ def collect_rollouts_via_api(
                 wdir = ensure_dir(artifacts_dir / task.id / f"rollout_{k:02d}")
                 (wdir / "main.py").write_text(completion, encoding="utf-8")
                 (ensure_dir(wdir / "tests") / "test_task.py").write_text(task.tests, encoding="utf-8")
-                t0 = time.time()
-                if verifier_backend == "docker":
-                    res = docker_verify(
-                        task.prompt,
-                        completion,
-                        str(wdir),
-                        timeout_s=int(cfg.rlm.verifier_timeout_s),
-                        image=cfg.rlm.docker_image,
-                        memory_mb=int(cfg.rlm.docker_memory_mb),
-                        cpus=float(cfg.rlm.docker_cpus),
-                        pids=int(cfg.rlm.docker_pids),
-                    )
-                else:
-                    res = pytest_verify(
-                        task.prompt,
-                        completion,
-                        str(wdir),
-                        timeout_s=int(cfg.rlm.verifier_timeout_s),
-                    )
-                latency_ms = (time.time() - t0) * 1000.0
-                passed = bool(getattr(res, "passed", False))
-                reward = float(getattr(res, "reward", 0.0))
+                passed, reward, latency_ms = _run_task_verifier(cfg, task.prompt, completion, wdir)
                 rollouts.append(
                     Rollout(
                         task_id=task.id,
@@ -1247,24 +1225,7 @@ def collect_rollouts_via_api(
                 (wdir / "main.py").write_text(completion, encoding="utf-8")
                 (ensure_dir(wdir / "tests") / "test_task.py").write_text(task.tests, encoding="utf-8")
                 
-                t0 = time.time()
-                if verifier_backend == "docker":
-                    res = docker_verify(
-                        task.prompt,
-                        completion,
-                        str(wdir),
-                        timeout_s=int(cfg.rlm.verifier_timeout_s),
-                        image=cfg.rlm.docker_image,
-                        memory_mb=int(cfg.rlm.docker_memory_mb),
-                        cpus=float(cfg.rlm.docker_cpus),
-                        pids=int(cfg.rlm.docker_pids),
-                    )
-                else:
-                    res = pytest_verify(task.prompt, completion, str(wdir), timeout_s=int(cfg.rlm.verifier_timeout_s))
-                latency_ms = (time.time() - t0) * 1000.0
-                
-                passed = bool(getattr(res, "passed", False))
-                reward = float(getattr(res, "reward", 0.0))
+                passed, reward, latency_ms = _run_task_verifier(cfg, task.prompt, completion, wdir)
                 
                 rollouts.append(Rollout(
                     task_id=task.id,

mlxsmith/sdk/__init__.py

@@ -259,6 +259,7 @@ def preference_forward_backward(
     kl_coeff: float = 0.0,
     train_on_prompt: bool = False,
     max_seq_len: Optional[int] = None,
+    delta: float = 0.0,
 ) -> Tuple[Any, Any | None]:
     """Execute preference-based forward/backward pass.
     
@@ -296,23 +297,38 @@ def preference_forward_backward(
             reference_backend=reference_backend,
             kl_coeff=kl_coeff,
             train_on_prompt=train_on_prompt,
+            delta=delta,
         )
 
     return backend.value_and_grad(loss_fn)
 
 
-def create_optimizer(backend: Any, *, lr: float, weight_decay: float = 0.0) -> Tuple[Any, Any]:
+def create_optimizer(
+    backend: Any,
+    *,
+    lr: float,
+    weight_decay: float = 0.0,
+    optimizer: Optional[str] = None,
+    optimizer_kwargs: Optional[dict] = None,
+) -> Tuple[Any, Any]:
     """Create optimizer for training.
     
     Args:
         backend: LLM backend instance
         lr: Learning rate
         weight_decay: Weight decay coefficient
+        optimizer: Optimizer name
+        optimizer_kwargs: Extra optimizer kwargs
         
     Returns:
         Tuple of (optimizer, parameters)
     """
-    return backend.optimizer_and_params(lr=lr, weight_decay=weight_decay)
+    return backend.optimizer_and_params(
+        lr=lr,
+        weight_decay=weight_decay,
+        optimizer=optimizer,
+        optimizer_kwargs=optimizer_kwargs,
+    )
 
 
 def optim_step(backend: Any, optimizer: Any, grads: Any) -> None:

mlxsmith/sdk/losses.py

@@ -70,6 +70,76 @@ def preference_diff(
     return (logp_c - logp_r) - ref_diff
 
 
+@register_loss("cpo")
+def cpo_loss(
+    backend,
+    chosen_ids: Sequence[int],
+    rejected_ids: Sequence[int],
+    *,
+    prompt_len_chosen: int,
+    prompt_len_rejected: int,
+    beta: float = 0.1,
+) -> Any:
+    mx = _require_mx(backend)
+    diff = preference_diff(
+        backend,
+        chosen_ids,
+        rejected_ids,
+        prompt_len_chosen=prompt_len_chosen,
+        prompt_len_rejected=prompt_len_rejected,
+        reference_backend=None,
+    )
+    scaled = _to_mx_scalar(mx, beta) * diff
+    return mx.log1p(mx.exp(-scaled))
+
+
+@register_loss("ipo")
+def ipo_loss(
+    backend,
+    chosen_ids: Sequence[int],
+    rejected_ids: Sequence[int],
+    *,
+    prompt_len_chosen: int,
+    prompt_len_rejected: int,
+    beta: float = 0.1,
+    reference_backend: Optional[Any] = None,
+) -> Any:
+    mx = _require_mx(backend)
+    diff = preference_diff(
+        backend,
+        chosen_ids,
+        rejected_ids,
+        prompt_len_chosen=prompt_len_chosen,
+        prompt_len_rejected=prompt_len_rejected,
+        reference_backend=reference_backend,
+    )
+    target = _to_mx_scalar(mx, 1.0 / (2.0 * float(beta))) if beta != 0 else _to_mx_scalar(mx, 0.0)
+    return (diff - target) ** 2
+
+
+@register_loss("hinge")
+def hinge_loss(
+    backend,
+    chosen_ids: Sequence[int],
+    rejected_ids: Sequence[int],
+    *,
+    prompt_len_chosen: int,
+    prompt_len_rejected: int,
+    delta: float = 0.0,
+    reference_backend: Optional[Any] = None,
+) -> Any:
+    mx = _require_mx(backend)
+    diff = preference_diff(
+        backend,
+        chosen_ids,
+        rejected_ids,
+        prompt_len_chosen=prompt_len_chosen,
+        prompt_len_rejected=prompt_len_rejected,
+        reference_backend=reference_backend,
+    )
+    return mx.maximum(_to_mx_scalar(mx, delta) - diff, _to_mx_scalar(mx, 0.0))
+
+
 @register_loss("dpo")
 def dpo_loss(
     backend,
@@ -149,8 +219,10 @@ def preference_loss(
     reference_backend: Optional[Any] = None,
     kl_coeff: float = 0.0,
     train_on_prompt: bool = False,
+    delta: float = 0.0,
 ) -> Any:
-    if algo.lower() == "orpo":
+    algo_l = algo.lower()
+    if algo_l == "orpo":
         return orpo_loss(
             backend,
             chosen_ids,
@@ -162,6 +234,35 @@ def preference_loss(
             kl_coeff=kl_coeff,
             train_on_prompt=train_on_prompt,
         )
+    if algo_l == "cpo":
+        return cpo_loss(
+            backend,
+            chosen_ids,
+            rejected_ids,
+            prompt_len_chosen=prompt_len_chosen,
+            prompt_len_rejected=prompt_len_rejected,
+            beta=beta,
+        )
+    if algo_l == "ipo":
+        return ipo_loss(
+            backend,
+            chosen_ids,
+            rejected_ids,
+            prompt_len_chosen=prompt_len_chosen,
+            prompt_len_rejected=prompt_len_rejected,
+            beta=beta,
+            reference_backend=reference_backend,
+        )
+    if algo_l == "hinge":
+        return hinge_loss(
+            backend,
+            chosen_ids,
+            rejected_ids,
+            prompt_len_chosen=prompt_len_chosen,
+            prompt_len_rejected=prompt_len_rejected,
+            delta=delta,
+            reference_backend=reference_backend,
+        )
     return dpo_loss(
         backend,
         chosen_ids,

mlxsmith/sdk/training_client.py

@@ -210,6 +210,7 @@ class TrainingClient:
                         beta=batch.extra.get("beta", 0.1),
                         reference_backend=batch.extra.get("reference_backend"),
                         kl_coeff=batch.extra.get("kl_coeff", 0.0),
+                        delta=batch.extra.get("delta", 0.0),
                         train_on_prompt=batch.train_on_prompt,
                         max_seq_len=batch.max_seq_len,
                     )
@@ -514,12 +515,20 @@ class TrainingClient:
     # Utility Methods
     # ========================================================================
     
-    def create_optimizer(self, lr: float = 1e-4, weight_decay: float = 0.0) -> APIFuture[Any]:
+    def create_optimizer(
+        self,
+        lr: float = 1e-4,
+        weight_decay: float = 0.0,
+        optimizer: Optional[str] = None,
+        optimizer_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> APIFuture[Any]:
         """Create optimizer for training.
         
         Args:
             lr: Learning rate
             weight_decay: Weight decay coefficient
+            optimizer: Optimizer name (e.g., adamw, adam, qhadam, muon)
+            optimizer_kwargs: Extra optimizer kwargs
             
         Returns:
             APIFuture resolving to optimizer instance
@@ -535,6 +544,8 @@ class TrainingClient:
                 self.backend,
                 lr=lr,
                 weight_decay=weight_decay,
+                optimizer=optimizer,
+                optimizer_kwargs=optimizer_kwargs,
             )
             return self.optimizer
         
@@ -588,10 +599,18 @@ class TrainingClient:
         if len(grads_list) == 1:
             return grads_list[0]
         
-        # Average gradients
-        # This is backend-specific; for now return first grad
-        # In practice, MLX would average the arrays
-        return grads_list[0]
+        from ..util import tree_add, tree_scale
+
+        agg = None
+        count = 0
+        for grads in grads_list:
+            if grads is None:
+                continue
+            agg = tree_add(agg, grads)
+            count += 1
+        if agg is None or count == 0:
+            return None
+        return tree_scale(agg, 1.0 / float(count))
     
     def shutdown(self) -> None:
         """Shutdown the client and its thread pool."""

mlxsmith/train/distill.py

@@ -109,7 +109,12 @@ def run_distill(
             )
             student.apply_lora_from_config(lora_cfg)
 
-        opt, _params = student.optimizer_and_params(lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
+        opt, _params = student.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 = int(cfg.train.iters)