Files
rustytorch/demos/rtx-distllm-demo/src/sample_data.rs
T
2026-03-04 00:08:42 +00:00

690 lines
19 KiB
Rust

//! 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
// ============================================================================
/// 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);
}
}
}