--- license: apache-2.0 language: - en library_name: pytorch pipeline_tag: text-generation tags: - causal-lm - small-language-model - mla - multi-head-latent-attention - muon - sft - instruct - chat datasets: - HuggingFaceTB/smoltalk base_model: Shiv-22/tinylm-checkpoints-v2 --- # TinyLM 275M — Instruct (SmolTalk SFT) An **instruction-tuned** 275M small language model. This is the **E3-full** base (Run D — MLA + Muon — continued-pretrained on a 7.34B-token distill-mixture) supervised-fine-tuned for one epoch on **[SmolTalk](https://huggingface.co/datasets/HuggingFaceTB/smoltalk)** with a ChatML-style template and prompt-loss masking. - **Source code:** https://github.com/shivnarainms22/TinyLM - **Base model:** [`Shiv-22/tinylm-checkpoints-v2`](https://huggingface.co/Shiv-22/tinylm-checkpoints-v2) → **`e3_full/step_06999.pt`** — the final E3-full checkpoint from the v2 continued-pretraining track. Note this is *not* [`Shiv-22/tinylm`](https://huggingface.co/Shiv-22/tinylm) (Run D); E3-full is Run D continued on a 7.34B-token distill mixture (see [`results/v2/E3full_vs_runD.md`](https://github.com/shivnarainms22/TinyLM/blob/main/results/v2/E3full_vs_runD.md)) - **SFT details & eval:** [`results/v3/sft_vs_base.md`](https://github.com/shivnarainms22/TinyLM/blob/main/results/v3/sft_vs_base.md) - **📄 What the project found (all four tracks, one page):** [`results/FINDINGS.md`](https://github.com/shivnarainms22/TinyLM/blob/main/results/FINDINGS.md) ## TinyLM model family | Repo | What it is | |---|---| | [`Shiv-22/tinylm`](https://huggingface.co/Shiv-22/tinylm) | **Base 275M** — Run D (MLA + Muon), ablation winner; the model for downstream use | | [`Shiv-22/tinylm-instruct`](https://huggingface.co/Shiv-22/tinylm-instruct) **← this repo** | **Instruct** — the base SmolTalk-SFT'd for chat (ChatML) | | [`Shiv-22/tinylm-checkpoints-v2`](https://huggingface.co/Shiv-22/tinylm-checkpoints-v2) | **All 4 ablation arms** (A/B/C/D) + the **E3-full** continued-pretraining checkpoints | | [`Shiv-22/tinylm-checkpoints`](https://huggingface.co/Shiv-22/tinylm-checkpoints) | **v1 historical** checkpoint (1B×21 tokens, pre data-fix) | Source & full results: [github.com/shivnarainms22/TinyLM](https://github.com/shivnarainms22/TinyLM) > **Scope.** This is a 275M research model, part of a training-methods portfolio. > It follows instructions and holds a chat format, but it is **not** a reliable > assistant — expect factual slips, weak multi-step reasoning, and occasional > repetition. See **Limitations**. --- ## Prompt format The model was trained on this exact text template (no new special tokens; vocab stays 32,000). Encode **without** a BOS token; the model closes its own turn with EOS. ``` <|user|> {your message} <|assistant|> ``` An optional system turn may precede it: `<|system|>\n{system}\n`. Generation continues after the final `<|assistant|>\n` and stops at EOS. The repo's `scripts/generate_sft_samples.py` builds this priming string exactly. ## Usage This is a **custom PyTorch model** (not a `transformers` architecture), so load it with the repo's code rather than `AutoModelForCausalLM`: ```bash git clone https://github.com/shivnarainms22/TinyLM && cd TinyLM pip install torch transformers huggingface_hub ``` ```python from huggingface_hub import hf_hub_download from transformers import AutoTokenizer import sys; sys.path.insert(0, "src") from tinylm.loader import load_from_checkpoint from tinylm.sft import render_chat, _ROLE_PREFIX ckpt = hf_hub_download("Shiv-22/tinylm-instruct", "tinylm_sft_smoltalk.pt") model = load_from_checkpoint(ckpt, device="cpu").eval() tok = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf") enc = lambda s: tok.encode(s, add_special_tokens=False) # no BOS ids, _ = render_chat([{"role": "user", "content": "Explain photosynthesis simply."}], enc, tok.eos_token_id) ids += enc(_ROLE_PREFIX["assistant"]) # prime the assistant turn # ... greedy-decode from `ids`, stop at eos — see scripts/generate_sft_samples.py ``` `scripts/generate_sft_samples.py --checkpoint ` does exactly this end to end. ## Training | | | |---|---| | Base | E3-full `step_06999` (275M, MLA + Muon) | | Data | [SmolTalk](https://huggingface.co/datasets/HuggingFaceTB/smoltalk), ChatML render, **loss on assistant turns only** | | Objective | Memory-safe chunked cross-entropy, `ignore_index=-100` | | Steps | 20,000 (~655M tokens, 1 epoch cap) | | LR | cosine 2e-5 → 2e-6, 3% warmup | | Hardware | 1× A100, 5h48m, final train loss ~1.1 | **Training logs:** SFT loss curve on Weights & Biases — [`tinylm-v3`](https://wandb.ai/shivnarainms22-northeastern-university/tinylm-v3). The base model's continued-pretraining runs are in [`tinylm-v2`](https://wandb.ai/shivnarainms22-northeastern-university/tinylm-v2). ## Evaluation — the alignment tax Instruction-tuning is evaluated honestly for its **cost** on the locked completion/MCQ suite (it cannot measure instruction-following, only what SFT trades away). Base = E3-full at 0-shot; both scored with the same harness. | Benchmark | Metric | E3-full base | Instruct (SFT) | Δ | |-----------|--------|-------------:|---------------:|------:| | LAMBADA | perplexity ↓ | 23.20 | **22.97** | −0.23 (preserved) | | HellaSwag | acc_norm | 0.4125 | 0.4107 | −0.0018 (flat) | | Winogrande | acc | 0.5146 | 0.5209 | +0.0063 (flat) | | ARC-Easy | acc | 0.5753 | 0.5501 | **−0.0252** (~2.5σ) | Language modeling is **preserved** (LAMBADA perplexity even nudges down — no catastrophic forgetting), commonsense is flat, and the one real cost is ~2.5 points on ARC-Easy multiple-choice — the expected tax of shifting the model from ranking canned options toward producing chat answers. ## Sample outputs Greedy decoding, unedited. Full set (8 prompts): [`results/v3/sft_samples.md`](https://github.com/shivnarainms22/TinyLM/blob/main/results/v3/sft_samples.md). **Explain photosynthesis in one simple paragraph.** > Photosynthesis is the process by which plants convert light energy into > chemical energy, which is used to produce food for the plant. The process > involves the absorption of light energy by chlorophyll in the leaves of a > plant, which then absorbs the energy and uses it to produce glucose, a type > of sugar. […] **Rewrite this sentence to be more polite: 'Send me the file now.'** > I'll send you the file now. **Give me two tips for staying focused while studying.** > First, establish a study routine that works for you. Set aside dedicated time > each day to study… Another key tip is to create a dedicated study space… turn > off any distractions like phones or TVs. The model reliably adopts the assistant role, stays on topic, and produces coherent, on-format answers — that is the instruction-following the benchmark table above cannot see. ## Limitations These are real and visible in the same sample set — shown, not hidden: - **Weak multi-step reasoning / arithmetic.** Asked a train's average speed (60 km / 2 h), it answered "120 km/h" then "≈1.25" — wrong. Consistent with the v2/v3 finding that reasoning at 275M is **capacity-bound**, and SFT does not add reasoning the base model lacks. - **Factual errors.** "good morning" → Spanish came back as *"¿Dónde está el día?"* (incorrect). It will state wrong facts confidently. - **Repetition on open-ended generation.** Under greedy decoding, creative prompts (e.g. a poem) can collapse into a repeated line. Sampling (`temperature 0.7`, `top_p 0.9`) mitigates this but trades determinism. - **No safety tuning.** Only SmolTalk SFT — no preference optimization (DPO/RLHF), no safety alignment. Do not deploy in user-facing or high-stakes settings. ## Intended use Research and education: studying small-model post-training, prompt-template effects, and the gap between benchmark scores and instruction-following. Not intended for production assistance.