diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/README.md b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/README.md
new file mode 100644
index 0000000000..eca67b9e3f
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/README.md
@@ -0,0 +1,339 @@
+# BESE + Mamba-3 SSD Hybrid
+
+**Author:** Omer Bese ([@mrbese](https://github.com/mrbese))
+**Date:** 2026-04-16
+**Track:** Non-record (SSM / State-space model submission)
+**val_bpb:** 1.3571 (INT6 + LZMA + sliding window eval with n-gram tilt)
+**Artifact size:** 7,614,888 bytes (48% of 16 MB limit)
+
+---
+
+## Overview
+
+This submission combines two experimental ideas requested by the challenge organizers:
+
+1. **State-space models** (specifically Mamba-3 SSD) — checking the "State-space models" bounty from the challenge README
+2. **Novel tokenizer** (BESE, a custom 288-vocab byte-level tokenizer) — testing whether sub-byte tokenization gives SSMs an advantage through 2x token density
+
+To our knowledge, this is the first submission to pair a custom byte-level tokenizer with a Mamba-3 architecture.
+
+## Architecture
+
+**Hybrid: 6 Mamba-3 SSD blocks + 2 Attention blocks (8 layers total)**
+
+```
+Layer 0: Mamba-3 SSD
+Layer 1: Mamba-3 SSD
+Layer 2: Attention (GQA, FlashAttention/SDPA)
+Layer 3: Mamba-3 SSD
+Layer 4: Mamba-3 SSD
+Layer 5: Attention (GQA, FlashAttention/SDPA)
+Layer 6: Mamba-3 SSD
+Layer 7: Mamba-3 SSD
+```
+
+### Model Configuration
+
+| Parameter | Value | Notes |
+|-----------|-------|-------|
+| `model_dim` | 512 | |
+| `num_layers` | 8 | 6 Mamba + 2 Attention |
+| `d_state` | 128 | SSM state dimension |
+| `expand` | 2 | d_inner = 1024 |
+| `headdim` | 64 | SSM head dimension |
+| `nheads` (SSM) | 16 | d_inner / headdim |
+| `ngroups` | 1 | All heads share B/C (reference Mamba-2 default) |
+| `chunk_size` | 64 | SSD chunk size |
+| `num_heads` (Attn) | 8 | |
+| `num_kv_heads` (Attn) | 4 | GQA |
+| `mlp_mult` | 3.0 | Attention block MLP |
+| `vocab_size` | 288 | BESE tokenizer |
+| **Total params** | **15,152,432** | |
+
+### Key Design Decisions
+
+**1. ngroups=1 (shared B/C across heads)**
+
+All 16 SSM heads share the same B (input-to-state) and C (state-to-output) projections, with only 1 group. This matches the reference Mamba-2 implementation and was confirmed optimal by PR #1644 ablations. Saves ~6.9M parameters vs per-head B/C (ngroups=16), which we reallocate to larger d_state.
+
+**2. No depth recurrence on SSM layers**
+
+PR #1355 measured a -69 mBPB penalty from depth recurrence on Mamba blocks. Unlike transformers where attention re-processes with updated context, SSM state from pass 1 does not inform pass 2 (initial_states=None). We disable depth recurrence entirely.
+
+**3. Two attention layers at positions [2, 5]**
+
+Following PR #1644's architecture, attention layers provide global token mixing at strategic points, dividing the SSM blocks into three equal segments. The SSM layers handle local sequential processing (O(n)), while attention provides periodic global information bottlenecks.
+
+**4. d_state=128 with ngroups=1**
+
+With shared B/C (ngroups=1), the projection cost is only 1 x d_state per position for B and C. Doubling d_state from 64 to 128 costs just ~400K extra parameters but doubles the SSM's memory bandwidth — how much past context each state vector can retain.
+
+## BESE Tokenizer
+
+BESE (Byte-Encoded Sub-byte Encoding) is a two-layer tokenizer:
+- **Layer 1:** 40 base tokens (digrams covering 95% of English byte pairs)
+- **Layer 2:** 248 BPE merges on top of the base tokens
+- **Total vocab:** 288 tokens
+
+Compared to SP1024 (the challenge default), BESE produces ~2x more tokens per byte of text. This means:
+- **Embedding table:** 288 x 512 = 147K params (vs SP1024's 1024 x 512 = 524K, or SP8192's 8192 x 512 = 4.2M)
+- **Saved parameters** go directly into model capacity
+- **Longer effective sequences** for the same token count
+
+### BPB Correctness Proof
+
+The competition rules require that any tokenizer change "prove with certainty that the val_bpb is correctly calculated." Because BESE is custom, we include the proof inline below rather than referencing prior submissions.
+
+**The invariant we maintain:**
+
+> For any input string `s`,
+> `sum(BYTES_PER_TOKEN[t] for t in encode(s)) == len(s.encode("utf-8"))`
+
+This means our BPB number is computed against the same denominator (UTF-8 bytes) as every SP1024 / SP8192 submission, so scores are directly comparable.
+
+**Per-token byte accounting** (defined in `bese_constants.py::build_bytes_per_token`):
+
+| Token category | Count | Bytes per token |
+|---|---|---|
+| Single-letter tokens (e, t, a, o, i, n, s, r, h, d, l) | 11 | 1 |
+| Group tokens (ufbz, cwvj, mykx, gpq) — prefix only | 4 | **0** |
+| Position tokens (P1–P4) — carry the actual character | 4 | 1 |
+| Punctuation (space, period, comma, newline, ?, quote, OTHER_PUNCT) | 7 | 1 |
+| Digit tokens (0–9) | 10 | 1 |
+| Special tokens (PAD, BOS, EOS, UNK) | 4 | 0 |
+| BPE merge tokens | 248 | recursive sum of constituents |
+
+**Why it holds — by case** (see `bese_fast_bpe.py::_text_to_base_tokens`):
+
+For every input character `ch`:
+
+1. **Mapped path:** If `lower(ch)` is in `ENCODE_TABLE` AND the mapped tokens' byte sum equals `len(ch.encode("utf-8"))`, we emit those tokens. The equality is checked at runtime — if it ever failed we'd fall through to the byte-fallback path. Group tokens contribute 0, position tokens contribute 1 → a (group, pos) pair encodes one ASCII character (1 byte) correctly. Single letters contribute 1 byte directly.
+
+2. **Byte-fallback path:** Otherwise (uppercase ASCII not normalized, non-ASCII, unknown punctuation), we emit `OTHER_PUNCT_ID` exactly `len(ch.encode("utf-8"))` times. Each `OTHER_PUNCT_ID` contributes 1 byte → total bytes match the UTF-8 byte count of the character exactly.
+
+In both branches the sum-of-bytes invariant is preserved per character, so it holds for the full string by linearity.
+
+**For BPE merges**, the byte count is computed transitively when merges are loaded (`bese_fast_bpe.py::compute_bytes_per_token`):
+
+```python
+bpt = np.zeros(self.vocab_size, dtype=np.int16)
+bpt[:BASE_VOCAB_SIZE] = BYTES_PER_TOKEN
+merge_bpt = {i: int(BYTES_PER_TOKEN[i]) for i in range(BASE_VOCAB_SIZE)}
+for pair, new_id in self.merges:
+ merge_bpt[new_id] = merge_bpt[pair[0]] + merge_bpt[pair[1]]
+ bpt[new_id] = merge_bpt[new_id]
+```
+
+So every merge inherits its byte count from its constituents — invariant preserved through the BPE layer.
+
+**Self-test** (run this from inside the records folder; runs in <1 s on CPU):
+
+```python
+import numpy as np
+from bese_fast_bpe import BeseFastBPE
+
+# Load tokenizer
+tok = BeseFastBPE.load("tokenizer.json")
+bpt = tok.compute_bytes_per_token()
+
+# Roundtrip + byte-count check across diverse inputs
+test_cases = [
+ "the quick brown fox jumps over the lazy dog",
+ "Hello, World! 1234567890",
+ "naïve café résumé — emojis: 🚀✨🌍",
+ "newlines\nand\ttabs and \"quotes\" and 'apostrophes'",
+ "MixedCase HTML tag
and JSON {\"a\": 1}",
+]
+for s in test_cases:
+ ids = tok.encode(s)
+ # 1. Lossless roundtrip on the supported character set
+ # 2. Byte-count invariant
+ assert sum(int(bpt[t]) for t in ids) == len(s.encode("utf-8")), \
+ f"byte invariant failed on: {s!r}"
+print("OK — BPB byte invariant holds for all test cases")
+```
+
+The same `bpt` table is what the training/eval code uses to compute val_bpb, so this self-test is checking the exact accounting used to score the submission.
+
+**val_bpb formula** (matches the upstream definition):
+
+```
+val_bpb = (sum_of_token_NLLs_in_nats / sum_of_bytes_per_token) / log(2)
+```
+
+where the numerator and denominator are summed over the same set of evaluated tokens. Because the per-token byte counts sum to the true UTF-8 byte length of the validation set, this is identical to the SP1024 / SP8192 BPB formula evaluated on the same FineWeb validation data.
+
+## Training
+
+- **Hardware:** 8x NVIDIA H100 80GB SXM (RunPod)
+- **Training time:** 600 seconds (wallclock cap)
+- **Steps completed:** 2,191
+- **Step average:** 274 ms/step
+- **Optimizer:** Muon (Newton-Schulz) for 2D matrices, AdamW for scalars and embeddings
+- **EMA decay:** 0.9965
+- **Warmdown:** 5000 iterations
+- **SWA:** Activated at step 1200
+- **Sequence length:** 2048 (train and eval)
+- **Batch tokens:** 786,432 per step (global)
+
+### Training Curve
+
+| Step | val_bpb |
+|------|---------|
+| 0 | 4.1571 |
+| 500 | 1.5460 |
+| 1000 | 1.4268 |
+| 1500 | 1.3806 |
+| 2000 | 1.3489 |
+| 2191 (final) | 1.3458 |
+
+## Evaluation
+
+| Stage | val_bpb | Notes |
+|-------|---------|-------|
+| Raw (post-EMA) | 1.3475 | Diagnostic |
+| INT6 roundtrip | 1.3809 | Quantized model |
+| **INT6 + Sliding Window + N-gram tilt** | **1.3571** | **Final submission score** |
+
+- **Quantization:** Mixed INT6 (6-bit) for MLP, attention, and Mamba projection weights. Scalar params (D, dt_bias, A_log, norms) stored as FP16.
+- **Compression:** LZMA preset 9
+- **Sliding window eval:** stride=64, full 2048 context per window
+- **N-gram tilt:** Pre-computed trigram prior from training data, applied as additive logit bias during sliding window eval
+
+> **Statistical-significance note.** The headline 1.3571 BPB is a **single-seed** result from the `dim=512, d_state=128` configuration (`train_log_run1.txt`). The two additional logs (`train_log_run2_d64.txt`, `train_log_run3_dim576.txt`) are **architecture ablations** (different `d_state` and `model_dim`), not seed replicates of the headline config. Three-seed validation of the headline config is in the *Ongoing Work* section, pending compute credits. We submit this as a non-record entry under the rule that allows in-progress and unoptimized solutions for novel ideas; it is not a leaderboard claim.
+
+### Artifact Size
+
+| Component | Bytes |
+|-----------|-------|
+| Compressed model (INT6 + LZMA) | 7,452,680 |
+| Code (train_gpt.py + mamba3_ssd.py + tokenizer) | 162,208 |
+| **Total** | **7,614,888** |
+| Budget remaining | 8,385,112 (52% unused) |
+
+## Additional Runs
+
+We ran three configurations to ablate the architecture. These are **architecture ablations**, not seed replicates of a single config:
+
+| Config | Params | Steps | Raw BPB | INT6 BPB | SW BPB | Artifact |
+|--------|--------|-------|---------|----------|--------|----------|
+| dim=512, d_state=64 | 14.8M | 2,482 | 1.3254 | 1.3445 | not completed | 7.96 MB |
+| **dim=512, d_state=128** (headline) | **15.2M** | **2,191** | **1.3458** | **1.3809** | **1.3571** | **7.56 MB** |
+| dim=576, d_state=128, mlp3.5 | 19.7M | 1,847 | 1.3415 | 1.4053 | not completed | 8.42 MB |
+
+Key findings:
+- **d_state=128 vs 64:** Slightly worse raw BPB (fewer steps) but sliding window eval works and n-gram tilt recovers the gap
+- **dim=576 (wider model):** Best per-step learning rate and best raw BPB (1.3415 at step 1847), but larger INT6 quantization gap (+60 mBPB). Suggests QAT would unlock significant gains for wider Mamba models.
+- **Artifact headroom:** Even the widest model uses only 8.42/16 MB, leaving substantial room for growth
+
+## SSD Implementation Notes
+
+We implemented Mamba-3 SSD in pure PyTorch (no custom CUDA/Triton kernels) using the chunked parallel formulation from the Mamba-2 paper. Key components:
+
+- **segsum:** Stable cumulative sum for decay computation via lower-triangular masking
+- **ssd_chunked:** Chunked parallel SSD with intra-chunk quadratic attention and inter-chunk state recurrence
+- **Causality fix:** We discovered and fixed a causality bug in the reference implementation's inter-chunk decay matrix (diagonal was 1, allowing each chunk to see its own state through Y_off). Fixed by shifting the column index in the einsum.
+
+We attempted integration with the official `mamba-ssm` Triton kernels (`mamba_chunk_scan_combined`), which worked on single GPU but caused segfaults under multi-GPU torchrun after ~100 steps. The pure PyTorch fallback is stable and provides correct results, though ~2-3x slower per step.
+
+## Code Structure
+
+The submission folder contains four Python files:
+
+| File | Bytes | Role |
+|---|---|---|
+| `train_gpt.py` | 108,597 | Self-contained training entry point — model definition, training loop, eval, quantization, compression |
+| `mamba3_ssd.py` | 20,912 | Mamba-3 SSD block + chunked parallel `ssd_chunked` algorithm |
+| `bese_fast_bpe.py` | 25,263 | BESE tokenizer encode/decode + BPE merge application |
+| `bese_constants.py` | 3,941 | BESE alphabet constants and `BYTES_PER_TOKEN` lookup table |
+
+**On the FAQ rule "all counted code should live in `train_gpt.py`":** the helper modules above are fully accounted for in `submission.json::code_bytes` (162,208 bytes total) and bundled with the submission, so the artifact-size accounting is honest. We split them out for readability — `mamba3_ssd.py` is the SSD algorithm, `bese_*.py` is the tokenizer — but they could be inlined into `train_gpt.py` mechanically with no logic change, and the artifact size would be identical. We left them separate to make the per-component contribution to artifact size legible to reviewers; happy to inline if preferred.
+
+`tokenizer.json` (3,495 bytes) holds the trained BPE merges and is loaded at training time only (it is *not* part of the artifact — the merges are baked into the model's input pipeline).
+
+## Comparison to Other SSM Submissions
+
+| Submission | BPB | Arch | Vocab | Artifact |
+|-----------|------|------|-------|----------|
+| PR #1479 GDN hybrid | 1.1450 | 8 GDN + 2 Attn | SP8192 | 13.83 MB |
+| PR #1245 Hymba | 1.1470 | 8L hybrid | SP8192 | ~15 MB |
+| PR #1644 Mamba-3 | 1.1473 | 5 SSM + 2 Attn | SP8192 | ~14 MB |
+| **This (BESE + Mamba-3)** | **1.3571** | **6 SSM + 2 Attn** | **BESE 288** | **7.56 MB** |
+
+The ~210 mBPB gap to the best SSM submissions is attributable to:
+- Byte-level prediction with 288 vocab (estimated ~30-50 mBPB penalty vs SP8192)
+- Pure PyTorch SSD without Triton kernels (fewer training steps, ~40-60 mBPB)
+- No test-time training (~30-50 mBPB based on other submissions)
+- No torch.compile (~20-30 mBPB)
+
+The unique contribution is demonstrating that byte-level tokenization + SSM is viable, achieving competitive artifact efficiency (half the 16 MB budget) while leaving substantial room for optimization.
+
+## Reproduction
+
+### Quick path — run the submitted artifact from this records folder
+
+The four Python files in this folder are self-contained. From an 8xH100 SXM RunPod pod with the official Parameter Golf template (PyTorch 2.x, CUDA 12.x):
+
+```bash
+# From the cloned upstream repo:
+cd parameter-golf
+
+# Install the two extra packages (everything else is in the template):
+pip install einops sentencepiece
+
+# Prepare the SP1024 cached FineWeb (used as the input pipeline before BESE re-encoding):
+python3 data/cached_challenge_fineweb.py --variant sp1024
+
+# Run training + eval directly from this records folder:
+cd records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid
+
+torchrun --standalone --nproc_per_node=8 train_gpt.py \
+ VOCAB_SIZE=288 \
+ TOKENIZER_PATH=./tokenizer.json \
+ DATA_PATH=../../../data/datasets/fineweb10B_sp1024/ \
+ RUN_ID=bese_mamba3_repro
+```
+
+This produces the same INT6+LZMA artifact and the `final_*` BPB numbers in `train_log_run1.txt`. Total wallclock on 8xH100 SXM is ~10 min training + ~7-8 min eval (sliding window + n-gram tilt is the eval bottleneck).
+
+### Full pipeline — rebuild the BESE shards from scratch
+
+If you want to reproduce the data-prep step (untimed) and not rely on cached BESE shards, the full pipeline lives on the author's fork:
+
+```bash
+cd /workspace
+git clone https://github.com/mrbese/parameter-golf-bese.git bese
+cd bese
+git checkout v7-mamba
+
+# Reuses /workspace/parameter-golf/data/datasets/fineweb10B_sp1024/ from above.
+pip install einops --break-system-packages
+python scripts/runpod_v7_mamba.py --num-gpus 8
+
+# Or with pre-existing BESE shards (cached on RunPod network volume):
+python scripts/runpod_v7_mamba.py --skip-shards --num-gpus 8
+```
+
+The fork's `scripts/runpod_v7_mamba.py` is just an orchestrator around the same `train_gpt.py` shipped in this records folder; it adds shard re-encoding from SP1024 → BESE base tokens, BPE training, and ngram-prior building. The records folder contains the artifacts of those steps (`tokenizer.json`, the BPE merges) so the quick path above can skip them.
+
+## Ongoing Work
+
+We have a pending compute credit request and plan to continue optimizing this submission. Planned next steps:
+
+- **Triton kernel integration**: Fix the multi-GPU segfault in `mamba_chunk_scan_combined` to get 2-3x faster steps (~150ms vs 274ms), enabling ~4,000 steps in 600s
+- **torch.compile**: Unblocked once Triton kernels are stable — additional ~15% step speedup
+- **Wider model with QAT**: dim=576 + mlp3.5 achieved 1.3415 raw BPB but has a +60 mBPB INT6 gap. Quantization-aware training should close this gap substantially
+- **Test-time training (TTT)**: Disabled in current runs to save credits. Other SSM submissions show ~30-50 mBPB improvement from TTT
+- **SP8192 + BESE comparison**: Direct ablation of tokenizer impact on the same Mamba architecture
+- **Three-seed statistical significance** for the headline `dim=512, d_state=128` configuration
+
+Conservative target with all optimizations: **1.17-1.20 BPB**, which would be competitive with the best SSM submissions while maintaining BESE's artifact efficiency advantage.
+
+## Acknowledgments
+
+Architecture decisions informed by:
+- PR #1644 by mradassaad (best Mamba-3 submission, exhaustive ablation study)
+- PR #1355 by mradassaad (SSM depth recurrence ablation)
+- PR #1245 by mkenney2 (Hymba hybrid architecture)
+- The Mamba-2 paper (Dao and Gu, 2024) for the SSD algorithm
+- mamba3-minimal (VikramKarLex) for the reference pure-PyTorch implementation
diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/bese_constants.py b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/bese_constants.py
new file mode 100644
index 0000000000..a54f1abbfb
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/bese_constants.py
@@ -0,0 +1,125 @@
+"""
+Shared BESE alphabet constants and lookup tables (40-token base vocabulary).
+
+Used by bese_tokenizer.py and bese_bpe_tokenizer.py — single source of truth.
+
+v3: Reordered groups by aggregate frequency (most common group → lowest ID).
+ Group 15 (ufbz, 6.57%) > Group 16 (cwvj, 6.35%) > Group 17 (mykx, 5.35%) > Group 18 (gpq, 4.0%).
+ The group token ID itself now encodes frequency tier information.
+
+v2: Promoted h, d, l to single-token letters (8→11 singles).
+ Regrouped remaining 15 letters into 4 groups (5→4 groups).
+ Reordered positions within groups by frequency (most freq → P1).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+# --- Special tokens ---
+PAD_ID = 0
+BOS_ID = 1
+EOS_ID = 2
+UNK_ID = 3
+
+# Single-token letters: 11 most frequent in English (e=4 ... l=14)
+SINGLE_LETTERS = "etaoinsrhdl"
+SINGLE_LETTER_START = 4
+
+# Key groups (group token = 0 bytes; position completes character)
+# Ordered by frequency within each group (most frequent → P1).
+# Letters that commonly co-occur in English bigrams are placed in
+# DIFFERENT groups to give BPE cleaner merge patterns.
+GROUP_START = 15
+GROUPS = [
+ "ufbz", # Group 15: u(2.8%) f(2.2%) b(1.5%) z(0.07%) — 6.57% aggregate
+ "cwvj", # Group 16: c(2.8%) w(2.4%) v(1.0%) j(0.15%) — 6.35% aggregate
+ "mykx", # Group 17: m(2.4%) y(2.0%) k(0.8%) x(0.15%) — 5.35% aggregate
+ "gpq", # Group 18: g(2.0%) p(1.9%) q(0.10%) — 4.00% aggregate
+]
+
+POS_START = 19
+
+SPACE_ID = 23
+PERIOD_ID = 24
+COMMA_ID = 25
+NEWLINE_ID = 26
+QUESTION_ID = 27
+QUOTE_ID = 28
+OTHER_PUNCT_ID = 29
+DIGIT_START = 30
+
+BASE_VOCAB_SIZE = 40
+VOCAB_SIZE = BASE_VOCAB_SIZE # alias for base-only tokenizer
+
+
+def build_encode_table() -> dict[str, list[int]]:
+ """Character -> base token id(s)."""
+ table: dict[str, list[int]] = {}
+ for i, ch in enumerate(SINGLE_LETTERS):
+ table[ch] = [SINGLE_LETTER_START + i]
+ for gi, group in enumerate(GROUPS):
+ group_token = GROUP_START + gi
+ for pi, ch in enumerate(group):
+ pos_token = POS_START + pi
+ table[ch] = [group_token, pos_token]
+ table[" "] = [SPACE_ID]
+ table["."] = [PERIOD_ID]
+ table[","] = [COMMA_ID]
+ table["\n"] = [NEWLINE_ID]
+ table["?"] = [QUESTION_ID]
+ for ch in ["'", '"', "\u2018", "\u2019", "\u201c", "\u201d"]:
+ table[ch] = [QUOTE_ID]
+ for d in range(10):
+ table[str(d)] = [DIGIT_START + d]
+ return table
+
+
+def build_decode_table() -> dict[tuple[int, ...], str]:
+ """Base token sequence -> single character (best-effort)."""
+ table: dict[tuple[int, ...], str] = {}
+ for i, ch in enumerate(SINGLE_LETTERS):
+ table[(SINGLE_LETTER_START + i,)] = ch
+ for gi, group in enumerate(GROUPS):
+ group_token = GROUP_START + gi
+ for pi, ch in enumerate(group):
+ pos_token = POS_START + pi
+ table[(group_token, pos_token)] = ch
+ table[(SPACE_ID,)] = " "
+ table[(PERIOD_ID,)] = "."
+ table[(COMMA_ID,)] = ","
+ table[(NEWLINE_ID,)] = "\n"
+ table[(QUESTION_ID,)] = "?"
+ table[(QUOTE_ID,)] = "'"
+ table[(OTHER_PUNCT_ID,)] = "?"
+ for d in range(10):
+ table[(DIGIT_START + d,)] = str(d)
+ return table
+
+
+def build_bytes_per_token() -> np.ndarray:
+ """UTF-8 bytes each base token represents (BPB-critical)."""
+ bpt = np.zeros(BASE_VOCAB_SIZE, dtype=np.int16)
+ for i in range(len(SINGLE_LETTERS)):
+ bpt[SINGLE_LETTER_START + i] = 1
+ for i in range(4):
+ bpt[POS_START + i] = 1
+ for tid in (
+ SPACE_ID,
+ PERIOD_ID,
+ COMMA_ID,
+ NEWLINE_ID,
+ QUESTION_ID,
+ QUOTE_ID,
+ OTHER_PUNCT_ID,
+ ):
+ bpt[tid] = 1
+ for d in range(10):
+ bpt[DIGIT_START + d] = 1
+ return bpt
+
+
+# Eager singletons (import cost is tiny; avoids recomputation)
+ENCODE_TABLE = build_encode_table()
+DECODE_TABLE = build_decode_table()
+BYTES_PER_TOKEN = build_bytes_per_token()
diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/bese_fast_bpe.py b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/bese_fast_bpe.py
new file mode 100644
index 0000000000..8f2e9a1ce8
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/bese_fast_bpe.py
@@ -0,0 +1,684 @@
+"""
+Fast BPE training and encoding for BESE tokenizer.
+
+Replaces the O(num_merges * total_tokens) pure-Python implementation with
+an O(total_tokens * log(total_tokens)) priority-queue based approach.
+
+The key insight: instead of scanning the entire corpus once per merge,
+we maintain a priority queue of pair counts and a doubly-linked list of
+tokens. Each merge only touches positions where the merged pair exists.
+
+On 10K FineWeb docs (~40M base tokens), this reduces BPE training from
+~22 minutes to ~30 seconds, and encoding from ~25 minutes to ~60 seconds.
+"""
+
+from __future__ import annotations
+
+import json
+import heapq
+import numpy as np
+from collections import defaultdict
+from pathlib import Path
+
+try:
+ from .bese_constants import (
+ BASE_VOCAB_SIZE,
+ BOS_ID,
+ BYTES_PER_TOKEN,
+ DECODE_TABLE,
+ EOS_ID,
+ ENCODE_TABLE,
+ GROUP_START,
+ GROUPS,
+ OTHER_PUNCT_ID,
+ PAD_ID,
+ SINGLE_LETTERS,
+ UNK_ID,
+ )
+except ImportError:
+ from bese_constants import (
+ BASE_VOCAB_SIZE,
+ BOS_ID,
+ BYTES_PER_TOKEN,
+ DECODE_TABLE,
+ EOS_ID,
+ ENCODE_TABLE,
+ GROUP_START,
+ GROUPS,
+ OTHER_PUNCT_ID,
+ PAD_ID,
+ SINGLE_LETTERS,
+ UNK_ID,
+ )
+
+
+def _text_to_base_tokens(text: str) -> list[int]:
+ """Convert text to BESE base token sequence."""
+ tokens = []
+ for ch in text:
+ lower = ch.lower()
+ if lower in ENCODE_TABLE:
+ utf8_len = len(ch.encode("utf-8"))
+ mapped = ENCODE_TABLE[lower]
+ mapped_bytes = sum(BYTES_PER_TOKEN[t] for t in mapped)
+ if utf8_len == mapped_bytes:
+ tokens.extend(mapped)
+ else:
+ tokens.extend([OTHER_PUNCT_ID] * utf8_len)
+ else:
+ utf8_len = len(ch.encode("utf-8"))
+ tokens.extend([OTHER_PUNCT_ID] * utf8_len)
+ return tokens
+
+
+# ---------------------------------------------------------------------------
+# Fast BPE training using indexed pair counting
+# ---------------------------------------------------------------------------
+
+class _Node:
+ """Doubly-linked list node for fast BPE merge operations."""
+ __slots__ = ("token", "prev", "next", "doc_id")
+
+ def __init__(self, token: int, doc_id: int):
+ self.token = token
+ self.prev = None
+ self.next = None
+ self.doc_id = doc_id
+
+
+def _encode_texts_worker(texts: list[str]) -> list[list[int]]:
+ """Worker function for parallel base-token encoding."""
+ return [_text_to_base_tokens(text) for text in texts]
+
+
+def train_bpe_merges_fast(texts: list[str], num_merges: int = 250, verbose: bool = True) -> list:
+ """
+ Learn BPE merges using an efficient indexed approach.
+
+ Instead of scanning all sequences for every merge, we:
+ 1. Build a doubly-linked list of all tokens (encoding parallelized across CPUs)
+ 2. Maintain a max-heap of pair counts for O(log n) best-pair lookup
+ 3. For each merge, update only the affected positions
+
+ This is O(total_tokens + num_merges * avg_pair_count * log(num_pairs)) instead of
+ O(num_merges * total_tokens).
+ """
+ import multiprocessing as mp
+
+ if verbose:
+ print(f"Encoding {len(texts)} texts with base BESE tokenizer...")
+
+ # Step 1: Parallel encode all texts to base tokens
+ n_workers = min(mp.cpu_count(), 128)
+ chunk_size = max(1, len(texts) // n_workers)
+ chunks = [texts[i:i + chunk_size] for i in range(0, len(texts), chunk_size)]
+
+ import time as _time
+ t_enc = _time.time()
+ with mp.Pool(n_workers) as pool:
+ encoded_chunks = pool.map(_encode_texts_worker, chunks)
+ all_encoded = [tokens for chunk in encoded_chunks for tokens in chunk]
+ if verbose:
+ print(f" Parallel encoding: {n_workers} workers, {_time.time() - t_enc:.1f}s")
+
+ # Step 1b: Build linked lists and pair index from encoded tokens
+ doc_heads = []
+ pair_positions = defaultdict(set)
+ all_nodes = []
+
+ total_base = 0
+ for doc_id, base_tokens in enumerate(all_encoded):
+ total_base += len(base_tokens)
+ if not base_tokens:
+ doc_heads.append(None)
+ continue
+
+ nodes = []
+ for t in base_tokens:
+ node = _Node(t, doc_id)
+ node_id = len(all_nodes)
+ if nodes:
+ prev_node = nodes[-1]
+ prev_node.next = node_id
+ node.prev = len(all_nodes) - 1
+ nodes.append(node)
+ all_nodes.append(node)
+
+ doc_heads.append(len(all_nodes) - len(nodes))
+
+ for i in range(len(nodes) - 1):
+ nid = len(all_nodes) - len(nodes) + i
+ pair = (nodes[i].token, nodes[i + 1].token)
+ pair_positions[pair].add(nid)
+
+ del all_encoded
+
+ if verbose:
+ print(f"Base tokens: {total_base:,}")
+ print(f"Unique pairs: {len(pair_positions):,}")
+ print(f"Learning {num_merges} BPE merges (heap mode)...")
+
+ # Step 2: Greedily merge most frequent pairs using a max-heap
+ merges = []
+ next_id = BASE_VOCAB_SIZE
+
+ pair_counts = {pair: len(positions) for pair, positions in pair_positions.items()}
+
+ # Build max-heap (negate counts for max-heap via min-heap)
+ heap = [(-count, pair) for pair, count in pair_counts.items() if count >= 2]
+ heapq.heapify(heap)
+
+ merge_num = 0
+ while merge_num < num_merges and heap:
+ # Pop best pair from heap (skip stale entries)
+ while heap:
+ neg_count, best_pair = heapq.heappop(heap)
+ current_count = pair_counts.get(best_pair, 0)
+ if current_count >= 2 and current_count == -neg_count:
+ best_count = current_count
+ break
+ else:
+ break
+
+ # Get all positions where this pair occurs and filter stale entries
+ raw_positions = pair_positions.get(best_pair, set())
+ positions = []
+ for nid in raw_positions:
+ node_a = all_nodes[nid]
+ if node_a.next is None:
+ continue
+ node_b = all_nodes[node_a.next]
+ if node_a.token != best_pair[0] or node_b.token != best_pair[1]:
+ continue
+ positions.append(nid)
+
+ # Remove the pair from tracking
+ if best_pair in pair_positions:
+ del pair_positions[best_pair]
+ if best_pair in pair_counts:
+ del pair_counts[best_pair]
+
+ # Re-check count after filtering stale positions
+ if len(positions) < 2:
+ continue # don't increment merge_num — this wasn't a real merge
+
+ new_id = next_id
+ merges.append((best_pair, new_id))
+
+ # Apply the merge at each valid position
+ for nid in positions:
+ node_a = all_nodes[nid]
+ node_b = all_nodes[node_a.next]
+
+ # Remove old pairs involving node_a and node_b from index
+ # Left neighbor pair: (prev, a)
+ if node_a.prev is not None:
+ prev_node = all_nodes[node_a.prev]
+ old_left_pair = (prev_node.token, node_a.token)
+ nid_prev = node_a.prev
+ pair_positions.get(old_left_pair, set()).discard(nid_prev)
+ if old_left_pair in pair_positions and not pair_positions[old_left_pair]:
+ del pair_positions[old_left_pair]
+ if old_left_pair in pair_counts:
+ pair_counts[old_left_pair] = max(pair_counts[old_left_pair] - 1, 0)
+ if pair_counts[old_left_pair] == 0:
+ del pair_counts[old_left_pair]
+
+ # Right neighbor pair: (b, next)
+ if node_b.next is not None:
+ next_node = all_nodes[node_b.next]
+ old_right_pair = (node_b.token, next_node.token)
+ b_nid = node_a.next
+ pair_positions.get(old_right_pair, set()).discard(b_nid)
+ if old_right_pair in pair_positions and not pair_positions[old_right_pair]:
+ del pair_positions[old_right_pair]
+ if old_right_pair in pair_counts:
+ pair_counts[old_right_pair] = max(pair_counts[old_right_pair] - 1, 0)
+ if pair_counts[old_right_pair] == 0:
+ del pair_counts[old_right_pair]
+
+ # Merge: node_a becomes new_id, node_b is unlinked
+ node_a.token = new_id
+ node_a.next = node_b.next
+ if node_b.next is not None:
+ all_nodes[node_b.next].prev = nid
+
+ # Add new pairs and push onto heap
+ if node_a.prev is not None:
+ prev_node = all_nodes[node_a.prev]
+ new_left_pair = (prev_node.token, new_id)
+ pair_positions.setdefault(new_left_pair, set()).add(node_a.prev)
+ new_count = pair_counts.get(new_left_pair, 0) + 1
+ pair_counts[new_left_pair] = new_count
+ heapq.heappush(heap, (-new_count, new_left_pair))
+
+ if node_a.next is not None:
+ next_node = all_nodes[node_a.next]
+ new_right_pair = (new_id, next_node.token)
+ pair_positions.setdefault(new_right_pair, set()).add(nid)
+ new_count = pair_counts.get(new_right_pair, 0) + 1
+ pair_counts[new_right_pair] = new_count
+ heapq.heappush(heap, (-new_count, new_right_pair))
+
+ next_id += 1
+ merge_num += 1
+
+ if verbose and (merge_num <= 20 or merge_num % 50 == 0 or merge_num == num_merges):
+ # Count remaining tokens (approximate)
+ print(
+ f" Merge {merge_num:4d}: ({best_pair[0]:3d},{best_pair[1]:3d}) -> {new_id:4d}"
+ f" count={best_count:6d}"
+ )
+
+ if verbose:
+ print(f"\nDone. Learned {len(merges)} merges.")
+ print(
+ f"Vocabulary: {BASE_VOCAB_SIZE} base + {len(merges)} merges = "
+ f"{BASE_VOCAB_SIZE + len(merges)} total"
+ )
+ return merges
+
+
+# ---------------------------------------------------------------------------
+# HuggingFace tokenizers backend (Rust, ~100x faster than pure Python)
+# ---------------------------------------------------------------------------
+
+def train_bpe_merges_hf(texts: list[str], num_merges: int = 1024, verbose: bool = True) -> list:
+ """
+ Train BPE merges using the HuggingFace `tokenizers` library (Rust backend).
+
+ ~100x faster than train_bpe_merges_fast: reduces 2-3 hours to ~2-5 minutes
+ on 32 cores for 100K FineWeb docs. Returns merges in the same format as
+ train_bpe_merges_fast: [((left_id, right_id), new_id), ...]
+
+ Requires: pip install tokenizers
+ """
+ try:
+ from tokenizers import Tokenizer, models, trainers, pre_tokenizers
+ except ImportError:
+ raise ImportError(
+ "HuggingFace tokenizers not installed. Run: pip install tokenizers\n"
+ "Falling back to pure-Python BPE is possible via train_bpe_merges_fast()."
+ )
+
+ import time as _time
+ import multiprocessing as mp
+ import tempfile, os
+
+ # Map base token IDs 0-(BASE_VOCAB_SIZE-1) to unique Unicode private-use chars.
+ # U+E000..U+E027 are valid single-codepoint, valid UTF-8, never in real text.
+ BASE_CHARS = [chr(0xE000 + i) for i in range(BASE_VOCAB_SIZE)]
+
+ # ------------------------------------------------------------------
+ # Step 1: Parallel encode texts → base token IDs
+ # ------------------------------------------------------------------
+ if verbose:
+ print(f"Encoding {len(texts)} texts with base BESE tokenizer...")
+
+ n_workers = min(mp.cpu_count(), 128)
+ chunk_size = max(1, len(texts) // n_workers)
+ chunks = [texts[i:i + chunk_size] for i in range(0, len(texts), chunk_size)]
+
+ t_enc = _time.time()
+ with mp.Pool(n_workers) as pool:
+ encoded_chunks = pool.map(_encode_texts_worker, chunks)
+ all_encoded = [doc for chunk in encoded_chunks for doc in chunk]
+ if verbose:
+ print(f" Parallel encoding: {n_workers} workers, {_time.time() - t_enc:.1f}s")
+
+ # ------------------------------------------------------------------
+ # Step 2: Write corpus to a temp file — one doc per line, no spaces.
+ # Each base token becomes one Unicode private-use char.
+ # Whitespace pre-tokenizer treats each line (= no-space sequence) as
+ # one "word", so BPE merges happen freely within each document.
+ # ------------------------------------------------------------------
+ if verbose:
+ print(" Writing corpus temp file...")
+ t_write = _time.time()
+ fd, tmp_path = tempfile.mkstemp(suffix=".txt")
+ try:
+ with os.fdopen(fd, "w", encoding="utf-8") as f:
+ for doc_tokens in all_encoded:
+ f.write("".join(BASE_CHARS[t] for t in doc_tokens) + "\n")
+ if verbose:
+ size_mb = os.path.getsize(tmp_path) / 1e6
+ print(f" Corpus file: {size_mb:.0f} MB ({_time.time() - t_write:.1f}s)")
+
+ # ------------------------------------------------------------------
+ # Step 3: Train BPE with HF tokenizers (Rust, multi-threaded)
+ # ------------------------------------------------------------------
+ if verbose:
+ print(f" Training BPE ({num_merges} merges) with HuggingFace tokenizers...")
+ t_bpe = _time.time()
+
+ tokenizer = Tokenizer(models.BPE())
+ # Whitespace splits only on Unicode whitespace — our private-use chars
+ # are never whitespace, so each doc line becomes exactly one "word".
+ tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
+
+ trainer = trainers.BpeTrainer(
+ vocab_size=BASE_VOCAB_SIZE + num_merges,
+ min_frequency=2,
+ initial_alphabet=BASE_CHARS,
+ special_tokens=[],
+ show_progress=verbose,
+ )
+ tokenizer.train([tmp_path], trainer)
+
+ if verbose:
+ print(f" HF BPE done in {_time.time() - t_bpe:.1f}s")
+ finally:
+ try:
+ os.unlink(tmp_path)
+ except OSError:
+ pass
+
+ # ------------------------------------------------------------------
+ # Step 4: Extract merges and convert char-strings → token ID pairs.
+ # Save the model to a temp dir to read merges.txt (one merge per line:
+ # "str_a str_b"). Replay the sequence to reconstruct integer IDs.
+ # ------------------------------------------------------------------
+ merge_dir = tempfile.mkdtemp()
+ try:
+ tokenizer.model.save(merge_dir)
+ merges_path = os.path.join(merge_dir, "merges.txt")
+ with open(merges_path, encoding="utf-8") as mf:
+ raw_lines = mf.read().splitlines()
+ finally:
+ import shutil
+ shutil.rmtree(merge_dir, ignore_errors=True)
+
+ # Filter header lines (start with #) and parse "str_a str_b" pairs
+ hf_merges = []
+ for line in raw_lines:
+ if line.startswith("#") or not line.strip():
+ continue
+ parts = line.split(" ", 1)
+ if len(parts) == 2:
+ hf_merges.append((parts[0], parts[1]))
+
+ str_to_id: dict[str, int] = {c: i for i, c in enumerate(BASE_CHARS)}
+ result: list = []
+
+ for str_a, str_b in hf_merges:
+ id_a = str_to_id.get(str_a)
+ id_b = str_to_id.get(str_b)
+ if id_a is None or id_b is None:
+ continue # shouldn't happen with a clean corpus
+ new_id = BASE_VOCAB_SIZE + len(result)
+ str_to_id[str_a + str_b] = new_id
+ result.append(((id_a, id_b), new_id))
+ if len(result) >= num_merges:
+ break
+
+ if verbose:
+ print(f"\nDone. Learned {len(result)} merges.")
+ print(
+ f"Vocabulary: {BASE_VOCAB_SIZE} base + {len(result)} merges = "
+ f"{BASE_VOCAB_SIZE + len(result)} total"
+ )
+ return result
+
+
+# ---------------------------------------------------------------------------
+# Fast encoding using merge priority (hash-map based)
+# ---------------------------------------------------------------------------
+
+def _build_merge_priority(merges):
+ """Build a priority lookup: (a, b) -> (priority, new_id).
+ Lower priority = should be applied first."""
+ return {pair: (i, new_id) for i, (pair, new_id) in enumerate(merges)}
+
+
+def encode_fast(text: str, merges: list, _merge_priority=None) -> np.ndarray:
+ """
+ Encode text to BESE+BPE tokens using hash-map based merging.
+
+ Instead of scanning for each merge in order (O(merges * tokens)),
+ we use a linked list + priority queue approach:
+ 1. Start with base tokens in a linked list
+ 2. Find all adjacent pairs and their merge priorities
+ 3. Apply merges in priority order (lowest first)
+ 4. After each merge, check if new pairs can be merged
+
+ This is O(tokens * log(tokens)) instead of O(merges * tokens).
+ """
+ base_tokens = _text_to_base_tokens(text)
+ if not base_tokens or not merges:
+ return np.array(base_tokens, dtype=np.uint16)
+
+ if _merge_priority is None:
+ _merge_priority = _build_merge_priority(merges)
+
+ # Build doubly linked list
+ n = len(base_tokens)
+ tokens = list(base_tokens)
+ prev_arr = list(range(-1, n - 1)) # prev[i] = i-1
+ next_arr = list(range(1, n + 1)) # next[i] = i+1, n = sentinel
+ next_arr[-1] = n # sentinel
+
+ # Priority queue: (priority, position_id, pair_at_creation)
+ # pair_at_creation is used to validate the entry is still current
+ heap = []
+
+ # Initialize heap with all mergeable pairs
+ for i in range(n - 1):
+ pair = (tokens[i], tokens[i + 1])
+ if pair in _merge_priority:
+ pri, _ = _merge_priority[pair]
+ heapq.heappush(heap, (pri, i, pair))
+
+ while heap:
+ pri, pos, pair_check = heapq.heappop(heap)
+
+ # Validate: position must still have this exact pair
+ if tokens[pos] != pair_check[0]:
+ continue
+ nxt = next_arr[pos]
+ if nxt >= n or tokens[nxt] != pair_check[1]:
+ continue
+
+ _, new_id = _merge_priority[pair_check]
+
+ # Merge: pos takes new_id, nxt is removed
+ tokens[pos] = new_id
+ next_arr[pos] = next_arr[nxt]
+ if next_arr[nxt] < n:
+ prev_arr[next_arr[nxt]] = pos
+
+ # Check new left pair
+ if prev_arr[pos] >= 0:
+ left = prev_arr[pos]
+ new_pair = (tokens[left], new_id)
+ if new_pair in _merge_priority:
+ lp, _ = _merge_priority[new_pair]
+ heapq.heappush(heap, (lp, left, new_pair))
+
+ # Check new right pair
+ if next_arr[pos] < n:
+ right = next_arr[pos]
+ new_pair = (new_id, tokens[right])
+ if new_pair in _merge_priority:
+ rp, _ = _merge_priority[new_pair]
+ heapq.heappush(heap, (rp, pos, new_pair))
+
+ # Collect result
+ result = []
+ pos = 0
+ while pos < n:
+ result.append(tokens[pos])
+ nxt = next_arr[pos]
+ if nxt <= pos:
+ break # safety
+ pos = nxt
+
+ # Handle case where linked list walk doesn't start from a valid head
+ if not result:
+ # Find first valid node
+ for i in range(n):
+ if prev_arr[i] < 0 or i == 0:
+ pos = i
+ while pos < n:
+ result.append(tokens[pos])
+ pos = next_arr[pos]
+ break
+
+ return np.array(result, dtype=np.uint16)
+
+
+class FastBESEBPETokenizer:
+ """Fast BESE+BPE tokenizer using indexed merge operations."""
+
+ def __init__(self, merges=None):
+ self.merges = merges or []
+ self.pad_id = PAD_ID
+ self.bos_id = BOS_ID
+ self.eos_id = EOS_ID
+ self.unk_id = UNK_ID
+ self._merge_map = {pair: new_id for pair, new_id in self.merges}
+ self._merge_priority = _build_merge_priority(self.merges)
+ self._bpt = self._build_bpt()
+ self._decode_chains = {new_id: pair for pair, new_id in self.merges}
+
+ @property
+ def vocab_size(self):
+ return BASE_VOCAB_SIZE + len(self.merges)
+
+ def _build_bpt(self):
+ bpt = np.zeros(self.vocab_size, dtype=np.int16)
+ bpt[:BASE_VOCAB_SIZE] = BYTES_PER_TOKEN
+ merge_bpt = {i: int(BYTES_PER_TOKEN[i]) for i in range(BASE_VOCAB_SIZE)}
+ for pair, new_id in self.merges:
+ merge_bpt[new_id] = merge_bpt[pair[0]] + merge_bpt[pair[1]]
+ bpt[new_id] = merge_bpt[new_id]
+ return bpt
+
+ def encode(self, text: str) -> np.ndarray:
+ return encode_fast(text, self.merges, self._merge_priority)
+
+ def encode_batch(self, texts: list[str]) -> list[np.ndarray]:
+ return [self.encode(text) for text in texts]
+
+ def decode_token_to_base(self, token_id: int) -> list[int]:
+ if token_id < BASE_VOCAB_SIZE:
+ return [token_id]
+ if token_id in self._decode_chains:
+ left, right = self._decode_chains[token_id]
+ return self.decode_token_to_base(left) + self.decode_token_to_base(right)
+ return [UNK_ID]
+
+ def decode(self, token_ids: list[int]) -> str:
+ base_tokens = []
+ for tid in token_ids:
+ base_tokens.extend(self.decode_token_to_base(tid))
+ result = []
+ i = 0
+ while i < len(base_tokens):
+ tid = base_tokens[i]
+ if GROUP_START <= tid < GROUP_START + len(GROUPS):
+ if i + 1 < len(base_tokens):
+ key = (tid, base_tokens[i + 1])
+ result.append(DECODE_TABLE.get(key, "?"))
+ i += 2
+ continue
+ key = (tid,)
+ if key in DECODE_TABLE:
+ result.append(DECODE_TABLE[key])
+ elif tid not in (PAD_ID, BOS_ID, EOS_ID, UNK_ID):
+ result.append("?")
+ i += 1
+ return "".join(result)
+
+ def get_bytes_per_token_lut(self) -> np.ndarray:
+ return self._bpt.copy()
+
+ def save(self, path):
+ """Save in same format as BESEBPETokenizer for compatibility."""
+ path = Path(path)
+ path.parent.mkdir(parents=True, exist_ok=True)
+ payload = {
+ "tokenizer_type": "bese_bpe",
+ "version": 2,
+ "base_vocab_size": BASE_VOCAB_SIZE,
+ "num_merges": len(self.merges),
+ "vocab_size": self.vocab_size,
+ "single_letters": SINGLE_LETTERS,
+ "groups": GROUPS,
+ "merges": [[list(pair), new_id] for pair, new_id in self.merges],
+ }
+ path.write_text(json.dumps(payload, separators=(",", ":")) + "\n", encoding="utf-8")
+
+ @classmethod
+ def load(cls, path):
+ path = Path(path)
+ payload = json.loads(path.read_text(encoding="utf-8"))
+ merges = [(tuple(pair), new_id) for pair, new_id in payload["merges"]]
+ return cls(merges=merges)
+
+ def build_luts_for_training(self, device=None):
+ """Build lookup tables compatible with train_gpt.py eval_val function."""
+ import torch
+
+ vs = self.vocab_size
+ has_leading_space = np.zeros(vs, dtype=np.bool_)
+ is_boundary = np.zeros(vs, dtype=np.bool_)
+ is_boundary[PAD_ID] = True
+ is_boundary[BOS_ID] = True
+ is_boundary[EOS_ID] = True
+ is_boundary[UNK_ID] = True
+ kwargs = {"device": device} if device is not None else {}
+ return (
+ torch.tensor(self._bpt.copy(), dtype=torch.int16, **kwargs),
+ torch.tensor(has_leading_space, dtype=torch.bool, **kwargs),
+ torch.tensor(is_boundary, dtype=torch.bool, **kwargs),
+ )
+
+
+# ---------------------------------------------------------------------------
+# Self-test
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+ import time
+
+ sample = [
+ "The cat sat on the mat.",
+ "Hello world!",
+ "The quick brown fox jumps over the lazy dog.",
+ "Parameter Golf is a challenge to train the best language model.",
+ "BESE uses a 40-token structured alphabet with BPE merges on top.",
+ ] * 200
+
+ print("=== Fast BPE Training ===")
+ t0 = time.time()
+ merges = train_bpe_merges_fast(sample, num_merges=100, verbose=True)
+ t1 = time.time()
+ print(f"Training took {t1-t0:.2f}s")
+
+ print("\n=== Fast Encoding ===")
+ tok = FastBESEBPETokenizer(merges=merges)
+
+ # Correctness check
+ test_texts = [
+ "The cat sat on the mat.",
+ "Hello world!",
+ "Testing 123 with special chars: é, ñ, ü",
+ "Multiple\nlines\nof\ntext.",
+ ]
+ for text in test_texts:
+ enc = tok.encode(text)
+ dec = tok.decode(enc.tolist())
+ bpt = tok.get_bytes_per_token_lut()
+ tb = int(sum(bpt[t] for t in enc))
+ ub = len(text.encode("utf-8"))
+ status = "OK" if tb == ub else "FAIL"
+ print(f' [{status}] "{text[:40]}..." -> {len(enc)} tokens, bytes {tb}/{ub}')
+
+ # Speed benchmark
+ print("\n=== Speed Benchmark ===")
+ big_text = " ".join(sample)
+ t0 = time.time()
+ for _ in range(10):
+ tok.encode(big_text)
+ t1 = time.time()
+ print(f"Encoded {len(big_text)} chars x 10 in {t1-t0:.2f}s ({len(big_text)*10/(t1-t0):.0f} chars/sec)")
diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/mamba3_ssd.py b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/mamba3_ssd.py
new file mode 100644
index 0000000000..3d528b87c4
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/mamba3_ssd.py
@@ -0,0 +1,573 @@
+"""
+Mamba-3 SSD (Structured State Space Duality) hybrid for Parameter Golf.
+
+Architecture: 6 Mamba-3 blocks + 2 Attention blocks (positions 2, 5).
+Uses fused Triton kernels from mamba-ssm when available, falls back to
+pure PyTorch (segsum/einsum) otherwise.
+
+Based on mamba3-minimal (https://github.com/VikramKarLex/mamba3-minimal).
+Informed by PR #1644 (best Mamba-3 at 1.1473 BPB) and PR #1355 ablations.
+"""
+from __future__ import annotations
+
+import math
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+try:
+ from einops import rearrange
+except ImportError:
+ raise ImportError("einops is required: pip install einops")
+
+# Fused Triton kernel from mamba-ssm — disabled due to multi-GPU segfault after ~100 steps.
+# The pure PyTorch fallback (ngroups=1, d_state=128) with the new architecture should
+# still be faster than v7 v1 (ngroups=16, d_state=64) due to fewer params + no depth recurrence.
+_HAS_MAMBA_KERNEL = False
+
+
+# ---------------------------------------------------------------------------
+# Core SSD helpers (pure PyTorch fallback)
+# ---------------------------------------------------------------------------
+
+def segsum(x: Tensor) -> Tensor:
+ """Stable cumulative sum for decay computation.
+
+ Computes a lower-triangular matrix where entry [i,j] = sum(x[j..i]).
+ Used for computing state decay within chunks.
+
+ Args:
+ x: (..., L) tensor
+ Returns:
+ (..., L, L) lower-triangular cumsum matrix
+ """
+ T = x.size(-1)
+ x = x[..., None].repeat(*([1] * (x.ndim - 1)), 1, T) # (..., T, T)
+ mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=-1)
+ x = x.masked_fill(~mask, 0)
+ x_segsum = torch.cumsum(x, dim=-2)
+ mask2 = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=0)
+ x_segsum = x_segsum.masked_fill(~mask2, -torch.inf)
+ return x_segsum
+
+
+def ssd_chunked(
+ x: Tensor,
+ A: Tensor,
+ B: Tensor,
+ C: Tensor,
+ chunk_size: int,
+ initial_states: Tensor | None = None,
+) -> Tensor:
+ """Structured State Space Duality — chunked parallel computation (fallback).
+
+ Args:
+ x: (batch, seq_len, heads, head_dim) — input after projection
+ A: (batch, seq_len, heads) — log decay rates (dt * A_param)
+ B: (batch, seq_len, ngroups, d_state) — input-to-state projection
+ C: (batch, seq_len, ngroups, d_state) — state-to-output projection
+ chunk_size: size of chunks for parallel computation
+ initial_states: optional (batch, heads, head_dim, d_state) initial hidden state
+
+ Returns:
+ (batch, seq_len, heads, head_dim) output
+ """
+ batch, seq_len, nheads, headdim = x.shape
+ ngroups = B.shape[2]
+ d_state = B.shape[-1]
+
+ # Broadcast B/C if ngroups < nheads
+ if ngroups < nheads:
+ repeat_factor = nheads // ngroups
+ B = B.repeat_interleave(repeat_factor, dim=2) # (b, l, nheads, d_state)
+ C = C.repeat_interleave(repeat_factor, dim=2)
+
+ # Pad sequence to multiple of chunk_size
+ pad_len = (chunk_size - seq_len % chunk_size) % chunk_size
+ if pad_len > 0:
+ x = F.pad(x, (0, 0, 0, 0, 0, pad_len))
+ A = F.pad(A, (0, 0, 0, pad_len))
+ B = F.pad(B, (0, 0, 0, 0, 0, pad_len))
+ C = F.pad(C, (0, 0, 0, 0, 0, pad_len))
+
+ # Reshape into chunks: (batch, n_chunks, chunk_size, ...)
+ x = rearrange(x, "b (c l) h p -> b c l h p", l=chunk_size)
+ A = rearrange(A, "b (c l) h -> b c l h", l=chunk_size)
+ B = rearrange(B, "b (c l) h n -> b c l h n", l=chunk_size)
+ C = rearrange(C, "b (c l) h n -> b c l h n", l=chunk_size)
+
+ # Transpose heads for einsum convenience: b h c l ...
+ A = rearrange(A, "b c l h -> b h c l")
+
+ # Step 1: Intra-chunk quadratic attention
+ L = torch.exp(segsum(A))
+
+ Y_diag = torch.einsum(
+ "bclhn, bcshn, bhcls, bcshp -> bclhp",
+ C, B, L, x
+ )
+
+ # Step 2: Per-chunk state accumulation
+ A_cumsum = torch.cumsum(A, dim=-1)
+ decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
+
+ states = torch.einsum(
+ "bclhn, bhcl, bclhp -> bchpn",
+ B, decay_states, x
+ )
+
+ # Step 3: Inter-chunk recurrence (with causality fix)
+ if initial_states is not None:
+ states = torch.cat([initial_states.unsqueeze(1), states], dim=1)
+
+ A_chunk_decay = A_cumsum[:, :, :, -1]
+ decay_chunk = torch.exp(segsum(F.pad(A_chunk_decay, (1, 0))))
+
+ if initial_states is not None:
+ c_init = states.shape[1]
+ new_states = torch.einsum(
+ "bhzc, bchpn -> bzhpn",
+ decay_chunk[:, :, 1:c_init + 1, :c_init],
+ states
+ )
+ else:
+ c = states.shape[1]
+ new_states = torch.einsum(
+ "bhzc, bchpn -> bzhpn",
+ decay_chunk[:, :, :c, 1:c + 1],
+ states
+ )
+
+ # Step 4: State-to-output
+ state_decay_out = torch.exp(A_cumsum)
+ Y_off = torch.einsum(
+ "bclhn, bchpn, bhcl -> bclhp",
+ C, new_states, state_decay_out
+ )
+
+ Y = Y_diag + Y_off
+ Y = rearrange(Y, "b c l h p -> b (c l) h p")
+
+ if pad_len > 0:
+ Y = Y[:, :seq_len]
+
+ return Y
+
+
+# ---------------------------------------------------------------------------
+# Mamba-3 Block
+# ---------------------------------------------------------------------------
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-5):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def forward(self, x: Tensor) -> Tensor:
+ return F.rms_norm(x, (x.size(-1),), self.weight, self.eps)
+
+
+class Mamba3Block(nn.Module):
+ """Mamba-3 block with SSD (Structured State Space Duality).
+
+ Uses fused Triton kernels from mamba-ssm if available (2-3x faster),
+ falls back to pure PyTorch (segsum/einsum) otherwise.
+
+ ngroups=1: all heads share B/C projections (matches reference Mamba-2,
+ confirmed optimal by PR #1644 ablations).
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ d_state: int = 128,
+ expand: int = 2,
+ headdim: int = 64,
+ chunk_size: int = 64,
+ ngroups: int = 1,
+ ):
+ super().__init__()
+ self.dim = dim
+ self.d_state = d_state
+ self.d_inner = expand * dim
+ self.nheads = self.d_inner // headdim
+ self.headdim = headdim
+ self.chunk_size = chunk_size
+ self.ngroups = ngroups
+
+ # Input projection: x → (z, x_proj, B, C, dt)
+ # z: gating signal (d_inner)
+ # x_proj: SSM input (d_inner)
+ # B: input-to-state (ngroups * d_state) — shared across heads
+ # C: state-to-output (ngroups * d_state) — shared across heads
+ # dt: timestep (nheads)
+ d_proj = self.d_inner * 2 + ngroups * d_state * 2 + self.nheads
+ self.in_proj = nn.Linear(dim, d_proj, bias=False)
+ self.out_proj = nn.Linear(self.d_inner, dim, bias=False)
+
+ # SSM parameters
+ self.D = nn.Parameter(torch.ones(self.nheads))
+ self.dt_bias = nn.Parameter(torch.randn(self.nheads) * 0.1)
+ self.A_log = nn.Parameter(torch.log(0.5 + torch.rand(self.nheads) * 0.5))
+
+ # Normalization for B and C projections
+ self.B_norm = RMSNorm(d_state)
+ self.C_norm = RMSNorm(d_state)
+
+ # Pre-normalization
+ self.norm = nn.LayerNorm(dim)
+
+ self._init_weights()
+
+ def _init_weights(self):
+ nn.init.orthogonal_(self.in_proj.weight, gain=1.0)
+ nn.init.zeros_(self.out_proj.weight)
+
+ def forward(self, u: Tensor) -> Tensor:
+ """
+ Args:
+ u: (batch, seq_len, dim)
+ Returns:
+ (batch, seq_len, dim)
+ """
+ batch, seq_len, dim = u.shape
+
+ # Normalize input
+ u_normed = self.norm(u)
+
+ # Project
+ proj = self.in_proj(u_normed)
+
+ # Split projections
+ d_inner = self.d_inner
+ nheads = self.nheads
+ ngroups = self.ngroups
+ d_state = self.d_state
+
+ z = proj[..., :d_inner]
+ x = proj[..., d_inner:d_inner * 2]
+ B_flat = proj[..., d_inner * 2:d_inner * 2 + ngroups * d_state]
+ C_flat = proj[..., d_inner * 2 + ngroups * d_state:d_inner * 2 + ngroups * d_state * 2]
+ dt = proj[..., -nheads:]
+
+ # Reshape into multi-head/group form
+ B = B_flat.reshape(batch, seq_len, ngroups, d_state)
+ C = C_flat.reshape(batch, seq_len, ngroups, d_state)
+ B = self.B_norm(B)
+ C = self.C_norm(C)
+ x = x.reshape(batch, seq_len, nheads, self.headdim)
+ z = z.reshape(batch, seq_len, nheads, self.headdim)
+
+ if _HAS_MAMBA_KERNEL:
+ # Fused Triton kernel: handles softplus(dt+dt_bias), dt*A, chunking,
+ # inter-chunk recurrence, D skip connection, SiLU gating — all in one kernel
+ A = -torch.exp(self.A_log.float()) # (nheads,) 1D, negative
+ y = mamba_chunk_scan_combined(
+ x.contiguous(),
+ dt.contiguous(),
+ A, B, C,
+ chunk_size=self.chunk_size,
+ D=self.D,
+ z=z.contiguous(),
+ dt_bias=self.dt_bias,
+ dt_softplus=True,
+ )
+ else:
+ # Pure PyTorch fallback
+ dt_proc = F.softplus(dt + self.dt_bias)
+ A = -torch.exp(self.A_log.float())
+ dA = dt_proc * A
+ y = ssd_chunked(x, dA, B, C, self.chunk_size)
+ y = y + x * self.D[None, None, :, None]
+ # Merge heads, gate with SiLU, output
+ y = y.reshape(batch, seq_len, d_inner)
+ z_flat = z.reshape(batch, seq_len, d_inner)
+ y = y * F.silu(z_flat)
+ return u + self.out_proj(y)
+
+ # Merge heads and output projection (kernel path)
+ y = y.reshape(batch, seq_len, d_inner)
+ return u + self.out_proj(y)
+
+
+# ---------------------------------------------------------------------------
+# Attention Block (reused from existing GPT, simplified)
+# ---------------------------------------------------------------------------
+
+class Rotary(nn.Module):
+ """RoPE (Rotary Position Embedding)."""
+ def __init__(self, dim: int, base: float = 10000.0, train_seq_len: int = 1024, rope_dims: int = 0):
+ super().__init__()
+ self.dim = dim
+ self.base = base
+ self.train_seq_len = train_seq_len
+ self.rope_dims = rope_dims if rope_dims > 0 else dim
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.rope_dims, 2, dtype=torch.float32) / self.rope_dims))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self._seq_len_cached = 0
+ self._cos_cached: Tensor | None = None
+ self._sin_cached: Tensor | None = None
+
+ def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]:
+ if (
+ self._cos_cached is None
+ or self._sin_cached is None
+ or self._seq_len_cached != seq_len
+ or self._cos_cached.device != device
+ ):
+ rd = self.rope_dims
+ if seq_len > self.train_seq_len:
+ scale = seq_len / self.train_seq_len
+ new_base = self.base * (scale ** (rd / (rd - 2)))
+ inv_freq = 1.0 / (new_base ** (torch.arange(0, rd, 2, dtype=torch.float32, device=device) / rd))
+ else:
+ inv_freq = self.inv_freq.to(device)
+ t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
+ freqs = torch.outer(t, inv_freq)
+ self._cos_cached = freqs.cos()[None, :, None, :]
+ self._sin_cached = freqs.sin()[None, :, None, :]
+ self._seq_len_cached = seq_len
+ return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype)
+
+
+def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor, rope_dims: int = 0) -> Tensor:
+ """Apply rotary position embeddings."""
+ if rope_dims > 0 and rope_dims < x.size(-1):
+ x_rope, x_pass = x[..., :rope_dims], x[..., rope_dims:]
+ half = rope_dims // 2
+ x1, x2 = x_rope[..., :half], x_rope[..., half:]
+ x_rope = torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1)
+ return torch.cat((x_rope, x_pass), dim=-1)
+ half = x.size(-1) // 2
+ x1, x2 = x[..., :half], x[..., half:]
+ return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1)
+
+
+class AttentionBlock(nn.Module):
+ """Standard causal self-attention block for the hybrid architecture."""
+
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int = 8,
+ num_kv_heads: int = 4,
+ mlp_mult: float = 3.0,
+ rope_base: float = 10000.0,
+ qk_gain_init: float = 5.25,
+ rope_dims: int = 16,
+ ):
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ self.num_kv_heads = num_kv_heads
+ self.head_dim = dim // num_heads
+ mlp_dim = int(dim * mlp_mult)
+
+ self.c_q = nn.Linear(dim, dim, bias=False)
+ self.c_k = nn.Linear(dim, num_kv_heads * self.head_dim, bias=False)
+ self.c_v = nn.Linear(dim, num_kv_heads * self.head_dim, bias=False)
+ self.c_proj = nn.Linear(dim, dim, bias=False)
+
+ self.mlp_fc = nn.Linear(dim, mlp_dim, bias=False)
+ self.mlp_proj = nn.Linear(mlp_dim, dim, bias=False)
+
+ self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32))
+
+ self.rope_dims = rope_dims
+ self.rotary = Rotary(self.head_dim, base=rope_base, train_seq_len=4096, rope_dims=rope_dims)
+
+ self.attn_norm = nn.LayerNorm(dim)
+ self.mlp_norm = nn.LayerNorm(dim)
+
+ self._init_weights()
+
+ def _init_weights(self):
+ for name, p in self.named_parameters():
+ if p.ndim == 2:
+ if "proj" in name:
+ nn.init.zeros_(p)
+ else:
+ nn.init.orthogonal_(p, gain=1.0)
+
+ def forward(self, x: Tensor) -> Tensor:
+ bsz, seqlen, dim = x.shape
+
+ h = self.attn_norm(x)
+ q = self.c_q(h).reshape(bsz, seqlen, self.num_heads, self.head_dim)
+ k = self.c_k(h).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim)
+ v = self.c_v(h).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim)
+
+ q = F.rms_norm(q, (q.size(-1),))
+ k = F.rms_norm(k, (k.size(-1),))
+ cos, sin = self.rotary(seqlen, x.device, q.dtype)
+ q = apply_rotary_emb(q, cos, sin, self.rope_dims)
+ k = apply_rotary_emb(k, cos, sin, self.rope_dims)
+ q = q * self.q_gain.to(dtype=q.dtype)[None, None, :, None]
+
+ try:
+ from flash_attn_interface import flash_attn_func as fa3
+ y = fa3(q, k, v, causal=True)
+ except ImportError:
+ q_t = q.transpose(1, 2)
+ rep = self.num_heads // self.num_kv_heads
+ k_t = k.transpose(1, 2).repeat_interleave(rep, dim=1)
+ v_t = v.transpose(1, 2).repeat_interleave(rep, dim=1)
+ y = F.scaled_dot_product_attention(q_t, k_t, v_t, is_causal=True)
+ y = y.transpose(1, 2)
+
+ y = y.reshape(bsz, seqlen, dim)
+ x = x + self.c_proj(y)
+
+ h = self.mlp_norm(x)
+ x = x + self.mlp_proj(F.silu(self.mlp_fc(h)))
+
+ return x
+
+
+# ---------------------------------------------------------------------------
+# Hybrid Model: 6 Mamba-3 + 2 Attention
+# ---------------------------------------------------------------------------
+
+class HybridMambaGPT(nn.Module):
+ """Mamba-3 + Attention hybrid for Parameter Golf.
+
+ Architecture:
+ - 6 Mamba-3 blocks (linear O(n) scaling, ngroups=1)
+ - 2 Attention blocks at configurable positions (default: 2, 5)
+ - No depth recurrence (hurts SSMs by -69 mBPB per PR #1355)
+ - BESE 288 vocab with tied embeddings
+ """
+
+ def __init__(
+ self,
+ vocab_size: int = 288,
+ num_layers: int = 8,
+ model_dim: int = 512,
+ d_state: int = 128,
+ expand: int = 2,
+ headdim: int = 64,
+ chunk_size: int = 64,
+ attn_pos: int | list[int] = None,
+ num_heads: int = 8,
+ num_kv_heads: int = 4,
+ mlp_mult: float = 3.0,
+ rope_base: float = 10000.0,
+ qk_gain_init: float = 5.25,
+ rope_dims: int = 16,
+ logit_softcap: float = 30.0,
+ tied_embed_init_std: float = 0.005,
+ ngroups: int = 1,
+ # Legacy depth recurrence params (kept for API compat, but disabled)
+ depth_recurrence_start: int = 0,
+ depth_recurrence_end: int = 0,
+ depth_recurrence_loops: int = 1,
+ depth_recurrence_activation_frac: float = 1.0,
+ ):
+ super().__init__()
+ self.vocab_size = vocab_size
+ self.model_dim = model_dim
+ self.num_layers = num_layers
+ self.logit_softcap = logit_softcap
+ self.tied_embed_init_std = tied_embed_init_std
+
+ # Attention positions: default [2, 5] for 8-layer model
+ if attn_pos is None:
+ self.attn_positions = [2, 5]
+ elif isinstance(attn_pos, int):
+ self.attn_positions = [attn_pos]
+ else:
+ self.attn_positions = list(attn_pos)
+ # Legacy single attn_pos for training script compat
+ self.attn_pos = self.attn_positions[0]
+
+ # Depth recurrence config (kept for API compat, defaults to disabled)
+ self._rec_start = depth_recurrence_start
+ self._rec_end = depth_recurrence_end
+ self._rec_target_loops = depth_recurrence_loops
+ self._rec_activation_frac = depth_recurrence_activation_frac
+ self._rec_loops = 1
+ self._training_progress = 0.0
+
+ # Token embedding (tied with lm_head)
+ self.tok_emb = nn.Embedding(vocab_size, model_dim)
+ nn.init.normal_(self.tok_emb.weight, mean=0.0, std=tied_embed_init_std)
+
+ # SmearGate for temporal smoothing
+ self.smear_gate = nn.Parameter(torch.zeros(model_dim, dtype=torch.float32))
+
+ # Build layers: Mamba at most positions, Attention at attn_positions
+ self.layers = nn.ModuleList()
+ for i in range(num_layers):
+ if i in self.attn_positions:
+ self.layers.append(AttentionBlock(
+ dim=model_dim,
+ num_heads=num_heads,
+ num_kv_heads=num_kv_heads,
+ mlp_mult=mlp_mult,
+ rope_base=rope_base,
+ qk_gain_init=qk_gain_init,
+ rope_dims=rope_dims,
+ ))
+ else:
+ self.layers.append(Mamba3Block(
+ dim=model_dim,
+ d_state=d_state,
+ expand=expand,
+ headdim=headdim,
+ chunk_size=chunk_size,
+ ngroups=ngroups,
+ ))
+
+ # Final norm
+ self.final_norm = nn.LayerNorm(model_dim)
+
+ # No separate lm_head — tied with tok_emb
+ self.lm_head = None
+
+ def _smear(self, x: Tensor) -> Tensor:
+ """Temporal smoothing gate."""
+ g = torch.sigmoid(self.smear_gate.to(dtype=x.dtype))[None, None, :]
+ x_prev = torch.cat([torch.zeros_like(x[:, :1]), x[:, :-1]], dim=1)
+ return (1 - g) * x + g * x_prev
+
+ def _run_layers(self, x: Tensor) -> Tensor:
+ """Run all layers sequentially (no depth recurrence for SSMs)."""
+ # Depth recurrence support (disabled by default for Mamba)
+ if self._training_progress >= self._rec_activation_frac:
+ self._rec_loops = self._rec_target_loops
+
+ for i, layer in enumerate(self.layers):
+ if self._rec_start <= i <= self._rec_end and self._rec_loops > 1:
+ for _ in range(self._rec_loops):
+ x = layer(x)
+ else:
+ x = layer(x)
+ return x
+
+ def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor:
+ """Forward pass with cross-entropy loss."""
+ x = self.tok_emb(input_ids)
+ x = F.rms_norm(x, (x.size(-1),))
+ x = self._smear(x)
+ x = self._run_layers(x)
+ x = self.final_norm(x)
+
+ logits = F.linear(x, self.tok_emb.weight) # tied embeddings
+ logits = self.logit_softcap * torch.tanh(logits / self.logit_softcap)
+
+ targets = target_ids.reshape(-1)
+ logits_flat = logits.reshape(-1, logits.size(-1))
+ return F.cross_entropy(logits_flat.float(), targets, reduction="mean")
+
+ def forward_logits(self, input_ids: Tensor) -> Tensor:
+ """Return logits without computing loss (for eval/TTT)."""
+ x = self.tok_emb(input_ids)
+ x = F.rms_norm(x, (x.size(-1),))
+ x = self._smear(x)
+ x = self._run_layers(x)
+ x = self.final_norm(x)
+
+ logits = F.linear(x, self.tok_emb.weight)
+ logits = self.logit_softcap * torch.tanh(logits / self.logit_softcap)
+ return logits
diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/requirements.txt b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/requirements.txt
new file mode 100644
index 0000000000..30ac27f76e
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/requirements.txt
@@ -0,0 +1,3 @@
+einops>=0.8.0
+sentencepiece>=0.2.0
+flash-attn>=2.7.0
diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/submission.json b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/submission.json
new file mode 100644
index 0000000000..b1b29ce20b
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/submission.json
@@ -0,0 +1,24 @@
+{
+ "name": "BESE + Mamba-3 SSD Hybrid",
+ "track": "non_record_16mb",
+ "author": "Omer Bese",
+ "github_id": "mrbese",
+ "val_bpb": 1.3571,
+ "date": "2026-04-16",
+ "summary": "BESE 288-vocab byte-level tokenizer + Mamba-3 SSD / Attention hybrid (6 Mamba + 2 Attention). First SSM submission using a custom sub-byte tokenizer. ngroups=1, d_state=128, no depth recurrence. INT6 + LZMA + sliding window eval with n-gram tilt.",
+ "hardware": "8xH100 80GB SXM",
+ "training_time_seconds": 600,
+ "compressed_model_bytes": 7452680,
+ "code_bytes": 162208,
+ "total_artifact_bytes": 7614888,
+ "vocab_size": 288,
+ "num_layers": 8,
+ "model_dim": 512,
+ "tokenizer": "BESE+BPE (custom 40-base + 248 merges, 288 vocab)",
+ "architecture": "Mamba-3 SSD hybrid (6 Mamba + 2 Attention at positions 2,5)",
+ "d_state": 128,
+ "expand": 2,
+ "ngroups": 1,
+ "headdim": 64,
+ "notes": "Non-record SSM track submission. First combination of a custom byte-level tokenizer (BESE) with Mamba-3 SSD. BESE's 288 vocab halves the embedding table vs SP1024 (saving ~3-4 MB), funding model capacity. Artifact is 7.56 MB — half the 16 MB limit. Architecture informed by PR #1644 (best Mamba-3) and PR #1355 (SSM ablations). Pure PyTorch SSD fallback used due to multi-GPU Triton kernel stability issues. Additional runs with dim=576 achieved 1.3415 raw BPB (pre-quantization) at 8.42 MB, demonstrating headroom for further optimization."
+}
diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/tokenizer.json b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/tokenizer.json
new file mode 100644
index 0000000000..431f12a5dd
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/tokenizer.json
@@ -0,0 +1 @@
+{"tokenizer_type":"bese_bpe","version":2,"base_vocab_size":40,"num_merges":248,"vocab_size":288,"single_letters":"etaoinsrhdl","groups":["ufbz","cwvj","mykx","gpq"],"merges":[[[4,23],40],[[16,19],41],[[15,19],42],[[23,5],43],[[17,19],44],[[18,20],45],[[15,20],46],[[18,19],47],[[23,6],48],[[8,9],49],[[17,20],50],[[16,20],51],[[43,12],52],[[10,23],53],[[15,21],54],[[4,11],55],[[5,23],56],[[7,9],57],[[13,23],58],[[16,21],59],[[7,11],60],[[4,9],61],[[52,40],62],[[7,42],63],[[29,29],64],[[6,9],65],[[11,4],66],[[17,21],67],[[50,23],68],[[6,5],69],[[49,47],70],[[25,23],71],[[8,5],72],[[24,23],73],[[7,23],74],[[6,11],75],[[48,9],76],[[10,5],77],[[6,14],78],[[7,46],79],[[4,10],80],[[8,41],81],[[7,44],82],[[5,12],83],[[23,41],84],[[64,64],85],[[8,57],86],[[23,10],87],[[76,58],88],[[8,14],89],[[8,10],90],[[4,14],91],[[23,51],92],[[43,74],93],[[4,13],94],[[23,45],95],[[4,58],96],[[46,60],97],[[23,44],98],[[41,12],99],[[11,7],100],[[8,53],101],[[23,54],102],[[6,10],103],[[50,63],104],[[23,12],105],[[14,4],106],[[61,5],107],[[7,51],108],[[23,13],109],[[4,41],110],[[7,14],111],[[6,41],112],[[42,10],113],[[23,46],114],[[59,40],115],[[42,11],116],[[17,22],117],[[4,53],118],[[8,44],119],[[6,44],120],[[6,56],121],[[8,13],122],[[59,55],123],[[8,47],124],[[4,44],125],[[23,49],126],[[42,9],127],[[4,5],128],[[72,12],129],[[8,11],130],[[16,22],131],[[69,86],132],[[51,12],133],[[70,23],134],[[6,13],135],[[14,7],136],[[14,40],137],[[23,79],138],[[7,45],139],[[14,68],140],[[29,23],141],[[12,6],142],[[8,46],143],[[6,50],144],[[48,14],145],[[4,117],146],[[43,7],147],[[89,14],148],[[40,83],149],[[5,55],150],[[42,14],151],[[8,56],152],[[48,11],153],[[23,47],154],[[66,10],155],[[6,47],156],[[7,10],157],[[9,4],158],[[6,54],159],[[52,121],160],[[9,7],161],[[51,129],162],[[63,11],163],[[42,5],164],[[11,8],165],[[76,13],166],[[40,5],167],[[23,14],168],[[6,67],169],[[40,10],170],[[61,56],171],[[23,104],172],[[41,82],173],[[15,22],174],[[124,12],175],[[18,21],176],[[42,45],177],[[45,100],178],[[52,4],179],[[6,53],180],[[40,79],181],[[4,75],182],[[45,14],183],[[11,82],184],[[6,49],185],[[78,14],186],[[176,42],187],[[23,66],188],[[10,12],189],[[42,44],190],[[63,9],191],[[41,57],192],[[41,5],193],[[65,13],194],[[10,73],195],[[7,7],196],[[10,71],197],[[52,101],198],[[6,45],199],[[4,77],200],[[8,59],201],[[55,23],202],[[4,73],203],[[10,48],204],[[24,62],205],[[55,40],206],[[110,5],207],[[8,77],208],[[54,4],209],[[6,68],210],[[4,59],211],[[4,45],212],[[65,58],213],[[4,71],214],[[53,79],215],[[23,77],216],[[45,55],217],[[153,40],218],[[7,5],219],[[63,14],220],[[64,53],221],[[80,10],222],[[32,30],223],[[92,12],224],[[8,115],225],[[40,51],226],[[61,41],227],[[91,14],228],[[42,13],229],[[7,83],230],[[23,29],231],[[65,5],232],[[43,11],233],[[65,41],234],[[47,12],235],[[23,97],236],[[62,10],237],[[45,60],238],[[48,53],239],[[61,13],240],[[25,88],241],[[40,41],242],[[48,23],243],[[63,5],244],[[40,44],245],[[54,42],246],[[8,45],247],[[112,67],248],[[40,49],249],[[145,14],250],[[8,78],251],[[7,123],252],[[75,5],253],[[92,129],254],[[63,56],255],[[23,57],256],[[7,41],257],[[46,184],258],[[40,13],259],[[47,11],260],[[142,115],261],[[51,60],262],[[41,14],263],[[10,4],264],[[7,13],265],[[10,88],266],[[161,56],267],[[63,235],268],[[220,58],269],[[9,108],270],[[53,49],271],[[14,8],272],[[4,6],273],[[104,11],274],[[4,46],275],[[13,55],276],[[10,74],277],[[40,6],278],[[42,41],279],[[48,5],280],[[48,54],281],[[11,6],282],[[48,10],283],[[84,12],284],[[56,79],285],[[112,5],286],[[4,123],287]]}
diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_gpt.py b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_gpt.py
new file mode 100644
index 0000000000..57038cad0e
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_gpt.py
@@ -0,0 +1,2253 @@
+from __future__ import annotations
+import copy
+import glob
+import io
+import lzma
+import math
+import os
+import random
+import subprocess
+import sys
+import time
+import uuid
+from pathlib import Path
+
+# Fix Triton JIT cache race under torchrun: give each rank its own cache dir.
+# Must be set BEFORE any Triton import (mamba-ssm triggers Triton on import).
+_local_rank = os.environ.get("LOCAL_RANK", "0")
+os.environ.setdefault("TRITON_CACHE_DIR", f"/tmp/triton_cache_rank{_local_rank}")
+
+import numpy as np
+import sentencepiece as spm
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import Tensor, nn
+from flash_attn_interface import flash_attn_func as flash_attn_3_func
+
+# BESE tokenizer support: add tokenizer root to sys.path if configured
+_BESE_TOK_ROOT = os.environ.get("BESE_TOKENIZER_ROOT", "")
+if _BESE_TOK_ROOT:
+ sys.path.insert(0, _BESE_TOK_ROOT)
+
+# Mamba-3 hybrid support
+_MODEL_TYPE = os.environ.get("MODEL_TYPE", "transformer")
+if _MODEL_TYPE == "mamba_hybrid":
+ from mamba3_ssd import HybridMambaGPT
+
+class Hyperparameters:
+ data_path = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024")
+ train_files = os.path.join(data_path, "fineweb_train_*.bin")
+ val_files = os.path.join(data_path, "fineweb_val_*.bin")
+ tokenizer_path = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model")
+ run_id = os.environ.get("RUN_ID", str(uuid.uuid4()))
+ seed = int(os.environ.get("SEED", 1337))
+ val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288))
+ val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 4000))
+ train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 500))
+ iterations = int(os.environ.get("ITERATIONS", 20000))
+ warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 3500))
+ warmup_steps = int(os.environ.get("WARMUP_STEPS", 20))
+ train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 786_432))
+ train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 2048))
+ eval_seq_len = int(os.environ.get("EVAL_SEQ_LEN", 2048))
+ max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0))
+ qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 5.0))
+ vocab_size = int(os.environ.get("VOCAB_SIZE", 1024))
+ num_layers = int(os.environ.get("NUM_LAYERS", 11))
+ num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4))
+ model_dim = int(os.environ.get("MODEL_DIM", 512))
+ num_heads = int(os.environ.get("NUM_HEADS", 8))
+ mlp_mult = float(os.environ.get("MLP_MULT", 3.0))
+ tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1")))
+ rope_base = float(os.environ.get("ROPE_BASE", 10000.0))
+ logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0))
+ embed_lr = float(os.environ.get("EMBED_LR", 0.6))
+ head_lr = float(os.environ.get("HEAD_LR", 0.008))
+ tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.035))
+ tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005))
+ matrix_lr = float(os.environ.get("MATRIX_LR", 0.022))
+ scalar_lr = float(os.environ.get("SCALAR_LR", 0.025))
+ muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.99))
+ muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5))
+ muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.92))
+ muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 1500))
+ beta1 = float(os.environ.get("BETA1", 0.9))
+ beta2 = float(os.environ.get("BETA2", 0.95))
+ adam_eps = float(os.environ.get("ADAM_EPS", 1e-8))
+ grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.3))
+ eval_stride = int(os.environ.get("EVAL_STRIDE", 64))
+ mtp_num_heads = int(os.environ.get("MTP_NUM_HEADS", 0))
+ mtp_loss_weight = float(os.environ.get("MTP_LOSS_WEIGHT", 0.2))
+ swa_enabled = bool(int(os.environ.get("SWA_ENABLED", "1")))
+ swa_every = int(os.environ.get("SWA_EVERY", 50))
+ lawa_enabled = bool(int(os.environ.get("LAWA_ENABLED", "0")))
+ lawa_k = int(os.environ.get("LAWA_K", 10))
+ lawa_freq = int(os.environ.get("LAWA_FREQ", 100))
+ muon_wd = float(os.environ.get("MUON_WD", 0.095))
+ adam_wd = float(os.environ.get("ADAM_WD", 0.095))
+ qat_enabled = bool(int(os.environ.get("QAT_ENABLED", "0")))
+ bigram_vocab_size = int(os.environ.get("BIGRAM_VOCAB_SIZE", 2048))
+ bigram_dim = int(os.environ.get("BIGRAM_DIM", 128))
+ xsa_last_n = int(os.environ.get("XSA_LAST_N", 4))
+ rope_dims = int(os.environ.get("ROPE_DIMS", 16))
+ ln_scale = bool(int(os.environ.get("LN_SCALE", "1")))
+ dtg_enabled = bool(int(os.environ.get("DTG_ENABLED", "0")))
+ late_qat_threshold = float(os.environ.get("LATE_QAT_THRESHOLD", 0.15))
+ ve_enabled = bool(int(os.environ.get("VE_ENABLED", "1")))
+ ve_dim = int(os.environ.get("VE_DIM", 128))
+ ve_layers = os.environ.get("VE_LAYERS", "9,10")
+ gated_attention = bool(int(os.environ.get("GATED_ATTENTION", "0")))
+ value_residual = bool(int(os.environ.get("VALUE_RESIDUAL", "0")))
+ # v5: N-gram tilt at eval time
+ ngram_tilt_enabled = bool(int(os.environ.get("NGRAM_TILT_ENABLED", "0")))
+ ngram_tilt_beta = float(os.environ.get("NGRAM_TILT_BETA", 0.3))
+ ngram_tilt_max_n = int(os.environ.get("NGRAM_TILT_MAX_N", 4))
+ ngram_prior_path = os.environ.get("NGRAM_PRIOR_PATH", "")
+
+ # v8: Noisy QAT (Gaussian noise calibrated to INT6 quantization error)
+ noisy_qat_enabled = bool(int(os.environ.get("NOISY_QAT_ENABLED", "0")))
+ noisy_qat_activation_frac = float(os.environ.get("NOISY_QAT_ACTIVATION_FRAC", 0.20))
+ noisy_qat_clip_range = int(os.environ.get("NOISY_QAT_CLIP_RANGE", 31))
+
+ # v8: Bigram prior (frozen log-prob matrix as logit bias)
+ bigram_prior_enabled = bool(int(os.environ.get("BIGRAM_PRIOR_ENABLED", "0")))
+ bigram_prior_path = os.environ.get("BIGRAM_PRIOR_PATH", "")
+
+ # v6.1: Legal TTT (test-time training) at eval
+ ttt_enabled = bool(int(os.environ.get("TTT_ENABLED", "0")))
+ ttt_lr = float(os.environ.get("TTT_LR", 0.005))
+ ttt_momentum = float(os.environ.get("TTT_MOMENTUM", 0.9))
+ ttt_epochs = int(os.environ.get("TTT_EPOCHS", 3))
+ ttt_grad_clip = float(os.environ.get("TTT_GRAD_CLIP", 1.0))
+ ttt_chunk_size = int(os.environ.get("TTT_CHUNK_SIZE", 32768))
+
+ # v5: Depth recurrence (loop middle layers multiple times)
+ depth_recurrence_start = int(os.environ.get("DEPTH_RECURRENCE_START", 3))
+ depth_recurrence_end = int(os.environ.get("DEPTH_RECURRENCE_END", 5)) # inclusive
+ depth_recurrence_loops = int(os.environ.get("DEPTH_RECURRENCE_LOOPS", 1)) # 1 = no recurrence
+ depth_recurrence_activation_frac = float(os.environ.get("DEPTH_RECURRENCE_ACTIVATION_FRAC", 0.35))
+
+ # v5: Parallel residuals (GPT-J style, attn + mlp in parallel for late layers)
+ parallel_residual_start = int(os.environ.get("PARALLEL_RESIDUAL_START", 999)) # 999 = disabled
+
+ # v5: EMA decay (was hardcoded to 0.997)
+ ema_decay = float(os.environ.get("EMA_DECAY", 0.997))
+
+ # v7: Mamba-3 hybrid
+ model_type = os.environ.get("MODEL_TYPE", "transformer")
+ d_state = int(os.environ.get("D_STATE", 128))
+ mamba_expand = int(os.environ.get("MAMBA_EXPAND", 2))
+ mamba_headdim = int(os.environ.get("MAMBA_HEADDIM", 64))
+ mamba_chunk_size = int(os.environ.get("MAMBA_CHUNK_SIZE", 64))
+ mamba_ngroups = int(os.environ.get("MAMBA_NGROUPS", 1))
+ # ATTN_LAYER_POS: single int or comma-separated list (e.g. "2,5")
+ _attn_pos_raw = os.environ.get("ATTN_LAYER_POS", "2,5")
+ attn_layer_pos = [int(x) for x in _attn_pos_raw.split(",")] if "," in _attn_pos_raw else int(_attn_pos_raw)
+
+# --- Batched Newton-Schulz orthogonalization ---
+
+def zeropower_via_newtonschulz5(G: Tensor, steps: int = 5, eps: float = 1e-7) -> Tensor:
+ """Batched Newton-Schulz orthogonalization. G: (B,M,N) or (M,N)."""
+ a, b, c = (3.4445, -4.7750, 2.0315)
+ was_2d = G.ndim == 2
+ if was_2d:
+ G = G.unsqueeze(0)
+ X = G.bfloat16()
+ transposed = X.size(-2) > X.size(-1)
+ if transposed:
+ X = X.mT
+ X = X / (X.norm(dim=(-2, -1), keepdim=True) + eps)
+ for _ in range(steps):
+ A = X @ X.mT
+ B = b * A + c * (A @ A)
+ X = a * X + B @ X
+ if transposed:
+ X = X.mT
+ if was_2d:
+ X = X.squeeze(0)
+ return X
+
+# v5.3: compile NS5 — fuses the 5-iteration bmm loop into a single CUDA kernel launch
+zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5, dynamic=True)
+
+# --- Parallel Muon optimizer ---
+
+class Muon(torch.optim.Optimizer):
+ """Parallel Muon: post-backward reduce-scatter -> local NS5 -> all-gather.
+
+ No DDP for bank params. After backward, this optimizer:
+ 1. Launches async reduce-scatter for all banks (biggest first)
+ 2. Returns control so Adam can step on small params while RS is in-flight
+ 3. Waits for each RS, runs local NS5 on the shard, launches async all-gather
+ 4. Each all-gather overlaps with next bank's NS5
+ """
+ def __init__(self, params, lr: float, momentum: float, backend_steps: int,
+ nesterov: bool = True, weight_decay: float = 0.0):
+ super().__init__(
+ params,
+ dict(lr=lr, momentum=momentum, backend_steps=backend_steps,
+ nesterov=nesterov, weight_decay=weight_decay),
+ )
+ self._built = False
+
+ def _build(self):
+ self._distributed = dist.is_available() and dist.is_initialized()
+ self._world_size = dist.get_world_size() if self._distributed else 1
+ self._rank = dist.get_rank() if self._distributed else 0
+ ws = self._world_size
+
+ self._bank_meta = []
+ for group in self.param_groups:
+ for p in group["params"]:
+ B = p.shape[0]
+ padded_B = ((B + ws - 1) // ws) * ws
+ shard_B = padded_B // ws
+ tail = p.shape[1:]
+ dev = p.device
+ self._bank_meta.append({
+ 'p': p,
+ 'B': B,
+ 'padded_grad': torch.zeros(padded_B, *tail, device=dev, dtype=torch.bfloat16),
+ 'shard': torch.zeros(shard_B, *tail, device=dev, dtype=torch.bfloat16),
+ 'shard_mom': torch.zeros(shard_B, *tail, device=dev, dtype=torch.bfloat16),
+ 'full_update': torch.zeros(padded_B, *tail, device=dev, dtype=torch.bfloat16),
+ 'scale': max(1, p.shape[-2] / p.shape[-1]) ** 0.5,
+ })
+ # Sort by size descending -- launch biggest reduce-scatters first
+ self._bank_meta.sort(key=lambda m: -m['p'].numel())
+ self._built = True
+
+ def launch_reduce_scatters(self):
+ """Phase 1: launch async reduce-scatter for all banks. Call right after backward."""
+ if not self._built:
+ self._build()
+ if not self._distributed:
+ return
+ self._rs_futures = []
+ for m in self._bank_meta:
+ p = m['p']
+ if p.grad is None:
+ self._rs_futures.append(None)
+ continue
+ pg = m['padded_grad']
+ pg[:m['B']].copy_(p.grad.bfloat16())
+ if pg.shape[0] > m['B']:
+ pg[m['B']:].zero_()
+ fut = dist.reduce_scatter_tensor(m['shard'], pg, op=dist.ReduceOp.AVG, async_op=True)
+ self._rs_futures.append(fut)
+
+ @torch.no_grad()
+ def step(self, closure=None):
+ """Phase 3: wait for RS, local NS5, all-gather. Call AFTER Adam steps."""
+ loss = None
+ if closure is not None:
+ with torch.enable_grad():
+ loss = closure()
+
+ if not self._built:
+ self._build()
+
+ for group in self.param_groups:
+ lr = group["lr"]
+ momentum = group["momentum"]
+ backend_steps = group["backend_steps"]
+ nesterov = group["nesterov"]
+ wd = group.get("weight_decay", 0.0)
+
+ prev_ag_handle = None
+ prev_m = None
+
+ sharded = self._distributed and hasattr(self, '_rs_futures')
+
+ for i, m in enumerate(self._bank_meta):
+ p = m['p']
+ if p.grad is None:
+ continue
+
+ if prev_ag_handle is not None:
+ prev_ag_handle.wait()
+ pp = prev_m['p']
+ upd = prev_m['full_update'][:prev_m['B']]
+ if wd > 0.0:
+ pp.data.mul_(1.0 - lr * wd)
+ pp.add_(upd.to(dtype=pp.dtype), alpha=-lr * prev_m['scale'])
+
+ if sharded and self._rs_futures[i] is not None:
+ self._rs_futures[i].wait()
+ g = m['shard']
+ buf = m['shard_mom']
+ else:
+ g = p.grad.bfloat16()
+ state = self.state[p]
+ if "momentum_buffer" not in state:
+ state["momentum_buffer"] = torch.zeros_like(g)
+ buf = state["momentum_buffer"]
+
+ buf.mul_(momentum).add_(g)
+ if nesterov:
+ update = g.add(buf, alpha=momentum)
+ else:
+ update = buf
+
+ update = zeropower_via_newtonschulz5(update, steps=backend_steps)
+
+ if sharded:
+ prev_ag_handle = dist.all_gather_into_tensor(
+ m['full_update'], update, async_op=True)
+ prev_m = m
+ else:
+ if wd > 0.0:
+ p.data.mul_(1.0 - lr * wd)
+ p.add_(update.to(dtype=p.dtype), alpha=-lr * m['scale'])
+
+ if prev_ag_handle is not None:
+ prev_ag_handle.wait()
+ pp = prev_m['p']
+ upd = prev_m['full_update'][:prev_m['B']]
+ if wd > 0.0:
+ pp.data.mul_(1.0 - lr * wd)
+ pp.add_(upd.to(dtype=pp.dtype), alpha=-lr * prev_m['scale'])
+
+ if hasattr(self, '_rs_futures'):
+ del self._rs_futures
+
+ return loss
+
+# --- Tokenizer evaluation helpers ---
+
+def build_sentencepiece_luts(
+ sp: spm.SentencePieceProcessor, vocab_size: int, device: torch.device
+) -> tuple[Tensor, Tensor, Tensor]:
+ sp_vocab_size = int(sp.vocab_size())
+ table_size = max(sp_vocab_size, vocab_size)
+ base_bytes_np = np.zeros((table_size,), dtype=np.int16)
+ has_leading_space_np = np.zeros((table_size,), dtype=np.bool_)
+ is_boundary_token_np = np.ones((table_size,), dtype=np.bool_)
+ for token_id in range(sp_vocab_size):
+ if sp.is_control(token_id) or sp.is_unknown(token_id) or sp.is_unused(token_id):
+ continue
+ is_boundary_token_np[token_id] = False
+ if sp.is_byte(token_id):
+ base_bytes_np[token_id] = 1
+ continue
+ piece = sp.id_to_piece(token_id)
+ if piece.startswith("\u2581"):
+ has_leading_space_np[token_id] = True
+ piece = piece[1:]
+ base_bytes_np[token_id] = len(piece.encode("utf-8"))
+ return (
+ torch.tensor(base_bytes_np, dtype=torch.int16, device=device),
+ torch.tensor(has_leading_space_np, dtype=torch.bool, device=device),
+ torch.tensor(is_boundary_token_np, dtype=torch.bool, device=device),
+ )
+def load_validation_tokens(pattern: str, seq_len: int) -> Tensor:
+ files = [Path(p) for p in sorted(glob.glob(pattern))]
+ if not files:
+ raise FileNotFoundError(f"No files found for pattern: {pattern}")
+ tokens = torch.cat([load_data_shard(file) for file in files]).contiguous()
+ usable = ((tokens.numel() - 1) // seq_len) * seq_len
+ if usable <= 0:
+ raise ValueError(f"Validation split is too short for TRAIN_SEQ_LEN={seq_len}")
+ return tokens[: usable + 1]
+def eval_val(
+ args: Hyperparameters,
+ model: nn.Module,
+ rank: int,
+ world_size: int,
+ device: torch.device,
+ grad_accum_steps: int,
+ val_tokens: Tensor,
+ base_bytes_lut: Tensor,
+ has_leading_space_lut: Tensor,
+ is_boundary_token_lut: Tensor,
+ eval_seq_len: int | None = None,
+) -> tuple[float, float]:
+ seq_len = eval_seq_len or args.train_seq_len
+ local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps)
+ if local_batch_tokens < seq_len:
+ raise ValueError(
+ "VAL_BATCH_SIZE must provide at least one sequence per rank; "
+ f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, "
+ f"GRAD_ACCUM_STEPS={grad_accum_steps}, seq_len={seq_len}"
+ )
+ local_batch_seqs = local_batch_tokens // seq_len
+ total_seqs = (val_tokens.numel() - 1) // seq_len
+ seq_start = (total_seqs * rank) // world_size
+ seq_end = (total_seqs * (rank + 1)) // world_size
+ val_loss_sum = torch.zeros((), device=device, dtype=torch.float64)
+ val_token_count = torch.zeros((), device=device, dtype=torch.float64)
+ val_byte_count = torch.zeros((), device=device, dtype=torch.float64)
+ model.eval()
+ with torch.inference_mode():
+ for batch_seq_start in range(seq_start, seq_end, local_batch_seqs):
+ batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end)
+ raw_start = batch_seq_start * seq_len
+ raw_end = batch_seq_end * seq_len + 1
+ local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True)
+ x = local[:-1].reshape(-1, seq_len)
+ y = local[1:].reshape(-1, seq_len)
+ with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+ batch_loss = model(x, y).detach()
+ batch_token_count = float(y.numel())
+ val_loss_sum += batch_loss.to(torch.float64) * batch_token_count
+ val_token_count += batch_token_count
+ prev_ids = x.reshape(-1)
+ tgt_ids = y.reshape(-1)
+ token_bytes = base_bytes_lut[tgt_ids].to(dtype=torch.int16)
+ token_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(dtype=torch.int16)
+ val_byte_count += token_bytes.to(torch.float64).sum()
+ if dist.is_available() and dist.is_initialized():
+ dist.all_reduce(val_loss_sum, op=dist.ReduceOp.SUM)
+ dist.all_reduce(val_token_count, op=dist.ReduceOp.SUM)
+ dist.all_reduce(val_byte_count, op=dist.ReduceOp.SUM)
+ val_loss = val_loss_sum / val_token_count
+ bits_per_token = val_loss.item() / math.log(2.0)
+ tokens_per_byte = val_token_count.item() / val_byte_count.item()
+ model.train()
+ return float(val_loss.item()), float(bits_per_token * tokens_per_byte)
+
+# --- N-gram tilt for eval-time logit boosting ---
+
+class NgramTilt:
+ """Causal n-gram prediction booster for eval time.
+
+ Builds a live n-gram table from ground-truth tokens seen so far (causal),
+ and optionally loads a pre-computed prior from training data.
+ Applies a small additive tilt to logits based on n-gram predictions.
+ """
+
+ def __init__(self, vocab_size: int, beta: float = 0.3, max_n: int = 4):
+ self.vocab_size = vocab_size
+ self.beta = beta
+ self.max_n = max_n
+ self.prior_table: dict | None = None
+ # Dense lookup tensors for vectorized tilt (populated by load_prior)
+ self.bigram_arr: "Tensor | None" = None # shape [V]
+ self.trigram_arr: "Tensor | None" = None # shape [V, V]
+
+ def load_prior(self, path: str) -> None:
+ """Load pre-computed n-gram table from compressed artifact."""
+ import pickle
+ import zlib
+ with open(path, 'rb') as f:
+ compressed = f.read()
+ self.prior_table = pickle.loads(zlib.decompress(compressed))
+ self._precompute_lookup_arrays()
+
+ def _precompute_lookup_arrays(self) -> None:
+ """Convert prior_table dicts into dense int16 tensors for vectorized eval.
+
+ bigram_arr[prev] = top-1 token for bigram prefix (prev,), or -1.
+ trigram_arr[prev2, prev1] = top-1 token for trigram prefix (prev2, prev1), or -1.
+ Lookup is O(1) tensor indexing instead of O(N) Python dict iteration.
+ """
+ if self.prior_table is None:
+ return
+ V = self.vocab_size
+ if 2 in self.prior_table:
+ arr = torch.full((V,), -1, dtype=torch.int16)
+ for (prev,), top_list in self.prior_table[2].items():
+ if top_list:
+ arr[prev] = top_list[0][0]
+ self.bigram_arr = arr
+ if 3 in self.prior_table:
+ arr = torch.full((V, V), -1, dtype=torch.int16)
+ for (prev2, prev1), top_list in self.prior_table[3].items():
+ if top_list:
+ arr[prev2, prev1] = top_list[0][0]
+ self.trigram_arr = arr
+
+
+# --- Quantization helpers ---
+
+CONTROL_TENSOR_NAME_PATTERNS = tuple(
+ pattern
+ for pattern in os.environ.get(
+ "CONTROL_TENSOR_NAME_PATTERNS",
+ "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,skip_weight,skip_weights,smear,dtg_gate,ve_layer_scales,ve_shared.scale,attn_gate,vr_lambda",
+ ).split(",")
+ if pattern
+)
+INT8_PER_ROW_SCALE_DTYPE = torch.float16
+INT8_CLIP_PERCENTILE = 99.99984
+INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0
+def quantize_float_tensor(t: Tensor) -> tuple[Tensor, Tensor]:
+ t32 = t.float()
+ if t32.ndim == 2:
+ clip_abs = (
+ torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1)
+ if t32.numel()
+ else torch.empty((t32.shape[0],), dtype=torch.float32)
+ )
+ clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None])
+ scale = (clip_abs / 127.0).clamp_min(1.0 / 127.0)
+ q = torch.clamp(torch.round(clipped / scale[:, None]), -127, 127).to(torch.int8).contiguous()
+ return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous()
+ clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0
+ scale = torch.tensor(clip_abs / 127.0 if clip_abs > 0 else 1.0, dtype=torch.float32)
+ q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -127, 127).to(torch.int8).contiguous()
+ return q, scale
+# --- Data loading ---
+
+def load_data_shard(file: Path) -> Tensor:
+ header_bytes = 256 * np.dtype(" None:
+ self.file_idx = (self.file_idx + 1) % len(self.files)
+ self.tokens = load_data_shard(self.files[self.file_idx])
+ self.pos = 0
+ def take(self, n: int) -> Tensor:
+ chunks: list[Tensor] = []
+ remaining = n
+ while remaining > 0:
+ avail = self.tokens.numel() - self.pos
+ if avail <= 0:
+ self._advance_file()
+ continue
+ k = min(remaining, avail)
+ chunks.append(self.tokens[self.pos : self.pos + k])
+ self.pos += k
+ remaining -= k
+ return chunks[0] if len(chunks) == 1 else torch.cat(chunks)
+class DistributedTokenLoader:
+ def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device):
+ self.rank = rank
+ self.world_size = world_size
+ self.device = device
+ self.stream = TokenStream(pattern)
+ def next_batch(self, global_tokens: int, seq_len: int, grad_accum_steps: int) -> tuple[Tensor, Tensor]:
+ local_tokens = global_tokens // (self.world_size * grad_accum_steps)
+ per_rank_span = local_tokens + 1
+ chunk = self.stream.take(per_rank_span * self.world_size)
+ start = self.rank * per_rank_span
+ local = chunk[start : start + per_rank_span].to(dtype=torch.int64)
+ x = local[:-1].reshape(-1, seq_len)
+ y = local[1:].reshape(-1, seq_len)
+ return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True)
+
+# --- Transformer modules ---
+
+class RMSNorm(nn.Module):
+ def __init__(self, eps: float | None = None):
+ super().__init__()
+ self.eps = eps
+ def forward(self, x: Tensor) -> Tensor:
+ return F.rms_norm(x, (x.size(-1),), eps=self.eps)
+class CastedLinear(nn.Linear):
+ _qat_enabled: bool = False
+ _noisy_qat_enabled: bool = False
+ _noisy_qat_clip_range: int = 31
+ def forward(self, x: Tensor) -> Tensor:
+ w = self.weight.to(x.dtype)
+ if self.training and w.ndim == 2:
+ if CastedLinear._noisy_qat_enabled:
+ # v8: Noisy QAT — inject Gaussian noise calibrated to INT6 quantization error.
+ # Quantization step delta = row_max / clip_range.
+ # Uniform rounding error has std = delta / sqrt(12).
+ with torch.no_grad():
+ w32 = self.weight.float()
+ clip_range = CastedLinear._noisy_qat_clip_range
+ row_max = w32.abs().amax(dim=1).clamp_min(1e-8)
+ delta = row_max / clip_range
+ noise_std = delta / (12.0 ** 0.5)
+ noise = torch.randn_like(w) * noise_std[:, None].to(w.dtype)
+ w = w + noise
+ elif CastedLinear._qat_enabled:
+ with torch.no_grad():
+ w32 = self.weight.float()
+ row_max = w32.abs().amax(dim=1)
+ scale = (row_max / 31.0).clamp_min(1.0 / 31.0)
+ w_q = (torch.clamp(torch.round(w32 / scale[:, None]), -32, 31) * scale[:, None]).to(x.dtype)
+ w = w + (w_q - w).detach()
+ bias = self.bias.to(x.dtype) if self.bias is not None else None
+ return F.linear(x, w, bias)
+def restore_low_dim_params_to_fp32(module: nn.Module) -> None:
+ with torch.no_grad():
+ for name, param in module.named_parameters():
+ if (param.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)) and param.dtype != torch.float32:
+ param.data = param.data.float()
+class Rotary(nn.Module):
+ def __init__(self, dim: int, base: float = 10000.0, train_seq_len: int = 1024, rope_dims: int = 0):
+ super().__init__()
+ self.dim = dim
+ self.base = base
+ self.train_seq_len = train_seq_len
+ self.rope_dims = rope_dims if rope_dims > 0 else dim
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.rope_dims, 2, dtype=torch.float32) / self.rope_dims))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self._seq_len_cached = 0
+ self._cos_cached: Tensor | None = None
+ self._sin_cached: Tensor | None = None
+ def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]:
+ if (
+ self._cos_cached is None
+ or self._sin_cached is None
+ or self._seq_len_cached != seq_len
+ or self._cos_cached.device != device
+ ):
+ rd = self.rope_dims
+ if seq_len > self.train_seq_len:
+ scale = seq_len / self.train_seq_len
+ new_base = self.base * (scale ** (rd / (rd - 2)))
+ inv_freq = 1.0 / (new_base ** (torch.arange(0, rd, 2, dtype=torch.float32, device=device) / rd))
+ else:
+ inv_freq = self.inv_freq.to(device)
+ t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
+ freqs = torch.outer(t, inv_freq)
+ self._cos_cached = freqs.cos()[None, :, None, :]
+ self._sin_cached = freqs.sin()[None, :, None, :]
+ self._seq_len_cached = seq_len
+ return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype)
+def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor, rope_dims: int = 0) -> Tensor:
+ if rope_dims > 0 and rope_dims < x.size(-1):
+ x_rope, x_pass = x[..., :rope_dims], x[..., rope_dims:]
+ half = rope_dims // 2
+ x1, x2 = x_rope[..., :half], x_rope[..., half:]
+ x_rope = torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1)
+ return torch.cat((x_rope, x_pass), dim=-1)
+ half = x.size(-1) // 2
+ x1, x2 = x[..., :half], x[..., half:]
+ return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1)
+
+class CausalSelfAttention(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int,
+ num_kv_heads: int,
+ rope_base: float,
+ qk_gain_init: float,
+ gated_attention: bool = False,
+ value_residual: bool = False,
+ ):
+ super().__init__()
+ if dim % num_heads != 0:
+ raise ValueError("model_dim must be divisible by num_heads")
+ if num_heads % num_kv_heads != 0:
+ raise ValueError("num_heads must be divisible by num_kv_heads")
+ self.num_heads = num_heads
+ self.num_kv_heads = num_kv_heads
+ self.head_dim = dim // num_heads
+ if self.head_dim % 2 != 0:
+ raise ValueError("head_dim must be even for RoPE")
+ # No CastedLinear -- weights come from banks
+ self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32))
+ self.rope_dims = 0 # set by GPT.__init__ for partial RoPE
+ self.rotary = Rotary(self.head_dim, base=rope_base, train_seq_len=1024)
+ self.use_xsa = False # set by GPT.__init__ for deep layers only
+ # Gated attention and value residual (non-banked small params)
+ self.gated_attention = gated_attention
+ if gated_attention:
+ self.attn_gate = nn.Linear(dim, num_heads, bias=True)
+ nn.init.zeros_(self.attn_gate.weight)
+ nn.init.constant_(self.attn_gate.bias, 4.0)
+ self.value_residual = value_residual
+ if value_residual:
+ self.vr_lambda = nn.Parameter(torch.tensor([0.5, 0.5], dtype=torch.float32))
+ def _xsa_efficient(self, y: Tensor, v: Tensor) -> Tensor:
+ """Efficient XSA: subtract self-value projection via GQA-aware reshape (no repeat_interleave).
+ y: [B, T, H, D], v: [B, T, Hkv, D]. H must be divisible by Hkv."""
+ B, T, H, D = y.shape
+ Hkv = v.size(-2)
+ group = H // Hkv
+ y_g = y.reshape(B, T, Hkv, group, D) # [B, T, Hkv, group, D]
+ vn = F.normalize(v, dim=-1).unsqueeze(-2) # [B, T, Hkv, 1, D] -- broadcast ready
+ proj = (y_g * vn).sum(dim=-1, keepdim=True) * vn
+ return (y_g - proj).reshape(B, T, H, D)
+ def forward(self, x: Tensor, q_w: Tensor, k_w: Tensor, v_w: Tensor, out_w: Tensor, v_embed: Tensor | None = None, v0: Tensor | None = None) -> tuple[Tensor, Tensor | None]:
+ bsz, seqlen, dim = x.shape
+ q = F.linear(x, q_w.to(x.dtype)).reshape(bsz, seqlen, self.num_heads, self.head_dim)
+ k = F.linear(x, k_w.to(x.dtype)).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim)
+ v = F.linear(x, v_w.to(x.dtype))
+ if v_embed is not None:
+ v = v + v_embed
+ v = v.reshape(bsz, seqlen, self.num_kv_heads, self.head_dim)
+ raw_v = v if self.value_residual else None
+ if self.value_residual and v0 is not None:
+ lam = self.vr_lambda.to(dtype=v.dtype)
+ v = lam[0] * v0 + lam[1] * v
+ q = F.rms_norm(q, (q.size(-1),))
+ k = F.rms_norm(k, (k.size(-1),))
+ cos, sin = self.rotary(seqlen, x.device, q.dtype)
+ q = apply_rotary_emb(q, cos, sin, self.rope_dims)
+ k = apply_rotary_emb(k, cos, sin, self.rope_dims)
+ q = q * self.q_gain.to(dtype=q.dtype)[None, None, :, None]
+ y = flash_attn_3_func(q, k, v, causal=True)
+ if self.use_xsa:
+ y = self._xsa_efficient(y, v)
+ if self.gated_attention:
+ # gate shape: (bsz, seqlen, num_heads) -> (bsz, seqlen, num_heads, 1) for B,T,H,D layout
+ gate = torch.sigmoid(self.attn_gate(x)).unsqueeze(-1)
+ y = y * gate
+ y = y.reshape(bsz, seqlen, dim)
+ return F.linear(y, out_w.to(x.dtype)), raw_v
+
+class SmearGate(nn.Module):
+ def __init__(self, dim: int):
+ super().__init__()
+ self.gate = nn.Parameter(torch.zeros(dim, dtype=torch.float32))
+ def forward(self, x: Tensor) -> Tensor:
+ g = torch.sigmoid(self.gate.to(dtype=x.dtype))[None, None, :]
+ x_prev = torch.cat([torch.zeros_like(x[:, :1]), x[:, :-1]], dim=1)
+ return (1 - g) * x + g * x_prev
+
+class BigramHashEmbedding(nn.Module):
+ def __init__(self, bigram_vocab_size: int, bigram_dim: int, model_dim: int):
+ super().__init__()
+ self.bigram_vocab_size = bigram_vocab_size
+ self.embed = nn.Embedding(bigram_vocab_size, bigram_dim)
+ nn.init.zeros_(self.embed.weight)
+ self.proj = CastedLinear(bigram_dim, model_dim, bias=False) if bigram_dim != model_dim else None
+ if self.proj is not None:
+ nn.init.zeros_(self.proj.weight)
+ self.scale = nn.Parameter(torch.tensor(0.05, dtype=torch.float32))
+ def bigram_hash(self, tokens: Tensor) -> Tensor:
+ t = tokens.to(torch.int32)
+ mod = self.bigram_vocab_size - 1
+ out = torch.empty_like(t)
+ out[..., 0] = mod
+ out[..., 1:] = torch.bitwise_xor(36313 * t[..., 1:], 27191 * t[..., :-1]) % mod
+ return out.long()
+ def forward(self, token_ids: Tensor) -> Tensor:
+ h = self.embed(self.bigram_hash(token_ids))
+ if self.proj is not None:
+ h = self.proj(h)
+ return h * self.scale.to(dtype=h.dtype)
+
+class ValueEmbedding(nn.Module):
+ """Reinject token identity into attention values at specific layers.
+ Each table maps vocab tokens to a low-dim embedding, projected to model_dim."""
+ def __init__(self, vocab_size: int, ve_dim: int, model_dim: int):
+ super().__init__()
+ self.embed = nn.Embedding(vocab_size, ve_dim)
+ nn.init.normal_(self.embed.weight, std=0.01)
+ self.proj = CastedLinear(ve_dim, model_dim, bias=False) if ve_dim != model_dim else None
+ if self.proj is not None:
+ nn.init.zeros_(self.proj.weight)
+ self.scale = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
+ def forward(self, token_ids: Tensor) -> Tensor:
+ h = self.embed(token_ids)
+ if self.proj is not None:
+ h = self.proj(h)
+ return h * self.scale.to(dtype=h.dtype)
+
+class MLP(nn.Module):
+ def __init__(self, dim: int, mlp_mult: float):
+ super().__init__()
+ # No CastedLinear -- weights come from banks
+ def forward(self, x: Tensor, up_w: Tensor, down_w: Tensor) -> Tensor:
+ x = F.leaky_relu(F.linear(x, up_w.to(x.dtype)), negative_slope=0.5)
+ return F.linear(x.square(), down_w.to(x.dtype))
+
+class Block(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int,
+ num_kv_heads: int,
+ mlp_mult: float,
+ rope_base: float,
+ qk_gain_init: float,
+ layer_idx: int = 0,
+ ln_scale: bool = False,
+ dtg: bool = False,
+ gated_attention: bool = False,
+ value_residual: bool = False,
+ parallel_residual: bool = False,
+ ):
+ super().__init__()
+ self.parallel_residual = parallel_residual
+ self.attn_norm = RMSNorm()
+ self.mlp_norm = RMSNorm()
+ self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init,
+ gated_attention=gated_attention, value_residual=value_residual)
+ self.mlp = MLP(dim, mlp_mult)
+ self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32))
+ self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32))
+ self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float())
+ self.ln_scale_factor = 1.0 / math.sqrt(layer_idx + 1) if ln_scale else 1.0
+ if dtg:
+ self.dtg_gate = nn.Linear(dim, 1, bias=True)
+ nn.init.zeros_(self.dtg_gate.weight)
+ nn.init.constant_(self.dtg_gate.bias, 2.0)
+ else:
+ self.dtg_gate = None
+ def forward(self, x: Tensor, x0: Tensor, q_w: Tensor, k_w: Tensor, v_w: Tensor, out_w: Tensor, up_w: Tensor, down_w: Tensor, v_embed: Tensor | None = None, v0: Tensor | None = None) -> tuple[Tensor, Tensor | None]:
+ mix = self.resid_mix.to(dtype=x.dtype)
+ x_in = mix[0][None, None, :] * x + mix[1][None, None, :] * x0
+ normed = self.attn_norm(x_in) * self.ln_scale_factor
+ attn_out, raw_v = self.attn(normed, q_w, k_w, v_w, out_w, v_embed=v_embed, v0=v0)
+ if self.parallel_residual:
+ # GPT-J style: attn and MLP both see the same normed input
+ mlp_out = self.mlp(normed, up_w, down_w)
+ x_out = x_in + self.attn_scale.to(dtype=x_in.dtype)[None, None, :] * attn_out + self.mlp_scale.to(dtype=x_in.dtype)[None, None, :] * mlp_out
+ else:
+ x_out = x_in + self.attn_scale.to(dtype=x_in.dtype)[None, None, :] * attn_out
+ x_out = x_out + self.mlp_scale.to(dtype=x_out.dtype)[None, None, :] * self.mlp(self.mlp_norm(x_out) * self.ln_scale_factor, up_w, down_w)
+ if self.dtg_gate is not None:
+ gate = torch.sigmoid(self.dtg_gate(x_in.detach()))
+ x_out = x_in + gate * (x_out - x_in)
+ return x_out, raw_v
+
+class GPT(nn.Module):
+ def __init__(
+ self,
+ vocab_size: int,
+ num_layers: int,
+ model_dim: int,
+ num_heads: int,
+ num_kv_heads: int,
+ mlp_mult: float,
+ tie_embeddings: bool,
+ tied_embed_init_std: float,
+ logit_softcap: float,
+ rope_base: float,
+ qk_gain_init: float,
+ mtp_num_heads: int = 0,
+ mtp_loss_weight: float = 0.1,
+ bigram_vocab_size: int = 0,
+ bigram_dim: int = 128,
+ xsa_last_n: int = 0,
+ rope_dims: int = 0,
+ ln_scale: bool = False,
+ dtg: bool = False,
+ ve_enabled: bool = False,
+ ve_dim: int = 128,
+ ve_layers: str = "9,10",
+ gated_attention: bool = False,
+ value_residual: bool = False,
+ parallel_residual_start: int = 999,
+ depth_recurrence_start: int = 3,
+ depth_recurrence_end: int = 5,
+ depth_recurrence_loops: int = 1,
+ depth_recurrence_activation_frac: float = 0.35,
+ ):
+ super().__init__()
+ self._ve_target_dim = num_kv_heads * (model_dim // num_heads) # kv_dim for value projection
+ if logit_softcap <= 0.0:
+ raise ValueError(f"logit_softcap must be positive, got {logit_softcap}")
+ self.tie_embeddings = tie_embeddings
+ self.tied_embed_init_std = tied_embed_init_std
+ self.logit_softcap = logit_softcap
+ self.value_residual = value_residual
+ self.mtp_num_heads = mtp_num_heads
+ self.mtp_loss_weight = mtp_loss_weight
+ # v5: depth recurrence
+ self.depth_recurrence_start = depth_recurrence_start
+ self.depth_recurrence_end = depth_recurrence_end
+ self.depth_recurrence_loops = depth_recurrence_loops
+ self.depth_recurrence_activation_frac = depth_recurrence_activation_frac
+ self._training_progress = 0.0 # set by training loop
+ # _rec_loops: controls how many times the recurrence zone runs.
+ # Starts at 1 (no recurrence); training loop sets to depth_recurrence_loops
+ # once _training_progress crosses the activation threshold.
+ # This avoids torch.compile recompilation every step — it only changes once.
+ self._rec_loops = 1
+ self.tok_emb = nn.Embedding(vocab_size, model_dim)
+ # v8: Frozen bigram log-probability prior (logit bias during training + inference)
+ self._bigram_prior_active = False
+ self.bigram_prior_scale = nn.Parameter(torch.tensor(0.5, dtype=torch.float32))
+ self.bigram = BigramHashEmbedding(bigram_vocab_size, bigram_dim, model_dim) if bigram_vocab_size > 0 else None
+ self.smear = SmearGate(model_dim)
+ self.num_encoder_layers = num_layers // 2
+ self.num_decoder_layers = num_layers - self.num_encoder_layers
+ self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers)
+ self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32))
+ # Parameter banks: contiguous 3D tensors for batched optimizer
+ head_dim = model_dim // num_heads
+ kv_dim = num_kv_heads * head_dim
+ mlp_dim = int(mlp_mult * model_dim)
+ self.num_layers = num_layers
+ self.qo_bank = nn.Parameter(torch.empty(2 * num_layers, model_dim, model_dim))
+ self.kv_bank = nn.Parameter(torch.empty(2 * num_layers, kv_dim, model_dim))
+ self.mlp_up_bank = nn.Parameter(torch.empty(num_layers, mlp_dim, model_dim))
+ self.mlp_down_bank = nn.Parameter(torch.empty(num_layers, model_dim, mlp_dim))
+ self.blocks = nn.ModuleList(
+ [
+ Block(
+ model_dim,
+ num_heads,
+ num_kv_heads,
+ mlp_mult,
+ rope_base,
+ qk_gain_init,
+ layer_idx=i,
+ ln_scale=ln_scale,
+ dtg=dtg,
+ gated_attention=gated_attention,
+ value_residual=value_residual,
+ parallel_residual=(i >= parallel_residual_start),
+ )
+ for i in range(num_layers)
+ ]
+ )
+ if rope_dims > 0:
+ head_dim = model_dim // num_heads
+ for block in self.blocks:
+ block.attn.rope_dims = rope_dims
+ block.attn.rotary = Rotary(head_dim, base=rope_base, train_seq_len=1024, rope_dims=rope_dims)
+ self.ve_layer_indices = [int(x) for x in ve_layers.split(",") if x.strip()] if ve_enabled else []
+ kv_dim_ve = self._ve_target_dim
+ if self.ve_layer_indices:
+ self.ve_shared = ValueEmbedding(vocab_size, ve_dim, kv_dim_ve)
+ self.ve_layer_scales = nn.ParameterList(
+ [nn.Parameter(torch.ones(1, dtype=torch.float32)) for _ in self.ve_layer_indices]
+ )
+ else:
+ self.ve_shared = None
+ self.ve_layer_scales = nn.ParameterList()
+ self.final_norm = RMSNorm()
+ self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False)
+ if self.lm_head is not None:
+ self.lm_head._zero_init = True
+ self.mtp_heads = nn.ModuleList(
+ [CastedLinear(model_dim, vocab_size, bias=False) for _ in range(mtp_num_heads)]
+ )
+ for head in self.mtp_heads:
+ head._zero_init = True
+ if xsa_last_n > 0:
+ for i in range(max(0, num_layers - xsa_last_n), num_layers):
+ self.blocks[i].attn.use_xsa = True
+ self._init_weights()
+ def _init_weights(self) -> None:
+ if self.tie_embeddings:
+ nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std)
+ n = self.num_layers
+ proj_scale = 1.0 / math.sqrt(2 * n)
+ # Init banks: orthogonal, with proj layers scaled down and out/down zero-init
+ for i in range(n):
+ nn.init.orthogonal_(self.qo_bank.data[i], gain=1.0) # Q
+ nn.init.zeros_(self.qo_bank.data[n + i]) # Out (zero init)
+ nn.init.orthogonal_(self.kv_bank.data[i], gain=1.0) # K
+ nn.init.orthogonal_(self.kv_bank.data[n + i], gain=1.0) # V
+ nn.init.orthogonal_(self.mlp_up_bank.data[i], gain=1.0) # MLP up
+ nn.init.zeros_(self.mlp_down_bank.data[i]) # MLP down (zero init)
+ # Scale proj layers (out_proj and mlp_down are "proj" layers)
+ self.qo_bank.data[n + i].mul_(proj_scale)
+ self.mlp_down_bank.data[i].mul_(proj_scale)
+ # Init remaining nn.Linear modules (bigram proj, mtp heads, lm_head)
+ for name, module in self.named_modules():
+ if isinstance(module, nn.Linear):
+ if getattr(module, "_zero_init", False):
+ nn.init.zeros_(module.weight)
+ elif module.weight.ndim == 2 and module.weight.shape[0] >= 64 and module.weight.shape[1] >= 64:
+ nn.init.orthogonal_(module.weight, gain=1.0)
+ def load_bigram_prior(self, path: str) -> None:
+ """Load frozen bigram log-prob matrix as a non-parameter buffer."""
+ mat = torch.load(path, map_location="cpu", weights_only=True) # [V, V] float32
+ # Move to model's device (load is CPU, model may already be on CUDA)
+ device = next(self.parameters()).device
+ self.register_buffer("bigram_prior_mat", mat.to(device))
+ self._bigram_prior_active = True
+
+ def _get_ve(self, layer_idx: int, input_ids: Tensor, ve_cache: dict | None = None) -> Tensor | None:
+ """Get value embedding for a specific layer using shared table + per-layer scale."""
+ if self.ve_shared is None or layer_idx not in self.ve_layer_indices:
+ return None
+ if ve_cache is not None and 've' not in ve_cache:
+ ve_cache['ve'] = self.ve_shared(input_ids)
+ ve_base = ve_cache['ve'] if ve_cache is not None else self.ve_shared(input_ids)
+ ve_idx = self.ve_layer_indices.index(layer_idx)
+ return ve_base * self.ve_layer_scales[ve_idx].to(dtype=ve_base.dtype)
+ def _run_layers(self, x: Tensor, x0: Tensor, input_ids: Tensor, ve_cache: dict) -> Tensor:
+ """Run all layers with U-Net skips and optional depth recurrence."""
+ n = self.num_layers
+ rec_start = self.depth_recurrence_start
+ rec_end = min(self.depth_recurrence_end, n - 1)
+ rec_loops = self._rec_loops # Set externally by training loop; avoids recompilation
+
+ v0 = None
+ skips: list[Tensor] = []
+
+ def run_block(i: int, x: Tensor, v0: Tensor | None) -> tuple[Tensor, Tensor | None]:
+ ve = self._get_ve(i, input_ids, ve_cache)
+ x_out, raw_v = self.blocks[i](x, x0,
+ self.qo_bank[i], self.kv_bank[i], self.kv_bank[n + i],
+ self.qo_bank[n + i], self.mlp_up_bank[i], self.mlp_down_bank[i],
+ v_embed=ve, v0=v0)
+ if v0 is None and raw_v is not None:
+ v0 = raw_v
+ return x_out, v0
+
+ # Section 1: Pre-recurrence layers
+ for i in range(min(rec_start, n)):
+ if i < self.num_encoder_layers:
+ x, v0 = run_block(i, x, v0)
+ skips.append(x)
+ else:
+ di = i - self.num_encoder_layers
+ if skips and di < self.num_skip_weights:
+ x = x + self.skip_weights[di].to(dtype=x.dtype)[None, None, :] * skips.pop()
+ x, v0 = run_block(i, x, v0)
+
+ # Section 2: Recurrence zone — run [rec_start..rec_end] rec_loops times
+ # First pass: push/pop skips as normal
+ # Additional passes: no skip operations
+ for loop_pass in range(rec_loops):
+ for i in range(rec_start, min(rec_end + 1, n)):
+ if loop_pass == 0:
+ # First pass: handle skips normally
+ if i < self.num_encoder_layers:
+ x, v0 = run_block(i, x, v0)
+ skips.append(x)
+ else:
+ di = i - self.num_encoder_layers
+ if skips and di < self.num_skip_weights:
+ x = x + self.skip_weights[di].to(dtype=x.dtype)[None, None, :] * skips.pop()
+ x, v0 = run_block(i, x, v0)
+ else:
+ # Additional passes: no skip operations
+ x, v0 = run_block(i, x, v0)
+
+ # Section 3: Post-recurrence layers
+ for i in range(min(rec_end + 1, n), n):
+ if i < self.num_encoder_layers:
+ x, v0 = run_block(i, x, v0)
+ skips.append(x)
+ else:
+ di = i - self.num_encoder_layers
+ if skips and di < self.num_skip_weights:
+ x = x + self.skip_weights[di].to(dtype=x.dtype)[None, None, :] * skips.pop()
+ x, v0 = run_block(i, x, v0)
+
+ return x
+
+ def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor:
+ x = self.tok_emb(input_ids)
+ if self.bigram is not None:
+ x = x + self.bigram(input_ids)
+ x = F.rms_norm(x, (x.size(-1),))
+ x = self.smear(x)
+ x0 = x
+ ve_cache: dict = {}
+ x = self._run_layers(x, x0, input_ids, ve_cache)
+ x = self.final_norm(x)
+ x_flat = x.reshape(-1, x.size(-1))
+ targets = target_ids.reshape(-1)
+ if self.tie_embeddings:
+ logits_proj = F.linear(x_flat, self.tok_emb.weight)
+ else:
+ if self.lm_head is None:
+ raise RuntimeError("lm_head is required when tie_embeddings=False")
+ logits_proj = self.lm_head(x_flat)
+ # v8: Bigram prior logit bias (before softcap)
+ if self._bigram_prior_active:
+ prev_tokens = input_ids.reshape(-1)
+ bias = self.bigram_prior_mat.to(logits_proj.device)[prev_tokens.long()]
+ logits_proj = logits_proj + self.bigram_prior_scale * bias.to(logits_proj.dtype)
+ logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap)
+ main_loss = F.cross_entropy(logits.float(), targets, reduction="mean")
+ if self.training and self.mtp_num_heads > 0 and self.mtp_loss_weight > 0.0:
+ _, seqlen, dim = x.shape
+ mtp_loss_sum = x.new_zeros(())
+ mtp_loss_count = 0
+ for k, mtp_head in enumerate(self.mtp_heads):
+ valid_t = seqlen - (k + 1)
+ if valid_t <= 0:
+ continue
+ mtp_hidden = x[:, :valid_t, :].reshape(-1, dim)
+ mtp_targets = target_ids[:, k + 1 :].reshape(-1)
+ mtp_logits_proj = mtp_head(mtp_hidden)
+ mtp_logits = self.logit_softcap * torch.tanh(mtp_logits_proj / self.logit_softcap)
+ mtp_loss_sum = mtp_loss_sum + F.cross_entropy(mtp_logits.float(), mtp_targets, reduction="mean")
+ mtp_loss_count += 1
+ if mtp_loss_count > 0:
+ main_loss = main_loss + self.mtp_loss_weight * (mtp_loss_sum / mtp_loss_count)
+ return main_loss
+ def forward_logits(self, input_ids: Tensor) -> Tensor:
+ """Return logits (bsz, seq_len, vocab) without computing loss."""
+ x = self.tok_emb(input_ids)
+ if self.bigram is not None:
+ x = x + self.bigram(input_ids)
+ x = F.rms_norm(x, (x.size(-1),))
+ x = self.smear(x)
+ x0 = x
+ ve_cache: dict = {}
+ x = self._run_layers(x, x0, input_ids, ve_cache)
+ x = self.final_norm(x)
+ if self.tie_embeddings:
+ logits_proj = F.linear(x, self.tok_emb.weight)
+ else:
+ logits_proj = self.lm_head(x)
+ # v8: Bigram prior logit bias (before softcap)
+ if self._bigram_prior_active:
+ bias = self.bigram_prior_mat.to(logits_proj.device)[input_ids.long()]
+ logits_proj = logits_proj + self.bigram_prior_scale * bias.to(logits_proj.dtype)
+ return self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap)
+
+# --- Sliding window evaluation ---
+
+def eval_val_sliding(
+ args: Hyperparameters,
+ base_model: nn.Module,
+ rank: int,
+ world_size: int,
+ device: torch.device,
+ val_tokens: Tensor,
+ base_bytes_lut: Tensor,
+ has_leading_space_lut: Tensor,
+ is_boundary_token_lut: Tensor,
+ stride: int,
+ batch_seqs: int = 32,
+ eval_seq_len: int | None = None,
+ ngram_tilt: "NgramTilt | None" = None,
+) -> tuple[float, float]:
+ """Sliding window evaluation: each token scored with maximum context."""
+ seq_len = eval_seq_len or args.train_seq_len
+ total_tokens = val_tokens.numel() - 1
+ window_starts = [ws for ws in range(0, total_tokens, stride)
+ if min(ws + seq_len, total_tokens) - ws >= 1]
+ total_windows = len(window_starts)
+ my_s = (total_windows * rank) // world_size
+ my_e = (total_windows * (rank + 1)) // world_size
+ my_windows = window_starts[my_s:my_e]
+ loss_sum = torch.zeros((), device=device, dtype=torch.float64)
+ token_count = torch.zeros((), device=device, dtype=torch.float64)
+ byte_count = torch.zeros((), device=device, dtype=torch.float64)
+ base_model.eval()
+ if isinstance(base_model, nn.Module) and hasattr(base_model, 'smear_gate'):
+ # Mamba hybrid — skip torch.compile (einops breaks fullgraph=True)
+ compiled_logits = base_model.forward_logits
+ else:
+ compiled_logits = torch.compile(base_model.forward_logits, dynamic=False, fullgraph=True)
+ with torch.inference_mode():
+ for bi in range(0, len(my_windows), batch_seqs):
+ batch_ws = my_windows[bi:bi + batch_seqs]
+ bsz = len(batch_ws)
+ x_batch = torch.zeros(bsz, seq_len, dtype=torch.int64, device=device)
+ y_batch = torch.zeros(bsz, seq_len, dtype=torch.int64, device=device)
+ wlens: list[int] = []
+ for i, ws in enumerate(batch_ws):
+ end = min(ws + seq_len, total_tokens)
+ wlen = end - ws
+ wlens.append(wlen)
+ chunk = val_tokens[ws:end + 1].to(dtype=torch.int64, device=device)
+ x_batch[i, :wlen] = chunk[:-1]
+ y_batch[i, :wlen] = chunk[1:]
+ with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+ logits = compiled_logits(x_batch)
+ # v5.3: apply n-gram prior tilt before scoring (vectorized via precomputed arrays)
+ if ngram_tilt is not None and ngram_tilt.prior_table is not None:
+ logits_f = logits.float()
+ delta = ngram_tilt.beta * 0.5
+ dev = logits.device
+ # Bigram: x[:, t-1] → top token → boost logits[:, t, top]
+ if ngram_tilt.bigram_arr is not None:
+ bg = ngram_tilt.bigram_arr.to(dev)
+ top_tok = bg[x_batch[:, :-1].long()] # [bsz, seq_len-1]
+ b_idx, t_idx = (top_tok >= 0).nonzero(as_tuple=True)
+ if b_idx.numel() > 0:
+ logits_f[b_idx, t_idx + 1, top_tok[b_idx, t_idx].long()] += delta
+ # Trigram: (x[:, t-2], x[:, t-1]) → top token → boost logits[:, t, top]
+ if ngram_tilt.trigram_arr is not None:
+ tg = ngram_tilt.trigram_arr.to(dev)
+ top_tok = tg[x_batch[:, :-2].long(), x_batch[:, 1:-1].long()] # [bsz, seq_len-2]
+ b_idx, t_idx = (top_tok >= 0).nonzero(as_tuple=True)
+ if b_idx.numel() > 0:
+ logits_f[b_idx, t_idx + 2, top_tok[b_idx, t_idx].long()] += delta
+ logits = logits_f
+ nll = F.cross_entropy(
+ logits.reshape(-1, logits.size(-1)).float(),
+ y_batch.reshape(-1),
+ reduction="none",
+ ).reshape(bsz, seq_len)
+ for i, ws in enumerate(batch_ws):
+ wlen = wlens[i]
+ s = 0 if ws == 0 else max(wlen - stride, 0)
+ scored_nll = nll[i, s:wlen].to(torch.float64)
+ loss_sum += scored_nll.sum()
+ token_count += float(wlen - s)
+ tgt = y_batch[i, s:wlen]
+ prev = x_batch[i, s:wlen]
+ tb = base_bytes_lut[tgt].to(torch.float64)
+ tb += (has_leading_space_lut[tgt] & ~is_boundary_token_lut[prev]).to(torch.float64)
+ byte_count += tb.sum()
+ if dist.is_available() and dist.is_initialized():
+ dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
+ dist.all_reduce(token_count, op=dist.ReduceOp.SUM)
+ dist.all_reduce(byte_count, op=dist.ReduceOp.SUM)
+ val_loss = (loss_sum / token_count).item()
+ bits_per_token = val_loss / math.log(2.0)
+ tokens_per_byte = token_count.item() / byte_count.item()
+ base_model.train()
+ return val_loss, bits_per_token * tokens_per_byte
+
+
+# --- Legal TTT (test-time training) sliding window evaluation ---
+
+def _run_ttt_sliding_window_eval(
+ model: GPT,
+ val_tokens: Tensor,
+ seq_len: int,
+ stride: int,
+ base_bytes_lut: Tensor,
+ has_leading_space_lut: Tensor,
+ is_boundary_token_lut: Tensor,
+ device: torch.device,
+ ttt_lr: float = 0.005,
+ ttt_momentum: float = 0.9,
+ ttt_epochs: int = 3,
+ ttt_grad_clip: float = 1.0,
+ ngram_tilt: "NgramTilt | None" = None,
+ chunk_size: int = 32768,
+) -> tuple[float, float]:
+ """
+ Legal TTT evaluation — FAST chunk-based version.
+
+ Score-first protocol:
+ 1. Score chunk with non-overlapping seq_len windows (ONE pass per window, no stride)
+ 2. Train on the SAME already-scored chunk with SGD
+ 3. Move to next chunk
+
+ Speed: ~(1 + ttt_epochs) × (chunk_size / seq_len) forward passes per chunk.
+ With chunk_size=32K, seq_len=2048, 1 epoch: ~32 passes/chunk × ~2300 chunks = ~74K passes.
+ At ~4ms/pass on H100: ~5 min. (vs ~85 min for the old stride-64 version)
+ """
+ model.eval()
+ original_state = copy.deepcopy(model.state_dict())
+
+ # Ensure all params are float32 for SGD (dequantized model may have mixed dtypes)
+ for p in model.parameters():
+ p.data = p.data.float()
+
+ ttt_optimizer = torch.optim.SGD(
+ model.parameters(), lr=ttt_lr, momentum=ttt_momentum
+ )
+
+ n_tokens = val_tokens.shape[0]
+ total_nll = 0.0
+ total_base_bytes = 0.0
+
+ # Non-overlapping chunks for TTT
+ chunk_starts = list(range(0, n_tokens - 1, chunk_size))
+ n_chunks = len(chunk_starts)
+
+ for chunk_idx, chunk_start in enumerate(chunk_starts):
+ chunk_end = min(chunk_start + chunk_size, n_tokens - 1)
+ actual_chunk_len = chunk_end - chunk_start
+ if actual_chunk_len < 2:
+ continue
+
+ # === PHASE 1: SCORE with non-overlapping seq_len windows ===
+ with torch.inference_mode():
+ pos = chunk_start
+ while pos < chunk_end:
+ end = min(pos + seq_len, chunk_end)
+ if end - pos < 2:
+ break
+
+ x = val_tokens[pos:end].unsqueeze(0).to(dtype=torch.int64, device=device)
+ y = val_tokens[pos + 1:end + 1].unsqueeze(0).to(dtype=torch.int64, device=device)
+
+ with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+ logits = model.forward_logits(x)
+
+ # Apply n-gram tilt
+ if ngram_tilt is not None and ngram_tilt.prior_table is not None:
+ logits_f = logits.float()
+ delta = ngram_tilt.beta * 0.5
+ dev = logits.device
+ if ngram_tilt.bigram_arr is not None:
+ bg = ngram_tilt.bigram_arr.to(dev)
+ top_tok = bg[x[:, :-1].long()]
+ b_idx, t_idx = (top_tok >= 0).nonzero(as_tuple=True)
+ if b_idx.numel() > 0:
+ logits_f[b_idx, t_idx + 1, top_tok[b_idx, t_idx].long()] += delta
+ if ngram_tilt.trigram_arr is not None:
+ tg = ngram_tilt.trigram_arr.to(dev)
+ top_tok = tg[x[:, :-2].long(), x[:, 1:-1].long()]
+ b_idx, t_idx = (top_tok >= 0).nonzero(as_tuple=True)
+ if b_idx.numel() > 0:
+ logits_f[b_idx, t_idx + 2, top_tok[b_idx, t_idx].long()] += delta
+ logits = logits_f
+
+ # Score ALL tokens in this window (single pass, no stride)
+ wlen = end - pos
+ nll = F.cross_entropy(
+ logits[0, :wlen].float(),
+ y[0, :wlen],
+ reduction="none",
+ )
+ total_nll += nll.to(torch.float64).sum().item()
+
+ tgt = y[0, :wlen]
+ prev = x[0, :wlen]
+ tb = base_bytes_lut[tgt].to(torch.float64)
+ tb += (has_leading_space_lut[tgt] & ~is_boundary_token_lut[prev]).to(torch.float64)
+ total_base_bytes += tb.sum().item()
+
+ pos += seq_len # non-overlapping — advance by full seq_len
+
+ # === PHASE 2: TRAIN on already-scored chunk ===
+ is_last = (chunk_idx >= n_chunks - 1)
+ if not is_last:
+ model.train()
+
+ # Cosine LR decay across chunks
+ cos_scale = 0.5 * (1.0 + math.cos(math.pi * chunk_idx / max(n_chunks - 1, 1)))
+ for pg in ttt_optimizer.param_groups:
+ pg['lr'] = ttt_lr * cos_scale
+
+ # Train in seq_len-sized sub-windows
+ for _epoch in range(ttt_epochs):
+ train_pos = chunk_start
+ while train_pos < chunk_end:
+ train_end = min(train_pos + seq_len, chunk_end)
+ if train_end - train_pos < 2:
+ break
+
+ x_t = val_tokens[train_pos:train_end].unsqueeze(0).to(dtype=torch.int64, device=device)
+ y_t = val_tokens[train_pos + 1:train_end + 1].unsqueeze(0).to(dtype=torch.int64, device=device)
+
+ ttt_optimizer.zero_grad()
+ with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+ loss = model(x_t, y_t)
+ loss.backward()
+ torch.nn.utils.clip_grad_norm_(model.parameters(), ttt_grad_clip)
+ ttt_optimizer.step()
+
+ train_pos += seq_len
+
+ model.eval()
+
+ # Progress logging every 100 chunks
+ if chunk_idx % 100 == 0 or chunk_idx == n_chunks - 1:
+ partial_bpb = (total_nll / max(total_base_bytes, 1e-9)) / math.log(2)
+ print(f" ttt_chunk:{chunk_idx}/{n_chunks} partial_bpb:{partial_bpb:.6f}", flush=True)
+
+ # Restore original weights (TTT is eval-only, don't persist changes)
+ model.load_state_dict(original_state)
+
+ if total_base_bytes < 1e-9:
+ return 0.0, 0.0
+ val_loss = total_nll / total_base_bytes
+ val_bpb = val_loss / math.log(2.0)
+ return val_loss, val_bpb
+
+
+# --- GPTQ-lite int6 quantization ---
+
+def _classify_param(name: str) -> str:
+ if "tok_emb" in name or "lm_head" in name:
+ return "embed"
+ if ".mlp." in name or "mlp_fc" in name or "mlp_proj" in name:
+ return "mlp"
+ if ".attn." in name or "c_q" in name or "c_k" in name or "c_v" in name or "c_proj" in name:
+ return "attn"
+ if "in_proj" in name or "out_proj" in name:
+ return "mamba"
+ if ".proj." in name and ".mlp." not in name:
+ return "attn"
+ return "other"
+def quantize_int6_per_row(t: Tensor, clip_range: int = 31, clip_percentiles: list[float] | None = None) -> tuple[Tensor, Tensor]:
+ t32 = t.float()
+ if clip_percentiles is None:
+ clip_percentiles = [0.9990, 0.9995, 0.9999, 0.99999, 1.0]
+ if t32.ndim == 2:
+ best_q, best_s, best_err = None, None, float('inf')
+ for pct in clip_percentiles:
+ if pct < 1.0:
+ row_clip = torch.quantile(t32.abs(), pct, dim=1)
+ else:
+ row_clip = t32.abs().amax(dim=1)
+ s = (row_clip / clip_range).clamp_min(1.0 / clip_range).to(torch.float16)
+ q = torch.clamp(torch.round(t32 / s.float()[:, None]), -clip_range, clip_range).to(torch.int8)
+ recon = q.float() * s.float()[:, None]
+ err = (t32 - recon).pow(2).mean().item()
+ if err < best_err:
+ best_q, best_s, best_err = q, s, err
+ return best_q, best_s
+ amax = t32.abs().max().item()
+ scale = torch.tensor(amax / clip_range if amax > 0 else 1.0, dtype=torch.float16)
+ q = torch.clamp(torch.round(t32 / scale.float()), -clip_range, clip_range).to(torch.int8)
+ return q, scale
+
+def _unbank_state_dict(sd: dict[str, Tensor], num_layers: int) -> dict[str, Tensor]:
+ """Convert 3D bank tensors into individual 2D tensors with standard names."""
+ out: dict[str, Tensor] = {}
+ n = num_layers
+ for name, tensor in sd.items():
+ if name == "qo_bank":
+ for i in range(n):
+ out[f"blocks.{i}.attn.c_q.weight"] = tensor[i]
+ out[f"blocks.{i}.attn.proj.weight"] = tensor[n + i]
+ elif name == "kv_bank":
+ for i in range(n):
+ out[f"blocks.{i}.attn.c_k.weight"] = tensor[i]
+ out[f"blocks.{i}.attn.c_v.weight"] = tensor[n + i]
+ elif name == "mlp_up_bank":
+ for i in range(n):
+ out[f"blocks.{i}.mlp.fc.weight"] = tensor[i]
+ elif name == "mlp_down_bank":
+ for i in range(n):
+ out[f"blocks.{i}.mlp.proj.weight"] = tensor[i]
+ else:
+ out[name] = tensor
+ return out
+
+def _rebank_state_dict(sd: dict[str, Tensor], num_layers: int, template_sd: dict[str, Tensor]) -> dict[str, Tensor]:
+ """Convert individual 2D tensors back into 3D bank tensors."""
+ out: dict[str, Tensor] = {}
+ n = num_layers
+ # Reconstruct banks from individual weight keys
+ qo_slices = [None] * (2 * n)
+ kv_slices = [None] * (2 * n)
+ up_slices = [None] * n
+ down_slices = [None] * n
+ consumed = set()
+ for i in range(n):
+ qk = f"blocks.{i}.attn.c_q.weight"
+ if qk in sd:
+ qo_slices[i] = sd[qk]
+ consumed.add(qk)
+ ok = f"blocks.{i}.attn.proj.weight"
+ if ok in sd:
+ qo_slices[n + i] = sd[ok]
+ consumed.add(ok)
+ kk = f"blocks.{i}.attn.c_k.weight"
+ if kk in sd:
+ kv_slices[i] = sd[kk]
+ consumed.add(kk)
+ vk = f"blocks.{i}.attn.c_v.weight"
+ if vk in sd:
+ kv_slices[n + i] = sd[vk]
+ consumed.add(vk)
+ fk = f"blocks.{i}.mlp.fc.weight"
+ if fk in sd:
+ up_slices[i] = sd[fk]
+ consumed.add(fk)
+ dk = f"blocks.{i}.mlp.proj.weight"
+ if dk in sd:
+ down_slices[i] = sd[dk]
+ consumed.add(dk)
+ out["qo_bank"] = torch.stack(qo_slices).to(dtype=template_sd["qo_bank"].dtype)
+ out["kv_bank"] = torch.stack(kv_slices).to(dtype=template_sd["kv_bank"].dtype)
+ out["mlp_up_bank"] = torch.stack(up_slices).to(dtype=template_sd["mlp_up_bank"].dtype)
+ out["mlp_down_bank"] = torch.stack(down_slices).to(dtype=template_sd["mlp_down_bank"].dtype)
+ for name, tensor in sd.items():
+ if name not in consumed:
+ out[name] = tensor
+ return out
+
+def mixed_quantize_int6(state_dict: dict[str, Tensor], int6_cats: set[str]):
+ num_layers_total = max(
+ (int(k.split(".")[1]) for k in state_dict if k.startswith("blocks.")),
+ default=0,
+ ) + 1
+ result: dict[str, Tensor] = {}
+ meta: dict[str, object] = {}
+ for name, tensor in state_dict.items():
+ t = tensor.detach().cpu().contiguous()
+ cat = _classify_param(name)
+ if not t.is_floating_point() or t.numel() <= 65536 or "bigram_prior" in name:
+ result[name] = t.to(torch.float16) if t.is_floating_point() else t
+ meta[name] = "passthrough"
+ continue
+ if any(p in name for p in CONTROL_TENSOR_NAME_PATTERNS):
+ result[name] = t.float()
+ meta[name] = "passthrough_ctrl"
+ continue
+ if cat in int6_cats and t.ndim >= 1:
+ # v6: per-layer adaptive GPTQ clipping
+ # MLP layers: tighter clipping preserves precision
+ # Attn layers: looser clipping improves compressibility
+ if cat == "mlp":
+ q, s = quantize_int6_per_row(t, clip_range=31, clip_percentiles=[0.9995, 0.9999, 1.0])
+ else: # attn
+ q, s = quantize_int6_per_row(t, clip_range=31, clip_percentiles=[0.999, 0.9995, 0.9999, 0.99999, 1.0])
+ result[name + ".q"] = q
+ result[name + ".scale"] = s
+ meta[name] = {"type": "int6"}
+ else:
+ q, s = quantize_float_tensor(t)
+ result[name + ".q"] = q
+ result[name + ".scale"] = s
+ meta[name] = {"type": "int8"}
+ return result, meta
+def dequantize_mixed_int6(result: dict[str, Tensor], meta: dict[str, object],
+ template_sd: dict[str, Tensor]) -> dict[str, Tensor]:
+ out: dict[str, Tensor] = {}
+ for name, orig in template_sd.items():
+ info = meta.get(name)
+ if info is None:
+ continue
+ orig_dtype = orig.dtype
+ if info in ("passthrough", "passthrough_ctrl"):
+ t = result[name]
+ if t.dtype == torch.float16 and orig_dtype in (torch.float32, torch.bfloat16):
+ t = t.to(orig_dtype)
+ out[name] = t
+ continue
+ q, s = result[name + ".q"], result[name + ".scale"]
+ if s.ndim > 0:
+ out[name] = (q.float() * s.float().view(q.shape[0], *([1] * (q.ndim - 1)))).to(orig_dtype)
+ else:
+ out[name] = (q.float() * float(s.item())).to(orig_dtype)
+ return out
+
+# --- Training ---
+
+def main() -> None:
+ code = Path(__file__).read_text(encoding="utf-8")
+ args = Hyperparameters()
+ # base_model.forward is not torch.compiled (NS5 is compiled separately at module level)
+ distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ
+ rank = int(os.environ.get("RANK", "0"))
+ world_size = int(os.environ.get("WORLD_SIZE", "1"))
+ local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+ if world_size <= 0:
+ raise ValueError(f"WORLD_SIZE must be positive, got {world_size}")
+ if 8 % world_size != 0:
+ raise ValueError(f"WORLD_SIZE={world_size} must divide 8 so grad_accum_steps stays integral")
+ grad_accum_steps = 8 // world_size
+ grad_scale = 1.0 / grad_accum_steps
+ if not torch.cuda.is_available():
+ raise RuntimeError("CUDA is required")
+ device = torch.device("cuda", local_rank)
+ torch.cuda.set_device(device)
+ if distributed:
+ dist.init_process_group(backend="nccl", device_id=device)
+ dist.barrier()
+ master_process = rank == 0
+ torch.backends.cuda.matmul.allow_tf32 = True
+ torch.backends.cudnn.allow_tf32 = True
+ from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp
+ enable_cudnn_sdp(False)
+ enable_flash_sdp(True)
+ enable_mem_efficient_sdp(False)
+ enable_math_sdp(False)
+ logfile = None
+ if master_process:
+ os.makedirs("logs", exist_ok=True)
+ logfile = f"logs/{args.run_id}.txt"
+ print(logfile)
+ def log0(msg: str, console: bool = True) -> None:
+ if not master_process:
+ return
+ if console:
+ print(msg)
+ if logfile is not None:
+ with open(logfile, "a", encoding="utf-8") as f:
+ print(msg, file=f)
+ log0(code, console=False)
+ log0("=" * 100, console=False)
+ log0(f"Running Python {sys.version}", console=False)
+ log0(f"Running PyTorch {torch.__version__}", console=False)
+ log0(
+ subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False).stdout,
+ console=False,
+ )
+ log0("=" * 100, console=False)
+ random.seed(args.seed)
+ np.random.seed(args.seed)
+ torch.manual_seed(args.seed)
+ torch.cuda.manual_seed_all(args.seed)
+ # --- Tokenizer loading: SentencePiece (.model) or BESE BPE (.json) ---
+ _use_bese = args.tokenizer_path.endswith(".json")
+ if _use_bese:
+ from bese_fast_bpe import FastBESEBPETokenizer
+ bese_tok = FastBESEBPETokenizer.load(args.tokenizer_path)
+ if bese_tok.vocab_size != args.vocab_size:
+ raise ValueError(
+ f"VOCAB_SIZE={args.vocab_size} does not match BESE vocab_size={bese_tok.vocab_size}"
+ )
+ dataset_dir = Path(args.data_path).resolve()
+ actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin")))
+ effective_eval_seq_len = args.eval_seq_len if args.eval_seq_len > 0 else args.train_seq_len
+ val_seq_len = max(args.train_seq_len, effective_eval_seq_len)
+ val_tokens = load_validation_tokens(args.val_files, val_seq_len)
+ base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = bese_tok.build_luts_for_training(device)
+ log0(f"val_bpb:enabled tokenizer_kind=bese_bpe tokenizer_path={args.tokenizer_path} vocab_size={bese_tok.vocab_size}")
+ else:
+ if not args.tokenizer_path.endswith(".model"):
+ raise ValueError(f"Tokenizer must be .model (SentencePiece) or .json (BESE BPE): {args.tokenizer_path}")
+ sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path)
+ if int(sp.vocab_size()) != args.vocab_size:
+ raise ValueError(
+ f"VOCAB_SIZE={args.vocab_size} does not match tokenizer vocab_size={int(sp.vocab_size())}"
+ )
+ dataset_dir = Path(args.data_path).resolve()
+ actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin")))
+ effective_eval_seq_len = args.eval_seq_len if args.eval_seq_len > 0 else args.train_seq_len
+ val_seq_len = max(args.train_seq_len, effective_eval_seq_len)
+ val_tokens = load_validation_tokens(args.val_files, val_seq_len)
+ base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts(
+ sp, args.vocab_size, device
+ )
+ log0(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}")
+ log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files}")
+ log0(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}")
+ CastedLinear._qat_enabled = args.qat_enabled
+ if args.model_type == "mamba_hybrid":
+ from mamba3_ssd import HybridMambaGPT
+ base_model = HybridMambaGPT(
+ vocab_size=args.vocab_size,
+ num_layers=args.num_layers,
+ model_dim=args.model_dim,
+ d_state=args.d_state,
+ expand=args.mamba_expand,
+ headdim=args.mamba_headdim,
+ chunk_size=args.mamba_chunk_size,
+ attn_pos=args.attn_layer_pos,
+ num_heads=args.num_heads,
+ num_kv_heads=args.num_kv_heads,
+ mlp_mult=args.mlp_mult,
+ rope_base=args.rope_base,
+ qk_gain_init=args.qk_gain_init,
+ rope_dims=args.rope_dims,
+ logit_softcap=args.logit_softcap,
+ tied_embed_init_std=args.tied_embed_init_std,
+ ngroups=args.mamba_ngroups,
+ depth_recurrence_start=args.depth_recurrence_start,
+ depth_recurrence_end=args.depth_recurrence_end,
+ depth_recurrence_loops=args.depth_recurrence_loops,
+ depth_recurrence_activation_frac=args.depth_recurrence_activation_frac,
+ ).to(device).bfloat16()
+ else:
+ base_model = GPT(
+ vocab_size=args.vocab_size,
+ num_layers=args.num_layers,
+ model_dim=args.model_dim,
+ num_heads=args.num_heads,
+ num_kv_heads=args.num_kv_heads,
+ mlp_mult=args.mlp_mult,
+ tie_embeddings=args.tie_embeddings,
+ tied_embed_init_std=args.tied_embed_init_std,
+ logit_softcap=args.logit_softcap,
+ rope_base=args.rope_base,
+ qk_gain_init=args.qk_gain_init,
+ mtp_num_heads=args.mtp_num_heads,
+ mtp_loss_weight=args.mtp_loss_weight,
+ bigram_vocab_size=args.bigram_vocab_size,
+ bigram_dim=args.bigram_dim,
+ xsa_last_n=args.xsa_last_n,
+ rope_dims=args.rope_dims,
+ ln_scale=args.ln_scale,
+ dtg=args.dtg_enabled,
+ ve_enabled=args.ve_enabled,
+ ve_dim=args.ve_dim,
+ ve_layers=args.ve_layers,
+ gated_attention=args.gated_attention,
+ value_residual=args.value_residual,
+ parallel_residual_start=args.parallel_residual_start,
+ depth_recurrence_start=args.depth_recurrence_start,
+ depth_recurrence_end=args.depth_recurrence_end,
+ depth_recurrence_loops=args.depth_recurrence_loops,
+ depth_recurrence_activation_frac=args.depth_recurrence_activation_frac,
+ ).to(device).bfloat16()
+ # Banks stay FP32 (like CastedLinear weights), cast to BF16 in forward
+ base_model.qo_bank.data = base_model.qo_bank.data.float()
+ base_model.kv_bank.data = base_model.kv_bank.data.float()
+ base_model.mlp_up_bank.data = base_model.mlp_up_bank.data.float()
+ base_model.mlp_down_bank.data = base_model.mlp_down_bank.data.float()
+ for module in base_model.modules():
+ if isinstance(module, CastedLinear):
+ module.float()
+ restore_low_dim_params_to_fp32(base_model)
+ # v8: Load bigram prior
+ if args.model_type != "mamba_hybrid" and args.bigram_prior_enabled and args.bigram_prior_path and os.path.exists(args.bigram_prior_path):
+ base_model.load_bigram_prior(args.bigram_prior_path)
+ log0(f"bigram_prior:loaded path:{args.bigram_prior_path}")
+ # No DDP -- Parallel Muon handles bank grad communication via reduce-scatter,
+ # and non-bank grads are manually all-reduced before Adam steps.
+ if args.model_type == "mamba_hybrid":
+ # Skip torch.compile — Triton kernel (mamba_chunk_scan_combined) provides
+ # the real speedup; dynamo tracing with custom Triton ops causes hangs
+ compiled_model = base_model
+ log0("torch.compile:skipped (mamba_hybrid, Triton kernel is the fast path)")
+ else:
+ compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True)
+ model = compiled_model
+
+ if args.model_type == "mamba_hybrid":
+ # v7: Mamba hybrid optimizer setup
+ # 2D weight matrices (ndim>=2, numel>4096) -> Muon (Newton-Schulz)
+ # 1D params (biases, norms, scales, D, dt_bias, A_log) -> Adam
+ # Token embedding -> Adam (separate LR)
+ token_lr = args.tied_embed_lr # Mamba always uses tied embeddings
+ matrix_params = []
+ scalar_params = []
+ tok_params = [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}]
+ for name, p in base_model.named_parameters():
+ if "tok_emb" in name:
+ continue # already in tok_params
+ elif p.ndim >= 2 and p.numel() > 4096:
+ matrix_params.append(p)
+ else:
+ scalar_params.append(p)
+ optimizer_tok = torch.optim.AdamW(
+ tok_params,
+ betas=(args.beta1, args.beta2),
+ eps=args.adam_eps,
+ weight_decay=args.adam_wd,
+ fused=True,
+ )
+ optimizer_muon = Muon(
+ matrix_params,
+ lr=args.matrix_lr,
+ momentum=args.muon_momentum,
+ backend_steps=args.muon_backend_steps,
+ weight_decay=args.muon_wd,
+ )
+ for group in optimizer_muon.param_groups:
+ group["base_lr"] = args.matrix_lr
+ optimizer_scalar = torch.optim.AdamW(
+ [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}],
+ betas=(args.beta1, args.beta2),
+ eps=args.adam_eps,
+ weight_decay=args.adam_wd,
+ fused=True,
+ )
+ replicated_params = list(optimizer_tok.param_groups[0]["params"])
+ replicated_params.extend(scalar_params)
+ optimizer_head = None # Mamba uses tied embeddings, no separate lm_head
+ optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar]
+ else:
+ # Optimizer split:
+ # - 4 parameter banks -> Muon (batched Newton-Schulz)
+ # - token embedding -> Adam
+ # - scalars/control tensors -> Adam
+ # - bigram proj, mtp heads, VE proj -> Adam (small matrix params not worth banking)
+ matrix_params = [
+ base_model.qo_bank, base_model.kv_bank,
+ base_model.mlp_up_bank, base_model.mlp_down_bank,
+ ]
+ block_named_params = list(base_model.blocks.named_parameters())
+ scalar_params = [
+ p
+ for name, p in block_named_params
+ if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)
+ ]
+ if base_model.skip_weights.numel() > 0:
+ scalar_params.append(base_model.skip_weights)
+ scalar_params.append(base_model.smear.gate)
+ if base_model._bigram_prior_active:
+ scalar_params.append(base_model.bigram_prior_scale)
+ if base_model.bigram is not None:
+ scalar_params.append(base_model.bigram.scale)
+ token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr
+ tok_params = [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}]
+ if base_model.bigram is not None:
+ tok_params.append({"params": [base_model.bigram.embed.weight], "lr": token_lr, "base_lr": token_lr})
+ if base_model.bigram.proj is not None:
+ scalar_params.append(base_model.bigram.proj.weight)
+ if base_model.ve_shared is not None:
+ tok_params.append({"params": [base_model.ve_shared.embed.weight], "lr": token_lr, "base_lr": token_lr})
+ if base_model.ve_shared.proj is not None:
+ scalar_params.append(base_model.ve_shared.proj.weight)
+ scalar_params.append(base_model.ve_shared.scale)
+ for s in base_model.ve_layer_scales:
+ scalar_params.append(s)
+ optimizer_tok = torch.optim.AdamW(
+ tok_params,
+ betas=(args.beta1, args.beta2),
+ eps=args.adam_eps,
+ weight_decay=args.adam_wd,
+ fused=True,
+ )
+ optimizer_muon = Muon(
+ matrix_params,
+ lr=args.matrix_lr,
+ momentum=args.muon_momentum,
+ backend_steps=args.muon_backend_steps,
+ weight_decay=args.muon_wd,
+ )
+ for group in optimizer_muon.param_groups:
+ group["base_lr"] = args.matrix_lr
+ optimizer_scalar = torch.optim.AdamW(
+ [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}],
+ betas=(args.beta1, args.beta2),
+ eps=args.adam_eps,
+ weight_decay=args.adam_wd,
+ fused=True,
+ )
+ # Non-bank params that need manual all-reduce (replicated across GPUs)
+ replicated_params = list(optimizer_tok.param_groups[0]["params"])
+ for pg in optimizer_tok.param_groups[1:]:
+ replicated_params.extend(pg["params"])
+ replicated_params.extend(scalar_params)
+
+ optimizer_head = None
+ if base_model.lm_head is not None:
+ optimizer_head = torch.optim.Adam(
+ [{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}],
+ betas=(args.beta1, args.beta2),
+ eps=args.adam_eps,
+ fused=True,
+ )
+ replicated_params.append(base_model.lm_head.weight)
+ optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar]
+ if optimizer_head is not None:
+ optimizers.append(optimizer_head)
+ n_params = sum(p.numel() for p in base_model.parameters())
+ mtp_params = sum(p.numel() for p in base_model.mtp_heads.parameters()) if hasattr(base_model, 'mtp_heads') else 0
+ log0(f"model_params:{n_params}")
+ if args.model_type != "mamba_hybrid":
+ log0(f"mtp_num_heads:{args.mtp_num_heads} mtp_loss_weight:{args.mtp_loss_weight} mtp_params:{mtp_params}")
+ xsa_layers = [i for i, b in enumerate(base_model.blocks) if b.attn.use_xsa]
+ log0(f"XSA:last_{args.xsa_last_n} active_layers:{xsa_layers}")
+ else:
+ log0(f"model_type:mamba_hybrid d_state:{args.d_state} expand:{args.mamba_expand} ngroups:{args.mamba_ngroups} attn_pos:{args.attn_layer_pos}")
+ log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}")
+ log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False")
+ log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}")
+ log0(
+ f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} "
+ f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} "
+ f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}"
+ )
+ log0(
+ f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} "
+ f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} "
+ f"max_wallclock_seconds:{args.max_wallclock_seconds:.3f}"
+ )
+ log0(f"seed:{args.seed}")
+ train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device)
+ def zero_grad_all() -> None:
+ for opt in optimizers:
+ opt.zero_grad(set_to_none=True)
+ max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None
+ def lr_mul(step: int, elapsed_ms: float) -> float:
+ if args.warmdown_iters <= 0:
+ return 1.0
+ if max_wallclock_ms is None:
+ warmdown_start = max(args.iterations - args.warmdown_iters, 0)
+ return max((args.iterations - step) / max(args.warmdown_iters, 1), 0.0) if warmdown_start <= step < args.iterations else 1.0
+ step_ms = elapsed_ms / max(step, 1)
+ warmdown_ms = args.warmdown_iters * step_ms
+ remaining_ms = max(max_wallclock_ms - elapsed_ms, 0.0)
+ return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0
+ if args.warmup_steps > 0:
+ initial_model_state = {name: tensor.detach().cpu().clone() for name, tensor in base_model.state_dict().items()}
+ initial_optimizer_states = [copy.deepcopy(opt.state_dict()) for opt in optimizers]
+ model.train()
+ for warmup_step in range(args.warmup_steps):
+ zero_grad_all()
+ for micro_step in range(grad_accum_steps):
+ x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps)
+ with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+ warmup_loss = model(x, y)
+ (warmup_loss * grad_scale).backward()
+ # All-reduce all grads for warmup (simple, not optimized)
+ if distributed:
+ for p in base_model.parameters():
+ if p.grad is not None:
+ dist.all_reduce(p.grad, op=dist.ReduceOp.AVG)
+ for opt in optimizers:
+ opt.step()
+ zero_grad_all()
+ if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps:
+ log0(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}")
+ base_model.load_state_dict(initial_model_state, strict=True)
+ for opt, state in zip(optimizers, initial_optimizer_states, strict=True):
+ opt.load_state_dict(state)
+ zero_grad_all()
+ train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device)
+ swa_state: dict[str, Tensor] | None = None
+ swa_count = 0
+ from collections import deque
+ lawa_queue: deque[dict[str, Tensor]] = deque(maxlen=args.lawa_k)
+ ema_state = {name: t.detach().float().clone() for name, t in base_model.state_dict().items()}
+ ema_decay = args.ema_decay
+ training_time_ms = 0.0
+ approx_training_time_ms = 0.0 # v5: init before training loop (used for recurrence activation)
+ stop_after_step: int | None = None
+ # v5.3: warm up NS5 compile before the 600s clock starts — first call triggers JIT
+ if torch.cuda.is_available():
+ _warmup_g = torch.randn(8, 8, device=device, dtype=torch.bfloat16)
+ zeropower_via_newtonschulz5(_warmup_g)
+ del _warmup_g
+ torch.cuda.synchronize()
+ t0 = time.perf_counter()
+ step = 0
+ while True:
+ last_step = step == args.iterations or (stop_after_step is not None and step >= stop_after_step)
+ should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0)
+ if should_validate:
+ torch.cuda.synchronize()
+ training_time_ms += 1000.0 * (time.perf_counter() - t0)
+ val_loss, val_bpb = eval_val(
+ args,
+ model,
+ rank,
+ world_size,
+ device,
+ grad_accum_steps,
+ val_tokens,
+ base_bytes_lut,
+ has_leading_space_lut,
+ is_boundary_token_lut,
+ )
+ log0(
+ f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} "
+ f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms"
+ )
+ torch.cuda.synchronize()
+ t0 = time.perf_counter()
+ if last_step:
+ if stop_after_step is not None and step < args.iterations:
+ log0(
+ f"stopping_early: wallclock_cap train_time:{training_time_ms:.0f}ms "
+ f"step:{step}/{args.iterations}"
+ )
+ break
+ elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0)
+ scale = lr_mul(step, elapsed_ms)
+ if args.late_qat_threshold > 0 and scale < args.late_qat_threshold and not CastedLinear._qat_enabled:
+ CastedLinear._qat_enabled = True
+ log0(f"late_qat:enabled step:{step} scale:{scale:.4f}")
+ # v8: Noisy QAT activation based on training progress fraction
+ if args.noisy_qat_enabled and not CastedLinear._noisy_qat_enabled:
+ training_progress = getattr(base_model, '_training_progress', 0.0)
+ if training_progress >= args.noisy_qat_activation_frac:
+ CastedLinear._noisy_qat_enabled = True
+ CastedLinear._noisy_qat_clip_range = args.noisy_qat_clip_range
+ log0(f"noisy_qat:enabled step:{step} progress:{training_progress:.3f} clip_range:{args.noisy_qat_clip_range}")
+ zero_grad_all()
+ # v5: update training progress for depth recurrence activation
+ # Use wallclock fraction (not step fraction) since training is wallclock-capped
+ if max_wallclock_ms is not None and max_wallclock_ms > 0:
+ base_model._training_progress = approx_training_time_ms / max_wallclock_ms
+ else:
+ base_model._training_progress = step / max(args.iterations, 1)
+ # Activate recurrence once threshold is crossed (changes _rec_loops once, one recompile)
+ if base_model._rec_loops == 1 and args.depth_recurrence_loops > 1 and base_model._training_progress >= args.depth_recurrence_activation_frac:
+ base_model._rec_loops = args.depth_recurrence_loops
+ log0(f"depth_recurrence:activated step:{step} progress:{base_model._training_progress:.3f} loops:{args.depth_recurrence_loops}")
+ train_loss = torch.zeros((), device=device)
+ for micro_step in range(grad_accum_steps):
+ x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps)
+ with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+ loss = model(x, y)
+ train_loss += loss.detach()
+ (loss * grad_scale).backward()
+ train_loss /= grad_accum_steps
+ frac = min(step / args.muon_momentum_warmup_steps, 1.0) if args.muon_momentum_warmup_steps > 0 else 1.0
+ muon_momentum = (1 - frac) * args.muon_momentum_warmup_start + frac * args.muon_momentum
+ for group in optimizer_muon.param_groups:
+ group["momentum"] = muon_momentum
+ for opt in optimizers:
+ for group in opt.param_groups:
+ group["lr"] = group["base_lr"] * scale
+ if args.grad_clip_norm > 0:
+ torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm)
+ # === 3-phase overlapped optimizer step ===
+ # Phase 1: Launch async reduce-scatter for banks (biggest first)
+ optimizer_muon.launch_reduce_scatters()
+ # Phase 2: All-reduce non-bank grads + step Adam (while bank RS is in-flight)
+ if distributed:
+ for p in replicated_params:
+ if p.grad is not None:
+ dist.all_reduce(p.grad, op=dist.ReduceOp.AVG)
+ optimizer_tok.step()
+ optimizer_scalar.step()
+ if optimizer_head is not None:
+ optimizer_head.step()
+ # Phase 3: Wait for RS, local NS5, all-gather (banks processed last)
+ optimizer_muon.step()
+ zero_grad_all()
+ # EMA update — v5.3: batched foreach ops instead of Python loop (~20 free steps)
+ with torch.no_grad():
+ _ema_vals = list(ema_state.values())
+ _model_vals = [t.detach().float() for t in base_model.state_dict().values()]
+ torch._foreach_mul_(_ema_vals, ema_decay)
+ torch._foreach_add_(_ema_vals, _model_vals, alpha=1.0 - ema_decay)
+ step += 1
+ approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0)
+ if args.swa_enabled and scale < 0.2 and step % args.swa_every == 0:
+ if swa_state is None:
+ swa_state = {name: t.detach().cpu().clone() for name, t in base_model.state_dict().items()}
+ swa_count = 1
+ log0(f"swa:start step:{step}")
+ else:
+ for name, t in base_model.state_dict().items():
+ swa_state[name] += t.detach().cpu()
+ swa_count += 1
+ if args.lawa_enabled and step % args.lawa_freq == 0:
+ lawa_queue.append({name: t.detach().cpu().clone() for name, t in base_model.state_dict().items()})
+ should_log_train = (
+ args.train_log_every > 0
+ and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None)
+ )
+ if should_log_train:
+ log0(
+ f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} "
+ f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms"
+ )
+ reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms
+ if distributed and max_wallclock_ms is not None:
+ reached_cap_tensor = torch.tensor(int(reached_cap), device=device)
+ dist.all_reduce(reached_cap_tensor, op=dist.ReduceOp.MAX)
+ reached_cap = bool(reached_cap_tensor.item())
+ if stop_after_step is None and reached_cap:
+ stop_after_step = step
+ log0(
+ f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB "
+ f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB"
+ )
+ # Apply weight averaging
+ if args.lawa_enabled and len(lawa_queue) > 1:
+ log0(f"lawa:applying LAWA averaging k={len(lawa_queue)}")
+ current_state = base_model.state_dict()
+ avg_state = {name: torch.zeros(t.shape, dtype=torch.float32, device='cpu') for name, t in current_state.items()}
+ for snap in lawa_queue:
+ for name in avg_state:
+ avg_state[name] += snap[name].float()
+ for name in avg_state:
+ avg_state[name] /= len(lawa_queue)
+ avg_state[name] = avg_state[name].to(dtype=current_state[name].dtype)
+ base_model.load_state_dict(avg_state, strict=True)
+ else:
+ log0("ema:applying EMA weights")
+ current_state = base_model.state_dict()
+ avg_state = {name: t.to(dtype=current_state[name].dtype) for name, t in ema_state.items()}
+ base_model.load_state_dict(avg_state, strict=True)
+ torch.cuda.synchronize()
+ t_diag = time.perf_counter()
+ diag_val_loss, diag_val_bpb = eval_val(
+ args, compiled_model, rank, world_size, device, grad_accum_steps,
+ val_tokens, base_bytes_lut, has_leading_space_lut, is_boundary_token_lut,
+ )
+ torch.cuda.synchronize()
+ log0(
+ f"DIAGNOSTIC post_ema val_loss:{diag_val_loss:.4f} val_bpb:{diag_val_bpb:.4f} "
+ f"eval_time:{1000.0 * (time.perf_counter() - t_diag):.0f}ms"
+ )
+ full_state_dict = base_model.state_dict()
+ export_sd = {k: v for k, v in full_state_dict.items()
+ if "mtp_heads" not in k and k != "bigram_prior_mat"}
+ excluded_mtp = sum(int(t.numel()) for k, t in full_state_dict.items() if "mtp_heads" in k)
+ if excluded_mtp > 0:
+ log0(f"export_excluding_mtp_params:{excluded_mtp}")
+ if master_process:
+ torch.save(export_sd, "final_model.pt")
+ model_bytes = os.path.getsize("final_model.pt")
+ code_bytes = len(code.encode("utf-8"))
+ log0(f"Serialized model: {model_bytes} bytes")
+ log0(f"Code size: {code_bytes} bytes")
+ # v8: Disable QAT noise before quantization
+ CastedLinear._noisy_qat_enabled = False
+ CastedLinear._qat_enabled = False
+ # Unbank 3D tensors into individual 2D tensors for quantization
+ sd_cpu = {k: v.detach().cpu() for k, v in export_sd.items()}
+ if args.model_type == "mamba_hybrid":
+ unbanked_sd = sd_cpu # Mamba model has no banks to unbank
+ else:
+ unbanked_sd = _unbank_state_dict(sd_cpu, args.num_layers)
+ int6_cats = {"mlp", "attn", "mamba"} if args.model_type == "mamba_hybrid" else {"mlp", "attn"}
+ quant_result, quant_meta = mixed_quantize_int6(unbanked_sd, int6_cats)
+ save_dict = {"w": quant_result, "m": quant_meta}
+ # Bundle n-gram prior table if it exists
+ ngram_path = os.environ.get("NGRAM_PRIOR_PATH", "")
+ if ngram_path and os.path.exists(ngram_path):
+ with open(ngram_path, "rb") as nf:
+ save_dict["ngram"] = nf.read()
+ log0(f"Bundled n-gram prior: {len(save_dict['ngram']):,} bytes from {ngram_path}")
+ quant_buf = io.BytesIO()
+ torch.save(save_dict, quant_buf)
+ quant_raw = quant_buf.getvalue()
+ quant_blob = lzma.compress(quant_raw, preset=9)
+ if master_process:
+ with open("final_model.int6.ptz", "wb") as f:
+ f.write(quant_blob)
+ quant_file_bytes = len(quant_blob)
+ code_bytes = len(code.encode("utf-8"))
+ log0(f"Serialized model int6+lzma: {quant_file_bytes} bytes")
+ log0(f"Total submission size int6+lzma: {quant_file_bytes + code_bytes} bytes")
+ if distributed:
+ dist.barrier()
+ with open("final_model.int6.ptz", "rb") as f:
+ quant_blob_disk = f.read()
+ quant_state = torch.load(
+ io.BytesIO(lzma.decompress(quant_blob_disk)),
+ map_location="cpu",
+ )
+ # Extract bundled n-gram prior to temp file for NgramTilt.load_prior()
+ if "ngram" in quant_state:
+ import tempfile
+ _ngram_tmp = tempfile.NamedTemporaryFile(suffix=".bin", delete=False)
+ _ngram_tmp.write(quant_state["ngram"])
+ _ngram_tmp.close()
+ os.environ["NGRAM_PRIOR_PATH"] = _ngram_tmp.name
+ log0(f"Extracted bundled n-gram prior: {len(quant_state['ngram']):,} bytes → {_ngram_tmp.name}")
+ del quant_state["ngram"]
+ deq_unbanked = dequantize_mixed_int6(quant_state["w"], quant_state["m"], unbanked_sd)
+ if args.model_type == "mamba_hybrid":
+ deq_state = deq_unbanked # Mamba model has no banks to rebank
+ else:
+ deq_state = _rebank_state_dict(deq_unbanked, args.num_layers, sd_cpu)
+ if args.model_type == "mamba_hybrid":
+ from mamba3_ssd import HybridMambaGPT
+ eval_model = HybridMambaGPT(
+ vocab_size=args.vocab_size,
+ num_layers=args.num_layers,
+ model_dim=args.model_dim,
+ d_state=args.d_state,
+ expand=args.mamba_expand,
+ headdim=args.mamba_headdim,
+ chunk_size=args.mamba_chunk_size,
+ attn_pos=args.attn_layer_pos,
+ num_heads=args.num_heads,
+ num_kv_heads=args.num_kv_heads,
+ mlp_mult=args.mlp_mult,
+ rope_base=args.rope_base,
+ qk_gain_init=args.qk_gain_init,
+ rope_dims=args.rope_dims,
+ logit_softcap=args.logit_softcap,
+ tied_embed_init_std=args.tied_embed_init_std,
+ ngroups=args.mamba_ngroups,
+ depth_recurrence_start=args.depth_recurrence_start,
+ depth_recurrence_end=args.depth_recurrence_end,
+ depth_recurrence_loops=args.depth_recurrence_loops,
+ depth_recurrence_activation_frac=args.depth_recurrence_activation_frac,
+ ).to(device).bfloat16()
+ eval_model._training_progress = 1.0
+ eval_model._rec_loops = args.depth_recurrence_loops
+ restore_low_dim_params_to_fp32(eval_model)
+ eval_model.load_state_dict(deq_state, strict=True)
+ else:
+ eval_model = GPT(
+ vocab_size=args.vocab_size, num_layers=args.num_layers, model_dim=args.model_dim,
+ num_heads=args.num_heads, num_kv_heads=args.num_kv_heads, mlp_mult=args.mlp_mult,
+ tie_embeddings=args.tie_embeddings, tied_embed_init_std=args.tied_embed_init_std,
+ logit_softcap=args.logit_softcap, rope_base=args.rope_base, qk_gain_init=args.qk_gain_init,
+ mtp_num_heads=0, mtp_loss_weight=0.0,
+ bigram_vocab_size=args.bigram_vocab_size, bigram_dim=args.bigram_dim,
+ xsa_last_n=args.xsa_last_n,
+ rope_dims=args.rope_dims, ln_scale=args.ln_scale, dtg=args.dtg_enabled,
+ ve_enabled=args.ve_enabled, ve_dim=args.ve_dim, ve_layers=args.ve_layers,
+ gated_attention=args.gated_attention, value_residual=args.value_residual,
+ parallel_residual_start=args.parallel_residual_start,
+ depth_recurrence_start=args.depth_recurrence_start,
+ depth_recurrence_end=args.depth_recurrence_end,
+ depth_recurrence_loops=args.depth_recurrence_loops,
+ depth_recurrence_activation_frac=args.depth_recurrence_activation_frac,
+ ).to(device).bfloat16()
+ eval_model._training_progress = 1.0 # v5: enable depth recurrence in eval
+ eval_model._rec_loops = args.depth_recurrence_loops
+ eval_model.qo_bank.data = eval_model.qo_bank.data.float()
+ eval_model.kv_bank.data = eval_model.kv_bank.data.float()
+ eval_model.mlp_up_bank.data = eval_model.mlp_up_bank.data.float()
+ eval_model.mlp_down_bank.data = eval_model.mlp_down_bank.data.float()
+ for m in eval_model.modules():
+ if isinstance(m, CastedLinear):
+ m.float()
+ restore_low_dim_params_to_fp32(eval_model)
+ eval_model.load_state_dict(deq_state, strict=True)
+ # v8: Load bigram prior on eval model (after state_dict — buffer excluded from artifact)
+ if args.bigram_prior_enabled and args.bigram_prior_path and os.path.exists(args.bigram_prior_path):
+ eval_model.load_bigram_prior(args.bigram_prior_path)
+ # v5.2: run INT6 eval in eager mode — torch.compile(dynamic=True) wraps INT6
+ # scale tensors as SymFloat proxies, causing AttributeError on .size() calls
+ torch.cuda.synchronize()
+ t_qeval = time.perf_counter()
+ q_val_loss, q_val_bpb = eval_val(
+ args, eval_model, rank, world_size, device, grad_accum_steps,
+ val_tokens, base_bytes_lut, has_leading_space_lut, is_boundary_token_lut,
+ eval_seq_len=effective_eval_seq_len,
+ )
+ torch.cuda.synchronize()
+ log0(
+ f"final_int6_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} "
+ f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms"
+ )
+ log0(f"final_int6_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}")
+ # v5.3: create n-gram tilt for sliding window eval (prior-only, no live updates)
+ _sw_ngram_tilt: NgramTilt | None = None
+ if args.ngram_tilt_enabled and args.ngram_prior_path:
+ _prior_path = os.environ.get("NGRAM_PRIOR_PATH", args.ngram_prior_path)
+ if os.path.exists(_prior_path):
+ _sw_ngram_tilt = NgramTilt(
+ args.vocab_size, beta=args.ngram_tilt_beta, max_n=args.ngram_tilt_max_n
+ )
+ _sw_ngram_tilt.load_prior(_prior_path)
+ log0(f"ngram_tilt: loaded prior for sliding window eval from {_prior_path}")
+ sw_seq_len = effective_eval_seq_len
+ if args.eval_stride > 0 and args.eval_stride < sw_seq_len:
+ torch.cuda.synchronize()
+ t_slide = time.perf_counter()
+ _sw_batch = 2 if args.model_type == "mamba_hybrid" else 32
+ sw_val_loss, sw_val_bpb = eval_val_sliding(
+ args, eval_model, rank, world_size, device,
+ val_tokens, base_bytes_lut, has_leading_space_lut, is_boundary_token_lut,
+ stride=args.eval_stride,
+ batch_seqs=_sw_batch,
+ eval_seq_len=sw_seq_len,
+ ngram_tilt=_sw_ngram_tilt,
+ )
+ torch.cuda.synchronize()
+ log0(
+ f"final_int6_sliding_window val_loss:{sw_val_loss:.4f} val_bpb:{sw_val_bpb:.4f} "
+ f"stride:{args.eval_stride} eval_time:{1000.0 * (time.perf_counter() - t_slide):.0f}ms"
+ )
+ log0(f"final_int6_sliding_window_exact val_loss:{sw_val_loss:.8f} val_bpb:{sw_val_bpb:.8f}")
+ log0(f"final_int6_lzma_roundtrip_exact val_loss:{sw_val_loss:.8f} val_bpb:{sw_val_bpb:.8f}")
+ if args.eval_stride != 64 and 64 < sw_seq_len:
+ torch.cuda.synchronize()
+ t_slide64 = time.perf_counter()
+ sw64_val_loss, sw64_val_bpb = eval_val_sliding(
+ args, eval_model, rank, world_size, device,
+ val_tokens, base_bytes_lut, has_leading_space_lut, is_boundary_token_lut,
+ stride=64,
+ batch_seqs=_sw_batch,
+ eval_seq_len=sw_seq_len,
+ ngram_tilt=_sw_ngram_tilt,
+ )
+ torch.cuda.synchronize()
+ log0(
+ f"final_int6_sliding_window_s64 val_loss:{sw64_val_loss:.4f} val_bpb:{sw64_val_bpb:.4f} "
+ f"stride:64 eval_time:{1000.0 * (time.perf_counter() - t_slide64):.0f}ms"
+ )
+ log0(f"final_int6_sliding_window_s64_exact val_loss:{sw64_val_loss:.8f} val_bpb:{sw64_val_bpb:.8f}")
+ log0(f"final_int6_lzma_roundtrip_exact val_loss:{sw64_val_loss:.8f} val_bpb:{sw64_val_bpb:.8f}")
+ # === Legal TTT eval (if enabled) ===
+ if args.ttt_enabled:
+ log0("ttt_eval:starting")
+ torch.cuda.synchronize()
+ t_ttt = time.perf_counter()
+ ttt_val_loss, ttt_val_bpb = _run_ttt_sliding_window_eval(
+ model=eval_model,
+ val_tokens=val_tokens,
+ seq_len=sw_seq_len,
+ stride=args.eval_stride,
+ base_bytes_lut=base_bytes_lut,
+ has_leading_space_lut=has_leading_space_lut,
+ is_boundary_token_lut=is_boundary_token_lut,
+ device=device,
+ ttt_lr=args.ttt_lr,
+ ttt_momentum=args.ttt_momentum,
+ ttt_epochs=args.ttt_epochs,
+ ttt_grad_clip=args.ttt_grad_clip,
+ ngram_tilt=_sw_ngram_tilt,
+ chunk_size=args.ttt_chunk_size,
+ )
+ torch.cuda.synchronize()
+ log0(
+ f"final_ttt_sliding_window val_loss:{ttt_val_loss:.4f} val_bpb:{ttt_val_bpb:.4f} "
+ f"eval_time:{1000.0 * (time.perf_counter() - t_ttt):.0f}ms"
+ )
+ log0(f"final_ttt_sliding_window_exact val_loss:{ttt_val_loss:.8f} val_bpb:{ttt_val_bpb:.8f}")
+ if distributed:
+ dist.destroy_process_group()
+if __name__ == "__main__":
+ main()
diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_log_run1.txt b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_log_run1.txt
new file mode 100644
index 0000000000..bc9d33618a
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_log_run1.txt
@@ -0,0 +1,191 @@
+
+========================================================================
+ BESE v7-mamba: Mamba-3 + Attention Hybrid
+========================================================================
+ Start time: 2026-04-16 06:41:11
+ Architecture: 6 Mamba-3 + 2 Attention (pos 2,5), dim=512, d_state=128, ngroups=1, BESE 288 vocab
+ Thesis: BESE's 2x token density is an advantage with O(n) SSMs
+
+========================================================================
+ Installing dependencies
+========================================================================
+
+ [pip-mamba-ssm] Running: /usr/local/bin/python -m pip install mamba-ssm causal-conv1d einops --quiet --no-build-isolation --break-system-packages
+WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.
+ [pip-mamba-ssm] Completed in 0.8s (0.0 min)
+ Found 44 training shards
+ Detected 8 GPUs
+
+========================================================================
+ PHASE 1: MAMBA HYBRID TRAINING (600s wallclock, 8 GPUs)
+========================================================================
+ Shards: 44 train, 1 val
+ Training env:
+ ADAM_WD=0.095
+ ATTN_LAYER_POS=2,5
+ BESE_TOKENIZER_ROOT=/workspace/bese/tokenizer
+ BIGRAM_PRIOR_ENABLED=0
+ DATA_PATH=/runpod-volume/bese_shards_v5
+ DEPTH_RECURRENCE_ACTIVATION_FRAC=1.0
+ DEPTH_RECURRENCE_END=0
+ DEPTH_RECURRENCE_LOOPS=1
+ DEPTH_RECURRENCE_START=0
+ D_STATE=128
+ EMA_DECAY=0.9965
+ EVAL_SEQ_LEN=2048
+ EVAL_STRIDE=64
+ LATE_QAT_THRESHOLD=0
+ MAMBA_CHUNK_SIZE=64
+ MAMBA_EXPAND=2
+ MAMBA_HEADDIM=64
+ MAMBA_NGROUPS=1
+ MATRIX_LR=0.026
+ MAX_WALLCLOCK_SECONDS=600
+ MLP_MULT=3.0
+ MODEL_DIM=512
+ MODEL_TYPE=mamba_hybrid
+ MUON_WD=0.095
+ NGRAM_PRIOR_PATH=/runpod-volume/artifacts/ngram_table_v6.bin
+ NGRAM_TILT_ENABLED=1
+ NGRAM_TILT_MAX_N=3
+ NOISY_QAT_ENABLED=0
+ NUM_HEADS=8
+ NUM_KV_HEADS=4
+ NUM_LAYERS=8
+ PYTHONPATH=/workspace/bese
+ QAT_ENABLED=0
+ QK_GAIN_INIT=5.25
+ RUN_ID=bese_v7_mamba
+ TOKENIZER_PATH=/runpod-volume/tokenizers/bese_bpe_248_v5.json
+ TRAIN_LOG_EVERY=100
+ TRAIN_SEQ_LEN=2048
+ TTT_CHUNK_SIZE=32768
+ TTT_ENABLED=0
+ TTT_EPOCHS=1
+ TTT_GRAD_CLIP=1.0
+ TTT_LR=0.005
+ TTT_MOMENTUM=0.9
+ VAL_LOSS_EVERY=500
+ VOCAB_SIZE=288
+ WARMDOWN_ITERS=5000
+
+ [torchrun] Running: torchrun --standalone --nproc_per_node=8 /workspace/bese/integration/train_gpt_bese.py
+W0416 06:41:14.972000 3190 torch/distributed/run.py:803]
+W0416 06:41:14.972000 3190 torch/distributed/run.py:803] *****************************************
+W0416 06:41:14.972000 3190 torch/distributed/run.py:803] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
+W0416 06:41:14.972000 3190 torch/distributed/run.py:803] *****************************************
+logs/bese_v7_mamba.txt
+val_bpb:enabled tokenizer_kind=bese_bpe tokenizer_path=/runpod-volume/tokenizers/bese_bpe_248_v5.json vocab_size=288
+train_loader:dataset:bese_shards_v5 train_shards:44
+val_loader:shards pattern=/runpod-volume/bese_shards_v5/fineweb_val_*.bin tokens:76632064
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+torch.compile:skipped (mamba_hybrid, Triton kernel is the fast path)
+model_params:15152432
+model_type:mamba_hybrid d_state:128 expand:2 ngroups:1 attn_pos:[2, 5]
+world_size:8 grad_accum_steps:1
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+sdp_backends:cudnn=False flash=True mem_efficient=False math=False
+attention_mode:gqa num_heads:8 num_kv_heads:4
+tie_embeddings:True embed_lr:0.035 head_lr:0.0 matrix_lr:0.026 scalar_lr:0.025
+train_batch_tokens:786432 train_seq_len:2048 iterations:20000 warmup_steps:20 max_wallclock_seconds:600.000
+seed:1337
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+warmup_step:1/20
+warmup_step:2/20
+warmup_step:3/20
+warmup_step:4/20
+warmup_step:5/20
+warmup_step:6/20
+warmup_step:7/20
+warmup_step:8/20
+warmup_step:9/20
+warmup_step:10/20
+warmup_step:11/20
+warmup_step:12/20
+warmup_step:13/20
+warmup_step:14/20
+warmup_step:15/20
+warmup_step:16/20
+warmup_step:17/20
+warmup_step:18/20
+warmup_step:19/20
+warmup_step:20/20
+step:0/20000 val_loss:5.6808 val_bpb:4.1571 train_time:0ms step_avg:0.01ms
+step:1/20000 train_loss:5.6067 train_time:343ms step_avg:343.47ms
+step:2/20000 train_loss:9.2367 train_time:647ms step_avg:323.62ms
+step:3/20000 train_loss:7.3239 train_time:927ms step_avg:309.13ms
+step:4/20000 train_loss:5.9848 train_time:1206ms step_avg:301.58ms
+step:5/20000 train_loss:6.0772 train_time:1511ms step_avg:302.26ms
+step:6/20000 train_loss:5.7907 train_time:1823ms step_avg:303.76ms
+step:7/20000 train_loss:5.9782 train_time:2134ms step_avg:304.81ms
+step:8/20000 train_loss:5.7510 train_time:2444ms step_avg:305.46ms
+step:9/20000 train_loss:5.5104 train_time:2725ms step_avg:302.81ms
+step:10/20000 train_loss:6.0598 train_time:3005ms step_avg:300.50ms
+step:100/20000 train_loss:2.8180 train_time:27511ms step_avg:275.11ms
+step:200/20000 train_loss:2.3938 train_time:54885ms step_avg:274.42ms
+step:300/20000 train_loss:2.2195 train_time:82293ms step_avg:274.31ms
+step:400/20000 train_loss:2.0719 train_time:109671ms step_avg:274.18ms
+step:500/20000 train_loss:1.9991 train_time:136924ms step_avg:273.85ms
+step:500/20000 val_loss:2.1127 val_bpb:1.5460 train_time:136962ms step_avg:273.92ms
+step:600/20000 train_loss:1.9919 train_time:164358ms step_avg:273.93ms
+step:700/20000 train_loss:1.9344 train_time:191760ms step_avg:273.94ms
+step:800/20000 train_loss:1.9071 train_time:219112ms step_avg:273.89ms
+step:900/20000 train_loss:1.9019 train_time:246474ms step_avg:273.86ms
+step:1000/20000 train_loss:1.9488 train_time:273699ms step_avg:273.70ms
+step:1000/20000 val_loss:1.9497 val_bpb:1.4268 train_time:273713ms step_avg:273.71ms
+step:1100/20000 train_loss:1.8819 train_time:301030ms step_avg:273.66ms
+swa:start step:1200
+step:1200/20000 train_loss:1.8882 train_time:328375ms step_avg:273.65ms
+step:1300/20000 train_loss:1.9037 train_time:355888ms step_avg:273.76ms
+step:1400/20000 train_loss:1.7942 train_time:383255ms step_avg:273.75ms
+step:1500/20000 train_loss:1.7417 train_time:410513ms step_avg:273.68ms
+step:1500/20000 val_loss:1.8866 val_bpb:1.3806 train_time:410542ms step_avg:273.69ms
+step:1600/20000 train_loss:1.8532 train_time:437902ms step_avg:273.69ms
+step:1700/20000 train_loss:1.8551 train_time:465397ms step_avg:273.76ms
+step:1800/20000 train_loss:1.7673 train_time:492787ms step_avg:273.77ms
+step:1900/20000 train_loss:1.7762 train_time:520052ms step_avg:273.71ms
+step:2000/20000 train_loss:1.9196 train_time:547511ms step_avg:273.76ms
+step:2000/20000 val_loss:1.8433 val_bpb:1.3489 train_time:547568ms step_avg:273.78ms
+step:2100/20000 train_loss:1.8175 train_time:574913ms step_avg:273.77ms
+step:2191/20000 val_loss:1.8391 val_bpb:1.3458 train_time:600028ms step_avg:273.86ms
+stopping_early: wallclock_cap train_time:600028ms step:2191/20000
+peak memory allocated: 41596 MiB reserved: 42438 MiB
+ema:applying EMA weights
+DIAGNOSTIC post_ema val_loss:1.8414 val_bpb:1.3475 eval_time:9444ms
+Serialized model: 30357515 bytes
+Code size: 108633 bytes
+Bundled n-gram prior: 298,802 bytes from /runpod-volume/artifacts/ngram_table_v6.bin
+Serialized model int6+lzma: 7452680 bytes
+Total submission size int6+lzma: 7561313 bytes
+Extracted bundled n-gram prior: 298,802 bytes → /tmp/tmpwuww5dk4.bin
+final_int6_roundtrip val_loss:1.8870 val_bpb:1.3809 eval_time:9447ms
+final_int6_roundtrip_exact val_loss:1.88698964 val_bpb:1.38086544
+ngram_tilt: loaded prior for sliding window eval from /tmp/tmpwuww5dk4.bin
+final_int6_sliding_window val_loss:1.8545 val_bpb:1.3571 stride:64 eval_time:401929ms
+final_int6_sliding_window_exact val_loss:1.85446154 val_bpb:1.35706229
+final_int6_lzma_roundtrip_exact val_loss:1.85446154 val_bpb:1.35706229
+ [torchrun] Completed in 1129.5s (18.8 min)
+
+========================================================================
+ RUN SUMMARY
+========================================================================
+ int6_bpb: 1.3809
+ model_bytes: 30357515
+ size_bytes: 7561313
+ sliding_bpb: 1.3571
+
+ Total wall time: 1130.5s (18.8 min)
+ Saved artifact to /runpod-volume/checkpoints/v7_mamba/final_model.int6.ptz
diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_log_run2_d64.txt b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_log_run2_d64.txt
new file mode 100644
index 0000000000..0f47409dcc
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_log_run2_d64.txt
@@ -0,0 +1,179 @@
+
+========================================================================
+ BESE v7-mamba: Mamba-3 + Attention Hybrid
+========================================================================
+ Start time: 2026-04-16 06:23:17
+ Architecture: 6 Mamba-3 + 2 Attention (pos 2,5), dim=512, d_state=128, ngroups=1, BESE 288 vocab
+ Thesis: BESE's 2x token density is an advantage with O(n) SSMs
+
+========================================================================
+ Installing dependencies
+========================================================================
+
+ [pip-mamba-ssm] Running: /usr/local/bin/python -m pip install mamba-ssm causal-conv1d einops --quiet --no-build-isolation --break-system-packages
+WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.
+ [pip-mamba-ssm] Completed in 0.5s (0.0 min)
+ Found 44 training shards
+ Detected 8 GPUs
+
+========================================================================
+ PHASE 1: MAMBA HYBRID TRAINING (600s wallclock, 8 GPUs)
+========================================================================
+ Shards: 44 train, 1 val
+ Training env:
+ ADAM_WD=0.095
+ ATTN_LAYER_POS=2,5
+ BESE_TOKENIZER_ROOT=/workspace/bese/tokenizer
+ BIGRAM_PRIOR_ENABLED=0
+ DATA_PATH=/runpod-volume/bese_shards_v5
+ DEPTH_RECURRENCE_ACTIVATION_FRAC=1.0
+ DEPTH_RECURRENCE_END=0
+ DEPTH_RECURRENCE_LOOPS=1
+ DEPTH_RECURRENCE_START=0
+ D_STATE=64
+ EMA_DECAY=0.9965
+ EVAL_SEQ_LEN=2048
+ EVAL_STRIDE=64
+ LATE_QAT_THRESHOLD=0
+ MAMBA_CHUNK_SIZE=64
+ MAMBA_EXPAND=2
+ MAMBA_HEADDIM=64
+ MAMBA_NGROUPS=1
+ MATRIX_LR=0.026
+ MAX_WALLCLOCK_SECONDS=600
+ MLP_MULT=3.0
+ MODEL_DIM=512
+ MODEL_TYPE=mamba_hybrid
+ MUON_WD=0.095
+ NGRAM_PRIOR_PATH=/runpod-volume/artifacts/ngram_table_v6.bin
+ NGRAM_TILT_ENABLED=1
+ NGRAM_TILT_MAX_N=3
+ NOISY_QAT_ENABLED=0
+ NUM_HEADS=8
+ NUM_KV_HEADS=4
+ NUM_LAYERS=8
+ PYTHONPATH=/workspace/bese
+ QAT_ENABLED=0
+ QK_GAIN_INIT=5.25
+ RUN_ID=bese_v7_mamba
+ TOKENIZER_PATH=/runpod-volume/tokenizers/bese_bpe_248_v5.json
+ TRAIN_LOG_EVERY=100
+ TRAIN_SEQ_LEN=2048
+ TTT_CHUNK_SIZE=32768
+ TTT_ENABLED=0
+ TTT_EPOCHS=1
+ TTT_GRAD_CLIP=1.0
+ TTT_LR=0.005
+ TTT_MOMENTUM=0.9
+ VAL_LOSS_EVERY=500
+ VOCAB_SIZE=288
+ WARMDOWN_ITERS=5000
+
+ [torchrun] Running: torchrun --standalone --nproc_per_node=8 /workspace/bese/integration/train_gpt_bese.py
+W0416 06:23:19.869000 393 torch/distributed/run.py:803]
+W0416 06:23:19.869000 393 torch/distributed/run.py:803] *****************************************
+W0416 06:23:19.869000 393 torch/distributed/run.py:803] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
+W0416 06:23:19.869000 393 torch/distributed/run.py:803] *****************************************
+logs/bese_v7_mamba.txt
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+val_bpb:enabled tokenizer_kind=bese_bpe tokenizer_path=/runpod-volume/tokenizers/bese_bpe_248_v5.json vocab_size=288
+train_loader:dataset:bese_shards_v5 train_shards:44
+val_loader:shards pattern=/runpod-volume/bese_shards_v5/fineweb_val_*.bin tokens:76632064
+torch.compile:skipped (mamba_hybrid, Triton kernel is the fast path)
+model_params:14758448
+model_type:mamba_hybrid d_state:64 expand:2 ngroups:1 attn_pos:[2, 5]
+world_size:8 grad_accum_steps:1
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+sdp_backends:cudnn=False flash=True mem_efficient=False math=False
+attention_mode:gqa num_heads:8 num_kv_heads:4
+tie_embeddings:True embed_lr:0.035 head_lr:0.0 matrix_lr:0.026 scalar_lr:0.025
+train_batch_tokens:786432 train_seq_len:2048 iterations:20000 warmup_steps:20 max_wallclock_seconds:600.000
+seed:1337
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+warmup_step:1/20
+warmup_step:2/20
+warmup_step:3/20
+warmup_step:4/20
+warmup_step:5/20
+warmup_step:6/20
+warmup_step:7/20
+warmup_step:8/20
+warmup_step:9/20
+warmup_step:10/20
+warmup_step:11/20
+warmup_step:12/20
+warmup_step:13/20
+warmup_step:14/20
+warmup_step:15/20
+warmup_step:16/20
+warmup_step:17/20
+warmup_step:18/20
+warmup_step:19/20
+warmup_step:20/20
+step:0/20000 val_loss:5.6808 val_bpb:4.1571 train_time:0ms step_avg:0.01ms
+step:1/20000 train_loss:5.6067 train_time:299ms step_avg:299.44ms
+step:2/20000 train_loss:9.2128 train_time:524ms step_avg:261.99ms
+step:3/20000 train_loss:7.0552 train_time:752ms step_avg:250.58ms
+step:4/20000 train_loss:6.3110 train_time:982ms step_avg:245.58ms
+step:5/20000 train_loss:6.0682 train_time:1218ms step_avg:243.68ms
+step:6/20000 train_loss:7.1301 train_time:1448ms step_avg:241.36ms
+step:7/20000 train_loss:5.8157 train_time:1679ms step_avg:239.89ms
+step:8/20000 train_loss:6.0502 train_time:1913ms step_avg:239.18ms
+step:9/20000 train_loss:5.6456 train_time:2142ms step_avg:237.97ms
+step:10/20000 train_loss:5.3274 train_time:2373ms step_avg:237.26ms
+step:100/20000 train_loss:2.7963 train_time:22209ms step_avg:222.09ms
+step:200/20000 train_loss:2.3271 train_time:44369ms step_avg:221.84ms
+step:300/20000 train_loss:2.1084 train_time:66756ms step_avg:222.52ms
+step:400/20000 train_loss:1.9847 train_time:88923ms step_avg:222.31ms
+step:500/20000 train_loss:1.9546 train_time:110963ms step_avg:221.93ms
+step:500/20000 val_loss:2.0728 val_bpb:1.5168 train_time:111014ms step_avg:222.03ms
+step:600/20000 train_loss:1.9583 train_time:133209ms step_avg:222.02ms
+step:700/20000 train_loss:1.9098 train_time:155370ms step_avg:221.96ms
+step:800/20000 train_loss:1.8891 train_time:177637ms step_avg:222.05ms
+step:900/20000 train_loss:1.8852 train_time:199850ms step_avg:222.06ms
+step:1000/20000 train_loss:1.9317 train_time:221906ms step_avg:221.91ms
+step:1000/20000 val_loss:1.9346 val_bpb:1.4157 train_time:221929ms step_avg:221.93ms
+step:1100/20000 train_loss:1.8665 train_time:244220ms step_avg:222.02ms
+step:1200/20000 train_loss:1.8726 train_time:270219ms step_avg:225.18ms
+step:1300/20000 train_loss:1.8927 train_time:296728ms step_avg:228.25ms
+step:1400/20000 train_loss:1.7816 train_time:324143ms step_avg:231.53ms
+step:1500/20000 train_loss:1.7280 train_time:346233ms step_avg:230.82ms
+step:1500/20000 val_loss:1.8749 val_bpb:1.3720 train_time:346284ms step_avg:230.86ms
+swa:start step:1600
+step:1600/20000 train_loss:1.8388 train_time:375188ms step_avg:234.49ms
+step:1700/20000 train_loss:1.8406 train_time:400637ms step_avg:235.67ms
+step:1800/20000 train_loss:1.7537 train_time:428145ms step_avg:237.86ms
+step:1900/20000 train_loss:1.7648 train_time:450224ms step_avg:236.96ms
+step:2000/20000 train_loss:1.9046 train_time:476517ms step_avg:238.26ms
+step:2000/20000 val_loss:1.8287 val_bpb:1.3382 train_time:476555ms step_avg:238.28ms
+step:2100/20000 train_loss:1.8045 train_time:502294ms step_avg:239.19ms
+step:2200/20000 train_loss:1.8343 train_time:528033ms step_avg:240.02ms
+step:2300/20000 train_loss:1.7125 train_time:555838ms step_avg:241.67ms
+step:2400/20000 train_loss:1.7194 train_time:577916ms step_avg:240.80ms
+step:2482/20000 val_loss:1.8112 val_bpb:1.3254 train_time:599898ms step_avg:241.70ms
+stopping_early: wallclock_cap train_time:599898ms step:2482/20000
+peak memory allocated: 32456 MiB reserved: 33996 MiB
+ema:applying EMA weights
+DIAGNOSTIC post_ema val_loss:1.8123 val_bpb:1.3262 eval_time:7471ms
+Serialized model: 29568011 bytes
+Code size: 108498 bytes
+Bundled n-gram prior: 298,802 bytes from /runpod-volume/artifacts/ngram_table_v6.bin
+Serialized model int6+lzma: 7850968 bytes
+Total submission size int6+lzma: 7959466 bytes
+Extracted bundled n-gram prior: 298,802 bytes → /tmp/tmpqzlx37jt.bin
+final_int6_roundtrip val_loss:1.8373 val_bpb:1.3445 eval_time:7464ms
+final_int6_roundtrip_exact val_loss:1.83732240 val_bpb:1.34451984
+ngram_tilt: loaded prior for sliding window eval from /tmp/tmpqzlx37jt.bin
diff --git a/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_log_run3_dim576.txt b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_log_run3_dim576.txt
new file mode 100644
index 0000000000..2efa8064ee
--- /dev/null
+++ b/records/track_non_record_16mb/2026-04-16_BESE_Mamba3_Hybrid/train_log_run3_dim576.txt
@@ -0,0 +1,172 @@
+
+========================================================================
+ BESE v7-mamba: Mamba-3 + Attention Hybrid
+========================================================================
+ Start time: 2026-04-16 07:03:45
+ Architecture: 6 Mamba-3 + 2 Attention (pos 2,5), dim=512, d_state=128, ngroups=1, BESE 288 vocab
+ Thesis: BESE's 2x token density is an advantage with O(n) SSMs
+
+========================================================================
+ Installing dependencies
+========================================================================
+
+ [pip-mamba-ssm] Running: /usr/local/bin/python -m pip install mamba-ssm causal-conv1d einops --quiet --no-build-isolation --break-system-packages
+WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.
+ [pip-mamba-ssm] Completed in 0.5s (0.0 min)
+ Found 44 training shards
+ Detected 8 GPUs
+
+========================================================================
+ PHASE 1: MAMBA HYBRID TRAINING (600s wallclock, 8 GPUs)
+========================================================================
+ Shards: 44 train, 1 val
+ Training env:
+ ADAM_WD=0.095
+ ATTN_LAYER_POS=2,5
+ BESE_TOKENIZER_ROOT=/workspace/bese/tokenizer
+ BIGRAM_PRIOR_ENABLED=0
+ DATA_PATH=/runpod-volume/bese_shards_v5
+ DEPTH_RECURRENCE_ACTIVATION_FRAC=1.0
+ DEPTH_RECURRENCE_END=0
+ DEPTH_RECURRENCE_LOOPS=1
+ DEPTH_RECURRENCE_START=0
+ D_STATE=128
+ EMA_DECAY=0.9965
+ EVAL_SEQ_LEN=2048
+ EVAL_STRIDE=64
+ LATE_QAT_THRESHOLD=0
+ MAMBA_CHUNK_SIZE=64
+ MAMBA_EXPAND=2
+ MAMBA_HEADDIM=64
+ MAMBA_NGROUPS=1
+ MATRIX_LR=0.026
+ MAX_WALLCLOCK_SECONDS=600
+ MLP_MULT=3.5
+ MODEL_DIM=576
+ MODEL_TYPE=mamba_hybrid
+ MUON_WD=0.095
+ NGRAM_PRIOR_PATH=/runpod-volume/artifacts/ngram_table_v6.bin
+ NGRAM_TILT_ENABLED=1
+ NGRAM_TILT_MAX_N=3
+ NOISY_QAT_ENABLED=0
+ NUM_HEADS=8
+ NUM_KV_HEADS=4
+ NUM_LAYERS=8
+ PYTHONPATH=/workspace/bese
+ QAT_ENABLED=0
+ QK_GAIN_INIT=5.25
+ RUN_ID=bese_v7_mamba
+ TOKENIZER_PATH=/runpod-volume/tokenizers/bese_bpe_248_v5.json
+ TRAIN_LOG_EVERY=100
+ TRAIN_SEQ_LEN=2048
+ TTT_CHUNK_SIZE=32768
+ TTT_ENABLED=0
+ TTT_EPOCHS=1
+ TTT_GRAD_CLIP=1.0
+ TTT_LR=0.005
+ TTT_MOMENTUM=0.9
+ VAL_LOSS_EVERY=500
+ VOCAB_SIZE=288
+ WARMDOWN_ITERS=5000
+
+ [torchrun] Running: torchrun --standalone --nproc_per_node=8 /workspace/bese/integration/train_gpt_bese.py
+W0416 07:03:48.426000 4539 torch/distributed/run.py:803]
+W0416 07:03:48.426000 4539 torch/distributed/run.py:803] *****************************************
+W0416 07:03:48.426000 4539 torch/distributed/run.py:803] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
+W0416 07:03:48.426000 4539 torch/distributed/run.py:803] *****************************************
+logs/bese_v7_mamba.txt
+val_bpb:enabled tokenizer_kind=bese_bpe tokenizer_path=/runpod-volume/tokenizers/bese_bpe_248_v5.json vocab_size=288
+train_loader:dataset:bese_shards_v5 train_shards:44
+val_loader:shards pattern=/runpod-volume/bese_shards_v5/fineweb_val_*.bin tokens:76632064
+torch.compile:skipped (mamba_hybrid, Triton kernel is the fast path)
+model_params:19707412
+model_type:mamba_hybrid d_state:128 expand:2 ngroups:1 attn_pos:[2, 5]
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+world_size:8 grad_accum_steps:1
+sdp_backends:cudnn=False flash=True mem_efficient=False math=False
+attention_mode:gqa num_heads:8 num_kv_heads:4
+tie_embeddings:True embed_lr:0.035 head_lr:0.0 matrix_lr:0.026 scalar_lr:0.025
+train_batch_tokens:786432 train_seq_len:2048 iterations:20000 warmup_steps:20 max_wallclock_seconds:600.000
+seed:1337
+/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2920: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /pytorch/aten/src/ATen/native/layer_norm.cpp:344.)
+ return torch.rms_norm(input, normalized_shape, weight, eps)
+warmup_step:1/20
+warmup_step:2/20
+warmup_step:3/20
+warmup_step:4/20
+warmup_step:5/20
+warmup_step:6/20
+warmup_step:7/20
+warmup_step:8/20
+warmup_step:9/20
+warmup_step:10/20
+warmup_step:11/20
+warmup_step:12/20
+warmup_step:13/20
+warmup_step:14/20
+warmup_step:15/20
+warmup_step:16/20
+warmup_step:17/20
+warmup_step:18/20
+warmup_step:19/20
+warmup_step:20/20
+step:0/20000 val_loss:5.6905 val_bpb:4.1642 train_time:0ms step_avg:0.01ms
+step:1/20000 train_loss:5.6097 train_time:402ms step_avg:401.59ms
+step:2/20000 train_loss:10.0167 train_time:757ms step_avg:378.64ms
+step:3/20000 train_loss:8.1614 train_time:1089ms step_avg:362.89ms
+step:4/20000 train_loss:6.6466 train_time:1451ms step_avg:362.66ms
+step:5/20000 train_loss:6.2076 train_time:1781ms step_avg:356.29ms
+step:6/20000 train_loss:5.9359 train_time:2112ms step_avg:351.93ms
+step:7/20000 train_loss:5.4867 train_time:2472ms step_avg:353.16ms
+step:8/20000 train_loss:7.4613 train_time:2810ms step_avg:351.28ms
+step:9/20000 train_loss:5.4067 train_time:3141ms step_avg:348.97ms
+step:10/20000 train_loss:4.8497 train_time:3511ms step_avg:351.05ms
+step:100/20000 train_loss:2.7812 train_time:32602ms step_avg:326.02ms
+step:200/20000 train_loss:2.3231 train_time:65086ms step_avg:325.43ms
+step:300/20000 train_loss:2.1042 train_time:97600ms step_avg:325.33ms
+step:400/20000 train_loss:1.9740 train_time:130093ms step_avg:325.23ms
+step:500/20000 train_loss:1.9478 train_time:162428ms step_avg:324.86ms
+step:500/20000 val_loss:2.0630 val_bpb:1.5097 train_time:162441ms step_avg:324.88ms
+step:600/20000 train_loss:1.9510 train_time:194879ms step_avg:324.80ms
+step:700/20000 train_loss:1.9020 train_time:227317ms step_avg:324.74ms
+step:800/20000 train_loss:1.8759 train_time:259769ms step_avg:324.71ms
+swa:start step:850
+step:900/20000 train_loss:1.8707 train_time:292294ms step_avg:324.77ms
+step:1000/20000 train_loss:1.9150 train_time:324727ms step_avg:324.73ms
+step:1000/20000 val_loss:1.9184 val_bpb:1.4038 train_time:324774ms step_avg:324.77ms
+step:1100/20000 train_loss:1.8442 train_time:357287ms step_avg:324.81ms
+step:1200/20000 train_loss:1.8574 train_time:389845ms step_avg:324.87ms
+step:1300/20000 train_loss:1.8761 train_time:422495ms step_avg:325.00ms
+step:1400/20000 train_loss:1.7627 train_time:455064ms step_avg:325.05ms
+step:1500/20000 train_loss:1.7030 train_time:487496ms step_avg:325.00ms
+step:1500/20000 val_loss:1.8534 val_bpb:1.3563 train_time:487528ms step_avg:325.02ms
+step:1600/20000 train_loss:1.8179 train_time:519979ms step_avg:324.99ms
+step:1700/20000 train_loss:1.8258 train_time:552485ms step_avg:324.99ms
+step:1800/20000 train_loss:1.7433 train_time:584978ms step_avg:324.99ms
+step:1847/20000 val_loss:1.8332 val_bpb:1.3415 train_time:600197ms step_avg:324.96ms
+stopping_early: wallclock_cap train_time:600197ms step:1847/20000
+peak memory allocated: 46850 MiB reserved: 47580 MiB
+ema:applying EMA weights
+DIAGNOSTIC post_ema val_loss:1.8383 val_bpb:1.3452 eval_time:11114ms
+Serialized model: 39470347 bytes
+Code size: 108633 bytes
+Bundled n-gram prior: 298,802 bytes from /runpod-volume/artifacts/ngram_table_v6.bin
+Serialized model int6+lzma: 8311720 bytes
+Total submission size int6+lzma: 8420353 bytes
+Extracted bundled n-gram prior: 298,802 bytes → /tmp/tmp7yqe1ver.bin
+final_int6_roundtrip val_loss:1.9204 val_bpb:1.4053 eval_time:11150ms
+final_int6_roundtrip_exact val_loss:1.92043110 val_bpb:1.40533731
+ngram_tilt: loaded prior for sliding window eval from /tmp/tmp7yqe1ver.bin