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
ui/presets.py
ADDED
|
@@ -0,0 +1,194 @@
|
|
|
1
|
+
"""
|
|
2
|
+
ui/presets.py
|
|
3
|
+
|
|
4
|
+
Default + preset configurations matching minimind's MoE defaults.
|
|
5
|
+
|
|
6
|
+
Each preset is a complete dict that can be:
|
|
7
|
+
1. Loaded into the Config tab sliders via the "Load Preset" button.
|
|
8
|
+
2. Used as the initial value when the user clicks "Reset to MiniMind".
|
|
9
|
+
3. Persisted to / loaded from JSON via the Save/Load Config buttons.
|
|
10
|
+
|
|
11
|
+
The single source of truth for the *minimind* preset is what
|
|
12
|
+
``MiniMindConfig(use_moe=True)`` actually produces — we mirror those values
|
|
13
|
+
here so the UI shows the same model the trainer would build.
|
|
14
|
+
"""
|
|
15
|
+
from __future__ import annotations
|
|
16
|
+
|
|
17
|
+
import json
|
|
18
|
+
import os
|
|
19
|
+
from typing import Dict, List
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
# ── MiniMind MoE — exact defaults from model/model_minimind.py ─────
|
|
23
|
+
# These values match MiniMindConfig(hidden_size=768, num_hidden_layers=8,
|
|
24
|
+
# use_moe=True) field-for-field, with Chronos-specific extensions left at
|
|
25
|
+
# their conservative defaults (no router-anchor by default).
|
|
26
|
+
MINIMIND_MOE_DEFAULTS: Dict = {
|
|
27
|
+
# MiniMind core
|
|
28
|
+
"hidden_size": 768,
|
|
29
|
+
"num_hidden_layers": 8,
|
|
30
|
+
"num_attention_heads": 8,
|
|
31
|
+
"num_key_value_heads": 4,
|
|
32
|
+
"vocab_size": 6400,
|
|
33
|
+
"intermediate_size": 0, # 0 = auto: ceil(H·π/64)·64 → 2432
|
|
34
|
+
"moe_intermediate_size": 0, # 0 = auto, same value
|
|
35
|
+
"max_seq_len": 512, # training sequence length (NOT max_position_embeddings)
|
|
36
|
+
"tie_word_embeddings": True,
|
|
37
|
+
|
|
38
|
+
# MoE
|
|
39
|
+
"num_experts": 4,
|
|
40
|
+
"num_experts_per_tok": 2, # M8c: top-2 doubles activated params at fixed total params
|
|
41
|
+
"num_shared_experts": 1, # Chronos extension; minimind has no shared
|
|
42
|
+
|
|
43
|
+
# Attention shape (Chronos hybrid)
|
|
44
|
+
"kv_latent_dim": 64,
|
|
45
|
+
"rope_dim": 32,
|
|
46
|
+
"sliding_window_size": 2048,
|
|
47
|
+
"use_hybrid_attention": True,
|
|
48
|
+
|
|
49
|
+
# Chronos prediction
|
|
50
|
+
"lookahead_steps": 2,
|
|
51
|
+
|
|
52
|
+
# Loss coefficients
|
|
53
|
+
"lambda_balance": 5e-4,
|
|
54
|
+
"lambda_temporal": 1e-3,
|
|
55
|
+
"lambda_lookahead": 0.1,
|
|
56
|
+
"lambda_router_anchor": 0.0,
|
|
57
|
+
|
|
58
|
+
# Hardware
|
|
59
|
+
"vram_budget_gb": 4.0,
|
|
60
|
+
"pinned_memory_max_fraction": 0.25,
|
|
61
|
+
"storage_format": "safetensors",
|
|
62
|
+
|
|
63
|
+
# Training
|
|
64
|
+
"learning_rate": 5e-4,
|
|
65
|
+
"batch_size": 16,
|
|
66
|
+
"accumulation_steps": 8,
|
|
67
|
+
"epochs": 2,
|
|
68
|
+
"save_interval": 1000,
|
|
69
|
+
"save_dir": "./out",
|
|
70
|
+
"dtype": "fp16",
|
|
71
|
+
}
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
# Order matches ui/tabs/config_tab.py::all_inputs (used by Apply-Preset
|
|
75
|
+
# to push values back into the slider widgets in the right slots).
|
|
76
|
+
CONFIG_INPUT_ORDER: List[str] = [
|
|
77
|
+
"hidden_size", "num_hidden_layers", "num_experts", "num_experts_per_tok",
|
|
78
|
+
"num_shared_experts", "lookahead_steps", "kv_latent_dim", "sliding_window_size",
|
|
79
|
+
"num_attention_heads", "num_key_value_heads", "rope_dim", "moe_intermediate_size",
|
|
80
|
+
"vocab_size", "dtype", "tie_word_embeddings",
|
|
81
|
+
"lambda_balance", "lambda_temporal", "lambda_lookahead", "lambda_router_anchor",
|
|
82
|
+
"vram_budget_gb", "pinned_memory_max_fraction", "use_hybrid_attention", "storage_format",
|
|
83
|
+
"learning_rate", "batch_size", "accumulation_steps", "max_seq_len",
|
|
84
|
+
"epochs", "save_interval", "save_dir",
|
|
85
|
+
]
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
# Named presets — keys are user-facing labels (also used in the dropdown).
|
|
89
|
+
PRESETS: Dict[str, Dict] = {
|
|
90
|
+
"MiniMind-MoE (default)": dict(MINIMIND_MOE_DEFAULTS),
|
|
91
|
+
|
|
92
|
+
# M8c: "Recommended-CN" — sized for the /Users/Fona/Downloads/Hybrid_LLM/Dataset
|
|
93
|
+
# 1.3 GB pretrain corpus. Previously the default 15M-param model
|
|
94
|
+
# (H=256, L=4, E=4, topk=1 → ~6M activated) was far below the size
|
|
95
|
+
# needed for coherent Chinese output (empirically ~100M activated).
|
|
96
|
+
# This preset targets ~120M total / ~60M activated params at vocab=6400.
|
|
97
|
+
"Recommended-CN (≈120M)": {
|
|
98
|
+
**MINIMIND_MOE_DEFAULTS,
|
|
99
|
+
"hidden_size": 512, "num_hidden_layers": 8, "num_experts": 8,
|
|
100
|
+
"num_experts_per_tok": 2, "num_shared_experts": 1,
|
|
101
|
+
"num_attention_heads": 8, "num_key_value_heads": 2,
|
|
102
|
+
"kv_latent_dim": 64, "rope_dim": 32, "sliding_window_size": 2048,
|
|
103
|
+
"max_seq_len": 512, "batch_size": 8, "accumulation_steps": 8,
|
|
104
|
+
"epochs": 1, "save_interval": 500,
|
|
105
|
+
"learning_rate": 3e-4,
|
|
106
|
+
},
|
|
107
|
+
|
|
108
|
+
"Tiny (smoke / CI)": {
|
|
109
|
+
**MINIMIND_MOE_DEFAULTS,
|
|
110
|
+
"hidden_size": 128, "num_hidden_layers": 2, "num_experts": 4,
|
|
111
|
+
"num_attention_heads": 4, "num_key_value_heads": 4,
|
|
112
|
+
"kv_latent_dim": 16, "rope_dim": 8, "sliding_window_size": 64,
|
|
113
|
+
"max_seq_len": 64, "batch_size": 2, "accumulation_steps": 1,
|
|
114
|
+
"epochs": 1, "save_interval": 100,
|
|
115
|
+
},
|
|
116
|
+
|
|
117
|
+
"Small (256H × 4L × 4E)": {
|
|
118
|
+
**MINIMIND_MOE_DEFAULTS,
|
|
119
|
+
"hidden_size": 256, "num_hidden_layers": 4, "num_experts": 4,
|
|
120
|
+
"max_seq_len": 256, "batch_size": 8,
|
|
121
|
+
},
|
|
122
|
+
|
|
123
|
+
"Medium (512H × 8L × 8E)": {
|
|
124
|
+
**MINIMIND_MOE_DEFAULTS,
|
|
125
|
+
"hidden_size": 512, "num_experts": 8,
|
|
126
|
+
"max_seq_len": 1024, "batch_size": 8,
|
|
127
|
+
},
|
|
128
|
+
|
|
129
|
+
"Large (1024H × 16L × 16E)": {
|
|
130
|
+
**MINIMIND_MOE_DEFAULTS,
|
|
131
|
+
"hidden_size": 1024, "num_hidden_layers": 16, "num_experts": 16,
|
|
132
|
+
"num_attention_heads": 16, "num_key_value_heads": 4,
|
|
133
|
+
"kv_latent_dim": 128, "rope_dim": 64,
|
|
134
|
+
"max_seq_len": 2048, "batch_size": 4, "accumulation_steps": 16,
|
|
135
|
+
"vram_budget_gb": 16.0,
|
|
136
|
+
},
|
|
137
|
+
}
|
|
138
|
+
|
|
139
|
+
|
|
140
|
+
def preset_names() -> List[str]:
|
|
141
|
+
return list(PRESETS.keys())
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
def get_preset(name: str) -> Dict:
|
|
145
|
+
"""Return a deep-ish copy of the named preset (avoids accidental mutation)."""
|
|
146
|
+
return dict(PRESETS.get(name, MINIMIND_MOE_DEFAULTS))
|
|
147
|
+
|
|
148
|
+
|
|
149
|
+
def values_in_input_order(cfg: Dict) -> List:
|
|
150
|
+
"""Return preset values in the same order as Config tab's all_inputs.
|
|
151
|
+
|
|
152
|
+
For keys missing from `cfg`, falls back to MINIMIND_MOE_DEFAULTS. If a
|
|
153
|
+
value is ``None`` (missing in both), it is replaced by a safe-default
|
|
154
|
+
per field type so Gradio's slider/number widgets don't reject it.
|
|
155
|
+
"""
|
|
156
|
+
# Hardcoded fallbacks for fields whose MiniMind default is None or absent.
|
|
157
|
+
# Values chosen to match the slider minima where applicable.
|
|
158
|
+
SAFE = {
|
|
159
|
+
"hidden_size": 768, "num_hidden_layers": 8,
|
|
160
|
+
"num_experts": 4, "num_experts_per_tok": 2, "num_shared_experts": 1,
|
|
161
|
+
"lookahead_steps": 2, "kv_latent_dim": 64, "sliding_window_size": 2048,
|
|
162
|
+
"num_attention_heads": 8, "num_key_value_heads": 4, "rope_dim": 32,
|
|
163
|
+
"moe_intermediate_size": 0, "vocab_size": 6400,
|
|
164
|
+
"dtype": "fp16", "tie_word_embeddings": True,
|
|
165
|
+
"lambda_balance": 5e-4, "lambda_temporal": 1e-3,
|
|
166
|
+
"lambda_lookahead": 0.1, "lambda_router_anchor": 0.0,
|
|
167
|
+
"vram_budget_gb": 4.0, "pinned_memory_max_fraction": 0.25,
|
|
168
|
+
"use_hybrid_attention": True, "storage_format": "safetensors",
|
|
169
|
+
"learning_rate": 5e-4, "batch_size": 16, "accumulation_steps": 8,
|
|
170
|
+
"max_seq_len": 512, "epochs": 2, "save_interval": 1000,
|
|
171
|
+
"save_dir": "./out",
|
|
172
|
+
}
|
|
173
|
+
out = []
|
|
174
|
+
for k in CONFIG_INPUT_ORDER:
|
|
175
|
+
v = cfg.get(k) if k in cfg else MINIMIND_MOE_DEFAULTS.get(k)
|
|
176
|
+
if v is None:
|
|
177
|
+
v = SAFE.get(k, 0)
|
|
178
|
+
out.append(v)
|
|
179
|
+
return out
|
|
180
|
+
|
|
181
|
+
|
|
182
|
+
def save_config(cfg: Dict, path: str) -> str:
|
|
183
|
+
"""Write cfg as pretty JSON. Returns the absolute path written."""
|
|
184
|
+
abs_path = os.path.abspath(path)
|
|
185
|
+
os.makedirs(os.path.dirname(abs_path) or ".", exist_ok=True)
|
|
186
|
+
with open(abs_path, "w") as f:
|
|
187
|
+
json.dump(cfg, f, indent=2, sort_keys=True)
|
|
188
|
+
return abs_path
|
|
189
|
+
|
|
190
|
+
|
|
191
|
+
def load_config(path: str) -> Dict:
|
|
192
|
+
"""Read JSON and return the dict. Raises FileNotFoundError if missing."""
|
|
193
|
+
with open(os.path.abspath(path)) as f:
|
|
194
|
+
return json.load(f)
|
ui/tabs/__init__.py
ADDED
ui/tabs/autotune_tab.py
ADDED
|
@@ -0,0 +1,210 @@
|
|
|
1
|
+
"""
|
|
2
|
+
ui/tabs/autotune_tab.py — Optuna search over Chronos hyperparameters.
|
|
3
|
+
|
|
4
|
+
Improvements over the M5 version:
|
|
5
|
+
|
|
6
|
+
- **Log accumulation**: events append to a rolling buffer, persisted to
|
|
7
|
+
``./auto_tune_cache/tune_log.txt`` so history survives a UI refresh.
|
|
8
|
+
- **Best-params history**: every trial's params + score appended to
|
|
9
|
+
``./auto_tune_cache/trials.jsonl``.
|
|
10
|
+
- **Expanded search space**: structural knobs (hidden_size, num_experts,
|
|
11
|
+
num_shared_experts, kv_latent_dim, λ_lookahead) can be ticked for search.
|
|
12
|
+
- **Apply Best → Config**: one button writes the best params back into
|
|
13
|
+
the Config tab's slider widgets.
|
|
14
|
+
"""
|
|
15
|
+
import json
|
|
16
|
+
import os
|
|
17
|
+
import time
|
|
18
|
+
|
|
19
|
+
from ui.gradio_compat import gr
|
|
20
|
+
|
|
21
|
+
import chronos.deps # noqa: F401
|
|
22
|
+
from ui.i18n import t, register_translatable
|
|
23
|
+
from chronos.tuning.chronos_auto_tuner import ChronosAutoTuner, ChronosSearchSpaceConfig
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
_tuner = ChronosAutoTuner()
|
|
27
|
+
_CACHE_DIR = os.path.abspath("./auto_tune_cache")
|
|
28
|
+
_LOG_FILE = os.path.join(_CACHE_DIR, "tune_log.txt")
|
|
29
|
+
_TRIALS_FILE = os.path.join(_CACHE_DIR, "trials.jsonl")
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _ensure_cache():
|
|
33
|
+
os.makedirs(_CACHE_DIR, exist_ok=True)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def _append_log(text: str):
|
|
37
|
+
_ensure_cache()
|
|
38
|
+
with open(_LOG_FILE, "a") as f:
|
|
39
|
+
f.write(text + ("\n" if not text.endswith("\n") else ""))
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def _read_log() -> str:
|
|
43
|
+
if not os.path.exists(_LOG_FILE):
|
|
44
|
+
return ""
|
|
45
|
+
with open(_LOG_FILE) as f:
|
|
46
|
+
return f.read()
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def _append_trial(record: dict):
|
|
50
|
+
_ensure_cache()
|
|
51
|
+
with open(_TRIALS_FILE, "a") as f:
|
|
52
|
+
f.write(json.dumps(record) + "\n")
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
# Map best-params keys to indices into the Config tab's all_inputs list.
|
|
56
|
+
# IMPORTANT: must stay in sync with ui/tabs/config_tab.py:all_inputs order.
|
|
57
|
+
PARAM_TO_CONFIG_IDX = {
|
|
58
|
+
"hidden_size": 0,
|
|
59
|
+
"num_hidden_layers": 1,
|
|
60
|
+
"num_experts": 2,
|
|
61
|
+
"num_experts_per_tok": 3,
|
|
62
|
+
"num_shared_experts": 4,
|
|
63
|
+
"lookahead_steps": 5,
|
|
64
|
+
"kv_latent_dim": 6,
|
|
65
|
+
"sliding_window_size": 7,
|
|
66
|
+
"lambda_balance": 15,
|
|
67
|
+
"lambda_temporal": 16,
|
|
68
|
+
"lambda_lookahead": 17,
|
|
69
|
+
"lambda_router_anchor": 18,
|
|
70
|
+
"learning_rate": 23,
|
|
71
|
+
"batch_size": 24,
|
|
72
|
+
}
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def build_autotune_tab(config_state: gr.State, config_inputs: list):
|
|
76
|
+
with gr.Tab(t("tab.autotune")) as tab:
|
|
77
|
+
register_translatable(tab, "tab.autotune")
|
|
78
|
+
|
|
79
|
+
gr.Markdown("### Optuna TPE — search over Chronos hyperparameters")
|
|
80
|
+
|
|
81
|
+
with gr.Row():
|
|
82
|
+
data_path = gr.Textbox(label=t("autotune.data_path"), scale=3,
|
|
83
|
+
placeholder="./tests/fixtures/tiny_pretrain.jsonl")
|
|
84
|
+
n_trials = gr.Slider(2, 200, value=20, step=1, label=t("autotune.n_trials"))
|
|
85
|
+
probe_steps = gr.Slider(10, 1000, value=80, step=10, label=t("autotune.probe_steps"))
|
|
86
|
+
register_translatable(data_path, "autotune.data_path")
|
|
87
|
+
register_translatable(n_trials, "autotune.n_trials")
|
|
88
|
+
register_translatable(probe_steps, "autotune.probe_steps")
|
|
89
|
+
|
|
90
|
+
gr.Markdown(t("autotune.tip"))
|
|
91
|
+
with gr.Row():
|
|
92
|
+
t_lb = gr.Checkbox(value=True, label="λ1 balance")
|
|
93
|
+
t_lt = gr.Checkbox(value=True, label="λ2 temporal")
|
|
94
|
+
t_ll = gr.Checkbox(value=False, label="λ lookahead")
|
|
95
|
+
t_la = gr.Checkbox(value=False, label="lookahead_steps")
|
|
96
|
+
t_lr = gr.Checkbox(value=True, label="learning_rate")
|
|
97
|
+
with gr.Row():
|
|
98
|
+
t_hs = gr.Checkbox(value=False, label="hidden_size")
|
|
99
|
+
t_ne = gr.Checkbox(value=False, label="num_experts")
|
|
100
|
+
t_ns = gr.Checkbox(value=False, label="num_shared_experts")
|
|
101
|
+
t_kv = gr.Checkbox(value=False, label="kv_latent_dim")
|
|
102
|
+
|
|
103
|
+
with gr.Row():
|
|
104
|
+
start_btn = gr.Button(t("autotune.start"), variant="primary")
|
|
105
|
+
stop_btn = gr.Button(t("autotune.stop"), variant="stop")
|
|
106
|
+
clear_btn = gr.Button(t("autotune.clear_log"))
|
|
107
|
+
apply_btn = gr.Button(t("autotune.apply_best"), variant="secondary")
|
|
108
|
+
register_translatable(start_btn, "autotune.start")
|
|
109
|
+
register_translatable(stop_btn, "autotune.stop")
|
|
110
|
+
register_translatable(clear_btn, "autotune.clear_log")
|
|
111
|
+
register_translatable(apply_btn, "autotune.apply_best")
|
|
112
|
+
|
|
113
|
+
status_box = gr.Textbox(value="idle", label="Status", interactive=False)
|
|
114
|
+
best_box = gr.JSON(label=t("autotune.best"))
|
|
115
|
+
log_box = gr.Textbox(label=t("autotune.log"), lines=20,
|
|
116
|
+
interactive=False, autoscroll=True,
|
|
117
|
+
value=_read_log())
|
|
118
|
+
register_translatable(log_box, "autotune.log")
|
|
119
|
+
register_translatable(best_box, "autotune.best")
|
|
120
|
+
|
|
121
|
+
gr.Markdown(f"Log persisted at `{_LOG_FILE}` (survives refresh).")
|
|
122
|
+
|
|
123
|
+
def start_tune(dpath, n_tr, p_steps,
|
|
124
|
+
tlb, tlt, tll, tla, tlr,
|
|
125
|
+
ths, tne, tns, tkv,
|
|
126
|
+
cfg):
|
|
127
|
+
if _tuner.is_running():
|
|
128
|
+
return "already running", _read_log(), {}
|
|
129
|
+
if not dpath or not os.path.exists(dpath):
|
|
130
|
+
return f"error: data_path not found: {dpath!r}", _read_log(), {}
|
|
131
|
+
|
|
132
|
+
_append_log(f"\n──── Auto-tune started @ {time.strftime('%Y-%m-%d %H:%M:%S')} ────")
|
|
133
|
+
|
|
134
|
+
ss = ChronosSearchSpaceConfig(
|
|
135
|
+
tune_lambda_balance=bool(tlb),
|
|
136
|
+
tune_lambda_temporal=bool(tlt),
|
|
137
|
+
tune_lambda_lookahead=bool(tll),
|
|
138
|
+
tune_lookahead_steps=bool(tla),
|
|
139
|
+
tune_lr=bool(tlr),
|
|
140
|
+
tune_hidden_size=bool(ths),
|
|
141
|
+
tune_num_experts=bool(tne),
|
|
142
|
+
tune_num_shared_experts=bool(tns),
|
|
143
|
+
tune_kv_latent_dim=bool(tkv),
|
|
144
|
+
)
|
|
145
|
+
save_dir = (cfg or {}).get("save_dir", "./out")
|
|
146
|
+
hidden_size = (cfg or {}).get("hidden_size", 256)
|
|
147
|
+
_tuner.start(
|
|
148
|
+
model_id=os.path.join(save_dir, f"chronos_{hidden_size}_moe.pth"),
|
|
149
|
+
dataset_path=dpath, train_ratio=0.95,
|
|
150
|
+
prompt_template="alpaca", think_mode="keep",
|
|
151
|
+
search_space=ss, n_trials=int(n_tr),
|
|
152
|
+
probe_steps=int(p_steps),
|
|
153
|
+
output_dir=_CACHE_DIR, seed=42,
|
|
154
|
+
)
|
|
155
|
+
return "running", _read_log(), {}
|
|
156
|
+
|
|
157
|
+
def stop_tune():
|
|
158
|
+
_tuner.stop()
|
|
159
|
+
_append_log(f"──── stopped @ {time.strftime('%Y-%m-%d %H:%M:%S')} ────")
|
|
160
|
+
return "stopped"
|
|
161
|
+
|
|
162
|
+
def clear_log_file():
|
|
163
|
+
if os.path.exists(_LOG_FILE):
|
|
164
|
+
os.remove(_LOG_FILE)
|
|
165
|
+
return ""
|
|
166
|
+
|
|
167
|
+
def poll_tune():
|
|
168
|
+
events = _tuner.poll()
|
|
169
|
+
for ev in events:
|
|
170
|
+
if ev.get("type") == "log":
|
|
171
|
+
_append_log(ev.get("line", ""))
|
|
172
|
+
elif ev.get("type") == "trial_done":
|
|
173
|
+
_append_trial(ev)
|
|
174
|
+
best = {}
|
|
175
|
+
if _tuner.result and _tuner.result.best_params:
|
|
176
|
+
best = _tuner.result.best_params
|
|
177
|
+
return _tuner.status, _read_log(), best
|
|
178
|
+
|
|
179
|
+
def apply_best(best):
|
|
180
|
+
"""Write best_params back into the Config tab's slider widgets.
|
|
181
|
+
Returns a list of gr.update for every config_input — untouched
|
|
182
|
+
inputs get gr.update() (no-op). None values are dropped so
|
|
183
|
+
Gradio's numeric widgets don't reject them."""
|
|
184
|
+
updates = [gr.update() for _ in config_inputs]
|
|
185
|
+
if not best:
|
|
186
|
+
return updates
|
|
187
|
+
for k, v in best.items():
|
|
188
|
+
if v is None:
|
|
189
|
+
continue
|
|
190
|
+
idx = PARAM_TO_CONFIG_IDX.get(k)
|
|
191
|
+
if idx is not None:
|
|
192
|
+
updates[idx] = gr.update(value=v)
|
|
193
|
+
return updates
|
|
194
|
+
|
|
195
|
+
start_btn.click(
|
|
196
|
+
fn=start_tune,
|
|
197
|
+
inputs=[data_path, n_trials, probe_steps,
|
|
198
|
+
t_lb, t_lt, t_ll, t_la, t_lr,
|
|
199
|
+
t_hs, t_ne, t_ns, t_kv,
|
|
200
|
+
config_state],
|
|
201
|
+
outputs=[status_box, log_box, best_box],
|
|
202
|
+
)
|
|
203
|
+
stop_btn.click(fn=stop_tune, outputs=[status_box])
|
|
204
|
+
clear_btn.click(fn=clear_log_file, outputs=[log_box])
|
|
205
|
+
apply_btn.click(fn=apply_best, inputs=[best_box], outputs=config_inputs)
|
|
206
|
+
|
|
207
|
+
timer = gr.Timer(value=2.0)
|
|
208
|
+
timer.tick(fn=poll_tune, outputs=[status_box, log_box, best_box])
|
|
209
|
+
|
|
210
|
+
return tab
|
ui/tabs/benchmark_tab.py
ADDED
|
@@ -0,0 +1,196 @@
|
|
|
1
|
+
"""
|
|
2
|
+
ui/tabs/benchmark_tab.py — MiniMind vs Chronos comparison view.
|
|
3
|
+
|
|
4
|
+
Renders the comparison_results.json as:
|
|
5
|
+
|
|
6
|
+
- a formatted Markdown table (parameters, RAM, tokens/sec)
|
|
7
|
+
- per-metric bar charts (gr.BarPlot) for visual comparison
|
|
8
|
+
- the raw JSON (collapsed; for power users)
|
|
9
|
+
- a streaming log of the comparison subprocess
|
|
10
|
+
"""
|
|
11
|
+
import json
|
|
12
|
+
import os
|
|
13
|
+
|
|
14
|
+
from ui.gradio_compat import gr
|
|
15
|
+
import pandas as pd
|
|
16
|
+
|
|
17
|
+
import chronos.deps # noqa: F401
|
|
18
|
+
from ui.i18n import t, register_translatable
|
|
19
|
+
|
|
20
|
+
RESULTS_FILE = os.path.join(os.path.dirname(__file__), '../../benchmark_results.json')
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
# Metrics we know how to chart. Order matters — drives the row order of
|
|
24
|
+
# the bar plots. Each entry: (display_name, json_key, unit_label, lower_is_better).
|
|
25
|
+
CHART_METRICS = [
|
|
26
|
+
("Parameters", "params_m", "M params", False),
|
|
27
|
+
("Model RAM (resident)", "ram_model_gb", "GB", True),
|
|
28
|
+
("Train RAM", "ram_train_gb", "GB", True),
|
|
29
|
+
("Inference RAM", "ram_infer_gb", "GB", True),
|
|
30
|
+
("Decode tokens/sec", "tokens_per_sec","tokens/s", False),
|
|
31
|
+
]
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _format_table(data: dict) -> str:
|
|
35
|
+
"""Render comparison results as a Markdown table."""
|
|
36
|
+
if not data:
|
|
37
|
+
return "_No benchmark data yet. Click **Run** to start._"
|
|
38
|
+
|
|
39
|
+
mm = data.get("minimind", {}) or {}
|
|
40
|
+
ch = data.get("chronos", {}) or {}
|
|
41
|
+
|
|
42
|
+
rows = ["| Metric | MiniMind | Chronos | Δ (Chronos − MiniMind) | Better |",
|
|
43
|
+
"|---|---|---|---|---|"]
|
|
44
|
+
for label, key, unit, lower_is_better in CHART_METRICS:
|
|
45
|
+
mv = mm.get(key)
|
|
46
|
+
cv = ch.get(key)
|
|
47
|
+
if mv is None and cv is None:
|
|
48
|
+
continue
|
|
49
|
+
mvs = "n/a" if mv is None else f"{mv:.4g} {unit}"
|
|
50
|
+
cvs = "n/a" if cv is None else f"{cv:.4g} {unit}"
|
|
51
|
+
if mv is None or cv is None:
|
|
52
|
+
delta = "n/a"; better = ""
|
|
53
|
+
else:
|
|
54
|
+
d = cv - mv
|
|
55
|
+
sign = "−" if d < 0 else ("+" if d > 0 else "±")
|
|
56
|
+
delta = f"{sign}{abs(d):.4g} {unit}"
|
|
57
|
+
if lower_is_better:
|
|
58
|
+
better = "**Chronos**" if cv < mv else ("MiniMind" if cv > mv else "tie")
|
|
59
|
+
else:
|
|
60
|
+
better = "**Chronos**" if cv > mv else ("MiniMind" if cv < mv else "tie")
|
|
61
|
+
rows.append(f"| {label} | {mvs} | {cvs} | {delta} | {better} |")
|
|
62
|
+
|
|
63
|
+
extra = []
|
|
64
|
+
if "kv_cache_type" in ch:
|
|
65
|
+
extra.append(f"- **Chronos KV cache**: `{ch['kv_cache_type']}`")
|
|
66
|
+
cs = ch.get("cache_stats") or {}
|
|
67
|
+
if cs:
|
|
68
|
+
extra.append("- **Cache stats**: " + ", ".join(
|
|
69
|
+
f"`{k}={v}`" for k, v in cs.items()
|
|
70
|
+
if k in ("vram_experts", "ram_experts", "num_clusters",
|
|
71
|
+
"storage_format", "hit_rate", "expert_size_kb",
|
|
72
|
+
"pinned_ram_used_gb")
|
|
73
|
+
))
|
|
74
|
+
if extra:
|
|
75
|
+
rows.append("")
|
|
76
|
+
rows.extend(extra)
|
|
77
|
+
|
|
78
|
+
return "\n".join(rows)
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def _to_chart_df(data: dict) -> pd.DataFrame:
|
|
82
|
+
"""Convert benchmark JSON into a long-format DataFrame for gr.BarPlot."""
|
|
83
|
+
rows = []
|
|
84
|
+
for model_key, label_key in (("minimind", "MiniMind"), ("chronos", "Chronos")):
|
|
85
|
+
d = data.get(model_key, {}) or {}
|
|
86
|
+
for label, key, unit, _lower in CHART_METRICS:
|
|
87
|
+
v = d.get(key)
|
|
88
|
+
if v is None:
|
|
89
|
+
continue
|
|
90
|
+
rows.append({
|
|
91
|
+
"metric": label,
|
|
92
|
+
"model": d.get("name", label_key),
|
|
93
|
+
"value": float(v),
|
|
94
|
+
"unit": unit,
|
|
95
|
+
})
|
|
96
|
+
return pd.DataFrame(rows) if rows else pd.DataFrame(
|
|
97
|
+
columns=["metric", "model", "value", "unit"]
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
def _empty_df():
|
|
102
|
+
return pd.DataFrame(columns=["metric", "model", "value", "unit"])
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def _load_existing():
|
|
106
|
+
p = os.path.abspath(RESULTS_FILE)
|
|
107
|
+
if os.path.exists(p):
|
|
108
|
+
try:
|
|
109
|
+
with open(p) as f:
|
|
110
|
+
d = json.load(f)
|
|
111
|
+
return d, _format_table(d), _to_chart_df(d), "Loaded previous results."
|
|
112
|
+
except Exception as e:
|
|
113
|
+
return {}, f"_Could not parse {p}: {e}_", _empty_df(), str(e)
|
|
114
|
+
return {}, _format_table({}), _empty_df(), "No results yet. Click Run."
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def build_benchmark_tab():
|
|
118
|
+
with gr.Tab(t("tab.benchmark")) as tab:
|
|
119
|
+
register_translatable(tab, "tab.benchmark")
|
|
120
|
+
|
|
121
|
+
run_btn = gr.Button(t("bench.run"), variant="primary")
|
|
122
|
+
register_translatable(run_btn, "bench.run")
|
|
123
|
+
|
|
124
|
+
# Formatted summary
|
|
125
|
+
summary_md = gr.Markdown(value=_format_table({}))
|
|
126
|
+
|
|
127
|
+
# Bar chart — one BarPlot per model so units stack cleanly per metric.
|
|
128
|
+
chart = gr.BarPlot(
|
|
129
|
+
value=_empty_df(),
|
|
130
|
+
x="metric", y="value", color="model",
|
|
131
|
+
title="MiniMind vs Chronos — comparison metrics",
|
|
132
|
+
tooltip=["model", "metric", "value", "unit"],
|
|
133
|
+
height=320,
|
|
134
|
+
)
|
|
135
|
+
|
|
136
|
+
with gr.Accordion("Raw JSON", open=False):
|
|
137
|
+
results_box = gr.JSON(label=t("bench.results"))
|
|
138
|
+
register_translatable(results_box, "bench.results")
|
|
139
|
+
|
|
140
|
+
log_box = gr.Textbox(
|
|
141
|
+
label=t("bench.log"), lines=20, interactive=False, autoscroll=True
|
|
142
|
+
)
|
|
143
|
+
register_translatable(log_box, "bench.log")
|
|
144
|
+
|
|
145
|
+
def run_benchmark():
|
|
146
|
+
import subprocess, sys
|
|
147
|
+
script = os.path.join(os.path.dirname(__file__), '../../benchmark_compare.py')
|
|
148
|
+
log_lines = []
|
|
149
|
+
yield (
|
|
150
|
+
{}, _format_table({}), _empty_df(),
|
|
151
|
+
"Running benchmark... (this takes ~2 minutes)\n",
|
|
152
|
+
)
|
|
153
|
+
try:
|
|
154
|
+
proc = subprocess.Popen(
|
|
155
|
+
[sys.executable, script],
|
|
156
|
+
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
|
157
|
+
text=True, bufsize=1,
|
|
158
|
+
)
|
|
159
|
+
for line in proc.stdout:
|
|
160
|
+
log_lines.append(line.rstrip())
|
|
161
|
+
yield (
|
|
162
|
+
{}, _format_table({}), _empty_df(),
|
|
163
|
+
"\n".join(log_lines[-50:]),
|
|
164
|
+
)
|
|
165
|
+
proc.wait()
|
|
166
|
+
p = os.path.abspath(RESULTS_FILE)
|
|
167
|
+
if os.path.exists(p):
|
|
168
|
+
with open(p) as f:
|
|
169
|
+
results = json.load(f)
|
|
170
|
+
yield (
|
|
171
|
+
results, _format_table(results), _to_chart_df(results),
|
|
172
|
+
"\n".join(log_lines),
|
|
173
|
+
)
|
|
174
|
+
else:
|
|
175
|
+
yield (
|
|
176
|
+
{}, "_Run finished but no results file was produced._",
|
|
177
|
+
_empty_df(), "\n".join(log_lines) + "\n[No results file found]",
|
|
178
|
+
)
|
|
179
|
+
except Exception as e:
|
|
180
|
+
yield (
|
|
181
|
+
{}, f"_error: {e}_", _empty_df(), f"Error: {e}",
|
|
182
|
+
)
|
|
183
|
+
|
|
184
|
+
run_btn.click(
|
|
185
|
+
fn=run_benchmark,
|
|
186
|
+
outputs=[results_box, summary_md, chart, log_box],
|
|
187
|
+
)
|
|
188
|
+
|
|
189
|
+
# Render existing results on tab open
|
|
190
|
+
d, md, df, log = _load_existing()
|
|
191
|
+
results_box.value = d
|
|
192
|
+
summary_md.value = md
|
|
193
|
+
chart.value = df
|
|
194
|
+
log_box.value = log
|
|
195
|
+
|
|
196
|
+
return tab
|