Promote Sign in Sign up

Repository

transformer-mamba-llm-research

🚀 Transformer-Mamba Hybrid LLM Research architectures for efficient language modeling & research. : Alternating Transformer and Mamba layers : Mixed precision, gradient scaling, an · License unknown · Install via pip · Maintenance unknown

vukrosic
Repository
License
Unknown

README

🚀 Transformer-Mamba Hybrid LLM Research

LLM combining Transformer attention and Mamba SSM architectures for efficient language modeling & research.

✨ Features

  • 🧑‍💻 Only 252 Lines of Code
  • 🔄 Hybrid Architecture: Alternating Transformer and Mamba layers
  • ⚡ Efficient Training: Mixed precision, gradient scaling, and optimized data loading
  • 🎯 Flexible Patterns: Configurable layer arrangements (e.g., "MMAMAMAM")
  • 🚀 Multi-GPU Support: Automatic DataParallel for multi GPUs (only 1 tested)
  • 💾 Easy Inference: Interactive chat mode and text generation

️ Architecture

The model alternates between:

  • Mamba SSM: State space model with convolution and simplified parallel processing
  • Transformer: Multi-head attention with causal masking

Quick Start

1. Install Dependencies

pip install -r requirements.txt

2. Train the Model

python hybrid_llm.py

3. Run Inference

python inference.py --interactive

or

python inference.py --prompt "The future of AI is"

📊 Model Configuration

  • Hidden Size: 384
  • Layers: 8 (configurable pattern)
  • Sequence Length: 512
  • Vocabulary: Auto-detected from tokenizer
  • Parameters: ~2.5M (configurable)

🎮 Usage Examples

Interactive Chat

python inference.py --model model.pt --interactive

Text Generation

python inference.py --model model.pt --prompt "The future of AI is"

📁 Project Structure

├── hybrid_llm.py    # Training script and model definition
├── inference.py     # Inference and chat interface
├── gpu_monitor.py  # GPU monitoring utilities
└── requirements.txt # Python dependencies

Customization

Modify HybridConfig in hybrid_llm.py to adjust:

  • Model dimensions
  • Layer patterns
  • Training parameters
  • Architecture choices

📈 Performance

  • Training: Optimized with AMP, gradient clipping, and efficient data loading
  • Inference: Fast generation with temperature and top-k sampling
  • Memory: Efficient attention and SSM implementations

Contributing

This is a research project - feel free to experiment with different architectures and configurations!

📄 License

See LICENSE file for details.