Gemma3 270M: Small Language Model Implementation from Scratch

This repository contains a PyTorch implementation of the Google's Gemma3 model with 270 million parameters, trained from scratch on the TinyStories dataset.

Github Repo Link

About

This model is based on Google DeepMind's Gemma3 architecture and was built from the ground up to explore training dynamics, architecture design, and generation quality of small LLMs. It includes advanced components such as:

Sliding Window Attention (512-token window)
Rotary Positional Embeddings (RoPE)
RMSNorm for stable training
Grouped Key-Value Attention (1 KV group)

Model Architecture

Parameters: 270M total (170M embedding + 100M transformer)
Layers: 18 transformer blocks
Attention Heads: 4 Query Heads, 1 KV Group
Hidden Dimension: 2048
Embedding Dimension: 640
Head Dimension: 256
Vocabulary Size: 50,257 (GPT-2 tokenizer)
Context Length: 32,768 tokens (trained with 128 block size)
Sliding Window: 512 tokens

Training Details

Dataset: TinyStories by Roneneldan
Steps: 150,000 steps (not epochs)
Batch Size: 32
Loss Function: Cross-Entropy
Optimizer: AdamW
LR Scheduler: Linear Warmup + Cosine Decay
Hardware: Single NVIDIA A100 GPU

Requirements

pip install git+https://huggingface.co/Shubhamw11/Gemma-270M-TinyStories

How to use

You can now load and use the model from the gemma3_tinystories library:

from gemma3_tinystories import HFGemma3Model, Gemma3Config
import tiktoken
import torch

config = Gemma3Config.from_pretrained("Shubhamw11/Gemma-270M-TinyStories")
model = HFGemma3Model.from_pretrained("Shubhamw11/Gemma-270M-TinyStories", config=config).model
tokenizer = tiktoken.get_encoding("gpt2")

Generate text


#define the device
device = "cuda" if torch.cuda.is_available() else "cpu"

input_text = "Once upon a time, there was a little"
context = torch.tensor(tokenizer.encode(input_text), dtype=torch.long).unsqueeze(0).to(device)
model.to(device)
response = model.generate(context, max_new_tokens=200, temperature=1.1, top_k=5)

print(tokenizer.decode(response.squeeze().tolist()))

Loading…

About

Sliding Window Attention (512-token window)

Rotary Positional Embeddings (RoPE)

RMSNorm for stable training

Grouped Key-Value Attention (1 KV group)

Model Architecture

Parameters: 270M total (170M embedding + 100M transformer)

Layers: 18 transformer blocks

Attention Heads: 4 Query Heads, 1 KV Group

Hidden Dimension: 2048

Embedding Dimension: 640

Head Dimension: 256

Vocabulary Size: 50,257 (GPT-2 tokenizer)

Context Length: 32,768 tokens (trained with 128 block size)

Sliding Window: 512 tokens

How to use

You can now load and use the model from the gemma3_tinystories library:

from gemma3_tinystories import HFGemma3Model, Gemma3Config import tiktoken import torch config = Gemma3Config.from_pretrained("Shubhamw11/Gemma-270M-TinyStories") model = HFGemma3Model.from_pretrained("Shubhamw11/Gemma-270M-TinyStories", config=config).model tokenizer = tiktoken.get_encoding("gpt2")

Generate text

#define the device device = "cuda" if torch.cuda.is_available() else "cpu" input_text = "Once upon a time, there was a little" context = torch.tensor(tokenizer.encode(input_text), dtype=torch.long).unsqueeze(0).to(device) model.to(device) response = model.generate(context, max_new_tokens=200, temperature=1.1, top_k=5) print(tokenizer.decode(response.squeeze().tolist()))