custom-models

# Custom Models Guide to implementing custom model architectures in LitGPT. ## Overview LitGPT's clean, single-file implementations make it easy to create custom architectures. You can extend the base `GPT` class or create entirely new models. **Use cases**: - Implementing new research architectures - Adapting models for specific domains - Experimenting with attention mechanisms - Adding custom layers or components ## Key Files and Classes ### Core Architecture (`litgpt/model.py`) **Main classes**: - `GPT`: Top-level model class - `Block`: Transformer block (attention + MLP) - `CausalSelfAttention`: Attention mechanism - `MLP`: Feed-forward network - `RMSNorm` / `LayerNorm`: Normalization layers **Configuration** (`litgpt/config.py`): - `Config`: Base configuration dataclass - Model-specific configs: `LlamaConfig`, `MistralConfig`, `PhiConfig`, etc. ## Custom Architecture Workflow ### Step 1: Define Configuration Create a `Config` dataclass with your model's hyperparameters: python from dataclasses import dataclass from litgpt.config import Config @dataclass class MyModelConfig(Config): """Configuration for my custom model.""" # Standard parameters name: str = "my-model-7b" block_size: int = 4096 vocab_size: int = 32000 n_layer: int = 32 n_head: int = 32 n_embd: int = 4096 # Custom parameters custom_param: float = 0.1 use_custom_attention: bool = True # Optional: override defaults rope_base: int = 10000 intermediate_size: int = 11008 ### Step 2: Implement Custom Components #### Option A: Custom Attention python from litgpt.model import CausalSelfAttention import torch import torch.nn as nn class CustomAttention(CausalSelfAttention): """Custom attention mechanism.""" def __init__(self, config): super().__init__(config) # Add custom components self.custom_proj = nn.Linear(config.n_embd, config.n_embd) self.custom_param = config.custom_param def forward(self, x, mask=None, input_pos=None): B, T, C = x.size() # Standard Q, K, V projections q = self.attn(x) k = self.attn(x) v = self.attn(x) # Custom modification q = q + self.custom_proj(x) * self.custom_param # Rest of attention computation q = q.view(B, T, self.n_head, self.head_size) k = k.view(B, T, self.n_query_groups, self.head_size) v = v.view(B, T, self.n_query_groups, self.head_size)