Project-Chronos 0.1.1__py3-none-any.whl
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.
- chronos/__init__.py +16 -0
- chronos/_backend_legacy.py +64 -0
- chronos/app.py +122 -0
- chronos/backend/__init__.py +40 -0
- chronos/backend/dispatcher.py +347 -0
- chronos/backend/ext/__init__.py +2 -0
- chronos/backend/ext/opencl.py +22 -0
- chronos/cli.py +62 -0
- chronos/data/__init__.py +0 -0
- chronos/data/flexible_dataset.py +309 -0
- chronos/deps.py +72 -0
- chronos/eval/__init__.py +2 -0
- chronos/eval/benchmark.py +159 -0
- chronos/eval/io_profiler.py +232 -0
- chronos/io/__init__.py +13 -0
- chronos/io/async_prefetcher.py +220 -0
- chronos/io/cluster_layout.py +317 -0
- chronos/io/expert_store.py +559 -0
- chronos/io/io_simulator.py +40 -0
- chronos/io/storage.py +237 -0
- chronos/mlx/__init__.py +32 -0
- chronos/mlx/attention.py +230 -0
- chronos/mlx/expert_store.py +155 -0
- chronos/mlx/inference.py +159 -0
- chronos/mlx/model.py +181 -0
- chronos/mlx/moe.py +102 -0
- chronos/model/__init__.py +6 -0
- chronos/model/config.py +54 -0
- chronos/model/hf_io.py +149 -0
- chronos/model/hybrid_attention.py +237 -0
- chronos/model/lookahead_router.py +55 -0
- chronos/model/model_chronos.py +271 -0
- chronos/model/moe_chronos.py +189 -0
- chronos/model/temporal_loss.py +124 -0
- chronos/router/__init__.py +12 -0
- chronos/router/expert_predictor.py +163 -0
- chronos/router/intent_classifier.py +193 -0
- chronos/router/prefill_scheduler.py +248 -0
- chronos/runtime/__init__.py +2 -0
- chronos/runtime/cache_manager.py +117 -0
- chronos/runtime/inference_engine.py +295 -0
- chronos/runtime/metrics.py +75 -0
- chronos/serving/__init__.py +13 -0
- chronos/serving/vllm_adapter.py +85 -0
- chronos/trainer/__init__.py +6 -0
- chronos/trainer/chronos_trainer.py +151 -0
- chronos/trainer/distill_trainer.py +193 -0
- chronos/trainer/dpo_trainer.py +190 -0
- chronos/trainer/grpo_trainer.py +288 -0
- chronos/trainer/loss_mixin.py +165 -0
- chronos/trainer/optim_utils.py +101 -0
- chronos/trainer/orpo_trainer.py +189 -0
- chronos/trainer/reward.py +130 -0
- chronos/trainer/sft_trainer.py +189 -0
- chronos/tuning/__init__.py +1 -0
- chronos/tuning/_base_tuner.py +284 -0
- chronos/tuning/chronos_auto_tuner.py +203 -0
- project_chronos-0.1.1.dist-info/METADATA +670 -0
- project_chronos-0.1.1.dist-info/RECORD +83 -0
- project_chronos-0.1.1.dist-info/WHEEL +5 -0
- project_chronos-0.1.1.dist-info/entry_points.txt +7 -0
- project_chronos-0.1.1.dist-info/licenses/LICENSE +177 -0
- project_chronos-0.1.1.dist-info/top_level.txt +8 -0
- train_chronos.py +133 -0
- train_chronos_distill.py +95 -0
- train_chronos_dpo.py +82 -0
- train_chronos_grpo.py +88 -0
- train_chronos_orpo.py +81 -0
- train_chronos_sft.py +85 -0
- ui/__init__.py +3 -0
- ui/estimator.py +204 -0
- ui/gradio_compat.py +72 -0
- ui/i18n.py +540 -0
- ui/presets.py +194 -0
- ui/tabs/__init__.py +3 -0
- ui/tabs/autotune_tab.py +210 -0
- ui/tabs/benchmark_tab.py +196 -0
- ui/tabs/config_tab.py +318 -0
- ui/tabs/data_tab.py +148 -0
- ui/tabs/inference_tab.py +496 -0
- ui/tabs/iomon_tab.py +37 -0
- ui/tabs/pipeline_tab.py +134 -0
- ui/tabs/train_tab.py +1688 -0
|
@@ -0,0 +1,670 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: Project_Chronos
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: On-device low-latency lookahead dual-layer MoE inference — Lookahead Routing + Async Prefetch + Hybrid Attention
|
|
5
|
+
License-Expression: Apache-2.0
|
|
6
|
+
Keywords: llm,moe,inference,lookahead-routing,hybrid-attention,edge-ai
|
|
7
|
+
Classifier: Development Status :: 3 - Alpha
|
|
8
|
+
Classifier: Intended Audience :: Science/Research
|
|
9
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
10
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
11
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
12
|
+
Requires-Python: >=3.10
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
License-File: LICENSE
|
|
15
|
+
Requires-Dist: torch>=2.4.0
|
|
16
|
+
Requires-Dist: transformers>=4.40.0
|
|
17
|
+
Requires-Dist: datasets>=2.18.0
|
|
18
|
+
Requires-Dist: safetensors>=0.4.0
|
|
19
|
+
Requires-Dist: optuna>=3.6.0
|
|
20
|
+
Requires-Dist: numpy>=1.24.0
|
|
21
|
+
Requires-Dist: psutil>=5.9.0
|
|
22
|
+
Requires-Dist: tqdm>=4.66.0
|
|
23
|
+
Requires-Dist: gradio>=4.0.0
|
|
24
|
+
Provides-Extra: cluster
|
|
25
|
+
Requires-Dist: python-louvain>=0.16; extra == "cluster"
|
|
26
|
+
Requires-Dist: networkx>=3.0; extra == "cluster"
|
|
27
|
+
Provides-Extra: mlx
|
|
28
|
+
Requires-Dist: mlx>=0.12.0; extra == "mlx"
|
|
29
|
+
Provides-Extra: dev
|
|
30
|
+
Requires-Dist: pytest>=7.0; extra == "dev"
|
|
31
|
+
Requires-Dist: black>=24.0; extra == "dev"
|
|
32
|
+
Requires-Dist: ruff>=0.4.0; extra == "dev"
|
|
33
|
+
Dynamic: license-file
|
|
34
|
+
Dynamic: requires-python
|
|
35
|
+
|
|
36
|
+
# Project Chronos (Experimental)
|
|
37
|
+
|
|
38
|
+
**A storage-aware MoE stack built for SSD+DRAM hybrid inference, with a full six-stage training pipeline.**
|
|
39
|
+
|
|
40
|
+
[](https://pypi.org/project/project-chronos/)
|
|
41
|
+
[](LICENSE)
|
|
42
|
+
[](https://python.org)
|
|
43
|
+
|
|
44
|
+
[中文文档](README_zh.md)
|
|
45
|
+
|
|
46
|
+
---
|
|
47
|
+
|
|
48
|
+
## Why reactive MoE offload breaks down
|
|
49
|
+
|
|
50
|
+
Mainstream MoE models such as Mixtral, DeepSeek-MoE, and Qwen-MoE make routing decisions **at decode time, token by token**. The model discovers which experts it needs only when the next token is already on the critical path. If those experts are not resident in VRAM, generation stalls while weights are moved from SSD to RAM to VRAM.
|
|
51
|
+
|
|
52
|
+
That is not a tuning issue. It is an architectural mismatch.
|
|
53
|
+
|
|
54
|
+
Most offload runtimes assume the model was originally designed for full VRAM residency, then try to patch storage pressure afterward. On consumer hardware, that usually means the decode loop ends up paying the IO bill over and over again.
|
|
55
|
+
|
|
56
|
+
---
|
|
57
|
+
|
|
58
|
+
## Chronos in one line
|
|
59
|
+
|
|
60
|
+
**Move IO into prefill. Move synchronization down to the expert-event level.**
|
|
61
|
+
|
|
62
|
+
```mermaid
|
|
63
|
+
sequenceDiagram
|
|
64
|
+
autonumber
|
|
65
|
+
participant Prompt as Prompt tokens
|
|
66
|
+
participant Intent as IntentClassifier
|
|
67
|
+
participant SSD as Clustered SSD cache
|
|
68
|
+
participant RAM as Pinned RAM
|
|
69
|
+
participant H2D as H2D stream
|
|
70
|
+
participant GPU as Decode compute stream
|
|
71
|
+
participant LR as LookaheadRouter
|
|
72
|
+
|
|
73
|
+
rect rgb(236, 248, 255)
|
|
74
|
+
Prompt->>Intent: prefill reads the whole prompt
|
|
75
|
+
Intent->>SSD: predict generation-level expert set
|
|
76
|
+
SSD->>RAM: mmap cluster-packed .ctsr files
|
|
77
|
+
RAM->>H2D: enqueue hot experts before first decode token
|
|
78
|
+
H2D-->>GPU: per-expert ready events
|
|
79
|
+
end
|
|
80
|
+
|
|
81
|
+
loop every decode token
|
|
82
|
+
GPU->>LR: block-0 hidden state
|
|
83
|
+
LR->>H2D: prefetch experts for t+1..t+K
|
|
84
|
+
GPU->>GPU: forward token t with resident experts
|
|
85
|
+
GPU->>H2D: wait only on the expert events it needs
|
|
86
|
+
end
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
Two things matter here:
|
|
90
|
+
|
|
91
|
+
1. **Prefill-time loading**: `PrefillScheduler` and `IntentClassifier` bulk-load predicted experts before the first decode token.
|
|
92
|
+
2. **Per-expert event sync (M3)**: `promote_to_vram(blocking=False)` records a `torch.cuda.Event` on the `_h2d_stream`, and the compute stream waits only on the experts it actually needs. No more `stream.synchronize()` for the whole system.
|
|
93
|
+
|
|
94
|
+
Under a simulated 30 ms SSD delay, the new path keeps **35 ms+ per token** of pipeline slack compared with the older blocking path.
|
|
95
|
+
|
|
96
|
+
---
|
|
97
|
+
|
|
98
|
+
## Three-tier storage (M1: cluster-aware safetensors)
|
|
99
|
+
|
|
100
|
+
```mermaid
|
|
101
|
+
flowchart TB
|
|
102
|
+
subgraph GPU["VRAM / Metal unified memory"]
|
|
103
|
+
Dense["Dense layers"]
|
|
104
|
+
Shared["Shared experts<br/>always resident"]
|
|
105
|
+
Hot["Predicted hot experts<br/>promoted before use"]
|
|
106
|
+
Events["Per-expert CUDA events<br/>no global stream sync"]
|
|
107
|
+
end
|
|
108
|
+
|
|
109
|
+
subgraph RAM["Pinned RAM staging tier"]
|
|
110
|
+
Buffer["Prefetched expert tensors"]
|
|
111
|
+
LRU["RAM LRU cache"]
|
|
112
|
+
end
|
|
113
|
+
|
|
114
|
+
subgraph Disk["NVMe SSD"]
|
|
115
|
+
C1["cluster_000.ctsr"]
|
|
116
|
+
C2["cluster_001.ctsr"]
|
|
117
|
+
Manifest["cluster_manifest.json<br/>expert -> cluster map"]
|
|
118
|
+
end
|
|
119
|
+
|
|
120
|
+
Manifest --> C1
|
|
121
|
+
Manifest --> C2
|
|
122
|
+
C1 -->|"sequential mmap read"| Buffer
|
|
123
|
+
C2 -->|"sequential mmap read"| Buffer
|
|
124
|
+
Buffer --> LRU
|
|
125
|
+
LRU -->|"non-blocking H2D"| Hot
|
|
126
|
+
Hot --> Events
|
|
127
|
+
Shared --> Events
|
|
128
|
+
Dense --> Events
|
|
129
|
+
```
|
|
130
|
+
|
|
131
|
+
`cluster_manifest.json` and `.ctsr` files are produced offline by Louvain clustering. At runtime, Chronos uses `safetensors.safe_open(...).mmap` to bring an entire expert cluster into RAM in one shot, turning random reads into mostly sequential reads.
|
|
132
|
+
|
|
133
|
+
### Cache misses degrade gracefully instead of stalling
|
|
134
|
+
|
|
135
|
+
Even the worst case does not hard-stop generation:
|
|
136
|
+
|
|
137
|
+
```python
|
|
138
|
+
# Pure tensor math: no Python branch, no graph break under torch.compile
|
|
139
|
+
output = avail[i] * expert_output + (1.0 - avail[i]) * shared_expert_output
|
|
140
|
+
```
|
|
141
|
+
|
|
142
|
+
The shared expert is always resident, so generation continues while the missing expert finishes loading in the background. Quality degrades smoothly and recovers automatically once the expert becomes available.
|
|
143
|
+
|
|
144
|
+
---
|
|
145
|
+
|
|
146
|
+
## Dual-layer routing with supervised lookahead (M2)
|
|
147
|
+
|
|
148
|
+
```mermaid
|
|
149
|
+
flowchart LR
|
|
150
|
+
Prompt["Prompt / prefill context"] --> IC["IntentClassifier<br/>generation-level expert prior"]
|
|
151
|
+
IC --> Budget["ExpertPredictor<br/>budgeted hot set"]
|
|
152
|
+
Budget --> Prefill["PrefillScheduler<br/>bulk preload before decode"]
|
|
153
|
+
|
|
154
|
+
Token["Decode token t"] --> Block0["Transformer block 0"]
|
|
155
|
+
Block0 --> LR["LookaheadRouter<br/>Q_t^(1..K)"]
|
|
156
|
+
LR --> Future["Future expert predictions<br/>t+1 ... t+K"]
|
|
157
|
+
Future --> Async["AsyncPrefetcher<br/>prefetch queue"]
|
|
158
|
+
|
|
159
|
+
Prefill --> Cache["CacheManager / ExpertStore"]
|
|
160
|
+
Async --> Cache
|
|
161
|
+
Cache --> MoE["ChronosMOEFeedForward<br/>resident expert or shared fallback"]
|
|
162
|
+
```
|
|
163
|
+
|
|
164
|
+
| | IntentClassifier (Layer 1) | LookaheadRouter (Layer 2) |
|
|
165
|
+
|---|---|---|
|
|
166
|
+
| **When it runs** | Once during prefill | Every token during decode |
|
|
167
|
+
| **Input** | Full prompt (up to 512 tokens) | Hidden state after Block 0 |
|
|
168
|
+
| **Output** | Expert set for the full generation | Expert IDs for t+1, t+2 |
|
|
169
|
+
| **Training target** | Supervised from real activation logs | **`L_lookahead` supervised by real router decisions at t+k** |
|
|
170
|
+
| **Parameter count** | ~10-15M | ~2M |
|
|
171
|
+
|
|
172
|
+
Before M2, the lookahead head was just a head with no real supervision. M2 adds a proper soft-target objective:
|
|
173
|
+
|
|
174
|
+
```math
|
|
175
|
+
L_{\mathrm{lookahead}}
|
|
176
|
+
= \frac{1}{|K_{\mathrm{valid}}|}
|
|
177
|
+
\sum_{k \in K_{\mathrm{valid}}}
|
|
178
|
+
\mathbb{E}_{b,t}
|
|
179
|
+
\left[
|
|
180
|
+
- \sum_{e=1}^{E}
|
|
181
|
+
\mathrm{sg}(P_{b,t+k,e}) \log Q_{b,t,e}^{(k)}
|
|
182
|
+
\right]
|
|
183
|
+
```
|
|
184
|
+
|
|
185
|
+
That turns the lookahead router into an actual predictor of future routing, instead of a best-effort heuristic.
|
|
186
|
+
|
|
187
|
+
---
|
|
188
|
+
|
|
189
|
+
## Full training stack: Stage 1 -> Stage 6
|
|
190
|
+
|
|
191
|
+
Each stage has its own entry script, and every one of them inherits the shared Chronos loss mixer: balance loss, temporal locality loss, lookahead loss, and, for alignment stages, a router KL anchor that keeps RL/DPO updates from destroying cache locality.
|
|
192
|
+
|
|
193
|
+
```mermaid
|
|
194
|
+
flowchart LR
|
|
195
|
+
P1["Stage 1<br/>Pretrain<br/>CE + Chronos mix"] --> P2["Stage 2<br/>SFT<br/>assistant-token CE + mix"]
|
|
196
|
+
P2 --> P3["Stage 3<br/>DPO<br/>preference loss + anchor"]
|
|
197
|
+
P2 --> P4["Stage 4<br/>ORPO<br/>NLL + odds-ratio + anchor"]
|
|
198
|
+
P4 --> P5["Stage 5<br/>GRPO<br/>rollout reward + KL + anchor"]
|
|
199
|
+
P5 --> P6["Stage 6<br/>Distill<br/>KD + CE + anchor"]
|
|
200
|
+
|
|
201
|
+
subgraph Shared["Shared Chronos training terms"]
|
|
202
|
+
B["raw load-balance aux"]
|
|
203
|
+
T["temporal locality"]
|
|
204
|
+
LA["supervised lookahead"]
|
|
205
|
+
A["router KL anchor<br/>alignment stages"]
|
|
206
|
+
end
|
|
207
|
+
|
|
208
|
+
Shared -. applied to .-> P1
|
|
209
|
+
Shared -. applied to .-> P2
|
|
210
|
+
Shared -. applied to .-> P3
|
|
211
|
+
Shared -. applied to .-> P4
|
|
212
|
+
Shared -. applied to .-> P5
|
|
213
|
+
Shared -. applied to .-> P6
|
|
214
|
+
```
|
|
215
|
+
|
|
216
|
+
| Stage | Script | Core objective | Router KL anchor (default lambda) |
|
|
217
|
+
|---|---|---|---|
|
|
218
|
+
| 1 Pretrain | `train_chronos.py` | CE + balance + temporal + lookahead | 0.0 (off) |
|
|
219
|
+
| 2 SFT | `train_chronos_sft.py` | SFT loss + shared Chronos mix | 0.01 (weak) |
|
|
220
|
+
| 3 DPO | `train_chronos_dpo.py` | DPO `log-sigma(beta * logits)` + mix | 0.10 (strong) |
|
|
221
|
+
| 4 ORPO | `train_chronos_orpo.py` | NLL + lambda * ORPO term | 0.10 |
|
|
222
|
+
| 5 GRPO | `train_chronos_grpo.py` | `PG * A - beta * KL` with `ToyReward` or pluggable `LMRewardModel` | 0.10 |
|
|
223
|
+
| 6 Distill | `train_chronos_distill.py` | `alpha * T^2 * KL(student || teacher) + (1 - alpha) * CE` | 0.05 |
|
|
224
|
+
|
|
225
|
+
The full six-stage comparison harness lives in `tools/compare_minimind_chronos_v3.py`.
|
|
226
|
+
|
|
227
|
+
---
|
|
228
|
+
|
|
229
|
+
## Backend dispatch (M5)
|
|
230
|
+
|
|
231
|
+
```mermaid
|
|
232
|
+
flowchart TD
|
|
233
|
+
Request["User request<br/>train / inference / WebUI / CLI"] --> Dispatcher["BackendDispatcher"]
|
|
234
|
+
Dispatcher --> Probe["Capability probes"]
|
|
235
|
+
Probe --> CUDA["CUDA<br/>training + inference"]
|
|
236
|
+
Probe --> MPS["MPS<br/>training + inference"]
|
|
237
|
+
Probe --> MLX["MLX<br/>Apple Silicon native path"]
|
|
238
|
+
Probe --> CPU["CPU<br/>portable fallback"]
|
|
239
|
+
Probe --> EXT["Extension hooks<br/>Vulkan / OpenCL"]
|
|
240
|
+
Dispatcher --> Choice["Priority + availability decision"]
|
|
241
|
+
Choice --> Runtime["Chronos runtime / trainer"]
|
|
242
|
+
```
|
|
243
|
+
|
|
244
|
+
```python
|
|
245
|
+
from chronos.backend import BackendDispatcher
|
|
246
|
+
|
|
247
|
+
d = BackendDispatcher()
|
|
248
|
+
d.available() # ['mlx', 'mps', 'cpu'] on Apple Silicon
|
|
249
|
+
# ['cuda', 'cpu'] on NVIDIA hosts
|
|
250
|
+
d.select() # choose the best available backend automatically
|
|
251
|
+
d.describe() # human-readable capability summary
|
|
252
|
+
```
|
|
253
|
+
|
|
254
|
+
- **First-class backends for training and inference**: `cpu`, `mps`, `cuda`, `mlx`
|
|
255
|
+
- **Inference-only / experimental**: `vulkan` when PyTorch was custom-built with `USE_VULKAN=ON`
|
|
256
|
+
- **Third-party extension hook**: `opencl`, via `chronos/backend/ext/opencl.py:PROBE()`
|
|
257
|
+
|
|
258
|
+
Honest note: upstream PyTorch does not ship a real OpenCL backend, and Vulkan support is still niche. Chronos provides a dispatcher seam so external integrations can plug in cleanly without touching core code.
|
|
259
|
+
|
|
260
|
+
---
|
|
261
|
+
|
|
262
|
+
## Hugging Face and vLLM compatibility (M5)
|
|
263
|
+
|
|
264
|
+
- `ChronosForCausalLM` subclasses `PreTrainedModel` and registers `AutoConfig` and `AutoModelForCausalLM`, so loading does **not** require `trust_remote_code`:
|
|
265
|
+
|
|
266
|
+
```python
|
|
267
|
+
from transformers import AutoModelForCausalLM
|
|
268
|
+
|
|
269
|
+
model = AutoModelForCausalLM.from_pretrained("./out_dir")
|
|
270
|
+
```
|
|
271
|
+
|
|
272
|
+
- `chronos.model.hf_io.save_chronos_pretrained` and `load_chronos_pretrained` emit standard `model.safetensors` + `config.json`, while also carrying `cluster_manifest.json` and `.ctsr` files for expert-cache layout. Roundtrip logit drift is `0.00e+00`.
|
|
273
|
+
|
|
274
|
+
- `chronos.serving.register_chronos_with_vllm()` registers Chronos with the vLLM `ModelRegistry` when vLLM is installed. If vLLM is absent, it prints an install hint and exits cleanly. Worker-side mask injection is documented in [docs/vllm_integration.md](docs/vllm_integration.md).
|
|
275
|
+
|
|
276
|
+
---
|
|
277
|
+
|
|
278
|
+
## Compared with existing offload stacks
|
|
279
|
+
|
|
280
|
+
| Feature | llama.cpp offload | vLLM offload | **Project Chronos** |
|
|
281
|
+
|---|---|---|---|
|
|
282
|
+
| Expert prediction | None | None | **Predictive (`IntentCLF` + `LookaheadRouter`)** |
|
|
283
|
+
| Lookahead training | n/a | n/a | **Supervised `L_lookahead` (M2)** |
|
|
284
|
+
| IO timing | During decode, blocking | During decode, blocking | **During prefill, async** |
|
|
285
|
+
| Decode pipeline | Synchronous | Synchronous | **Dual-stream + per-expert events (M3)** |
|
|
286
|
+
| Cache miss behavior | Hard stall | Hard stall | **Soft gating, zero hard stall** |
|
|
287
|
+
| Disk format | GGUF | safetensors | **Cluster-packed safetensors (`.ctsr`)** |
|
|
288
|
+
| Training integration | Post-hoc patch | Post-hoc patch | **Native six-stage stack + router KL anchor** |
|
|
289
|
+
| Backend dispatch | Compile-time fixed | CUDA only | **`cpu` / `mps` / `cuda` / `mlx` + extension hooks** |
|
|
290
|
+
| Apple Silicon support | Partial | No | **Full MLX backend** |
|
|
291
|
+
| Hugging Face compatibility | GGUF only | Yes | **Yes, with expert-cache metadata** |
|
|
292
|
+
| vLLM compatibility | n/a | Native | **Optional adapter** |
|
|
293
|
+
|
|
294
|
+
---
|
|
295
|
+
|
|
296
|
+
## Objective
|
|
297
|
+
|
|
298
|
+
```math
|
|
299
|
+
L_{\mathrm{total}} =
|
|
300
|
+
L_{\mathrm{base}}
|
|
301
|
+
+ \lambda_{\mathrm{bal}} L_{\mathrm{aux}}
|
|
302
|
+
+ \lambda_{\mathrm{tmp}} L_{\mathrm{temporal}}
|
|
303
|
+
+ \lambda_{\mathrm{LA}} L_{\mathrm{lookahead}}
|
|
304
|
+
+ \lambda_{\mathrm{anc}} L_{\mathrm{routerKL}}
|
|
305
|
+
```
|
|
306
|
+
|
|
307
|
+
```math
|
|
308
|
+
L_{\mathrm{aux}} = E \sum_{e=1}^{E} load_e \cdot \overline{p}_e
|
|
309
|
+
```
|
|
310
|
+
|
|
311
|
+
```math
|
|
312
|
+
L_{\mathrm{temporal}} =
|
|
313
|
+
\mathbb{E}_{b,t}
|
|
314
|
+
\left[
|
|
315
|
+
\left\| P_{b,t,:} - P_{b,t-1,:} \right\|_2^2
|
|
316
|
+
\right]
|
|
317
|
+
```
|
|
318
|
+
|
|
319
|
+
```math
|
|
320
|
+
L_{\mathrm{routerKL}} =
|
|
321
|
+
D_{\mathrm{KL}}
|
|
322
|
+
\left(
|
|
323
|
+
\pi_{\theta}^{\mathrm{router}}
|
|
324
|
+
\|
|
|
325
|
+
\pi_{\mathrm{ref}}^{\mathrm{router}}
|
|
326
|
+
\right)
|
|
327
|
+
```
|
|
328
|
+
|
|
329
|
+
- `L_base`: stage-specific objective (`CE`, `DPO`, `ORPO`, `GRPO`, or distillation).
|
|
330
|
+
- `L_aux`: the unscaled MoE load-balance auxiliary term; Chronos applies `lambda_bal` once in `chronos_loss_term`.
|
|
331
|
+
- `L_temporal`: encourages adjacent tokens to reuse similar expert distributions.
|
|
332
|
+
- `L_lookahead`: soft-target cross entropy from the real future router distribution to the lookahead prediction. `sg(...)` means stop-gradient.
|
|
333
|
+
- `L_routerKL`: keeps alignment-stage updates from destroying the routing layout captured at stage start.
|
|
334
|
+
|
|
335
|
+
All lambda terms are searchable with Optuna TPE, together with structural hyperparameters such as `hidden_size`, `num_experts`, and `kv_latent_dim`.
|
|
336
|
+
|
|
337
|
+
---
|
|
338
|
+
|
|
339
|
+
## Installation
|
|
340
|
+
|
|
341
|
+
```bash
|
|
342
|
+
pip install project-chronos
|
|
343
|
+
```
|
|
344
|
+
|
|
345
|
+
Or from source:
|
|
346
|
+
|
|
347
|
+
```bash
|
|
348
|
+
git clone https://github.com/FonaTech/Project_Chronos
|
|
349
|
+
cd Project_Chronos
|
|
350
|
+
pip install -e ".[dev]"
|
|
351
|
+
```
|
|
352
|
+
|
|
353
|
+
**MLX (Apple Silicon):**
|
|
354
|
+
|
|
355
|
+
```bash
|
|
356
|
+
pip install "project-chronos[mlx]"
|
|
357
|
+
```
|
|
358
|
+
|
|
359
|
+
**vLLM serving (optional, Linux + CUDA only):**
|
|
360
|
+
|
|
361
|
+
```bash
|
|
362
|
+
pip install vllm
|
|
363
|
+
```
|
|
364
|
+
|
|
365
|
+
> **minimind dependency**: Project Chronos uses [minimind](https://github.com/jingyaogong/minimind) as its MoE kernel.
|
|
366
|
+
> If it is not found locally, Chronos clones it automatically into `~/.cache/chronos/minimind-master/` on first import.
|
|
367
|
+
> minimind is licensed under **Apache-2.0**. See [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md) for attribution details.
|
|
368
|
+
|
|
369
|
+
**Requirements**: Python 3.10+, PyTorch 2.4+
|
|
370
|
+
|
|
371
|
+
---
|
|
372
|
+
|
|
373
|
+
## Quick start
|
|
374
|
+
|
|
375
|
+
### Web UI (M6: 7 tabs, 4 languages)
|
|
376
|
+
|
|
377
|
+
```bash
|
|
378
|
+
chronos-ui
|
|
379
|
+
# or
|
|
380
|
+
python chronos_app.py
|
|
381
|
+
```
|
|
382
|
+
|
|
383
|
+
Tabs included:
|
|
384
|
+
|
|
385
|
+
- `Config` with a live parameter / memory estimator merged in from the old Designer
|
|
386
|
+
- `Train` with its own `data_path`
|
|
387
|
+
- `6-Stage Pipeline` with per-stage dataset paths
|
|
388
|
+
- `Inference`
|
|
389
|
+
- `Benchmark` with Markdown table + bar plot
|
|
390
|
+
- `Auto-Tune` with persistent logs and one-click `Apply Best -> Config`
|
|
391
|
+
- `IO Monitor`
|
|
392
|
+
|
|
393
|
+
Built-in i18n: `zh-Hans`, `zh-Hant`, `en`, `ja`
|
|
394
|
+
|
|
395
|
+
### Stage 1: pretrain
|
|
396
|
+
|
|
397
|
+
```bash
|
|
398
|
+
python train_chronos.py \
|
|
399
|
+
--data_path ./tests/fixtures/tiny_pretrain.jsonl \
|
|
400
|
+
--hidden_size 256 --num_hidden_layers 4 --num_experts 4 \
|
|
401
|
+
--epochs 1 --device cpu --save_dir ./out
|
|
402
|
+
```
|
|
403
|
+
|
|
404
|
+
### Stage 2-5: alignment chain
|
|
405
|
+
|
|
406
|
+
```bash
|
|
407
|
+
python train_chronos_sft.py --data_path ./tests/fixtures/tiny_sft.jsonl --from_weight chronos --save_dir ./out --device cpu
|
|
408
|
+
python train_chronos_dpo.py --data_path ./tests/fixtures/tiny_dpo.jsonl --from_weight sft --save_dir ./out --device cpu
|
|
409
|
+
python train_chronos_orpo.py --data_path ./tests/fixtures/tiny_dpo.jsonl --from_weight sft --save_dir ./out --device cpu
|
|
410
|
+
python train_chronos_grpo.py --data_path ./tests/fixtures/tiny_grpo.jsonl --from_weight orpo --save_dir ./out --device cpu \
|
|
411
|
+
--reward toy # or lm:/path/to/reward-model
|
|
412
|
+
```
|
|
413
|
+
|
|
414
|
+
### Stage 6: distillation
|
|
415
|
+
|
|
416
|
+
```bash
|
|
417
|
+
python train_chronos_distill.py \
|
|
418
|
+
--data_path ./tests/fixtures/tiny_sft.jsonl \
|
|
419
|
+
--teacher_path ./out/sft_192_moe.pth \
|
|
420
|
+
--from_weight grpo --save_dir ./out --device cpu \
|
|
421
|
+
--alpha 0.7 --temperature 4.0
|
|
422
|
+
```
|
|
423
|
+
|
|
424
|
+
### End-to-end comparison (minimind vs Chronos)
|
|
425
|
+
|
|
426
|
+
```bash
|
|
427
|
+
python tools/compare_minimind_chronos_v3.py \
|
|
428
|
+
--pretrain_steps 150 --align_steps 30 --distill_steps 30 \
|
|
429
|
+
--simulated_ssd_ms 30 --device cpu \
|
|
430
|
+
--output results/compare_results_v3.json
|
|
431
|
+
```
|
|
432
|
+
|
|
433
|
+
This emits per-stage loss, HF roundtrip logit delta, tokens/sec, active-expert ratio, resident-expert bytes, M3 pipeline slack, and the backend inventory on the current host.
|
|
434
|
+
|
|
435
|
+
### Cluster-pack expert storage for sequential SSD reads
|
|
436
|
+
|
|
437
|
+
```python
|
|
438
|
+
from chronos.io.cluster_layout import (
|
|
439
|
+
collect_activation_log,
|
|
440
|
+
build_cooccurrence_matrix,
|
|
441
|
+
try_louvain_clustering,
|
|
442
|
+
repack_expert_weights_safetensors,
|
|
443
|
+
)
|
|
444
|
+
|
|
445
|
+
log = collect_activation_log(model, calib_loader, "cpu", max_batches=50)
|
|
446
|
+
clusters = try_louvain_clustering(build_cooccurrence_matrix(log, num_experts))
|
|
447
|
+
repack_expert_weights_safetensors(model, clusters, "./expert_cache_clustered")
|
|
448
|
+
```
|
|
449
|
+
|
|
450
|
+
### Search lambdas and structure hyperparameters automatically
|
|
451
|
+
|
|
452
|
+
```python
|
|
453
|
+
from chronos.tuning.chronos_auto_tuner import ChronosAutoTuner, ChronosSearchSpaceConfig
|
|
454
|
+
|
|
455
|
+
tuner = ChronosAutoTuner()
|
|
456
|
+
tuner.start(
|
|
457
|
+
model_id="./out/chronos_256_moe.pth",
|
|
458
|
+
dataset_path="./dataset/train.jsonl",
|
|
459
|
+
search_space=ChronosSearchSpaceConfig(
|
|
460
|
+
tune_lambda_balance=True,
|
|
461
|
+
tune_lambda_temporal=True,
|
|
462
|
+
tune_lambda_lookahead=True,
|
|
463
|
+
tune_lookahead_steps=True,
|
|
464
|
+
tune_hidden_size=True,
|
|
465
|
+
tune_num_experts=True,
|
|
466
|
+
tune_num_shared_experts=True,
|
|
467
|
+
tune_kv_latent_dim=True,
|
|
468
|
+
),
|
|
469
|
+
n_trials=20,
|
|
470
|
+
)
|
|
471
|
+
```
|
|
472
|
+
|
|
473
|
+
---
|
|
474
|
+
|
|
475
|
+
## Project layout
|
|
476
|
+
|
|
477
|
+
```text
|
|
478
|
+
Project_Chronos/
|
|
479
|
+
├── chronos/
|
|
480
|
+
│ ├── deps.py # Auto-download minimind if missing
|
|
481
|
+
│ ├── __init__.py # AutoConfig / AutoModelForCausalLM registration
|
|
482
|
+
│ ├── model/
|
|
483
|
+
│ │ ├── config.py # ChronosConfig
|
|
484
|
+
│ │ ├── hybrid_attention.py # MLAAttention + SlidingWindowAttention
|
|
485
|
+
│ │ ├── lookahead_router.py # Per-token lookahead predictor
|
|
486
|
+
│ │ ├── moe_chronos.py # ChronosMOEFeedForward + shared experts + soft gating
|
|
487
|
+
│ │ ├── model_chronos.py # ChronosForCausalLM
|
|
488
|
+
│ │ ├── temporal_loss.py # Temporal locality + lookahead losses
|
|
489
|
+
│ │ └── hf_io.py # save/load_chronos_pretrained + HF registration
|
|
490
|
+
│ ├── io/
|
|
491
|
+
│ │ ├── expert_store.py # Three-tier storage + per-expert events
|
|
492
|
+
│ │ ├── async_prefetcher.py # Async prefetch engine
|
|
493
|
+
│ │ ├── storage.py # ClusterStorage: .ctsr safetensors + manifest
|
|
494
|
+
│ │ ├── cluster_layout.py # Co-occurrence clustering + repacking
|
|
495
|
+
│ │ └── io_simulator.py # CHRONOS_SIM_SSD_MS test hook
|
|
496
|
+
│ ├── router/
|
|
497
|
+
│ │ ├── intent_classifier.py # Prompt-level expert predictor
|
|
498
|
+
│ │ ├── expert_predictor.py # IntentVector -> ExpertSet
|
|
499
|
+
│ │ └── prefill_scheduler.py # Prefill-time expert preloader
|
|
500
|
+
│ ├── mlx/
|
|
501
|
+
│ │ ├── attention.py / moe.py / model.py / expert_store.py / inference.py
|
|
502
|
+
│ ├── runtime/
|
|
503
|
+
│ │ ├── cache_manager.py # prefetch_for_next_step / ensure_resident
|
|
504
|
+
│ │ ├── inference_engine.py # End-to-end inference engine
|
|
505
|
+
│ │ └── metrics.py # MetricsBus for the IO Monitor
|
|
506
|
+
│ ├── trainer/
|
|
507
|
+
│ │ ├── loss_mixin.py # chronos_loss_term + router_kl_anchor
|
|
508
|
+
│ │ ├── chronos_trainer.py # Pretrain
|
|
509
|
+
│ │ ├── sft_trainer.py # Stage 2
|
|
510
|
+
│ │ ├── dpo_trainer.py # Stage 3
|
|
511
|
+
│ │ ├── orpo_trainer.py # Stage 4
|
|
512
|
+
│ │ ├── grpo_trainer.py # Stage 5
|
|
513
|
+
│ │ ├── distill_trainer.py # Stage 6
|
|
514
|
+
│ │ └── reward.py # ToyReward / LMRewardModel / build_reward_fn
|
|
515
|
+
│ ├── tuning/
|
|
516
|
+
│ │ └── chronos_auto_tuner.py # Optuna lambda + architecture search
|
|
517
|
+
│ ├── eval/
|
|
518
|
+
│ │ ├── io_profiler.py # Lookahead validation
|
|
519
|
+
│ │ └── benchmark.py # End-to-end benchmarking
|
|
520
|
+
│ ├── data/
|
|
521
|
+
│ │ └── flexible_dataset.py # Flexible JSONL dataset loader
|
|
522
|
+
│ ├── backend/
|
|
523
|
+
│ │ ├── __init__.py # BackendDispatcher (cpu/mps/cuda/mlx)
|
|
524
|
+
│ │ ├── dispatcher.py # Capability probing + priority logic
|
|
525
|
+
│ │ └── ext/opencl.py # Third-party OpenCL extension hook
|
|
526
|
+
│ ├── _backend_legacy.py # Backward-compatible older APIs
|
|
527
|
+
│ ├── serving/
|
|
528
|
+
│ │ ├── __init__.py
|
|
529
|
+
│ │ └── vllm_adapter.py # Optional vLLM registration
|
|
530
|
+
│ └── cli.py # Unified CLI
|
|
531
|
+
├── ui/ # Gradio Web UI (zh-Hans / zh-Hant / en / ja)
|
|
532
|
+
│ ├── i18n.py
|
|
533
|
+
│ ├── estimator.py # Live parameter / memory estimator
|
|
534
|
+
│ └── tabs/
|
|
535
|
+
│ ├── config_tab.py # Config + Designer merged together
|
|
536
|
+
│ ├── train_tab.py # Owns data_path
|
|
537
|
+
│ ├── pipeline_tab.py # Per-stage datasets across all 6 stages
|
|
538
|
+
│ ├── inference_tab.py
|
|
539
|
+
│ ├── benchmark_tab.py # Markdown table + gr.BarPlot
|
|
540
|
+
│ ├── autotune_tab.py # Persistent logs + Apply Best -> Config
|
|
541
|
+
│ └── iomon_tab.py # MetricsBus dashboard
|
|
542
|
+
├── chronos_app.py # Web UI entry
|
|
543
|
+
├── train_chronos.py # Stage 1 entry
|
|
544
|
+
├── train_chronos_sft.py # Stage 2 entry
|
|
545
|
+
├── train_chronos_dpo.py # Stage 3 entry
|
|
546
|
+
├── train_chronos_orpo.py # Stage 4 entry
|
|
547
|
+
├── train_chronos_grpo.py # Stage 5 entry
|
|
548
|
+
├── train_chronos_distill.py # Stage 6 entry
|
|
549
|
+
├── tools/
|
|
550
|
+
│ ├── compare_minimind_chronos.py
|
|
551
|
+
│ ├── compare_minimind_chronos_v2.py
|
|
552
|
+
│ └── compare_minimind_chronos_v3.py
|
|
553
|
+
├── tests/
|
|
554
|
+
│ ├── test_smoke.py
|
|
555
|
+
│ ├── test_smoke_cuda.py
|
|
556
|
+
│ └── fixtures/
|
|
557
|
+
├── docs/
|
|
558
|
+
│ └── vllm_integration.md
|
|
559
|
+
├── pyproject.toml
|
|
560
|
+
└── README.md / README_zh.md / THIRD_PARTY_NOTICES.md
|
|
561
|
+
```
|
|
562
|
+
|
|
563
|
+
---
|
|
564
|
+
|
|
565
|
+
## Roadmap
|
|
566
|
+
|
|
567
|
+
```mermaid
|
|
568
|
+
timeline
|
|
569
|
+
title Project Chronos delivery map
|
|
570
|
+
Phase 1
|
|
571
|
+
: LookaheadRouter
|
|
572
|
+
: Temporal locality regularization
|
|
573
|
+
: Router-probability collection path
|
|
574
|
+
Phase 2
|
|
575
|
+
: Async IO engine
|
|
576
|
+
: Three-tier SSD/RAM/VRAM storage
|
|
577
|
+
: Co-activation clustering
|
|
578
|
+
Phase 3
|
|
579
|
+
: Hybrid MLA + SlidingWindow attention
|
|
580
|
+
: PrefillScheduler
|
|
581
|
+
: Dual-layer routing
|
|
582
|
+
Phase 4
|
|
583
|
+
: Native MLX backend
|
|
584
|
+
: Web UI and CLI
|
|
585
|
+
: Optuna search
|
|
586
|
+
: Open-source release
|
|
587
|
+
M1-M3
|
|
588
|
+
: Cluster-aware safetensors storage
|
|
589
|
+
: Supervised lookahead loss
|
|
590
|
+
: Dual-stream decode with per-expert events
|
|
591
|
+
M4-M6
|
|
592
|
+
: SFT / DPO / ORPO / GRPO trainers
|
|
593
|
+
: Router KL anchor
|
|
594
|
+
: HF IO, vLLM adapter, multi-backend dispatch
|
|
595
|
+
: Stage 6 distillation and pluggable rewards
|
|
596
|
+
: Web UI v2, benchmark plots, IO Monitor
|
|
597
|
+
Next
|
|
598
|
+
: Train IntentClassifier on large activation corpora
|
|
599
|
+
: Benchmark 7B+ checkpoints
|
|
600
|
+
: Inject masks on the vLLM worker path
|
|
601
|
+
: Ship real Vulkan / OpenCL kernels
|
|
602
|
+
```
|
|
603
|
+
|
|
604
|
+
```mermaid
|
|
605
|
+
mindmap
|
|
606
|
+
root((Chronos innovation surface))
|
|
607
|
+
Predictive routing
|
|
608
|
+
IntentClassifier
|
|
609
|
+
Prompt-level hot expert prior
|
|
610
|
+
Budgeted expert-set prediction
|
|
611
|
+
LookaheadRouter
|
|
612
|
+
Per-token future routing
|
|
613
|
+
Soft-target CE from real future routers
|
|
614
|
+
Storage-aware MoE
|
|
615
|
+
Clustered safetensors
|
|
616
|
+
.ctsr packed expert clusters
|
|
617
|
+
manifest-driven mmap
|
|
618
|
+
Three-tier cache
|
|
619
|
+
NVMe SSD
|
|
620
|
+
Pinned RAM
|
|
621
|
+
VRAM / unified memory
|
|
622
|
+
Soft fallback
|
|
623
|
+
Shared experts always resident
|
|
624
|
+
No hard stall on cache miss
|
|
625
|
+
Decode pipeline
|
|
626
|
+
Prefill-time expert loading
|
|
627
|
+
AsyncPrefetcher queue
|
|
628
|
+
H2D stream
|
|
629
|
+
Per-expert CUDA events
|
|
630
|
+
Training stack
|
|
631
|
+
Pretrain
|
|
632
|
+
SFT
|
|
633
|
+
DPO
|
|
634
|
+
ORPO
|
|
635
|
+
GRPO
|
|
636
|
+
Distill
|
|
637
|
+
Router KL anchor
|
|
638
|
+
Deployment
|
|
639
|
+
HF safetensors IO
|
|
640
|
+
AutoModel registration
|
|
641
|
+
vLLM adapter
|
|
642
|
+
CPU / CUDA / MPS / MLX dispatch
|
|
643
|
+
Web UI and CLI
|
|
644
|
+
```
|
|
645
|
+
|
|
646
|
+
---
|
|
647
|
+
|
|
648
|
+
## Citation
|
|
649
|
+
|
|
650
|
+
```bibtex
|
|
651
|
+
@misc{chronos2026,
|
|
652
|
+
title = {Project Chronos: Prefill-Time Expert Loading and Dual-Layer Routing
|
|
653
|
+
for Zero-Stall On-Device MoE Inference},
|
|
654
|
+
author = {Fona and Project Chronos Contributors},
|
|
655
|
+
year = {2026},
|
|
656
|
+
url = {https://github.com/FonaTech/Project_Chronos}
|
|
657
|
+
}
|
|
658
|
+
```
|
|
659
|
+
|
|
660
|
+
---
|
|
661
|
+
|
|
662
|
+
## Third-party attribution
|
|
663
|
+
|
|
664
|
+
Project Chronos builds on **jingyaogong**'s [minimind](https://github.com/jingyaogong/minimind), licensed under **Apache-2.0**. Full attribution lives in [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md).
|
|
665
|
+
|
|
666
|
+
---
|
|
667
|
+
|
|
668
|
+
## License
|
|
669
|
+
|
|
670
|
+
Apache 2.0 - see [LICENSE](LICENSE)
|