rustytorch/demos/rtx-distllm-demo/src/sample_data.rs

//! Sample data and configurations for DistributedLLM demo.
//!
//! This module provides pre-configured model, cluster, and inference
//! settings for demonstration and testing purposes.

use distllm_shared::{
    ClusterConfig, CommBackend, DataType, GenerationConfig, InferenceRequest, KVCacheConfig,
    ModelConfig, NetworkTopology, NodeConfig, ParallelismConfig, PipelineSchedule, SamplingParams,
};

// ============================================================================
// Model Configurations
// ============================================================================

/// A small model configuration used for **real tensor compute** demo runs.
///
/// The larger configs below (`llama_70b_config`, `llama_405b_config`, ...)
/// describe genuinely trillion-parameter-scale models: their weight tensors
/// are far too large to actually allocate (would require hundreds of GB of
/// RAM per shard). They remain useful for cluster-planning /
/// memory-estimation code paths that only read `ModelConfig` fields and
/// never allocate real weights. Anywhere this demo performs *real*
/// `Tensor::randn` weight allocation and real attention compute
/// (`ModelShard::load`, `DistributedLLM::generate`, `run_demo`), this small
/// config (or something similarly sized) should be used instead, so the
/// real compute stays representative without exhausting host memory.
#[must_use]
pub fn tiny_realcompute_config() -> ModelConfig {
    ModelConfig {
        name: "tiny-realcompute-demo".to_string(),
        num_params: 0.05,
        num_layers: 8,
        hidden_dim: 256,
        num_heads: 8,
        num_kv_heads: 8,
        intermediate_dim: 512,
        vocab_size: 32000,
        max_seq_len: 4096,
        head_dim: 32,
        rope_theta: 10000.0,
        dtype: DataType::Float32,
    }
}

/// LLaMA 7B model configuration.
#[must_use]
pub fn llama_7b_config() -> ModelConfig {
    ModelConfig {
        name: "llama-7b".to_string(),
        num_params: 7.0,
        num_layers: 32,
        hidden_dim: 4096,
        num_heads: 32,
        num_kv_heads: 32,
        intermediate_dim: 11008,
        vocab_size: 32000,
        max_seq_len: 4096,
        head_dim: 128,
        rope_theta: 10000.0,
        dtype: DataType::BFloat16,
    }
}

/// LLaMA 13B model configuration.
#[must_use]
pub fn llama_13b_config() -> ModelConfig {
    ModelConfig {
        name: "llama-13b".to_string(),
        num_params: 13.0,
        num_layers: 40,
        hidden_dim: 5120,
        num_heads: 40,
        num_kv_heads: 40,
        intermediate_dim: 13824,
        vocab_size: 32000,
        max_seq_len: 4096,
        head_dim: 128,
        rope_theta: 10000.0,
        dtype: DataType::BFloat16,
    }
}

/// LLaMA 70B model configuration.
#[must_use]
pub fn llama_70b_config() -> ModelConfig {
    ModelConfig {
        name: "llama-70b".to_string(),
        num_params: 70.0,
        num_layers: 80,
        hidden_dim: 8192,
        num_heads: 64,
        num_kv_heads: 8, // GQA
        intermediate_dim: 28672,
        vocab_size: 32000,
        max_seq_len: 4096,
        head_dim: 128,
        rope_theta: 10000.0,
        dtype: DataType::BFloat16,
    }
}

/// LLaMA 405B model configuration (requires large cluster).
#[must_use]
pub fn llama_405b_config() -> ModelConfig {
    ModelConfig {
        name: "llama-405b".to_string(),
        num_params: 405.0,
        num_layers: 126,
        hidden_dim: 16384,
        num_heads: 128,
        num_kv_heads: 8, // GQA
        intermediate_dim: 53248,
        vocab_size: 128256,
        max_seq_len: 131072,
        head_dim: 128,
        rope_theta: 500000.0,
        dtype: DataType::BFloat16,
    }
}

/// Mixtral 8x7B MoE model configuration.
#[must_use]
pub fn mixtral_8x7b_config() -> ModelConfig {
    ModelConfig {
        name: "mixtral-8x7b".to_string(),
        num_params: 46.7, // Active params per forward pass
        num_layers: 32,
        hidden_dim: 4096,
        num_heads: 32,
        num_kv_heads: 8, // GQA
        intermediate_dim: 14336,
        vocab_size: 32000,
        max_seq_len: 32768,
        head_dim: 128,
        rope_theta: 1000000.0,
        dtype: DataType::BFloat16,
    }
}

/// GPT-4 scale model configuration (hypothetical).
#[must_use]
pub fn gpt4_scale_config() -> ModelConfig {
    ModelConfig {
        name: "gpt4-scale".to_string(),
        num_params: 1760.0, // 1.76T params
        num_layers: 120,
        hidden_dim: 25600,
        num_heads: 200,
        num_kv_heads: 25, // GQA
        intermediate_dim: 102400,
        vocab_size: 128000,
        max_seq_len: 128000,
        head_dim: 128,
        rope_theta: 500000.0,
        dtype: DataType::BFloat16,
    }
}

// ============================================================================
// Cluster Configurations
// ============================================================================

/// Single Mac Studio configuration.
#[must_use]
pub fn single_node_cluster() -> ClusterConfig {
    ClusterConfig {
        name: "single-mac-studio".to_string(),
        nodes: vec![NodeConfig {
            id: 0,
            hostname: "localhost".to_string(),
            memory_gb: 192.0,
            cpu_cores: 24,
            gpu_type: None,
            gpu_memory_gb: None,
            network: CommBackend::PCIe,
        }],
        topology: NetworkTopology::Ring,
        total_memory_gb: 192.0,
    }
}

/// 2-node Thunderbolt 5 cluster.
#[must_use]
pub fn two_node_cluster() -> ClusterConfig {
    ClusterConfig {
        name: "tb5-2node-cluster".to_string(),
        nodes: (0..2)
            .map(|i| NodeConfig {
                id: i,
                hostname: format!("mac-studio-{}", i),
                memory_gb: 192.0,
                cpu_cores: 24,
                gpu_type: None,
                gpu_memory_gb: None,
                network: CommBackend::Thunderbolt5,
            })
            .collect(),
        topology: NetworkTopology::Ring,
        total_memory_gb: 384.0,
    }
}

/// 4-node Thunderbolt 5 cluster.
#[must_use]
pub fn four_node_cluster() -> ClusterConfig {
    ClusterConfig {
        name: "tb5-4node-cluster".to_string(),
        nodes: (0..4)
            .map(|i| NodeConfig {
                id: i,
                hostname: format!("mac-studio-{}", i),
                memory_gb: 192.0,
                cpu_cores: 24,
                gpu_type: None,
                gpu_memory_gb: None,
                network: CommBackend::Thunderbolt5,
            })
            .collect(),
        topology: NetworkTopology::Ring,
        total_memory_gb: 768.0,
    }
}

/// 8-node Thunderbolt 5 cluster.
#[must_use]
pub fn eight_node_cluster() -> ClusterConfig {
    ClusterConfig {
        name: "tb5-8node-cluster".to_string(),
        nodes: (0..8)
            .map(|i| NodeConfig {
                id: i,
                hostname: format!("mac-studio-{}", i),
                memory_gb: 192.0,
                cpu_cores: 24,
                gpu_type: None,
                gpu_memory_gb: None,
                network: CommBackend::Thunderbolt5,
            })
            .collect(),
        topology: NetworkTopology::Mesh,
        total_memory_gb: 1536.0,
    }
}

/// 16-node heterogeneous cluster.
#[must_use]
pub fn sixteen_node_cluster() -> ClusterConfig {
    ClusterConfig {
        name: "hybrid-16node-cluster".to_string(),
        nodes: (0..16)
            .map(|i| NodeConfig {
                id: i,
                hostname: format!("node-{}", i),
                memory_gb: if i < 8 { 192.0 } else { 96.0 },
                cpu_cores: if i < 8 { 24 } else { 12 },
                gpu_type: None,
                gpu_memory_gb: None,
                network: if i < 8 {
                    CommBackend::Thunderbolt5
                } else {
                    CommBackend::Ethernet
                },
            })
            .collect(),
        topology: NetworkTopology::Tree,
        total_memory_gb: 2304.0,
    }
}

// ============================================================================
// Parallelism Configurations
// ============================================================================

/// Single-node parallelism (no distribution).
#[must_use]
pub fn single_node_parallel_config() -> ParallelismConfig {
    ParallelismConfig {
        tensor_parallel: 1,
        pipeline_parallel: 1,
        world_size: 1,
        num_micro_batches: 1,
        pipeline_schedule: PipelineSchedule::GPipe,
        comm_backend: CommBackend::PCIe,
    }
}

/// 2-node pipeline parallelism.
#[must_use]
pub fn two_node_parallel_config() -> ParallelismConfig {
    ParallelismConfig {
        tensor_parallel: 1,
        pipeline_parallel: 2,
        world_size: 2,
        num_micro_batches: 4,
        pipeline_schedule: PipelineSchedule::OneForwardOneBackward,
        comm_backend: CommBackend::Thunderbolt5,
    }
}

/// 4-node pipeline parallelism.
#[must_use]
pub fn four_node_parallel_config() -> ParallelismConfig {
    ParallelismConfig {
        tensor_parallel: 1,
        pipeline_parallel: 4,
        world_size: 4,
        num_micro_batches: 8,
        pipeline_schedule: PipelineSchedule::OneForwardOneBackward,
        comm_backend: CommBackend::Thunderbolt5,
    }
}

/// 4-node hybrid parallelism (2 TP x 2 PP).
#[must_use]
pub fn four_node_hybrid_config() -> ParallelismConfig {
    ParallelismConfig {
        tensor_parallel: 2,
        pipeline_parallel: 2,
        world_size: 4,
        num_micro_batches: 4,
        pipeline_schedule: PipelineSchedule::InterleavedOneForwardOneBackward,
        comm_backend: CommBackend::Thunderbolt5,
    }
}

/// 8-node pipeline parallelism.
#[must_use]
pub fn eight_node_parallel_config() -> ParallelismConfig {
    ParallelismConfig {
        tensor_parallel: 1,
        pipeline_parallel: 8,
        world_size: 8,
        num_micro_batches: 16,
        pipeline_schedule: PipelineSchedule::ZeroBubble,
        comm_backend: CommBackend::Thunderbolt5,
    }
}

/// 8-node hybrid parallelism (4 TP x 2 PP).
#[must_use]
pub fn eight_node_hybrid_config() -> ParallelismConfig {
    ParallelismConfig {
        tensor_parallel: 4,
        pipeline_parallel: 2,
        world_size: 8,
        num_micro_batches: 8,
        pipeline_schedule: PipelineSchedule::InterleavedOneForwardOneBackward,
        comm_backend: CommBackend::Thunderbolt5,
    }
}

// ============================================================================
// Inference Requests
// ============================================================================

/// Simple chat inference request.
#[must_use]
pub fn chat_inference_request() -> InferenceRequest {
    InferenceRequest {
        prompt: "You are a helpful AI assistant.\n\nUser: Explain the theory of relativity in simple terms.\n\nAssistant:".to_string(),
        max_tokens: 512,
        temperature: 0.7,
        top_p: 0.9,
        top_k: 50,
        frequency_penalty: 0.0,
        presence_penalty: 0.0,
        stop_sequences: vec!["User:".to_string()],
        seed: Some(42),
        stream: true,
    }
}

/// Code completion request.
#[must_use]
pub fn code_completion_request() -> InferenceRequest {
    InferenceRequest {
        prompt: "```rust\n// Implement a parallel merge sort using Rayon\nfn parallel_merge_sort<T: Ord + Clone + Send>(arr: &mut [T]) {\n".to_string(),
        max_tokens: 1024,
        temperature: 0.2,
        top_p: 0.95,
        top_k: 40,
        frequency_penalty: 0.0,
        presence_penalty: 0.0,
        stop_sequences: vec!["```".to_string()],
        seed: Some(123),
        stream: false,
    }
}

/// Long-form content generation request.
#[must_use]
pub fn long_generation_request() -> InferenceRequest {
    InferenceRequest {
        prompt: "Write a detailed technical blog post about distributed systems:".to_string(),
        max_tokens: 4096,
        temperature: 0.8,
        top_p: 0.9,
        top_k: 100,
        frequency_penalty: 0.3,
        presence_penalty: 0.3,
        stop_sequences: vec![],
        seed: None,
        stream: true,
    }
}

/// Summarization request.
#[must_use]
pub fn summarization_request() -> InferenceRequest {
    InferenceRequest {
        prompt: "Summarize the following article in 3 bullet points:\n\n[Article text would go here]\n\nSummary:".to_string(),
        max_tokens: 256,
        temperature: 0.3,
        top_p: 0.9,
        top_k: 20,
        frequency_penalty: 0.0,
        presence_penalty: 0.0,
        stop_sequences: vec!["\n\n".to_string()],
        seed: Some(456),
        stream: false,
    }
}

// ============================================================================
// Generation Configurations
// ============================================================================

/// Creative generation configuration.
#[must_use]
pub fn creative_generation_config() -> GenerationConfig {
    GenerationConfig {
        max_new_tokens: 1024,
        min_new_tokens: 50,
        do_sample: true,
        sampling: SamplingParams {
            temperature: 0.9,
            top_p: 0.95,
            top_k: 100,
            typical_p: 0.95,
            eta_cutoff: 0.0,
            epsilon_cutoff: 0.0,
        },
        num_beams: 1,
        early_stopping: false,
        repetition_penalty: 1.1,
        length_penalty: 1.0,
        num_return_sequences: 1,
    }
}

/// Precise generation configuration.
#[must_use]
pub fn precise_generation_config() -> GenerationConfig {
    GenerationConfig {
        max_new_tokens: 512,
        min_new_tokens: 1,
        do_sample: false,
        sampling: SamplingParams {
            temperature: 0.0,
            top_p: 1.0,
            top_k: 1,
            typical_p: 1.0,
            eta_cutoff: 0.0,
            epsilon_cutoff: 0.0,
        },
        num_beams: 4,
        early_stopping: true,
        repetition_penalty: 1.0,
        length_penalty: 0.8,
        num_return_sequences: 1,
    }
}

/// Balanced generation configuration.
#[must_use]
pub fn balanced_generation_config() -> GenerationConfig {
    GenerationConfig {
        max_new_tokens: 512,
        min_new_tokens: 10,
        do_sample: true,
        sampling: SamplingParams {
            temperature: 0.7,
            top_p: 0.9,
            top_k: 50,
            typical_p: 1.0,
            eta_cutoff: 0.0,
            epsilon_cutoff: 0.0,
        },
        num_beams: 1,
        early_stopping: false,
        repetition_penalty: 1.0,
        length_penalty: 1.0,
        num_return_sequences: 1,
    }
}

// ============================================================================
// KV Cache Configurations
// ============================================================================

/// KV cache for LLaMA 70B.
#[must_use]
pub fn llama_70b_kv_cache() -> KVCacheConfig {
    KVCacheConfig {
        max_seq_len: 4096,
        max_batch_size: 1,
        num_layers: 80,
        num_kv_heads: 8,
        head_dim: 128,
        dtype: DataType::BFloat16,
        paged_attention: true,
        block_size: 16,
    }
}

/// KV cache for LLaMA 405B.
#[must_use]
pub fn llama_405b_kv_cache() -> KVCacheConfig {
    KVCacheConfig {
        max_seq_len: 8192,
        max_batch_size: 1,
        num_layers: 126,
        num_kv_heads: 8,
        head_dim: 128,
        dtype: DataType::BFloat16,
        paged_attention: true,
        block_size: 32,
    }
}

/// KV cache for long context.
#[must_use]
pub fn long_context_kv_cache() -> KVCacheConfig {
    KVCacheConfig {
        max_seq_len: 131072,
        max_batch_size: 1,
        num_layers: 126,
        num_kv_heads: 8,
        head_dim: 128,
        dtype: DataType::FP8, // Use FP8 for memory efficiency
        paged_attention: true,
        block_size: 64,
    }
}

// ============================================================================
// Benchmark Configurations
// ============================================================================

/// Configuration for throughput benchmarking.
#[derive(Debug, Clone)]
pub struct BenchmarkConfig {
    /// Number of warmup iterations.
    pub warmup_iters: usize,
    /// Number of benchmark iterations.
    pub bench_iters: usize,
    /// Prompt lengths to test.
    pub prompt_lengths: Vec<usize>,
    /// Generation lengths to test.
    pub gen_lengths: Vec<usize>,
    /// Batch sizes to test.
    pub batch_sizes: Vec<usize>,
}

impl Default for BenchmarkConfig {
    fn default() -> Self {
        Self {
            warmup_iters: 3,
            bench_iters: 10,
            prompt_lengths: vec![128, 512, 1024, 2048],
            gen_lengths: vec![64, 128, 256, 512],
            batch_sizes: vec![1],
        }
    }
}

/// Throughput benchmark configuration.
#[must_use]
pub fn throughput_benchmark() -> BenchmarkConfig {
    BenchmarkConfig {
        warmup_iters: 5,
        bench_iters: 20,
        prompt_lengths: vec![128, 512, 2048],
        gen_lengths: vec![128, 512],
        batch_sizes: vec![1, 2, 4],
    }
}

/// Latency benchmark configuration.
#[must_use]
pub fn latency_benchmark() -> BenchmarkConfig {
    BenchmarkConfig {
        warmup_iters: 10,
        bench_iters: 50,
        prompt_lengths: vec![64, 256, 1024],
        gen_lengths: vec![1, 32, 128],
        batch_sizes: vec![1],
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_llama_7b_config() {
        let config = llama_7b_config();
        assert_eq!(config.num_layers, 32);
        assert_eq!(config.num_heads, 32);
    }

    #[test]
    fn test_llama_70b_config() {
        let config = llama_70b_config();
        assert_eq!(config.num_layers, 80);
        assert_eq!(config.num_kv_heads, 8); // GQA
    }

    #[test]
    fn test_llama_405b_config() {
        let config = llama_405b_config();
        assert_eq!(config.num_layers, 126);
        assert!(config.num_params > 400.0);
    }

    #[test]
    fn test_four_node_cluster() {
        let cluster = four_node_cluster();
        assert_eq!(cluster.nodes.len(), 4);
        assert!((cluster.total_memory_gb - 768.0).abs() < 1.0);
    }

    #[test]
    fn test_eight_node_cluster() {
        let cluster = eight_node_cluster();
        assert_eq!(cluster.nodes.len(), 8);
        assert_eq!(cluster.topology, NetworkTopology::Mesh);
    }

    #[test]
    fn test_parallel_configs() {
        let single = single_node_parallel_config();
        assert_eq!(single.world_size, 1);

        let four = four_node_parallel_config();
        assert_eq!(four.pipeline_parallel, 4);

        let hybrid = four_node_hybrid_config();
        assert_eq!(
            hybrid.tensor_parallel * hybrid.pipeline_parallel,
            hybrid.world_size
        );
    }

    #[test]
    fn test_inference_requests() {
        let chat = chat_inference_request();
        assert!(chat.stream);
        assert!(!chat.stop_sequences.is_empty());

        let code = code_completion_request();
        assert!(code.temperature < 0.5);
    }

    #[test]
    fn test_generation_configs() {
        let creative = creative_generation_config();
        assert!(creative.sampling.temperature > 0.8);
        assert!(creative.do_sample);

        let precise = precise_generation_config();
        assert!(!precise.do_sample);
        assert!(precise.num_beams > 1);
    }

    #[test]
    fn test_kv_cache_configs() {
        let cache_70b = llama_70b_kv_cache();
        assert_eq!(cache_70b.num_layers, 80);
        assert!(cache_70b.paged_attention);

        let cache_405b = llama_405b_kv_cache();
        assert_eq!(cache_405b.num_layers, 126);
    }

    #[test]
    fn test_benchmark_configs() {
        let throughput = throughput_benchmark();
        assert!(throughput.bench_iters > throughput.warmup_iters);

        let latency = latency_benchmark();
        assert!(!latency.gen_lengths.is_empty());
    }

    #[test]
    fn test_model_memory_estimates() {
        let llama_70b = llama_70b_config();
        let mem = llama_70b.estimated_memory_gb();
        // 70B * 2 bytes = 140 GB
        assert!((mem - 140.0).abs() < 5.0);

        let llama_405b = llama_405b_config();
        let mem_405b = llama_405b.estimated_memory_gb();
        // 405B * 2 bytes = 810 GB
        assert!((mem_405b - 810.0).abs() < 10.0);
    }

    #[test]
    fn test_cluster_bandwidth() {
        let cluster = four_node_cluster();
        for node in &cluster.nodes {
            assert_eq!(node.network.bandwidth_gbps(), 80.0);
        }
    }
}