🫐 Train Your Own Small Language Model

A minimal toolkit for training and using small language models with the Muon optimizer.

YouTube video: https://youtu.be/Dw0b0Kc9Kpk

🚀 Quick Start

Option 1: Google Colab (No Setup Required)

Click the badge above to run everything in your browser with free GPU access!

Option 2: Local Setup

# Clone and setup
git clone https://github.com/vukrosic/build-and-release-your-own-llm
cd build-and-release-your-own-llm
python setup.py  # Installs requirements and creates .env file

🎯 Four Ways to Use This Project

1. 🚀 Quick Start - Use My Pre-trained Model

Want to try text generation immediately?

# Install dependencies
pip install -r requirements.txt

# Run inference with my pre-trained model
python inference.py

The script will:

Show available checkpoints from vukrosic/blueberry-1
Download the model automatically
Let you generate text interactively

No setup required! The model downloads automatically.

2. 🏗️ Train Your Own Model (Single GPU)

Want to train from scratch on one GPU?

# Install dependencies
pip install -r requirements.txt

# Start training (takes ~20 minutes on GPU)
python train_llm.py

# Use your trained model
python inference.py

Your model will be saved in checkpoints/ and you can resume training anytime.

3. ⚡ Distributed Training (Multi-GPU)

Want to train faster with multiple GPUs?

# Install dependencies
pip install -r requirements.txt

# For 2x RTX 4090 or similar high-end GPUs:
python run_4090.py

# Monitor GPU utilization (in another terminal):
python gpu_monitor.py

# If you experience GPU imbalance, run diagnostics:
python fix_gpu_imbalance.py

Distributed Training Features:

Automatic GPU detection and memory optimization
Load balancing across multiple GPUs
Real-time monitoring of GPU utilization
Fallback options if distributed training fails
Memory-safe settings for RTX 4090 class GPUs

GPU Utilization Monitoring:
The gpu_monitor.py script shows real-time GPU stats:

[14:23:45] GPU Status:
  GPU 0: 🟢  87% util | 18234/24564MB (74.2%) | 67°C | NVIDIA GeForce RTX 4090
  GPU 1: 🟢  89% util | 18891/24564MB (76.9%) | 69°C | NVIDIA GeForce RTX 4090
  📊 Average: 88.0% | Range: 87%-89% | Imbalance: 2%

Want to share your model on Hugging Face?

# 1. Copy environment template
cp .env.example .env

# 2. Edit .env file:
# HF_REPO_NAME=your-username/your-model-name
# HF_TOKEN=hf_your_token_here
# PUSH_TO_HUB=true

# 3. Train (single or multi-GPU)
python train_llm.py          # Single GPU
# OR
python run_4090.py           # Multi-GPU

Get your HF token from: https://huggingface.co/settings/tokens

📁 Project Structure

├── train_llm.py              # Single GPU training with Muon optimizer
├── train_distributed_llm.py  # Multi-GPU distributed training
├── run_4090.py               # Optimized launcher for RTX 4090 class GPUs
├── gpu_monitor.py            # Real-time GPU utilization monitoring
├── fix_gpu_imbalance.py      # Diagnostic tool for GPU load balancing
├── inference.py              # Text generation and model loading
├── upload_to_hf.py           # Upload checkpoints to Hugging Face
├── example_usage.py          # Example workflow script
├── setup.py                  # Easy setup script
├── requirements.txt          # Python dependencies
├── .env.example             # Environment variables template
└── README.md                # This file

🎯 What You Get

21M parameter transformer model (384d, 6 layers, 8 heads)
Muon optimizer for efficient training
Automatic checkpointing every 5000 steps
Resume training from any checkpoint
Interactive text generation
Hugging Face integration (optional)

📊 Expected Results

Training time: ~16-20 minutes on modern GPU
Final perplexity: ~1.06
Model size: ~21M parameters
Memory usage: ~4-6GB GPU

🔧 Customization

Change Model Size

Edit train_llm.py:

@dataclass
class ModelConfig:
    d_model: int = 512      # Bigger model (was 384)
    n_layers: int = 8       # More layers (was 6)
    max_steps: int = 5000  # Train longer for better results (20000)

Use Your Own Data

Edit the dataset loading in train_llm.py:

# Replace this line:
dataset = load_dataset("HuggingFaceTB/smollm-corpus", "cosmopedia-v2", split="train", streaming=True)

# With your dataset:
dataset = load_dataset("your-dataset-name", split="train", streaming=True)

Adjust Training Speed

batch_size: int = 16        # Smaller = less memory
gradient_accumulation_steps: int = 8  # Larger = same effective batch size

📊 Understanding the Output

During Training

Training: 67%|██████▋   | 20000/30000 [12:34<06:15, 26.6it/s, loss=1.234, acc=0.876, ppl=3.4, lr=8.5e-03]

loss: Lower is better (target: ~1.0)
acc: Accuracy (target: ~98%)
ppl: Perplexity (target: ~1.1)
lr: Learning rate (automatically scheduled)

During Inference

Prompt: The future of AI is
Generated text: The future of AI is bright and full of possibilities. Machine learning algorithms continue to evolve...

🚨 Common Issues

"CUDA out of memory"

# In train_llm.py, reduce batch size:
batch_size: int = 12  # or even 8

"No checkpoints found"

Make sure you've run training first:

python train_llm.py  # Wait for it to complete
python inference.py  # Now this will work

"HF upload failed"

Check your token permissions:

Go to https://huggingface.co/settings/tokens
Make sure token has "Write" permission
Update your .env file

⚡ Distributed Training Guide

Understanding the Scripts

run_4090.py - The main launcher that:

Sets optimal memory settings for RTX 4090 class GPUs
Clears GPU memory to prevent fragmentation
Calls run_novita_4090_training() from train_distributed_llm.py
Provides fallback options if distributed training fails

train_distributed_llm.py - Contains the distributed training logic:

run_novita_4090_training() - Main entry point for multi-GPU training
launch_distributed_ddp() - PyTorch native DDP implementation (recommended)
launch_distributed() - Custom distributed training with Muon optimizer
main_ddp() - DDP training loop with standard optimizers
main() - Custom distributed training loop

Training Flow

run_4090.py → calls → run_novita_4090_training()
run_novita_4090_training() tries in order:
- PyTorch DDP (most reliable)
- Custom distributed training
- Single GPU fallback

GPU Utilization Troubleshooting

If you see uneven GPU utilization (e.g., 90% vs 19%):

# 1. Run diagnostics
python fix_gpu_imbalance.py

# 2. Monitor in real-time
python gpu_monitor.py

# 3. Try the balanced launcher
python run_balanced.py  # Created by fix_gpu_imbalance.py

Common Causes of GPU Imbalance:

Data loading bottlenecks
Memory fragmentation
Poor batch distribution
CPU affinity issues

Solutions Applied:

Rank 0 loads data, broadcasts to other GPUs
Improved distributed samplers with proper padding
Memory optimization and periodic cache clearing
CPU affinity optimization for data loading

Note on GPU Monitoring:
Different monitoring tools may show different utilization percentages:

Novita UI: May show instantaneous or averaged values
gpu_monitor.py: Shows real-time utilization every 2 seconds
nvidia-smi: Shows current utilization at query time

If you see discrepancies (e.g., Novita shows 90%/19% but gpu_monitor.py shows 90%/90%), this usually means:

The imbalance is intermittent (happening between measurements)
Different measurement intervals are being used
One tool is showing peak values while another shows averages

Recommendation: Run gpu_monitor.py for several minutes to see the true utilization pattern over time.

Advanced Configuration

Edit train_distributed_llm.py to customize:

# At the top of the file:
NUM_GPUS = 2                    # Number of GPUs to use
BASE_BATCH_SIZE = 6             # Batch size per GPU
BASE_LR = 0.01                  # Base learning rate
GPU_IDS = [0, 1]                # Specific GPU IDs to use

Performance Expectations

Single GPU (RTX 4090):

Training time: ~16-20 minutes
Memory usage: ~18-20GB
Utilization: ~85-95%

Dual GPU (2x RTX 4090):

Training time: ~8-12 minutes
Memory usage: ~18-20GB per GPU
Utilization: ~80-90% each (balanced)

🎉 What's Next?

Experiment with prompts - Try different starting texts
Adjust generation parameters - Change temperature and top_k in inference.py
Train on your data - Replace the dataset with your own text
Scale up - Increase model size for better performance
Multi-GPU training - Use distributed training for faster results
Monitor performance - Use gpu_monitor.py to optimize utilization
Share your model - Upload to Hugging Face for others to use

📦 Checkpoint Management

Automatic Checkpointing

The training script now saves checkpoints every 5000 steps in the checkpoints/ directory:

checkpoints/
├── checkpoint_step_5000/
│   ├── model.pt          # Model weights and optimizer state
│   ├── config.json       # Model configuration
│   └── tokenizer files   # Tokenizer configuration
├── checkpoint_step_10000/
└── checkpoint_step_15000/

Upload to Hugging Face

Share your trained models with the community:

# Set your Hugging Face token
export HF_TOKEN="hf_your_token_here"

# List available checkpoints
python upload_to_hf.py --list

# Upload latest checkpoint
python upload_to_hf.py --repo-name username/my-awesome-model

# Upload specific checkpoint
python upload_to_hf.py --repo-name username/my-model --checkpoint checkpoints/checkpoint_step_10000

# Create private repository
python upload_to_hf.py --repo-name username/my-model --private

Get your token from: https://huggingface.co/settings/tokens

Example Workflow

# Run the complete example
python example_usage.py

# Or step by step:
python train_llm.py                    # Train model (saves checkpoints)
python upload_to_hf.py --list          # See available checkpoints  
python upload_to_hf.py --repo-name username/model  # Upload to HF

💡 Pro Tips

Resume training: The script automatically detects checkpoints
Monitor GPU usage: Use nvidia-smi to check memory usage
Save compute: Use smaller models for experimentation
Better results: More training steps = better model (usually)
Checkpoint frequency: Adjust save_every in ModelConfig for different intervals
Share early: Upload intermediate checkpoints to track training progress

Happy training! 🚀

README