Cukinator commited on
Commit
be12dfa
·
verified ·
1 Parent(s): d7ef1d4

Rebuild model card from reproducible CPU audit; include evidence and limitations

Browse files
README.md CHANGED
@@ -13,7 +13,7 @@ tags:
13
 
14
  1,346 checkpoint files across 34 run folders, totaling **215.178 GB** (200.400 GiB). The inventory includes **21 `_final.pt` files** and **1325 step files**. Counts and sizes are derived from the pinned Hub inventory, not estimates from nominal architecture sizes.
15
 
16
- For evaluated FP32 exports, see [Cukinator/cpu1-ablations-final](https://huggingface.co/Cukinator/cpu1-ablations-final). Source code: [Cukinator/1.58bits](https://github.com/igna-s/New-1.58).
17
 
18
  ## What the audit found
19
 
@@ -89,7 +89,7 @@ from huggingface_hub import hf_hub_download
89
  torch.set_num_threads(1)
90
  path = hf_hub_download("Cukinator/cpu1-ablation-checkpoints", "run_04/checkpoint_run_04_final.pt",
91
  revision="fb402d025e843ed00fb990966bf2430c32eb958c", local_dir="D:/cpu1/active", cache_dir="D:/cpu1/hf/hub")
92
- from train_ablation import build_ablation_model, load_ablation_checkpoint
93
  state, config = load_ablation_checkpoint(path)
94
  config["use_gradient_checkpointing"] = False
95
  model = build_ablation_model(config)
@@ -101,14 +101,14 @@ with torch.inference_mode():
101
  print([tuple(head.shape) for head in logits])
102
  ```
103
 
104
- This forward uses zero target bytes in the local decoder; it is a shape smoke test. For scored likelihood, pass `byte_targets` as shown in the audit. For byte generation with retained Transformer context and one recurrent prefill, use `generation()` in `eval/audit_readmes.py`. Generation remains experimental; inspect the saved examples before choosing a checkpoint.
105
 
106
  ## Resume training
107
 
108
  Use an intermediate checkpoint, not an inference-only packed final. The current source CLI names the argument `--resume_ckpt`:
109
 
110
  ```powershell
111
- python train_ablation.py --run run_04 --resume_ckpt D:/cpu1/active/run_04/checkpoint_run_04_step640.pt
112
  ```
113
 
114
  Download the exact step file from its run folder first. Resume requires matching source/configuration, dependencies, and suitable training hardware. This audit did not run a training-resume experiment.
@@ -128,7 +128,7 @@ This card was rebuilt from a fresh CPU audit on **2026-10-06**. All 34 published
128
  - Shared sample JSON SHA-256: `11d143c06478ff7bbe364c541a58e6a2e76504c4cb0595b294f46e98f3cfe798`.
129
  - Byte targets are the next patch; incomplete final patches are discarded. The first patch is context only. LocalByteDecoder receives ground-truth target bytes and shifts them internally for causal teacher forcing within a patch. Cross-entropy is summed over all targets and divided by their total count.
130
  - Byte loss is nats/byte; `BPB = NLL / ln(2)` and byte PPL is `exp(NLL)`. A uniform byte predictor gives BPB=8 and PPL=256. BPE loss uses Qwen2.5-3B token IDs with IDs outside the configured vocabulary mapped to 0, matching the dataset adapter. Its PPL is over that lossy mapped token stream and cannot be compared with byte PPL or interpreted as the likelihood of the original text.
131
- - All BPE models shared one resolved Qwen tokenizer in this run. Its resolved revision is recorded in `run_01.json` and `environment.json`; reuse that tokenizer revision/cache for strict reproduction. The current evaluator resolves the tokenizer Hub name at runtime rather than hard-coding that revision.
132
  - Causality: change the last 16 of 32 input steps and compare all prefix logits, tolerance `1e-5`. Recurrent chunk equivalence: compare a 128-step forward against four 32-step chunks across all output heads, tolerance `1e-4`. Failed checks remain visible.
133
  - DeleteGate is disabled consistently: phase 1, elimination rate 0. The audit does not measure token elimination, phase-2 quality, or the benefit of DeleteGate.
134
  - Forward timing: 10 warm-up steps, then 100 measured single-input forwards. `steps/s = 100 / total measured seconds`; p50 and p95 are per-step latency. One byte-model step consumes a patch of 4 bytes; one BPE step consumes one token. Local decoder inputs are zeros in this microbenchmark. Transformer timing has one input step and no persistent attention context. This is not a context-scaling test or a generation benchmark.
@@ -141,14 +141,14 @@ Export repository revision audited: `65951289352c1a036b12ba6d888af037e7c90b13`.
141
 
142
  From the source checkout, install PyTorch plus `huggingface_hub`, `transformers`, `tokenizers`, `pyarrow`, `psutil`, `numpy`, and `datasets`. Pin the versions from `evaluation/environment.json` when reproducing this run.
143
 
144
- The audit additions are bundled in `evaluation/source/`. Copy those Python files into the checkout’s `eval/` directory when using a source revision that predates this audit. The base Git commit is recorded in `evaluation/environment.json`.
145
 
146
  ```powershell
147
- python eval/audit_readmes.py --storage D:/cpu1-readme-audit --threads 1
148
- python eval/audit_readmes.py --storage D:/cpu1-readme-audit --cards
149
  ```
150
 
151
- Model downloads, tokenizer files, Hugging Face caches, and temporary files go to D:. Only one model checkpoint is kept in the task storage at a time; it is removed after its result is saved. Interrupted runs resume from matching evaluator hashes. Raw reports remain in `resultados/readme_audit/`.
152
 
153
  ## Fresh diagnostic results
154
 
@@ -310,7 +310,7 @@ These configurations did not reproduce a full-sequence forward when carrying sta
310
 
311
  ### Recurrent implementation diagnostic
312
 
313
- The source’s `AblationMLGRU._parallel_scan` was tested independently against its defining serial recurrence with random candidate/forget tensors, seed 19, shape `[1,32,8]`, FP32. The maximum difference was **19.100506**; splitting the same scan into two 16-step chunks differed by **11.471270** (tolerance `1e-4`). See [`implementation_diagnostics.json`](https://huggingface.co/Cukinator/cpu1-ablation-checkpoints/blob/main/evaluation/implementation_diagnostics.json) and the bundled `audit_primitives.py`.
314
 
315
  The implementation replaces log-cumulative-sum-exp with a cumulative-maximum normalization followed by a cumulative sum; accumulated terms are not rescaled when that maximum changes. The observed outputs therefore do not satisfy the intended recurrence. This is a source implementation defect; disabling or repairing the scan would create a different inference protocol and was not silently applied to these results. The trained-checkpoint audit uses the source as supplied.
316
 
 
13
 
14
  1,346 checkpoint files across 34 run folders, totaling **215.178 GB** (200.400 GiB). The inventory includes **21 `_final.pt` files** and **1325 step files**. Counts and sizes are derived from the pinned Hub inventory, not estimates from nominal architecture sizes.
15
 
16
+ For evaluated FP32 exports, see [Cukinator/cpu1-ablations-final](https://huggingface.co/Cukinator/cpu1-ablations-final). Source code: [igna-s/New-1.58](https://github.com/igna-s/New-1.58).
17
 
18
  ## What the audit found
19
 
 
89
  torch.set_num_threads(1)
90
  path = hf_hub_download("Cukinator/cpu1-ablation-checkpoints", "run_04/checkpoint_run_04_final.pt",
91
  revision="fb402d025e843ed00fb990966bf2430c32eb958c", local_dir="D:/cpu1/active", cache_dir="D:/cpu1/hf/hub")
92
+ from training.ablation import build_ablation_model, load_ablation_checkpoint
93
  state, config = load_ablation_checkpoint(path)
94
  config["use_gradient_checkpointing"] = False
95
  model = build_ablation_model(config)
 
101
  print([tuple(head.shape) for head in logits])
102
  ```
103
 
104
+ This forward uses zero target bytes in the local decoder; it is a shape smoke test. For scored likelihood, pass `byte_targets` as shown in the audit. For byte generation with retained Transformer context and one recurrent prefill, use `generation()` in `evaluation/audit.py`. Generation remains experimental; inspect the saved examples before choosing a checkpoint.
105
 
106
  ## Resume training
107
 
108
  Use an intermediate checkpoint, not an inference-only packed final. The current source CLI names the argument `--resume_ckpt`:
109
 
110
  ```powershell
111
+ python -m training.ablation --run run_04 --resume_ckpt D:/cpu1/active/run_04/checkpoint_run_04_step640.pt
112
  ```
113
 
114
  Download the exact step file from its run folder first. Resume requires matching source/configuration, dependencies, and suitable training hardware. This audit did not run a training-resume experiment.
 
128
  - Shared sample JSON SHA-256: `11d143c06478ff7bbe364c541a58e6a2e76504c4cb0595b294f46e98f3cfe798`.
129
  - Byte targets are the next patch; incomplete final patches are discarded. The first patch is context only. LocalByteDecoder receives ground-truth target bytes and shifts them internally for causal teacher forcing within a patch. Cross-entropy is summed over all targets and divided by their total count.
130
  - Byte loss is nats/byte; `BPB = NLL / ln(2)` and byte PPL is `exp(NLL)`. A uniform byte predictor gives BPB=8 and PPL=256. BPE loss uses Qwen2.5-3B token IDs with IDs outside the configured vocabulary mapped to 0, matching the dataset adapter. Its PPL is over that lossy mapped token stream and cannot be compared with byte PPL or interpreted as the likelihood of the original text.
131
+ - BPE models shared one Qwen tokenizer, resolved at runtime. The original audit did not capture its resolved revision (null in the records); that missing pin limits exact BPE reproduction. Inspect tokenizer_revision in each result when reproducing later runs.
132
  - Causality: change the last 16 of 32 input steps and compare all prefix logits, tolerance `1e-5`. Recurrent chunk equivalence: compare a 128-step forward against four 32-step chunks across all output heads, tolerance `1e-4`. Failed checks remain visible.
133
  - DeleteGate is disabled consistently: phase 1, elimination rate 0. The audit does not measure token elimination, phase-2 quality, or the benefit of DeleteGate.
134
  - Forward timing: 10 warm-up steps, then 100 measured single-input forwards. `steps/s = 100 / total measured seconds`; p50 and p95 are per-step latency. One byte-model step consumes a patch of 4 bytes; one BPE step consumes one token. Local decoder inputs are zeros in this microbenchmark. Transformer timing has one input step and no persistent attention context. This is not a context-scaling test or a generation benchmark.
 
141
 
142
  From the source checkout, install PyTorch plus `huggingface_hub`, `transformers`, `tokenizers`, `pyarrow`, `psutil`, `numpy`, and `datasets`. Pin the versions from `evaluation/environment.json` when reproducing this run.
143
 
144
+ The audit additions are bundled in `evaluation/source/`. Copy those Python files into the checkout’s `evaluation/` directory when using a source revision that predates this audit. The base Git commit is recorded in `evaluation/environment.json`.
145
 
146
  ```powershell
147
+ python -m evaluation.audit --storage D:/cpu1-readme-audit --threads 1
148
+ python -m evaluation.audit --storage D:/cpu1-readme-audit --cards
149
  ```
150
 
151
+ Model downloads, tokenizer files, Hugging Face caches, and temporary files go to D:. Only one model checkpoint is kept in the task storage at a time; it is removed after its result is saved. Interrupted runs resume from matching evaluator hashes. Raw reports remain in `outputs/audit/`.
152
 
153
  ## Fresh diagnostic results
154
 
 
310
 
311
  ### Recurrent implementation diagnostic
312
 
313
+ The source’s `AblationMLGRU._parallel_scan` was tested independently against its defining serial recurrence with random candidate/forget tensors, seed 19, shape `[1,32,8]`, FP32. The maximum difference was **19.100506**; splitting the same scan into two 16-step chunks differed by **11.471270** (tolerance `1e-4`). See [`implementation_diagnostics.json`](https://huggingface.co/Cukinator/cpu1-ablation-checkpoints/blob/main/evaluation/implementation_diagnostics.json) and the bundled `evaluation/diagnostics.py`.
314
 
315
  The implementation replaces log-cumulative-sum-exp with a cumulative-maximum normalization followed by a cumulative sum; accumulated terms are not rescaled when that maximum changes. The observed outputs therefore do not satisfy the intended recurrence. This is a source implementation defect; disabling or repairing the scan would create a different inference protocol and was not silently applied to these results. The trained-checkpoint audit uses the source as supplied.
316
 
evaluation/publication.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "Cukinator/cpu1-ablations-final": {
3
- "oid": "d38678e7c5d18577f36edade64495237c72a44b5",
4
- "url": "https://huggingface.co/Cukinator/cpu1-ablations-final/commit/d38678e7c5d18577f36edade64495237c72a44b5"
5
  }
6
  }
 
1
  {
2
  "Cukinator/cpu1-ablations-final": {
3
+ "oid": "304d4a9f7414da8682aae5825daa16e40d74713a",
4
+ "url": "https://huggingface.co/Cukinator/cpu1-ablations-final/commit/304d4a9f7414da8682aae5825daa16e40d74713a"
5
  }
6
  }
evaluation/source/cpu1/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ """Shared CPU-1 model components and checkpoint codecs."""
evaluation/source/cpu1/components.py ADDED
The diff for this file is too large to render. See raw diff
 
evaluation/source/evaluation/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ """Reproducible audit and checkpoint tooling."""
evaluation/source/evaluation/audit.py ADDED
@@ -0,0 +1,299 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Reproducible Hub audit. Download and evaluate exactly one model at a time.
2
+
3
+ python -m evaluation.audit --storage D:/cpu1-readme-audit --threads 1
4
+ python -m evaluation.audit --storage D:/cpu1-readme-audit --cards
5
+ python -m evaluation.audit --storage D:/cpu1-readme-audit --upload
6
+ """
7
+ import argparse
8
+ import hashlib
9
+ import json
10
+ import math
11
+ import os
12
+ from pathlib import Path
13
+ import platform
14
+ import statistics
15
+ import sys
16
+ import time
17
+ import urllib.request
18
+ from datetime import datetime, timezone
19
+
20
+ ROOT = Path(__file__).resolve().parents[1]
21
+ sys.path.insert(0, str(ROOT))
22
+ OUT = ROOT / 'outputs' / 'audit'
23
+ FINAL = 'Cukinator/cpu1-ablations-final'
24
+ ARCHIVE = 'Cukinator/cpu1-ablation-checkpoints'
25
+ FINAL_REV = '65951289352c1a036b12ba6d888af037e7c90b13'
26
+ ARCHIVE_REV = 'fb402d025e843ed00fb990966bf2430c32eb958c'
27
+ EVIDENCE_REV = 'd38678e7c5d18577f36edade64495237c72a44b5'
28
+ DATA = 'Salesforce/wikitext'
29
+ DATA_REV = 'b08601e04326c79dfdd32d625aee71d232d685c3'
30
+ DATA_FILE = 'wikitext-2-raw-v1/test-00000-of-00001.parquet'
31
+ PROMPTS = ['The history of artificial intelligence', 'Once upon a time', 'The purpose of science is']
32
+
33
+
34
+ def save(path, value):
35
+ def finite_json(obj):
36
+ if isinstance(obj, float) and not math.isfinite(obj):
37
+ return None
38
+ if isinstance(obj, dict):
39
+ return {k: finite_json(v) for k, v in obj.items()}
40
+ if isinstance(obj, (list, tuple)):
41
+ return [finite_json(v) for v in obj]
42
+ return obj
43
+ path.parent.mkdir(parents=True, exist_ok=True)
44
+ temporary = path.with_suffix('.json.tmp')
45
+ temporary.write_text(json.dumps(finite_json(value), indent=2, ensure_ascii=False, allow_nan=False) + '\n', encoding='utf-8')
46
+ temporary.replace(path)
47
+
48
+
49
+ def sha(path):
50
+ h = hashlib.sha256()
51
+ with path.open('rb') as f:
52
+ for block in iter(lambda: f.read(1024 * 1024), b''):
53
+ h.update(block)
54
+ return h.hexdigest()
55
+
56
+
57
+ def restore_evidence(storage):
58
+ """Restore published reports only; never download model weights."""
59
+ import shutil
60
+ from huggingface_hub import snapshot_download
61
+ folder = Path(snapshot_download(
62
+ FINAL, revision=EVIDENCE_REV,
63
+ allow_patterns=['evaluation/*.json', 'evaluation/CORPUS_LICENSE.txt'],
64
+ local_dir=str(storage/'evidence'), cache_dir=str(storage/'hf'/'hub'),
65
+ ))/'evaluation'
66
+ OUT.mkdir(parents=True, exist_ok=True)
67
+ for source in sorted(folder.glob('*')):
68
+ if source.suffix == '.json' or source.name == 'CORPUS_LICENSE.txt':
69
+ shutil.copyfile(source, OUT/source.name)
70
+ print(f'Restored published evidence to {OUT}; no model weights downloaded.', flush=True)
71
+
72
+
73
+ def prepare(storage):
74
+ OUT.mkdir(parents=True, exist_ok=True)
75
+ invpath = OUT / 'inventory.json'
76
+ if not invpath.exists():
77
+ inventory = {r: json.load(urllib.request.urlopen(f'https://huggingface.co/api/models/{r}/revision/{revision}?blobs=true')) for r, revision in ((FINAL, FINAL_REV), (ARCHIVE, ARCHIVE_REV))}
78
+ save(invpath, inventory)
79
+ corpuspath = OUT / 'corpus.json'
80
+ if not corpuspath.exists():
81
+ import pyarrow.parquet as pq
82
+ dest = storage / 'data' / 'test.parquet'
83
+ dest.parent.mkdir(parents=True, exist_ok=True)
84
+ urllib.request.urlretrieve(f'https://huggingface.co/datasets/{DATA}/resolve/{DATA_REV}/{DATA_FILE}', dest)
85
+ texts = pq.read_table(dest, columns=['text']).column('text').to_pylist()
86
+ # First 32 non-heading paragraphs with at least 260 UTF-8 bytes.
87
+ samples = [t.encode('utf-8')[:260].decode('utf-8', errors='ignore') for t in texts if len(t.encode('utf-8')) >= 260 and not t.strip().startswith('=')][:32]
88
+ assert len(samples) == 32
89
+ save(corpuspath, {'dataset': DATA, 'revision': DATA_REV, 'file': DATA_FILE, 'parquet_sha256': sha(dest), 'selection': 'first 32 non-heading paragraphs >=260 UTF-8 bytes; truncate to 260 bytes, discard incomplete trailing UTF-8', 'texts': samples})
90
+ return json.loads(invpath.read_text(encoding='utf-8')), json.loads(corpuspath.read_text(encoding='utf-8'))
91
+
92
+
93
+ def loss_eval(model, config, corpus, tokenizer):
94
+ import torch
95
+ import torch.nn.functional as F
96
+ byte = config['tokenization'] in ('byte', 'byte_dynamic')
97
+ p = config.get('patch_size', 4)
98
+ total, count, mapped, all_ids = 0., 0, 0, 0
99
+ with torch.inference_mode():
100
+ for text in corpus['texts']:
101
+ if byte:
102
+ raw = list(text.encode('utf-8'))
103
+ data = torch.tensor(raw[:len(raw)//p*p]).reshape(1, -1, p)
104
+ x, y = data[:, :-1], data[:, 1:]
105
+ logits = model(x, byte_targets=y, training_phase=1, target_elimination=0.)[0]
106
+ if isinstance(logits, (tuple, list)):
107
+ logits = torch.stack(logits, dim=2)
108
+ else:
109
+ ids = torch.tensor(tokenizer.encode(text, add_special_tokens=False)).reshape(1, -1)
110
+ cap = config['bpe_vocab_size']
111
+ mapped += int((ids >= cap).sum())
112
+ all_ids += ids.numel()
113
+ ids = torch.where(ids < cap, ids, 0)
114
+ x, y = ids[:, :-1], ids[:, 1:]
115
+ logits = model(x, training_phase=1, target_elimination=0.)[0]
116
+ loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1), reduction='sum')
117
+ if not torch.isfinite(loss):
118
+ raise ValueError('nonfinite loss')
119
+ total += loss.item()
120
+ count += y.numel()
121
+ avg = total / count
122
+ return {'mean_nll': avg, 'targets': count, 'unit': 'byte' if byte else 'mapped_qwen_token', 'perplexity': math.exp(avg), 'bits_per_byte': avg/math.log(2) if byte else None, 'mapped_ids_fraction': mapped/all_ids if all_ids else None, 'source': 'shared_wikitext_test_sample', 'teacher_forcing': True, 'delete_gate': 'disabled (phase 1, elimination 0)'}
123
+
124
+
125
+ def generation(model, config, prompt, length=128):
126
+ """Keep Transformer context; recurrent models carry state exactly once."""
127
+ import torch
128
+ from training.ablation import _global_pass, _sample_logits
129
+ p = config.get('patch_size', 4)
130
+ recurrent = config['token_mixer'] == 'mlgru'
131
+ raw = prompt.encode('utf-8')
132
+ raw += b'\x00' * ((-len(raw)) % p)
133
+ context = torch.tensor(list(raw)).reshape(1, -1, p)
134
+ states = [None] * len(model.blocks)
135
+ output = bytearray()
136
+ with torch.inference_mode():
137
+ h, states = _global_pass(model, context, states, 0.)
138
+ while len(output) < length:
139
+ new = torch.zeros(1, 1, p, dtype=torch.long)
140
+ for k in range(p):
141
+ logits = model.local_decoder(h[:, -1:], new)[k][0, 0] if model.use_local_decoder else model.heads[k](h[:, -1:])[0, 0]
142
+ value = _sample_logits(logits, 1., 40, .9)
143
+ new[0, 0, k] = value
144
+ output.append(value)
145
+ if len(output) >= length:
146
+ break
147
+ if recurrent:
148
+ h, states = _global_pass(model, new, states, 0.)
149
+ else:
150
+ context = torch.cat([context, new], dim=1)
151
+ window = config.get('window_size', 512)
152
+ h, _ = _global_pass(model, context[:, -window:], [None]*len(model.blocks), 0.)
153
+ return bytes(output[:length])
154
+
155
+
156
+ def benchmark(engine):
157
+ """Isolated single-input forward, not generated bytes/s or context scaling."""
158
+ import torch
159
+ dummy = torch.zeros((1, 1, engine.info.patch_size) if engine.info.is_byte_level else (1, 1), dtype=torch.long)
160
+ times = []
161
+ with torch.inference_mode():
162
+ for _ in range(10):
163
+ _, engine.hidden_states = engine._ablation_step(engine._model, dummy)
164
+ engine.reset_state()
165
+ for _ in range(100):
166
+ t = time.perf_counter()
167
+ _, engine.hidden_states = engine._ablation_step(engine._model, dummy)
168
+ times.append(time.perf_counter()-t)
169
+ engine.reset_state()
170
+ values = sorted(times)
171
+ return {'forward_steps_per_second': len(times)/sum(times), 'p50_ms': statistics.median(times)*1000, 'p95_ms': values[94]*1000, 'jitter_cv': statistics.stdev(times)/statistics.mean(times), 'warmup_steps': 10, 'measured_steps': 100, 'context_steps': 1, 'backend': 'PyTorch FP32; quantization simulated for ternary modules; no C++ kernels', 'meaning': 'one input patch/token forward, zero local decoder targets; excludes autoregressive local sampling; Transformer has no persistent context'}
172
+
173
+
174
+ def causality(engine):
175
+ import torch
176
+ from evaluation.correctness import _get_logits_list, _make_dummy_input
177
+ torch.manual_seed(101)
178
+ x = _make_dummy_input(engine.info, 1, 32)
179
+ changed = x.clone()
180
+ changed[:, 16:] = (changed[:, 16:] + 71) % engine.info.vocab_size
181
+ with torch.inference_mode():
182
+ a = _get_logits_list(engine._model(x)[0], engine.info)
183
+ b = _get_logits_list(engine._model(changed)[0], engine.info)
184
+ diff = max((u[:, :16]-v[:, :16]).abs().max().item() for u,v in zip(a,b))
185
+ return {'status': ('PASS' if diff < 1e-5 else 'FAIL') if math.isfinite(diff) else 'ERROR', 'max_prefix_logit_diff': diff, 'tolerance': 1e-5, 'protocol': '32 input steps; change steps 16:32; compare all heads at steps 0:16 with identical zero local targets'}
186
+
187
+
188
+ def evaluate(storage, inventory, corpus, args):
189
+ import torch
190
+ import psutil
191
+ from huggingface_hub import hf_hub_download
192
+ from evaluation.checkpoints import load_run
193
+ from evaluation.correctness import run_correctness
194
+ from transformers import AutoTokenizer
195
+ torch.set_num_threads(args.threads)
196
+ torch.set_num_interop_threads(1)
197
+ revision = inventory[FINAL]['sha']
198
+ runs = [f['rfilename'].split('/')[0] for f in inventory[FINAL]['siblings'] if f['rfilename'].endswith('/model.pt')]
199
+ if args.runs:
200
+ runs = [r for r in runs if r in args.runs]
201
+ import subprocess
202
+ source = subprocess.run(['git', 'rev-parse', 'HEAD'], cwd=ROOT, capture_output=True, text=True)
203
+ source_commit = source.stdout.strip() if source.returncode == 0 else None
204
+ tokenizer = None
205
+ for index, run in enumerate(runs, 1):
206
+ engine = model = raw = None
207
+ result_path = OUT / f'{run}.json'
208
+ if result_path.exists() and not args.force and json.loads(result_path.read_text(encoding='utf-8')).get('suite_sha256') == sha(Path(__file__)):
209
+ print(f'[{index}/{len(runs)}] existing {run}', flush=True)
210
+ continue
211
+ # Only this task-owned file may be removed. No cache of model blobs.
212
+ model_path = storage / 'active' / run / 'model.pt'
213
+ for stale in (storage / 'active').glob('*/model.pt'):
214
+ target = stale.resolve()
215
+ if not target.is_relative_to((storage / 'active').resolve()):
216
+ raise ValueError('unsafe cleanup path')
217
+ target.unlink()
218
+ print(f'[{index}/{len(runs)}] downloading {run} to {model_path}', flush=True)
219
+ started = time.perf_counter()
220
+ hf_hub_download(FINAL, f'{run}/model.pt', revision=revision, local_dir=str(storage/'active'), cache_dir=str(storage/'hf'/'hub'))
221
+ record = {'run': run, 'revision': revision, 'file': f'{run}/model.pt', 'sha256': sha(model_path), 'file_bytes': model_path.stat().st_size, 'date_utc': datetime.now(timezone.utc).isoformat(), 'threads': args.threads, 'torch': torch.__version__, 'platform': platform.platform(), 'processor': platform.processor(), 'corpus_sha256': sha(OUT/'corpus.json'), 'suite_sha256': sha(Path(__file__)), 'seed': 42, 'source_commit': source_commit}
222
+ try:
223
+ torch.manual_seed(42)
224
+ engine, raw, config, _, _ = load_run(model_path, run)
225
+ model = engine._model
226
+ record.update({'config': config, 'step': raw.get('step'), 'stored_training_metrics': raw.get('metrics', {}), 'parameters': sum(p.numel() for p in model.parameters()), 'parameter_bytes': sum(p.numel()*p.element_size() for p in model.parameters()), 'format': raw.get('format'), 'strict_load': True})
227
+ record['correctness'] = run_correctness(engine)
228
+ record['correctness']['causality'] = causality(engine)
229
+ if not engine.info.is_byte_level and tokenizer is None:
230
+ tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-3B', cache_dir=str(storage/'hf'/'hub'), trust_remote_code=False)
231
+ record['tokenizer_revision'] = tokenizer.init_kwargs.get('_commit_hash')
232
+ record['loss'] = loss_eval(model, config, corpus, tokenizer)
233
+ record['benchmark'] = benchmark(engine)
234
+ if engine.info.is_byte_level:
235
+ samples = []
236
+ for i, prompt in enumerate(PROMPTS):
237
+ torch.manual_seed(42+i)
238
+ t = time.perf_counter()
239
+ generated = generation(model, config, prompt)
240
+ elapsed = time.perf_counter()-t
241
+ try:
242
+ generated.decode('utf-8', errors='strict')
243
+ valid = True
244
+ except UnicodeDecodeError:
245
+ valid = False
246
+ samples.append({'prompt': prompt, 'seed': 42+i, 'hex': generated.hex(), 'text_replacement_decode': generated.decode('utf-8', errors='replace'), 'bytes': len(generated), 'utf8_valid': valid, 'ascii_printable_fraction': sum(32<=b<127 or b in (9,10,13) for b in generated)/len(generated), 'repeated_8byte_fraction': sum(generated[i:i+8]==generated[i-8:i] for i in range(8,len(generated)-7))/max(1,len(generated)-15), 'seconds': elapsed, 'generated_bytes_per_second': len(generated)/elapsed})
247
+ record['generation'] = samples
248
+ else:
249
+ record['generation'] = {'status': 'N/A', 'reason': 'truncated Qwen vocabulary; token mapping is lossy, no comparable byte generation test'}
250
+ record['process_rss_bytes'] = psutil.Process().memory_info().rss
251
+ except Exception as exc:
252
+ record['error'] = f'{type(exc).__name__}: {exc}'
253
+ print(record['error'], flush=True)
254
+ record['total_seconds'] = time.perf_counter()-started
255
+ record['evaluation_complete'] = True
256
+ save(result_path, record)
257
+ print(f'{run}: loss={record.get("loss",{}).get("mean_nll")} elapsed={record["total_seconds"]:.1f}s', flush=True)
258
+ model_path.unlink()
259
+ engine = model = raw = None
260
+ import gc
261
+ gc.collect()
262
+
263
+
264
+ def main():
265
+ parser = argparse.ArgumentParser(description=__doc__)
266
+ parser.add_argument('--storage', type=Path, default=Path('D:/cpu1-readme-audit'))
267
+ parser.add_argument('--threads', type=int, default=1)
268
+ parser.add_argument('--runs', nargs='+')
269
+ parser.add_argument('--force', action='store_true')
270
+ parser.add_argument('--cards', action='store_true')
271
+ parser.add_argument('--upload', action='store_true')
272
+ parser.add_argument('--restore-evidence', action='store_true', help='Download only the published audit JSONs and corpus license')
273
+ args = parser.parse_args()
274
+ if args.storage.resolve().drive.upper() != 'D:':
275
+ raise ValueError('This audit requires storage on D:')
276
+ args.storage.mkdir(parents=True, exist_ok=True)
277
+ for name in ('HF_HOME','HF_HUB_CACHE','HF_DATASETS_CACHE','TRANSFORMERS_CACHE','TEMP','TMP'):
278
+ dest = args.storage / ('temp' if name in ('TEMP','TMP') else name.lower())
279
+ dest.mkdir(parents=True, exist_ok=True)
280
+ os.environ[name] = str(dest)
281
+ os.environ['HF_HUB_DISABLE_XET'] = '1'
282
+ os.environ['HF_HUB_DISABLE_PROGRESS_BARS'] = '1'
283
+ if args.restore_evidence:
284
+ restore_evidence(args.storage)
285
+ return
286
+ inventory, corpus = prepare(args.storage)
287
+ if args.cards or args.upload:
288
+ from evaluation.cards import build_cards, upload_cards
289
+ build_cards(inventory, corpus)
290
+ if args.upload:
291
+ upload_cards(inventory)
292
+ else:
293
+ evaluate(args.storage, inventory, corpus, args)
294
+
295
+
296
+ if __name__ == '__main__':
297
+ if hasattr(sys.stdout, 'reconfigure'):
298
+ sys.stdout.reconfigure(encoding='utf-8')
299
+ main()
evaluation/source/evaluation/cards.py ADDED
@@ -0,0 +1,299 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Generate Hub model cards and the dataset template from audited evidence."""
2
+ from collections import defaultdict
3
+ from datetime import datetime
4
+ import json
5
+ from pathlib import Path
6
+ import re
7
+ import sys
8
+
9
+ from evaluation.audit import ROOT, OUT, FINAL, ARCHIVE, DATA, DATA_REV, DATA_FILE, PROMPTS, sha
10
+
11
+
12
+ def link(repo, filename, revision='main'):
13
+ return f'https://huggingface.co/{repo}/blob/{revision}/{filename}'
14
+
15
+
16
+ def manifest_rows(inventory):
17
+ folders = defaultdict(list)
18
+ for f in inventory[ARCHIVE]['siblings']:
19
+ if f['rfilename'].endswith('.pt'):
20
+ folders[f['rfilename'].split('/')[0]].append(f)
21
+ rows = ['| Folder | Step files | Highest filename step | Final file MB | Total GB |', '|---|---:|---:|---:|---:|']
22
+ for run, files in sorted(folders.items()):
23
+ steps = [int(m.group(1)) for f in files if (m := re.search(r'_step(\d+)\.pt$', f['rfilename']))]
24
+ finals = [f for f in files if f['rfilename'].endswith('_final.pt')]
25
+ final = f'{sum(f["size"] for f in finals)/1e6:.2f}' if finals else '—'
26
+ rows.append(f'| `{run}` | {len(steps)} | {max(steps) if steps else "—"} | {final} | {sum(f["size"] for f in files)/1e9:.3f} |')
27
+ return rows
28
+
29
+
30
+ def audit_summary(records):
31
+ checks = [r.get('correctness', {}) for r in records]
32
+ causal = [c.get('causality', {}) for c in checks]
33
+ recurrent = [c['state_prop'] for c in checks if c.get('state_prop', {}).get('status') != 'N/A']
34
+ return [
35
+ '## What the audit found', '',
36
+ f'The shared audit covers {len(records)} FP32 exports on 32 short WikiText test passages, with CPU execution, one thread, and a fixed protocol. It does not evaluate every intermediate archive checkpoint.', '',
37
+ '| Check | Outcome |', '|---|---|',
38
+ f'| Strict checkpoint loading | {sum(r.get("strict_load", False) for r in records)}/{len(records)} loaded |',
39
+ f'| Numerical evaluation | {sum(bool(r.get("error")) for r in records)} models produced an evaluation error |',
40
+ f'| Prefix causality | {sum(c.get("status") == "FAIL" for c in causal)}/{sum(c.get("status") in ("PASS", "FAIL") for c in causal)} finite-output models failed |',
41
+ f'| Recurrent chunk equivalence | {sum(c.get("status") == "FAIL" for c in recurrent)}/{len(recurrent)} recurrent models failed |', '',
42
+ 'Full results, generation examples, and the recurrent implementation diagnostic follow below. Loss units depend on tokenization. Loss from a model that fails causality cannot establish autoregressive prediction quality; printable output alone does not establish useful language generation.', '',
43
+ ]
44
+
45
+
46
+ def protocol(inventory, corpus, records):
47
+ return [
48
+ '## Evaluation protocol', '',
49
+ f'This card was rebuilt from a fresh CPU audit on **2026-10-06**. All {len(records)} published FP32 exports were attempted; {sum(r.get("strict_load",False) for r in records)} loaded with zero missing or unexpected state keys. Numerical and loading errors remain visible. Raw JSON, the exact text sample, and the repository inventories are included in [`evaluation/`](evaluation/). Previous `resultados/ablation_*.json` reports are historical and were not mixed into these tables.', '',
50
+ '- Hardware: Intel Core i7-8750H (6 cores / 12 logical processors), Windows 11; CPU execution with **1 PyTorch thread**, batch size 1, FP32 parameters. PyTorch ' + records[0]['torch'] + '. No custom C/C++ kernels were used.',
51
+ '- Test text: [`Salesforce/wikitext`, `wikitext-2-raw-v1`, test split](https://huggingface.co/datasets/Salesforce/wikitext). First 32 non-heading paragraphs with at least 260 UTF-8 bytes, each truncated to 260 bytes without an incomplete trailing character. This is a small diagnostic sample, not the complete benchmark. No exhaustive training-overlap audit was performed.',
52
+ '- Corpus attribution: WikiText by its dataset authors, based on English Wikipedia contributions. The dataset metadata lists CC-BY-SA-3.0 and GFDL; the included corpus excerpt retains those upstream terms, separately from the model repository license. See `evaluation/CORPUS_LICENSE.txt`.',
53
+ f'- Dataset revision: `{DATA_REV}`; exact file: `{DATA_FILE}`; parquet SHA-256: `{corpus["parquet_sha256"]}`.',
54
+ f'- Shared sample JSON SHA-256: `{sha(OUT/"corpus.json")}`.',
55
+ '- Byte targets are the next patch; incomplete final patches are discarded. The first patch is context only. LocalByteDecoder receives ground-truth target bytes and shifts them internally for causal teacher forcing within a patch. Cross-entropy is summed over all targets and divided by their total count.',
56
+ '- Byte loss is nats/byte; `BPB = NLL / ln(2)` and byte PPL is `exp(NLL)`. A uniform byte predictor gives BPB=8 and PPL=256. BPE loss uses Qwen2.5-3B token IDs with IDs outside the configured vocabulary mapped to 0, matching the dataset adapter. Its PPL is over that lossy mapped token stream and cannot be compared with byte PPL or interpreted as the likelihood of the original text.',
57
+ '- BPE models shared one Qwen tokenizer, resolved at runtime. The original audit did not capture its resolved revision (null in the records); that missing pin limits exact BPE reproduction. Inspect tokenizer_revision in each result when reproducing later runs.',
58
+ '- Causality: change the last 16 of 32 input steps and compare all prefix logits, tolerance `1e-5`. Recurrent chunk equivalence: compare a 128-step forward against four 32-step chunks across all output heads, tolerance `1e-4`. Failed checks remain visible.',
59
+ '- DeleteGate is disabled consistently: phase 1, elimination rate 0. The audit does not measure token elimination, phase-2 quality, or the benefit of DeleteGate.',
60
+ '- Forward timing: 10 warm-up steps, then 100 measured single-input forwards. `steps/s = 100 / total measured seconds`; p50 and p95 are per-step latency. One byte-model step consumes a patch of 4 bytes; one BPE step consumes one token. Local decoder inputs are zeros in this microbenchmark. Transformer timing has one input step and no persistent attention context. This is not a context-scaling test or a generation benchmark.',
61
+ '- Byte generation: three prompts, 128 new bytes per prompt, seeds 42/43/44, temperature 1.0, top-k 40, top-p 0.9. Recurrent models carry state; Transformers reprocess the growing context. Generated bytes/s includes prompt prefill and local sampling. UTF-8 validity is tested on the complete continuation. Samples are retained as hex and decoded text. BPE generation is N/A because vocabulary truncation prevents an equivalent byte-output comparison.',
62
+ '- Printable ASCII fraction, UTF-8 validity, and repetition are diagnostics; they do not measure factual accuracy, instruction following, or useful language quality. RAM columns report parameter storage, not a measured peak process footprint. Ternary models still hold FP32 parameters and simulate quantization in PyTorch.', '',
63
+ f'Export repository revision audited: `{inventory[FINAL]["sha"]}`. Archive inventory revision: `{inventory[ARCHIVE]["sha"]}`. Each result records its checkpoint SHA-256, actual parameter count, saved step, configuration, thread count, time, sample hash, and evaluator hash.', '',
64
+ '### Reproduce', '',
65
+ 'From the source checkout, install PyTorch plus `huggingface_hub`, `transformers`, `tokenizers`, `pyarrow`, `psutil`, `numpy`, and `datasets`. Pin the versions from `evaluation/environment.json` when reproducing this run.', '',
66
+ 'The audit additions are bundled in `evaluation/source/`. Copy those Python files into the checkout’s `evaluation/` directory when using a source revision that predates this audit. The base Git commit is recorded in `evaluation/environment.json`.', '',
67
+ '```powershell',
68
+ 'python -m evaluation.audit --storage D:/cpu1-readme-audit --threads 1',
69
+ 'python -m evaluation.audit --storage D:/cpu1-readme-audit --cards',
70
+ '```', '',
71
+ 'Model downloads, tokenizer files, Hugging Face caches, and temporary files go to D:. Only one model checkpoint is kept in the task storage at a time; it is removed after its result is saved. Interrupted runs resume from matching evaluator hashes. Raw reports remain in `outputs/audit/`.', '',
72
+ ]
73
+
74
+
75
+ def tables(records):
76
+ lines = ['## Fresh diagnostic results', '',
77
+ '**Do not rank models across tokenizations or treat a failed-causality loss as autoregressive language-model perplexity.** Saved training metrics are retained separately in the JSONs; the numbers below are recomputed on the shared test sample.', '']
78
+ for byte in (True, False):
79
+ lines += ['### Byte models' if byte else '### BPE models (mapped Qwen token stream)', '',
80
+ '| Run | Params M | Saved step | File MB | NLL | BPB | PPL | Causal | State chunks |',
81
+ '|---|---:|---:|---:|---:|---:|---:|---|---|']
82
+ for r in records:
83
+ if 'config' not in r or (r['config']['tokenization'] in ('byte','byte_dynamic')) != byte:
84
+ continue
85
+ lm = r.get('loss')
86
+ cc = r['correctness']
87
+ if lm is None:
88
+ lines.append(f'| [`{r["run"]}`](evaluation/{r["run"]}.json) | {r["parameters"]/1e6:.3f} | {r["step"]} | {r["file_bytes"]/1e6:.2f} | ERROR | — | — | {cc["causality"]["status"]} | {cc["state_prop"]["status"]} |')
89
+ continue
90
+ bpb = f'{lm["bits_per_byte"]:.4f}' if byte else '—'
91
+ lines.append(f'| [`{r["run"]}`](evaluation/{r["run"]}.json) | {r["parameters"]/1e6:.3f} | {r["step"]} | {r["file_bytes"]/1e6:.2f} | {lm["mean_nll"]:.4f} | {bpb} | {lm["perplexity"]:.3f} | {cc["causality"]["status"]} | {cc["state_prop"]["status"]} |')
92
+ if not byte:
93
+ lines += ['', '| BPE run | IDs mapped to 0 | Vocabulary cap |', '|---|---:|---:|']
94
+ for r in records:
95
+ if r.get('config',{}).get('tokenization') == 'bpe' and 'loss' in r:
96
+ lines.append(f'| `{r["run"]}` | {r["loss"]["mapped_ids_fraction"]*100:.2f}% | {r["config"]["bpe_vocab_size"]:,} |')
97
+ lines += ['']
98
+ lines += ['### CPU timing and byte generation', '', '| Run | Param MB (FP32) | Forward steps/s | p50 ms | p95 ms | Generated bytes/s | Valid UTF-8 samples |', '|---|---:|---:|---:|---:|---:|---:|']
99
+ for r in records:
100
+ bm = r.get('benchmark')
101
+ gen = r.get('generation')
102
+ if bm is None:
103
+ lines.append(f'| `{r["run"]}` | {r.get("parameter_bytes",0)/1e6:.2f} | ERROR | — | — | N/A | N/A |')
104
+ continue
105
+ if isinstance(gen, list):
106
+ speed = f'{sum(s["bytes"] for s in gen)/sum(s["seconds"] for s in gen):.2f}'
107
+ utf = f'{sum(s["utf8_valid"] for s in gen)}/{len(gen)}'
108
+ else:
109
+ speed = utf = 'N/A'
110
+ lines.append(f'| `{r["run"]}` | {r["parameter_bytes"]/1e6:.2f} | {bm["forward_steps_per_second"]:.2f} | {bm["p50_ms"]:.2f} | {bm["p95_ms"]:.2f} | {speed} | {utf} |')
111
+ lines += ['', '### Example continuations', '',
112
+ 'The following fixed examples use the first audit prompt, `The history of artificial intelligence`, seed 42, and the sampling settings above. They illustrate why low diagnostic loss or valid UTF-8 alone does not establish useful text generation. Escapes preserve the recorded characters; complete byte sequences and all three prompts are in the JSON reports.', '']
113
+ for name in ('run_04_v3','run_14_v3'):
114
+ r = next((r for r in records if r['run']==name),None)
115
+ if r and isinstance(r.get('generation'),list):
116
+ lines += [f'`{name}`:', '', '```text', json.dumps(r['generation'][0]['text_replacement_decode'],ensure_ascii=True), '```', '']
117
+ lines += ['', '### Correctness checks', '', '| Run | Shapes | Probabilities | Determinism | State chunks | Ternary projection | State size | Causality |', '|---|---|---|---|---|---|---|---|']
118
+ keys = ['forward_shapes','probs','determinism','state_prop','ternary_weights','state_size','causality']
119
+ for r in records:
120
+ lines.append('| `' + r['run'] + '` | ' + ' | '.join(r.get('correctness',{}).get(k,{}).get('status','ERROR') for k in keys) + ' |')
121
+ failed_causal = [r['run'] for r in records if r.get('correctness',{}).get('causality',{}).get('status')=='FAIL']
122
+ failed_state = [r['run'] for r in records if r.get('correctness',{}).get('state_prop',{}).get('status')=='FAIL']
123
+ lines += ['', f'**Causality failures: {len(failed_causal)}/{len(records)}.** ' + (', '.join(f'`{r}`' for r in failed_causal) or 'None.'), '',
124
+ 'Changing future inputs changed prefix logits in these models. Some saved configurations enable noncausal Bolmo cross-patch operations; the test detects the effect without isolating its individual source. Their full-sequence losses are descriptive reconstruction scores and cannot establish causal prediction quality.', '',
125
+ f'**State chunk failures: {len(failed_state)}.** ' + (', '.join(f'`{r}`' for r in failed_state) or 'None.'), '',
126
+ 'These configurations did not reproduce a full-sequence forward when carrying state between chunks at the stated tolerance. An O(1) state tensor alone does not establish correct streaming equivalence. No fix to the trained architectures or weights was made for this audit.', '']
127
+ diagnostic = OUT/'implementation_diagnostics.json'
128
+ if diagnostic.exists():
129
+ d = json.loads(diagnostic.read_text(encoding='utf-8'))
130
+ lines += ['### Recurrent implementation diagnostic', '',
131
+ f'The source’s `AblationMLGRU._parallel_scan` was tested independently against its defining serial recurrence with random candidate/forget tensors, seed 19, shape `[1,32,8]`, FP32. The maximum difference was **{d["serial_max_diff"]:.6f}**; splitting the same scan into two 16-step chunks differed by **{d["chunks_max_diff"]:.6f}** (tolerance `1e-4`). See [`implementation_diagnostics.json`](evaluation/implementation_diagnostics.json) and the bundled `evaluation/diagnostics.py`.', '',
132
+ 'The implementation replaces log-cumulative-sum-exp with a cumulative-maximum normalization followed by a cumulative sum; accumulated terms are not rescaled when that maximum changes. The observed outputs therefore do not satisfy the intended recurrence. This is a source implementation defect; disabling or repairing the scan would create a different inference protocol and was not silently applied to these results. The trained-checkpoint audit uses the source as supplied.', '']
133
+ errors = [r for r in records if r.get('error')]
134
+ if errors:
135
+ lines += ['### Evaluation errors', '', 'Scores and speed are unavailable when model loading or numerical evaluation fails. Nonfinite values are represented as JSON null; they are never substituted with valid-looking numbers.', '', '| Run | Error |', '|---|---|']
136
+ lines += [f'| `{r["run"]}` | {r["error"].replace("|","/")} |' for r in errors]
137
+ lines += ['']
138
+ return lines
139
+
140
+
141
+ def architecture(records):
142
+ lines = ['## Architecture recorded in each checkpoint', '',
143
+ 'Run names identify experiments; dimensions and features below come from the saved configuration. FP16 is the training quantization label, while this audit executes the loaded tensors in FP32. “50M” and “10M” in older cards were design families, not exact model sizes.', '',
144
+ '| Run | Mixer | Tokens | Quant. label | Width / layers / FFN | Local decoder | FP residual | Bolmo | DeleteGate | PFNet hidden | Channel decay |',
145
+ '|---|---|---|---|---|---|---|---|---|---:|---|']
146
+ for r in records:
147
+ if 'config' not in r:
148
+ continue
149
+ c = r['config']
150
+ yes = lambda k: 'yes' if c.get(k, False) else 'no'
151
+ lines.append(f'| `{r["run"]}` | {c["token_mixer"]} | {c["tokenization"]} | {c["quantization"]} | {c["d_model"]} / {c["n_layers"]} / {c["d_ff"]} | {yes("use_local_decoder") if c["tokenization"]!="bpe" else "N/A"} | {yes("fp_residual")} | {yes("use_bolmo_gating")} | {yes("use_delete_gate")} | {c.get("pfnet_hidden",0)} | {yes("mlgru_per_channel_decay")} |')
152
+ lines += ['', 'These experiments change parameter counts, representations, training signal, budgets, saved steps, and sometimes multiple architectural settings. The audit is not a controlled causal estimate of the benefit of an individual component. `chinchilla_tokens_per_param` is a configured budget; it does not prove that a run completed that budget. Actual saved steps are reported separately.', '']
153
+ return lines
154
+
155
+
156
+ def quickstart(repo, revision, compact=False):
157
+ filename = 'run_04/checkpoint_run_04_final.pt' if compact else 'run_04/model.pt'
158
+ loader = ['from training.ablation import build_ablation_model, load_ablation_checkpoint', f'state, config = load_ablation_checkpoint(path)'] if compact else ['from training.ablation import build_ablation_model', 'payload = torch.load(path, map_location="cpu", weights_only=False)', 'config, state = payload["config"], payload["state_dict"]']
159
+ return ['## Load one checkpoint', '', 'Clone the [source repository](https://github.com/igna-s/New-1.58) and run this from its root. These are custom PyTorch checkpoints; `AutoModelForCausalLM.from_pretrained()` does not load them. The examples use a trusted checkpoint source and strict state loading.', '', '```python', 'import os', 'os.environ["HF_HOME"] = "D:/cpu1/hf"', 'os.environ["HF_HUB_DISABLE_XET"] = "1"', 'import torch', 'from huggingface_hub import hf_hub_download', 'torch.set_num_threads(1)', f'path = hf_hub_download("{repo}", "{filename}",', f' revision="{revision}", local_dir="D:/cpu1/active", cache_dir="D:/cpu1/hf/hub")', *loader, 'config["use_gradient_checkpointing"] = False', 'model = build_ablation_model(config)', 'model.load_state_dict(state, strict=True)', 'model.eval()', 'x = torch.tensor(list(b"The history of AI. ")).reshape(1, -1, 4)', 'with torch.inference_mode():', ' logits, states, extra = model(x, training_phase=1, target_elimination=0.0)', 'print([tuple(head.shape) for head in logits])', '```', '', 'This forward uses zero target bytes in the local decoder; it is a shape smoke test. For scored likelihood, pass `byte_targets` as shown in the audit. For byte generation with retained Transformer context and one recurrent prefill, use `generation()` in `evaluation/audit.py`. Generation remains experimental; inspect the saved examples before choosing a checkpoint.', '']
160
+
161
+
162
+ def build_cards(inventory, corpus):
163
+ names = sorted(f['rfilename'].split('/')[0] for f in inventory[FINAL]['siblings'] if f['rfilename'].endswith('/model.pt'))
164
+ records = [json.loads((OUT/f'{r}.json').read_text(encoding='utf-8')) for r in names]
165
+ for r in records:
166
+ if not r.get('evaluation_complete') or r['revision'] != inventory[FINAL]['sha']:
167
+ raise ValueError(f'Incomplete audit for {r["run"]}')
168
+ if len({r['suite_sha256'] for r in records}) != 1 or len({r['corpus_sha256'] for r in records}) != 1:
169
+ raise ValueError('Mixed audit protocols')
170
+ header = ['---', 'license: apache-2.0', 'language:', '- en', 'tags:', '- cpu-1', '- ablation', '- ternary', '- research', '---', '']
171
+ shared = protocol(inventory, corpus, records) + tables(records) + architecture(records)
172
+ final_files = [f for f in inventory[FINAL]['siblings'] if f['rfilename'].endswith('/model.pt')]
173
+ final_card = header + ['# CPU-1 ablations: FP32 checkpoint exports', '',
174
+ f'{len(final_files)} custom PyTorch model exports, {sum(f["size"] for f in final_files)/1e9:.3f} GB of checkpoint files (decimal units). This is a research ablation collection, not an instruction-tuned assistant. It includes Transformer and MLGRU variants with byte patches or truncated BPE vocabularies, and floating-point or simulated ternary projections.', '',
175
+ f'Source: [igna-s/New-1.58](https://github.com/igna-s/New-1.58). Training archive: [{ARCHIVE}](https://huggingface.co/{ARCHIVE}). Training dataset companion: [Cukinator/cpu1-ablation-dataset](https://huggingface.co/datasets/Cukinator/cpu1-ablation-dataset).', '',
176
+ *audit_summary(records),
177
+ '## Payload format', '', '`<run>/model.pt` contains `format="fp32"`, `state_dict`, `config`, `step`, `run_name`, `metrics`, and source metadata. The exported parameters occupy four bytes each. “1.58-bit” describes the ternary value alphabet; these files are not stored as 1.58-bit tensors and this PyTorch path does not use packed-weight inference.', '',
178
+ 'Historical `metrics` may be empty or come from different training validation protocols. They are preserved as metadata, not relabeled as freshly measured test results.', '',
179
+ *quickstart(FINAL, inventory[FINAL]['sha']), *shared,
180
+ '## License', '', 'The Hub metadata declares Apache-2.0 for the model repository. The source checkout does not include a standalone LICENSE file. Upstream datasets and tokenizer assets retain their own licenses.', '']
181
+ archived = [f for f in inventory[ARCHIVE]['siblings'] if f['rfilename'].endswith('.pt')]
182
+ finals = [f for f in archived if f['rfilename'].endswith('_final.pt')]
183
+ archive_card = header + ['# CPU-1 ablations: training checkpoint archive', '',
184
+ f'{len(archived):,} checkpoint files across {len({f["rfilename"].split("/")[0] for f in archived})} run folders, totaling **{sum(f["size"] for f in archived)/1e9:.3f} GB** ({sum(f["size"] for f in archived)/2**30:.3f} GiB). The inventory includes **{len(finals)} `_final.pt` files** and **{len(archived)-len(finals)} step files**. Counts and sizes are derived from the pinned Hub inventory, not estimates from nominal architecture sizes.', '',
185
+ f'For evaluated FP32 exports, see [{FINAL}](https://huggingface.co/{FINAL}). Source code: [igna-s/New-1.58](https://github.com/igna-s/New-1.58).', '',
186
+ *audit_summary(records),
187
+ '## What is stored', '',
188
+ '- `checkpoint_<run>_final.pt`: inference-oriented final checkpoints. Ternary runs use `compact_2bit`; floating-point baselines use floating-point checkpoint storage. File names alone do not identify the serialization format: inspect `format` inside the payload.',
189
+ '- `checkpoint_<run>_step<N>.pt`: intermediate training state; includes saved configuration and model state, and may include optimizer state. `_phase2_step<N>.pt` names distinguish phase-2 files where present.',
190
+ '- The 2-bit codec packs a ternary alphabet into two bits per stored value plus scales and floating-point tensors. It is not a guarantee of 1.58 effective bits per total model parameter. Resume step files and packed finals have different purposes.',
191
+ '- A folder without `_final.pt` is an intermediate-checkpoint archive. The highest filename step establishes the highest stored step; it does not establish convergence, completed training budget, or an early-stop reason.', '',
192
+ '## Archive inventory', '', *manifest_rows(inventory), '',
193
+ 'Sizes use decimal MB and GB. Step counts include phase-2 step files. Highest step is parsed from filenames and does not merge training phases into a continuous training counter.', '',
194
+ *quickstart(ARCHIVE, inventory[ARCHIVE]['sha'], compact=True),
195
+ '## Resume training', '', 'Use an intermediate checkpoint, not an inference-only packed final. The current source CLI names the argument `--resume_ckpt`:', '',
196
+ '```powershell', 'python -m training.ablation --run run_04 --resume_ckpt D:/cpu1/active/run_04/checkpoint_run_04_step640.pt', '```', '',
197
+ 'Download the exact step file from its run folder first. Resume requires matching source/configuration, dependencies, and suitable training hardware. This audit did not run a training-resume experiment.', '',
198
+ '## Scope of the measured results', '',
199
+ f'The results below were measured on the {len(records)} published exports in [{FINAL}](https://huggingface.co/{FINAL}), **not by evaluating every intermediate archive checkpoint**. Matching run names do not prove tensor identity between a packed final, an intermediate step, and an export. No packed/FP32 numerical-parity claim is made by this card.', '',
200
+ *shared, '## License', '', 'Hub metadata declares Apache-2.0. Dataset and tokenizer licenses are separate. The source checkout does not include a standalone LICENSE file.', '']
201
+ directory = ROOT / 'docs' / 'model_cards'
202
+ for stem, card in [('cpu1-ablations-final', final_card), ('cpu1-ablation-checkpoints', archive_card)]:
203
+ content = '\n'.join(card)
204
+ repo = FINAL if stem == 'cpu1-ablations-final' else ARCHIVE
205
+ content = content.replace('](evaluation/)', f'](https://huggingface.co/{repo}/tree/main/evaluation)')
206
+ content = content.replace('](evaluation/', f'](https://huggingface.co/{repo}/blob/main/evaluation/')
207
+ (directory/(stem+'.md')).write_text(content, encoding='utf-8')
208
+ save_environment(records)
209
+ (OUT/'CORPUS_LICENSE.txt').write_text('The text excerpts in corpus.json originate from WikiText (Salesforce/wikitext), based on English Wikipedia.\nDataset authors and contributors: see https://huggingface.co/datasets/Salesforce/wikitext .\nUpstream dataset metadata declares CC-BY-SA-3.0 and GFDL.\nCC-BY-SA-3.0: https://creativecommons.org/licenses/by-sa/3.0/\nGFDL: https://www.gnu.org/licenses/fdl-1.3.html\nChanges: selection of 32 test paragraphs and truncation to 260 UTF-8 bytes.\nThe corpus remains under its upstream terms; the model repository Apache-2.0 metadata does not relicense these excerpts.\n',encoding='utf-8')
210
+ dataset_card()
211
+ print('Generated both model cards and the dataset template; source README preserved.', flush=True)
212
+
213
+
214
+ def dataset_card():
215
+ """Update the local companion template; publication scope is the two model repos."""
216
+ path = OUT/'dataset_inventory.json'
217
+ if not path.exists():
218
+ return
219
+ v = json.loads(path.read_text(encoding='utf-8'))
220
+ lines = ['---', 'license: mit', 'task_categories:', '- text-generation', 'language:', '- en', 'configs:', '- config_name: byte_marginalized', ' data_files:', ' - split: train', ' path: byte_marginalized/shard_*.parquet', '- config_name: bpe_tokenized', ' data_files:', ' - split: train', ' path: bpe_tokenized/shard_*.parquet', '---', '',
221
+ '# CPU-1 ablation training dataset', '',
222
+ 'Companion dataset for the CPU-1 distillation experiments. The extraction source config uses Qwen/Qwen2.5-3B as teacher over FineWeb text. The byte and BPE subsets have different schemas and different shard counts; they must not be assumed to be aligned row by row.', '',
223
+ '## Verified Hub inventory: 2026-10-06', '', f'Revision: `{v["sha"]}`. The following counts and compressed sizes are from the Hub file listing, not from a full scan of the dataset rows.', '', '| Subset | Parquet files | Compressed GB | Shard range | Missing IDs within range |', '|---|---:|---:|---|---|']
224
+ for subset in ('byte_marginalized','bpe_tokenized'):
225
+ fs = [f for f in v['siblings'] if f['rfilename'].startswith(subset+'/') and f['rfilename'].endswith('.parquet')]
226
+ ids = sorted(int(re.search(r'shard_(\d+)\.parquet$',f['rfilename']).group(1)) for f in fs)
227
+ missing = sorted(set(range(min(ids),max(ids)+1))-set(ids))
228
+ lines.append(f'| `{subset}` | {len(fs)} | {sum(f["size"] for f in fs)/1e9:.3f} | {min(ids):04d}–{max(ids):04d} | '+', '.join(f'{n:04d}' for n in missing)+' |')
229
+ lines += ['', 'Sizes use decimal GB. Auxiliary `stream_checkpoint.json` and `stream_carry.pkl` files are extraction bookkeeping, not dataset examples. The file listing does not establish document totals, row totals, teacher coverage, split independence, or whether extraction has finished. Older shard and token totals must not be treated as current.', '',
230
+ '## Schema expected by the source adapters', '',
231
+ '| Subset | Fields | Interpretation |', '|---|---|---|',
232
+ '| Byte | `patches`, `targets`, `teacher_probs`, `teacher_mask`; optional `teacher_hidden`, `teacher_dim` | Flattened byte patches and next-patch targets; teacher signal and mask. Hidden arrays may be compacted to teacher-covered positions; follow the adapter. |',
233
+ '| BPE | `input_ids`, `labels`, `teacher_probs_bpe`, `teacher_ids_bpe`, `teacher_mask_bpe` | Qwen token IDs and shifted labels; flattened teacher top-k values and their original IDs. |', '',
234
+ 'This table describes `data_pipeline/prepare.py` and the readers in `cpu1/components.py`; the October audit did not validate every row or measure teacher-signal quality. Check actual row lengths, optional fields and `teacher_dim` rather than assuming a fixed hidden width or identical schemas in all shards.', '',
235
+ 'The extraction source currently configures four-byte patches, 5,000-step sequences, and top-128 BPE teacher entries. Those are writer settings, not a claim that every stored row matches them. Byte teacher signals are first-byte marginalizations of BPE candidates at mapped boundaries; they are not full autoregressive distributions for all four positions. Training uses the coverage masks.', '',
236
+ 'The BPE adapter maps IDs outside the configured student vocabulary to ID 0, and filters teacher IDs accordingly. This is a lossy representation. Its token perplexity cannot be compared directly with byte perplexity or with an untruncated Qwen language model.', '',
237
+ '## Stream one subset', '', '```python', 'import os', 'os.environ["HF_HOME"] = "D:/cpu1/hf"', 'os.environ["HF_DATASETS_CACHE"] = "D:/cpu1/datasets"', 'from datasets import load_dataset', 'dataset = load_dataset(', ' "Cukinator/cpu1-ablation-dataset",', ' "byte_marginalized",', f' revision="{v["sha"]}",', ' split="train", streaming=True,', ' cache_dir="D:/cpu1/datasets",', ')', 'row = next(iter(dataset))', 'print(row.keys())', '```', '',
238
+ 'Only a training split is declared by the Hub card. If you need held-out evaluation, create and document a deterministic independent split or use a separate test corpus. Never silently substitute the training split for validation. Do not zip the byte and BPE streams by shard index or row index.', '',
239
+ '## Model evaluation', '',
240
+ 'The new model-card audit uses a separate 32-passage WikiText test sample and loads no teacher model. Its corpus, procedure and raw results are in the [FP32 model repository](https://huggingface.co/Cukinator/cpu1-ablations-final). That diagnostic does not evaluate this dataset’s teacher coverage, entropy, hidden-state alignment, or full training distribution.', '',
241
+ '## Related artifacts', '', '- [Source and adapters](https://github.com/igna-s/New-1.58)', '- [FP32 exports and common evaluation](https://huggingface.co/Cukinator/cpu1-ablations-final)', '- [Packed finals and intermediate checkpoints](https://huggingface.co/Cukinator/cpu1-ablation-checkpoints)', '',
242
+ '## License', '', 'The dataset Hub metadata declares MIT. FineWeb source data and Qwen tokenizer/teacher assets retain their own upstream licenses. The repository label does not replace those upstream terms.', '']
243
+ (ROOT/'docs'/'model_cards'/'cpu1-ablation-dataset.md').write_text('\n'.join(lines), encoding='utf-8')
244
+
245
+
246
+ def save_environment(records):
247
+ existing = OUT/'environment.json'
248
+ if existing.exists():
249
+ metadata = json.loads(existing.read_text(encoding='utf-8'))
250
+ if metadata.get('audit_script_sha256') == records[0]['suite_sha256']:
251
+ return # Rendering historical evidence must preserve its provenance.
252
+ from evaluation.audit import save
253
+ from importlib.metadata import version
254
+ packages = {name: version(name) for name in ('torch','transformers','tokenizers','huggingface_hub','datasets','numpy','pyarrow','psutil')}
255
+ tokenizer_revision = next((r['tokenizer_revision'] for r in records if r.get('tokenizer_revision')), None)
256
+ save(OUT/'environment.json', {'python': sys.version, 'packages': packages, 'hardware': records[0].get('processor', 'CPU'), 'platform': records[0]['platform'], 'threads': records[0]['threads'], 'source_commit': records[0].get('source_commit'), 'source_changes': records[0].get('source_changes'), 'audit_script_sha256': records[0]['suite_sha256'], 'qwen_tokenizer_resolved_revision': tokenizer_revision})
257
+
258
+
259
+ def upload_cards(inventory):
260
+ from huggingface_hub import HfApi, CommitOperationAdd
261
+ from evaluation.audit import save
262
+ import os
263
+ token = os.environ.get('HF_TOKEN')
264
+ if not token:
265
+ print('Provide the authorized Hugging Face token on stdin (not saved).', flush=True)
266
+ token = sys.stdin.readline().strip()
267
+ api = HfApi(token=token)
268
+ identity = api.whoami()
269
+ if identity['name'].lower() != 'cukinator':
270
+ raise ValueError('Token does not belong to the authorized repository owner')
271
+ commits = {}
272
+ for repo, stem in [(FINAL,'cpu1-ablations-final'), (ARCHIVE,'cpu1-ablation-checkpoints')]:
273
+ current = api.model_info(repo, files_metadata=True)
274
+ parent = current.sha
275
+ # A README-only publication changes the tip without changing audited weights.
276
+ # Pin the new commit to that tip, but require the exact audited model inventory.
277
+ expected = {f['rfilename']: (f['size'], f['lfs']['sha256'])
278
+ for f in inventory[repo]['siblings'] if f['rfilename'].endswith('.pt')}
279
+ actual = {f.rfilename: (f.size, f.lfs.sha256 if f.lfs else None)
280
+ for f in current.siblings if f.rfilename.endswith('.pt')}
281
+ if actual != expected:
282
+ raise ValueError(f'{repo} model files changed since the audit; rerun evaluation before publishing')
283
+ files = [('README.md', ROOT/'docs'/'model_cards'/f'{stem}.md')]
284
+ files += [(f'evaluation/{p.name}', p) for p in sorted(OUT.glob('*.json'))]
285
+ files += [('evaluation/CORPUS_LICENSE.txt', OUT/'CORPUS_LICENSE.txt')]
286
+ source_dirs = ('cpu1', 'training', 'evaluation', 'tests')
287
+ files += [(f'evaluation/source/{p.relative_to(ROOT).as_posix()}', p)
288
+ for directory in source_dirs for p in sorted((ROOT/directory).glob('*.py'))]
289
+ files += [(f'evaluation/source/requirements/{p.name}', p)
290
+ for p in sorted((ROOT/'requirements').glob('*.txt'))]
291
+ commit = api.create_commit(repo_id=repo, repo_type='model', parent_commit=parent, commit_message='Rebuild model card from reproducible CPU audit; include evidence and limitations', operations=[CommitOperationAdd(path_in_repo=name,path_or_fileobj=str(path)) for name,path in files])
292
+ commits[repo] = {'oid': commit.oid, 'url': commit.commit_url}
293
+ # Verify the actual published README bytes against the local card.
294
+ from huggingface_hub import hf_hub_download
295
+ readme = hf_hub_download(repo, 'README.md', revision=commit.oid, cache_dir=os.environ['HF_HUB_CACHE'], token=token)
296
+ if Path(readme).read_bytes() != files[0][1].read_bytes():
297
+ raise RuntimeError('Published card verification failed')
298
+ print(f'Published and verified {repo}: {commit.oid}', flush=True)
299
+ save(OUT/'publication.json', commits)
evaluation/source/evaluation/checkpoints.py ADDED
@@ -0,0 +1,115 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Strict loading of custom CPU-1 checkpoints for the common audit."""
2
+ from dataclasses import dataclass
3
+ from pathlib import Path
4
+
5
+ import torch
6
+
7
+
8
+ @dataclass
9
+ class AblationInfo:
10
+ """
11
+ Architecture metadata used by the checkpoint audit.
12
+ """
13
+ run_name: str
14
+ model_class: str # config['name']
15
+ tok: str # byte / byte_dynamic / bpe
16
+ quant: str # ternary / fp16
17
+ is_byte_level: bool
18
+ is_recurrent: bool # True → MLGRU (estado persistente), False → Transformer
19
+ is_ternary: bool
20
+ use_local_dec: bool # LocalByteDecoder
21
+ patch_size: int # bytes por patch
22
+ vocab_size: int # 256 (byte) / bpe_vocab_size
23
+ n_heads: int # patch_size para byte, 1 para BPE
24
+ d_model: int
25
+ n_layers: int
26
+
27
+ @staticmethod
28
+ def from_config(run_name: str, config: dict) -> "AblationInfo":
29
+ tok = config.get("tokenization", "byte")
30
+ quant = config.get("quantization", "fp16")
31
+ mixer = config.get("token_mixer", "mlgru")
32
+ patch = int(config.get("patch_size", 4))
33
+ is_byte = tok in ("byte", "byte_dynamic")
34
+ use_lbd = is_byte and config.get("use_local_decoder", True)
35
+ vocab = 256 if is_byte else int(config.get("bpe_vocab_size", 16384))
36
+ # n_heads: siempre patch_size para byte (LocalDecoder se desempaqueta
37
+ # en _get_logits_list de evaluation.correctness a P tensores [B,S,256])
38
+ n_heads = patch if is_byte else 1
39
+ return AblationInfo(
40
+ run_name = run_name,
41
+ model_class = config.get("name", run_name),
42
+ tok = tok,
43
+ quant = quant,
44
+ is_byte_level = is_byte,
45
+ is_recurrent = (mixer == "mlgru"),
46
+ is_ternary = (quant == "ternary"),
47
+ use_local_dec = use_lbd,
48
+ patch_size = patch,
49
+ vocab_size = vocab,
50
+ n_heads = n_heads,
51
+ d_model = int(config.get("d_model", 512)),
52
+ n_layers = int(config.get("n_layers", 12)),
53
+ )
54
+
55
+
56
+ class AblationEngine:
57
+ """Loaded model and recurrent carry for correctness checks and timing."""
58
+
59
+ def __init__(self, model, config: dict, info: AblationInfo):
60
+ self._model = model
61
+ self.config = config
62
+ self.info = info
63
+ self.hidden_states = None
64
+
65
+ def _get_model(self):
66
+ return self._model
67
+
68
+ def reset_state(self):
69
+ self.hidden_states = None
70
+
71
+ def _ablation_step(self, model, dummy: torch.Tensor):
72
+ result = model(
73
+ dummy,
74
+ hidden_states=self.hidden_states if self.info.is_recurrent else None,
75
+ training_phase=1,
76
+ target_elimination=0.0,
77
+ )
78
+ return result[0], result[1] if self.info.is_recurrent else None
79
+
80
+
81
+ def load_run(ckpt_path: Path, run_name: str) -> tuple:
82
+ """
83
+ Carga el checkpoint y construye el modelo. Soporta tres formatos:
84
+
85
+ compact_2bit → load_ablation_checkpoint desempaqueta el ternary
86
+ bf16 → load_ablation_checkpoint castea a fp32
87
+ fp32 → carga directa desde raw['state_dict']
88
+ (formato de cpu1-ablations-final y export_ablations.py)
89
+
90
+ Devuelve (engine, raw_ckpt, config, missing_keys, unexpected_keys).
91
+ """
92
+ from training.ablation import load_ablation_checkpoint, build_ablation_model
93
+
94
+ raw = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
95
+
96
+ fmt = raw.get("format", "bf16")
97
+ if fmt == "fp32":
98
+ # Descargado de Cukinator/cpu1-ablations-final o exportado localmente.
99
+ # Claves: state_dict, config, metrics, step, run_name, format, source_repo
100
+ config = raw["config"]
101
+ state = raw["state_dict"]
102
+ else:
103
+ # compact_2bit o bf16: load_ablation_checkpoint lo desempaqueta / castea.
104
+ state, config = load_ablation_checkpoint(str(ckpt_path))
105
+
106
+ config["use_gradient_checkpointing"] = False
107
+ model = build_ablation_model(config)
108
+ missing, unexpected = model.load_state_dict(state, strict=True)
109
+ if missing or unexpected:
110
+ raise RuntimeError(f"Checkpoint incompatible: missing={missing}, unexpected={unexpected}")
111
+ model.eval()
112
+
113
+ info = AblationInfo.from_config(run_name, config)
114
+ engine = AblationEngine(model, config, info)
115
+ return engine, raw, config, missing, unexpected
evaluation/source/evaluation/correctness.py ADDED
@@ -0,0 +1,239 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Eval — Correctness: verifica propiedades del modelo.
3
+
4
+ Checks adaptivos segun arquitectura (engine.info):
5
+ forward_shapes — Forward pass sin crash para varias shapes
6
+ probs — Distribuciones de probabilidad validas
7
+ determinism — Mismo input → mismo output
8
+ state_prop — Estado MLGRU se propaga correctamente (solo recurrente)
9
+ ternary_weights — Pesos BitLinear son {-1, 0, 1} (solo ternary)
10
+ state_size — Estado MLGRU < 1MB (solo recurrente)
11
+
12
+ Checks use the model loaded by evaluation.checkpoints.
13
+ """
14
+ import torch
15
+
16
+ from evaluation.helpers import get_model, BitLinear
17
+
18
+
19
+ def _make_dummy_input(info, batch: int = 1, seq_len: int = 64):
20
+ """Crear input dummy segun tipo de modelo."""
21
+ if info.is_byte_level:
22
+ return torch.randint(0, 256, (batch, seq_len, info.patch_size), dtype=torch.long)
23
+ else:
24
+ return torch.randint(0, info.vocab_size, (batch, seq_len), dtype=torch.long)
25
+
26
+
27
+ def _unpack_forward(result):
28
+ """Extraer logits y states de cualquier modelo.
29
+ CPU1ModelV2: (logits_tuple, states)
30
+ AblationModel: (logits_or_tuple, states, extra_losses)
31
+ """
32
+ if len(result) == 2:
33
+ return result[0], result[1]
34
+ else:
35
+ return result[0], result[1]
36
+
37
+
38
+ def _get_logits_list(logits, info):
39
+ """Normalizar logits a lista de tensores [B, S, V] siempre.
40
+
41
+ Casos:
42
+ - tuple/list de heads (byte sin LocalDecoder): devuelve lista tal cual.
43
+ - LocalByteDecoder: [B, S, P, 256] → lista de P tensores [B, S, 256].
44
+ - tensor 3D (BPE o produccion): devuelve lista de 1 elemento.
45
+ """
46
+ if isinstance(logits, (tuple, list)):
47
+ return list(logits)
48
+ if logits.dim() == 4:
49
+ # LocalByteDecoder: [B, S, P, 256] → P × [B, S, 256]
50
+ return [logits[:, :, i, :] for i in range(logits.shape[2])]
51
+ return [logits]
52
+
53
+
54
+ def check_forward_shapes(engine) -> dict:
55
+ """Forward sin crash para varias combinaciones batch × seq."""
56
+ info = engine.info
57
+ model = engine._get_model()
58
+ configs = [(b, s) for b in [1, 2] for s in [1, 64, 256]]
59
+ errors = []
60
+
61
+ for batch, seq_len in configs:
62
+ x = _make_dummy_input(info, batch, seq_len)
63
+ with torch.inference_mode():
64
+ result = model(x)
65
+ logits, _ = _unpack_forward(result)
66
+
67
+ logits_list = _get_logits_list(logits, info)
68
+ expected_heads = info.n_heads
69
+ expected_vocab = info.vocab_size
70
+
71
+ if len(logits_list) != expected_heads:
72
+ errors.append(f"Esperado {expected_heads} heads, obtenido {len(logits_list)}")
73
+ continue
74
+
75
+ for i, l in enumerate(logits_list):
76
+ expected_shape = (batch, seq_len, expected_vocab)
77
+ if l.shape != expected_shape:
78
+ errors.append(f"Head {i}: shape {l.shape} != {expected_shape}")
79
+ if torch.isnan(l).any():
80
+ errors.append(f"NaN en logits batch={batch} seq={seq_len} head={i}")
81
+ if torch.isinf(l).any():
82
+ errors.append(f"Inf en logits batch={batch} seq={seq_len} head={i}")
83
+
84
+ assert not errors, "\n".join(errors)
85
+ return {'configs_tested': len(configs), 'heads': info.n_heads}
86
+
87
+
88
+ def check_probs(engine) -> dict:
89
+ """Las heads producen distribuciones validas (suman a 1)."""
90
+ info = engine.info
91
+ model = engine._get_model()
92
+ x = _make_dummy_input(info, batch=1, seq_len=128)
93
+
94
+ with torch.inference_mode():
95
+ logits, _ = _unpack_forward(model(x))
96
+
97
+ logits_list = _get_logits_list(logits, info)
98
+ max_err = 0.0
99
+ for i, head_logits in enumerate(logits_list):
100
+ probs = torch.softmax(head_logits, dim=-1)
101
+ sums = probs.sum(dim=-1)
102
+ err = (sums - 1.0).abs().max().item()
103
+ max_err = max(max_err, err)
104
+ assert err < 1e-4, f"Head {i}: suma de probs error={err:.6f}"
105
+
106
+ return {'max_error': max_err}
107
+
108
+
109
+ def check_determinism(engine) -> dict:
110
+ """Mismo input con mismo estado da output identico."""
111
+ info = engine.info
112
+ model = engine._get_model()
113
+ torch.manual_seed(42)
114
+ x = _make_dummy_input(info, batch=1, seq_len=64)
115
+
116
+ with torch.inference_mode():
117
+ logits1, _ = _unpack_forward(model(x, hidden_states=None))
118
+ logits2, _ = _unpack_forward(model(x, hidden_states=None))
119
+
120
+ list1 = _get_logits_list(logits1, info)
121
+ list2 = _get_logits_list(logits2, info)
122
+
123
+ for i in range(len(list1)):
124
+ diff = (list1[i] - list2[i]).abs().max().item()
125
+ assert diff == 0.0, f"Head {i}: output no deterministico, diff={diff}"
126
+
127
+ return {'deterministic': True}
128
+
129
+
130
+ def check_state_propagation(engine) -> dict:
131
+ """Procesar en chunks da el mismo resultado que todo junto (solo recurrente)."""
132
+ info = engine.info
133
+ if not info.is_recurrent:
134
+ return {'skipped': True, 'reason': 'modelo no recurrente (Transformer)'}
135
+
136
+ model = engine._get_model()
137
+ torch.manual_seed(0)
138
+ x = _make_dummy_input(info, batch=1, seq_len=128)
139
+
140
+ with torch.inference_mode():
141
+ logits_full, _ = _unpack_forward(model(x))
142
+
143
+ states = None
144
+ chunk_outputs = []
145
+ with torch.inference_mode():
146
+ for i in range(0, 128, 32):
147
+ chunk = x[:, i:i + 32] if x.dim() == 2 else x[:, i:i + 32, :]
148
+ result = model(chunk, states)
149
+ logits_chunk, states = _unpack_forward(result)
150
+ chunk_outputs.append(_get_logits_list(logits_chunk, info))
151
+
152
+ full_heads = _get_logits_list(logits_full, info)
153
+ diff = max((full_heads[k] - torch.cat([chunk[k] for chunk in chunk_outputs], dim=1)).abs().max().item()
154
+ for k in range(len(full_heads)))
155
+ assert diff < 1e-4, f"State propagation error: {diff:.6f}"
156
+ return {'max_diff': diff}
157
+
158
+
159
+ def check_ternary_weights(engine) -> dict:
160
+ """Pesos BitLinear quantizados son {-1, 0, 1} (solo ternary)."""
161
+ info = engine.info
162
+ if not info.is_ternary:
163
+ return {'skipped': True, 'reason': 'modelo FP16 (sin BitLinear)'}
164
+
165
+ model = engine._get_model()
166
+ n_checked = 0
167
+
168
+ for name, module in model.named_modules():
169
+ if not isinstance(module, BitLinear):
170
+ continue
171
+
172
+ if hasattr(module, 'weight_T_ternary'):
173
+ w = module.weight_T_ternary
174
+ unique = set(w.unique().tolist())
175
+ invalid = unique - {-1.0, 0.0, 1.0}
176
+ assert not invalid, f"{name}: valores no ternarios en weight_T_ternary: {invalid}"
177
+ elif module.weight.numel() > 0:
178
+ w = module.weight.data
179
+ w_scale = w.abs().mean()
180
+ w_q = (w / (w_scale + 1e-8)).round().clamp(-1, 1)
181
+ unique = set(w_q.unique().tolist())
182
+ invalid = unique - {-1.0, 0.0, 1.0}
183
+ assert not invalid, f"{name}: valores no ternarios: {invalid}"
184
+
185
+ n_checked += 1
186
+
187
+ assert n_checked > 0, "No se encontraron modulos BitLinear"
188
+ return {'modules_checked': n_checked}
189
+
190
+
191
+ def check_state_size(engine) -> dict:
192
+ """Estado MLGRU < 1MB (solo recurrente)."""
193
+ info = engine.info
194
+ if not info.is_recurrent:
195
+ return {'skipped': True, 'reason': 'modelo no recurrente (Transformer)'}
196
+
197
+ model = engine._get_model()
198
+ x = _make_dummy_input(info, batch=1, seq_len=1)
199
+
200
+ with torch.inference_mode():
201
+ _, states = _unpack_forward(model(x))
202
+
203
+ # Filtrar None states (pueden existir si algun bloque es transformer)
204
+ valid_states = [s for s in states if s is not None]
205
+ state_bytes = sum(s.numel() * s.element_size() for s in valid_states)
206
+ state_kb = state_bytes / 1024
207
+
208
+ assert state_kb < 1024, f"Estado MLGRU demasiado grande: {state_kb:.1f} KB"
209
+ return {'state_kb': round(state_kb, 1), 'batch_size': 1, 'dtype': 'float32'}
210
+
211
+
212
+ # =========================================================================
213
+ # Entry point
214
+ # =========================================================================
215
+
216
+ def run_correctness(engine) -> dict:
217
+ """Ejecutar todos los checks de correctness. 1 engine → 6 checks."""
218
+ from evaluation.helpers import run_check
219
+
220
+ print(f"\n--- Correctness ({engine.info.model_class}) ---")
221
+ results = {}
222
+ checks = [
223
+ ('forward_shapes', check_forward_shapes),
224
+ ('probs', check_probs),
225
+ ('determinism', check_determinism),
226
+ ('state_prop', check_state_propagation),
227
+ ('ternary_weights', check_ternary_weights),
228
+ ('state_size', check_state_size),
229
+ ]
230
+ for name, fn in checks:
231
+ r = run_check(name, fn, engine)
232
+ results[name] = r
233
+ status = r['status']
234
+ extra = ''
235
+ if r.get('result', {}).get('skipped'):
236
+ extra = f" (skip: {r['result']['reason']})"
237
+ print(f" {status:5s} {name}{extra}")
238
+
239
+ return results
evaluation/source/evaluation/diagnostics.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Check the recurrent scan against its defining equation without model files."""
2
+ from pathlib import Path
3
+ import sys
4
+
5
+ ROOT = Path(__file__).resolve().parents[1]
6
+ sys.path.insert(0, str(ROOT))
7
+
8
+
9
+ def main():
10
+ import torch
11
+ from cpu1.components import AblationMLGRU
12
+ from evaluation.audit import OUT, save, sha
13
+ torch.set_num_threads(1)
14
+ torch.manual_seed(19)
15
+ c = torch.randn(1, 32, 8)
16
+ f = torch.sigmoid(torch.randn(1, 32, 8))
17
+ initial = torch.randn(1, 8)
18
+ actual, _ = AblationMLGRU._parallel_scan(c, f, initial)
19
+ h = initial.clone()
20
+ steps = []
21
+ for i in range(32):
22
+ h = f[:, i]*h + (1-f[:, i])*c[:, i]
23
+ steps.append(h)
24
+ reference = torch.stack(steps, dim=1)
25
+ first, carry = AblationMLGRU._parallel_scan(c[:, :16], f[:, :16], initial)
26
+ second, _ = AblationMLGRU._parallel_scan(c[:, 16:], f[:, 16:], carry)
27
+ serial_diff = (actual-reference).abs().max().item()
28
+ chunk_diff = (actual-torch.cat([first, second], dim=1)).abs().max().item()
29
+ result = {'seed':19, 'shape':[1,32,8], 'torch':torch.__version__, 'dtype':'float32', 'reference':'h[t] = f[t] * h[t-1] + (1-f[t]) * c[t]', 'serial_max_diff':serial_diff, 'chunks_max_diff':chunk_diff, 'tolerance':1e-4, 'status':'PASS' if max(serial_diff,chunk_diff)<1e-4 else 'FAIL', 'components_sha256':sha(ROOT/'cpu1'/'components.py'), 'script_sha256':sha(Path(__file__))}
30
+ save(OUT/'implementation_diagnostics.json', result)
31
+ print(result)
32
+
33
+
34
+ if __name__ == '__main__':
35
+ main()
evaluation/source/evaluation/export.py ADDED
@@ -0,0 +1,252 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Desempaqueta los checkpoints compactos (2-bit ternary) del ablation study,
3
+ guarda versiones float32 listas para usar y las sube a HuggingFace.
4
+
5
+ Repo destino: Cukinator/cpu1-ablations-final
6
+ Modelos locales: models/ablations_final/{run_name}/
7
+
8
+ Uso:
9
+ python -m evaluation.export
10
+ python -m evaluation.export --runs run_02 run_03
11
+ python -m evaluation.export --skip-upload # solo exporta localmente
12
+ python -m evaluation.export --skip-export # solo sube los ya exportados
13
+ """
14
+
15
+ import os
16
+ import sys
17
+ import json
18
+ import time
19
+ import argparse
20
+ from pathlib import Path
21
+
22
+ ROOT = Path(__file__).parent.parent
23
+ sys.path.insert(0, str(ROOT))
24
+
25
+ HF_REPO_SRC = "Cukinator/cpu1-ablation-checkpoints"
26
+ HF_REPO_DEST = "Cukinator/cpu1-ablations-final"
27
+ MODELS_DIR = ROOT / "local" / "models" / "ablation_models"
28
+ EXPORT_DIR = ROOT / "local" / "models" / "ablations_final"
29
+
30
+ AVAILABLE_RUNS = [
31
+ "run_01",
32
+ "run_02",
33
+ "run_02a_byte_only_heads",
34
+ "run_03",
35
+ "run_04",
36
+ "run_05",
37
+ "run_05b_kernel_strict",
38
+ "run_06",
39
+ "run_07",
40
+ "run_08",
41
+ "run_09",
42
+ "run_10",
43
+ "run_13",
44
+ "run_14",
45
+ "run_15",
46
+ "run_16",
47
+ # Round-2 re-runs at 15 tok/param
48
+ "run_04_r2",
49
+ "run_05_r2",
50
+ "run_07_r2",
51
+ "run_08_r2",
52
+ "run_09_r2",
53
+ "run_10_r2",
54
+ "run_13_r2",
55
+ "run_14_r2",
56
+ "run_15_r2",
57
+ "run_16_r2",
58
+ # Round-3 cold-start rescue (bf16 AMP + lr_scale=2 + CE-only + 50 tok/p)
59
+ "run_04_r3",
60
+ "run_07_r3",
61
+ "run_14_r3",
62
+ "run_15_r3",
63
+ ]
64
+
65
+ USAGE_README = (Path(__file__).parent.parent / "docs" / "model_cards"
66
+ / "cpu1-ablations-final.md").read_text(encoding="utf-8")
67
+
68
+
69
+ def get_hf_token() -> str | None:
70
+ token = os.environ.get("HF_TOKEN")
71
+ if not token:
72
+ env_path = ROOT / ".env"
73
+ if env_path.exists():
74
+ for line in env_path.read_text().splitlines():
75
+ if line.startswith("HF_TOKEN"):
76
+ token = line.split("=", 1)[1].strip().strip('"')
77
+ break
78
+ return token or None
79
+
80
+
81
+ def export_run(run_name: str) -> Path | None:
82
+ """Desempaqueta el checkpoint compacto y guarda un .pt float32."""
83
+ import torch
84
+ from training.ablation import load_ablation_checkpoint
85
+
86
+ src = MODELS_DIR / run_name / f"checkpoint_{run_name}_final.pt"
87
+ if not src.exists():
88
+ print(f" [missing] {src}")
89
+ return None
90
+
91
+ out_dir = EXPORT_DIR / run_name
92
+ out_dir.mkdir(parents=True, exist_ok=True)
93
+ out_path = out_dir / "model.pt"
94
+
95
+ if out_path.exists():
96
+ print(f" [skip] {run_name} ya exportado ({out_path.stat().st_size / 1e6:.1f} MB)")
97
+ return out_path
98
+
99
+ print(f" Exportando {run_name}...")
100
+ t0 = time.perf_counter()
101
+
102
+ try:
103
+ state, config = load_ablation_checkpoint(str(src))
104
+ except Exception as e:
105
+ print(f" [ERROR] load: {e}")
106
+ return None
107
+
108
+ # Métricas del checkpoint original
109
+ import torch as _torch
110
+ raw = _torch.load(str(src), map_location="cpu", weights_only=False)
111
+ metrics = raw.get("metrics", {})
112
+ step = raw.get("step", "?")
113
+
114
+ payload = {
115
+ "state_dict": state,
116
+ "config": config,
117
+ "metrics": metrics,
118
+ "step": step,
119
+ "run_name": run_name,
120
+ "format": "fp32",
121
+ "source_repo": HF_REPO_SRC,
122
+ }
123
+
124
+ import torch as _torch
125
+ _torch.save(payload, out_path)
126
+ elapsed = time.perf_counter() - t0
127
+ size_mb = out_path.stat().st_size / 1e6
128
+ print(f" [ok] {run_name} → {out_path} ({size_mb:.1f} MB, {elapsed:.1f}s)")
129
+ return out_path
130
+
131
+
132
+ def upload_run(run_name: str, token: str | None) -> bool:
133
+ """Sube models/ablations_final/{run_name}/ al repo HF."""
134
+ from huggingface_hub import HfApi
135
+
136
+ out_path = EXPORT_DIR / run_name / "model.pt"
137
+ if not out_path.exists():
138
+ print(f" [skip upload] {run_name} — archivo no encontrado")
139
+ return False
140
+
141
+ api = HfApi(token=token)
142
+
143
+ print(f" Subiendo {run_name} ({out_path.stat().st_size / 1e6:.1f} MB)...")
144
+ t0 = time.perf_counter()
145
+ try:
146
+ api.upload_file(
147
+ path_or_fileobj=str(out_path),
148
+ path_in_repo=f"{run_name}/model.pt",
149
+ repo_id=HF_REPO_DEST,
150
+ repo_type="model",
151
+ commit_message=f"Add {run_name} (float32, unpacked)",
152
+ )
153
+ elapsed = time.perf_counter() - t0
154
+ print(f" [ok] {run_name} subido en {elapsed:.1f}s")
155
+ return True
156
+ except Exception as e:
157
+ print(f" [ERROR] upload {run_name}: {e}")
158
+ return False
159
+
160
+
161
+ def ensure_repo(token: str | None):
162
+ from huggingface_hub import HfApi
163
+ api = HfApi(token=token)
164
+ try:
165
+ api.repo_info(repo_id=HF_REPO_DEST, repo_type="model")
166
+ print(f" Repo existente: {HF_REPO_DEST}")
167
+ except Exception:
168
+ print(f" Creando repo: {HF_REPO_DEST} ...")
169
+ api.create_repo(
170
+ repo_id=HF_REPO_DEST,
171
+ repo_type="model",
172
+ private=False,
173
+ exist_ok=True,
174
+ )
175
+ print(f" [ok] Repo creado")
176
+
177
+
178
+ def upload_readme(token: str | None):
179
+ from huggingface_hub import HfApi
180
+ api = HfApi(token=token)
181
+ try:
182
+ import tempfile, os
183
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".md",
184
+ delete=False, encoding="utf-8") as f:
185
+ f.write(USAGE_README)
186
+ tmp = f.name
187
+ api.upload_file(
188
+ path_or_fileobj=tmp,
189
+ path_in_repo="README.md",
190
+ repo_id=HF_REPO_DEST,
191
+ repo_type="model",
192
+ commit_message="Add README",
193
+ )
194
+ os.unlink(tmp)
195
+ print(" [ok] README.md subido")
196
+ except Exception as e:
197
+ print(f" [warn] README upload: {e}")
198
+
199
+
200
+ def main():
201
+ parser = argparse.ArgumentParser(
202
+ description="Exporta y sube checkpoints ablation desempaquetados"
203
+ )
204
+ parser.add_argument("--runs", nargs="+", default=None)
205
+ parser.add_argument("--skip-upload", action="store_true",
206
+ help="Solo exporta localmente, no sube a HF")
207
+ parser.add_argument("--skip-export", action="store_true",
208
+ help="Solo sube los .pt ya exportados (no re-exporta)")
209
+ args = parser.parse_args()
210
+
211
+ runs = args.runs or AVAILABLE_RUNS
212
+ token = get_hf_token()
213
+
214
+ print(f"\n{'=' * 65}")
215
+ print(f" Export ablations -- {HF_REPO_SRC} -> {HF_REPO_DEST}")
216
+ print(f" Runs: {runs}")
217
+ print(f" HF_TOKEN: {'****' + token[-4:] if token else 'None'}")
218
+ print(f" Local export dir: {EXPORT_DIR}")
219
+ print(f"{'=' * 65}\n")
220
+
221
+ EXPORT_DIR.mkdir(parents=True, exist_ok=True)
222
+
223
+ # — Fase 1: exportar localmente —
224
+ if not args.skip_export:
225
+ print("## Exportando checkpoints (2-bit → float32)\n")
226
+ for run in runs:
227
+ export_run(run)
228
+
229
+ # — Fase 2: subir a HF —
230
+ if not args.skip_upload:
231
+ if not token:
232
+ print("\n[ERROR] No hay HF_TOKEN — no se puede subir.")
233
+ return
234
+
235
+ print("\n## Subiendo a HuggingFace\n")
236
+ ensure_repo(token)
237
+ upload_readme(token)
238
+
239
+ ok = 0
240
+ for run in runs:
241
+ if upload_run(run, token):
242
+ ok += 1
243
+
244
+ print(f"\n{'=' * 65}")
245
+ print(f" Listo — {ok}/{len(runs)} modelos subidos a {HF_REPO_DEST}")
246
+ print(f"{'=' * 65}\n")
247
+ else:
248
+ print(f"\nExportacion local completa → {EXPORT_DIR}")
249
+
250
+
251
+ if __name__ == "__main__":
252
+ main()
evaluation/source/evaluation/helpers.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Shared access and result handling for the active checkpoint audit."""
2
+ import os
3
+ import time
4
+
5
+ os.environ.setdefault('KMP_AFFINITY', 'granularity=fine,compact,1,0')
6
+ os.environ.setdefault('KMP_BLOCKTIME', '0')
7
+
8
+ from cpu1.components import BitLinear
9
+
10
+
11
+ def get_model(engine):
12
+ return engine._get_model()
13
+
14
+
15
+ def run_check(name: str, fn, *args, **kwargs) -> dict:
16
+ """Ejecutar un check individual con manejo de errores y timing."""
17
+ try:
18
+ t0 = time.perf_counter()
19
+ result = fn(*args, **kwargs)
20
+ elapsed = time.perf_counter() - t0
21
+ status = 'N/A' if isinstance(result, dict) and result.get('skipped') else 'PASS'
22
+ return {'status': status, 'time_s': round(elapsed, 2), 'result': result}
23
+ except AssertionError as e:
24
+ return {'status': 'FAIL', 'error': str(e)}
25
+ except Exception as e:
26
+ return {'status': 'ERROR', 'error': f'{type(e).__name__}: {e}'}
evaluation/source/evaluation/verify_checkpoints.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Verify the two published loading examples, sequentially, in D: storage."""
2
+ import gc
3
+ import hashlib
4
+ import os
5
+ from pathlib import Path
6
+ import sys
7
+
8
+ ROOT = Path(__file__).resolve().parents[1]
9
+ sys.path.insert(0, str(ROOT))
10
+
11
+
12
+ def digest(state):
13
+ result = hashlib.sha256()
14
+ for name, value in sorted(state.items()):
15
+ result.update(name.encode())
16
+ result.update(str((value.dtype, tuple(value.shape))).encode())
17
+ result.update(value.contiguous().numpy().tobytes())
18
+ return result.hexdigest()
19
+
20
+
21
+ def main():
22
+ import json
23
+ from evaluation.audit import FINAL, ARCHIVE, OUT, save, sha
24
+ storage = Path('D:/cpu1-readme-audit')
25
+ assert not list((storage/'active').glob('*/model.pt')), 'Finish the model audit first'
26
+ os.environ['HF_HOME'] = str(storage/'hf_home')
27
+ os.environ['HF_HUB_CACHE'] = str(storage/'hf_hub_cache')
28
+ os.environ['HF_HUB_DISABLE_XET'] = '1'
29
+ os.environ['HF_HUB_DISABLE_PROGRESS_BARS'] = '1'
30
+ import torch
31
+ from huggingface_hub import hf_hub_download
32
+ from training.ablation import build_ablation_model, load_ablation_checkpoint
33
+ torch.set_num_threads(1)
34
+ inventory = json.loads((OUT/'inventory.json').read_text(encoding='utf-8'))
35
+ records = []
36
+ for repo, filename in [(ARCHIVE,'run_04/checkpoint_run_04_final.pt'),(FINAL,'run_04/model.pt')]:
37
+ path = Path(hf_hub_download(repo,filename,revision=inventory[repo]['sha'],local_dir=str(storage/'smoke'),cache_dir=str(storage/'hf_hub_cache')))
38
+ raw = torch.load(path,map_location='cpu',weights_only=False)
39
+ fmt = raw.get('format')
40
+ if repo == FINAL:
41
+ state, config = raw['state_dict'], raw['config']
42
+ else:
43
+ state, config = load_ablation_checkpoint(str(path))
44
+ fingerprint = digest(state)
45
+ config['use_gradient_checkpointing'] = False
46
+ model = build_ablation_model(config)
47
+ model.load_state_dict(state,strict=True)
48
+ model.eval()
49
+ x = torch.tensor(list(b'The history of AI. ')).reshape(1,-1,4)
50
+ with torch.inference_mode():
51
+ logits, states, extra = model(x,training_phase=1,target_elimination=0.)
52
+ assert all(torch.isfinite(head).all() for head in logits)
53
+ shapes = [list(head.shape) for head in logits]
54
+ assert shapes == [[1,5,256]]*4
55
+ records.append({'repo':repo,'revision':inventory[repo]['sha'],'file':filename,'format':fmt,'sha256':sha(path),'state_sha256':fingerprint,'strict_load':True,'finite_logits':True,'shapes':shapes})
56
+ print(f'Example verified: {repo} ({fmt})',flush=True)
57
+ del raw, state, model, logits, states, extra
58
+ gc.collect()
59
+ path.unlink()
60
+ save(OUT/'loading_examples.json',{'examples':records,'run_04_state_dicts_identical':records[0]['state_sha256']==records[1]['state_sha256'],'scope':'only the run_04 pair; does not establish parity across other models or intermediate steps'})
61
+
62
+
63
+ if __name__ == '__main__':
64
+ main()
evaluation/source/requirements/audit.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ -r training.txt
2
+ psutil>=5.9.0
evaluation/source/requirements/data.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ -r training.txt
2
+ psutil>=5.9.0
3
+ tqdm>=4.66.0
4
+ # Optional teacher quantization: install bitsandbytes or torchao when needed.
evaluation/source/requirements/training.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ # Install the PyTorch build for your hardware separately.
2
+ transformers>=4.47.0
3
+ datasets>=3.2.0
4
+ huggingface_hub>=0.27.0
5
+ tokenizers>=0.21.0
6
+ numpy>=1.26.0
7
+ pyarrow>=18.0.0
evaluation/source/tests/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ """Offline regression and integration checks."""
evaluation/source/tests/test_audit.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Offline regression checks for the scoring protocol, not trained-model quality."""
2
+ import json
3
+ import math
4
+ from pathlib import Path
5
+ import tempfile
6
+ import unittest
7
+
8
+ import torch
9
+
10
+ from evaluation.audit import loss_eval, save
11
+ from evaluation.helpers import run_check
12
+ from cpu1.components import LocalByteDecoder
13
+
14
+
15
+ class AuditTests(unittest.TestCase):
16
+ @classmethod
17
+ def setUpClass(cls):
18
+ torch.set_num_threads(1)
19
+
20
+ def test_local_teacher_forcing_does_not_look_ahead(self):
21
+ torch.manual_seed(4)
22
+ decoder = LocalByteDecoder(16, 4, 16, 4, use_ternary=False).eval()
23
+ hidden = torch.randn(1, 2, 16)
24
+ targets = torch.tensor([[[5, 6, 7, 8], [9, 10, 11, 12]]])
25
+ future = targets.clone()
26
+ future[:, :, 2:] += 20
27
+ with torch.inference_mode():
28
+ a, b = decoder(hidden, targets), decoder(hidden, future)
29
+ # Head 2 sees only ground-truth bytes 0 and 1; head 3 sees byte 2.
30
+ for k in range(3):
31
+ torch.testing.assert_close(a[k], b[k], rtol=0, atol=0)
32
+ self.assertGreater((a[3]-b[3]).abs().max().item(), 0)
33
+
34
+ def test_loss_uses_shifted_targets_and_target_weighting(self):
35
+ class Model:
36
+ def __init__(self):
37
+ self.seen = []
38
+
39
+ def __call__(self, x, byte_targets, **kwargs):
40
+ self.seen.append(byte_targets.clone())
41
+ logits = torch.zeros(*byte_targets.shape, 256)
42
+ # Distinct losses per passage and unequal target counts.
43
+ logits[..., 0] = len(self.seen)
44
+ return tuple(logits[:, :, k] for k in range(4)), None, {}
45
+ model = Model()
46
+ result = loss_eval(model, {'tokenization':'byte', 'patch_size':4}, {'texts':['abcdefgh','ijklmnopqrst']}, None)
47
+ self.assertEqual(result['targets'], 12)
48
+ self.assertEqual(model.seen[0].flatten().tolist(), list(b'efgh'))
49
+ self.assertEqual(model.seen[1].flatten().tolist(), list(b'mnopqrst'))
50
+ expected = (4*math.log(255+math.exp(1))+8*math.log(255+math.exp(2)))/12
51
+ self.assertAlmostEqual(result['mean_nll'], expected, places=5)
52
+
53
+ def test_skips_are_not_passes(self):
54
+ result = run_check('irrelevant', lambda: {'skipped':True,'reason':'not recurrent'})
55
+ self.assertEqual(result['status'], 'N/A')
56
+
57
+ def test_nonfinite_evidence_is_valid_json(self):
58
+ with tempfile.TemporaryDirectory() as directory:
59
+ path = Path(directory)/'record.json'
60
+ save(path, {'loss':float('nan'), 'nested':[float('inf'), 2.]})
61
+ self.assertEqual(json.loads(path.read_text()), {'loss':None,'nested':[None,2.]})
62
+
63
+
64
+ if __name__ == '__main__':
65
+ torch.set_num_threads(1)
66
+ unittest.main()
evaluation/source/tests/test_checkpoints.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Checkpoint compatibility checks for the public package layout."""
2
+ from pathlib import Path
3
+ import tempfile
4
+ import unittest
5
+
6
+ import torch
7
+
8
+ from cpu1.components import pack_model_compact
9
+ from evaluation.checkpoints import load_run
10
+ from training.ablation import build_ablation_model, get_run_config
11
+
12
+
13
+ class CheckpointTests(unittest.TestCase):
14
+ @classmethod
15
+ def setUpClass(cls):
16
+ torch.set_num_threads(1)
17
+ cls.storage = Path('D:/cpu1-readme-audit/test-checkpoints')
18
+ cls.storage.mkdir(parents=True, exist_ok=True)
19
+
20
+ def config(self):
21
+ cfg = get_run_config('run_04')
22
+ cfg.update(d_model=16, n_layers=1, d_ff=32, n_heads=2,
23
+ head_dim=8, window_size=8, local_decoder_dim=16,
24
+ local_decoder_heads=2, use_gradient_checkpointing=False)
25
+ return cfg
26
+
27
+ def test_fp32_and_compact_payloads_load_in_new_package(self):
28
+ config = self.config()
29
+ model = build_ablation_model(config).eval()
30
+ x = torch.tensor(list(b'abcdefgh')).reshape(1, 2, 4)
31
+ for format_name in ('fp32', 'compact_2bit'):
32
+ with self.subTest(format=format_name):
33
+ with tempfile.TemporaryDirectory(dir=self.storage) as directory:
34
+ payload = {'config': config, 'format': format_name}
35
+ if format_name == 'fp32':
36
+ payload['state_dict'] = model.state_dict()
37
+ else:
38
+ payload['compact'] = pack_model_compact(model)
39
+ path = Path(directory)/'model.pt'
40
+ torch.save(payload, path)
41
+ engine, _, loaded_config, missing, extra = load_run(path, 'test')
42
+ self.assertEqual(missing, [])
43
+ self.assertEqual(extra, [])
44
+ self.assertEqual(loaded_config['d_model'], 16)
45
+ with torch.inference_mode():
46
+ heads = engine._get_model()(x, byte_targets=x)[0]
47
+ self.assertEqual([tuple(h.shape) for h in heads], [(1, 2, 256)]*4)
48
+ self.assertTrue(all(torch.isfinite(h).all() for h in heads))
49
+
50
+ def test_incompatible_payload_is_rejected(self):
51
+ config = self.config()
52
+ state = build_ablation_model(config).state_dict()
53
+ state.pop(next(iter(state)))
54
+ with tempfile.TemporaryDirectory(dir=self.storage) as directory:
55
+ path = Path(directory)/'model.pt'
56
+ torch.save({'config': config, 'format': 'fp32', 'state_dict': state}, path)
57
+ with self.assertRaises(RuntimeError):
58
+ load_run(path, 'incompatible')
59
+
60
+
61
+ if __name__ == '__main__':
62
+ unittest.main()
evaluation/source/training/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ """Training entry points for the CPU-1 research models."""
evaluation/source/training/ablation.py ADDED
The diff for this file is too large to render. See raw diff
 
evaluation/source/training/ablation_amd.py ADDED
The diff for this file is too large to render. See raw diff
 
evaluation/source/training/final.py ADDED
@@ -0,0 +1,686 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Training entry point for the production CPU-1 model (CPU1ModelV2).
3
+
4
+ Two modes:
5
+ 1. Phase 1 — distillation training of CPU1ModelV2 (Qwen teacher → CPU-1).
6
+ python -m training.final --dry-run --total-steps 10
7
+ python -m training.final --data-path Cukinator/cpu1-dataset
8
+
9
+ 2. Phase 2 — DeleteGate training of CPU1ModelV2WithGates from a Phase-1 ckpt.
10
+ python -m training.final --phase2 --phase1-ckpt checkpoint_final.pt
11
+
12
+ Inference is implemented in `inference/` (C/C++ kernels).
13
+ The full evaluation suite lives in `evaluation/`.
14
+ """
15
+
16
+ import os
17
+ import math
18
+ import time
19
+ import argparse
20
+ import contextlib
21
+ import inspect
22
+ from pathlib import Path
23
+
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from torch import Tensor
28
+ from torch.utils.data import DataLoader
29
+
30
+ from cpu1.components import (
31
+ # Constants
32
+ D_MODEL, N_LAYERS, D_FF, FP_RANK, PFNET_HIDDEN, VOCAB_SIZE, PATCH_SIZE,
33
+ # Primitives / blocks
34
+ RMSNorm, BitLinear, BitEmbedding, FPResidual, PFNet,
35
+ BytePatchEmbeddingV2, BitHead, LocalByteDecoder, CPU1BlockV2,
36
+ BoundaryPredictor, DeleteGate, EmbeddingAligner,
37
+ # Training infrastructure
38
+ DistillationDataset, FineWebStreamingDataset, distillation_loss,
39
+ get_lr, build_param_groups, save_checkpoint, load_checkpoint,
40
+ )
41
+
42
+ _FINEWEB_REPO = 'HuggingFaceFW/fineweb'
43
+ _FINEWEB_CONFIG = {
44
+ 'tokenization': 'byte',
45
+ 'seq_len': 5000,
46
+ 'patch_size': PATCH_SIZE,
47
+ 'seed': 42,
48
+ 'training_signal': 'ce_only',
49
+ }
50
+
51
+
52
+ # =============================================================================
53
+ # 1. CPU1ModelV2 — production architecture
54
+ # =============================================================================
55
+
56
+ class CPU1ModelV2(nn.Module):
57
+ """CPU-1 v2: ternary weights, MatMul-free token mixing, MegaByte byte head.
58
+
59
+ Inference state per-layer:
60
+ - MLGRU h: d_model × 2 bytes (28 KB total at 14 layers × d=1024) → L2
61
+ - PFNet weights: 14 × 16 KB = 224 KB → L2
62
+ """
63
+
64
+ def __init__(self, config: dict | None = None):
65
+ super().__init__()
66
+ config = config or {}
67
+ d_model = config.get('d_model', D_MODEL)
68
+ n_layers = config.get('n_layers', N_LAYERS)
69
+ d_ff = config.get('d_ff', D_FF)
70
+ fp_rank = config.get('fp_rank', FP_RANK)
71
+ pfnet_hidden = config.get('pfnet_hidden', PFNET_HIDDEN)
72
+
73
+ self.d_model = d_model
74
+ self.n_layers = n_layers
75
+
76
+ self.patch_embedding = BytePatchEmbeddingV2(d_model)
77
+ self.blocks = nn.ModuleList([
78
+ CPU1BlockV2(d_model, d_ff, fp_rank, pfnet_hidden) for _ in range(n_layers)
79
+ ])
80
+ self.norm = RMSNorm(d_model)
81
+
82
+ # Default head: MegaByte LocalByteDecoder (autoregressive intra-patch
83
+ # chain). Pass legacy_heads=True to load older 4-independent-BitHead
84
+ # checkpoints (keys byte_heads.{0..3}.*).
85
+ if config.get('legacy_heads', False):
86
+ self.byte_heads = nn.ModuleList([
87
+ BitHead(d_model, VOCAB_SIZE) for _ in range(4)
88
+ ])
89
+ self.local_decoder = None
90
+ else:
91
+ self.local_decoder = LocalByteDecoder(
92
+ d_model=d_model,
93
+ patch_size=PATCH_SIZE,
94
+ d_local=config.get('local_decoder_dim', 128),
95
+ n_heads=config.get('local_decoder_heads', 4),
96
+ ff_ratio=config.get('local_decoder_ff_ratio', 4.0),
97
+ )
98
+ self.byte_heads = None
99
+
100
+ self.use_gradient_checkpointing = config.get(
101
+ 'use_gradient_checkpointing', False
102
+ )
103
+
104
+ def _forward_global(
105
+ self,
106
+ patches: Tensor,
107
+ hidden_states: list | None = None,
108
+ ) -> tuple[Tensor, list]:
109
+ """Embedding → all MLGRU blocks → final norm. Returns (post-norm, states)."""
110
+ x = self.patch_embedding(patches)
111
+
112
+ if hidden_states is None:
113
+ hidden_states = [None] * len(self.blocks)
114
+
115
+ new_states: list = []
116
+
117
+ if self.use_gradient_checkpointing and self.training:
118
+ from torch.utils.checkpoint import checkpoint
119
+ for i, block in enumerate(self.blocks):
120
+ def _block_fn(x_in, h_in, _block=block):
121
+ return _block(x_in, h_in)
122
+ x, h_new = checkpoint(_block_fn, x, hidden_states[i],
123
+ use_reentrant=False)
124
+ new_states.append(h_new)
125
+ else:
126
+ for i, block in enumerate(self.blocks):
127
+ # Touch a layer-i+2 weight pointer to hint at memory prefetch.
128
+ if i + 2 < len(self.blocks):
129
+ _ = self.blocks[i + 2].token_mixer.forget_proj.weight.data_ptr()
130
+ x, h_new = block(x, hidden_states[i])
131
+ new_states.append(h_new)
132
+
133
+ return self.norm(x), new_states
134
+
135
+ def forward(
136
+ self,
137
+ patches: Tensor,
138
+ hidden_states: list | None = None,
139
+ byte_targets: Tensor | None = None,
140
+ ) -> tuple[tuple, list]:
141
+ """Returns (logits_tuple, new_hidden_states).
142
+
143
+ byte_targets is required when using LocalByteDecoder (teacher forcing).
144
+ Falls back to zeros for benchmarks / dry runs.
145
+ """
146
+ x, new_states = self._forward_global(patches, hidden_states)
147
+
148
+ if self.local_decoder is not None:
149
+ if byte_targets is None:
150
+ B, S = x.shape[0], x.shape[1]
151
+ byte_targets = torch.zeros(
152
+ B, S, PATCH_SIZE, dtype=torch.long, device=x.device
153
+ )
154
+ logits = self.local_decoder(x, byte_targets)
155
+ else:
156
+ logits = tuple(head(x) for head in self.byte_heads)
157
+
158
+ return logits, new_states
159
+
160
+
161
+ # =============================================================================
162
+ # 2. CPU1ModelV2WithGates — Phase-2 architecture (DeleteGate + optional aux)
163
+ # =============================================================================
164
+
165
+ class CPU1ModelV2WithGates(nn.Module):
166
+ """CPU-1 v2 + BoundaryPredictor / DeleteGate / EmbeddingAligner gates.
167
+
168
+ Identical to CPU1ModelV2 when all use_* flags are False — i.e. it's a
169
+ superset that adds the auxiliary modules used in Phase 2 training.
170
+ """
171
+
172
+ def __init__(self, config: dict):
173
+ super().__init__()
174
+ d_model = config.get('d_model', D_MODEL)
175
+ n_layers = config.get('n_layers', N_LAYERS)
176
+ d_ff = config.get('d_ff', D_FF)
177
+ fp_rank = config.get('fp_rank', FP_RANK)
178
+ pfnet_hidden = config.get('pfnet_hidden', PFNET_HIDDEN)
179
+
180
+ self.use_boundary = config.get('use_boundary_predictor', False)
181
+ self.use_delete_gate = config.get('use_delete_gate', False)
182
+ self.use_emb_align = config.get('use_embedding_alignment', False)
183
+ self.delete_gate_pos = n_layers // 2
184
+
185
+ self.patch_embedding = BytePatchEmbeddingV2(d_model)
186
+ self.blocks = nn.ModuleList([
187
+ CPU1BlockV2(d_model, d_ff, fp_rank, pfnet_hidden)
188
+ for _ in range(n_layers)
189
+ ])
190
+ self.norm = RMSNorm(d_model)
191
+
192
+ if config.get('legacy_heads', False):
193
+ self.byte_heads = nn.ModuleList([
194
+ BitHead(d_model, VOCAB_SIZE) for _ in range(4)
195
+ ])
196
+ self.local_decoder = None
197
+ else:
198
+ self.local_decoder = LocalByteDecoder(
199
+ d_model=d_model,
200
+ patch_size=PATCH_SIZE,
201
+ d_local=config.get('local_decoder_dim', 128),
202
+ n_heads=config.get('local_decoder_heads', 4),
203
+ ff_ratio=config.get('local_decoder_ff_ratio', 4.0),
204
+ )
205
+ self.byte_heads = None
206
+
207
+ if self.use_boundary:
208
+ self.boundary_predictor = BoundaryPredictor(embed_dim=8, hidden_dim=32)
209
+
210
+ if self.use_delete_gate:
211
+ self.delete_gate = DeleteGate(d_model=d_model, elimination_rate=0.0)
212
+
213
+ if self.use_emb_align:
214
+ self.embedding_aligner = EmbeddingAligner(
215
+ student_dim=d_model,
216
+ teacher_dim=config.get('teacher_dim', 3584),
217
+ )
218
+
219
+ self.use_gradient_checkpointing = config.get(
220
+ 'use_gradient_checkpointing', False
221
+ )
222
+
223
+ def forward(
224
+ self,
225
+ patches: Tensor,
226
+ hidden_states: list | None = None,
227
+ training_phase: int = 1,
228
+ target_elimination: float = 0.0,
229
+ teacher_hidden: Tensor | None = None,
230
+ teacher_mask: Tensor | None = None,
231
+ byte_targets: Tensor | None = None,
232
+ **kwargs,
233
+ ) -> tuple:
234
+ """Returns (logits_tuple, new_hidden_states, extra_losses_dict)."""
235
+ x = self.patch_embedding(patches)
236
+
237
+ if hidden_states is None:
238
+ hidden_states = [None] * len(self.blocks)
239
+
240
+ new_states: list = []
241
+ extra_losses: dict = {}
242
+
243
+ if self.use_gradient_checkpointing and self.training:
244
+ from torch.utils.checkpoint import checkpoint
245
+ for i, block in enumerate(self.blocks):
246
+ def _block_fn(x_in, h_in, _block=block):
247
+ return _block(x_in, h_in)
248
+ x, h_new = checkpoint(_block_fn, x, hidden_states[i],
249
+ use_reentrant=False)
250
+ new_states.append(h_new)
251
+
252
+ if (self.use_delete_gate
253
+ and i == self.delete_gate_pos
254
+ and training_phase == 2):
255
+ x, _ = self.delete_gate.forward_train(
256
+ x, target_rate=target_elimination
257
+ )
258
+ extra_losses['delete_gate'] = self.delete_gate.last_elimination_loss
259
+ else:
260
+ for i, block in enumerate(self.blocks):
261
+ if i + 2 < len(self.blocks):
262
+ _ = self.blocks[i + 2].token_mixer.forget_proj.weight.data_ptr()
263
+ x, h_new = block(x, hidden_states[i])
264
+ new_states.append(h_new)
265
+
266
+ if (self.use_delete_gate
267
+ and i == self.delete_gate_pos
268
+ and training_phase == 2):
269
+ x, _ = self.delete_gate.forward_train(
270
+ x, target_rate=target_elimination
271
+ )
272
+ extra_losses['delete_gate'] = self.delete_gate.last_elimination_loss
273
+
274
+ x_pre_norm = x
275
+ x = self.norm(x)
276
+
277
+ if self.local_decoder is not None:
278
+ if byte_targets is None:
279
+ B, S = x.shape[0], x.shape[1]
280
+ byte_targets = torch.zeros(
281
+ B, S, PATCH_SIZE, dtype=torch.long, device=x.device
282
+ )
283
+ logits = self.local_decoder(x, byte_targets)
284
+ else:
285
+ logits = tuple(head(x) for head in self.byte_heads)
286
+
287
+ if (self.use_emb_align and training_phase == 1
288
+ and teacher_hidden is not None and teacher_mask is not None):
289
+ extra_losses['emb_align'] = self.embedding_aligner.alignment_loss(
290
+ x_pre_norm, teacher_hidden, teacher_mask,
291
+ )
292
+
293
+ return logits, new_states, extra_losses
294
+
295
+
296
+ # =============================================================================
297
+ # 3. Phase 1 training — distillation against a first-byte teacher
298
+ # =============================================================================
299
+
300
+ def train(args: argparse.Namespace):
301
+ """Phase 1 training loop with multi-GPU DataParallel."""
302
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
303
+ n_gpus = torch.cuda.device_count() if device.type == 'cuda' else 0
304
+
305
+ if device.type == 'cuda':
306
+ torch.backends.cudnn.benchmark = True
307
+ torch.set_float32_matmul_precision('high')
308
+ torch.backends.cuda.matmul.allow_tf32 = True
309
+ torch.backends.cudnn.allow_tf32 = True
310
+
311
+ print(f"Device: {device} ({n_gpus} GPU{'s' if n_gpus != 1 else ''})")
312
+
313
+ model = CPU1ModelV2({
314
+ 'use_gradient_checkpointing': args.gradient_checkpointing,
315
+ }).to(device)
316
+
317
+ n_params = sum(p.numel() for p in model.parameters())
318
+ print(f"Parámetros: {n_params:,}")
319
+
320
+ if n_gpus > 1:
321
+ model = torch.nn.DataParallel(model)
322
+ print(f"[Multi-GPU] DataParallel on {n_gpus} GPUs")
323
+
324
+ param_groups = build_param_groups(model)
325
+ optimizer = torch.optim.AdamW(
326
+ param_groups,
327
+ lr=args.max_lr,
328
+ betas=(0.9, 0.95),
329
+ eps=1e-8,
330
+ )
331
+
332
+ use_amp = device.type == 'cuda'
333
+ scaler = torch.amp.GradScaler(device.type, enabled=use_amp)
334
+
335
+ if args.dry_run:
336
+ dataset = DistillationDataset(
337
+ dry_run_samples=args.dry_run_samples,
338
+ seq_len=args.seq_len,
339
+ )
340
+ elif args.data_path is None or args.data_path == _FINEWEB_REPO:
341
+ fw_cfg = {**_FINEWEB_CONFIG, 'seq_len': args.seq_len}
342
+ print(f"[fineweb] Streaming HuggingFaceFW/fineweb sample-10BT (seq_len={args.seq_len})")
343
+ dataset = FineWebStreamingDataset(fw_cfg, split='train')
344
+ else:
345
+ dataset = DistillationDataset(
346
+ data_path=args.data_path,
347
+ seq_len=args.seq_len,
348
+ )
349
+
350
+ use_workers = args.num_workers > 0 and not args.dry_run
351
+ total_batch = args.batch_per_gpu * max(n_gpus, 1)
352
+ dataloader = DataLoader(
353
+ dataset,
354
+ batch_size=total_batch,
355
+ shuffle=True,
356
+ num_workers=args.num_workers if not args.dry_run else 0,
357
+ pin_memory=(device.type == 'cuda'),
358
+ drop_last=False,
359
+ persistent_workers=use_workers,
360
+ prefetch_factor=2 if use_workers else None,
361
+ )
362
+
363
+ start_step = 0
364
+ if args.resume_from:
365
+ model_for_load = model.module if hasattr(model, 'module') else model
366
+ start_step = load_checkpoint(args.resume_from, model_for_load, optimizer, scaler)
367
+ print(f"Resumiendo desde step {start_step}")
368
+
369
+ grad_accum = args.gradient_accumulation_steps
370
+ total_steps = args.total_steps
371
+ effective_batch = total_batch * grad_accum
372
+
373
+ print(f"Batch efectivo: {effective_batch} "
374
+ f"({args.batch_per_gpu}/gpu × {max(n_gpus, 1)} GPUs × {grad_accum} accum)")
375
+ print(f"Steps: {start_step} → {total_steps}")
376
+ print(f"Seq len: {args.seq_len} patches = {args.seq_len * PATCH_SIZE} bytes")
377
+
378
+ model.train()
379
+ data_iter = iter(dataloader)
380
+ epoch = 0
381
+ log_timer = time.time()
382
+ metrics: dict[str, float] = {}
383
+ last_loss = 0.0
384
+
385
+ for step in range(start_step, total_steps):
386
+ lr = get_lr(step, args.warmup_steps, total_steps, args.max_lr, args.min_lr)
387
+ for pg in optimizer.param_groups:
388
+ pg['lr'] = lr * pg.get('lr_scale', 1.0)
389
+
390
+ optimizer.zero_grad(set_to_none=True)
391
+ should_log = (step % args.log_every == 0)
392
+
393
+ for micro_step in range(grad_accum):
394
+ try:
395
+ batch = next(data_iter)
396
+ except StopIteration:
397
+ epoch += 1
398
+ data_iter = iter(dataloader)
399
+ batch = next(data_iter)
400
+
401
+ patches = batch['patches'].to(device, non_blocking=True)
402
+ targets = batch['targets'].to(device, non_blocking=True)
403
+ B, S = patches.shape[0], patches.shape[1]
404
+ _tp = batch.get('teacher_probs')
405
+ _tm = batch.get('teacher_mask')
406
+ teacher_probs = (
407
+ _tp.to(device, non_blocking=True) if _tp is not None
408
+ else torch.zeros(B, S, VOCAB_SIZE, device=device)
409
+ )
410
+ teacher_mask = (
411
+ _tm.to(device, non_blocking=True) if _tm is not None
412
+ else torch.zeros(B, S, dtype=torch.bool, device=device)
413
+ )
414
+ teacher_hidden = batch.get('teacher_hidden')
415
+ if teacher_hidden is not None:
416
+ teacher_hidden = teacher_hidden.to(device, non_blocking=True)
417
+
418
+ is_last_micro = (micro_step == grad_accum - 1)
419
+ amp_ctx = (torch.amp.autocast('cuda', dtype=torch.float16)
420
+ if use_amp else contextlib.nullcontext())
421
+ with amp_ctx:
422
+ # Detect WithGates models that take teacher_hidden in their forward
423
+ fwd = model.module.forward if hasattr(model, 'module') else model.forward
424
+ sig_params = inspect.signature(fwd).parameters
425
+ if 'teacher_hidden' in sig_params:
426
+ logits, _, _ = model(
427
+ patches, byte_targets=targets,
428
+ teacher_hidden=teacher_hidden, teacher_mask=teacher_mask,
429
+ )
430
+ else:
431
+ logits, _ = model(patches, byte_targets=targets)
432
+ loss, metrics = distillation_loss(
433
+ logits, targets, teacher_probs, teacher_mask,
434
+ temperature=args.temperature,
435
+ compute_metrics=(should_log and is_last_micro),
436
+ )
437
+ if is_last_micro:
438
+ last_loss = loss.item()
439
+ loss = loss / grad_accum
440
+
441
+ scaler.scale(loss).backward()
442
+
443
+ scaler.unscale_(optimizer)
444
+ grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
445
+
446
+ scaler.step(optimizer)
447
+ scaler.update()
448
+
449
+ if should_log:
450
+ elapsed = time.time() - log_timer
451
+ steps_since_log = max(args.log_every, 1)
452
+ patches_per_sec = (
453
+ total_batch * args.seq_len * grad_accum * steps_since_log
454
+ ) / max(elapsed, 1e-6)
455
+
456
+ parts = [
457
+ f"Step {step:6d}/{total_steps}",
458
+ f"Loss: {metrics.get('loss_total', 0):.4f}",
459
+ f"H0-KL: {metrics.get('loss_head0_kl', 0):.4f}",
460
+ f"H0-CE: {metrics.get('loss_head0_ce', 0):.4f}",
461
+ f"H123: {metrics.get('loss_heads_123_ce', 0):.4f}",
462
+ f"Tcov: {metrics.get('teacher_coverage', 0):.1%}",
463
+ f"‖g‖: {grad_norm:.2f}",
464
+ f"LR: {lr:.2e}",
465
+ f"P/s: {patches_per_sec:.0f}",
466
+ ]
467
+ if use_amp:
468
+ gpu_gb = torch.cuda.max_memory_allocated(device) / (1024 ** 3)
469
+ parts.append(f"VRAM: {gpu_gb:.1f}G")
470
+
471
+ print(" | ".join(parts))
472
+ log_timer = time.time()
473
+
474
+ if step > 0 and step % args.save_every == 0:
475
+ ckpt_path = str(Path(args.output_dir) / f"checkpoint_{step}.pt")
476
+ save_checkpoint(
477
+ model, optimizer, scaler, step,
478
+ last_loss,
479
+ ckpt_path,
480
+ upload_to_hf=args.upload_hf,
481
+ )
482
+ print(f" Checkpoint: {ckpt_path}")
483
+
484
+ ckpt_path = str(Path(args.output_dir) / "checkpoint_final.pt")
485
+ save_checkpoint(
486
+ model, optimizer, scaler, total_steps,
487
+ last_loss,
488
+ ckpt_path,
489
+ upload_to_hf=args.upload_hf,
490
+ )
491
+ print(f"Training completo. Final: {ckpt_path}")
492
+
493
+
494
+ # =============================================================================
495
+ # 4. Phase 2 training — DeleteGate elimination on top of a Phase-1 checkpoint
496
+ # =============================================================================
497
+
498
+ def train_phase2(
499
+ model,
500
+ dataset,
501
+ config: dict,
502
+ checkpoint_phase1: str,
503
+ target_elimination: float = 0.4,
504
+ total_steps: int = 5000,
505
+ lr: float = 1e-4,
506
+ device: torch.device | None = None,
507
+ ) -> dict:
508
+ """Phase 2: load Phase-1 weights, activate DeleteGate, freeze EmbAligner,
509
+ and train with CE + elimination regularization (no KL — sequence shortens).
510
+ """
511
+ if device is None:
512
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
513
+
514
+ ckpt = torch.load(checkpoint_phase1, map_location='cpu', weights_only=False)
515
+ model.load_state_dict(ckpt['model_state_dict'], strict=False)
516
+ model = model.to(device)
517
+
518
+ if hasattr(model, 'delete_gate'):
519
+ model.delete_gate.set_elimination_rate(target_elimination)
520
+
521
+ if hasattr(model, 'embedding_aligner'):
522
+ for p in model.embedding_aligner.parameters():
523
+ p.requires_grad = False
524
+
525
+ n_gpus = torch.cuda.device_count() if device.type == 'cuda' else 0
526
+ if n_gpus > 1:
527
+ model = torch.nn.DataParallel(model)
528
+ print(f"[Phase2 Multi-GPU] DataParallel on {n_gpus} GPUs")
529
+
530
+ def phase2_loss(logits, targets, extra_losses):
531
+ ce = sum(
532
+ F.cross_entropy(
533
+ logits[i].reshape(-1, 256),
534
+ targets[:, :, i].long().reshape(-1)
535
+ ) for i in range(4)
536
+ ) / 4
537
+
538
+ elim_loss = extra_losses.get(
539
+ 'delete_gate', torch.tensor(0.0, device=device)
540
+ )
541
+ if elim_loss.dim() > 0:
542
+ elim_loss = elim_loss.mean()
543
+
544
+ return ce + 0.1 * elim_loss
545
+
546
+ model_params = model.module if hasattr(model, 'module') else model
547
+ optimizer = torch.optim.AdamW(
548
+ [p for p in model_params.parameters() if p.requires_grad],
549
+ lr=lr, betas=(0.9, 0.95),
550
+ )
551
+
552
+ use_amp = device.type == 'cuda'
553
+ scaler = torch.amp.GradScaler() if use_amp else None
554
+
555
+ total_batch = 4 * max(n_gpus, 1)
556
+ dataloader = torch.utils.data.DataLoader(
557
+ dataset, batch_size=total_batch, shuffle=True, drop_last=True,
558
+ )
559
+
560
+ model.train()
561
+ step = 0
562
+ t_start = time.time()
563
+
564
+ while step < total_steps:
565
+ for batch in dataloader:
566
+ if step >= total_steps:
567
+ break
568
+
569
+ patches = batch['patches'].to(device, non_blocking=True)
570
+ targets = batch['targets'].to(device, non_blocking=True)
571
+
572
+ amp_ctx = (torch.amp.autocast(device_type='cuda', dtype=torch.float16)
573
+ if use_amp else contextlib.nullcontext())
574
+ with amp_ctx:
575
+ logits, _, extra_losses = model(
576
+ patches, training_phase=2,
577
+ target_elimination=target_elimination,
578
+ )
579
+ loss = phase2_loss(logits, targets, extra_losses)
580
+
581
+ optimizer.zero_grad(set_to_none=True)
582
+ if use_amp:
583
+ scaler.scale(loss).backward()
584
+ scaler.unscale_(optimizer)
585
+ torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
586
+ scaler.step(optimizer)
587
+ scaler.update()
588
+ else:
589
+ loss.backward()
590
+ torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
591
+ optimizer.step()
592
+
593
+ step += 1
594
+ if step % 500 == 0:
595
+ elapsed = time.time() - t_start
596
+ print(f"[Phase2] Step {step}/{total_steps} | "
597
+ f"Loss: {loss.item():.4f} | Time: {elapsed:.0f}s")
598
+
599
+ return {'phase': 2, 'target_elimination': target_elimination, 'steps': step}
600
+
601
+
602
+ # =============================================================================
603
+ # 5. CLI
604
+ # =============================================================================
605
+
606
+ def parse_args() -> argparse.Namespace:
607
+ p = argparse.ArgumentParser(description="Training CPU-1 con distillation loss")
608
+
609
+ p.add_argument('--data-path', type=str, default=None,
610
+ help="Ruta a HF dataset, directorio, o .parquet. "
611
+ "Por defecto usa HuggingFaceFW/fineweb (CE-only, sin teacher).")
612
+ p.add_argument('--dry-run', action='store_true',
613
+ help="Usar datos aleatorios (testing)")
614
+ p.add_argument('--dry-run-samples', type=int, default=1000)
615
+ p.add_argument('--seq-len', type=int, default=5000,
616
+ help="Patches por secuencia")
617
+
618
+ p.add_argument('--batch-per-gpu', type=int, default=4)
619
+ p.add_argument('--gradient-accumulation-steps', type=int, default=2)
620
+ p.add_argument('--total-steps', type=int, default=100_000)
621
+ p.add_argument('--warmup-steps', type=int, default=2000)
622
+ p.add_argument('--max-lr', type=float, default=3e-4)
623
+ p.add_argument('--min-lr', type=float, default=3e-5)
624
+ p.add_argument('--temperature', type=float, default=4.0)
625
+ p.add_argument('--gradient-checkpointing', action='store_true')
626
+
627
+ p.add_argument('--output-dir', type=str, default='./checkpoints')
628
+ p.add_argument('--save-every', type=int, default=500)
629
+ p.add_argument('--log-every', type=int, default=100)
630
+ p.add_argument('--upload-hf', action='store_true',
631
+ help="Subir checkpoints a Cukinator/cpu1-checkpoints")
632
+
633
+ p.add_argument('--resume-from', type=str, default=None,
634
+ help="Ruta a checkpoint .pt para resumir")
635
+ p.add_argument('--num-workers', type=int, default=4)
636
+
637
+ # Phase 2 controls
638
+ p.add_argument('--phase2', action='store_true',
639
+ help="Run DeleteGate Phase 2 starting from --phase1-ckpt")
640
+ p.add_argument('--phase1-ckpt', type=str, default=None,
641
+ help="Phase 1 checkpoint to start Phase 2 from")
642
+ p.add_argument('--target-elimination', type=float, default=0.4,
643
+ help="Phase 2 fraction of states to eliminate")
644
+
645
+ return p.parse_args()
646
+
647
+
648
+ if __name__ == "__main__":
649
+ args = parse_args()
650
+
651
+ if args.phase2:
652
+ if not args.phase1_ckpt:
653
+ raise SystemExit("--phase2 requires --phase1-ckpt")
654
+
655
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
656
+ cfg = {
657
+ 'use_delete_gate': True,
658
+ 'use_gradient_checkpointing': args.gradient_checkpointing,
659
+ }
660
+ model = CPU1ModelV2WithGates(cfg)
661
+
662
+ if args.dry_run:
663
+ dataset = DistillationDataset(
664
+ dry_run_samples=args.dry_run_samples,
665
+ seq_len=args.seq_len,
666
+ )
667
+ else:
668
+ dataset = DistillationDataset(
669
+ data_path=args.data_path,
670
+ seq_len=args.seq_len,
671
+ )
672
+
673
+ result = train_phase2(
674
+ model=model,
675
+ dataset=dataset,
676
+ config=cfg,
677
+ checkpoint_phase1=args.phase1_ckpt,
678
+ target_elimination=args.target_elimination,
679
+ total_steps=args.total_steps,
680
+ lr=args.max_lr,
681
+ device=device,
682
+ )
683
+ print(f"[Phase2] Done: {result}")
684
+ else:
685
+ os.makedirs(args.output_dir, exist_ok=True)
686
+ train(args)