OpenLanguageModel: Readable and Composable Small-Language-Model Pretraining for Education and Research

# Readable and Composable Small-Language-Model Pretraining for Education and Research
Source: [https://arxiv.org/html/2607.16669](https://arxiv.org/html/2607.16669)
###### Abstract

OpenLanguageModel \(OLM\) is an open\-source PyTorch library for building and pretraining small language models while keeping their machinery visible\. In OLM, model code reads like the architecture: components are ordinary modules, whileBlock,Residual,Repeat, andParalleldescribe how they are wired\. The resulting model can move unchanged from a teaching notebook to a complete pretraining run or a research ablation\. OLM connects this readable model layer to tokenizers, local and streaming datasets, optimization, mixed precision, callbacks, checkpoints, and hardware\-aware CPU, single\-GPU, and single\-node multi\-GPU execution\. We demonstrate the full path by tracing GPT\-2 from diagram to code, launching a FineWeb\-Edu training script, replacing one attention component, and lettingAutoTrainerconfigure the available machine\. The package includes 27 presets across nine familiar model families and documentation that progresses from LM fundamentals to architecture research\. Validation shows close agreement with independent reference implementations, 90\.6% four\-GPU weak\-scaling efficiency for a 348M\-parameter workload, compact architecture edits, and positive early usability results\. OLM is MIT\-licensed and available through PyPI, GitHub, and its documentation site\.

## 1Introduction

Small language models \(SLMs\) put the complete language\-modelling workflow within reach of a classroom, a single workstation, or a focused research project\. They are large enough to exercise real tokenization, data streaming, optimization, checkpointing, and distributed training, yet small enough that one person can still understand and change the whole model\.

OpenLanguageModel \(OLM\) brings those two goals—understanding the model and running it seriously—into the same PyTorch library\. Its model definitions follow the diagrams used to explain transformers\. Embeddings, attention, normalization, feed\-forwards, residual paths, repeated layers, and branches remain visible as ordinary modules\. The same object then connects to a complete SLM pretraining stack, carrying a lesson directly into an experiment\.

This design serves students and educators who want to build an LM from first principles, practitioners who want a readable pretraining recipe, and researchers who want local component experiments with the surrounding system held steady\. Users can start from a named reference model, assemble a model from components, insert a handwrittennn\.Module, or bring the resulting model to their own PyTorch loop\.

OLM keeps three things together: architecture code that mirrors the model, end\-to\-end hardware\-aware SLM pretraining, and a learning path that grows into research\. The paper follows that workflow directly\. We first read a model, train it, change one idea, and scale the run; we then open the system design and give compact validation of correctness, customization, scaling, and usability\.

## 2A Guided Tour of OLM

A first\-time user can install OLM withpip install openlanguagemodel\. The demonstration then follows one model from source code to a saved checkpoint\.

#### Read the model\.

Figure[1](https://arxiv.org/html/2607.16669#S2.F1)places a GPT\-2 diagram beside its OLM definition\. The correspondence is literal: token and positional embeddings feed a repeated stack of residual attention and feed\-forward blocks, followed by an output head\. Opening a Llama\-style definition reveals the same structural vocabulary with RMSNorm, grouped\-query attention, RoPE, and SwiGLU in place of GPT\-2’s components\. The full architecture is visible at a glance in the definition\.

![Refer to caption](https://arxiv.org/html/2607.16669v1/figures/gpt2-architecture-code.png)

Figure 1:GPT\-2 as a diagram and as an OLM model\. The order, nesting, residual paths, and layer repetition are visible in the executable PyTorch definition\.
#### Train the same object\.

Figure[2](https://arxiv.org/html/2607.16669#S2.F2)A shows a compact pretraining script: choose a tokenizer, stream FineWeb\-Edu, construct a familiar model, and pass ordinary PyTorch objects toAutoTrainer\. The trainer supplies the training loop, mixed precision, accumulation, scheduling, metrics, and checkpoints\. A local text dataset or a user\-written loop can slot into the same workflow while the model stays unchanged\.

A\.Connect a model to complete pretraining

```
tok = HFTokenizer("gpt2")
data = FineWebEduDataset(
    tok, context_length=1024, streaming=True)
loader = DataLoader(data, batch_size=8)

model = GPT2Model(
    vocab_size=tok.vocab_size,
    embed_dim=768, num_layers=12,
    num_heads=12, max_seq_len=1024)

trainer = AutoTrainer(
    model, AdamW, loader,
    device="auto", context_length=1024,
    learning_rate=3e-4, grad_accum_steps=8)
trainer.train(epochs=1, max_steps=1000)
```

B\.Make one local architecture change

```
# Baseline component for this run:
attention = GroupedQueryAttention(
    d_model, heads, kv_heads, context)

# For the ALiBi ablation, replace only
# the assignment above with this one:
# attention = MultiHeadAttentionwithALiBi(
#     d_model, heads, max_seq_len=context)

layer = Block([
    Residual(Block([
        RMSNorm(d_model), attention])),
    Residual(Block([
        RMSNorm(d_model), SwiGLUFFN(d_model)])),
])
```

Figure 2:Two moments from the OLM demo\. The left panel connects a readable model to the training stack\. The right panel changes the attention rule while preserving the surrounding layer and training path\. Imports are omitted for space\.
#### Change one idea\.

The right side of Figure[2](https://arxiv.org/html/2607.16669#S2.F2)makes the research workflow concrete\. The attention object is one line in an otherwise unchanged transformer layer\. Replacing grouped\-query attention with ALiBi attention leaves the residual wiring, normalization, feed\-forward path, data pipeline, optimizer, and trainer intact\. A new attention rule can also subclassAttentionBaseor enter the block as any ordinarynn\.Module\.

#### Fit the run to the machine\.

Withdevice="auto", OLM reports the detected devices, available memory, estimated model footprint, and selected execution strategy\. The same call runs on CPU or one GPU and configures DDP or FSDP for multiple GPUs on one node\. During training, the user sees loss, perplexity, learning rate, and throughput; callbacks can add validation or experiment tracking\. The demo ends by saving, reloading, and sampling from the model\.

## 3The System Behind the Demo

### 3\.1Architecture as an Object Graph

OLM represents both components and wiring withtorch\.nn\.Moduleobjects\.Blockstores an orderedModuleList;Residualcomputesx\+f​\(x\)x\+f\(x\);Repeatcreates independently parameterized layers from a factory; andParallelsends one input through multiple branches before a selectable merge\. Because the combinators share the same module\-in/module\-out interface as attention or feed\-forward layers, they nest naturally\. A parallel residual path can be written asResidual\(Parallel\(\[attention, ffn\]\)\); a heterogeneous stack is simply an explicit list of different layers\.

The object graph is the model\. It can be printed, traversed, placed on a device, optimized, saved, or embedded inside a larger handwritten module using normal PyTorch tools\. The trainer consumes this graph directly, so source, inspection, and execution share one representation\. Custom control flow can still use a conventionalforwardmethod\.

### 3\.2Reusable Components and Reference Models

OLM exposes attention, positional encoding, embeddings, normalization, feed\-forward layers, output heads, and structural combinators as independent imports\. Configurable implementations of[GPT\-2](https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/openai/gpt2.py),[Llama 2](https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/meta/llama2.py),[Llama 3](https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/meta/llama3.py),[Qwen 2\.5](https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/alibaba/qwen2.py),[Phi\-3](https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/microsoft/phi3.py),[Phi\-4](https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/microsoft/phi4.py),[Gemma 2](https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/google/gemma2.py),[OLMo](https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/allenai/olmo.py), and[OPT](https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/facebook/opt.py)use those same pieces, with 27 named presets for familiar parameterizations\. Family blocks and model skeletons remain separate from preset dimensions, making a reference model both a ready starting point and readable example code\.

The component layer includes explicit educational implementations as well as efficient paths backed by PyTorch scaled dot\-product attention\. Modern details such as RoPE, ALiBi, grouped\-query attention, RMSNorm, SwiGLU, alternating local attention, and mixture\-of\-experts feed\-forwards remain visible where they enter the model\.

### 3\.3A Connected Pretraining Stack

The founding team has trained more than 40 SLMs below one billion parameters across its research projects, and that repeated use has shaped OLM’s training path\. The library provides local and streaming datasets, Hugging Face tokenizer integration, data loading, optimizers, losses, warmup and decay schedules, mixed precision, gradient accumulation and clipping, callbacks, throughput logging, and checkpoint save/reload\.

AutoTrainerinspects CPU and CUDA devices, GPU count and memory, and the estimated model footprint, then chooses the base trainer, DDP, or FSDP according to a balanced, speed, or memory\-oriented policy\. Users can select a strategy directly or keep the model in their own loop\. Architecture and training stay separate, so a component experiment continues to use the same tested data and optimization path\.

### 3\.4Documentation as Part of the System

The documentation at[https://openlanguagemodel\.github\.io/openlanguagemodel/](https://openlanguagemodel.github.io/openlanguagemodel/)begins with tokens, embeddings, attention, and next\-token prediction; continues to a runnable small model with training, generation, and save/reload; and then opens the block system, modern architectures, distributed training, and the API reference\. Figure[3](https://arxiv.org/html/2607.16669#S3.F3)summarizes the three entry points\. The code introduced in the beginner path is the same code used later in the research path\.

Learntokens→\\rightarrowembeddings→\\rightarrowattention→\\rightarrowfirst LMTrainlocal/FineWeb\-Edu data→\\rightarrowAutoTrainer→\\rightarrowcheckpointResearchreplace attention, normalization, FFN, or wiringSametorch\.nn\.Modulegraph, data interfaces, and training tools

Figure 3:Three entry points, one code path\. Learning, pretraining, and research use the same inspectable modules and connected training stack\.

## 4Compact System Validation

The evaluation checks the capabilities shown in the tour: model coverage and reference fidelity, local architecture changes, single\-node scaling, and the experience of early users\. Table[1](https://arxiv.org/html/2607.16669#S4.T1)gives the short version; protocols and packaged records are described in Appendix[A](https://arxiv.org/html/2607.16669#A1)\.

Table 1:Validation summary\.#### Coverage and reference fidelity\.

Reduced configurations of all nine model families passed forward, analytical parameter\-count, tied\-embedding, and bitwise checkpoint round\-trip checks\. All 27 named presets passed configuration and parameter\-formula checks, and the local suite passed 149 tests\. For a numerical check, we copied weights from independent Hugging Face implementations into four\-layer FP32 GPT\-2, Llama 3, and Qwen 2\.5 variants\. Across three seeds per family, the maximum absolute logit difference was3\.576×10−73\.576\\times 10^\{\-7\}, the maximum loss difference was4\.768×10−74\.768\\times 10^\{\-7\}, and selected gradient cosine similarity remained above0\.99999999999980\.9999999999998\.

#### Local architecture changes\.

We implemented six downstream edits in new user scripts while keeping framework, training, and data files unchanged: parallel residuals, heterogeneous blocks, selected\-layer MoE, RoPE\-to\-ALiBi replacement, cross\-family block mixing, and an N\-way learned merge\. On the four tasks completed in OLM, LitGPT, and Pico, the new scripts contained 117, 171, and 195 lines respectively\. Across all six OLM–LitGPT tasks, the totals were 167 and 256 lines\. This small code surface reflects reuse of exposed components and combinators\.

#### Single\-node scaling\.

A 348,447,744\-parameter Llama\-3\-style model was weak\-scaled on one, two, and four A100\-SXM4\-80GB GPUs with BF16, AdamW, PyTorch SDPA, and 2,048 tokens per GPU\. Mean throughput across three runs was 17,785, 32,426, and 64,427 tokens/s\. The four\-GPU run retained 90\.6% efficiency relative to ideal scaling from one GPU, while peak allocated memory stayed at 10\.47 GB per GPU on two and four GPUs\.

#### Early user experience\.

In a de\-identified survey of 20 OLM users, the mean System Usability Scale score was 73\.8 \(95% bootstrap CI 69\.4–78\.4\)\. All 20 respondents agreed that the architecture was understandable and that OLM felt natural with ordinary PyTorch; 19 agreed that architecture changes were localized\. Among 17 respondents able to compare effort, 14 found architecture\-level changes easier in OLM\.

## 5Positioning and Availability

OLM builds on PyTorch’s imperative module system\[[5](https://arxiv.org/html/2607.16669#bib.bib5)\]and complements several strong language\-model ecosystems\. Transformers offers broad access to pretrained models\[[6](https://arxiv.org/html/2607.16669#bib.bib6)\]; LitGPT provides compact training recipes\[[4](https://arxiv.org/html/2607.16669#bib.bib4)\]; LMFlow emphasizes foundation\-model finetuning and inference\[[2](https://arxiv.org/html/2607.16669#bib.bib2)\]; and Pico supports hypothesis\-driven SLM research\[[3](https://arxiv.org/html/2607.16669#bib.bib3)\]\. OLM’s focus is the path that connects these concerns: readable architecture code, guided learning, complete SLM pretraining, and local research edits in one library\.

As of this writing, OLM v2\.2\.0 is the latest release; the project is actively developed by seven contributors and distributed under the MIT license\. The package is installable from[https://pypi\.org/project/openlanguagemodel/](https://pypi.org/project/openlanguagemodel/); source is hosted at[https://github\.com/openlanguagemodel/openlanguagemodel](https://github.com/openlanguagemodel/openlanguagemodel); and the website at[https://openlanguagemodel\.github\.io/openlanguagemodel/](https://openlanguagemodel.github.io/openlanguagemodel/)provides the beginner course, first\-model tutorial, training guides, architecture guide, reference\-model source, and generated API documentation\. These public entry points make the demonstrated system available for teaching, research, and extension\.

## 6Conclusion

OLM starts from a simple promise: the code used to explain a language model should still be useful when it is time to train and change that model\. A student can follow the arrows in GPT\-2, a practitioner can connect the same object to streaming data and hardware\-aware training, and a researcher can replace the attention line while keeping the rest of the run intact\.

The demonstration shows that path end to end\. OLM combines readable PyTorch composition, a complete SLM pretraining suite, familiar reference models, and documentation that grows with the user\. The MIT\-licensed package, source, and tutorials are available now for teaching, experiments, and new contributions\.

## Limitations

The validation is scoped to the system shown in the demo\. Numerical parity uses reduced models, the GPU study uses synthetic tokens and DDP on one A100 node, and edit size is a proxy for code locality rather than authoring time\. Future evaluation will extend to full\-scale quality, FSDP, and multi\-node training\. The 20\-person survey is an early convenience sample of recent users, so its scores describe that group rather than the broader LM community\.

## Ethical Considerations

OLM is general\-purpose research software\. Its readable, open\-source design supports inspection, while users remain responsible for dataset rights, bias and harmful\-output evaluation, access controls, and computational resources\. Only de\-identified aggregate results from the user study are included in the submission; the row\-level workbook remains private\.

## Acknowledgements

Created by members of FAIRC, with contributions from Leap AI Club at Plaksha University\.

## References

- Brooke \[1996\]John Brooke\.SUS: A “quick and dirty” usability scale\.In Patrick W\. Jordan, Bruce Thomas, Ian L\. McClelland, and Bernard Weerdmeester, editors,*Usability Evaluation in Industry*, pages 189–194\. Taylor & Francis, London, United Kingdom, 1996\.doi:10\.1201/9781498710411\-35\.
- Diao et al\. \[2024\]Shizhe Diao, Rui Pan, Hanze Dong, KaShun Shum, Jipeng Zhang, Wei Xiong, and Tong Zhang\.LMFlow: An extensible toolkit for finetuning and inference of large foundation models\.In*Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies \(Volume 3: System Demonstrations\)*, pages 116–127, Mexico City, Mexico, 2024\. Association for Computational Linguistics\.doi:10\.18653/v1/2024\.naacl\-demo\.12\.URL[https://aclanthology\.org/2024\.naacl\-demo\.12/](https://aclanthology.org/2024.naacl-demo.12/)\.
- Diehl Martinez et al\. \[2025\]Richard Diehl Martinez, David Demitri Africa, Yuval Weiss, Suchir Salhan, Ryan Daniels, and Paula Buttery\.Pico: A modular framework for hypothesis\-driven small language model research\.In*Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing: System Demonstrations*, pages 295–306, Suzhou, China, 2025\. Association for Computational Linguistics\.doi:10\.18653/v1/2025\.emnlp\-demos\.22\.URL[https://aclanthology\.org/2025\.emnlp\-demos\.22/](https://aclanthology.org/2025.emnlp-demos.22/)\.
- Lightning AI \[2023\]Lightning AI\.LitGPT\.[https://github\.com/Lightning\-AI/litgpt](https://github.com/Lightning-AI/litgpt), 2023\.
- Paszke et al\. \[2019\]Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala\.PyTorch: An imperative style, high\-performance deep learning library\.In*Advances in Neural Information Processing Systems*, volume 32, 2019\.URL[https://proceedings\.neurips\.cc/paper/2019/hash/bdbca288fee7f92f2bfa9f7012727740\-Abstract\.html](https://proceedings.neurips.cc/paper/2019/hash/bdbca288fee7f92f2bfa9f7012727740-Abstract.html)\.
- Wolf et al\. \[2020\]Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush\.Transformers: State\-of\-the\-art natural language processing\.In*Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations*, pages 38–45, Online, 2020\. Association for Computational Linguistics\.doi:10\.18653/v1/2020\.emnlp\-demos\.6\.URL[https://aclanthology\.org/2020\.emnlp\-demos\.6/](https://aclanthology.org/2020.emnlp-demos.6/)\.

## Appendix AEvaluation Details

Parity used deterministic FP32 CPU execution and seeds 11, 22, and 33\. Scaling used one 2,048\-token sequence per GPU, 50 warm\-up and 200 measured steps, and three replicates\. Edit counts cover new downstream\-script lines\. SUS uses de\-identified aggregates, the standard transform\[[1](https://arxiv.org/html/2607.16669#bib.bib1)\], and 10,000 bootstrap resamples \(seed 20260716\)\.