0.3.2a3: compile segment-mode cluster locality, recover from a failed compile trace

CHANGELOG.md

@@ -2,6 +2,22 @@
 
 ## [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).
+
+- `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%.
+
+### 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.

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.

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):

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)

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))

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)
@@ -251,26 +277,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 +309,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)
 
 
@@ -922,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 = []
@@ -1382,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

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"