613 lines
18 KiB
Rust
613 lines
18 KiB
Rust
//! Parallelism strategies for distributed LLM inference.
|
|
//!
|
|
//! This module implements tensor parallelism and pipeline parallelism
|
|
//! for distributing large language models across multiple nodes.
|
|
|
|
use distllm_shared::LayerAssignment;
|
|
|
|
// ============================================================================
|
|
// Tensor Parallelism
|
|
// ============================================================================
|
|
|
|
/// Tensor parallelism for splitting linear layers across nodes.
|
|
///
|
|
/// In tensor parallelism, each attention head and FFN dimension is split
|
|
/// across multiple devices. This requires all-reduce after each layer.
|
|
#[derive(Debug, Clone)]
|
|
pub struct TensorParallel {
|
|
/// Tensor parallel world size.
|
|
pub world_size: usize,
|
|
/// This node's rank in tensor parallel group.
|
|
pub rank: usize,
|
|
/// Shard size for attention heads.
|
|
pub head_shard_size: usize,
|
|
/// Shard size for FFN hidden dimension.
|
|
pub ffn_shard_size: usize,
|
|
}
|
|
|
|
impl TensorParallel {
|
|
/// Create a new tensor parallel configuration.
|
|
#[must_use]
|
|
pub fn new(world_size: usize, rank: usize) -> Self {
|
|
Self {
|
|
world_size,
|
|
rank,
|
|
head_shard_size: 0,
|
|
ffn_shard_size: 0,
|
|
}
|
|
}
|
|
|
|
/// Configure for a specific model.
|
|
pub fn configure(&mut self, num_heads: usize, ffn_dim: usize) {
|
|
self.head_shard_size = num_heads / self.world_size;
|
|
self.ffn_shard_size = ffn_dim / self.world_size;
|
|
}
|
|
|
|
/// Get the attention head range for this rank.
|
|
#[must_use]
|
|
pub fn head_range(&self, total_heads: usize) -> (usize, usize) {
|
|
let heads_per_rank = total_heads / self.world_size;
|
|
let start = self.rank * heads_per_rank;
|
|
let end = if self.rank == self.world_size - 1 {
|
|
total_heads
|
|
} else {
|
|
(self.rank + 1) * heads_per_rank
|
|
};
|
|
(start, end)
|
|
}
|
|
|
|
/// Get the FFN dimension range for this rank.
|
|
#[must_use]
|
|
pub fn ffn_range(&self, total_dim: usize) -> (usize, usize) {
|
|
let dim_per_rank = total_dim / self.world_size;
|
|
let start = self.rank * dim_per_rank;
|
|
let end = if self.rank == self.world_size - 1 {
|
|
total_dim
|
|
} else {
|
|
(self.rank + 1) * dim_per_rank
|
|
};
|
|
(start, end)
|
|
}
|
|
|
|
/// Split a linear layer weight matrix column-wise.
|
|
#[must_use]
|
|
pub fn column_split(&self, weight_shape: (usize, usize)) -> (usize, usize) {
|
|
let (rows, cols) = weight_shape;
|
|
let cols_per_rank = cols / self.world_size;
|
|
(rows, cols_per_rank)
|
|
}
|
|
|
|
/// Split a linear layer weight matrix row-wise.
|
|
#[must_use]
|
|
pub fn row_split(&self, weight_shape: (usize, usize)) -> (usize, usize) {
|
|
let (rows, cols) = weight_shape;
|
|
let rows_per_rank = rows / self.world_size;
|
|
(rows_per_rank, cols)
|
|
}
|
|
|
|
/// Calculate communication volume for all-reduce in MB.
|
|
#[must_use]
|
|
pub fn all_reduce_volume_mb(
|
|
&self,
|
|
hidden_dim: usize,
|
|
seq_len: usize,
|
|
dtype_bytes: usize,
|
|
) -> f64 {
|
|
let elements = hidden_dim * seq_len;
|
|
let bytes = elements * dtype_bytes;
|
|
bytes as f64 / (1024.0 * 1024.0)
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Pipeline Parallelism
|
|
// ============================================================================
|
|
|
|
/// Pipeline parallelism for layer-wise distribution.
|
|
///
|
|
/// In pipeline parallelism, consecutive layers are assigned to different
|
|
/// stages. This requires point-to-point communication between stages.
|
|
#[derive(Debug, Clone)]
|
|
pub struct PipelineParallel {
|
|
/// Number of pipeline stages.
|
|
pub num_stages: usize,
|
|
/// This node's stage index.
|
|
pub stage: usize,
|
|
/// Layers assigned to each stage.
|
|
pub layers_per_stage: Vec<usize>,
|
|
}
|
|
|
|
impl PipelineParallel {
|
|
/// Create a new pipeline parallel configuration.
|
|
#[must_use]
|
|
pub fn new(num_stages: usize, stage: usize) -> Self {
|
|
Self {
|
|
num_stages,
|
|
stage,
|
|
layers_per_stage: vec![],
|
|
}
|
|
}
|
|
|
|
/// Configure for a specific number of layers.
|
|
pub fn configure(&mut self, num_layers: usize) {
|
|
let base_layers = num_layers / self.num_stages;
|
|
let extra = num_layers % self.num_stages;
|
|
|
|
self.layers_per_stage = (0..self.num_stages)
|
|
.map(|s| base_layers + usize::from(s < extra))
|
|
.collect();
|
|
}
|
|
|
|
/// Get the layer range for this stage.
|
|
#[must_use]
|
|
pub fn layer_range(&self) -> (usize, usize) {
|
|
if self.layers_per_stage.is_empty() {
|
|
return (0, 0);
|
|
}
|
|
|
|
let start: usize = self.layers_per_stage[..self.stage].iter().sum();
|
|
let end = start + self.layers_per_stage[self.stage];
|
|
(start, end)
|
|
}
|
|
|
|
/// Get the previous stage (for receiving activations).
|
|
#[must_use]
|
|
pub fn prev_stage(&self) -> Option<usize> {
|
|
if self.stage > 0 {
|
|
Some(self.stage - 1)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Get the next stage (for sending activations).
|
|
#[must_use]
|
|
pub fn next_stage(&self) -> Option<usize> {
|
|
if self.stage < self.num_stages - 1 {
|
|
Some(self.stage + 1)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Check if this is the first stage.
|
|
#[must_use]
|
|
pub fn is_first_stage(&self) -> bool {
|
|
self.stage == 0
|
|
}
|
|
|
|
/// Check if this is the last stage.
|
|
#[must_use]
|
|
pub fn is_last_stage(&self) -> bool {
|
|
self.stage == self.num_stages - 1
|
|
}
|
|
|
|
/// Calculate activation transfer size in MB.
|
|
#[must_use]
|
|
pub fn activation_size_mb(
|
|
&self,
|
|
hidden_dim: usize,
|
|
seq_len: usize,
|
|
batch_size: usize,
|
|
dtype_bytes: usize,
|
|
) -> f64 {
|
|
let elements = hidden_dim * seq_len * batch_size;
|
|
let bytes = elements * dtype_bytes;
|
|
bytes as f64 / (1024.0 * 1024.0)
|
|
}
|
|
|
|
/// Calculate pipeline bubble ratio.
|
|
#[must_use]
|
|
pub fn bubble_ratio(&self, num_micro_batches: usize) -> f64 {
|
|
let p = self.num_stages as f64;
|
|
let m = num_micro_batches as f64;
|
|
(p - 1.0) / (p + m - 1.0)
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Layer Partitioner
|
|
// ============================================================================
|
|
|
|
/// Partitioner for optimal layer assignment across nodes.
|
|
#[derive(Debug)]
|
|
pub struct LayerPartitioner {
|
|
/// Total number of layers.
|
|
pub num_layers: usize,
|
|
/// Total number of nodes.
|
|
pub num_nodes: usize,
|
|
/// Tensor parallelism degree.
|
|
pub tensor_parallel: usize,
|
|
/// Pipeline parallelism degree.
|
|
pub pipeline_parallel: usize,
|
|
/// Layer memory costs (optional).
|
|
pub layer_costs: Option<Vec<f64>>,
|
|
}
|
|
|
|
impl LayerPartitioner {
|
|
/// Create a new layer partitioner.
|
|
#[must_use]
|
|
pub fn new(
|
|
num_layers: usize,
|
|
num_nodes: usize,
|
|
tensor_parallel: usize,
|
|
pipeline_parallel: usize,
|
|
) -> Self {
|
|
Self {
|
|
num_layers,
|
|
num_nodes,
|
|
tensor_parallel,
|
|
pipeline_parallel,
|
|
layer_costs: None,
|
|
}
|
|
}
|
|
|
|
/// Set custom layer memory costs.
|
|
pub fn set_layer_costs(&mut self, costs: Vec<f64>) {
|
|
self.layer_costs = Some(costs);
|
|
}
|
|
|
|
/// Partition layers across nodes.
|
|
#[must_use]
|
|
pub fn partition(&self) -> Vec<LayerAssignment> {
|
|
let mut assignments = Vec::with_capacity(self.num_layers);
|
|
|
|
// Calculate layers per pipeline stage
|
|
let layers_per_stage = self.num_layers / self.pipeline_parallel;
|
|
let extra_layers = self.num_layers % self.pipeline_parallel;
|
|
|
|
let mut layer_idx = 0;
|
|
for stage in 0..self.pipeline_parallel {
|
|
let stage_layers = layers_per_stage + usize::from(stage < extra_layers);
|
|
|
|
for local_layer in 0..stage_layers {
|
|
// For tensor parallel, each layer is replicated across TP ranks
|
|
// but we assign to the primary node (TP rank 0)
|
|
let node_id = stage * self.tensor_parallel;
|
|
|
|
let memory_mb = self
|
|
.layer_costs
|
|
.as_ref()
|
|
.map_or(1024.0, |c| c.get(layer_idx).copied().unwrap_or(1024.0));
|
|
|
|
assignments.push(LayerAssignment {
|
|
layer_id: layer_idx,
|
|
node_id,
|
|
memory_mb,
|
|
pipeline_stage: stage,
|
|
tensor_rank: 0,
|
|
});
|
|
|
|
layer_idx += 1;
|
|
let _ = local_layer; // Suppress unused warning
|
|
}
|
|
}
|
|
|
|
assignments
|
|
}
|
|
|
|
/// Partition with load balancing based on layer costs.
|
|
#[must_use]
|
|
pub fn partition_balanced(&self) -> Vec<LayerAssignment> {
|
|
if self.layer_costs.is_none() {
|
|
return self.partition();
|
|
}
|
|
|
|
let costs = self.layer_costs.as_ref().unwrap();
|
|
let total_cost: f64 = costs.iter().sum();
|
|
let target_cost_per_stage = total_cost / self.pipeline_parallel as f64;
|
|
|
|
let mut assignments = Vec::with_capacity(self.num_layers);
|
|
let mut current_stage = 0;
|
|
let mut current_stage_cost = 0.0;
|
|
|
|
for (layer_idx, &cost) in costs.iter().enumerate() {
|
|
let node_id = current_stage * self.tensor_parallel;
|
|
|
|
assignments.push(LayerAssignment {
|
|
layer_id: layer_idx,
|
|
node_id,
|
|
memory_mb: cost,
|
|
pipeline_stage: current_stage,
|
|
tensor_rank: 0,
|
|
});
|
|
|
|
current_stage_cost += cost;
|
|
|
|
// Move to next stage if we've exceeded target and not on last stage
|
|
if current_stage_cost >= target_cost_per_stage
|
|
&& current_stage < self.pipeline_parallel - 1
|
|
{
|
|
current_stage += 1;
|
|
current_stage_cost = 0.0;
|
|
}
|
|
}
|
|
|
|
assignments
|
|
}
|
|
|
|
/// Get memory usage per node.
|
|
#[must_use]
|
|
pub fn memory_per_node(&self, assignments: &[LayerAssignment]) -> Vec<f64> {
|
|
let mut memory = vec![0.0; self.num_nodes];
|
|
for a in assignments {
|
|
if a.node_id < memory.len() {
|
|
memory[a.node_id] += a.memory_mb;
|
|
}
|
|
}
|
|
memory
|
|
}
|
|
|
|
/// Calculate load imbalance ratio.
|
|
#[must_use]
|
|
pub fn imbalance_ratio(&self, assignments: &[LayerAssignment]) -> f64 {
|
|
let memory = self.memory_per_node(assignments);
|
|
let max_mem = memory.iter().copied().fold(0.0, f64::max);
|
|
let min_mem = memory.iter().copied().fold(f64::MAX, f64::min);
|
|
let avg_mem = memory.iter().sum::<f64>() / memory.len() as f64;
|
|
|
|
if avg_mem > 0.0 {
|
|
(max_mem - min_mem) / avg_mem
|
|
} else {
|
|
0.0
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Micro-batch Schedule
|
|
// ============================================================================
|
|
|
|
/// Micro-batch schedule for pipeline execution.
|
|
#[derive(Debug, Clone)]
|
|
pub struct MicroBatchSchedule {
|
|
/// Number of micro-batches.
|
|
pub num_micro_batches: usize,
|
|
/// Number of pipeline stages.
|
|
pub num_stages: usize,
|
|
/// Schedule entries (stage, micro_batch, is_forward).
|
|
pub entries: Vec<(usize, usize, bool)>,
|
|
}
|
|
|
|
impl MicroBatchSchedule {
|
|
/// Create a GPipe schedule (all forwards then all backwards).
|
|
#[must_use]
|
|
pub fn gpipe(num_stages: usize, num_micro_batches: usize) -> Self {
|
|
let mut entries = Vec::new();
|
|
|
|
// Forward passes
|
|
for mb in 0..num_micro_batches {
|
|
for stage in 0..num_stages {
|
|
entries.push((stage, mb, true));
|
|
}
|
|
}
|
|
|
|
// Backward passes
|
|
for mb in (0..num_micro_batches).rev() {
|
|
for stage in (0..num_stages).rev() {
|
|
entries.push((stage, mb, false));
|
|
}
|
|
}
|
|
|
|
Self {
|
|
num_micro_batches,
|
|
num_stages,
|
|
entries,
|
|
}
|
|
}
|
|
|
|
/// Create a 1F1B schedule (interleaved forward/backward).
|
|
#[must_use]
|
|
pub fn one_f_one_b(num_stages: usize, num_micro_batches: usize) -> Self {
|
|
let mut entries = Vec::new();
|
|
|
|
// Warmup: forward passes for first (num_stages) micro-batches
|
|
for mb in 0..num_stages.min(num_micro_batches) {
|
|
for stage in 0..num_stages {
|
|
entries.push((stage, mb, true));
|
|
}
|
|
}
|
|
|
|
// Steady state: alternate 1F1B
|
|
for mb in num_stages..num_micro_batches {
|
|
for stage in 0..num_stages {
|
|
entries.push((stage, mb, true));
|
|
}
|
|
let backward_mb = mb - num_stages;
|
|
for stage in (0..num_stages).rev() {
|
|
entries.push((stage, backward_mb, false));
|
|
}
|
|
}
|
|
|
|
// Cooldown: remaining backward passes
|
|
for mb in (num_micro_batches.saturating_sub(num_stages)..num_micro_batches).rev() {
|
|
for stage in (0..num_stages).rev() {
|
|
entries.push((stage, mb, false));
|
|
}
|
|
}
|
|
|
|
Self {
|
|
num_micro_batches,
|
|
num_stages,
|
|
entries,
|
|
}
|
|
}
|
|
|
|
/// Get schedule for a specific stage.
|
|
#[must_use]
|
|
pub fn stage_schedule(&self, stage: usize) -> Vec<(usize, bool)> {
|
|
self.entries
|
|
.iter()
|
|
.filter(|(s, _, _)| *s == stage)
|
|
.map(|(_, mb, fwd)| (*mb, *fwd))
|
|
.collect()
|
|
}
|
|
|
|
/// Calculate bubble time slots.
|
|
#[must_use]
|
|
pub fn bubble_slots(&self) -> usize {
|
|
// In ideal case with no bubbles, each stage would execute
|
|
// 2 * num_micro_batches operations (fwd + bwd)
|
|
let ideal = self.num_stages * self.num_micro_batches * 2;
|
|
let actual = self.entries.len();
|
|
actual.saturating_sub(ideal) / self.num_stages
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_tensor_parallel_creation() {
|
|
let tp = TensorParallel::new(4, 1);
|
|
assert_eq!(tp.world_size, 4);
|
|
assert_eq!(tp.rank, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_parallel_head_range() {
|
|
let tp = TensorParallel::new(4, 1);
|
|
let (start, end) = tp.head_range(64);
|
|
assert_eq!(start, 16);
|
|
assert_eq!(end, 32);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_parallel_ffn_range() {
|
|
let tp = TensorParallel::new(4, 2);
|
|
let (start, end) = tp.ffn_range(28672);
|
|
assert_eq!(start, 14336);
|
|
assert_eq!(end, 21504);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_parallel_column_split() {
|
|
let tp = TensorParallel::new(4, 0);
|
|
let (rows, cols) = tp.column_split((4096, 16384));
|
|
assert_eq!(rows, 4096);
|
|
assert_eq!(cols, 4096);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pipeline_parallel_creation() {
|
|
let pp = PipelineParallel::new(4, 1);
|
|
assert_eq!(pp.num_stages, 4);
|
|
assert_eq!(pp.stage, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pipeline_parallel_configure() {
|
|
let mut pp = PipelineParallel::new(4, 0);
|
|
pp.configure(80);
|
|
assert_eq!(pp.layers_per_stage.iter().sum::<usize>(), 80);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pipeline_parallel_layer_range() {
|
|
let mut pp = PipelineParallel::new(4, 1);
|
|
pp.configure(80);
|
|
let (start, end) = pp.layer_range();
|
|
assert_eq!(start, 20);
|
|
assert_eq!(end, 40);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pipeline_parallel_neighbors() {
|
|
let pp = PipelineParallel::new(4, 1);
|
|
assert_eq!(pp.prev_stage(), Some(0));
|
|
assert_eq!(pp.next_stage(), Some(2));
|
|
|
|
let first = PipelineParallel::new(4, 0);
|
|
assert!(first.is_first_stage());
|
|
assert_eq!(first.prev_stage(), None);
|
|
|
|
let last = PipelineParallel::new(4, 3);
|
|
assert!(last.is_last_stage());
|
|
assert_eq!(last.next_stage(), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pipeline_bubble_ratio() {
|
|
let pp = PipelineParallel::new(4, 0);
|
|
let ratio = pp.bubble_ratio(16);
|
|
// (4-1)/(4+16-1) = 3/19 = 0.157...
|
|
assert!((ratio - 0.158).abs() < 0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_partitioner() {
|
|
let partitioner = LayerPartitioner::new(80, 4, 1, 4);
|
|
let assignments = partitioner.partition();
|
|
assert_eq!(assignments.len(), 80);
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_partitioner_stages() {
|
|
let partitioner = LayerPartitioner::new(80, 4, 1, 4);
|
|
let assignments = partitioner.partition();
|
|
|
|
// First 20 layers should be stage 0
|
|
assert!(assignments[..20].iter().all(|a| a.pipeline_stage == 0));
|
|
// Next 20 should be stage 1
|
|
assert!(assignments[20..40].iter().all(|a| a.pipeline_stage == 1));
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_partitioner_balanced() {
|
|
let mut partitioner = LayerPartitioner::new(8, 4, 1, 4);
|
|
partitioner.set_layer_costs(vec![1.0, 2.0, 3.0, 4.0, 4.0, 3.0, 2.0, 1.0]);
|
|
let assignments = partitioner.partition_balanced();
|
|
assert_eq!(assignments.len(), 8);
|
|
}
|
|
|
|
#[test]
|
|
fn test_memory_per_node() {
|
|
let partitioner = LayerPartitioner::new(8, 2, 1, 2);
|
|
let assignments = partitioner.partition();
|
|
let memory = partitioner.memory_per_node(&assignments);
|
|
assert_eq!(memory.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gpipe_schedule() {
|
|
let schedule = MicroBatchSchedule::gpipe(4, 8);
|
|
assert!(!schedule.entries.is_empty());
|
|
|
|
// All forwards should come before backwards
|
|
let first_backward = schedule
|
|
.entries
|
|
.iter()
|
|
.position(|(_, _, fwd)| !fwd)
|
|
.unwrap();
|
|
let last_forward = schedule
|
|
.entries
|
|
.iter()
|
|
.rposition(|(_, _, fwd)| *fwd)
|
|
.unwrap();
|
|
assert!(last_forward < first_backward);
|
|
}
|
|
|
|
#[test]
|
|
fn test_1f1b_schedule() {
|
|
let schedule = MicroBatchSchedule::one_f_one_b(4, 8);
|
|
assert!(!schedule.entries.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_stage_schedule() {
|
|
let schedule = MicroBatchSchedule::gpipe(4, 4);
|
|
let stage_0 = schedule.stage_schedule(0);
|
|
assert!(!stage_0.is_empty());
|
|
// Stage 0 should have 4 forwards and 4 backwards
|
|
let forwards = stage_0.iter().filter(|(_, fwd)| *fwd).count();
|
|
let backwards = stage_0.iter().filter(|(_, fwd)| !*fwd).count();
|
|
assert_eq!(forwards, 4);
|
|
assert_eq!(backwards, 4);
|
|
}
|
|
}
|