Model Architecture

The goal of LiteGPT-25M is not to achieve state-of-the-art performance, but to provide a clean and understandable implementation of a modern Llama-style language model that can be trained from scratch and extended with further techniques in future experiments.

Overview

Model type: Decoder-only Transformer
Parameters: ~25M
Context length: 512
Vocabulary size: 16,384 (custom ByteLevel BPE)
Attention: GQA Flash-Attention (8 query / 4 key-value heads)
Positional Encoding: RoPE
Normalization: RMSNorm
Activation: SwiGLU

Architecture Diagram

Input Tokens [B, T]
        │
        ▼
┌─────────────────────┐
│  Token Embeddings   │
│  [vocab, d_model]   │
└─────────────────────┘
           │
           │
           ▼
╔══════════════════════════════╗
║ Transformer Block × 8        ║
║                              ║
║       RMSNorm                ║
║           │                  ║
║           ▼                  ║
║     GQA Flash Attention      ║
║           │                  ║
║           ▼                  ║
║       Residual Add           ║
║           │                  ║
║           ▼                  ║
║       RMSNorm                ║
║           │                  ║
║           ▼                  ║
║        SwiGLU MLP            ║
║           │                  ║
║           ▼                  ║
║      Residual Add            ║
╚══════════════════════════════╝
            │
            ▼
┌─────────────────────┐
│   Final RMSNorm     │
└─────────────────────┘
            │
            ▼
┌─────────────────────┐
│      LM Head        │
└─────────────────────┘
            │
            ▼
      Logits [B,T,V]

Configuration

Parameter	Value
batch_size	64
grad_accum_steps	2
n_layers	8
d_model	448
n_heads	8
n_kv_heads	4
head_dim	56
ffn_dim	1152
context_length	512
vocab_size	16384

Transformer Block

Attention

Grouped Query Attention (GQA) with 8 query heads and 4 key-value heads
FlashAttention for memory-efficient computation
RoPE for positional encoding
Causal masking

Feed Forward Network (SwiGLU)

FFN(x) = W2(SiLU(W1(x)) ⊙ W3(x))

Expansion ratio: ~2.57×

Residual Connections

x = x + Attention(x)

x = x + FFN(x)

Normalization

RMSNorm (pre-attention and pre-FFN)

Parameter Count

Component	Params
Token Embeddings	(n_vocab x d_model) 16384 x 448 = 7,340,032
Attention (GQA)	[(448x448) + (448x224) + (448x224) + (448x448)] x 8 = 4,816,896
SwiGLU FFN	[(448x1152) + (448x1152) + (1152x448)] x 8 = 12,386,304
RMSNorm	(448 x 2) x 8 = 7,168
Final RMSNorm	448
LM Head	weight tied with token embeddings
Total	~24.6M

Design Decisions

This model incorporates a subset of modern LLM architectural improvements while remaining small enough to train from scratch on a single NVIDIA A5000 GPU.

Modern Decoder Architecture

The model uses:

Learned token embeddings
Rotary Positional Embeddings (RoPE)
Grouped Query Attention (GQA)
FlashAttention
SwiGLU feed-forward networks
RMSNorm
Causal masking
Weight tying between token embeddings and LM head

Modern Features With Minimal Complexity

Several architectural improvements from models such as Llama, Mistral, and Gemma are included because they provide meaningful gains in training efficiency, model quality, or inference performance without substantially increasing implementation complexity.

Key improvements over GPT-2 include:

RoPE for improved positional encoding and length extrapolation
GQA for reduced KV cache size and improved inference efficiency
FlashAttention for faster and more memory-efficient attention computation
SwiGLU for improved feed-forward expressiveness
RMSNorm for simpler and more efficient normalization

Small Scale Training

The model is designed to train on a single NVIDIA A5000. Model size, context length, and batch size are selected to balance training speed, memory usage, and model quality within a constrained compute budget.

Dataset

Metric	Value
Total Tokens	~500M
Tokenizer	Custom ByteLevel BPE
Vocabulary Size	16,384
Context Length	512
Train Split	90%
Validation Split	10%
Storage Format	uint16

Dataset Composition

Dataset	Tokens	Percentage
FineWeb	300M	60%
TinyStories	200M	40%

Data Pipeline

FineWeb (300M)
TinyStories (200M)
        │
        ▼
┌──────────────────┐
│ Dataset Mixing   │
└────────┬─────────┘
         │
         ▼
┌──────────────────┐
│ ByteLevel BPE    │
│ (vocab 16,384)   │
└────────┬─────────┘
         │
         ▼
      Token IDs
         │
         ▼
┌──────────────────┐
│ Append EOT Token │
└────────┬─────────┘
         │
         ▼
┌──────────────────┐
│ Random 90/10     │
│ Train / Val      │
└────────┬─────────┘
         │
    ┌────┴────┐
    ▼         ▼
train.bin  val.bin
(uint16)   (uint16)

Tokenizer

Type: ByteLevel BPE (trained on the corpus)
Vocabulary size: 16,384
Special tokens: <|endoftext|>
Tokens stored as uint16

Training

Hardware

GPU: RTX A5000 x1
Platform: Runpod

Hyperparameters

Parameter	Value
Batch Size	64
Gradient Accumulation	2
Effective Batch Size	128
Sequence Length	512
Learning Rate	6e-4
Min Learning Rate	6e-5
Weight Decay	0.1
Warmup Steps	4000
Max Steps	40000
Eval Interval	500
Eval Iters	100
Grad Clip	1.0

Tokens Seen

64 x 512 x 40000 = 1,310,720,000 tokens ~ 1.3B tokens

Optimizer

AdamW (β1 = 0.9, β2 = 0.95)

Learning Rate Schedule

Cosine Decay with linear warmup (4000 steps warmup, 36000 steps decay)

Mixed Precision

BF16 / FP16

Loss Function

Cross Entropy Loss

Results

Metric	values
Train Loss	2.6853755712509155
Val Loss	2.8251807928085326
Perplexity	16.863993599262944
tokens_per_sec	239081.02266364102
best_val_loss	2.8036279249191285

LiteGPT Evaluation Suite

Language modeling

Settings: temperature=0.8, top_k=50, max_tokens=100

Prompt: The capital of France is

Output: The capital of France is well known for its great geographical experience. The historical setting is not one of a sort of the greatest of all the times, but it's a central part of the history of the country. It's a historical town that has seen a huge shift inside this famous city of the same name. The city is known as "The Dolceus" and also known as the "The Dolceus." It's a place where a lot of people live and live,

Prompt: Once upon a time, there lived

Output: Once upon a time, there lived a little girl named Lily. She was very excited because she was going on a trip to see Grandma. She helped her mom pack their luggage with clothes and toys. Then, they got in the car and drove to Grandma's house. As they drove, Lily saw something strange. It was a big, gray cloud that seemed to float away. She pointed and asked her mom what it was. Her mom said it was a mild cloud, but it was a very hot day.

Prompt: The theory of evolution states that

Output: The theory of evolution states that a single model is responsible for some type of behavior, not all animals, not all human beings. It is not the same concept to decide the same. It does not matter if it is the same in a species or a family, or the exact same animal. When it is no longer possible to choose one species, it is only a matter of determining the kind of a single species. |We can try to make a simple concept based on a model, name or model (or model as

Prompt: Machine learning is

Output: Machine learning is part of the same, but the reason is that each semester is different, and the number of students attending is different. This is the reason for course selection and inclusion among those at lower-level learning. So, we need to make sure we understand what the learning system is, how it is being presented and what we use to help others find their own kindergarten to finish it. So, we will get to see how each is different. We will see what our teachers and

Reasoning

Settings: temperature=0.4, top_k=40, max_tokens=100

Prompt: If John has 5 apples and gives 2 to Alice, then John has

Output: If John has 5 apples and gives 2 to Alice, then John has 5 apples and 1 to Bob, then Alice is 3 to Bob, then Alice is 3 to Bob and Bob is 3 to Bob. A few days ago, I made a couple of apples and added them to the bottom of the apple. I think I did a pretty good job with the apple, but I didn't have a lot of apples, so I just added the apples and added the apples. I made a few more apples, but I think I would add a lot more apples

Prompt: The next number in the sequence is: 2, 4, 8, 16,

Output: The next number in the sequence is: 2, 4, 8, 16, 17, 17, 17, 17, 18, 18, 18, 18, 18, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,

Prompt: Earth is to Solar System as electron is to

Output: Earth is to Solar System as electron is to be able to generate a high-level of energy and solar power. The solar system is designed to generate a high-level of energy, which is then used to generate a high-level of energy. Solar systems can generate electricity, which is then used to generate a high-level of energy. Solar systems are used to generate electricity, which is then used to generate electricity. Solar systems are used to generate electricity. The solar system is used to generate electricity which is then

World knowledge

Settings: temperature=0.4, top_k=40, max_tokens=100

Prompt: Python is a programming language that

Output: Python is a programming language that is used to communicate to the audience. The language is translated into English. The language is translated into English. The language is translated into English. The language is translated into English. The language is translated into English. The language is translated into English. The language is translated into English. The language is translated into English. The language is translated into English. The language is translated into English. The language is translated into English. The language is translated into

Prompt: The Internet works because

Output: The Internet works because it’s a platform that helps people with disabilities. What is a mobile app? A mobile app is a platform that allows people to use it to work, work, and play. It allows people to use it to work, and it allows them to use it to work. What is a mobile app? A mobile app is a platform that allows people to use it to work, work, and play. It is a platform that allows people to use it to work, work

Prompt: The largest planet in our solar system is

Output: The largest planet in our solar system is the Earth's solar system. It's a great solar system for solar panels and solar panels. It's also a great solar power for solar panels and solar panels. The solar system is a great solar system for solar panels and solar panels. It's also a great solar system for solar panels and solar panels. The solar system is a great solar power system for solar panels and solar panels. It's also a great solar power system for solar panels and solar panels. The solar system is

Coding

Settings: temperature=0.2, top_k=20, max_tokens=150

Prompt: def factorial(n):

Output: def factorial(n): a) the use of the term "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "condition" in the context of a "cond

Prompt: class LinkedList:

Output: class LinkedList: - LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: LinkedList: Linked

Prompt: for i in range(10):

Output: for i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(10): - i in range(

Prompt: #include int main() {

Output: #include int main() {\iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS/iOS

Prompt: function quickSort(arr) {

Output: function quickSort(arr) {iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone(s) iphone

HTML

Settings: temperature=0.2, top_k=20, max_tokens=150

### Prompt: <!DOCTYPE html> <html>
- **Output:** <!DOCTYPE html> <html><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a><a>

Markdown

Settings: temperature=0.8, top_k=50, max_tokens=100

Prompt: # Neural Networks

Output: # Neural Networks. There are several networks that provide the same network as the other networks. At the same time, a network can provide a network with the same network as a wireless network (SMT). What's the difference between VoIP networks? The VoIP network has an IP address that connects both networks (i.e. IP address and IP address) and the Internet address. The network is the network's IP address that connects to a certain network network. This connects

Story completion

Settings: temperature=0.9, top_k=100, max_tokens=200

Prompt: The dragon opened its eyes and

Output: The dragon opened its eyes and looked around. It saw something shiny and colorful. It jumped and grabbed the object with its sharp claws. The dragon flew away and the dragon breathed a sigh of joy. It flew back up into the sky and found its owner, who was so proud of it. He said the dragon had been an amazing hunter.

Prompt: The detective walked into the abandoned warehouse and

Output: The detective walked into the abandoned warehouse and turned right, halfway next to the counter—the warehouse floor was all finished. I wasn’t sure what the guy was doing in front of the counter, but I was having a very slight, deliberate grimacing moment. This might be the first time I saw an account from the FBI for a meeting with a customer, or someone from the FBI’s office. This was my first time seeing the FBI’s staffer work at an incident—one that is usually a business or a place of business. My goal is to solve the problem, and there’s not as much fun as I think. A lot of people that I’m not interested in hearing from have been in the FBI’s management room, and it’s been very, very good. When I needed to tell them where it was going, I was really drawn to a number of good habits, but still had to be careful. There’s been some

Dialogue

Settings: temperature=0.9, top_k=100, max_tokens=200

Prompt: Alice: Hello! Bob:

Output: Alice: Hello! Bob: I'm sorry. Wow. You're a little bit of an understatement to me. Maybe. I know how to keep you up by yourself, even though I need to share my stories. I know, I found myself having trouble finding relationships, since I'm usually with my best friend or coworkers. I know I've been in relationship situations since I last met my friend from college and I understand I can help. I know I have found the best relationship possible for a couple of weeks now, but it's actually all about finding love. The two of you and your best friend, Mildred P. Stewart (Pennsylvania Girl). That's where the two are: a close relationship, a relationship, and an intercourse. I love it when I'm with her, and her love. These two have just been on a good track. But lately, I actually have found several relationships. "To be truly love means to be happy,

TinyStories style

Settings: temperature=0.9, top_k=100, max_tokens=200

Prompt: Tom had a little red ball.

Output: Tom had a little red ball. Lily had a red ball and Tom had a blue ball. They both loved to play with their toys. But they both wanted the red ball. They both pulled it from each other. "Give me the red ball!" Lily shouted. "It's mine!" "No, it's mine!" Tom yelled. "It's mine!" They pulled so hard that they separate the red ball from each other. They both fell down and started to cry. Mom heard the noise and came to see what was wrong. She saw the red ball on the floor and the tears on Lily's face. She was not angry, but she was sad. "Why are you fighting?" Mom asked. "You both have the red ball, but you both have the same toy." "That is boring!" Lily said. "You both need to share and take turns." "Okay!" Mom said. "But you can both have the red ball

Prompt: Lily loved going to the park because

Output: Lily loved going to the park because it was always free. One day, she saw a boy crying and her knee started to hurt. 'Hello there, little girl,' the boy said, pointing to his knee. 'I'm sad because my knee hurts,' he said. The man smiled and said, 'Don't worry, I have a band-aid in my bag and I'll take it out to make you feel better.' The boy went back to his house and his mother put the band-aid on his knee. The man smiled and said, 'I'm happy to see this band-aid and I'm glad so you can help me!' The boy thanked his mother and helped the man get back in his shoes. The man said goodbye and the boy smiled as he skipped away.

Long-form continuation

Settings: temperature=0.8, top_k=50, max_tokens=100

Prompt: The Great Wall of China is a series of fortifications that were built across the historical northern borders of ancient Chinese states and Imperial China as protection against various nomadic groups from the Eurasian Steppe. Several walls were built from the 7th century BC, with selective stretches later joined together by Qin Shi Huang, the first emperor of China. Little of the Qin wall remains. Later on, many successive dynasties built and maintained multiple stretches of border walls. The best-known sections of the wall were built by the Ming dynasty.

Output: The Great Wall of China is a series of fortifications that were built across the historical northern borders of ancient Chinese states and Imperial China as protection against various nomadic groups from the Eurasian Steppe. Several walls were built from the 7th century BC, with selective stretches later joined together by Qin Shi Huang, the first emperor of China. Little of the Qin wall remains. Later on, many successive dynasties built and maintained multiple stretches of border walls. The best-known sections of the wall were built by the Ming dynasty. The Chinese Army was established in 1899, and became the first American Army occupied by the army. The city has a large population of 100,000 people; the entire city of China is the largest in the world. The city grew by 31 percent, with a total of 322,000 square feet. The most recent World War II military shrank the city's population, while the largest, and the largest, is the most populous city in China. In 1901,

Downloads last month: 257