0.3.1a3: segment clustering, MLX multimodal (all stages), tool calls, cache limit

CHANGELOG.md

@@ -2,6 +2,63 @@
 
 ## [Unreleased]
 
+## [0.3.1a3] — 2026-11-08 (beta)
+
+This release fixes the router masks and the lazy-expert decode. It also adds a document-level cluster objective, MLX multimodal training in all six stages, tool-call SFT records and tool-call parsing in `chronos serve`, a kernel-level MLX cache limit and progress-based ETAs.
+
+### Added
+- `cluster_locality_mode="segment"`: the cluster objective treats each document as one unit.
+  - Documents are found from the tokens: `<|endoftext|>` separators, and long documents are cut every `cluster_locality_window` tokens. No labels or domain tags are used.
+  - The loss is `1 − I(document; cluster) / log C`, pooled over real tokens and all virtual layers.
+  - The default stays `"window"`, the 0.3.1a2 objective.
+  - Torch and MLX give the same value and gradient. The option is on the Web UI config page (advanced) in four languages.
+- MLX multimodal training. `chronos.mlx.multimodal` ports the image patch encoder, the audio conv encoder and the gated cross-attention fusion.
+  - Pre-encoded `*_embeds` inputs and the per-modality and per-token masks work as in Torch.
+  - Logits match Torch within 2e-4 in FP32 tests.
+  - MLX checkpoints save and load the fusion weights, and `freeze_modality_encoders` works on MLX.
+  - `run_mlx_stage` takes media in all six stages: pretrain, SFT, DPO, ORPO, GRPO and distill.
+  - DPO and ORPO: the chosen and rejected answers condition on the prompt's media.
+  - GRPO: every rollout conditions on the prompt's media. The rollout encodes the media once and reuses it at every decode step, with identical logits. The policy update re-encodes the media, so the encoders get gradients.
+  - Distill: the student gets the media. A text-only teacher is accepted, as on PyTorch.
+- Multimodal preference and GRPO data.
+  - `MultimodalStreamingDPODataset` and `preference_collate` read `chosen`/`rejected` records that carry `image`, `audio` or `*_embeds` fields, the same fields SFT uses.
+  - `load_grpo_prompt_records` keeps each prompt's record so the media can be resolved per step.
+  - `record_media` is the shared resolver.
+  - The Web UI selects these automatically when the dataset contains media.
+  - The PyTorch DPO, ORPO and GRPO trainers are text-only. With media, `auto` picks MLX, and an explicit PyTorch backend raises an error instead of dropping the media.
+- Tool calls in the chat template.
+  - `tools` render as one `<tools>` system block.
+  - The assistant's `tool_calls` render as `<tool_call>{json}</tool_call>` lines, and `tool` turns carry the results.
+  - Conversations without tools render exactly as before.
+- `chronos.data.chat_records` provides `render_chat`, `encode_chat` and `write_sft_records`.
+  - Supervised tokens: the assistant text, its tool calls and `<|im_end|>`.
+  - Context tokens: system and user turns, tool definitions and tool results.
+  - Records are written as `.bin/.idx/.plen/.lmask`.
+- `SFTRecords` reads `.lmask` (one uint8 per token). When present it replaces the `.plen` prompt rule. Older readers still use `.plen`.
+- `chronos serve` supports tool calls end to end.
+  - It passes `tools` and the assistant's `tool_calls` from `/api/chat` and `/v1/chat/completions` into the template.
+  - When the request has `tools`, it parses `<tool_call>` blocks in the output into structured calls. `/api/chat` returns Ollama `message.tool_calls`. `/v1/chat/completions` returns OpenAI `tool_calls` with string arguments and `finish_reason: "tool_calls"`.
+  - Streaming parses the calls as they are generated, and no client sees part of a tag.
+  - An invalid block stays in the content. Requests without `tools` are unchanged.
+  - `chronos.serving.tool_calls` holds the parser.
+- `chronos.mlx.memory`: every MLX model sets the MLX buffer-cache limit before it allocates weights.
+  - The limit is `min(RAM/8, 16 GiB)`.
+  - A recipe (`extra.mlx_cache_limit_gb`) or `CHRONOS_MLX_CACHE_LIMIT_GB` can lower it. Larger values are clamped to the limit.
+  - `LongRun` logs the limit and records it in `runtime.json`.
+  - Only reclaimable cache is bounded. Weights, activations and gradients are not.
+- Monitor and Web UI ETAs follow measured progress.
+  - A phase ends at the first of three points: the wall deadline, the step total, or `target_passes`.
+  - The pass rate comes from row timestamps over the last 30 minutes, so evaluation and checkpoint time count.
+  - Each phase shows its estimated end and which of the three points sets it.
+  - Later phases also get the current phase's unused time.
+  - A stopped or failed run shows its remaining work and its state, but no end time.
+
+### Fixed
+- Router auxiliary losses and routing telemetry counted right-padding in SFT batches. Balance, entropy, lookahead and cluster terms now use only real tokens. Packed pretraining is unchanged.
+- Training telemetry now sums every micro-batch of a step. Before, it sampled only the last one.
+- `SFTRecords` with `.lmask`: under NumPy 2, the -100 context labels of uint16 shards wrapped to 65436. Labels are now built in int32.
+- MLX lazy-expert decode used the per-expert loop, which sums in a different order from the full-memory grouped kernel. In BF16 on the GPU this flipped near-tied greedy tokens: agreement was 31.25 % at every budget. The lazy path now runs the same grouped kernel over the live experts. Lazy and full greedy now agree at every budget (measured: 0.278, 0.5, 0.778 and 1.0 of experts, planned and window residency).
+
 ## [0.3.1a2] — 2026-11-07 (beta)
 
 Kernel/extension split: the `chronos` kernel trains general models and holds no translation-specific code. The en↔zh demo moved to `demos/translation` (`chronos-translate`), which imports the kernel and is not part of the `Project_Chronos` wheel. The extension is not published on PyPI; install it from a checkout with `pip install -e demos/translation` (see `demos/translation/README.md`).

README.md

@@ -286,6 +286,14 @@ 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.
 
+Memory, media and tools on MLX:
+
+- Every MLX model caps the MLX buffer cache at `min(RAM/8, 16 GiB)` before it allocates weights (`chronos.mlx.memory`). A recipe's `extra.mlx_cache_limit_gb` or `CHRONOS_MLX_CACHE_LIMIT_GB` may lower the cap, never raise it. The cap bounds reclaimable cache only; weights, activations and gradients are not limited.
+- Image and audio encoders and the gated cross-attention fusion run on MLX in all six stages, with the same weights and numerics as PyTorch (`chronos.mlx.multimodal`). In DPO and ORPO, both answers of a pair condition on the prompt's media. In GRPO, every rollout conditions on it. Preference and GRPO records carry media in the same `image`, `audio` and `*_embeds` fields as SFT. The PyTorch DPO, ORPO and GRPO trainers are text-only, so `auto` routes media data in those stages to MLX.
+- Tool calls:
+  - Training: the chat template renders `tools`, assistant `tool_calls` and `tool` results. `chronos.data.chat_records.encode_chat` gives each token a loss flag (assistant text and calls are learned; everything else is context), and `write_sft_records` writes `.lmask` shards that `SFTRecords` reads.
+  - Serving: when a request carries `tools`, `chronos serve` parses the model's `<tool_call>` output into Ollama `message.tool_calls` or OpenAI `tool_calls`, in streaming and non-streaming responses.
+
 The full six-stage comparison harness lives in `tools/compare_minimind_chronos_v3.py`.
 
 ---
@@ -565,7 +573,7 @@ Bundles belong to the translation extension (`demos/translation`). Build the two
 
 ```bash
 chronos-translate bundle build --run-dir runs/enzh_50m_v1 --which best --pack \
-    --wheel dist/project_chronos-0.3.1a2-py3-none-any.whl --wheel dist/chronos_translate-0.3.1a2-py3-none-any.whl
+    --wheel dist/project_chronos-0.3.1a3-py3-none-any.whl --wheel dist/chronos_translate-0.3.1a3-py3-none-any.whl
 ```
 
 | Bundle path | Content |

README_zh.md

@@ -296,6 +296,14 @@ checkpoint 保存事件、stop 事件、verify 结果都会同步到界面。Sto
 边界协作式生效；定期保存会写出可继续训练的 `.pth` 与同名 `.config.json`，
 因此 MLX 阶段可以继续接 PyTorch、Export、Diagnose 链路。
 
+MLX 上的内存、多模态与工具调用：
+
+- 每个 MLX 模型在分配权重前把 MLX 缓冲区缓存上限设为 `min(内存/8, 16 GiB)`（`chronos.mlx.memory`）。配方的 `extra.mlx_cache_limit_gb` 或 `CHRONOS_MLX_CACHE_LIMIT_GB` 只能调低，不能调高。上限只约束可回收缓存，不限制权重、激活和梯度。
+- 图像、音频编码器和门控交叉注意力融合在 MLX 上覆盖全部 6 个阶段，权重与数值和 PyTorch 一致（`chronos.mlx.multimodal`）。DPO/ORPO 的一对回答以 prompt 的媒体为共同条件，GRPO 的每条 rollout 也以该媒体为条件。偏好数据和 GRPO 数据的媒体字段与 SFT 相同：`image`、`audio`、`*_embeds`。PyTorch 的 DPO/ORPO/GRPO 只支持文本，所以这三个阶段遇到带媒体的数据时，`auto` 会选 MLX。
+- 工具调用：
+  - 训练：chat 模板支持 `tools`、assistant 的 `tool_calls` 和 `tool` 结果。`chronos.data.chat_records.encode_chat` 为每个 token 标出是否计 loss（assistant 文本和调用参与学习，其余为上下文），`write_sft_records` 写出 `SFTRecords` 可读的 `.lmask` 分片。
+  - 推理：请求带 `tools` 时，`chronos serve` 把模型输出的 `<tool_call>` 解析成 Ollama 的 `message.tool_calls` 或 OpenAI 的 `tool_calls`，流式与非流式都支持。
+
 完整 6 阶段端到端对比见 `tools/compare_minimind_chronos_v3.py`。
 
 ---
@@ -569,7 +577,7 @@ OLLAMA_HOST=127.0.0.1:11435 ollama run <模型名> "你好"
 
 ```bash
 chronos-translate bundle build --run-dir runs/enzh_50m_v1 --which best --pack \
-    --wheel dist/project_chronos-0.3.1a2-py3-none-any.whl --wheel dist/chronos_translate-0.3.1a2-py3-none-any.whl
+    --wheel dist/project_chronos-0.3.1a3-py3-none-any.whl --wheel dist/chronos_translate-0.3.1a3-py3-none-any.whl
 ```
 
 | 包内路径 | 内容 |

chronos/__init__.py

@@ -1,5 +1,5 @@
 # Project Chronos: On-device low-latency lookahead dual-layer MoE inference
-__version__ = "0.3.1a2"
+__version__ = "0.3.1a3"
 
 # Bootstrap minimind dependency before any submodule imports
 from chronos.deps import ensure_minimind as _ensure_minimind

chronos/config_schema.py

@@ -37,7 +37,7 @@
     "lookahead_steps", "layer_aware_lookahead", "layer_aware_cache",
     "lambda_balance", "lambda_temporal", "lambda_lookahead",
     "lambda_lookahead_topk", "lambda_lookahead_union", "lambda_router_locality",
-    "lambda_cluster_locality", "cluster_locality_clusters", "cluster_locality_window",
+    "lambda_cluster_locality", "cluster_locality_clusters", "cluster_locality_window", "cluster_locality_mode",
     "lambda_router_anchor", "fallback_mask_prob", "offload_miss_policy_default",
     "use_hybrid_attention", "kv_latent_dim", "rope_dim", "rope_layout", "sliding_window_size",
     "kv_cache_dtype", "kv_quant_group_size", "mla_decode", "mla_cache_rope", "router_scope",

chronos/data/chat_records.py

@@ -0,0 +1,142 @@
+"""Multi-turn chat records (tool calls, tool results) for SFT shards.
+
+``render_chat`` reproduces ``chronos.tokenizer.model.CHAT_TEMPLATE`` in
+Python and returns the text as ``(segment, supervised)`` pieces, so each
+token's loss flag is known without guessing from offsets:
+
+* learned: an assistant turn's content, its ``<tool_call>{json}</tool_call>``
+  lines and the closing ``<|im_end|>\\n``;
+* context: system/tool definitions, user turns, ``tool`` results and every
+  ``<|im_start|>role\\n`` header (the generation prompt the model is given).
+
+``encode_chat`` tokenizes the pieces one by one and concatenates the ids, the
+same way the translation shards join prompt and reply encodings, so every
+piece boundary is a token boundary.  ``write_sft_records`` writes the
+``.bin/.idx/.plen/.lmask`` layout read by ``chronos.data.shards.SFTRecords``.
+
+A record whose last turn is not an assistant turn ends with context tokens;
+the trailing context is dropped (it carries no loss and would read as
+padding to the router's real-token mask, which ends at the last supervised
+token).
+"""
+from __future__ import annotations
+
+import json
+import os
+
+import numpy as np
+
+ROLES = ("system", "user", "assistant", "tool")
+
+
+def _tool_json(value) -> str:
+    # ``tojson`` as transformers overrides it for chat templates: key order
+    # kept, no ASCII/HTML escaping, default ", " / ": " separators.
+    return json.dumps(value, ensure_ascii=False)
+
+
+def _call_parts(call: dict) -> dict:
+    fn = call.get("function") if isinstance(call, dict) and call.get("function") else call
+    if not isinstance(fn, dict) or "name" not in fn:
+        raise ValueError(f"tool call needs a name: {call!r}")
+    args = fn.get("arguments", {})
+    return {"name": fn["name"], "arguments": args}
+
+
+def render_chat(messages, tools=None, *, add_generation_prompt: bool = False) -> list[tuple[str, bool]]:
+    """``[(text, supervised)]`` pieces whose concatenation is the template text."""
+    pieces: list[tuple[str, bool]] = []
+    messages = list(messages or [])
+    for m in messages:
+        if m.get("role") not in ROLES:
+            raise ValueError(f"unknown chat role {m.get('role')!r}; expected one of {ROLES}")
+    skip_first = bool(tools) and bool(messages) and messages[0]["role"] == "system"
+    if tools:
+        head = "<|im_start|>system\n"
+        if skip_first:
+            head += (messages[0].get("content") or "") + "\n\n"
+        head += "<tools>\n" + "".join(_tool_json(t) + "\n" for t in tools) + "</tools><|im_end|>\n"
+        pieces.append((head, False))
+    for i, m in enumerate(messages):
+        if skip_first and i == 0:
+            continue
+        role, text = m["role"], m.get("content") or ""
+        pieces.append((f"<|im_start|>{role}\n", False))
+        if role != "assistant":
+            pieces.append((text + "<|im_end|>\n", False))
+            continue
+        body = text
+        for j, call in enumerate(m.get("tool_calls") or []):
+            if text or j:
+                body += "\n"
+            body += "<tool_call>" + _tool_json(_call_parts(call)) + "</tool_call>"
+        pieces.append((body + "<|im_end|>\n", True))
+    if add_generation_prompt:
+        pieces.append(("<|im_start|>assistant\n", False))
+    return [(t, s) for t, s in pieces if t]
+
+
+def _ids(tokenizer, text: str) -> list[int]:
+    """Token ids from a ``ChronosTokenizer``, HF ``tokenizers.Tokenizer`` or
+    ``transformers`` tokenizer (no BOS/EOS added)."""
+    try:
+        out = tokenizer.encode(text, add_special_tokens=False)
+    except TypeError:
+        out = tokenizer.encode(text)
+    return list(getattr(out, "ids", out))
+
+
+def encode_chat(tokenizer, messages, tools=None) -> tuple[np.ndarray, np.ndarray]:
+    """``(ids int64, supervised uint8)`` for one conversation; trailing context
+    after the last assistant turn is dropped.  Raises if nothing is learned."""
+    ids: list[int] = []
+    mask: list[int] = []
+    for text, supervised in render_chat(messages, tools):
+        piece = _ids(tokenizer, text)
+        ids.extend(piece)
+        mask.extend([int(supervised)] * len(piece))
+    if 1 not in mask:
+        raise ValueError("conversation has no assistant turn to learn")
+    end = len(mask) - mask[::-1].index(1)
+    return np.asarray(ids[:end], np.int64), np.asarray(mask[:end], np.uint8)
+
+
+def write_sft_records(prefix: str, records, *, dtype=np.uint16, max_len: int | None = None) -> dict:
+    """Write ``(ids, supervised)`` records to ``<prefix>.{bin,idx,plen,lmask}``.
+
+    ``.plen`` is each record's first supervised offset (older readers apply
+    the prompt rule); ``.lmask`` carries the exact per-token flags.  Records
+    longer than ``max_len`` are skipped and counted.  Files are written to
+    ``.tmp`` names and renamed, so a reader never sees a partial set.
+    """
+    limit = np.iinfo(dtype).max
+    tokens, masks, offsets, plens, skipped = [], [], [0], [], 0
+    for ids, mask in records:
+        ids, mask = np.asarray(ids), np.asarray(mask, np.uint8)
+        if ids.shape != mask.shape or ids.ndim != 1:
+            raise ValueError("each record needs 1-D ids and a mask of the same length")
+        if max_len is not None and ids.size > max_len:
+            skipped += 1
+            continue
+        if ids.size and (int(ids.max()) > limit or int(ids.min()) < 0):
+            raise ValueError(f"token id outside {np.dtype(dtype).name}")
+        hits = np.flatnonzero(mask)
+        if not hits.size:
+            raise ValueError("record has no supervised token")
+        tokens.append(ids.astype(dtype)); masks.append(mask)
+        offsets.append(offsets[-1] + ids.size); plens.append(int(hits[0]))
+    arrays = {
+        ".bin": np.concatenate(tokens) if tokens else np.zeros(0, dtype),
+        ".idx": np.asarray(offsets, np.int64),
+        ".plen": np.asarray(plens, np.int32),
+        ".lmask": np.concatenate(masks) if masks else np.zeros(0, np.uint8),
+    }
+    os.makedirs(os.path.dirname(os.path.abspath(prefix)), exist_ok=True)
+    for suffix, arr in arrays.items():
+        tmp = prefix + suffix + ".tmp"
+        with open(tmp, "wb") as fh:
+            arr.tofile(fh)
+    for suffix in (".bin", ".idx", ".lmask", ".plen"):   # .plen last: readers key on it
+        os.replace(prefix + suffix + ".tmp", prefix + suffix)
+    return {"records": len(plens), "skipped_too_long": skipped, "tokens": int(offsets[-1]),
+            "supervised_tokens": int(arrays[".lmask"].sum())}