my-LLM

Overview

my-LLM is a compact language model developed from scratch as an educational initiative, featuring approximately 30 million parameters in a decoder-only transformer architecture. It is trained on the TinyStories dataset, leveraging Kaggle's free T4 GPUs.

This project aims to provide an in-depth understanding of transformer models by manually implementing key components, including the model architecture, training loop, data pipeline, and tokenizer. A series of controlled experiments will follow, culminating in a detailed technical write-up of findings and insights.

Key Features

• Built from the ground up: The implementation relies solely on PyTorch tensor operations, with no use of nn.Transformer or high-level model classes from Hugging Face.
• Efficient architecture: The multi-head attention mechanism is vectorized for efficient batch processing, adhering to industry practices.
• Modular design: Each component is encapsulated within its own nn.Module, ensuring easy testing and integration against established tensor shape contracts.
• Configurability and reproducibility: Hyperparameters and configurations are managed via versioned YAML files, while a fixed random seed ensures reproducibility during development.
• Budget-friendly: The model is designed to run effectively within a reasonable compute budget, with 30 million parameters providing a balance of speed and narrative coherence.

Objectives

• Develop a complete decoder-only transformer without high-level abstractions.
• Train the model to generate coherent narratives in the style of TinyStories on a single T4 GPU.
• Conduct ablations to evaluate which design choices significantly impact performance, such as transformer normalization techniques and learning rate schedules.
• Explore the model's scaling characteristics through experiments with varying size, data, and compute resources.
• Document findings in a comprehensive technical blog.

Implementation Status

The following functionalities have been successfully implemented and tested:

• BPE tokenizer with encoding/decoding verification
• Data exploration, including story count and character statistics
• YAML configuration loader
• Comprehensive transformer architecture (including embeddings, multi-head attention, feed-forward networks)
• Training loop with advanced features like AdamW optimizer and learning rate scheduling
• Kaggle integration for seamless training usage
• Baseline model training completed: 30M parameters with measurable evaluation metrics

Project Directory Structure

my-LLM/
├── configs/              # YAML experiment configurations
├── notebooks/            # Jupyter Notebooks for exploration and training
├── scripts/              # Utilities for datasets and model generation
├── src/                 # Core source code for model, training, and tokenizer
└── requirements.txt      # List of dependencies

Experimental Insights

Performance metrics guide the model’s evolution:

• Significant understanding of normalization techniques (RMSNorm vs LayerNorm) has been documented, showing both methods yield comparable results, yet LayerNorm offers faster computation.
• Ablations on warmup lengths affirm that under certain conditions, warmup steps do not substantially impact training stability.
• Learning rate adjustments reveal peak performance occurs around 1e-3, emphasizing the importance of hyperparameter tuning.
• A comprehensive scaling study suggests a fixed computational budget impacts model performance, necessitating further exploration of optimization techniques.

For a detailed exploration of model capabilities and architectural choices, the accompanying write-up titled What a 30M-Parameter Language Model Taught Me About Training LLMs is available.