Add files via upload

Author: FonaTech

pyproject.toml

@@ -19,7 +19,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 dependencies = [
-    "torch>=2.1.0",
+    "torch>=2.4.0",
     "transformers>=4.40.0",
     "datasets>=2.18.0",
     "safetensors>=0.4.0",

requirements.txt

@@ -1,9 +1,8 @@
-torch>=2.1.0
+torch>=2.4.0
 transformers>=4.40.0
 datasets>=2.18.0
 tokenizers>=0.19.0
 safetensors>=0.4.0
 optuna>=3.6.0
 numpy>=1.24.0
 tqdm>=4.66.0
-gradio>=6.0.0

ui/tabs/config_tab.py

@@ -12,15 +12,13 @@
 """
 from ui.gradio_compat import gr
 
-from chronos.model.config import ChronosConfig
 from ui.i18n import t, register_translatable
 from ui.estimator import (
     ArchConfig, total_params, active_params, memory_footprint,
     estimated_decode_tps, fmt_bytes, fmt_params,
 )
 from ui.presets import (
-    MINIMIND_MOE_DEFAULTS, PRESETS, preset_names, get_preset,
-    values_in_input_order, save_config, load_config, CONFIG_INPUT_ORDER,
+    preset_names, get_preset, values_in_input_order, save_config, load_config,
 )
 
 
@@ -88,7 +86,9 @@ def build_config_tab():
     data_path has been moved to the Train tab, so the tuple now has 3
     elements. Callers must be updated.
     """
-    config_state = gr.State(ChronosConfig().__dict__.copy())
+    initial_preset = "Recommended-CN (≈120M)"
+    initial_cfg = get_preset(initial_preset)
+    config_state = gr.State(dict(initial_cfg))
 
     with gr.Tab(t("tab.config")) as tab:
         register_translatable(tab, "tab.config")
@@ -97,7 +97,7 @@ def build_config_tab():
         with gr.Row():
             preset_dd = gr.Dropdown(
                 choices=preset_names(),
-                value="Recommended-CN (≈120M)",
+                value=initial_preset,
                 label=t("config.preset"),
                 scale=2,
             )
@@ -123,7 +123,7 @@ def build_config_tab():
             with gr.Column(scale=3):
                 gr.Markdown(f"### {t('config.arch')}")
 
-                D = MINIMIND_MOE_DEFAULTS  # MiniMind-MoE starting values
+                D = initial_cfg
 
                 with gr.Row():
                     hidden_size     = gr.Slider(*RANGES["hidden_size"],       value=D["hidden_size"],         label=t("config.hidden_size"))
@@ -211,7 +211,7 @@ def build_config_tab():
                     register_translatable(save_interval, "config.save_interval")
                     register_translatable(save_dir,      "config.save_dir")
 
-                config_display = gr.JSON(label="Current Config (saved to config_state)", value={})
+                config_display = gr.JSON(label="Current Config (saved to config_state)", value=initial_cfg)
 
             with gr.Column(scale=1, min_width=260):
                 gr.Markdown(f"### 🧬 {t('designer.title')}")
@@ -282,11 +282,11 @@ def update_config(*vals):
         # ── Preset / Save / Load wiring ───────────────────────
         def apply_preset(name):
             cfg = get_preset(name)
-            return [gr.update(value=v) for v in values_in_input_order(cfg)]
+            return [dict(cfg), dict(cfg)] + [gr.update(value=v) for v in values_in_input_order(cfg)]
 
         def reset_minimind():
             cfg = get_preset("MiniMind-MoE (default)")
-            return [gr.update(value=v) for v in values_in_input_order(cfg)]
+            return [dict(cfg), dict(cfg)] + [gr.update(value=v) for v in values_in_input_order(cfg)]
 
         def do_save(cfg, path):
             try:
@@ -299,21 +299,20 @@ def do_load(path):
             try:
                 cfg = load_config(path)
                 vals = values_in_input_order(cfg)
-                return [f"✅ Loaded from `{path}`"] + [gr.update(value=v) for v in vals]
+                return [f"✅ Loaded from `{path}`", dict(cfg), dict(cfg)] + [gr.update(value=v) for v in vals]
             except FileNotFoundError:
-                return [f"❌ Not found: `{path}`"] + [gr.update() for _ in all_inputs]
+                return [f"❌ Not found: `{path}`", gr.update(), gr.update()] + [gr.update() for _ in all_inputs]
             except Exception as e:
-                return [f"❌ Load failed: {e}"] + [gr.update() for _ in all_inputs]
+                return [f"❌ Load failed: {e}", gr.update(), gr.update()] + [gr.update() for _ in all_inputs]
 
-        apply_preset_btn.click(fn=apply_preset, inputs=[preset_dd], outputs=all_inputs)
+        apply_preset_btn.click(fn=apply_preset, inputs=[preset_dd], outputs=[config_state, config_display] + all_inputs)
         # Selecting a preset in the dropdown should sync immediately;
         # the explicit "Load Preset" button stays as a re-apply affordance.
-        preset_dd.change(fn=apply_preset, inputs=[preset_dd], outputs=all_inputs)
-        reset_btn.click(fn=reset_minimind, outputs=all_inputs)
+        preset_dd.change(fn=apply_preset, inputs=[preset_dd], outputs=[config_state, config_display] + all_inputs)
+        reset_btn.click(fn=reset_minimind, outputs=[config_state, config_display] + all_inputs)
         save_btn.click(fn=do_save, inputs=[config_state, cfg_path], outputs=[save_status])
-        load_btn.click(fn=do_load, inputs=[cfg_path], outputs=[save_status] + all_inputs)
+        load_btn.click(fn=do_load, inputs=[cfg_path], outputs=[save_status, config_state, config_display] + all_inputs)
 
-        _initial_cfg = ChronosConfig().__dict__
-        total_box.value, active_box.value, vram_box.value, ssd_box.value, kv_box.value, tps_box.value = _estimate(_initial_cfg)
+        total_box.value, active_box.value, vram_box.value, ssd_box.value, kv_box.value, tps_box.value = _estimate(initial_cfg)
 
     return config_state, all_inputs, save_dir

ui/tabs/train_tab.py

@@ -2,6 +2,7 @@
 ui/tabs/train_tab.py — Full training loop: Pretrain / SFT / DPO / ORPO / GRPO / Distill
 """
 import os
+import glob
 import time
 import threading
 import queue
@@ -47,6 +48,18 @@
     "grpo": "./tests/fixtures/tiny_grpo.jsonl",
     "distill": "./tests/fixtures/tiny_sft.jsonl",
 }
+CHECKPOINT_TOPOLOGY_KEYS = (
+    "hidden_size",
+    "num_hidden_layers",
+    "num_experts",
+    "moe_intermediate_size",
+    "vocab_size",
+    "lookahead_steps",
+    "kv_latent_dim",
+    "rope_dim",
+    "num_attention_heads",
+    "num_key_value_heads",
+)
 STAGE_HELP_TEXT = {
     "pretrain": {
         "zh-Hans": "从通用语料继续预训练。`init_weight` 可留空；若存在同名 checkpoint，将按当前拓扑尝试恢复。",
@@ -394,7 +407,36 @@ def _sniff_checkpoint(path: str) -> dict:
             total = int(lookahead_proj.shape[0])
             n_exp = int(out["num_experts"])
             if n_exp > 0 and total % n_exp == 0:
-                out["lookahead_steps"] = total // n_exp
+                # The router predicts current-token routing plus K future
+                # steps, so the saved output rows are (lookahead_steps + 1)
+                # * num_experts. Do not report the +1 as user topology.
+                out["lookahead_steps"] = max(0, total // n_exp - 1)
+
+        qnope = sd.get("model.layers.0.self_attn.q_nope_proj.weight")
+        qrope = sd.get("model.layers.0.self_attn.q_rope_proj.weight")
+        kvdown = sd.get("model.layers.0.self_attn.kv_down_proj.weight")
+        vproj = sd.get("model.layers.0.self_attn.v_proj.weight")
+        if qnope is not None and qrope is not None and out.get("hidden_size"):
+            candidates = [
+                n for n in range(1, out["hidden_size"] + 1)
+                if out["hidden_size"] % n == 0
+                and int(qrope.shape[0]) % n == 0
+                and int(qnope.shape[0]) % n == 0
+            ]
+            if candidates:
+                preferred = 8 if 8 in candidates else max(candidates)
+                head_dim = out["hidden_size"] // preferred
+                rope_dim = int(qrope.shape[0]) // preferred
+                nope_dim = int(qnope.shape[0]) // preferred
+                if rope_dim + nope_dim == head_dim:
+                    out["num_attention_heads"] = preferred
+                    out["rope_dim"] = rope_dim
+        if kvdown is not None:
+            out["kv_latent_dim"] = int(kvdown.shape[0])
+        if vproj is not None and out.get("kv_latent_dim"):
+            head_dim = out.get("hidden_size", 0) // max(out.get("num_attention_heads", 8), 1)
+            if head_dim > 0:
+                out["num_key_value_heads"] = max(1, int(vproj.shape[0]) // head_dim)
 
         return out
     except Exception:
@@ -477,16 +519,42 @@ def _default_init_path(self, save_dir: str, mode: str, hidden_size: int) -> str:
             return ""
         return os.path.join(save_dir, f"{upstream}_{hidden_size}_moe.pth")
 
+    def _resolve_default_init_path(self, save_dir: str, mode: str, hidden_size: int) -> str:
+        exact = self._default_init_path(save_dir, mode, hidden_size)
+        if os.path.exists(exact):
+            return exact
+        upstream = STAGE_DEFAULT_INIT.get(mode)
+        if not upstream:
+            return exact
+        candidates = [
+            p for p in glob.glob(os.path.join(save_dir, f"{upstream}_*_moe.pth"))
+            if os.path.isfile(p)
+        ]
+        if not candidates:
+            return exact
+        return max(candidates, key=os.path.getmtime)
+
     def _topology_mismatches(self, sniffed: dict, model_cfg_kwargs: dict) -> list[str]:
         mismatches = []
-        for k in [
-            "hidden_size", "num_hidden_layers", "num_experts",
-            "moe_intermediate_size", "vocab_size", "lookahead_steps",
-        ]:
+        for k in CHECKPOINT_TOPOLOGY_KEYS:
             if k in sniffed and k in model_cfg_kwargs and int(sniffed[k]) != int(model_cfg_kwargs[k]):
                 mismatches.append(f"{k}: ckpt={sniffed[k]} != ui={model_cfg_kwargs[k]}")
         return mismatches
 
+    def _adopt_checkpoint_topology(self, model_cfg_kwargs: dict, sniffed: dict) -> dict:
+        adopted = {}
+        for key in CHECKPOINT_TOPOLOGY_KEYS:
+            if key not in sniffed:
+                continue
+            current = model_cfg_kwargs.get(key)
+            value = sniffed[key]
+            if current is None or int(current) != int(value):
+                adopted[key] = value
+            model_cfg_kwargs[key] = value
+        if "moe_intermediate_size" in sniffed:
+            model_cfg_kwargs["intermediate_size"] = sniffed["moe_intermediate_size"]
+        return adopted
+
     def _build_stage_args(self, cfg: dict, mode: str, save_dir: str, hidden_size: int):
         reward_spec = (cfg.get("reward_spec") or "toy").strip() or "toy"
         teacher_path = (cfg.get("teacher_path") or "").strip()
@@ -681,6 +749,24 @@ def _run(self, cfg: dict, mode: str):
                 if opt_key in AUTO_SENTINEL_KEYS and val == 0:
                     continue
                 model_cfg_kwargs[opt_key] = val
+            save_dir = cfg.get("save_dir", "./out")
+            init_weight = (cfg.get("init_weight") or "").strip()
+            load_path = ""
+            init_sniffed = {}
+            if mode != "pretrain":
+                load_path = init_weight or self._resolve_default_init_path(
+                    save_dir, mode, int(model_cfg_kwargs["hidden_size"])
+                )
+                if os.path.exists(load_path):
+                    init_sniffed = _sniff_checkpoint(load_path)
+                    adopted = self._adopt_checkpoint_topology(model_cfg_kwargs, init_sniffed)
+                    if adopted:
+                        summary = ", ".join(f"{k}={v}" for k, v in adopted.items())
+                        self._put(
+                            f"[{mode.upper()}] Adopted topology from init checkpoint "
+                            f"{load_path}: {summary}"
+                        )
+
             model_cfg = ChronosConfig(**model_cfg_kwargs)
             model = ChronosForCausalLM(model_cfg)
             params_m = sum(p.numel() for p in model.parameters()) / 1e6
@@ -689,14 +775,12 @@ def _run(self, cfg: dict, mode: str):
                       f"E={model_cfg.num_experts}, ffn={model_cfg.intermediate_size}, "
                       f"vocab={model_cfg.vocab_size})")
 
-            save_dir = cfg.get("save_dir", "./out")
             ckp_path = self._stage_checkpoint_path(save_dir, mode, model_cfg.hidden_size)
 
             # Init-weight resolution:
             #   - pretrain: optionally resume from its own checkpoint if present
             #   - other stages: require an upstream weight (explicit init_weight
             #     or the stage's default predecessor checkpoint).
-            init_weight = (cfg.get("init_weight") or "").strip()
             if mode == "pretrain":
                 resume_path = init_weight or ckp_path
                 if os.path.exists(resume_path):
@@ -714,7 +798,6 @@ def _run(self, cfg: dict, mode: str):
                 else:
                     self._put("Pretraining from random init")
             else:
-                load_path = init_weight or self._default_init_path(save_dir, mode, model_cfg.hidden_size)
                 if not os.path.exists(load_path):
                     raise FileNotFoundError(
                         f"[{mode.upper()}] requires an upstream checkpoint to initialize from. "
@@ -723,7 +806,7 @@ def _run(self, cfg: dict, mode: str):
                         f"or run the prior stage first."
                     )
 
-                sniffed = _sniff_checkpoint(load_path)
+                sniffed = init_sniffed or _sniff_checkpoint(load_path)
                 mismatch_hints = self._topology_mismatches(sniffed, model_cfg_kwargs)
                 if mismatch_hints:
                     raise RuntimeError(