| --- |
| license: mit |
| datasets: |
| - roneneldan/TinyStories |
| language: |
| - en |
| metrics: |
| - perplexity |
| pipeline_tag: text-generation |
| tags: |
| - slm |
| - transformer |
| - attention |
| - optimization |
| - pytorch |
| - tinystories |
| - educational |
| --- |
| # Model Card for Helium-Nano (45M) |
|
|
| **Helium-Nano** is a 45-million parameter Small Language Model (SLM) trained on the TinyStories dataset. It demonstrates how a highly optimized custom Transformer architecture can achieve coherent English storytelling capabilities with minimal compute resources. The model was trained in under 1 hour on a single Nvidia L4 GPU, achieving a throughput of **409k tokens/second** via PyTorch 2.0 compile and architectural optimizations. |
|
|
| ## Model Details |
|
|
| ### Model Description |
|
|
| Helium-Nano is a decoder-only Transformer designed to investigate training dynamics and scaling laws in low-resource environments. Despite its small size, it produces grammatically correct and narratively consistent short stories. |
|
|
| The primary goal of this model was engineering efficiency. By implementing **BFloat16 mixed precision**, **Flash Attention principles**, **Torch.compile (Inductor)**, and **Float32-optimized Rotary Embeddings (RoPE)**, the training pipeline achieved a 16x speedup over standard eager-mode baselines. |
|
|
| - **Developed by:** Debmalya/batmanLovesAI |
| - **Model type:** Decoder-only Transformer (Custom Architecture) |
| - **Language(s) (NLP):** English |
| - **License:** MIT |
| - **Finetuned from model:** N/A (Trained from scratch) |
|
|
| ### Model Sources |
|
|
| - **Repository:** [Link to Github Repo](https://github.com/DebmalyaSen34/HeliumLM) |
| - **Dataset Paper:** [TinyStories: How Small Can Language Models Be?](https://arxiv.org/abs/2305.07759) |
| - **Optimization Techniques:** [Small Language Models: Architectures, Techniques, Evaluation, Problems and Future Adaptation](https://arxiv.org/abs/2505.19529) |
|
|
| ## Uses |
|
|
| ### Direct Use |
|
|
| - **Story Generation:** Generating simple, coherent short stories suitable for early childhood reading levels. |
| - **Educational:** A lightweight baseline for experimenting with model interpretation, quantization, or fine-tuning on consumer hardware. |
| - **Performance Benchmarking:** Testing inference speeds of small transformers on various hardware. |
|
|
| ### Out-of-Scope Use |
|
|
| - **Factual Queries:** The model is trained on fiction; it has no world knowledge and will hallucinate facts. |
| - **Reasoning/Math:** The model is not capable of complex logic or arithmetic. |
| - **Harmful Content:** While the dataset is heavily filtered, users should not attempt to generate toxic or biased content. |
|
|
| ## Bias, Risks, and Limitations |
|
|
| - **Dataset Bias:** The model reflects the vocabulary and concepts found in the TinyStories dataset, which focuses on simple, positive narratives using a limited vocabulary (approx 3-year-old level). |
| - **Repetition:** Like many SLMs, the model may enter repetitive loops if the temperature is too low or repetition penalty is not applied during inference. |
| - **Hallucinations:** The model prioritizes grammatical structure over semantic logic. |
|
|
| ## How to Get Started with the Model |
|
|
| Since this uses a custom architecture, you need to instantiate the model class before loading weights. |
|
|
| ```python |
| import torch |
| from tokenizers import Tokenizer |
| # Assuming TinySLM class is defined in your local files |
| |
| # 1. Load Tokenizer |
| tokenizer = Tokenizer.from_file("tokenizer.json") |
| |
| # 2. Initialize Model |
| config = { |
| "vocab_size": 32000, |
| "d_model": 512, |
| "n_head": 8, |
| "n_layers": 10, |
| "max_seq_len": 512 |
| } |
| model = TinySLM(config) |
| |
| # 3. Load Weights |
| state_dict = torch.load("helium_nano_45m.pt", map_location="cpu") |
| model.load_state_dict(state_dict) |
| model.eval() |
| |
| # 4. Generate |
| prompt = "Once upon a time, there was a little" |
| # ... inference code ... |
