mlxsmith 0.1.2__tar.gz → 0.1.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (96) hide show
  1. {mlxsmith-0.1.2/src/mlxsmith.egg-info → mlxsmith-0.1.3}/PKG-INFO +27 -1
  2. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/README.md +23 -0
  3. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/pyproject.toml +3 -1
  4. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/bench.py +12 -2
  5. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/cli.py +187 -1
  6. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/config_models.py +15 -1
  7. mlxsmith-0.1.3/src/mlxsmith/integrations/__init__.py +19 -0
  8. mlxsmith-0.1.3/src/mlxsmith/integrations/mlx_lm_lora.py +117 -0
  9. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/llm/backend.py +8 -1
  10. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/llm/mlx_lm_backend.py +59 -2
  11. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/llm/mock_backend.py +8 -1
  12. mlxsmith-0.1.3/src/mlxsmith/optim/__init__.py +3 -0
  13. mlxsmith-0.1.3/src/mlxsmith/optim/muon.py +93 -0
  14. mlxsmith-0.1.3/src/mlxsmith/orchestrator/daemon.py +116 -0
  15. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/orchestrator/trainer_worker.py +4 -0
  16. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/rlm/loop.py +53 -92
  17. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/sdk/__init__.py +18 -2
  18. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/sdk/losses.py +102 -1
  19. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/sdk/training_client.py +24 -5
  20. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/train/distill.py +6 -1
  21. mlxsmith-0.1.3/src/mlxsmith/train/online_dpo.py +249 -0
  22. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/train/pref.py +31 -29
  23. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/train/rft.py +123 -38
  24. mlxsmith-0.1.3/src/mlxsmith/train/self_verify.py +199 -0
  25. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/train/sft.py +13 -2
  26. mlxsmith-0.1.3/src/mlxsmith/verifiers/llm_judge.py +278 -0
  27. mlxsmith-0.1.3/src/mlxsmith/verifiers/prime.py +127 -0
  28. {mlxsmith-0.1.2 → mlxsmith-0.1.3/src/mlxsmith.egg-info}/PKG-INFO +27 -1
  29. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith.egg-info/SOURCES.txt +11 -0
  30. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith.egg-info/requires.txt +4 -0
  31. mlxsmith-0.1.3/tests/test_lora_integration.py +47 -0
  32. mlxsmith-0.1.3/tests/test_online_dpo_self_verify.py +61 -0
  33. mlxsmith-0.1.3/tests/test_pref_variants.py +31 -0
  34. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/tests/test_sdk.py +17 -0
  35. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/tests/test_training_smoke.py +6 -0
  36. mlxsmith-0.1.3/tests/test_verifiers.py +73 -0
  37. mlxsmith-0.1.2/src/mlxsmith/orchestrator/daemon.py +0 -449
  38. mlxsmith-0.1.2/tests/test_verifiers.py +0 -25
  39. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/LICENSE +0 -0
  40. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/setup.cfg +0 -0
  41. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/__init__.py +0 -0
  42. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/accel/__init__.py +0 -0
  43. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/accel/base.py +0 -0
  44. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/accel/none.py +0 -0
  45. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/adapters.py +0 -0
  46. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/api/__init__.py +0 -0
  47. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/api/handlers.py +0 -0
  48. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/api/schemas.py +0 -0
  49. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/auth.py +0 -0
  50. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/config.py +0 -0
  51. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/data.py +0 -0
  52. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/envs/__init__.py +0 -0
  53. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/envs/system.py +0 -0
  54. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/envs/token_env.py +0 -0
  55. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/eval.py +0 -0
  56. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/infer.py +0 -0
  57. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/llm/__init__.py +0 -0
  58. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/llm/interface.py +0 -0
  59. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/llm/registry.py +0 -0
  60. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/models.py +0 -0
  61. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/orchestrator/__init__.py +0 -0
  62. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/orchestrator/inference_worker.py +0 -0
  63. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/orchestrator/queue.py +0 -0
  64. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/rlm/__init__.py +0 -0
  65. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/rlm/corpus.py +0 -0
  66. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/rlm/gating.py +0 -0
  67. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/rlm/generate.py +0 -0
  68. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/rlm/history.py +0 -0
  69. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/rlm/inference.py +0 -0
  70. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/rlm/mutate.py +0 -0
  71. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/rlm/trainer.py +0 -0
  72. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/rlm/weights.py +0 -0
  73. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/runs.py +0 -0
  74. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/sdk/future.py +0 -0
  75. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/sdk/sampling_client.py +0 -0
  76. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/server.py +0 -0
  77. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/train/__init__.py +0 -0
  78. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/train/lora.py +0 -0
  79. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/util.py +0 -0
  80. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/verifiers/__init__.py +0 -0
  81. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/verifiers/compose.py +0 -0
  82. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/verifiers/docker_verifier.py +0 -0
  83. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/verifiers/jsonschema.py +0 -0
  84. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/verifiers/pytest_verifier.py +0 -0
  85. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/verifiers/regex.py +0 -0
  86. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith/verifiers/types.py +0 -0
  87. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith.egg-info/dependency_links.txt +0 -0
  88. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith.egg-info/entry_points.txt +0 -0
  89. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/src/mlxsmith.egg-info/top_level.txt +0 -0
  90. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/tests/test_api.py +0 -0
  91. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/tests/test_auth.py +0 -0
  92. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/tests/test_config.py +0 -0
  93. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/tests/test_data.py +0 -0
  94. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/tests/test_rlm.py +0 -0
  95. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/tests/test_rlm_mutation.py +0 -0
  96. {mlxsmith-0.1.2 → mlxsmith-0.1.3}/tests/test_runs.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: mlxsmith
3
- Version: 0.1.2
3
+ Version: 0.1.3
4
4
  Summary: Apple Silicon MLX fine-tuning toolkit — SFT, DPO/ORPO, GRPO, distillation, and OpenAI-compatible serving.
5
5
  Author-email: Shannon Labs <hmbown@gmail.com>
6
6
  License: MIT
@@ -36,6 +36,8 @@ Provides-Extra: llm
36
36
  Requires-Dist: mlx-lm>=0.30.5; extra == "llm"
37
37
  Requires-Dist: transformers>=5.0.0; extra == "llm"
38
38
  Requires-Dist: datasets>=3.0.0; extra == "llm"
39
+ Provides-Extra: lora
40
+ Requires-Dist: mlx-lm-lora>=1.0.0; extra == "lora"
39
41
  Provides-Extra: serve
40
42
  Requires-Dist: fastapi>=0.128.0; extra == "serve"
41
43
  Requires-Dist: uvicorn>=0.40.0; extra == "serve"
@@ -46,6 +48,7 @@ Requires-Dist: ruff>=0.14.0; extra == "dev"
46
48
  Provides-Extra: all
47
49
  Requires-Dist: mlx>=0.30.4; extra == "all"
48
50
  Requires-Dist: mlx-lm>=0.30.5; extra == "all"
51
+ Requires-Dist: mlx-lm-lora>=1.0.0; extra == "all"
49
52
  Requires-Dist: transformers>=5.0.0; extra == "all"
50
53
  Requires-Dist: datasets>=3.0.0; extra == "all"
51
54
  Requires-Dist: fastapi>=0.128.0; extra == "all"
@@ -74,6 +77,9 @@ pip install mlxsmith
74
77
  # Apple Silicon training + serving
75
78
  pip install "mlxsmith[mlx,llm,serve]"
76
79
 
80
+ # mlx-lm-lora passthrough (advanced training methods)
81
+ pip install "mlxsmith[lora]"
82
+
77
83
  # Everything
78
84
  pip install "mlxsmith[all]"
79
85
  ```
@@ -131,6 +137,22 @@ mlxsmith distill --teacher large-model --student small-model --mode opd
131
137
  mlxsmith pipeline
132
138
  ```
133
139
 
140
+ ### mlx-lm-lora parity (all methods)
141
+
142
+ Use the passthrough to access mlx-lm-lora features (DPO variants, GRPO variants,
143
+ PPO, synthetic datasets, judge training, etc.):
144
+
145
+ ```bash
146
+ # Train with mlx-lm-lora directly
147
+ mlxsmith lora train --model Qwen/Qwen3-4B-Instruct-2507 --data data/prefs --train-mode dpo -- --beta 0.1
148
+
149
+ # Generate synthetic datasets
150
+ mlxsmith lora synthetic prompts -- --model mlx-community/Qwen3-4B-Instruct-2507-4bit --num-samples 1000
151
+
152
+ # Train judge model
153
+ mlxsmith lora judge -- --model mlx-community/Qwen3-4B-Instruct-2507-4bit --data data/prefs
154
+ ```
155
+
134
156
  ## Serving
135
157
 
136
158
  OpenAI-compatible `/v1/chat/completions` endpoint.
@@ -202,6 +224,7 @@ Built-in verifiers for eval, RFT, and preference tuning:
202
224
  - **pytest** — sandboxed test execution
203
225
  - **docker** — containerized verification
204
226
  - **compose** — multi-verifier composition (AND/OR/weighted)
227
+ - **llm_judge** — LLM-based self-verification / ThinkPRM-style verifier
205
228
 
206
229
  See `docs/VERIFIERS.md` for the verifier API.
207
230
 
@@ -230,6 +253,9 @@ mlxsmith config env # show environment variable mapping
230
253
 
231
254
  Config sources (in priority order): CLI flags > environment variables (`MLXSMITH__SECTION__KEY`) > config file > defaults.
232
255
 
256
+ Training optimizers are configurable via `train.optimizer` and `train.optimizer_kwargs`
257
+ (for example `adamw`, `adam`, `qhadam`, `muon` when available in MLX).
258
+
233
259
  ## SDK (programmatic API)
234
260
 
235
261
  For building custom training loops:
@@ -19,6 +19,9 @@ pip install mlxsmith
19
19
  # Apple Silicon training + serving
20
20
  pip install "mlxsmith[mlx,llm,serve]"
21
21
 
22
+ # mlx-lm-lora passthrough (advanced training methods)
23
+ pip install "mlxsmith[lora]"
24
+
22
25
  # Everything
23
26
  pip install "mlxsmith[all]"
24
27
  ```
@@ -76,6 +79,22 @@ mlxsmith distill --teacher large-model --student small-model --mode opd
76
79
  mlxsmith pipeline
77
80
  ```
78
81
 
82
+ ### mlx-lm-lora parity (all methods)
83
+
84
+ Use the passthrough to access mlx-lm-lora features (DPO variants, GRPO variants,
85
+ PPO, synthetic datasets, judge training, etc.):
86
+
87
+ ```bash
88
+ # Train with mlx-lm-lora directly
89
+ mlxsmith lora train --model Qwen/Qwen3-4B-Instruct-2507 --data data/prefs --train-mode dpo -- --beta 0.1
90
+
91
+ # Generate synthetic datasets
92
+ mlxsmith lora synthetic prompts -- --model mlx-community/Qwen3-4B-Instruct-2507-4bit --num-samples 1000
93
+
94
+ # Train judge model
95
+ mlxsmith lora judge -- --model mlx-community/Qwen3-4B-Instruct-2507-4bit --data data/prefs
96
+ ```
97
+
79
98
  ## Serving
80
99
 
81
100
  OpenAI-compatible `/v1/chat/completions` endpoint.
@@ -147,6 +166,7 @@ Built-in verifiers for eval, RFT, and preference tuning:
147
166
  - **pytest** — sandboxed test execution
148
167
  - **docker** — containerized verification
149
168
  - **compose** — multi-verifier composition (AND/OR/weighted)
169
+ - **llm_judge** — LLM-based self-verification / ThinkPRM-style verifier
150
170
 
151
171
  See `docs/VERIFIERS.md` for the verifier API.
152
172
 
@@ -175,6 +195,9 @@ mlxsmith config env # show environment variable mapping
175
195
 
176
196
  Config sources (in priority order): CLI flags > environment variables (`MLXSMITH__SECTION__KEY`) > config file > defaults.
177
197
 
198
+ Training optimizers are configurable via `train.optimizer` and `train.optimizer_kwargs`
199
+ (for example `adamw`, `adam`, `qhadam`, `muon` when available in MLX).
200
+
178
201
  ## SDK (programmatic API)
179
202
 
180
203
  For building custom training loops:
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "mlxsmith"
7
- version = "0.1.2"
7
+ version = "0.1.3"
8
8
  description = "Apple Silicon MLX fine-tuning toolkit — SFT, DPO/ORPO, GRPO, distillation, and OpenAI-compatible serving."
9
9
  readme = {file = "README.md", content-type = "text/markdown"}
10
10
  requires-python = ">=3.10"
@@ -47,6 +47,7 @@ llm = [
47
47
  "transformers>=5.0.0",
48
48
  "datasets>=3.0.0",
49
49
  ]
50
+ lora = ["mlx-lm-lora>=1.0.0"]
50
51
  serve = [
51
52
  "fastapi>=0.128.0",
52
53
  "uvicorn>=0.40.0",
@@ -56,6 +57,7 @@ dev = ["pytest>=9.0.0", "ruff>=0.14.0"]
56
57
  all = [
57
58
  "mlx>=0.30.4",
58
59
  "mlx-lm>=0.30.5",
60
+ "mlx-lm-lora>=1.0.0",
59
61
  "transformers>=5.0.0",
60
62
  "datasets>=3.0.0",
61
63
  "fastapi>=0.128.0",
@@ -44,7 +44,12 @@ def run_bench(
44
44
  mode = (mode or "inference").lower()
45
45
 
46
46
  if mode == "trainer":
47
- opt, _params = llm.optimizer_and_params(lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
47
+ opt, _params = llm.optimizer_and_params(
48
+ lr=cfg.train.lr,
49
+ weight_decay=cfg.train.weight_decay,
50
+ optimizer=cfg.train.optimizer,
51
+ optimizer_kwargs=cfg.train.optimizer_kwargs,
52
+ )
48
53
  prompt_ids = llm.encode(prompt)
49
54
  ids = llm.encode(prompt + " " + "x" * max_tokens)
50
55
  for i in range(max(1, reps)):
@@ -59,7 +64,12 @@ def run_bench(
59
64
  elapsed = max(time.time() - t0, 1e-6)
60
65
  results.append({"rep": i, "steps": steps, "time_s": elapsed, "steps_per_s": steps / elapsed})
61
66
  elif mode == "end_to_end":
62
- opt, _params = llm.optimizer_and_params(lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
67
+ opt, _params = llm.optimizer_and_params(
68
+ lr=cfg.train.lr,
69
+ weight_decay=cfg.train.weight_decay,
70
+ optimizer=cfg.train.optimizer,
71
+ optimizer_kwargs=cfg.train.optimizer_kwargs,
72
+ )
63
73
  for i in range(max(1, reps)):
64
74
  t0 = time.time()
65
75
  gen = llm.generate(prompt, max_new_tokens=max_tokens, temperature=0.0)
@@ -24,6 +24,8 @@ from .train.sft import run_sft
24
24
  from .train.pref import run_pref
25
25
  from .train.rft import run_rft
26
26
  from .train.distill import run_distill
27
+ from .train.online_dpo import run_online_dpo
28
+ from .train.self_verify import run_self_verify
27
29
  from .eval import run_eval
28
30
  from .bench import run_bench
29
31
  from .rlm import run_rlm, run_rlm_orchestrated
@@ -40,6 +42,13 @@ from .envs import (
40
42
  resolve_env_path as resolve_env_path_plugin,
41
43
  load_manifest as load_env_manifest,
42
44
  )
45
+ from .integrations.mlx_lm_lora import (
46
+ build_train_command as build_mlx_lm_lora_train_command,
47
+ build_synthetic_command as build_mlx_lm_lora_synth_command,
48
+ build_judge_command as build_mlx_lm_lora_judge_command,
49
+ build_reward_functions_command as build_mlx_lm_lora_reward_functions_command,
50
+ run_command as run_mlx_lm_lora_command,
51
+ )
43
52
 
44
53
  app = typer.Typer(
45
54
  add_completion=False,
@@ -65,6 +74,9 @@ def init(path: str = typer.Argument(..., help="Project directory to create")):
65
74
  (p / "verifiers" / "regex.py").write_text(_sample_verifier_regex(), encoding="utf-8")
66
75
  (p / "verifiers" / "pytest.py").write_text(_sample_verifier_pytest(), encoding="utf-8")
67
76
  (p / "verifiers" / "jsonschema.py").write_text(_sample_verifier_jsonschema(), encoding="utf-8")
77
+ (p / "verifiers" / "llm_judge.py").write_text(_sample_verifier_llm_judge(), encoding="utf-8")
78
+ (p / "verifiers" / "rubrics").mkdir(parents=True, exist_ok=True)
79
+ (p / "verifiers" / "rubrics" / "coding.txt").write_text(_sample_judge_rubric(), encoding="utf-8")
68
80
  (p / "eval" / "suites" / "coding.yaml").write_text(_sample_eval_suite(), encoding="utf-8")
69
81
  console.print(f"[green]Initialized[/green] {p.resolve()}")
70
82
 
@@ -341,14 +353,19 @@ def pref(
341
353
  data: str = typer.Option("data/prefs", "--data"),
342
354
  model: str = typer.Option(..., "--model", help="Base adapter or model path (e.g., runs/sft_0001/adapter)"),
343
355
  accel: Optional[str] = typer.Option(None, "--accel", help="Override accel.backend"),
344
- algo: Optional[str] = typer.Option(None, "--algo", help="Override pref.algo (dpo|orpo|grpo)"),
356
+ algo: Optional[str] = typer.Option(None, "--algo", help="Override pref.algo (legacy)"),
357
+ loss_type: Optional[str] = typer.Option(None, "--loss-type", help="dpo|cpo|orpo|ipo|hinge"),
345
358
  ):
346
359
  root = project_root_from_cwd()
360
+ overrides = {}
361
+ if loss_type is not None:
362
+ overrides["pref.loss_type"] = loss_type
347
363
  cfg = get_config(
348
364
  config_path=config,
349
365
  root=root,
350
366
  accel_backend=accel,
351
367
  algo=algo,
368
+ **overrides,
352
369
  )
353
370
  data_dir = root / data
354
371
  run = run_pref(root, cfg, data_dir, Path(model), cfg.accel.backend)
@@ -363,13 +380,27 @@ def rft(
363
380
  model: str = typer.Option(..., "--model"),
364
381
  accel: Optional[str] = typer.Option(None, "--accel", help="Override accel.backend"),
365
382
  rollouts: Optional[int] = typer.Option(None, "--rollouts", help="Override rft.rollouts"),
383
+ loss_type: Optional[str] = typer.Option(None, "--loss-type", help="grpo|dr_grpo|dapo"),
384
+ epsilon_low: Optional[float] = typer.Option(None, "--epsilon-low"),
385
+ epsilon_high: Optional[float] = typer.Option(None, "--epsilon-high"),
386
+ token_level_loss: Optional[bool] = typer.Option(None, "--token-level-loss/--sequence-level-loss"),
366
387
  ):
367
388
  root = project_root_from_cwd()
389
+ overrides = {}
390
+ if loss_type is not None:
391
+ overrides["rft.loss_type"] = loss_type
392
+ if epsilon_low is not None:
393
+ overrides["rft.epsilon_low"] = epsilon_low
394
+ if epsilon_high is not None:
395
+ overrides["rft.epsilon_high"] = epsilon_high
396
+ if token_level_loss is not None:
397
+ overrides["rft.token_level_loss"] = token_level_loss
368
398
  cfg = get_config(
369
399
  config_path=config,
370
400
  root=root,
371
401
  accel_backend=accel,
372
402
  rollouts=rollouts,
403
+ **overrides,
373
404
  )
374
405
  run = run_rft(root, cfg, root / env, root / verifier, Path(model), cfg.accel.backend)
375
406
  console.print(f"[bold]Run:[/bold] {run.run_dir}")
@@ -437,6 +468,142 @@ def distill(
437
468
  console.print(f"[bold]Run:[/bold] {run.run_dir}")
438
469
 
439
470
 
471
+ @app.command("online-dpo")
472
+ def online_dpo(
473
+ data: str = typer.Option(..., "--data", help="JSONL with prompts"),
474
+ model: str = typer.Option(..., "--model"),
475
+ judge_model: Optional[str] = typer.Option(None, "--judge-model"),
476
+ judge_backend: str = typer.Option("mlx-lm", "--judge-backend"),
477
+ rubric: Optional[str] = typer.Option(None, "--rubric"),
478
+ group_size: Optional[int] = typer.Option(None, "--group-size"),
479
+ max_new_tokens: Optional[int] = typer.Option(None, "--max-new-tokens"),
480
+ temperature: Optional[float] = typer.Option(None, "--temperature"),
481
+ config: str = typer.Option("mlxsmith.yaml", "-c", "--config", help="Config file path"),
482
+ accel: Optional[str] = typer.Option(None, "--accel", help="Override accel.backend"),
483
+ ):
484
+ root = project_root_from_cwd()
485
+ cfg = get_config(config_path=config, root=root, accel_backend=accel)
486
+ run = run_online_dpo(
487
+ root,
488
+ cfg,
489
+ Path(data),
490
+ model,
491
+ cfg.accel.backend,
492
+ judge_model=judge_model,
493
+ judge_backend=judge_backend,
494
+ rubric=rubric,
495
+ group_size=group_size,
496
+ max_new_tokens=max_new_tokens,
497
+ temperature=temperature,
498
+ )
499
+ console.print(f"[bold]Run:[/bold] {run.run_dir}")
500
+
501
+
502
+ @app.command("self-verify")
503
+ def self_verify(
504
+ data: str = typer.Option(..., "--data", help="JSONL with prompts"),
505
+ model: str = typer.Option(..., "--model"),
506
+ verifier_model: Optional[str] = typer.Option(None, "--verifier-model"),
507
+ verifier_backend: str = typer.Option("mlx-lm", "--verifier-backend"),
508
+ rubric: Optional[str] = typer.Option(None, "--rubric"),
509
+ max_new_tokens: Optional[int] = typer.Option(None, "--max-new-tokens"),
510
+ temperature: Optional[float] = typer.Option(None, "--temperature"),
511
+ config: str = typer.Option("mlxsmith.yaml", "-c", "--config", help="Config file path"),
512
+ accel: Optional[str] = typer.Option(None, "--accel", help="Override accel.backend"),
513
+ ):
514
+ root = project_root_from_cwd()
515
+ cfg = get_config(config_path=config, root=root, accel_backend=accel)
516
+ run = run_self_verify(
517
+ root,
518
+ cfg,
519
+ Path(data),
520
+ model,
521
+ cfg.accel.backend,
522
+ verifier_model=verifier_model,
523
+ verifier_backend=verifier_backend,
524
+ rubric=rubric,
525
+ max_new_tokens=max_new_tokens,
526
+ temperature=temperature,
527
+ )
528
+ console.print(f"[bold]Run:[/bold] {run.run_dir}")
529
+
530
+
531
+ lora_app = typer.Typer(help="mlx-lm-lora passthrough commands")
532
+ app.add_typer(lora_app, name="lora")
533
+
534
+
535
+ @lora_app.command(
536
+ "train",
537
+ context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
538
+ )
539
+ def lora_train(
540
+ ctx: typer.Context,
541
+ config: Optional[str] = typer.Option(None, "--config", help="mlx-lm-lora config path"),
542
+ model: Optional[str] = typer.Option(None, "--model", help="Model id or path"),
543
+ data: Optional[str] = typer.Option(None, "--data", help="Dataset path or HF dataset"),
544
+ train_mode: Optional[str] = typer.Option(None, "--train-mode", help="sft|dpo|orpo|grpo|ppo|..."),
545
+ train_type: Optional[str] = typer.Option(None, "--train-type", help="lora|dora|full"),
546
+ dry_run: bool = typer.Option(False, "--dry-run"),
547
+ ):
548
+ """Run mlx-lm-lora training with passthrough args.
549
+
550
+ Use `--` to pass through any additional mlx-lm-lora flags.
551
+ """
552
+ root = project_root_from_cwd()
553
+ cmd = build_mlx_lm_lora_train_command(
554
+ config=config,
555
+ model=model,
556
+ data=data,
557
+ train_mode=train_mode,
558
+ train_type=train_type,
559
+ extra_args=list(ctx.args),
560
+ )
561
+ run_mlx_lm_lora_command(cmd, dry_run=dry_run, cwd=root)
562
+
563
+
564
+ @lora_app.command(
565
+ "synthetic",
566
+ context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
567
+ )
568
+ def lora_synthetic(
569
+ ctx: typer.Context,
570
+ kind: str = typer.Argument(..., help="prompts|sft|dpo"),
571
+ dry_run: bool = typer.Option(False, "--dry-run"),
572
+ ):
573
+ """Run mlx-lm-lora synthetic dataset generation."""
574
+ root = project_root_from_cwd()
575
+ cmd = build_mlx_lm_lora_synth_command(kind, extra_args=list(ctx.args))
576
+ run_mlx_lm_lora_command(cmd, dry_run=dry_run, cwd=root)
577
+
578
+
579
+ @lora_app.command(
580
+ "judge",
581
+ context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
582
+ )
583
+ def lora_judge(
584
+ ctx: typer.Context,
585
+ dry_run: bool = typer.Option(False, "--dry-run"),
586
+ ):
587
+ """Run mlx-lm-lora judge model training."""
588
+ root = project_root_from_cwd()
589
+ cmd = build_mlx_lm_lora_judge_command(extra_args=list(ctx.args))
590
+ run_mlx_lm_lora_command(cmd, dry_run=dry_run, cwd=root)
591
+
592
+
593
+ @lora_app.command(
594
+ "reward-functions",
595
+ context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
596
+ )
597
+ def lora_reward_functions(
598
+ ctx: typer.Context,
599
+ dry_run: bool = typer.Option(False, "--dry-run"),
600
+ ):
601
+ """List mlx-lm-lora reward functions."""
602
+ root = project_root_from_cwd()
603
+ cmd = build_mlx_lm_lora_reward_functions_command(extra_args=list(ctx.args))
604
+ run_mlx_lm_lora_command(cmd, dry_run=dry_run, cwd=root)
605
+
606
+
440
607
  @app.command()
441
608
  def eval(
442
609
  suite: str = typer.Option("eval/suites/coding.yaml", "--suite"),
@@ -933,6 +1100,25 @@ def verify(prompt: str, completion: str, workdir: str, **kwargs):
933
1100
  """
934
1101
 
935
1102
 
1103
+ def _sample_verifier_llm_judge() -> str:
1104
+ return """from mlxsmith.verifiers.llm_judge import verify as _verify
1105
+
1106
+ def verify(prompt: str, completion: str, workdir: str, **kwargs):
1107
+ # Pass model=... or set MLXSMITH_JUDGE_MODEL for the judge model id.
1108
+ return _verify(prompt, completion, workdir, **kwargs)
1109
+ """
1110
+
1111
+
1112
+ def _sample_judge_rubric() -> str:
1113
+ return """Score from 0.0 to 1.0.
1114
+ - 1.0: Correct, complete, and safe.
1115
+ - 0.7: Mostly correct with small issues.
1116
+ - 0.4: Partial correctness or unclear reasoning.
1117
+ - 0.0: Incorrect or unsafe.
1118
+ Return JSON only.
1119
+ """
1120
+
1121
+
936
1122
  def _sample_eval_suite() -> str:
937
1123
  return """name: coding-eval-sample
938
1124
  notes: |
@@ -47,6 +47,8 @@ class TrainConfig(BaseModel):
47
47
  grad_accum: int = 8
48
48
  lr: float = 2e-4
49
49
  weight_decay: float = 0.0
50
+ optimizer: str = "adamw"
51
+ optimizer_kwargs: Dict[str, Any] = Field(default_factory=dict)
50
52
  iters: int = 1000
51
53
  save_every: int = 100
52
54
  eval_every: int = 100
@@ -61,6 +63,11 @@ class TrainConfig(BaseModel):
61
63
  raise ValueError("value must be non-negative")
62
64
  return v
63
65
 
66
+ @field_validator("optimizer")
67
+ @classmethod
68
+ def normalize_optimizer(cls, v: str) -> str:
69
+ return v.strip().lower()
70
+
64
71
 
65
72
  class LoraConfig(BaseModel):
66
73
  """LoRA/DoRA adapter configuration."""
@@ -89,11 +96,13 @@ class LoraConfig(BaseModel):
89
96
 
90
97
 
91
98
  class PrefConfig(BaseModel):
92
- """Preference tuning configuration (DPO, ORPO, GRPO)."""
99
+ """Preference tuning configuration (DPO variants)."""
93
100
 
94
101
  algo: Literal["dpo", "orpo", "grpo"] = "dpo"
102
+ loss_type: Literal["dpo", "cpo", "orpo", "ipo", "hinge"] = "dpo"
95
103
  beta: float = 0.1
96
104
  kl_coeff: float = 0.0
105
+ delta: float = 0.0
97
106
  reference_model: Optional[str] = None
98
107
 
99
108
 
@@ -101,12 +110,16 @@ class RftConfig(BaseModel):
101
110
  """Reinforcement fine-tuning configuration."""
102
111
 
103
112
  algo: Literal["grpo"] = "grpo"
113
+ loss_type: Literal["grpo", "dr_grpo", "dapo"] = "grpo"
104
114
  rollouts: int = 8
105
115
  kl_coeff: float = 0.02
106
116
  max_steps_per_task: int = 1
107
117
  temperature: float = 0.8
108
118
  max_new_tokens: int = 256
109
119
  normalize_advantage: bool = True
120
+ epsilon_low: float = 0.2
121
+ epsilon_high: float = 0.2
122
+ token_level_loss: bool = False
110
123
  reference_model: Optional[str] = None
111
124
 
112
125
 
@@ -164,6 +177,7 @@ CLI_ALIASES: dict[str, tuple[str, ...]] = {
164
177
  "lr": ("train", "lr"),
165
178
  "batch_size": ("train", "batch_size"),
166
179
  "iters": ("train", "iters"),
180
+ "optimizer": ("train", "optimizer"),
167
181
  "model_id": ("model", "id"),
168
182
  "accel_backend": ("accel", "backend"),
169
183
  "host": ("serve", "host"),
@@ -0,0 +1,19 @@
1
+ """External integrations for mlxsmith."""
2
+
3
+ from .mlx_lm_lora import (
4
+ build_train_command as build_mlx_lm_lora_train_command,
5
+ build_synthetic_command as build_mlx_lm_lora_synth_command,
6
+ build_judge_command as build_mlx_lm_lora_judge_command,
7
+ build_reward_functions_command as build_mlx_lm_lora_reward_functions_command,
8
+ run_command as run_mlx_lm_lora_command,
9
+ ensure_available as ensure_mlx_lm_lora_available,
10
+ )
11
+
12
+ __all__ = [
13
+ "build_mlx_lm_lora_train_command",
14
+ "build_mlx_lm_lora_synth_command",
15
+ "build_mlx_lm_lora_judge_command",
16
+ "build_mlx_lm_lora_reward_functions_command",
17
+ "run_mlx_lm_lora_command",
18
+ "ensure_mlx_lm_lora_available",
19
+ ]
@@ -0,0 +1,117 @@
1
+ """Passthrough helpers for mlx-lm-lora CLI integration."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import importlib.util
6
+ import os
7
+ import shlex
8
+ import subprocess
9
+ import sys
10
+ from pathlib import Path
11
+ from typing import Optional, Sequence
12
+
13
+ from rich.console import Console
14
+
15
+ console = Console()
16
+
17
+
18
+ def ensure_available() -> None:
19
+ if importlib.util.find_spec("mlx_lm_lora") is None:
20
+ raise RuntimeError(
21
+ "mlx-lm-lora is not installed. Install with: pip install 'mlxsmith[lora]' or 'mlx-lm-lora'"
22
+ )
23
+
24
+
25
+ def _flag_present(args: Sequence[str], *flags: str) -> bool:
26
+ for flag in flags:
27
+ if flag in args:
28
+ return True
29
+ if flag.startswith("--"):
30
+ prefix = flag + "="
31
+ if any(a.startswith(prefix) for a in args):
32
+ return True
33
+ return False
34
+
35
+
36
+ def _append_flag(cmd: list[str], args: Sequence[str], flag: str, value: Optional[str]) -> None:
37
+ if value is None:
38
+ return
39
+ if _flag_present(args, flag):
40
+ return
41
+ cmd.extend([flag, value])
42
+
43
+
44
+ def _base_python_cmd(module: str) -> list[str]:
45
+ return [sys.executable, "-m", module]
46
+
47
+
48
+ def build_train_command(
49
+ *,
50
+ config: Optional[str] = None,
51
+ model: Optional[str] = None,
52
+ data: Optional[str] = None,
53
+ train_mode: Optional[str] = None,
54
+ train_type: Optional[str] = None,
55
+ extra_args: Sequence[str] = (),
56
+ ) -> list[str]:
57
+ args = list(extra_args)
58
+ cmd = _base_python_cmd("mlx_lm_lora.train")
59
+ _append_flag(cmd, args, "--config", config)
60
+ _append_flag(cmd, args, "--model", model)
61
+ _append_flag(cmd, args, "--data", data)
62
+ _append_flag(cmd, args, "--train-mode", train_mode)
63
+ _append_flag(cmd, args, "--train-type", train_type)
64
+ cmd.extend(args)
65
+ return cmd
66
+
67
+
68
+ def build_synthetic_command(
69
+ kind: str,
70
+ *,
71
+ extra_args: Sequence[str] = (),
72
+ ) -> list[str]:
73
+ kind = kind.strip().lower()
74
+ module = {
75
+ "prompts": "mlx_lm_lora.synthetic_prompts",
76
+ "sft": "mlx_lm_lora.synthetic_sft",
77
+ "dpo": "mlx_lm_lora.synthetic_dpo",
78
+ }.get(kind)
79
+ if module is None:
80
+ raise ValueError(f"Unknown synthetic kind: {kind}")
81
+ cmd = _base_python_cmd(module)
82
+ cmd.extend(list(extra_args))
83
+ return cmd
84
+
85
+
86
+ def build_judge_command(*, extra_args: Sequence[str] = ()) -> list[str]:
87
+ cmd = _base_python_cmd("mlx_lm_lora.train_judge")
88
+ cmd.extend(list(extra_args))
89
+ return cmd
90
+
91
+
92
+ def build_reward_functions_command(*, extra_args: Sequence[str] = ()) -> list[str]:
93
+ cmd = _base_python_cmd("mlx_lm_lora.train")
94
+ cmd.append("--list-reward-functions")
95
+ cmd.extend(list(extra_args))
96
+ return cmd
97
+
98
+
99
+ def run_command(
100
+ cmd: Sequence[str],
101
+ *,
102
+ dry_run: bool = False,
103
+ cwd: Optional[Path] = None,
104
+ env: Optional[dict] = None,
105
+ ) -> int:
106
+ if dry_run:
107
+ console.print("[cyan]mlx-lm-lora cmd[/cyan]", shlex.join(list(cmd)))
108
+ return 0
109
+ ensure_available()
110
+ run_env = os.environ.copy()
111
+ if env:
112
+ run_env.update(env)
113
+ console.print("[cyan]mlx-lm-lora cmd[/cyan]", shlex.join(list(cmd)))
114
+ result = subprocess.run(list(cmd), cwd=str(cwd) if cwd else None, env=run_env, check=False)
115
+ if result.returncode != 0:
116
+ raise RuntimeError(f"mlx-lm-lora failed with exit code {result.returncode}")
117
+ return result.returncode
@@ -112,7 +112,14 @@ class LLMBackend(Protocol):
112
112
  def value_and_grad(self, loss_fn) -> tuple[Any, Any | None]:
113
113
  """Return (loss, grads) using backend autograd when available."""
114
114
 
115
- def optimizer_and_params(self, *, lr: float, weight_decay: float = 0.0) -> tuple[Any, Any]:
115
+ def optimizer_and_params(
116
+ self,
117
+ *,
118
+ lr: float,
119
+ weight_decay: float = 0.0,
120
+ optimizer: str | None = None,
121
+ optimizer_kwargs: dict | None = None,
122
+ ) -> tuple[Any, Any]:
116
123
  """Return (optimizer, trainable_params_tree)."""
117
124
 
118
125
  def apply_grads(self, optimizer: Any, grads: Any) -> None: