New models chosen and rag docs updated

components/model_selector.py

@@ -10,22 +10,22 @@ from dash import html, dcc
 
 # Available models organized by family
 AVAILABLE_MODELS = [
-    # LLaMA-like models (Qwen)
-    {"label": "Qwen2.5-0.5B", "value": "Qwen/Qwen2.5-0.5B"},
-    # {"label": "Qwen2.5-1.5B", "value": "Qwen/Qwen2.5-1.5B"},
-    
-    # GPT-2 family
-    {"label": "GPT-2 (124M)", "value": "gpt2"}
-    # {"label": "GPT-2 Medium (355M)", "value": "gpt2-medium"},
-    # {"label": "GPT-2 Large (774M)", "value": "gpt2-large"},
-    
-    # # OPT family
-    # {"label": "OPT-125M", "value": "facebook/opt-125m"},
-    # {"label": "OPT-350M", "value": "facebook/opt-350m"},
-    
-    # # GPT-NeoX family (Pythia)
-    # {"label": "Pythia-70M", "value": "EleutherAI/pythia-70m"},
-    # {"label": "Pythia-160M", "value": "EleutherAI/pythia-160m"},
+    # GPT-2 family (OpenAI) — absolute positional encoding, LayerNorm, GELU
+    {"label": "GPT-2 (124M)", "value": "gpt2"},
+    {"label": "GPT-2 Medium (355M)", "value": "gpt2-medium"},
+
+    # GPT-Neo (EleutherAI) — absolute PE, LayerNorm, GELU
+    {"label": "GPT-Neo 125M", "value": "EleutherAI/gpt-neo-125M"},
+
+    # Pythia (EleutherAI) — rotary PE, LayerNorm, GELU, parallel attn+MLP
+    {"label": "Pythia-160M", "value": "EleutherAI/pythia-160m"},
+    {"label": "Pythia-410M", "value": "EleutherAI/pythia-410m"},
+
+    # OPT (Meta) — absolute PE, LayerNorm, ReLU activation
+    {"label": "OPT-125M", "value": "facebook/opt-125m"},
+
+    # Qwen2.5 (Alibaba) — rotary PE, RMSNorm, SiLU activation
+    {"label": "Qwen2.5-0.5B (494M)", "value": "Qwen/Qwen2.5-0.5B"},
 ]
 
 def create_model_selector():

rag_docs/README.md

@@ -37,8 +37,10 @@ This folder contains documents used by the AI chatbot for Retrieval-Augmented Ge
 
 ### Category 3: Model-Specific Documentation
 - `gpt2_overview.md` - GPT-2 architecture, why it's a good starter, variants
-- `llama_overview.md` - LLaMA/Qwen/Mistral architecture, RoPE, GQA differences
-- `opt_overview.md` - OPT architecture, comparison with GPT-2
+- `gpt_neo_overview.md` - GPT-Neo architecture, local attention, comparison with GPT-2
+- `pythia_overview.md` - Pythia architecture, RoPE, parallel attn+MLP, interpretability focus
+- `opt_overview.md` - OPT architecture, ReLU activation, comparison with GPT-2
+- `qwen_overview.md` - Qwen2.5 (LLaMA-like) architecture, RMSNorm, SiLU, GQA
 
 ### Category 4: Guided Experiments (Step-by-Step)
 - `experiment_first_analysis.md` - Your first analysis with GPT-2

rag_docs/gpt2_overview.md

@@ -39,14 +39,14 @@ When analyzing GPT-2, you'll typically see:
 
 ## GPT-2 Variants
 
-The dashboard supports all GPT-2 sizes, though only the small variant is in the default dropdown:
+The dashboard includes GPT-2 Small and Medium in the dropdown. Larger variants can be loaded by typing the model ID:
 
 - **GPT-2 Small** (124M params, 12 layers) -- in dropdown as "GPT-2 (124M)"
-- **GPT-2 Medium** (355M params, 24 layers) -- enter `gpt2-medium` in the dropdown
+- **GPT-2 Medium** (355M params, 24 layers) -- in dropdown as "GPT-2 Medium (355M)"
 - **GPT-2 Large** (774M params, 36 layers) -- enter `gpt2-large`
 - **GPT-2 XL** (1.5B params, 48 layers) -- enter `gpt2-xl`
 
-Larger variants have more layers and heads but use more memory and are slower.
+Comparing GPT-2 Small and Medium is a great way to see how attention heads specialize as models scale up within the same architecture.
 
 ## HuggingFace Model IDs
 

rag_docs/gpt_neo_overview.md

@@ -0,0 +1,49 @@
+# GPT-Neo Overview
+
+## What Is GPT-Neo?
+
+GPT-Neo is a family of open-source language models created by EleutherAI in 2021. It was one of the first serious open-source alternatives to OpenAI's GPT-3, designed as an open replication of the GPT architecture. GPT-Neo is the predecessor to EleutherAI's later Pythia suite.
+
+## Architecture Details
+
+| Property | GPT-Neo 125M |
+|----------|-------------|
+| Parameters | 85M (actual) |
+| Layers | 12 |
+| Attention Heads | 12 per layer |
+| Hidden Dimension | 768 |
+| Vocabulary Size | 50,257 |
+| Positional Encoding | Learned absolute |
+| Normalization | LayerNorm |
+| Activation Function | GELU |
+
+## Key Architectural Feature: Local Attention
+
+GPT-Neo's most distinctive feature is its **alternating local and global attention** pattern:
+- **Even layers**: Use local attention with a window of 256 tokens (each token only attends to nearby tokens)
+- **Odd layers**: Use standard global attention (each token can attend to all previous tokens)
+
+This alternating pattern is a significant architectural difference from GPT-2 (which uses global attention in every layer) and creates interesting attention visualization patterns in the dashboard.
+
+## Why GPT-Neo Is Useful for Learning
+
+- **Same size as GPT-2**: At 125M parameters with 12 layers and 12 heads, it's directly comparable to GPT-2
+- **Same positional encoding**: Uses learned absolute positions like GPT-2, so attention pattern differences are due to architecture, not PE
+- **Different attention pattern**: The local attention in even layers creates visually distinct attention maps — great for understanding how attention scope affects behavior
+- **Bridge to Pythia**: GPT-Neo → GPT-NeoX → Pythia is an evolutionary chain. Comparing GPT-Neo (absolute PE) with Pythia (rotary PE) shows the effect of positional encoding
+
+## Comparing GPT-Neo and GPT-2
+
+| Feature | GPT-2 (124M) | GPT-Neo 125M |
+|---------|-------------|--------------|
+| Layers × Heads | 12 × 12 | 12 × 12 |
+| Hidden Dim | 768 | 768 |
+| Attention | All global | Alternating local/global |
+| PE Type | Learned absolute | Learned absolute |
+| Training Data | WebText | The Pile |
+| Creator | OpenAI | EleutherAI |
+
+## HuggingFace Model IDs
+
+- `EleutherAI/gpt-neo-125M` (in dropdown)
+- `EleutherAI/gpt-neo-1.3B`, `EleutherAI/gpt-neo-2.7B` (larger, enter manually)

rag_docs/llama_overview.md

@@ -1,46 +0,0 @@
-# LLaMA Overview
-
-## What Is LLaMA?
-
-LLaMA (Large Language Model Meta AI) is a family of open-weight language models developed by Meta. First released in 2023, LLaMA models introduced several architectural improvements over GPT-2 and became the foundation for many other models (Mistral, Qwen, etc.). In the dashboard, models labeled "LLaMA-like" share this architecture.
-
-## Architectural Differences from GPT-2
-
-LLaMA models use several key innovations:
-
-### RoPE (Rotary Position Embeddings)
-Instead of GPT-2's learned absolute position embeddings, LLaMA uses **Rotary Position Embeddings (RoPE)**. RoPE encodes position information by rotating the query and key vectors in attention. This means:
-- The model can generalize better to different sequence lengths
-- Position information is baked into the attention computation itself
-- Attention patterns may look different from GPT-2 because of how positions are encoded
-
-### RMSNorm Instead of LayerNorm
-LLaMA uses **RMSNorm** (Root Mean Square Normalization) instead of the standard LayerNorm used in GPT-2. RMSNorm is simpler and slightly faster -- it only normalizes the magnitude of the vectors without centering them first.
-
-### SiLU Activation
-Where GPT-2 uses GELU activation in the MLP, LLaMA uses **SiLU** (Sigmoid Linear Unit, also called "Swish"). This is a smooth activation function that tends to produce slightly different MLP behavior.
-
-### Grouped-Query Attention (GQA)
-Larger LLaMA variants use **Grouped-Query Attention**, where multiple query heads share the same key and value heads. This reduces memory usage and speeds up inference without significantly hurting quality. This means the number of key/value heads may be smaller than the number of query heads.
-
-## Models Using LLaMA Architecture
-
-The dashboard's "llama_like" family includes:
-- **Meta LLaMA**: LLaMA 2 (7B, 13B, 70B), LLaMA 3 (1B, 3B, 8B, 70B)
-- **Qwen**: Qwen2, Qwen2.5 (0.5B to 72B) -- available in the dashboard dropdown as "Qwen2.5-0.5B"
-- **Mistral**: Mistral-7B, Mixtral-8x7B
-
-## What to Expect in the Dashboard
-
-When using a LLaMA-like model (such as Qwen2.5-0.5B):
-
-- **More layers and heads**: Even the small Qwen2.5-0.5B has 24 layers and 14 heads, compared to GPT-2's 12 layers and 12 heads
-- **Different attention patterns**: RoPE-based attention may show different positional patterns compared to GPT-2
-- **Different tokenizer**: LLaMA-family models use a different BPE vocabulary, so the same text may tokenize differently
-- **Comparing with GPT-2**: Running the same prompt on both GPT-2 and a LLaMA-like model is a great way to see how architecture affects predictions
-
-## HuggingFace Model IDs
-
-- `Qwen/Qwen2.5-0.5B` (in default dropdown)
-- `meta-llama/Llama-3.2-1B`, `meta-llama/Llama-3.1-8B`
-- `mistralai/Mistral-7B-v0.3`

rag_docs/model_selector_guide.md

@@ -2,39 +2,43 @@
 
 ## How to Choose a Model
 
-The dashboard supports several transformer model families. You can select a model from the dropdown menu in the generator section at the top of the page.
+The dashboard supports seven transformer models from five architecture families. Select a model from the dropdown menu in the generator section at the top of the page.
 
 ### Available Models
 
-Currently, the dashboard offers:
+| Model | Family | Params | Layers × Heads | Key Feature |
+|-------|--------|--------|----------------|-------------|
+| **GPT-2 (124M)** | GPT-2 | 85M | 12 × 12 | The MI classic — start here |
+| **GPT-2 Medium (355M)** | GPT-2 | 302M | 24 × 16 | Scale comparison within GPT-2 |
+| **GPT-Neo 125M** | GPT-Neo | 85M | 12 × 12 | Local attention in alternating layers |
+| **Pythia-160M** | Pythia | 85M | 12 × 12 | Rotary PE, parallel attn+MLP |
+| **Pythia-410M** | Pythia | 302M | 24 × 16 | Larger Pythia for scale comparison |
+| **OPT-125M** | OPT | 85M | 12 × 12 | ReLU activation (unique contrast) |
+| **Qwen2.5-0.5B (494M)** | Qwen2 | 391M | 24 × 14 | Modern: RMSNorm, SiLU, rotary PE |
 
-- **GPT-2 (124M)**: OpenAI's GPT-2 small model. 12 layers, 12 attention heads, 768-dimensional embeddings. This is the best model to start with -- it's small, fast, and well-studied.
-- **Qwen2.5-0.5B**: A LLaMA-like model from Alibaba's Qwen family. 24 layers, 14 attention heads, 896-dimensional embeddings. Slightly larger and uses different architectural features (RoPE, SiLU activation).
+### Architecture Comparisons
 
-You can also enter a custom **HuggingFace model ID** in the dropdown (type it in). The dashboard supports GPT-2, LLaMA, OPT, GPT-NeoX, BLOOM, Falcon, and MPT model families.
+These models were chosen to highlight specific architectural differences:
+
+- **Positional encoding**: GPT-2, GPT-Neo, and OPT use absolute positions. Pythia and Qwen use rotary (RoPE). Comparing the same prompt across both types shows how PE affects attention patterns.
+- **Activation function**: Most models use GELU, but OPT uses ReLU and Qwen uses SiLU. This affects MLP behavior.
+- **Normalization**: GPT-2, Neo, Pythia, and OPT use LayerNorm. Qwen uses RMSNorm.
+- **Attention scope**: GPT-Neo alternates between local (256-token window) and global attention, unlike all other models.
 
 ### What Happens When You Load a Model
 
 1. The model is downloaded from HuggingFace (this may take a moment the first time)
 2. The dashboard **auto-detects** the model's architecture family
 3. Internal hooks are automatically configured to capture attention patterns, MLP activations, and other data
-4. The layer and head dropdowns in the sidebar and ablation panel are populated based on the model's structure
-
-### Auto-Detection
-
-The dashboard has a registry that maps model names to their architecture family. When it recognizes a model, it automatically configures:
-- Which internal modules to hook for attention capture
-- Which normalization parameters to track
-- The correct patterns for extracting layer outputs
-
-If you enter an unknown model, the sidebar's configuration dropdowns may need manual adjustment.
+4. The layer and head dropdowns are populated based on the model's structure
 
 ### Tips for Choosing
 
 - **Start with GPT-2**: It's small, fast, and the most widely studied. Most educational resources reference GPT-2.
-- **Try Qwen2.5-0.5B for comparison**: It uses a different architecture (LLaMA-style). Comparing results between GPT-2 and Qwen can highlight how architectural differences affect attention patterns.
-- **Larger models are slower**: Models with more parameters take longer to load and analyze. Stick to small models for interactive exploration.
-- **Memory matters**: Larger models require more RAM. If the dashboard becomes unresponsive, try a smaller model.
+- **Compare same-size models**: GPT-2, GPT-Neo 125M, Pythia-160M, and OPT-125M all have 12 layers and 12 heads — differences in their attention patterns come from architecture, not scale.
+- **Compare scale**: GPT-2 vs GPT-2 Medium (or Pythia-160M vs Pythia-410M) shows how more layers and heads change behavior.
+- **Try Qwen2.5 for a modern perspective**: It uses an entirely different design philosophy from GPT-2.
+- **Memory matters**: All dropdown models are small enough for interactive exploration. Larger models can be entered manually but may be slow.
 
 ### Generation Settings
 

rag_docs/opt_overview.md

@@ -18,7 +18,8 @@ OPT's architecture is close to GPT-2 but has some differences:
 
 ### Key Differences from GPT-2
 
-- **Learned positional embeddings**: Like GPT-2, OPT uses learned absolute position embeddings (unlike LLaMA's RoPE)
+- **ReLU activation**: OPT uses **ReLU** instead of GPT-2's GELU. This is the only model in the dashboard with ReLU, making it useful for comparing how activation functions affect MLP behavior.
+- **Learned positional embeddings**: Like GPT-2, OPT uses learned absolute position embeddings (unlike Pythia's or Qwen's RoPE)
 - **LayerNorm placement**: OPT uses pre-norm LayerNorm (applied before each sublayer), which is slightly different from GPT-2's original arrangement
 - **Larger variants available**: OPT scales up to 175 billion parameters, though only smaller variants are practical for interactive use
 
@@ -36,13 +37,9 @@ When using OPT models:
 - **OPT-125M is very similar to GPT-2**: Same number of layers (12), heads (12), and hidden dimension (768). You'll see similar attention patterns and predictions.
 - **Different module paths**: The dashboard auto-detects OPT's internal structure (e.g., `model.decoder.layers.N.self_attn`), so hooking works automatically.
 - **Tokenization**: OPT's tokenizer is very similar to GPT-2's, so the same text usually produces similar (but not identical) token sequences.
-- **Good for comparison**: Running the same prompt on GPT-2 and OPT-125M can show how similar architectures with different training data produce different predictions.
+- **Good for comparison**: Running the same prompt on GPT-2 and OPT-125M can show how similar architectures with different training data and activation functions produce different predictions.
 
 ## HuggingFace Model IDs
 
-- `facebook/opt-125m`
-- `facebook/opt-350m`
-- `facebook/opt-1.3b`
-- `facebook/opt-2.7b`
-
-Note: OPT models are not in the default dropdown but can be loaded by typing the model ID directly.
+- `facebook/opt-125m` (in dropdown)
+- `facebook/opt-350m`, `facebook/opt-1.3b` (larger, enter manually)

rag_docs/pythia_overview.md

@@ -0,0 +1,53 @@
+# Pythia Overview
+
+## What Is Pythia?
+
+Pythia is a suite of language models created by EleutherAI specifically to support interpretability research. Released in 2023, the Pythia models were trained on the deduplicated Pile dataset with full public training data and checkpoints, making them uniquely transparent.
+
+## Architecture Details
+
+| Property | Pythia-160M | Pythia-410M |
+|----------|-------------|-------------|
+| Parameters | 85M (actual) | 302M (actual) |
+| Layers | 12 | 24 |
+| Attention Heads | 12 per layer | 16 per layer |
+| Hidden Dimension | 768 | 1024 |
+| Vocabulary Size | 50,304 | 50,304 |
+| Positional Encoding | Rotary (RoPE) |  Rotary (RoPE) |
+| Normalization | LayerNorm | LayerNorm |
+| Activation Function | GELU | GELU |
+| Parallel Attn+MLP | Yes | Yes |
+
+## Key Architectural Features
+
+### Rotary Position Embeddings (RoPE)
+Unlike GPT-2's learned absolute positions, Pythia uses **Rotary Position Embeddings**. RoPE encodes position by rotating the query and key vectors, which means:
+- Position information is baked into the attention computation
+- The model can potentially generalize to longer sequences
+- Attention patterns may look different from GPT-2
+
+### Parallel Attention + MLP
+Pythia computes attention and MLP in **parallel** within each layer (rather than sequentially). This is a design choice that slightly changes how information flows through the network.
+
+### GPT-NeoX Architecture
+Pythia uses the GPT-NeoX architecture (`GPTNeoXForCausalLM`), which is EleutherAI's evolution of GPT-Neo. The module paths use `gpt_neox.layers.{N}.attention` instead of GPT-2's `transformer.h.{N}.attn`.
+
+## Why Pythia Matters for Interpretability
+
+- **Open training data**: The Pile dataset is fully public, so you can trace what the model learned
+- **Training checkpoints**: Hundreds of checkpoints are available, letting researchers study how features develop during training
+- **Designed for research**: EleutherAI specifically built Pythia to be a good subject for mechanistic interpretability
+- **Comparison with GPT-Neo**: Same creator but different architecture (rotary vs absolute PE), enabling clean ablation studies of positional encoding
+
+## Comparing Pythia and GPT-2
+
+Running the same prompt on both is highly educational:
+- **Same size, different PE**: Pythia-160M and GPT-2 both have 12 layers and 12 heads, but use different positional encodings
+- **Parallel vs sequential**: Pythia processes attention and MLP in parallel; GPT-2 does them sequentially
+- **Different training data**: GPT-2 was trained on WebText; Pythia on The Pile
+
+## HuggingFace Model IDs
+
+- `EleutherAI/pythia-160m` (in dropdown)
+- `EleutherAI/pythia-410m` (in dropdown)
+- `EleutherAI/pythia-1b`, `EleutherAI/pythia-1.4b` (larger, enter manually)

rag_docs/qwen_overview.md

@@ -0,0 +1,55 @@
+# Qwen2.5 Overview
+
+## What Is Qwen2.5?
+
+Qwen2.5 is a family of language models developed by Alibaba's Qwen team, released in 2024. It uses a modern LLaMA-like architecture with several improvements over older models like GPT-2. The dashboard includes **Qwen2.5-0.5B** as a representative of this modern architecture.
+
+## Architecture Details
+
+| Property | Qwen2.5-0.5B |
+|----------|-------------|
+| Parameters | 494M |
+| Layers | 24 |
+| Attention Heads | 14 per layer |
+| Hidden Dimension | 896 |
+| Vocabulary Size | 151,936 |
+| Positional Encoding | Rotary (RoPE) |
+| Normalization | RMSNorm |
+| Activation Function | SiLU |
+
+## Key Architectural Differences from GPT-2
+
+### RoPE (Rotary Position Embeddings)
+Instead of GPT-2's learned absolute position embeddings, Qwen uses **Rotary Position Embeddings (RoPE)**. RoPE encodes position by rotating query and key vectors in attention. This means:
+- Better generalization to different sequence lengths
+- Position information baked into the attention computation
+- Attention patterns may look different from GPT-2
+
+### RMSNorm Instead of LayerNorm
+Qwen uses **RMSNorm** (Root Mean Square Normalization) instead of standard LayerNorm. RMSNorm only normalizes vector magnitude without centering, making it simpler and slightly faster.
+
+### SiLU Activation
+Where GPT-2 uses GELU in the MLP, Qwen uses **SiLU** (Sigmoid Linear Unit, also called "Swish"), a smooth activation function common in modern architectures.
+
+### Grouped-Query Attention (GQA)
+Qwen uses **Grouped-Query Attention**, where multiple query heads share key and value heads. This reduces memory usage while maintaining quality. In the attention visualization, you may notice this affects how attention patterns are structured.
+
+## Why Include Qwen2.5?
+
+- **Modern architecture**: Represents the latest generation of LLM design (post-LLaMA innovations)
+- **Contrast with GPT-2**: Every major architectural choice is different — PE type, normalization, activation, attention variant
+- **Good size**: 494M parameters with 24 layers provides enough depth for interesting head specialization while staying fast
+- **Open access**: No gated access or license agreement required (unlike LLaMA or Gemma)
+
+## What to Expect in the Dashboard
+
+When using Qwen2.5-0.5B:
+- **More layers**: 24 layers vs GPT-2's 12 means more attention patterns to explore
+- **Different tokenizer**: Qwen has a much larger vocabulary (152K vs 50K), so the same text tokenizes differently
+- **Different attention patterns**: RoPE-based attention creates different positional patterns from GPT-2
+- **Comparing with GPT-2**: Running the same prompt on both is the best way to see how architecture affects predictions and attention
+
+## HuggingFace Model IDs
+
+- `Qwen/Qwen2.5-0.5B` (in dropdown)
+- `Qwen/Qwen2.5-1.5B`, `Qwen/Qwen2.5-3B` (larger, enter manually)

rag_docs/recommended_starting_points.md

@@ -55,10 +55,14 @@ If you're new to the dashboard, follow this path:
 5. **Experiment: Comparing Heads** -- Systematically compare head categories
 6. **Experiment: Beam Search** -- Explore alternative generation paths
 
-## After the Basics
+## After the Basics: Cross-Model Comparisons
 
-Once you've completed the guided experiments:
-- **Compare models**: Run the same prompt on GPT-2 and Qwen2.5-0.5B to see architectural differences
+Once you've completed the guided experiments, try comparing models to see how architecture affects behavior:
+
+- **GPT-2 vs GPT-Neo 125M**: Same size and PE type, but GPT-Neo alternates local/global attention — see how attention scope matters
+- **GPT-2 vs Pythia-160M**: Same size but different positional encoding (absolute vs rotary) — see how RoPE changes attention patterns
+- **GPT-2 vs OPT-125M**: Same size but OPT uses ReLU instead of GELU — compare MLP behavior
+- **GPT-2 vs GPT-2 Medium**: Same architecture at different scales — see how head specialization changes with more layers
+- **Pythia-160M vs Qwen2.5-0.5B**: Both use rotary PE but different normalization (LayerNorm vs RMSNorm) and activation (GELU vs SiLU)
 - **Try longer prompts**: See how attention patterns change with more context
 - **Combine techniques**: Use attribution to find important tokens, then ablate heads to find the components that process those tokens
-- **Explore edge cases**: Try prompts in other languages, code snippets, or mathematical expressions

scripts/analyze_heads.py

@@ -48,6 +48,7 @@ HF_TO_TL_NAME = {
     "openai-community/gpt2-large": "gpt2-large",
     "gpt2-xl": "gpt2-xl",
     "openai-community/gpt2-xl": "gpt2-xl",
+    "EleutherAI/gpt-neo-125M": "gpt-neo-125M",
     "EleutherAI/pythia-70m": "pythia-70m",
     "EleutherAI/pythia-160m": "pythia-160m",
     "EleutherAI/pythia-410m": "pythia-410m",
@@ -56,6 +57,7 @@ HF_TO_TL_NAME = {
     "facebook/opt-125m": "opt-125m",
     "facebook/opt-350m": "opt-350m",
     "facebook/opt-1.3b": "opt-1.3b",
+    "Qwen/Qwen2.5-0.5B": "qwen2.5-0.5b",
 }
 
 # Default models to analyze
@@ -64,10 +66,11 @@ DEFAULT_MODELS = ["gpt2"]
 ALL_PRIORITY_MODELS = [
     "gpt2",
     "gpt2-medium",
-    "EleutherAI/pythia-70m",
+    "EleutherAI/gpt-neo-125M",
     "EleutherAI/pythia-160m",
     "EleutherAI/pythia-410m",
     "facebook/opt-125m",
+    "Qwen/Qwen2.5-0.5B",
 ]
 
 # ============================================================================

utils/model_config.py

@@ -34,6 +34,18 @@ MODEL_FAMILIES: Dict[str, Dict[str, Any]] = {
         "norm_type": "layernorm",
     },
     
+    # GPT-Neo (EleutherAI) — similar to GPT-2 but with local attention
+    "gpt_neo": {
+        "description": "GPT-Neo architecture (EleutherAI)",
+        "templates": {
+            "attention_pattern": "transformer.h.{N}.attn.attention",
+            "mlp_pattern": "transformer.h.{N}.mlp",
+            "block_pattern": "transformer.h.{N}",
+        },
+        "norm_parameter": "transformer.ln_f.weight",
+        "norm_type": "layernorm",
+    },
+    
     # OPT
     "opt": {
         "description": "OPT architecture",
@@ -143,7 +155,12 @@ MODEL_TO_FAMILY: Dict[str, str] = {
     "facebook/opt-13b": "opt",
     "facebook/opt-30b": "opt",
     
-    # GPT-NeoX models
+    # GPT-Neo models (EleutherAI)
+    "EleutherAI/gpt-neo-125M": "gpt_neo",
+    "EleutherAI/gpt-neo-1.3B": "gpt_neo",
+    "EleutherAI/gpt-neo-2.7B": "gpt_neo",
+    
+    # GPT-NeoX / Pythia models (EleutherAI)
     "EleutherAI/gpt-neox-20b": "gpt_neox",
     "EleutherAI/pythia-70m": "gpt_neox",
     "EleutherAI/pythia-160m": "gpt_neox",