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

mlxsmith/bench.py

@@ -44,7 +44,12 @@ def run_bench(
     mode = (mode or "inference").lower()
 
     if mode == "trainer":
-        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,
+        )
         prompt_ids = llm.encode(prompt)
         ids = llm.encode(prompt + " " + "x" * max_tokens)
         for i in range(max(1, reps)):
@@ -59,7 +64,12 @@ def run_bench(
             elapsed = max(time.time() - t0, 1e-6)
             results.append({"rep": i, "steps": steps, "time_s": elapsed, "steps_per_s": steps / elapsed})
     elif mode == "end_to_end":
-        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,
+        )
         for i in range(max(1, reps)):
             t0 = time.time()
             gen = llm.generate(prompt, max_new_tokens=max_tokens, temperature=0.0)

mlxsmith/cli.py

@@ -24,6 +24,8 @@ from .train.sft import run_sft
 from .train.pref import run_pref
 from .train.rft import run_rft
 from .train.distill import run_distill
+from .train.online_dpo import run_online_dpo
+from .train.self_verify import run_self_verify
 from .eval import run_eval
 from .bench import run_bench
 from .rlm import run_rlm, run_rlm_orchestrated
@@ -40,6 +42,13 @@ from .envs import (
     resolve_env_path as resolve_env_path_plugin,
     load_manifest as load_env_manifest,
 )
+from .integrations.mlx_lm_lora import (
+    build_train_command as build_mlx_lm_lora_train_command,
+    build_synthetic_command as build_mlx_lm_lora_synth_command,
+    build_judge_command as build_mlx_lm_lora_judge_command,
+    build_reward_functions_command as build_mlx_lm_lora_reward_functions_command,
+    run_command as run_mlx_lm_lora_command,
+)
 
 app = typer.Typer(
     add_completion=False,
@@ -65,6 +74,9 @@ def init(path: str = typer.Argument(..., help="Project directory to create")):
     (p / "verifiers" / "regex.py").write_text(_sample_verifier_regex(), encoding="utf-8")
     (p / "verifiers" / "pytest.py").write_text(_sample_verifier_pytest(), encoding="utf-8")
     (p / "verifiers" / "jsonschema.py").write_text(_sample_verifier_jsonschema(), encoding="utf-8")
+    (p / "verifiers" / "llm_judge.py").write_text(_sample_verifier_llm_judge(), encoding="utf-8")
+    (p / "verifiers" / "rubrics").mkdir(parents=True, exist_ok=True)
+    (p / "verifiers" / "rubrics" / "coding.txt").write_text(_sample_judge_rubric(), encoding="utf-8")
     (p / "eval" / "suites" / "coding.yaml").write_text(_sample_eval_suite(), encoding="utf-8")
     console.print(f"[green]Initialized[/green] {p.resolve()}")
 
@@ -341,14 +353,19 @@ def pref(
     data: str = typer.Option("data/prefs", "--data"),
     model: str = typer.Option(..., "--model", help="Base adapter or model path (e.g., runs/sft_0001/adapter)"),
     accel: Optional[str] = typer.Option(None, "--accel", help="Override accel.backend"),
-    algo: Optional[str] = typer.Option(None, "--algo", help="Override pref.algo (dpo|orpo|grpo)"),
+    algo: Optional[str] = typer.Option(None, "--algo", help="Override pref.algo (legacy)"),
+    loss_type: Optional[str] = typer.Option(None, "--loss-type", help="dpo|cpo|orpo|ipo|hinge"),
 ):
     root = project_root_from_cwd()
+    overrides = {}
+    if loss_type is not None:
+        overrides["pref.loss_type"] = loss_type
     cfg = get_config(
         config_path=config,
         root=root,
         accel_backend=accel,
         algo=algo,
+        **overrides,
     )
     data_dir = root / data
     run = run_pref(root, cfg, data_dir, Path(model), cfg.accel.backend)
@@ -363,13 +380,27 @@ def rft(
     model: str = typer.Option(..., "--model"),
     accel: Optional[str] = typer.Option(None, "--accel", help="Override accel.backend"),
     rollouts: Optional[int] = typer.Option(None, "--rollouts", help="Override rft.rollouts"),
+    loss_type: Optional[str] = typer.Option(None, "--loss-type", help="grpo|dr_grpo|dapo"),
+    epsilon_low: Optional[float] = typer.Option(None, "--epsilon-low"),
+    epsilon_high: Optional[float] = typer.Option(None, "--epsilon-high"),
+    token_level_loss: Optional[bool] = typer.Option(None, "--token-level-loss/--sequence-level-loss"),
 ):
     root = project_root_from_cwd()
+    overrides = {}
+    if loss_type is not None:
+        overrides["rft.loss_type"] = loss_type
+    if epsilon_low is not None:
+        overrides["rft.epsilon_low"] = epsilon_low
+    if epsilon_high is not None:
+        overrides["rft.epsilon_high"] = epsilon_high
+    if token_level_loss is not None:
+        overrides["rft.token_level_loss"] = token_level_loss
     cfg = get_config(
         config_path=config,
         root=root,
         accel_backend=accel,
         rollouts=rollouts,
+        **overrides,
     )
     run = run_rft(root, cfg, root / env, root / verifier, Path(model), cfg.accel.backend)
     console.print(f"[bold]Run:[/bold] {run.run_dir}")
@@ -437,6 +468,142 @@ def distill(
     console.print(f"[bold]Run:[/bold] {run.run_dir}")
 
 
+@app.command("online-dpo")
+def online_dpo(
+    data: str = typer.Option(..., "--data", help="JSONL with prompts"),
+    model: str = typer.Option(..., "--model"),
+    judge_model: Optional[str] = typer.Option(None, "--judge-model"),
+    judge_backend: str = typer.Option("mlx-lm", "--judge-backend"),
+    rubric: Optional[str] = typer.Option(None, "--rubric"),
+    group_size: Optional[int] = typer.Option(None, "--group-size"),
+    max_new_tokens: Optional[int] = typer.Option(None, "--max-new-tokens"),
+    temperature: Optional[float] = typer.Option(None, "--temperature"),
+    config: str = typer.Option("mlxsmith.yaml", "-c", "--config", help="Config file path"),
+    accel: Optional[str] = typer.Option(None, "--accel", help="Override accel.backend"),
+):
+    root = project_root_from_cwd()
+    cfg = get_config(config_path=config, root=root, accel_backend=accel)
+    run = run_online_dpo(
+        root,
+        cfg,
+        Path(data),
+        model,
+        cfg.accel.backend,
+        judge_model=judge_model,
+        judge_backend=judge_backend,
+        rubric=rubric,
+        group_size=group_size,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+    )
+    console.print(f"[bold]Run:[/bold] {run.run_dir}")
+
+
+@app.command("self-verify")
+def self_verify(
+    data: str = typer.Option(..., "--data", help="JSONL with prompts"),
+    model: str = typer.Option(..., "--model"),
+    verifier_model: Optional[str] = typer.Option(None, "--verifier-model"),
+    verifier_backend: str = typer.Option("mlx-lm", "--verifier-backend"),
+    rubric: Optional[str] = typer.Option(None, "--rubric"),
+    max_new_tokens: Optional[int] = typer.Option(None, "--max-new-tokens"),
+    temperature: Optional[float] = typer.Option(None, "--temperature"),
+    config: str = typer.Option("mlxsmith.yaml", "-c", "--config", help="Config file path"),
+    accel: Optional[str] = typer.Option(None, "--accel", help="Override accel.backend"),
+):
+    root = project_root_from_cwd()
+    cfg = get_config(config_path=config, root=root, accel_backend=accel)
+    run = run_self_verify(
+        root,
+        cfg,
+        Path(data),
+        model,
+        cfg.accel.backend,
+        verifier_model=verifier_model,
+        verifier_backend=verifier_backend,
+        rubric=rubric,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+    )
+    console.print(f"[bold]Run:[/bold] {run.run_dir}")
+
+
+lora_app = typer.Typer(help="mlx-lm-lora passthrough commands")
+app.add_typer(lora_app, name="lora")
+
+
+@lora_app.command(
+    "train",
+    context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
+)
+def lora_train(
+    ctx: typer.Context,
+    config: Optional[str] = typer.Option(None, "--config", help="mlx-lm-lora config path"),
+    model: Optional[str] = typer.Option(None, "--model", help="Model id or path"),
+    data: Optional[str] = typer.Option(None, "--data", help="Dataset path or HF dataset"),
+    train_mode: Optional[str] = typer.Option(None, "--train-mode", help="sft|dpo|orpo|grpo|ppo|..."),
+    train_type: Optional[str] = typer.Option(None, "--train-type", help="lora|dora|full"),
+    dry_run: bool = typer.Option(False, "--dry-run"),
+):
+    """Run mlx-lm-lora training with passthrough args.
+
+    Use `--` to pass through any additional mlx-lm-lora flags.
+    """
+    root = project_root_from_cwd()
+    cmd = build_mlx_lm_lora_train_command(
+        config=config,
+        model=model,
+        data=data,
+        train_mode=train_mode,
+        train_type=train_type,
+        extra_args=list(ctx.args),
+    )
+    run_mlx_lm_lora_command(cmd, dry_run=dry_run, cwd=root)
+
+
+@lora_app.command(
+    "synthetic",
+    context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
+)
+def lora_synthetic(
+    ctx: typer.Context,
+    kind: str = typer.Argument(..., help="prompts|sft|dpo"),
+    dry_run: bool = typer.Option(False, "--dry-run"),
+):
+    """Run mlx-lm-lora synthetic dataset generation."""
+    root = project_root_from_cwd()
+    cmd = build_mlx_lm_lora_synth_command(kind, extra_args=list(ctx.args))
+    run_mlx_lm_lora_command(cmd, dry_run=dry_run, cwd=root)
+
+
+@lora_app.command(
+    "judge",
+    context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
+)
+def lora_judge(
+    ctx: typer.Context,
+    dry_run: bool = typer.Option(False, "--dry-run"),
+):
+    """Run mlx-lm-lora judge model training."""
+    root = project_root_from_cwd()
+    cmd = build_mlx_lm_lora_judge_command(extra_args=list(ctx.args))
+    run_mlx_lm_lora_command(cmd, dry_run=dry_run, cwd=root)
+
+
+@lora_app.command(
+    "reward-functions",
+    context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
+)
+def lora_reward_functions(
+    ctx: typer.Context,
+    dry_run: bool = typer.Option(False, "--dry-run"),
+):
+    """List mlx-lm-lora reward functions."""
+    root = project_root_from_cwd()
+    cmd = build_mlx_lm_lora_reward_functions_command(extra_args=list(ctx.args))
+    run_mlx_lm_lora_command(cmd, dry_run=dry_run, cwd=root)
+
+
 @app.command()
 def eval(
     suite: str = typer.Option("eval/suites/coding.yaml", "--suite"),
@@ -933,6 +1100,25 @@ def verify(prompt: str, completion: str, workdir: str, **kwargs):
 """
 
 
+def _sample_verifier_llm_judge() -> str:
+    return """from mlxsmith.verifiers.llm_judge import verify as _verify
+
+def verify(prompt: str, completion: str, workdir: str, **kwargs):
+    # Pass model=... or set MLXSMITH_JUDGE_MODEL for the judge model id.
+    return _verify(prompt, completion, workdir, **kwargs)
+"""
+
+
+def _sample_judge_rubric() -> str:
+    return """Score from 0.0 to 1.0.
+- 1.0: Correct, complete, and safe.
+- 0.7: Mostly correct with small issues.
+- 0.4: Partial correctness or unclear reasoning.
+- 0.0: Incorrect or unsafe.
+Return JSON only.
+"""
+
+
 def _sample_eval_suite() -> str:
     return """name: coding-eval-sample
 notes: |

mlxsmith/config_models.py

@@ -47,6 +47,8 @@ class TrainConfig(BaseModel):
     grad_accum: int = 8
     lr: float = 2e-4
     weight_decay: float = 0.0
+    optimizer: str = "adamw"
+    optimizer_kwargs: Dict[str, Any] = Field(default_factory=dict)
     iters: int = 1000
     save_every: int = 100
     eval_every: int = 100
@@ -61,6 +63,11 @@ class TrainConfig(BaseModel):
             raise ValueError("value must be non-negative")
         return v
 
+    @field_validator("optimizer")
+    @classmethod
+    def normalize_optimizer(cls, v: str) -> str:
+        return v.strip().lower()
+
 
 class LoraConfig(BaseModel):
     """LoRA/DoRA adapter configuration."""
@@ -89,11 +96,13 @@ class LoraConfig(BaseModel):
 
 
 class PrefConfig(BaseModel):
-    """Preference tuning configuration (DPO, ORPO, GRPO)."""
+    """Preference tuning configuration (DPO variants)."""
     
     algo: Literal["dpo", "orpo", "grpo"] = "dpo"
+    loss_type: Literal["dpo", "cpo", "orpo", "ipo", "hinge"] = "dpo"
     beta: float = 0.1
     kl_coeff: float = 0.0
+    delta: float = 0.0
     reference_model: Optional[str] = None
 
 
@@ -101,12 +110,16 @@ class RftConfig(BaseModel):
     """Reinforcement fine-tuning configuration."""
     
     algo: Literal["grpo"] = "grpo"
+    loss_type: Literal["grpo", "dr_grpo", "dapo"] = "grpo"
     rollouts: int = 8
     kl_coeff: float = 0.02
     max_steps_per_task: int = 1
     temperature: float = 0.8
     max_new_tokens: int = 256
     normalize_advantage: bool = True
+    epsilon_low: float = 0.2
+    epsilon_high: float = 0.2
+    token_level_loss: bool = False
     reference_model: Optional[str] = None
 
 
@@ -164,6 +177,7 @@ CLI_ALIASES: dict[str, tuple[str, ...]] = {
     "lr": ("train", "lr"),
     "batch_size": ("train", "batch_size"),
     "iters": ("train", "iters"),
+    "optimizer": ("train", "optimizer"),
     "model_id": ("model", "id"),
     "accel_backend": ("accel", "backend"),
     "host": ("serve", "host"),

mlxsmith/integrations/__init__.py

@@ -0,0 +1,19 @@
+"""External integrations for mlxsmith."""
+
+from .mlx_lm_lora import (
+    build_train_command as build_mlx_lm_lora_train_command,
+    build_synthetic_command as build_mlx_lm_lora_synth_command,
+    build_judge_command as build_mlx_lm_lora_judge_command,
+    build_reward_functions_command as build_mlx_lm_lora_reward_functions_command,
+    run_command as run_mlx_lm_lora_command,
+    ensure_available as ensure_mlx_lm_lora_available,
+)
+
+__all__ = [
+    "build_mlx_lm_lora_train_command",
+    "build_mlx_lm_lora_synth_command",
+    "build_mlx_lm_lora_judge_command",
+    "build_mlx_lm_lora_reward_functions_command",
+    "run_mlx_lm_lora_command",
+    "ensure_mlx_lm_lora_available",
+]

mlxsmith/integrations/mlx_lm_lora.py

@@ -0,0 +1,117 @@
+"""Passthrough helpers for mlx-lm-lora CLI integration."""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+import shlex
+import subprocess
+import sys
+from pathlib import Path
+from typing import Optional, Sequence
+
+from rich.console import Console
+
+console = Console()
+
+
+def ensure_available() -> None:
+    if importlib.util.find_spec("mlx_lm_lora") is None:
+        raise RuntimeError(
+            "mlx-lm-lora is not installed. Install with: pip install 'mlxsmith[lora]' or 'mlx-lm-lora'"
+        )
+
+
+def _flag_present(args: Sequence[str], *flags: str) -> bool:
+    for flag in flags:
+        if flag in args:
+            return True
+        if flag.startswith("--"):
+            prefix = flag + "="
+            if any(a.startswith(prefix) for a in args):
+                return True
+    return False
+
+
+def _append_flag(cmd: list[str], args: Sequence[str], flag: str, value: Optional[str]) -> None:
+    if value is None:
+        return
+    if _flag_present(args, flag):
+        return
+    cmd.extend([flag, value])
+
+
+def _base_python_cmd(module: str) -> list[str]:
+    return [sys.executable, "-m", module]
+
+
+def build_train_command(
+    *,
+    config: Optional[str] = None,
+    model: Optional[str] = None,
+    data: Optional[str] = None,
+    train_mode: Optional[str] = None,
+    train_type: Optional[str] = None,
+    extra_args: Sequence[str] = (),
+) -> list[str]:
+    args = list(extra_args)
+    cmd = _base_python_cmd("mlx_lm_lora.train")
+    _append_flag(cmd, args, "--config", config)
+    _append_flag(cmd, args, "--model", model)
+    _append_flag(cmd, args, "--data", data)
+    _append_flag(cmd, args, "--train-mode", train_mode)
+    _append_flag(cmd, args, "--train-type", train_type)
+    cmd.extend(args)
+    return cmd
+
+
+def build_synthetic_command(
+    kind: str,
+    *,
+    extra_args: Sequence[str] = (),
+) -> list[str]:
+    kind = kind.strip().lower()
+    module = {
+        "prompts": "mlx_lm_lora.synthetic_prompts",
+        "sft": "mlx_lm_lora.synthetic_sft",
+        "dpo": "mlx_lm_lora.synthetic_dpo",
+    }.get(kind)
+    if module is None:
+        raise ValueError(f"Unknown synthetic kind: {kind}")
+    cmd = _base_python_cmd(module)
+    cmd.extend(list(extra_args))
+    return cmd
+
+
+def build_judge_command(*, extra_args: Sequence[str] = ()) -> list[str]:
+    cmd = _base_python_cmd("mlx_lm_lora.train_judge")
+    cmd.extend(list(extra_args))
+    return cmd
+
+
+def build_reward_functions_command(*, extra_args: Sequence[str] = ()) -> list[str]:
+    cmd = _base_python_cmd("mlx_lm_lora.train")
+    cmd.append("--list-reward-functions")
+    cmd.extend(list(extra_args))
+    return cmd
+
+
+def run_command(
+    cmd: Sequence[str],
+    *,
+    dry_run: bool = False,
+    cwd: Optional[Path] = None,
+    env: Optional[dict] = None,
+) -> int:
+    if dry_run:
+        console.print("[cyan]mlx-lm-lora cmd[/cyan]", shlex.join(list(cmd)))
+        return 0
+    ensure_available()
+    run_env = os.environ.copy()
+    if env:
+        run_env.update(env)
+    console.print("[cyan]mlx-lm-lora cmd[/cyan]", shlex.join(list(cmd)))
+    result = subprocess.run(list(cmd), cwd=str(cwd) if cwd else None, env=run_env, check=False)
+    if result.returncode != 0:
+        raise RuntimeError(f"mlx-lm-lora failed with exit code {result.returncode}")
+    return result.returncode

mlxsmith/llm/backend.py

@@ -112,7 +112,14 @@ class LLMBackend(Protocol):
     def value_and_grad(self, loss_fn) -> tuple[Any, Any | None]:
         """Return (loss, grads) using backend autograd when available."""
 
-    def optimizer_and_params(self, *, lr: float, weight_decay: float = 0.0) -> tuple[Any, Any]:
+    def optimizer_and_params(
+        self,
+        *,
+        lr: float,
+        weight_decay: float = 0.0,
+        optimizer: str | None = None,
+        optimizer_kwargs: dict | None = None,
+    ) -> tuple[Any, Any]:
         """Return (optimizer, trainable_params_tree)."""
 
     def apply_grads(self, optimizer: Any, grads: Any) -> None:

mlxsmith/llm/mlx_lm_backend.py

@@ -467,7 +467,54 @@ class MlxLMBackend:
             return vag(self.model, loss_fn)(self.model)
         return loss_fn(self.model), None
 
-    def optimizer_and_params(self, *, lr: float, weight_decay: float = 0.0) -> tuple[Any, Any]:
+    def optimizer_and_params(
+        self,
+        *,
+        lr: float,
+        weight_decay: float = 0.0,
+        optimizer: str | None = None,
+        optimizer_kwargs: dict | None = None,
+    ) -> tuple[Any, Any]:
+        return self._optimizer_and_params(
+            lr=lr,
+            weight_decay=weight_decay,
+            optimizer=optimizer,
+            optimizer_kwargs=optimizer_kwargs,
+        )
+
+    def _resolve_optimizer(self, name: str):
+        assert self.optim is not None
+        name = name.strip().lower()
+        if name in ("muon", "muonclip", "muon_clip"):
+            from ..optim.muon import Muon, MuonClip
+
+            return MuonClip if name in ("muonclip", "muon_clip") else Muon
+        mapping = {
+            "adamw": ["AdamW", "Adamw", "adamw"],
+            "adam": ["Adam", "adam"],
+            "sgd": ["SGD", "Sgd", "sgd"],
+            "rmsprop": ["RMSprop", "RmsProp", "rmsprop"],
+            "qhadam": ["QHAdam", "Qhadam", "qhadam"],
+            "muon": ["Muon", "muon", "MuonW", "muonw"],
+        }
+        candidates = mapping.get(name, [name, name.capitalize(), name.upper()])
+        for cand in candidates:
+            opt_cls = getattr(self.optim, cand, None)
+            if opt_cls is not None:
+                return opt_cls
+        raise RuntimeError(
+            f"Optimizer '{name}' is not available in mlx.optimizers. "
+            "Install a build that provides it or choose a supported optimizer."
+        )
+
+    def _optimizer_and_params(
+        self,
+        *,
+        lr: float,
+        weight_decay: float = 0.0,
+        optimizer: str | None = None,
+        optimizer_kwargs: dict | None = None,
+    ) -> tuple[Any, Any]:
         assert self.optim is not None
         if self.model is None:
             raise RuntimeError("Backend not loaded")
@@ -483,7 +530,17 @@ class MlxLMBackend:
         if not params:
             params = getattr(self.model, "parameters", lambda: self.model)()
 
-        opt = self.optim.AdamW(learning_rate=lr, weight_decay=weight_decay)
+        opt_name = optimizer or "adamw"
+        opt_cls = self._resolve_optimizer(opt_name)
+        kwargs = dict(optimizer_kwargs or {})
+        if "learning_rate" not in kwargs and "lr" not in kwargs:
+            kwargs["learning_rate"] = lr
+            kwargs["lr"] = lr
+        if "weight_decay" not in kwargs:
+            kwargs["weight_decay"] = weight_decay
+        opt = self._call_with_supported_kwargs(opt_cls, **kwargs)
+        if opt is None:
+            opt = opt_cls()
         opt.init(params)
         return opt, params
 

mlxsmith/llm/mock_backend.py

@@ -200,7 +200,14 @@ class MockBackend:
     def value_and_grad(self, loss_fn):
         return loss_fn(self.model), None
 
-    def optimizer_and_params(self, *, lr: float, weight_decay: float = 0.0):
+    def optimizer_and_params(
+        self,
+        *,
+        lr: float,
+        weight_decay: float = 0.0,
+        optimizer: str | None = None,
+        optimizer_kwargs: dict | None = None,
+    ):
         return object(), {}
 
     def apply_grads(self, optimizer: Any, grads: Any) -> None:

mlxsmith/optim/__init__.py

@@ -0,0 +1,3 @@
+from .muon import Muon, MuonClip
+
+__all__ = ["Muon", "MuonClip"]

mlxsmith/optim/muon.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from ..util import tree_map
+
+
+class Muon:
+    """Muon optimizer wrapper (AdamW with Newton-Schulz orthogonalization for 2D grads)."""
+
+    def __init__(
+        self,
+        learning_rate: Optional[float] = None,
+        lr: Optional[float] = None,
+        weight_decay: float = 0.0,
+        clip: Optional[float] = None,
+        ns_iters: int = 5,
+        a: float = 3.4445,
+        b: float = -4.7750,
+        c: float = 2.0315,
+    ) -> None:
+        import mlx.optimizers as optim  # type: ignore
+
+        self.learning_rate = float(lr if lr is not None else learning_rate if learning_rate is not None else 1e-4)
+        self.weight_decay = float(weight_decay)
+        self.clip = clip
+        self.ns_iters = int(ns_iters)
+        self.a = float(a)
+        self.b = float(b)
+        self.c = float(c)
+        self._base = optim.AdamW(learning_rate=self.learning_rate, weight_decay=self.weight_decay)
+        self.state: dict[str, Any] = {}
+
+    def init(self, params: Any) -> None:
+        self._base.init(params)
+        self.state = getattr(self._base, "state", {})
+
+    def _clip_grad(self, g):
+        if self.clip is None or self.clip <= 0:
+            return g
+        import mlx.core as mx  # type: ignore
+
+        norm = mx.sqrt((g * g).sum())
+        scale = mx.minimum(mx.array(1.0), mx.array(float(self.clip)) / mx.maximum(norm, mx.array(1e-8)))
+        return g * scale
+
+    def _orthogonalize(self, g):
+        import mlx.core as mx  # type: ignore
+
+        g = self._clip_grad(g)
+        if getattr(g, "ndim", None) != 2:
+            return g
+        in_dim = int(g.shape[1])
+        eye = mx.eye(in_dim, dtype=g.dtype)
+        out = g
+        for _ in range(self.ns_iters):
+            gtg = mx.matmul(out.T, out)
+            gtg2 = mx.matmul(gtg, gtg)
+            out = mx.matmul(out, self.a * eye + self.b * gtg + self.c * gtg2)
+        return out
+
+    def update(self, model: Any, grads: Any) -> None:
+        if grads is None:
+            return
+        transformed = tree_map(self._orthogonalize, grads)
+        self._base.update(model, transformed)
+        self.state = getattr(self._base, "state", self.state)
+
+
+class MuonClip(Muon):
+    """Muon with explicit gradient clipping before orthogonalization."""
+
+    def __init__(
+        self,
+        learning_rate: Optional[float] = None,
+        lr: Optional[float] = None,
+        weight_decay: float = 0.0,
+        clip: Optional[float] = 1.0,
+        ns_iters: int = 5,
+        a: float = 3.4445,
+        b: float = -4.7750,
+        c: float = 2.0315,
+    ) -> None:
+        super().__init__(
+            learning_rate=learning_rate,
+            lr=lr,
+            weight_decay=weight_decay,
+            clip=clip,
+            ns_iters=ns_iters,
+            a=a,
+            b=b,
+            c=c,
+        )