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.
Files changed (83) hide show
  1. chronos/__init__.py +16 -0
  2. chronos/_backend_legacy.py +64 -0
  3. chronos/app.py +122 -0
  4. chronos/backend/__init__.py +40 -0
  5. chronos/backend/dispatcher.py +347 -0
  6. chronos/backend/ext/__init__.py +2 -0
  7. chronos/backend/ext/opencl.py +22 -0
  8. chronos/cli.py +62 -0
  9. chronos/data/__init__.py +0 -0
  10. chronos/data/flexible_dataset.py +309 -0
  11. chronos/deps.py +72 -0
  12. chronos/eval/__init__.py +2 -0
  13. chronos/eval/benchmark.py +159 -0
  14. chronos/eval/io_profiler.py +232 -0
  15. chronos/io/__init__.py +13 -0
  16. chronos/io/async_prefetcher.py +220 -0
  17. chronos/io/cluster_layout.py +317 -0
  18. chronos/io/expert_store.py +559 -0
  19. chronos/io/io_simulator.py +40 -0
  20. chronos/io/storage.py +237 -0
  21. chronos/mlx/__init__.py +32 -0
  22. chronos/mlx/attention.py +230 -0
  23. chronos/mlx/expert_store.py +155 -0
  24. chronos/mlx/inference.py +159 -0
  25. chronos/mlx/model.py +181 -0
  26. chronos/mlx/moe.py +102 -0
  27. chronos/model/__init__.py +6 -0
  28. chronos/model/config.py +54 -0
  29. chronos/model/hf_io.py +149 -0
  30. chronos/model/hybrid_attention.py +237 -0
  31. chronos/model/lookahead_router.py +55 -0
  32. chronos/model/model_chronos.py +271 -0
  33. chronos/model/moe_chronos.py +189 -0
  34. chronos/model/temporal_loss.py +124 -0
  35. chronos/router/__init__.py +12 -0
  36. chronos/router/expert_predictor.py +163 -0
  37. chronos/router/intent_classifier.py +193 -0
  38. chronos/router/prefill_scheduler.py +248 -0
  39. chronos/runtime/__init__.py +2 -0
  40. chronos/runtime/cache_manager.py +117 -0
  41. chronos/runtime/inference_engine.py +295 -0
  42. chronos/runtime/metrics.py +75 -0
  43. chronos/serving/__init__.py +13 -0
  44. chronos/serving/vllm_adapter.py +85 -0
  45. chronos/trainer/__init__.py +6 -0
  46. chronos/trainer/chronos_trainer.py +151 -0
  47. chronos/trainer/distill_trainer.py +193 -0
  48. chronos/trainer/dpo_trainer.py +190 -0
  49. chronos/trainer/grpo_trainer.py +288 -0
  50. chronos/trainer/loss_mixin.py +165 -0
  51. chronos/trainer/optim_utils.py +101 -0
  52. chronos/trainer/orpo_trainer.py +189 -0
  53. chronos/trainer/reward.py +130 -0
  54. chronos/trainer/sft_trainer.py +189 -0
  55. chronos/tuning/__init__.py +1 -0
  56. chronos/tuning/_base_tuner.py +284 -0
  57. chronos/tuning/chronos_auto_tuner.py +203 -0
  58. project_chronos-0.1.1.dist-info/METADATA +670 -0
  59. project_chronos-0.1.1.dist-info/RECORD +83 -0
  60. project_chronos-0.1.1.dist-info/WHEEL +5 -0
  61. project_chronos-0.1.1.dist-info/entry_points.txt +7 -0
  62. project_chronos-0.1.1.dist-info/licenses/LICENSE +177 -0
  63. project_chronos-0.1.1.dist-info/top_level.txt +8 -0
  64. train_chronos.py +133 -0
  65. train_chronos_distill.py +95 -0
  66. train_chronos_dpo.py +82 -0
  67. train_chronos_grpo.py +88 -0
  68. train_chronos_orpo.py +81 -0
  69. train_chronos_sft.py +85 -0
  70. ui/__init__.py +3 -0
  71. ui/estimator.py +204 -0
  72. ui/gradio_compat.py +72 -0
  73. ui/i18n.py +540 -0
  74. ui/presets.py +194 -0
  75. ui/tabs/__init__.py +3 -0
  76. ui/tabs/autotune_tab.py +210 -0
  77. ui/tabs/benchmark_tab.py +196 -0
  78. ui/tabs/config_tab.py +318 -0
  79. ui/tabs/data_tab.py +148 -0
  80. ui/tabs/inference_tab.py +496 -0
  81. ui/tabs/iomon_tab.py +37 -0
  82. ui/tabs/pipeline_tab.py +134 -0
  83. 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
+ [![PyPI](https://img.shields.io/pypi/v/project-chronos)](https://pypi.org/project/project-chronos/)
41
+ [![License](https://img.shields.io/badge/license-Apache%202.0-blue)](LICENSE)
42
+ [![Python](https://img.shields.io/badge/python-3.10%2B-blue)](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)