Upload Atom2.7m model

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+bg.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,162 @@
+---
+license: apache-2.0
+language:
+- en
+pipeline_tag: text-generation
+tags:
+- causal-lm
+- gpt
+- small-language-model
+- arithmetic
+- custom-tokenizer
+- custom-code
+- safetensors
+- lm-evaluation-harness
+datasets:
+- openbmb/Ultra-FineWeb
+- HuggingFaceFW/fineweb-edu
+- HuggingFaceTB/finemath
+- HuggingFaceTB/smollm-corpus
+---
+
+![bg](bg.png)
+# Atom2.7m
+
+Atom2.7m is a small decoder-only causal language model trained with a general byte-level BPE tokenizer plus arithmetic-specific digit features. The model has 2,738,880 parameters and uses custom code for both the model and the tokenizer path.
+
+## Model Details
+
+- Architecture: decoder-only GPT
+- Parameters: 2,738,880
+- Layers: 5
+- Hidden size: 192
+- Attention heads: 4
+- KV heads: 2
+- Context length: 512
+- Vocabulary size: 4,096
+- Token embeddings: tied input/output embeddings
+- Arithmetic feature embeddings:
+  - `place_vocab_size`: 66
+  - `role_vocab_size`: 12
+
+## Tokenizer
+
+This model should not be evaluated or used with a plain Hugging Face tokenizer path alone. It uses a custom fusion tokenizer implemented in `tokenizer_utils.py`.
+
+The tokenizer keeps byte-level BPE for ordinary text, but treats arithmetic sensitive spans specially:
+
+- digits `0`-`9` are atomic and never BPE-merged
+- digit spans are emitted least-significant-digit first
+- `+ - * / = ( )` are isolated atomic tokens
+- whitespace is isolated from text
+- `place_ids` are assigned to every digit run
+- `role_ids` are assigned only for strict integer equation spans
+
+The model expects aligned `input_ids`, `place_ids`, and `role_ids`.
+
+## Usage
+
+```python
+from pathlib import Path
+
+import torch
+from transformers import AutoModelForCausalLM
+
+from tokenizer_utils import load_tokenizer
+
+model_dir = Path(".")
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_dir,
+    trust_remote_code=True,
+).eval()
+tokenizer = load_tokenizer(model_dir)
+
+text = "12 + 34 ="
+encoding = tokenizer.encode(text)
+
+input_ids = torch.tensor([encoding.input_ids])
+place_ids = torch.tensor([encoding.place_ids])
+role_ids = torch.tensor([encoding.role_ids])
+
+with torch.no_grad():
+    outputs = model(
+        input_ids=input_ids,
+        place_ids=place_ids,
+        role_ids=role_ids,
+    )
+```
+
+For correct results, do not rely on `pipeline("text-generation")` unless it is wrapped to provide `place_ids` and `role_ids`.
+
+## Evaluation
+
+### ArithMark 2.0
+
+Use the included fusion-aware benchmark script:
+
+```bash
+python benchmark_fusion_arithmark.py \
+  --checkpoint . \
+  --tokenizer-dir . \
+  --data-path arithmark_2.0.jsonl \
+  --batch-size 64 \
+  --device cuda \
+  --output benchmark_results/fusion_arithmark_2.0_results.json
+```
+
+### lm-evaluation-harness
+
+Use the included launcher so the `atom2.7m` model wrapper is registered:
+
+```bash
+python lm_eval_fusion run \
+  --model atom2.7m \
+  --model_args pretrained=.,tokenizer_dir=. \
+  --tasks hellaswag,arc_easy,arc_challenge,piqa \
+  --device cuda:0 \
+  --batch_size auto \
+  --output_path benchmark_results/lm_eval
+```
+
+The wrapper uses `tokenizer_utils.load_tokenizer()` and forwards `place_ids` and `role_ids` to the model.
+
+## Results
+
+| Benchmark | Metric | Value |
+| --- | --- | ---: |
+| ArithMark 2.0 | acc | 0.6380 |
+| arc_challenge | acc_norm | 0.2261 |
+| arc_easy | acc_norm | 0.3270 |
+| hellaswag | acc_norm | 0.2733 |
+| piqa | acc_norm | 0.5305 |
+
+## Training Data
+
+The pretraining mixture targeted about 3.5B tokens:
+
+- Ultra-FineWeb: 900M
+- FineWeb-Edu: 900M
+- FineMath: 450M
+- Cosmopedia-v2: 337.5M
+- UltraData-Math-L2-preview: 337.5M
+- Ultra-FineWeb-L3-en-QA-Synthetic: 225M
+- Synthetic-Arithmetic: 350M
+
+Synthetic-Arithmetic is AtomCalc-style canonical integer equation data. The training curriculum is included as `pretraining_curriculum.json`.
+
+## Limitations
+
+- This is a very small model and should be treated as an experimental research artifact.
+- The custom tokenizer makes plain `AutoTokenizer` or default `lm_eval --model hf` unsuitable for final reported numbers.
+- Numeric text is represented least-significant-digit first internally.
+- Role annotations intentionally target strict integer equations, not broad math prose, decimals, rationals, or QA formats.
+
+## Files
+
+- `model.safetensors`: model weights
+- `config.json`, `config.py`, `configuration_gpt.py`, `model.py`: custom model code
+- `tokenizer.json`, `tokenizer_utils.py`: tokenizer files and fusion wrapper
+- `benchmark_fusion_arithmark.py`: ArithMark evaluation
+- `lm_eval_fusion.py`, `lm_eval_fusion`: lm-eval custom model wrapper
+- `pretraining_curriculum.json`: training curriculum

benchmark_fusion_arithmark.py

@@ -0,0 +1,291 @@
+"""Score an Atom2.7m checkpoint on ArithMark 2.0."""
+
+from __future__ import annotations
+
+import argparse
+from collections import Counter
+from contextlib import nullcontext
+import json
+from pathlib import Path
+import re
+import urllib.request
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM
+
+from tokenizer_utils import SPECIAL_TOKENS, FusionTokenizer, load_tokenizer
+
+
+DATA_URL = (
+    "https://huggingface.co/datasets/AxiomicLabs/Arithmark-2.0/"
+    "resolve/main/arithmark_2.0.jsonl"
+)
+PAD_ID = SPECIAL_TOKENS.index("<|pad|>")
+
+
+def ensure_data(path: Path) -> Path:
+    if path.exists():
+        return path
+    path.parent.mkdir(parents=True, exist_ok=True)
+    urllib.request.urlretrieve(DATA_URL, path)
+    return path
+
+
+def load_examples(path: Path, *, max_examples: int = 0) -> list[dict]:
+    examples = []
+    with path.open("r", encoding="utf-8") as handle:
+        for line in handle:
+            if not line.strip():
+                continue
+            examples.append(json.loads(line))
+            if max_examples > 0 and len(examples) >= max_examples:
+                break
+    return examples
+
+
+def _encoded_choice(
+    tokenizer: FusionTokenizer,
+    context: str,
+    ending: str,
+) -> tuple[list[int], list[int], list[int], int]:
+    context_encoding = tokenizer.encode(context)
+    full_encoding = tokenizer.encode(context + ending)
+    continuation_length = len(full_encoding.input_ids) - len(context_encoding.input_ids)
+    return (
+        full_encoding.input_ids,
+        full_encoding.place_ids,
+        full_encoding.role_ids,
+        continuation_length,
+    )
+
+
+@torch.inference_mode()
+def evaluate(
+    model: torch.nn.Module,
+    tokenizer: FusionTokenizer,
+    examples: list[dict],
+    *,
+    device: torch.device,
+    batch_size: int,
+    dump_failures: bool = False,
+    failure_operator_count: int | None = None,
+    max_failures: int = 100,
+) -> dict:
+    correct = 0
+    total = 0
+    by_operator_count: dict[str, list[int]] = {}
+    by_topic: dict[str, list[int]] = {}
+    failures: list[dict] = []
+    failure_summary: Counter[tuple[str, str, str]] = Counter()
+    model.eval()
+
+    for start in range(0, len(examples), batch_size):
+        batch_examples = examples[start : start + batch_size]
+        encoded = []
+        offsets = []
+        for example in batch_examples:
+            flat_start = len(encoded)
+            encoded.extend(
+                _encoded_choice(tokenizer, example["ctx"], ending)
+                for ending in example["endings"]
+            )
+            offsets.append((flat_start, len(example["endings"])))
+
+        max_length = max(len(item[0]) for item in encoded)
+        input_ids = torch.full(
+            (len(encoded), max_length),
+            PAD_ID,
+            dtype=torch.long,
+            device=device,
+        )
+        place_ids = torch.zeros_like(input_ids)
+        role_ids = torch.zeros_like(input_ids)
+        attention_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+        lengths = []
+        continuation_lengths = []
+        for row, (ids, places, roles, continuation_length) in enumerate(encoded):
+            length = len(ids)
+            input_ids[row, :length] = torch.tensor(ids, device=device)
+            place_ids[row, :length] = torch.tensor(places, device=device)
+            role_ids[row, :length] = torch.tensor(roles, device=device)
+            attention_mask[row, :length] = True
+            lengths.append(length)
+            continuation_lengths.append(continuation_length)
+
+        autocast = (
+            torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+            if device.type == "cuda"
+            else nullcontext()
+        )
+        with autocast:
+            logits = model(
+                input_ids=input_ids,
+                place_ids=place_ids,
+                role_ids=role_ids,
+                attention_mask=attention_mask,
+            ).logits
+        log_probs = F.log_softmax(logits.float(), dim=-1)
+
+        for example_index, example in enumerate(batch_examples):
+            flat_start, choice_count = offsets[example_index]
+            likelihoods = []
+            for choice_index in range(choice_count):
+                row = flat_start + choice_index
+                length = lengths[row]
+                continuation_length = continuation_lengths[row]
+                continuation_start = length - continuation_length
+                likelihood = 0.0
+                for position in range(continuation_start, length):
+                    likelihood += float(
+                        log_probs[row, position - 1, input_ids[row, position]].item()
+                    )
+                likelihoods.append(likelihood)
+
+            prediction = max(range(choice_count), key=likelihoods.__getitem__)
+            label = int(example["label"])
+            matched = prediction == label
+            correct += int(matched)
+            total += 1
+            metadata = example.get("metadata", {})
+            operator_count = str(metadata.get("operator_count", "unknown"))
+            topic = str(metadata.get("topic", "unknown"))
+            for grouped, key in (
+                (by_operator_count, operator_count),
+                (by_topic, topic),
+            ):
+                group = grouped.setdefault(key, [0, 0])
+                group[0] += int(matched)
+                group[1] += 1
+
+            if not matched and dump_failures:
+                op_count_int = None
+                try:
+                    op_count_int = int(operator_count)
+                except ValueError:
+                    pass
+                if failure_operator_count is None or op_count_int == failure_operator_count:
+                    context = str(example["ctx"]).strip()
+                    expression = context[:-1].strip() if context.endswith("=") else context
+                    operands = [int(value) for value in re.findall(r"\d+", expression)]
+                    operator = "".join(re.findall(r"[+\-*/]", expression))
+                    predicted_answer = str(example["endings"][prediction]).strip()
+                    correct_answer = str(example["endings"][label]).strip()
+                    width = max((len(str(value)) for value in operands), default=0)
+                    failure_summary[(topic, operator, f"width={width}")] += 1
+                    if len(failures) < max_failures:
+                        failures.append(
+                            {
+                                "ctx": context,
+                                "topic": topic,
+                                "operator_count": operator_count,
+                                "operator": operator,
+                                "operands": operands,
+                                "max_operand_digits": width,
+                                "correct_answer": correct_answer,
+                                "predicted_answer": predicted_answer,
+                                "choices": [str(value).strip() for value in example["endings"]],
+                                "choice_scores": [round(value, 4) for value in likelihoods],
+                                "score_margin_correct_minus_predicted": round(
+                                    likelihoods[label] - likelihoods[prediction],
+                                    4,
+                                ),
+                            }
+                        )
+
+    results = {
+        "benchmark": "arithmark_2.0",
+        "model_type": "atom2.7m",
+        "accuracy": correct / max(total, 1),
+        "correct": correct,
+        "total": total,
+        "by_operator_count": {
+            key: {
+                "accuracy": values[0] / max(values[1], 1),
+                "correct": values[0],
+                "total": values[1],
+            }
+            for key, values in sorted(by_operator_count.items())
+        },
+        "by_topic": {
+            key: {
+                "accuracy": values[0] / max(values[1], 1),
+                "correct": values[0],
+                "total": values[1],
+            }
+            for key, values in sorted(by_topic.items())
+        },
+    }
+    if dump_failures:
+        results["failure_summary"] = {
+            "|".join(key): value
+            for key, value in failure_summary.most_common()
+        }
+        results["failures"] = failures
+    return results
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", type=Path, default=Path("outputs/fusion_run/final_model"))
+    parser.add_argument("--tokenizer-dir", type=Path, default=Path("tokenizer_4k"))
+    parser.add_argument("--data-path", type=Path, default=Path("arithmark_2.0.jsonl"))
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--device", default="auto")
+    parser.add_argument("--output", type=Path)
+    parser.add_argument(
+        "--max-examples",
+        type=int,
+        default=0,
+        help="Evaluate only the first N examples. Default evaluates all examples.",
+    )
+    parser.add_argument(
+        "--dump-failures",
+        action="store_true",
+        help="Include incorrectly scored examples and grouped failure summary.",
+    )
+    parser.add_argument(
+        "--failure-operator-count",
+        type=int,
+        default=None,
+        help="Only dump failures with this operator count, e.g. 1 for easy examples.",
+    )
+    parser.add_argument("--max-failures", type=int, default=100)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+
+    data_path = ensure_data(args.data_path)
+    examples = load_examples(data_path, max_examples=args.max_examples)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.checkpoint,
+        trust_remote_code=True,
+    ).to(device)
+    tokenizer = load_tokenizer(args.tokenizer_dir)
+    results = evaluate(
+        model,
+        tokenizer,
+        examples,
+        device=device,
+        batch_size=args.batch_size,
+        dump_failures=args.dump_failures,
+        failure_operator_count=args.failure_operator_count,
+        max_failures=args.max_failures,
+    )
+    print(json.dumps(results, indent=2, sort_keys=True))
+    if args.output:
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(
+            json.dumps(results, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+
+
+if __name__ == "__main__":
+    main()

config.json

@@ -0,0 +1,29 @@
+{
+  "architectures": [
+    "GPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.GPTConfig",
+    "AutoModelForCausalLM": "model.GPTForCausalLM"
+  },
+  "block_size": 512,
+  "dtype": "float32",
+  "head_dim": 48,
+  "hidden_size": 192,
+  "intermediate_size": 480,
+  "labels_are_shifted": true,
+  "max_position_embeddings": 512,
+  "model_type": "gpt",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 2,
+  "place_vocab_size": 66,
+  "rms_norm_eps": 1e-06,
+  "role_vocab_size": 12,
+  "rope_theta": 5000.0,
+  "transformers_version": "4.57.6",
+  "use_place_embeddings": true,
+  "use_role_embeddings": true,
+  "vocab_size": 4096,
+  "xsa_projection": true
+}

config.py

@@ -0,0 +1,233 @@
+"""Environment-driven training configuration."""
+
+from __future__ import annotations
+
+import os
+import math
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from transformers import PretrainedConfig
+
+
+DEFAULT_VOCAB_SIZE = 4096
+DEFAULT_HIDDEN_SIZE = 192
+DEFAULT_NUM_HIDDEN_LAYERS = 5
+DEFAULT_NUM_ATTENTION_HEADS = 4
+DEFAULT_NUM_KEY_VALUE_HEADS = 2
+DEFAULT_HEAD_DIM = DEFAULT_HIDDEN_SIZE // DEFAULT_NUM_ATTENTION_HEADS
+DEFAULT_INTERMEDIATE_SIZE = DEFAULT_HIDDEN_SIZE * 5 // 2
+DEFAULT_BLOCK_SIZE = 512
+DEFAULT_ROPE_THETA = 5000.0
+DEFAULT_PLACE_VOCAB_SIZE = 66
+DEFAULT_ROLE_VOCAB_SIZE = 12
+
+
+class GPTConfig(PretrainedConfig):
+    model_type = "gpt"
+
+    def __init__(
+        self,
+        vocab_size: int = DEFAULT_VOCAB_SIZE,
+        hidden_size: int = DEFAULT_HIDDEN_SIZE,
+        num_hidden_layers: int = DEFAULT_NUM_HIDDEN_LAYERS,
+        num_attention_heads: int = DEFAULT_NUM_ATTENTION_HEADS,
+        num_key_value_heads: int | None = DEFAULT_NUM_KEY_VALUE_HEADS,
+        intermediate_size: int | None = DEFAULT_INTERMEDIATE_SIZE,
+        head_dim: int | None = None,
+        block_size: int = DEFAULT_BLOCK_SIZE,
+        rope_theta: float = DEFAULT_ROPE_THETA,
+        rms_norm_eps: float = 1e-6,
+        xsa_projection: bool = True,
+        tie_word_embeddings: bool = True,
+        labels_are_shifted: bool = False,
+        use_place_embeddings: bool = True,
+        use_role_embeddings: bool = True,
+        place_vocab_size: int = DEFAULT_PLACE_VOCAB_SIZE,
+        role_vocab_size: int = DEFAULT_ROLE_VOCAB_SIZE,
+        **kwargs,
+    ):
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        if head_dim is None:
+            if hidden_size % num_attention_heads != 0:
+                raise ValueError("hidden_size must be divisible by num_attention_heads")
+            head_dim = hidden_size // num_attention_heads
+        if intermediate_size is None:
+            intermediate_size = hidden_size * 4
+        if num_attention_heads % num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        if head_dim % 2 != 0:
+            raise ValueError("head_dim must be even for RoPE")
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.vocab_size = int(vocab_size)
+        self.hidden_size = int(hidden_size)
+        self.num_hidden_layers = int(num_hidden_layers)
+        self.num_attention_heads = int(num_attention_heads)
+        self.num_key_value_heads = int(num_key_value_heads)
+        self.intermediate_size = int(intermediate_size)
+        self.head_dim = int(head_dim)
+        self.block_size = int(block_size)
+        self.max_position_embeddings = int(block_size)
+        self.rope_theta = float(rope_theta)
+        self.rms_norm_eps = float(rms_norm_eps)
+        self.xsa_projection = bool(xsa_projection)
+        self.labels_are_shifted = bool(labels_are_shifted)
+        self.use_place_embeddings = bool(use_place_embeddings)
+        self.use_role_embeddings = bool(use_role_embeddings)
+        self.place_vocab_size = int(place_vocab_size)
+        self.role_vocab_size = int(role_vocab_size)
+
+
+def _bool_env(name: str, default: bool) -> bool:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    return raw.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _path_env(name: str, default: str) -> str:
+    return str(Path(os.environ.get(name, default)).expanduser())
+
+
+@dataclass
+class Hyperparameters:
+    data_dir: str = field(default_factory=lambda: _path_env("DATA_DIR", "."))
+    tokenized_dir: str = field(default_factory=lambda: _path_env("TOKENIZED_DIR", "tokenized2"))
+    tokenizer_dir: str = field(default_factory=lambda: _path_env("TOKENIZER_DIR", "tokenizer_4k"))
+    tokenizer_path: str = field(default_factory=lambda: os.environ.get("TOKENIZER_PATH", ""))
+    curriculum_path: str = field(default_factory=lambda: os.environ.get("CURRICULUM_PATH", ""))
+    mix_weights_path: str = field(default_factory=lambda: os.environ.get("MIX_WEIGHTS_PATH", ""))
+    run_id: str = field(default_factory=lambda: os.environ.get("RUN_ID", str(uuid.uuid4())))
+    seed: int = field(default_factory=lambda: int(os.environ.get("SEED", "1337")))
+    rank: int = field(init=False)
+
+    iterations: int = field(default_factory=lambda: int(os.environ.get("ITERATIONS", "10000")))
+    requested_train_tokens: int = field(init=False)
+    train_tokens: int = field(init=False)
+    decay_start_frac: float = field(default_factory=lambda: float(os.environ.get("DECAY_START_FRAC", "0.7")))
+    warmup_steps: int = field(default_factory=lambda: int(os.environ.get("WARMUP_STEPS", "0")))
+    lr_warmup_steps: int = field(default_factory=lambda: int(os.environ.get("LR_WARMUP_STEPS", "500")))
+    train_batch_tokens: int = field(default_factory=lambda: int(os.environ.get("TRAIN_BATCH_TOKENS", str(DEFAULT_BLOCK_SIZE * 512))))
+    train_seq_len: int = field(init=False)
+    eval_seq_len: int = field(init=False)
+    grad_accum_steps: int = field(default_factory=lambda: int(os.environ.get("GRAD_ACCUM_STEPS", "4")))
+    train_log_every: int = field(default_factory=lambda: int(os.environ.get("TRAIN_LOG_EVERY", "100")))
+    train_log_dense_steps: int = field(default_factory=lambda: int(os.environ.get("TRAIN_LOG_DENSE_STEPS", "100")))
+    train_log_ramp_steps: int = field(
+        default_factory=lambda: int(
+            os.environ.get(
+                "TRAIN_LOG_RAMP_STEPS",
+                os.environ.get("TRAIN_LOG_FIRST_STEPS", "500"),
+            )
+        )
+    )
+
+    val_batch_tokens: int = field(default_factory=lambda: int(os.environ.get("VAL_BATCH_TOKENS", str(DEFAULT_BLOCK_SIZE * 256))))
+    val_loss_every: int = field(default_factory=lambda: int(os.environ.get("VAL_LOSS_EVERY", "1000")))
+    val_max_tokens: int = field(default_factory=lambda: int(os.environ.get("VAL_MAX_TOKENS", "10_000_000")))
+
+    vocab_size: int = field(default_factory=lambda: int(os.environ.get("VOCAB_SIZE", str(DEFAULT_VOCAB_SIZE))))
+    hidden_size: int = field(default_factory=lambda: int(os.environ.get("HIDDEN_SIZE", os.environ.get("MODEL_DIM", str(DEFAULT_HIDDEN_SIZE)))))
+    num_hidden_layers: int = field(default_factory=lambda: int(os.environ.get("NUM_HIDDEN_LAYERS", os.environ.get("NUM_LAYERS", str(DEFAULT_NUM_HIDDEN_LAYERS)))))
+    num_attention_heads: int = field(default_factory=lambda: int(os.environ.get("NUM_ATTENTION_HEADS", os.environ.get("NUM_HEADS", str(DEFAULT_NUM_ATTENTION_HEADS)))))
+    num_key_value_heads: int = field(default_factory=lambda: int(os.environ.get("NUM_KEY_VALUE_HEADS", os.environ.get("NUM_KV_HEADS", str(DEFAULT_NUM_KEY_VALUE_HEADS)))))
+    head_dim: int = field(init=False)
+    intermediate_size: int = field(default_factory=lambda: int(os.environ.get("INTERMEDIATE_SIZE", os.environ.get("INTERMEDIATE", str(DEFAULT_INTERMEDIATE_SIZE)))))
+    block_size: int = field(default_factory=lambda: int(os.environ.get("BLOCK_SIZE", str(DEFAULT_BLOCK_SIZE))))
+    rope_theta: float = field(default_factory=lambda: float(os.environ.get("ROPE_THETA", os.environ.get("ROPE_BASE", str(DEFAULT_ROPE_THETA)))))
+    rms_norm_eps: float = field(default_factory=lambda: float(os.environ.get("RMS_NORM_EPS", "1e-6")))
+    xsa_projection: bool = field(default_factory=lambda: _bool_env("XSA_PROJECTION", True))
+    tie_word_embeddings: bool = field(default_factory=lambda: _bool_env("TIE_WORD_EMBEDDINGS", _bool_env("TIE_EMBEDDINGS", True)))
+    use_place_embeddings: bool = field(default_factory=lambda: _bool_env("USE_PLACE_EMBEDDINGS", True))
+    use_role_embeddings: bool = field(default_factory=lambda: _bool_env("USE_ROLE_EMBEDDINGS", True))
+    place_vocab_size: int = field(default_factory=lambda: int(os.environ.get("PLACE_VOCAB_SIZE", str(DEFAULT_PLACE_VOCAB_SIZE))))
+    role_vocab_size: int = field(default_factory=lambda: int(os.environ.get("ROLE_VOCAB_SIZE", str(DEFAULT_ROLE_VOCAB_SIZE))))
+
+    min_lr: float = field(default_factory=lambda: float(os.environ.get("MIN_LR", "0.0")))
+    lr: float = field(default_factory=lambda: float(os.environ.get("LR", "0.004")))
+    beta1: float = field(default_factory=lambda: float(os.environ.get("BETA1", "0.9")))
+    beta2: float = field(default_factory=lambda: float(os.environ.get("BETA2", "0.95")))
+    adam_eps: float = field(default_factory=lambda: float(os.environ.get("ADAM_EPS", "1e-8")))
+    weight_decay: float = field(default_factory=lambda: float(os.environ.get("WEIGHT_DECAY", "0.005")))
+
+    compile_model: bool = field(default_factory=lambda: _bool_env("COMPILE_MODEL", True))
+    autocast: bool = field(default_factory=lambda: _bool_env("AUTOCAST", True))
+    bf16: bool = field(default_factory=lambda: _bool_env("BF16", True))
+    device: str = field(default_factory=lambda: os.environ.get("DEVICE", "auto"))
+
+    output_dir: str = field(default_factory=lambda: _path_env("OUTPUT_DIR", "outputs"))
+    checkpoint_name: str = field(default_factory=lambda: os.environ.get("CHECKPOINT_NAME", "final_model"))
+    logfile: str = field(init=False)
+    model_path: str = field(init=False)
+    is_main_process: bool = True
+    train_files: str = field(init=False)
+    val_files: str = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.rank = int(os.environ.get("RANK", "0"))
+        if self.rank < 0:
+            raise ValueError("RANK must be non-negative")
+        self.is_main_process = self.rank == 0
+        self.head_dim = int(os.environ.get("HEAD_DIM", str(self.hidden_size // self.num_attention_heads)))
+        self.train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", str(self.block_size)))
+        self.eval_seq_len = int(os.environ.get("EVAL_SEQ_LEN", os.environ.get("TRAIN_SEQ_LEN", str(self.train_seq_len))))
+        token_alignment = self.train_seq_len * self.grad_accum_steps
+        if self.train_batch_tokens % token_alignment != 0:
+            raise ValueError(
+                "TRAIN_BATCH_TOKENS must be divisible by TRAIN_SEQ_LEN * GRAD_ACCUM_STEPS"
+            )
+        requested_train_tokens = int(os.environ.get("TRAIN_TOKENS", "0"))
+        self.requested_train_tokens = requested_train_tokens or self.iterations * self.train_batch_tokens
+        if self.requested_train_tokens <= 0:
+            raise ValueError("TRAIN_TOKENS must be positive")
+        self.train_tokens = self.requested_train_tokens - (self.requested_train_tokens % token_alignment)
+        if self.train_tokens <= 0:
+            raise ValueError(
+                "TRAIN_TOKENS must provide at least TRAIN_SEQ_LEN * GRAD_ACCUM_STEPS tokens"
+            )
+        self.iterations = math.ceil(self.train_tokens / self.train_batch_tokens)
+        tokenized = Path(self.tokenized_dir)
+        self.train_files = os.environ.get("TRAIN_FILES", str(tokenized / "*" / "shard_*.bin"))
+        self.val_files = os.environ.get("VAL_FILES", os.environ.get("TRAIN_FILES", self.train_files))
+        explicit_legacy_mix = bool(os.environ.get("MIX_WEIGHTS_PATH"))
+        if not self.curriculum_path and not explicit_legacy_mix:
+            tokenized_curriculum = tokenized / "curriculum.json"
+            default_curriculum = Path("pretraining_curriculum.json")
+            if tokenized_curriculum.exists():
+                self.curriculum_path = str(tokenized_curriculum)
+            elif default_curriculum.exists():
+                self.curriculum_path = str(default_curriculum)
+        if not self.mix_weights_path and not self.curriculum_path:
+            mix_weights_path = tokenized / "mix_weights.json"
+            self.mix_weights_path = str(mix_weights_path) if mix_weights_path.exists() else ""
+        if not self.tokenizer_path:
+            tok_dir = Path(self.tokenizer_dir)
+            json_path = tok_dir / "tokenizer.json"
+            self.tokenizer_path = str(json_path if json_path.exists() else tok_dir)
+        out = Path(self.output_dir)
+        self.logfile = os.environ.get("LOGFILE", str(out / "logs" / f"{self.run_id}.txt"))
+        self.model_path = os.environ.get("MODEL_PATH", str(out / self.checkpoint_name))
+
+    def to_model_config(self) -> GPTConfig:
+        return GPTConfig(
+            vocab_size=self.vocab_size,
+            hidden_size=self.hidden_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            num_key_value_heads=self.num_key_value_heads,
+            head_dim=self.head_dim,
+            intermediate_size=self.intermediate_size,
+            block_size=self.block_size,
+            rope_theta=self.rope_theta,
+            rms_norm_eps=self.rms_norm_eps,
+            xsa_projection=self.xsa_projection,
+            tie_word_embeddings=self.tie_word_embeddings,
+            use_place_embeddings=self.use_place_embeddings,
+            use_role_embeddings=self.use_role_embeddings,
+            place_vocab_size=self.place_vocab_size,
+            role_vocab_size=self.role_vocab_size,
+            labels_are_shifted=True,
+        )

configuration_gpt.py

@@ -0,0 +1,35 @@
+"""Exports for the GPT model configuration.
+
+New code should import these from :mod:`GPT.config`.
+"""
+
+from .config import (
+    DEFAULT_BLOCK_SIZE,
+    DEFAULT_HEAD_DIM,
+    DEFAULT_HIDDEN_SIZE,
+    DEFAULT_INTERMEDIATE_SIZE,
+    DEFAULT_NUM_ATTENTION_HEADS,
+    DEFAULT_NUM_HIDDEN_LAYERS,
+    DEFAULT_NUM_KEY_VALUE_HEADS,
+    DEFAULT_PLACE_VOCAB_SIZE,
+    DEFAULT_ROPE_THETA,
+    DEFAULT_ROLE_VOCAB_SIZE,
+    DEFAULT_VOCAB_SIZE,
+    GPTConfig,
+)
+
+
+__all__ = [
+    "DEFAULT_BLOCK_SIZE",
+    "DEFAULT_HEAD_DIM",
+    "DEFAULT_HIDDEN_SIZE",
+    "DEFAULT_INTERMEDIATE_SIZE",
+    "DEFAULT_NUM_ATTENTION_HEADS",
+    "DEFAULT_NUM_HIDDEN_LAYERS",
+    "DEFAULT_NUM_KEY_VALUE_HEADS",
+    "DEFAULT_PLACE_VOCAB_SIZE",
+    "DEFAULT_ROPE_THETA",
+    "DEFAULT_ROLE_VOCAB_SIZE",
+    "DEFAULT_VOCAB_SIZE",
+    "GPTConfig",
+]

lm_eval_fusion

@@ -0,0 +1,9 @@
+#!/usr/bin/env python
+"""Run lm-eval with the local Atom2.7m model registered."""
+
+import lm_eval_fusion  # noqa: F401
+from lm_eval.__main__ import cli_evaluate
+
+
+if __name__ == "__main__":
+    cli_evaluate()

lm_eval_fusion.py

@@ -0,0 +1,299 @@
+"""lm-eval wrapper for Atom2.7m checkpoints.
+
+The standard ``hf`` lm-eval model does not use the fusion tokenizer wrapper and
+does not pass arithmetic feature streams. This model keeps lm-eval's
+log-likelihood interface while encoding with ``tokenizer_utils.load_tokenizer``
+and forwarding ``place_ids`` and ``role_ids``.
+"""
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from lm_eval.api.model import LM
+from lm_eval.api.registry import register_model
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM
+
+from tokenizer_utils import EOT_ID, FusionTokenizer, load_tokenizer
+
+
+def _parse_bool(value: Any, default: bool = False) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, bool):
+        return value
+    return str(value).strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _parse_batch_size(value: int | str | None, max_batch_size: int | None) -> int:
+    if value is None:
+        return 1
+    if isinstance(value, int):
+        return value
+    text = str(value).strip().lower()
+    if text == "auto" or text.startswith("auto:"):
+        return int(max_batch_size or 64)
+    return int(text)
+
+
+def _dtype_from_name(value: str | torch.dtype | None) -> torch.dtype | None:
+    if value is None or value == "auto":
+        return None
+    if isinstance(value, torch.dtype):
+        return value
+    normalized = str(value).replace("torch.", "").lower()
+    if normalized in {"bf16", "bfloat16"}:
+        return torch.bfloat16
+    if normalized in {"fp16", "float16", "half"}:
+        return torch.float16
+    if normalized in {"fp32", "float32", "float"}:
+        return torch.float32
+    raise ValueError(f"Unsupported dtype: {value!r}")
+
+
+@register_model("atom2.7m")
+class FusionGPTLM(LM):
+    """Fusion-tokenizer GPT adapter for lm-eval log-likelihood tasks."""
+
+    def __init__(
+        self,
+        pretrained: str = "outputs/fusion_run/final_model",
+        tokenizer_dir: str = "tokenizer_4k",
+        batch_size: int | str | None = 1,
+        max_batch_size: int | None = 64,
+        max_length: int | None = None,
+        device: str | None = "cuda",
+        dtype: str | torch.dtype | None = "auto",
+        mixed_precision_dtype: str | torch.dtype | None = "auto",
+        trust_remote_code: bool | str | None = None,
+        **_: Any,
+    ) -> None:
+        super().__init__()
+        del trust_remote_code
+        if device is None or device == "auto":
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self._device = torch.device(device)
+        self.batch_size = _parse_batch_size(batch_size, max_batch_size)
+        self.tokenizer: FusionTokenizer = load_tokenizer(Path(tokenizer_dir))
+        self.model = AutoModelForCausalLM.from_pretrained(
+            Path(pretrained),
+            trust_remote_code=True,
+        ).to(self.device)
+        model_dtype = _dtype_from_name(dtype)
+        if model_dtype is not None:
+            self.model = self.model.to(dtype=model_dtype)
+        if mixed_precision_dtype == "auto":
+            self.mixed_precision_dtype = (
+                torch.bfloat16 if self.device.type == "cuda" else None
+            )
+        else:
+            self.mixed_precision_dtype = _dtype_from_name(mixed_precision_dtype)
+        self.model.eval()
+        self.max_length = int(
+            max_length
+            or getattr(self.model.config, "block_size", None)
+            or getattr(self.model.config, "max_position_embeddings", 512)
+        )
+
+    @property
+    def eot_token_id(self) -> int:
+        return EOT_ID
+
+    def tok_encode(
+        self,
+        string: str,
+        add_special_tokens: bool | None = None,
+        left_truncate_len: int | None = None,
+        **_: Any,
+    ) -> list[int]:
+        del add_special_tokens
+        ids = self.tokenizer.encode(string).input_ids
+        if left_truncate_len is not None:
+            ids = ids[-left_truncate_len:]
+        return ids
+
+    def tok_decode(self, tokens, skip_special_tokens: bool = True) -> str:
+        if isinstance(tokens, int):
+            tokens = [tokens]
+        return self.tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
+
+    def _encode_request(
+        self,
+        context: str,
+        continuation: str,
+    ) -> tuple[list[int], list[int], list[int], list[int], int]:
+        if context == "":
+            continuation_encoding = self.tokenizer.encode(continuation)
+            ids = [self.eot_token_id] + continuation_encoding.input_ids
+            place_ids = [0] + continuation_encoding.place_ids
+            role_ids = [0] + continuation_encoding.role_ids
+            context_len = 1
+            continuation_ids = continuation_encoding.input_ids
+        else:
+            n_spaces = len(context) - len(context.rstrip())
+            if n_spaces > 0:
+                continuation = context[-n_spaces:] + continuation
+                context = context[:-n_spaces]
+            full_encoding = self.tokenizer.encode(context + continuation)
+            context_encoding = self.tokenizer.encode(context)
+            ids = full_encoding.input_ids
+            place_ids = full_encoding.place_ids
+            role_ids = full_encoding.role_ids
+            context_len = len(context_encoding.input_ids)
+            continuation_ids = ids[context_len:]
+
+        if not continuation_ids:
+            raise ValueError("Continuation encoded to zero tokens")
+        return ids, place_ids, role_ids, continuation_ids, context_len
+
+    def loglikelihood(
+        self,
+        requests: list["Instance"],
+        disable_tqdm: bool = False,
+    ) -> list[tuple[float, bool]]:
+        encoded = [
+            self._encode_request(context, continuation)
+            for context, continuation in tqdm(
+                [req.args for req in requests],
+                desc="Fusion tokenizing inputs",
+                disable=disable_tqdm,
+            )
+        ]
+        results: list[tuple[float, bool]] = []
+        for start in tqdm(
+            range(0, len(encoded), self.batch_size),
+            desc="Running fusion loglikelihood requests",
+            disable=disable_tqdm or self.rank != 0,
+        ):
+            batch = encoded[start : start + self.batch_size]
+            rows = []
+            row_places = []
+            row_roles = []
+            row_targets = []
+            row_score_slices = []
+            for ids, place_ids, role_ids, continuation_ids, context_len in batch:
+                window_start = max(0, len(ids) - (self.max_length + 1))
+                window_ids = ids[window_start:]
+                window_places = place_ids[window_start:]
+                window_roles = role_ids[window_start:]
+                input_ids = window_ids[:-1]
+                targets = window_ids[1:]
+                full_score_start = context_len - 1
+                full_score_end = len(ids) - 1
+                score_start = max(full_score_start, window_start) - window_start
+                score_end = full_score_end - window_start
+                if score_end <= score_start:
+                    raise ValueError("No continuation tokens remain after truncation")
+                scored_continuation_ids = continuation_ids[-(score_end - score_start) :]
+                rows.append(input_ids)
+                row_places.append(window_places[:-1])
+                row_roles.append(window_roles[:-1])
+                row_targets.append(targets)
+                row_score_slices.append((score_start, score_end, scored_continuation_ids))
+
+            max_len = max(len(row) for row in rows)
+            input_tensor = torch.full(
+                (len(rows), max_len),
+                self.eot_token_id,
+                dtype=torch.long,
+                device=self.device,
+            )
+            place_tensor = torch.zeros_like(input_tensor)
+            role_tensor = torch.zeros_like(input_tensor)
+            attention_mask = torch.zeros_like(input_tensor, dtype=torch.bool)
+            target_tensor = torch.full_like(input_tensor, self.eot_token_id)
+            for row, (ids, places, roles, targets) in enumerate(
+                zip(rows, row_places, row_roles, row_targets, strict=True)
+            ):
+                length = len(ids)
+                input_tensor[row, :length] = torch.tensor(ids, device=self.device)
+                place_tensor[row, :length] = torch.tensor(places, device=self.device)
+                role_tensor[row, :length] = torch.tensor(roles, device=self.device)
+                target_tensor[row, :length] = torch.tensor(targets, device=self.device)
+                attention_mask[row, :length] = True
+
+            autocast = (
+                torch.autocast(
+                    device_type=self.device.type,
+                    dtype=self.mixed_precision_dtype,
+                    enabled=self.mixed_precision_dtype is not None,
+                )
+                if self.device.type == "cuda"
+                else nullcontext()
+            )
+            with torch.inference_mode(), autocast:
+                logits = self.model(
+                    input_ids=input_tensor,
+                    place_ids=place_tensor,
+                    role_ids=role_tensor,
+                    attention_mask=attention_mask,
+                ).logits
+                log_probs = F.log_softmax(logits.float(), dim=-1)
+
+            for row, (score_start, score_end, continuation_ids) in enumerate(row_score_slices):
+                row_log_probs = log_probs[row, score_start:score_end]
+                row_targets_for_score = target_tensor[row, score_start:score_end]
+                token_log_probs = torch.gather(
+                    row_log_probs,
+                    1,
+                    row_targets_for_score.unsqueeze(-1),
+                ).squeeze(-1)
+                greedy = torch.equal(
+                    row_log_probs.argmax(dim=-1),
+                    torch.tensor(continuation_ids, dtype=torch.long, device=self.device),
+                )
+                results.append((float(token_log_probs.sum().item()), bool(greedy)))
+
+        return results
+
+    def loglikelihood_rolling(
+        self,
+        requests: list["Instance"],
+        disable_tqdm: bool = False,
+    ) -> list[float]:
+        results = []
+        for (text,) in tqdm(
+            [req.args for req in requests],
+            desc="Running fusion rolling loglikelihood",
+            disable=disable_tqdm or self.rank != 0,
+        ):
+            encoding = self.tokenizer.encode(text)
+            ids = encoding.input_ids
+            places = encoding.place_ids
+            roles = encoding.role_ids
+            total = 0.0
+            start = 0
+            while start < len(ids):
+                end = min(len(ids), start + self.max_length)
+                prefix = [self.eot_token_id] if start == 0 else ids[start - 1 : start]
+                chunk_ids = prefix + ids[start:end]
+                chunk_places = [0] + places[start:end] if start == 0 else places[start - 1 : end]
+                chunk_roles = [0] + roles[start:end] if start == 0 else roles[start - 1 : end]
+                input_ids = torch.tensor([chunk_ids[:-1]], dtype=torch.long, device=self.device)
+                place_ids = torch.tensor([chunk_places[:-1]], dtype=torch.long, device=self.device)
+                role_ids = torch.tensor([chunk_roles[:-1]], dtype=torch.long, device=self.device)
+                targets = torch.tensor(chunk_ids[1:], dtype=torch.long, device=self.device)
+                with torch.inference_mode():
+                    logits = self.model(
+                        input_ids=input_ids,
+                        place_ids=place_ids,
+                        role_ids=role_ids,
+                    ).logits[0]
+                    log_probs = F.log_softmax(logits.float(), dim=-1)
+                total += float(
+                    torch.gather(log_probs, 1, targets.unsqueeze(-1)).sum().item()
+                )
+                start = end
+            results.append(total)
+        return results
+
+    def generate_until(self, requests, disable_tqdm: bool = False) -> list[str]:
+        raise NotImplementedError(
+            "FusionGPTLM currently supports loglikelihood tasks. "
+            "Use tasks with multiple-choice/loglikelihood output."
+        )

model.py

@@ -0,0 +1,335 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+from transformers import PreTrainedModel
+from transformers.cache_utils import DynamicCache
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .config import GPTConfig
+
+
+CONTROL_TENSOR_NAME_PATTERNS = (
+    "scale",
+    "gate",
+    "gain",
+    "norm",
+    "ln_",
+    "rms",
+)
+
+
+class CastedLinear(nn.Linear):
+    """Store linear params in FP32, cast to activation dtype for matmul."""
+
+    def forward(self, x: Tensor) -> Tensor:
+        weight = self.weight.to(dtype=x.dtype)
+        bias = self.bias.to(dtype=x.dtype) if self.bias is not None else None
+        return F.linear(x, weight, bias)
+
+
+def restore_fp32_params(model: nn.Module) -> None:
+    """Keep linear weights and control params in FP32 after dtype conversion."""
+    for module in model.modules():
+        if isinstance(module, CastedLinear):
+            module.float()
+    for name, param in model.named_parameters():
+        if (
+            param.ndim < 2
+            or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)
+        ) and param.dtype != torch.float32:
+            param.data = param.data.float()
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x.float() * rms).to(dtype=x.dtype) * self.weight.to(dtype=x.dtype)
+
+
+def build_rope_inv_freq(head_dim, theta=2500.0):
+    return 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+
+
+def precompute_rope_cos_sin(head_dim, seq_len, theta=2500.0):
+    freqs = build_rope_inv_freq(head_dim, theta)
+    t = torch.arange(seq_len, dtype=torch.float32)
+    freqs = torch.outer(t, freqs)
+    return freqs.cos(), freqs.sin()
+
+
+def _apply_rope(x, cos, sin):
+    x_float = x.float()
+    x_pair = x_float.reshape(*x_float.shape[:-1], -1, 2)
+    even = x_pair[..., 0]
+    odd = x_pair[..., 1]
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    x_rot = torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1)
+    return x_rot.flatten(-2).type_as(x)
+
+
+def apply_rotary_emb(q, k, freqs_cis):
+    cos, sin = freqs_cis
+    return _apply_rope(q, cos, sin), _apply_rope(k, cos, sin)
+
+class GPTAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.num_attention_heads
+        self.n_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.n_rep = self.n_head // self.n_kv_heads
+        self.xsa_projection = config.xsa_projection
+
+        self.q_proj = CastedLinear(config.hidden_size, self.n_head * self.head_dim, bias=False)
+        self.k_proj = CastedLinear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = CastedLinear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.o_proj = CastedLinear(self.n_head * self.head_dim, config.hidden_size, bias=False)
+
+    def _xsa_efficient(self, y: Tensor, v_current: Tensor) -> Tensor:
+        B, H, T, D = y.shape
+        Hkv = v_current.size(1)
+        group = H // Hkv
+
+        y_g = y.reshape(B, Hkv, group, T, D)
+        v_n = F.normalize(v_current, dim=-1).unsqueeze(2)
+
+        proj = (y_g * v_n).sum(dim=-1, keepdim=True) * v_n
+        return (y_g - proj).reshape(B, H, T, D)
+
+    def forward(self, x, freqs_cis, past_key_value=None, use_cache=False, attention_mask=None):
+        B, T, _ = x.size()
+
+        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        k_current = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v_current = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+
+        q, k_current = apply_rotary_emb(q, k_current, freqs_cis)
+
+        if past_key_value is not None:
+            k, v = past_key_value.update(k_current, v_current, self.layer_idx)
+        else:
+            k, v = k_current, v_current
+
+        S = k.size(2)
+
+        is_causal = past_key_value is None or past_key_value.get_seq_length(self.layer_idx) == T
+
+        attn_mask = None
+        if attention_mask is not None:
+            key_pad = attention_mask.to(torch.bool)[:, None, None, :]
+
+            if is_causal and T > 1:
+                causal = torch.ones(T, S, dtype=torch.bool, device=x.device).tril(diagonal=S - T)
+                attn_mask = key_pad & causal[None, None, :, :]
+            else:
+                attn_mask = key_pad.expand(B, 1, T, S)
+
+            is_causal = False
+
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            is_causal=is_causal,
+            enable_gqa=(self.n_kv_heads != self.n_head),
+        )
+
+        if self.xsa_projection:
+            y = self._xsa_efficient(y, v_current)
+
+        y = y.transpose(1, 2).contiguous().view(B, T, self.n_head * self.head_dim)
+        return self.o_proj(y)
+
+
+class GPTMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.w_gate = CastedLinear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w_up = CastedLinear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w_down = CastedLinear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
+
+
+class GPTBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.ln_1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attn = GPTAttention(config, layer_idx)
+        self.ln_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = GPTMLP(config)
+
+    def forward(self, x, freqs_cis, past_key_value=None, use_cache=False, attention_mask=None):
+        x = x + self.attn(self.ln_1(x), freqs_cis, past_key_value, use_cache, attention_mask=attention_mask)
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class GPTPreTrainedModel(PreTrainedModel):
+    config_class = GPTConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        std = self.config.hidden_size ** -0.5
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+
+
+class GPTForCausalLM(GPTPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.transformer = nn.ModuleDict(dict(
+            wte=nn.Embedding(config.vocab_size, config.hidden_size),
+            h=nn.ModuleList([GPTBlock(config, i) for i in range(config.num_hidden_layers)]),
+            ln_f=RMSNorm(config.hidden_size, eps=config.rms_norm_eps),
+        ))
+        self.lm_head = CastedLinear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.transformer["wte"].weight
+        if getattr(config, "use_place_embeddings", True):
+            self.place_embeddings = nn.Embedding(
+                config.place_vocab_size,
+                config.hidden_size,
+                padding_idx=0,
+            )
+        else:
+            self.place_embeddings = None
+        if getattr(config, "use_role_embeddings", True):
+            self.role_embeddings = nn.Embedding(
+                config.role_vocab_size,
+                config.hidden_size,
+                padding_idx=0,
+            )
+        else:
+            self.role_embeddings = None
+        self._freqs_cis_cache = None
+        self.post_init()
+        with torch.no_grad():
+            if self.place_embeddings is not None:
+                self.place_embeddings.weight[0].zero_()
+            if self.role_embeddings is not None:
+                self.role_embeddings.weight[0].zero_()
+        restore_fp32_params(self)
+
+    def _apply(self, fn):
+        module = super()._apply(fn)
+        restore_fp32_params(self)
+        return module
+
+    def get_input_embeddings(self):
+        return self.transformer["wte"]
+
+    def set_input_embeddings(self, value):
+        self.transformer["wte"] = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def embed_tokens(self, input_ids, *, place_ids=None, role_ids=None, **kwargs):
+        embeddings = self.transformer["wte"](input_ids)
+        if self.place_embeddings is not None:
+            if place_ids is None:
+                place_ids = torch.zeros_like(input_ids)
+            if place_ids.shape != input_ids.shape:
+                raise ValueError("place_ids must match input_ids shape")
+            embeddings = embeddings + self.place_embeddings(place_ids)
+        if self.role_embeddings is not None:
+            if role_ids is None:
+                role_ids = torch.zeros_like(input_ids)
+            if role_ids.shape != input_ids.shape:
+                raise ValueError("role_ids must match input_ids shape")
+            embeddings = embeddings + self.role_embeddings(role_ids)
+        return embeddings
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
+        if past_key_values is not None and past_key_values.get_seq_length() > 0:
+            input_ids = input_ids[:, -1:]
+            if kwargs.get("place_ids") is not None:
+                kwargs["place_ids"] = kwargs["place_ids"][:, -1:]
+            if kwargs.get("role_ids") is not None:
+                kwargs["role_ids"] = kwargs["role_ids"][:, -1:]
+        return {
+            "input_ids": input_ids,
+            "place_ids": kwargs.get("place_ids"),
+            "role_ids": kwargs.get("role_ids"),
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+            "use_cache": True,
+        }
+
+    def _get_freqs_cis(self, seq_len, device):
+        cache = self._freqs_cis_cache
+        if cache is None or cache[0].device != device or cache[0].size(0) < seq_len:
+            cache = tuple(
+                tensor.to(device)
+                for tensor in precompute_rope_cos_sin(self.config.head_dim, seq_len, self.config.rope_theta)
+            )
+            if torch.is_inference_mode_enabled():
+                return cache[0][:seq_len], cache[1][:seq_len]
+            self._freqs_cis_cache = cache
+        return cache[0][:seq_len], cache[1][:seq_len]
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        labels=None,
+        past_key_values: Optional[DynamicCache] = None,
+        use_cache=False,
+        **kwargs,
+    ):
+        B, T = input_ids.size()
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
+        x = self.embed_tokens(input_ids, **kwargs)
+        cos, sin = self._get_freqs_cis(past_len + T, input_ids.device)
+        freqs_cis = (
+            cos[past_len:past_len + T],
+            sin[past_len:past_len + T],
+        )
+
+        for block in self.transformer["h"]:
+            x = block(x, freqs_cis, past_key_values if use_cache else None, use_cache, attention_mask=attention_mask)
+
+        x = self.transformer["ln_f"](x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            if getattr(self.config, "labels_are_shifted", False):
+                loss = F.cross_entropy(logits.float().reshape(-1, logits.size(-1)), labels.reshape(-1))
+            else:
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+                loss = F.cross_entropy(shift_logits.float().view(-1, shift_logits.size(-1)), shift_labels.reshape(-1))
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values if use_cache else None,
+        )

model.safetensors

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

pretraining_curriculum.json

@@ -0,0 +1,46 @@
+{
+  "version": 1,
+  "transition_fraction": 0.1,
+  "stages": [
+    {
+      "name": "early",
+      "start": 0.0,
+      "end": 0.4,
+      "weights": {
+        "Ultra-FineWeb": 0.5,
+        "FineWeb-Edu": 0.38,
+        "FineMath": 0.05,
+        "Cosmopedia-v2": 0.05,
+        "UltraData-Math-L2-preview": 0.02
+      }
+    },
+    {
+      "name": "mid",
+      "start": 0.4,
+      "end": 0.8,
+      "weights": {
+        "Ultra-FineWeb": 0.12,
+        "FineWeb-Edu": 0.22,
+        "FineMath": 0.18,
+        "Cosmopedia-v2": 0.13,
+        "UltraData-Math-L2-preview": 0.12,
+        "Ultra-FineWeb-L3-en-QA-Synthetic": 0.05,
+        "Synthetic-Arithmetic": 0.18
+      }
+    },
+    {
+      "name": "late",
+      "start": 0.8,
+      "end": 1.0,
+      "weights": {
+        "Ultra-FineWeb": 0.105,
+        "FineWeb-Edu": 0.21,
+        "FineMath": 0.14,
+        "Cosmopedia-v2": 0.14,
+        "UltraData-Math-L2-preview": 0.105,
+        "Ultra-FineWeb-L3-en-QA-Synthetic": 0.2,
+        "Synthetic-Arithmetic": 0.1
+      }
+    }
+  ]
+}

requirements.txt

@@ -0,0 +1,6 @@
+torch
+transformers
+tokenizers
+safetensors
+tqdm
+lm-eval

special_tokens_map.json

@@ -0,0 +1,9 @@
+{
+  "additional_special_tokens": [
+    "<|endoftext|>"
+  ],
+  "bos_token": "<|bos|>",
+  "eos_token": "<|eos|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|unk|>"
+}

tokenizer_config.json

@@ -0,0 +1,11 @@
+{
+  "additional_special_tokens": [
+    "<|endoftext|>"
+  ],
+  "bos_token": "<|bos|>",
+  "eos_token": "<|eos|>",
+  "model_max_length": 512,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "GPT2TokenizerFast",
+  "unk_token": "<|unk|>"
+}

tokenizer_utils.py

@@ -0,0 +1,328 @@
+"""Shared construction and loading helpers for the project's tokenizer."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import json
+from pathlib import Path
+import re
+from typing import Any, Iterable
+
+
+SPECIAL_TOKENS = [
+    "<|pad|>",
+    "<|bos|>",
+    "<|eos|>",
+    "<|unk|>",
+    "<|endoftext|>",
+]
+EOT_ID = SPECIAL_TOKENS.index("<|endoftext|>")
+ARITHMETIC_TOKENS = ("+", "-", "*", "/", "=", "(", ")")
+MAX_PLACE_ID = 64
+PLACE_OVERFLOW_ID = MAX_PLACE_ID + 1
+PLACE_VOCAB_SIZE = PLACE_OVERFLOW_ID + 1
+RESULT_ROLE_ID = 10
+SPACE_ROLE_ID = 11
+ROLE_VOCAB_SIZE = SPACE_ROLE_ID + 1
+MAX_OPERAND_ROLES = 9
+
+
+@dataclass(frozen=True)
+class FusionEncoding:
+    ids: list[int]
+    place_ids: list[int]
+    role_ids: list[int]
+    tokens: list[str] = field(default_factory=list)
+
+    @property
+    def input_ids(self) -> list[int]:
+        return self.ids
+
+    def __len__(self) -> int:
+        return len(self.ids)
+
+    def __iter__(self):
+        return iter(self.ids)
+
+    def __post_init__(self) -> None:
+        if not (len(self.ids) == len(self.place_ids) == len(self.role_ids)):
+            raise ValueError("Fusion tokenizer streams must have equal length")
+
+
+def build_tokenizer() -> Any:
+    """Build a byte-level BPE tokenizer with explicit lossless boundaries."""
+    from tokenizers import Regex, Tokenizer, decoders, models, pre_tokenizers
+
+    tokenizer = Tokenizer(models.BPE(unk_token="<|unk|>"))
+    tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
+        [
+            pre_tokenizers.Split(
+                Regex(r"\s+|\d|[+\-*/=()]|[^\s\d+\-*/=()]+"),
+                behavior="isolated",
+            ),
+            pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
+        ]
+    )
+    tokenizer.decoder = decoders.ByteLevel()
+    return tokenizer
+
+
+class FusionTokenizer:
+    """Runtime wrapper adding LSD-first digit streams to a trained BPE tokenizer."""
+
+    _digit_span_re = re.compile(r"\d+")
+
+    def __init__(self, tokenizer: Any):
+        self.tokenizer = tokenizer
+        self._digit_token_ids = frozenset(
+            token_id
+            for digit in "0123456789"
+            if (token_id := self.tokenizer.token_to_id(digit)) is not None
+        )
+        self._digit_id_to_text = {
+            int(self.tokenizer.token_to_id(digit)): digit
+            for digit in "0123456789"
+            if self.tokenizer.token_to_id(digit) is not None
+        }
+        self._equals_id = self.tokenizer.token_to_id("=")
+        self._special_token_ids = frozenset(
+            token_id
+            for token in SPECIAL_TOKENS
+            if (token_id := self.tokenizer.token_to_id(token)) is not None
+        )
+        if len(self._digit_token_ids) != 10:
+            raise ValueError("Tokenizer vocabulary must contain atomic digit tokens 0-9")
+        if self._equals_id is None:
+            raise ValueError("Tokenizer vocabulary must contain an atomic '=' token")
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.tokenizer, name)
+
+    @property
+    def digit_token_ids(self) -> frozenset[int]:
+        return self._digit_token_ids
+
+    @property
+    def special_token_ids(self) -> frozenset[int]:
+        return self._special_token_ids
+
+    def get_vocab_size(self, with_added_tokens: bool = True) -> int:
+        return int(self.tokenizer.get_vocab_size(with_added_tokens=with_added_tokens))
+
+    def get_vocab(self, with_added_tokens: bool = True) -> dict[str, int]:
+        return self.tokenizer.get_vocab(with_added_tokens=with_added_tokens)
+
+    def token_to_id(self, token: str) -> int | None:
+        return self.tokenizer.token_to_id(token)
+
+    def id_to_token(self, token_id: int) -> str | None:
+        return self.tokenizer.id_to_token(int(token_id))
+
+    @classmethod
+    def _reverse_digit_spans(cls, text: str) -> str:
+        return cls._digit_span_re.sub(lambda match: match.group(0)[::-1], text)
+
+    def _decode_token_piece(self, token_id: int) -> str:
+        return self.tokenizer.decode([int(token_id)], skip_special_tokens=False)
+
+    @staticmethod
+    def _is_equation_whitespace(piece: str) -> bool:
+        return bool(piece) and piece.isspace() and "\n" not in piece and "\r" not in piece
+
+    def _is_equation_piece(self, token_id: int, piece: str) -> bool:
+        if token_id in self._special_token_ids:
+            return False
+        if token_id in self._digit_token_ids:
+            return True
+        if self._is_equation_whitespace(piece):
+            return True
+        return len(piece) == 1 and piece in set(ARITHMETIC_TOKENS)
+
+    def _annotate_equation_span(
+        self,
+        ids: list[int],
+        pieces: list[str],
+        start: int,
+        end: int,
+        role_ids: list[int],
+    ) -> None:
+        equals_positions = [
+            index
+            for index in range(start, end)
+            if ids[index] == self._equals_id
+        ]
+        if len(equals_positions) != 1:
+            return
+        equals_position = equals_positions[0]
+
+        digit_runs: list[tuple[int, int]] = []
+        index = start
+        while index < end:
+            if ids[index] not in self._digit_token_ids:
+                index += 1
+                continue
+            run_start = index
+            while index < end and ids[index] in self._digit_token_ids:
+                index += 1
+            digit_runs.append((run_start, index))
+
+        operand_runs = [(a, b) for a, b in digit_runs if b <= equals_position]
+        result_runs = [(a, b) for a, b in digit_runs if a > equals_position]
+        if not operand_runs or not result_runs or len(operand_runs) > MAX_OPERAND_ROLES:
+            return
+
+        for index in range(start, end):
+            if self._is_equation_whitespace(pieces[index]):
+                role_ids[index] = SPACE_ROLE_ID
+
+        for role, (run_start, run_end) in enumerate(operand_runs, start=1):
+            for index in range(run_start, run_end):
+                role_ids[index] = role
+        for run_start, run_end in result_runs:
+            for index in range(run_start, run_end):
+                role_ids[index] = RESULT_ROLE_ID
+
+    def annotate_ids(self, ids: Iterable[int]) -> tuple[list[int], list[int]]:
+        input_ids = [int(token_id) for token_id in ids]
+        place_ids = [0] * len(input_ids)
+        role_ids = [0] * len(input_ids)
+        pieces = [self._decode_token_piece(token_id) for token_id in input_ids]
+
+        index = 0
+        while index < len(input_ids):
+            if input_ids[index] not in self._digit_token_ids:
+                index += 1
+                continue
+            run_start = index
+            while index < len(input_ids) and input_ids[index] in self._digit_token_ids:
+                offset = index - run_start + 1
+                place_ids[index] = min(offset, PLACE_OVERFLOW_ID)
+                index += 1
+
+        span_start: int | None = None
+        for index, (token_id, piece) in enumerate(zip(input_ids, pieces, strict=True)):
+            if self._is_equation_piece(token_id, piece):
+                if span_start is None:
+                    span_start = index
+                continue
+            if span_start is not None:
+                self._annotate_equation_span(input_ids, pieces, span_start, index, role_ids)
+                span_start = None
+        if span_start is not None:
+            self._annotate_equation_span(input_ids, pieces, span_start, len(input_ids), role_ids)
+
+        return place_ids, role_ids
+
+    def encode(self, text: str, *args, **kwargs) -> FusionEncoding:
+        transformed = self._reverse_digit_spans(text)
+        encoding = self.tokenizer.encode(transformed, *args, **kwargs)
+        ids = [int(token_id) for token_id in encoding.ids]
+        place_ids, role_ids = self.annotate_ids(ids)
+        return FusionEncoding(
+            ids=ids,
+            place_ids=place_ids,
+            role_ids=role_ids,
+            tokens=list(getattr(encoding, "tokens", [])),
+        )
+
+    def encode_batch(self, texts: list[str], *args, **kwargs) -> list[FusionEncoding]:
+        return [self.encode(text, *args, **kwargs) for text in texts]
+
+    def decode(
+        self,
+        token_ids: Iterable[int],
+        skip_special_tokens: bool = True,
+    ) -> str:
+        pieces: list[str] = []
+        text_ids: list[int] = []
+        digit_buffer: list[str] = []
+
+        def flush_text() -> None:
+            if text_ids:
+                pieces.append(
+                    self.tokenizer.decode(
+                        text_ids,
+                        skip_special_tokens=skip_special_tokens,
+                    )
+                )
+                text_ids.clear()
+
+        def flush_digits() -> None:
+            if digit_buffer:
+                pieces.extend(reversed(digit_buffer))
+                digit_buffer.clear()
+
+        for raw_id in token_ids:
+            token_id = int(raw_id)
+            if token_id in self._digit_token_ids:
+                flush_text()
+                digit_buffer.append(self._digit_id_to_text[token_id])
+                continue
+
+            flush_digits()
+            text_ids.append(token_id)
+
+        flush_text()
+        flush_digits()
+        return "".join(pieces)
+
+
+def build_trainer(vocab_size: int, min_frequency: int) -> Any:
+    from tokenizers import pre_tokenizers, trainers
+
+    return trainers.BpeTrainer(
+        vocab_size=vocab_size,
+        min_frequency=min_frequency,
+        special_tokens=SPECIAL_TOKENS,
+        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+    )
+
+
+def tokenizer_files(tokenizer_dir: Path) -> tuple[Path, Path, Path]:
+    return (
+        tokenizer_dir / "tokenizer.json",
+        tokenizer_dir / "vocab.json",
+        tokenizer_dir / "merges.txt",
+    )
+
+
+def validate_tokenizer(tokenizer_dir: Path) -> None:
+    tokenizer_json, vocab_path, merges_path = tokenizer_files(tokenizer_dir)
+    if not tokenizer_json.exists():
+        raise FileNotFoundError(
+            f"Missing {tokenizer_json}. Retrain with train_tokenizer.py so the "
+            "whitespace and digit boundary rules are preserved."
+        )
+    if vocab_path.exists():
+        with vocab_path.open("r", encoding="utf-8") as f:
+            vocab = json.load(f)
+    else:
+        with tokenizer_json.open("r", encoding="utf-8") as f:
+            tokenizer_data = json.load(f)
+        vocab = tokenizer_data.get("model", {}).get("vocab")
+        if not isinstance(vocab, dict):
+            raise FileNotFoundError(f"Missing vocab.json and no embedded vocab in {tokenizer_json}")
+
+    max_id = max(vocab.values())
+    if max_id > 65_535:
+        raise ValueError(f"Tokenizer max id {max_id} does not fit in uint16")
+    if vocab.get("<|endoftext|>") != EOT_ID:
+        raise ValueError(
+            f"Expected <|endoftext|> id {EOT_ID}, "
+            f"got {vocab.get('<|endoftext|>')}"
+        )
+    missing = [
+        token
+        for token in (*[str(value) for value in range(10)], *ARITHMETIC_TOKENS)
+        if token not in vocab
+    ]
+    if missing:
+        raise ValueError(f"Tokenizer missing required atomic tokens: {missing}")
+
+
+def load_tokenizer(tokenizer_dir: Path) -> Any:
+    from tokenizers import Tokenizer
+
+    validate_tokenizer(tokenizer_dir)
+    tokenizer_json, _, _ = tokenizer_files(tokenizer_dir)
+    return FusionTokenizer(Tokenizer.from_file(str(tokenizer_json)))