Skip to content

Latest commit

 

History

5 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 

Repository files navigation

CogAI4Sci - Custom Llama 3.2 Implementation

A high-performance, from-scratch implementation of Llama 3.2 (1B/3B) with instant startup times and optimized inference.

🚀 Features

  • From-scratch implementation of Llama 3.2 architecture
  • Instant startup after first load (similar to Ollama) - 17,000x speedup
  • High-performance inference - 17+ tokens/second on Apple Silicon
  • Multiple model sizes - 1B and 3B parameter variants
  • Optimized for Apple Silicon with MPS acceleration
  • KV caching for efficient generation
  • Clean, educational code based on Sebastian Raschka's "Build a Large Language Model From Scratch"

📊 Performance

Metric Performance
Startup Time 0.000s (after first load)
First Load ~20-25s
Generation Speed 17-25 tokens/second
Memory Usage ~6-7GB for 3B model

🔧 Installation

# Clone the repository
git clone https://github.com/yourusername/cogai4sci.git
cd cogai4sci

# Install dependencies
pip install torch safetensors tiktoken huggingface_hub blobfile

🏃 Quick Start

FastLLM - Instant Startup Version

from fast_llm import FastLLM

# Initialize model - first time loads model, subsequent times are instant!
model = FastLLM("meta-llama/Llama-3.2-1B-Instruct", use_float16=True)

# Generate text with streaming
prompt = "What is the future of AI?"
print(prompt, end="")

for token in model.generate(
    prompt, 
    max_new_tokens=100,
    temperature=0.8,
    top_k=40,
    use_kv_cache=True
):
    print(token, end="", flush=True)

Jupyter Notebook Usage

Open comprehensive_llama_demo.ipynb for a complete interactive guide with all features, benchmarks, and educational content.

📁 Project Structure

cogai4sci/
├── fast_llm.py                    # Main FastLLM class with instant startup
├── comprehensive_llama_demo.ipynb # Complete interactive demo notebook
└── README.md                      # This documentation

🔬 Key Implementation

FastLLM (fast_llm.py)

  • Class-level model caching for instant subsequent loads (similar to Ollama)
  • Optimized weight loading for both 1B and 3B models
  • Streaming token generation with proper tokenization
  • KV cache optimization for faster inference
  • Complete architecture including RoPE, GroupedQueryAttention, FeedForward
  • Proper dtype handling for Apple Silicon compatibility
  • Comprehensive tokenizer with special token support

🧠 Architecture Details

Model Architecture

  • Transformer-based with grouped query attention
  • RoPE (Rotary Position Embedding) with proper scaling
  • SwiGLU activation in feed-forward networks
  • RMSNorm for layer normalization
  • Weight tying between embedding and output layers

Optimizations

  • Mixed precision (float16/bfloat16) for memory efficiency
  • KV caching for faster autoregressive generation
  • Repetition penalty with sliding window
  • Top-k sampling for better text quality
  • MPS acceleration on Apple Silicon

🔧 Technical Fixes

Issues Resolved

  1. Italian text generation bug - Fixed tokenizer chat mode triggering
  2. 3B model loading - Added support for multi-file safetensors
  3. Performance optimization - Achieved 17+ tokens/second
  4. Memory efficiency - Reduced memory usage with proper dtype handling
  5. Startup time - Implemented model caching for instant subsequent loads

📈 Benchmarks

All benchmarks, tests, and demonstrations are included in the comprehensive notebook:

jupyter notebook comprehensive_llama_demo.ipynb

🎯 Use Cases

  • Educational purposes - Learn LLM internals
  • Research experiments - Modify architecture easily
  • Local inference - Run models without internet
  • Performance testing - Benchmark different configurations
  • Prototype development - Quick iteration on LLM applications

🤝 Contributing

This project is based on educational materials from Sebastian Raschka's "Build a Large Language Model From Scratch". Contributions welcome!

📚 References

📄 License

This project is released under the Apache 2.0 License. See LICENSE for details.

⚡ Quick Commands

# Complete interactive demo with all features
jupyter notebook comprehensive_llama_demo.ipynb

# Quick Python usage
python -c "from fast_llm import FastLLM; model = FastLLM('meta-llama/Llama-3.2-1B-Instruct'); print(''.join(model.generate('Hello world', max_new_tokens=20)))"

Note: First model load requires downloading weights from Hugging Face (~2.5GB for 3B model). Subsequent loads are instant thanks to our caching system!

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages