From e32660dacdeb540d9dd783a0541fd935d626b43b Mon Sep 17 00:00:00 2001 From: Fona <62028332+FonaTech@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:15:28 +0800 Subject: [PATCH 1/2] 0.3.2a3: compile segment-mode cluster locality, recover from a failed compile trace - cluster_locality_mode="segment" now traces under mx.compile: document ids come from segment_ids_mx(input_ids, chunk=window) and the loss uses static shapes (B*S segment slots, unused slots weigh 0). Matches the NumPy ids and the Torch loss. - CompiledStep snapshots weights, RNG key and side outputs before a first trace and restores them when the trace fails, so the eager fallback no longer trains on tracers. - tests/test_segment_compile.py (5 cases; 4 fail on 0.3.2a2). Co-Authored-By: Claude Sonnet 5.5 --- CHANGELOG.md | 14 +++ chronos/__init__.py | 2 +- chronos/expert_clusters.py | 25 +++-- chronos/mlx/training/compiled.py | 37 +++++++- chronos/mlx/training/trainer.py | 38 ++++---- pyproject.toml | 2 +- tests/test_segment_compile.py | 157 +++++++++++++++++++++++++++++++ 7 files changed, 248 insertions(+), 27 deletions(-) create mode 100644 tests/test_segment_compile.py diff --git a/CHANGELOG.md b/CHANGELOG.md index baf6fec..c6c76ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,20 @@ ## [Unreleased] +## [0.3.2a3] — 2026-11-10 (beta) + +Fixes two kernel bugs found while benchmarking the 50M and 30M multilingual configs on the real corpus. Nothing new is switched on by default. + +### Fixed +- `cluster_locality_mode="segment"` could not be compiled. The loss found document ids with NumPy (`segment_ids`, `np.unique`) on a traced array, so `CHRONOS_MLX_COMPILE=on` raised "Attempting to eval an array during function transformations", and the default `auto` mode crashed on the next `mx.eval`. The loss now uses `segment_ids_mx(input_ids, chunk=window)`, which has a static shape (`row * S + id` addresses `B * S` segment slots and unused slots weigh 0), and runs inside the compiled step. It matches the NumPy ids on 300 random batches and the Torch loss within 1.7e-5 on 60. +- A compile trace that failed part way left tracers in the model weights, the RNG key and the router outputs, because `mx.compile` swaps them back only on success. The eager fallback then trained on them and the next `mx.eval` raised "Attempting to eval an array without a primitive". `CompiledStep` now snapshots those on a first trace and restores them on failure; the RNG is re-seeded from the saved key words, so the random stream continues exactly (checked on 2,000 random keys). + +### Measured +- Segment mode, compiled, real corpus, 10 steps, document masking on, steady tokens/s: 50M 24,856 at ctx 2048 and 20,416 at 4096; 30M 22,454 and 16,416. The eager fallback cost 15-17%. + +### Tests +- `tests/test_segment_compile.py` (5 cases) fails on 0.3.2a2 (4 of 5) and passes here: ids equal the NumPy rule with chunking, loss equals the Torch reference eagerly and compiled, a one-segment batch gives 0, a compiled segment-mode step matches the eager gradients, and a failed trace leaves weights, RNG and side outputs usable. + ## [0.3.2a2] — 2026-11-09 (beta) Adds narrower recomputation, two lookahead options and web UI fields for every long-context switch, and fixes a crash found by training on the real corpus. Every new switch is off by default. `reports/long_context_20261109/real_corpus_ab.json` holds the measurements below. diff --git a/chronos/__init__.py b/chronos/__init__.py index 61ba8e2..6dcbeb7 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.2a2" +__version__ = "0.3.2a3" # Register with HuggingFace AutoXXX factories so # AutoModelForCausalLM.from_pretrained("path/to/chronos") works out of the box. diff --git a/chronos/expert_clusters.py b/chronos/expert_clusters.py index 7c116de..37b1e36 100644 --- a/chronos/expert_clusters.py +++ b/chronos/expert_clusters.py @@ -226,15 +226,28 @@ def segment_ids(input_ids, chunk: int = 0, boundary_id: int = SEGMENT_BOUNDARY_I return np.cumsum(new, axis=1) -def segment_ids_mx(input_ids, boundary_id: int = SEGMENT_BOUNDARY_ID): - """``segment_ids`` (chunk = 0) for an MLX array, traceable under ``mx.compile``. +def segment_ids_mx(input_ids, boundary_id: int = SEGMENT_BOUNDARY_ID, chunk: int = 0): + """``segment_ids`` for an MLX array, traceable under ``mx.compile`` (no NumPy, no data-dependent shape). Same rule: a separator token closes the document it ends, so the id of position t is the - number of separators strictly before t. Returns an int32 ``[B, S]`` array.""" + number of separators strictly before t. ``chunk`` > 0 cuts a document into ``chunk``-token parts, + as ``segment_ids`` does. Returns int32 ``[B, S]`` ids that increase along each row and stay below + ``S``, so ``row * S + id`` is a dense batch-wide segment number with the static bound ``B * S``.""" import mlx.core as mx - is_bnd = (input_ids == int(boundary_id)).astype(mx.int32) - before = mx.cumsum(is_bnd, axis=1) - is_bnd - return mx.stop_gradient(before) + ids = input_ids if isinstance(input_ids, mx.array) else mx.array(input_ids) + is_bnd = (ids == int(boundary_id)).astype(mx.int32) + doc = mx.cumsum(is_bnd, axis=1) - is_bnd + if int(chunk or 0) <= 0: + return mx.stop_gradient(doc) + B, S = ids.shape + pos = mx.broadcast_to(mx.arange(S, dtype=mx.int32)[None, :], (B, S)) + doc_changed = doc[:, 1:] != doc[:, :-1] + first = mx.concatenate([mx.ones((B, 1), dtype=mx.bool_), doc_changed], axis=1) + start = mx.cummax(mx.where(first, pos, 0), axis=1) + part = (pos - start) // int(chunk) + new = doc_changed | (part[:, 1:] != part[:, :-1]) + seg = mx.concatenate([mx.zeros((B, 1), dtype=mx.int32), mx.cumsum(new.astype(mx.int32), axis=1)], axis=1) + return mx.stop_gradient(seg) def segment_index(segments): diff --git a/chronos/mlx/training/compiled.py b/chronos/mlx/training/compiled.py index b4422af..d500080 100644 --- a/chronos/mlx/training/compiled.py +++ b/chronos/mlx/training/compiled.py @@ -77,6 +77,33 @@ def _side_set(model, side: dict) -> None: object.__setattr__(m, "_shared_topk_idx", topk) +def _state_slots(node, out: list) -> list: + """``(container, key, array)`` for every array leaf of a module's state tree.""" + for key, value in list(node.items() if isinstance(node, dict) else enumerate(node)): + if isinstance(value, mx.array): + out.append((node, key, value)) + elif isinstance(value, (dict, list)): + _state_slots(value, out) + return out + + +def _capture(model) -> dict: + """Weights, RNG key and side outputs as they are before a first trace.""" + return {"slots": _state_slots(model.state, []), + "rng": [int(w) for w in mx.random.state[0].tolist()], + "side": _side_get(model)} + + +def _restore(model, snap: dict) -> None: + """Undo what a failed trace left behind: the saved arrays go back into the module tree, the + RNG is re-seeded from the saved key words (the stream continues exactly), side outputs are reset.""" + for node, key, value in snap["slots"]: + node[key] = value + hi, lo = snap["rng"] + mx.random.seed((hi << 32) | lo) + _side_set(model, snap["side"]) + + def _static(value): """Hashable stand-in for a non-array argument, or ``None`` when unsupported.""" if value is None or isinstance(value, (bool, int, float, str)): @@ -206,7 +233,8 @@ def __call__(self, *args): with self._pin(depth): return self.eager(*args) hit = self._graphs.get(key) - if hit is None: + fresh = hit is None + if fresh: self._prewarm(args) hit = self._build(key, args) self._graphs[key] = hit @@ -217,11 +245,16 @@ def __call__(self, *args): self._graphs.move_to_end(key) fn, meta = hit arrays = [a for a in args if isinstance(a, mx.array)] + # A trace that fails half way leaves tracers in the weights, the RNG key and the side outputs + # (mx.compile swaps them back only on success), so snapshot before tracing and restore on failure. + snapshot = _capture(self.model) if fresh else None try: with self._pin(depth): loss, grads, flat = fn(*arrays) - except Exception as exc: # trace-time failure: never lose the step + except Exception as exc: # trace-time failure: never lose the step self._graphs.pop(key, None) + if snapshot is not None: + _restore(self.model, snapshot) if self.mode == "on": raise self._failed.add(key) diff --git a/chronos/mlx/training/trainer.py b/chronos/mlx/training/trainer.py index 582b49e..bc9bff3 100644 --- a/chronos/mlx/training/trainer.py +++ b/chronos/mlx/training/trainer.py @@ -251,26 +251,29 @@ def entropy(p): def _segment_cluster_loss(router_4d: mx.array, membership: mx.array, segments, - token_mask: mx.array | None = None) -> mx.array: + 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 + ``chronos.model.temporal_loss.segment_cluster_loss``), written with static shapes so it traces + under ``mx.compile``. + + ``segments`` ``[B, S]`` comes from ``chronos.expert_clusters.segment_ids`` (NumPy) or + ``segment_ids_mx``: the token stream's own document separators, no label or task token. Ids stay + below ``S``, so ``row * S + id`` indexes ``B * S`` segment slots; unused slots have weight 0. Each + segment's cluster mass is pooled over its real tokens and every dispatch; segments weigh by + real-token count. A batch with a single segment gives 0, as in the Torch twin.""" B, S, V, E = router_4d.shape C = int(membership.shape[1]) - inv, n = segment_index(segments) - if C <= 1 or n <= 1: + if C <= 1: return mx.zeros(()) + seg = (segments if isinstance(segments, mx.array) else mx.array(segments)).astype(mx.int32) + slots = B * S + idx = (seg + (mx.arange(B, dtype=mx.int32) * S)[:, None]).reshape(-1) 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) + else token_mask.reshape(B * S).astype(mx.float32)) + total = mx.zeros((slots, C), mx.float32).at[idx].add(mass * m[:, None]) + count = mx.zeros((slots,), mx.float32).at[idx].add(m) + present = mx.zeros((slots,), mx.float32).at[idx].add(mx.ones((B * S,), mx.float32)) # padding counts ps = total / mx.maximum(count, 1.0)[:, None] weight = count / mx.maximum(count.sum(), 1.0) @@ -280,17 +283,18 @@ def entropy(p): marginal = (ps * weight[:, None]).sum(axis=0) mutual = entropy(marginal) - (entropy(ps) * weight).sum() - return 1.0 - mutual / math.log(C) + loss = 1.0 - mutual / math.log(C) + return mx.where((present > 0).sum() > 1, loss, mx.zeros(())) 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 + from chronos.expert_clusters import cluster_locality_mode, segment_ids_mx 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 _segment_cluster_loss(router_4d, membership, segment_ids_mx(input_ids, chunk=window), token_mask) return _cluster_locality_loss(router_4d, membership, window, token_mask) diff --git a/pyproject.toml b/pyproject.toml index b07567a..fd14572 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "Project_Chronos" -version = "0.3.2a2" +version = "0.3.2a3" 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_segment_compile.py b/tests/test_segment_compile.py new file mode 100644 index 0000000..8aac4c9 --- /dev/null +++ b/tests/test_segment_compile.py @@ -0,0 +1,157 @@ +"""Segment-mode cluster locality under mx.compile, and recovery from a failed trace (0.3.2a3).""" +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') +nn = pytest.importorskip('mlx.nn') +from mlx.utils import tree_flatten # noqa: E402 + + +def _config(**kw): + from chronos.model.config import ChronosConfig + base = dict(architecture_version=3, hidden_size=64, num_hidden_layers=4, prelude_layers=1, coda_layers=1, + recurrent_depth=3, min_recurrent_depth=2, num_attention_heads=4, num_key_value_heads=2, + kv_latent_dim=16, rope_dim=8, num_experts=6, num_experts_per_tok=2, num_shared_experts=1, + moe_intermediate_size=32, vocab_size=97, max_position_embeddings=128, sliding_window_size=8, + dropout=0.0, fallback_mask_prob=0.0, adaptive_halting=True, lookahead_steps=2, + layer_aware_lookahead=True, lambda_cluster_locality=0.5, cluster_locality_window=8, + cluster_locality_mode='segment') + base.update(kw) + return ChronosConfig(**base) + + +def _model(cfg): + from chronos.mlx.model import ChronosMLXModel + mx.random.seed(0) + model = ChronosMLXModel(cfg) + mx.eval(model.parameters()) + model.train() + return model + + +def _batch(cfg, b=2, s=40): + mx.random.seed(1) + ids = mx.random.randint(3, cfg.vocab_size, (b, s)) + ids = mx.where(mx.arange(s)[None, :] % 11 == 10, 0, ids) # several documents per row + labels = mx.random.randint(3, cfg.vocab_size, (b, s)) + return ids, labels + + +def _live(*trees) -> bool: + try: + mx.eval(*trees) + return True + except Exception: + return False + + +def test_segment_ids_mx_equals_numpy_with_chunks(): + from chronos.expert_clusters import segment_ids, segment_ids_mx + rng = np.random.default_rng(7) + for _ in range(120): + b, s = int(rng.integers(1, 4)), int(rng.integers(2, 70)) + ids = rng.integers(0, 6, (b, s)) + chunk = int(rng.choice([0, 1, 3, 8, 256])) + assert np.array_equal(np.array(segment_ids_mx(mx.array(ids), chunk=chunk)), segment_ids(ids, chunk)) + + +def test_static_segment_loss_matches_torch_eager_and_compiled(): + torch = pytest.importorskip('torch') + from chronos.expert_clusters import segment_ids, segment_ids_mx + from chronos.mlx.training.trainer import _segment_cluster_loss + from chronos.model.temporal_loss import segment_cluster_loss + rng = np.random.default_rng(3) + mem = mx.array(np.eye(6, 3, dtype=np.float32) + np.eye(6, 3, k=-3, dtype=np.float32)) + for _ in range(25): + b, s = int(rng.integers(1, 4)), int(rng.integers(4, 48)) + ids = rng.integers(0, 5, (b, s)) + r = mx.softmax(mx.array(rng.normal(size=(b, s, 3, 6)).astype(np.float32)), axis=-1) + mask = mx.array(rng.random((b, s)) > 0.2) + chunk = int(rng.choice([0, 4, 16])) + want = float(segment_cluster_loss(torch.tensor(np.array(r)), torch.tensor(np.array(mem)), + segment_ids(ids, chunk), torch.tensor(np.array(mask)))) + eager = float(_segment_cluster_loss(r, mem, mx.array(segment_ids(ids, chunk)), mask)) + traced = float(mx.compile(lambda rr, ii, mm: _segment_cluster_loss( + rr, mem, segment_ids_mx(ii, chunk=chunk), mm))(r, mx.array(ids), mask)) + assert eager == pytest.approx(want, abs=1e-4) + assert traced == pytest.approx(want, abs=1e-4) + + +def test_single_segment_batch_has_no_locality_loss(): + from chronos.mlx.training.trainer import _segment_cluster_loss + r = mx.softmax(mx.random.normal((2, 8, 2, 6)), axis=-1) + mem = mx.array(np.eye(6, 3, dtype=np.float32) + np.eye(6, 3, k=-3, dtype=np.float32)) + one_row = mx.zeros((1, 8), dtype=mx.int32) + assert float(_segment_cluster_loss(r[:1], mem, one_row)) == 0.0 + + +def test_segment_mode_runs_in_the_compiled_step_and_matches_eager(): + from chronos.mlx.training.compiled import make_grad_fn + from chronos.mlx.training.trainer import _ce_stage_loss, _collect_router_probs + cfg = _config() + ids, labels = _batch(cfg) + model = _model(cfg) + eager_loss, eager_grads = nn.value_and_grad(model, _ce_stage_loss)(model, ids, labels) + mx.eval(eager_loss, eager_grads) + want = dict(tree_flatten(eager_grads)) + model2 = _model(cfg) + step = make_grad_fn(model2, _ce_stage_loss, mode='on', log=lambda m: None) + loss, grads = step(model2, ids, labels) + mx.eval(loss, grads) + assert step.stats['compiled'] == 1 and step.stats['eager'] == 0 + assert abs(loss.item() - eager_loss.item()) < 1e-4 + for name, g in dict(tree_flatten(grads)).items(): + denom = float(mx.maximum(mx.abs(want[name]).max(), 1e-12).item()) + assert float((mx.abs(g - want[name]).max()).item()) / denom < 1e-3, name + assert _live(_collect_router_probs(model2)) + assert model2._last_loss_components['cluster_locality_loss'].item() > 0.0 + + +def _failing_loss(model, ids, labels): + from chronos.mlx.training.trainer import _ce_stage_loss + loss = _ce_stage_loss(model, ids, labels) + np.asarray(ids) # a host read of a traced array: always fails while tracing + return loss + + +def test_failed_trace_leaves_weights_rng_and_side_outputs_usable(): + """A trace that raises used to leave tracers in the weights, the RNG key and the router outputs; + the eager fallback then ran on them and the next mx.eval raised 'array without a primitive'.""" + from chronos.mlx.training.compiled import CompiledStep + from chronos.mlx.training.trainer import _collect_router_probs + cfg = _config(dropout=0.1, fallback_mask_prob=0.2) + ids, labels = _batch(cfg) + lines: list = [] + model = _model(cfg) + mx.random.seed(11) + ref_loss, ref_grads = nn.value_and_grad(model, _failing_loss_eager)(model, ids, labels) + mx.eval(ref_loss, ref_grads) + ref_next = mx.random.normal((4,)) + mx.eval(ref_next) + model2 = _model(cfg) + step = CompiledStep(model2, _failing_loss, mode='auto', log=lines.append) + mx.random.seed(11) + loss, grads = step(model2, ids, labels) + assert any('compile failed' in line for line in lines) + assert _live(model2.parameters(), mx.random.state, loss, grads) + assert _live(_collect_router_probs(model2), list(model2._last_loss_components.values())) + assert abs(loss.item() - ref_loss.item()) < 1e-6 + nxt = mx.random.normal((4,)) + mx.eval(nxt) + assert np.array_equal(np.array(nxt), np.array(ref_next)) # the RNG stream continues exactly + again_loss, _ = step(model2, ids, labels) # a second step takes the eager path + assert _live(again_loss) and step.stats['eager'] == 2 + + +def _failing_loss_eager(model, ids, labels): + from chronos.mlx.training.trainer import _ce_stage_loss + return _ce_stage_loss(model, ids, labels) From 4faef7776e032494d41bfd41801d2f6e0bd4d710 Mon Sep 17 00:00:00 2001 From: Fona <62028332+FonaTech@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:56:10 +0800 Subject: [PATCH 2/2] Reject token ids beyond vocab_size before MLX reads the embedding table MLX does not bounds-check gathers, so an id >= vocab_size read neighbouring memory and crashed the process at random. run_mlx_stage (every batch, GRPO prompts) and LongRun now raise a ValueError naming the id range. test_mlx_training_six_stage_loss_paths crashed about one run in two on CPU, on clean dev as well: its GRPO stage tokenizes with the bundled 6400-token tokenizer while the model had vocab_size=64 (ids up to 6170). The test now uses vocab_size=6400 (25/25 runs pass). Co-Authored-By: Claude Sonnet 5.5 --- CHANGELOG.md | 2 ++ chronos/mlx/training/longrun.py | 3 +- chronos/mlx/training/trainer.py | 28 +++++++++++++++ tests/test_smoke.py | 2 +- tests/test_vocab_range_guard.py | 61 +++++++++++++++++++++++++++++++++ 5 files changed, 94 insertions(+), 2 deletions(-) create mode 100644 tests/test_vocab_range_guard.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c6c76ac..d7739cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,8 @@ Fixes two kernel bugs found while benchmarking the 50M and 30M multilingual conf - `cluster_locality_mode="segment"` could not be compiled. The loss found document ids with NumPy (`segment_ids`, `np.unique`) on a traced array, so `CHRONOS_MLX_COMPILE=on` raised "Attempting to eval an array during function transformations", and the default `auto` mode crashed on the next `mx.eval`. The loss now uses `segment_ids_mx(input_ids, chunk=window)`, which has a static shape (`row * S + id` addresses `B * S` segment slots and unused slots weigh 0), and runs inside the compiled step. It matches the NumPy ids on 300 random batches and the Torch loss within 1.7e-5 on 60. - A compile trace that failed part way left tracers in the model weights, the RNG key and the router outputs, because `mx.compile` swaps them back only on success. The eager fallback then trained on them and the next `mx.eval` raised "Attempting to eval an array without a primitive". `CompiledStep` now snapshots those on a first trace and restores them on failure; the RNG is re-seeded from the saved key words, so the random stream continues exactly (checked on 2,000 random keys). +- `run_mlx_stage` and LongRun now check every batch's token ids against the model's `vocab_size` and raise a `ValueError` that names the range. MLX does not bounds-check embedding gathers, so an id beyond the table read neighbouring memory and crashed the process at random. `tests/test_smoke.py::test_mlx_training_six_stage_loss_paths` hit this about one run in two on CPU: its GRPO stage tokenizes with the bundled 6,400-token tokenizer while the model had `vocab_size=64`. The test now uses 6400 (25 of 25 runs pass), and `tests/test_vocab_range_guard.py` covers the check. + ### Measured - Segment mode, compiled, real corpus, 10 steps, document masking on, steady tokens/s: 50M 24,856 at ctx 2048 and 20,416 at 4096; 30M 22,454 and 16,416. The eager fallback cost 15-17%. diff --git a/chronos/mlx/training/longrun.py b/chronos/mlx/training/longrun.py index e6603d2..7b9fad1 100644 --- a/chronos/mlx/training/longrun.py +++ b/chronos/mlx/training/longrun.py @@ -514,7 +514,7 @@ def _validate(self, source, phase, n_batches=8) -> float: def run_phase(self, pi: int, deadline_s: float, max_steps: int | None = None): """Train phase ``pi`` until its step total, its deadline, or ``max_steps`` more optimizer updates (used by tests and staged runs).""" - from chronos.mlx.training.trainer import _add_grads, _collect_router_probs, _tree_all_finite + from chronos.mlx.training.trainer import _add_grads, _check_ids_in_vocab, _collect_router_probs, _tree_all_finite phase = self.spec.phases[pi] source = self._source(phase) if self.state["rng"].get(phase.name): @@ -578,6 +578,7 @@ def run_phase(self, pi: int, deadline_s: float, max_steps: int | None = None): routes = RoutingCounts() if will_log or want_tensors else None for _ in range(accum): ids, labels = source.batch(phase.micro_tokens) + _check_ids_in_vocab(ids, int(self.model.config.vocab_size), "longrun batch") 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)) diff --git a/chronos/mlx/training/trainer.py b/chronos/mlx/training/trainer.py index bc9bff3..bbd51ef 100644 --- a/chronos/mlx/training/trainer.py +++ b/chronos/mlx/training/trainer.py @@ -47,6 +47,32 @@ def _ids_to_mx(x) -> mx.array: return mx.array(x).astype(mx.int32) +def _check_ids_in_vocab(ids, vocab_size: int, where: str) -> None: + """Raise instead of reading outside the embedding table. + + MLX does not bounds-check gathers, so an id >= vocab_size reads whatever memory follows the table + and crashes the process some of the time. Float arrays are not token ids and pass.""" + import numpy as np + arr = ids.detach().cpu().numpy() if isinstance(ids, torch.Tensor) else np.asarray(ids) + if arr.dtype.kind not in "iu" or arr.size == 0: + return + low, high = int(arr.min()), int(arr.max()) + if low < 0 or high >= int(vocab_size): + raise ValueError(f"{where}: token ids span [{low}, {high}] but the model's vocab_size is {int(vocab_size)}; " + "use the tokenizer the model was trained with") + + +def _check_batch_ids(batch, vocab_size: int, stage: str) -> None: + """Range-check the input ids of one stage batch (a tuple, a multimodal dict or a DPO pair).""" + if isinstance(batch, dict): + for key in ("input_ids", "x_chosen", "x_rejected"): + if key in batch: + _check_ids_in_vocab(batch[key], vocab_size, f"{stage} batch[{key!r}]") + elif isinstance(batch, (tuple, list)) and batch: + _check_ids_in_vocab(batch[0], vocab_size, f"{stage} batch ids") + + + def _f32(x: mx.array) -> mx.array: """FP32 view for loss math; a no-op when logits were born FP32.""" return x if x.dtype == mx.float32 else x.astype(mx.float32) @@ -926,6 +952,7 @@ def _build_grpo_batch(model, prompt, tokenizer, reward_fn, args) -> dict: temperature = float(getattr(args, "temperature", 1.0) if args is not None else 1.0) enc = tokenizer(rendered, return_tensors="pt", truncation=True, max_length=max_seq_len) prompt_ids = _ids_to_mx(enc.input_ids) + _check_ids_in_vocab(prompt_ids, int(model.config.vocab_size), "grpo prompt (tokenized with the tokenizer bundled in chronos.deps)") prompt_len = int(prompt_ids.shape[1]) trajectories = [] rewards = [] @@ -1386,6 +1413,7 @@ def evaluate_validation(epoch): for batch in data_iter: if _should_stop(): break + _check_batch_ids(batch, int(model.config.vocab_size), stage) if router_ref is None and stage in {"sft", "dpo", "orpo", "distill"}: router_ref = _capture_router_ref(model, batch, stage) train_batch = batch diff --git a/tests/test_smoke.py b/tests/test_smoke.py index 46c9aaf..b38970e 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -2233,7 +2233,7 @@ def test_mlx_training_six_stage_loss_paths(tmp_path): num_experts_per_tok=1, num_shared_experts=1, num_attention_heads=2, num_key_value_heads=2, kv_latent_dim=8, rope_dim=4, max_seq_len=16, - vocab_size=64, use_moe=True, + vocab_size=6400, use_moe=True, lambda_balance=5e-4, lambda_temporal=1e-3, lambda_lookahead=0.1, lambda_lookahead_topk=0.05) diff --git a/tests/test_vocab_range_guard.py b/tests/test_vocab_range_guard.py new file mode 100644 index 0000000..800b159 --- /dev/null +++ b/tests/test_vocab_range_guard.py @@ -0,0 +1,61 @@ +"""Out-of-range token ids are rejected before they reach an MLX embedding gather (0.3.2a3).""" +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") +torch = pytest.importorskip("torch") + + +def test_check_ids_in_vocab_accepts_valid_and_rejects_out_of_range(): + from chronos.mlx.training.trainer import _check_ids_in_vocab + _check_ids_in_vocab(np.array([[0, 5, 63]]), 64, "x") + _check_ids_in_vocab(torch.tensor([[0, 63]]), 64, "x") + _check_ids_in_vocab(mx.array([[0, 63]]), 64, "x") + _check_ids_in_vocab(np.zeros((0, 3), dtype=np.int64), 64, "empty") + _check_ids_in_vocab(np.array([[0.5]]), 64, "float ids are not token ids") + for bad in (np.array([[0, 64]]), torch.tensor([[-1, 3]]), mx.array([[6170]])): + with pytest.raises(ValueError, match="vocab_size is 64"): + _check_ids_in_vocab(bad, 64, "x") + + +def test_check_batch_ids_covers_tuples_dicts_and_dpo_pairs(): + from chronos.mlx.training.trainer import _check_batch_ids + ok = torch.tensor([[1, 2, 3]]) + _check_batch_ids((ok, ok), 10, "sft") + _check_batch_ids({"input_ids": ok, "labels": ok}, 10, "sft") + _check_batch_ids({"x_chosen": ok, "x_rejected": ok}, 10, "dpo") + with pytest.raises(ValueError, match="sft batch ids"): + _check_batch_ids((torch.tensor([[1, 10]]), ok), 10, "sft") + with pytest.raises(ValueError, match="x_rejected"): + _check_batch_ids({"x_chosen": ok, "x_rejected": torch.tensor([[10]])}, 10, "dpo") + + +def test_run_mlx_stage_refuses_a_batch_larger_than_the_vocab(tmp_path): + from torch.utils.data import DataLoader, TensorDataset + from chronos.mlx.training import run_mlx_stage + from chronos.model.config import ChronosConfig + cfg = ChronosConfig(hidden_size=32, num_hidden_layers=1, num_experts=2, num_experts_per_tok=1, + num_shared_experts=1, num_attention_heads=2, num_key_value_heads=2, kv_latent_dim=8, + rope_dim=4, max_seq_len=16, vocab_size=64, use_moe=True) + + class Args: + learning_rate = 1e-4 + weight_decay = 0.0 + steps = 1 + epochs = 1 + max_seq_len = 8 + + ids = torch.randint(0, 64, (2, 8), dtype=torch.long) + ids[0, 3] = 6170 # an id from a larger vocabulary + with pytest.raises(ValueError, match="vocab_size is 64"): + run_mlx_stage(stage="pretrain", config=cfg, checkpoint_path=None, save_dir=str(tmp_path), + loader=DataLoader(TensorDataset(ids, ids.clone()), batch_size=1), args=Args())