CoolFace
Modelpublic

Se00n00/TinyLM2-50M-Base

sourceHugging Facemitupdated 17d agoView on Hugging Face
0likes270downloads
Model Card

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 lmeval to run them <img src="losschart.png">

Model Architecture & Hyperparameters

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

HyperparameterValueDescription
ArchitectureALiBi Decoder-Only TransformerAutoregressive Decoder-Only Transformer
Total Parameters~50.96M (53,430,272)Compact and ultra-fast for edge & local CPU/GPU
inference
`vocab_size`50,271Includes special chat tags (`<SYSTEM>, <USER>, <ASSISTANT>`)
`hidden_size` (`d_model`)512Model hidden dimension
`intermediate_size` (`ff_hidden_d`)819SwiGLU Gated Feedforward hidden dimension
`num_hidden_layers`12Number of Transformer block layers
`num_attention_heads`8Attention heads (Head dim = 64)
`max_position_embeddings`2,048Maximum context sequence length
NormalizationRMSNorm (eps=1e-8)Scale normalization for accelerated throughput
Activation FunctionSwiGLU (SiLU)Gated Feedforward activation
Positional EncodingALiBiAttention with Linear Biases
Tie Word EmbeddingsTrueTied 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.

PropertyValue
Tokenizer TypeGPT2Tokenizer (Byte-Level BPE)
Vocabulary Size50,271
Special Tokens`<START> <END> <UNK>`
Chat Control Tokens`<SYSTEM> <USER> <ASSISTANT>`
Extra Special Tokens`<THINK> </THINK> <AVAILABLE_TOOLS> </AVAILABLE_TOOLS> <TOOL_CALLS> </TOOL_CALLS> <TOOL_RESULTS> </TOOL_RESULTS>`
Chat TemplateNative Jinja2 support via tokenizer.apply_chat_template()

## Training Configuration

ParameterValue
Pipeline ProcessPre-Training (PT)
DatasetHuggingFaceFW/fineweb-edu (sample-100BT), Salesforce/wikitext (wikitext-2-v1)
Total Tokens4,000,000,000 (4B)
Epochs1
Learning Rate1e-4
Learning Rate ScheduleCosine (warmup_ratio=0.01)
Micro-Batch Size2 per device
Gradient Accumulation16 steps
Effective Batch Size32 × 2,048 tokens
OptimizerAdamW (weight_decay=0.1)
Max Sequence Length2,048 tokens
Precisionfloat16
HardwareNVIDIA Tesla T4 x 2 GPU

## Inference

python
  # 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. …