BarunLM-35M: 35M-parameter base language model trained on 5.7B tokens

.gitattributes

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

CITATION.cff

@@ -0,0 +1,12 @@
+cff-version: 1.2.0
+message: "If you use BarunLM-35M, please cite this release."
+title: "BarunLM-35M"
+type: software
+version: 1.0.0
+date-released: 2026-07-31
+authors:
+  - family-names: Singh
+    given-names: Harshal
+repository-code: "https://github.com/harrrshall/barunlm-35m"
+url: "https://huggingface.co/harrrshall/BarunLM-35M"
+license: Apache-2.0

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 Harshal Singh
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,26 @@
+BarunLM-35M
+Copyright 2026 Harshal Singh
+
+The BarunLM source code and released model weights are licensed under the
+Apache License 2.0. Training datasets are not redistributed with this release
+and retain their own licenses and terms:
+
+- HuggingFaceTB/smollm-corpus (FineWeb-Edu and Cosmopedia v2), revision
+  3ba9d605774198c5868892d7a8deda78031a781f: ODC-By 1.0. FineWeb-derived
+  material is also subject to the applicable Common Crawl terms.
+- HuggingFaceTB/finemath (FineMath-4+), revision
+  e92b25a616738fe95dc186b64dfb19f9c8525594: ODC-By 1.0 and Common Crawl
+  terms of use.
+- mlfoundations/dclm-baseline-1.0-parquet, revision
+  817d6752765f6a41261085171dd546b104f60626: CC BY 4.0. The dataset card
+  describes DCLM-Baseline as a research baseline.
+- code-search-net/code_search_net (Python), revision
+  bd0cf261e357a3eb5c8fba490d23ec1a1cd59555: source repositories retain
+  their individual licenses.
+- codeparrot/codeparrot-clean, revision
+  35a59fb025bc0a102f7d96eac09d145b896d487b: source files retain the
+  repository-level licenses recorded by the dataset.
+
+Users are responsible for reviewing the upstream dataset cards and terms for
+their intended use. This notice is informational and does not replace those
+terms.

README.md

@@ -0,0 +1,205 @@
+---
+license: apache-2.0
+language:
+  - en
+pipeline_tag: text-generation
+library_name: pytorch
+tags:
+  - base-model
+  - causal-lm
+  - small-language-model
+  - text-generation
+datasets:
+  - HuggingFaceTB/smollm-corpus
+  - HuggingFaceTB/finemath
+  - mlfoundations/dclm-baseline-1.0-parquet
+  - code-search-net/code_search_net
+  - codeparrot/codeparrot-clean
+---
+
+# BarunLM-35M
+
+BarunLM-35M is a compact, decoder-only base language model built for strong
+capability per parameter. With exactly **35,072,768 parameters**, it reaches
+**41.01%** on a fixed, decontaminated nine-task zero-shot suite and exceeds
+[LFM2.5-230M-Base](https://huggingface.co/LiquidAI/LFM2.5-230M-Base) while
+using **6.55× fewer parameters**.
+
+[Download the model](https://huggingface.co/harrrshall/BarunLM-35M) ·
+[Inspect benchmark data](benchmark_results.json) · [License](LICENSE)
+
+## Highlights
+
+- **Parameter-efficient hybrid attention.** Three local-attention layers feed
+  each full-attention layer, combining a 256-token working window with periodic
+  global information exchange.
+- **Selective residual routing.** A learned convex selector every four layers
+  chooses between the block-group input and its transformed state, providing a
+  lightweight path for preserving useful representations.
+- **Stable small-model optimization.** Grouped-query attention, partial RoPE,
+  QK normalization, gated attention outputs, and bounded SwiGLU are integrated
+  into a single 35M-parameter design.
+- **Capacity-aligned pretraining.** The model was trained on 5.70B tokens—about
+  162.5 tokens per parameter—from a curated blend of educational web text,
+  synthetic exposition, mathematics, general web text, and code.
+- **Contamination-aware evaluation.** Reported scores exclude 1,854 samples
+  identified by a correctness-blind exact 13-token scan over the complete
+  training history.
+
+## Results
+
+All rows below were evaluated zero-shot with LM Evaluation Harness 0.4.12 on
+ARC-Challenge, ARC-Easy, BoolQ, HellaSwag, LAMBADA OpenAI, OpenBookQA, PIQA,
+SciQ, and WinoGrande. The reported value is the unweighted macro average after
+applying the same frozen decontamination decisions to every model.
+
+| Model | Loaded parameters | Macro accuracy | BarunLM lead |
+|---|---:|---:|---:|
+| **BarunLM-35M** | **35.1M** | **41.01%** | — |
+| [LFM2.5-230M-Base](https://huggingface.co/LiquidAI/LFM2.5-230M-Base) | 229.7M | 39.20% | **+1.81 pp** |
+| [Pythia-160M](https://huggingface.co/EleutherAI/pythia-160m-deduped) | 162.3M | 37.35% | **+3.66 pp** |
+| [Stentor-30M](https://huggingface.co/StentorLabs/Stentor-30M) | 30.4M | 36.46% | **+4.55 pp** |
+| [TinyStories-33M](https://huggingface.co/roneneldan/TinyStories-33M)¹ | 68.5M | 33.16% | **+7.85 pp** |
+| [Pythia-70M](https://huggingface.co/EleutherAI/pythia-70m-deduped) | 70.4M | 31.71% | **+9.30 pp** |
+
+The paired 10,000-resample bootstrap interval for the BarunLM minus LFM2.5
+macro difference is **[+0.92, +2.71] percentage points**. Exact revisions,
+task scores, confidence intervals, sample counts, and evidence hashes are in
+[`benchmark_results.json`](benchmark_results.json).
+
+¹ TinyStories is included as a narrow-domain diagnostic rather than a
+general-purpose peer. Parameter counts are computed from the loaded models,
+not inferred from repository names.
+
+These results establish parameter efficiency on this evaluation suite; they
+do not imply universal superiority across tasks or deployment settings.
+
+## Architecture
+
+| Component | Configuration |
+|---|---|
+| Parameters | 35,072,768 |
+| Layers / width | 12 / 448 |
+| Attention | 7 query heads, 1 key/value head |
+| Attention rhythm | 3 local layers, then 1 full layer |
+| Local window | 256 tokens |
+| Position encoding | 50% partial RoPE |
+| Feed-forward width | 1,228 |
+| Residual selection | Every 4 layers |
+| Vocabulary | 16,384 byte-level BPE tokens |
+| Context length | 2,048 tokens |
+| Embeddings | Input/output weights tied |
+
+The design treats global attention as a periodic communication layer rather
+than a cost paid at every depth. Local layers concentrate computation on nearby
+structure, while the global layer propagates information across the sequence.
+The residual selector then gives each four-layer group a learned preservation
+path with negligible parameter overhead.
+
+## Quick start
+
+BarunLM uses a small native PyTorch implementation. It is a completion model,
+not a chat model, and does not require `trust_remote_code`.
+
+For the shortest path from this model repository:
+
+```bash
+hf download harrrshall/BarunLM-35M --local-dir BarunLM-35M
+cd BarunLM-35M
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+python generate.py \
+  --prompt "The future of efficient language models is" \
+  --max-new-tokens 48 \
+  --temperature 0.8
+```
+
+The maintained source package and tests are available in the
+[GitHub repository](https://github.com/harrrshall/barunlm-35m):
+
+```bash
+git clone https://github.com/harrrshall/barunlm-35m.git
+cd barunlm-35m
+python -m venv .venv
+source .venv/bin/activate
+pip install -e .
+python examples/generate.py \
+  --prompt "The future of efficient language models is" \
+  --max-new-tokens 48 \
+  --temperature 0.8
+```
+
+Both examples verify the three required model artifacts and select CUDA
+automatically when available. For deterministic greedy decoding, pass
+`--temperature 0`.
+
+To download the release without running it:
+
+```bash
+hf download harrrshall/BarunLM-35M --local-dir BarunLM-35M
+sha256sum -c BarunLM-35M/SHA256SUMS
+```
+
+## Training
+
+BarunLM-35M was pretrained on **5,699,985,408 realized tokens** at sequence
+length 2,048. The final 4B-token continuation used Muon with a peak learning
+rate of `1e-4`, weight decay `0.1`, batch size 48, and 40,690 optimizer steps
+on one NVIDIA H200.
+
+The complete corpus combines the following pinned public sources:
+
+- [FineWeb-Edu and Cosmopedia v2](https://huggingface.co/datasets/HuggingFaceTB/smollm-corpus)
+- [FineMath-4+](https://huggingface.co/datasets/HuggingFaceTB/finemath)
+- [DCLM-Baseline](https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0-parquet)
+- [CodeSearchNet Python](https://huggingface.co/datasets/code-search-net/code_search_net)
+- [CodeParrot Clean](https://huggingface.co/datasets/codeparrot/codeparrot-clean)
+
+Documents were deduplicated across stages before admission. Source revisions
+and upstream licensing notices are recorded in [`NOTICE`](NOTICE); no training
+data is redistributed in this repository or the model release.
+
+## Release integrity
+
+| Artifact | SHA-256 |
+|---|---|
+| `model.safetensors` | `f2a7c88b9f2c2e3584809081407ab136795d82e30e89b730e007781c45d01447` |
+| `barun_config.json` | `9b3a1d71baa95a198744d250f9629231738d942570b8685c44307fd83dd33565` |
+| `tokenizer.json` | `70ded9605fccd09c2340ca7e225361eab0ae8b4dbbb0d6e26343ab5183979db6` |
+
+The checkpoint contains 35,072,768 unique trainable parameters. All stored
+floating-point tensors passed an independent finiteness and shape audit.
+
+## Intended use and limitations
+
+BarunLM-35M is intended for research on compact language models, controlled
+text-generation experiments, education, and local prototyping.
+
+- It is a base model, not instruction-tuned, and should be prompted as a text
+  continuation model.
+- Its 2,048-token context and 35M-parameter capacity limit factual recall,
+  multi-step reasoning, instruction following, and long-context synthesis.
+- Training and evaluation are English-centric.
+- Outputs may be inaccurate, biased, unsafe, repetitive, or fabricated.
+- The model has not undergone a comprehensive safety evaluation and should not
+  be used for medical, legal, financial, or other high-stakes decisions.
+
+Users should evaluate the model in their own domain and apply appropriate
+guardrails before deployment.
+
+## Acknowledgements
+
+BarunLM builds on ideas developed across the open language-model ecosystem,
+including [grouped-query attention](https://arxiv.org/abs/2305.13245),
+[rotary position embeddings](https://arxiv.org/abs/2104.09864), and
+[SwiGLU](https://arxiv.org/abs/2002.05202). We thank the teams behind PyTorch,
+Hugging Face, LM Evaluation Harness, FineWeb-Edu, Cosmopedia, FineMath, DCLM,
+CodeSearchNet, and CodeParrot for the open infrastructure and data that made
+this release possible.
+
+## License
+
+The source code and released weights are available under the
+[Apache License 2.0](LICENSE). Upstream datasets retain their own licenses and
+terms as described in [`NOTICE`](NOTICE).

SHA256SUMS

@@ -0,0 +1,3 @@
+f2a7c88b9f2c2e3584809081407ab136795d82e30e89b730e007781c45d01447  model.safetensors
+9b3a1d71baa95a198744d250f9629231738d942570b8685c44307fd83dd33565  barun_config.json
+70ded9605fccd09c2340ca7e225361eab0ae8b4dbbb0d6e26343ab5183979db6  tokenizer.json

barun_config.json

@@ -0,0 +1,22 @@
+{
+  "activation_clip": 10.0,
+  "attention_gate": true,
+  "dim": 448,
+  "dropout": 0.0,
+  "ffn_dim": 1228,
+  "full_attention_every": 4,
+  "local_window": 256,
+  "max_seq_len": 2048,
+  "mtp_loss_weight": 0.0,
+  "mtp_offset": 2,
+  "n_heads": 7,
+  "n_kv_heads": 1,
+  "n_layers": 12,
+  "norm_eps": 1e-06,
+  "qk_norm": true,
+  "residual_select_every": 4,
+  "rope_fraction": 0.5,
+  "rope_theta": 10000.0,
+  "tie_embeddings": true,
+  "vocab_size": 16384
+}

barunlm/__init__.py

@@ -0,0 +1,4 @@
+from .config import BarunConfig
+from .model import BarunLM
+
+__all__ = ["BarunConfig", "BarunLM"]

barunlm/config.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+
+@dataclass(slots=True)
+class BarunConfig:
+    vocab_size: int = 16_384
+    dim: int = 448
+    n_layers: int = 12
+    n_heads: int = 7
+    n_kv_heads: int = 1
+    ffn_dim: int = 1_228
+    max_seq_len: int = 2_048
+    rope_theta: float = 10_000.0
+    rope_fraction: float = 0.5
+    local_window: int = 256
+    full_attention_every: int = 4
+    attention_gate: bool = True
+    qk_norm: bool = True
+    residual_select_every: int = 4
+    activation_clip: float = 10.0
+    dropout: float = 0.0
+    norm_eps: float = 1e-6
+    tie_embeddings: bool = True
+    mtp_offset: int = 2
+    mtp_loss_weight: float = 0.0
+
+    def __post_init__(self) -> None:
+        if self.dim % self.n_heads:
+            raise ValueError("dim must be divisible by n_heads")
+        if self.n_heads % self.n_kv_heads:
+            raise ValueError("n_heads must be divisible by n_kv_heads")
+        head_dim = self.dim // self.n_heads
+        rope_dim = int(head_dim * self.rope_fraction)
+        if rope_dim < 2 or rope_dim % 2:
+            raise ValueError("rope_fraction must yield a positive, even rotary dimension")
+        if not 0 <= self.dropout < 1:
+            raise ValueError("dropout must be in [0, 1)")
+        if self.local_window < 1:
+            raise ValueError("local_window must be positive")
+        if self.full_attention_every < 1:
+            raise ValueError("full_attention_every must be positive")
+        if self.residual_select_every < 0:
+            raise ValueError("residual_select_every cannot be negative")
+        if self.mtp_offset < 1:
+            raise ValueError("mtp_offset must be positive")
+
+    @property
+    def head_dim(self) -> int:
+        return self.dim // self.n_heads
+
+    @property
+    def rope_dim(self) -> int:
+        value = int(self.head_dim * self.rope_fraction)
+        return value - value % 2
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> BarunConfig:
+        with Path(path).open() as handle:
+            return cls(**json.load(handle))
+
+    def save_json(self, path: str | Path) -> None:
+        target = Path(path)
+        target.parent.mkdir(parents=True, exist_ok=True)
+        with target.open("w") as handle:
+            json.dump(asdict(self), handle, indent=2, sort_keys=True)
+            handle.write("\n")

barunlm/model.py

@@ -0,0 +1,430 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import lru_cache
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from .config import BarunConfig
+
+try:
+    from flash_attn import flash_attn_func
+except ImportError:  # pragma: no cover - exercised on GPU installations only
+    flash_attn_func = None
+
+try:
+    from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+except ImportError:  # pragma: no cover - compatibility fallback for older PyTorch
+    create_block_mask = None
+    flex_attention = None
+
+compiled_flex_attention = (
+    torch.compile(flex_attention, dynamic=False) if flex_attention is not None else None
+)
+
+
+@dataclass
+class BarunOutput:
+    logits: Tensor
+    loss: Tensor | None = None
+    causal_loss: Tensor | None = None
+    mtp_loss: Tensor | None = None
+    past_key_values: list[tuple[Tensor, Tensor]] | None = None
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6, affine: bool = True) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim)) if affine else None
+
+    def forward(self, x: Tensor) -> Tensor:
+        dtype = x.dtype
+        y = x.float() * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + self.eps)
+        if self.weight is not None:
+            y = y * self.weight.float()
+        return y.to(dtype)
+
+
+def _rotate_half(x: Tensor) -> Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class PartialRotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, max_seq_len: int, theta: float) -> None:
+        super().__init__()
+        inv_freq = theta ** (-torch.arange(0, dim, 2, dtype=torch.float32) / dim)
+        positions = torch.arange(max_seq_len, dtype=torch.float32)
+        freqs = torch.outer(positions, inv_freq)
+        angles = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos", angles.cos()[None, None, :, :], persistent=False)
+        self.register_buffer("sin", angles.sin()[None, None, :, :], persistent=False)
+        self.dim = dim
+
+    def forward(self, q: Tensor, k: Tensor, position_offset: int = 0) -> tuple[Tensor, Tensor]:
+        seq_len = q.shape[-2]
+        end = position_offset + seq_len
+        if end > self.cos.shape[-2]:
+            raise ValueError(
+                f"rotary position {end} exceeds configured maximum {self.cos.shape[-2]}"
+            )
+        cos = self.cos[:, :, position_offset:end].to(device=q.device, dtype=q.dtype)
+        sin = self.sin[:, :, position_offset:end].to(device=q.device, dtype=q.dtype)
+        q_rot, q_pass = q[..., : self.dim], q[..., self.dim :]
+        k_rot, k_pass = k[..., : self.dim], k[..., self.dim :]
+        q = torch.cat((q_rot * cos + _rotate_half(q_rot) * sin, q_pass), dim=-1)
+        k = torch.cat((k_rot * cos + _rotate_half(k_rot) * sin, k_pass), dim=-1)
+        return q, k
+
+
+@lru_cache(maxsize=32)
+def _local_causal_mask(seq_len: int, window: int, device_type: str) -> Tensor:
+    device = torch.device(device_type)
+    row = torch.arange(seq_len, device=device)[:, None]
+    col = torch.arange(seq_len, device=device)[None, :]
+    allowed = (col <= row) & (col > row - window)
+    mask = torch.zeros((seq_len, seq_len), device=device, dtype=torch.float32)
+    return mask.masked_fill(~allowed, float("-inf"))
+
+
+@lru_cache(maxsize=32)
+def _local_bidirectional_mask(seq_len: int, window: int, device_type: str) -> Tensor:
+    """Window constraint for custom masks such as bidirectional PrefixLM prompts."""
+    device = torch.device(device_type)
+    row = torch.arange(seq_len, device=device)[:, None]
+    col = torch.arange(seq_len, device=device)[None, :]
+    allowed = (col - row).abs() < window
+    mask = torch.zeros((seq_len, seq_len), device=device, dtype=torch.float32)
+    return mask.masked_fill(~allowed, float("-inf"))
+
+
+@lru_cache(maxsize=32)
+def _local_block_mask(seq_len: int, window: int, device: str):
+    if create_block_mask is None:
+        return None
+
+    def causal_window(batch, head, query_index, key_index):
+        del batch, head
+        return (query_index >= key_index) & (query_index - key_index < window)
+
+    return create_block_mask(
+        causal_window,
+        B=None,
+        H=None,
+        Q_LEN=seq_len,
+        KV_LEN=seq_len,
+        device=device,
+        _compile=True,
+    )
+
+
+class GroupedAttention(nn.Module):
+    def __init__(self, config: BarunConfig, *, is_full: bool) -> None:
+        super().__init__()
+        self.config = config
+        self.is_full = is_full
+        self.q_proj = nn.Linear(config.dim, config.n_heads * config.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.dim, config.n_kv_heads * config.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.dim, config.n_kv_heads * config.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.n_heads * config.head_dim, config.dim, bias=False)
+        self.g_proj = (
+            nn.Linear(config.dim, config.n_heads * config.head_dim, bias=False)
+            if config.attention_gate
+            else None
+        )
+        self.q_norm = RMSNorm(config.head_dim, config.norm_eps) if config.qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(config.head_dim, config.norm_eps) if config.qk_norm else nn.Identity()
+        self.rope = PartialRotaryEmbedding(config.rope_dim, config.max_seq_len, config.rope_theta)
+
+    def _torch_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Tensor | None) -> Tensor:
+        use_flex = (
+            not self.is_full
+            and mask is None
+            and compiled_flex_attention is not None
+            and q.is_cuda
+            and self.config.dropout == 0
+        )
+        if use_flex:
+            block_mask = _local_block_mask(q.shape[-2], self.config.local_window, str(q.device))
+            return compiled_flex_attention(q, k, v, block_mask=block_mask, enable_gqa=True)
+        repeat = self.config.n_heads // self.config.n_kv_heads
+        if repeat > 1:
+            k = k.repeat_interleave(repeat, dim=1)
+            v = v.repeat_interleave(repeat, dim=1)
+        if mask is not None:
+            mask = mask.to(dtype=q.dtype)
+            if not self.is_full:
+                local_mask = _local_bidirectional_mask(
+                    q.shape[-2], self.config.local_window, q.device.type
+                ).to(dtype=q.dtype)
+                mask = mask + local_mask
+            return F.scaled_dot_product_attention(
+                q, k, v, attn_mask=mask, dropout_p=self.config.dropout if self.training else 0.0
+            )
+        if self.is_full:
+            return F.scaled_dot_product_attention(
+                q, k, v, is_causal=True, dropout_p=self.config.dropout if self.training else 0.0
+            )
+        local_mask = _local_causal_mask(q.shape[-2], self.config.local_window, q.device.type)
+        return F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=local_mask,
+            dropout_p=self.config.dropout if self.training else 0.0,
+        )
+
+    def _project(self, x: Tensor, position_offset: int) -> tuple[Tensor, Tensor, Tensor]:
+        batch, seq_len, _ = x.shape
+        q = self.q_proj(x).view(batch, seq_len, self.config.n_heads, self.config.head_dim)
+        k = self.k_proj(x).view(batch, seq_len, self.config.n_kv_heads, self.config.head_dim)
+        v = self.v_proj(x).view(batch, seq_len, self.config.n_kv_heads, self.config.head_dim)
+        q = self.q_norm(q).transpose(1, 2)
+        k = self.k_norm(k).transpose(1, 2)
+        v = v.transpose(1, 2)
+        q, k = self.rope(q, k, position_offset)
+        return q, k, v
+
+    def _finish(self, x: Tensor, out: Tensor) -> Tensor:
+        batch, seq_len, _ = x.shape
+        out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
+        if self.g_proj is not None:
+            out = out * torch.sigmoid(self.g_proj(x))
+        return self.o_proj(out)
+
+    def forward(self, x: Tensor, attention_mask: Tensor | None = None) -> Tensor:
+        q, k, v = self._project(x, position_offset=0)
+
+        use_flash = flash_attn_func is not None and x.is_cuda and attention_mask is None
+        if use_flash:
+            window = (-1, -1) if self.is_full else (self.config.local_window - 1, 0)
+            out = flash_attn_func(
+                q.transpose(1, 2),
+                k.transpose(1, 2),
+                v.transpose(1, 2),
+                dropout_p=self.config.dropout if self.training else 0.0,
+                causal=True,
+                window_size=window,
+            ).transpose(1, 2)
+        else:
+            out = self._torch_attention(q, k, v, attention_mask)
+        return self._finish(x, out)
+
+    def forward_cached(
+        self,
+        x: Tensor,
+        cache: tuple[Tensor, Tensor] | None,
+        position_offset: int,
+        attention_mask: Tensor | None = None,
+    ) -> tuple[Tensor, tuple[Tensor, Tensor]]:
+        q, new_k, new_v = self._project(x, position_offset)
+        if cache is None:
+            k, v = new_k, new_v
+            out = self._torch_attention(q, k, v, mask=attention_mask)
+        else:
+            k = torch.cat((cache[0], new_k), dim=-2)
+            v = torch.cat((cache[1], new_v), dim=-2)
+            if not self.is_full:
+                k = k[:, :, -self.config.local_window :]
+                v = v[:, :, -self.config.local_window :]
+            repeat = self.config.n_heads // self.config.n_kv_heads
+            expanded_k = k.repeat_interleave(repeat, dim=1) if repeat > 1 else k
+            expanded_v = v.repeat_interleave(repeat, dim=1) if repeat > 1 else v
+            if attention_mask is not None:
+                attention_mask = attention_mask[..., -k.shape[-2] :].to(dtype=q.dtype)
+            out = F.scaled_dot_product_attention(
+                q,
+                expanded_k,
+                expanded_v,
+                attn_mask=attention_mask,
+                is_causal=False,
+            )
+        if not self.is_full:
+            k = k[:, :, -self.config.local_window :]
+            v = v[:, :, -self.config.local_window :]
+        return self._finish(x, out), (k, v)
+
+
+class BoundedSwiGLU(nn.Module):
+    def __init__(self, config: BarunConfig) -> None:
+        super().__init__()
+        self.gate_up = nn.Linear(config.dim, 2 * config.ffn_dim, bias=False)
+        self.down = nn.Linear(config.ffn_dim, config.dim, bias=False)
+        self.clip = config.activation_clip
+
+    def forward(self, x: Tensor) -> Tensor:
+        gate, up = self.gate_up(x).chunk(2, dim=-1)
+        if self.clip > 0:
+            gate = gate.clamp(max=self.clip)
+            up = up.clamp(min=-self.clip, max=self.clip)
+        return self.down(F.silu(gate) * up)
+
+
+class BarunBlock(nn.Module):
+    def __init__(self, config: BarunConfig, *, is_full: bool) -> None:
+        super().__init__()
+        self.attn_norm = RMSNorm(config.dim, config.norm_eps)
+        self.attn = GroupedAttention(config, is_full=is_full)
+        self.ffn_norm = RMSNorm(config.dim, config.norm_eps)
+        self.ffn = BoundedSwiGLU(config)
+        self.dropout = config.dropout
+
+    def forward(self, x: Tensor, attention_mask: Tensor | None = None) -> Tensor:
+        x = x + F.dropout(self.attn(self.attn_norm(x), attention_mask), self.dropout, self.training)
+        x = x + F.dropout(self.ffn(self.ffn_norm(x)), self.dropout, self.training)
+        return x
+
+    def forward_cached(
+        self,
+        x: Tensor,
+        cache: tuple[Tensor, Tensor] | None,
+        position_offset: int,
+        attention_mask: Tensor | None = None,
+    ) -> tuple[Tensor, tuple[Tensor, Tensor]]:
+        attention, new_cache = self.attn.forward_cached(
+            self.attn_norm(x), cache, position_offset, attention_mask
+        )
+        x = x + attention
+        x = x + self.ffn(self.ffn_norm(x))
+        return x, new_cache
+
+
+class ResidualSelector(nn.Module):
+    """Convexly select between a group's input checkpoint and its transformed output."""
+
+    def __init__(self, config: BarunConfig) -> None:
+        super().__init__()
+        self.norm = RMSNorm(config.dim, config.norm_eps, affine=False)
+        self.score = nn.Linear(config.dim, 1, bias=False)
+        self.last_mean_weights: Tensor | None = None
+
+    def forward(self, checkpoint: Tensor, current: Tensor) -> Tensor:
+        candidates = torch.stack((checkpoint, current), dim=-2)
+        weights = self.score(self.norm(candidates)).softmax(dim=-2)
+        if not torch.compiler.is_compiling():
+            self.last_mean_weights = weights.detach().float().mean(dim=(0, 1, 3))
+        return (weights * candidates).sum(dim=-2)
+
+
+class BarunLM(nn.Module):
+    def __init__(self, config: BarunConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.embedding = nn.Embedding(config.vocab_size, config.dim)
+        self.layers = nn.ModuleList(
+            BarunBlock(config, is_full=(index + 1) % config.full_attention_every == 0)
+            for index in range(config.n_layers)
+        )
+        selector_count = (
+            config.n_layers // config.residual_select_every if config.residual_select_every else 0
+        )
+        self.selectors = nn.ModuleList(ResidualSelector(config) for _ in range(selector_count))
+        self.final_norm = RMSNorm(config.dim, config.norm_eps)
+        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
+        if config.mtp_loss_weight > 0:
+            self.mtp_norm = RMSNorm(config.dim, config.norm_eps)
+            self.mtp_proj = nn.Linear(config.dim, config.dim, bias=False)
+        else:
+            self.mtp_norm = None
+            self.mtp_proj = None
+        if config.tie_embeddings:
+            self.lm_head.weight = self.embedding.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear | nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        for name, parameter in module.named_parameters(recurse=False):
+            if name == "weight" and isinstance(module, RMSNorm) and parameter is not None:
+                nn.init.ones_(parameter)
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        labels: Tensor | None = None,
+        attention_mask: Tensor | None = None,
+        past_key_values: list[tuple[Tensor, Tensor]] | None = None,
+        use_cache: bool = False,
+        position_offset: int = 0,
+    ) -> BarunOutput:
+        if position_offset + input_ids.shape[1] > self.config.max_seq_len:
+            raise ValueError(f"sequence length exceeds max_seq_len={self.config.max_seq_len}")
+        if past_key_values is not None and len(past_key_values) != len(self.layers):
+            raise ValueError("past_key_values must have one entry per layer")
+        x = self.embedding(input_ids)
+        checkpoint = x
+        selector_index = 0
+        new_past_key_values = [] if use_cache else None
+        for index, layer in enumerate(self.layers):
+            if use_cache:
+                layer_cache = past_key_values[index] if past_key_values is not None else None
+                x, new_cache = layer.forward_cached(
+                    x, layer_cache, position_offset, attention_mask
+                )
+                new_past_key_values.append(new_cache)
+            else:
+                x = layer(x, attention_mask)
+            stride = self.config.residual_select_every
+            if stride and (index + 1) % stride == 0:
+                x = self.selectors[selector_index](checkpoint, x)
+                checkpoint = x
+                selector_index += 1
+        hidden = self.final_norm(x)
+        logits = self.lm_head(hidden)
+        if labels is None:
+            return BarunOutput(logits=logits, past_key_values=new_past_key_values)
+
+        causal_loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), labels.reshape(-1))
+        mtp_loss = None
+        loss = causal_loss
+        offset = self.config.mtp_offset
+        if self.mtp_proj is not None and hidden.shape[1] > offset:
+            mtp_hidden = self.mtp_norm(hidden[:, :-offset] + self.mtp_proj(hidden[:, :-offset]))
+            mtp_logits = self.lm_head(mtp_hidden)
+            mtp_labels = input_ids[:, offset:].clone()
+            target_mask = labels[:, offset - 1 : -1] == -100
+            mtp_labels.masked_fill_(target_mask, -100)
+            mtp_loss = F.cross_entropy(
+                mtp_logits.reshape(-1, mtp_logits.shape[-1]), mtp_labels.reshape(-1)
+            )
+            loss = loss + self.config.mtp_loss_weight * mtp_loss
+        return BarunOutput(
+            logits=logits,
+            loss=loss,
+            causal_loss=causal_loss,
+            mtp_loss=mtp_loss,
+            past_key_values=new_past_key_values,
+        )
+
+    @torch.no_grad()
+    def generate(self, input_ids: Tensor, max_new_tokens: int, temperature: float = 0.8) -> Tensor:
+        self.eval()
+        if input_ids.shape[1] + max_new_tokens > self.config.max_seq_len:
+            raise ValueError("prompt plus generation exceeds max_seq_len")
+        output = self(input_ids, use_cache=True)
+        past_key_values = output.past_key_values
+        for generated in range(max_new_tokens):
+            logits = output.logits[:, -1]
+            if temperature <= 0:
+                next_token = logits.argmax(dim=-1, keepdim=True)
+            else:
+                probs = (logits / temperature).softmax(dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            input_ids = torch.cat((input_ids, next_token), dim=1)
+            if generated + 1 < max_new_tokens:
+                output = self(
+                    next_token,
+                    past_key_values=past_key_values,
+                    use_cache=True,
+                    position_offset=input_ids.shape[1] - 1,
+                )
+                past_key_values = output.past_key_values
+        return input_ids
+
+    def parameter_counts(self) -> dict[str, int]:
+        total = sum(parameter.numel() for parameter in self.parameters())
+        embedding = self.embedding.weight.numel()
+        return {"total": total, "non_embedding": total - embedding, "embedding": embedding}

benchmark_results.json

@@ -0,0 +1,94 @@
+{
+  "schema_version": 1,
+  "model": {
+    "id": "harrrshall/BarunLM-35M",
+    "parameters": 35072768,
+    "training_tokens": 5699985408,
+    "checkpoint_sha256": "f2a7c88b9f2c2e3584809081407ab136795d82e30e89b730e007781c45d01447"
+  },
+  "protocol": {
+    "harness": "lm-eval==0.4.12",
+    "num_fewshot": 0,
+    "dtype": "bfloat16",
+    "seed": 1234,
+    "bootstrap_iterations": 10000,
+    "tasks": [
+      "arc_challenge",
+      "arc_easy",
+      "boolq",
+      "hellaswag",
+      "lambada_openai",
+      "openbookqa",
+      "piqa",
+      "sciq",
+      "winogrande"
+    ],
+    "decontamination": {
+      "method": "correctness-blind exact 13-token scan over the complete training history",
+      "training_tokens_scanned": 5700000000,
+      "excluded_samples": 1854
+    }
+  },
+  "results": {
+    "macro_accuracy": 0.4101019372910188,
+    "task_scores": {
+      "arc_challenge": 0.222990492653414,
+      "arc_easy": 0.38862761864044465,
+      "boolq": 0.5593303791235844,
+      "hellaswag": 0.2899662473151273,
+      "lambada_openai": 0.2450731707317073,
+      "openbookqa": 0.28,
+      "piqa": 0.5926950747094631,
+      "sciq": 0.5992115637319316,
+      "winogrande": 0.5130228887134964
+    }
+  },
+  "comparisons": [
+    {
+      "id": "LiquidAI/LFM2.5-230M-Base",
+      "revision": "9d2be5519834990d30996f878b6771cccbd24f2c",
+      "parameters": 229693184,
+      "macro_accuracy": 0.39197443061078113,
+      "barunlm_lead": 0.01812750668023767,
+      "barunlm_lead_ci95": [0.00919175854904696, 0.02710881312805468]
+    },
+    {
+      "id": "EleutherAI/pythia-160m-deduped",
+      "revision": "582159a2dfe3e712a8d47ae83dec95ae3bde8e7e",
+      "parameters": 162322944,
+      "macro_accuracy": 0.3734593023210671,
+      "barunlm_lead": 0.036642634969951675,
+      "barunlm_lead_ci95": [0.02735687681514485, 0.04598902006059502]
+    },
+    {
+      "id": "StentorLabs/Stentor-30M",
+      "revision": "6a69098f55961cc0df3a0751e44082616e4afdfa",
+      "parameters": 30419712,
+      "macro_accuracy": 0.3645857462123749,
+      "barunlm_lead": 0.045516191078643894,
+      "barunlm_lead_ci95": [0.03729788428965733, 0.05363598444691119]
+    },
+    {
+      "id": "roneneldan/TinyStories-33M",
+      "revision": "2ad0a164221b7c4d21cac7c46aec74f6f98dbfc8",
+      "parameters": 68514048,
+      "macro_accuracy": 0.33156687160177734,
+      "barunlm_lead": 0.07853506568924143,
+      "barunlm_lead_ci95": [0.0686233699582622, 0.08816152027826075],
+      "scope": "narrow-domain diagnostic"
+    },
+    {
+      "id": "EleutherAI/pythia-70m-deduped",
+      "revision": "e93a9faa9c77e5d09219f6c868bfc7a1bd65593c",
+      "parameters": 70426624,
+      "macro_accuracy": 0.3171108992671423,
+      "barunlm_lead": 0.09299103802387654,
+      "barunlm_lead_ci95": [0.08290362117614143, 0.10295299346141393]
+    }
+  ],
+  "evidence_sha256": {
+    "candidate_results": "0e0137bd990a6c0c0d8e12a53cd58dc7ab1faec928cecc5df630572deb13248e",
+    "peer_summary": "16ec0e53ac64e2bca9f4cf28be561ab8172d8a35bcc82cf6f810a288942054f4",
+    "lfm_pair_summary": "f3ba7d8c1ba50bc1fe34aaa546d3c4746a97f4526147a9d53127328fe2c60ee7"
+  }
+}

generate.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import argparse
+import hashlib
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_model
+from tokenizers import Tokenizer
+
+from barunlm import BarunConfig, BarunLM
+
+ROOT = Path(__file__).resolve().parent
+EXPECTED_SHA256 = {
+    "model.safetensors": "f2a7c88b9f2c2e3584809081407ab136795d82e30e89b730e007781c45d01447",
+    "barun_config.json": "9b3a1d71baa95a198744d250f9629231738d942570b8685c44307fd83dd33565",
+    "tokenizer.json": "70ded9605fccd09c2340ca7e225361eab0ae8b4dbbb0d6e26343ab5183979db6",
+}
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate text with BarunLM-35M.")
+    parser.add_argument("--prompt", default="The future of small language models is")
+    parser.add_argument("--max-new-tokens", type=int, default=48)
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
+    parser.add_argument("--verify-only", action="store_true")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    for name, expected in EXPECTED_SHA256.items():
+        actual = sha256(ROOT / name)
+        if actual != expected:
+            raise RuntimeError(f"{name} SHA-256 mismatch: expected {expected}, got {actual}")
+
+    config = BarunConfig.from_json(ROOT / "barun_config.json")
+    model = BarunLM(config)
+    missing, unexpected = load_model(model, ROOT / "model.safetensors", strict=False)
+    if missing or unexpected:
+        raise RuntimeError(f"checkpoint mismatch: missing={missing}, unexpected={unexpected}")
+    count = model.parameter_counts()["total"]
+    if count != 35_072_768:
+        raise RuntimeError(f"unexpected parameter count: {count}")
+    print(f"verified parameters={count} checkpoint_sha256={EXPECTED_SHA256['model.safetensors']}")
+    if args.verify_only:
+        return
+
+    if args.max_new_tokens < 1:
+        raise ValueError("--max-new-tokens must be positive")
+    if args.temperature < 0:
+        raise ValueError("--temperature cannot be negative")
+    device_name = "cuda" if args.device == "auto" and torch.cuda.is_available() else args.device
+    if device_name == "auto":
+        device_name = "cpu"
+    if device_name == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA was requested but is unavailable")
+    device = torch.device(device_name)
+    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+    model.to(device=device, dtype=dtype).eval()
+
+    tokenizer = Tokenizer.from_file(str(ROOT / "tokenizer.json"))
+    prompt_ids = tokenizer.encode(args.prompt, add_special_tokens=False).ids
+    if not prompt_ids:
+        raise ValueError("prompt must encode to at least one token")
+    if len(prompt_ids) + args.max_new_tokens > config.max_seq_len:
+        raise ValueError("prompt and continuation exceed the 2,048-token context")
+
+    input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
+    with torch.inference_mode():
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=args.max_new_tokens,
+            temperature=args.temperature,
+        )
+    print(tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=True))
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

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

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.4
+safetensors>=0.4.5
+tokenizers>=0.20,<1