mlx-forge 0.2.0__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 (127) hide show
  1. mlx_forge-0.2.0/LICENSE +21 -0
  2. mlx_forge-0.2.0/PKG-INFO +246 -0
  3. mlx_forge-0.2.0/README.md +201 -0
  4. mlx_forge-0.2.0/mlx_forge/__init__.py +456 -0
  5. mlx_forge-0.2.0/mlx_forge/_version.py +1 -0
  6. mlx_forge-0.2.0/mlx_forge/adapters/__init__.py +0 -0
  7. mlx_forge-0.2.0/mlx_forge/adapters/lora.py +287 -0
  8. mlx_forge-0.2.0/mlx_forge/adapters/targeting.py +162 -0
  9. mlx_forge-0.2.0/mlx_forge/cli/__init__.py +0 -0
  10. mlx_forge-0.2.0/mlx_forge/cli/data_cmd.py +220 -0
  11. mlx_forge-0.2.0/mlx_forge/cli/generate_cmd.py +99 -0
  12. mlx_forge-0.2.0/mlx_forge/cli/main.py +254 -0
  13. mlx_forge-0.2.0/mlx_forge/cli/prepare_cmd.py +23 -0
  14. mlx_forge-0.2.0/mlx_forge/cli/studio_cmd.py +21 -0
  15. mlx_forge-0.2.0/mlx_forge/cli/train_cmd.py +11 -0
  16. mlx_forge-0.2.0/mlx_forge/config.py +176 -0
  17. mlx_forge-0.2.0/mlx_forge/data/__init__.py +0 -0
  18. mlx_forge-0.2.0/mlx_forge/data/backend.py +157 -0
  19. mlx_forge-0.2.0/mlx_forge/data/batching.py +238 -0
  20. mlx_forge-0.2.0/mlx_forge/data/catalog.py +370 -0
  21. mlx_forge-0.2.0/mlx_forge/data/converter.py +204 -0
  22. mlx_forge-0.2.0/mlx_forge/data/formats.py +191 -0
  23. mlx_forge-0.2.0/mlx_forge/data/mixing.py +61 -0
  24. mlx_forge-0.2.0/mlx_forge/data/packing.py +93 -0
  25. mlx_forge-0.2.0/mlx_forge/data/preprocessing.py +198 -0
  26. mlx_forge-0.2.0/mlx_forge/data/registry.py +221 -0
  27. mlx_forge-0.2.0/mlx_forge/data/validate.py +227 -0
  28. mlx_forge-0.2.0/mlx_forge/inference/__init__.py +1 -0
  29. mlx_forge-0.2.0/mlx_forge/inference/cache.py +122 -0
  30. mlx_forge-0.2.0/mlx_forge/inference/engine.py +232 -0
  31. mlx_forge-0.2.0/mlx_forge/inference/sampling.py +88 -0
  32. mlx_forge-0.2.0/mlx_forge/logging/__init__.py +0 -0
  33. mlx_forge-0.2.0/mlx_forge/logging/metrics.py +52 -0
  34. mlx_forge-0.2.0/mlx_forge/losses/__init__.py +6 -0
  35. mlx_forge-0.2.0/mlx_forge/losses/dpo.py +98 -0
  36. mlx_forge-0.2.0/mlx_forge/losses/sft.py +77 -0
  37. mlx_forge-0.2.0/mlx_forge/manifest.py +164 -0
  38. mlx_forge-0.2.0/mlx_forge/models/__init__.py +0 -0
  39. mlx_forge-0.2.0/mlx_forge/models/_base/__init__.py +12 -0
  40. mlx_forge-0.2.0/mlx_forge/models/_base/activations.py +25 -0
  41. mlx_forge-0.2.0/mlx_forge/models/_base/args.py +38 -0
  42. mlx_forge-0.2.0/mlx_forge/models/_base/attention.py +101 -0
  43. mlx_forge-0.2.0/mlx_forge/models/_base/rope.py +276 -0
  44. mlx_forge-0.2.0/mlx_forge/models/architectures/__init__.py +4 -0
  45. mlx_forge-0.2.0/mlx_forge/models/architectures/gemma.py +352 -0
  46. mlx_forge-0.2.0/mlx_forge/models/architectures/llama.py +236 -0
  47. mlx_forge-0.2.0/mlx_forge/models/architectures/phi3.py +261 -0
  48. mlx_forge-0.2.0/mlx_forge/models/architectures/phi4.py +221 -0
  49. mlx_forge-0.2.0/mlx_forge/models/architectures/qwen2.py +221 -0
  50. mlx_forge-0.2.0/mlx_forge/models/architectures/qwen3.py +231 -0
  51. mlx_forge-0.2.0/mlx_forge/models/architectures/qwen3_5.py +752 -0
  52. mlx_forge-0.2.0/mlx_forge/models/loader.py +143 -0
  53. mlx_forge-0.2.0/mlx_forge/models/memory.py +376 -0
  54. mlx_forge-0.2.0/mlx_forge/models/quantize.py +39 -0
  55. mlx_forge-0.2.0/mlx_forge/models/registry.py +108 -0
  56. mlx_forge-0.2.0/mlx_forge/models/resolve.py +205 -0
  57. mlx_forge-0.2.0/mlx_forge/recipes/__init__.py +5 -0
  58. mlx_forge-0.2.0/mlx_forge/recipes/auto_config.py +104 -0
  59. mlx_forge-0.2.0/mlx_forge/recipes/built_in/chat_sft.yaml +42 -0
  60. mlx_forge-0.2.0/mlx_forge/recipes/built_in/instruction_sft.yaml +42 -0
  61. mlx_forge-0.2.0/mlx_forge/recipes/built_in/preference_dpo.yaml +46 -0
  62. mlx_forge-0.2.0/mlx_forge/recipes/built_in/writing_style.yaml +42 -0
  63. mlx_forge-0.2.0/mlx_forge/recipes/registry.py +90 -0
  64. mlx_forge-0.2.0/mlx_forge/studio/__init__.py +4 -0
  65. mlx_forge-0.2.0/mlx_forge/studio/api/__init__.py +1 -0
  66. mlx_forge-0.2.0/mlx_forge/studio/api/config_schema.py +152 -0
  67. mlx_forge-0.2.0/mlx_forge/studio/api/data_library.py +76 -0
  68. mlx_forge-0.2.0/mlx_forge/studio/api/datasets.py +44 -0
  69. mlx_forge-0.2.0/mlx_forge/studio/api/inference.py +62 -0
  70. mlx_forge-0.2.0/mlx_forge/studio/api/memory.py +58 -0
  71. mlx_forge-0.2.0/mlx_forge/studio/api/models.py +46 -0
  72. mlx_forge-0.2.0/mlx_forge/studio/api/queue.py +63 -0
  73. mlx_forge-0.2.0/mlx_forge/studio/api/recipes.py +67 -0
  74. mlx_forge-0.2.0/mlx_forge/studio/api/runs.py +73 -0
  75. mlx_forge-0.2.0/mlx_forge/studio/api/training.py +48 -0
  76. mlx_forge-0.2.0/mlx_forge/studio/frontend/assets/index-DfE9wCUu.js +46 -0
  77. mlx_forge-0.2.0/mlx_forge/studio/frontend/assets/index-DoKRRrtV.css +1 -0
  78. mlx_forge-0.2.0/mlx_forge/studio/frontend/index.html +14 -0
  79. mlx_forge-0.2.0/mlx_forge/studio/server.py +210 -0
  80. mlx_forge-0.2.0/mlx_forge/studio/services/__init__.py +1 -0
  81. mlx_forge-0.2.0/mlx_forge/studio/services/data_library_service.py +46 -0
  82. mlx_forge-0.2.0/mlx_forge/studio/services/dataset_service.py +73 -0
  83. mlx_forge-0.2.0/mlx_forge/studio/services/memory_service.py +71 -0
  84. mlx_forge-0.2.0/mlx_forge/studio/services/metrics_watcher.py +56 -0
  85. mlx_forge-0.2.0/mlx_forge/studio/services/model_library_service.py +77 -0
  86. mlx_forge-0.2.0/mlx_forge/studio/services/model_service.py +107 -0
  87. mlx_forge-0.2.0/mlx_forge/studio/services/queue_service.py +178 -0
  88. mlx_forge-0.2.0/mlx_forge/studio/services/recipe_service.py +47 -0
  89. mlx_forge-0.2.0/mlx_forge/studio/services/run_service.py +242 -0
  90. mlx_forge-0.2.0/mlx_forge/studio/services/training_service.py +113 -0
  91. mlx_forge-0.2.0/mlx_forge/trainer/__init__.py +0 -0
  92. mlx_forge-0.2.0/mlx_forge/trainer/callbacks.py +150 -0
  93. mlx_forge-0.2.0/mlx_forge/trainer/checkpoint.py +187 -0
  94. mlx_forge-0.2.0/mlx_forge/trainer/dpo_trainer.py +118 -0
  95. mlx_forge-0.2.0/mlx_forge/trainer/optimizer.py +123 -0
  96. mlx_forge-0.2.0/mlx_forge/trainer/state.py +20 -0
  97. mlx_forge-0.2.0/mlx_forge/trainer/trainer.py +319 -0
  98. mlx_forge-0.2.0/mlx_forge.egg-info/PKG-INFO +246 -0
  99. mlx_forge-0.2.0/mlx_forge.egg-info/SOURCES.txt +125 -0
  100. mlx_forge-0.2.0/mlx_forge.egg-info/dependency_links.txt +1 -0
  101. mlx_forge-0.2.0/mlx_forge.egg-info/entry_points.txt +2 -0
  102. mlx_forge-0.2.0/mlx_forge.egg-info/requires.txt +25 -0
  103. mlx_forge-0.2.0/mlx_forge.egg-info/top_level.txt +1 -0
  104. mlx_forge-0.2.0/pyproject.toml +85 -0
  105. mlx_forge-0.2.0/setup.cfg +4 -0
  106. mlx_forge-0.2.0/tests/test_adapters.py +203 -0
  107. mlx_forge-0.2.0/tests/test_backend.py +278 -0
  108. mlx_forge-0.2.0/tests/test_catalog.py +290 -0
  109. mlx_forge-0.2.0/tests/test_config.py +246 -0
  110. mlx_forge-0.2.0/tests/test_data.py +217 -0
  111. mlx_forge-0.2.0/tests/test_data_validate.py +139 -0
  112. mlx_forge-0.2.0/tests/test_integration.py +288 -0
  113. mlx_forge-0.2.0/tests/test_labels.py +479 -0
  114. mlx_forge-0.2.0/tests/test_m10_performance.py +667 -0
  115. mlx_forge-0.2.0/tests/test_m11_studio.py +604 -0
  116. mlx_forge-0.2.0/tests/test_m12_frontend.py +206 -0
  117. mlx_forge-0.2.0/tests/test_m13_integration.py +708 -0
  118. mlx_forge-0.2.0/tests/test_m9_foundation.py +711 -0
  119. mlx_forge-0.2.0/tests/test_mixing.py +81 -0
  120. mlx_forge-0.2.0/tests/test_model_loading.py +323 -0
  121. mlx_forge-0.2.0/tests/test_qwen3_5.py +1019 -0
  122. mlx_forge-0.2.0/tests/test_resolve.py +223 -0
  123. mlx_forge-0.2.0/tests/test_trainer_infra.py +294 -0
  124. mlx_forge-0.2.0/tests/test_v2_dpo.py +566 -0
  125. mlx_forge-0.2.0/tests/test_v2_memory.py +212 -0
  126. mlx_forge-0.2.0/tests/test_v2_recipes.py +170 -0
  127. mlx_forge-0.2.0/tests/test_v2_studio.py +588 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Jiekai Wang
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,246 @@
1
+ Metadata-Version: 2.4
2
+ Name: mlx-forge
3
+ Version: 0.2.0
4
+ Summary: Fine-tune, experiment with, and run LLMs locally on your Mac
5
+ Author: Jiekai Wang
6
+ License-Expression: MIT
7
+ Project-URL: Repository, https://github.com/moyuan5989/mlx-forge
8
+ Project-URL: Issues, https://github.com/moyuan5989/mlx-forge/issues
9
+ Keywords: mlx,lora,fine-tuning,apple-silicon,llm,qlora,dpo
10
+ Classifier: Development Status :: 4 - Beta
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: Intended Audience :: Science/Research
13
+ Classifier: Operating System :: MacOS
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
19
+ Requires-Python: >=3.10
20
+ Description-Content-Type: text/markdown
21
+ License-File: LICENSE
22
+ Requires-Dist: mlx>=0.18.0
23
+ Requires-Dist: pydantic>=2.0
24
+ Requires-Dist: pyyaml>=6.0
25
+ Requires-Dist: numpy>=1.24.0
26
+ Requires-Dist: transformers>=4.35.0
27
+ Requires-Dist: safetensors>=0.4.0
28
+ Requires-Dist: huggingface-hub>=0.20.0
29
+ Requires-Dist: datasets>=2.16.0
30
+ Requires-Dist: fastapi>=0.104.0
31
+ Requires-Dist: uvicorn[standard]>=0.24.0
32
+ Requires-Dist: websockets>=12.0
33
+ Provides-Extra: wandb
34
+ Requires-Dist: wandb>=0.16.0; extra == "wandb"
35
+ Provides-Extra: dev
36
+ Requires-Dist: pytest>=7.0; extra == "dev"
37
+ Requires-Dist: pytest-timeout>=2.0; extra == "dev"
38
+ Requires-Dist: pytest-asyncio>=0.23.0; extra == "dev"
39
+ Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
40
+ Requires-Dist: httpx>=0.25.0; extra == "dev"
41
+ Requires-Dist: ruff>=0.4.0; extra == "dev"
42
+ Provides-Extra: all
43
+ Requires-Dist: mlx-forge[wandb]; extra == "all"
44
+ Dynamic: license-file
45
+
46
+ # MLX Forge
47
+
48
+ **Fine-tune LLMs on your Mac with MLX. No cloud, no CUDA required.**
49
+
50
+ [![PyPI](https://img.shields.io/pypi/v/mlx-forge)](https://pypi.org/project/mlx-forge/)
51
+ [![Python](https://img.shields.io/pypi/pyversions/mlx-forge)](https://pypi.org/project/mlx-forge/)
52
+ [![License](https://img.shields.io/badge/license-MIT-blue)](LICENSE)
53
+ [![Tests](https://img.shields.io/github/actions/workflow/status/moyuan5989/mlx-forge/test.yml?label=tests)](https://github.com/moyuan5989/mlx-forge/actions)
54
+
55
+ MLX Forge is a complete LLM fine-tuning toolkit that runs entirely on your Mac. Pick a model, upload your data, and start training — all from a browser-based UI. Supports LoRA, QLoRA, DPO, 18+ models, and 20+ curated datasets out of the box.
56
+
57
+ ```bash
58
+ pip install mlx-forge
59
+ mlx-forge studio
60
+ ```
61
+
62
+ <p align="center">
63
+ <img src="assets/studio-new-training.png" alt="MLX Forge Studio — New Training" width="800">
64
+ </p>
65
+
66
+ ## Why MLX Forge?
67
+
68
+ - **One command to start** — `pip install mlx-forge && mlx-forge studio`.
69
+ - **Browser-based Studio UI** — Guided training wizard, real-time loss charts, model library with memory estimates, interactive playground.
70
+ - **Runs on Apple Silicon** — Built on [MLX](https://github.com/ml-explore/mlx). Your data stays on your machine.
71
+ - **Production training features** — QLoRA (67% memory reduction), sequence packing (2-5x speedup), gradient checkpointing, DPO alignment, compiled training loop.
72
+
73
+ ## Quick Start
74
+
75
+ ### Studio UI (recommended)
76
+
77
+ ```bash
78
+ mlx-forge studio
79
+ # Opens at http://127.0.0.1:8741
80
+ ```
81
+
82
+ Pick a recipe, choose a model, upload your data, and start training — all from the browser.
83
+
84
+ ### CLI
85
+
86
+ ```bash
87
+ # Browse and download a dataset
88
+ mlx-forge data catalog
89
+ mlx-forge data download alpaca-cleaned --max-samples 5000
90
+
91
+ # Train
92
+ mlx-forge train --config train.yaml
93
+ ```
94
+
95
+ Models are downloaded from Hugging Face on first run and cached locally. All subsequent runs work offline.
96
+
97
+ ## Studio UI
98
+
99
+ <p align="center">
100
+ <img src="assets/studio-model-library.png" alt="MLX Forge Studio — Model Library" width="800">
101
+ </p>
102
+
103
+ - **New Training** — Guided wizard: pick a recipe (chat, instruction, DPO, writing style), choose a model, configure, and launch
104
+ - **Model Library** — Browse 18+ models with memory estimates for your hardware
105
+ - **Experiments** — Compare runs, view loss curves in real time
106
+ - **Datasets** — Manage your training data
107
+ - **Playground** — Chat with your fine-tuned models interactively
108
+
109
+ ## Supported Models
110
+
111
+ 18 curated models in the Studio library, all tested on Apple Silicon:
112
+
113
+ | Architecture | Models | Sizes |
114
+ |-------------|--------|-------|
115
+ | Qwen | Qwen 2.5, Qwen 3, Qwen 3.5 | 0.5B - 8B |
116
+ | Gemma | Gemma 2, Gemma 3 | 1B - 9B |
117
+ | Llama | Llama 3.1 | 8B |
118
+ | Phi | Phi-3 Mini, Phi-4 Mini | 3.8B |
119
+ | DeepSeek | DeepSeek-R1-Distill (Qwen-based) | 1.5B - 7B |
120
+ | Mistral | Mistral (uses Llama architecture) | 7B |
121
+
122
+ Any HF model using a supported architecture will work — the table above shows the curated models with pre-computed memory estimates in Studio.
123
+
124
+ ## Features
125
+
126
+ **Training**
127
+ - LoRA and QLoRA (4-bit) fine-tuning with 67% memory reduction
128
+ - DPO (Direct Preference Optimization) for alignment
129
+ - Sequence packing for 2-5x speedup on short sequences
130
+ - Gradient checkpointing for 40-60% memory savings
131
+ - Compiled training loop with gradient accumulation
132
+ - Cosine, linear, step, and exponential LR schedules with warmup
133
+ - Resume from any checkpoint
134
+
135
+ **Data**
136
+ - 20+ curated datasets across 7 categories (general, code, math, conversation, reasoning, safety, domain)
137
+ - Auto-detection of chat, completions, text, and preference formats
138
+ - Multi-dataset mixing with weighted sampling
139
+ - Data validation with train/val overlap detection
140
+
141
+ ## CLI Reference
142
+
143
+ | Command | Description |
144
+ |---------|-------------|
145
+ | `mlx-forge studio` | Launch the Studio UI |
146
+ | `mlx-forge train --config FILE` | Run LoRA/QLoRA/DPO training |
147
+ | `mlx-forge generate --model MODEL` | Generate text or interactive chat |
148
+ | `mlx-forge prepare --data FILE --model MODEL` | Pre-tokenize a dataset |
149
+ | `mlx-forge data catalog` | Browse 20+ curated datasets |
150
+ | `mlx-forge data download DATASET` | Download a dataset from the catalog |
151
+ | `mlx-forge data import FILE --name NAME` | Import a local JSONL file |
152
+ | `mlx-forge data validate FILE` | Validate JSONL data |
153
+ | `mlx-forge data inspect NAME` | Preview dataset samples |
154
+ | `mlx-forge data stats NAME` | Show dataset statistics |
155
+
156
+ ## Configuration
157
+
158
+ ```yaml
159
+ schema_version: 1
160
+
161
+ model:
162
+ path: "Qwen/Qwen3-0.6B" # HF model ID or local path
163
+ quantization: # Optional: QLoRA (67% memory savings)
164
+ bits: 4
165
+ group_size: 64
166
+
167
+ adapter:
168
+ preset: "attention-qv" # attention-qv | attention-all | mlp | all-linear
169
+ rank: 16
170
+ scale: 32.0
171
+
172
+ data:
173
+ train: "./train.jsonl"
174
+ valid: "./val.jsonl"
175
+ packing: false # Sequence packing (2-5x speedup)
176
+ max_seq_length: 2048
177
+
178
+ training:
179
+ optimizer: adamw # adam | adamw | sgd | adafactor
180
+ learning_rate: 1.0e-5
181
+ num_iters: 1000
182
+ batch_size: 4
183
+ gradient_checkpointing: false # 40-60% memory savings
184
+ # training_type: dpo # For DPO training
185
+ # dpo_beta: 0.1
186
+
187
+ runtime:
188
+ seed: 42
189
+ ```
190
+
191
+ ## Data Formats
192
+
193
+ MLX Forge auto-detects four JSONL formats:
194
+
195
+ **Chat** — Multi-turn conversations (loss on assistant turns only):
196
+ ```json
197
+ {"messages": [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi!"}]}
198
+ ```
199
+
200
+ **Completions** — Prompt-completion pairs:
201
+ ```json
202
+ {"prompt": "Translate to French: Hello", "completion": "Bonjour"}
203
+ ```
204
+
205
+ **Text** — Raw text for continued pretraining:
206
+ ```json
207
+ {"text": "The quick brown fox jumps over the lazy dog."}
208
+ ```
209
+
210
+ **Preference** — For DPO alignment training:
211
+ ```json
212
+ {"chosen": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "good"}], "rejected": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "bad"}]}
213
+ ```
214
+
215
+ ## Library API
216
+
217
+ All CLI commands are backed by Python functions:
218
+
219
+ ```python
220
+ from mlx_forge import prepare, train
221
+ from mlx_forge.config import TrainingConfig
222
+
223
+ # Train from a config file
224
+ config = TrainingConfig.from_yaml("train.yaml")
225
+ result = train(config=config)
226
+ print(f"Best val loss: {result.best_val_loss:.4f}")
227
+ ```
228
+
229
+ ```python
230
+ from mlx_forge import generate
231
+
232
+ # Generate text with a fine-tuned adapter
233
+ generate(
234
+ model="Qwen/Qwen3-0.6B",
235
+ adapter="~/.mlxforge/runs/my-run/checkpoints/best",
236
+ prompt="Explain quantum computing in simple terms.",
237
+ )
238
+ ```
239
+
240
+ ## Contributing
241
+
242
+ See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, coding standards, and how to submit changes.
243
+
244
+ ## License
245
+
246
+ [MIT](LICENSE)
@@ -0,0 +1,201 @@
1
+ # MLX Forge
2
+
3
+ **Fine-tune LLMs on your Mac with MLX. No cloud, no CUDA required.**
4
+
5
+ [![PyPI](https://img.shields.io/pypi/v/mlx-forge)](https://pypi.org/project/mlx-forge/)
6
+ [![Python](https://img.shields.io/pypi/pyversions/mlx-forge)](https://pypi.org/project/mlx-forge/)
7
+ [![License](https://img.shields.io/badge/license-MIT-blue)](LICENSE)
8
+ [![Tests](https://img.shields.io/github/actions/workflow/status/moyuan5989/mlx-forge/test.yml?label=tests)](https://github.com/moyuan5989/mlx-forge/actions)
9
+
10
+ MLX Forge is a complete LLM fine-tuning toolkit that runs entirely on your Mac. Pick a model, upload your data, and start training — all from a browser-based UI. Supports LoRA, QLoRA, DPO, 18+ models, and 20+ curated datasets out of the box.
11
+
12
+ ```bash
13
+ pip install mlx-forge
14
+ mlx-forge studio
15
+ ```
16
+
17
+ <p align="center">
18
+ <img src="assets/studio-new-training.png" alt="MLX Forge Studio — New Training" width="800">
19
+ </p>
20
+
21
+ ## Why MLX Forge?
22
+
23
+ - **One command to start** — `pip install mlx-forge && mlx-forge studio`.
24
+ - **Browser-based Studio UI** — Guided training wizard, real-time loss charts, model library with memory estimates, interactive playground.
25
+ - **Runs on Apple Silicon** — Built on [MLX](https://github.com/ml-explore/mlx). Your data stays on your machine.
26
+ - **Production training features** — QLoRA (67% memory reduction), sequence packing (2-5x speedup), gradient checkpointing, DPO alignment, compiled training loop.
27
+
28
+ ## Quick Start
29
+
30
+ ### Studio UI (recommended)
31
+
32
+ ```bash
33
+ mlx-forge studio
34
+ # Opens at http://127.0.0.1:8741
35
+ ```
36
+
37
+ Pick a recipe, choose a model, upload your data, and start training — all from the browser.
38
+
39
+ ### CLI
40
+
41
+ ```bash
42
+ # Browse and download a dataset
43
+ mlx-forge data catalog
44
+ mlx-forge data download alpaca-cleaned --max-samples 5000
45
+
46
+ # Train
47
+ mlx-forge train --config train.yaml
48
+ ```
49
+
50
+ Models are downloaded from Hugging Face on first run and cached locally. All subsequent runs work offline.
51
+
52
+ ## Studio UI
53
+
54
+ <p align="center">
55
+ <img src="assets/studio-model-library.png" alt="MLX Forge Studio — Model Library" width="800">
56
+ </p>
57
+
58
+ - **New Training** — Guided wizard: pick a recipe (chat, instruction, DPO, writing style), choose a model, configure, and launch
59
+ - **Model Library** — Browse 18+ models with memory estimates for your hardware
60
+ - **Experiments** — Compare runs, view loss curves in real time
61
+ - **Datasets** — Manage your training data
62
+ - **Playground** — Chat with your fine-tuned models interactively
63
+
64
+ ## Supported Models
65
+
66
+ 18 curated models in the Studio library, all tested on Apple Silicon:
67
+
68
+ | Architecture | Models | Sizes |
69
+ |-------------|--------|-------|
70
+ | Qwen | Qwen 2.5, Qwen 3, Qwen 3.5 | 0.5B - 8B |
71
+ | Gemma | Gemma 2, Gemma 3 | 1B - 9B |
72
+ | Llama | Llama 3.1 | 8B |
73
+ | Phi | Phi-3 Mini, Phi-4 Mini | 3.8B |
74
+ | DeepSeek | DeepSeek-R1-Distill (Qwen-based) | 1.5B - 7B |
75
+ | Mistral | Mistral (uses Llama architecture) | 7B |
76
+
77
+ Any HF model using a supported architecture will work — the table above shows the curated models with pre-computed memory estimates in Studio.
78
+
79
+ ## Features
80
+
81
+ **Training**
82
+ - LoRA and QLoRA (4-bit) fine-tuning with 67% memory reduction
83
+ - DPO (Direct Preference Optimization) for alignment
84
+ - Sequence packing for 2-5x speedup on short sequences
85
+ - Gradient checkpointing for 40-60% memory savings
86
+ - Compiled training loop with gradient accumulation
87
+ - Cosine, linear, step, and exponential LR schedules with warmup
88
+ - Resume from any checkpoint
89
+
90
+ **Data**
91
+ - 20+ curated datasets across 7 categories (general, code, math, conversation, reasoning, safety, domain)
92
+ - Auto-detection of chat, completions, text, and preference formats
93
+ - Multi-dataset mixing with weighted sampling
94
+ - Data validation with train/val overlap detection
95
+
96
+ ## CLI Reference
97
+
98
+ | Command | Description |
99
+ |---------|-------------|
100
+ | `mlx-forge studio` | Launch the Studio UI |
101
+ | `mlx-forge train --config FILE` | Run LoRA/QLoRA/DPO training |
102
+ | `mlx-forge generate --model MODEL` | Generate text or interactive chat |
103
+ | `mlx-forge prepare --data FILE --model MODEL` | Pre-tokenize a dataset |
104
+ | `mlx-forge data catalog` | Browse 20+ curated datasets |
105
+ | `mlx-forge data download DATASET` | Download a dataset from the catalog |
106
+ | `mlx-forge data import FILE --name NAME` | Import a local JSONL file |
107
+ | `mlx-forge data validate FILE` | Validate JSONL data |
108
+ | `mlx-forge data inspect NAME` | Preview dataset samples |
109
+ | `mlx-forge data stats NAME` | Show dataset statistics |
110
+
111
+ ## Configuration
112
+
113
+ ```yaml
114
+ schema_version: 1
115
+
116
+ model:
117
+ path: "Qwen/Qwen3-0.6B" # HF model ID or local path
118
+ quantization: # Optional: QLoRA (67% memory savings)
119
+ bits: 4
120
+ group_size: 64
121
+
122
+ adapter:
123
+ preset: "attention-qv" # attention-qv | attention-all | mlp | all-linear
124
+ rank: 16
125
+ scale: 32.0
126
+
127
+ data:
128
+ train: "./train.jsonl"
129
+ valid: "./val.jsonl"
130
+ packing: false # Sequence packing (2-5x speedup)
131
+ max_seq_length: 2048
132
+
133
+ training:
134
+ optimizer: adamw # adam | adamw | sgd | adafactor
135
+ learning_rate: 1.0e-5
136
+ num_iters: 1000
137
+ batch_size: 4
138
+ gradient_checkpointing: false # 40-60% memory savings
139
+ # training_type: dpo # For DPO training
140
+ # dpo_beta: 0.1
141
+
142
+ runtime:
143
+ seed: 42
144
+ ```
145
+
146
+ ## Data Formats
147
+
148
+ MLX Forge auto-detects four JSONL formats:
149
+
150
+ **Chat** — Multi-turn conversations (loss on assistant turns only):
151
+ ```json
152
+ {"messages": [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi!"}]}
153
+ ```
154
+
155
+ **Completions** — Prompt-completion pairs:
156
+ ```json
157
+ {"prompt": "Translate to French: Hello", "completion": "Bonjour"}
158
+ ```
159
+
160
+ **Text** — Raw text for continued pretraining:
161
+ ```json
162
+ {"text": "The quick brown fox jumps over the lazy dog."}
163
+ ```
164
+
165
+ **Preference** — For DPO alignment training:
166
+ ```json
167
+ {"chosen": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "good"}], "rejected": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "bad"}]}
168
+ ```
169
+
170
+ ## Library API
171
+
172
+ All CLI commands are backed by Python functions:
173
+
174
+ ```python
175
+ from mlx_forge import prepare, train
176
+ from mlx_forge.config import TrainingConfig
177
+
178
+ # Train from a config file
179
+ config = TrainingConfig.from_yaml("train.yaml")
180
+ result = train(config=config)
181
+ print(f"Best val loss: {result.best_val_loss:.4f}")
182
+ ```
183
+
184
+ ```python
185
+ from mlx_forge import generate
186
+
187
+ # Generate text with a fine-tuned adapter
188
+ generate(
189
+ model="Qwen/Qwen3-0.6B",
190
+ adapter="~/.mlxforge/runs/my-run/checkpoints/best",
191
+ prompt="Explain quantum computing in simple terms.",
192
+ )
193
+ ```
194
+
195
+ ## Contributing
196
+
197
+ See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, coding standards, and how to submit changes.
198
+
199
+ ## License
200
+
201
+ [MIT](LICENSE)