Add tokenizer, inference code, model card, and 20-query report

20_query_inference_report.md

@@ -0,0 +1,90 @@
+# Tiny LLM 75M — 20 Query Inference Report
+
+Repo: `razor5050/tinyllm-75m-openwebtext-chat`
+Checkpoint: `final.pt`
+SFT step: `781`
+Parameters: `75,074,112`
+State load: missing `0`, unexpected `0`
+Elapsed: `218.86s`
+
+## Questions and answers
+
+### 1. What is photosynthesis?
+
+Photosynthesis is a chemical process in which a molecule is formed and then into a form called a molecular structure. It is a chemical that is composed of the elements that are
+
+### 2. Why do leaves look green?
+
+The word "green" is actually a word that is often used in various forms of environmental and environmental sciences. It is often used in various fields, such as agriculture, engineering,
+
+### 3. What is gravity?
+
+Gravity is a phenomenon known as gravity. It occurs when a substance of gravity is formed in a shape or shape that is not present in a substance, and the substance is not
+
+### 4. Name one planet in our solar system.
+
+The solar system was named from the first solar system, named by the name "Solar".
+
+### 5. What is the largest ocean on Earth?
+
+The largest ocean on Earth is the Nusa-Eibagan, the largest ocean on Earth. It is a ocean in the southern Peninsula of the United States, and
+
+### 6. What gas do humans breathe in to survive?
+
+A gas in water, the gas's core, is a fluid that is released through the water's chemical reactions, which can cause blood death and heart disease. A gas in water
+
+### 7. What does the heart do?
+
+The heart is a system of organs that are designed to be designed to be designed to perform functions that perform functions that perform functions that perform functions that perform functions that perform functions that perform
+
+### 8. What is a volcano?
+
+A volcano is a volcano that consists of a range of extreme events, such as volcanic eruptions, earthquakes, and earthquakes. It is a volcanic eruption that occurs within a range of
+
+### 9. What is the Moon?
+
+The Moon is a vast and beautiful space that is packed with diverse cultures, diverse cultures, and diverse cultures. It is a vast, and well-designed space that can be accessed
+
+### 10. Why does rain fall from clouds?
+
+The wind is a powerful force that can create storms and create storms. By rain fall, the wind causes the wind to move, which is the way the air is released from the
+
+### 11. What is electricity?
+
+Electototototototototototototototototototototototototototototototototototot
+
+### 12. Name one continent.
+
+A continent a continent in the America, United States, and United States.
+
+### 13. What is a river?
+
+A river is a river that is a river that flows through the river. It is a river that is flowing through the river, making it a river that flows through rivers. It
+
+### 14. What animal is known as the king of the jungle?
+
+The king of the jungle is the king of the jungle. This species is known as the "King of the jungle." The king of the jungle is the king of the jungle.
+
+### 15. What is the purpose of roots in a plant?
+
+A plant is a plant that is a plant that is composed of a plant that is composed of a plant that is composed of a plant that is composed of a plant that is composed
+
+### 16. What is a computer?
+
+A computer is a computer that is designed to monitor and monitor any specific data, including any relevant data, data, or data, or any relevant data, or any relevant data in
+
+### 17. What is the internet?
+
+The Internet is a digital movement that spans various regions, each having its own unique and complex connections. The internet is a phenomenon that occurs through various methods, including internet-specific communication
+
+### 18. What is a vaccine?
+
+A vaccine is an vaccine designed to promote the growth of the immune system, which regulates the immune system's activity, such as the immune system, immune cells, and immune cells
+
+### 19. What country is Tokyo in?
+
+Japan is a world where Tokyo is known for its stunning water quality and its culture. It is a hub for the world's largest expesticive waste waste agency, providing a vital
+
+### 20. What is climate change?
+
+Climate change is often associated with a significant increase in global temperature, associated with a rapid acceleration in the warming cycle, with a greater amount of heat and heat being added to the global

20_query_inference_results.json

@@ -0,0 +1,113 @@
+{
+  "meta": {
+    "repo": "razor5050/tinyllm-75m-openwebtext-chat",
+    "checkpoint": "final.pt",
+    "step": 781,
+    "parameters": 75074112,
+    "missing_keys": 0,
+    "unexpected_keys": 0
+  },
+  "elapsed_seconds": 218.86,
+  "results": [
+    {
+      "n": 1,
+      "question": "What is photosynthesis?",
+      "answer": "Photosynthesis is a chemical process in which a molecule is formed and then into a form called a molecular structure. It is a chemical that is composed of the elements that are"
+    },
+    {
+      "n": 2,
+      "question": "Why do leaves look green?",
+      "answer": "The word \"green\" is actually a word that is often used in various forms of environmental and environmental sciences. It is often used in various fields, such as agriculture, engineering,"
+    },
+    {
+      "n": 3,
+      "question": "What is gravity?",
+      "answer": "Gravity is a phenomenon known as gravity. It occurs when a substance of gravity is formed in a shape or shape that is not present in a substance, and the substance is not"
+    },
+    {
+      "n": 4,
+      "question": "Name one planet in our solar system.",
+      "answer": "The solar system was named from the first solar system, named by the name \"Solar\"."
+    },
+    {
+      "n": 5,
+      "question": "What is the largest ocean on Earth?",
+      "answer": "The largest ocean on Earth is the Nusa-Eibagan, the largest ocean on Earth. It is a ocean in the southern Peninsula of the United States, and"
+    },
+    {
+      "n": 6,
+      "question": "What gas do humans breathe in to survive?",
+      "answer": "A gas in water, the gas's core, is a fluid that is released through the water's chemical reactions, which can cause blood death and heart disease. A gas in water"
+    },
+    {
+      "n": 7,
+      "question": "What does the heart do?",
+      "answer": "The heart is a system of organs that are designed to be designed to be designed to perform functions that perform functions that perform functions that perform functions that perform functions that perform functions that perform"
+    },
+    {
+      "n": 8,
+      "question": "What is a volcano?",
+      "answer": "A volcano is a volcano that consists of a range of extreme events, such as volcanic eruptions, earthquakes, and earthquakes. It is a volcanic eruption that occurs within a range of"
+    },
+    {
+      "n": 9,
+      "question": "What is the Moon?",
+      "answer": "The Moon is a vast and beautiful space that is packed with diverse cultures, diverse cultures, and diverse cultures. It is a vast, and well-designed space that can be accessed"
+    },
+    {
+      "n": 10,
+      "question": "Why does rain fall from clouds?",
+      "answer": "The wind is a powerful force that can create storms and create storms. By rain fall, the wind causes the wind to move, which is the way the air is released from the"
+    },
+    {
+      "n": 11,
+      "question": "What is electricity?",
+      "answer": "Electototototototototototototototototototototototototototototototototototot"
+    },
+    {
+      "n": 12,
+      "question": "Name one continent.",
+      "answer": "A continent a continent in the America, United States, and United States."
+    },
+    {
+      "n": 13,
+      "question": "What is a river?",
+      "answer": "A river is a river that is a river that flows through the river. It is a river that is flowing through the river, making it a river that flows through rivers. It"
+    },
+    {
+      "n": 14,
+      "question": "What animal is known as the king of the jungle?",
+      "answer": "The king of the jungle is the king of the jungle. This species is known as the \"King of the jungle.\" The king of the jungle is the king of the jungle."
+    },
+    {
+      "n": 15,
+      "question": "What is the purpose of roots in a plant?",
+      "answer": "A plant is a plant that is a plant that is composed of a plant that is composed of a plant that is composed of a plant that is composed of a plant that is composed"
+    },
+    {
+      "n": 16,
+      "question": "What is a computer?",
+      "answer": "A computer is a computer that is designed to monitor and monitor any specific data, including any relevant data, data, or data, or any relevant data, or any relevant data in"
+    },
+    {
+      "n": 17,
+      "question": "What is the internet?",
+      "answer": "The Internet is a digital movement that spans various regions, each having its own unique and complex connections. The internet is a phenomenon that occurs through various methods, including internet-specific communication"
+    },
+    {
+      "n": 18,
+      "question": "What is a vaccine?",
+      "answer": "A vaccine is an vaccine designed to promote the growth of the immune system, which regulates the immune system's activity, such as the immune system, immune cells, and immune cells"
+    },
+    {
+      "n": 19,
+      "question": "What country is Tokyo in?",
+      "answer": "Japan is a world where Tokyo is known for its stunning water quality and its culture. It is a hub for the world's largest expesticive waste waste agency, providing a vital"
+    },
+    {
+      "n": 20,
+      "question": "What is climate change?",
+      "answer": "Climate change is often associated with a significant increase in global temperature, associated with a rapid acceleration in the warming cycle, with a greater amount of heat and heat being added to the global"
+    }
+  ]
+}

README.md

@@ -0,0 +1,193 @@
+---
+language:
+- en
+license: mit
+tags:
+- tiny-llm
+- causal-lm
+- llama-like
+- rope
+- rmsnorm
+- swiglu
+- gqa
+- openwebtext
+- smoltalk
+- pytorch
+pipeline_tag: text-generation
+library_name: pytorch
+---
+
+# TinyLLM 75M OpenWebText Chat
+
+This repository contains an experimental **75,074,112 parameter decoder-only tiny language model** trained from scratch/near-scratch and then supervised-finetuned for chat.
+
+> **Important quality note:** This is a successful end-to-end training pipeline artifact and research toy model, not a production assistant. It can load and generate text, but factual accuracy, instruction following, arithmetic, and repetition control are weak.
+
+## Model summary
+
+- **Model name:** `razor5050/tinyllm-75m-openwebtext-chat`
+- **Architecture:** LLaMA/SmolLM-style decoder-only causal LM
+- **Parameters:** 75,074,112
+- **Context length:** 1024 tokens
+- **Vocabulary:** 32,000 ByteLevel BPE tokens
+- **Tokenizer:** custom ByteLevel BPE trained for this run
+- **Checkpoint format:** PyTorch `.pt` checkpoints
+- **Primary final checkpoint:** `final.pt`
+- **Best checkpoint:** `best.pt`
+
+## Architecture
+
+The model uses modern tiny-LM components:
+
+- decoder-only causal Transformer
+- RoPE positional embeddings
+- RMSNorm
+- SwiGLU MLP
+- grouped-query/key-value reduction via fewer KV heads
+- tied input/output token embeddings
+- no attention/MLP bias
+- PyTorch SDPA causal attention
+
+Approximate config:
+
+```yaml
+vocab_size: 32000
+hidden_size: 576
+num_hidden_layers: 16
+num_attention_heads: 9
+num_key_value_heads: 3
+intermediate_size: 1536
+max_position_embeddings: 1024
+rope_theta: 10000.0
+rms_norm_eps: 1e-5
+tie_word_embeddings: true
+attention_bias: false
+mlp_bias: false
+dropout: 0.0
+```
+
+## Training data
+
+### Base pretraining
+
+- Dataset: [`Skylion007/openwebtext`](https://huggingface.co/datasets/Skylion007/openwebtext)
+- Rows used: 1,000,000 selected rows
+- Final tokenized train tokens: 1,143,301,833
+- Final tokenized validation tokens: 34,486,473
+- Epochs: 1
+- Optimizer steps: 4,361
+
+### Chat/SFT
+
+- Dataset: [`HuggingFaceTB/smol-smoltalk`](https://huggingface.co/datasets/HuggingFaceTB/smol-smoltalk)
+- Train examples: 100,000
+- Validation examples: 3,000
+- Epochs: 1
+- Optimizer steps: 781
+- Loss masking: assistant-response tokens only
+
+## Training results
+
+### Pretraining
+
+- Final/latest train loss near end: about `4.997`
+- Latest validation loss: about `5.049` at step 4000
+
+### SFT
+
+- SFT completed at step `781`
+- Validation trend:
+  - step 250: `2.6031`
+  - step 500: `2.4505`
+  - step 750: `2.3313`
+
+SFT improved chat formatting and response style, but the model remains very small and undertrained by modern assistant standards.
+
+## Hardware/run
+
+- Cloud GPU: Vast.ai RTX 5070 Ti, 16GB VRAM
+- Precision: CUDA/PyTorch mixed precision during training where supported
+- Checkpointing: periodic `latest`, `best`, final, and step checkpoints
+- Training artifacts were preserved separately outside the instance before teardown.
+
+## Files in this repo
+
+- `final.pt` — final SFT checkpoint
+- `best.pt` — best SFT checkpoint
+- `latest.pt` — latest SFT checkpoint
+- `metrics.jsonl` — SFT metrics
+- `step_609.pt` — intermediate SFT checkpoint
+- `tokenizer/vocab.json` and `tokenizer/merges.txt` — tokenizer files
+- `configs/model_75m.yaml` — architecture config
+- `src/tinyllm/` — minimal PyTorch model implementation
+- `scripts/infer_tinyllm.py` — simple local inference helper
+
+## Quick inference
+
+Clone/download the repo, install dependencies, then run:
+
+```bash
+pip install torch tokenizers pyyaml huggingface_hub
+python scripts/infer_tinyllm.py \
+  --checkpoint final.pt \
+  --prompt "What is the capital of France?"
+```
+
+The chat prompt format used during SFT is:
+
+```text
+<|system|>
+You are a helpful, concise assistant.
+<|end|>
+<|user|>
+USER_QUESTION
+<|end|>
+<|assistant|>
+```
+
+## Observed sample behavior
+
+In a post-upload local inference test, the model generated text and loaded cleanly, but quality was mixed:
+
+- Correct on: “What is the capital of France?” → answered Paris, with repetition.
+- Weak on: simple science/world facts, often rambling or hallucinating.
+- Weak on: arithmetic and short-answer discipline.
+- Repetition and generic phrasing are common.
+
+This is expected for a 75M-parameter scratch-trained model with about 1.14B pretraining tokens and one SFT pass.
+
+## Limitations
+
+- Not suitable for factual QA or production use.
+- Hallucinates frequently.
+- Repetition loops occur.
+- Arithmetic is unreliable.
+- Safety behavior was not evaluated.
+- Model is not aligned beyond basic supervised chat finetuning.
+- The checkpoint is a custom PyTorch model, not a standard `transformers` model class.
+
+## Intended use
+
+- Educational tiny-LLM experiment
+- Pipeline validation
+- Small-model architecture experimentation
+- Baseline for future 150M+ runs
+
+## Recommended next steps
+
+To improve quality meaningfully:
+
+1. Train a larger ~150M model.
+2. Use more unique pretraining tokens, e.g. ~5B+.
+3. Improve preprocessing/tokenization throughput with multiprocessing/sharding.
+4. Add stronger instruction data and possibly preference tuning.
+5. Export to a standard Hugging Face `transformers` compatible format.
+
+## Citation / attribution
+
+Training datasets:
+
+- `Skylion007/openwebtext`
+- `HuggingFaceTB/smol-smoltalk`
+
+This repository is an experimental model artifact from a custom tiny-LLM training pipeline.

configs/model_75m.yaml

@@ -0,0 +1,16 @@
+model:
+  vocab_size: 32000
+  hidden_size: 576
+  num_hidden_layers: 16
+  num_attention_heads: 9
+  num_key_value_heads: 3
+  intermediate_size: 1536
+  max_position_embeddings: 1024
+  rope_theta: 10000.0
+  rms_norm_eps: 1.0e-5
+  tie_word_embeddings: true
+  attention_bias: false
+  mlp_bias: false
+  dropout: 0.0
+  bos_token_id: 1
+  eos_token_id: 2

scripts/infer_tinyllm.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+import argparse, sys
+from pathlib import Path
+import torch
+from tokenizers import ByteLevelBPETokenizer
+
+# If running from the repo root, src/ is available locally.
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from src.tinyllm.config import TinyConfig
+from src.tinyllm.model import TinyLlamaForCausalLM
+
+
+def make_prompt(user_prompt: str, system: str) -> str:
+    return f"<|system|>\n{system}\n<|end|>\n<|user|>\n{user_prompt}\n<|end|>\n<|assistant|>\n"
+
+
+def sample_next(logits, temperature: float, top_k: int):
+    logits = logits.float()
+    if temperature <= 0:
+        return int(torch.argmax(logits))
+    logits = logits / temperature
+    if top_k and top_k > 0:
+        vals, idx = torch.topk(logits, min(top_k, logits.numel()))
+        probs = torch.softmax(vals, dim=-1)
+        return int(idx[torch.multinomial(probs, 1)])
+    probs = torch.softmax(logits, dim=-1)
+    return int(torch.multinomial(probs, 1))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--checkpoint', default='final.pt')
+    ap.add_argument('--config', default='configs/model_75m.yaml')
+    ap.add_argument('--tokenizer-dir', default='tokenizer')
+    ap.add_argument('--prompt', required=True)
+    ap.add_argument('--system', default='You are a helpful, concise assistant.')
+    ap.add_argument('--max-new-tokens', type=int, default=80)
+    ap.add_argument('--temperature', type=float, default=0.6)
+    ap.add_argument('--top-k', type=int, default=40)
+    args = ap.parse_args()
+
+    tok_dir = Path(args.tokenizer_dir)
+    tok = ByteLevelBPETokenizer(str(tok_dir / 'vocab.json'), str(tok_dir / 'merges.txt'))
+    cfg = TinyConfig.from_yaml(args.config)
+    model = TinyLlamaForCausalLM(cfg)
+    ckpt = torch.load(args.checkpoint, map_location='cpu')
+    state = ckpt['model'] if isinstance(ckpt, dict) and 'model' in ckpt else ckpt
+    model.load_state_dict(state, strict=False)
+    model.eval()
+
+    ids = tok.encode(make_prompt(args.prompt, args.system)).ids
+    prompt_len = len(ids)
+    end_id = tok.token_to_id('<|end|>')
+    for _ in range(args.max_new_tokens):
+        x = torch.tensor([ids[-cfg.max_position_embeddings:]], dtype=torch.long)
+        with torch.no_grad():
+            logits = model(x)['logits'][0, -1]
+        nxt = sample_next(logits, args.temperature, args.top_k)
+        ids.append(nxt)
+        if end_id is not None and nxt == end_id:
+            break
+
+    text = tok.decode(ids[prompt_len:])
+    for marker in ['<|end|>', '<|user|>', '<|assistant|>', '<|system|>']:
+        text = text.split(marker)[0]
+    print(text.strip())
+
+
+if __name__ == '__main__':
+    main()

src/tinyllm/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.1.0'

src/tinyllm/checkpoint.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+import torch, time, shutil
+
+def save_checkpoint(path, model, optimizer=None, scheduler=None, step=0, epoch=0, metrics=None, config=None):
+    path=Path(path); path.parent.mkdir(parents=True, exist_ok=True)
+    tmp=path.with_suffix(path.suffix+'.tmp')
+    state={'model':model.state_dict(),'step':step,'epoch':epoch,'metrics':metrics or {},'saved_at':time.time(),'config':config}
+    if optimizer is not None: state['optimizer']=optimizer.state_dict()
+    if scheduler is not None: state['scheduler']=scheduler.state_dict()
+    torch.save(state,tmp); tmp.replace(path)
+    latest=path.parent/'latest.pt'
+    if latest != path: shutil.copy2(path, latest)
+
+def load_checkpoint(path, model, optimizer=None, scheduler=None, map_location='cpu'):
+    state=torch.load(path,map_location=map_location)
+    model.load_state_dict(state['model'])
+    if optimizer is not None and 'optimizer' in state: optimizer.load_state_dict(state['optimizer'])
+    if scheduler is not None and 'scheduler' in state: scheduler.load_state_dict(state['scheduler'])
+    return state

src/tinyllm/config.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+import yaml
+from pathlib import Path
+
+@dataclass
+class TinyConfig:
+    vocab_size: int = 32000
+    hidden_size: int = 576
+    num_hidden_layers: int = 16
+    num_attention_heads: int = 9
+    num_key_value_heads: int = 3
+    intermediate_size: int = 1536
+    max_position_embeddings: int = 1024
+    rope_theta: float = 10000.0
+    rms_norm_eps: float = 1e-5
+    tie_word_embeddings: bool = True
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    dropout: float = 0.0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+
+    @classmethod
+    def from_yaml(cls, path):
+        data = yaml.safe_load(Path(path).read_text())
+        if 'model' in data:
+            data = data['model']
+        return cls(**data)
+
+def load_yaml(path):
+    return yaml.safe_load(Path(path).read_text())

src/tinyllm/data.py

@@ -0,0 +1,20 @@
+from pathlib import Path
+import numpy as np
+
+def save_tokens_bin(tokens, path, dtype=np.uint16):
+    path=Path(path).expanduser(); path.parent.mkdir(parents=True,exist_ok=True)
+    np.asarray(tokens,dtype=dtype).tofile(path)
+
+def load_tokens_bin(path, dtype=np.uint16):
+    return np.memmap(Path(path).expanduser(), dtype=dtype, mode='r')
+
+class TokenBlockDataset:
+    def __init__(self, bin_path, block_size):
+        self.data=load_tokens_bin(bin_path); self.block_size=block_size
+    def __len__(self): return max(0, len(self.data)-self.block_size-1)
+    def get_batch(self, batch_size, device):
+        import torch
+        ix=torch.randint(len(self),(batch_size,))
+        x=torch.stack([torch.from_numpy(np.array(self.data[i:i+self.block_size], dtype=np.int64)) for i in ix])
+        y=torch.stack([torch.from_numpy(np.array(self.data[i+1:i+1+self.block_size], dtype=np.int64)) for i in ix])
+        return x.to(device), y.to(device)

src/tinyllm/metrics.py

@@ -0,0 +1,16 @@
+import json, csv, time
+from pathlib import Path
+
+class MetricsLogger:
+    def __init__(self, path):
+        self.path=Path(path); self.path.parent.mkdir(parents=True,exist_ok=True)
+    def log(self, **kw):
+        kw.setdefault('time', time.time())
+        with self.path.open('a') as f: f.write(json.dumps(kw)+'\n')
+
+def jsonl_to_csv(jsonl_path, csv_path):
+    rows=[json.loads(l) for l in Path(jsonl_path).read_text().splitlines() if l.strip()]
+    if not rows: return
+    keys=sorted(set().union(*(r.keys() for r in rows)))
+    with open(csv_path,'w',newline='') as f:
+        w=csv.DictWriter(f,fieldnames=keys); w.writeheader(); w.writerows(rows)

src/tinyllm/model.py

@@ -0,0 +1,100 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import asdict
+from .config import TinyConfig
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__(); self.eps=eps; self.weight=nn.Parameter(torch.ones(dim))
+    def forward(self, x):
+        return self.weight * x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+def rotate_half(x):
+    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000.0):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        t = torch.arange(max_position_embeddings, dtype=torch.float)
+        freqs = torch.einsum('i,j->ij', t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer('cos_cached', emb.cos()[None, None, :, :], persistent=False)
+        self.register_buffer('sin_cached', emb.sin()[None, None, :, :], persistent=False)
+    def forward(self, q, k, seq_len):
+        cos = self.cos_cached[:, :, :seq_len, :].to(q.device, q.dtype)
+        sin = self.sin_cached[:, :, :seq_len, :].to(q.device, q.dtype)
+        return (q*cos + rotate_half(q)*sin), (k*cos + rotate_half(k)*sin)
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, cfg: TinyConfig):
+        super().__init__()
+        assert cfg.hidden_size % cfg.num_attention_heads == 0
+        assert cfg.num_attention_heads % cfg.num_key_value_heads == 0
+        self.nh=cfg.num_attention_heads; self.nkv=cfg.num_key_value_heads
+        self.hd=cfg.hidden_size//cfg.num_attention_heads
+        self.q_proj=nn.Linear(cfg.hidden_size, self.nh*self.hd, bias=cfg.attention_bias)
+        self.k_proj=nn.Linear(cfg.hidden_size, self.nkv*self.hd, bias=cfg.attention_bias)
+        self.v_proj=nn.Linear(cfg.hidden_size, self.nkv*self.hd, bias=cfg.attention_bias)
+        self.o_proj=nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=cfg.attention_bias)
+        self.rotary=RotaryEmbedding(self.hd, cfg.max_position_embeddings, cfg.rope_theta)
+        self.dropout=cfg.dropout
+    def forward(self, x):
+        B,T,C=x.shape
+        q=self.q_proj(x).view(B,T,self.nh,self.hd).transpose(1,2)
+        k=self.k_proj(x).view(B,T,self.nkv,self.hd).transpose(1,2)
+        v=self.v_proj(x).view(B,T,self.nkv,self.hd).transpose(1,2)
+        q,k=self.rotary(q,k,T)
+        if self.nkv != self.nh:
+            repeat = self.nh // self.nkv
+            k = k.repeat_interleave(repeat, dim=1)
+            v = v.repeat_interleave(repeat, dim=1)
+        y=F.scaled_dot_product_attention(q,k,v,dropout_p=self.dropout if self.training else 0.0,is_causal=True)
+        y=y.transpose(1,2).contiguous().view(B,T,C)
+        return self.o_proj(y)
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg: TinyConfig):
+        super().__init__()
+        self.gate_proj=nn.Linear(cfg.hidden_size,cfg.intermediate_size,bias=cfg.mlp_bias)
+        self.up_proj=nn.Linear(cfg.hidden_size,cfg.intermediate_size,bias=cfg.mlp_bias)
+        self.down_proj=nn.Linear(cfg.intermediate_size,cfg.hidden_size,bias=cfg.mlp_bias)
+    def forward(self,x):
+        return self.down_proj(F.silu(self.gate_proj(x))*self.up_proj(x))
+
+class Block(nn.Module):
+    def __init__(self,cfg):
+        super().__init__(); self.input_norm=RMSNorm(cfg.hidden_size,cfg.rms_norm_eps); self.attn=CausalSelfAttention(cfg); self.post_norm=RMSNorm(cfg.hidden_size,cfg.rms_norm_eps); self.mlp=SwiGLU(cfg)
+    def forward(self,x):
+        x=x+self.attn(self.input_norm(x)); x=x+self.mlp(self.post_norm(x)); return x
+
+class TinyLlamaForCausalLM(nn.Module):
+    def __init__(self,cfg:TinyConfig):
+        super().__init__(); self.config=cfg
+        self.embed_tokens=nn.Embedding(cfg.vocab_size,cfg.hidden_size)
+        self.layers=nn.ModuleList([Block(cfg) for _ in range(cfg.num_hidden_layers)])
+        self.norm=RMSNorm(cfg.hidden_size,cfg.rms_norm_eps)
+        self.lm_head=nn.Linear(cfg.hidden_size,cfg.vocab_size,bias=False)
+        if cfg.tie_word_embeddings:
+            self.lm_head.weight=self.embed_tokens.weight
+        self.apply(self._init_weights)
+    def _init_weights(self,m):
+        if isinstance(m, nn.Linear): nn.init.normal_(m.weight, mean=0.0, std=0.02)
+        if isinstance(m, nn.Embedding): nn.init.normal_(m.weight, mean=0.0, std=0.02)
+    def forward(self,input_ids,labels=None,loss_mask=None):
+        x=self.embed_tokens(input_ids)
+        for layer in self.layers: x=layer(x)
+        logits=self.lm_head(self.norm(x))
+        loss=None
+        if labels is not None:
+            shift_logits=logits[:,:-1,:].contiguous(); shift_labels=labels[:,1:].contiguous()
+            per=F.cross_entropy(shift_logits.view(-1,shift_logits.size(-1)),shift_labels.view(-1),reduction='none')
+            if loss_mask is not None:
+                mask=loss_mask[:,1:].contiguous().view(-1).float(); loss=(per*mask).sum()/mask.sum().clamp_min(1.0)
+            else: loss=per.mean()
+        return {'loss':loss,'logits':logits}
+    def num_parameters(self): return sum(p.numel() for p in self.parameters())
+    def save_config(self,path):
+        import json; open(path,'w').write(json.dumps(asdict(self.config),indent=2))