SmolMixtral - Mixtral Inspired Model

• So, I trained a MoE based a 124M (8x12M) architecture I coded from ground up.
• Trained on TinyStories dataset from HuggingFace consisting of 1M texts for a total of 14000 steps

Examples

📊 Training Results & Model Weights

• Hugging Face Model: YuvrajSingh9886/SmolMixtral
• WandB Checkpoints: Check the WandB report above for additional trained model checkpoints

Features

• Flash Attention: Efficient attention mechanism with memory optimization
• Mixture of Experts (MoE): 8 experts with top-2 routing and noisy top-k support
• SWiGLU Activation: Advanced activation function in expert layers
• Rotary Positional Embeddings: Position encoding for sequence understanding
• Liger Kernels: Optimized kernels for faster training (optional)
• Distributed Training: Support for multi-GPU training with DDP
• Advanced Optimizer: AdamW optimizer with custom learning rate scheduling
• Gradio Interface: Interactive web interface for text generation

Model Architecture

Default Configuration

• Embedding Dimensions: 512
• Decoder Layers: 8
• Attention Heads: 8
• MoE Experts: 8 (top-2 routing)
• Block Size: 1024 tokens
• Vocabulary Size: Based on Llama-2-7b tokenizer (~32,000 tokens)
• Batch Size: 16

Full Parameter List

Model Architecture Parameters

• epochs: Number of training epochs (default: 4)
• block_size: Maximum sequence length (default: 1024)
• batch_size: Training batch size (default: 16)
• embeddings_dims: Model embedding dimensions (default: 512)
• no_of_heads: Number of attention heads (default: 8)
• no_of_decoder_layers: Number of decoder layers (default: 8)
• attn_dropout: Attention dropout rate (default: 0.1)
• dropout: General dropout rate (default: 0.1)

Mixture of Experts (MoE) Parameters

• experts: Number of MoE experts (default: 8)
• top_experts: Number of experts to route to (default: 2)
• noisy_topk: Use noisy top-k routing (default: False)

Training Hyperparameters

• max_lr: Maximum learning rate (default: 6e-4)
• weight_decay_optim: Weight decay for optimizer (default: 0.01)
• beta_1: Beta1 for optimizer (default: 0.9)
• beta_2: Beta2 for optimizer (default: 0.95)
• eps: Epsilon for optimizer (default: 1e-8)
• clip: Gradient clipping value (default: 1.0)

System Configuration

• device: Device to use (default: 'cuda:9')
• use_checkpointing: Use gradient checkpointing (default: False)
• use_liger: Use Liger kernels for optimization (default: True)
• use_flash_attention: Use Flash Attention (default: True)
• use_compile: Use torch.compile (default: True)

Data Configuration

• vocab_size: Vocabulary size (default: based on tokenizer + 768)
• val_epochs: Validation frequency (default: 2)

Quick Start

Installation

Important: Hugging Face Token Setup

1. Get a Hugging Face Token:
   • Go to Hugging Face Settings
   • Create a new token with "Read" permissions
   • Accept the Llama-2 license at meta-llama/Llama-2-7b-hf
2. Set your token in config.py:
   TOKEN = 'your_token_here'

Using Pre-trained Weights

1. Download Model Weights:
   • Option 1: Download from Hugging Face - YuvrajSingh9886/SmolMixtral
   • Option 2: Visit the WandB Training Report for additional checkpoints
   • Place downloaded files in the checkpoints/ directory
2. Load Pre-trained Model for Inference:
   bash

   1# Using the Gradio web interface 2cd gradio 3python app.py 4 5# Or use in your own code 6python inference.py

Training Examples

Basic Training (Single GPU)

Training with Custom Parameters

Multi-GPU Distributed Training

Inference with Gradio

File Structure

SmolMixtral/
├── config.py          # Model configuration and hyperparameters
├── model.py           # Model architecture (Mixtral, MoE, Attention, etc.)
├── data.py           # Data loading and preparation
├── inference.py      # Inference functions and text generation
├── trainer.py        # Main training loop with DDP support
├── install.sh        # Setup script
├── requirements.txt  # Python dependencies
├── model_summary.py  # Model architecture summary
├── gradio/
│   └── app.py        # Gradio web interface
├── checkpoints/      # Model checkpoints
├── generated_data/   # Generated text outputs
├── images/           # Project images
└── old/             # Original files

Training Features

• Gradient Accumulation: Configurable batch size scaling
• Learning Rate Scheduling: Cosine decay with warmup
• Gradient Clipping: Prevents gradient explosion
• Wandb Integration: Experiment tracking and logging
• Checkpointing: Regular model checkpoints during training
• Loss Calculation: Optimized cross-entropy with padding token handling
• Distributed Training: Multi-GPU support with DDP
• Memory Optimization: Gradient checkpointing support

Generation Methods

1. Top-k Sampling: Traditional sampling with temperature control

Advanced Usage

Configuration

Custom Dataset Training

Monitoring and Logging

Performance Tips

1. Use Liger Kernels: Keep use_liger = True for optimized operations
2. Flash Attention: Keep use_flash_attention = True for memory efficiency
3. Gradient Checkpointing: Use use_checkpointing = True for memory-constrained setups
4. Batch Size Tuning: Start with smaller batch sizes and increase gradually
5. Block Size: Larger block sizes improve quality but require more memory

Troubleshooting

Common Issues

Authentication Error (401)

Out of Memory (OOM)

Slow Training

Contributing

Requirements

• Python 3.8+
• PyTorch 2.0+
• Transformers
• Datasets
• Gradio
• Wandb
• Liger-kernel (optional)

License