| | --- |
| | language: en |
| | license: mit |
| | tags: |
| | - tiny |
| | - language-model |
| | - causal-lm |
| | - pytorch |
| | datasets: |
| | - roneneldan/TinyStories |
| | - Skylion007/openwebtext |
| | pipeline_tag: text-generation |
| | library_name: transformers |
| | --- |
| | |
| | # TinyLM |
| |
|
| | A 3.4M parameter causal language model trained from scratch, for experimentation. |
| |
|
| | ## Architecture |
| |
|
| | | Hyperparameter | Value | |
| | |---|---| |
| | | Parameters | 3.403.968 | |
| | | Layers | 4 | |
| | | Hidden size | 64 | |
| | | Attention heads | 4 | |
| | | FFN dim | 192 | |
| | | Embedding rank | 32 | |
| | | Context length | 256 | |
| | | Tokenizer | GPT-2 (50257 vocab) | |
| |
|
| | Uses a **factored (low-rank) embedding** to keep the vocab projection from eating the entire parameter budget, with weight tying on the output head. |
| |
|
| | ## Training |
| |
|
| | | | | |
| | |---|---| |
| | | Datasets | Skylion007/openwebtext (10k samples), roneneldan/TinyStories (10k samples) | |
| | | Optimizer | AdamW (lr=3e-3, weight_decay=0.01) | |
| | | Scheduler | Cosine annealing with warm restarts | |
| | | Mixed precision | fp16 (torch.cuda.amp) | |
| | | Hardware | Nvidia P100 | |
| | |
| | ## Usage |
| | ```python |
| | from huggingface_hub import snapshot_download |
| | import importlib.util |
| | import torch |
| | |
| | # Download files |
| | snapshot_download(repo_id="Fu01978/TinyLM", local_dir="./tinylm") |
| |
|
| | # Load via script |
| | spec = importlib.util.spec_from_file_location("modeling_tinylm", "./tinylm/modeling_tinylm.py") |
| | module = importlib.util.module_from_spec(spec) |
| | spec.loader.exec_module(module) |
| |
|
| | model, tokenizer, config = module.load_tinylm("./tinylm") |
| | model.eval() |
| | |
| | # Generate |
| | output = module.generate(model, tokenizer, "Once upon a time, ") |
| | print(output) |
| | ``` |
