osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
884 lines
31 KiB
Rust
884 lines
31 KiB
Rust
//! DistributedLLM: Trillion-Parameter Inference across Thunderbolt 5 Mac Cluster
|
|
//!
|
|
//! This demo showcases distributed large language model inference using
|
|
//! tensor parallelism and pipeline parallelism across a cluster of Mac
|
|
//! systems connected via Thunderbolt 5.
|
|
//!
|
|
//! ## What is real vs. simulated
|
|
//!
|
|
//! - **Real**: per-layer projection weight allocation (`ModelShard::load`),
|
|
//! Q/K/V tensor construction, matmul + softmax attention compute
|
|
//! (`DistributedAttention::forward`, via `rtx_tensor::Tensor::
|
|
//! scaled_dot_product_attention`), and all wall-clock timings measured
|
|
//! with `std::time::Instant`. Weight/KV memory sizes are computed from
|
|
//! actual tensor element counts and dtype byte widths, not guessed
|
|
//! constants.
|
|
//! - **Simulated**: the network transport layer (`DistributedCommunicator`
|
|
//! models cluster bandwidth/latency but performs no actual inter-process
|
|
//! communication), the tokenizer (word-based, not a real subword
|
|
//! vocabulary), and end-of-sequence detection (a fixed probability, not a
|
|
//! learned distribution). Flash-attention memory savings remain a
|
|
//! documented closed-form estimate rather than a real tiled kernel.
|
|
//! - Model sharding across multiple nodes
|
|
//! - Distributed attention with sharded KV cache
|
|
//! - Pipeline parallelism with micro-batching
|
|
//! - High-bandwidth inter-node communication
|
|
|
|
pub mod attention;
|
|
pub mod parallel;
|
|
pub mod sample_data;
|
|
|
|
use distllm_shared::{
|
|
ClusterConfig, FinishReason, GenerationConfig, InferenceRequest, InferenceResult,
|
|
KVCacheConfig, LayerAssignment, ModelConfig, ParallelismConfig, PerformanceMetrics,
|
|
};
|
|
|
|
use attention::{DistributedAttention, KVCache};
|
|
use parallel::{LayerPartitioner, PipelineParallel, TensorParallel};
|
|
use rtx_tensor::{Device, Tensor};
|
|
|
|
// ============================================================================
|
|
// Model Shard
|
|
// ============================================================================
|
|
|
|
/// Real per-layer projection weights for one transformer layer.
|
|
///
|
|
/// Shapes are sized (and, under tensor parallelism, sharded) from the
|
|
/// model's `ModelConfig` the same way a real transformer layer would be:
|
|
/// Q/K/V/O attention projections plus a gated MLP (SwiGLU-style
|
|
/// gate/up/down projections).
|
|
#[derive(Debug)]
|
|
pub struct LayerWeights {
|
|
/// Query projection: `[hidden_dim, sharded_num_heads * head_dim]`.
|
|
pub q_proj: Tensor,
|
|
/// Key projection: `[hidden_dim, sharded_num_kv_heads * head_dim]`.
|
|
pub k_proj: Tensor,
|
|
/// Value projection: `[hidden_dim, sharded_num_kv_heads * head_dim]`.
|
|
pub v_proj: Tensor,
|
|
/// Output projection: `[sharded_num_heads * head_dim, hidden_dim]`.
|
|
pub o_proj: Tensor,
|
|
/// MLP gate projection: `[hidden_dim, sharded_intermediate_dim]`.
|
|
pub gate_proj: Tensor,
|
|
/// MLP up projection: `[hidden_dim, sharded_intermediate_dim]`.
|
|
pub up_proj: Tensor,
|
|
/// MLP down projection: `[sharded_intermediate_dim, hidden_dim]`.
|
|
pub down_proj: Tensor,
|
|
}
|
|
|
|
impl LayerWeights {
|
|
/// Allocate real (randomly initialized) weight tensors for one layer,
|
|
/// sized according to `model_config` and sharded `tensor_world`-ways
|
|
/// for tensor parallelism (column/row split approximation).
|
|
fn allocate(
|
|
model_config: &ModelConfig,
|
|
tensor_world: usize,
|
|
device: &Device,
|
|
) -> Result<Self, String> {
|
|
let hidden_dim = model_config.hidden_dim;
|
|
let sharded_heads = (model_config.num_heads / tensor_world.max(1)).max(1);
|
|
let sharded_kv_heads = (model_config.num_kv_heads / tensor_world.max(1)).max(1);
|
|
let sharded_intermediate = (model_config.intermediate_dim / tensor_world.max(1)).max(1);
|
|
let head_dim = model_config.head_dim;
|
|
|
|
let q_dim = sharded_heads * head_dim;
|
|
let kv_dim = sharded_kv_heads * head_dim;
|
|
|
|
Ok(Self {
|
|
q_proj: Tensor::randn(&[hidden_dim, q_dim], device).map_err(|e| e.to_string())?,
|
|
k_proj: Tensor::randn(&[hidden_dim, kv_dim], device).map_err(|e| e.to_string())?,
|
|
v_proj: Tensor::randn(&[hidden_dim, kv_dim], device).map_err(|e| e.to_string())?,
|
|
o_proj: Tensor::randn(&[q_dim, hidden_dim], device).map_err(|e| e.to_string())?,
|
|
gate_proj: Tensor::randn(&[hidden_dim, sharded_intermediate], device)
|
|
.map_err(|e| e.to_string())?,
|
|
up_proj: Tensor::randn(&[hidden_dim, sharded_intermediate], device)
|
|
.map_err(|e| e.to_string())?,
|
|
down_proj: Tensor::randn(&[sharded_intermediate, hidden_dim], device)
|
|
.map_err(|e| e.to_string())?,
|
|
})
|
|
}
|
|
|
|
/// Total bytes occupied by this layer's weights, computed from actual
|
|
/// tensor element counts and the tensor's dtype byte width (real
|
|
/// allocation size, not an estimate).
|
|
fn size_bytes(&self) -> usize {
|
|
let dtype_bytes = self.q_proj.dtype().size_bytes();
|
|
[
|
|
&self.q_proj,
|
|
&self.k_proj,
|
|
&self.v_proj,
|
|
&self.o_proj,
|
|
&self.gate_proj,
|
|
&self.up_proj,
|
|
&self.down_proj,
|
|
]
|
|
.iter()
|
|
.map(|t| t.numel() * dtype_bytes)
|
|
.sum()
|
|
}
|
|
}
|
|
|
|
/// A shard of the model assigned to this node.
|
|
#[derive(Debug)]
|
|
pub struct ModelShard {
|
|
/// Shard ID.
|
|
pub id: usize,
|
|
/// Layer range (start, end exclusive).
|
|
pub layer_range: (usize, usize),
|
|
/// Tensor parallel rank.
|
|
pub tensor_rank: usize,
|
|
/// Tensor parallel world size.
|
|
pub tensor_world: usize,
|
|
/// Real per-layer projection weight tensors allocated on load.
|
|
pub layer_weights: Vec<LayerWeights>,
|
|
/// Total weight size in MB, computed from actual tensor byte sizes.
|
|
pub weights_mb: f64,
|
|
/// Whether this shard is loaded.
|
|
pub loaded: bool,
|
|
}
|
|
|
|
impl ModelShard {
|
|
/// Create a new model shard.
|
|
#[must_use]
|
|
pub fn new(id: usize, layer_start: usize, layer_end: usize) -> Self {
|
|
Self {
|
|
id,
|
|
layer_range: (layer_start, layer_end),
|
|
tensor_rank: 0,
|
|
tensor_world: 1,
|
|
layer_weights: Vec::new(),
|
|
weights_mb: 0.0,
|
|
loaded: false,
|
|
}
|
|
}
|
|
|
|
/// Get the number of layers in this shard.
|
|
#[must_use]
|
|
pub fn num_layers(&self) -> usize {
|
|
self.layer_range.1 - self.layer_range.0
|
|
}
|
|
|
|
/// Allocate real weight tensors for every layer in this shard's range,
|
|
/// sized (and tensor-parallel sharded) from `model_config`. `weights_mb`
|
|
/// is then computed by summing the actual allocated tensor byte sizes.
|
|
///
|
|
/// # Errors
|
|
/// Returns an error string if tensor allocation fails (e.g. invalid
|
|
/// shape derived from a degenerate model config).
|
|
pub fn load(&mut self, model_config: &ModelConfig) -> Result<(), String> {
|
|
let device = Device::cpu();
|
|
let tensor_world = self.tensor_world.max(1);
|
|
|
|
self.layer_weights = (0..self.num_layers())
|
|
.map(|_| LayerWeights::allocate(model_config, tensor_world, &device))
|
|
.collect::<Result<Vec<_>, _>>()?;
|
|
|
|
let total_bytes: usize = self
|
|
.layer_weights
|
|
.iter()
|
|
.map(LayerWeights::size_bytes)
|
|
.sum();
|
|
self.weights_mb = total_bytes as f64 / (1024.0 * 1024.0);
|
|
self.loaded = true;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Distributed Communicator
|
|
// ============================================================================
|
|
|
|
/// Communicator for distributed operations.
|
|
///
|
|
/// This is a **simulated network layer**: it performs no real inter-process
|
|
/// or inter-node communication. It models realistic transfer timing (ring
|
|
/// all-reduce cost, point-to-point send/recv cost) given the cluster's
|
|
/// configured bandwidth/latency (`ClusterConfig`) and the *real* byte size
|
|
/// of the data being "transferred" (now computed from actual tensor sizes
|
|
/// rather than invented numbers).
|
|
#[derive(Debug)]
|
|
pub struct DistributedCommunicator {
|
|
/// This node's rank.
|
|
pub rank: usize,
|
|
/// Total world size.
|
|
pub world_size: usize,
|
|
/// Cluster configuration.
|
|
pub cluster: ClusterConfig,
|
|
/// Communication latency in microseconds.
|
|
pub latency_us: f64,
|
|
/// Bandwidth in GB/s.
|
|
pub bandwidth_gbps: f64,
|
|
/// Simulated message count.
|
|
message_count: usize,
|
|
}
|
|
|
|
impl DistributedCommunicator {
|
|
/// Create a new communicator.
|
|
#[must_use]
|
|
pub fn new(rank: usize, cluster: ClusterConfig) -> Self {
|
|
let world_size = cluster.num_nodes();
|
|
let bandwidth_gbps = if !cluster.nodes.is_empty() {
|
|
cluster.nodes[0].network.bandwidth_gbps()
|
|
} else {
|
|
80.0 // Default to TB5
|
|
};
|
|
|
|
Self {
|
|
rank,
|
|
world_size,
|
|
cluster,
|
|
latency_us: 10.0, // 10us latency
|
|
bandwidth_gbps,
|
|
message_count: 0,
|
|
}
|
|
}
|
|
|
|
/// Simulate an all-reduce operation.
|
|
pub fn all_reduce(&mut self, data_size_mb: f64) -> f64 {
|
|
self.message_count += 1;
|
|
// Ring all-reduce: 2 * (n-1) / n * data_size
|
|
let transfer_size =
|
|
2.0 * (self.world_size - 1) as f64 / self.world_size as f64 * data_size_mb;
|
|
let transfer_time_ms = transfer_size * 1000.0 / self.bandwidth_gbps;
|
|
let latency_ms = self.latency_us / 1000.0 * 2.0 * (self.world_size - 1) as f64;
|
|
transfer_time_ms + latency_ms
|
|
}
|
|
|
|
/// Simulate a point-to-point send.
|
|
pub fn send(&mut self, _dest: usize, data_size_mb: f64) -> f64 {
|
|
self.message_count += 1;
|
|
let transfer_time_ms = data_size_mb * 1000.0 / self.bandwidth_gbps;
|
|
transfer_time_ms + self.latency_us / 1000.0
|
|
}
|
|
|
|
/// Simulate a point-to-point receive.
|
|
pub fn recv(&mut self, _source: usize, data_size_mb: f64) -> f64 {
|
|
self.message_count += 1;
|
|
let transfer_time_ms = data_size_mb * 1000.0 / self.bandwidth_gbps;
|
|
transfer_time_ms + self.latency_us / 1000.0
|
|
}
|
|
|
|
/// Get total message count.
|
|
#[must_use]
|
|
pub fn message_count(&self) -> usize {
|
|
self.message_count
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Pipeline Scheduler
|
|
// ============================================================================
|
|
|
|
/// Scheduler for pipeline parallel execution.
|
|
#[derive(Debug)]
|
|
pub struct PipelineScheduler {
|
|
/// Number of pipeline stages.
|
|
pub num_stages: usize,
|
|
/// Number of micro-batches.
|
|
pub num_micro_batches: usize,
|
|
/// Current micro-batch index.
|
|
current_micro_batch: usize,
|
|
/// Schedule type.
|
|
pub schedule: distllm_shared::PipelineSchedule,
|
|
}
|
|
|
|
impl PipelineScheduler {
|
|
/// Create a new pipeline scheduler.
|
|
#[must_use]
|
|
pub fn new(num_stages: usize, num_micro_batches: usize) -> Self {
|
|
Self {
|
|
num_stages,
|
|
num_micro_batches,
|
|
current_micro_batch: 0,
|
|
schedule: distllm_shared::PipelineSchedule::GPipe,
|
|
}
|
|
}
|
|
|
|
/// Get the next stage to execute for the current micro-batch.
|
|
#[must_use]
|
|
pub fn next_stage(&self) -> Option<usize> {
|
|
if self.current_micro_batch < self.num_micro_batches {
|
|
Some(self.current_micro_batch % self.num_stages)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Advance to the next micro-batch.
|
|
pub fn advance(&mut self) {
|
|
self.current_micro_batch += 1;
|
|
}
|
|
|
|
/// Reset the scheduler.
|
|
pub fn reset(&mut self) {
|
|
self.current_micro_batch = 0;
|
|
}
|
|
|
|
/// Calculate pipeline bubble ratio.
|
|
#[must_use]
|
|
pub fn bubble_ratio(&self) -> f64 {
|
|
let total_slots = self.num_stages * self.num_micro_batches;
|
|
let bubble_slots = self.num_stages - 1;
|
|
bubble_slots as f64 / total_slots as f64
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Distributed LLM System
|
|
// ============================================================================
|
|
|
|
/// Main distributed LLM inference system.
|
|
#[derive(Debug)]
|
|
pub struct DistributedLLM {
|
|
/// Model configuration.
|
|
pub model_config: ModelConfig,
|
|
/// Parallelism configuration.
|
|
pub parallel_config: ParallelismConfig,
|
|
/// Model shards on this node.
|
|
pub shards: Vec<ModelShard>,
|
|
/// Distributed communicator.
|
|
pub communicator: DistributedCommunicator,
|
|
/// Pipeline scheduler.
|
|
pub scheduler: PipelineScheduler,
|
|
/// Layer assignments.
|
|
pub layer_assignments: Vec<LayerAssignment>,
|
|
/// Distributed attention module.
|
|
pub attention: DistributedAttention,
|
|
/// KV cache.
|
|
pub kv_cache: KVCache,
|
|
/// Tensor parallel module.
|
|
pub tensor_parallel: TensorParallel,
|
|
/// Pipeline parallel module.
|
|
pub pipeline_parallel: PipelineParallel,
|
|
/// Whether model is loaded.
|
|
pub loaded: bool,
|
|
/// RNG state for sampling.
|
|
rng_state: u64,
|
|
/// Measured prefill latency (ms) from the most recent `generate()` call.
|
|
last_prefill_latency_ms: f64,
|
|
/// Measured average per-token decode latency (ms) from the most recent
|
|
/// `generate()` call.
|
|
last_decode_latency_ms: f64,
|
|
/// Measured tokens/sec throughput from the most recent `generate()` call.
|
|
last_tokens_per_second: f64,
|
|
/// Accumulated real attention compute time (ms) across all layers/shards
|
|
/// in the most recent `generate()` call.
|
|
last_attention_compute_ms: f64,
|
|
/// Accumulated simulated all-reduce time (ms) across all layers/shards
|
|
/// in the most recent `generate()` call.
|
|
last_allreduce_time_ms: f64,
|
|
/// Total wall-clock time (ms) of the most recent `generate()` call.
|
|
last_total_time_ms: f64,
|
|
}
|
|
|
|
impl DistributedLLM {
|
|
/// Create a new distributed LLM system.
|
|
#[must_use]
|
|
pub fn new(
|
|
model_config: ModelConfig,
|
|
parallel_config: ParallelismConfig,
|
|
cluster: ClusterConfig,
|
|
rank: usize,
|
|
) -> Self {
|
|
let communicator = DistributedCommunicator::new(rank, cluster);
|
|
let scheduler = PipelineScheduler::new(
|
|
parallel_config.pipeline_parallel,
|
|
parallel_config.num_micro_batches,
|
|
);
|
|
|
|
let kv_cache_config = KVCacheConfig {
|
|
max_seq_len: model_config.max_seq_len,
|
|
num_layers: model_config.num_layers,
|
|
num_kv_heads: model_config.num_kv_heads / parallel_config.tensor_parallel,
|
|
head_dim: model_config.head_dim,
|
|
dtype: model_config.dtype,
|
|
..Default::default()
|
|
};
|
|
|
|
let attention = DistributedAttention::new(
|
|
model_config.num_heads / parallel_config.tensor_parallel,
|
|
model_config.num_kv_heads / parallel_config.tensor_parallel,
|
|
model_config.head_dim,
|
|
parallel_config.tensor_parallel,
|
|
);
|
|
|
|
let kv_cache = KVCache::new(kv_cache_config);
|
|
|
|
let tensor_parallel = TensorParallel::new(
|
|
parallel_config.tensor_parallel,
|
|
rank % parallel_config.tensor_parallel,
|
|
);
|
|
|
|
let pipeline_parallel = PipelineParallel::new(
|
|
parallel_config.pipeline_parallel,
|
|
rank / parallel_config.tensor_parallel,
|
|
);
|
|
|
|
Self {
|
|
model_config,
|
|
parallel_config,
|
|
shards: Vec::new(),
|
|
communicator,
|
|
scheduler,
|
|
layer_assignments: Vec::new(),
|
|
attention,
|
|
kv_cache,
|
|
tensor_parallel,
|
|
pipeline_parallel,
|
|
loaded: false,
|
|
rng_state: 42,
|
|
last_prefill_latency_ms: 0.0,
|
|
last_decode_latency_ms: 0.0,
|
|
last_tokens_per_second: 0.0,
|
|
last_attention_compute_ms: 0.0,
|
|
last_allreduce_time_ms: 0.0,
|
|
last_total_time_ms: 0.0,
|
|
}
|
|
}
|
|
|
|
/// Load the model, distributing layers across nodes.
|
|
pub fn load_model(&mut self) {
|
|
// Use layer partitioner to assign layers
|
|
let partitioner = LayerPartitioner::new(
|
|
self.model_config.num_layers,
|
|
self.parallel_config.world_size,
|
|
self.parallel_config.tensor_parallel,
|
|
self.parallel_config.pipeline_parallel,
|
|
);
|
|
|
|
self.layer_assignments = partitioner.partition();
|
|
|
|
// Create shards for this node's layers
|
|
let my_layers: Vec<_> = self
|
|
.layer_assignments
|
|
.iter()
|
|
.filter(|a| a.node_id == self.communicator.rank)
|
|
.collect();
|
|
|
|
if !my_layers.is_empty() {
|
|
let start = my_layers.first().map_or(0, |a| a.layer_id);
|
|
let end = my_layers.last().map_or(0, |a| a.layer_id + 1);
|
|
|
|
let mut shard = ModelShard::new(self.communicator.rank, start, end);
|
|
shard.tensor_rank = self.tensor_parallel.rank;
|
|
shard.tensor_world = self.tensor_parallel.world_size;
|
|
shard
|
|
.load(&self.model_config)
|
|
.expect("real tensor weight allocation failed (CPU device should always succeed)");
|
|
self.shards.push(shard);
|
|
}
|
|
|
|
self.loaded = true;
|
|
}
|
|
|
|
/// Generate text from a prompt.
|
|
pub fn generate(
|
|
&mut self,
|
|
request: &InferenceRequest,
|
|
config: &GenerationConfig,
|
|
) -> InferenceResult {
|
|
let start = std::time::Instant::now();
|
|
|
|
if !self.loaded {
|
|
self.load_model();
|
|
}
|
|
|
|
// Tokenize prompt (simulated)
|
|
let prompt_tokens = self.tokenize(&request.prompt);
|
|
let num_prompt_tokens = prompt_tokens.len();
|
|
|
|
// Reset per-call accumulators (real measurements, not constants).
|
|
self.last_attention_compute_ms = 0.0;
|
|
self.last_allreduce_time_ms = 0.0;
|
|
|
|
// Prefill phase: runs a real matmul/softmax attention forward pass
|
|
// per layer via `DistributedAttention::forward`.
|
|
let prefill_start = std::time::Instant::now();
|
|
self.prefill(&prompt_tokens);
|
|
let prefill_time = prefill_start.elapsed().as_secs_f64() * 1000.0;
|
|
self.last_prefill_latency_ms = prefill_time;
|
|
|
|
// Decode phase
|
|
let mut generated_tokens = Vec::new();
|
|
let mut token_latencies = Vec::new();
|
|
let max_tokens = config.max_new_tokens.min(request.max_tokens);
|
|
|
|
let ttft = prefill_time; // Time to first token
|
|
|
|
for _ in 0..max_tokens {
|
|
let token_start = std::time::Instant::now();
|
|
|
|
// Generate next token
|
|
let (token, finished) = self.decode_step(&config.sampling);
|
|
|
|
let token_time = token_start.elapsed().as_secs_f64() * 1000.0;
|
|
token_latencies.push(token_time);
|
|
|
|
generated_tokens.push(token);
|
|
|
|
if finished {
|
|
break;
|
|
}
|
|
}
|
|
|
|
let total_time = start.elapsed().as_secs_f64() * 1000.0;
|
|
let num_generated = generated_tokens.len();
|
|
let tokens_per_second = num_generated as f64 / (total_time / 1000.0);
|
|
|
|
self.last_tokens_per_second = tokens_per_second;
|
|
self.last_total_time_ms = total_time;
|
|
self.last_decode_latency_ms = if token_latencies.is_empty() {
|
|
0.0
|
|
} else {
|
|
token_latencies.iter().sum::<f64>() / token_latencies.len() as f64
|
|
};
|
|
|
|
// Detokenize
|
|
let text = self.detokenize(&generated_tokens);
|
|
|
|
let finish_reason = if num_generated >= max_tokens {
|
|
FinishReason::MaxTokens
|
|
} else {
|
|
FinishReason::EndOfSequence
|
|
};
|
|
|
|
InferenceResult {
|
|
tokens: generated_tokens,
|
|
text,
|
|
latency_ms: total_time,
|
|
tokens_per_second,
|
|
time_to_first_token_ms: ttft,
|
|
token_latencies_ms: token_latencies,
|
|
prompt_tokens: num_prompt_tokens,
|
|
completion_tokens: num_generated,
|
|
finish_reason,
|
|
}
|
|
}
|
|
|
|
/// Tokenize input text.
|
|
///
|
|
/// Simplified simulation: a real subword tokenizer/vocabulary is out of
|
|
/// scope for this demo, so tokens are assigned by whitespace-split word
|
|
/// position rather than via a learned vocabulary.
|
|
fn tokenize(&self, text: &str) -> Vec<u32> {
|
|
text.split_whitespace()
|
|
.enumerate()
|
|
.map(|(i, _)| i as u32 + 1)
|
|
.collect()
|
|
}
|
|
|
|
/// Detokenize tokens to text (simulated; see [`Self::tokenize`]).
|
|
fn detokenize(&self, tokens: &[String]) -> String {
|
|
tokens.join(" ")
|
|
}
|
|
|
|
/// Run prefill phase (process prompt).
|
|
///
|
|
/// Runs a real attention forward pass (Q/K/V allocation, matmul,
|
|
/// softmax) for every layer assigned to this shard, and accumulates the
|
|
/// measured compute time. The tensor-parallel all-reduce remains a
|
|
/// simulated network operation (see [`DistributedCommunicator`]), now
|
|
/// sized from the real hidden-state byte count.
|
|
fn prefill(&mut self, tokens: &[u32]) {
|
|
let seq_len = tokens.len();
|
|
|
|
for shard in &self.shards {
|
|
for layer_id in shard.layer_range.0..shard.layer_range.1 {
|
|
let (_output, elapsed) = self
|
|
.attention
|
|
.forward(seq_len, layer_id)
|
|
.expect("real attention forward pass failed");
|
|
self.last_attention_compute_ms += elapsed.as_secs_f64() * 1000.0;
|
|
|
|
self.kv_cache.update(layer_id, seq_len);
|
|
}
|
|
}
|
|
|
|
// All-reduce for tensor parallel (simulated network transfer, real byte size).
|
|
if self.tensor_parallel.world_size > 1 {
|
|
let hidden_size_mb =
|
|
(self.model_config.hidden_dim * seq_len * 2) as f64 / (1024.0 * 1024.0);
|
|
self.last_allreduce_time_ms += self.communicator.all_reduce(hidden_size_mb);
|
|
}
|
|
}
|
|
|
|
/// Run one decode step (generate one token).
|
|
///
|
|
/// Runs a real single-token attention forward pass per layer, same as
|
|
/// [`Self::prefill`] but with `seq_len = 1`.
|
|
fn decode_step(&mut self, sampling: &distllm_shared::SamplingParams) -> (String, bool) {
|
|
for shard in &self.shards {
|
|
for layer_id in shard.layer_range.0..shard.layer_range.1 {
|
|
let (_output, elapsed) = self
|
|
.attention
|
|
.forward(1, layer_id)
|
|
.expect("real attention forward pass failed");
|
|
self.last_attention_compute_ms += elapsed.as_secs_f64() * 1000.0;
|
|
|
|
self.kv_cache.update(layer_id, 1);
|
|
}
|
|
}
|
|
|
|
// All-reduce for tensor parallel (simulated network transfer, real byte size).
|
|
if self.tensor_parallel.world_size > 1 {
|
|
let hidden_size_mb = (self.model_config.hidden_dim * 2) as f64 / (1024.0 * 1024.0);
|
|
self.last_allreduce_time_ms += self.communicator.all_reduce(hidden_size_mb);
|
|
}
|
|
|
|
// Sample token (simplified simulation: word-list sampling, not a real
|
|
// learned vocabulary distribution).
|
|
let token = self.sample_token(sampling);
|
|
|
|
// Check for end of sequence (simulated: fixed probability, not a
|
|
// learned EOS distribution).
|
|
let is_eos = self.random() < 0.02; // 2% chance of EOS
|
|
|
|
(token, is_eos)
|
|
}
|
|
|
|
/// Sample a token based on sampling parameters.
|
|
fn sample_token(&mut self, _sampling: &distllm_shared::SamplingParams) -> String {
|
|
// Simulated token sampling
|
|
let words = [
|
|
"the", "a", "an", "is", "was", "are", "were", "be", "been", "being", "have", "has",
|
|
"had", "do", "does", "did", "will", "would", "could", "should", "may", "might", "must",
|
|
"shall", "can", "need", "and", "but", "or", "if", "when", "while", "as", "because",
|
|
"although", "this", "that", "these", "those", "it", "they", "we", "you", "I", "which",
|
|
"who", "what", "where", "how", "why", "when",
|
|
];
|
|
|
|
let idx = (self.random() * words.len() as f64) as usize;
|
|
words[idx.min(words.len() - 1)].to_string()
|
|
}
|
|
|
|
/// Get performance metrics.
|
|
///
|
|
/// `memory_utilization` is derived from actual allocated tensor byte
|
|
/// sizes (`ModelShard::weights_mb`). Latency/throughput/utilization
|
|
/// figures are derived from real `Instant`-measured timings accumulated
|
|
/// during the most recent `generate()` call (zero before the first
|
|
/// call). `compute_utilization` and `network_utilization` are the
|
|
/// measured attention-compute and all-reduce shares of total wall-clock
|
|
/// time, respectively — real ratios, not fixed constants.
|
|
#[must_use]
|
|
pub fn get_metrics(&self) -> PerformanceMetrics {
|
|
let memory_used: f64 = self.shards.iter().map(|s| s.weights_mb).sum();
|
|
let total_memory = self.communicator.cluster.total_memory() * 1024.0; // Convert to MB
|
|
|
|
let (compute_utilization, network_utilization) = if self.last_total_time_ms > 0.0 {
|
|
(
|
|
(self.last_attention_compute_ms / self.last_total_time_ms).min(1.0),
|
|
(self.last_allreduce_time_ms / self.last_total_time_ms).min(1.0),
|
|
)
|
|
} else {
|
|
(0.0, 0.0)
|
|
};
|
|
|
|
PerformanceMetrics {
|
|
prefill_latency_ms: self.last_prefill_latency_ms,
|
|
decode_latency_ms: self.last_decode_latency_ms,
|
|
tokens_per_second: self.last_tokens_per_second,
|
|
memory_utilization: memory_used / total_memory,
|
|
compute_utilization,
|
|
network_utilization,
|
|
bubble_ratio: self.scheduler.bubble_ratio(),
|
|
allreduce_time_ms: self.last_allreduce_time_ms,
|
|
}
|
|
}
|
|
|
|
/// Random number generator.
|
|
fn random(&mut self) -> f64 {
|
|
self.rng_state = self
|
|
.rng_state
|
|
.wrapping_mul(6364136223846793005)
|
|
.wrapping_add(1442695040888963407);
|
|
(self.rng_state >> 11) as f64 / (1u64 << 53) as f64
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Run Demo
|
|
// ============================================================================
|
|
|
|
/// Run the distributed LLM demo.
|
|
///
|
|
/// Uses [`sample_data::tiny_realcompute_config`] rather than a real 70B/405B
|
|
/// config: this demo now performs *real* `Tensor::randn` weight allocation
|
|
/// and real attention matmul/softmax compute (see module docs), and a
|
|
/// genuine 70B-parameter model's weights cannot be allocated on typical
|
|
/// demo/CI hardware. The cluster topology and parallelism strategy are
|
|
/// still the full 4-node Thunderbolt 5 configuration.
|
|
#[must_use]
|
|
pub fn run_demo() -> InferenceResult {
|
|
let model_config = sample_data::tiny_realcompute_config();
|
|
let parallel_config = sample_data::four_node_parallel_config();
|
|
let cluster = sample_data::four_node_cluster();
|
|
|
|
let mut llm = DistributedLLM::new(model_config, parallel_config, cluster, 0);
|
|
llm.load_model();
|
|
|
|
let request = sample_data::chat_inference_request();
|
|
let gen_config = GenerationConfig {
|
|
max_new_tokens: 128,
|
|
do_sample: true,
|
|
..Default::default()
|
|
};
|
|
|
|
llm.generate(&request, &gen_config)
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_model_shard_creation() {
|
|
let shard = ModelShard::new(0, 0, 20);
|
|
assert_eq!(shard.num_layers(), 20);
|
|
assert!(!shard.loaded);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_shard_load() {
|
|
let mut shard = ModelShard::new(0, 0, 4);
|
|
let model_config = sample_data::tiny_realcompute_config();
|
|
shard
|
|
.load(&model_config)
|
|
.expect("weight allocation should succeed");
|
|
assert!(shard.loaded);
|
|
assert_eq!(shard.layer_weights.len(), 4);
|
|
// Real tensor-derived byte size must be positive.
|
|
assert!(shard.weights_mb > 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_communicator_creation() {
|
|
let cluster = sample_data::four_node_cluster();
|
|
let comm = DistributedCommunicator::new(0, cluster);
|
|
assert_eq!(comm.rank, 0);
|
|
assert_eq!(comm.world_size, 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_all_reduce_time() {
|
|
let cluster = sample_data::four_node_cluster();
|
|
let mut comm = DistributedCommunicator::new(0, cluster);
|
|
let time = comm.all_reduce(100.0);
|
|
assert!(time > 0.0);
|
|
assert_eq!(comm.message_count(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pipeline_scheduler() {
|
|
let mut scheduler = PipelineScheduler::new(4, 8);
|
|
assert!(scheduler.next_stage().is_some());
|
|
scheduler.advance();
|
|
assert_eq!(scheduler.current_micro_batch, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bubble_ratio() {
|
|
let scheduler = PipelineScheduler::new(4, 16);
|
|
let ratio = scheduler.bubble_ratio();
|
|
// 3 bubble slots out of 64 total = 0.046875
|
|
assert!(ratio < 0.1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_distributed_llm_creation() {
|
|
let model_config = sample_data::llama_70b_config();
|
|
let parallel_config = sample_data::four_node_parallel_config();
|
|
let cluster = sample_data::four_node_cluster();
|
|
|
|
let llm = DistributedLLM::new(model_config, parallel_config, cluster, 0);
|
|
assert!(!llm.loaded);
|
|
}
|
|
|
|
#[test]
|
|
fn test_load_model() {
|
|
// Real weight allocation happens here, so use the tiny compute
|
|
// config rather than a genuine 70B-scale model (see
|
|
// `sample_data::tiny_realcompute_config` docs).
|
|
let model_config = sample_data::tiny_realcompute_config();
|
|
let parallel_config = sample_data::four_node_parallel_config();
|
|
let cluster = sample_data::four_node_cluster();
|
|
|
|
let mut llm = DistributedLLM::new(model_config, parallel_config, cluster, 0);
|
|
llm.load_model();
|
|
|
|
assert!(llm.loaded);
|
|
assert!(!llm.layer_assignments.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate() {
|
|
let model_config = sample_data::tiny_realcompute_config();
|
|
let parallel_config = ParallelismConfig {
|
|
tensor_parallel: 1,
|
|
pipeline_parallel: 1,
|
|
world_size: 1,
|
|
..Default::default()
|
|
};
|
|
let cluster = distllm_shared::ClusterConfig::default();
|
|
|
|
let mut llm = DistributedLLM::new(model_config, parallel_config, cluster, 0);
|
|
|
|
let request = InferenceRequest {
|
|
prompt: "Hello world".to_string(),
|
|
max_tokens: 10,
|
|
..Default::default()
|
|
};
|
|
|
|
let config = GenerationConfig {
|
|
max_new_tokens: 10,
|
|
..Default::default()
|
|
};
|
|
|
|
let result = llm.generate(&request, &config);
|
|
|
|
assert!(!result.tokens.is_empty());
|
|
assert!(result.latency_ms > 0.0);
|
|
assert!(result.tokens_per_second > 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_metrics() {
|
|
let model_config = sample_data::tiny_realcompute_config();
|
|
let parallel_config = ParallelismConfig {
|
|
tensor_parallel: 1,
|
|
pipeline_parallel: 1,
|
|
world_size: 1,
|
|
..Default::default()
|
|
};
|
|
let cluster = distllm_shared::ClusterConfig::default();
|
|
|
|
let mut llm = DistributedLLM::new(model_config, parallel_config, cluster, 0);
|
|
llm.load_model();
|
|
|
|
let request = InferenceRequest {
|
|
prompt: "Hello world".to_string(),
|
|
max_tokens: 4,
|
|
..Default::default()
|
|
};
|
|
let config = GenerationConfig {
|
|
max_new_tokens: 4,
|
|
..Default::default()
|
|
};
|
|
let _ = llm.generate(&request, &config);
|
|
|
|
let metrics = llm.get_metrics();
|
|
assert!(metrics.tokens_per_second > 0.0);
|
|
assert!(metrics.memory_utilization > 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_run_demo() {
|
|
let result = run_demo();
|
|
assert!(!result.tokens.is_empty());
|
|
assert!(result.completion_tokens > 0);
|
|
}
|
|
}
|