diff --git a/CHANGELOG.md b/CHANGELOG.md index f4602e8..1ffc7e4 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/chronos/__init__.py b/chronos/__init__.py index 22af209..914085a 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.2a0" +__version__ = "0.3.2a1" # Register with HuggingFace AutoXXX factories so # AutoModelForCausalLM.from_pretrained("path/to/chronos") works out of the box. diff --git a/chronos/config_schema.py b/chronos/config_schema.py index c4301bb..a6fd961 100644 --- a/chronos/config_schema.py +++ b/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", diff --git a/chronos/data/shards.py b/chronos/data/shards.py index 34a9ae1..2de910f 100644 --- a/chronos/data/shards.py +++ b/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 diff --git a/chronos/mlx/attention.py b/chronos/mlx/attention.py index 7d7b281..0bc12a2 100644 --- a/chronos/mlx/attention.py +++ b/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,6 +304,7 @@ 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: @@ -293,7 +312,12 @@ def __call__( 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) diff --git a/chronos/mlx/blockwise_attention.py b/chronos/mlx/blockwise_attention.py new file mode 100644 index 0000000..5b2b4f2 --- /dev/null +++ b/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 diff --git a/chronos/mlx/chunked_ce.py b/chronos/mlx/chunked_ce.py new file mode 100644 index 0000000..5c60e37 --- /dev/null +++ b/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 diff --git a/chronos/mlx/model.py b/chronos/mlx/model.py index 4300518..1c03fbb 100644 --- a/chronos/mlx/model.py +++ b/chronos/mlx/model.py @@ -68,6 +68,18 @@ def __call__(self, x: mx.array) -> mx.array: +def _recompute_block(cfg, index: int, block) -> bool: + """Whether physical block ``index`` is an ``mx.checkpoint`` region under + ``config.recompute_policy`` ("all" | "mla" | "swa" | "alternate").""" + policy = getattr(cfg, "recompute_policy", "all") + if policy == "all": + return True + if policy == "alternate": + return index % 2 == 0 + is_mla = isinstance(block.self_attn, MLAAttentionMLX) + return is_mla if policy == "mla" else not is_mla + + class ChronosMLXBlock(nn.Module): def __init__(self, layer_id: int, config): super().__init__() @@ -160,24 +172,38 @@ def __call__( 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 + return_hidden: bool = False, # skip the LM head (chunked losses) + segments=None, # [B, S] document ids: attention stays in a document ): - """``logits_dtype=mx.float32`` makes the logits FP32 at birth: the + """``segments`` ``[B, S]`` (non-decreasing document ids of a packed row, + as ``chronos.expert_clusters.segment_ids`` gives) keeps every attention + layer inside the query's own document (document masking for varlen + packing). Rows longer than one block take the blockwise path + (``attention_block_size``); shorter ones get an explicit mask. + + ``return_hidden=True`` returns ``(hidden [B,S,H], lookahead, cache)`` + without the LM head, so a loss never materializes [B,S,V] logits + (``chronos.mlx.chunked_ce``). + + ``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 the logits) instead of rounding logits to the compute dtype and casting the large [B,S,V] tensor afterwards.""" + self._set_segments(segments) moes = [layer.mlp for layer in self.layers if isinstance(layer.mlp, ChronosMLXMOE)] for moe in moes: moe.reset_iteration_routing() moe.open_pack_scope() try: return self._forward(input_ids, cache, available_expert_masks, logits_dtype, recurrent_depth, - modality_inputs, modality_mask) + modality_inputs, modality_mask, return_hidden) finally: for moe in moes: moe.close_pack_scope() + self._set_segments(None) def _forward(self, input_ids, cache, available_expert_masks, logits_dtype, depth_override=None, - modality_inputs=None, modality_mask=None): + modality_inputs=None, modality_mask=None, return_hidden=False): B, S = input_ids.shape x = self.dropout(self.embed_tokens(input_ids)) # [B, S, H] object.__setattr__(self, "last_modality_states", {}) @@ -206,7 +232,7 @@ def _forward(self, input_ids, cache, available_expert_masks, logits_dtype, depth def run(i, hidden, slot=None, **kw): avail = available_expert_masks[i] if available_expert_masks is not None else None c = cache[i if slot is None else slot] if "kv_override" not in kw else None - if recompute and c is None: + if recompute and c is None and _recompute_block(cfg, i, self.layers[i]): return call_block(self.layers[i], hidden, mask=mask, cache=None, available_expert_mask=avail, **kw) return self.layers[i](hidden, mask=mask, cache=c, available_expert_mask=avail, **kw) @@ -298,13 +324,25 @@ def run(i, hidden, slot=None, **kw): lookahead_probs = self.lookahead_router(x) x = self.norm(x) - head = self.embed_tokens.weight if self.lm_head is None else self.lm_head.weight + if return_hidden: + return x, lookahead_probs, new_caches + head = self.head_weight() if logits_dtype is not None and logits_dtype != head.dtype: logits = x.astype(logits_dtype) @ head.astype(logits_dtype).T else: logits = self.embed_tokens.as_linear(x) if self.lm_head is None else self.lm_head(x) return logits, lookahead_probs, new_caches + def _set_segments(self, segments) -> None: + """Give every attention layer this forward's document ids (or clear them).""" + seg = None if segments is None else mx.stop_gradient(mx.array(segments).astype(mx.int32)) + for layer in self.layers: + object.__setattr__(layer.self_attn, "_segments", seg) + + def head_weight(self) -> mx.array: + """The LM head matrix ``[V, H]`` (the embedding when tied).""" + return self.embed_tokens.weight if self.lm_head is None else self.lm_head.weight + @staticmethod def from_chronos_pytorch(pt_model, config, include_experts: bool = True, dtype: str | None = None): """ diff --git a/chronos/mlx/moe.py b/chronos/mlx/moe.py index a1d641c..8e99044 100644 --- a/chronos/mlx/moe.py +++ b/chronos/mlx/moe.py @@ -79,6 +79,7 @@ def __init__(self, config): self.fallback_mask_prob = float(getattr(config, "fallback_mask_prob", 0.0) or 0.0) # Sorted gather_mm dispatch (requires mlx >= 0.32.4 on M5-class GPUs). self.sparse_dispatch = bool(getattr(config, "moe_sparse_dispatch", True)) and gather_mm_supported() + self.chunk_tokens = int(getattr(config, "moe_chunk_tokens", 0) or 0) self.runtime_on_demand_loader = None self.runtime_touch_expert = None @@ -126,6 +127,39 @@ def _expert_stacks(self): def last_router_probs(self, value): object.__setattr__(self, "_last_router_probs", value) + def _grouped_chunked(self, x_flat, topk_idx, topk_weight, fallback): + """Sparse dispatch over the routed experts, in token chunks when + ``config.moe_chunk_tokens`` is set and the call is a training forward. + + Each chunk runs under ``mx.checkpoint``: its ``[chunk*K, I]`` gate/up + activations are recomputed in the backward pass instead of kept for the + whole row. Routing, weights and the fallback mask are inputs, so every + token sees the same experts and the result equals one dispatch (up to + summation order inside gather_mm).""" + stacks = self._expert_stacks() + n = x_flat.shape[0] + chunk = self.chunk_tokens + if not self.training or chunk <= 0 or n <= chunk: + return grouped_swiglu(x_flat, topk_idx, topk_weight, *stacks, fallback=fallback) + avail, shared = (None, None) if fallback is None else fallback + outs = [] + prev = None + for start in range(0, n, chunk): + end = min(n, start + chunk) + xs = x_flat[start:end] + if prev is not None: + xs = mx.depends(xs, prev) # one chunk's activations live at a time + if fallback is None: + def part(xc, ic, wc, g, u, d): + return grouped_swiglu(xc, ic, wc, g, u, d) + prev = mx.checkpoint(part)(xs, topk_idx[start:end], topk_weight[start:end], *stacks) + else: + def part(xc, ic, wc, sc, g, u, d): + return grouped_swiglu(xc, ic, wc, g, u, d, fallback=(avail, sc)) + prev = mx.checkpoint(part)(xs, topk_idx[start:end], topk_weight[start:end], shared[start:end], *stacks) + outs.append(prev) + return mx.concatenate(outs, axis=0) + def _shared_out(self, x: mx.array) -> mx.array: return sum(e(x) for e in self.shared_experts) / len(self.shared_experts) @@ -237,8 +271,7 @@ def __call__( # Cache-miss simulation (PyTorch parity): an unavailable routed # expert is served by the shared expert, as at offload inference. fallback = (draw(lambda: training_availability(self.num_experts, self.fallback_mask_prob)), shared) - y = grouped_swiglu(x_flat, topk_idx, topk_weight, *self._expert_stacks(), - fallback=fallback) + y = self._grouped_chunked(x_flat, topk_idx, topk_weight, fallback) if shared is not None: y = y + shared elif available_expert_mask is None: diff --git a/chronos/mlx/training/longrun.py b/chronos/mlx/training/longrun.py index b4730ea..e6603d2 100644 --- a/chronos/mlx/training/longrun.py +++ b/chronos/mlx/training/longrun.py @@ -47,6 +47,12 @@ class PhaseSpec: min_lr_ratio: float = 0.1 weight_decay: float = 0.1 grad_clip: float = 1.0 + # Length curriculum: ``[[progress, seq_len], ...]`` (progress = elapsed fraction + # of the phase, ascending). The window length steps up as the phase advances; + # micro_tokens stays fixed, so rows shrink as rows get longer. None = seq_len. + seq_len_curriculum: list | None = None + # SFT: join short records into full rows (``seq_len`` tokens) with separators. + pack_sft: bool = False @dataclass @@ -231,7 +237,23 @@ def recurrent_metrics(state: dict | None) -> dict: return out +def curriculum_seq_len(phase, progress: float) -> int: + """Window length for ``progress`` in [0, 1] under ``phase.seq_len_curriculum``.""" + stages = getattr(phase, "seq_len_curriculum", None) + if not stages: + return int(phase.seq_len) + length = int(stages[0][1]) + for start, n in stages: + if progress >= float(start): + length = int(n) + return length + + class PretrainSource: + def set_seq_len(self, n: int) -> None: + self.stream.seq_len = int(n) + self.valid.seq_len = int(n) + def __init__(self, shards_dir, seq_len, seed): from chronos.data.shards import PackedTokenStream self.stream = PackedTokenStream(os.path.join(shards_dir, "pretrain", "train"), seq_len, seed=seed) @@ -258,7 +280,8 @@ def set_rng_state(self, state): class SFTSource: - def __init__(self, shards_dir, seed): + def __init__(self, shards_dir, seed, pack_len: int = 0): + self.pack_len = int(pack_len or 0) from chronos.data.shards import SFTRecords self.records = SFTRecords(os.path.join(shards_dir, "sft", "train"), seed=seed) self.valid = SFTRecords(os.path.join(shards_dir, "sft", "valid"), seed=seed + 1) @@ -268,12 +291,16 @@ def _shift(ids, labels): return ids[:, :-1], labels[:, 1:] # labels already -100 outside replies def batch(self, micro_tokens: int): + if self.pack_len: + return self._shift(*self.records.packed_batch(micro_tokens, self.pack_len)) return self._shift(*self.records.batch(micro_tokens)) def valid_batches(self, n: int, micro_tokens: int): state = self.valid.rng.bit_generator.state self.valid.rng = np.random.default_rng(0) - out = [self._shift(*self.valid.batch(micro_tokens)) for _ in range(n)] + draw = (lambda: self.valid.packed_batch(micro_tokens, self.pack_len)) if self.pack_len else \ + (lambda: self.valid.batch(micro_tokens)) + out = [self._shift(*draw()) for _ in range(n)] self.valid.rng.bit_generator.state = state return out @@ -441,14 +468,9 @@ def _tensor_telemetry(self, phase, step: int, grads, before, accum: int, lr: flo 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 + from chronos.mlx.training.trainer import _chronos_regularized_loss, draft_exit_kl, draft_exit_logits, lm_ce draft = draft_exit_logits(model, ids, model.config) # None unless draft_exit_kl_coef > 0 - logits, lookahead, _ = model(ids, logits_dtype=mx.float32) - lg = logits.reshape(-1, logits.shape[-1]) - lb = labels.reshape(-1) - mask = lb != -100 - ce = nn.losses.cross_entropy(lg, mx.where(mask, lb, 0), reduction="none") - ce = (ce * mask).sum() / mx.maximum(mask.sum(), 1) + ce, lookahead, logits = lm_ce(model, ids, labels, aligned=True, need_logits=draft is not None) # 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}, @@ -456,8 +478,8 @@ def _loss_fn(self, model, ids, labels, token_mask=None): def _source(self, phase: PhaseSpec): if phase.name == "pretrain": - return PretrainSource(self.spec.shards_dir, phase.seq_len, self.spec.seed) - return SFTSource(self.spec.shards_dir, self.spec.seed) + return PretrainSource(self.spec.shards_dir, max([int(phase.seq_len)] + [int(n) for _, n in (phase.seq_len_curriculum or [])]), self.spec.seed) + return SFTSource(self.spec.shards_dir, self.spec.seed, pack_len=phase.seq_len if phase.pack_sft else 0) def _micro_step(self, grad_fn, ids, labels): mask = token_mask_for(labels) @@ -535,6 +557,8 @@ def run_phase(self, pi: int, deadline_s: float, max_steps: int | None = None): progress = (self.state["phase_wall_s"] / self.state["phase_budget_s"]) if time_based else \ step / max(total or 1_000_000, 1) controller.set_training_progress(int(progress * 1_000_000), 1_000_000) + if phase.name == "pretrain" and phase.seq_len_curriculum: + source.set_seq_len(curriculum_seq_len(phase, progress)) # MLX's global RNG (cache-miss simulation, dropout) cannot be # snapshotted, so make each step's randomness a pure function of # (seed, phase, step, rollbacks): resume is then bit-exact. diff --git a/chronos/mlx/training/trainer.py b/chronos/mlx/training/trainer.py index 0e0631d..0fb76bd 100644 --- a/chronos/mlx/training/trainer.py +++ b/chronos/mlx/training/trainer.py @@ -63,6 +63,15 @@ def _masked_lm_loss(logits: mx.array, labels: mx.array) -> mx.array: return (losses * mask.astype(losses.dtype)).sum() / mx.maximum(mask.astype(mx.float32).sum(), 1.0) +def _masked_lm_loss_aligned(logits: mx.array, labels: mx.array) -> mx.array: + """Mean CE of labels already aligned with the logits (no shift); -100 = ignored.""" + flat_logits = _f32(logits.reshape(-1, logits.shape[-1])) + flat_labels = labels.reshape(-1) + mask = flat_labels != -100 + losses = nn.losses.cross_entropy(flat_logits, mx.where(mask, flat_labels, 0), reduction="none") + return (losses * mask).sum() / mx.maximum(mask.sum(), 1) + + def _gather_token_logprobs(logits: mx.array, labels: mx.array) -> mx.array: labels = labels.astype(mx.int32) logits = _f32(logits) @@ -119,12 +128,31 @@ def _masked_mean(values: mx.array, mask: mx.array | None) -> mx.array: 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: +_PAIR_SEGMENTS: list = [None] # [B, S] document ids of the batch being regularized +_LOOKAHEAD_SEGMENTS: list = [None] # document ids for the lookahead terms (own switch) + + +def _pair_mask(token_mask: mx.array | None, k: int = 1, groups: int = 1) -> mx.array | None: + """``[B, S-k]``: positions t and t+k are both real tokens of one document. + + Documents come from ``_PAIR_SEGMENTS`` (set by ``_chronos_regularized_loss`` + when ``config.reset_router_at_documents``): temporal, locality and + lookahead terms then never pair tokens across a packed-row boundary. + ``groups`` > 1: the caller flattened that many routing groups into the + batch axis (``B*G`` rows), so the document ids are repeated to match.""" + seg = _PAIR_SEGMENTS[0] + if seg is not None and groups > 1: + seg = mx.repeat(seg, groups, axis=0) + if token_mask is None and seg is None: return None - m = token_mask.astype(mx.float32) - return m[:, :-k] * m[:, k:] + m = None + if token_mask is not None: + t = token_mask.astype(mx.float32) + m = t[:, :-k] * t[:, k:] + if seg is not None and seg.shape[1] > k: + same = (seg[:, :-k] == seg[:, k:]).astype(mx.float32) + m = same if m is None else m * same + return m def _group_mask(token_mask: mx.array | None, groups: int) -> mx.array | None: @@ -325,7 +353,7 @@ def _lookahead_supervision_loss( teacher_k = teacher[:, k:, :] pred_k = lookahead_probs[:, :-k, k, :].astype(mx.float32) total = total + _masked_mean(-(teacher_k * mx.log(mx.maximum(pred_k, 1e-9))).sum(axis=-1), - _pair_mask(token_mask, k)) + _pair_mask(token_mask, k, groups)) terms += 1 return total / max(terms, 1) @@ -360,7 +388,8 @@ 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 + _masked_mean(-mx.log(mx.maximum(hit_mass / float(top_k), 1e-9)), _pair_mask(token_mask, k)) + total = total + _masked_mean(-mx.log(mx.maximum(hit_mass / float(top_k), 1e-9)), + _pair_mask(token_mask, k, groups)) terms += 1 return total / max(terms, 1) @@ -394,7 +423,7 @@ def _lookahead_union_loss( mass = (pred_union * union_target).sum(axis=-1) target_count = mx.maximum(union_target.sum(axis=-1), 1.0) # 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)) + return _masked_mean(-mx.log(mx.maximum(mass / target_count, 1e-9)), _pair_mask(token_mask, K, groups)) def _offload_training_metrics( @@ -537,6 +566,27 @@ def _chronos_regularized_loss(model, base_loss, lookahead_probs, config, router_ 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).""" + segments = look_segments = None + reset_router = bool(getattr(config, "reset_router_at_documents", False)) + reset_look = getattr(config, "reset_lookahead_at_documents", None) + reset_look = reset_router if reset_look is None else bool(reset_look) + if input_ids is not None and (reset_router or reset_look): + from chronos.expert_clusters import segment_ids + docs = mx.array(segment_ids(input_ids)) + segments = docs if reset_router else None + look_segments = docs if reset_look else None + _PAIR_SEGMENTS[0] = segments + _LOOKAHEAD_SEGMENTS[0] = look_segments + try: + return _regularized_terms(model, base_loss, lookahead_probs, config, router_ref, metrics, + draft_kl, token_mask, input_ids) + finally: + _PAIR_SEGMENTS[0] = None + _LOOKAHEAD_SEGMENTS[0] = None + + +def _regularized_terms(model, base_loss, lookahead_probs, config, router_ref, metrics, draft_kl, + token_mask, input_ids): object.__setattr__(model, "_last_lookahead_probs", lookahead_probs) terms = {"primary_loss": base_loss.astype(mx.float32)} def add(name, coefficient, value): @@ -563,15 +613,20 @@ def add(name, coefficient, value): mean = router.mean(axis=2) 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, token_mask)) - add("lookahead_topk_loss", config.lambda_lookahead_topk, - _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, token_mask)) + saved_pairs = _PAIR_SEGMENTS[0] + _PAIR_SEGMENTS[0] = _LOOKAHEAD_SEGMENTS[0] # lookahead keeps its own document rule + try: + 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, token_mask)) + add("lookahead_topk_loss", config.lambda_lookahead_topk, + _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, token_mask)) + finally: + _PAIR_SEGMENTS[0] = saved_pairs 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) @@ -581,11 +636,37 @@ def add(name, coefficient, value): return loss +def lm_ce(model, input_ids, labels, *, aligned: bool = False, need_logits: bool = False, **forward_kw): + """Forward + LM cross-entropy: ``(ce, lookahead, logits_or_None)``. + + ``config.ce_chunk_size`` > 0 (and no caller needing the full logits) + computes the loss in row chunks from the final hidden state + (``chronos.mlx.chunked_ce``): no ``[B, S, V]`` tensor exists. The value + equals the full FP32 path. ``aligned``: labels already sit on the + logits (no causal shift). ``config.document_masking`` keeps attention inside + each packed document (``segments`` from the separator tokens).""" + from chronos.mlx.chunked_ce import ce_chunk_for, chunked_lm_loss, chunked_shifted_lm_loss + chunk = 0 if need_logits else ce_chunk_for(model.config, int(input_ids.shape[0]) * int(input_ids.shape[1])) + if getattr(model.config, "document_masking", False) and "segments" not in forward_kw: + from chronos.expert_clusters import segment_ids + # Varlen packing: queries attend only inside their own document; the + # separator token closes the document it ends (segment_ids rule). + forward_kw["segments"] = mx.array(segment_ids(input_ids)) + if chunk: + hidden, lookahead, _cache = model(input_ids, return_hidden=True, **forward_kw) + loss_fn = chunked_lm_loss if aligned else chunked_shifted_lm_loss + return loss_fn(hidden, model.head_weight(), labels, chunk=chunk), lookahead, None + logits, lookahead, _cache = model(input_ids, logits_dtype=mx.float32, **forward_kw) + if aligned: + return _masked_lm_loss_aligned(logits, labels), lookahead, logits + return _masked_lm_loss(logits, labels), lookahead, logits + + 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) + base, lookahead, logits = lm_ce(model, input_ids, labels, + need_logits=draft is not None, **(media or {})) 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, token_mask=real_token_mask(labels), input_ids=input_ids) diff --git a/chronos/model/config.py b/chronos/model/config.py index 8968338..d7f88f9 100644 --- a/chronos/model/config.py +++ b/chronos/model/config.py @@ -91,6 +91,39 @@ def __init__(self, **kwargs): self.optimizer_page_dir: str = kwargs.pop("optimizer_page_dir", "./optimizer_pages") self.page_size_mb: int = kwargs.pop("page_size_mb", 64) self.activation_checkpointing: bool = kwargs.pop("activation_checkpointing", False) + # Long-context training (0.3.2a1). ce_chunk_size > 0 computes the LM + # loss in row chunks from the final hidden state, never building the + # [B, S, V] logits (0 = full logits, the historical path). + self.ce_chunk_size: int = int(kwargs.pop("ce_chunk_size", 0) or 0) + # attention_block_size > 0 runs MLX attention over long multi-row + # queries in blocks that read only their visible keys (no [S, T] mask; + # SWA costs O(S * window)). Exact; 0 = one fused call (historical). + self.attention_block_size: int = int(kwargs.pop("attention_block_size", 0) or 0) + # Packed rows: when True, the temporal / router-locality / lookahead + # regularizers drop token pairs that straddle a document separator, + # so no routing penalty links two unrelated documents. + self.reset_router_at_documents: bool = bool(kwargs.pop("reset_router_at_documents", False)) + # Packed rows: attention stays inside each document (separator-delimited). + self.document_masking: bool = bool(kwargs.pop("document_masking", False)) + # Which blocks are recomputed when activation_checkpointing is on: + # "all", "mla" (global-attention blocks), "swa", "alternate" (every second + # block). Selective policies trade some memory for less recompute time. + self.recompute_policy: str = str(kwargs.pop("recompute_policy", "all") or "all").lower() + if self.recompute_policy not in ("all", "mla", "swa", "alternate"): + raise ValueError(f"unknown recompute_policy {self.recompute_policy!r}") + # MoE sparse dispatch in chunks of this many tokens, each recomputed in the + # backward pass, so [N*K, I] expert activations never exist for a whole row + # (0 = one dispatch, the historical path). + self.moe_chunk_tokens: int = int(kwargs.pop("moe_chunk_tokens", 0) or 0) + # Lookahead supervision's own document rule: None follows + # reset_router_at_documents; True/False decouples it from the temporal and + # router-locality terms. + _lrd = kwargs.pop("reset_lookahead_at_documents", None) + self.reset_lookahead_at_documents = None if _lrd is None else bool(_lrd) + # Long-context RoPE for training: {"type": "yarn"|"linear"|"ntk", "factor": f, + # "original_max_position_embeddings": n, ...}; applied after the base init. + _rts = kwargs.pop("rope_train_scaling", None) + self.rope_train_scaling = dict(_rts) if _rts else None # Multimodal encoders are native kernel inputs, but are instantiated # only for modalities discovered/enabled by the data pipeline. @@ -235,6 +268,8 @@ def __init__(self, **kwargs): # Force MoE on kwargs.setdefault("use_moe", True) super().__init__(**kwargs) + if self.rope_train_scaling is not None: + self.rope_scaling = dict(self.rope_train_scaling) from chronos.config_schema import normalize_model_dimensions resolved = normalize_model_dimensions(vars(self)) diff --git a/chronos/model/kernel.py b/chronos/model/kernel.py index bb980de..e94c194 100644 --- a/chronos/model/kernel.py +++ b/chronos/model/kernel.py @@ -84,9 +84,16 @@ def __init__(self, hidden_size: int = 768, num_hidden_layers: int = 8, use_moe: def _inv_freq(dim: int, base: float, scaling: dict | None, end: int) -> tuple[torch.Tensor, float]: cpu = torch.device("cpu") # explicit: callers may run under ``with torch.device('meta')`` + kind = str((scaling or {}).get("type", "yarn")).lower() + if scaling is not None and kind == "ntk": + # NTK-aware: stretch the base so the lowest frequency spans the longer window. + base = base * float(scaling.get("factor", 1.0)) ** (dim / max(dim - 2, 1)) freqs = 1.0 / (base ** (torch.arange(0, dim, 2, device=cpu)[: dim // 2].float() / dim)) attn_factor = 1.0 - if scaling is not None: + if scaling is not None and kind == "linear": + freqs = freqs / float(scaling.get("factor", 1.0)) + attn_factor = float(scaling.get("attention_factor", 1.0)) + elif scaling is not None and kind == "yarn": # YaRN: f'(i) = f(i) * ((1 - g) + g / s), g a linear ramp over the band. orig_max = scaling.get("original_max_position_embeddings", 2048) factor = scaling.get("factor", 16) diff --git a/chronos/model/model_chronos.py b/chronos/model/model_chronos.py index c0d06fb..4d045d8 100644 --- a/chronos/model/model_chronos.py +++ b/chronos/model/model_chronos.py @@ -9,7 +9,7 @@ from .precision_ops import RMSNorm, lm_head_logits from .lookahead_router import LookaheadRouter from .moe_chronos import ChronosMOEFeedForward -from .hybrid_attention import make_attention +from .hybrid_attention import MLAAttention, make_attention from .recurrent import RecurrentController, recurrent_auxiliary_state from .outputs import ChronosCausalLMOutput from .multimodal import MultimodalFusion @@ -153,6 +153,17 @@ def __init__(self, config: ChronosConfig): config.head_dim, config.max_position_embeddings, config.rope_theta, config.rope_scaling, "cpu") self.last_recurrent_state = None + def _recompute_block(self, layer) -> bool: + """Whether ``layer`` is recomputed under ``config.recompute_policy`` + ("all" | "mla" | "swa" | "alternate"); MLX parity.""" + policy = getattr(self.config, "recompute_policy", "all") + if policy == "all": + return True + if policy == "alternate": + return next(i for i, other in enumerate(self.layers) if other is layer) % 2 == 0 + is_mla = isinstance(layer.self_attn, MLAAttention) + return is_mla if policy == "mla" else not is_mla + def _run_block( self, layer, hidden_states, position_embeddings, *, past_key_value=None, use_cache=False, attention_mask=None, available_expert_mask=None, @@ -164,6 +175,7 @@ def _run_block( if ( bool(getattr(self.config, "activation_checkpointing", False)) and self.training and not use_cache and not return_shared_kv + and self._recompute_block(layer) ): from torch.utils.checkpoint import checkpoint diff --git a/pyproject.toml b/pyproject.toml index 1663d21..a9c4a09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "Project_Chronos" -version = "0.3.2a0" +version = "0.3.2a1" 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_long_context.py b/tests/test_long_context.py new file mode 100644 index 0000000..d6a611e --- /dev/null +++ b/tests/test_long_context.py @@ -0,0 +1,365 @@ +"""Long-context training kernels (0.3.2a1): exactness against the reference paths. + +Chunked cross-entropy must give the loss and gradients of the full-logit +FP32 path, both for the stage loss (causal shift) and for LongRun's +pre-aligned labels, with -100 ignored. +""" +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) + +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) + 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)) + labels = mx.random.randint(3, cfg.vocab_size, (b, s)) + labels = mx.where(mx.arange(s)[None, :] < 7, -100, labels) # a masked prompt + return ids, labels + + +def _grads(model, loss_fn, ids, labels): + loss, grads = nn.value_and_grad(model, loss_fn)(model, ids, labels) + mx.eval(loss, grads) + return loss, dict(tree_flatten(grads)) + + +def _rel(a, b): + return float((mx.abs(a - b).max() / mx.maximum(mx.abs(a).max(), 1e-12)).item()) + + +@pytest.mark.parametrize("chunk", [1, 16, 33]) +def test_chunked_ce_equals_full_logits_stage_loss(chunk): + from chronos.mlx.training.trainer import _ce_stage_loss + full_cfg, chunk_cfg = _config(), _config(ce_chunk_size=chunk) + ids, labels = _batch(full_cfg) + a = _grads(_model(full_cfg), _ce_stage_loss, ids, labels) + b = _grads(_model(chunk_cfg), _ce_stage_loss, ids, labels) + assert abs(a[0].item() - b[0].item()) < 1e-5 + assert a[1].keys() == b[1].keys() + for name, g in a[1].items(): + assert _rel(g, b[1][name]) < 1e-4, name + + +def test_chunked_ce_aligned_labels_match_reference(): + from chronos.mlx.chunked_ce import chunked_lm_loss + mx.random.seed(2) + h = mx.random.normal((2, 50, 32)) + w = mx.random.normal((97, 32)) * 0.1 + lab = mx.random.randint(0, 97, (2, 50)) + lab = mx.where(mx.random.uniform(shape=lab.shape) < 0.3, -100, lab) + lg = (h @ w.T).reshape(-1, 97) + m = lab.reshape(-1) != -100 + ref = (nn.losses.cross_entropy(lg, mx.where(m, lab.reshape(-1), 0), reduction="none") * m).sum() / m.sum() + for chunk in (1, 7, 64, 1000): + assert abs(chunked_lm_loss(h, w, lab, chunk=chunk).item() - ref.item()) < 1e-5 + + +def test_all_ignored_labels_give_zero_loss(): + from chronos.mlx.chunked_ce import chunked_lm_loss + h = mx.ones((1, 9, 8)) + w = mx.ones((5, 8)) + assert chunked_lm_loss(h, w, mx.full((1, 9), -100), chunk=4).item() == 0.0 + + +def test_ce_chunk_field_round_trips_through_the_schema(): + from chronos.config_schema import MODEL_FIELDS + assert "ce_chunk_size" in MODEL_FIELDS + assert _config(ce_chunk_size=4096).ce_chunk_size == 4096 + assert _config().ce_chunk_size == 0 + + +@pytest.mark.parametrize("case", [{}, dict(recurrent_attention_mode="per_iteration"), dict(use_hybrid_attention=False), + dict(sliding_window_size=64)]) +def test_blockwise_attention_model_equals_single_kernel(case): + from chronos.mlx.training.trainer import _ce_stage_loss + ref_cfg, blk_cfg = _config(**case), _config(attention_block_size=8, **case) + ids, labels = _batch(ref_cfg) + a = _grads(_model(ref_cfg), _ce_stage_loss, ids, labels) + b = _grads(_model(blk_cfg), _ce_stage_loss, ids, labels) + assert abs(a[0].item() - b[0].item()) < 1e-5 + for name, g in a[1].items(): + assert _rel(g, b[1][name]) < 1e-4, name + + +def test_blockwise_attention_keeps_decode_identical(): + from chronos.mlx.model import ChronosMLXModel + ref_cfg, blk_cfg = _config(), _config(attention_block_size=8) + ids, _ = _batch(ref_cfg, b=1, s=30) + outs = [] + for cfg in (ref_cfg, blk_cfg): + mx.random.seed(0) + m = ChronosMLXModel(cfg) + m.eval() + logits, _, cache = m(ids[:, :24]) + step, _, _ = m(ids[:, 24:25], cache=cache) + mx.eval(logits, step) + outs.append((logits, step)) + assert _rel(outs[0][0], outs[1][0]) < 1e-5 + assert _rel(outs[0][1], outs[1][1]) < 1e-5 + + +@pytest.mark.parametrize("block", [0, 8]) +def test_document_masking_isolates_packed_documents(block): + """With document masking, a document's logits inside a packed row equal + the logits of that document alone (attention never crosses the separator).""" + from chronos.mlx.model import ChronosMLXModel + cfg = _config(attention_block_size=block, sliding_window_size=64, recurrent_depth=1, min_recurrent_depth=1, + rope_layout="per_head", + adaptive_halting=False) + mx.random.seed(0) + model = ChronosMLXModel(cfg) + model.eval() + from chronos.expert_clusters import segment_ids + import numpy as np + first = list(range(10, 23)) + [0] # document 1 ends with its separator + second = list(range(30, 49)) + row = mx.array([first + second]) + seg = mx.array(segment_ids(np.array(row))) + packed, _, _ = model(row, segments=seg) + alone, _, _ = model(mx.array([second])) + mx.eval(packed, alone) + # Document masking removes every cross-document term. What remains is + # the absolute-position offset: RoPE is relative in exact arithmetic, + # but FP32 cos/sin tables at shifted rows round differently (measured + # 8.8e-4 for the same document shifted by 14 positions, dense model). + assert _rel(alone[0], packed[0, len(first):]) < 2e-3 + leaky, _, _ = model(row) + assert _rel(alone[0], leaky[0, len(first):]) > 1e-3 # without masking the context differs + + +def test_router_pairs_do_not_cross_documents(): + from chronos.mlx.training import trainer as T + seg = mx.array([[0, 0, 0, 1, 1]]) + T._PAIR_SEGMENTS[0] = seg + try: + m = T._pair_mask(None, 1) + finally: + T._PAIR_SEGMENTS[0] = None + assert m.tolist() == [[1.0, 1.0, 0.0, 1.0]] + assert T._pair_mask(None, 1) is None + + + +# ---- 0.3.2a1 part 2: MoE chunks, selective recompute, RoPE, lookahead, curriculum, packing ---- + +def _grad_equal(a, b, tol=1e-4): + assert abs(a[0].item() - b[0].item()) < 1e-5 + assert a[1].keys() == b[1].keys() + for name, g in a[1].items(): + assert _rel(g, b[1][name]) < tol, name + + +@pytest.mark.parametrize("chunk", [7, 16]) +@pytest.mark.parametrize("miss", [0.0, 0.4]) +def test_moe_chunked_dispatch_equals_single_dispatch(chunk, miss): + from chronos.mlx.training.trainer import _ce_stage_loss + ref_cfg = _config(fallback_mask_prob=miss) + chk_cfg = _config(moe_chunk_tokens=chunk, fallback_mask_prob=miss) + ids, labels = _batch(ref_cfg) + mx.random.seed(5) + a = _grads(_model(ref_cfg), _ce_stage_loss, ids, labels) + mx.random.seed(5) + b = _grads(_model(chk_cfg), _ce_stage_loss, ids, labels) + _grad_equal(a, b) + + +@pytest.mark.parametrize("policy", ["all", "mla", "swa", "alternate"]) +def test_selective_recompute_keeps_loss_and_gradients(policy): + from chronos.mlx.training.trainer import _ce_stage_loss + ref_cfg = _config() + rc_cfg = _config(activation_checkpointing=True, recompute_policy=policy) + ids, labels = _batch(ref_cfg) + a = _grads(_model(ref_cfg), _ce_stage_loss, ids, labels) + b = _grads(_model(rc_cfg), _ce_stage_loss, ids, labels) + _grad_equal(a, b, tol=5e-4) + + +def test_selective_recompute_policy_picks_the_blocks(): + from chronos.mlx.model import _recompute_block + model = _model(_config()) + picked = {} + for policy in ("all", "mla", "swa", "alternate"): + model.config.recompute_policy = policy + picked[policy] = [_recompute_block(model.config, i, b) for i, b in enumerate(model.layers)] + assert all(picked["all"]) + assert picked["mla"] == [not v for v in picked["swa"]] + assert any(picked["mla"]) and not all(picked["mla"]) + assert picked["alternate"] == [i % 2 == 0 for i in range(len(model.layers))] + with pytest.raises(ValueError): + _config(recompute_policy="nope") + + +@pytest.mark.parametrize("kind", ["linear", "ntk", "yarn"]) +def test_rope_scaling_tables_match_between_mlx_and_torch(kind): + pytest.importorskip("torch") + import numpy as np + from chronos.mlx.attention import _rope_freqs + from chronos.model.kernel import precompute_freqs_cis + sc = {"type": kind, "factor": 4, "original_max_position_embeddings": 64, "beta_fast": 32, "beta_slow": 1} + c_t, s_t = precompute_freqs_cis(32, 256, 1e6, sc) + c_m, s_m = _rope_freqs(32, max_len=256, base=1e6, rope_scaling=sc) + assert np.abs(np.array(c_m) - c_t.numpy()).max() < 2e-5 + assert np.abs(np.array(s_m) - s_t.numpy()).max() < 2e-5 + base_c, _ = precompute_freqs_cis(32, 256, 1e6, None) + assert np.abs(base_c.numpy() - c_t.numpy()).max() > 1e-3 # scaling really changes the table + + +def test_linear_rope_scaling_compresses_positions(): + import numpy as np + from chronos.mlx.attention import _rope_freqs + plain, _ = _rope_freqs(16, max_len=64, base=1e4, rope_scaling=None) + lin, _ = _rope_freqs(16, max_len=64, base=1e4, rope_scaling={"type": "linear", "factor": 2}) + assert np.abs(np.array(lin[8]) - np.array(plain[4])).max() < 1e-5 + + +def test_rope_train_scaling_sets_rope_scaling_and_round_trips(): + from chronos.config_schema import MODEL_FIELDS + cfg = _config(rope_train_scaling={"type": "ntk", "factor": 8}) + assert cfg.rope_scaling == {"type": "ntk", "factor": 8} + assert _config().rope_scaling is None + new = {"rope_train_scaling", "recompute_policy", "moe_chunk_tokens", "reset_lookahead_at_documents"} + assert new <= set(MODEL_FIELDS) + + +def test_lookahead_document_rule_is_decoupled_from_router_rule(): + from chronos.mlx.training import trainer as T + base = dict(reset_router_at_documents=True, lambda_lookahead=1.0, lambda_temporal=1.0) + ids, labels = _batch(_config()) + ids = mx.where(mx.arange(ids.shape[1])[None, :] % 9 == 8, 0, ids) # separators inside each row + cases = {"follow": {}, "off": dict(reset_lookahead_at_documents=False), + "on": dict(reset_lookahead_at_documents=True)} + vals = {} + for name, extra in cases.items(): + model = _model(_config(**base, **extra)) + T._ce_stage_loss(model, ids, labels) + vals[name] = {k: float(v.item()) for k, v in model._last_loss_components.items()} + assert vals["follow"]["lookahead_loss"] == pytest.approx(vals["on"]["lookahead_loss"], rel=1e-6) + assert vals["off"]["lookahead_loss"] != pytest.approx(vals["on"]["lookahead_loss"], rel=1e-6) + assert vals["off"]["temporal_loss"] == pytest.approx(vals["on"]["temporal_loss"], rel=1e-6) + assert T._LOOKAHEAD_SEGMENTS[0] is None and T._PAIR_SEGMENTS[0] is None + + +def test_length_curriculum_steps_up_with_progress(): + from chronos.mlx.training.longrun import PhaseSpec, curriculum_seq_len + kw = dict(name="pretrain", budget_fraction=1.0, peak_lr=1e-3, tokens_per_step=1024) + ph = PhaseSpec(seq_len=512, seq_len_curriculum=[[0.0, 512], [0.3, 2048], [0.7, 8192]], **kw) + got = [curriculum_seq_len(ph, p) for p in (0.0, 0.29, 0.3, 0.69, 0.7, 1.0)] + assert got == [512, 512, 2048, 2048, 8192, 8192] + assert curriculum_seq_len(PhaseSpec(seq_len=256, **kw), 0.9) == 256 + + +def _write_sft(folder, records, with_lmask): + import numpy as np + os.makedirs(folder, exist_ok=True) + flat = np.concatenate([np.asarray(r, np.uint16) for r, _ in records]) + off = np.concatenate([[0], np.cumsum([len(r) for r, _ in records])]).astype(np.int64) + flat.tofile(os.path.join(folder, "a.bin")) + off.tofile(os.path.join(folder, "a.idx")) + np.asarray([p for _, p in records], np.int32).tofile(os.path.join(folder, "a.plen")) + if with_lmask: + mask = np.concatenate([np.arange(len(r)) >= p for r, p in records]).astype(np.uint8) + mask.tofile(os.path.join(folder, "a.lmask")) + + +@pytest.mark.parametrize("with_lmask", [False, True]) +def test_sft_packing_fills_rows_with_separated_records(tmp_path, with_lmask): + import numpy as np + from chronos.data.shards import SFTRecords + from chronos.expert_clusters import segment_ids + recs = [([5, 6, 7, 8], 2), ([9, 10, 11], 1), ([12, 13, 14, 15, 16], 3)] + _write_sft(str(tmp_path), recs, with_lmask) + rec = SFTRecords(str(tmp_path), seed=0) + ids, labels = rec.packed_batch(token_budget=64, seq_len=15) + assert ids.shape == labels.shape == (4, 16) + assert not np.any(ids[:, 0] == 0) # a row starts with a record + assert set(ids[ids != 0].tolist()) <= {5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + assert (ids == 0).sum() <= ids.shape[0] * 5 # at most one separator per 3+ record tokens + assert not np.any((labels != -100) & (ids == 0)) # separators carry no loss + allowed = {7, 8, 10, 11, 14, 15, 16} # reply tokens of the three records + assert set(labels[labels != -100].tolist()) <= allowed + assert segment_ids(ids).max() >= 1 # several documents per row + + +def test_longrun_pretrain_source_follows_the_curriculum(tmp_path): + import numpy as np + from chronos.mlx.training.longrun import PretrainSource + for split in ("train", "valid"): + d = tmp_path / "pretrain" / split + d.mkdir(parents=True) + (np.arange(5000, dtype=np.uint16) % 90).tofile(d / "a.bin") + src = PretrainSource(str(tmp_path), 64, seed=0) + assert src.batch(256)[0].shape == (4, 64) + src.set_seq_len(128) + assert src.batch(256)[0].shape == (2, 128) + + +# ---- Torch parity (0.3.2a1): selective recompute on the Torch side ---- + +@pytest.mark.parametrize("policy", ["all", "mla", "swa", "alternate"]) +def test_torch_selective_recompute_policy_keeps_gradients(policy): + torch = pytest.importorskip("torch") + from chronos.model.model_chronos import ChronosForCausalLM + + def run(**kw): + torch.manual_seed(0) + cfg = _config(**kw) + model = ChronosForCausalLM(cfg) + model.train() + ids = torch.randint(3, cfg.vocab_size, (2, 24)) + out = model(ids, labels=ids) + out.loss.backward() + grads = {n: p.grad.clone() for n, p in model.named_parameters() if p.grad is not None} + return out.loss.item(), grads + + base_loss, base = run() + loss, grads = run(activation_checkpointing=True, recompute_policy=policy) + assert abs(loss - base_loss) < 1e-5 + assert base.keys() == grads.keys() + for name, g in base.items(): + denom = max(float(g.abs().max()), 1e-12) + assert float((g - grads[name]).abs().max()) / denom < 1e-3, name + + +def test_torch_recompute_policy_selects_expected_layers(): + pytest.importorskip("torch") + from chronos.model.hybrid_attention import MLAAttention + from chronos.model.model_chronos import ChronosModel + model = ChronosModel(_config()) + is_mla = [isinstance(layer.self_attn, MLAAttention) for layer in model.layers] + alternate = [i % 2 == 0 for i in range(len(is_mla))] + expect = {"all": [True] * len(is_mla), "mla": is_mla, "swa": [not v for v in is_mla], "alternate": alternate} + for policy, want in expect.items(): + model.config.recompute_policy = policy + assert [model._recompute_block(layer) for layer in model.layers] == want