This repository contains implementation of the Gemma 3 270M architecture, featuring modern transformer optimizations including sliding window attention, rotary position embedding (RoPE), and root mean square normalization (RMSNorm). The model achieves efficient performance while maintaining the core capabilities of larger language models.
Model Architecture
Core Specifications
Parameters: 270M total (170M embedding + 100M transformer)
Layers: 18 transformer blocks
Attention Heads: 4 query heads, 1 key-value 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
Key Architectural Features
Hybrid Attention Pattern
15 Sliding Window Attention layers: Efficient local attention with 512-token window
3 Full Attention layers: Global context at positions 6, 12, and 18
Reduces computational complexity from O(n²) to O(n×W) for most layers
Advanced Components
RoPE (Rotary Position Embedding): Preserves token semantics while encoding positional information
RMSNorm: Stable normalization with zero-centered weights and (1+w) scaling
1# Load and tokenize dataset2from datasets import load_dataset
3ds = load_dataset("roneneldan/TinyStories")45# Initialize model6model = Gemma3Model(GEMMA3_CONFIG_270M)78# Train for 60,000 iterations9# (See full training loop in the notebook)
Inference
python
1# Load trained model2model = Gemma3Model(GEMMA3_CONFIG_270M)3model.load_state_dict(torch.load("best_model_params.pt"))45# Generate text6sentence ="Once upon a time there was a pumpkin."7context = torch.tensor(enc.encode_ordinary(sentence)).unsqueeze(0)8output = model.generate(context, max_new_tokens=200)9generated_text = enc.decode(output.squeeze().tolist())
Implementation Details
RoPE Implementation
Dual frequency bases: 10,000 (local) and 1,000,000 (global)
Preserves semantic meaning while encoding positional relationships
Applied to both query and key vectors before attention computation
Sliding Window Attention
Causal mask combined with sliding window constraint
Tokens attend to recent W=512 tokens plus all previous tokens for full attention layers
Significant memory and computation savings for long sequences
RMSNorm Features
Zero-centered weight initialization
(1 + weight) scaling during forward pass
Float32 computation with dtype preservation
Optional bias parameters
Training Results
The model was successfully trained for 60,000 iterations with:
Converging training and validation loss curves
Stable gradient norms with clipping
Effective learning rate scheduling
Generated coherent stories in the TinyStories domain
File Structure
├── gemma_3_270_m_slm_from_scratch.py # Main implementation
├── train.bin # Processed training data
├── validation.bin # Processed validation data
├── best_model_params.pt # Best model checkpoint
└── README.md # This file
Key Innovations Implemented
Efficient Architecture: Hybrid attention pattern balances performance and computational cost
Modern Optimizations: RoPE, RMSNorm, and grouped query attention
Training Stability: Careful hyperparameter tuning and gradient management
Memory Efficiency: Quantization-aware training ready and memory-mapped data loading
Performance Characteristics
Memory Usage: ~550MB RAM for inference
Training Time: Approximately 60,000 iterations on GPU
Generation Speed: Fast inference suitable for edge deployment
Specialization Ready: Architecture optimized for task-specific fine-tuning
