Ex0bit commited on
Commit
73dc7d2
·
1 Parent(s): 6b4c197

Clarify README: MLX is the training engine, ANE is experimental

Browse files

- Remove ANE bridge build from Setup (not required)
- Make MLX self-test the primary Quick Validation step
- Mark all ANE files as [Experimental] in project structure
- Add "Note on ANE Code" section explaining ANE is not used for training
- Architecture section now explicitly states pure MLX autograd

Files changed (1) hide show
  1. README.md +23 -24
README.md CHANGED
@@ -85,34 +85,33 @@ A system for just-in-time (JIT) LoRA training that modifies a running language m
85
 
86
  ## Architecture
87
 
 
 
88
  ```
89
  User → React Frontend → Express Proxy → Neural Daemon (FastAPI, :8766)
90
 
91
- MLX Inference + LoRA Adapter
92
 
93
  SSE Token Stream → Frontend → TTS
94
 
95
- [After response] Background LoRA Training
96
 
97
- Updated adapter for next query
98
  ```
99
 
100
  ## Project Structure
101
 
102
  ```
103
  ├── src/
104
- │ ├── mlx_lora_trainer.py # Core training engine — LoRALinear, autograd, early stopping
105
- │ ├── neural_daemon.py # FastAPI daemon — inference, training orchestration, SSE
106
  │ ├── neural_config.py # Hyperparameter configuration
107
  │ ├── neural_data.py # Training data manager — rolling + replay buffers
108
- │ ├── ane_bridge_py.py # Python ctypes wrapper for ANE bridge
109
- │ ├── ane_lora_trainer.py # ANE training engine (requires ANE bridge)
110
- │ ├── ane_mil_lora.py # ANE kernel generators for LoRA forward/backward
111
  │ ├── export_to_lms.py # GGUF export for LM Studio
112
- ── bridge/ # ANE C bridge (from github.com/maderix/ANE, MIT)
113
- ├── ane_bridge.h # C API header
114
- ├── ane_bridge.m # Objective-C implementation
115
- └── Makefile # Build: `make` libane_bridge.dylib
116
  ├── tests/
117
  │ ├── test_daemon_e2e.py # Experiment 1 — 4 fictional facts
118
  │ ├── test_deep_e2e.py # Experiment 2 — 41 facts, 10 domains, 70 test cases
@@ -145,21 +144,13 @@ User → React Frontend → Express Proxy → Neural Daemon (FastAPI, :8766)
145
  ```bash
146
  git clone https://github.com/eelbaz/jit-lora.git
147
  cd jit-lora
148
- pip install -r requirements.txt
149
-
150
- # Build the ANE bridge (requires Xcode Command Line Tools)
151
- cd src/bridge && make && cd ../..
152
  ```
153
 
154
- The ANE bridge (`src/bridge/`) provides direct access to Apple Neural Engine hardware via private APIs. It is based on [maderix/ANE](https://github.com/maderix/ANE) (MIT License). Requires macOS 15+ on Apple Silicon.
155
-
156
  ### Quick Validation
157
 
158
  ```bash
159
- # Verify ANE bridge works
160
- python3 src/ane_bridge_py.py
161
-
162
- # Verify MLX training engine
163
  python3 src/mlx_lora_trainer.py
164
  ```
165
 
@@ -174,11 +165,19 @@ curl -X POST http://localhost:8766/activate \
174
  -H "Content-Type: application/json" \
175
  -d '{"hf_repo":"Qwen/Qwen3.5-2B-Base"}'
176
 
177
- python3 tests/test_daemon_e2e.py # 4 facts, 20s
178
- python3 tests/test_deep_e2e.py # 41 facts, 121s
179
  python3 tests/test_statistical_e2e.py # 35+ facts, 3 trials, ~4 min
180
  ```
181
 
 
 
 
 
 
 
 
 
182
  ## Citation
183
 
184
  ```bibtex
 
85
 
86
  ## Architecture
87
 
88
+ The training engine is **pure MLX** — `nn.value_and_grad()` for real autograd, Adam optimizer, cosine LR with early stopping. LoRA adapters are injected in-place into the model, so `mlx_lm.stream_generate()` automatically uses the updated weights with no special handling.
89
+
90
  ```
91
  User → React Frontend → Express Proxy → Neural Daemon (FastAPI, :8766)
92
 
93
+ MLX Inference with in-place LoRA adapter
94
 
95
  SSE Token Stream → Frontend → TTS
96
 
97
+ [After response] MLX LoRA backprop (background)
98
 
99
+ Updated adapter weights for next query
100
  ```
101
 
102
  ## Project Structure
103
 
104
  ```
105
  ├── src/
106
+ │ ├── mlx_lora_trainer.py # Training engine — LoRALinear, nn.value_and_grad, Adam, early stopping
107
+ │ ├── neural_daemon.py # FastAPI daemon — inference, training orchestration, SSE streaming
108
  │ ├── neural_config.py # Hyperparameter configuration
109
  │ ├── neural_data.py # Training data manager — rolling + replay buffers
 
 
 
110
  │ ├── export_to_lms.py # GGUF export for LM Studio
111
+ ── ane_bridge_py.py # [Experimental] Python ctypes wrapper for ANE bridge
112
+ ├── ane_lora_trainer.py # [Experimental] ANE training engine (not used — see note below)
113
+ ├── ane_mil_lora.py # [Experimental] ANE kernel generators for LoRA forward/backward
114
+ └── bridge/ # [Experimental] ANE C bridge (from github.com/maderix/ANE, MIT)
115
  ├── tests/
116
  │ ├── test_daemon_e2e.py # Experiment 1 — 4 fictional facts
117
  │ ├── test_deep_e2e.py # Experiment 2 — 41 facts, 10 domains, 70 test cases
 
144
  ```bash
145
  git clone https://github.com/eelbaz/jit-lora.git
146
  cd jit-lora
147
+ pip install mlx mlx-lm fastapi uvicorn requests numpy
 
 
 
148
  ```
149
 
 
 
150
  ### Quick Validation
151
 
152
  ```bash
153
+ # Verify MLX training engine (downloads Qwen2.5-0.5B, trains 5 steps, ~30s)
 
 
 
154
  python3 src/mlx_lora_trainer.py
155
  ```
156
 
 
165
  -H "Content-Type: application/json" \
166
  -d '{"hf_repo":"Qwen/Qwen3.5-2B-Base"}'
167
 
168
+ python3 tests/test_daemon_e2e.py # 4 facts, ~20s
169
+ python3 tests/test_deep_e2e.py # 41 facts, ~121s
170
  python3 tests/test_statistical_e2e.py # 35+ facts, 3 trials, ~4 min
171
  ```
172
 
173
+ ## Note on ANE Code
174
+
175
+ The `ane_*.py` files and `bridge/` directory are **experimental and not used for training**. The initial approach attempted to run LoRA kernels directly on Apple's Neural Engine via the private `AppleNeuralEngine.framework`. While the forward kernels compile and run, ANE produces IOSurface-backed tensors that are opaque to any autograd system — making gradient-based training impossible through ANE alone.
176
+
177
+ All training in this project uses **MLX autograd on GPU**. The ANE code remains in the repo for a potential future hybrid inference path (see Section 8.2 of the paper), where ANE could accelerate LoRA forward passes during multi-agent inference while the GPU handles the base model. This path is speculative and has not been benchmarked.
178
+
179
+ If you're interested in ANE internals, the bridge is based on [maderix/ANE](https://github.com/maderix/ANE) (MIT License) and requires macOS 15+ on Apple Silicon. Build with `cd src/bridge && make`. But this is **not required** to run any of the experiments or use the training system.
180
+
181
  ## Citation
182
 
183
  ```bibtex