Rebuild model card from reproducible CPU audit; include evidence and limitations
Browse files- README.md +10 -10
- evaluation/publication.json +2 -2
- evaluation/source/cpu1/__init__.py +1 -0
- evaluation/source/cpu1/components.py +0 -0
- evaluation/source/evaluation/__init__.py +1 -0
- evaluation/source/evaluation/audit.py +299 -0
- evaluation/source/evaluation/cards.py +299 -0
- evaluation/source/evaluation/checkpoints.py +115 -0
- evaluation/source/evaluation/correctness.py +239 -0
- evaluation/source/evaluation/diagnostics.py +35 -0
- evaluation/source/evaluation/export.py +252 -0
- evaluation/source/evaluation/helpers.py +26 -0
- evaluation/source/evaluation/verify_checkpoints.py +64 -0
- evaluation/source/requirements/audit.txt +2 -0
- evaluation/source/requirements/data.txt +4 -0
- evaluation/source/requirements/training.txt +7 -0
- evaluation/source/tests/__init__.py +1 -0
- evaluation/source/tests/test_audit.py +66 -0
- evaluation/source/tests/test_checkpoints.py +62 -0
- evaluation/source/training/__init__.py +1 -0
- evaluation/source/training/ablation.py +0 -0
- evaluation/source/training/ablation_amd.py +0 -0
- evaluation/source/training/final.py +686 -0
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: [
|
| 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
|
| 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 `
|
| 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
|
| 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 |
-
-
|
| 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 `
|
| 145 |
|
| 146 |
```powershell
|
| 147 |
-
python
|
| 148 |
-
python
|
| 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 `
|
| 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 `
|
| 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": "
|
| 4 |
-
"url": "https://huggingface.co/Cukinator/cpu1-ablations-final/commit/
|
| 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)
|