osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
871 lines
24 KiB
Rust
871 lines
24 KiB
Rust
//! Shared types for the DistributedLLM trillion-parameter inference demo.
|
|
//!
|
|
//! This crate provides IPC types for distributed large language model inference
|
|
//! across Thunderbolt 5 Mac clusters. Supports tensor parallelism, pipeline
|
|
//! parallelism, and hybrid strategies for models up to trillion parameters.
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
// ============================================================================
|
|
// Model Configuration
|
|
// ============================================================================
|
|
|
|
/// Configuration for a large language model.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct ModelConfig {
|
|
/// Model name (e.g., "llama-70b", "llama-405b").
|
|
pub name: String,
|
|
/// Total number of parameters (in billions).
|
|
pub num_params: f64,
|
|
/// Number of transformer layers.
|
|
pub num_layers: usize,
|
|
/// Hidden dimension size.
|
|
pub hidden_dim: usize,
|
|
/// Number of attention heads.
|
|
pub num_heads: usize,
|
|
/// Number of key-value heads (for GQA).
|
|
pub num_kv_heads: usize,
|
|
/// Intermediate dimension (FFN).
|
|
pub intermediate_dim: usize,
|
|
/// Vocabulary size.
|
|
pub vocab_size: usize,
|
|
/// Maximum sequence length.
|
|
pub max_seq_len: usize,
|
|
/// Head dimension.
|
|
pub head_dim: usize,
|
|
/// RoPE theta for positional encoding.
|
|
pub rope_theta: f64,
|
|
/// Data type for weights.
|
|
pub dtype: DataType,
|
|
}
|
|
|
|
impl Default for ModelConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
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,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl ModelConfig {
|
|
/// Calculate estimated memory required in GB.
|
|
#[must_use]
|
|
pub fn estimated_memory_gb(&self) -> f64 {
|
|
let bytes_per_param = match self.dtype {
|
|
DataType::Float32 => 4.0,
|
|
DataType::Float16 | DataType::BFloat16 => 2.0,
|
|
DataType::Int8 | DataType::FP8 => 1.0,
|
|
DataType::Int4 => 0.5,
|
|
};
|
|
self.num_params * bytes_per_param
|
|
}
|
|
|
|
/// Calculate memory per layer in MB.
|
|
#[must_use]
|
|
pub fn memory_per_layer_mb(&self) -> f64 {
|
|
self.estimated_memory_gb() * 1024.0 / self.num_layers as f64
|
|
}
|
|
}
|
|
|
|
/// Data types for model weights.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
|
pub enum DataType {
|
|
Float32,
|
|
Float16,
|
|
#[default]
|
|
BFloat16,
|
|
Int8,
|
|
Int4,
|
|
FP8,
|
|
}
|
|
|
|
// ============================================================================
|
|
// Parallelism Configuration
|
|
// ============================================================================
|
|
|
|
/// Configuration for distributed parallelism.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct ParallelismConfig {
|
|
/// Tensor parallelism degree (splits attention heads/FFN).
|
|
pub tensor_parallel: usize,
|
|
/// Pipeline parallelism degree (splits layers).
|
|
pub pipeline_parallel: usize,
|
|
/// Total world size (number of nodes).
|
|
pub world_size: usize,
|
|
/// Number of micro-batches for pipeline.
|
|
pub num_micro_batches: usize,
|
|
/// Pipeline schedule type.
|
|
pub pipeline_schedule: PipelineSchedule,
|
|
/// Communication backend.
|
|
pub comm_backend: CommBackend,
|
|
}
|
|
|
|
impl Default for ParallelismConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
tensor_parallel: 1,
|
|
pipeline_parallel: 1,
|
|
world_size: 1,
|
|
num_micro_batches: 1,
|
|
pipeline_schedule: PipelineSchedule::GPipe,
|
|
comm_backend: CommBackend::Thunderbolt5,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl ParallelismConfig {
|
|
/// Calculate data parallelism degree.
|
|
#[must_use]
|
|
pub fn data_parallel(&self) -> usize {
|
|
self.world_size / (self.tensor_parallel * self.pipeline_parallel)
|
|
}
|
|
|
|
/// Validate parallelism configuration.
|
|
#[must_use]
|
|
pub fn is_valid(&self) -> bool {
|
|
self.world_size > 0
|
|
&& self.tensor_parallel > 0
|
|
&& self.pipeline_parallel > 0
|
|
&& self
|
|
.world_size
|
|
.is_multiple_of(self.tensor_parallel * self.pipeline_parallel)
|
|
}
|
|
}
|
|
|
|
/// Pipeline parallelism schedule.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
|
pub enum PipelineSchedule {
|
|
/// GPipe: all-forward then all-backward.
|
|
#[default]
|
|
GPipe,
|
|
/// 1F1B: interleaved forward-backward.
|
|
OneForwardOneBackward,
|
|
/// Interleaved 1F1B with virtual stages.
|
|
InterleavedOneForwardOneBackward,
|
|
/// Zero bubble schedule.
|
|
ZeroBubble,
|
|
}
|
|
|
|
/// Communication backend for distributed inference.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
|
pub enum CommBackend {
|
|
/// Thunderbolt 5 (80 Gbps bidirectional).
|
|
#[default]
|
|
Thunderbolt5,
|
|
/// Thunderbolt 4 (40 Gbps bidirectional).
|
|
Thunderbolt4,
|
|
/// Ethernet (10/25/100 GbE).
|
|
Ethernet,
|
|
/// InfiniBand.
|
|
InfiniBand,
|
|
/// PCIe direct.
|
|
PCIe,
|
|
/// NVLink (for GPU clusters).
|
|
NVLink,
|
|
}
|
|
|
|
impl CommBackend {
|
|
/// Get theoretical bandwidth in GB/s.
|
|
#[must_use]
|
|
pub fn bandwidth_gbps(&self) -> f64 {
|
|
match self {
|
|
Self::Thunderbolt5 => 80.0,
|
|
Self::Thunderbolt4 => 40.0,
|
|
Self::Ethernet => 100.0,
|
|
Self::InfiniBand => 400.0,
|
|
Self::PCIe => 64.0,
|
|
Self::NVLink => 900.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Inference Request/Response
|
|
// ============================================================================
|
|
|
|
/// Request for text generation inference.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct InferenceRequest {
|
|
/// Input prompt text.
|
|
pub prompt: String,
|
|
/// Maximum tokens to generate.
|
|
pub max_tokens: usize,
|
|
/// Sampling temperature.
|
|
pub temperature: f32,
|
|
/// Top-p (nucleus) sampling parameter.
|
|
pub top_p: f32,
|
|
/// Top-k sampling parameter.
|
|
pub top_k: usize,
|
|
/// Frequency penalty.
|
|
pub frequency_penalty: f32,
|
|
/// Presence penalty.
|
|
pub presence_penalty: f32,
|
|
/// Stop sequences.
|
|
pub stop_sequences: Vec<String>,
|
|
/// Random seed for reproducibility.
|
|
pub seed: Option<u64>,
|
|
/// Whether to stream tokens.
|
|
pub stream: bool,
|
|
}
|
|
|
|
impl Default for InferenceRequest {
|
|
fn default() -> Self {
|
|
Self {
|
|
prompt: String::new(),
|
|
max_tokens: 256,
|
|
temperature: 0.7,
|
|
top_p: 0.9,
|
|
top_k: 50,
|
|
frequency_penalty: 0.0,
|
|
presence_penalty: 0.0,
|
|
stop_sequences: vec![],
|
|
seed: None,
|
|
stream: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Result from text generation inference.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct InferenceResult {
|
|
/// Generated tokens (as strings).
|
|
pub tokens: Vec<String>,
|
|
/// Generated text.
|
|
pub text: String,
|
|
/// Total latency in milliseconds.
|
|
pub latency_ms: f64,
|
|
/// Tokens per second throughput.
|
|
pub tokens_per_second: f64,
|
|
/// Time to first token in milliseconds.
|
|
pub time_to_first_token_ms: f64,
|
|
/// Per-token latencies.
|
|
pub token_latencies_ms: Vec<f64>,
|
|
/// Number of prompt tokens.
|
|
pub prompt_tokens: usize,
|
|
/// Number of generated tokens.
|
|
pub completion_tokens: usize,
|
|
/// Finish reason.
|
|
pub finish_reason: FinishReason,
|
|
}
|
|
|
|
impl Default for InferenceResult {
|
|
fn default() -> Self {
|
|
Self {
|
|
tokens: vec![],
|
|
text: String::new(),
|
|
latency_ms: 0.0,
|
|
tokens_per_second: 0.0,
|
|
time_to_first_token_ms: 0.0,
|
|
token_latencies_ms: vec![],
|
|
prompt_tokens: 0,
|
|
completion_tokens: 0,
|
|
finish_reason: FinishReason::EndOfSequence,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Reason for finishing generation.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
|
pub enum FinishReason {
|
|
/// Reached end of sequence token.
|
|
#[default]
|
|
EndOfSequence,
|
|
/// Reached maximum token limit.
|
|
MaxTokens,
|
|
/// Hit a stop sequence.
|
|
StopSequence,
|
|
/// Generation was cancelled.
|
|
Cancelled,
|
|
}
|
|
|
|
// ============================================================================
|
|
// Layer Assignment
|
|
// ============================================================================
|
|
|
|
/// Assignment of a layer to a node.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct LayerAssignment {
|
|
/// Layer index.
|
|
pub layer_id: usize,
|
|
/// Node ID where layer is placed.
|
|
pub node_id: usize,
|
|
/// Memory required in MB.
|
|
pub memory_mb: f64,
|
|
/// Pipeline stage index.
|
|
pub pipeline_stage: usize,
|
|
/// Tensor parallel rank within stage.
|
|
pub tensor_rank: usize,
|
|
}
|
|
|
|
impl LayerAssignment {
|
|
/// Create a new layer assignment.
|
|
#[must_use]
|
|
pub fn new(layer_id: usize, node_id: usize, memory_mb: f64) -> Self {
|
|
Self {
|
|
layer_id,
|
|
node_id,
|
|
memory_mb,
|
|
pipeline_stage: 0,
|
|
tensor_rank: 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Generation Configuration
|
|
// ============================================================================
|
|
|
|
/// Configuration for text generation.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct GenerationConfig {
|
|
/// Maximum new tokens to generate.
|
|
pub max_new_tokens: usize,
|
|
/// Minimum new tokens to generate.
|
|
pub min_new_tokens: usize,
|
|
/// Whether to use sampling (vs greedy).
|
|
pub do_sample: bool,
|
|
/// Sampling parameters.
|
|
pub sampling: SamplingParams,
|
|
/// Number of beams for beam search.
|
|
pub num_beams: usize,
|
|
/// Early stopping for beam search.
|
|
pub early_stopping: bool,
|
|
/// Repetition penalty.
|
|
pub repetition_penalty: f32,
|
|
/// Length penalty for beam search.
|
|
pub length_penalty: f32,
|
|
/// Number of return sequences.
|
|
pub num_return_sequences: usize,
|
|
}
|
|
|
|
impl Default for GenerationConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_new_tokens: 256,
|
|
min_new_tokens: 1,
|
|
do_sample: true,
|
|
sampling: SamplingParams::default(),
|
|
num_beams: 1,
|
|
early_stopping: false,
|
|
repetition_penalty: 1.0,
|
|
length_penalty: 1.0,
|
|
num_return_sequences: 1,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Sampling parameters for generation.
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
|
|
pub struct SamplingParams {
|
|
/// Temperature for softmax.
|
|
pub temperature: f32,
|
|
/// Top-p (nucleus) sampling.
|
|
pub top_p: f32,
|
|
/// Top-k sampling.
|
|
pub top_k: usize,
|
|
/// Typical-p sampling.
|
|
pub typical_p: f32,
|
|
/// Eta cutoff for entropy-based sampling.
|
|
pub eta_cutoff: f32,
|
|
/// Epsilon cutoff for probability-based sampling.
|
|
pub epsilon_cutoff: f32,
|
|
}
|
|
|
|
impl Default for SamplingParams {
|
|
fn default() -> Self {
|
|
Self {
|
|
temperature: 0.7,
|
|
top_p: 0.9,
|
|
top_k: 50,
|
|
typical_p: 1.0,
|
|
eta_cutoff: 0.0,
|
|
epsilon_cutoff: 0.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Cluster Configuration
|
|
// ============================================================================
|
|
|
|
/// Configuration for a compute cluster.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ClusterConfig {
|
|
/// Cluster name.
|
|
pub name: String,
|
|
/// List of nodes in the cluster.
|
|
pub nodes: Vec<NodeConfig>,
|
|
/// Network topology.
|
|
pub topology: NetworkTopology,
|
|
/// Total cluster memory in GB.
|
|
pub total_memory_gb: f64,
|
|
}
|
|
|
|
impl Default for ClusterConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
name: "default-cluster".to_string(),
|
|
nodes: vec![NodeConfig::default()],
|
|
topology: NetworkTopology::Ring,
|
|
total_memory_gb: 192.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl ClusterConfig {
|
|
/// Get total number of nodes.
|
|
#[must_use]
|
|
pub fn num_nodes(&self) -> usize {
|
|
self.nodes.len()
|
|
}
|
|
|
|
/// Get total memory across all nodes.
|
|
#[must_use]
|
|
pub fn total_memory(&self) -> f64 {
|
|
self.nodes.iter().map(|n| n.memory_gb).sum()
|
|
}
|
|
}
|
|
|
|
/// Configuration for a single compute node.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct NodeConfig {
|
|
/// Node ID.
|
|
pub id: usize,
|
|
/// Node hostname.
|
|
pub hostname: String,
|
|
/// Memory available in GB.
|
|
pub memory_gb: f64,
|
|
/// Number of CPU cores.
|
|
pub cpu_cores: usize,
|
|
/// GPU type if available.
|
|
pub gpu_type: Option<String>,
|
|
/// GPU memory in GB if available.
|
|
pub gpu_memory_gb: Option<f64>,
|
|
/// Network interface.
|
|
pub network: CommBackend,
|
|
}
|
|
|
|
impl Default for NodeConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
id: 0,
|
|
hostname: "localhost".to_string(),
|
|
memory_gb: 192.0,
|
|
cpu_cores: 12,
|
|
gpu_type: None,
|
|
gpu_memory_gb: None,
|
|
network: CommBackend::Thunderbolt5,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Network topology for cluster.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
|
pub enum NetworkTopology {
|
|
/// Ring topology.
|
|
#[default]
|
|
Ring,
|
|
/// Fully connected mesh.
|
|
Mesh,
|
|
/// Star topology with central switch.
|
|
Star,
|
|
/// Tree/hierarchical topology.
|
|
Tree,
|
|
/// Daisy chain.
|
|
DaisyChain,
|
|
}
|
|
|
|
// ============================================================================
|
|
// KV Cache Configuration
|
|
// ============================================================================
|
|
|
|
/// Configuration for KV cache.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct KVCacheConfig {
|
|
/// Maximum sequence length to cache.
|
|
pub max_seq_len: usize,
|
|
/// Maximum batch size.
|
|
pub max_batch_size: usize,
|
|
/// Number of layers.
|
|
pub num_layers: usize,
|
|
/// Number of KV heads.
|
|
pub num_kv_heads: usize,
|
|
/// Head dimension.
|
|
pub head_dim: usize,
|
|
/// Data type for cache.
|
|
pub dtype: DataType,
|
|
/// Whether to use paged attention.
|
|
pub paged_attention: bool,
|
|
/// Block size for paged attention.
|
|
pub block_size: usize,
|
|
}
|
|
|
|
impl Default for KVCacheConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_seq_len: 4096,
|
|
max_batch_size: 1,
|
|
num_layers: 32,
|
|
num_kv_heads: 32,
|
|
head_dim: 128,
|
|
dtype: DataType::BFloat16,
|
|
paged_attention: true,
|
|
block_size: 16,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl KVCacheConfig {
|
|
/// Calculate total cache memory in GB.
|
|
#[must_use]
|
|
pub fn total_memory_gb(&self) -> f64 {
|
|
let bytes_per_elem = match self.dtype {
|
|
DataType::Float32 => 4,
|
|
DataType::Float16 | DataType::BFloat16 => 2,
|
|
DataType::Int8 | DataType::FP8 => 1,
|
|
DataType::Int4 => 1, // Rounded up
|
|
};
|
|
let total_elements = self.max_batch_size
|
|
* self.max_seq_len
|
|
* self.num_layers
|
|
* self.num_kv_heads
|
|
* self.head_dim
|
|
* 2; // K and V
|
|
(total_elements * bytes_per_elem) as f64 / (1024.0 * 1024.0 * 1024.0)
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Performance Metrics
|
|
// ============================================================================
|
|
|
|
/// Performance metrics for inference.
|
|
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
|
pub struct PerformanceMetrics {
|
|
/// Prefill latency in ms.
|
|
pub prefill_latency_ms: f64,
|
|
/// Decode latency per token in ms.
|
|
pub decode_latency_ms: f64,
|
|
/// Tokens per second.
|
|
pub tokens_per_second: f64,
|
|
/// Memory utilization (0-1).
|
|
pub memory_utilization: f64,
|
|
/// Compute utilization (0-1).
|
|
pub compute_utilization: f64,
|
|
/// Network bandwidth utilization (0-1).
|
|
pub network_utilization: f64,
|
|
/// Pipeline bubble ratio.
|
|
pub bubble_ratio: f64,
|
|
/// All-reduce time in ms.
|
|
pub allreduce_time_ms: f64,
|
|
}
|
|
|
|
// ============================================================================
|
|
// Sample Data Functions
|
|
// ============================================================================
|
|
|
|
/// Create a 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,
|
|
intermediate_dim: 28672,
|
|
vocab_size: 32000,
|
|
max_seq_len: 4096,
|
|
head_dim: 128,
|
|
rope_theta: 10000.0,
|
|
dtype: DataType::BFloat16,
|
|
}
|
|
}
|
|
|
|
/// Create a LLaMA 405B model configuration.
|
|
#[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,
|
|
intermediate_dim: 53248,
|
|
vocab_size: 128256,
|
|
max_seq_len: 131072,
|
|
head_dim: 128,
|
|
rope_theta: 500000.0,
|
|
dtype: DataType::BFloat16,
|
|
}
|
|
}
|
|
|
|
/// Create a 4-node Thunderbolt 5 cluster configuration.
|
|
#[must_use]
|
|
pub fn four_node_cluster_config() -> 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,
|
|
}
|
|
}
|
|
|
|
/// Create an 8-node Thunderbolt 5 cluster configuration.
|
|
#[must_use]
|
|
pub fn eight_node_cluster_config() -> 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,
|
|
}
|
|
}
|
|
|
|
/// Create a sample chat inference request.
|
|
#[must_use]
|
|
pub fn chat_request() -> InferenceRequest {
|
|
InferenceRequest {
|
|
prompt: "You are a helpful AI assistant.\n\nUser: Explain quantum computing 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,
|
|
}
|
|
}
|
|
|
|
/// Create a sample completion inference request.
|
|
#[must_use]
|
|
pub fn completion_request() -> InferenceRequest {
|
|
InferenceRequest {
|
|
prompt: "The Rust programming language is known for".to_string(),
|
|
max_tokens: 256,
|
|
temperature: 0.8,
|
|
top_p: 0.95,
|
|
top_k: 40,
|
|
frequency_penalty: 0.1,
|
|
presence_penalty: 0.1,
|
|
stop_sequences: vec![],
|
|
seed: None,
|
|
stream: false,
|
|
}
|
|
}
|
|
|
|
/// Create sample layer assignments for a model across nodes.
|
|
#[must_use]
|
|
pub fn sample_layer_assignments(num_layers: usize, num_nodes: usize) -> Vec<LayerAssignment> {
|
|
let layers_per_node = num_layers.div_ceil(num_nodes);
|
|
(0..num_layers)
|
|
.map(|layer_id| {
|
|
let node_id = layer_id / layers_per_node;
|
|
LayerAssignment {
|
|
layer_id,
|
|
node_id: node_id.min(num_nodes - 1),
|
|
memory_mb: 1024.0,
|
|
pipeline_stage: node_id.min(num_nodes - 1),
|
|
tensor_rank: 0,
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_model_config_default() {
|
|
let config = ModelConfig::default();
|
|
assert_eq!(config.name, "llama-7b");
|
|
assert_eq!(config.num_layers, 32);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_memory_estimation() {
|
|
let config = llama_70b_config();
|
|
let memory = config.estimated_memory_gb();
|
|
// 70B params * 2 bytes (BF16) = 140 GB
|
|
assert!((memory - 140.0).abs() < 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parallelism_config() {
|
|
let config = ParallelismConfig {
|
|
tensor_parallel: 4,
|
|
pipeline_parallel: 2,
|
|
world_size: 8,
|
|
..Default::default()
|
|
};
|
|
assert!(config.is_valid());
|
|
assert_eq!(config.data_parallel(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parallelism_invalid() {
|
|
let config = ParallelismConfig {
|
|
tensor_parallel: 3,
|
|
pipeline_parallel: 2,
|
|
world_size: 8,
|
|
..Default::default()
|
|
};
|
|
assert!(!config.is_valid());
|
|
}
|
|
|
|
#[test]
|
|
fn test_inference_request_default() {
|
|
let request = InferenceRequest::default();
|
|
assert_eq!(request.max_tokens, 256);
|
|
assert!((request.temperature - 0.7).abs() < 0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_assignment() {
|
|
let assignment = LayerAssignment::new(5, 1, 512.0);
|
|
assert_eq!(assignment.layer_id, 5);
|
|
assert_eq!(assignment.node_id, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cluster_config() {
|
|
let cluster = four_node_cluster_config();
|
|
assert_eq!(cluster.num_nodes(), 4);
|
|
assert!((cluster.total_memory() - 768.0).abs() < 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_kv_cache_memory() {
|
|
let config = KVCacheConfig {
|
|
max_seq_len: 4096,
|
|
max_batch_size: 1,
|
|
num_layers: 80,
|
|
num_kv_heads: 8,
|
|
head_dim: 128,
|
|
dtype: DataType::BFloat16,
|
|
..Default::default()
|
|
};
|
|
let memory = config.total_memory_gb();
|
|
// Should be reasonable for LLaMA 70B
|
|
assert!(memory > 0.0 && memory < 100.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_llama_70b_config() {
|
|
let config = llama_70b_config();
|
|
assert_eq!(config.num_layers, 80);
|
|
assert_eq!(config.num_heads, 64);
|
|
assert_eq!(config.num_kv_heads, 8);
|
|
}
|
|
|
|
#[test]
|
|
fn test_llama_405b_config() {
|
|
let config = llama_405b_config();
|
|
assert_eq!(config.num_layers, 126);
|
|
assert!((config.num_params - 405.0).abs() < 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sample_layer_assignments() {
|
|
let assignments = sample_layer_assignments(80, 4);
|
|
assert_eq!(assignments.len(), 80);
|
|
// First 20 layers should be on node 0
|
|
assert_eq!(assignments[0].node_id, 0);
|
|
assert_eq!(assignments[19].node_id, 0);
|
|
// Next 20 on node 1
|
|
assert_eq!(assignments[20].node_id, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_chat_request() {
|
|
let request = chat_request();
|
|
assert!(request.prompt.contains("User:"));
|
|
assert!(request.stream);
|
|
}
|
|
|
|
#[test]
|
|
fn test_completion_request() {
|
|
let request = completion_request();
|
|
assert!(request.prompt.contains("Rust"));
|
|
assert!(!request.stream);
|
|
}
|
|
|
|
#[test]
|
|
fn test_comm_backend_bandwidth() {
|
|
assert_eq!(CommBackend::Thunderbolt5.bandwidth_gbps(), 80.0);
|
|
assert_eq!(CommBackend::NVLink.bandwidth_gbps(), 900.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_serialization() {
|
|
let config = llama_70b_config();
|
|
let json = serde_json::to_string(&config).unwrap();
|
|
let parsed: ModelConfig = serde_json::from_str(&json).unwrap();
|
|
assert_eq!(config, parsed);
|
|
}
|
|
|
|
#[test]
|
|
fn test_generation_config() {
|
|
let config = GenerationConfig::default();
|
|
assert!(config.do_sample);
|
|
assert_eq!(config.num_beams, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sampling_params() {
|
|
let params = SamplingParams::default();
|
|
assert!((params.temperature - 0.7).abs() < 0.01);
|
|
assert_eq!(params.top_k, 50);
|
|
}
|
|
|
|
#[test]
|
|
fn test_performance_metrics() {
|
|
let metrics = PerformanceMetrics {
|
|
tokens_per_second: 50.0,
|
|
memory_utilization: 0.85,
|
|
..Default::default()
|
|
};
|
|
assert!(metrics.tokens_per_second > 0.0);
|
|
}
|
|
}
|