# MAP-NEO Mini: A DIY LLM from Scratch

This repository demonstrates a complete, end-to-end journey of building, extending, and deploying a custom 253 M-parameter language model on modest hardware (RTX 5070, 8 GB VRAM). It covers data preparation, model training, context window extension, interactive inference, and GPU optimization.

***

## 🚀 Project Overview

- **Model**: MAP-NEO Mini (253 M parameters)  
- **Architecture**: Rotary embeddings, RMSNorm, SwiGLU, Flash Attention  
- **Hardware**: Intel i5 CPU, 16 GB RAM → NVIDIA RTX 4000 (20 GB) → RTX 5070 (8 GB)  
- **Data**: RefinedWeb (100 K high-quality web docs, 41 M tokens)  
- **Context Window**: Extended from 1,024 → 16,384 tokens  
- **Training**: Mixed precision (bf16), gradient checkpointing, gradient accumulation  
- **Fine-Tuning**: Planned conversational instruction tuning with UltraChat  

***

## 📂 Repository Structure

```
AI/
├─ checkpoints/                  # Model checkpoints & configs
│  ├─ checkpoint_step_149999.pt  # Last pre-training checkpoint
│  ├─ extended_context_model.pt  # 8K context model
│  └─ model_config.json          # Config for extended model
├─ data/                         # Raw and processed data
│  ├─ shards/                    # Raw JSONL shards
│  ├─ processed/                 # Filtered JSONL
│  └─ tokens/                    # Packed token sequences
├─ clean_conversational_neo/     # Conversational training scripts
├─ configs/                      # training_config.json, data_config.json
├─ logs/                         # TensorBoard logs
├─ notebooks/                    # Exploratory Jupyter notebooks
├─ advanced_generate.py          # Advanced inference & context tests
├─ conversation_data_prep.py     # Prepares chat data for fine-tuning
├─ data_prep.py                  # RefinedWeb download & preprocessing
├─ debug_downloaded_data.py      # Inspect raw data quality
├─ extend_context.py             # Script to extend model context window
├─ finetune_neo.py               # Base fine-tuning script
├─ generate_text.py              # Simple generation utility
├─ interactive_chat.py           # Interactive chat interface
├─ model_neo.py                  # Model & config definitions
├─ requirements.txt              # Python dependencies
├─ run_training.py               # Orchestrates data prep → training
├─ scale_data.py                 # Utilities for sampling & scaling datasets
├─ setup_project.py              # Initial setup (venv, downloads)
├─ test_conversational_neo.py    # Tests on small conversational model
└─ train_neo.py                  # Main pre-training script
```

***

## 🛠️ Setup & Installation

1. **Clone** this repo.  
2. **Create virtual environment** (Python 3.10+):  
   ```bash
   python -m venv .venv
   source .venv/bin/activate  # or .venv\Scripts\activate
   pip install --upgrade pip
   pip install -r requirements.txt
   ```
3. **Install GPU drivers** and **CUDA** (if using RTX).  
4. **Optional**: `pip install tensorboard pynvml` for logging and GPU monitoring.

***

## 📊 Data Preparation

- **Dataset**: `tiiuae/falcon-refinedweb`  
- **Script**: `data_prep.py`  
  - Downloads 100 K docs, filters for quality (200–10,000 chars, English only)  
  - Tokenizes with GPT-2 BPE, packs into sequences of length 1,024  
- **Output**:  
  - Raw shards: `data/shards/refinedweb_sample_raw.jsonl`  
  - Filtered: `data/processed/refinedweb_filtered.jsonl`  
  - Packed tokens: `data/tokens/packed_1024.txt`  

```bash
python data_prep.py --num_docs 100000 --seq_length 1024 --tokenizer gpt2 --output_dir data
```

***

## 🏋️ Pre-Training

- **Script**: `train_neo.py`  
- **Config**:  
  ```python
  batch_size = 1
  gradient_accumulation_steps = 32
  max_steps = 150000
  warmup_steps = 3750
  mixed_precision = "bf16"
  gradient_checkpointing = True
  ```
- **Accelerator** handles mixed-precision, gradient accumulation, checkpointing.  
- **Resume** from any checkpoint: set `resume_from_checkpoint` in `TrainingConfig`.

```bash
python train_neo.py        # fresh
python train_neo.py --resume checkpoints/checkpoint_step_7500.pt  # resume
```

- **Speed**: ~10 it/s → 4 hours for 150K steps on RTX 5070.

***

## 🔧 Context Extension

- **Script**: `extend_context.py`  
- Extends `config.max_seq_len` → 16,384 and interpolates position embeddings.  
- **Output**: `checkpoints/extended_context_model_16k.pt`

```bash
python extend_context.py --new_max_len 16384
```

***

## 🤖 Inference & Testing

### **Simple Generation**  
`advanced_generate.py` tests fixed prompts and long context usage with VRAM monitoring.  

### **Interactive Chat**  
`interactive_chat.py` provides a full chat interface:  
- `/help`, `/params`, `/memory`, `/context`, `/clear`, `/save`, `/load`, `/multi`, `/system`, `/exit`  
- Real-time GPU usage and context tracking  
- Customizable sampling parameters  

```bash
python interactive_chat.py
```

***

## 📈 Fine-Tuning Plan

- **Dataset Recommendation**: `openbmb/UltraChat` (1.5 M dialogs) + `BAAI/Infinity-Instruct` + `vicgalle/alpaca-gpt4`  
- **Script**: `finetune_neo.py` (will be extended for conversational data)  
- **Goal**: Transform base model → instruction-following chat assistant  

```bash
python finetune_neo.py \
  --base_model checkpoints/extended_context_model_16k.pt \
  --dataset /path/to/UltraChat \
  --epochs 3 --lr 5e-6 --batch_size 1
```

***

## 🔑 Key Lessons & Tips

- **Quality > Quantity**: RefinedWeb quality cut training steps by 25%  
- **Memory Efficiency**: Achieved 3.6 K tokens at ~1.3 GB VRAM  
- **Batch Size Tradeoff**: 1 vs 2 batch size critical for VRAM overflow  
- **Cache Clearing**: `torch.cuda.empty_cache()` essential for long context tests  
- **Resume Training**: Checkpointing during pre-training saved 10+ hours  
- **Conversational Fine-Tuning**: Final step to transform base model into chat assistant

***

## 📂 Next Steps

1. **Review and run conversational fine-tuning** on UltraChat.  
2. **Evaluate** on standardized benchmarks (perplexity, MMLU, HellaSwag).  
3. **Quantize** or **prune** for faster inference on edge devices.  
4. **Deploy** with FastAPI + SSE for streaming responses.  
5. **Document** model card and share results.

***

Thank you for following this detailed project! Your model is now a **powerful, efficient LLM** ready for conversational fine-tuning and deployment. Good luck!