rustytorch/crates/training/rtx-transformers/src/architectures/gpt/config.rs

//! GPT configuration
//!
//! This module defines the configuration structure for GPT models,
//! including various preset configurations for different GPT variants.

use crate::architectures::TransformerConfig;
use crate::{Result, TransformerError};
use serde::{Deserialize, Serialize};

/// GPT-specific configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GPTConfig {
    /// Base transformer configuration
    pub base: TransformerConfig,
    /// Vocabulary size
    pub vocab_size: usize,
    /// Maximum sequence length
    pub max_sequence_length: usize,
    /// Number of transformer layers
    pub num_layers: usize,
    /// Hidden dimension
    pub hidden_size: usize,
    /// Number of attention heads
    pub num_heads: usize,
    /// Feed-forward dimension
    pub intermediate_size: usize,
    /// Dropout probability
    pub dropout: f64,
    /// Whether to use bias in linear layers
    pub use_bias: bool,
    /// Activation function
    pub activation: String,
    /// Layer norm epsilon
    pub layer_norm_eps: f64,
    /// Initializer range for weights
    pub initializer_range: f64,
}

impl Default for GPTConfig {
    fn default() -> Self {
        Self {
            base: TransformerConfig::default(),
            vocab_size: 50257, // GPT-2 vocab size
            max_sequence_length: 1024,
            num_layers: 12,
            hidden_size: 768,
            num_heads: 12,
            intermediate_size: 3072,
            dropout: 0.1,
            use_bias: true,
            activation: "gelu".to_string(),
            layer_norm_eps: 1e-5,
            initializer_range: 0.02,
        }
    }
}

impl GPTConfig {
    /// Create GPT-2 small configuration (117M parameters)
    #[must_use]
    pub fn gpt2_small() -> Self {
        Self::default()
    }

    /// Create GPT-2 medium configuration (345M parameters)
    #[must_use]
    pub fn gpt2_medium() -> Self {
        let mut config = Self::default();
        config.num_layers = 24;
        config.hidden_size = 1024;
        config.num_heads = 16;
        config.intermediate_size = 4096;
        config
    }

    /// Create GPT-2 large configuration (774M parameters)
    #[must_use]
    pub fn gpt2_large() -> Self {
        let mut config = Self::default();
        config.num_layers = 36;
        config.hidden_size = 1280;
        config.num_heads = 20;
        config.intermediate_size = 5120;
        config
    }

    /// Create GPT-2 XL configuration (1.5B parameters)
    #[must_use]
    pub fn gpt2_xl() -> Self {
        let mut config = Self::default();
        config.num_layers = 48;
        config.hidden_size = 1600;
        config.num_heads = 25;
        config.intermediate_size = 6400;
        config
    }

    /// Create GPT-3 style configuration
    pub fn gpt3(size: &str) -> Result<Self> {
        let mut config = Self::default();
        match size {
            "small" => {
                config.num_layers = 12;
                config.hidden_size = 768;
                config.num_heads = 12;
                config.intermediate_size = 3072;
            }
            "medium" => {
                config.num_layers = 24;
                config.hidden_size = 1024;
                config.num_heads = 16;
                config.intermediate_size = 4096;
            }
            "large" => {
                config.num_layers = 24;
                config.hidden_size = 1536;
                config.num_heads = 16;
                config.intermediate_size = 6144;
            }
            "xl" => {
                config.num_layers = 24;
                config.hidden_size = 2048;
                config.num_heads = 24;
                config.intermediate_size = 8192;
            }
            "2.7b" => {
                config.num_layers = 32;
                config.hidden_size = 2560;
                config.num_heads = 32;
                config.intermediate_size = 10240;
            }
            "6.7b" => {
                config.num_layers = 32;
                config.hidden_size = 4096;
                config.num_heads = 32;
                config.intermediate_size = 16384;
            }
            "13b" => {
                config.num_layers = 40;
                config.hidden_size = 5140;
                config.num_heads = 40;
                config.intermediate_size = 20560;
            }
            "175b" => {
                config.num_layers = 96;
                config.hidden_size = 12288;
                config.num_heads = 96;
                config.intermediate_size = 49152;
                config.max_sequence_length = 2048;
            }
            _ => {
                return Err(TransformerError::config(format!(
                    "Unsupported GPT-3 size: {size}"
                )));
            }
        }
        config.vocab_size = 50257; // GPT-3 uses same tokenizer as GPT-2
        Ok(config)
    }
}