TinyLM2-50M

TinyLM2-50M is a compact decoder-only Transformer language model designed for efficient instruction following and conversational AI. The model has approximately 50M parameters and has been pre-trained on 4 billion tokens using the ALiBi decoder-only architecture, providing a lightweight platform for researching and understanding language-model behavior while remaining suitable for local inference and resource-constrained environments.

Evaluation

All evaluations are zero-shot unless stated otherwise, and i used lm_eval to run them

Model Architecture & Hyperparameters

TinyLM2-50M is built on a custom ALiBi Decoder-Only Transformer architecture with pre-normalization and gated feedforward networks:

Hyperparameter Value Description
Architecture ALiBi Decoder-Only Transformer Autoregressive Decoder-Only Transformer
Total Parameters ~50.96M (53,430,272) Compact and ultra-fast for edge & local CPU/GPU
inference
vocab_size 50,271 Includes special chat tags (`<
hidden_size (d_model) 512 Model hidden dimension
intermediate_size (ff_hidden_d) 819 SwiGLU Gated Feedforward hidden dimension
num_hidden_layers 12 Number of Transformer block layers
num_attention_heads 8 Attention heads (Head dim = 64)
max_position_embeddings 2,048 Maximum context sequence length
Normalization RMSNorm (eps=1e-8) Scale normalization for accelerated throughput
Activation Function SwiGLU (SiLU) Gated Feedforward activation
Positional Encoding ALiBi Attention with Linear Biases
Tie Word Embeddings True Tied input embedding and LM head projection

Tokenizer & Chat Template

The model uses a custom Byte-Level BPE Tokenizer equipped with special tokens and a pre-configured Jinja2 chat_template for multi-turn conversations.

Property Value
Tokenizer Type GPT2Tokenizer (Byte-Level BPE)
Vocabulary Size 50,271
Special Tokens `<
Chat Control Tokens `<
Extra Special Tokens `<
Chat Template Native Jinja2 support via tokenizer.apply_chat_template()

Training Configuration

Parameter Value
Pipeline Process Pre-Training (PT)
Dataset HuggingFaceFW/fineweb-edu (sample-100BT), Salesforce/wikitext (wikitext-2-v1)
Total Tokens 4,000,000,000 (4B)
Epochs 1
Learning Rate 1e-4
Learning Rate Schedule Cosine (warmup_ratio=0.01)
Micro-Batch Size 2 per device
Gradient Accumulation 16 steps
Effective Batch Size 32 × 2,048 tokens
Optimizer AdamW (weight_decay=0.1)
Max Sequence Length 2,048 tokens
Precision float16
Hardware NVIDIA Tesla T4 x 2 GPU

Inference

# pip install torch transformers
                                                                             
import torch                                                                                            
from transformers import pipeline                                                                       
                                                                                                        
pipe = pipeline(                                                                                        
  "text-generation",                                                                                  
  model="Se00n00/TinyLM2-50M",
  trust_remote_code = True                                                                            
)                                                                                                       
                                                                                                        
messages = [                                                                                            
  {"role": "system", "content": "You are a helpful AI assistant."},                                   
  {"role": "user", "content": "Explain artificial intelligence in simple terms."}                     
]                                                                                                       
                                                                                                        
prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)       
                                                                                                        
result = pipe(                                                                                          
  prompt,                                                                                             
  max_new_tokens=120,                                                                                 
  do_sample=True,                                                                                     
  temperature=0.5,                                                                                    
  top_k=40,                                                                                           
  top_p=0.9                                                                                           
)                                                                                                       
                                                                                                        
print(result[0]['generated_text'])                                                                  

──────

Sample Outputs

Raw next-token continuation (no chat template, temperature 0.8, top-p 0.9):

Prompt: Photosynthesis is the process by which

Plants take up oxygen and use it to make energy. Plant leaves, flowers, and even fruit can also convert carbon dioxide (CO2) into sugars. As you grow, plants are able to store the carbon dioxide from their leaves and convert it to sugars and carbohydrates. This is known as photosynthesis, and it is a process that converts the carbon dioxide into a form of energy. …

Downloads last month
226
Safetensors
Model size
79.2M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Datasets used to train Se00n00/TinyLM2-50M-Base

Collection including Se00n00/TinyLM2-50M-Base

Evaluation results