Publish trained base checkpoint, evaluation confidence intervals and TensorBoard logs

README.md

@@ -0,0 +1,227 @@
+---
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- pretrained
+- from-scratch
+- tiny-llm-ablation
+- custom_code
+- tensorboard
+- causal-lm
+datasets:
+- HuggingFaceFW/fineweb-edu
+model-index:
+- name: q-51M-base
+  results:
+  - task:
+      type: text-generation
+      name: Zero-shot continuation likelihood
+    dataset:
+      type: Rowan/hellaswag
+      name: HellaSwag
+      config: default
+      split: validation
+      args:
+        num_few_shot: 0
+    metrics:
+    - type: acc_norm
+      name: acc_norm (fraction; lm-eval 0.4.12)
+      value: 0.29177454690300736
+      args:
+        ci95_low: 0.28296457902614003
+        ci95_high: 0.3007437626650362
+        ci_method: Wilson
+  - task:
+      type: text-generation
+      name: Zero-shot continuation likelihood
+    dataset:
+      type: allenai/ai2_arc
+      name: ARC-Easy
+      config: ARC-Easy
+      split: test
+      args:
+        num_few_shot: 0
+    metrics:
+    - type: acc_norm
+      name: acc_norm (fraction; lm-eval 0.4.12)
+      value: 0.43308080808080807
+      args:
+        ci95_low: 0.41328092357316454
+        ci95_high: 0.45309672993759426
+        ci_method: Wilson
+  - task:
+      type: text-generation
+      name: Zero-shot continuation likelihood
+    dataset:
+      type: allenai/ai2_arc
+      name: ARC-Challenge
+      config: ARC-Challenge
+      split: test
+      args:
+        num_few_shot: 0
+    metrics:
+    - type: acc_norm
+      name: acc_norm (fraction; lm-eval 0.4.12)
+      value: 0.24232081911262798
+      args:
+        ci95_low: 0.21865689491479706
+        ci95_high: 0.2676684124053075
+        ci_method: Wilson
+  - task:
+      type: text-generation
+      name: Zero-shot continuation likelihood
+    dataset:
+      type: baber/piqa
+      name: PIQA
+      config: default
+      split: validation
+      args:
+        num_few_shot: 0
+    metrics:
+    - type: acc_norm
+      name: acc_norm (fraction; lm-eval 0.4.12)
+      value: 0.5990206746463548
+      args:
+        ci95_low: 0.5764309418644076
+        ci95_high: 0.6211973601030507
+        ci_method: Wilson
+  - task:
+      type: text-generation
+      name: Zero-shot continuation likelihood
+    dataset:
+      type: allenai/winogrande
+      name: WinoGrande
+      config: winogrande_xl
+      split: validation
+      args:
+        num_few_shot: 0
+    metrics:
+    - type: acc
+      name: acc (fraction; lm-eval 0.4.12)
+      value: 0.500394632991318
+      args:
+        ci95_low: 0.47290359623559514
+        ci95_high: 0.5278832839791312
+        ci_method: Wilson
+  - task:
+      type: text-generation
+      name: Zero-shot continuation likelihood
+    dataset:
+      type: allenai/openbookqa
+      name: OpenBookQA
+      config: main
+      split: test
+      args:
+        num_few_shot: 0
+    metrics:
+    - type: acc_norm
+      name: acc_norm (fraction; lm-eval 0.4.12)
+      value: 0.282
+      args:
+        ci95_low: 0.2443364332272433
+        ci95_high: 0.32298777921399346
+        ci_method: Wilson
+  - task:
+      type: text-generation
+      name: Zero-shot continuation likelihood
+    dataset:
+      type: aps/super_glue
+      name: BoolQ
+      config: boolq
+      split: validation
+      args:
+        num_few_shot: 0
+    metrics:
+    - type: acc
+      name: acc (fraction; lm-eval 0.4.12)
+      value: 0.598776758409786
+      args:
+        ci95_low: 0.5818706845463824
+        ci95_high: 0.615451027004275
+        ci_method: Wilson
+  - task:
+      type: text-generation
+      name: Zero-shot continuation likelihood
+    dataset:
+      type: EleutherAI/lambada_openai
+      name: LAMBADA OpenAI
+      config: default
+      split: test
+      args:
+        num_few_shot: 0
+    metrics:
+    - type: acc
+      name: acc (fraction; lm-eval 0.4.12)
+      value: 0.20861633999611875
+      args:
+        ci95_low: 0.1977414700784266
+        ci95_high: 0.2199253277124155
+        ci_method: Wilson
+---
+
+# Q-51M Base
+
+A **50,878,208-parameter** English base model in the **Tiny llm ablation** experiment. Trained from scratch on exactly **3,932,160,000 source tokens** over **15,000 optimizer steps**. The token count measures processed input blocks, not unique text or supervised target tokens.
+
+## Architecture and references
+
+10 decoder layers, width 512, SwiGLU width 1792, 8 query / 2 KV heads, RoPE, per-head QK normalization, gated residual branches, tied embeddings; context 2048.
+
+Architecture and unchanged 32,768-token tokenizer: [q-project/Q-50M-Base](https://huggingface.co/q-project/Q-50M-Base). This checkpoint was trained from random initialization; it is not a fine-tune of the reference weights.
+
+## Training
+
+- Data: [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu), `sample-10BT`, streamed from local Parquet shards; shuffle buffer 100,000.
+- Objective: Causal next-token cross-entropy.
+- Batch: 32 sequences × 4 accumulation × 2048 tokens = 262,144 source tokens per step.
+- Fused AdamW; peak LR 0.001, betas (0.9, 0.95), weight decay 0.1 (no decay for bias/norm/1D parameters), gradient clipping 1.0. Linear warmup for 150 steps, then cosine decay to 10% of peak LR.
+- BF16 compute on one RTX 5070 Ti (16 GB), seed 2026; checkpoints retain FP32 weights. [Exact training configuration](training_config.json).
+
+## Evaluation
+
+Full official task splits, **zero-shot**, **lm-eval 0.4.12**, no chat template, BF16 on RTX 5070 Ti, maximum context 2048. Accuracy is a percentage. **± is one standard error; the separate bracketed column is the 95% Wilson confidence interval.** Intervals describe finite evaluation-sample uncertainty, not variation across training seeds; no multiple-comparison correction is applied.
+
+| Dataset | Split | Examples | Metric | Score ± SE (%) | 95% CI (%) |
+|---|---|---:|---|---:|---:|
+| [HellaSwag](https://huggingface.co/datasets/Rowan/hellaswag) | validation | 10,042 | `acc_norm` | 29.18 ± 0.45 | [28.30, 30.07] |
+| [ARC-Easy](https://huggingface.co/datasets/allenai/ai2_arc) | test | 2,376 | `acc_norm` | 43.31 ± 1.02 | [41.33, 45.31] |
+| [ARC-Challenge](https://huggingface.co/datasets/allenai/ai2_arc) | test | 1,172 | `acc_norm` | 24.23 ± 1.25 | [21.87, 26.77] |
+| [PIQA](https://huggingface.co/datasets/baber/piqa) | validation | 1,838 | `acc_norm` | 59.90 ± 1.14 | [57.64, 62.12] |
+| [WinoGrande](https://huggingface.co/datasets/allenai/winogrande) | validation | 1,267 | `acc` | 50.04 ± 1.41 | [47.29, 52.79] |
+| [OpenBookQA](https://huggingface.co/datasets/allenai/openbookqa) | test | 500 | `acc_norm` | 28.20 ± 2.01 | [24.43, 32.30] |
+| [BoolQ](https://huggingface.co/datasets/aps/super_glue) | validation | 3,270 | `acc` | 59.88 ± 0.86 | [58.19, 61.55] |
+| [LAMBADA OpenAI](https://huggingface.co/datasets/EleutherAI/lambada_openai) | test | 5,153 | `acc` | 20.86 ± 0.57 | [19.77, 21.99] |
+
+Q50M uses autoregressive continuation likelihood. LAMBADA accuracy requires the complete final-word token sequence. `acc_norm` is harness length-normalized option scoring; raw accuracy is also stored in [results.json](evaluation/results.json).
+
+**[WikiText-2 raw test](https://huggingface.co/datasets/Salesforce/wikitext), conditional continuation:** CPU FP32 re-evaluation on 291 nonoverlapping blocks (512 prefix + 512 scored suffix tokens), 148,992 scored tokens; 335 tail tokens excluded. NLL **3.560862**, 95% CI **[3.521807, 3.599507]**; token PPL **35.194**, 95% CI **[33.846, 36.580]**. Percentile block bootstrap, 10,000 resamples, seed 2026; exponentiate NLL endpoints for PPL. Blocks are the resampling unit; this does not model all within-document dependence. This is not standard rolling AR or word PPL. The earlier BF16 point is retained separately in TensorBoard, with no borrowed FP32 interval.
+
+The metadata contains author-reported `model-index` scores. The evaluated dataset repositories had no registered `eval.yaml` on 2026-09-20, so no `.eval_results` leaderboard entry or verified badge is claimed. [Machine-readable results and provenance](evaluation/results.json).
+
+## Usage
+
+Install `requirements.txt` (tested with Transformers 5.17.0 / PyTorch 2.11.0). Custom model code is included; `trust_remote_code=True` is required. This example runs on CPU.
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+repo = "d0rj/q-51M-base"
+tokenizer = AutoTokenizer.from_pretrained(repo)
+model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True).eval()
+inputs = tokenizer("The capital of France is", return_tensors="pt")
+output = model.generate(**inputs, max_new_tokens=64, do_sample=False)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+To reproduce the core evaluation from a downloaded repository, install `evaluation/requirements.txt` and run:
+
+```bash
+python evaluation/run_core.py --device cuda:0 --dtype bfloat16 --batch-size 8 --output evaluation-rerun
+```
+
+## TensorBoard and limitations
+
+[TensorBoard event files](tensorboard/) include training telemetry and `eval/<task>/<metric>` at step 15,000, plus separate CI bounds. Recovered text-log telemetry covers steps 12,020–15,000 only (150 points); earlier training telemetry is unavailable. `train/loss` was corrected for the historical gradient-accumulation logging scale; `recovery/loss_as_printed` preserves the printed values.
+
+These are small English continuation models, not instruction-tuned assistants. Equal source-token budgets do not imply equal target-token supervision or FLOPs. Benchmark contamination was not audited; results are from one training seed. Reference-model scores from different prompts, tokenizers or corpora are not directly interchangeable.

config.json

@@ -0,0 +1,34 @@
+{
+  "architectures": [
+    "Q50MForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "attention_scalar_gate": true,
+  "bos_token_id": 1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "gate_multiplier": 2.0,
+  "head_dim": 64,
+  "hidden_size": 512,
+  "intermediate_size": 1792,
+  "max_position_embeddings": 2048,
+  "mlp_scalar_gate": true,
+  "model_type": "q50m",
+  "nope_every_n": null,
+  "nope_layers": [],
+  "num_attention_heads": 8,
+  "num_hidden_layers": 10,
+  "num_key_value_heads": 2,
+  "pad_token_id": 0,
+  "qk_norm": true,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "vocab_size": 32768,
+  "auto_map": {
+    "AutoConfig": "configuration_q50m.Q50MConfig",
+    "AutoModelForCausalLM": "modeling_q50m.Q50MForCausalLM"
+  },
+  "_name_or_path": "d0rj/q-51M-base"
+}

configuration_q50m.py

@@ -0,0 +1,68 @@
+"""Configuration for the Q-50M architecture (replica of q-project/Q-50M-Base).
+
+512-wide, 10-layer decoder with GQA (8 q / 2 kv heads), per-head QK-Norm,
+content-dependent scalar residual gates on both branches, tied embeddings.
+"""
+
+from transformers import PretrainedConfig
+
+
+class Q50MConfig(PretrainedConfig):
+    model_type = "q50m"
+
+    def __init__(
+        self,
+        vocab_size: int = 32768,
+        hidden_size: int = 512,
+        intermediate_size: int = 1792,
+        num_hidden_layers: int = 10,
+        num_attention_heads: int = 8,
+        num_key_value_heads: int = 2,
+        head_dim: int = 64,
+        max_position_embeddings: int = 2048,
+        rope_theta: float = 10000.0,
+        rms_norm_eps: float = 1e-5,
+        qk_norm: bool = True,
+        attention_scalar_gate: bool = True,
+        mlp_scalar_gate: bool = True,
+        gate_multiplier: float = 2.0,
+        nope_every_n: int | None = None,
+        nope_layers: list[int] | None = None,
+        attention_dropout: float = 0.0,
+        tie_word_embeddings: bool = True,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        pad_token_id: int = 0,
+        **kwargs,
+    ):
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.qk_norm = qk_norm
+        self.attention_scalar_gate = attention_scalar_gate
+        self.mlp_scalar_gate = mlp_scalar_gate
+        self.gate_multiplier = gate_multiplier
+        self.nope_every_n = nope_every_n
+        if nope_layers is not None:
+            self.nope_layers = list(nope_layers)
+        elif nope_every_n:
+            self.nope_layers = [
+                i for i in range(num_hidden_layers) if (i + 1) % nope_every_n == 0
+            ]
+        else:
+            self.nope_layers = []
+        self.attention_dropout = attention_dropout

evaluation/adapters.py

@@ -0,0 +1,164 @@
+"""HF loading and an explicitly experimental masked-model adapter.
+
+Only imported for an actual evaluation, never for planning.
+"""
+
+import torch
+from lm_eval.models.huggingface import HFLM
+from transformers import AutoTokenizer
+from tqdm import tqdm
+
+
+class UL2HFLM(HFLM):
+    """S-denoiser continuation scoring; only answer text contributes to NLL.
+
+    Source: S + context + sentinel_0 + EOS.
+    Decoder input: BOS + sentinel_0 + answer[:-1]. Controls are conditioned
+    on, never scored; the softmax still includes the complete vocabulary.
+    """
+
+    def _encode_pair(self, context, continuation):
+        # Match causal HFLM's text boundary, without tokenizer-added controls.
+        spaces = len(context) - len(context.rstrip())
+        if spaces:
+            continuation = context[-spaces:] + continuation
+            context = context[:-spaces]
+        whole = self.tok_encode(context + continuation, add_special_tokens=False)
+        prefix = self.tok_encode(context, add_special_tokens=False)
+        return prefix, whole[len(prefix):]
+
+    @torch.inference_mode()
+    def _loglikelihood_tokens(self, requests, disable_tqdm=False, override_bs=None):
+        contract = self.model.config.ul2
+        span = contract["sentinel_ids"][0]
+        bos = self.model.config.decoder_start_token_id
+        batch_size = override_bs or self.batch_size
+        results = [None] * len(requests)
+        ordered = sorted(enumerate(requests), key=lambda x: -(len(x[1][1]) + len(x[1][2])))
+        for start in tqdm(range(0, len(ordered), batch_size), disable=disable_tqdm, desc="UL2 likelihood"):
+            batch = ordered[start:start + batch_size]
+            sources, decoders = [], []
+            for _, (_, context, answer) in batch:
+                if not answer or len(answer) >= self.max_length:
+                    raise ValueError("UL2 scoring requires a nonempty answer shorter than the context limit")
+                budget = min(self.max_length + 1 - len(answer), self.model.config.max_position_embeddings - 3)
+                context = context[-budget:]
+                sources.append([contract["mode_ids"]["S"], *context, span, contract["eos_id"]])
+                decoders.append([bos, span, *answer[:-1]])
+
+            def pad(rows):
+                ids = torch.full((len(rows), max(map(len, rows))), contract["pad_id"], dtype=torch.long, device=self.device)
+                mask = torch.zeros_like(ids)
+                for i, row in enumerate(rows):
+                    ids[i, :len(row)] = torch.tensor(row, device=self.device)
+                    mask[i, :len(row)] = 1
+                return ids, mask
+
+            source, source_mask = pad(sources)
+            decoder, decoder_mask = pad(decoders)
+            hidden = self.model.model(input_ids=source, attention_mask=source_mask,
+                                      decoder_input_ids=decoder, decoder_attention_mask=decoder_mask,
+                                      use_cache=False).last_hidden_state
+            # Project only scored positions and bound temporary vocabulary logits.
+            for row, (index, (key, _, answer)) in enumerate(batch):
+                score, greedy = 0.0, True
+                for offset in range(0, len(answer), 256):
+                    target = torch.tensor(answer[offset:offset + 256], device=self.device)
+                    logits = self.model.lm_head(hidden[row, 1 + offset:1 + offset + len(target)]).float()
+                    score += torch.log_softmax(logits, -1).gather(1, target[:, None]).sum().item()
+                    greedy = greedy and bool((logits.argmax(-1) == target).all())
+                results[index] = (score, greedy)
+                if key is not None:
+                    self.cache_hook.add_partial("loglikelihood", key, results[index])
+        return results
+
+    def loglikelihood_rolling(self, requests, disable_tqdm=False):
+        raise NotImplementedError("UL2 continuation likelihood is not rolling autoregressive perplexity")
+
+    def generate_until(self, requests, disable_tqdm=False):
+        raise NotImplementedError("UL2 benchmark adapter currently supports likelihood tasks only")
+
+
+class DiffusionHFLM(HFLM):
+    """Continuation PLL, with other answer tokens visible.
+
+    The bool is full-token reconstruction accuracy, NOT AR exact match.
+    Generation uses one masked next-token slot at a time, not the model's
+    unconditional parallel denoiser. Both are experimental protocols.
+    """
+
+    @torch.inference_mode()
+    def _loglikelihood_tokens(self, requests, disable_tqdm=False, override_bs=None):
+        results = []
+        for cache_key, context, continuation in requests:
+            if not continuation:
+                results.append((0.0, True))
+                continue
+            if len(continuation) >= self.max_length:
+                raise ValueError("Diffusion PLL needs room for context and the entire continuation")
+            context = context[-(self.max_length - len(continuation)):]
+            tokens = context + continuation
+            base = torch.tensor([tokens], dtype=torch.long, device=self.device)
+            timestep = torch.tensor([1 / len(tokens)], device=self.device)
+            score, greedy = 0.0, True
+            for pos in range(len(context), len(tokens)):
+                masked = base.clone()
+                masked[0, pos] = self.model.config.mask_token_id
+                logits = self.model(input_ids=masked, timesteps=timestep).logits[0, pos].float()
+                score += torch.log_softmax(logits, dim=-1)[tokens[pos]].item()
+                greedy = greedy and logits.argmax().item() == tokens[pos]
+            result = (score, greedy)
+            results.append(result)
+            if cache_key is not None:
+                self.cache_hook.add_partial("loglikelihood", cache_key, result)
+        return results
+
+    def loglikelihood_rolling(self, requests, disable_tqdm=False):
+        raise NotImplementedError("Diffusion PLL must not be reported as autoregressive perplexity")
+
+    @torch.inference_mode()
+    def _model_generate(self, context, max_length, stop, **generation_kwargs):
+        if context.shape[0] != 1:
+            raise ValueError("Diffusion generation requires batch size 1")
+        if generation_kwargs.get("do_sample", False):
+            raise ValueError("Diffusion adapter implements greedy generation only")
+        sequence = context
+        prompt_length = context.shape[1]
+        while sequence.shape[1] < min(max_length, self.max_length):
+            slot = torch.full((1, 1), self.model.config.mask_token_id,
+                              dtype=torch.long, device=self.device)
+            masked = torch.cat((sequence, slot), dim=1)
+            timestep = torch.tensor([1 / masked.shape[1]], device=self.device)
+            logits = self.model(input_ids=masked, timesteps=timestep).logits[:, -1]
+            predicted = logits.argmax(dim=-1, keepdim=True)
+            sequence = torch.cat((sequence, predicted), dim=1)
+            text = self.tok_decode(sequence[0, prompt_length:].tolist())
+            if predicted.item() == self.eot_token_id or any(s and s in text for s in stop):
+                break
+        return sequence
+
+
+def load_model(plan, args):
+    common = dict(batch_size=args.batch_size, device=args.device, dtype=args.dtype,
+                  max_length=args.max_length, trust_remote_code=args.trust_remote_code,
+                  revision=args.revision)
+    tokenizer = AutoTokenizer.from_pretrained(
+        plan["tokenizer"], trust_remote_code=args.trust_remote_code,
+        revision=args.tokenizer_revision or args.revision,
+    )
+    if plan["local_model"] is None:
+        return HFLM(pretrained=plan["checkpoint"], tokenizer=tokenizer, **common)
+
+    from tiny_llm.models import MODEL_REGISTRY
+
+    entry = MODEL_REGISTRY[plan["local_model"]]
+    dtype = args.dtype if args.dtype == "auto" else getattr(torch, args.dtype)
+    # Checkpoints have no auto_map or tokenizer. Never load train_state.pt.
+    config = entry["config_class"].from_pretrained(plan["checkpoint"])
+    model = entry["model_class"].from_pretrained(plan["checkpoint"], config=config, dtype=dtype)
+    model.to(args.device).eval()
+    backend = "seq2seq" if plan["kind"] == "seq2seq" else "causal"
+    adapter = DiffusionHFLM if plan["kind"] == "diffusion" else HFLM
+    if plan["kind"] == "seq2seq" and getattr(config, "ul2", None):
+        adapter = UL2HFLM
+    return adapter(pretrained=model, tokenizer=tokenizer, backend=backend, **common)

evaluation/requirements.txt

@@ -0,0 +1,2 @@
+lm_eval==0.4.12
+tensorboard>=2.18

evaluation/results.json

@@ -0,0 +1,554 @@
+{
+  "model": "d0rj/q-51M-base",
+  "step": 15000,
+  "source_tokens": 3932160000,
+  "core_protocol": "autoregressive likelihood",
+  "core_device": "RTX 5070 Ti",
+  "core_dtype": "bfloat16",
+  "harness": "lm-eval 0.4.12",
+  "date": "2026-09-20",
+  "tasks": {
+    "hellaswag": {
+      "samples": 10042,
+      "shots": 0,
+      "dataset": "Rowan/hellaswag",
+      "config": "default",
+      "split": "validation",
+      "metrics": {
+        "acc": {
+          "value": 0.27633937462656843,
+          "stderr": 0.004462727543055927,
+          "ci95": [
+            0.2676798053353836,
+            0.28516999640463647
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        },
+        "acc_norm": {
+          "value": 0.29177454690300736,
+          "stderr": 0.004536500714147901,
+          "ci95": [
+            0.28296457902614003,
+            0.3007437626650362
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        }
+      }
+    },
+    "arc_easy": {
+      "samples": 2376,
+      "shots": 0,
+      "dataset": "allenai/ai2_arc",
+      "config": "ARC-Easy",
+      "split": "test",
+      "metrics": {
+        "acc": {
+          "value": 0.49074074074074076,
+          "stderr": 0.010258024147860667,
+          "ci95": [
+            0.47067078447639055,
+            0.5108405889653497
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        },
+        "acc_norm": {
+          "value": 0.43308080808080807,
+          "stderr": 0.010167478013701743,
+          "ci95": [
+            0.41328092357316454,
+            0.45309672993759426
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        }
+      }
+    },
+    "arc_challenge": {
+      "samples": 1172,
+      "shots": 0,
+      "dataset": "allenai/ai2_arc",
+      "config": "ARC-Challenge",
+      "split": "test",
+      "metrics": {
+        "acc": {
+          "value": 0.197098976109215,
+          "stderr": 0.011625047669880649,
+          "ci95": [
+            0.17532930720089182,
+            0.22084779246346944
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        },
+        "acc_norm": {
+          "value": 0.24232081911262798,
+          "stderr": 0.012521593295800071,
+          "ci95": [
+            0.21865689491479706,
+            0.2676684124053075
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        }
+      }
+    },
+    "piqa": {
+      "samples": 1838,
+      "shots": 0,
+      "dataset": "baber/piqa",
+      "config": "default",
+      "split": "validation",
+      "metrics": {
+        "acc": {
+          "value": 0.6115342763873776,
+          "stderr": 0.011371877593210428,
+          "ci95": [
+            0.5890412805423275,
+            0.6335620266188204
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        },
+        "acc_norm": {
+          "value": 0.5990206746463548,
+          "stderr": 0.011434766962108364,
+          "ci95": [
+            0.5764309418644076,
+            0.6211973601030507
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        }
+      }
+    },
+    "winogrande": {
+      "samples": 1267,
+      "shots": 0,
+      "dataset": "allenai/winogrande",
+      "config": "winogrande_xl",
+      "split": "validation",
+      "metrics": {
+        "acc": {
+          "value": 0.500394632991318,
+          "stderr": 0.01405248130604961,
+          "ci95": [
+            0.47290359623559514,
+            0.5278832839791312
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        }
+      }
+    },
+    "openbookqa": {
+      "samples": 500,
+      "shots": 0,
+      "dataset": "allenai/openbookqa",
+      "config": "main",
+      "split": "test",
+      "metrics": {
+        "acc": {
+          "value": 0.152,
+          "stderr": 0.016071982367911835,
+          "ci95": [
+            0.12319239267380164,
+            0.18611414828743875
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        },
+        "acc_norm": {
+          "value": 0.282,
+          "stderr": 0.020143572847290726,
+          "ci95": [
+            0.2443364332272433,
+            0.32298777921399346
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        }
+      }
+    },
+    "boolq": {
+      "samples": 3270,
+      "shots": 0,
+      "dataset": "aps/super_glue",
+      "config": "boolq",
+      "split": "validation",
+      "metrics": {
+        "acc": {
+          "value": 0.598776758409786,
+          "stderr": 0.008572708337179049,
+          "ci95": [
+            0.5818706845463824,
+            0.615451027004275
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        }
+      }
+    },
+    "lambada_openai": {
+      "samples": 5153,
+      "shots": 0,
+      "dataset": "EleutherAI/lambada_openai",
+      "config": "default",
+      "split": "test",
+      "metrics": {
+        "acc": {
+          "value": 0.20861633999611875,
+          "stderr": 0.005660825573438258,
+          "ci95": [
+            0.1977414700784266,
+            0.2199253277124155
+          ],
+          "ci_method": "Wilson score, z=1.959963984540054"
+        },
+        "perplexity": {
+          "value": 189.6865878818963,
+          "stderr": 8.94540405067453
+        }
+      }
+    }
+  },
+  "continuation_cpu_fp32": {
+    "nll": 3.560861705095088,
+    "nll_ci95": [
+      3.521807046434314,
+      3.599506885284411
+    ],
+    "ppl": 35.19351050996088,
+    "ppl_ci95": [
+      33.84553368820155,
+      36.580191764630115
+    ],
+    "blocks": 291,
+    "scored_tokens": 148992,
+    "block_nll": [
+      2.8818949460983276,
+      3.559372901916504,
+      3.4997223615646362,
+      3.5533182621002197,
+      3.904844641685486,
+      3.7352592945098877,
+      4.410188436508179,
+      3.7979665994644165,
+      3.6401803493499756,
+      3.706244707107544,
+      3.3963332176208496,
+      2.9677618741989136,
+      2.9415470361709595,
+      3.381311774253845,
+      3.761952519416809,
+      3.502508759498596,
+      3.827163577079773,
+      3.726546883583069,
+      3.598921537399292,
+      3.8391060829162598,
+      3.615709662437439,
+      3.3927921056747437,
+      3.9658381938934326,
+      2.9928760528564453,
+      3.289688467979431,
+      3.522820830345154,
+      3.514780282974243,
+      3.4957292079925537,
+      3.4607547521591187,
+      3.3068007230758667,
+      3.829126238822937,
+      4.061469554901123,
+      3.32765793800354,
+      3.3051648139953613,
+      3.4504024982452393,
+      3.987992525100708,
+      3.617965340614319,
+      3.45237398147583,
+      3.2098472118377686,
+      3.6943342685699463,
+      3.2918165922164917,
+      3.7236536741256714,
+      3.3088908195495605,
+      3.5559070110321045,
+      3.3810545206069946,
+      3.626940131187439,
+      3.8755948543548584,
+      3.4037437438964844,
+      3.8654167652130127,
+      3.5353585481643677,
+      3.423677682876587,
+      3.527420997619629,
+      3.8018072843551636,
+      4.020987868309021,
+      3.4471275806427,
+      4.067643642425537,
+      3.1591018438339233,
+      3.1642297506332397,
+      3.5973926782608032,
+      3.252170205116272,
+      3.200021743774414,
+      3.070353388786316,
+      3.260322332382202,
+      2.5585895776748657,
+      3.14900541305542,
+      3.241077184677124,
+      3.1020973920822144,
+      2.980721950531006,
+      3.3583728075027466,
+      3.282616138458252,
+      3.4956382513046265,
+      3.161021590232849,
+      3.7117087841033936,
+      3.4659478664398193,
+      3.7045702934265137,
+      3.560722827911377,
+      3.566222310066223,
+      3.6698633432388306,
+      3.529252290725708,
+      3.6588650941848755,
+      3.5175918340682983,
+      3.6442832946777344,
+      3.546086549758911,
+      3.9669642448425293,
+      3.9136862754821777,
+      4.107044577598572,
+      3.8092015981674194,
+      3.8360373973846436,
+      3.784490466117859,
+      3.9585541486740112,
+      3.542818546295166,
+      3.608612895011902,
+      3.712354898452759,
+      3.7949299812316895,
+      3.5544456243515015,
+      3.6868759393692017,
+      3.4533783197402954,
+      3.4440205097198486,
+      3.385461211204529,
+      4.006747007369995,
+      3.7272919416427612,
+      3.518463134765625,
+      2.864761471748352,
+      3.4449726343154907,
+      3.4412227869033813,
+      3.110851764678955,
+      3.3665395975112915,
+      3.103492498397827,
+      3.2595232725143433,
+      3.3449599742889404,
+      3.251083731651306,
+      3.348357915878296,
+      2.8886704444885254,
+      3.0338516235351562,
+      3.6125348806381226,
+      3.265676736831665,
+      3.208403468132019,
+      3.8483210802078247,
+      3.346174716949463,
+      3.436376690864563,
+      3.76150643825531,
+      3.8037009239196777,
+      3.6503586769104004,
+      3.6716541051864624,
+      3.9631524085998535,
+      3.73912513256073,
+      3.523168683052063,
+      3.6097822189331055,
+      3.807469367980957,
+      3.660905122756958,
+      3.8983651399612427,
+      3.7528629302978516,
+      3.823133945465088,
+      3.454434037208557,
+      3.8696566820144653,
+      3.581449508666992,
+      3.627530574798584,
+      3.8549742698669434,
+      3.481117367744446,
+      3.6938220262527466,
+      3.5973401069641113,
+      3.131093144416809,
+      4.145915865898132,
+      3.846521258354187,
+      3.906494617462158,
+      3.995554804801941,
+      4.2677435874938965,
+      4.066015124320984,
+      3.2945024967193604,
+      3.526572823524475,
+      3.762774705886841,
+      3.6110243797302246,
+      3.0586884021759033,
+      3.4789986610412598,
+      3.098861813545227,
+      3.476531982421875,
+      3.6977832317352295,
+      3.6894044876098633,
+      4.153103709220886,
+      3.3070541620254517,
+      3.663380265235901,
+      3.6190026998519897,
+      2.991021156311035,
+      2.888830542564392,
+      2.809903621673584,
+      3.1375738382339478,
+      3.0203697681427,
+      3.3764541149139404,
+      3.0283684730529785,
+      3.8419175148010254,
+      3.703011989593506,
+      3.9893038272857666,
+      3.885919213294983,
+      4.4654247760772705,
+      4.350051641464233,
+      3.898257851600647,
+      4.068492889404297,
+      3.401093602180481,
+      3.329573631286621,
+      4.586219072341919,
+      3.7894513607025146,
+      4.155319690704346,
+      4.0569844245910645,
+      4.123142957687378,
+      4.2685511112213135,
+      4.168790817260742,
+      3.9254218339920044,
+      4.219369173049927,
+      4.145432949066162,
+      4.055754661560059,
+      4.449324369430542,
+      4.017173767089844,
+      3.2817542552948,
+      4.302147626876831,
+      4.085442543029785,
+      3.6487646102905273,
+      3.7623459100723267,
+      3.4463300704956055,
+      3.7141886949539185,
+      3.693334460258484,
+      3.477225184440613,
+      3.063503861427307,
+      3.466049075126648,
+      3.2781829833984375,
+      3.5349220037460327,
+      3.6245510578155518,
+      3.305224299430847,
+      3.6646769046783447,
+      3.295943021774292,
+      3.564467668533325,
+      3.447028160095215,
+      3.268126368522644,
+      3.388076663017273,
+      3.4262938499450684,
+      3.2643288373947144,
+      3.563586711883545,
+      3.6324812173843384,
+      3.3838117122650146,
+      2.9912376403808594,
+      3.1752302646636963,
+      3.2856974601745605,
+      3.4931013584136963,
+      3.2308322191238403,
+      3.2233749628067017,
+      3.502491593360901,
+      3.303251266479492,
+      3.0688891410827637,
+      3.631072521209717,
+      3.84873628616333,
+      3.3207255601882935,
+      3.7225327491760254,
+      3.5921976566314697,
+      3.480934739112854,
+      3.669029474258423,
+      3.7562562227249146,
+      3.764595627784729,
+      3.822733163833618,
+      3.6460816860198975,
+      3.160399556159973,
+      3.726219892501831,
+      3.462749719619751,
+      3.500929832458496,
+      3.1595499515533447,
+      3.0313973426818848,
+      3.2162559032440186,
+      3.2418947219848633,
+      3.2680046558380127,
+      3.580800771713257,
+      4.050579786300659,
+      3.4574050903320312,
+      3.9652373790740967,
+      4.038872718811035,
+      3.5344061851501465,
+      3.4817302227020264,
+      3.8452590703964233,
+      3.525113582611084,
+      3.5390002727508545,
+      3.4374747276306152,
+      3.7511621713638306,
+      3.556244730949402,
+      3.4953012466430664,
+      3.5543426275253296,
+      3.1843537092208862,
+      3.2016040086746216,
+      3.458970069885254,
+      3.6360448598861694,
+      3.3213725090026855,
+      3.729948401451111,
+      3.7068639993667603,
+      2.9791568517684937,
+      3.8019036054611206,
+      3.64823317527771,
+      2.9146808385849,
+      3.088519811630249,
+      3.0383111238479614,
+      2.6207664012908936,
+      3.47467839717865,
+      3.4298750162124634,
+      3.7091829776763916,
+      3.6953502893447876,
+      3.5971832275390625,
+      3.733664035797119,
+      3.306243062019348,
+      3.5431629419326782,
+      3.622551202774048,
+      3.268967390060425,
+      3.5540285110473633,
+      3.8727654218673706,
+      4.464851140975952,
+      4.037501811981201,
+      3.530550479888916
+    ]
+  },
+  "continuation_ci_method": "percentile bootstrap of nonoverlapping 1024-token blocks; conditional suffix NLL; exponentiate interval endpoints for PPL",
+  "continuation_bootstrap_samples": 10000,
+  "continuation_bootstrap_seed": 2026,
+  "corpus_sha256": "aca2f46735043bcfd0a44eca981d04627b9cdf74c4c9a04bf0856d04066f58fc",
+  "weights_sha256": "aeab3483f72ca0f9937730614c9cabbcbfa2bdcf8327936d7e6f9794d464c1cf",
+  "source_sha256": {
+    "src/tiny_llm/__init__.py": "8c7f075d51006a53fe1d645ef5a84d65db58ff4251e7544074fe29ee7549c14c",
+    "src/tiny_llm/benchmarks.py": "5d9f461366389bd1f926c6a6468192f3e2313c87caf5b1de27f407d563a858a4",
+    "src/tiny_llm/checkpoint.py": "ac9df5f8c981a47eb1f1e59d0687ff41a1f6c38ca74ced50f71ac446a58d6ef5",
+    "src/tiny_llm/continuation.py": "548b7e76401cd545af3e7ff1aedf041fc96a330c8c9cc80266e0e20830f0e95d",
+    "src/tiny_llm/data.py": "973e4ff3e7b6e80c74bd99173f89cb8ce94c6ef3a21ff54acf1a42eab9654d89",
+    "src/tiny_llm/evaluation/__init__.py": "e49e80f7be3ab6e200f0db14b028e465de9caa8f6f705430eabce0c41b8322ba",
+    "src/tiny_llm/evaluation/adapters.py": "5a881b82bf6938ca953be751de021889f2c51a559df88f537da1e6ce4dfc6e15",
+    "src/tiny_llm/evaluation/catalog.py": "5681c9433c67890644488f7f73280e18662f28a5de8c18374443bd7c84ba4be9",
+    "src/tiny_llm/evaluation/cli.py": "b4875211ed8390c72ee96e2cb5f0abc02eb9ea7c5e4828b7c3eea7d6ad682f11",
+    "src/tiny_llm/losses.py": "592a05287369f6e5077b3faca404a9392de4772cf06389fb8aa7e36924e74f86",
+    "src/tiny_llm/models/__init__.py": "55f88bbd8cf8bfd70e5e17c1d0050299412528385c3714bb712623f40b0a681d",
+    "src/tiny_llm/models/diffusion/__init__.py": "c92eaf5986d604d8aaf1894a862a217cefa32ab25359175b88110e45a71c1b25",
+    "src/tiny_llm/models/diffusion/configuration_diffusion_lm.py": "8f796fb3bb105dfedc5d14d939ffc1c0cb3a119904fe90ca0feb76ffee9fc9a2",
+    "src/tiny_llm/models/diffusion/modeling_diffusion_lm.py": "71162bf7fbe92cbf08cf3bf64f59fc11098c5c302787fc1c7fe9fcb14c9479e5",
+    "src/tiny_llm/models/looped/__init__.py": "4b46b426021c59ed313cb2f49afc1926ad2793d771ef90117e3fef0462957cff",
+    "src/tiny_llm/models/looped/configuration_looped_lm.py": "b8d99c52b98e0a97a4cae76ca4c20affd8f4d019ec38e0a9beefd5fad864dc25",
+    "src/tiny_llm/models/looped/modeling_looped_lm.py": "9bdb3c5e3aaae663d09e98055f95df1adbc42e0a4bf5dffb7f6c39242af41b9e",
+    "src/tiny_llm/models/moe_t5/__init__.py": "18e33fe38b2049521f573512ac8796afa209e1b25c2052784e1c927a7160d66c",
+    "src/tiny_llm/models/moe_t5/configuration_alicet5_moe.py": "4647b1e9ab9ce07887bffc1abc3beb3e691962a380b4cc237aae376490bb56b7",
+    "src/tiny_llm/models/moe_t5/modeling_alicet5_moe.py": "c7bbba2f01dc889cc604c7bc161354397239b28503d05f7e00665658be84751d",
+    "src/tiny_llm/models/prefixlm/__init__.py": "c94c3c861acba8d2147a3a6c46bccfc71ac8adc8d70757bede009598029d31ea",
+    "src/tiny_llm/models/prefixlm/configuration_prefix_lm.py": "20b3a62ae3628b5a64566ef19d2947c90f10c6159daad00dd13d18fc5910c42c",
+    "src/tiny_llm/models/prefixlm/modeling_prefix_lm.py": "f62a38211df175696f64945173d1e41d9f7fcf6152a7648020da8c91788e98fd",
+    "src/tiny_llm/models/q50m/__init__.py": "9d0dfec6fa0cfa6967e2f17a73d8fbeca69ba93b4a93651eda4d95626d5f5eb5",
+    "src/tiny_llm/models/q50m/configuration_q50m.py": "796806b2d5d982b633796bc7de53c3073b1b6e52d259f44ef798b63f55eaea92",
+    "src/tiny_llm/models/q50m/modeling_q50m.py": "b7a712050167fad24471193b0ca7baf4ca7e5ebdf8036c91b506d3fff5b49c8b",
+    "src/tiny_llm/telemetry.py": "f5a35de68da1a2304159fd3613fa3773bc368982ed18712791c847285b895c09",
+    "src/tiny_llm/ul2.py": "f4d948d89d5e106965612971dd473a54b57719f811761f18457c6221d763b876"
+  },
+  "packages": {
+    "lm_eval": "0.4.12",
+    "torch": "2.11.0+cu128",
+    "transformers": "5.17.0",
+    "datasets": "5.0.1",
+    "accelerate": "1.15.0"
+  }
+}

evaluation/run_core.py

@@ -0,0 +1,36 @@
+"""Reproduce the released model's eight-task likelihood evaluation."""
+import argparse
+import json
+import importlib.metadata
+from pathlib import Path
+import torch
+from transformers import AutoTokenizer,AutoModelForCausalLM,AutoModelForSeq2SeqLM
+from lm_eval import evaluator,tasks
+from lm_eval.models.huggingface import HFLM
+from adapters import UL2HFLM
+
+def main():
+    if importlib.metadata.version('lm_eval') != '0.4.12':
+        raise RuntimeError('This reproduction protocol requires lm_eval==0.4.12')
+    p=argparse.ArgumentParser(description=__doc__)
+    p.add_argument('--device',default='cpu')
+    p.add_argument('--dtype',default='float32',choices=['float32','bfloat16'])
+    p.add_argument('--batch-size',type=int,default=1)
+    p.add_argument('--output',type=Path,required=True)
+    p.add_argument('--limit',type=int,help='Smoke only; not a full benchmark')
+    a=p.parse_args(); a.output.mkdir(parents=True,exist_ok=False)
+    torch.set_num_threads(4)
+    root=Path(__file__).resolve().parents[1]
+    config=json.loads((root/'config.json').read_text()); ul2=bool(config.get('ul2'))
+    cls=AutoModelForSeq2SeqLM if ul2 else AutoModelForCausalLM
+    model=cls.from_pretrained(root,trust_remote_code=True,dtype=getattr(torch,a.dtype)).to(a.device).eval()
+    tok=AutoTokenizer.from_pretrained(root)
+    adapter=(UL2HFLM if ul2 else HFLM)(pretrained=model,tokenizer=tok,backend='seq2seq' if ul2 else 'causal',device=a.device,batch_size=a.batch_size,max_length=2048)
+    mapping={'hellaswag':'hellaswag','arc_easy':'arc','arc_challenge':'arc','piqa':'piqa','winogrande':'winogrande','openbookqa':'openbookqa','boolq':'super_glue/boolq','lambada_openai':'lambada'}
+    manager=tasks.TaskManager(include_defaults=False,include_path=sorted({Path(tasks.__file__).parent/v for v in mapping.values()}))
+    for name in mapping:
+        result=evaluator.simple_evaluate(model=adapter,tasks=[name],num_fewshot=0,limit=a.limit,bootstrap_iters=1000,log_samples=False,task_manager=manager,random_seed=1234,numpy_random_seed=1234,torch_random_seed=1234,fewshot_random_seed=1234,apply_chat_template=False)
+        def fallback(x): return x.item() if hasattr(x,'item') else str(x)
+        (a.output/(name+'.json')).write_text(json.dumps(result,indent=2,default=fallback))
+
+if __name__=='__main__': main()

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "do_sample": false
+}

model.safetensors

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

modeling_q50m.py

@@ -0,0 +1,445 @@
+"""Modeling for the Q-50M architecture (replica of q-project/Q-50M-Base).
+
+Self-contained implementation of the exact architecture described in the
+model card: Mistral-shaped 512-wide 10-layer decoder, GQA 8q/2kv, per-head
+QK-Norm, RoPE in every layer (configurable NoPE), scalar content-dependent
+gates (gate = multiplier * sigmoid(Wx)) on attention and MLP residual
+branches, SwiGLU MLP, tied input/output embeddings, final RMSNorm.
+
+Follows the transformers v5 modeling pattern (DynamicCache, create_causal_mask,
+GradientCheckpointingLayer, GenerationMixin).
+"""
+
+from collections.abc import Callable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs
+
+from .configuration_q50m import Q50MConfig
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_pos_emb(q, k, cos, sin):
+    # cos/sin: (batch, seq, head_dim)
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin
+
+
+class Q50MRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states.to(input_dtype)).to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Q50MRotaryEmbedding(nn.Module):
+    def __init__(self, config: Q50MConfig, device=None):
+        super().__init__()
+        self.config = config
+        inv_freq = 1.0 / (
+            config.rope_theta
+            ** (torch.arange(0, config.head_dim, 2, dtype=torch.float32, device=device) / config.head_dim)
+        )
+        self.inv_freq = nn.Buffer(inv_freq, persistent=True)
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        inv_freq_expanded = (
+            self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1)
+            .to(dtype=torch.float32, device=x.device)
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos().to(dtype=x.dtype), emb.sin().to(dtype=x.dtype)
+
+
+class Q50MScalarGate(nn.Module):
+    """Content-dependent scalar gate: output = branch * multiplier * sigmoid(W x)."""
+
+    def __init__(self, hidden_size: int, multiplier: float = 2.0):
+        super().__init__()
+        self.projection = nn.Linear(hidden_size, 1, bias=False)
+        self.multiplier = multiplier
+
+    def forward(self, branch: torch.Tensor, residual_input: torch.Tensor) -> torch.Tensor:
+        return branch * self.multiplier * torch.sigmoid(self.projection(residual_input))
+
+    def reset_parameters(self):
+        nn.init.zeros_(self.projection.weight)
+
+
+def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def _eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = _repeat_kv(key, module.num_key_value_groups)
+    value_states = _repeat_kv(value, module.num_key_value_groups)
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, attn_weights
+
+
+def _sdpa_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    """Local SDPA wrapper. HF's generic SDPA interface falls back to the math
+    backend for GQA shapes (q heads != kv heads), which materializes
+    (batch, heads, seq, seq) fp32 attention scores. Expanding kv first keeps
+    the fused flash / memory-efficient kernels eligible."""
+    n_rep = getattr(module, "num_key_value_groups", 1)
+    if n_rep > 1:
+        key = _repeat_kv(key, n_rep)
+        value = _repeat_kv(value, n_rep)
+    is_causal = (
+        attention_mask is None
+        and getattr(module, "is_causal", False)
+        and query.shape[2] > 1
+    )
+    attn_output = nn.functional.scaled_dot_product_attention(
+        query,
+        key,
+        value,
+        attn_mask=attention_mask,
+        dropout_p=dropout,
+        is_causal=is_causal,
+        scale=scaling,
+    )
+    return attn_output.transpose(1, 2).contiguous(), None
+
+
+class Q50MAttention(nn.Module):
+    def __init__(self, config: Q50MConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = config.head_dim
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_kv_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        self.use_rope = layer_idx not in config.nope_layers
+
+        inner = self.num_heads * self.head_dim
+        kv_inner = self.num_kv_heads * self.head_dim
+        self.q_proj = nn.Linear(config.hidden_size, inner, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, kv_inner, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, kv_inner, bias=False)
+        self.o_proj = nn.Linear(inner, config.hidden_size, bias=False)
+        self.q_norm = Q50MRMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else None
+        self.k_norm = Q50MRMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else None
+        self.output_gate = (
+            Q50MScalarGate(config.hidden_size, config.gate_multiplier)
+            if config.attention_scalar_gate
+            else None
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        if self.q_norm is not None:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        if self.use_rope:
+            cos, sin = position_embeddings
+            q, k = _apply_rotary_pos_emb(q, k, cos, sin)
+
+        if past_key_values is not None:
+            k, v = past_key_values.update(k, v, self.layer_idx)
+
+        if self.config._attn_implementation == "sdpa":
+            attention_interface: Callable = _sdpa_attention_forward
+        else:
+            attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+                self.config._attn_implementation, _eager_attention_forward
+            )
+        attn_output, attn_weights = attention_interface(
+            self,
+            q, k, v,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        if self.output_gate is not None:
+            attn_output = self.output_gate(attn_output, hidden_states)
+        return attn_output, attn_weights
+
+
+class Q50MMLP(nn.Module):
+    def __init__(self, config: Q50MConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.output_gate = (
+            Q50MScalarGate(config.hidden_size, config.gate_multiplier)
+            if config.mlp_scalar_gate
+            else None
+        )
+
+    def forward(self, hidden_states):
+        output = self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+        if self.output_gate is not None:
+            output = self.output_gate(output, hidden_states)
+        return output
+
+
+class Q50MDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Q50MConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = Q50MAttention(config, layer_idx)
+        self.mlp = Q50MMLP(config)
+        self.input_layernorm = Q50MRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Q50MRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class Q50MPreTrainedModel(PreTrainedModel):
+    config_class = Q50MConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Q50MDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_sdpa = True
+    _supports_flash_attn = False
+    _supports_flex_attn = False
+
+
+class Q50MModel(Q50MPreTrainedModel):
+    def __init__(self, config: Q50MConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Q50MDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = Q50MRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Q50MRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        standard_positions = position_ids is None
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.unsqueeze(0)
+
+        # Unpadded, uncached SDPA can use its native causal flag, including
+        # under torch.compile where HF otherwise materializes a dense mask.
+        causal_mask = None if (
+            attention_mask is None and past_key_values is None and standard_positions
+            and self.config._attn_implementation == "sdpa"
+        ) else create_causal_mask(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if self.gradient_checkpointing and self.training:
+                hidden_states = self._gradient_checkpointing_func(
+                    decoder_layer.forward,
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    use_cache,
+                    position_embeddings,
+                )
+            else:
+                hidden_states = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    position_embeddings=position_embeddings,
+                    position_ids=position_ids,
+                    past_key_values=past_key_values,
+                    use_cache=use_cache,
+                    **kwargs,
+                )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+class Q50MForCausalLM(Q50MPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: Q50MConfig):
+        super().__init__(config)
+        self.model = Q50MModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+        )

release-manifest.json

@@ -0,0 +1,43 @@
+{
+  "repo": "d0rj/q-51M-base",
+  "files_sha256": {
+    "README.md": "70275eaece2a4b14b996e9e53fdfad2ac6045c24f2bca89f4891ba240ec76cc6",
+    "config.json": "c5c22bdc342dc8ba2dfadfad98733a697892ff696bc4a6a49fe07efc0096f9ad",
+    "configuration_q50m.py": "e1fcd06df2efb41474bb18d3954f54c04dc7733350c600afced9ab9dcf61d5c7",
+    "evaluation/adapters.py": "5a881b82bf6938ca953be751de021889f2c51a559df88f537da1e6ce4dfc6e15",
+    "evaluation/requirements.txt": "953e2117755df78f80177d62b34ff1d7a297efb5edc87a442864385424f416d3",
+    "evaluation/results.json": "662314441ea584474f234bf20ad0c37100291ab4fb6ba80176b8da0a71159799",
+    "evaluation/run_core.py": "217376c792abb0cbd4729580c78725388b83e5df0d126dcfe7e198568f6ed28f",
+    "generation_config.json": "baaf419957543384a7ed26a4d70753267fa60aea90bac2c55bb2232d70f9602a",
+    "model.safetensors": "aeab3483f72ca0f9937730614c9cabbcbfa2bdcf8327936d7e6f9794d464c1cf",
+    "modeling_q50m.py": "d8f48321584187276d991a8ef5792d3048029af5a319e16c3fb6800587546431",
+    "requirements.txt": "ac09bc9a52f1464d18f19032046f29b2faefc89c5716deca900ead1ee8aa572e",
+    "tensorboard/events.out.tfevents.1789820315.DESKTOP-OU3A33R.3028.0": "d65557d13fbf25ec08e48f2d845ea6c4cb38129c738cd0b64208019b7223cfca",
+    "tensorboard/events.out.tfevents.1789905190.DESKTOP-OU3A33R.12878.0.eval-core-20260920": "9f517d66584b507239662826fc95fe04d18f6ced59d3c7359d08968b6dfb7a3f",
+    "tokenizer.json": "6155a4212832ddd656acd1cd45ab4acb327dbb080e3acdbec84f086adf17414e",
+    "tokenizer_config.json": "cccfc292e0846e8ab8aa04dafae32483a831bcf1df6f15100f82d306b4a9f4c6",
+    "training_config.json": "39a3b2780f731260fe969aa0c458477f2aede7552bac925f69398fcaa7e405b1"
+  },
+  "validation": {
+    "logits_bitwise_equal": true,
+    "cached_generation_matches_uncached": true,
+    "parameters": 50878208,
+    "generated_ids": [
+      [
+        463,
+        2934,
+        283,
+        4226,
+        314,
+        265,
+        2046,
+        283,
+        6270,
+        15,
+        493,
+        314,
+        265
+      ]
+    ]
+  }
+}

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.11.0
+transformers==5.17.0
+tokenizers>=0.22
+safetensors>=0.4

tensorboard/events.out.tfevents.1789820315.DESKTOP-OU3A33R.3028.0

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

tensorboard/events.out.tfevents.1789905190.DESKTOP-OU3A33R.12878.0.eval-core-20260920

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

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "is_local": true,
+  "model_max_length": 2048,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}

training_config.json

@@ -0,0 +1,35 @@
+{
+  "step": 15000,
+  "source_tokens": 3932160000,
+  "args": {
+    "model": "q50m",
+    "max_steps": 15000,
+    "batch_size": 32,
+    "grad_accum": 4,
+    "seq_len": 2048,
+    "lr": 0.001,
+    "min_lr_frac": 0.1,
+    "weight_decay": 0.1,
+    "warmup_frac": 0.01,
+    "betas": [
+      0.9,
+      0.95
+    ],
+    "grad_clip": 1.0,
+    "dtype": "bf16",
+    "device": "auto",
+    "seed": 2026,
+    "text_file": null,
+    "dataset_name": "parquet",
+    "log_every": 20,
+    "save_steps": 100,
+    "max_host_memory_gib": 12.0,
+    "shuffle_buffer": 100000,
+    "compile": true,
+    "ce_chunk": 8192,
+    "grad_ckpt": true,
+    "resume": true,
+    "dataset": "HuggingFaceFW/fineweb-edu",
+    "dataset_subset": "sample-10BT"
+  }
+}