Add files via upload

Author: FonaTech

README.md

@@ -195,6 +195,21 @@ Training dtype and resource policy:
 - Native MLX training pushes UI logs, scalar readouts, and chart points every `log_interval` steps, and Web UI Stop is checked at each batch boundary.
 - The Web UI writes a warning-only `<checkpoint>.verify.json` after each stage. It checks no-mask vs all-available MoE parity and, on Apple Silicon, MLX prefill logits against the PyTorch CPU baseline.
 
+Native MLX training is a separate Apple Silicon backend, not `torch.to("mlx")`.
+The six-stage MLX trainer mirrors the PyTorch loss stack in `chronos.mlx.*`:
+masked CE, DPO/ORPO/GRPO/distillation losses, load balance, temporal locality,
+lookahead soft-target supervision, lookahead top-k hit loss, and router KL
+anchor. Numerically sensitive pieces run in FP32 even when model weights use
+BF16/FP16, which is why `auto` prefers BF16 on MLX/MPS: FP16 has too little
+exponent range for router softmax, CE, and Adam moments on small unstable
+training runs.
+
+The MLX Web UI trainer reports the same fields as CPU/MPS training:
+step/loss/steps-per-second/ETA, checkpoint-save events, stop events, and
+verify results. Stop is cooperative at batch boundaries and save intervals
+write a valid `.pth` plus sibling `.config.json`, so MLX stages can resume or
+feed the PyTorch/export pipeline.
+
 The full six-stage comparison harness lives in `tools/compare_minimind_chronos_v3.py`.
 
 ---
@@ -231,6 +246,32 @@ d.describe()    # human-readable capability summary
 
 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.
 
+### MLX lazy/offload runtime
+
+MLX uses Apple unified memory, so Chronos treats "VRAM" and "RAM" as logical
+tiers rather than physical buses. The native lazy runtime still enforces the
+same offload contract as CUDA/MPS/CPU:
+
+- **Hot slots** hold only the execution-budget experts materialized as live
+  MLX modules.
+- **Warm cache** holds a bounded prediction buffer loaded from `.ctsr`
+  safetensors, never a hidden copy of every expert.
+- **Cold storage** is the checkpoint/export reader or per-expert cluster cache.
+  In lazy mode, live experts are replaced by placeholders after cache creation;
+  Chronos does not keep a `_saved_live` full-expert repair cache.
+- Lookahead predictions are queued into the warm cache and only promoted when
+  ready. A true miss synchronously materializes the selected expert only; it
+  does not full-load the model.
+- MLX attention dynamically grows RoPE lookup tables past
+  `max_position_embeddings`, so long prompts plus decode do not fail at token
+  257/513/etc.
+
+Inference compare/sweep reports real hot/warm counts, resident hit rate,
+prediction hit rate, sync SSD loads, MLX active/cache/peak memory, process RSS,
+prefill/decode time, and tokens/sec. A lazy run is not considered
+offload-ready unless deterministic exact-lazy output matches full-DRAM output
+and fallback weight stays zero.
+
 ---
 
 ## Hugging Face and vLLM compatibility (M5)
@@ -361,7 +402,7 @@ Tabs included:
 - `6-Stage Pipeline` with per-stage dataset paths
 - `Inference`
 - `Export` for FP16/Q8_0 safetensors and GGUF deployment artifacts
-- `Benchmark` with Markdown table + bar plot
+- `Benchmark` with Markdown table + comparison plots
 - `Auto-Tune` with persistent logs and one-click `Apply Best -> Config`
 - `IO Monitor`
 

README_zh.md

@@ -189,6 +189,19 @@ flowchart LR
 - MLX 原生训练会按 `log_interval` 同步 UI 日志、标量读数和图表点；Web UI Stop 会在每个 batch 边界检查并停止。
 - Web UI 每阶段保存后会写出 warning-only 的 `<checkpoint>.verify.json`，检查 no-mask vs all-available MoE 等价性，并在 Apple Silicon 上对比 MLX prefill logits 与 PyTorch CPU baseline。
 
+MLX 训练是 Apple Silicon 原生后端，不是 `torch.to("mlx")`。六阶段 MLX
+trainer 在 `chronos.mlx.*` 中复刻 PyTorch 训练栈：masked CE、DPO/ORPO/GRPO/
+蒸馏损失、load balance、temporal locality、lookahead soft-target 监督、
+lookahead top-k 命中损失、router KL anchor。即使权重使用 BF16/FP16，路由
+softmax、CE、Adam moments 等数值敏感路径仍以 FP32 计算；因此 `auto` 在
+MLX/MPS 上默认优先 BF16。FP16 的指数范围太小，在小模型/小 batch/路由分布
+剧烈波动时更容易把 loss 推成 NaN。
+
+Web UI 的 MLX 训练日志与 CPU/MPS 对齐：step、loss、steps/s、ETA、
+checkpoint 保存事件、stop 事件、verify 结果都会同步到界面。Stop 在 batch
+边界协作式生效；定期保存会写出可继续训练的 `.pth` 与同名 `.config.json`，
+因此 MLX 阶段可以继续接 PyTorch、Export、Diagnose 链路。
+
 完整 6 阶段端到端对比见 `tools/compare_minimind_chronos_v3.py`。
 
 ---
@@ -224,6 +237,28 @@ d.describe()       # 人类可读的能力总览
 
 诚实声明：上游 PyTorch 没有 OpenCL 后端、Vulkan 也仅在自定义构建中可用。Chronos 提供 dispatcher 接缝，使第三方插件无需改核心代码即可接入。
 
+### MLX lazy/offload 运行时
+
+MLX 使用 Apple 统一内存，因此 Chronos 在 MLX 上的 “VRAM/RAM” 是逻辑分层，
+不是 PCIe 式物理分层。但 lazy runtime 仍严格遵守与 CUDA/MPS/CPU 一致的
+offload contract：
+
+- **Hot slots**：只保存执行预算内、已经 materialize 成 MLX live module 的专家。
+- **Warm cache**：只保存有界预测缓冲，从 `.ctsr` safetensors 中按专家/簇加载；
+  不允许暗中保留全量专家副本。
+- **Cold storage**：checkpoint/export reader 或 per-expert cluster cache。Lazy
+  模式创建 cache 后会把 live experts 替换为 placeholder；不会保留 `_saved_live`
+  这种全量专家修复缓存。
+- Lookahead 预测只进入 warm cache 队列，只有 ready 的专家才会 promote。真实 miss
+  只同步 materialize 当前被选中的专家，不会 full-load 整个模型。
+- MLX attention 会按需扩容 RoPE lookup table，因此 prompt 长度加 decode 超过
+  `max_position_embeddings` 时不会在第 257/513 等 token 崩溃。
+
+Inference compare/sweep 会报告真实 hot/warm 数量、resident hit、prediction hit、
+sync SSD loads、MLX active/cache/peak memory、进程 RSS、prefill/decode 时间与
+tokens/s。只有确定性 exact-lazy 输出与 full-DRAM 一致、且 fallback weight 为 0，
+才应该把该 checkpoint 标记为 offload-ready。
+
 ---
 
 ## HuggingFace / vLLM 兼容（M5）
@@ -344,7 +379,7 @@ chronos-ui
 python chronos_app.py
 ```
 
-包含：⚙️ Config（含右侧实时参数估算面板，合并了 Designer）/ 🏋️ Train（拥有 data_path）/ 🧪 6-Stage Pipeline（每阶段独立数据路径）/ 💬 Inference（含懒加载 vs 全量 DRAM 对比）/ 📦 Export（FP16/Q8_0 safetensors 与 GGUF）/ 📊 Benchmark（Markdown 表 + BarPlot）/ 🔬 Auto-Tune（持久化日志 + 一键 Apply Best → Config）/ 📡 IO Monitor。i18n 支持 zh-Hans / zh-Hant / en / ja。
+包含：⚙️ Config（含右侧实时参数估算面板，合并了 Designer）/ 🏋️ Train（拥有 data_path）/ 🧪 6-Stage Pipeline（每阶段独立数据路径）/ 💬 Inference（含懒加载 vs 全量 DRAM 对比）/ 📦 Export（FP16/Q8_0 safetensors 与 GGUF）/ 📊 Benchmark（Markdown 表 + 对比图）/ 🔬 Auto-Tune（持久化日志 + 一键 Apply Best → Config）/ 📡 IO Monitor。i18n 支持 zh-Hans / zh-Hant / en / ja。
 
 ### 部署导出
 

chronos/mlx/attention.py

@@ -59,7 +59,39 @@ def _apply_rope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
     return (x * c + rotated_half * s).astype(x.dtype)
 
 
-class MLAAttentionMLX(nn.Module):
+class _RopeCacheMixin:
+    def _init_rope_cache(self, config, dim: int) -> None:
+        rope_len = int(getattr(config, "max_position_embeddings", 4096))
+        object.__setattr__(self, "_rope_dim", int(dim))
+        object.__setattr__(self, "_rope_base", float(getattr(config, "rope_theta", 1_000_000.0)))
+        object.__setattr__(self, "_rope_scaling", getattr(config, "rope_scaling", None))
+        cos, sin = _rope_freqs(
+            int(dim),
+            max_len=rope_len,
+            base=self._rope_base,
+            rope_scaling=self._rope_scaling,
+        )
+        object.__setattr__(self, "_cos", cos)
+        object.__setattr__(self, "_sin", sin)
+
+    def _ensure_rope_cache(self, needed_len: int) -> None:
+        current_len = int(self._cos.shape[0])
+        if int(needed_len) <= current_len:
+            return
+        # Grow geometrically to avoid reallocating every decode step past
+        # max_position_embeddings.
+        new_len = max(int(needed_len), current_len * 2)
+        cos, sin = _rope_freqs(
+            int(self._rope_dim),
+            max_len=new_len,
+            base=float(self._rope_base),
+            rope_scaling=self._rope_scaling,
+        )
+        object.__setattr__(self, "_cos", cos)
+        object.__setattr__(self, "_sin", sin)
+
+
+class MLAAttentionMLX(_RopeCacheMixin, nn.Module):
     """
     Multi-head Latent Attention for MLX.
 
@@ -94,17 +126,7 @@ def __init__(self, config):
         self.v_proj       = nn.Linear(self.kv_latent_dim, self.n_kv_heads * self.head_dim, bias=False)
         self.o_proj       = nn.Linear(self.n_heads * self.head_dim, H, bias=False)
 
-        rope_len = int(getattr(config, "max_position_embeddings", 4096))
-        cos, sin = _rope_freqs(
-            self.head_dim,
-            max_len=rope_len,
-            base=float(getattr(config, "rope_theta", 1_000_000.0)),
-            rope_scaling=getattr(config, "rope_scaling", None),
-        )
-        # Bypass mlx.nn.Module.__setattr__; these are runtime lookup tables,
-        # not trainable checkpoint parameters.
-        object.__setattr__(self, "_cos", cos)
-        object.__setattr__(self, "_sin", sin)
+        self._init_rope_cache(config, self.head_dim)
 
     def __call__(
         self,
@@ -134,6 +156,7 @@ def __call__(
             k_rope_past = None
 
         T_full = c_kv_full.shape[1]
+        self._ensure_rope_cache(T_full)
 
         # ── Keys / Values (recomputed from cache) ─────────────────
         k_nope = self.k_nope_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.nope_dim)
@@ -172,7 +195,7 @@ def __call__(
         return self.o_proj(out), new_cache
 
 
-class SlidingWindowAttentionMLX(nn.Module):
+class SlidingWindowAttentionMLX(_RopeCacheMixin, nn.Module):
     """
     Standard multi-head attention with a sliding KV window for MLX.
     KV cache is capped at window_size tokens — O(1) memory per step.
@@ -193,15 +216,7 @@ def __init__(self, config):
         self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
         self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
 
-        rope_len = int(getattr(config, "max_position_embeddings", 4096))
-        cos, sin = _rope_freqs(
-            self.head_dim,
-            max_len=rope_len,
-            base=float(getattr(config, "rope_theta", 1_000_000.0)),
-            rope_scaling=getattr(config, "rope_scaling", None),
-        )
-        object.__setattr__(self, "_cos", cos)
-        object.__setattr__(self, "_sin", sin)
+        self._init_rope_cache(config, self.head_dim)
 
     def __call__(
         self,
@@ -216,6 +231,7 @@ def __call__(
         v_new = self.v_proj(x).reshape(B, S, self.n_kv_heads, self.head_dim)
 
         past_len = cache[0].shape[1] if cache is not None else 0
+        self._ensure_rope_cache(past_len + S)
         q = _apply_rope(q, self._cos[past_len:past_len + S], self._sin[past_len:past_len + S])
         k_new = _apply_rope(k_new, self._cos[past_len:past_len + S], self._sin[past_len:past_len + S])