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)
-
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
1from gemma3_tinystories import HFGemma3Model, Gemma3Config
2import tiktoken
3import torch
4
5config = Gemma3Config.from_pretrained("Shubhamw11/Gemma-270M-TinyStories")
6model = HFGemma3Model.from_pretrained("Shubhamw11/Gemma-270M-TinyStories", config=config).model
7tokenizer = tiktoken.get_encoding("gpt2")
1
2#define the device
3device = "cuda" if torch.cuda.is_available() else "cpu"
4
5input_text = "Once upon a time, there was a little"
6context = torch.tensor(tokenizer.encode(input_text), dtype=torch.long).unsqueeze(0).to(device)
7model.to(device)
8response = model.generate(context, max_new_tokens=200, temperature=1.1, top_k=5)
9
10print(tokenizer.decode(response.squeeze().tolist()))