0.3.2a1: long-context MLX training

CHANGELOG.md

@@ -2,6 +2,35 @@
 
 ## [Unreleased]
 
+## [0.3.2a1] — 2026-11-09 (beta)
+
+This release lets the MLX trainer run long contexts on one Mac. Every new switch is off by default, so existing configs, checkpoints and recipes behave as before.
+
+### Added
+- `attention_block_size`: blockwise causal attention (`chronos/mlx/blockwise_attention.py`). Each block of queries reads only the keys it can see, under `mx.checkpoint`, and a block that ends at its last visible key uses the kernel's `"causal"` mode with no mask array. Value difference 0, gradient difference ≤ 8.5e-7 against the full-mask path.
+- `ce_chunk_size`: chunked cross-entropy from the final hidden state, so no `[B, S, V]` logits exist (`chronos/mlx/chunked_ce.py`). Loss difference 0, gradient relative difference ≤ 5e-6.
+- `moe_chunk_tokens`: MoE sparse dispatch in token chunks, each recomputed in the backward pass, including the cache-miss simulation.
+- `recompute_policy` (`all`, `mla`, `swa`, `alternate`) chooses which blocks `activation_checkpointing` recomputes, on both the MLX and the Torch side.
+- `rope_train_scaling` (`yarn`, `linear`, `ntk`): training-time RoPE scaling; MLX and Torch tables agree to 2e-5.
+- `document_masking` and `reset_router_at_documents`: attention and the temporal and locality losses stay inside each packed document. `reset_lookahead_at_documents` gives the lookahead terms their own switch (default: follow the router switch).
+- `PhaseSpec.seq_len_curriculum` (`[[progress, seq_len], ...]`) steps the pretrain window length up during a phase, and `PhaseSpec.pack_sft` joins short SFT records into full rows separated by the document separator (`SFTRecords.packed_batch`).
+
+### Fixed
+- The lookahead terms with several routing groups paired document ids of the wrong batch size when `reset_router_at_documents` was on.
+
+### Measured
+Full MLX training step (hidden 512, 8 layers, 16 experts top-2, bf16, batch 1, M5 Max). "long" = `attention_block_size=1024`, `ce_chunk_size=1024`, `moe_chunk_tokens=8192`; "+recompute" adds `activation_checkpointing`.
+
+| Context | Default: tokens/s, extra peak | long: tokens/s, extra peak | long + recompute: tokens/s, extra peak |
+|---|---|---|---|
+| 2k | 20.5k, 2.96 GiB | 23.0k, 2.16 GiB | 19.4k, 1.50 GiB |
+| 8k | 9.9k, 21.5 GiB | 23.3k, 5.55 GiB | 19.3k, 2.87 GiB |
+| 16k | 4.3k, 76.2 GiB | 18.1k, 6.76 GiB | 15.6k, 4.76 GiB |
+| 32k | not run (exceeds RAM) | 11.7k, 12.7 GiB | 9.7k, 8.3 GiB |
+| 64k | not run (exceeds RAM) | 6.6k, 24.5 GiB | 5.9k, 15.4 GiB |
+
+A single MLA layer at 64k needs 148.6 GiB extra without blocks and 2.5 GiB with them.
+
 ## [0.3.2a0] — 2026-11-08 (beta)
 
 This release replaces the minimind kernel with a native one, compiles the MLX training step, adds MLX activation recomputation, and keeps tool-call tags that a custom tokenizer registers as special tokens. README "Native kernel (0.3.2)" has the migration reasons and the measured comparison.

chronos/__init__.py

@@ -1,5 +1,5 @@
 # Project Chronos: On-device low-latency lookahead dual-layer MoE inference
-__version__ = "0.3.2a0"
+__version__ = "0.3.2a1"
 
 # Register with HuggingFace AutoXXX factories so
 # AutoModelForCausalLM.from_pretrained("path/to/chronos") works out of the box.

chronos/config_schema.py

@@ -45,7 +45,9 @@
     "vram_budget_gb", "pinned_memory_max_fraction", "prefetch_depth",
     "recommended_resident_experts", "offload_working_set_experts", "storage_format",
     "cluster_manifest_path", "training_paging", "optimizer_paging",
-    "optimizer_page_dir", "page_size_mb", "activation_checkpointing",
+    "optimizer_page_dir", "page_size_mb", "activation_checkpointing", "ce_chunk_size", "attention_block_size",
+    "reset_router_at_documents", "document_masking", "recompute_policy", "moe_chunk_tokens",
+    "reset_lookahead_at_documents", "rope_train_scaling",
     "modalities", "auto_detect_modalities", "image_size", "image_patch_size",
     "max_image_tokens", "audio_sample_rate", "max_audio_seconds", "max_audio_tokens",
     "modality_dropout", "freeze_modality_encoders", "image_encoder_type",

chronos/data/shards.py

@@ -130,6 +130,39 @@ def batch(self, token_budget: int, max_rows: int = 256, width_multiple: int = 64
                 labels[r, p:n] = seq[p:]
         return ids, labels
 
+    def packed_batch(self, token_budget: int, seq_len: int, sep_id: int = 0):
+        """Rows of exactly ``seq_len + 1`` tokens, each filled with whole records.
+
+        Records are taken in random order and joined by one ``sep_id`` token
+        (the document separator that ``segment_ids`` cuts on); a record that
+        does not fit is cut at the row end, a record longer than a row is
+        truncated.  Labels keep the loss mask of each record; separators and
+        the row's unused tail carry no loss.  Compared with ``batch`` there is no
+        length bucketing and almost no padding, so short SFT data wastes no compute."""
+        rows = max(1, int(token_budget) // int(seq_len))
+        width = int(seq_len) + 1
+        ids = np.full((rows, width), self.pad_id, dtype=np.int32)
+        labels = np.full((rows, width), -100, dtype=np.int32)
+        for r in range(rows):
+            pos = 0
+            while pos < width:
+                si, ri, n = self.index[int(self.rng.integers(0, len(self.index)))]
+                off = self.offsets[si][ri]
+                take = min(int(n), width - pos)
+                seq = np.asarray(self.tokens[si][off:off + take], dtype=np.int32)
+                ids[r, pos:pos + take] = seq
+                lm = self.lmasks[si]
+                if lm is not None:
+                    keep = np.asarray(lm[off:off + take]).astype(bool)
+                else:
+                    keep = np.arange(take) >= int(self.plens[si][ri])
+                labels[r, pos:pos + take] = np.where(keep, seq, -100)
+                pos += take
+                if pos < width:
+                    ids[r, pos] = sep_id   # closes the document; carries no loss
+                    pos += 1
+        return ids, labels
+
     def supervised_tokens(self) -> int:
         """Tokens that carry loss over every record (``.lmask`` or ``.plen`` rule)."""
         total = 0

chronos/mlx/attention.py

@@ -21,6 +21,7 @@
 from chronos.mlx.checkpointing import Dropout as TapeDropout   # replayed under recomputation
 
 from chronos.mlx.kv_quant import kv_read, kv_store
+from chronos.mlx.blockwise_attention import blockwise_attention, use_blockwise
 
 
 def _rope_freqs(
@@ -35,9 +36,16 @@ def _rope_freqs(
     whole last dimension. MLX must mirror that exactly for checkpoint parity.
     """
     half = max(dim // 2, 1)
+    kind = str((rope_scaling or {}).get("type", "yarn")).lower()
+    if rope_scaling is not None and kind == "ntk":
+        # NTK-aware: stretch the base so the lowest frequency spans the longer window.
+        base = base * float(rope_scaling.get("factor", 1.0)) ** (dim / max(dim - 2, 1))
     theta = 1.0 / (base ** (mx.arange(0, dim, 2, dtype=mx.float32)[:half] / dim))
     attn_factor = 1.0
-    if rope_scaling is not None:
+    if rope_scaling is not None and kind == "linear":
+        theta = theta / float(rope_scaling.get("factor", 1.0))
+        attn_factor = float(rope_scaling.get("attention_factor", 1.0))
+    elif rope_scaling is not None and kind == "yarn":
         orig_max = float(rope_scaling.get("original_max_position_embeddings", 2048))
         factor = float(rope_scaling.get("factor", 16))
         beta_fast = float(rope_scaling.get("beta_fast", 32.0))
@@ -122,6 +130,15 @@ def _merge_mask(allowed: mx.array | None, mask: mx.array | None, S: int, T: int)
     return m if allowed is None else (allowed & m)
 
 
+def _segment_mask(allowed, segments, S: int, T: int):
+    """AND a document mask (``segments`` ``[B, T]`` ids, queries = last S) into ``allowed``."""
+    if segments is None:
+        return allowed
+    seg = segments[:, -T:]
+    same = seg[:, None, -S:, None] == seg[:, None, None, :]
+    return same if allowed is None else (allowed & same)
+
+
 class _RopeCacheMixin:
     """FP32 RoPE tables, built once and never converted per call.
 
@@ -218,6 +235,7 @@ def __init__(self, config):
         self.dropout_p = float(config.dropout)
 
         self._init_rope_cache(config, self.head_dim, rotated_dim=self.rope_dim)
+        object.__setattr__(self, "_block_cfg", config)
         # Inference cache/decode policy (no parameters; see ChronosConfig).
         object.__setattr__(self, "mla_decode", str(getattr(config, "mla_decode", "expanded") or "expanded").lower())
         object.__setattr__(self, "cache_rope", bool(getattr(config, "mla_cache_rope", False)))
@@ -286,14 +304,20 @@ def __call__(
         q_rope = _apply_rope(q_rope, self._cos[past_len:past_len + S], self._sin[past_len:past_len + S])
         scale = 1.0 / math.sqrt(self.head_dim)
         allowed = _merge_mask(_causal_allowed(S, T_full), mask, S, T_full)
+        allowed = _segment_mask(allowed, getattr(self, "_segments", None), S, T_full)
         if self._use_absorbed(S, T_full):
             out = self._absorbed(q_nope, q_rope, c_kv_full, k_rope, scale, allowed)
         else:
             k_nope = self.k_nope_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.nope_dim)
             v = self.v_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.head_dim)
             q = mx.concatenate([q_nope, q_rope], axis=-1).transpose(0, 2, 1, 3)   # [B, nH, S, d]
             k = mx.concatenate([k_nope, k_rope], axis=-1).transpose(0, 2, 1, 3)   # [B, nKV, T, d]
-            out = self._attend(q, k, v.transpose(0, 2, 1, 3), scale, allowed)
+            blk = use_blockwise(self._block_cfg, S, T_full, self.training) if mask is None else 0
+            if blk:
+                out = blockwise_attention(q, k, v.transpose(0, 2, 1, 3), scale=scale, window=None, block=blk,
+                                     segments=getattr(self, "_segments", None), past=T_full - S)
+            else:
+                out = self._attend(q, k, v.transpose(0, 2, 1, 3), scale, allowed)
             out = out.transpose(0, 2, 1, 3)
         out = self.resid_dropout(self.o_proj(out.reshape(B, S, -1)))
         if return_shared_kv:
@@ -346,6 +370,7 @@ def __init__(self, config):
         self.dropout_p = float(config.dropout)
 
         self._init_rope_cache(config, self.head_dim)
+        object.__setattr__(self, "_block_cfg", config)
         object.__setattr__(self, "kv_int8", str(getattr(config, "kv_cache_dtype", "auto") or "auto").lower() == "int8")
         object.__setattr__(self, "kv_group", int(getattr(config, "kv_quant_group_size", 32) or 32))
         # Extra retained entries so self-speculative decoding can roll back up
@@ -394,10 +419,17 @@ def __call__(
         allowed = k_positions <= q_positions
         if self.window_size is not None:              # None: dense (use_hybrid_attention=False)
             allowed = allowed & (k_positions > q_positions - self.window_size)
+        allowed = _segment_mask(allowed, getattr(self, "_segments", None), S, T_full)
 
         scale = 1.0 / math.sqrt(self.head_dim)
-        out = self._attend(q.transpose(0, 2, 1, 3), k.transpose(0, 2, 1, 3), v.transpose(0, 2, 1, 3),
-                           scale, _merge_mask(allowed, mask, S, T_full))
+        blk = use_blockwise(self._block_cfg, S, T_full, self.training) if mask is None else 0
+        if blk:
+            out = blockwise_attention(q.transpose(0, 2, 1, 3), k.transpose(0, 2, 1, 3), v.transpose(0, 2, 1, 3),
+                                     scale=scale, window=self.window_size, block=blk,
+                                     segments=getattr(self, "_segments", None), past=cached_len)
+        else:
+            out = self._attend(q.transpose(0, 2, 1, 3), k.transpose(0, 2, 1, 3), v.transpose(0, 2, 1, 3),
+                               scale, _merge_mask(allowed, mask, S, T_full))
         out = out.transpose(0, 2, 1, 3).reshape(B, S, -1)
 
         # Cache stores [B, T, nKV, d] (pre-repeat, pre-transpose)

chronos/mlx/blockwise_attention.py

@@ -0,0 +1,94 @@
+"""Blockwise causal attention for long sequences (MLX): no ``[S, T]`` mask.
+
+The fused SDPA kernel with an explicit boolean mask needs the ``[S, T]``
+mask and, in its VJP, ``[nH, S, T]`` scores, so memory grows as S**2
+(measured: one SWA layer at S=32768 peaked at 40 GiB extra).
+
+``blockwise_attention`` splits the queries into blocks of ``block`` rows.
+Block b only reads the keys it can see: ``[b*block - window + 1, (b+1)*block)``
+for a sliding window, ``[0, (b+1)*block)`` for full causal attention.  Each
+block runs the fused kernel with a small ``[block, span]`` mask, under
+``mx.checkpoint`` so its scores are recomputed in the backward pass.  A
+sliding-window layer then costs O(S * (window + block)), and the full
+causal layer stores O(block * S) at a time instead of O(S**2).
+
+``segments`` ``[B, S]`` (document ids, non-decreasing along a row) adds
+document masking: a query attends only to keys of its own document.
+
+Exactness: every query sees exactly the keys of the reference mask, and
+softmax runs over that set inside the fused kernel, so the output equals
+the full-mask path up to floating-point summation order.
+"""
+import mlx.core as mx
+
+
+def _block_mask(q0, q1, k0, k1, window, seg_q, seg_k):
+    """Boolean ``[(B,) 1, q1-q0, k1-k0]`` mask for absolute query rows q0..q1-1."""
+    q_pos = mx.arange(q0, q1)[:, None]
+    k_pos = mx.arange(k0, k1)[None, :]
+    allowed = k_pos <= q_pos
+    if window is not None:
+        allowed = allowed & (k_pos > q_pos - int(window))
+    if seg_q is None:
+        return allowed
+    same = seg_q[:, None, :, None] == seg_k[:, None, None, :]
+    return allowed[None, None] & same
+
+
+def _mask_for(q0, q1, k0, k1, window, seg_q, seg_k):
+    """Array mask for one block, or ``"causal"`` for a square diagonal span
+    where the plain causal rule applies (no mask array in the kernel)."""
+    # MLX aligns "causal" to the bottom-right, so a block whose last query is
+    # the last visible key needs no mask array whatever its key span.
+    if seg_q is None and k1 == q1 and (window is None or window >= q1 - k0):
+        return "causal"
+    return mx.stop_gradient(_block_mask(q0, q1, k0, k1, window, seg_q, seg_k))
+
+
+def _attend_block(q, k, v, mask, scale):
+    return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask)
+
+
+_attend_block_ckpt = mx.checkpoint(_attend_block)
+
+
+def blockwise_attention(q, k, v, *, scale, window=None, block=1024, segments=None, past=0, checkpoint=True):
+    """Causal (optionally windowed / document-masked) attention in query blocks.
+
+    ``q [B, nH, S, d]``, ``k``/``v`` ``[B, nKV, T, d]`` with ``T = past + S``
+    (``past`` cached keys precede the queries).  Returns ``[B, nH, S, dv]``.
+    """
+    S, T = q.shape[2], k.shape[2]
+    block = max(1, int(block))
+    run = _attend_block_ckpt if checkpoint else _attend_block
+    seg = None if segments is None else mx.stop_gradient(segments)
+    outs = []
+    prev = None
+    for start in range(0, S, block):
+        end = min(S, start + block)
+        qa0, qa1 = past + start, past + end
+        k0 = 0 if window is None else max(0, qa0 - int(window) + 1)
+        k1 = min(T, qa1)
+        seg_q = None if seg is None else seg[:, qa0:qa1]
+        seg_k = None if seg is None else seg[:, k0:k1]
+        mask = _mask_for(qa0, qa1, k0, k1, window, seg_q, seg_k)
+        qb = q[:, :, start:end]
+        if prev is not None:
+            qb = mx.depends(qb, prev)   # one block's scores live at a time
+        prev = run(qb, k[:, :, k0:k1], v[:, :, k0:k1], mask, scale)
+        outs.append(prev)
+    return outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=2)
+
+
+def use_blockwise(config, S, T, training):
+    """Rows per attention block from ``config.attention_block_size`` (0 = off).
+
+    Only for multi-row queries without attention dropout, and only when the
+    sequence is longer than one block (short sequences keep the exact
+    single-kernel path)."""
+    size = int(getattr(config, "attention_block_size", 0) or 0)
+    if size <= 0 or S <= size:
+        return 0
+    if training and float(getattr(config, "dropout", 0.0) or 0.0) > 0.0:
+        return 0
+    return size

chronos/mlx/chunked_ce.py

@@ -0,0 +1,64 @@
+"""Chunked cross-entropy for MLX: the LM-head loss without ``[B, S, V]`` logits.
+
+For long contexts the FP32 logits dominate peak memory (S=65536, V=16384:
+4 GiB per row).  ``chunked_lm_loss`` walks the flattened hidden states in
+``chunk`` rows.  Each chunk computes ``hidden @ W.T`` in FP32, its
+logsumexp and the label logit, inside ``mx.checkpoint``, so the backward
+pass recomputes that chunk's logits instead of keeping them.  Peak logit
+memory is one chunk (``chunk * V * 4`` bytes).
+
+The value equals the mean cross-entropy over labels that are not
+``ignore_index``, in FP32, as the full path with ``logits_dtype=mx.float32``.
+"""
+import mlx.core as mx
+
+DEFAULT_CHUNK = 4096
+
+
+def _chunk_sum(h, w, lab, keep):
+    """Summed CE of one chunk: rows ``h [n, H]``, head ``w [V, H]``, FP32."""
+    logits = h.astype(mx.float32) @ w.astype(mx.float32).T
+    lse = mx.logsumexp(logits, axis=-1)
+    picked = mx.take_along_axis(logits, mx.stop_gradient(lab)[:, None], axis=-1)[:, 0]
+    return ((lse - picked) * mx.stop_gradient(keep)).sum()
+
+
+_chunk_sum_ckpt = mx.checkpoint(_chunk_sum)
+
+
+def chunked_ce_sum(hidden, weight, labels, ignore_index=-100, chunk=DEFAULT_CHUNK):
+    """``(sum of CE, number of supervised labels)`` over ``hidden [..., H]`` / ``labels [...]``."""
+    H = hidden.shape[-1]
+    h = hidden.reshape(-1, H)
+    lb = labels.reshape(-1)
+    keep = (lb != ignore_index).astype(mx.float32)
+    safe = mx.where(lb != ignore_index, lb, 0).astype(mx.int32)
+    n = h.shape[0]
+    step = max(1, int(chunk))
+    total = mx.zeros((), dtype=mx.float32)
+    for start in range(0, n, step):
+        end = min(n, start + step)
+        rows = h[start:end]
+        if start:
+            # Serialize the chunks: chunk k+1 starts after chunk k's sum exists,
+            # so the scheduler never holds several chunks' logits at once.
+            rows = mx.depends(rows, total)
+        total = total + _chunk_sum_ckpt(rows, weight, safe[start:end], keep[start:end])
+    return total, keep.sum()
+
+
+def chunked_lm_loss(hidden, weight, labels, ignore_index=-100, chunk=DEFAULT_CHUNK):
+    """Mean CE of labels aligned with ``hidden`` (no shift), chunked."""
+    total, count = chunked_ce_sum(hidden, weight, labels, ignore_index, chunk)
+    return total / mx.maximum(count, 1.0)
+
+
+def chunked_shifted_lm_loss(hidden, weight, labels, ignore_index=-100, chunk=DEFAULT_CHUNK):
+    """Mean CE where position t predicts ``labels[:, t + 1]`` (the causal shift)."""
+    return chunked_lm_loss(hidden[:, :-1, :], weight, labels[:, 1:], ignore_index, chunk)
+
+
+def ce_chunk_for(config, seq_rows: int) -> int:
+    """Rows per CE chunk from ``config.ce_chunk_size``; 0 = no chunking (full logits)."""
+    size = int(getattr(config, "ce_chunk_size", 0) or 0)
+    return 0 if size <= 0 or size >= int(seq_rows) else size