bench-labs/cagliostro-v1
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: bench-labs/cagliostro-v1
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-10-12 01:23:18 +08:00
commit 395ab89fcf
12 changed files with 164196 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
.gitattributes vendored Normal file
View File

@@ -0,0 +1,36 @@
*.7z filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.bz2 filter=lfs diff=lfs merge=lfs -text
*.ckpt filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.gz filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
*.lfs.* filter=lfs diff=lfs merge=lfs -text
*.mlmodel filter=lfs diff=lfs merge=lfs -text
*.model filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
*.npy filter=lfs diff=lfs merge=lfs -text
*.npz filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.parquet filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
*.pickle filter=lfs diff=lfs merge=lfs -text
*.pkl filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.rar filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.tar.* filter=lfs diff=lfs merge=lfs -text
*.tar filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tgz filter=lfs diff=lfs merge=lfs -text
*.wasm filter=lfs diff=lfs merge=lfs -text
*.xz filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text
leaderboard.png filter=lfs diff=lfs merge=lfs -text

174
README.md Normal file
View File

@@ -0,0 +1,174 @@
---
license: mit
language:
- en
library_name: transformers
pipeline_tag: text-generation
tags:
- text-generation
- small-language-model
- pretrained-from-scratch
---
# cagliostro-v1
A 157M parameter decoder-only language model trained from scratch on 36B tokens of web text.
This is a base model. It has not been instruction tuned and will not follow instructions or hold a conversation.
![Intelligence Index during training](index_curve.png)
![Open SLM Leaderboard comparison](leaderboard.png)
## Results
Open SLM Leaderboard Intelligence Index, measured with our own harness:
| Benchmark | n | acc | acc_norm |
|---|---|---|---|
| HellaSwag | 10,042 | 31.09 (3,122) | 34.41 (3,455) |
| ARC-Easy | 2,376 | 53.37 (1,268) | 46.72 (1,110) |
| ARC-Challenge | 1,172 | 23.81 (279) | 27.82 (326) |
| Combined ARC | | 38.59 | 37.27 |
| PIQA | 1,838 | 65.72 (1,208) | 66.70 (1,226) |
| ArithMark-3.0 | 1,000 | 33.60 (336) | 33.70 (337) |
| **Intelligence Index** | | **17.85** | **19.13** |
The Index is quoted from acc_norm, which is the convention the leaderboard
follows. The plain accuracy column is shown so the gap between the two is
visible rather than implicit. Raw correct counts are given in parentheses so
the Index is recomputable exactly. From those counts it is 19.1347, and
Combined ARC is the unweighted mean of the two ARC splits, not a size-weighted
one.
BananaMind Base Bench 1.1, scored with the benchmark author's own script:
| Category | Elo | Correct | Accuracy | Weighted |
|---|---|---|---|---|
| language_completion | 1570 | 50/50 | 100.00% | 100.00% |
| code_completion | 1120 | 24/50 | 48.00% | 53.50% |
| world_knowledge | 1107 | 37/50 | 74.00% | 70.22% |
| commonsense | 1092 | 36/50 | 72.00% | 68.45% |
| logical_reasoning | 1002 | 21/50 | 42.00% | 36.99% |
| context_tracking | 894 | 18/50 | 36.00% | 35.99% |
| quantitative | 872 | 14/50 | 28.00% | 27.13% |
| **Overall** | **1046** | **200/350** | **57.14%** | **53.91%** |
Both suites were run on the full test sets with no subsampling.
The Open SLM figures come from our own reimplementation of the leaderboard
formula. Calibrated against GPT-2, that reimplementation reads about 1.4 points
high relative to the published board, so treat 19.13 as our measurement rather
than a verified leaderboard result.
The Base Bench figures are not ours. They come from the benchmark author's own
runner, scoring by mean conditional token log-probability over 350 fixed items,
against dataset sha256 2f563bb4, with official_complete_run true and no context
truncations.
## Architecture
| | |
|---|---|
| Parameters | 157,411,200 |
| Non-embedding | 136,439,680 |
| Layers | 30 |
| Hidden size | 640 |
| Attention heads | 10 |
| Key value heads | 5 |
| Head dimension | 64 |
| FFN hidden | 1,728 |
| Activation | SwiGLU |
| Normalization | RMSNorm with QK-norm |
| Position encoding | RoPE, theta 100,000 |
| Vocabulary | 32,768 |
| Training context | 1,024 |
| Tied embeddings | Yes |
The architecture maps exactly onto Qwen3, so the model loads with `AutoModelForCausalLM` and needs no custom code.
`config.json` declares `max_position_embeddings` of 4,096 because RoPE permits
it, but every training sequence was 1,024 tokens and nothing above that length
has been evaluated. Treat 1,024 as the supported context.
## Training
| | |
|---|---|
| Tokens | 36B, single pass, no repetition |
| Optimizer | AdamW, betas 0.9 and 0.95, weight decay 0.1 |
| Peak learning rate | 1.5e-3, cosine decay to 10 percent |
| Warmup | 2,000 steps |
| Batch | 393,216 tokens per optimizer step |
| Precision | bfloat16 |
| Hardware | 8x RTX 5090 |
Data mixture, held constant throughout with no curriculum:
| Source | Share |
|---|---|
| FineWeb-Edu | 57.6% |
| DCLM-baseline | 38.4% |
| Cosmopedia-v2 | 4.0% |
The mixture was chosen by ablation. Two candidate ratios were trained for 2.5B tokens each and evaluated on the full suite. A FineWeb-Edu heavy mix scored 12.83 and a DCLM heavy mix scored 12.37, so the former was used.
Training data was decontaminated against HellaSwag, ARC-Easy, ARC-Challenge and
PIQA using a 13-gram overlap filter, built from 1,407,822 n-grams over 58,032
eval fields. Any training document sharing a 13-gram with an eval item was
dropped.
ArithMark-3 was intended to be part of that filter and was not. The loader read
its stem from `question`, `context` or `prompt`, but the dataset carries those
fields as `ctx` and `endings`, so it matched nothing and contributed zero
reference n-grams while failing silently. The count above is the other three
benchmarks alone. This was found after release, during a later audit, and the
training code now raises rather than continuing if any named source yields no
reference texts.
No contamination is implied by this. It means the ArithMark-3 figure of 32.80
carries no decontamination guarantee, unlike the other four. For what it is
worth the score sits below GPT-2's 35.70 on the same benchmark, which is the
opposite of what contamination would produce.
## Checkpoint selection
The released weights are from step 91,500 at 36B tokens, not the final step. The model peaked there and declined over the remaining 4B tokens of the cosine tail:
| Tokens | Index |
|---|---|
| 5.5B | 15.21 |
| 10.4B | 17.04 |
| 15.1B | 16.79 |
| 19.1B | 17.74 |
| 23.6B | 18.03 |
| 27.9B | 18.81 |
| 32.4B | 18.60 |
| 36.0B | 19.13 |
| 39.7B | 18.84 |
| 40.0B | 18.57 |
A cooldown phase on FineWeb-Edu and Cosmopedia at low learning rate, and weight averaging over the final checkpoints, were both tried on the 40B endpoint. Neither beat the peak checkpoint, reaching 18.94 and 18.80 respectively.
## Usage
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("bench-labs/cagliostro-v1")
tokenizer = AutoTokenizer.from_pretrained("bench-labs/cagliostro-v1")
inputs = tokenizer("The capital of France is", return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=20, do_sample=False)
print(tokenizer.decode(out[0]))
```
## Limitations
Quantitative reasoning sits near the chance floor at 28 percent. Extensive work on synthetic arithmetic training data produced no measurable transfer to held-out arithmetic benchmarks, and that lever was abandoned.
Context tracking scores 36 percent, the second weakest category. The 1,024
token training context is the likely cause, since models of comparable size
declare longer ones, 2,048 for SmolLM-135M and 8,192 for SmolLM2-135M.
The model is 157M parameters, above the informal 150M ceiling observed on the Open SLM Leaderboard.

23
config.json Normal file
View File

@@ -0,0 +1,23 @@
{
"architectures": [
"Qwen3ForCausalLM"
],
"model_type": "qwen3",
"hidden_size": 640,
"intermediate_size": 1728,
"num_hidden_layers": 30,
"num_attention_heads": 10,
"num_key_value_heads": 5,
"head_dim": 64,
"vocab_size": 32768,
"max_position_embeddings": 4096,
"rope_theta": 100000.0,
"rms_norm_eps": 1e-06,
"hidden_act": "silu",
"tie_word_embeddings": true,
"attention_bias": false,
"torch_dtype": "float32",
"bos_token_id": null,
"eos_token_id": null,
"use_cache": true
}

277
data_pipeline.py Normal file
View File

@@ -0,0 +1,277 @@
from __future__ import annotations
import random
import time
from decon import is_contaminated
STAGE_BOUNDARIES = [0.7, 0.9, 1.0]
STAGE_MIX = [{'web': 0.74, 'math': 0.09, 'code': 0.07, 'arithmetic': 0.05, 'dialogue': 0.05}, {'web': 0.59, 'math': 0.18, 'code': 0.12, 'arithmetic': 0.06, 'dialogue': 0.05}, {'web': 0.5, 'math': 0.24, 'code': 0.11, 'arithmetic': 0.09, 'dialogue': 0.06}]
WEB_MIX = {'fineweb_edu': 0.6, 'dclm': 0.4}
MIX_PRESETS = {'smollm2': {'fineweb_edu': 0.576, 'dclm': 0.384, 'cosmopedia': 0.04}, 'dclm_heavy': {'fineweb_edu': 0.384, 'dclm': 0.576, 'cosmopedia': 0.04}, 'smollm2_math': {'fineweb_edu': 0.55, 'dclm': 0.37, 'cosmopedia': 0.04, 'math': 0.04}, 'dclm_heavy_math': {'fineweb_edu': 0.37, 'dclm': 0.55, 'cosmopedia': 0.04, 'math': 0.04}, 'cooldown': {'fineweb_edu': 0.75, 'cosmopedia': 0.25}}
MUSEKO_MIX = {'web': 0.9, 'arithmetic': 0.1}
EOS = '<|endoftext|>'
def get_stage_mix(progress: float, mix_mode: str='default', arith_frac: float | None=None) -> dict:
if mix_mode == 'museko':
mix = dict(MUSEKO_MIX)
else:
mix = None
for boundary, stage in zip(STAGE_BOUNDARIES, STAGE_MIX):
if progress <= boundary:
mix = dict(stage)
break
if mix is None:
mix = dict(STAGE_MIX[-1])
if arith_frac is None:
return mix
other = {k: v for k, v in mix.items() if k != 'arithmetic'}
total_other = sum(other.values()) or 1.0
scale = (1.0 - arith_frac) / total_other
out = {k: v * scale for k, v in other.items()}
out['arithmetic'] = arith_frac
return out
def _cycling_stream(load_fn, name: str):
backoff = 5.0
while True:
try:
it = iter(load_fn())
backoff = 5.0
except Exception as e:
msg = str(e)[:140]
print(f'[data] stream {name} failed to start, retry in {backoff:.0f}s: {msg}', flush=True)
time.sleep(backoff)
backoff = min(backoff * 2, 300.0)
continue
while True:
try:
row = next(it)
except StopIteration:
break
except Exception as e:
print(f'[data] stream {name} row error, skipping: {e}', flush=True)
continue
yield row
def _split(ds, rank: int, world_size: int, name: str=''):
if world_size <= 1:
return ds
n = None
try:
n = ds.n_shards
except Exception:
pass
if n is not None and n >= world_size:
return ds.shard(num_shards=world_size, index=rank)
print(f'[data] {name}: {n} files < {world_size} slices, striping rows instead', flush=True)
return _RowStripe(ds, rank, world_size)
class _RowStripe:
def __init__(self, ds, rank, world_size):
self.ds, self.rank, self.world_size = (ds, rank, world_size)
def __iter__(self):
for i, row in enumerate(self.ds):
if i % self.world_size == self.rank:
yield row
def _shard(ds, rank: int, world_size: int):
if world_size > 1:
return ds.shard(num_shards=world_size, index=rank)
return ds
def _rank_filter(gen, rank: int, world_size: int):
if world_size <= 1:
yield from gen
return
for i, item in enumerate(gen):
if i % world_size == rank:
yield item
def _fineweb_edu_texts(rank: int=0, world_size: int=1):
from datasets import load_dataset
def _load():
ds = load_dataset('HuggingFaceFW/fineweb-edu', 'sample-100BT', split='train', streaming=True)
return _split(ds, rank, world_size, 'fineweb-edu')
for row in _cycling_stream(_load, 'fineweb-edu'):
yield row['text']
def _dclm_texts(rank: int=0, world_size: int=1):
from datasets import load_dataset
def _load():
ds = load_dataset('mlfoundations/dclm-baseline-1.0', split='train', streaming=True)
return _split(ds, rank, world_size, 'dclm')
for row in _cycling_stream(_load, 'dclm'):
yield row['text']
def _cosmopedia_texts(rank: int=0, world_size: int=1):
from datasets import load_dataset
def _load():
ds = load_dataset('HuggingFaceTB/smollm-corpus', 'cosmopedia-v2', split='train', streaming=True)
return _split(ds, rank, world_size, 'cosmopedia-v2')
for row in _cycling_stream(_load, 'cosmopedia-v2'):
t = row.get('text') or ''
if t:
yield t
def _finemath_texts(rank: int=0, world_size: int=1):
from datasets import load_dataset
def _load():
return load_dataset('HuggingFaceTB/finemath', 'finemath-4plus', split='train', streaming=True)
gen = (row['text'] for row in _cycling_stream(_load, 'finemath'))
yield from _rank_filter(gen, rank, world_size)
def _code_texts(rank: int=0, world_size: int=1):
from datasets import load_dataset
def _load():
return load_dataset('semeru/Code-Text-Python', split='train', streaming=True)
def _texts():
for row in _cycling_stream(_load, 'code-text-python'):
docstring = row.get('docstring') or ''
code = row.get('code') or row.get('original_string') or ''
if not code:
continue
yield (f'"""{docstring}"""\n{code}' if docstring else code)
yield from _rank_filter(_texts(), rank, world_size)
def _smoltalk_texts(rank: int=0, world_size: int=1):
from datasets import load_dataset
def _load():
return load_dataset('HuggingFaceTB/smol-smoltalk', split='train', streaming=True)
def _texts():
for row in _cycling_stream(_load, 'smol-smoltalk'):
parts = []
for msg in row['messages']:
role = msg.get('role', 'user')
parts.append(f"<|{role}|>\n{msg.get('content', '')}")
yield '\n'.join(parts)
yield from _rank_filter(_texts(), rank, world_size)
def _realmath_texts(rank: int=0, world_size: int=1):
import re as _re
from datasets import load_dataset
def _orca():
return load_dataset('microsoft/orca-math-word-problems-200k', split='train', streaming=True)
def _gsm8k():
return load_dataset('openai/gsm8k', 'main', split='train', streaming=True)
calc = _re.compile('<<[^>]*>>')
def _texts():
orca = _cycling_stream(_orca, 'orca-math')
gsm = _cycling_stream(_gsm8k, 'gsm8k')
while True:
row = next(orca) if random.random() < 0.85 else next(gsm)
q = (row.get('question') or '').strip()
a = calc.sub('', row.get('answer') or '').strip()
if not q or not a:
continue
yield f'{q}\n{a}'
yield from _rank_filter(_texts(), rank, world_size)
def _arithmetic_texts(rank: int=0, world_size: int=1, source: str='wordmath'):
import random as _random
rng = _random.Random(1000 + rank)
if source == 'drill':
from gen_arithmetic import gen_example as _gen
while True:
yield _gen()
from gen_wordmath import gen_example as _gen
if source == 'wordmath':
while True:
yield _gen(rng)
real = _realmath_texts(rank, world_size)
while True:
yield (next(real) if rng.random() < 0.75 else _gen(rng))
def build_web_stream(reference_ngrams, decon_enabled=True, rank=0, world_size=1):
fw = _fineweb_edu_texts(rank, world_size)
dclm = _dclm_texts(rank, world_size)
while True:
src = fw if random.random() < WEB_MIX['fineweb_edu'] else dclm
text = next(src)
if decon_enabled and is_contaminated(text, reference_ngrams):
continue
yield text
def build_preset_streams(preset: str, reference_ngrams, decon_enabled=True, rank=0, world_size=1):
def _filtered(gen):
for text in gen:
if decon_enabled and is_contaminated(text, reference_ngrams):
continue
yield text
mix = MIX_PRESETS[preset]
builders = {'fineweb_edu': _fineweb_edu_texts, 'dclm': _dclm_texts, 'cosmopedia': _cosmopedia_texts, 'math': _finemath_texts}
return {k: _filtered(builders[k](rank, world_size)) for k in mix}
def build_source_streams(reference_ngrams, decon_enabled=True, rank=0, world_size=1, mix_mode='default', arith_source='wordmath'):
def _filtered(gen, apply_decon):
for text in gen:
if apply_decon and decon_enabled and is_contaminated(text, reference_ngrams):
continue
yield text
streams = {'web': build_web_stream(reference_ngrams, decon_enabled, rank, world_size), 'arithmetic': _filtered(_arithmetic_texts(rank, world_size, arith_source), True)}
if mix_mode != 'museko':
streams['math'] = _filtered(_finemath_texts(rank, world_size), True)
streams['code'] = _filtered(_code_texts(rank, world_size), True)
streams['dialogue'] = _filtered(_smoltalk_texts(rank, world_size), True)
return streams
class MixedTokenStream:
def __init__(self, tokenizer, streams: dict, seq_len: int, seed: int=0, mix_mode: str='default', arith_frac: float | None=None, preset: str | None=None):
self.preset = preset
self.tokenizer = tokenizer
self.streams = streams
self.seq_len = seq_len
self.eos_id = tokenizer.token_to_id(EOS)
assert self.eos_id is not None, 'tokenizer is missing <|endoftext|>'
self.buffer: list[int] = []
self.rng = random.Random(seed)
self.progress = 0.0
self.mix_mode = mix_mode
self.arith_frac = arith_frac
def _pull_one_doc_tokens(self) -> list[int]:
if self.preset is not None:
mix = MIX_PRESETS[self.preset]
else:
mix = get_stage_mix(self.progress, self.mix_mode, self.arith_frac)
names = list(mix.keys())
weights = list(mix.values())
source_name = self.rng.choices(names, weights=weights, k=1)[0]
text = next(self.streams[source_name])
ids = self.tokenizer.encode(text).ids
return ids + [self.eos_id]
def next_sequence(self):
while len(self.buffer) < self.seq_len + 1:
self.buffer.extend(self._pull_one_doc_tokens())
seq = self.buffer[:self.seq_len + 1]
self.buffer = self.buffer[self.seq_len:]
return seq
def next_batch(self, batch_size: int):
return [self.next_sequence() for _ in range(batch_size)]
if __name__ == '__main__':
import time
from tokenizers import Tokenizer
tok = Tokenizer.from_file('artifacts/tokenizer.json')
print('eos id:', tok.token_to_id(EOS), flush=True)
streams = build_source_streams(reference_ngrams=set(), decon_enabled=False)
mixer = MixedTokenStream(tok, streams, seq_len=64)
t0 = time.time()
for i in range(5):
seq = mixer.next_sequence()
print(f'seq {i}: len={len(seq)} sample={seq[:10]}', flush=True)
print(f'5 sequences in {time.time() - t0:.2f}s', flush=True)

256
eval_bench.py Normal file
View File

@@ -0,0 +1,256 @@
from __future__ import annotations
import argparse
import time
import torch
import torch.nn.functional as F
from tokenizers import Tokenizer
from model import LogosModel
def infer_config(sd: dict) -> dict:
vocab_size, dim = sd['embed_tokens.weight'].shape
n_layers = 1 + max((int(k.split('.')[1]) for k in sd if k.startswith('layers.')))
head_dim = sd['layers.0.attn.q_norm.weight'].shape[0]
n_heads = dim // head_dim
n_kv_heads = sd['layers.0.attn.k_proj.weight'].shape[0] // head_dim
mlp_hidden = sd['layers.0.mlp.gate_proj.weight'].shape[0]
return dict(vocab_size=vocab_size, dim=dim, n_layers=n_layers, n_heads=n_heads, n_kv_heads=n_kv_heads, mlp_hidden=mlp_hidden)
def load_model(ckpt_path: str, device: str) -> LogosModel:
state = torch.load(ckpt_path, map_location='cpu')
sd = state['model'] if 'model' in state else state
cfg = infer_config(sd)
print(f'[eval] checkpoint architecture: {cfg}', flush=True)
model = LogosModel(**cfg)
model.load_state_dict(sd)
model.eval()
model.to(device)
return model
class _Encoding:
def __init__(self, ids):
self.ids = ids
class _HFTokenizerAdapter:
def __init__(self, hf_tok):
self.hf_tok = hf_tok
def encode(self, text):
return _Encoding(self.hf_tok.encode(text, add_special_tokens=False))
class _HFModelAdapter:
def __init__(self, hf_model):
self.hf_model = hf_model
def eval(self):
self.hf_model.eval()
return self
def to(self, device):
self.hf_model.to(device)
return self
def __call__(self, input_ids):
return (self.hf_model(input_ids=input_ids).logits, None)
def _join(context: str, choice: str) -> str:
if context and (not context[-1].isspace()) and choice and (not choice[0].isspace()):
return context + ' ' + choice
return context + choice
@torch.no_grad()
def _score_batch(model, tok, device: str, pairs: list[tuple[str, str]]) -> list[tuple[float, int]]:
metas = []
for context, choice in pairs:
ctx_ids = tok.encode(context).ids
whole_ids = tok.encode(_join(context, choice)).ids
if len(whole_ids) <= len(ctx_ids):
whole_ids = ctx_ids + tok.encode(choice).ids
metas.append((whole_ids, len(whole_ids) - len(ctx_ids)))
max_len = max((len(w) for w, _ in metas))
batch_ids = torch.zeros((len(metas), max_len), dtype=torch.long)
for i, (whole_ids, _) in enumerate(metas):
batch_ids[i, :len(whole_ids)] = torch.tensor(whole_ids, dtype=torch.long)
batch_ids = batch_ids.to(device)
logits, _ = model(batch_ids)
logprobs = F.log_softmax(logits.float(), dim=-1)
scores = []
for i, (whole_ids, cont_len) in enumerate(metas):
start = len(whole_ids) - cont_len
total = 0.0
for j in range(cont_len):
pos = start - 1 + j
tid = whole_ids[start + j]
total += logprobs[i, pos, tid].item()
scores.append((total, max(cont_len, 1)))
return scores
def batched_score(model, tok, device, pairs: list[tuple[str, str]], batch_size: int, tag: str='') -> list[tuple[float, int]]:
scores: list[tuple[float, int]] = [(0.0, 1)] * len(pairs)
n_batches = (len(pairs) + batch_size - 1) // batch_size
t0 = time.time()
for bi, start in enumerate(range(0, len(pairs), batch_size)):
chunk = pairs[start:start + batch_size]
chunk_scores = _score_batch(model, tok, device, chunk)
scores[start:start + len(chunk)] = chunk_scores
if tag and (bi % 10 == 0 or bi == n_batches - 1):
elapsed = time.time() - t0
done = start + len(chunk)
rate = done / elapsed if elapsed > 0 else 0
print(f'[eval] {tag}: {done}/{len(pairs)} pairs scored ({rate:.1f}/s, {elapsed:.0f}s elapsed)', flush=True)
return scores
def _pick(ex_scores, normalized: bool) -> int:
if normalized:
return max(range(len(ex_scores)), key=lambda k: ex_scores[k][0] / ex_scores[k][1])
return max(range(len(ex_scores)), key=lambda k: ex_scores[k][0])
def eval_hellaswag_style(model, tok, device, ds, ctx_key, endings_key, label_key, limit=None, batch_size=32, tag=''):
if limit:
ds = ds.select(range(min(limit, len(ds))))
pairs, counts = ([], [])
for row in ds:
endings = row[endings_key]
for e in endings:
pairs.append((row[ctx_key], e))
counts.append(len(endings))
scores = batched_score(model, tok, device, pairs, batch_size, tag)
raw_correct, norm_correct, idx = (0, 0, 0)
for i, row in enumerate(ds):
n = counts[i]
ex_scores = scores[idx:idx + n]
idx += n
gold = int(row[label_key])
raw_correct += int(_pick(ex_scores, False) == gold)
norm_correct += int(_pick(ex_scores, True) == gold)
return (norm_correct / len(ds), len(ds), raw_correct / len(ds))
def eval_arc(model, tok, device, config, limit=None, batch_size=32, tag=''):
from datasets import load_dataset
ds = load_dataset('allenai/ai2_arc', config, split='test')
if limit:
ds = ds.select(range(min(limit, len(ds))))
kept_rows = []
pairs, counts = ([], [])
for row in ds:
choices = row['choices']['text']
labels = row['choices']['label']
answer = row['answerKey']
if answer not in labels or not choices:
continue
kept_rows.append((row, labels.index(answer)))
for c in choices:
pairs.append((row['question'], c))
counts.append(len(choices))
scores = batched_score(model, tok, device, pairs, batch_size, tag)
raw_correct, norm_correct, idx = (0, 0, 0)
for (row, gold_idx), n in zip(kept_rows, counts):
ex_scores = scores[idx:idx + n]
idx += n
raw_correct += int(_pick(ex_scores, False) == gold_idx)
norm_correct += int(_pick(ex_scores, True) == gold_idx)
return (norm_correct / len(kept_rows), len(kept_rows), raw_correct / len(kept_rows))
def eval_piqa(model, tok, device, limit=None, batch_size=32, tag=''):
from datasets import load_dataset
ds = load_dataset('ybisk/piqa', split='validation', revision='refs/convert/parquet')
if limit:
ds = ds.select(range(min(limit, len(ds))))
pairs = []
for row in ds:
pairs.append((row['goal'], row['sol1']))
pairs.append((row['goal'], row['sol2']))
scores = batched_score(model, tok, device, pairs, batch_size, tag)
raw_correct, norm_correct = (0, 0)
for i, row in enumerate(ds):
pair = [scores[2 * i], scores[2 * i + 1]]
gold = int(row['label'])
raw_correct += int(_pick(pair, False) == gold)
norm_correct += int(_pick(pair, True) == gold)
return (norm_correct / len(ds), len(ds), raw_correct / len(ds))
def N(score_pct: float, chance: float) -> float:
return 100 * (score_pct - chance) / (100 - chance)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--ckpt', default=None)
parser.add_argument('--hf_model', default=None, help='if set, evaluate a HF transformers reference model instead of a LogosModel checkpoint -- used to calibrate the harness against a model with known published scores')
parser.add_argument('--device', default='cpu')
parser.add_argument('--threads', type=int, default=24)
parser.add_argument('--batch_size', type=int, default=32)
parser.add_argument('--limit', type=int, default=None, help='cap per-benchmark examples (debug/speed)')
args = parser.parse_args()
if not args.ckpt and (not args.hf_model):
parser.error('one of --ckpt or --hf_model is required')
torch.set_num_threads(args.threads)
if args.hf_model:
from transformers import AutoModelForCausalLM, AutoTokenizer
print(f'[eval] loading HF reference model {args.hf_model}...', flush=True)
tok = _HFTokenizerAdapter(AutoTokenizer.from_pretrained(args.hf_model))
model = _HFModelAdapter(AutoModelForCausalLM.from_pretrained(args.hf_model)).eval().to(args.device)
ckpt_label = args.hf_model
else:
tok = Tokenizer.from_file('artifacts/tokenizer.json')
print(f'[eval] loading checkpoint {args.ckpt}...', flush=True)
model = load_model(args.ckpt, args.device)
ckpt_label = args.ckpt
results = {}
from datasets import load_dataset
t0 = time.time()
print('[eval] running HellaSwag...', flush=True)
hs_ds = load_dataset('Rowan/hellaswag', split='validation')
results['hellaswag'] = eval_hellaswag_style(model, tok, args.device, hs_ds, 'ctx', 'endings', 'label', args.limit, args.batch_size, tag='hellaswag')
print(f"[eval] HellaSwag acc_norm: {results['hellaswag'][0]:.4f} (n={results['hellaswag'][1]}, {time.time() - t0:.0f}s)", flush=True)
t0 = time.time()
print('[eval] running ARC-Easy...', flush=True)
results['arc_easy'] = eval_arc(model, tok, args.device, 'ARC-Easy', args.limit, args.batch_size, tag='arc_easy')
print(f"[eval] ARC-Easy acc_norm: {results['arc_easy'][0]:.4f} (n={results['arc_easy'][1]}, {time.time() - t0:.0f}s)", flush=True)
t0 = time.time()
print('[eval] running ARC-Challenge...', flush=True)
results['arc_challenge'] = eval_arc(model, tok, args.device, 'ARC-Challenge', args.limit, args.batch_size, tag='arc_challenge')
print(f"[eval] ARC-Challenge acc_norm: {results['arc_challenge'][0]:.4f} (n={results['arc_challenge'][1]}, {time.time() - t0:.0f}s)", flush=True)
t0 = time.time()
print('[eval] running PIQA...', flush=True)
results['piqa'] = eval_piqa(model, tok, args.device, args.limit, args.batch_size, tag='piqa')
print(f"[eval] PIQA acc_norm: {results['piqa'][0]:.4f} (n={results['piqa'][1]}, {time.time() - t0:.0f}s)", flush=True)
t0 = time.time()
print('[eval] running ArithMark-3.0...', flush=True)
am_ds = load_dataset('AxiomicLabs/ArithMark-3.0', split='train')
results['arithmark3'] = eval_hellaswag_style(model, tok, args.device, am_ds, 'ctx', 'endings', 'label', args.limit, args.batch_size, tag='arithmark3')
print(f"[eval] ArithMark-3.0 acc_norm: {results['arithmark3'][0]:.4f} (n={results['arithmark3'][1]}, {time.time() - t0:.0f}s)", flush=True)
def index_from(hs, arc_e, arc_c, piqa, am):
return (N(hs * 100, 25) + N((arc_e + arc_c) / 2 * 100, 25) + N(piqa * 100, 50) + 0.65 * N(am * 100, 25)) / 3.65
print('', flush=True)
print('[eval] ==== acc vs acc_norm ====', flush=True)
print(f"[eval] {'task':<16}{'acc':>9}{'acc_norm':>11}", flush=True)
for key, label in [('hellaswag', 'HellaSwag'), ('arc_easy', 'ARC-Easy'), ('arc_challenge', 'ARC-Chall'), ('piqa', 'PIQA'), ('arithmark3', 'ArithMark-3')]:
norm, _n, raw = results[key]
print(f'[eval] {label:<16}{raw * 100:>8.2f}%{norm * 100:>10.2f}%', flush=True)
idx_raw = index_from(results['hellaswag'][2], results['arc_easy'][2], results['arc_challenge'][2], results['piqa'][2], results['arithmark3'][2])
idx_norm = index_from(results['hellaswag'][0], results['arc_easy'][0], results['arc_challenge'][0], results['piqa'][0], results['arithmark3'][0])
print(f'[eval] Index from acc: {idx_raw:.2f}', flush=True)
print(f'[eval] Index from acc_norm: {idx_norm:.2f}', flush=True)
hs_acc, hs_n = (results['hellaswag'][0], results['hellaswag'][1])
e_acc, e_n = (results['arc_easy'][0], results['arc_easy'][1])
c_acc, c_n = (results['arc_challenge'][0], results['arc_challenge'][1])
piqa_acc, piqa_n = (results['piqa'][0], results['piqa'][1])
am_acc, am_n = (results['arithmark3'][0], results['arithmark3'][1])
combined_arc = (e_acc + c_acc) / 2
idx = (N(hs_acc * 100, 25) + N(combined_arc * 100, 25) + N(piqa_acc * 100, 50) + 0.65 * N(am_acc * 100, 25)) / 3.65
print('', flush=True)
print('[eval] ==== SUMMARY ====', flush=True)
print(f'[eval] checkpoint: {ckpt_label}', flush=True)
print(f'[eval] HellaSwag: {hs_acc * 100:.2f}% (n={hs_n})', flush=True)
print(f'[eval] ARC-Easy: {e_acc * 100:.2f}% (n={e_n})', flush=True)
print(f'[eval] ARC-Challenge: {c_acc * 100:.2f}% (n={c_n})', flush=True)
print(f'[eval] Combined ARC: {combined_arc * 100:.2f}%', flush=True)
print(f'[eval] PIQA: {piqa_acc * 100:.2f}% (n={piqa_n})', flush=True)
print(f'[eval] ArithMark-3.0: {am_acc * 100:.2f}% (n={am_n})', flush=True)
print(f'[eval] Intelligence Index (approx): {idx:.2f}', flush=True)
print(f'[eval] target (GPT-X2.5-135M): 25.17', flush=True)
if __name__ == '__main__':
main()

BIN
index_curve.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 78 KiB

3
leaderboard.png Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:0452cb6ff569e4fea144309618d65c2e0b1f002e13e1d53d63174d500f5fea2c
size 113941

170
model.py Normal file
View File

@@ -0,0 +1,170 @@
from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-06):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x):
dtype = x.dtype
x = x.float()
x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
return x.to(dtype) * self.weight
def rope_freqs(seq_len, head_dim, theta=100000.0, device=None):
inv_freq = 1.0 / theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
t = torch.arange(seq_len, device=device).float()
freqs = torch.outer(t, inv_freq)
return torch.cat([freqs, freqs], dim=-1)
def rotate_half(x):
x1, x2 = x.chunk(2, dim=-1)
return torch.cat([-x2, x1], dim=-1)
def apply_rope(x, cos, sin):
return x * cos + rotate_half(x) * sin
class Attention(nn.Module):
def __init__(self, dim, n_heads, n_kv_heads, head_dim, qk_norm_eps=1e-06):
super().__init__()
if n_heads % n_kv_heads != 0:
raise ValueError(f'n_heads ({n_heads}) must be divisible by n_kv_heads ({n_kv_heads}); valid choices for {n_heads} heads: {[k for k in range(1, n_heads + 1) if n_heads % k == 0]}')
self.n_heads = n_heads
self.n_kv_heads = n_kv_heads
self.head_dim = head_dim
self.n_rep = n_heads // n_kv_heads
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
self.q_norm = RMSNorm(head_dim, eps=qk_norm_eps)
self.k_norm = RMSNorm(head_dim, eps=qk_norm_eps)
def forward(self, x, cos, sin):
B, T, _ = x.shape
q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
q = self.q_norm(q)
k = self.k_norm(k)
q = apply_rope(q, cos, sin)
k = apply_rope(k, cos, sin)
q = q.transpose(1, 2)
k = k.transpose(1, 2).repeat_interleave(self.n_rep, dim=1)
v = v.transpose(1, 2).repeat_interleave(self.n_rep, dim=1)
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
out = out.transpose(1, 2).contiguous().view(B, T, -1)
return self.o_proj(out)
class SwiGLU(nn.Module):
def __init__(self, dim, hidden):
super().__init__()
self.gate_proj = nn.Linear(dim, hidden, bias=False)
self.up_proj = nn.Linear(dim, hidden, bias=False)
self.down_proj = nn.Linear(hidden, dim, bias=False)
def forward(self, x):
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
class DecoderLayer(nn.Module):
def __init__(self, dim, n_heads, n_kv_heads, head_dim, mlp_hidden, norm_eps=1e-06):
super().__init__()
self.attn_norm = RMSNorm(dim, eps=norm_eps)
self.attn = Attention(dim, n_heads, n_kv_heads, head_dim)
self.mlp_norm = RMSNorm(dim, eps=norm_eps)
self.mlp = SwiGLU(dim, mlp_hidden)
def forward(self, x, cos, sin):
x = x + self.attn(self.attn_norm(x), cos, sin)
x = x + self.mlp(self.mlp_norm(x))
return x
class LogosModel(nn.Module):
def __init__(self, vocab_size=32768, dim=640, n_layers=30, n_heads=10, n_kv_heads=5, mlp_hidden=1728, max_seq_len=4096, rope_theta=100000.0, norm_eps=1e-06, tie_embeddings=True):
super().__init__()
self.dim = dim
self.n_layers = n_layers
self.head_dim = dim // n_heads
self.max_seq_len = max_seq_len
self.rope_theta = rope_theta
self.embed_tokens = nn.Embedding(vocab_size, dim)
self.layers = nn.ModuleList([DecoderLayer(dim, n_heads, n_kv_heads, self.head_dim, mlp_hidden, norm_eps) for _ in range(n_layers)])
self.norm_out = RMSNorm(dim, eps=norm_eps)
self.lm_head = nn.Linear(dim, vocab_size, bias=False)
if tie_embeddings:
self.lm_head.weight = self.embed_tokens.weight
self.apply(self._init_weights)
for layer in self.layers:
nn.init.normal_(layer.mlp.down_proj.weight, mean=0.0, std=0.02 / math.sqrt(2 * n_layers))
nn.init.normal_(layer.attn.o_proj.weight, mean=0.0, std=0.02 / math.sqrt(2 * n_layers))
cos, sin = self._build_rope_cache(max_seq_len)
self.register_buffer('rope_cos', cos, persistent=False)
self.register_buffer('rope_sin', sin, persistent=False)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, mean=0.0, std=0.02)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Embedding):
nn.init.normal_(m.weight, mean=0.0, std=0.02)
def _build_rope_cache(self, seq_len):
freqs = rope_freqs(seq_len, self.head_dim, self.rope_theta)
return (freqs.cos()[None, :, None, :], freqs.sin()[None, :, None, :])
def forward(self, input_ids, labels=None, loss_chunk_size=2048):
B, T = input_ids.shape
x = self.embed_tokens(input_ids)
cos = self.rope_cos[:, :T].to(x.dtype)
sin = self.rope_sin[:, :T].to(x.dtype)
for layer in self.layers:
x = layer(x, cos, sin)
x = self.norm_out(x)
if labels is None:
return (self.lm_head(x), None)
shift_x = x[:, :-1].reshape(-1, x.size(-1))
shift_labels = labels[:, 1:].reshape(-1)
n = shift_x.size(0)
total_loss = x.new_zeros((), dtype=torch.float32)
total_count = x.new_zeros((), dtype=torch.float32)
for start in range(0, n, loss_chunk_size):
end = min(start + loss_chunk_size, n)
chunk_labels = shift_labels[start:end]
valid = chunk_labels != -100
count = valid.sum()
if count == 0:
continue
chunk_logits = self.lm_head(shift_x[start:end]).float()
chunk_loss = F.cross_entropy(chunk_logits, chunk_labels, ignore_index=-100, reduction='sum')
total_loss = total_loss + chunk_loss
total_count = total_count + count
loss = total_loss / total_count.clamp(min=1)
return (None, loss)
def num_params(self, exclude_embeddings=False):
n = sum((p.numel() for p in self.parameters()))
if exclude_embeddings:
n -= self.embed_tokens.weight.numel()
if self.lm_head.weight is not self.embed_tokens.weight:
n -= self.lm_head.weight.numel()
return n
if __name__ == '__main__':
m = LogosModel()
print(f'total params: {m.num_params():,}')
print(f'non-embedding params: {m.num_params(exclude_embeddings=True):,}')
x = torch.randint(0, 32768, (2, 128))
logits, loss = m(x, labels=x)
print('loss:', loss.item() if loss is not None else None)
logits, _ = m(x)
print('logits shape (no labels):', logits.shape)

3
model.safetensors Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:8ba65c6de38014ad12a773ad2dea43d221579e1c958b6d57c899ae809031448d
size 629681616

344
pretrain_logos.py Normal file
View File

@@ -0,0 +1,344 @@
from __future__ import annotations
import argparse
import contextlib
import datetime
import math
import os
import queue
import threading
import time
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from tokenizers import Tokenizer
from model import LogosModel
from adamuon import build_optimizer, CombinedOptimizer
from decon import build_reference_ngrams, load_eval_reference_texts
from data_pipeline import build_source_streams, build_preset_streams, MixedTokenStream
SEQ_LEN = 1024
BATCH_SIZE = 96
MUON_LR = 0.02
ADAM_LR = 0.0003
WEIGHT_DECAY = 0.1
WARMUP_STEPS = 2000
GRAD_CLIP = 1.0
ADAMW_BASELINE_PEAK_LR = 0.0006
class PlainAdamWOptimizer:
def __init__(self, model, lr: float, weight_decay: float):
self.opt = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=weight_decay)
def zero_grad(self):
self.opt.zero_grad()
def step(self):
self.opt.step()
def set_lr(self, lr: float):
for g in self.opt.param_groups:
g['lr'] = lr
def state_dict(self):
return self.opt.state_dict()
def load_state_dict(self, sd):
self.opt.load_state_dict(sd)
def cosine_schedule(step: int, total_steps: int, warmup_steps: int, min_ratio: float=0.1) -> float:
if step < warmup_steps:
return step / max(1, warmup_steps)
progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
progress = min(progress, 1.0)
cosine = 0.5 * (1 + math.cos(math.pi * progress))
return min_ratio + (1 - min_ratio) * cosine
def is_distributed() -> bool:
return 'WORLD_SIZE' in os.environ and int(os.environ['WORLD_SIZE']) > 1
def _worker_loop(out_q, stop_ev, progress_val, cfg):
import numpy as np
from tokenizers import Tokenizer as _Tok
from data_pipeline import build_source_streams as _build, MixedTokenStream as _Mix
tok = _Tok.from_file(cfg['tokenizer_path'])
if cfg['decon']:
from decon import build_reference_ngrams, load_eval_reference_texts
ref = build_reference_ngrams(load_eval_reference_texts())
else:
ref = set()
if cfg.get('preset'):
from data_pipeline import build_preset_streams as _bp
streams = _bp(cfg['preset'], ref, decon_enabled=cfg['decon'], rank=cfg['rank'], world_size=cfg['world_size'])
else:
streams = _build(ref, decon_enabled=cfg['decon'], rank=cfg['rank'], world_size=cfg['world_size'], mix_mode=cfg['data_mix'], arith_source=cfg['arith_source'])
mixer = _Mix(tok, streams, seq_len=cfg['seq_len'], seed=cfg['seed'], mix_mode=cfg['data_mix'], arith_frac=cfg['arith_frac'], preset=cfg.get('preset'))
while not stop_ev.is_set():
mixer.progress = progress_val.value
batch = np.asarray(mixer.next_batch(cfg['batch_size']), dtype=np.int64)
while not stop_ev.is_set():
try:
out_q.put(batch, timeout=1.0)
break
except Exception:
continue
class ProcessBatchPrefetcher:
def __init__(self, cfg, n_workers: int, queue_depth: int=6):
import multiprocessing as mp
self.ctx = mp.get_context('spawn')
self.q = self.ctx.Queue(maxsize=queue_depth)
self.stop_ev = self.ctx.Event()
self.progress = self.ctx.Value('d', 0.0)
self.procs = []
for i in range(n_workers):
wcfg = dict(cfg)
wcfg['seed'] = cfg['seed'] * 1000 + i
wcfg['rank'] = cfg['rank'] * n_workers + i
wcfg['world_size'] = cfg['world_size'] * n_workers
p = self.ctx.Process(target=_worker_loop, args=(self.q, self.stop_ev, self.progress, wcfg), daemon=True)
p.start()
self.procs.append(p)
def set_progress(self, value: float):
self.progress.value = value
def get(self):
return self.q.get()
def stop(self):
self.stop_ev.set()
for p in self.procs:
p.join(timeout=10)
if p.is_alive():
p.terminate()
class BatchPrefetcher:
def __init__(self, mixer, batch_size: int, queue_depth: int=4):
self.mixer = mixer
self.batch_size = batch_size
self.q: queue.Queue = queue.Queue(maxsize=queue_depth)
self._stop = threading.Event()
self._thread = threading.Thread(target=self._run, daemon=True)
def start(self) -> 'BatchPrefetcher':
self._thread.start()
return self
def _run(self):
while not self._stop.is_set():
self.q.put(self.mixer.next_batch(self.batch_size))
def get(self):
return self.q.get()
def stop(self):
self._stop.set()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--total_tokens', type=float, default=22000000000.0)
parser.add_argument('--ckpt_dir', type=str, default='/workspace/ckpt')
parser.add_argument('--ckpt_every', type=int, default=5000)
parser.add_argument('--log_every', type=int, default=20)
parser.add_argument('--resume', action='store_true')
parser.add_argument('--decon', action='store_true', default=True)
parser.add_argument('--no-decon', dest='decon', action='store_false')
parser.add_argument('--smoke_steps', type=int, default=0, help='if >0, stop after this many steps (trial run)')
parser.add_argument('--dim', type=int, default=640)
parser.add_argument('--n_layers', type=int, default=30)
parser.add_argument('--n_heads', type=int, default=10)
parser.add_argument('--mlp_hidden', type=int, default=1728)
parser.add_argument('--n_kv_heads', type=int, default=5, help='must divide n_heads; = n_heads for full MHA')
parser.add_argument('--optimizer', choices=['adamuon', 'adamw'], default='adamw', help='adamuon is retained only for reproducing the earlier ablation')
parser.add_argument('--lr_min_ratio', type=float, default=0.1, help='floor of the cosine as a fraction of peak. 0.0 decays to zero, which is what a cooldown phase wants; the 0.1 default is for the main run. Note the floor multiplies --lr, so on a resumed run near 100%% progress the effective LR is --lr * lr_min_ratio.')
parser.add_argument('--lr', type=float, default=ADAMW_BASELINE_PEAK_LR, help='AdamW peak LR. 6e-4 matches Museko-125M; BananaMind-2-Pro used 1.5e-3 at this scale')
parser.add_argument('--arith_source', choices=['wordmath', 'drill', 'mixed'], default='mixed', help="wordmath = big-endian grade-school word problems (matches ArithMark-3's format); drill = the old reversed little-endian digit drills, kept for the A/B ablation")
parser.add_argument('--grad_accum', type=int, default=1, help='micro-batches per optimizer step; allreduce runs only on the last one')
parser.add_argument('--prefetch', choices=['process', 'thread'], default='process', help='process = data pipeline in worker processes (escapes the GIL); thread = the old single-thread prefetcher, kept for comparison')
parser.add_argument('--workers', type=int, default=3, help='data worker processes per rank')
parser.add_argument('--arith_frac', type=float, default=None, help='pin the arithmetic slice to this fraction of the mix instead of using the stage ramp -- used to concentrate signal in short ablation runs')
parser.add_argument('--warmup_steps', type=int, default=WARMUP_STEPS, help='override for short ablation runs where the default 2000 would eat the whole run')
parser.add_argument('--batch_size', type=int, default=BATCH_SIZE, help='override the module default, which was tuned for a 96GB card, not smaller ones')
parser.add_argument('--data_dir', default=None, help='read pre-tokenised uint16 shards from here instead of streaming. Streaming sustained ~520k tok/s against an 843k compute ceiling and kept tripping HF rate limits; a memmap has neither problem.')
parser.add_argument('--preset', choices=['smollm2', 'dclm_heavy', 'smollm2_math', 'dclm_heavy_math', 'cooldown'], default=None, help="single-stage data preset; overrides the staged curriculum entirely. The staged mix scored Index 11.73 at 9.14B (below GPT-2's 13.58) because 97%% of the gap to GPT-X2.5 sits in HellaSwag/PIQA/ARC, which need web text, not the code/arithmetic/dialogue the curriculum was spending tokens on.")
parser.add_argument('--data_mix', choices=['default', 'museko'], default='default', help="museko = flat web+10%% arithmetic only, matching Museko-125M's recipe, for data-mix ablation")
args = parser.parse_args()
world_size = int(os.environ.get('WORLD_SIZE', 1))
rank = int(os.environ.get('RANK', 0))
local_rank = int(os.environ.get('LOCAL_RANK', 0))
ddp = is_distributed()
is_main = rank == 0
if ddp:
dist.init_process_group(backend='nccl', timeout=datetime.timedelta(minutes=90))
torch.cuda.set_device(local_rank)
device = f'cuda:{local_rank}'
else:
device = 'cuda'
def log(msg):
if is_main:
print(msg, flush=True)
if is_main:
os.makedirs(args.ckpt_dir, exist_ok=True)
tokens_per_step_global = args.batch_size * SEQ_LEN * world_size * args.grad_accum
total_steps = int(args.total_tokens // tokens_per_step_global)
log(f'[train] world_size={world_size} total_steps={total_steps:,} tokens_per_step_global={tokens_per_step_global:,} target_tokens={args.total_tokens:,.0f}')
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
log('[train] loading tokenizer...')
tok = Tokenizer.from_file('artifacts/tokenizer.json')
if args.decon:
log('[train] building decontamination reference set (every rank builds its own copy -- eval sets are small, cheap enough to not bother broadcasting)...')
ref_texts = load_eval_reference_texts()
reference_ngrams = build_reference_ngrams(ref_texts)
log(f'[train] decon reference: {len(reference_ngrams):,} n-grams from {len(ref_texts):,} eval fields')
else:
reference_ngrams = set()
log('[train] decontamination DISABLED (--no-decon)')
log(f'[train] building + prewarming data streams for rank {rank}/{world_size} (cold-start latency on first connection per source, one-time cost)...')
if args.preset:
streams = build_preset_streams(args.preset, reference_ngrams, decon_enabled=args.decon, rank=rank, world_size=world_size)
else:
streams = build_source_streams(reference_ngrams, decon_enabled=args.decon, rank=rank, world_size=world_size, mix_mode=args.data_mix, arith_source=args.arith_source)
mixer = MixedTokenStream(tok, streams, seq_len=SEQ_LEN, seed=rank, mix_mode=args.data_mix, arith_frac=args.arith_frac, preset=args.preset)
if args.prefetch == 'process':
log('[train] skipping in-process prewarm (workers own the streams)')
else:
for name, gen in streams.items():
t0 = time.time()
try:
_ = next(gen)
except StopIteration:
pass
if is_main:
print(f"[train] prewarmed '{name}' in {time.time() - t0:.1f}s", flush=True)
if ddp:
dist.barrier()
log(f'[train] building model (dim={args.dim} layers={args.n_layers} heads={args.n_heads} kv_heads={args.n_kv_heads} mlp={args.mlp_hidden}, optimizer={args.optimizer}, arith_source={args.arith_source} arith_frac={args.arith_frac})...')
model = LogosModel(dim=args.dim, n_layers=args.n_layers, n_heads=args.n_heads, n_kv_heads=args.n_kv_heads, mlp_hidden=args.mlp_hidden).to(device).to(torch.bfloat16)
log(f'[train] params: {model.num_params():,} (non-embedding {model.num_params(exclude_embeddings=True):,})')
if args.optimizer == 'adamw':
opt = PlainAdamWOptimizer(model, lr=args.lr, weight_decay=WEIGHT_DECAY)
else:
muon_opt, adam_opt = build_optimizer(model, muon_lr=MUON_LR, adam_lr=ADAM_LR, weight_decay=WEIGHT_DECAY)
opt = CombinedOptimizer(muon_opt, adam_opt)
model_ddp = None
if ddp:
model_ddp = DDP(model, device_ids=[local_rank])
model_c = torch.compile(model_ddp)
else:
model_c = torch.compile(model)
start_step = 0
ckpt_path = os.path.join(args.ckpt_dir, 'latest.pt')
if args.resume and os.path.exists(ckpt_path):
log(f'[train] resuming from {ckpt_path}')
state = torch.load(ckpt_path, map_location=device)
model.load_state_dict(state['model'])
opt.load_state_dict(state['opt'])
mixer.buffer = state.get('mixer_buffer', [])
if 'tokens_done' in state:
tokens_done = state['tokens_done']
else:
tps = state.get('tokens_per_step', tokens_per_step_global)
tokens_done = state['step'] * tps
log(f'[train] old-format checkpoint, inferred tokens from step x {tps:,}')
start_step = int(tokens_done // tokens_per_step_global)
log(f'[train] resumed at {tokens_done / 1000000000.0:.3f}B tokens -> step {start_step:,}/{total_steps:,} (this config: {tokens_per_step_global:,} tokens/step)')
if state.get('tokens_per_step') not in (None, tokens_per_step_global):
log(f"[train] NOTE: checkpoint was written at {state['tokens_per_step']:,} tokens/step; rescaled cleanly to the current geometry")
mixer.progress = start_step / total_steps
if args.data_dir:
from bin_data import BinaryMixStream
from data_pipeline import MIX_PRESETS
binstream = BinaryMixStream(args.data_dir, MIX_PRESETS[args.preset], seq_len=SEQ_LEN, seed=rank + 1)
if is_main:
for k, v in binstream.token_counts().items():
log(f'[train] {k}: {v:,} tokens on disk')
prefetcher = None
elif args.prefetch == 'process':
log(f'[train] starting {args.workers} data worker processes (each does its own cold-start, one-time)...')
prefetcher = ProcessBatchPrefetcher(dict(tokenizer_path=os.path.abspath('artifacts/tokenizer.json'), decon=args.decon, rank=rank, world_size=world_size, data_mix=args.data_mix, arith_source=args.arith_source, arith_frac=args.arith_frac, seq_len=SEQ_LEN, batch_size=args.batch_size, seed=rank + 1, preset=args.preset), n_workers=args.workers)
prefetcher.set_progress(start_step / max(1, total_steps))
del streams
else:
prefetcher = BatchPrefetcher(mixer, args.batch_size).start()
first_batch = None
if args.data_dir:
first_batch = binstream.next_batch(args.batch_size)
else:
t_warm = time.time()
first_batch = prefetcher.get()
log(f'[train] first batch ready after {time.time() - t_warm:.0f}s')
if ddp:
dist.barrier()
log('[train] all ranks have data, starting')
running_loss = 0.0
running_count = 0
t_last = time.time()
tok_since_log = 0
for step in range(start_step, total_steps):
if args.data_dir:
pass
elif args.prefetch == 'process':
prefetcher.set_progress(step / total_steps)
else:
mixer.progress = step / total_steps
lr_scale = cosine_schedule(step, total_steps, args.warmup_steps, args.lr_min_ratio)
if args.optimizer == 'adamw':
opt.set_lr(args.lr * lr_scale)
else:
opt.set_lr(MUON_LR * lr_scale, ADAM_LR * lr_scale)
opt.zero_grad()
loss_val = None
for micro in range(args.grad_accum):
if first_batch is not None:
batch, first_batch = (first_batch, None)
else:
batch = binstream.next_batch(args.batch_size) if args.data_dir else prefetcher.get()
if isinstance(batch, list):
x = torch.tensor(batch, dtype=torch.long, device=device)
else:
x = torch.from_numpy(batch).to(device, non_blocking=True)
is_last = micro == args.grad_accum - 1
sync_ctx = contextlib.nullcontext() if model_ddp is None or is_last else model_ddp.no_sync()
with sync_ctx:
_, loss = model_c(x, labels=x)
(loss / args.grad_accum).backward()
loss_val = loss.detach() if loss_val is None else loss_val + loss.detach()
loss_val = loss_val / args.grad_accum
torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
opt.step()
if ddp:
dist.all_reduce(loss_val, op=dist.ReduceOp.AVG)
running_loss += loss_val.item()
running_count += 1
tok_since_log += tokens_per_step_global
if step % args.log_every == 0:
dt = time.time() - t_last
tok_per_sec = tok_since_log / dt if dt > 0 else 0.0
avg_loss = running_loss / max(1, running_count)
eta_hours = (total_steps - step) * tokens_per_step_global / max(tok_per_sec, 1) / 3600
log(f'[train] step {step:,}/{total_steps:,} loss={avg_loss:.4f} lr_scale={lr_scale:.3f} tok/s={tok_per_sec:,.0f} eta={eta_hours:.1f}h stage_progress={mixer.progress:.3f}')
running_loss = 0.0
running_count = 0
t_last = time.time()
tok_since_log = 0
if is_main and step > 0 and (step % args.ckpt_every == 0):
meta = {'step': step, 'tokens_done': (step + 1) * tokens_per_step_global, 'tokens_per_step': tokens_per_step_global, 'total_tokens': args.total_tokens, 'world_size': world_size, 'batch_size': args.batch_size, 'grad_accum': args.grad_accum}
torch.save({'model': model.state_dict(), 'opt': opt.state_dict(), 'mixer_buffer': mixer.buffer, **meta}, ckpt_path)
torch.save({'model': model.state_dict(), **meta}, os.path.join(args.ckpt_dir, f'step{step}.pt'))
log(f'[train] checkpoint saved at step {step:,}')
if args.smoke_steps and step - start_step >= args.smoke_steps:
log(f'[train] smoke test: stopping after {args.smoke_steps} steps')
break
if prefetcher is not None:
prefetcher.stop()
if is_main:
torch.save({'model': model.state_dict(), 'opt': opt.state_dict(), 'step': step, 'tokens_done': (step + 1) * tokens_per_step_global, 'tokens_per_step': tokens_per_step_global, 'total_tokens': args.total_tokens, 'world_size': world_size, 'batch_size': args.batch_size, 'grad_accum': args.grad_accum, 'mixer_buffer': mixer.buffer}, ckpt_path)
log('[train] done')
if ddp:
dist.destroy_process_group()
if __name__ == '__main__':
main()

162902
tokenizer.json Normal file
View File

File diff suppressed because it is too large Load Diff

8
tokenizer_config.json Normal file
View File

@@ -0,0 +1,8 @@
{
"tokenizer_class": "PreTrainedTokenizerFast",
"model_max_length": 4096,
"eos_token": "<|endoftext|>",
"bos_token": null,
"pad_token": "<|endoftext|>",
"clean_up_tokenization_spaces": false
}