From a6477f7131ae7e458d182b93f7c7d556dc3d87d9 Mon Sep 17 00:00:00 2001 From: Fona <62028332+FonaTech@users.noreply.github.com> Date: Sun, 8 Nov 2026 13:23:46 +0800 Subject: [PATCH 1/3] 0.3.1a3: segment clustering, MLX multimodal, tool-call SFT, MLX cache limit - Router: padding mask, telemetry over all micro-batches, lazy decode uses the grouped kernel (fidelity 1.0 at every budget). - Cluster locality: cluster_locality_mode="segment" learns domain clusters per document from content alone (default stays "window"). - MLX: cache limit min(RAM/8, 16 GiB) applied in the kernel; recipe/env can only lower it. Multimodal fusion port with Torch parity and checkpoint round-trip; pretrain/SFT accept multimodal batches. - Tools: chat template renders /; chat_records writes per-token .lmask supervision read by SFTRecords; serve passes tools. - Monitor/WebUI: progress-based ETA per phase; stopped runs show their state instead of a clock ETA. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 42 +++ README.md | 8 +- README_zh.md | 8 +- chronos/__init__.py | 2 +- chronos/config_schema.py | 2 +- chronos/data/chat_records.py | 142 +++++++++++ chronos/data/shards.py | 29 ++- chronos/expert_clusters.py | 54 ++++ chronos/mlx/memory.py | 94 +++++++ chronos/mlx/model.py | 30 ++- chronos/mlx/moe.py | 59 +++++ chronos/mlx/multimodal.py | 284 +++++++++++++++++++++ chronos/mlx/training/inspect_checkpoint.py | 11 +- chronos/mlx/training/io.py | 4 + chronos/mlx/training/longrun.py | 123 +++++++-- chronos/mlx/training/telemetry.py | 40 ++- chronos/mlx/training/trainer.py | 221 +++++++++++++--- chronos/model/config.py | 5 + chronos/model/temporal_loss.py | 36 +++ chronos/monitor/snapshot.py | 104 +++++++- chronos/monitor/static/app.js | 13 +- chronos/serving/ollama_server.py | 30 ++- chronos/serving/runtime.py | 7 +- chronos/tokenizer/model.py | 21 +- chronos/trainer/chronos_trainer.py | 2 +- chronos/trainer/loss_mixin.py | 32 ++- chronos/trainer/sft_trainer.py | 2 +- chronos/webui/i18n/en.json | 9 + chronos/webui/i18n/ja.json | 9 + chronos/webui/i18n/zh-Hans.json | 9 + chronos/webui/i18n/zh-Hant.json | 9 + chronos/webui/schema.py | 2 + chronos/webui/static/pages_monitor.js | 14 +- chronos/workbench/training.py | 9 +- demos/translation/README.md | 10 +- demos/translation/README_zh.md | 8 +- demos/translation/pyproject.toml | 8 +- pyproject.toml | 2 +- tests/test_chat_records.py | 182 +++++++++++++ tests/test_mlx_cache_limit.py | 59 +++++ tests/test_mlx_multimodal.py | 153 +++++++++++ tests/test_moe_dispatch.py | 56 ++++ tests/test_monitor.py | 37 +++ tests/test_router_mask.py | 172 +++++++++++++ tests/test_segment_clusters.py | 190 ++++++++++++++ 45 files changed, 2205 insertions(+), 138 deletions(-) create mode 100644 chronos/data/chat_records.py create mode 100644 chronos/mlx/memory.py create mode 100644 chronos/mlx/multimodal.py create mode 100644 tests/test_chat_records.py create mode 100644 tests/test_mlx_cache_limit.py create mode 100644 tests/test_mlx_multimodal.py create mode 100644 tests/test_router_mask.py create mode 100644 tests/test_segment_clusters.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e8dce09..91b494c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,48 @@ ## [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, tool-call SFT records, 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` accepts `multimodal_collate` batches in pretrain and SFT. Preference and distillation stages with media still run on PyTorch. +- Tool calls in the chat template. + - `tools` render as one `` system block. + - The assistant's `tool_calls` render as `{json}` 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` passes `tools` and the assistant's `tool_calls` from `/api/chat` and `/v1/chat/completions` into the template. It does not parse `` output back into structured `tool_calls` yet. +- `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. +- 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`). diff --git a/README.md b/README.md index 7acc334..a3761fc 100644 --- a/README.md +++ b/README.md @@ -286,6 +286,12 @@ 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 pretrain and SFT, with the same weights and numerics as PyTorch (`chronos.mlx.multimodal`). DPO, ORPO, GRPO and distillation with media still use PyTorch. +- Tool-call SFT: 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. + The full six-stage comparison harness lives in `tools/compare_minimind_chronos_v3.py`. --- @@ -565,7 +571,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 | diff --git a/README_zh.md b/README_zh.md index d150e35..68d733f 100644 --- a/README_zh.md +++ b/README_zh.md @@ -296,6 +296,12 @@ 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 上支持预训练和 SFT，权重与数值和 PyTorch 一致（`chronos.mlx.multimodal`）。带媒体的 DPO、ORPO、GRPO 和蒸馏仍走 PyTorch。 +- 工具调用 SFT：chat 模板支持 `tools`、assistant 的 `tool_calls` 和 `tool` 结果。`chronos.data.chat_records.encode_chat` 为每个 token 标出是否计 loss（assistant 文本和调用参与学习，其余为上下文），`write_sft_records` 写出 `SFTRecords` 可读的 `.lmask` 分片。 + 完整 6 阶段端到端对比见 `tools/compare_minimind_chronos_v3.py`。 --- @@ -569,7 +575,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 ``` | 包内路径 | 内容 | diff --git a/chronos/__init__.py b/chronos/__init__.py index 2c5d388..a5591e4 100644 --- a/chronos/__init__.py +++ b/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 diff --git a/chronos/config_schema.py b/chronos/config_schema.py index 33d11f2..c4301bb 100644 --- a/chronos/config_schema.py +++ b/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", diff --git a/chronos/data/chat_records.py b/chronos/data/chat_records.py new file mode 100644 index 0000000..2963128 --- /dev/null +++ b/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 ``{json}`` + 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 += "\n" + "".join(_tool_json(t) + "\n" for t in 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_json(_call_parts(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 ``.{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())} diff --git a/chronos/data/shards.py b/chronos/data/shards.py index 0e586de..141c809 100644 --- a/chronos/data/shards.py +++ b/chronos/data/shards.py @@ -8,6 +8,12 @@ token array (``.bin``), int64 record offsets (``.idx``, n+1 entries) and int32 prompt lengths (``.plen``). Only tokens after the prompt are supervised, so the loss covers the assistant reply alone. +* optional ``sft/.../.lmask``: one uint8 per token of ``.bin`` (1 = + supervised). Multi-turn records (tool calls, tool results, several + assistant turns) need it: the assistant's text and ``tool_calls`` are + learned, user/system/tool-result tokens are context. When present it + replaces the ``.plen`` rule for that stem; ``.plen`` stays the first + supervised offset so older readers keep working. An optional ``shards_manifest.json`` may record ``"dtype": "uint16"|"uint32"``. """ @@ -67,12 +73,19 @@ class SFTRecords: def __init__(self, folder: str, dtype=np.uint16, seed: int = 0, pad_id: int = 0): stems = sorted(f[:-5] for f in os.listdir(folder) if f.endswith(".plen")) self.tokens, self.offsets, self.plens, self.index = [], [], [], [] + self.lmasks = [] for si, stem in enumerate(stems): base = os.path.join(folder, stem) tok = np.memmap(base + ".bin", dtype=dtype, mode="r") off = np.fromfile(base + ".idx", dtype=np.int64) plen = np.fromfile(base + ".plen", dtype=np.int32) + lmask = None + if os.path.exists(base + ".lmask"): + lmask = np.memmap(base + ".lmask", dtype=np.uint8, mode="r") + if lmask.size != tok.size: + raise ValueError(f"{base}.lmask has {lmask.size} entries for {tok.size} tokens") self.tokens.append(tok); self.offsets.append(off); self.plens.append(plen) + self.lmasks.append(lmask) self.index.append(np.stack([np.full(len(plen), si), np.arange(len(plen)), np.diff(off)], axis=1)) self.index = np.concatenate(self.index) if self.index else np.zeros((0, 3), np.int64) @@ -107,6 +120,18 @@ def batch(self, token_budget: int, max_rows: int = 256, width_multiple: int = 64 off = self.offsets[si][ri] seq = self.tokens[si][off:off + n] ids[r, :n] = seq - p = int(self.plens[si][ri]) - labels[r, p:n] = seq[p:] + lm = self.lmasks[si] + if lm is not None: + keep = np.asarray(lm[off:off + n]).astype(bool) + labels[r, :n] = np.where(keep, seq, -100) + else: + p = int(self.plens[si][ri]) + labels[r, p:n] = seq[p:] return ids, labels + + def supervised_tokens(self) -> int: + """Tokens that carry loss over every record (``.lmask`` or ``.plen`` rule).""" + total = 0 + for off, plen, lm in zip(self.offsets, self.plens, self.lmasks): + total += int(np.asarray(lm).sum()) if lm is not None else int(off[-1]) - int(plen.sum()) + return total diff --git a/chronos/expert_clusters.py b/chronos/expert_clusters.py index 7bf6aa2..ffea145 100644 --- a/chronos/expert_clusters.py +++ b/chronos/expert_clusters.py @@ -180,3 +180,57 @@ def cluster_locality_config(model_config: dict, shards_dir: str | None, *, enabl wsrc = "default (no SFT record lengths in the data)" return {"values": values, "source": {"lambda_cluster_locality": why, "cluster_locality_window": wsrc, "cluster_locality_clusters": "user value" if clusters else "ceil(E / 2k)"}} + + +# ---- cluster objective: unit of "one stream" -------------------------------- +# "window": fixed ``cluster_locality_window``-token windows (the original form). +# "segment": the documents of the batch, found from the token stream itself +# (``SEGMENT_BOUNDARY_ID`` separators of packed pretraining, one record per SFT +# row); documents longer than ``cluster_locality_window`` are cut into parts. +CLUSTER_MODES = ("window", "segment") +SEGMENT_BOUNDARY_ID = 0 # ``<|endoftext|>`` in Chronos tokenizers: document separator / padding + + +def cluster_locality_mode(config) -> str: + mode = str(getattr(config, "cluster_locality_mode", "window") or "window") + if mode not in CLUSTER_MODES: + raise ValueError(f"cluster_locality_mode must be one of {CLUSTER_MODES}, got {mode!r}") + return mode + + +def segment_ids(input_ids, chunk: int = 0, boundary_id: int = SEGMENT_BOUNDARY_ID): + """``[B, S]`` int64 document index of every position, from the tokens alone. + + A separator ``boundary_id`` closes the document it ends (so does each pad + of a right-padded row; ``token_mask`` removes those from the loss). + ``chunk`` > 0 cuts a document into ``chunk``-token parts. Ids increase + along each row; rows are told apart by ``segment_index``. No label, + prompt template or task token is consulted: what a document contains is + left for the router to discover.""" + import numpy as np + ids = np.asarray(input_ids) + B, S = ids.shape + is_bnd = ids == int(boundary_id) + doc = np.zeros((B, S), np.int64) + if S > 1: + doc[:, 1:] = np.cumsum(is_bnd[:, :-1], axis=1) + if int(chunk or 0) <= 0: + return doc + pos = np.arange(S)[None, :] + first = np.ones((B, S), bool) + first[:, 1:] = doc[:, 1:] != doc[:, :-1] + start = np.maximum.accumulate(np.where(first, pos, 0), axis=1) + key = doc * S + (pos - start) // int(chunk) + new = np.zeros((B, S), np.int64) + new[:, 1:] = key[:, 1:] != key[:, :-1] + return np.cumsum(new, axis=1) + + +def segment_index(segments): + """``(inverse [B*S] int32, n)``: dense batch-wide segment number of each position.""" + import numpy as np + seg = np.asarray(segments, np.int64) + B, S = seg.shape + flat = (seg + np.arange(B, dtype=np.int64)[:, None] * (S + 1)).reshape(-1) + uniq, inv = np.unique(flat, return_inverse=True) + return inv.astype(np.int32).reshape(-1), int(uniq.size) diff --git a/chronos/mlx/memory.py b/chronos/mlx/memory.py new file mode 100644 index 0000000..f380b89 --- /dev/null +++ b/chronos/mlx/memory.py @@ -0,0 +1,94 @@ +"""MLX buffer-cache limit, enforced by the kernel for every MLX model. + +MLX keeps freed Metal buffers in a cache for reuse. Its default cache +limit equals the memory limit, so on a unified-memory Mac the cache can grow +until macOS pages and training slows down (measured on a 64 GB machine: +footprint 55.8 → 13.6 GB and ~2.4k → ~5.2k tok/s after an 8 GiB limit). +The limit bounds only reclaimable cache; weights, activations and gradients +are active memory and are not affected. + +Default: ``min(physical memory / 8, 16 GiB)``. A caller value (recipe +``extra.mlx_cache_limit_gb``, ``CHRONOS_MLX_CACHE_LIMIT_GB``) may lower it, +never raise it above that ceiling. Values are GiB (×1024³); metrics report +decimal GB (8 GiB shows as 8.59 GB). +""" +from __future__ import annotations + +import os + +GIB = 1024 ** 3 +CEILING_GIB = 16.0 +FRACTION = 1.0 / 8.0 +ENV = "CHRONOS_MLX_CACHE_LIMIT_GB" + +_applied: dict | None = None + + +def physical_memory_bytes() -> int: + """Total physical memory (MLX device info, then sysctl/sysconf).""" + try: + import mlx.core as mx + info = mx.device_info() if hasattr(mx, "device_info") else mx.metal.device_info() + size = int(info.get("memory_size") or 0) + if size > 0: + return size + except Exception: # noqa: BLE001 - fall back to the OS + pass + try: + return int(os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) + except (ValueError, OSError, AttributeError): + return 0 + + +def ceiling_bytes(total: int | None = None) -> int: + """``min(total / 8, 16 GiB)``; 16 GiB when the total is unknown.""" + total = physical_memory_bytes() if total is None else int(total) + cap = int(CEILING_GIB * GIB) + return cap if total <= 0 else min(int(total * FRACTION), cap) + + +def resolve_cache_limit(requested_gb=None, total: int | None = None) -> tuple[int, str]: + """``(bytes, source)``: the limit to apply. ``requested_gb`` (GiB, > 0) or + the environment variable may lower the ceiling; larger values are clamped.""" + ceiling = ceiling_bytes(total) + raw, source = requested_gb, "recipe" + if raw in (None, ""): + raw, source = os.environ.get(ENV), "env" + if raw in (None, ""): + return ceiling, "default min(RAM/8, 16 GiB)" + try: + value = float(raw) + except (TypeError, ValueError): + raise ValueError(f"mlx_cache_limit_gb must be a positive number of GiB, got {raw!r}") from None + if not value > 0: + raise ValueError(f"mlx_cache_limit_gb must be > 0 GiB, got {raw!r}") + want = int(value * GIB) + if want > ceiling: + return ceiling, f"{source} {value:g} GiB clamped to min(RAM/8, 16 GiB)" + return want, source + + +def apply_cache_limit(requested_gb=None, *, log=None, force: bool = False) -> dict | None: + """Set the MLX cache limit once per process (again with ``force`` or a new + request). Returns ``{"limit_gb", "previous_gb", "ceiling_gb", "source"}`` + in GiB, or ``None`` without MLX.""" + global _applied + try: + import mlx.core as mx + except ModuleNotFoundError: + return None + if _applied is not None and not force and requested_gb in (None, "", _applied.get("requested")): + return _applied + limit, source = resolve_cache_limit(requested_gb) + previous = int(mx.set_cache_limit(limit)) + _applied = {"limit_gb": round(limit / GIB, 3), "previous_gb": round(previous / GIB, 3), + "ceiling_gb": round(ceiling_bytes() / GIB, 3), "source": source, + "requested": requested_gb} + if log is not None: + log(f"[run] MLX cache limit: {limit / GIB:g} GiB (previous {previous / GIB:.1f} GiB; {source})") + return _applied + + +def applied() -> dict | None: + """The last applied limit (``None`` before the first MLX model).""" + return None if _applied is None else {k: v for k, v in _applied.items() if k != "requested"} diff --git a/chronos/mlx/model.py b/chronos/mlx/model.py index 62bfb85..54a68bb 100644 --- a/chronos/mlx/model.py +++ b/chronos/mlx/model.py @@ -125,11 +125,19 @@ class ChronosMLXModel(nn.Module): def __init__(self, config): super().__init__() + # Kernel-level bound on MLX's free-buffer cache (min(RAM/8, 16 GiB)), + # applied once per process before any weight is allocated; a caller + # that set a lower limit first (recipe/env) keeps it. + from chronos.mlx.memory import apply_cache_limit + apply_cache_limit() self.config = config - if config.modalities: - raise ValueError("MLX modality encoders are not implemented; use the PyTorch MPS/CUDA backend for multimodal checkpoints") + # Image/audio encoders + gated cross-attention (Torch parity), only + # for the modalities the config enables. + from chronos.mlx.multimodal import MultimodalFusionMLX + self.multimodal_fusion = MultimodalFusionMLX(config) if config.modalities else None self.recurrent_controller = RecurrentControllerMLX(config) object.__setattr__(self, "last_recurrent_state", None) + object.__setattr__(self, "last_modality_states", {}) self.dropout = nn.Dropout(config.dropout) self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = [ChronosMLXBlock(i, config) @@ -147,6 +155,8 @@ def __call__( available_expert_masks=None, # list of [E] float masks, or None logits_dtype=None, # e.g. mx.float32 for losses recurrent_depth=None, # draft depth override (shared_first/first_only) + modality_inputs=None, # {"image"|"audio"|"*_embeds"|"*_mask": array} + modality_mask=None, # [B] or [B, 2] (image, audio) presence ): """``logits_dtype=mx.float32`` makes the logits FP32 at birth: the final hidden [B,S,H] and the head [V,H] are upcast (H/V the size of @@ -157,14 +167,23 @@ def __call__( moe.reset_iteration_routing() moe.open_pack_scope() try: - return self._forward(input_ids, cache, available_expert_masks, logits_dtype, recurrent_depth) + return self._forward(input_ids, cache, available_expert_masks, logits_dtype, recurrent_depth, + modality_inputs, modality_mask) finally: for moe in moes: moe.close_pack_scope() - def _forward(self, input_ids, cache, available_expert_masks, logits_dtype, depth_override=None): + def _forward(self, input_ids, cache, available_expert_masks, logits_dtype, depth_override=None, + modality_inputs=None, modality_mask=None): B, S = input_ids.shape x = self.dropout(self.embed_tokens(input_ids)) # [B, S, H] + object.__setattr__(self, "last_modality_states", {}) + if self.multimodal_fusion is not None: + x, states = self.multimodal_fusion(x, modality_inputs, modality_mask) + object.__setattr__(self, "last_modality_states", states) + elif modality_inputs: + raise ValueError("modality_inputs were provided, but this checkpoint has no enabled modalities; " + "set config.modalities before constructing the model") # Each attention layer builds its own boolean causal / window mask # (fused SDPA); no additive mask is materialized in the compute dtype. @@ -380,6 +399,9 @@ def _assign_norm(norm, sd, key, missing): proj_layers = mlx_model.lookahead_router.proj.layers _assign_linear(proj_layers[0], sd, "model.lookahead_router.proj.0.weight", missing) _assign_linear(proj_layers[2], sd, "model.lookahead_router.proj.2.weight", missing) + if mlx_model.multimodal_fusion is not None: + from chronos.mlx.multimodal import load_torch_fusion_state + load_torch_fusion_state(mlx_model.multimodal_fusion, sd, _to_mx, missing) if missing: raise RuntimeError( diff --git a/chronos/mlx/moe.py b/chronos/mlx/moe.py index 9bd1620..03b3051 100644 --- a/chronos/mlx/moe.py +++ b/chronos/mlx/moe.py @@ -118,6 +118,62 @@ def last_router_probs(self, value): def _shared_out(self, x: mx.array) -> mx.array: return sum(e(x) for e in self.shared_experts) / len(self.shared_experts) + def _resolve_live(self, i: int, python_avail: set) -> bool: + """Materialize selected expert ``i`` under the miss policy; True = real weights. + + Same decisions as the per-expert loop: a miss under ``sync_on_demand`` + loads synchronously (or raises); other policies fall back to the shared expert.""" + policy = getattr(self, "runtime_miss_policy", "") + loader = getattr(self, "runtime_on_demand_loader", None) + sync = loader is not None and policy == "sync_on_demand" + live = i in python_avail or (sync and bool(loader(int(i)))) + if live and isinstance(self.experts[i], LazyFeedForwardMLX): + if sync: + loader(int(i)) + if isinstance(self.experts[i], LazyFeedForwardMLX): + if policy == "sync_on_demand": + raise RuntimeError(f"MLX lazy expert {i} was selected but not materialized") + return False + if not live: + if policy == "sync_on_demand": + raise RuntimeError(f"MLX lazy expert {i} could not be loaded on demand") + return False + touch = getattr(self, "runtime_touch_expert", None) + if touch is not None: + touch(int(i)) + return True + + def _lazy_grouped(self, x_flat, topk_idx, topk_weight, selected_host, python_avail, shared_out): + """Lazy/offload MoE through ``grouped_swiglu`` over the live experts only. + + The full-memory path runs the same kernel, so a lazy decode reproduces + the full decode instead of re-rounding through the per-expert loop + (bf16 on the GPU: a different summation order flips near-tied greedy + tokens). Routed ids are remapped into a compact stack of the live + experts; a miss maps to a spare slot whose output the shared expert + replaces (the loop path's fallback).""" + selected = [i for i, s in enumerate(selected_host) if s] + live, experts = [], [] + for i in selected: + if self._resolve_live(i, python_avail): + # Take the module now: loading the next selected expert may + # evict this one from the hot set (its weights stay valid). + live.append(i) + experts.append(self.experts[i]) + if not live: + # Every routed slot missed: each is served by the shared output. + return shared_out + shared_out * topk_weight.sum(axis=-1, keepdims=True).astype(shared_out.dtype) + remap = [len(live)] * self.num_experts + for j, i in enumerate(live): + remap[i] = j + fallback = None + if len(live) < len(selected): + experts.append(experts[0]) # spare slot; its output is replaced + fallback = (mx.arange(len(experts)) < len(live), shared_out) + local = mx.array(remap, dtype=mx.int32)[topk_idx] + routed = grouped_swiglu(x_flat, local, topk_weight, *stack_expert_weights(experts), fallback=fallback) + return routed + shared_out + def __call__( self, x: mx.array, @@ -195,6 +251,9 @@ def __call__( selected_host = hit.tolist() # Read by residency planners from the runtime hooks below. object.__setattr__(self, "_selected_host", selected_host) + if self.sparse_dispatch and len(self.shared_experts) > 0 and _sparse_ok(x_flat.dtype): + y = self._lazy_grouped(x_flat, topk_idx, topk_weight, selected_host, python_avail, shared_out) + return y.reshape(B, S, H) for i, expert in enumerate(self.experts): sel = (topk_idx == i) diff --git a/chronos/mlx/multimodal.py b/chronos/mlx/multimodal.py new file mode 100644 index 0000000..14f93c6 --- /dev/null +++ b/chronos/mlx/multimodal.py @@ -0,0 +1,284 @@ +"""MLX twin of ``chronos.model.multimodal``: image/audio encoders and the gated +cross-attention that fuses modality tokens into the text states. + +Same parameters and numerics as the Torch module (parity-tested), so a +checkpoint trains on either backend. Layout differences are handled at +conversion time (``torch_fusion_state`` / ``load_torch_fusion_state``): MLX +convolutions are channels-last (Conv2d weight ``[O,kh,kw,I]``, Conv1d +``[O,k,I]``) and the Torch ``MultiheadAttention`` packs q/k/v into one +``in_proj``. The fused output only feeds the token embeddings, so routing, +recurrence, KV caches and the cluster objectives are unchanged. +""" +from __future__ import annotations + +import mlx.core as mx +import mlx.nn as nn + +MODALITY_KEYS = ("image", "audio", "image_embeds", "audio_embeds", + "image_mask", "audio_mask", "image_embeds_mask", "audio_embeds_mask") + + +def _mx(value, dtype=None): + """torch/numpy/list → mx.array (no copy for mx input).""" + if value is None or isinstance(value, mx.array): + return value if value is None or dtype is None else value.astype(dtype) + try: + import torch + if isinstance(value, torch.Tensor): + t = value.detach().cpu() + if t.dtype == torch.bfloat16: + t = t.float() + value = t.numpy() + except ModuleNotFoundError: + pass + arr = mx.array(value) + return arr if dtype is None else arr.astype(dtype) + + +class PatchImageEncoderMLX(nn.Module): + """``[B,3,H,W]`` (or ``[B,N,3,H,W]``) images → ``[B,T,H]`` patch tokens.""" + + def __init__(self, hidden_size: int, patch_size: int = 16, max_tokens: int = 196): + super().__init__() + self.patch = nn.Conv2d(3, hidden_size, patch_size, stride=patch_size, bias=False) + self.max_tokens = int(max_tokens) + self.norm = nn.LayerNorm(hidden_size) + + def __call__(self, images: mx.array) -> mx.array: + b = images.shape[0] + if images.ndim == 5: + n, c, h, w = images.shape[1:] + images = images.reshape(b * n, c, h, w) + x = self.patch(images.transpose(0, 2, 3, 1).astype(self.patch.weight.dtype)) # NHWC + x = x.reshape(x.shape[0], -1, x.shape[-1]) # row-major patches (Torch flatten(2)) + x = x.reshape(b, -1, x.shape[-1]) + return self.norm(x[:, : self.max_tokens]) + + +class ConvAudioEncoderMLX(nn.Module): + """``[B,T]`` (or ``[B,1,T]``) waveform → ``[B,T',H]`` frame tokens.""" + + def __init__(self, hidden_size: int, max_tokens: int = 128): + super().__init__() + mid = max(32, hidden_size // 2) + self.net = nn.Sequential( + nn.Conv1d(1, mid, 400, stride=160, padding=120), nn.SiLU(), + nn.Conv1d(mid, hidden_size, 5, stride=2, padding=2), nn.SiLU(), + ) + self.max_tokens = int(max_tokens) + self.norm = nn.LayerNorm(hidden_size) + + def __call__(self, audio: mx.array) -> mx.array: + if audio.ndim == 3: + audio = audio[:, 0] + x = self.net(audio[:, :, None].astype(self.net.layers[0].weight.dtype)) # NLC + return self.norm(x[:, : self.max_tokens]) + + +def _validity(present, b: int, n: int, active_col: mx.array, name: str) -> mx.array: + """``[B, n]`` bool: keys a row may attend (Torch ``~key_padding_mask``).""" + if present is None: + return mx.broadcast_to(active_col[:, None], (b, n)) + present = _mx(present).astype(mx.bool_) + if present.ndim == 2: + if present.shape[0] != b: + raise ValueError(f"{name} modality mask batch dimension mismatch") + if present.shape[1] >= n: + return present[:, :n] + return mx.concatenate([present, mx.zeros((b, n - present.shape[1]), dtype=mx.bool_)], axis=1) + present = present.reshape(-1) + if present.size == 1: + present = mx.broadcast_to(present, (b,)) + if present.size != b: + raise ValueError(f"{name} modality mask batch dimension mismatch") + return mx.broadcast_to(mx.logical_and(active_col, present)[:, None], (b, n)) + + +class MultimodalFusionMLX(nn.Module): + """Gated cross-attention of text states to image/audio tokens.""" + + def __init__(self, config): + super().__init__() + h = int(config.hidden_size) + heads = max(1, min(int(config.modality_fusion_heads), int(config.num_attention_heads))) + while h % heads: + heads -= 1 + enabled = set(getattr(config, "modalities", []) or []) + self.image_encoder = (PatchImageEncoderMLX(h, config.image_patch_size, config.max_image_tokens) + if "image" in enabled else None) + self.audio_encoder = (ConvAudioEncoderMLX(h, config.max_audio_tokens) + if "audio" in enabled else None) + self.hidden_size, self.num_heads = h, heads + self.image_embedding = mx.zeros((1, 1, h)) + self.audio_embedding = mx.zeros((1, 1, h)) + self.q_proj = nn.Linear(h, h) + self.k_proj = nn.Linear(h, h) + self.v_proj = nn.Linear(h, h) + self.out_proj = nn.Linear(h, h) + self.gate = mx.zeros(()) + self.attn_dropout = nn.Dropout(float(config.dropout)) + self.dropout = nn.Dropout(float(config.modality_dropout)) + self.freeze_encoders = bool(config.freeze_modality_encoders) + if self.freeze_encoders: + for encoder in (self.image_encoder, self.audio_encoder): + if encoder is not None: + encoder.freeze() + + def _fit_width(self, x: mx.array) -> mx.array: + width = x.shape[-1] + if width == self.hidden_size: + return x + if width > self.hidden_size: + return x[..., : self.hidden_size] + return mx.pad(x, [(0, 0)] * (x.ndim - 1) + [(0, self.hidden_size - width)]) + + def _embeds(self, value, text_dtype, offset): + x = _mx(value, text_dtype) + x = x[:, None] if x.ndim == 2 else x + return self._fit_width(x) + offset.astype(text_dtype) + + def __call__(self, text_states: mx.array, modality_inputs=None, modality_mask=None): + if not modality_inputs: + return text_states, {} + b, dtype = text_states.shape[0], text_states.dtype + parts, states = [], {} # parts: (modality, tokens, validity input) + for name, encoder, offset in (("image", self.image_encoder, self.image_embedding), + ("audio", self.audio_encoder, self.audio_embedding)): + embeds = modality_inputs.get(f"{name}_embeds") + if embeds is not None: + tokens = self._embeds(embeds, dtype, offset) + parts.append((name, tokens, modality_inputs.get(f"{name}_embeds_mask"))) + states[name] = tokens + raw = modality_inputs.get(name) + if raw is not None and embeds is None: + if encoder is None: + raise ValueError(f"{name} input provided but {name} modality is disabled") + tokens = encoder(_mx(raw, mx.float32)) + if self.freeze_encoders: + tokens = mx.stop_gradient(tokens) + tokens = (tokens + offset).astype(dtype) + parts.append((name, tokens, modality_inputs.get(f"{name}_mask"))) + states[name] = mx.concatenate([states[name], tokens], axis=1) if name in states else tokens + if not parts: + return text_states, {} + present = {m for m, _, _ in parts} + if modality_mask is None: + active = mx.broadcast_to(mx.array([["image" in present, "audio" in present]]), (b, 2)) + else: + raw_mask = _mx(modality_mask, mx.float32) + if raw_mask.ndim == 1: + col = raw_mask.reshape(-1, 1) if raw_mask.size != 1 else mx.broadcast_to(raw_mask.reshape(1, 1), (b, 1)) + active = mx.broadcast_to(col > 0, (b, 2)) + elif raw_mask.shape[-1] == 2: + active = raw_mask > 0 + else: + raise ValueError("modality_mask must be [B] or [B,2] (image,audio)") + memory = self.dropout(mx.concatenate([t for _, t, _ in parts], axis=1)) + valid = mx.concatenate([_validity(p, b, t.shape[1], active[:, 0 if m == "image" else 1], m) + for m, t, p in parts], axis=1) + cols = sorted({0 if m == "image" else 1 for m in present}) + no_media = mx.logical_not(mx.any(active[:, cols], axis=1)) + # Rows without media keep one key (no all-masked softmax) and get a zero update. + first = mx.arange(valid.shape[1])[None, :] == 0 + valid = mx.logical_or(valid, mx.logical_and(no_media[:, None], first)) + fused = self._attend(text_states, memory, valid) + fused = mx.where(no_media[:, None, None], mx.zeros_like(fused), fused) + return text_states + mx.sigmoid(self.gate).astype(dtype) * fused, states + + def _attend(self, x: mx.array, memory: mx.array, valid: mx.array) -> mx.array: + b, s, h = x.shape + t, nh = memory.shape[1], self.num_heads + hd = h // nh + q = self.q_proj(x).reshape(b, s, nh, hd).transpose(0, 2, 1, 3) + k = self.k_proj(memory).reshape(b, t, nh, hd).transpose(0, 2, 1, 3) + v = self.v_proj(memory).reshape(b, t, nh, hd).transpose(0, 2, 1, 3) + scores = (q.astype(mx.float32) @ k.astype(mx.float32).transpose(0, 1, 3, 2)) * (hd ** -0.5) + scores = mx.where(valid[:, None, None, :], scores, mx.array(-mx.inf, dtype=mx.float32)) + probs = self.attn_dropout(mx.softmax(scores, axis=-1)) + out = (probs @ v.astype(mx.float32)).astype(x.dtype) + return self.out_proj(out.transpose(0, 2, 1, 3).reshape(b, s, h)) + + +# ---- Torch ↔ MLX state (prefix ``model.multimodal_fusion.``) --------------- +PREFIX = "model.multimodal_fusion." + + +def _pairs(fusion: MultimodalFusionMLX): + """``(torch key, getter, setter, layout)``; layout: None | 'conv2d' | 'conv1d'.""" + out = [("gate", fusion, "gate", None), + ("modality_embeddings.image", fusion, "image_embedding", None), + ("modality_embeddings.audio", fusion, "audio_embedding", None), + ("attn.out_proj.weight", fusion.out_proj, "weight", None), + ("attn.out_proj.bias", fusion.out_proj, "bias", None)] + if fusion.image_encoder is not None: + enc = fusion.image_encoder + out += [("image_encoder.patch.weight", enc.patch, "weight", "conv2d"), + ("image_encoder.norm.weight", enc.norm, "weight", None), + ("image_encoder.norm.bias", enc.norm, "bias", None)] + if fusion.audio_encoder is not None: + enc = fusion.audio_encoder + for i in (0, 2): + conv = enc.net.layers[i] + out += [(f"audio_encoder.net.{i}.weight", conv, "weight", "conv1d"), + (f"audio_encoder.net.{i}.bias", conv, "bias", None)] + out += [("audio_encoder.norm.weight", enc.norm, "weight", None), + ("audio_encoder.norm.bias", enc.norm, "bias", None)] + return out + + +def _to_mlx_layout(arr: mx.array, layout): + if layout == "conv2d": + return arr.transpose(0, 2, 3, 1) # [O,I,kh,kw] → [O,kh,kw,I] + if layout == "conv1d": + return arr.transpose(0, 2, 1) # [O,I,k] → [O,k,I] + return arr + + +def _to_torch_layout(arr: mx.array, layout): + if layout == "conv2d": + return arr.transpose(0, 3, 1, 2) + if layout == "conv1d": + return arr.transpose(0, 2, 1) + return arr + + +def load_torch_fusion_state(fusion: MultimodalFusionMLX, sd, convert, missing: list) -> None: + """Copy ``model.multimodal_fusion.*`` tensors from a Torch state dict. + ``convert(tensor) -> mx.array`` applies the target dtype.""" + for key, obj, attr, layout in _pairs(fusion): + full = PREFIX + key + if full not in sd: + missing.append(full) + continue + value = _to_mlx_layout(convert(sd[full]), layout) + if tuple(value.shape) != tuple(getattr(obj, attr).shape): + raise ValueError(f"MLX checkpoint shape mismatch for {full}: {tuple(value.shape)} != {tuple(getattr(obj, attr).shape)}") + setattr(obj, attr, value) + w, b = PREFIX + "attn.in_proj_weight", PREFIX + "attn.in_proj_bias" + if w not in sd or b not in sd: + missing.extend(k for k in (w, b) if k not in sd) + return + weight, bias, h = convert(sd[w]), convert(sd[b]), fusion.hidden_size + for i, proj in enumerate((fusion.q_proj, fusion.k_proj, fusion.v_proj)): + proj.weight = weight[i * h:(i + 1) * h] + proj.bias = bias[i * h:(i + 1) * h] + + +def torch_fusion_state(fusion: MultimodalFusionMLX, to_torch) -> dict: + """``model.multimodal_fusion.*`` Torch tensors (``to_torch(mx.array)``).""" + sd = {PREFIX + key: to_torch(mx.contiguous(_to_torch_layout(getattr(obj, attr), layout))) + for key, obj, attr, layout in _pairs(fusion)} + sd[PREFIX + "attn.in_proj_weight"] = to_torch(mx.concatenate( + [fusion.q_proj.weight, fusion.k_proj.weight, fusion.v_proj.weight], axis=0)) + sd[PREFIX + "attn.in_proj_bias"] = to_torch(mx.concatenate( + [fusion.q_proj.bias, fusion.k_proj.bias, fusion.v_proj.bias], axis=0)) + return sd + + +def batch_modalities(batch): + """``(ids, labels, modality_inputs | None, modality_mask | None)`` from a + ``(ids, labels)`` pair or a ``multimodal_collate`` dict.""" + if isinstance(batch, dict): + media = batch.get("modality_inputs") or None + return batch["input_ids"], batch["labels"], media, batch.get("modality_mask") + return batch[0], batch[1], None, None diff --git a/chronos/mlx/training/inspect_checkpoint.py b/chronos/mlx/training/inspect_checkpoint.py index eb4183e..a087074 100644 --- a/chronos/mlx/training/inspect_checkpoint.py +++ b/chronos/mlx/training/inspect_checkpoint.py @@ -55,11 +55,13 @@ def _phase(spec: dict, state: dict) -> dict: def held_out(spec, config, model, phase, n_batches: int = 8) -> dict: - from chronos.mlx.training.longrun import PretrainSource, SFTSource, recurrent_metrics, routing_metrics + from chronos.mlx.training.longrun import (PretrainSource, RoutingCounts, SFTSource, recurrent_metrics, + token_mask_for) from chronos.mlx.training.trainer import _collect_router_probs source = (PretrainSource(spec["shards_dir"], phase["seq_len"], spec["seed"]) if phase["name"] == "pretrain" else SFTSource(spec["shards_dir"], spec["seed"])) total, count, extra = 0.0, 0, {} + routes, top_k = RoutingCounts(), int(config.num_experts_per_tok) for i, (ids, labels) in enumerate(source.valid_batches(n_batches, phase["micro_tokens"])): logits, _, _ = model(mx.array(ids, dtype=mx.int32), logits_dtype=mx.float32) lb = mx.array(labels, dtype=mx.int32).reshape(-1) @@ -69,11 +71,12 @@ def held_out(spec, config, model, phase, n_batches: int = 8) -> dict: s, n = (ce * mask).sum(), mask.sum() mx.eval(s, n) total, count = total + float(s.item()), count + int(n.item()) + routes.add(_collect_router_probs(model), top_k, token_mask_for(labels)) # real tokens, every batch if i == 0: - extra = {**routing_metrics(_collect_router_probs(model), int(config.num_experts_per_tok)), - **recurrent_metrics(getattr(model, "last_recurrent_state", None))} + extra = recurrent_metrics(getattr(model, "last_recurrent_state", None)) ce = total / max(count, 1) - return {"valid_ce": round(ce, 4), "valid_ppl": round(math.exp(min(ce, 30)), 2), "valid_tokens": count, **extra} + return {"valid_ce": round(ce, 4), "valid_ppl": round(math.exp(min(ce, 30)), 2), "valid_tokens": count, + **routes.summary(top_k), **extra} def greedy(model, tok, prompt_ids: list[int], stop_ids: set[int], limit: int) -> list[int]: diff --git a/chronos/mlx/training/io.py b/chronos/mlx/training/io.py index 1a1eb98..c5e0dc6 100644 --- a/chronos/mlx/training/io.py +++ b/chronos/mlx/training/io.py @@ -69,4 +69,8 @@ def mlx_state_to_torch(mlx_model, config) -> dict[str, torch.Tensor]: proj_layers = mlx_model.lookahead_router.proj.layers sd["model.lookahead_router.proj.0.weight"] = _torch_tensor(proj_layers[0].weight) sd["model.lookahead_router.proj.2.weight"] = _torch_tensor(proj_layers[2].weight) + fusion = getattr(mlx_model, "multimodal_fusion", None) + if fusion is not None: + from chronos.mlx.multimodal import torch_fusion_state + sd.update(torch_fusion_state(fusion, _torch_tensor)) return sd diff --git a/chronos/mlx/training/longrun.py b/chronos/mlx/training/longrun.py index 5405637..b4f1d63 100644 --- a/chronos/mlx/training/longrun.py +++ b/chronos/mlx/training/longrun.py @@ -142,22 +142,76 @@ def _optimizer_restore(optimizer, model, snap: dict) -> None: optimizer.state = tree_unflatten(list(current.items())) -def routing_metrics(router: mx.array | None, top_k: int) -> dict: - """Load balance and confidence over every dispatch ``[B, S, V, E]``.""" - if router is None: - return {} - E = router.shape[-1] - flat = router.reshape(-1, router.shape[2], E).astype(mx.float32) # [N, V, E] - top = mx.argpartition(-flat, kth=top_k - 1, axis=-1)[..., :top_k] - load = (top[..., None] == mx.arange(E)).astype(mx.float32).sum(-2).mean(0) / top_k # [V, E] - importance = flat.mean(0) - entropy = -(flat * mx.log(mx.maximum(flat, 1e-9))).sum(-1).mean() - cv = (load.std(-1) / mx.maximum(load.mean(-1), 1e-9)).mean() - dead = (load.max(0) < (0.05 / E)).sum() - mx.eval(load, importance, entropy, cv, dead) - return {"route_entropy": float(entropy.item()), "route_entropy_max": float(math.log(E)), - "load_cv": float(cv.item()), "dead_experts": int(dead.item()), - "max_expert_load": float(load.max().item() * E)} +def token_mask_for(labels) -> np.ndarray | None: + """Real-token mask ``[B, S]`` of a batch whose labels are aligned with the + inputs (``PretrainSource``/``SFTSource`` pairs, ``labels[t]`` = token t+1). + A row's real tokens end one past its last supervised label (that input + holds the record's final token); later positions are right-padding. + Prompt tokens stay real. ``None`` when nothing is padded (packed + pretraining), so packed batches keep the exact unmasked loss. Same rule + as ``trainer.real_token_mask(labels, aligned=True)``.""" + labels = np.asarray(labels) + if labels.ndim != 2 or labels.size == 0 or not (labels[:, -1] == -100).any(): + return None + pos = np.arange(labels.shape[1]) + last = np.where(labels != -100, pos[None, :], -1).max(axis=1) + mask = pos[None, :] <= (last + 1)[:, None] + return None if mask.all() else mask + + +class RoutingCounts: + """Top-k dispatch counts, router probability mass and entropy per virtual + layer, summed over every micro-batch of a step (real tokens only). + + Training telemetry used to read the last micro-batch alone, padding + included; a step's full batch is what the router losses see.""" + + def __init__(self): + self.load = self.importance = self.entropy = None + self.tokens = 0 + + def add(self, router: mx.array | None, top_k: int, token_mask=None) -> None: + if router is None: + return + E, V = router.shape[-1], router.shape[2] + flat = router.reshape(-1, V, E).astype(mx.float32) # [N, V, E] + w = mx.ones((flat.shape[0],)) if token_mask is None else \ + mx.array(np.asarray(token_mask).reshape(-1), dtype=mx.float32) + top = mx.argpartition(-flat, kth=top_k - 1, axis=-1)[..., :top_k] + hits = (top[..., None] == mx.arange(E)).astype(mx.float32).sum(-2) # [N, V, E] + ent = -(flat * mx.log(mx.maximum(flat, 1e-9))).sum(-1) # [N, V] + load, imp = (hits * w[:, None, None]).sum(0), (flat * w[:, None, None]).sum(0) + ent = (ent * w[:, None]).sum(0) + n = w.sum() + mx.eval(load, imp, ent, n) + if self.load is None: + self.load, self.importance, self.entropy = load, imp, ent + else: + self.load, self.importance, self.entropy = self.load + load, self.importance + imp, self.entropy + ent + self.tokens += int(n.item()) + + def summary(self, top_k: int) -> dict: + if self.load is None or self.tokens <= 0: + return {} + E = self.load.shape[-1] + load = self.load / (self.tokens * top_k) # [V, E] dispatch share + cv = load.std(-1) / mx.maximum(load.mean(-1), 1e-9) # [V] + entropy = (self.entropy / self.tokens).mean() + dead = (load.max(0) < (0.05 / E)).sum() + mx.eval(load, cv, entropy, dead) + return {"route_entropy": float(entropy.item()), "route_entropy_max": float(math.log(E)), + "load_cv": float(cv.mean().item()), "load_cv_max": float(cv.max().item()), + "dead_experts": int(dead.item()), + "max_expert_load": float(load.max().item() * E), "min_expert_load": float(load.min().item() * E), + "route_tokens": int(self.tokens)} + + +def routing_metrics(router: mx.array | None, top_k: int, token_mask=None) -> dict: + """Load balance and confidence over every real-token dispatch ``[B, S, V, E]`` + (``token_mask`` ``[B, S]``, 1 = real; ``None`` = all positions).""" + acc = RoutingCounts() + acc.add(router, top_k, token_mask) + return acc.summary(top_k) def recurrent_metrics(state: dict | None) -> dict: @@ -261,6 +315,10 @@ def __init__(self, spec: RunSpec, log=print): self.spec, self.log = spec, log os.makedirs(spec.run_dir, exist_ok=True) self.config = ChronosConfig(**spec.model_config) + # MLX free-buffer cache bound, set before the model exists: recipe + # ``extra.mlx_cache_limit_gb`` (GiB) may lower min(RAM/8, 16 GiB). + from chronos.mlx.memory import apply_cache_limit + self.cache_limit = apply_cache_limit((spec.extra or {}).get("mlx_cache_limit_gb"), log=log, force=True) import torch torch.manual_seed(spec.seed) # reproducible initialization (built in Torch, then converted) mx.random.seed(spec.seed) @@ -285,6 +343,9 @@ def __init__(self, spec: RunSpec, log=print): self.heartbeat = telemetry.Heartbeat(spec.run_dir) self.precision = getattr(self.model, "_chronos_precision", None) self.runtime = telemetry.runtime_info(self.model, spec, self.precision) + if self.cache_limit is not None: + self.runtime["mlx_cache_limit_gb"] = self.cache_limit["limit_gb"] + self.runtime["mlx_cache_limit"] = {k: v for k, v in self.cache_limit.items() if k != "requested"} try: telemetry._atomic_json(os.path.join(spec.run_dir, "runtime.json"), self.runtime) except OSError: @@ -354,7 +415,8 @@ def _beat(self, state: str | None = None, **fields): rollbacks=int(st.get("rollbacks", 0)), lr_scale=float(st.get("lr_scale", 1.0)), best_valid=st.get("best_valid"), **fields) - def _tensor_telemetry(self, phase, step: int, grads, before, accum: int, lr: float, norm: float): + def _tensor_telemetry(self, phase, step: int, grads, before, accum: int, lr: float, norm: float, + routes: RoutingCounts | None = None): """One ``parameters.jsonl`` row: per-tensor weight RMS/|max|, pre-clip gradient RMS, relative update, per-group norms, expert grid, routing.""" tel = self.telemetry @@ -365,14 +427,16 @@ def _tensor_telemetry(self, phase, step: int, grads, before, accum: int, lr: flo grad_scale=1.0 / max(1, accum)) row = {"time": time.strftime("%Y-%m-%dT%H:%M:%S"), "phase": phase.name, "step": step, "lr": lr, "grad_norm": norm, **stats, "experts": tel.expert_table(stats["tensors"]), - "routing": tel.routing_snapshot(self.model)} + "routing": tel.routing_snapshot(self.model, counts=routes)} row["telemetry_s"] = round(time.monotonic() - t0, 3) with open(self.parameters_path, "a", encoding="utf-8") as fh: fh.write(json.dumps(row, ensure_ascii=False) + "\n") return row # ---- training ---------------------------------------------------------- - def _loss_fn(self, model, ids, labels): + def _loss_fn(self, model, ids, labels, token_mask=None): + """CE on supervised labels + Chronos regularizers; ``token_mask`` ``[B, S]`` + (1 = real token) keeps right-padding out of the router regularizers.""" from chronos.mlx.training.trainer import _chronos_regularized_loss, draft_exit_kl, draft_exit_logits draft = draft_exit_logits(model, ids, model.config) # None unless draft_exit_kl_coef > 0 logits, lookahead, _ = model(ids, logits_dtype=mx.float32) @@ -384,7 +448,7 @@ def _loss_fn(self, model, ids, labels): # Labels here are pre-aligned with the logits (no shift). draft_kl = None if draft is None else draft_exit_kl(draft, logits, labels != -100) return _chronos_regularized_loss(model, ce, lookahead, model.config, metrics={"ce": ce, "primary": ce}, - draft_kl=draft_kl) + draft_kl=draft_kl, token_mask=token_mask, input_ids=ids) def _source(self, phase: PhaseSpec): if phase.name == "pretrain": @@ -392,7 +456,9 @@ def _source(self, phase: PhaseSpec): return SFTSource(self.spec.shards_dir, self.spec.seed) def _micro_step(self, grad_fn, ids, labels): - loss, grads = grad_fn(self.model, mx.array(ids, dtype=mx.int32), mx.array(labels, dtype=mx.int32)) + mask = token_mask_for(labels) + loss, grads = grad_fn(self.model, mx.array(ids, dtype=mx.int32), mx.array(labels, dtype=mx.int32), + None if mask is None else mx.array(mask)) return loss, grads def _apply(self, grads, n_micro, phase, lr): @@ -474,9 +540,16 @@ def run_phase(self, pi: int, deadline_s: float, max_steps: int | None = None): lr *= float(self.state.get("lr_scale", 1.0)) t0 = time.monotonic() acc, tokens, processed, loss_sum = None, 0, 0, mx.zeros(()) + # Routing telemetry covers the whole step (every micro-batch, real + # tokens only), counted only on steps that are logged or sampled. + will_log = (step + 1) % self.spec.log_every_steps == 0 or step < 3 + want_tensors = tensor_every > 0 and time.monotonic() - last_tensors >= tensor_every + routes = RoutingCounts() if will_log or want_tensors else None for _ in range(accum): ids, labels = source.batch(phase.micro_tokens) loss, grads = self._micro_step(grad_fn, ids, labels) + if routes is not None: + routes.add(_collect_router_probs(self.model), top_k, token_mask_for(labels)) acc = grads if acc is None else _add_grads(acc, grads) loss_sum = loss_sum + loss tokens += int((labels != -100).sum()) # supervised tokens @@ -484,7 +557,6 @@ def run_phase(self, pi: int, deadline_s: float, max_steps: int | None = None): mx.eval(acc, loss_sum) # bound the lazy graph to one micro-batch # Tensor telemetry keeps references to the pre-update weights (MLX # arrays are immutable: no copy) and the FP32 accumulated grads. - want_tensors = tensor_every > 0 and time.monotonic() - last_tensors >= tensor_every before = self.model.parameters() if want_tensors else None norm = self._apply(acc, accum, phase, lr) mx.eval(self.model.parameters(), self.optimizer.state, norm) @@ -513,7 +585,8 @@ def run_phase(self, pi: int, deadline_s: float, max_steps: int | None = None): self.state["phase_step"] = step + 1 if want_tensors: try: - self._tensor_telemetry(phase, step + 1, acc, before, accum, lr, float(norm.item())) + self._tensor_telemetry(phase, step + 1, acc, before, accum, lr, float(norm.item()), + routes=routes) except Exception as exc: # noqa: BLE001 - telemetry must never stop training self.log(f"[telemetry] tensor stats skipped: {exc!r}") last_tensors = time.monotonic() @@ -533,7 +606,7 @@ def run_phase(self, pi: int, deadline_s: float, max_steps: int | None = None): self.log(f"[run] {phase.name}: {tps:,.0f} tok/s → ~{self.state['phase_total']} steps " f"(~{self.state['phase_total'] * tokens / 1e9:.2f}B tokens)") self.state["phase_tokens_per_s"][phase.name] = tps - if (step + 1) % self.spec.log_every_steps == 0 or step < 3: + if will_log: comps = {k: round(float(v.item()), 5) for k, v in getattr(self.model, "_last_loss_components", {}).items()} rec = {"phase": phase.name, "step": step + 1, "total": self.state["phase_total"], "lr": lr, "grad_norm": float(norm.item()), "loss": float(loss_sum.item()) / accum, @@ -546,7 +619,7 @@ def run_phase(self, pi: int, deadline_s: float, max_steps: int | None = None): "phase_wall_s": self.state.get("phase_wall_s"), "phase_budget_s": self.state.get("phase_budget_s"), "lr_scale": float(self.state.get("lr_scale", 1.0)), "rollbacks": int(self.state.get("rollbacks", 0)), **comps, **recurrent_metrics(getattr(self.model, "last_recurrent_state", None)), - **routing_metrics(_collect_router_probs(self.model), top_k)} + **(routes.summary(top_k) if routes is not None else {})} self._emit(rec) self.log(f"[{phase.name}] {step + 1}/{self.state['phase_total']} loss={rec['loss']:.4f} " f"ce={comps.get('primary_loss', 0):.4f} gn={rec['grad_norm']:.2f} lr={lr:.2e} " diff --git a/chronos/mlx/training/telemetry.py b/chronos/mlx/training/telemetry.py index 673bd8e..b381dd7 100644 --- a/chronos/mlx/training/telemetry.py +++ b/chronos/mlx/training/telemetry.py @@ -144,10 +144,13 @@ def grid(i): return {"layers": layers, "weight_rms": grid(1), "grad_rms": grid(2)} if layers else {} -def routing_snapshot(model) -> dict | None: - """Per virtual layer: load share, mean router probability and entropy of - the last forward's dispatches (``[V, E]``), with ``L·`` - labels so recurrent iterations of one physical layer sit side by side.""" +def routing_snapshot(model, counts=None) -> dict | None: + """Per virtual layer: load share, mean router probability and entropy + (``[V, E]``), with ``L·`` labels so recurrent iterations + of one physical layer sit side by side. + + ``counts`` (``longrun.RoutingCounts``) = the whole step's real tokens; + without it, the last forward's dispatches (every position).""" from chronos.mlx.moe import ChronosMLXMOE probs, labels = [], [] for li, layer in enumerate(getattr(model, "layers", [])): @@ -156,21 +159,30 @@ def routing_snapshot(model) -> dict | None: for it, p in enumerate(moe.iteration_router_probs): probs.append(p) labels.append(f"L{li}·{it}") - if not probs: - return None top_k = int(model.config.num_experts_per_tok) - router = mx.stack(probs, axis=2) # [B, S, V, E] - E, V = router.shape[-1], router.shape[2] - flat = router.reshape(-1, V, E).astype(mx.float32) - top = mx.argpartition(-flat, kth=top_k - 1, axis=-1)[..., :top_k] - load = (top[..., None] == mx.arange(E)).astype(mx.float32).sum(-2).mean(0) / top_k - importance = flat.mean(0) - entropy = -(flat * mx.log(mx.maximum(flat, 1e-9))).sum(-1).mean(0) + if counts is not None and counts.load is not None and counts.tokens > 0: + n = float(counts.tokens) + load, importance, entropy = counts.load / (n * top_k), counts.importance / n, counts.entropy / n + E = int(load.shape[-1]) + if len(labels) != int(load.shape[0]): + labels = [f"V{v}" for v in range(int(load.shape[0]))] + source = "all micro-batches of the step, real tokens" + else: + if not probs: + return None + router = mx.stack(probs, axis=2) # [B, S, V, E] + E, V = router.shape[-1], router.shape[2] + flat = router.reshape(-1, V, E).astype(mx.float32) + top = mx.argpartition(-flat, kth=top_k - 1, axis=-1)[..., :top_k] + load = (top[..., None] == mx.arange(E)).astype(mx.float32).sum(-2).mean(0) / top_k + importance = flat.mean(0) + entropy = -(flat * mx.log(mx.maximum(flat, 1e-9))).sum(-1).mean(0) + source = "last micro-batch of the step" mx.eval(load, importance, entropy) r4 = lambda grid: [[round(v, 4) for v in row] for row in grid] return {"labels": labels, "experts": E, "top_k": top_k, "load": r4(load.tolist()), "importance": r4(importance.tolist()), "entropy": [round(v, 4) for v in entropy.tolist()], - "entropy_max": round(math.log(E), 4), "source": "last micro-batch of the step"} + "entropy_max": round(math.log(E), 4), "source": source} def code_fingerprint(root: str | None = None) -> dict: diff --git a/chronos/mlx/training/trainer.py b/chronos/mlx/training/trainer.py index 805d215..eadeaa5 100644 --- a/chronos/mlx/training/trainer.py +++ b/chronos/mlx/training/trainer.py @@ -25,6 +25,7 @@ from chronos.mlx.model import ChronosMLXModel from chronos.mlx.moe import ChronosMLXMOE + from chronos.mlx.multimodal import batch_modalities from chronos.mlx.training.io import mlx_state_to_torch from chronos.mlx.training.optim import MixedPrecisionAdamW except ModuleNotFoundError: @@ -34,6 +35,7 @@ tree_map = None ChronosMLXModel = None ChronosMLXMOE = None + batch_modalities = None mlx_state_to_torch = None @@ -109,14 +111,38 @@ def _first_iteration_router_probs(model) -> mx.array | None: return mx.stack(probs, axis=2) if probs else None -def _temporal_locality_loss(router_mean: mx.array) -> mx.array: +def _masked_mean(values: mx.array, mask: mx.array | None) -> mx.array: + """Mean of ``values`` over positions where ``mask`` (same shape, 1 = real) is set.""" + if mask is None: + return values.mean() + m = mask.astype(mx.float32) + return (values * m).sum() / mx.maximum(m.sum(), 1.0) + + +def _pair_mask(token_mask: mx.array | None, k: int = 1) -> mx.array | None: + """``[B, S-k]``: position t and t+k are both real tokens.""" + if token_mask is None: + return None + m = token_mask.astype(mx.float32) + return m[:, :-k] * m[:, k:] + + +def _group_mask(token_mask: mx.array | None, groups: int) -> mx.array | None: + """Repeat a ``[B, S]`` mask for the ``B*G`` rows of ``_flatten_lookahead_groups``.""" + if token_mask is None or groups <= 1: + return token_mask + return mx.repeat(token_mask, groups, axis=0) + + +def _temporal_locality_loss(router_mean: mx.array, token_mask: mx.array | None = None) -> mx.array: if router_mean.shape[1] <= 1: return mx.zeros(()) diff = router_mean[:, 1:, :] - router_mean[:, :-1, :] - return (diff * diff).sum(axis=-1).mean() + return _masked_mean((diff * diff).sum(axis=-1), _pair_mask(token_mask)) -def _router_locality_loss(router_mean: mx.array, num_experts_per_tok: int) -> mx.array: +def _router_locality_loss(router_mean: mx.array, num_experts_per_tok: int, + token_mask: mx.array | None = None) -> mx.array: if router_mean.shape[1] <= 1: return mx.zeros(()) B, _S, E = router_mean.shape @@ -126,7 +152,7 @@ def _router_locality_loss(router_mean: mx.array, num_experts_per_tok: int) -> mx target_ids = mx.argpartition(-prev, kth=top_k - 1, axis=-1)[:, :, :top_k] target_mask = (target_ids[..., None] == mx.arange(E)).astype(cur.dtype).sum(axis=-2) mass = (cur * target_mask).sum(axis=-1) - return (-mx.log(mx.maximum(mass / float(top_k), 1e-9))).mean() + return _masked_mean(-mx.log(mx.maximum(mass / float(top_k), 1e-9)), _pair_mask(token_mask)) _MEMBERSHIP_CACHE: dict = {} @@ -144,8 +170,15 @@ def _cluster_membership(config) -> mx.array: return hit -def _cluster_locality_loss(router_4d: mx.array, membership: mx.array, window: int = 0) -> mx.array: - """MLX twin of ``chronos.model.temporal_loss.cluster_locality_loss``: ``1 − I(window; cluster)/log C``.""" +def _cluster_locality_loss(router_4d: mx.array, membership: mx.array, window: int = 0, + token_mask: mx.array | None = None) -> mx.array: + """MLX twin of ``chronos.model.temporal_loss.cluster_locality_loss``: ``1 − I(window; cluster)/log C``. + + ``token_mask`` ``[B, S]`` (1 = real token) drops right-padding: a window's + cluster distribution averages its real tokens only, and windows count in + proportion to them, so pad runs are neither a "domain" of their own nor + dilute the windows they share. ``None`` = every position is real. + """ B, S, V, E = router_4d.shape C = int(membership.shape[1]) w = S if int(window or 0) <= 0 else min(int(window), S) @@ -153,26 +186,94 @@ def _cluster_locality_loss(router_4d: mx.array, membership: mx.array, window: in if C <= 1 or n <= 0: return mx.zeros(()) mass = (router_4d.astype(mx.float32).reshape(-1, E) @ membership).reshape(B, S, V, C) - pw = mass[:, : n * w].reshape(B, n, w, V, C).mean(axis=(2, 3)).reshape(B * n, C) + mass = mass[:, : n * w].reshape(B, n, w, V, C) + + def entropy(p): + p = mx.maximum(p, 1e-9) + return -(p * mx.log(p)).sum(axis=-1) + + if token_mask is None: + pw = mass.mean(axis=(2, 3)).reshape(B * n, C) + mutual = entropy(pw.mean(axis=0)) - entropy(pw).mean() + return 1.0 - mutual / math.log(C) + m = token_mask[:, : n * w].astype(mx.float32).reshape(B, n, w) + count = m.sum(axis=2) # [B, n] real tokens per window + pw = ((mass * m[..., None, None]).sum(axis=2).mean(axis=2) + / mx.maximum(count, 1.0)[..., None]).reshape(B * n, C) + weight = (count / mx.maximum(count.sum(), 1.0)).reshape(B * n) # windows ∝ real tokens + marginal = (pw * weight[:, None]).sum(axis=0) + mutual = entropy(marginal) - (entropy(pw) * weight).sum() + return 1.0 - mutual / math.log(C) + + +def _segment_cluster_loss(router_4d: mx.array, membership: mx.array, segments, + token_mask: mx.array | None = None) -> mx.array: + """``1 − I(segment; cluster)/log C`` with documents as the unit (MLX twin of + ``chronos.model.temporal_loss.segment_cluster_loss``). + + ``segments`` ``[B, S]`` comes from ``chronos.expert_clusters.segment_ids``: + the token stream's own document separators, no label or task token. Each + segment's cluster mass is pooled over its real tokens and every dispatch; + segments weigh by real-token count.""" + from chronos.expert_clusters import segment_index + B, S, V, E = router_4d.shape + C = int(membership.shape[1]) + inv, n = segment_index(segments) + if C <= 1 or n <= 1: + return mx.zeros(()) + mass = (router_4d.astype(mx.float32).reshape(-1, E) @ membership).reshape(B * S, V, C).mean(axis=1) + m = (mx.ones((B * S,), mx.float32) if token_mask is None + else token_mask.reshape(B * S).astype(mx.float32)) + idx = mx.array(inv) + total = mx.zeros((n, C), mx.float32).at[idx].add(mass * m[:, None]) + count = mx.zeros((n,), mx.float32).at[idx].add(m) + ps = total / mx.maximum(count, 1.0)[:, None] + weight = count / mx.maximum(count.sum(), 1.0) def entropy(p): p = mx.maximum(p, 1e-9) return -(p * mx.log(p)).sum(axis=-1) - mutual = entropy(pw.mean(axis=0)) - entropy(pw).mean() + marginal = (ps * weight[:, None]).sum(axis=0) + mutual = entropy(marginal) - (entropy(ps) * weight).sum() return 1.0 - mutual / math.log(C) -def _load_balance_loss(router_4d: mx.array, num_experts_per_tok: int) -> mx.array: +def _cluster_objective(router_4d: mx.array, config, token_mask=None, input_ids=None) -> mx.array: + """The configured cluster-locality term: ``window`` (fixed token windows) + or ``segment`` (documents found from ``input_ids``; window form without them).""" + from chronos.expert_clusters import cluster_locality_mode, segment_ids + window = int(getattr(config, "cluster_locality_window", 0) or 0) + membership = _cluster_membership(config) + if cluster_locality_mode(config) == "segment" and input_ids is not None: + return _segment_cluster_loss(router_4d, membership, segment_ids(input_ids, window), token_mask) + return _cluster_locality_loss(router_4d, membership, window, token_mask) + + +def _load_balance_loss(router_4d: mx.array, num_experts_per_tok: int, + token_mask: mx.array | None = None) -> mx.array: + """Switch-style ``E · Σ load·importance`` per virtual layer, summed. + + ``token_mask`` ``[B, S]`` (1 = real token) restricts load and importance + to real tokens; ``None`` = every position counts (packed pretraining).""" B, S, L, E = router_4d.shape top_k = max(1, min(int(num_experts_per_tok or 1), E)) total = mx.zeros(()) + w = None + if token_mask is not None: + w = token_mask.reshape(B * S).astype(router_4d.dtype) + w = w / mx.maximum(w.sum(), 1.0) for li in range(L): scores = router_4d[:, :, li, :].reshape(B * S, E) topk_idx = mx.argpartition(-scores, kth=top_k - 1, axis=-1)[:, :top_k] selected = (topk_idx[..., None] == mx.arange(E)).astype(scores.dtype) - load = selected.mean(axis=0) # [K, E], matching the PyTorch trainer - total = total + (load * scores.mean(axis=0)).sum() * E + if w is None: + load = selected.mean(axis=0) # [K, E], matching the PyTorch trainer + importance = scores.mean(axis=0) + else: + load = (selected * w[:, None, None]).sum(axis=0) + importance = (scores * w[:, None]).sum(axis=0) + total = total + (load * importance).sum() * E return total @@ -202,8 +303,11 @@ def _lookahead_supervision_loss( lookahead_probs: mx.array | None, teacher_router_probs: mx.array | None, lookahead_steps: int, + token_mask: mx.array | None = None, ) -> mx.array: + groups = lookahead_probs.shape[3] if lookahead_probs is not None and lookahead_probs.ndim == 5 else 1 lookahead_probs, teacher_router_probs = _flatten_lookahead_groups(lookahead_probs, teacher_router_probs) + token_mask = _group_mask(token_mask, groups) if lookahead_probs is None: return mx.zeros(()) if teacher_router_probs is None: @@ -220,7 +324,8 @@ def _lookahead_supervision_loss( continue teacher_k = teacher[:, k:, :] pred_k = lookahead_probs[:, :-k, k, :].astype(mx.float32) - total = total + (-(teacher_k * mx.log(mx.maximum(pred_k, 1e-9))).sum(axis=-1)).mean() + total = total + _masked_mean(-(teacher_k * mx.log(mx.maximum(pred_k, 1e-9))).sum(axis=-1), + _pair_mask(token_mask, k)) terms += 1 return total / max(terms, 1) @@ -230,8 +335,11 @@ def _lookahead_topk_hit_loss( teacher_router_probs: mx.array | None, lookahead_steps: int, num_experts_per_tok: int, + token_mask: mx.array | None = None, ) -> mx.array: + groups = lookahead_probs.shape[3] if lookahead_probs is not None and lookahead_probs.ndim == 5 else 1 lookahead_probs, teacher_router_probs = _flatten_lookahead_groups(lookahead_probs, teacher_router_probs) + token_mask = _group_mask(token_mask, groups) if lookahead_probs is None: return mx.zeros(()) if teacher_router_probs is None: @@ -252,7 +360,7 @@ def _lookahead_topk_hit_loss( target_ids = mx.argpartition(-teacher_k, kth=top_k - 1, axis=-1)[:, :, :top_k] target_mask = (target_ids[..., None] == mx.arange(E)).astype(pred_k.dtype).sum(axis=-2) hit_mass = (pred_k * target_mask).sum(axis=-1) - total = total + (-mx.log(mx.maximum(hit_mass / float(top_k), 1e-9))).mean() + total = total + _masked_mean(-mx.log(mx.maximum(hit_mass / float(top_k), 1e-9)), _pair_mask(token_mask, k)) terms += 1 return total / max(terms, 1) @@ -262,8 +370,11 @@ def _lookahead_union_loss( teacher_router_probs: mx.array | None, lookahead_steps: int, num_experts_per_tok: int, + token_mask: mx.array | None = None, ) -> mx.array: + groups = lookahead_probs.shape[3] if lookahead_probs is not None and lookahead_probs.ndim == 5 else 1 lookahead_probs, teacher_router_probs = _flatten_lookahead_groups(lookahead_probs, teacher_router_probs) + token_mask = _group_mask(token_mask, groups) if lookahead_probs is None: return mx.zeros(()) if teacher_router_probs is None: @@ -282,7 +393,8 @@ def _lookahead_union_loss( pred_union = lookahead_probs[:, :S - K, 1:K + 1, :].astype(mx.float32).max(axis=2) mass = (pred_union * union_target).sum(axis=-1) target_count = mx.maximum(union_target.sum(axis=-1), 1.0) - return (-mx.log(mx.maximum(mass / target_count, 1e-9))).mean() + # Union over t+1..t+K: position t counts when it and t+K are real (right padding). + return _masked_mean(-mx.log(mx.maximum(mass / target_count, 1e-9)), _pair_mask(token_mask, K)) def _offload_training_metrics( @@ -369,7 +481,7 @@ def draft_exit_active(model, config) -> bool: return 1 <= int(getattr(config, "draft_exit_depth", 1)) < int(config.recurrent_depth) -def draft_exit_logits(model, input_ids, config): +def draft_exit_logits(model, input_ids, config, media=None): """Draft-depth FP32 logits, or None when the KL hook is off. Run *before* the full-depth forward: the full pass then owns the routing @@ -378,7 +490,7 @@ def draft_exit_logits(model, input_ids, config): if not draft_exit_active(model, config): return None logits, _lookahead, _cache = model(input_ids, logits_dtype=mx.float32, - recurrent_depth=int(config.draft_exit_depth)) + recurrent_depth=int(config.draft_exit_depth), **(media or {})) return logits @@ -395,6 +507,22 @@ def draft_exit_kl(draft_logits: mx.array, full_logits: mx.array, keep: mx.array return (kl * keep).sum() / mx.maximum(keep.sum(), 1.0) +def real_token_mask(labels: mx.array, *, aligned: bool = False) -> mx.array: + """``[B, S]`` bool: input positions holding real tokens of a right-padded batch. + + Padding carries label -100 after each row's last supervised token, so a + row's real tokens end at that token: ``labels[t]`` is token ``t`` itself + (``aligned=False``, HF convention, the loss shifts) or token ``t+1`` + (``aligned=True``, labels pre-shifted onto the logits). Prompt tokens + (-100 before the reply) stay real: they are routed context. Packed + pretraining (no -100) gives an all-true mask. + """ + labels = mx.array(labels) if not isinstance(labels, mx.array) else labels + pos = mx.arange(labels.shape[1]) + last = mx.max(mx.where(labels != -100, pos[None, :], -1), axis=1) + return pos[None, :] <= (last + (1 if aligned else 0))[:, None] + + def _shifted_keep(labels: mx.array) -> mx.array: """Logit t is supervised when label t+1 is not -100 (CE shift).""" nxt = labels[:, 1:] != -100 @@ -402,8 +530,13 @@ def _shifted_keep(labels: mx.array) -> mx.array: def _chronos_regularized_loss(model, base_loss, lookahead_probs, config, router_ref=None, metrics=None, - draft_kl=None): - """Return total loss and retain the actual weighted components for logging.""" + draft_kl=None, token_mask=None, input_ids=None): + """Return total loss and retain the actual weighted components for logging. + + ``token_mask`` ``[B, S]`` (1 = real token, 0 = right padding) keeps padding + out of every router regularizer; ``None`` = all positions are real. + ``input_ids`` ``[B, S]`` lets ``cluster_locality_mode="segment"`` find the + documents of the batch (without it the window form is used).""" object.__setattr__(model, "_last_lookahead_probs", lookahead_probs) terms = {"primary_loss": base_loss.astype(mx.float32)} def add(name, coefficient, value): @@ -422,24 +555,23 @@ def add(name, coefficient, value): if router is not None: router = router.astype(mx.float32) top_k = config.num_experts_per_tok - add("balance_loss", config.lambda_balance, _load_balance_loss(router, top_k)) + add("balance_loss", config.lambda_balance, _load_balance_loss(router, top_k, token_mask)) if float(getattr(config, "lambda_cluster_locality", 0.0) or 0.0) > 0.0: add("cluster_locality_loss", config.lambda_cluster_locality, - _cluster_locality_loss(router, _cluster_membership(config), - int(getattr(config, "cluster_locality_window", 0) or 0))) + _cluster_objective(router, config, token_mask, input_ids)) if router.shape[1] > 1: mean = router.mean(axis=2) - add("temporal_loss", config.lambda_temporal, _temporal_locality_loss(mean)) - add("locality_loss", config.lambda_router_locality, _router_locality_loss(mean, top_k)) + add("temporal_loss", config.lambda_temporal, _temporal_locality_loss(mean, token_mask)) + add("locality_loss", config.lambda_router_locality, _router_locality_loss(mean, top_k, token_mask)) if lookahead_probs is not None and config.lambda_lookahead > 0 and config.lookahead_steps > 0: teacher = _lookahead_teacher(_first_iteration_router_probs(model).astype(mx.float32), lookahead_probs, config) add("lookahead_loss", config.lambda_lookahead, - _lookahead_supervision_loss(lookahead_probs, teacher, config.lookahead_steps)) + _lookahead_supervision_loss(lookahead_probs, teacher, config.lookahead_steps, token_mask)) add("lookahead_topk_loss", config.lambda_lookahead_topk, - _lookahead_topk_hit_loss(lookahead_probs, teacher, config.lookahead_steps, top_k)) + _lookahead_topk_hit_loss(lookahead_probs, teacher, config.lookahead_steps, top_k, token_mask)) add("lookahead_union_loss", config.lambda_lookahead_union, - _lookahead_union_loss(lookahead_probs, teacher, config.lookahead_steps, top_k)) + _lookahead_union_loss(lookahead_probs, teacher, config.lookahead_steps, top_k, token_mask)) add("anchor_loss", 1, _router_kl_anchor(mean, router_ref, config.lambda_router_anchor)) if draft_kl is not None: add("draft_exit_kl", getattr(config, "draft_exit_kl_coef", 0.0), draft_kl) @@ -449,13 +581,14 @@ def add(name, coefficient, value): return loss -def _ce_stage_loss(model, input_ids, labels): - draft = draft_exit_logits(model, input_ids, model.config) - logits, lookahead, _cache = model(input_ids, logits_dtype=mx.float32) +def _ce_stage_loss(model, input_ids, labels, modality_inputs=None, modality_mask=None): + media = {"modality_inputs": modality_inputs, "modality_mask": modality_mask} if modality_inputs else None + draft = draft_exit_logits(model, input_ids, model.config, media) + logits, lookahead, _cache = model(input_ids, logits_dtype=mx.float32, **(media or {})) base = _masked_lm_loss(logits, labels) draft_kl = None if draft is None else draft_exit_kl(draft, logits, _shifted_keep(labels)) return _chronos_regularized_loss(model, base, lookahead, model.config, metrics={"ce": base, "primary": base}, - draft_kl=draft_kl) + draft_kl=draft_kl, token_mask=real_token_mask(labels), input_ids=input_ids) def _dpo_stage_loss(model, reference, batch, beta: float = 0.1, router_ref=None): @@ -528,7 +661,9 @@ def _distill_stage_loss( kd = (temperature * temperature) * (kd_tok * mask).sum() / mx.maximum(mask.sum(), 1.0) ce = _masked_lm_loss(logits, labels) base = alpha * kd + (1.0 - alpha) * ce - return _chronos_regularized_loss(model, base, lookahead, model.config, router_ref=router_ref, metrics={"primary": kd, "secondary": ce, "ce": ce}) + return _chronos_regularized_loss(model, base, lookahead, model.config, router_ref=router_ref, + metrics={"primary": kd, "secondary": ce, "ce": ce}, + token_mask=real_token_mask(labels)) def _fit_teacher_input(ids: mx.array, teacher_vocab_size: int) -> mx.array: @@ -558,13 +693,15 @@ def _grpo_stage_loss(model, reference, batch, beta: float = 0.04, router_ref=Non def _capture_router_ref(model, batch, stage: str) -> mx.array | None: if stage in {"pretrain", "grpo"}: return None + media, media_mask = None, None if stage in {"sft", "distill"}: - ids = _ids_to_mx(batch[0]) + raw_ids, _labels, media, media_mask = batch_modalities(batch) + ids = _ids_to_mx(raw_ids) elif stage in {"dpo", "orpo"}: ids = mx.concatenate([_ids_to_mx(batch["x_chosen"]), _ids_to_mx(batch["x_rejected"])], axis=0) else: return None - _logits, _lookahead, _cache = model(ids) + _logits, _lookahead, _cache = model(ids, modality_inputs=media, modality_mask=media_mask) router_4d = _collect_router_probs(model) if router_4d is None: return None @@ -925,14 +1062,19 @@ def _save_checkpoint(reason: str) -> None: def loss_fn(m, batch): if stage in {"pretrain", "sft"}: - ids, labels = batch + ids, labels, media, media_mask = batch_modalities(batch) + if media: + return _ce_stage_loss(m, _ids_to_mx(ids), _ids_to_mx(labels), + modality_inputs=media, modality_mask=media_mask) return _ce_stage_loss(m, _ids_to_mx(ids), _ids_to_mx(labels)) if stage == "dpo": return _dpo_stage_loss(m, reference, batch, beta=beta, router_ref=router_ref) if stage == "orpo": return _orpo_stage_loss(m, batch, beta=beta, lambda_or=lambda_or, router_ref=router_ref) if stage == "distill": - ids, labels = batch + ids, labels, media, _media_mask = batch_modalities(batch) + if media: + raise ValueError("MLX distill does not take modality inputs yet; use pretrain/sft or the PyTorch backend") ids = _ids_to_mx(ids) labels = _ids_to_mx(labels) teacher_vocab = int(getattr(getattr(teacher, "config", config), "vocab_size", config.vocab_size)) @@ -1019,9 +1161,11 @@ def evaluate_validation(epoch): was_training = model.training model.eval() try: - for ids, labels in validation_loader: + for vbatch in validation_loader: + ids, labels, media, media_mask = batch_modalities(vbatch) labels = _ids_to_mx(labels) - logits, _, _ = model(_ids_to_mx(ids), logits_dtype=mx.float32) + logits, _, _ = model(_ids_to_mx(ids), logits_dtype=mx.float32, + modality_inputs=media, modality_mask=media_mask) ce = _masked_lm_loss(logits, labels) count = mx.sum(labels[:, 1:] != -100) mx.eval(ce, count) @@ -1063,7 +1207,8 @@ def evaluate_validation(epoch): # supervised-token count are materialised by one device sync. finite_check = _tree_finite_check((loss, grads)) candidate = _add_grads(pending_grads, grads) - n_supervised = mx.sum(_ids_to_mx(train_batch[1])[:, 1:] != -100) if supervised else None + n_supervised = (mx.sum(_ids_to_mx(batch_modalities(train_batch)[1])[:, 1:] != -100) + if supervised else None) mx.eval(loss, finite_check, candidate, *([n_supervised] if supervised else [])) steps += 1 finite = bool(finite_check.item()) diff --git a/chronos/model/config.py b/chronos/model/config.py index b964770..c9cdfaf 100644 --- a/chronos/model/config.py +++ b/chronos/model/config.py @@ -208,6 +208,11 @@ def __init__(self, **kwargs): self.lambda_cluster_locality: float = kwargs.pop("lambda_cluster_locality", 0.0) self.cluster_locality_clusters: int = kwargs.pop("cluster_locality_clusters", 0) self.cluster_locality_window: int = kwargs.pop("cluster_locality_window", 32) + # Unit of the constraint: "window" (fixed windows above) or "segment" + # (the batch's documents, read from the tokens' own separators; a + # document longer than the window is cut into window-sized parts). + # Neither consults labels or task tokens. + self.cluster_locality_mode: str = str(kwargs.pop("cluster_locality_mode", "window") or "window") # M4: weight on the router KL anchor (alignment stages only). # Prevents SFT/DPO/ORPO/GRPO gradients from drifting routing away # from the pretrained distribution that the cluster layout was diff --git a/chronos/model/temporal_loss.py b/chronos/model/temporal_loss.py index 61bdc97..b24141d 100644 --- a/chronos/model/temporal_loss.py +++ b/chronos/model/temporal_loss.py @@ -116,6 +116,42 @@ def cluster_locality_loss( return 1.0 - mutual / log_c +def segment_cluster_loss(router_4d: torch.Tensor, membership: torch.Tensor, segments, + token_mask: torch.Tensor | None = None) -> torch.Tensor: + """``1 − I(segment; cluster)/log C`` with the batch's documents as streams. + + ``segments`` ``[B, S]`` is ``chronos.expert_clusters.segment_ids(input_ids, + window)``: documents found from the tokens' own separators, long ones cut + into window-sized parts. Each segment's cluster mass is averaged over its + real tokens (``token_mask``, default all) and every dispatch; segments + weigh by real-token count. A fixed window that straddles two packed + documents rewards routing both into one cluster, so the split it learns + can follow position instead of content; with documents as the unit, the + way to keep each one inside a cluster while the batch spreads over all + clusters is to send similar documents to the same cluster. + """ + from chronos.expert_clusters import segment_index + if router_4d is None or router_4d.dim() != 4: + return torch.tensor(0.0) + B, S, V, E = router_4d.shape + C = int(membership.shape[1]) + inv, n = segment_index(segments) + if C <= 1 or n <= 1: + return router_4d.new_zeros(()) + mass = torch.einsum("bsve,ec->bsc", router_4d.float(), + membership.to(router_4d.device, torch.float32)).reshape(B * S, C) / V + m = (torch.ones(B * S, device=router_4d.device) if token_mask is None + else token_mask.reshape(B * S).to(router_4d.device, torch.float32)) + idx = torch.as_tensor(inv, device=router_4d.device, dtype=torch.long) + total = mass.new_zeros(n, C).index_add(0, idx, mass * m[:, None]) + count = mass.new_zeros(n).index_add(0, idx, m) + ps = total / count.clamp_min(1.0)[:, None] + weight = count / count.sum().clamp_min(1.0) + marginal = (ps * weight[:, None]).sum(dim=0) + mutual = _entropy(marginal) - (_entropy(ps) * weight).sum() + return 1.0 - mutual / math.log(C) + + def cluster_routing_metrics(router_4d: torch.Tensor, membership: torch.Tensor, top_k: int, window: int = 0) -> dict[str, float]: """Detached diagnostics of the cluster constraint (normalised entropies, clusters touched).""" diff --git a/chronos/monitor/snapshot.py b/chronos/monitor/snapshot.py index 3c22b13..44db490 100644 --- a/chronos/monitor/snapshot.py +++ b/chronos/monitor/snapshot.py @@ -165,7 +165,62 @@ def _status(self, path: str, spec: dict | None, rows: list[dict]) -> dict: # ---- progress / ETA ------------------------------------------------------ @staticmethod - def _progress(spec: dict | None, status: dict, latest: dict | None, rows: list[dict]) -> dict: + def _rate(prows: list[dict], key: str, horizon_s: float = 1800.0): + """Units of ``key`` per wall second over the latest rows (≤ ``horizon_s``). + + Wall time between log rows includes validation, checkpoints and + telemetry, so the rate already prices them in.""" + pts = [] + for r in prows: + v, ts = r.get(key), r.get("time") + if v is None or ts is None: + continue + try: + pts.append((datetime.fromisoformat(str(ts)).timestamp(), float(v))) + except ValueError: + continue + if len(pts) < 2: + return None + t_end, v_end = pts[-1] + first = next((p for p in pts if t_end - p[0] <= horizon_s), pts[0]) + if first is pts[-1]: + first = pts[-2] + dt, dv = t_end - first[0], v_end - first[1] + return dv / dt if dt > 0 and dv > 0 else None + + @classmethod + def _phase_eta(cls, phase: dict, prows: list[dict], hb: dict, active: bool, spec: dict | None) -> dict: + """Remaining wall seconds of one phase, from measured progress. + + A phase ends at the first of: its step total (step schedule), its data + target (``extra.target_passes``), or its wall-clock deadline. Each + measured candidate is projected at the recent rate; the smallest wins + and is reported with its basis. The deadline always bounds it.""" + last = prows[-1] if prows else {} + cands = {} + if active and hb.get("phase_wall_s") is not None and hb.get("phase_budget_s"): + cands["deadline"] = max(0.0, float(hb["phase_budget_s"]) - float(hb["phase_wall_s"])) + if (spec or {}).get("schedule", "time") != "time" and last.get("total"): + rate = cls._rate(prows, "step") + if rate: + cands["steps"] = max(0, int(last["total"]) - int(last.get("step") or 0)) / rate + target = (((spec or {}).get("extra") or {}).get("target_passes") or {}).get(phase.get("name")) + if target and last.get("passes") is not None: + rate = cls._rate(prows, "passes") + if rate: + cands["target_passes"] = max(0.0, float(target) - float(last["passes"])) / rate + if not cands and last.get("total") and last.get("step_s"): + cands["steps"] = max(0, int(last["total"]) - int(last.get("step") or 0)) * float(last["step_s"]) + if not cands: + return {} + basis = min(cands, key=cands.get) + out = {"remaining_s": cands[basis], "basis": basis} + if "deadline" in cands and basis != "deadline": + out["deadline_remaining_s"] = cands["deadline"] + return out + + @classmethod + def _progress(cls, spec: dict | None, status: dict, latest: dict | None, rows: list[dict]) -> dict: phases = (spec or {}).get("phases") or [] hb = status.get("heartbeat") or {} budget_h = float((spec or {}).get("budget_hours") or 0.0) @@ -174,6 +229,11 @@ def _progress(spec: dict | None, status: dict, latest: dict | None, rows: list[d current = hb.get("phase") or (latest or {}).get("phase") names = [p.get("name") for p in phases] index = names.index(current) if current in names else 0 + finished = status.get("state") == "finished" + # A clock ETA needs a live process; a stopped or failed run keeps its + # remaining time (work left at the measured rate) but no end time. + moving = status.get("state") in ("running", "starting", "evaluating", "checkpointing") + now = time.time() out = [] for i, phase in enumerate(phases): prows = [r for r in rows if r.get("phase") == phase.get("name") and "event" not in r] @@ -181,20 +241,48 @@ def _progress(spec: dict | None, status: dict, latest: dict | None, rows: list[d entry = {"name": phase.get("name"), "budget_h": round(train_s * float(phase.get("budget_fraction", 0)) / 3600, 2), "step": last.get("step", 0), "total": last.get("total"), "state": "complete" if i < index else ("active" if i == index else "pending")} - if i == index and status.get("state") == "finished": + if i == index and finished: entry["state"] = "complete" + if last.get("passes") is not None: + entry["passes"] = last["passes"] out.append(entry) result = {"phases": out, "current": current, "budget_hours": budget_h, "train_budget_hours": round(train_s / 3600, 2)} if hb.get("phase_wall_s") is not None and hb.get("phase_budget_s"): result["phase_fraction"] = min(1.0, float(hb["phase_wall_s"]) / float(hb["phase_budget_s"])) - result["phase_remaining_s"] = max(0.0, float(hb["phase_budget_s"]) - float(hb["phase_wall_s"])) elif latest and latest.get("total"): result["phase_fraction"] = min(1.0, latest["step"] / max(1, latest["total"])) - result["phase_remaining_s"] = max(0, latest["total"] - latest["step"]) * float(latest.get("step_s") or 0) - if result.get("phase_remaining_s") is not None: - later = sum(e["budget_h"] for e in out[index + 1:]) * 3600 - result["run_remaining_s"] = result["phase_remaining_s"] + later + budget_h * 3600 * reserve - result["eta"] = datetime.fromtimestamp(time.time() + result["run_remaining_s"]).isoformat(timespec="minutes") + if not finished and out: + cur_rows = [r for r in rows if r.get("phase") == current and "event" not in r] + eta = cls._phase_eta(phases[index], cur_rows, hb, True, spec) + if eta: + result["phase_remaining_s"] = eta["remaining_s"] + result["eta_basis"] = eta["basis"] + if "deadline_remaining_s" in eta: + result["phase_deadline_remaining_s"] = eta["deadline_remaining_s"] + out[index]["remaining_s"] = round(eta["remaining_s"], 1) + if moving: + out[index]["eta"] = datetime.fromtimestamp(now + eta["remaining_s"]).isoformat(timespec="minutes") + out[index]["eta_basis"] = eta["basis"] + # A phase that ends before its deadline hands the rest of its + # time to the next one (deadlines are fixed from the run start), + # so later phases are bounded by their deadlines; the run ends + # at the last one plus the final-eval reserve. + cursor = eta["remaining_s"] + slack = max(0.0, eta.get("deadline_remaining_s", eta["remaining_s"]) - eta["remaining_s"]) + for j in range(index + 1, len(out)): + span = out[j]["budget_h"] * 3600 + slack + target = (((spec or {}).get("extra") or {}).get("target_passes") or {}).get(out[j]["name"]) + out[j]["remaining_s"] = round(cursor + span, 1) + if moving: + out[j]["eta"] = datetime.fromtimestamp(now + cursor + span).isoformat(timespec="minutes") + out[j]["eta_basis"] = "deadline" if not target else "deadline (target_passes may end it earlier)" + cursor += span + slack = 0.0 + result["run_remaining_s"] = cursor + budget_h * 3600 * reserve + if moving: + result["eta"] = datetime.fromtimestamp(now + result["run_remaining_s"]).isoformat(timespec="minutes") + else: + result["eta_paused"] = status.get("state") if hb.get("run_wall_s") is not None: result["run_wall_s"] = hb["run_wall_s"] return result diff --git a/chronos/monitor/static/app.js b/chronos/monitor/static/app.js index 2fbd3f2..20bff24 100644 --- a/chronos/monitor/static/app.js +++ b/chronos/monitor/static/app.js @@ -18,7 +18,7 @@ const STR = { n_stale: "最近 {age} 没有新的日志行（正常间隔约 {int}）。训练进程可能卡住或被挂起。", n_derived: "此运行由旧版代码启动，没有 status.json / parameters.jsonl。状态由指标流和进程表推断；张量统计、路由热图和评测探针会在下次以新代码恢复（resume）后出现。", n_error: "运行报告错误：", n_paused: "训练进程处于 SIGSTOP 挂起状态，恢复后会继续。", n_none: "没有找到运行目录。用 chronos longrun 启动训练后，这里会自动出现。", - phase: "阶段", step: "步", of: "/", budget: "预算", remaining: "剩余", eta: "预计完成", wall: "已用", elapsed: "已运行", + phase: "阶段", step: "步", of: "/", budget: "预算", remaining: "剩余", eta: "预计完成", basis_deadline: "按时间预算", basis_steps: "按步数进度", basis_target_passes: "按数据轮次进度", phase_eta: "本阶段预计 {t}", wall: "已用", elapsed: "已运行", ph_complete: "完成", ph_active: "进行中", ph_pending: "未开始", pretrain: "预训练", sft: "SFT", k_loss: "训练损失", k_ce: "主损失 (CE)", k_valid: "最佳验证 CE", k_tps: "吞吐", k_eta: "预计完成", k_mem: "内存", k_grad: "梯度范数", k_depth: "平均循环深度", k_lr: "学习率", k_route: "路由熵", @@ -60,7 +60,7 @@ const STR = { n_stale: "No new log line for {age} (usual interval about {int}). The trainer may be stuck or suspended.", n_derived: "This run was started by older code and writes no status.json / parameters.jsonl. Status is derived from the metric stream and the process table; tensor stats, routing heatmaps and evaluation probes appear after the next resume with the new code.", n_error: "The run reported an error: ", n_paused: "The trainer is suspended (SIGSTOP) and will continue when resumed.", n_none: "No run directories found. Runs started with chronos longrun show up here.", - phase: "Phase", step: "step", of: "/", budget: "budget", remaining: "remaining", eta: "ETA", wall: "used", elapsed: "Elapsed", + phase: "Phase", step: "step", of: "/", budget: "budget", remaining: "remaining", eta: "ETA", basis_deadline: "by time budget", basis_steps: "by step progress", basis_target_passes: "by data-pass progress", phase_eta: "phase ends ~{t}", wall: "used", elapsed: "Elapsed", ph_complete: "complete", ph_active: "active", ph_pending: "pending", pretrain: "Pretrain", sft: "SFT", k_loss: "Training loss", k_ce: "Primary loss (CE)", k_valid: "Best valid CE", k_tps: "Throughput", k_eta: "ETA", k_mem: "Memory", k_grad: "Grad norm", k_depth: "Mean recurrent depth", k_lr: "Learning rate", k_route: "Routing entropy", @@ -310,7 +310,8 @@ function renderProgress() { const p = snap.progress || {}; const latest = snap.latest || {}; const pct = isNum(p.phase_fraction) ? ` · ${(p.phase_fraction * 100).toFixed(1)}%` : ""; - const right = [p.eta ? `${tr("eta")} ${fmtTime(p.eta)}` : null, + const basis = (b) => (b ? tr("basis_" + String(b).split(" ")[0]) : ""); + const right = [p.eta ? `${tr("eta")} ${fmtTime(p.eta)}${p.eta_basis ? ` (${basis(p.eta_basis)})` : ""}` : null, isNum(p.run_remaining_s) ? `${tr("remaining")} ${fmtDur(p.run_remaining_s)}` : null, isNum(p.run_wall_s) ? `${tr("wall")} ${fmtDur(p.run_wall_s)}` : null, isNum(p.budget_hours) && p.budget_hours ? `${tr("budget")} ${p.budget_hours} h` : null].filter(Boolean).join(" · "); @@ -330,7 +331,8 @@ function renderProgress() { h("div", { class: "track", role: "progressbar", "aria-label": phaseName(ph.name), "aria-valuemin": "0", "aria-valuemax": "100", "aria-valuenow": String(Math.round(frac * 100)) }, fill), h("b", { text: phaseName(ph.name) }), ` ${tr("ph_" + ph.state)} · ${ph.budget_h} h`, - ph.total ? ` · ${fmtInt(ph.step)}/${fmtInt(ph.total)}` : "")); + ph.total ? ` · ${fmtInt(ph.step)}/${fmtInt(ph.total)}` : "", + ph.eta && ph.state !== "complete" ? ` · ${tr("phase_eta", { t: fmtTime(ph.eta) })}` : "")); } box.append(phases); } @@ -369,7 +371,8 @@ function renderTiles() { tile(tr("k_valid"), best ? fmt(best.valid_ce) : "—", best ? `${tr("sub_at", { s: fmtInt(best.step) })} · ${tr("sub_ppl", { v: fmt(Math.exp(best.valid_ce), 3) })}` : null); tile(tr("k_tps"), isNum(L.tokens_per_s) ? `${fmtCount(L.tokens_per_s)} tok/s` : "—", isNum(L.processed_tokens_per_s) ? tr("sub_proc", { v: fmtCount(L.processed_tokens_per_s) }) : null, "tokens_per_s"); - tile(tr("k_eta"), p.eta ? fmtTime(p.eta) : "—", isNum(p.run_remaining_s) ? tr("sub_left", { v: fmtDur(p.run_remaining_s) }) : null); + tile(tr("k_eta"), p.eta ? fmtTime(p.eta) : (p.eta_paused ? tr(`st_${p.eta_paused}`) : "—"), + isNum(p.run_remaining_s) ? tr("sub_left", { v: fmtDur(p.run_remaining_s) }) : null); const osMem = isNum(mem.footprint_gb); // live process footprint has no history; MLX peak does tile(tr("k_mem"), osMem ? fmtGB(mem.footprint_gb) : fmtGB(L.peak_mem_gb), `${osMem ? tr("sub_footprint") + " · " : ""}${tr("sub_peak", { v: fmtGB(L.peak_mem_gb), a: fmtGB(L.active_mem_gb) })}`, osMem ? null : "peak_mem_gb"); diff --git a/chronos/serving/ollama_server.py b/chronos/serving/ollama_server.py index e5eb920..2d44006 100644 --- a/chronos/serving/ollama_server.py +++ b/chronos/serving/ollama_server.py @@ -42,6 +42,28 @@ def _digest(path: str) -> str: return hashlib.sha256(f"{path}:{st.st_size}:{st.st_mtime_ns}".encode()).hexdigest() +def _chat_messages(body: dict) -> list[dict]: + """Request messages as template input: role, content, and (when present) + the assistant's ``tool_calls`` (Ollama/OpenAI shapes; OpenAI string + arguments are parsed) so tool conversations render as they were trained.""" + out = [] + for m in body.get("messages") or []: + msg = {"role": m.get("role", "user"), "content": m.get("content") or ""} + calls = [] + for call in m.get("tool_calls") or []: + fn = dict(call.get("function") or call) + if isinstance(fn.get("arguments"), str): + try: + fn["arguments"] = json.loads(fn["arguments"]) + except ValueError: + pass + calls.append({"function": {"name": fn.get("name", ""), "arguments": fn.get("arguments", {})}}) + if calls: + msg["tool_calls"] = calls + out.append(msg) + return out + + def _options(body: dict) -> dict: opts = dict(body.get("options") or {}) return { @@ -238,8 +260,7 @@ def collect(items): async def chat(request: Request): body = await request.json() name = (body.get("model") or "").split(":")[0] - messages = [{"role": m.get("role", "user"), "content": m.get("content", "")} - for m in body.get("messages") or []] + messages = _chat_messages(body) if _unload_if_requested(body, name): return {"model": body.get("model"), "created_at": _now(), "message": {"role": "assistant", "content": ""}, "done": True, "done_reason": "unload"} @@ -247,7 +268,7 @@ async def chat(request: Request): if not messages: return {"model": body.get("model"), "created_at": _now(), "message": {"role": "assistant", "content": ""}, "done": True, "done_reason": "load"} - ids = runtime.encode_chat(lm, messages) + ids = runtime.encode_chat(lm, messages, body.get("tools")) def shape(text, final): if final is not None: @@ -271,8 +292,7 @@ async def openai_chat(request: Request): body = await request.json() name = (body.get("model") or "").split(":")[0] lm, load_s = _resolve(name) - messages = [{"role": m.get("role", "user"), "content": m.get("content", "")} for m in body.get("messages") or []] - ids = runtime.encode_chat(lm, messages) + ids = runtime.encode_chat(lm, _chat_messages(body), body.get("tools")) opts = _options({**body, "options": {"num_predict": body.get("max_tokens") or body.get("max_completion_tokens") or 256, "temperature": body.get("temperature", 0.0), "top_p": body.get("top_p", 1.0), "seed": body.get("seed"), "stop": body.get("stop"), diff --git a/chronos/serving/runtime.py b/chronos/serving/runtime.py index 3aa6c37..0a52226 100644 --- a/chronos/serving/runtime.py +++ b/chronos/serving/runtime.py @@ -245,8 +245,11 @@ def fwd(ids, cache, depth): return out return fwd - def encode_chat(self, lm: LoadedModel, messages: list[dict]) -> list[int]: - text = lm.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + def encode_chat(self, lm: LoadedModel, messages: list[dict], tools: list | None = None) -> list[int]: + # ``tools`` reach the template only when given, so templates without + # tool support render exactly as before. + extra = {"tools": tools} if tools else {} + text = lm.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, **extra) return lm.tokenizer(text, add_special_tokens=False)["input_ids"] def encode_prompt(self, lm: LoadedModel, prompt: str, system: str | None = None, raw: bool = False): diff --git a/chronos/tokenizer/model.py b/chronos/tokenizer/model.py index 09051da..6d87555 100644 --- a/chronos/tokenizer/model.py +++ b/chronos/tokenizer/model.py @@ -21,9 +21,28 @@ _BYTE_RE = re.compile(r"^<0x([0-9A-F]{2})>$") _BYTE_SPLIT_RE = re.compile(r"(<0x[0-9A-F]{2}>)") +# Plain turns render exactly as before (``<|im_start|>role\ncontent<|im_end|>\n``). +# Tools: the definitions go into one system block, an assistant turn's +# ``tool_calls`` follow its text as ``{json}`` lines, and +# ``tool`` turns carry the results. ``chronos.data.chat_records.render_chat`` +# produces the same text in Python (with per-turn spans for loss masks). CHAT_TEMPLATE = ( + "{% if tools %}{{ '<|im_start|>system\n' }}" + "{% if messages and messages[0]['role'] == 'system' %}{{ messages[0]['content'] + '\n\n' }}{% endif %}" + "{{ '\n' }}{% for tool in tools %}{{ tool | tojson + '\n' }}{% endfor %}" + "{{ '<|im_end|>\n' }}{% endif %}" "{% for message in messages %}" - "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n' }}" + "{% if not (tools and loop.first and message['role'] == 'system') %}" + "{{ '<|im_start|>' + message['role'] + '\n' + (message['content'] or '') }}" + "{% if message['role'] == 'assistant' and message.get('tool_calls') %}" + "{% set text = message['content'] or '' %}" + "{% for call in message['tool_calls'] %}" + "{% set fn = call['function'] if call.get('function') else call %}" + "{% if text or not loop.first %}{{ '\n' }}{% endif %}" + "{{ '' + {'name': fn['name'], 'arguments': fn['arguments']} | tojson + '' }}" + "{% endfor %}{% endif %}" + "{{ '<|im_end|>' + '\n' }}" + "{% endif %}" "{% endfor %}" "{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}" ) diff --git a/chronos/trainer/chronos_trainer.py b/chronos/trainer/chronos_trainer.py index d80cfbc..0cdf69d 100644 --- a/chronos/trainer/chronos_trainer.py +++ b/chronos/trainer/chronos_trainer.py @@ -107,7 +107,7 @@ def train_step(self, input_ids, labels, step, total_steps, *, modality_inputs=No loss = chronos_loss_term( self.model, ce_loss, lookahead_probs, self.config, - aux_loss=aux_loss, + aux_loss=aux_loss, input_ids=input_ids, labels=labels, ) # Standalone lookahead loss value for logging diff --git a/chronos/trainer/loss_mixin.py b/chronos/trainer/loss_mixin.py index 55774f8..691cf9f 100644 --- a/chronos/trainer/loss_mixin.py +++ b/chronos/trainer/loss_mixin.py @@ -34,6 +34,7 @@ lookahead_union_loss, router_locality_loss, router_offload_metrics, + segment_cluster_loss, ) @@ -62,6 +63,16 @@ def collect_router_probs(model, *, with_grad: bool = False) -> Optional[torch.Te return torch.stack(probs, dim=2) +def real_token_mask_torch(labels: torch.Tensor) -> torch.Tensor: + """``[B, S]`` bool real-token mask of a right-padded batch (HF labels: the + loss shifts). A row's real tokens end at its last supervised label; later + positions are padding. Prompt tokens stay real. Same rule as the MLX + ``trainer.real_token_mask(labels)``.""" + pos = torch.arange(labels.shape[1], device=labels.device) + last = torch.where(labels != -100, pos[None, :], torch.full_like(labels, -1)).max(dim=1).values + return pos[None, :] <= last[:, None] + + def cluster_membership_tensor(config, device=None) -> torch.Tensor: """``[E, C]`` one-hot of ``chronos.expert_clusters.training_clusters(config)``.""" from chronos.expert_clusters import cluster_membership, training_clusters @@ -115,6 +126,8 @@ def chronos_loss_term( config, *, aux_loss: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Combine `base_loss` with the Chronos-specific regularizers: @@ -127,6 +140,10 @@ def chronos_loss_term( The function is safe to call from any stage: it gracefully degrades when the necessary tensors are unavailable (e.g. lookahead_probs is None during a model that doesn't expose it). + + ``input_ids`` (and ``labels``, for the padding mask) let + ``cluster_locality_mode="segment"`` find the documents of the batch; + without them the window form is used. """ loss = base_loss @@ -163,11 +180,16 @@ def chronos_loss_term( ) lambda_cluster = float(getattr(config, "lambda_cluster_locality", 0.0) or 0.0) if lambda_cluster > 0.0: - loss = loss + lambda_cluster * cluster_locality_loss( - router_4d, - cluster_membership_tensor(config, router_4d.device), - int(getattr(config, "cluster_locality_window", 0) or 0), - ) + from chronos.expert_clusters import cluster_locality_mode, segment_ids + membership = cluster_membership_tensor(config, router_4d.device) + window = int(getattr(config, "cluster_locality_window", 0) or 0) + if cluster_locality_mode(config) == "segment" and input_ids is not None: + ids = input_ids.detach().cpu().numpy() + mask = None if labels is None else real_token_mask_torch(labels) + cluster_term = segment_cluster_loss(router_4d, membership, segment_ids(ids, window), mask) + else: + cluster_term = cluster_locality_loss(router_4d, membership, window) + loss = loss + lambda_cluster * cluster_term lambda_lookahead = float(getattr(config, "lambda_lookahead", 0.0)) lookahead_steps = int(getattr(config, "lookahead_steps", 0)) diff --git a/chronos/trainer/sft_trainer.py b/chronos/trainer/sft_trainer.py index 537bbb9..5f430cc 100644 --- a/chronos/trainer/sft_trainer.py +++ b/chronos/trainer/sft_trainer.py @@ -126,7 +126,7 @@ def train_step(self, input_ids, labels, step, total_steps, *, modality_inputs=No loss = chronos_loss_term( self.model, ce_loss, lookahead_probs, self.config, - aux_loss=aux_loss, + aux_loss=aux_loss, input_ids=input_ids, labels=labels, ) # Router KL anchor diff --git a/chronos/webui/i18n/en.json b/chronos/webui/i18n/en.json index 6878759..cdeca8b 100644 --- a/chronos/webui/i18n/en.json +++ b/chronos/webui/i18n/en.json @@ -22,6 +22,8 @@ "choice.auto": "Auto", "choice.balanced": "Balanced", "choice.cluster": "By cluster", + "choice.cluster_locality_mode.segment": "Documents (segment)", + "choice.cluster_locality_mode.window": "Fixed windows (window)", "choice.compare": "Compare lazy vs full", "choice.expanded": "Expanded (reference)", "choice.expert": "By expert", @@ -42,6 +44,7 @@ "choice.router_scope.per_iteration": "Route every iteration (per_iteration)", "choice.router_scope.shared": "Reuse iteration 0 routing (shared)", "choice.scheduled": "Scheduled", + "choice.segment": "By document", "choice.shared": "Shared", "choice.sync_on_demand": "Load and wait (sync_on_demand)", "choice.tune.search.lr": "Learning rate", @@ -139,6 +142,8 @@ "field.beta.label": "β (DPO/GRPO)", "field.cluster_locality_clusters.help": "Number of expert clusters for the locality term. 0 chooses it from the expert count.", "field.cluster_locality_clusters.label": "Clusters", + "field.cluster_locality_mode.help": "What one stream is for the cluster constraint. window: fixed windows of the length above. segment: each document of the batch (split at the data's own separators, long ones cut to the window), so similar documents learn to share a cluster without labels or task tags.", + "field.cluster_locality_mode.label": "Locality unit", "field.cluster_locality_window.help": "Segment length (tokens) the cluster constraint applies to. Best results matched the typical record length.", "field.cluster_locality_window.label": "Locality window", "field.coda_layers.help": "Layers that run once after the loop, before the output head.", @@ -586,6 +591,9 @@ "mon.chart.throughput": "Tokens per second", "mon.checkpoints": "Checkpoints", "mon.eta": "Time left", + "mon.eta_basis.deadline": "by time budget", + "mon.eta_basis.steps": "by step progress", + "mon.eta_basis.target_passes": "by data-pass progress", "mon.event": "Event", "mon.events": "Events", "mon.grad_norm": "Gradient norm", @@ -596,6 +604,7 @@ "mon.lr": "Learning rate", "mon.memory": "Memory", "mon.no_runs": "No runs yet. Start one on the Train page.", + "mon.phase_eta": "ends ~{t}", "mon.phase_note": "Steps restart in every stage, so charts show one stage at a time.", "mon.progress": "Progress", "mon.reference": "Reference", diff --git a/chronos/webui/i18n/ja.json b/chronos/webui/i18n/ja.json index 4bfd1cb..0480eba 100644 --- a/chronos/webui/i18n/ja.json +++ b/chronos/webui/i18n/ja.json @@ -22,6 +22,8 @@ "choice.auto": "自動", "choice.balanced": "バランス", "choice.cluster": "クラスタ単位", + "choice.cluster_locality_mode.segment": "文書単位（segment）", + "choice.cluster_locality_mode.window": "固定ウィンドウ（window）", "choice.compare": "遅延ロードと全メモリを比較", "choice.expanded": "展開（参照）", "choice.expert": "エキスパート単位", @@ -42,6 +44,7 @@ "choice.router_scope.per_iteration": "反復ごとにルーティング（per_iteration）", "choice.router_scope.shared": "初回のルーティングを再利用（shared）", "choice.scheduled": "スケジュール", + "choice.segment": "文書単位", "choice.shared": "共有", "choice.sync_on_demand": "読み込んで待機（sync_on_demand）", "choice.tune.search.lr": "学習率", @@ -139,6 +142,8 @@ "field.beta.label": "β（DPO/GRPO）", "field.cluster_locality_clusters.help": "局所性項で使うエキスパートクラスタの数。0 ならエキスパート数から決めます。", "field.cluster_locality_clusters.label": "クラスタ数", + "field.cluster_locality_mode.help": "クラスタ制約で一続きの入力とみなす単位。window: 上の長さの固定ウィンドウ。segment: バッチ内の各文書（データ自身の区切りで分割し、長い文書はウィンドウ長で再分割）。ラベルやタスクタグなしで、似た文書が同じクラスタを使うよう自発的に学習します。", + "field.cluster_locality_mode.label": "局所性の単位", "field.cluster_locality_window.help": "クラスタ制約をかける区間の長さ（トークン）。典型的なレコード長に合わせたときに最良の結果が出ました。", "field.cluster_locality_window.label": "局所性ウィンドウ", "field.coda_layers.help": "ループの後、出力ヘッドの前に 1 回だけ実行する層。", @@ -586,6 +591,9 @@ "mon.chart.throughput": "毎秒トークン数", "mon.checkpoints": "チェックポイント", "mon.eta": "残り時間", + "mon.eta_basis.deadline": "時間予算から", + "mon.eta_basis.steps": "ステップ進捗から", + "mon.eta_basis.target_passes": "データ周回の進捗から", "mon.event": "イベント", "mon.events": "イベント", "mon.grad_norm": "勾配ノルム", @@ -596,6 +604,7 @@ "mon.lr": "学習率", "mon.memory": "メモリ", "mon.no_runs": "まだ実行がありません。「学習」ページから開始してください。", + "mon.phase_eta": "{t} 頃に完了", "mon.phase_note": "ステップは段階ごとに数え直すため、グラフは一度に 1 つの段階だけを表示します。", "mon.progress": "進捗", "mon.reference": "参照訳", diff --git a/chronos/webui/i18n/zh-Hans.json b/chronos/webui/i18n/zh-Hans.json index 97e2cbc..06a2bf8 100644 --- a/chronos/webui/i18n/zh-Hans.json +++ b/chronos/webui/i18n/zh-Hans.json @@ -22,6 +22,8 @@ "choice.auto": "自动", "choice.balanced": "均衡", "choice.cluster": "按簇", + "choice.cluster_locality_mode.segment": "按文档（segment）", + "choice.cluster_locality_mode.window": "固定窗口（window）", "choice.compare": "懒加载与全内存对比", "choice.expanded": "展开（参考）", "choice.expert": "按专家", @@ -42,6 +44,7 @@ "choice.router_scope.per_iteration": "每轮重新路由（per_iteration）", "choice.router_scope.shared": "复用首轮路由（shared）", "choice.scheduled": "按调度", + "choice.segment": "按文档", "choice.shared": "共享", "choice.sync_on_demand": "同步按需加载（sync_on_demand）", "choice.tune.search.lr": "学习率", @@ -139,6 +142,8 @@ "field.beta.label": "β（DPO/GRPO）", "field.cluster_locality_clusters.help": "局部性约束使用的专家簇数。0 表示按专家数自动选择。", "field.cluster_locality_clusters.label": "簇数", + "field.cluster_locality_mode.help": "簇约束以什么作为一段连续输入。window：按上面的长度切固定窗口。segment：批次中的每篇文档（按数据自身的分隔符切分，过长的按窗口再切），相似文档会自发学到同一个簇，不依赖标签或任务标记。", + "field.cluster_locality_mode.label": "局部性单元", "field.cluster_locality_window.help": "簇约束作用的片段长度（token）。与典型记录长度一致时效果最好。", "field.cluster_locality_window.label": "局部性窗口", "field.coda_layers.help": "循环之后、输出头之前只运行一次的层。", @@ -586,6 +591,9 @@ "mon.chart.throughput": "每秒 token 数", "mon.checkpoints": "Checkpoint", "mon.eta": "剩余时间", + "mon.eta_basis.deadline": "按时间预算", + "mon.eta_basis.steps": "按步数进度", + "mon.eta_basis.target_passes": "按数据轮次进度", "mon.event": "事件", "mon.events": "事件", "mon.grad_norm": "梯度范数", @@ -596,6 +604,7 @@ "mon.lr": "学习率", "mon.memory": "内存", "mon.no_runs": "还没有运行记录。请在训练页面启动。", + "mon.phase_eta": "预计 {t} 完成", "mon.phase_note": "每个阶段的步数从头计数，因此图表一次只显示一个阶段。", "mon.progress": "进度", "mon.reference": "参考译文", diff --git a/chronos/webui/i18n/zh-Hant.json b/chronos/webui/i18n/zh-Hant.json index 1a98dc7..56f3ebb 100644 --- a/chronos/webui/i18n/zh-Hant.json +++ b/chronos/webui/i18n/zh-Hant.json @@ -22,6 +22,8 @@ "choice.auto": "自動", "choice.balanced": "平衡", "choice.cluster": "按叢集", + "choice.cluster_locality_mode.segment": "按文件（segment）", + "choice.cluster_locality_mode.window": "固定窗口（window）", "choice.compare": "延遲載入與全記憶體比較", "choice.expanded": "展開（參考）", "choice.expert": "按專家", @@ -42,6 +44,7 @@ "choice.router_scope.per_iteration": "每輪重新路由（per_iteration）", "choice.router_scope.shared": "沿用首輪路由（shared）", "choice.scheduled": "依排程", + "choice.segment": "按文件", "choice.shared": "共用", "choice.sync_on_demand": "同步按需載入（sync_on_demand）", "choice.tune.search.lr": "學習率", @@ -139,6 +142,8 @@ "field.beta.label": "β（DPO/GRPO）", "field.cluster_locality_clusters.help": "區域性約束使用的專家叢集數。0 表示按專家數自動選擇。", "field.cluster_locality_clusters.label": "叢集數", + "field.cluster_locality_mode.help": "簇約束以什麼作為一段連續輸入。window：按上面的長度切固定窗口。segment：批次中的每篇文件（按資料自身的分隔符切分，過長的按窗口再切），相似文件會自發學到同一個簇，不依賴標籤或任務標記。", + "field.cluster_locality_mode.label": "局部性單元", "field.cluster_locality_window.help": "叢集約束作用的片段長度（token）。與典型記錄長度一致時效果最好。", "field.cluster_locality_window.label": "區域性視窗", "field.coda_layers.help": "循環之後、輸出頭之前只執行一次的層。", @@ -586,6 +591,9 @@ "mon.chart.throughput": "每秒 token 數", "mon.checkpoints": "Checkpoint", "mon.eta": "剩餘時間", + "mon.eta_basis.deadline": "按時間預算", + "mon.eta_basis.steps": "按步數進度", + "mon.eta_basis.target_passes": "按資料輪次進度", "mon.event": "事件", "mon.events": "事件", "mon.grad_norm": "梯度範數", @@ -596,6 +604,7 @@ "mon.lr": "學習率", "mon.memory": "記憶體", "mon.no_runs": "尚無執行記錄。請在「訓練」頁開始一次。", + "mon.phase_eta": "預計 {t} 完成", "mon.phase_note": "每個階段的步數重新計數，因此圖表一次只顯示一個階段。", "mon.progress": "進度", "mon.reference": "參考譯文", diff --git a/chronos/webui/schema.py b/chronos/webui/schema.py index 5f041be..2aae9e2 100644 --- a/chronos/webui/schema.py +++ b/chronos/webui/schema.py @@ -136,6 +136,8 @@ def _f(key, group, type_, default=None, **kw) -> Field: recommended=0, source="rule", ref="clusters_auto"), _f("cluster_locality_window", "objectives", "int", 32, min=0, max=65536, step=1, recommended="data", source="rule", ref="window_median_record"), + _f("cluster_locality_mode", "objectives", "choice", "window", choices=["window", "segment"], + source="code", advanced=True), # ---- memory / lazy experts --------------------------------------------- _f("vram_budget_gb", "memory", "float", 4.0, min=0.01, max=256.0, step=0.01), _f("pinned_memory_max_fraction", "memory", "float", 0.25, min=0.01, max=0.95, step=0.01), diff --git a/chronos/webui/static/pages_monitor.js b/chronos/webui/static/pages_monitor.js index 1fb90da..9b881a1 100644 --- a/chronos/webui/static/pages_monitor.js +++ b/chronos/webui/static/pages_monitor.js @@ -46,6 +46,12 @@ export async function renderMonitor(main, page) { }); } +// How the remaining time was projected (measured progress or the time budget). +function etaBasis(basis) { + const key = `mon.eta_basis.${String(basis).split(" ")[0]}`; + return has(key) ? t(key) : String(basis); +} + function paint(body, snap) { const st = snap.status || {}; const pr = snap.progress || {}; @@ -58,7 +64,9 @@ function paint(body, snap) { h("span", { class: "phase-name", text: phaseLabel(p.name) }), h("div", { class: "progress", role: "progressbar", "aria-valuemin": 0, "aria-valuemax": p.total || 0, "aria-valuenow": p.step || 0, "aria-label": phaseLabel(p.name) }, bar(p.total ? Math.min(1, p.step / p.total) : (p.state === "complete" ? 1 : 0))), - h("span", { class: "note", text: `${int(p.step)} / ${p.total ? int(p.total) : "–"} · ${t(`phase.${p.state}`)}` }))); + h("span", { class: "note", text: [`${int(p.step)} / ${p.total ? int(p.total) : "–"} · ${t(`phase.${p.state}`)}`, + p.eta && p.state !== "complete" ? t("mon.phase_eta", { t: dateTime(p.eta) }) : null, + p.eta_basis && p.state !== "complete" ? etaBasis(p.eta_basis) : null].filter(Boolean).join(" · ") }))); // Cards with nothing to show return null; replaceChildren would print it as "null". body.replaceChildren(...[ h("div", { class: "row between" }, @@ -74,7 +82,9 @@ function paint(body, snap) { tile(t("mon.throughput"), num(latest.tokens_per_s, 0), latest.tokens_seen ? `${params(latest.tokens_seen)} ${t("mon.tokens_seen")}` : null), tile(t("mon.memory"), mem.footprint_gb !== undefined ? `${num(mem.footprint_gb, 1)} GB` : (latest.peak_mem_gb !== undefined ? `${num(latest.peak_mem_gb, 1)} GB` : "–"), mem.state === "measured" ? t("label.measured") : null), - tile(t("mon.eta"), duration(pr.run_remaining_s), pr.eta ? dateTime(pr.eta) : null), + tile(t("mon.eta"), duration(pr.run_remaining_s), + pr.eta ? [dateTime(pr.eta), pr.eta_basis ? etaBasis(pr.eta_basis) : null].filter(Boolean).join(" · ") + : (pr.eta_paused ? t(`state.${pr.eta_paused}`) : null)), tile(t("mon.lr"), num(latest.lr, 6), latest.grad_norm !== undefined ? `${t("mon.grad_norm")} ${num(latest.grad_norm, 3)}` : null)), phaseBar.length ? card(t("mon.progress"), ...phaseBar) : null, probeSection(snap), diff --git a/chronos/workbench/training.py b/chronos/workbench/training.py index d152007..e94db12 100644 --- a/chronos/workbench/training.py +++ b/chronos/workbench/training.py @@ -1281,13 +1281,15 @@ def _run(self, cfg: dict, mode: str): requested_backend = (cfg.get("train_backend") or "auto").strip().lower() or "auto" resolved_backend, device = resolve_training_device(requested_backend) - if resolved_backend == "mlx" and model_cfg.modalities: + if resolved_backend == "mlx" and model_cfg.modalities and mode not in {"pretrain", "sft"}: + # MLX trains the image/audio encoders and fusion in pretrain/SFT; + # the preference and distillation stages stay on PyTorch. if requested_backend != "auto": - raise ValueError("MLX 尚未实现多模态编码器训练，请选择 auto 或 PyTorch MPS/CUDA 后端。") + raise ValueError(f"MLX 多模态训练目前支持 pretrain/sft；{mode} 请选择 auto 或 PyTorch MPS/CUDA 后端。") supported = set(training_available()) target = next(name for name in ("cuda", "xpu", "mps", "cpu") if name in supported) resolved_backend, device = resolve_training_device(target) - self._put(f"[data] Modalities {model_cfg.modalities}: using {resolved_backend} for encoder training.") + self._put(f"[data] Modalities {model_cfg.modalities}: using {resolved_backend} for {mode}.") cfg["train_backend"] = requested_backend cfg["device"] = device cfg["dtype"] = _normalize_dtype_for_trainer(cfg.get("dtype")) @@ -1871,6 +1873,7 @@ def _eval_val() -> float: ) loss = chronos_loss_term( model, out.loss, lp, model_cfg, aux_loss=out.aux_loss, + input_ids=input_ids, labels=labels, ) # Logging-only: detached probs are fine here. router_4d_det = collect_router_probs(model, with_grad=False) diff --git a/demos/translation/README.md b/demos/translation/README.md index 1fc9f56..d3bdf9c 100644 --- a/demos/translation/README.md +++ b/demos/translation/README.md @@ -25,18 +25,18 @@ while 0.3.1 is in beta) or from the same checkout. ```sh git clone https://github.com/FonaTech/Project_Chronos.git && cd Project_Chronos -pip install -e '.[mlx,serve]' # the kernel (or: pip install --pre 'Project_Chronos[mlx,serve]==0.3.1a2') +pip install -e '.[mlx,serve]' # the kernel (or: pip install --pre 'Project_Chronos[mlx,serve]==0.3.1a3') pip install -e demos/translation # this extension (sacrebleu, sacremoses, pyarrow pinned) chronos-translate --version ``` -The extension pins `Project_Chronos==0.3.1a2`. Use the checkout's tag -(`git checkout v0.3.1a2`) so both sides match. A model bundle needs the two +The extension pins `Project_Chronos==0.3.1a3`. Use the checkout's tag +(`git checkout v0.3.1a3`) so both sides match. A model bundle needs the two as wheels; build them from the checkout: ```sh -python -m build --outdir dist . # dist/project_chronos-0.3.1a2-py3-none-any.whl -python -m build --outdir dist demos/translation # dist/chronos_translate-0.3.1a2-py3-none-any.whl +python -m build --outdir dist . # dist/project_chronos-0.3.1a3-py3-none-any.whl +python -m build --outdir dist demos/translation # dist/chronos_translate-0.3.1a3-py3-none-any.whl ``` The extension's tests run from the repo root: `pytest demos/translation/tests`. diff --git a/demos/translation/README_zh.md b/demos/translation/README_zh.md index 3ac1e05..87bf02b 100644 --- a/demos/translation/README_zh.md +++ b/demos/translation/README_zh.md @@ -23,17 +23,17 @@ Chronos 内核（`chronos/`）用于训练通用模型，不含任何翻译特 ```sh git clone https://github.com/FonaTech/Project_Chronos.git && cd Project_Chronos -pip install -e '.[mlx,serve]' # 内核（或：pip install --pre 'Project_Chronos[mlx,serve]==0.3.1a2'） +pip install -e '.[mlx,serve]' # 内核（或：pip install --pre 'Project_Chronos[mlx,serve]==0.3.1a3'） pip install -e demos/translation # 本扩展（sacrebleu、sacremoses、pyarrow 固定版本） chronos-translate --version ``` -扩展固定依赖 `Project_Chronos==0.3.1a2`。请检出对应标签（`git checkout v0.3.1a2`），使两边版本一致。 +扩展固定依赖 `Project_Chronos==0.3.1a3`。请检出对应标签（`git checkout v0.3.1a3`），使两边版本一致。 模型包需要这两个 wheel，可以从检出目录构建： ```sh -python -m build --outdir dist . # dist/project_chronos-0.3.1a2-py3-none-any.whl -python -m build --outdir dist demos/translation # dist/chronos_translate-0.3.1a2-py3-none-any.whl +python -m build --outdir dist . # dist/project_chronos-0.3.1a3-py3-none-any.whl +python -m build --outdir dist demos/translation # dist/chronos_translate-0.3.1a3-py3-none-any.whl ``` 扩展测试在仓库根目录运行：`pytest demos/translation/tests`。 diff --git a/demos/translation/pyproject.toml b/demos/translation/pyproject.toml index 4f96781..8b93640 100644 --- a/demos/translation/pyproject.toml +++ b/demos/translation/pyproject.toml @@ -4,21 +4,21 @@ build-backend = "setuptools.build_meta" [project] name = "chronos-translate" -version = "0.3.1a2" +version = "0.3.1a3" description = "English-Chinese translation demo for the Chronos kernel: corpus pipeline, BLEU/chrF runs, model bundles" readme = "README.md" license = "Apache-2.0" requires-python = ">=3.10" dependencies = [ - "Project_Chronos==0.3.1a2", + "Project_Chronos==0.3.1a3", "sacrebleu==2.6.0", "sacremoses==0.2.0", "pyarrow==23.0.1", ] [project.optional-dependencies] -mlx = ["Project_Chronos[mlx]==0.3.1a2"] -serve = ["Project_Chronos[serve]==0.3.1a2"] +mlx = ["Project_Chronos[mlx]==0.3.1a3"] +serve = ["Project_Chronos[serve]==0.3.1a3"] [project.scripts] chronos-translate = "chronos_translate.cli:main" diff --git a/pyproject.toml b/pyproject.toml index a51e710..e35c8d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "Project_Chronos" -version = "0.3.1a2" +version = "0.3.1a3" description = "Storage-aware recurrent-depth MoE: lookahead expert prefetch, SSD-lazy experts, hybrid MLA/SWA attention, self-speculative decoding" readme = "README.md" license = "Apache-2.0" diff --git a/tests/test_chat_records.py b/tests/test_chat_records.py new file mode 100644 index 0000000..c617d7c --- /dev/null +++ b/tests/test_chat_records.py @@ -0,0 +1,182 @@ +"""Tool-call chat template, multi-turn loss masks and ``.lmask`` SFT shards.""" +from __future__ import annotations + +import os +import sys + +import numpy as np +import pytest + +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +from chronos.data.chat_records import encode_chat, render_chat, write_sft_records +from chronos.tokenizer.model import CHAT_TEMPLATE + +OLD_TEMPLATE = ( + "{% for message in messages %}" + "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n' }}" + "{% endfor %}" + "{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}" +) + +TOOLS = [{"type": "function", "function": { + "name": "get_weather", "description": "查询天气", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}] + +PLAIN = [{"role": "system", "content": "你是翻译助手。"}, + {"role": "user", "content": "翻译成中文：Hello"}, + {"role": "assistant", "content": "你好"}] + +TOOL_CHAT = [ + {"role": "system", "content": "You can call tools."}, + {"role": "user", "content": "北京和上海今天天气？"}, + {"role": "assistant", "content": None, "tool_calls": [ + {"type": "function", "function": {"name": "get_weather", "arguments": {"city": "北京"}}}, + {"name": "get_weather", "arguments": {"city": "上海"}}]}, + {"role": "tool", "content": "{\"temp\": 12}"}, + {"role": "tool", "content": "{\"temp\": 18}"}, + {"role": "assistant", "content": "北京 12°C，上海 18°C。"}, +] + + +def _jinja(template, messages, tools=None, add_generation_prompt=False): + pytest.importorskip("transformers") + from transformers.utils.chat_template_utils import render_jinja_template + rendered, _ = render_jinja_template(conversations=[messages], tools=tools, chat_template=template, + add_generation_prompt=add_generation_prompt) + return rendered[0] + + +@pytest.mark.parametrize("gen", [False, True]) +def test_template_without_tools_is_unchanged(gen): + for msgs in (PLAIN, PLAIN[1:], PLAIN[1:2]): + assert _jinja(CHAT_TEMPLATE, msgs, add_generation_prompt=gen) == _jinja(OLD_TEMPLATE, msgs, add_generation_prompt=gen) + + +@pytest.mark.parametrize("messages,tools", [ + (PLAIN, None), (TOOL_CHAT, TOOLS), (TOOL_CHAT[1:], TOOLS), (PLAIN, TOOLS), + ([{"role": "user", "content": "x"}, + {"role": "assistant", "content": "先查一下。", "tool_calls": [{"name": "f", "arguments": "{\"a\": 1}"}]}], TOOLS), +]) +def test_python_render_matches_jinja_template(messages, tools): + text = "".join(t for t, _ in render_chat(messages, tools)) + assert text == _jinja(CHAT_TEMPLATE, messages, tools) + gen = "".join(t for t, _ in render_chat(messages, tools, add_generation_prompt=True)) + assert gen == _jinja(CHAT_TEMPLATE, messages, tools, add_generation_prompt=True) + + +def test_tool_call_text_layout(): + text = "".join(t for t, _ in render_chat(TOOL_CHAT, TOOLS)) + assert text.startswith("<|im_start|>system\nYou can call tools.\n\n\n{\"type\": \"function\"") + assert ('<|im_start|>assistant\n{"name": "get_weather", "arguments": {"city": "北京"}}\n' + '{"name": "get_weather", "arguments": {"city": "上海"}}<|im_end|>\n') in text + assert text.count("<|im_start|>tool\n") == 2 and text.count("<|im_start|>system") == 1 + + +def test_supervision_covers_assistant_turns_only(): + pieces = render_chat(TOOL_CHAT, TOOLS) + learned = [t for t, s in pieces if s] + assert len(learned) == 2 and all(t.endswith("<|im_end|>\n") for t in learned) + assert learned[0].startswith("") and learned[1] == "北京 12°C，上海 18°C。<|im_end|>\n" + context = "".join(t for t, s in pieces if not s) + assert "" in context and "<|im_start|>tool\n{\"temp\": 12}<|im_end|>" in context + assert "<|im_start|>assistant\n" in context # role header is the prompt, not a target + with pytest.raises(ValueError): + render_chat([{"role": "robot", "content": "x"}]) + + +@pytest.fixture(scope="module") +def tok(tmp_path_factory): + from chronos.tokenizer import build_tokenizer + import json + root = tmp_path_factory.mktemp("chat_tok") + with open(root / "d.jsonl", "w", encoding="utf-8") as fh: + for i in range(200): + fh.write(json.dumps({"text": f"天气 weather 北京 上海 call {i} tool_call get_weather city"}, + ensure_ascii=False) + "\n") + return build_tokenizer([str(root / "d.jsonl")], str(root / "out"), vocab_size=700, log=lambda *_: None) + + +def test_encode_matches_full_text_and_masks_exactly(tok): + ids, mask = encode_chat(tok, TOOL_CHAT, TOOLS) + text = "".join(t for t, _ in render_chat(TOOL_CHAT, TOOLS)) + assert tok.decode(ids.tolist()) == text # pieces join to the template text + learned = tok.decode(ids[mask == 1].tolist()) + assert learned.startswith("") and learned.endswith("。<|im_end|>\n") + assert "<|im_start|>" not in learned and "temp" not in learned + im_end = tok.token_id("<|im_end|>") + assert mask[-1] == 1 and int((ids[mask == 1] == im_end).sum()) == 2 + # Trailing context (a tool result with no reply) is dropped from the record. + ids2, mask2 = encode_chat(tok, TOOL_CHAT[:4], TOOLS) + assert mask2[-1] == 1 and tok.decode(ids2.tolist()).endswith("<|im_end|>\n") + with pytest.raises(ValueError): + encode_chat(tok, TOOL_CHAT[:2], TOOLS) + + +def test_lmask_shards_batch_and_router_mask(tok, tmp_path): + import importlib.util + import torch + from chronos.data.shards import SFTRecords + from chronos.trainer.loss_mixin import real_token_mask_torch + token_mask_for = None + if importlib.util.find_spec("mlx") is not None: + from chronos.mlx.training.longrun import token_mask_for + recs = [encode_chat(tok, TOOL_CHAT, TOOLS), encode_chat(tok, PLAIN, None), encode_chat(tok, TOOL_CHAT[1:4], TOOLS)] + stats = write_sft_records(str(tmp_path / "sft" / "part0"), recs, max_len=10_000) + assert stats["records"] == 3 and stats["supervised_tokens"] == sum(int(m.sum()) for _, m in recs) + folder = str(tmp_path / "sft") + assert sorted(os.listdir(folder)) == ["part0.bin", "part0.idx", "part0.lmask", "part0.plen"] + plen = np.fromfile(os.path.join(folder, "part0.plen"), np.int32) + assert plen.tolist() == [int(np.flatnonzero(m)[0]) for _, m in recs] + sft = SFTRecords(folder, seed=0) + assert sft.supervised_tokens() == stats["supervised_tokens"] + by_len = {len(i): (i, m) for i, m in recs} + seen = set() + for _ in range(30): + ids, labels = sft.batch(token_budget=1 << 16, width_multiple=16) + for r in range(ids.shape[0]): + n = int((ids[r] != 0).sum()) + ref_ids, ref_mask = by_len[n] + np.testing.assert_array_equal(ids[r, :n], ref_ids) + np.testing.assert_array_equal(labels[r, :n], np.where(ref_mask == 1, ref_ids, -100)) + assert (labels[r, n:] == -100).all() + seen.add(n) + # The router's real-token mask covers every record token (records end supervised). + np.testing.assert_array_equal(real_token_mask_torch(torch.as_tensor(labels)).numpy(), ids != 0) + if token_mask_for is not None: # MLX SFTSource shift: labels[t] = token t+1 + mask = token_mask_for(labels[:, 1:]) + real = ids[:, :-1] != 0 + np.testing.assert_array_equal(real if mask is None else mask, real) + assert len(seen) >= 2 + + +def test_lmask_size_mismatch_and_bad_records(tok, tmp_path): + from chronos.data.shards import SFTRecords + ids, mask = encode_chat(tok, PLAIN, None) + write_sft_records(str(tmp_path / "s" / "p"), [(ids, mask)]) + with open(tmp_path / "s" / "p.lmask", "ab") as fh: + fh.write(b"\x01") + with pytest.raises(ValueError, match="lmask"): + SFTRecords(str(tmp_path / "s")) + with pytest.raises(ValueError): + write_sft_records(str(tmp_path / "t" / "p"), [(ids, np.zeros_like(mask))]) + stats = write_sft_records(str(tmp_path / "u" / "p"), [(ids, mask)], max_len=len(ids) - 1) + assert stats["records"] == 0 and stats["skipped_too_long"] == 1 + + +def test_server_messages_keep_tool_calls_and_template_renders_them(): + from chronos.serving.ollama_server import _chat_messages + body = {"messages": [ + {"role": "user", "content": "北京天气？"}, + {"role": "assistant", "content": "", "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "get_weather", "arguments": "{\"city\": \"北京\"}"}}]}, + {"role": "tool", "content": "{\"temp\": 12}"}, + {"role": "assistant", "content": None, "tool_calls": [{"function": {"name": "f", "arguments": {"x": 1}}}]}]} + msgs = _chat_messages(body) + assert msgs[1]["tool_calls"] == [{"function": {"name": "get_weather", "arguments": {"city": "北京"}}}] + assert msgs[3]["content"] == "" and "tool_calls" not in msgs[0] and "tool_calls" not in msgs[2] + served = _jinja(CHAT_TEMPLATE, msgs, TOOLS, add_generation_prompt=True) + trained = "".join(t for t, _ in render_chat(msgs, TOOLS, add_generation_prompt=True)) + assert served == trained and '{"name": "get_weather", "arguments": {"city": "北京"}}' in served diff --git a/tests/test_mlx_cache_limit.py b/tests/test_mlx_cache_limit.py new file mode 100644 index 0000000..20a7cdf --- /dev/null +++ b/tests/test_mlx_cache_limit.py @@ -0,0 +1,59 @@ +"""Kernel-level MLX cache limit: min(RAM/8, 16 GiB); callers may only lower it.""" +from __future__ import annotations + +import os +import sys + +import pytest + +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +from chronos.mlx import memory # noqa: E402 + +GIB = 1024 ** 3 + + +@pytest.mark.parametrize("total_gib, want", [(16, 2), (64, 8), (128, 16), (512, 16)]) +def test_ceiling_is_an_eighth_of_ram_capped_at_16_gib(total_gib, want): + assert memory.ceiling_bytes(total_gib * GIB) == want * GIB + + +def test_requests_lower_but_never_raise_the_ceiling(monkeypatch): + monkeypatch.delenv(memory.ENV, raising=False) + total = 64 * GIB + assert memory.resolve_cache_limit(None, total) == (8 * GIB, "default min(RAM/8, 16 GiB)") + assert memory.resolve_cache_limit(4, total) == (4 * GIB, "recipe") + limit, src = memory.resolve_cache_limit(12, total) + assert limit == 8 * GIB and "clamped" in src + monkeypatch.setenv(memory.ENV, "2.5") + assert memory.resolve_cache_limit(None, total) == (int(2.5 * GIB), "env") + for bad in (0, -1, "x"): + with pytest.raises(ValueError): + memory.resolve_cache_limit(bad, total) + + +def test_model_construction_applies_the_limit_and_longrun_records_it(tmp_path, monkeypatch): + mx = pytest.importorskip("mlx.core") + monkeypatch.delenv(memory.ENV, raising=False) + monkeypatch.setattr(memory, "_applied", None) + before = mx.set_cache_limit(0) + mx.set_cache_limit(before) + try: + from chronos.mlx.model import ChronosMLXModel + from chronos.model.config import ChronosConfig + cfg = ChronosConfig(hidden_size=32, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, + vocab_size=48, num_experts=4, num_experts_per_tok=2, moe_intermediate_size=32, + max_position_embeddings=32, prelude_layers=0, coda_layers=0) + ChronosMLXModel(cfg) + got = memory.applied() + assert got is not None and got["limit_gb"] == pytest.approx(memory.ceiling_bytes() / GIB, abs=1e-3) + assert mx.set_cache_limit(int(got["limit_gb"] * GIB)) == int(got["limit_gb"] * GIB) + # A second model in the same process does not reset a caller's lower value. + memory.apply_cache_limit(1, force=True) + ChronosMLXModel(cfg) + assert memory.applied()["limit_gb"] == pytest.approx(1.0) + finally: + mx.set_cache_limit(before) + memory._applied = None diff --git a/tests/test_mlx_multimodal.py b/tests/test_mlx_multimodal.py new file mode 100644 index 0000000..62fcea2 --- /dev/null +++ b/tests/test_mlx_multimodal.py @@ -0,0 +1,153 @@ +"""MLX multimodal fusion: Torch parity, checkpoint round-trip, training step.""" +from __future__ import annotations + +import os +import sys + +import numpy as np +import pytest +import torch + +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +mx = pytest.importorskip("mlx.core") + + +def _cfg(modalities=("image", "audio"), **kw): + from chronos.model.config import ChronosConfig + base = dict(hidden_size=32, num_hidden_layers=3, num_attention_heads=4, num_key_value_heads=2, + vocab_size=48, num_experts=8, num_experts_per_tok=2, num_shared_experts=1, + moe_intermediate_size=32, max_position_embeddings=64, recurrent_depth=2, + prelude_layers=1, coda_layers=1, dropout=0.0, fallback_mask_prob=0.0, + modalities=list(modalities), image_patch_size=8, max_image_tokens=12, + max_audio_tokens=10, modality_dropout=0.0) + base.update(kw) + return ChronosConfig(**base) + + +def _torch_model(c, seed=0): + from chronos.model.model_chronos import ChronosForCausalLM + torch.manual_seed(seed) + pt = ChronosForCausalLM(c).eval() + with torch.no_grad(): # gate 0 → sigmoid 0.5; randomize so parity is tested + fusion = pt.model.multimodal_fusion + fusion.gate.fill_(0.7) + for p in fusion.modality_embeddings.values(): + p.normal_(0, 0.1) + fusion.attn.in_proj_bias.normal_(0, 0.1) + fusion.attn.out_proj.bias.normal_(0, 0.1) + return pt + + +def _media(b=3, seed=1): + g = torch.Generator().manual_seed(seed) + return {"image": torch.randn(b, 3, 16, 24, generator=g), "audio": torch.randn(b, 4000, generator=g)} + + +def _mlx(pt, c): + from chronos.mlx.model import ChronosMLXModel + m = ChronosMLXModel.from_chronos_pytorch(pt, c) + m.eval() + return m + + +CASES = { + "image_audio": lambda b: (_media(b), None), + "image_only": lambda b: ({"image": _media(b)["image"]}, None), + "mixed_rows": lambda b: (_media(b), torch.tensor([[1., 0.], [0., 1.], [0., 0.]])[:b]), + "scalar_mask": lambda b: (_media(b), torch.tensor([1., 0., 1.])[:b]), + "embeds_token_mask": lambda b: ({"image_embeds": torch.randn(b, 5, 20), + "image_embeds_mask": torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 1], [0] * 5]).bool()[:b], + "audio_embeds": torch.randn(b, 4, 40)}, + torch.tensor([[1., 1.], [1., 0.], [0., 0.]])[:b]), + "frame_mask": lambda b: ({**_media(b), "audio_mask": torch.tensor([True, False, True])[:b]}, None), +} + + +@pytest.mark.parametrize("case", sorted(CASES)) +def test_fusion_and_logits_match_torch(case): + c = _cfg() + pt = _torch_model(c) + ids = torch.randint(3, 48, (3, 9)) + media, mask = CASES[case](3) + with mx.stream(mx.cpu): + m = _mlx(pt, c) + text = torch.randn(3, 9, 32) + ref, ref_states = pt.model.multimodal_fusion(text, media, mask) + got, got_states = m.multimodal_fusion(mx.array(text.numpy()), media, mask) + np.testing.assert_allclose(np.array(got), ref.detach().numpy(), atol=2e-5) + assert sorted(got_states) == sorted(ref_states) + for k in ref_states: + np.testing.assert_allclose(np.array(got_states[k]), ref_states[k].detach().numpy(), atol=2e-5) + with torch.no_grad(): + ref_logits = pt(ids, modality_inputs=media, modality_mask=mask).logits + logits, _, _ = m(mx.array(ids.numpy()), modality_inputs=media, modality_mask=mask, logits_dtype=mx.float32) + np.testing.assert_allclose(np.array(logits), ref_logits.numpy(), atol=2e-4) + if case == "mixed_rows": # the row without media is text-only + with torch.no_grad(): + alone = pt(ids[2:], modality_inputs=None).logits + np.testing.assert_allclose(ref_logits[2:].numpy(), alone.numpy(), atol=1e-5) + + +def test_checkpoint_round_trip_and_text_only_unchanged(): + from chronos.mlx.training.io import mlx_state_to_torch + from chronos.model.model_chronos import ChronosForCausalLM + c = _cfg() + pt = _torch_model(c, seed=4) + with mx.stream(mx.cpu): + m = _mlx(pt, c) + sd = mlx_state_to_torch(m, c) + ref = pt.state_dict() + fusion_keys = [k for k in ref if k.startswith("model.multimodal_fusion.")] + assert fusion_keys and all(k in sd for k in fusion_keys) + back = ChronosForCausalLM(c) + back.load_state_dict(sd, strict=True) + for k in fusion_keys: + torch.testing.assert_close(back.state_dict()[k], ref[k], rtol=0, atol=0) + with mx.stream(mx.cpu): # text-only batches skip the fusion block + ids = torch.randint(3, 48, (2, 7)) + logits, _, _ = m(mx.array(ids.numpy()), logits_dtype=mx.float32) + with torch.no_grad(): + np.testing.assert_allclose(np.array(logits), pt(ids).logits.numpy(), atol=2e-4) + with pytest.raises(ValueError, match="no enabled modalities"): + from chronos.mlx.model import ChronosMLXModel + ChronosMLXModel(_cfg(modalities=()))(mx.array(ids.numpy()), modality_inputs=_media(2)) + + +def test_ce_stage_trains_encoders_and_fusion_from_collated_batches(): + from chronos.data.flexible_dataset import multimodal_collate + from chronos.mlx.training.trainer import _ce_stage_loss + import mlx.nn as nn + from mlx.utils import tree_flatten + c = _cfg() + pt = _torch_model(c, seed=5) + ids = torch.randint(3, 48, (9,)) + items = [{"input_ids": ids, "labels": ids, "modality_inputs": {"image": torch.randn(3, 16, 24)}, + "modality_mask": torch.tensor([1., 0.])}, + (ids, ids), + {"input_ids": ids, "labels": ids, "modality_inputs": {"audio": torch.randn(4000)}, + "modality_mask": torch.tensor([0., 1.])}] + batch = multimodal_collate(items) + assert set(batch["modality_inputs"]) >= {"image", "audio"} + from chronos.mlx.multimodal import batch_modalities + with mx.stream(mx.cpu): + m = _mlx(pt, c) + m.train() + + def loss_fn(model, b): + i, lab, media, mask = batch_modalities(b) + return _ce_stage_loss(model, mx.array(i.numpy()), mx.array(lab.numpy()), media, mask) + + loss, grads = nn.value_and_grad(m, loss_fn)(m, batch) + flat = dict(tree_flatten(grads)) + assert np.isfinite(float(loss)) + for key in ("multimodal_fusion.gate", "multimodal_fusion.q_proj.weight", + "multimodal_fusion.image_encoder.patch.weight", + "multimodal_fusion.audio_encoder.net.layers.0.weight"): + assert float(mx.abs(flat[key]).sum()) > 0, key + frozen = _mlx(_torch_model(_cfg(freeze_modality_encoders=True), seed=5), _cfg(freeze_modality_encoders=True)) + names = {k for k, _ in tree_flatten(frozen.trainable_parameters())} + assert not any(".image_encoder." in k or ".audio_encoder." in k for k in names) + assert "multimodal_fusion.q_proj.weight" in names diff --git a/tests/test_moe_dispatch.py b/tests/test_moe_dispatch.py index 739bb6e..11ce4c3 100644 --- a/tests/test_moe_dispatch.py +++ b/tests/test_moe_dispatch.py @@ -165,3 +165,59 @@ def loss(*p): assert _rel(a, b) < 1e-4 # Unavailable experts receive no gradient (they were not executed). assert mx.abs(sg[0][1]).max().item() == 0.0 and mx.abs(sg[0][4]).max().item() == 0.0 + + +def _lazy_moe(dtype, seed=5): + from mlx.utils import tree_map + from chronos.mlx.moe import ChronosMLXMOE + from chronos.mlx.moe_kernels import gather_mm_supported + if not gather_mm_supported(): + pytest.skip("mlx < 0.32.4") + mx.random.seed(seed) + moe = ChronosMLXMOE(moe_cfg()) + moe.update(tree_map(lambda a: a.astype(dtype), moe.parameters())) + moe.eval() + return moe, (mx.random.normal((2, 9, 48)) * 0.5).astype(dtype) + + +@pytest.mark.parametrize("dtype_name,stream_name", [("float32", "cpu"), ("bfloat16", "gpu")]) +def test_mlx_lazy_set_mask_runs_the_full_memory_kernel(dtype_name, stream_name): + """All experts live: the lazy (python-set) path equals the full-memory sparse path + bit for bit, so lazy greedy decoding can no longer drift from the full decode.""" + if stream_name == "gpu" and not mx.metal.is_available(): + pytest.skip("no Metal GPU") + device = mx.cpu if stream_name == "cpu" else mx.gpu + previous = mx.default_device() + mx.set_default_device(device) + try: + moe, x = _lazy_moe(getattr(mx, dtype_name)) + touched = [] + moe.runtime_touch_expert = touched.append + full = moe(x) + lazy = moe(x, available_expert_mask=set(range(moe.num_experts))) + mx.eval(full, lazy) + assert mx.array_equal(full, lazy).item() + assert touched and set(touched) <= set(range(moe.num_experts)) + finally: + mx.set_default_device(previous) + + +def test_mlx_lazy_misses_fall_back_to_shared_expert_like_the_loop_path(): + """Missing experts are served by the shared expert, same as the per-expert loop.""" + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + moe, x = _lazy_moe(mx.float32) + moe.runtime_miss_policy = "fallback_diagnostic" + for avail in ({0, 2, 3, 5}, {1}, set()): + grouped = moe(x, available_expert_mask=avail) + moe.sparse_dispatch = False + loop = moe(x, available_expert_mask=avail) + moe.sparse_dispatch = True + mx.eval(grouped, loop) + assert _rel(grouped, loop) < 1e-5, avail + moe.runtime_miss_policy = "sync_on_demand" + with pytest.raises(RuntimeError): + mx.eval(moe(x, available_expert_mask={0})) + finally: + mx.set_default_device(previous) diff --git a/tests/test_monitor.py b/tests/test_monitor.py index 580d891..01dd17a 100644 --- a/tests/test_monitor.py +++ b/tests/test_monitor.py @@ -3,6 +3,7 @@ import json import os +import time import shutil import subprocess import sys @@ -302,3 +303,39 @@ def test_monitor_static_assets_parse(): pytest.skip("node not installed") result = subprocess.run([node, "--check", os.path.join(static, "app.js")], capture_output=True, text=True) assert result.returncode == 0, result.stderr + + +def test_eta_follows_measured_data_passes_not_the_budget(tmp_path): + """A phase that ends at ``target_passes`` gets its ETA from the measured pass rate.""" + from chronos.monitor import RunMonitor + path = tmp_path / "passes" + os.makedirs(path / "checkpoints") + spec = {"run_dir": str(path), "budget_hours": 10.0, "eval_reserve_fraction": 0.0, "schedule": "time", + "phases": [{"name": "pretrain", "budget_fraction": 0.5}, {"name": "sft", "budget_fraction": 0.5}], + "extra": {"target_passes": {"pretrain": 2.0}}} + (path / "run_spec.json").write_text(json.dumps(spec)) + t0 = datetime.now() - timedelta(seconds=600) + rows = [{"phase": "pretrain", "step": i, "total": 10_000, "passes": 0.001 * i, "step_s": 1.0, + "time": (t0 + timedelta(seconds=i)).isoformat()} for i in range(1, 601)] + _write_jsonl(str(path / "metrics.jsonl"), rows) + hb = {"state": "running", "phase": "pretrain", "phase_wall_s": 600.0, "phase_budget_s": 18_000.0, + "updated": time.time(), "pid": None} + (path / "status.json").write_text(json.dumps(hb)) + prog = RunMonitor(str(tmp_path)).snapshot("passes")["progress"] + # 1.4 passes left at 0.001 passes/s → 1400 s, well inside the 17 400 s deadline. + assert prog["eta_basis"] == "target_passes" + assert prog["phase_remaining_s"] == pytest.approx(1400, rel=0.01) + assert prog["phase_deadline_remaining_s"] == pytest.approx(17_400) + pre, sft = prog["phases"] + assert pre["eta_basis"] == "target_passes" and pre["remaining_s"] == pytest.approx(1400, rel=0.01) + # Unused pretrain time moves to SFT, whose deadline is fixed from the run start. + assert sft["remaining_s"] == pytest.approx(17_400 + 18_000, rel=0.01) + assert prog["run_remaining_s"] == pytest.approx(sft["remaining_s"], rel=0.01) + # No live process: remaining work is still reported, but no clock end time. + assert "eta" not in prog and "eta" not in pre and prog["eta_paused"] == "stopped" + hb["pid"] = os.getpid() # a live writer + (path / "status.json").write_text(json.dumps(hb)) + live = RunMonitor(str(tmp_path)).snapshot("passes")["progress"] + end = datetime.fromisoformat(live["eta"]) + assert abs((end - datetime.now()).total_seconds() - live["run_remaining_s"]) < 120 + assert "eta" in live["phases"][0] and "eta_paused" not in live diff --git a/tests/test_router_mask.py b/tests/test_router_mask.py new file mode 100644 index 0000000..b2fe14a --- /dev/null +++ b/tests/test_router_mask.py @@ -0,0 +1,172 @@ +"""Router regularizers and routing telemetry ignore right-padding.""" +from __future__ import annotations + +import os +import sys + +import numpy as np +import pytest + +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +mx = pytest.importorskip("mlx.core") + + +def _sft_like(lengths, prompt, width, pad=0): + """Aligned (ids, labels) as ``SFTSource`` yields them: labels[t] = token t+1.""" + rng = np.random.default_rng(0) + raw_ids = np.full((len(lengths), width + 1), pad, np.int32) + raw_lab = np.full((len(lengths), width + 1), -100, np.int32) + for r, n in enumerate(lengths): + seq = rng.integers(3, 50, n) + raw_ids[r, :n] = seq + raw_lab[r, prompt:n] = seq[prompt:] + return raw_ids[:, :-1], raw_lab[:, 1:] + + +def _router(B, S, V, E, seed=0): + logits = np.random.default_rng(seed).normal(size=(B, S, V, E)).astype(np.float32) + return mx.softmax(mx.array(logits), axis=-1) + + +def test_token_mask_marks_every_real_input_and_no_padding(): + from chronos.mlx.training.longrun import token_mask_for + from chronos.mlx.training.trainer import real_token_mask + lengths = [7, 12, 3] + ids, labels = _sft_like(lengths, prompt=2, width=16) + mask = token_mask_for(labels) + assert mask is not None and mask.shape == ids.shape + np.testing.assert_array_equal(mask.sum(1), lengths) # last real token included + np.testing.assert_array_equal(mask, ids != 0) # records hold no pad id + np.testing.assert_array_equal(np.array(real_token_mask(mx.array(labels), aligned=True)), mask) + # HF convention (labels[t] = token t, the loss shifts): same real positions. + raw = np.where(np.arange(ids.shape[1])[None, :] >= 2, np.where(ids != 0, ids, -100), -100) + np.testing.assert_array_equal(np.array(real_token_mask(mx.array(raw))), ids != 0) + + +def test_token_mask_is_none_for_packed_batches(): + from chronos.mlx.training.longrun import token_mask_for + packed = np.random.default_rng(1).integers(3, 50, (4, 32)) + assert token_mask_for(packed) is None + ids, labels = _sft_like([17, 17], prompt=3, width=16) # rows fill the width + assert token_mask_for(labels) is None + + +def _losses(router, lookahead, teacher, membership, mask, window): + from chronos.mlx.training import trainer as T + mean = router.mean(axis=2) + out = {"balance": T._load_balance_loss(router, 2, mask), + "cluster": T._cluster_locality_loss(router, membership, window, mask), + "temporal": T._temporal_locality_loss(mean, mask), + "locality": T._router_locality_loss(mean, 2, mask), + "lookahead": T._lookahead_supervision_loss(lookahead, teacher, 2, mask), + "lookahead_topk": T._lookahead_topk_hit_loss(lookahead, teacher, 2, 2, mask), + "lookahead_union": T._lookahead_union_loss(lookahead, teacher, 2, 2, mask)} + return {k: float(v.item()) for k, v in out.items()} + + +def _membership(E, C): + m = np.zeros((E, C), np.float32) + m[np.arange(E), np.arange(E) % C] = 1 + return mx.array(m) + + +@pytest.mark.parametrize("grouped", [False, True]) +def test_masked_router_losses_match_unmasked_on_real_tokens(grouped): + """All-true mask == no mask; a padded row == the same row truncated.""" + B, S, V, E, C, w, n = 1, 32, 3, 6, 3, 8, 16 + router = _router(B, S, V, E) + member = _membership(E, C) + G = 2 if grouped else 1 + la = mx.softmax(mx.array(np.random.default_rng(2).normal(size=(B, S, 3, G, E)).astype(np.float32)), -1) + teacher = router[:, :, :G, :] if grouped else router.mean(axis=2) + if not grouped: + la = la[:, :, :, 0, :] + ones = mx.ones((B, S), dtype=mx.bool_) + full = _losses(router, la, teacher, member, None, w) + assert full == pytest.approx(_losses(router, la, teacher, member, ones, w), rel=1e-5, abs=1e-6) + + mask = mx.array(np.arange(S)[None, :] < n) + padded = _losses(router, la, teacher, member, mask, w) + trunc = _losses(router[:, :n], la[:, :n], teacher[:, :n], member, None, w) + assert padded == pytest.approx(trunc, rel=1e-5, abs=1e-6) + + +def test_padding_content_does_not_change_masked_losses(): + B, S, V, E, C = 3, 24, 2, 6, 3 + lengths = np.array([24, 10, 5]) + mask = mx.array(np.arange(S)[None, :] < lengths[:, None]) + a, b = _router(B, S, V, E, seed=3), _router(B, S, V, E, seed=4) + keep = mask[:, :, None, None] + mixed = mx.where(keep, a, b) # same real tokens, other padding + member = _membership(E, C) + ta, tm = a.mean(axis=2), mixed.mean(axis=2) + la = mx.softmax(mx.array(np.random.default_rng(5).normal(size=(B, S, 3, E)).astype(np.float32)), -1) + lm = mx.where(mask[:, :, None, None], la, mx.softmax(mx.ones_like(la), -1)) + got_a = _losses(a, la, ta, member, mask, 8) + got_m = _losses(mixed, lm, tm, member, mask, 8) + assert got_a == pytest.approx(got_m, rel=1e-5, abs=1e-6) + # Without the mask, padding moves the balance loss (the bug this guards). + from chronos.mlx.training.trainer import _load_balance_loss + assert float(_load_balance_loss(a, 2).item()) != pytest.approx(float(_load_balance_loss(mixed, 2).item())) + + +def test_routing_counts_sum_micro_batches_and_skip_padding(): + from chronos.mlx.training.longrun import RoutingCounts, routing_metrics + r1, r2 = _router(2, 8, 3, 5, seed=6), _router(2, 8, 3, 5, seed=7) + m1 = np.arange(8)[None, :] < np.array([[8], [3]]) + m2 = np.ones((2, 8), bool) + acc = RoutingCounts() + acc.add(r1, 2, m1) + acc.add(r2, 2, m2) + real = mx.array(np.concatenate([np.array(r1).reshape(-1, 3, 5)[m1.reshape(-1)], np.array(r2).reshape(-1, 3, 5)])) + want = routing_metrics(real.reshape(1, -1, 3, 5), 2) + got = acc.summary(2) + assert got["route_tokens"] == 11 + 16 == want["route_tokens"] + for k in ("route_entropy", "load_cv", "load_cv_max", "max_expert_load", "min_expert_load", "dead_experts"): + assert got[k] == pytest.approx(want[k], rel=1e-5, abs=1e-6), k + + +def _tiny_sft_shards(root, lengths, prompt=4): + rng = np.random.default_rng(0) + for split in ("train", "valid"): + d = os.path.join(root, "sft", split) + os.makedirs(d) + seqs = [rng.integers(3, 120, n).astype(np.uint16) for n in lengths] + np.concatenate(seqs).tofile(os.path.join(d, "a.bin")) + np.concatenate([[0], np.cumsum(lengths)]).astype(np.int64).tofile(os.path.join(d, "a.idx")) + np.full(len(lengths), prompt, np.int32).tofile(os.path.join(d, "a.plen")) + + +def test_sft_step_telemetry_covers_every_micro_batch_real_tokens_only(tmp_path): + import json + import time + from chronos.mlx.training.longrun import LongRun, PhaseSpec, RunSpec, SFTSource, token_mask_for + shards = str(tmp_path / "shards") + _tiny_sft_shards(shards, [9, 11, 13, 20, 22, 26] * 8) + model = dict(architecture_version=3, hidden_size=64, num_hidden_layers=4, num_attention_heads=4, + num_key_value_heads=2, vocab_size=128, num_experts=4, num_experts_per_tok=2, + num_shared_experts=1, moe_intermediate_size=64, max_position_embeddings=128, + recurrent_depth=2, prelude_layers=1, coda_layers=1, dropout=0.0, fallback_mask_prob=0.0) + phase = PhaseSpec(name="sft", budget_fraction=1.0, peak_lr=1e-3, tokens_per_step=256, micro_tokens=128) + spec = RunSpec(run_dir=str(tmp_path / "run"), shards_dir=shards, tokenizer_dir=shards, model_config=model, + phases=[phase], budget_hours=1.0, validate_every_s=1e9, checkpoint_every_s=1e9, + dtype="float32") + with mx.stream(mx.cpu): + run = LongRun(spec, log=lambda *_: None) + run.state["phase_total"] = 4 + run.run_phase(0, time.monotonic() + 3600, max_steps=2) + rows = [json.loads(x) for x in open(os.path.join(spec.run_dir, "metrics.jsonl"))] + steps = [r for r in rows if r.get("phase") == "sft" and "step" in r and "route_tokens" in r] + assert [r["step"] for r in steps] == [1, 2] + replay = SFTSource(shards, spec.seed) + for r in steps: # accum = 256 // 128 = 2 micro-batches per step + real = processed = 0 + for _ in range(2): + ids, labels = replay.batch(128) + mask = token_mask_for(labels) + real += int(mask.sum()) if mask is not None else ids.size + processed += ids.size + assert r["route_tokens"] == real < processed diff --git a/tests/test_segment_clusters.py b/tests/test_segment_clusters.py new file mode 100644 index 0000000..e652511 --- /dev/null +++ b/tests/test_segment_clusters.py @@ -0,0 +1,190 @@ +"""Segment-mode cluster objective: documents found from the tokens, no labels.""" +from __future__ import annotations + + +import os +import sys + +import numpy as np +import pytest +import torch + +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +from chronos.expert_clusters import cluster_locality_mode, segment_ids, segment_index # noqa: E402 + + +def test_segment_ids_split_at_separators_and_chunk_long_documents(): + ids = np.array([[5, 6, 0, 7, 8, 9, 0, 4], + [5, 6, 7, 8, 9, 10, 11, 12]]) + np.testing.assert_array_equal(segment_ids(ids), [[0, 0, 0, 1, 1, 1, 1, 2], + [0, 0, 0, 0, 0, 0, 0, 0]]) + # A separator closes its document; chunk=3 cuts documents into 3-token parts. + np.testing.assert_array_equal(segment_ids(ids, 3), [[0, 0, 0, 1, 1, 1, 2, 3], + [0, 0, 0, 1, 1, 1, 2, 2]]) + inv, n = segment_index(segment_ids(ids, 3)) + assert n == 7 and inv.shape == (16,) and inv.dtype == np.int32 + assert len(set(inv[:8]) & set(inv[8:])) == 0 # rows never share a segment + + +def test_mode_validation(): + class C: + cluster_locality_mode = "segment" + assert cluster_locality_mode(C) == "segment" + assert cluster_locality_mode(object()) == "window" # field absent: original form + C.cluster_locality_mode = "label" + with pytest.raises(ValueError): + cluster_locality_mode(C) + + +def _docs_batch(doc_len=5, n_docs=6, C=3, S=None): + """Packed rows of documents, each routed entirely into its own cluster ``d % C``.""" + S = S or doc_len * n_docs + ids = np.full((1, S), 7, np.int64) + ids[0, doc_len - 1::doc_len] = 0 + seg = segment_ids(ids) + cluster = (seg % C)[0] + return ids, cluster + + +def _router_from_clusters(cluster, E=6, C=3, V=2, sharp=12.0): + """[1, S, V, E] probabilities concentrated on cluster ``cluster[s]`` (2 experts per cluster).""" + S = len(cluster) + logits = np.zeros((1, S, V, E), np.float32) + for s, c in enumerate(cluster): + logits[0, s, :, 2 * c:2 * c + 2] = sharp + p = np.exp(logits - logits.max(-1, keepdims=True)) + return p / p.sum(-1, keepdims=True) + + +def _membership(E=6, C=3): + m = np.zeros((E, C), np.float32) + for e in range(E): + m[e, e // (E // C)] = 1 + return m + + +def test_document_routing_scores_zero_and_window_misalignment_is_not_penalised(): + """Each document in one cluster, documents spread evenly: segment loss ≈ 0. + The 8-token window form charges the same routing because its windows + straddle documents (doc length 5).""" + from chronos.model.temporal_loss import cluster_locality_loss, segment_cluster_loss + ids, cluster = _docs_batch() + r = torch.tensor(_router_from_clusters(cluster)) + mem = torch.tensor(_membership()) + seg = segment_cluster_loss(r, mem, segment_ids(ids)).item() + win = cluster_locality_loss(r[:, :24], mem, 8).item() + assert seg < 0.01 + assert win > 0.3 + uniform = torch.full_like(r, 1.0 / r.shape[-1]) + assert segment_cluster_loss(uniform, mem, segment_ids(ids)).item() == pytest.approx(1.0, abs=1e-5) + + +def test_padding_is_neither_a_document_nor_weight(): + from chronos.model.temporal_loss import segment_cluster_loss + ids, cluster = _docs_batch(n_docs=4) + r = _router_from_clusters(cluster) + pad = 9 + ids_p = np.concatenate([ids, np.zeros((1, pad), np.int64)], axis=1) + noise = np.random.default_rng(0).dirichlet(np.ones(6), size=(1, pad, 2)).astype(np.float32) + r_p = np.concatenate([r, noise], axis=1) + mask = torch.tensor(np.arange(ids_p.shape[1])[None, :] < ids.shape[1]) + mem = torch.tensor(_membership()) + a = segment_cluster_loss(torch.tensor(r), mem, segment_ids(ids)).item() + b = segment_cluster_loss(torch.tensor(r_p), mem, segment_ids(ids_p), mask).item() + assert b == pytest.approx(a, abs=1e-6) + + +def test_mlx_matches_torch_and_carries_gradient(): + mx = pytest.importorskip("mlx.core") + from chronos.mlx.training.trainer import _segment_cluster_loss + from chronos.model.temporal_loss import segment_cluster_loss + rng = np.random.default_rng(4) + ids = rng.integers(3, 40, (3, 24)) + ids[ids % 7 == 0] = 0 # random document boundaries + logits = rng.normal(size=(3, 24, 4, 6)).astype(np.float32) + mask = np.ones((3, 24), bool) + mask[1, 17:] = False + seg = segment_ids(ids, 8) + mem = _membership() + lt = torch.tensor(logits, requires_grad=True) + t = segment_cluster_loss(torch.softmax(lt, -1), torch.tensor(mem), seg, torch.tensor(mask)) + t.backward() + with mx.stream(mx.cpu): + f = lambda x: _segment_cluster_loss(mx.softmax(x, axis=-1), mx.array(mem), seg, mx.array(mask)) # noqa: E731 + val, grad = mx.value_and_grad(f)(mx.array(logits)) + mx.eval(val, grad) + assert float(val) == pytest.approx(t.item(), abs=1e-5) + np.testing.assert_allclose(np.array(grad), lt.grad.numpy(), atol=1e-5) + assert np.abs(np.array(grad)).sum() > 0 + assert np.abs(np.array(grad)[1, 17:]).max() == 0.0 # padding receives no gradient + + +def _tiny_cfg(mode, lam=0.5): + from chronos.model.config import ChronosConfig + return ChronosConfig(hidden_size=32, num_hidden_layers=3, num_attention_heads=4, num_key_value_heads=2, + vocab_size=48, num_experts=8, num_experts_per_tok=2, num_shared_experts=1, + moe_intermediate_size=32, max_position_embeddings=64, recurrent_depth=2, + prelude_layers=1, coda_layers=1, dropout=0.0, fallback_mask_prob=0.0, + lambda_cluster_locality=lam, cluster_locality_window=8, cluster_locality_mode=mode) + + +def test_segment_mode_reaches_both_trainers_and_differs_from_window(): + mx = pytest.importorskip("mlx.core") + from chronos.mlx.model import ChronosMLXModel + from chronos.mlx.training.trainer import _ce_stage_loss + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.trainer.loss_mixin import chronos_loss_term + ids = torch.randint(3, 48, (2, 16)) + ids[:, 5] = 0 + ids[1, 11] = 0 + vals = {} + for mode in ("window", "segment"): + torch.manual_seed(3) + c = _tiny_cfg(mode) + assert c.cluster_locality_mode == mode + pt = ChronosForCausalLM(c) + out = pt(ids) + base = out.logits.sum() * 0.0 + vals[("torch", mode)] = chronos_loss_term(pt, base, None, c, input_ids=ids, labels=ids).item() + no_ids = chronos_loss_term(pt, base, None, c).item() # without ids: window form + vals[("torch_noids", mode)] = no_ids + with mx.stream(mx.cpu): + m = ChronosMLXModel.from_chronos_pytorch(pt, c) + m.train() + x = mx.array(ids.numpy(), dtype=mx.int32) + _ce_stage_loss(m, x, x) + vals[("mlx", mode)] = float(m._last_loss_components["cluster_locality_loss"]) / 0.5 + assert vals[("torch", "segment")] != pytest.approx(vals[("torch", "window")], abs=1e-6) + assert vals[("torch_noids", "segment")] == pytest.approx(vals[("torch", "window")], abs=1e-6) + for mode in ("window", "segment"): + assert 0.0 <= vals[("mlx", mode)] <= 1.0 + + +def test_segment_ids_ignore_everything_but_the_separator(): + """Task/prompt tokens are ordinary tokens: changing them never moves a boundary.""" + rng = np.random.default_rng(2) + ids = rng.integers(1, 64, (4, 40)) + ids[:, [9, 23]] = 0 + other = ids.copy() + nz = other != 0 + other[nz] = rng.integers(1, 64, int(nz.sum())) + np.testing.assert_array_equal(segment_ids(ids, 16), segment_ids(other, 16)) + + +def test_schema_and_catalogs_know_the_field(): + import json + from chronos.config_schema import MODEL_FIELDS + from chronos.webui.schema import FIELDS + assert "cluster_locality_mode" in MODEL_FIELDS + field = FIELDS["cluster_locality_mode"] + assert tuple(field.choices) == ("window", "segment") and field.default == "window" + root = os.path.join(_repo_root, "chronos", "webui", "i18n") + for lang in ("en", "zh-Hans", "zh-Hant", "ja"): + cat = json.load(open(os.path.join(root, f"{lang}.json"), encoding="utf-8")) + for k in ("field.cluster_locality_mode.label", "field.cluster_locality_mode.help", "choice.window", + "choice.segment", "choice.cluster_locality_mode.window", "choice.cluster_locality_mode.segment"): + assert cat.get(k), (lang, k) + From cf485b260f82dbb1097ac8af6b835fdf46285996 Mon Sep 17 00:00:00 2001 From: Fona <62028332+FonaTech@users.noreply.github.com> Date: Sun, 8 Nov 2026 13:31:59 +0800 Subject: [PATCH 2/3] shards: keep -100 labels under NumPy 2 for uint16 .lmask records np.where(mask, uint16_seq, -100) stays uint16 under NEP 50 and wraps -100 to 65436; cast the sequence to int32 first. Co-Authored-By: Claude Opus 5.5 --- chronos/data/shards.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/chronos/data/shards.py b/chronos/data/shards.py index 141c809..34a9ae1 100644 --- a/chronos/data/shards.py +++ b/chronos/data/shards.py @@ -123,7 +123,8 @@ def batch(self, token_budget: int, max_rows: int = 256, width_multiple: int = 64 lm = self.lmasks[si] if lm is not None: keep = np.asarray(lm[off:off + n]).astype(bool) - labels[r, :n] = np.where(keep, seq, -100) + # int32 first: under NumPy 2 (NEP 50) np.where keeps uint16 and wraps -100 → 65436. + labels[r, :n] = np.where(keep, np.asarray(seq, dtype=np.int32), -100) else: p = int(self.plens[si][ri]) labels[r, p:n] = seq[p:] From e98372ea67658c599daac7080c07df40d279d4f4 Mon Sep 17 00:00:00 2001 From: Fona <62028332+FonaTech@users.noreply.github.com> Date: Sun, 8 Nov 2026 14:10:08 +0800 Subject: [PATCH 3/3] serve: parse tool calls; MLX: multimodal DPO/ORPO/GRPO/distill Tool-call inference - chronos.serving.tool_calls: streaming parser that never leaks a partial tag; invalid blocks stay content. - /api/chat returns Ollama message.tool_calls; /v1/chat/completions returns OpenAI tool_calls (string arguments, finish_reason "tool_calls"), streamed or not. Requests without tools are unchanged. Multimodal preference stages on MLX - DPO/ORPO: chosen and rejected condition on the prompt's media. - GRPO: rollouts condition on the prompt's media (encoded once, same logits); the policy update re-encodes so encoders train. - Distill: media reach the student; a text-only teacher is accepted. - Data: MultimodalStreamingDPODataset + preference_collate, load_grpo_prompt_records, shared record_media resolver. - Workbench: media data in DPO/ORPO/GRPO selects MLX under auto; an explicit PyTorch backend raises instead of dropping the media. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 21 ++- README.md | 6 +- README_zh.md | 6 +- chronos/data/flexible_dataset.py | 83 +++++++++ chronos/mlx/training/trainer.py | 110 ++++++++++-- chronos/serving/ollama_server.py | 78 +++++++-- chronos/serving/tool_calls.py | 130 ++++++++++++++ chronos/trainer/grpo_trainer.py | 25 +++ chronos/workbench/training.py | 59 +++++-- tests/test_mlx_multimodal_stages.py | 261 ++++++++++++++++++++++++++++ tests/test_tool_call_serving.py | 140 +++++++++++++++ 11 files changed, 866 insertions(+), 53 deletions(-) create mode 100644 chronos/serving/tool_calls.py create mode 100644 tests/test_mlx_multimodal_stages.py create mode 100644 tests/test_tool_call_serving.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 91b494c..237adc1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ## [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, tool-call SFT records, a kernel-level MLX cache limit and progress-based ETAs. +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. @@ -16,7 +16,16 @@ This release fixes the router masks and the lazy-expert decode. It also adds a d - 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` accepts `multimodal_collate` batches in pretrain and SFT. Preference and distillation stages with media still run on PyTorch. + - `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 `` system block. - The assistant's `tool_calls` render as `{json}` lines, and `tool` turns carry the results. @@ -26,7 +35,12 @@ This release fixes the router masks and the lazy-expert decode. It also adds a d - 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` passes `tools` and the assistant's `tool_calls` from `/api/chat` and `/v1/chat/completions` into the template. It does not parse `` output back into structured `tool_calls` yet. +- `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 `` 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. @@ -42,6 +56,7 @@ This release fixes the router masks and the lazy-expert decode. It also adds a d ### 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) diff --git a/README.md b/README.md index a3761fc..1726d3d 100644 --- a/README.md +++ b/README.md @@ -289,8 +289,10 @@ 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 pretrain and SFT, with the same weights and numerics as PyTorch (`chronos.mlx.multimodal`). DPO, ORPO, GRPO and distillation with media still use PyTorch. -- Tool-call SFT: 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. +- 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 `` 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`. diff --git a/README_zh.md b/README_zh.md index 68d733f..b1a9a11 100644 --- a/README_zh.md +++ b/README_zh.md @@ -299,8 +299,10 @@ checkpoint 保存事件、stop 事件、verify 结果都会同步到界面。Sto MLX 上的内存、多模态与工具调用： - 每个 MLX 模型在分配权重前把 MLX 缓冲区缓存上限设为 `min(内存/8, 16 GiB)`（`chronos.mlx.memory`）。配方的 `extra.mlx_cache_limit_gb` 或 `CHRONOS_MLX_CACHE_LIMIT_GB` 只能调低，不能调高。上限只约束可回收缓存，不限制权重、激活和梯度。 -- 图像、音频编码器和门控交叉注意力融合在 MLX 上支持预训练和 SFT，权重与数值和 PyTorch 一致（`chronos.mlx.multimodal`）。带媒体的 DPO、ORPO、GRPO 和蒸馏仍走 PyTorch。 -- 工具调用 SFT：chat 模板支持 `tools`、assistant 的 `tool_calls` 和 `tool` 结果。`chronos.data.chat_records.encode_chat` 为每个 token 标出是否计 loss（assistant 文本和调用参与学习，其余为上下文），`write_sft_records` 写出 `SFTRecords` 可读的 `.lmask` 分片。 +- 图像、音频编码器和门控交叉注意力融合在 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` 把模型输出的 `` 解析成 Ollama 的 `message.tool_calls` 或 OpenAI 的 `tool_calls`，流式与非流式都支持。 完整 6 阶段端到端对比见 `tools/compare_minimind_chronos_v3.py`。 diff --git a/chronos/data/flexible_dataset.py b/chronos/data/flexible_dataset.py index d93a5c7..2eb7793 100644 --- a/chronos/data/flexible_dataset.py +++ b/chronos/data/flexible_dataset.py @@ -303,6 +303,48 @@ def _media_value(value): return value +def record_media(record: dict, *, image_size: int = 224, audio_sample_rate: int = 16000, + max_audio_seconds: float = 30.0, media_error_policy: str = "warn"): + """``(modality_inputs, modality_mask[2])`` for one record, or ``None``. + + Pre-encoded ``*_embeds`` take precedence over a path/value of the same + modality. Used by the preference (DPO/ORPO) and GRPO prompt loaders; the + pair or the rollout group shares the record's media.""" + if not detect_modalities(record): + return None + normalized = normalize_multimodal_record(record) + media = {} + for name in ("image_embeds", "audio_embeds"): + value = normalized["modalities"].get(name) + if value is not None: + try: + media[name] = torch.as_tensor(value, dtype=torch.float32) + except Exception as exc: + if media_error_policy == "raise": + raise ValueError(f"invalid {name}: {exc}") from exc + warnings.warn(f"invalid {name}: {exc}", RuntimeWarning, stacklevel=2) + for modality, loader in ( + ("image", lambda v: _load_image(v, image_size)), + ("audio", lambda v: _load_audio(v, audio_sample_rate, max_audio_seconds)), + ): + if f"{modality}_embeds" in media: + continue + value = _media_value(normalized["modalities"].get(modality)) + if value is None: + continue + try: + media[modality] = loader(value) + except Exception as exc: + if media_error_policy == "raise": + raise + warnings.warn(str(exc), RuntimeWarning, stacklevel=2) + if not media: + return None + mask = torch.tensor([float("image" in media or "image_embeds" in media), + float("audio" in media or "audio_embeds" in media)]) + return media, mask + + class MediaResolver: """Lazy media resolver shared by JSONL and API serving paths.""" @@ -764,3 +806,44 @@ def __getitem__(self, index): "y_rejected": torch.tensor(r_ids[1:], dtype=torch.long), "mask_rejected":torch.tensor(r_mask[1:], dtype=torch.long), } + + +class MultimodalStreamingDPODataset(StreamingDPODataset): + """Preference pairs with the prompt's image/audio. + + Media fields sit on the record (``image``/``audio``/``*_embeds`` or a + ``modalities`` dict, as for SFT); chosen and rejected answer the same + multimodal prompt, so both sides of the pair condition on the same media. + Text-only records keep the ``StreamingDPODataset`` dict.""" + + def __init__(self, data_path, tokenizer, max_length=4096, *, image_size=224, + audio_sample_rate=16000, max_audio_seconds=30.0, media_error_policy="warn"): + super().__init__(data_path, tokenizer, max_length=max_length) + self.media_kwargs = dict(image_size=int(image_size), audio_sample_rate=int(audio_sample_rate), + max_audio_seconds=float(max_audio_seconds), + media_error_policy=str(media_error_policy)) + self.modality_summary = dataset_modality_summary(data_path) + + def __getitem__(self, index): + item = super().__getitem__(index) + found = record_media(self._read_record(index), **self.media_kwargs) + if found is not None: + item["modality_inputs"], item["modality_mask"] = found + return item + + +PREFERENCE_KEYS = ("x_chosen", "y_chosen", "mask_chosen", "x_rejected", "y_rejected", "mask_rejected") + + +def preference_collate(batch): + """Stack DPO/ORPO pairs; media (one set per pair) is batched with the + ``multimodal_collate`` rules, so mixed text/media pairs keep exact masks.""" + out = {key: torch.stack([item[key] for item in batch]) for key in PREFERENCE_KEYS} + if any(item.get("modality_inputs") for item in batch): + stub = torch.zeros(1, dtype=torch.long) + media = multimodal_collate([ + {"input_ids": stub, "labels": stub, "modality_inputs": item["modality_inputs"], + "modality_mask": item.get("modality_mask")} if item.get("modality_inputs") else (stub, stub) + for item in batch]) + out["modality_inputs"], out["modality_mask"] = media.get("modality_inputs", {}), media["modality_mask"] + return out diff --git a/chronos/mlx/training/trainer.py b/chronos/mlx/training/trainer.py index eadeaa5..dc4b0b6 100644 --- a/chronos/mlx/training/trainer.py +++ b/chronos/mlx/training/trainer.py @@ -591,6 +591,28 @@ def _ce_stage_loss(model, input_ids, labels, modality_inputs=None, modality_mask draft_kl=draft_kl, token_mask=real_token_mask(labels), input_ids=input_ids) +def _tile_media(batch, copies: int) -> dict: + """Model kwargs with the batch's media repeated ``copies`` times along the + batch axis (chosen‖rejected pairs, GRPO groups); ``{}`` for text-only.""" + media = batch.get("modality_inputs") if isinstance(batch, dict) else None + if not media: + return {} + from chronos.mlx.multimodal import _mx + tiled = {k: mx.concatenate([_mx(v)] * copies, axis=0) if copies > 1 else _mx(v) + for k, v in media.items() if v is not None} + mask = batch.get("modality_mask") + if mask is not None: + mask = _mx(mask, mx.float32) + mask = mx.concatenate([mask] * copies, axis=0) if copies > 1 else mask + return {"modality_inputs": tiled, "modality_mask": mask} + + +def _frozen_media(frozen, media: dict) -> dict: + """A text-only reference/teacher is a valid source for a multimodal + student (Torch distill parity): pass media only if it can fuse them.""" + return media if media and getattr(frozen, "multimodal_fusion", None) is not None else {} + + def _dpo_stage_loss(model, reference, batch, beta: float = 0.1, router_ref=None): x_c = _ids_to_mx(batch["x_chosen"]) x_r = _ids_to_mx(batch["x_rejected"]) @@ -599,12 +621,13 @@ def _dpo_stage_loss(model, reference, batch, beta: float = 0.1, router_ref=None) m_c = _ids_to_mx(batch["mask_chosen"]) m_r = _ids_to_mx(batch["mask_rejected"]) x = mx.concatenate([x_c, x_r], axis=0) - logits, lookahead, _cache = model(x, logits_dtype=mx.float32) + media = _tile_media(batch, 2) # chosen and rejected answer the same prompt media + logits, lookahead, _cache = model(x, logits_dtype=mx.float32, **media) B = x_c.shape[0] lp_c = _aligned_sequence_logprob(logits[:B], y_c, m_c) lp_r = _aligned_sequence_logprob(logits[B:], y_r, m_r) if reference is not None: - ref_logits, _ref_lp, _ref_cache = reference(x, logits_dtype=mx.float32) + ref_logits, _ref_lp, _ref_cache = reference(x, logits_dtype=mx.float32, **_frozen_media(reference, media)) ref_c = mx.stop_gradient(_aligned_sequence_logprob(ref_logits[:B], y_c, m_c)) ref_r = mx.stop_gradient(_aligned_sequence_logprob(ref_logits[B:], y_r, m_r)) base = -mx.mean(nn.log_sigmoid(beta * ((lp_c - lp_r) - (ref_c - ref_r)))) @@ -621,7 +644,7 @@ def _orpo_stage_loss(model, batch, beta: float = 0.1, lambda_or: float = 0.1, ro m_c = _ids_to_mx(batch["mask_chosen"]) m_r = _ids_to_mx(batch["mask_rejected"]) x = mx.concatenate([x_c, x_r], axis=0) - logits, lookahead, _cache = model(x, logits_dtype=mx.float32) + logits, lookahead, _cache = model(x, logits_dtype=mx.float32, **_tile_media(batch, 2)) B = x_c.shape[0] lp_c = _aligned_sequence_logprob(logits[:B], y_c, m_c, average=True) lp_r = _aligned_sequence_logprob(logits[B:], y_r, m_r, average=True) @@ -645,9 +668,13 @@ def _distill_stage_loss( alpha: float = 0.7, temperature: float = 4.0, router_ref=None, + modality_inputs=None, + modality_mask=None, ): - logits, lookahead, _cache = model(input_ids, logits_dtype=mx.float32) - t_logits, _tlp, _tcache = teacher(teacher_input_ids if teacher_input_ids is not None else input_ids, logits_dtype=mx.float32) + media = {"modality_inputs": modality_inputs, "modality_mask": modality_mask} if modality_inputs else {} + logits, lookahead, _cache = model(input_ids, logits_dtype=mx.float32, **media) + t_logits, _tlp, _tcache = teacher(teacher_input_ids if teacher_input_ids is not None else input_ids, + logits_dtype=mx.float32, **_frozen_media(teacher, media)) V = min(logits.shape[-1], t_logits.shape[-1]) S = min(logits.shape[1], t_logits.shape[1]) logits_kd = logits[:, :S, :V] @@ -674,7 +701,8 @@ def _grpo_stage_loss(model, reference, batch, beta: float = 0.04, router_ref=Non ids = _ids_to_mx(batch["ids"]) response_mask = mx.array(batch["response_mask"]).astype(mx.float32) advantages = mx.array(batch["advantages"]).astype(mx.float32) - logits, lookahead, _cache = model(ids, logits_dtype=mx.float32) + media = _tile_media(batch, int(ids.shape[0])) # one prompt's media for every trajectory + logits, lookahead, _cache = model(ids, logits_dtype=mx.float32, **media) lp_tok = _shifted_token_logprobs(logits, ids) mask = response_mask[:, 1:] seq_lp = (lp_tok * mask).sum(axis=1) / mx.maximum(mask.sum(axis=1), 1.0) @@ -682,7 +710,7 @@ def _grpo_stage_loss(model, reference, batch, beta: float = 0.04, router_ref=Non kl = mx.zeros((), dtype=logits.dtype) if reference is not None: - ref_logits, _ref_lp, _ref_cache = reference(ids, logits_dtype=mx.float32) + ref_logits, _ref_lp, _ref_cache = reference(ids, logits_dtype=mx.float32, **_frozen_media(reference, media)) ref_lp_tok = mx.stop_gradient(_shifted_token_logprobs(ref_logits, ids)) log_ratio = ref_lp_tok - lp_tok kl = ((mx.exp(log_ratio) - log_ratio - 1.0) * mask).sum() / mx.maximum(mask.sum(), 1.0) @@ -699,6 +727,8 @@ def _capture_router_ref(model, batch, stage: str) -> mx.array | None: ids = _ids_to_mx(raw_ids) elif stage in {"dpo", "orpo"}: ids = mx.concatenate([_ids_to_mx(batch["x_chosen"]), _ids_to_mx(batch["x_rejected"])], axis=0) + tiled = _tile_media(batch, 2) + media, media_mask = tiled.get("modality_inputs"), tiled.get("modality_mask") else: return None _logits, _lookahead, _cache = model(ids, modality_inputs=media, modality_mask=media_mask) @@ -717,13 +747,38 @@ def _sample_next(logits: mx.array, temperature: float) -> mx.array: return mx.random.categorical(mx.log(probs + 1e-9)).astype(mx.int32) -def _rollout_one(model, prompt_ids: mx.array, max_gen_len: int, temperature: float, eos_token_id: int | None) -> tuple[mx.array, list[int]]: +def _encoded_media(model, media: dict) -> dict: + """Rollout kwargs with raw image/audio encoded once into ``*_embeds``. + + Every decode step re-applies the fusion (Torch ``generate`` parity); the + encoder output is constant, so it is computed once and passed through + the pre-encoded path (same tokens, same offsets, same validity).""" + if not media: + return {} + fusion = getattr(model, "multimodal_fusion", None) + inputs = dict(media["modality_inputs"]) + if fusion is not None: + from chronos.mlx.multimodal import _mx + for name, encoder in (("image", fusion.image_encoder), ("audio", fusion.audio_encoder)): + raw = inputs.get(name) + if raw is None or encoder is None or inputs.get(f"{name}_embeds") is not None: + continue + inputs.pop(name) + inputs[f"{name}_embeds"] = mx.stop_gradient(encoder(_mx(raw, mx.float32))) + mask = inputs.pop(f"{name}_mask", None) + if mask is not None: + inputs[f"{name}_embeds_mask"] = mask + return {"modality_inputs": inputs, "modality_mask": media.get("modality_mask")} + + +def _rollout_one(model, prompt_ids: mx.array, max_gen_len: int, temperature: float, eos_token_id: int | None, + media: dict | None = None) -> tuple[mx.array, list[int]]: ids = prompt_ids cache = None generated: list[int] = [] for _ in range(max(1, int(max_gen_len))): token_in = ids if cache is None else ids[:, -1:] - logits, _lookahead, cache = model(token_in, cache=cache) + logits, _lookahead, cache = model(token_in, cache=cache, **(media or {})) next_id = _sample_next(logits[:, -1, :], temperature).reshape(1, 1) mx.eval(next_id) tok = int(next_id.item()) @@ -735,7 +790,28 @@ def _rollout_one(model, prompt_ids: mx.array, max_gen_len: int, temperature: flo return ids, generated -def _build_grpo_batch(model, prompt: str, tokenizer, reward_fn, args) -> dict: +def _grpo_prompt_media(item, args): + """``(prompt_text, media)`` for a GRPO prompt: a string, or the + ``{"prompt", "record"}`` dict of ``load_grpo_prompt_records`` (its + image/audio resolved here, once per step, with the dataset loaders).""" + if not isinstance(item, dict): + return str(item), None + from chronos.data.flexible_dataset import record_media + get = (lambda k, d: getattr(args, k, d)) if args is not None else (lambda _k, d: d) + found = record_media(item.get("record") or {}, image_size=int(get("image_size", 224)), + audio_sample_rate=int(get("audio_sample_rate", 16000)), + max_audio_seconds=float(get("max_audio_seconds", 30.0)), + media_error_policy=str(get("media_error_policy", "warn"))) + if found is None: + return str(item["prompt"]), None + media, mask = found + batched = {k: v.unsqueeze(0) for k, v in media.items()} # one prompt → batch of 1 + return str(item["prompt"]), {"modality_inputs": batched, "modality_mask": mask.unsqueeze(0)} + + +def _build_grpo_batch(model, prompt, tokenizer, reward_fn, args) -> dict: + prompt, media = _grpo_prompt_media(prompt, args) + rollout_media = _encoded_media(model, media) rendered = tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], tokenize=False, @@ -757,6 +833,7 @@ def _build_grpo_batch(model, prompt: str, tokenizer, reward_fn, args) -> dict: max_gen_len=max_gen_len, temperature=temperature, eos_token_id=tokenizer.eos_token_id, + media=rollout_media, ) trajectories.append(full_ids) rewards.append(float(reward_fn(prompt, response_tokens, tokenizer))) @@ -775,11 +852,14 @@ def _build_grpo_batch(model, prompt: str, tokenizer, reward_fn, args) -> dict: denom = mx.maximum(mx.sqrt(((adv * adv).mean())), 1e-6) advantages = adv / denom mx.eval(advantages) - return { + out = { "ids": ids_rows, "response_mask": mask_rows, "advantages": advantages.tolist(), } + if media: # the policy update re-encodes (encoder gradients), tiled per trajectory + out.update(media) + return out @dataclass @@ -1072,9 +1152,7 @@ def loss_fn(m, batch): if stage == "orpo": return _orpo_stage_loss(m, batch, beta=beta, lambda_or=lambda_or, router_ref=router_ref) if stage == "distill": - ids, labels, media, _media_mask = batch_modalities(batch) - if media: - raise ValueError("MLX distill does not take modality inputs yet; use pretrain/sft or the PyTorch backend") + ids, labels, media, media_mask = batch_modalities(batch) ids = _ids_to_mx(ids) labels = _ids_to_mx(labels) teacher_vocab = int(getattr(getattr(teacher, "config", config), "vocab_size", config.vocab_size)) @@ -1087,6 +1165,8 @@ def loss_fn(m, batch): alpha=alpha, temperature=temperature, router_ref=router_ref, + modality_inputs=media, + modality_mask=media_mask, ) if stage == "grpo": return _grpo_stage_loss(m, reference, batch, beta=beta, router_ref=router_ref) @@ -1195,7 +1275,7 @@ def evaluate_validation(epoch): break train_batch = _build_grpo_batch( model, - str(batch), + batch if isinstance(batch, dict) else str(batch), tokenizer, reward_fn, args, diff --git a/chronos/serving/ollama_server.py b/chronos/serving/ollama_server.py index 2d44006..7790ad3 100644 --- a/chronos/serving/ollama_server.py +++ b/chronos/serving/ollama_server.py @@ -30,6 +30,8 @@ except ImportError: # pragma: no cover - serving extra not installed FastAPI = HTTPException = Request = JSONResponse = PlainTextResponse = StreamingResponse = None +from chronos.serving.tool_calls import ToolCallStream, ollama_tool_calls, openai_tool_calls, split_tool_calls + OLLAMA_COMPAT_VERSION = "0.40.14" @@ -219,7 +221,9 @@ def events(): final["error"] = item["error"] yield shape(None, final) else: - yield shape(item["text"], None) + chunk = shape(item["text"], None) + if chunk is not None: # text held back (possible start of a tag) + yield chunk return events def _respond(body, events, collect): @@ -269,16 +273,35 @@ async def chat(request: Request): return {"model": body.get("model"), "created_at": _now(), "message": {"role": "assistant", "content": ""}, "done": True, "done_reason": "load"} ids = runtime.encode_chat(lm, messages, body.get("tools")) + # With ``tools`` the model may answer with blocks: parse them + # into message.tool_calls (Ollama shape) as they stream. + parser = ToolCallStream() if body.get("tools") else None + + def message(text, calls=()): + msg = {"role": "assistant", "content": text} + if calls: + msg["tool_calls"] = ollama_tool_calls(list(calls)) + return msg def shape(text, final): if final is not None: - return {**final, "message": {"role": "assistant", "content": ""}} - return {"model": body.get("model"), "created_at": _now(), - "message": {"role": "assistant", "content": text}, "done": False} + tail = parser.finish()[0] if parser is not None else "" + return {**final, "message": message(tail)} + if parser is None: + return {"model": body.get("model"), "created_at": _now(), "message": message(text), "done": False} + content, calls = parser.feed(text) + if not content and not calls: + return None + return {"model": body.get("model"), "created_at": _now(), "message": message(content, calls), + "done": False} def collect(items): - text = "".join(i["message"]["content"] for i in items[:-1]) - return {**items[-1], "message": {"role": "assistant", "content": text}} + text = "".join(i["message"]["content"] for i in items) + calls = [c for i in items for c in i["message"].get("tool_calls", [])] + msg = {"role": "assistant", "content": text.strip() if calls else text} + if calls: + msg["tool_calls"] = calls + return {**items[-1], "message": msg} return _respond(body, _run(body, ids, lm, load_s, shape), collect) @app.get("/v1/models") @@ -298,24 +321,49 @@ async def openai_chat(request: Request): "seed": body.get("seed"), "stop": body.get("stop"), **(body.get("chronos") or {})}}) cid, created = f"chatcmpl-{uuid.uuid4().hex}", int(time.time()) + with_tools = bool(body.get("tools")) + + def finish_reason(final, n_calls): + return "tool_calls" if n_calls else final.get("done_reason", "stop") + if body.get("stream"): def sse(): + parser = ToolCallStream() if with_tools else None + n_calls = 0 + + def chunk(delta, reason=None, extra=None): + out = {"id": cid, "object": "chat.completion.chunk", "created": created, "model": name, + "choices": [{"index": 0, "delta": delta, "finish_reason": reason}], **(extra or {})} + return f"data: {json.dumps(out, ensure_ascii=False)}\n\n" for item in runtime.generate_stream(name, ids, **opts): if item.get("done"): - chunk = {"id": cid, "object": "chat.completion.chunk", "created": created, "model": name, - "choices": [{"index": 0, "delta": {}, "finish_reason": item.get("done_reason", "stop")}], - "chronos": item} - yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n" + tail = parser.finish()[0] if parser is not None else "" + if tail: + yield chunk({"content": tail}) + yield chunk({}, finish_reason(item, n_calls), {"chronos": item}) + "data: [DONE]\n\n" + elif parser is None: + yield chunk({"content": item["text"]}) else: - chunk = {"id": cid, "object": "chat.completion.chunk", "created": created, "model": name, - "choices": [{"index": 0, "delta": {"content": item["text"]}, "finish_reason": None}]} - yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" + content, calls = parser.feed(item["text"]) + if content: + yield chunk({"content": content}) + if calls: + yield chunk({"tool_calls": openai_tool_calls(calls, start=n_calls)}) + n_calls += len(calls) return StreamingResponse(sse(), media_type="text/event-stream") items = list(runtime.generate_stream(name, ids, **opts)) final = items[-1] + text = "".join(i["text"] for i in items[:-1]) + msg = {"role": "assistant", "content": text} + calls = [] + if with_tools: + content, calls = split_tool_calls(text) + if calls: + msg = {"role": "assistant", "content": content or None, + "tool_calls": [{k: v for k, v in c.items() if k != "index"} # index: stream deltas only + for c in openai_tool_calls(calls)]} return {"id": cid, "object": "chat.completion", "created": created, "model": name, - "choices": [{"index": 0, "message": {"role": "assistant", "content": "".join(i["text"] for i in items[:-1])}, - "finish_reason": final.get("done_reason", "stop")}], + "choices": [{"index": 0, "message": msg, "finish_reason": finish_reason(final, len(calls))}], "usage": {"prompt_tokens": final.get("prompt_tokens", 0), "completion_tokens": final.get("completion_tokens", 0), "total_tokens": final.get("prompt_tokens", 0) + final.get("completion_tokens", 0)}, "chronos": final} diff --git a/chronos/serving/tool_calls.py b/chronos/serving/tool_calls.py new file mode 100644 index 0000000..b79e7b7 --- /dev/null +++ b/chronos/serving/tool_calls.py @@ -0,0 +1,130 @@ +"""Parse ``{json}`` blocks out of generated text. + +The chat template (``chronos.tokenizer.model.CHAT_TEMPLATE``) teaches the +model to emit tool calls as ``{"name": ..., "arguments": {...}} +`` lines after (optional) assistant text. ``ToolCallStream`` +consumes decoded text pieces as they stream and splits them into visible +content and parsed calls, holding back only what could still be the start +of a tag, so clients never see half a tag. + +A block whose JSON does not parse, or that has no ``name``, is returned as +ordinary content: nothing the model wrote is dropped. +""" +from __future__ import annotations + +import json +import uuid + +OPEN, CLOSE = "", "" + + +def _prefix_hold(text: str, tag: str) -> int: + """Length of the longest suffix of ``text`` that is a proper prefix of ``tag``.""" + for n in range(min(len(text), len(tag) - 1), 0, -1): + if tag.startswith(text[-n:]): + return n + return 0 + + +def parse_call(body: str) -> dict | None: + """``{"name": str, "arguments": dict}`` or ``None`` for an invalid block.""" + try: + obj = json.loads(body) + except (TypeError, ValueError): + return None + if not isinstance(obj, dict): + return None + fn = obj.get("function") if isinstance(obj.get("function"), dict) else obj + name = fn.get("name") + if not isinstance(name, str) or not name: + return None + args = fn.get("arguments", fn.get("parameters", {})) + if isinstance(args, str): + try: + args = json.loads(args) if args.strip() else {} + except ValueError: + return None + if not isinstance(args, dict): + return None + return {"name": name, "arguments": args} + + +class ToolCallStream: + """Incremental splitter: ``feed(piece) -> (content, calls)``, then ``finish()``. + + Whitespace between consecutive calls (the template's ``"\\n"`` joiner) is + not reported as content once a call has been parsed.""" + + def __init__(self): + self._buf = "" # undecided text (outside a block) + self._block = None # text inside an open block, or None + self._pending_ws = "" # whitespace after a call, emitted only if real text follows + self.calls: list[dict] = [] + + def _emit_text(self, text: str) -> str: + if not text: + return "" + if self.calls and not text.strip(): + self._pending_ws += text + return "" + out, self._pending_ws = self._pending_ws + text, "" + return out + + def feed(self, piece: str) -> tuple[str, list[dict]]: + content, calls = [], [] + self._buf += piece or "" + while True: + if self._block is None: + i = self._buf.find(OPEN) + if i < 0: + keep = _prefix_hold(self._buf, OPEN) + content.append(self._emit_text(self._buf[: len(self._buf) - keep])) + self._buf = self._buf[len(self._buf) - keep:] + break + content.append(self._emit_text(self._buf[:i])) + self._block, self._buf = "", self._buf[i + len(OPEN):] + else: + j = self._buf.find(CLOSE) + if j < 0: + keep = _prefix_hold(self._buf, CLOSE) + self._block += self._buf[: len(self._buf) - keep] + self._buf = self._buf[len(self._buf) - keep:] + break + body, self._buf = self._block + self._buf[:j], self._buf[j + len(CLOSE):] + self._block = None + call = parse_call(body.strip()) + if call is None: + content.append(self._emit_text(OPEN + body + CLOSE)) + else: + self._pending_ws = "" + self.calls.append(call) + calls.append(call) + return "".join(content), calls + + def finish(self) -> tuple[str, list[dict]]: + """Flush held text; an unterminated block is returned as content.""" + tail = self._buf if self._block is None else OPEN + self._block + self._buf + self._buf, self._block = "", None + text = tail if not (self.calls and not tail.strip()) else "" + self._pending_ws = "" + return text, [] + + +def split_tool_calls(text: str) -> tuple[str, list[dict]]: + """One-shot form of ``ToolCallStream``.""" + stream = ToolCallStream() + content, calls = stream.feed(text) + tail, _ = stream.finish() + return (content + tail).strip() if calls else content + tail, calls + + +def ollama_tool_calls(calls: list[dict]) -> list[dict]: + """Ollama ``message.tool_calls`` (arguments stay an object).""" + return [{"function": {"name": c["name"], "arguments": c["arguments"]}} for c in calls] + + +def openai_tool_calls(calls: list[dict], *, start: int = 0) -> list[dict]: + """OpenAI ``tool_calls`` (arguments are a JSON string, each call has an id).""" + return [{"index": start + k, "id": f"call_{uuid.uuid4().hex[:24]}", "type": "function", + "function": {"name": c["name"], "arguments": json.dumps(c["arguments"], ensure_ascii=False)}} + for k, c in enumerate(calls)] diff --git a/chronos/trainer/grpo_trainer.py b/chronos/trainer/grpo_trainer.py index f89cecc..d15fdff 100644 --- a/chronos/trainer/grpo_trainer.py +++ b/chronos/trainer/grpo_trainer.py @@ -315,3 +315,28 @@ def load_grpo_prompts(data_path: str, max_prompts: int = None) -> List[str]: if max_prompts and len(prompts) >= max_prompts: break return prompts + + +def load_grpo_prompt_records(data_path: str, max_prompts: int = None) -> list: + """Like ``load_grpo_prompts``, but a record that carries image/audio + (``image``/``audio``/``*_embeds`` or a ``modalities`` dict) is returned as + ``{"prompt": str, "record": dict}`` so the trainer can resolve its media + lazily, once per step; text-only records stay plain strings.""" + import json + from chronos.data.flexible_dataset import detect_modalities + out = [] + with open(data_path, encoding="utf-8") as f: + for line in f: + if not line.strip(): + continue + row = json.loads(line) + msgs = row.get("conversations") or row.get("messages") or [] + prompt = next((m.get("content") for m in msgs if m.get("role") == "user"), None) + if prompt is None and isinstance(row.get("prompt"), str): + prompt = row["prompt"] + if prompt is None: + continue + out.append({"prompt": prompt, "record": row} if detect_modalities(row) else prompt) + if max_prompts and len(out) >= max_prompts: + break + return out diff --git a/chronos/workbench/training.py b/chronos/workbench/training.py index e94db12..0760b35 100644 --- a/chronos/workbench/training.py +++ b/chronos/workbench/training.py @@ -532,6 +532,18 @@ def _coerce_max_steps(value) -> int | None: return steps if steps > 0 else None +def _stage_collate(mode: str, modalities, multimodal_collate, preference_collate): + """DataLoader collate for a stage: media batches need the multimodal rules + (CE stages) or the pair rules (DPO/ORPO); text-only keeps the default.""" + if not modalities: + return None + if mode in {"pretrain", "sft", "distill"}: + return multimodal_collate + if mode in {"dpo", "orpo"}: + return preference_collate + return None + + def _planned_train_steps(data_iter, epochs: int, max_steps=None) -> int: """Return the real UI progress denominator in micro/batch steps. @@ -1064,6 +1076,10 @@ def _build_stage_args(self, cfg: dict, mode: str, save_dir: str, hidden_size: in num_experts=int(cfg.get("num_experts", 4)), steps=int(cfg.get("max_steps", 0)) or None, weight_decay=float(cfg.get("weight_decay", 0.01)), + image_size=int(cfg.get("image_size", 224)), + audio_sample_rate=int(cfg.get("audio_sample_rate", 16000)), + max_audio_seconds=float(cfg.get("max_audio_seconds", 30.0)), + media_error_policy=cfg.get("media_error_policy", "warn"), ) def _record_stage_metric(self, mode: str, step: int, result, tps: float = 0.0): @@ -1281,15 +1297,15 @@ def _run(self, cfg: dict, mode: str): requested_backend = (cfg.get("train_backend") or "auto").strip().lower() or "auto" resolved_backend, device = resolve_training_device(requested_backend) - if resolved_backend == "mlx" and model_cfg.modalities and mode not in {"pretrain", "sft"}: - # MLX trains the image/audio encoders and fusion in pretrain/SFT; - # the preference and distillation stages stay on PyTorch. - if requested_backend != "auto": - raise ValueError(f"MLX 多模态训练目前支持 pretrain/sft；{mode} 请选择 auto 或 PyTorch MPS/CUDA 后端。") - supported = set(training_available()) - target = next(name for name in ("cuda", "xpu", "mps", "cpu") if name in supported) - resolved_backend, device = resolve_training_device(target) - self._put(f"[data] Modalities {model_cfg.modalities}: using {resolved_backend} for {mode}.") + if resolved_backend != "mlx" and model_cfg.modalities and mode in {"dpo", "orpo", "grpo"}: + # MLX conditions DPO/ORPO pairs and GRPO rollouts on the prompt's + # image/audio; the PyTorch preference trainers are text-only and + # would silently drop the media. + if requested_backend != "auto" or "mlx" not in set(training_available()): + raise ValueError(f"多模态 {mode.upper()} 训练需要 MLX 后端（PyTorch 的 {mode.upper()} 只支持文本）。" + "请选择 auto 或 MLX。") + resolved_backend, device = resolve_training_device("mlx") + self._put(f"[data] Modalities {model_cfg.modalities}: using mlx for {mode}.") cfg["train_backend"] = requested_backend cfg["device"] = device cfg["dtype"] = _normalize_dtype_for_trainer(cfg.get("dtype")) @@ -1339,8 +1355,8 @@ def _run(self, cfg: dict, mode: str): tokenizer = self._load_tokenizer() from chronos.data.flexible_dataset import ( FlexibleDataset, StreamingSFTDataset, MultimodalStreamingDataset, - MultimodalStreamingSFTDataset, StreamingDPODataset, - dataset_modality_summary, multimodal_collate, + MultimodalStreamingSFTDataset, StreamingDPODataset, MultimodalStreamingDPODataset, + dataset_modality_summary, multimodal_collate, preference_collate, ) modality_summary = dataset_modality_summary(data_path) detected_modalities = modality_summary.get("detected_modalities", []) @@ -1367,7 +1383,16 @@ def _run(self, cfg: dict, mode: str): if detected_modalities else StreamingSFTDataset(data_path, tokenizer, max_length=max_seq_len) ) elif mode in {"dpo", "orpo"}: - full_ds = StreamingDPODataset(data_path, tokenizer, max_length=max_seq_len) + if cfg.get("auto_detect_modalities", True) and detected_modalities: + model_cfg_kwargs["modalities"] = detected_modalities + full_ds = ( + MultimodalStreamingDPODataset(data_path, tokenizer, max_length=max_seq_len, + image_size=int(cfg.get("image_size", 224)), + audio_sample_rate=int(cfg.get("audio_sample_rate", 16000)), + max_audio_seconds=float(cfg.get("max_audio_seconds", 30.0)), + media_error_policy=cfg.get("media_error_policy", "warn")) + if detected_modalities else StreamingDPODataset(data_path, tokenizer, max_length=max_seq_len) + ) else: full_ds = None # Deterministic val split: every 1/val_ratio-th record goes to val @@ -1388,7 +1413,7 @@ def _run(self, cfg: dict, mode: str): DataLoader( train_ds, batch_size=batch_size, - collate_fn=multimodal_collate if detected_modalities and mode in {"pretrain", "sft", "distill"} else None, + collate_fn=_stage_collate(mode, detected_modalities, multimodal_collate, preference_collate), **dataloader_kwargs(device=device, num_workers=cfg.get("num_workers", "auto"), shuffle=True), ) if train_ds is not None else None @@ -1397,7 +1422,7 @@ def _run(self, cfg: dict, mode: str): DataLoader( val_ds, batch_size=batch_size, - collate_fn=multimodal_collate if detected_modalities and mode in {"pretrain", "sft", "distill"} else None, + collate_fn=_stage_collate(mode, detected_modalities, multimodal_collate, preference_collate), **dataloader_kwargs(device=device, num_workers=cfg.get("num_workers", "auto"), shuffle=False), ) if val_ds is not None else None @@ -1411,8 +1436,10 @@ def _run(self, cfg: dict, mode: str): workers = None if mode in {"grpo"}: tokenizer = self._load_tokenizer() - from chronos.trainer.grpo_trainer import load_grpo_prompts - prompts = load_grpo_prompts(data_path, max_prompts=None) + from chronos.trainer.grpo_trainer import load_grpo_prompt_records, load_grpo_prompts + # MLX resolves each record's image/audio per step; Torch GRPO is text-only. + prompts = (load_grpo_prompt_records(data_path, max_prompts=None) if resolved_backend == "mlx" + else load_grpo_prompts(data_path, max_prompts=None)) if not prompts: raise ValueError(f"[{mode.upper()}] no prompts found in {data_path}") else: diff --git a/tests/test_mlx_multimodal_stages.py b/tests/test_mlx_multimodal_stages.py new file mode 100644 index 0000000..598b218 --- /dev/null +++ b/tests/test_mlx_multimodal_stages.py @@ -0,0 +1,261 @@ +"""Multimodal DPO / ORPO / GRPO / distill on MLX: data path, media effect, +tiling parity, rollout encoder reuse, and end-to-end runs through +``run_mlx_stage``.""" +from __future__ import annotations + +import json +import os +import sys + +import numpy as np +import pytest +import torch + +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +EMB = 20 # pre-encoded feature width (fusion fits it to hidden_size) + + +def _wav(path, n=4000, rate=16000): + import wave + pcm = (np.sin(np.arange(n) * 0.05) * 12000).astype(" 0 and int(item["mask_rejected"].sum()) > 0 + assert "modality_inputs" not in items[1] + batch = preference_collate(items) + assert batch["x_chosen"].shape == (3, 23) + assert batch["modality_mask"].tolist() == [[1.0, 0.0], [0.0, 0.0], [0.0, 1.0]] + media = batch["modality_inputs"] + assert media["image_embeds"].shape == (3, 5, EMB) + assert media["image_embeds_mask"].tolist()[1] == [False] * 5 + assert media["audio"].shape == (3, 4000) and media["audio_mask"].tolist() == [False, False, True] + assert "modality_inputs" not in preference_collate([items[1], items[1]]) + + +def test_grpo_prompt_records_keep_media_lazily(data): + from chronos.trainer.grpo_trainer import load_grpo_prompt_records, load_grpo_prompts + _tok, path, _ = data + recs = load_grpo_prompt_records(path) + assert [isinstance(r, dict) for r in recs] == [True, False, True] + assert [r["prompt"] if isinstance(r, dict) else r for r in recs] == load_grpo_prompts(path) + + +# ---- MLX ------------------------------------------------------------------- + +def _cfg(**kw): + from chronos.model.config import ChronosConfig + base = dict(hidden_size=32, num_hidden_layers=3, num_attention_heads=4, num_key_value_heads=2, + vocab_size=600, num_experts=8, num_experts_per_tok=2, num_shared_experts=1, + moe_intermediate_size=32, max_position_embeddings=64, recurrent_depth=2, + prelude_layers=1, coda_layers=1, dropout=0.0, fallback_mask_prob=0.0, + modalities=["image", "audio"], image_patch_size=8, max_image_tokens=12, + max_audio_tokens=10, modality_dropout=0.0) + base.update(kw) + return ChronosConfig(**base) + + +def _torch_model(c, seed=0): + from chronos.model.model_chronos import ChronosForCausalLM + torch.manual_seed(seed) + pt = ChronosForCausalLM(c).eval() + with torch.no_grad(): # gate 0 → sigmoid 0.5 everywhere; randomize so media matter + fusion = pt.model.multimodal_fusion + fusion.gate.fill_(1.5) + for p in fusion.modality_embeddings.values(): + p.normal_(0, 0.5) + return pt + + +def _mlx_model(c, seed=0): + from chronos.mlx.model import ChronosMLXModel + m = ChronosMLXModel.from_chronos_pytorch(_torch_model(c, seed), c) + m.eval() + return m + + +def _pair_batch(data): + from chronos.data.flexible_dataset import MultimodalStreamingDPODataset, preference_collate + tok, path, _ = data + ds = MultimodalStreamingDPODataset(path, tok, max_length=24) + return preference_collate([ds[i] for i in range(3)]) + + +def _strip_media(batch): + return {k: v for k, v in batch.items() if k not in ("modality_inputs", "modality_mask")} + + +def test_dpo_orpo_tiling_matches_per_side_forward_and_media_change_the_loss(data): + mx = pytest.importorskip("mlx.core") + from chronos.mlx.training.trainer import _dpo_stage_loss, _orpo_stage_loss, _tile_media, _ids_to_mx + batch = _pair_batch(data) + with mx.stream(mx.cpu): + m = _mlx_model(_cfg()) + tiled = _tile_media(batch, 2) + x = mx.concatenate([_ids_to_mx(batch["x_chosen"]), _ids_to_mx(batch["x_rejected"])], axis=0) + both, _, _ = m(x, logits_dtype=mx.float32, **tiled) + once = _tile_media(batch, 1) + alone_c, _, _ = m(_ids_to_mx(batch["x_chosen"]), logits_dtype=mx.float32, **once) + alone_r, _, _ = m(_ids_to_mx(batch["x_rejected"]), logits_dtype=mx.float32, **once) + np.testing.assert_allclose(np.array(both), np.concatenate([np.array(alone_c), np.array(alone_r)]), atol=1e-5) + text, _, _ = m(x, logits_dtype=mx.float32) + rows = np.abs(np.array(both) - np.array(text)).max(axis=(1, 2)) + assert rows[[0, 2, 3, 5]].min() > 1e-3 # media rows (pairs 0 and 2, both sides) differ + np.testing.assert_allclose(rows[[1, 4]], 0, atol=1e-5) # the text-only pair is untouched + for fn in (lambda b: _dpo_stage_loss(m, m, b), lambda b: _orpo_stage_loss(m, b)): + with_media, without = float(fn(batch)), float(fn(_strip_media(batch))) + assert np.isfinite(with_media) and abs(with_media - without) > 1e-6 + + +def test_rollout_encodes_media_once_with_identical_logits(): + mx = pytest.importorskip("mlx.core") + from chronos.mlx.training.trainer import _encoded_media + g = torch.Generator().manual_seed(3) + media = {"modality_inputs": {"image": torch.randn(1, 3, 16, 24, generator=g), + "audio": torch.randn(1, 4000, generator=g)}, + "modality_mask": torch.tensor([[1.0, 1.0]])} + ids = mx.array(np.random.default_rng(1).integers(3, 600, (1, 7))) + with mx.stream(mx.cpu): + m = _mlx_model(_cfg(), seed=2) + enc = _encoded_media(m, media) + assert set(enc["modality_inputs"]) == {"image_embeds", "audio_embeds"} + raw, _, cache = m(ids[:, :-1], logits_dtype=mx.float32, **media) + pre, _, cache2 = m(ids[:, :-1], logits_dtype=mx.float32, **enc) + np.testing.assert_allclose(np.array(pre), np.array(raw), atol=1e-5) + step_raw, _, _ = m(ids[:, -1:], cache=cache, logits_dtype=mx.float32, **media) + step_pre, _, _ = m(ids[:, -1:], cache=cache2, logits_dtype=mx.float32, **enc) + np.testing.assert_allclose(np.array(step_pre), np.array(step_raw), atol=1e-5) + assert _encoded_media(m, None) == {} + + +class _Args: + learning_rate = 1e-3 + weight_decay = 0.0 + steps = 2 + epochs = 1 + max_seq_len = 24 + max_gen_len = 4 + num_generations = 2 + temperature = 1.0 + log_interval = 1 + save_interval = 100 + accumulation_steps = 1 + grad_clip = 1.0 + dtype = "float32" + reward = "toy" + + +@pytest.mark.parametrize("stage", ["dpo", "orpo", "grpo", "distill"]) +def test_run_mlx_stage_trains_fusion_with_media(stage, data, tmp_path, monkeypatch): + mx = pytest.importorskip("mlx.core") + from torch.utils.data import DataLoader + from chronos.data.flexible_dataset import (MultimodalStreamingDPODataset, MultimodalStreamingSFTDataset, + multimodal_collate, preference_collate) + from chronos.model.checkpoint import load_checkpoint_state_dict, save_state_dict_with_config + from chronos.trainer.grpo_trainer import load_grpo_prompt_records + from chronos.mlx.training import run_mlx_stage + import chronos.deps + tok, path, _ = data + c = _cfg() + ckpt = str(tmp_path / "init.pth") + pt = _torch_model(c, seed=7) + save_state_dict_with_config(pt, ckpt, c, stage="sft") + before = {k: v.clone() for k, v in pt.state_dict().items() if k.startswith("model.multimodal_fusion.")} + loader = prompts = None + if stage in {"dpo", "orpo"}: + loader = DataLoader(MultimodalStreamingDPODataset(path, tok, max_length=24), batch_size=2, + collate_fn=preference_collate) + elif stage == "distill": + loader = DataLoader(MultimodalStreamingSFTDataset(path, tok, max_length=24), batch_size=2, + collate_fn=multimodal_collate) + else: + prompts = [r for r in load_grpo_prompt_records(path) if isinstance(r, dict)] + monkeypatch.setattr(chronos.deps, "get_tokenizer_path", lambda: str(tmp_path.parent)) + import transformers + monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", classmethod(lambda cls, *_a, **_k: tok)) + events = [] + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + result = run_mlx_stage(stage=stage, config=c, checkpoint_path=ckpt, save_dir=str(tmp_path / "out"), + loader=loader, prompts=prompts, teacher_path=ckpt if stage == "distill" else None, + args=_Args(), progress_callback=events.append) + finally: + mx.set_default_device(previous) + assert not result.failure_reason and result.steps == 2 and result.checkpoint_saved, result + assert all(np.isfinite(e["loss"]) for e in events if "loss" in e) + after = load_checkpoint_state_dict(result.checkpoint_path, map_location="cpu") + moved = [k for k, v in before.items() if not torch.equal(after[k].float(), v.float())] + assert any("q_proj" in k or "in_proj" in k for k in moved), moved # fusion received gradients + if stage == "grpo": # raw audio prompt: the encoder trains too + assert any("audio_encoder" in k for k in moved), moved + + +def test_workbench_stage_collate_routes_media_batches(): + from chronos.data.flexible_dataset import multimodal_collate, preference_collate + from chronos.workbench.training import _stage_collate + pick = lambda mode, mods: _stage_collate(mode, mods, multimodal_collate, preference_collate) + assert pick("sft", []) is None and pick("dpo", []) is None + assert pick("pretrain", ["image"]) is multimodal_collate and pick("distill", ["audio"]) is multimodal_collate + assert pick("dpo", ["image"]) is preference_collate and pick("orpo", ["audio"]) is preference_collate + assert pick("grpo", ["image"]) is None # GRPO iterates prompt records, not a DataLoader diff --git a/tests/test_tool_call_serving.py b/tests/test_tool_call_serving.py new file mode 100644 index 0000000..8dedf5e --- /dev/null +++ b/tests/test_tool_call_serving.py @@ -0,0 +1,140 @@ +"""Tool-call inference: parsing and Ollama/OpenAI response shapes.""" +from __future__ import annotations + +import json +import os +import sys + +import pytest + +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +from chronos.serving.tool_calls import ToolCallStream, parse_call, split_tool_calls # noqa: E402 + +TOOLS = [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}] +CALL_TEXT = ('Let me check.\n{"name": "get_weather", "arguments": {"city": "北京"}}\n' + '{"name": "get_time", "arguments": {}}') + + +def _pieces(text, size): + return [text[i:i + size] for i in range(0, len(text), size)] + + +@pytest.mark.parametrize("size", [1, 2, 3, 7, 1000]) +def test_stream_splits_any_chunking_identically(size): + s = ToolCallStream() + content, calls = "", [] + for piece in _pieces(CALL_TEXT, size): + c, k = s.feed(piece) + assert "{"arguments": {}}')[1] == [] # no name + assert split_tool_calls("a {\"name\": \"f\"")[0] == "a {\"name\": \"f\"" # unterminated + assert parse_call('{"function": {"name": "f", "arguments": "{\\"a\\": 1}"}}') == {"name": "f", "arguments": {"a": 1}} + assert parse_call('{"name": "f", "arguments": [1]}') is None + assert parse_call('{"name": "f"}') == {"name": "f", "arguments": {}} + + +def test_round_trip_with_training_template(): + """What the template teaches the model to write is exactly what the server parses.""" + from chronos.data.chat_records import render_chat + calls = [{"function": {"name": "get_weather", "arguments": {"city": "北京"}}}, {"name": "get_time", "arguments": {}}] + msgs = [{"role": "user", "content": "天气?"}, {"role": "assistant", "content": "Let me check.", "tool_calls": calls}] + learned = "".join(t for t, sup in render_chat(msgs, TOOLS) if sup).removesuffix("<|im_end|>\n") + assert learned == CALL_TEXT + assert split_tool_calls(learned) == ("Let me check.", [{"name": "get_weather", "arguments": {"city": "北京"}}, + {"name": "get_time", "arguments": {}}]) + + +class _Runtime: + """Stub runtime: streams fixed text in small pieces.""" + + def __init__(self, text, piece=3): + self.text, self.piece, self.tools_seen = text, piece, [] + + def available(self): + return {"m": __file__} + + def get(self, name): + return type("LM", (), {"name": name})() + + def encode_chat(self, lm, messages, tools=None): + self.tools_seen.append(tools) + return [1, 2, 3] + + def generate_stream(self, name, ids, **_opts): + for p in _pieces(self.text, self.piece): + yield {"text": p} + yield {"done": True, "done_reason": "stop", "prompt_tokens": 3, "completion_tokens": 9} + + +def _client(text, piece=3): + pytest.importorskip("fastapi") + from fastapi.testclient import TestClient + from chronos.serving.ollama_server import create_app + rt = _Runtime(text, piece) + return TestClient(create_app(rt)), rt + + +def test_ollama_chat_returns_tool_calls_stream_and_non_stream(): + body = {"model": "m", "messages": [{"role": "user", "content": "天气?"}], "tools": TOOLS} + client, rt = _client(CALL_TEXT) + msg = client.post("/api/chat", json={**body, "stream": False}).json()["message"] + assert msg["content"] == "Let me check." + assert msg["tool_calls"] == [{"function": {"name": "get_weather", "arguments": {"city": "北京"}}}, + {"function": {"name": "get_time", "arguments": {}}}] + assert rt.tools_seen[-1] == TOOLS + lines = [json.loads(x) for x in client.post("/api/chat", json=body).text.splitlines() if x.strip()] + assert lines[-1]["done"] is True + streamed = [c for line in lines for c in line["message"].get("tool_calls", [])] + assert streamed == msg["tool_calls"] + assert "".join(line["message"]["content"] for line in lines).strip() == "Let me check." + assert not any("