Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,216 @@
+//! GPT decoder-only transformer implementation
+//!
+//! This module provides a complete GPT architecture with proper causal attention,
+//! embeddings, and language modeling heads organized into logical components for
+//! maintainability.
+//!
+//! ## Architecture Overview
+//!
+//! GPT (Generative Pre-trained Transformer) consists of:
+//! - Token and positional embeddings
+//! - Multi-head causal self-attention layers
+//! - Position-wise feed-forward networks
+//! - Layer normalization and residual connections
+//! - Language modeling head for next token prediction
+//!
+//! ## Usage
+//!
+//! ```rust,ignore
+//! use crate::architectures::gpt_old::{GPTConfig, GPTLMHeadModel};
+//! use rtx_tensor::Device;
+//!
+//! let config = GPTConfig::gpt2_small();
+//! let device = Device::cuda(0).unwrap_or(Device::default());
+//! let model = GPTLMHeadModel::new(config, &device)?;
+//! ```
+
+// Re-export from the main gpt module for backward compatibility
+pub use super::gpt::{
+    FeedForward, GPTBlock, GPTConfig, GPTLMHeadModel, GPTModel, LayerNorm, MultiHeadAttention,
+    PositionalEmbedding, TokenEmbedding,
+};
+
+// Legacy compatibility - keep some key types at this level for backward compatibility
+use crate::{Result, TransformerError};
+use rtx_tensor::{Device, Tensor};
+
+// Legacy functions for backward compatibility
+
+/// Create a simple GPT model for testing purposes
+pub fn create_simple_gpt_model(config: GPTConfig, device: &Device) -> Result<GPTModel> {
+    GPTModel::new(config, device)
+}
+
+/// Legacy helper to create language modeling model
+pub fn create_gpt_for_lm(config: GPTConfig, device: &Device) -> Result<GPTLMHeadModel> {
+    GPTLMHeadModel::new(config, device)
+}
+
+// Additional convenience functions for common GPT operations
+
+/// Load GPT configuration from preset
+pub fn load_gpt_config(model_name: &str) -> Result<GPTConfig> {
+    match model_name {
+        "gpt2" | "gpt2-small" => Ok(GPTConfig::gpt2_small()),
+        "gpt2-medium" => Ok(GPTConfig::gpt2_medium()),
+        "gpt2-large" => Ok(GPTConfig::gpt2_large()),
+        "gpt2-xl" => Ok(GPTConfig::gpt2_xl()),
+        _ => {
+            // Try parsing as GPT-3 size
+            if let Some(size) = model_name.strip_prefix("gpt3-") {
+                GPTConfig::gpt3(size)
+            } else {
+                Err(TransformerError::config(format!(
+                    "Unknown GPT model: {model_name}"
+                )))
+            }
+        }
+    }
+}
+
+/// Validate GPT configuration parameters
+pub fn validate_gpt_config(config: &GPTConfig) -> Result<()> {
+    if !config.hidden_size.is_multiple_of(config.num_heads) {
+        return Err(TransformerError::config(format!(
+            "Hidden size ({}) must be divisible by number of heads ({})",
+            config.hidden_size, config.num_heads
+        )));
+    }
+
+    if config.num_layers == 0 {
+        return Err(TransformerError::config(
+            "Number of layers must be greater than 0".to_string(),
+        ));
+    }
+
+    if config.vocab_size == 0 {
+        return Err(TransformerError::config(
+            "Vocabulary size must be greater than 0".to_string(),
+        ));
+    }
+
+    Ok(())
+}
+
+/// Estimate the number of parameters in a GPT model
+#[must_use] 
+pub fn estimate_gpt_parameters(config: &GPTConfig) -> usize {
+    // Token embeddings
+    let token_embedding_params = config.vocab_size * config.hidden_size;
+
+    // Position embeddings
+    let pos_embedding_params = config.max_sequence_length * config.hidden_size;
+
+    // Transformer layers
+    let attention_params_per_layer = 4 * config.hidden_size * config.hidden_size; // Q, K, V, O
+    let ff_params_per_layer = 2 * config.hidden_size * config.intermediate_size; // Up and down
+    let ln_params_per_layer = 4 * config.hidden_size; // 2 layer norms per layer
+    let params_per_layer = attention_params_per_layer + ff_params_per_layer + ln_params_per_layer;
+    let total_layer_params = config.num_layers * params_per_layer;
+
+    // Final layer norm
+    let final_ln_params = config.hidden_size;
+
+    // Language modeling head (tied embeddings don't add extra params)
+    let lm_head_params = if true {
+        // assuming tied embeddings
+        0
+    } else {
+        config.hidden_size * config.vocab_size
+    };
+
+    token_embedding_params
+        + pos_embedding_params
+        + total_layer_params
+        + final_ln_params
+        + lm_head_params
+}
+
+/// Create position IDs for a given sequence length
+pub fn create_position_ids(seq_len: usize, device: &Device) -> Result<Tensor> {
+    let position_ids_data: Vec<f32> = (0..seq_len).map(|i| i as f32).collect();
+    Tensor::from_data(position_ids_data, [seq_len], device)
+        .map_err(|e| TransformerError::tensor_op(e.to_string()))
+}
+
+/// Generate text with basic sampling (simplified implementation)
+pub fn generate_text(
+    model: &GPTLMHeadModel,
+    input_ids: &Tensor,
+    max_new_tokens: usize,
+    temperature: f32,
+) -> Result<Tensor> {
+    let device = input_ids.device();
+    let mut current_ids = input_ids.clone();
+
+    for _ in 0..max_new_tokens {
+        // Get logits for current sequence
+        let logits = model.forward(&current_ids, None)?;
+
+        // Extract logits for the last position
+        let batch_size = logits.shape().dims()[0];
+        let seq_len = logits.shape().dims()[1];
+        let vocab_size = logits.shape().dims()[2];
+
+        let logits_data = logits.to_cpu()?;
+
+        // For simplicity, just take the most likely token (greedy decoding)
+        let mut next_tokens = Vec::new();
+        for batch_idx in 0..batch_size {
+            let start_idx = (batch_idx * seq_len + (seq_len - 1)) * vocab_size;
+            let last_logits = &logits_data[start_idx..start_idx + vocab_size];
+
+            let max_idx = last_logits
+                .iter()
+                .enumerate()
+                .max_by(|(_, a), (_, b)| a.total_cmp(b))
+                .map_or(0, |(idx, _)| idx);
+
+            next_tokens.push(max_idx as f32);
+        }
+
+        // Create next token tensor
+        let next_token_tensor = Tensor::from_data(next_tokens, [batch_size, 1], device)?;
+
+        // Concatenate with current sequence
+        current_ids = Tensor::cat(&[current_ids, next_token_tensor], 1)?;
+    }
+
+    Ok(current_ids)
+}
+
+// Tests for backward compatibility
+#[cfg(all(test, feature = "disabled_tests"))]
+mod tests {
+    use super::*;
+    use rtx_tensor::Device;
+
+    #[test]
+    fn test_gpt_config_creation() {
+        let config = GPTConfig::gpt2_small();
+        assert_eq!(config.hidden_size, 768);
+        assert_eq!(config.num_heads, 12);
+        assert_eq!(config.num_layers, 12);
+    }
+
+    #[test]
+    fn test_gpt_config_validation() {
+        let config = GPTConfig::gpt2_small();
+        assert!(validate_gpt_config(&config).is_ok());
+    }
+
+    #[test]
+    fn test_parameter_estimation() {
+        let config = GPTConfig::gpt2_small();
+        let param_count = estimate_gpt_parameters(&config);
+        // GPT-2 small should have approximately 117M parameters
+        assert!(param_count > 100_000_000 && param_count < 130_000_000);
+    }
+
+    #[test]
+    fn test_load_gpt_config() {
+        assert!(load_gpt_config("gpt2").is_ok());
+        assert!(load_gpt_config("gpt2-medium").is_ok());
+        assert!(load_gpt_config("unknown-model").is_err());
+    }
+}