1019 lines
33 KiB
Rust
1019 lines
33 KiB
Rust
//! Automatic Model Partitioning for Distributed Training
|
|
//!
|
|
//! This module provides automatic model partitioning capabilities that
|
|
//! analyze model architecture and generate optimal sharding plans for
|
|
//! tensor parallelism (TP), pipeline parallelism (PP), and FSDP.
|
|
//!
|
|
//! # Features
|
|
//!
|
|
//! - Memory-aware partitioning to fit models on available GPU memory
|
|
//! - Compute-balanced distribution for optimal throughput
|
|
//! - Communication-minimal placement to reduce data movement
|
|
//! - Hybrid strategies combining multiple optimization goals
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_distributed::auto_partition::{AutoPartitioner, AutoPartitionConfig, ModelInfo};
|
|
//!
|
|
//! let config = AutoPartitionConfig::default();
|
|
//! let partitioner = AutoPartitioner::new(config);
|
|
//!
|
|
//! let model_info = ModelInfo::from_layer_sizes(&[
|
|
//! ("embedding", 1_000_000_000),
|
|
//! ("transformer.0", 500_000_000),
|
|
//! ("transformer.1", 500_000_000),
|
|
//! ("lm_head", 500_000_000),
|
|
//! ]);
|
|
//!
|
|
//! let plan = partitioner.partition(&model_info)?;
|
|
//! println!("{:?}", plan);
|
|
//! ```
|
|
|
|
use crate::error::{DistributedError, Result};
|
|
use crate::model_sharding::ShardingStrategy;
|
|
use parking_lot::RwLock;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
// =============================================================================
|
|
// Configuration
|
|
// =============================================================================
|
|
|
|
/// Strategy for automatic partitioning
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum AutoPartitionStrategy {
|
|
/// Minimize peak memory usage
|
|
MemoryOptimal,
|
|
/// Balance compute across devices
|
|
ComputeBalanced,
|
|
/// Minimize cross-device communication
|
|
CommunicationMinimal,
|
|
/// Balance all factors with weighted importance
|
|
Hybrid,
|
|
/// Prefer tensor parallelism for large layers
|
|
TensorParallelFirst,
|
|
/// Prefer pipeline parallelism for sequential models
|
|
PipelineParallelFirst,
|
|
}
|
|
|
|
impl Default for AutoPartitionStrategy {
|
|
fn default() -> Self {
|
|
Self::Hybrid
|
|
}
|
|
}
|
|
|
|
/// Configuration for the auto partitioner
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct AutoPartitionConfig {
|
|
/// Partitioning strategy
|
|
pub strategy: AutoPartitionStrategy,
|
|
/// Memory margin (0.1 = leave 10% headroom)
|
|
pub memory_margin: f32,
|
|
/// Maximum allowed compute imbalance ratio (1.2 = 20% max difference)
|
|
pub balance_threshold: f32,
|
|
/// Prefer tensor parallelism for large layers
|
|
pub prefer_tp_for_large_layers: bool,
|
|
/// Prefer pipeline parallelism for sequential structures
|
|
pub prefer_pp_for_sequential: bool,
|
|
/// Minimum layer size (bytes) to consider for tensor parallelism
|
|
pub min_tp_layer_size: usize,
|
|
/// Number of devices available
|
|
pub num_devices: usize,
|
|
/// Per-device memory budget (bytes)
|
|
pub device_memory: usize,
|
|
/// Weight for memory optimization in hybrid mode
|
|
pub memory_weight: f32,
|
|
/// Weight for compute balance in hybrid mode
|
|
pub compute_weight: f32,
|
|
/// Weight for communication minimization in hybrid mode
|
|
pub communication_weight: f32,
|
|
}
|
|
|
|
impl Default for AutoPartitionConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
strategy: AutoPartitionStrategy::Hybrid,
|
|
memory_margin: 0.1,
|
|
balance_threshold: 1.2,
|
|
prefer_tp_for_large_layers: true,
|
|
prefer_pp_for_sequential: true,
|
|
min_tp_layer_size: 100 * 1024 * 1024, // 100 MB
|
|
num_devices: 8,
|
|
device_memory: 80 * 1024 * 1024 * 1024, // 80 GB (A100/H100)
|
|
memory_weight: 0.4,
|
|
compute_weight: 0.3,
|
|
communication_weight: 0.3,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl AutoPartitionConfig {
|
|
/// Create config for memory-constrained environments
|
|
pub fn memory_optimized(num_devices: usize, device_memory: usize) -> Self {
|
|
Self {
|
|
strategy: AutoPartitionStrategy::MemoryOptimal,
|
|
num_devices,
|
|
device_memory,
|
|
memory_margin: 0.15,
|
|
memory_weight: 0.6,
|
|
compute_weight: 0.2,
|
|
communication_weight: 0.2,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
/// Create config for throughput-focused training
|
|
pub fn throughput_optimized(num_devices: usize) -> Self {
|
|
Self {
|
|
strategy: AutoPartitionStrategy::ComputeBalanced,
|
|
num_devices,
|
|
memory_weight: 0.2,
|
|
compute_weight: 0.6,
|
|
communication_weight: 0.2,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
/// Create config for communication-sensitive environments
|
|
pub fn communication_optimized(num_devices: usize) -> Self {
|
|
Self {
|
|
strategy: AutoPartitionStrategy::CommunicationMinimal,
|
|
num_devices,
|
|
memory_weight: 0.2,
|
|
compute_weight: 0.2,
|
|
communication_weight: 0.6,
|
|
..Default::default()
|
|
}
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Model Information
|
|
// =============================================================================
|
|
|
|
/// Information about a model layer
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct LayerInfo {
|
|
/// Layer name/path
|
|
pub name: String,
|
|
/// Parameter count
|
|
pub param_count: usize,
|
|
/// Memory size in bytes (including activations estimate)
|
|
pub memory_bytes: usize,
|
|
/// Estimated FLOPs per forward pass
|
|
pub flops: usize,
|
|
/// Layer type
|
|
pub layer_type: LayerType,
|
|
/// Whether this layer is shardable across tensor dimension
|
|
pub is_tp_shardable: bool,
|
|
/// Connections to other layers (for graph analysis)
|
|
pub connections: Vec<String>,
|
|
}
|
|
|
|
/// Types of model layers
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum LayerType {
|
|
/// Embedding layer
|
|
Embedding,
|
|
/// Linear/Dense layer
|
|
Linear,
|
|
/// Attention layer
|
|
Attention,
|
|
/// LayerNorm/BatchNorm
|
|
Normalization,
|
|
/// Activation function
|
|
Activation,
|
|
/// Convolutional layer
|
|
Convolution,
|
|
/// Pooling layer
|
|
Pooling,
|
|
/// Residual connection
|
|
Residual,
|
|
/// Other/Unknown
|
|
Other,
|
|
}
|
|
|
|
impl LayerInfo {
|
|
/// Create from basic size information
|
|
pub fn new(name: impl Into<String>, param_count: usize, layer_type: LayerType) -> Self {
|
|
// Estimate 4 bytes per param + 3x for gradients and optimizer states
|
|
let memory_bytes = param_count * 4 * 4;
|
|
// Rough FLOP estimate based on layer type
|
|
let flops = match layer_type {
|
|
LayerType::Linear => param_count * 2,
|
|
LayerType::Attention => param_count * 4,
|
|
_ => param_count,
|
|
};
|
|
|
|
Self {
|
|
name: name.into(),
|
|
param_count,
|
|
memory_bytes,
|
|
flops,
|
|
layer_type,
|
|
is_tp_shardable: matches!(
|
|
layer_type,
|
|
LayerType::Linear | LayerType::Embedding | LayerType::Attention
|
|
),
|
|
connections: Vec::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Complete model information for partitioning
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ModelInfo {
|
|
/// Model name
|
|
pub name: String,
|
|
/// Total parameter count
|
|
pub total_params: usize,
|
|
/// Total memory footprint
|
|
pub total_memory: usize,
|
|
/// Layers in execution order
|
|
pub layers: Vec<LayerInfo>,
|
|
/// Is the model purely sequential?
|
|
pub is_sequential: bool,
|
|
/// Detected transformer blocks
|
|
pub transformer_blocks: Vec<TransformerBlock>,
|
|
}
|
|
|
|
/// Detected transformer block structure
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TransformerBlock {
|
|
/// Block index
|
|
pub index: usize,
|
|
/// Layer names in this block
|
|
pub layer_names: Vec<String>,
|
|
/// Block memory size
|
|
pub memory_bytes: usize,
|
|
/// Block FLOPs
|
|
pub flops: usize,
|
|
}
|
|
|
|
impl ModelInfo {
|
|
/// Create from layer sizes
|
|
pub fn from_layer_sizes(layers: &[(&str, usize)]) -> Self {
|
|
let layer_infos: Vec<LayerInfo> = layers
|
|
.iter()
|
|
.map(|(name, size)| {
|
|
let layer_type = if name.contains("embed") {
|
|
LayerType::Embedding
|
|
} else if name.contains("attention") || name.contains("attn") {
|
|
LayerType::Attention
|
|
} else if name.contains("norm") {
|
|
LayerType::Normalization
|
|
} else {
|
|
LayerType::Linear
|
|
};
|
|
LayerInfo::new(*name, *size / 4, layer_type) // Convert bytes to param count
|
|
})
|
|
.collect();
|
|
|
|
let total_params = layer_infos.iter().map(|l| l.param_count).sum();
|
|
let total_memory = layer_infos.iter().map(|l| l.memory_bytes).sum();
|
|
|
|
Self {
|
|
name: "model".to_string(),
|
|
total_params,
|
|
total_memory,
|
|
layers: layer_infos,
|
|
is_sequential: true,
|
|
transformer_blocks: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Create from layer info
|
|
pub fn from_layers(name: impl Into<String>, layers: Vec<LayerInfo>) -> Self {
|
|
let total_params = layers.iter().map(|l| l.param_count).sum();
|
|
let total_memory = layers.iter().map(|l| l.memory_bytes).sum();
|
|
|
|
Self {
|
|
name: name.into(),
|
|
total_params,
|
|
total_memory,
|
|
layers,
|
|
is_sequential: true,
|
|
transformer_blocks: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Detect transformer blocks in the model
|
|
pub fn detect_transformer_blocks(&mut self) {
|
|
let mut blocks = Vec::new();
|
|
let mut current_block_layers = Vec::new();
|
|
let mut current_block_start = 0;
|
|
|
|
for (i, layer) in self.layers.iter().enumerate() {
|
|
if layer.name.contains("transformer") || layer.name.contains("block") {
|
|
// Extract block index from name
|
|
if let Some(idx) = extract_block_index(&layer.name) {
|
|
if !current_block_layers.is_empty() && idx != current_block_start {
|
|
// Save previous block
|
|
blocks.push(TransformerBlock {
|
|
index: current_block_start,
|
|
layer_names: current_block_layers.clone(),
|
|
memory_bytes: current_block_layers
|
|
.iter()
|
|
.filter_map(|name| self.layers.iter().find(|l| &l.name == name))
|
|
.map(|l| l.memory_bytes)
|
|
.sum(),
|
|
flops: current_block_layers
|
|
.iter()
|
|
.filter_map(|name| self.layers.iter().find(|l| &l.name == name))
|
|
.map(|l| l.flops)
|
|
.sum(),
|
|
});
|
|
current_block_layers.clear();
|
|
}
|
|
current_block_start = idx;
|
|
}
|
|
current_block_layers.push(layer.name.clone());
|
|
}
|
|
|
|
// Save last block
|
|
if i == self.layers.len() - 1 && !current_block_layers.is_empty() {
|
|
blocks.push(TransformerBlock {
|
|
index: current_block_start,
|
|
layer_names: current_block_layers.clone(),
|
|
memory_bytes: current_block_layers
|
|
.iter()
|
|
.filter_map(|name| self.layers.iter().find(|l| &l.name == name))
|
|
.map(|l| l.memory_bytes)
|
|
.sum(),
|
|
flops: current_block_layers
|
|
.iter()
|
|
.filter_map(|name| self.layers.iter().find(|l| &l.name == name))
|
|
.map(|l| l.flops)
|
|
.sum(),
|
|
});
|
|
}
|
|
}
|
|
|
|
self.transformer_blocks = blocks;
|
|
}
|
|
}
|
|
|
|
/// Extract block index from layer name (e.g., "transformer.12.attention" -> 12)
|
|
fn extract_block_index(name: &str) -> Option<usize> {
|
|
name.split('.').find_map(|s| s.parse::<usize>().ok())
|
|
}
|
|
|
|
// =============================================================================
|
|
// Sharding Plan
|
|
// =============================================================================
|
|
|
|
/// How a layer should be partitioned
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct LayerPartition {
|
|
/// Layer name
|
|
pub layer_name: String,
|
|
/// Sharding strategy for this layer
|
|
pub strategy: ShardingStrategy,
|
|
/// Device assignments (for PP)
|
|
pub device_ids: Vec<usize>,
|
|
/// TP dimension (for TP sharding)
|
|
pub tp_size: usize,
|
|
/// PP stage (for PP sharding)
|
|
pub pp_stage: usize,
|
|
}
|
|
|
|
/// Complete sharding plan for a model
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ShardingPlan {
|
|
/// Model name
|
|
pub model_name: String,
|
|
/// Per-layer partitioning
|
|
pub layer_partitions: Vec<LayerPartition>,
|
|
/// Total number of pipeline stages
|
|
pub num_pp_stages: usize,
|
|
/// Tensor parallel size
|
|
pub tp_size: usize,
|
|
/// Data parallel size (derived from devices / (tp * pp))
|
|
pub dp_size: usize,
|
|
/// Estimated peak memory per device
|
|
pub estimated_memory_per_device: usize,
|
|
/// Estimated compute imbalance factor
|
|
pub compute_imbalance: f32,
|
|
/// Estimated communication volume
|
|
pub communication_volume: usize,
|
|
/// Score from optimization
|
|
pub optimization_score: f64,
|
|
}
|
|
|
|
impl ShardingPlan {
|
|
/// Create a new sharding plan
|
|
pub fn new(model_name: impl Into<String>) -> Self {
|
|
Self {
|
|
model_name: model_name.into(),
|
|
layer_partitions: Vec::new(),
|
|
num_pp_stages: 1,
|
|
tp_size: 1,
|
|
dp_size: 1,
|
|
estimated_memory_per_device: 0,
|
|
compute_imbalance: 1.0,
|
|
communication_volume: 0,
|
|
optimization_score: 0.0,
|
|
}
|
|
}
|
|
|
|
/// Add a layer partition
|
|
pub fn add_partition(&mut self, partition: LayerPartition) {
|
|
self.layer_partitions.push(partition);
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Layer Profile
|
|
// =============================================================================
|
|
|
|
/// Profile of a layer's resource requirements
|
|
#[derive(Debug, Clone)]
|
|
struct LayerProfile {
|
|
/// Layer info
|
|
info: LayerInfo,
|
|
/// Memory required (bytes)
|
|
memory: usize,
|
|
/// Compute cost (FLOPs)
|
|
compute: usize,
|
|
/// Communication cost if sharded
|
|
comm_cost: usize,
|
|
/// Can this layer be sharded with TP?
|
|
can_tp_shard: bool,
|
|
}
|
|
|
|
// =============================================================================
|
|
// Auto Partitioner
|
|
// =============================================================================
|
|
|
|
/// Automatic model partitioner
|
|
pub struct AutoPartitioner {
|
|
/// Configuration
|
|
config: AutoPartitionConfig,
|
|
/// Per-device memory budgets
|
|
memory_budget: RwLock<HashMap<usize, usize>>,
|
|
/// Cached layer profiles
|
|
layer_profiles: RwLock<Vec<LayerProfile>>,
|
|
}
|
|
|
|
impl AutoPartitioner {
|
|
/// Create a new auto partitioner
|
|
pub fn new(config: AutoPartitionConfig) -> Self {
|
|
let mut memory_budget = HashMap::new();
|
|
let effective_memory =
|
|
(config.device_memory as f32 * (1.0 - config.memory_margin)) as usize;
|
|
for i in 0..config.num_devices {
|
|
memory_budget.insert(i, effective_memory);
|
|
}
|
|
|
|
Self {
|
|
config,
|
|
memory_budget: RwLock::new(memory_budget),
|
|
layer_profiles: RwLock::new(Vec::new()),
|
|
}
|
|
}
|
|
|
|
/// Analyze model and generate optimal sharding plan
|
|
pub fn partition(&self, model_info: &ModelInfo) -> Result<ShardingPlan> {
|
|
// Profile all layers
|
|
let profiles = self.profile_layers(model_info);
|
|
*self.layer_profiles.write() = profiles.clone();
|
|
|
|
// Check if model fits on single device
|
|
if model_info.total_memory <= self.effective_memory_per_device() {
|
|
return self.create_single_device_plan(model_info);
|
|
}
|
|
|
|
// Choose partitioning approach based on strategy
|
|
let plan = match self.config.strategy {
|
|
AutoPartitionStrategy::MemoryOptimal => {
|
|
self.partition_memory_optimal(model_info, &profiles)?
|
|
}
|
|
AutoPartitionStrategy::ComputeBalanced => {
|
|
self.partition_compute_balanced(model_info, &profiles)?
|
|
}
|
|
AutoPartitionStrategy::CommunicationMinimal => {
|
|
self.partition_comm_minimal(model_info, &profiles)?
|
|
}
|
|
AutoPartitionStrategy::Hybrid => self.partition_hybrid(model_info, &profiles)?,
|
|
AutoPartitionStrategy::TensorParallelFirst => {
|
|
self.partition_tp_first(model_info, &profiles)?
|
|
}
|
|
AutoPartitionStrategy::PipelineParallelFirst => {
|
|
self.partition_pp_first(model_info, &profiles)?
|
|
}
|
|
};
|
|
|
|
Ok(plan)
|
|
}
|
|
|
|
/// Get effective memory per device after margin
|
|
fn effective_memory_per_device(&self) -> usize {
|
|
(self.config.device_memory as f32 * (1.0 - self.config.memory_margin)) as usize
|
|
}
|
|
|
|
/// Profile layers for resource requirements
|
|
fn profile_layers(&self, model: &ModelInfo) -> Vec<LayerProfile> {
|
|
model
|
|
.layers
|
|
.iter()
|
|
.map(|layer| {
|
|
let comm_cost = if layer.is_tp_shardable {
|
|
// Communication cost proportional to layer size
|
|
layer.memory_bytes / self.config.num_devices
|
|
} else {
|
|
0
|
|
};
|
|
|
|
LayerProfile {
|
|
info: layer.clone(),
|
|
memory: layer.memory_bytes,
|
|
compute: layer.flops,
|
|
comm_cost,
|
|
can_tp_shard: layer.is_tp_shardable,
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Create plan for single-device execution
|
|
fn create_single_device_plan(&self, model: &ModelInfo) -> Result<ShardingPlan> {
|
|
let mut plan = ShardingPlan::new(&model.name);
|
|
plan.num_pp_stages = 1;
|
|
plan.tp_size = 1;
|
|
plan.dp_size = self.config.num_devices;
|
|
plan.estimated_memory_per_device = model.total_memory;
|
|
plan.compute_imbalance = 1.0;
|
|
|
|
for layer in &model.layers {
|
|
plan.add_partition(LayerPartition {
|
|
layer_name: layer.name.clone(),
|
|
strategy: ShardingStrategy::Replicated,
|
|
device_ids: vec![0],
|
|
tp_size: 1,
|
|
pp_stage: 0,
|
|
});
|
|
}
|
|
|
|
Ok(plan)
|
|
}
|
|
|
|
/// Memory-optimal partitioning (minimize peak memory)
|
|
fn partition_memory_optimal(
|
|
&self,
|
|
model: &ModelInfo,
|
|
profiles: &[LayerProfile],
|
|
) -> Result<ShardingPlan> {
|
|
let effective_memory = self.effective_memory_per_device();
|
|
|
|
// Calculate required TP degree to fit model
|
|
let mut tp_size = 1;
|
|
while model.total_memory / tp_size > effective_memory && tp_size < self.config.num_devices {
|
|
tp_size *= 2;
|
|
}
|
|
|
|
// If still doesn't fit, add PP
|
|
let mut pp_stages = 1;
|
|
let memory_per_tp = model.total_memory / tp_size;
|
|
if memory_per_tp > effective_memory {
|
|
pp_stages = (memory_per_tp + effective_memory - 1) / effective_memory;
|
|
pp_stages = pp_stages.min(self.config.num_devices / tp_size);
|
|
}
|
|
|
|
self.create_tp_pp_plan(model, profiles, tp_size, pp_stages)
|
|
}
|
|
|
|
/// Compute-balanced partitioning
|
|
fn partition_compute_balanced(
|
|
&self,
|
|
model: &ModelInfo,
|
|
profiles: &[LayerProfile],
|
|
) -> Result<ShardingPlan> {
|
|
let total_compute: usize = profiles.iter().map(|p| p.compute).sum();
|
|
let _target_per_device = total_compute / self.config.num_devices;
|
|
|
|
// Greedy assignment to balance compute
|
|
let mut device_compute = vec![0usize; self.config.num_devices];
|
|
let mut assignments = Vec::new();
|
|
|
|
for profile in profiles {
|
|
// Find device with least compute
|
|
let min_device = device_compute
|
|
.iter()
|
|
.enumerate()
|
|
.min_by_key(|(_, c)| *c)
|
|
.map_or(0, |(i, _)| i);
|
|
|
|
assignments.push((profile.info.name.clone(), min_device));
|
|
device_compute[min_device] += profile.compute;
|
|
}
|
|
|
|
// Convert to sharding plan
|
|
let mut plan = ShardingPlan::new(&model.name);
|
|
plan.num_pp_stages = 1;
|
|
plan.tp_size = 1;
|
|
plan.dp_size = self.config.num_devices;
|
|
|
|
for (layer_name, device) in assignments {
|
|
plan.add_partition(LayerPartition {
|
|
layer_name,
|
|
strategy: ShardingStrategy::Replicated,
|
|
device_ids: vec![device],
|
|
tp_size: 1,
|
|
pp_stage: device,
|
|
});
|
|
}
|
|
|
|
// Calculate imbalance
|
|
let max_compute = device_compute.iter().max().copied().unwrap_or(0);
|
|
let min_compute = device_compute.iter().min().copied().unwrap_or(1);
|
|
plan.compute_imbalance = if min_compute > 0 {
|
|
max_compute as f32 / min_compute as f32
|
|
} else {
|
|
f32::MAX
|
|
};
|
|
|
|
Ok(plan)
|
|
}
|
|
|
|
/// Communication-minimal partitioning
|
|
fn partition_comm_minimal(
|
|
&self,
|
|
model: &ModelInfo,
|
|
profiles: &[LayerProfile],
|
|
) -> Result<ShardingPlan> {
|
|
// Prefer PP over TP to minimize communication
|
|
let effective_memory = self.effective_memory_per_device();
|
|
let memory_per_device = model.total_memory / self.config.num_devices;
|
|
|
|
if memory_per_device <= effective_memory {
|
|
// Pure PP is possible
|
|
self.create_pp_only_plan(model, profiles)
|
|
} else {
|
|
// Need some TP
|
|
let tp_size = (model.total_memory + effective_memory - 1) / effective_memory;
|
|
let tp_size = tp_size.min(self.config.num_devices).next_power_of_two();
|
|
self.create_tp_pp_plan(model, profiles, tp_size, 1)
|
|
}
|
|
}
|
|
|
|
/// Hybrid partitioning (balance all factors)
|
|
fn partition_hybrid(
|
|
&self,
|
|
model: &ModelInfo,
|
|
profiles: &[LayerProfile],
|
|
) -> Result<ShardingPlan> {
|
|
let effective_memory = self.effective_memory_per_device();
|
|
|
|
// Try different TP/PP combinations and score them
|
|
let mut best_plan = None;
|
|
let mut best_score = f64::MIN;
|
|
|
|
for tp in [1, 2, 4, 8]
|
|
.iter()
|
|
.filter(|&&t| t <= self.config.num_devices)
|
|
{
|
|
let max_pp = self.config.num_devices / tp;
|
|
for pp in 1..=max_pp {
|
|
let memory_per_device = model.total_memory / (tp * pp);
|
|
if memory_per_device > effective_memory {
|
|
continue;
|
|
}
|
|
|
|
if let Ok(plan) = self.create_tp_pp_plan(model, profiles, *tp, pp) {
|
|
let score = self.score_plan(&plan);
|
|
if score > best_score {
|
|
best_score = score;
|
|
best_plan = Some(plan);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
best_plan.ok_or_else(|| {
|
|
DistributedError::configuration("Could not find valid partitioning scheme")
|
|
})
|
|
}
|
|
|
|
/// Tensor-parallel first partitioning
|
|
fn partition_tp_first(
|
|
&self,
|
|
model: &ModelInfo,
|
|
profiles: &[LayerProfile],
|
|
) -> Result<ShardingPlan> {
|
|
let effective_memory = self.effective_memory_per_device();
|
|
|
|
// Maximize TP before using PP
|
|
let mut tp_size = self.config.num_devices;
|
|
while tp_size > 1 {
|
|
let memory_per_device = model.total_memory / tp_size;
|
|
if memory_per_device <= effective_memory {
|
|
break;
|
|
}
|
|
tp_size /= 2;
|
|
}
|
|
|
|
if tp_size == 0 {
|
|
tp_size = 1;
|
|
}
|
|
|
|
let pp_stages = self.config.num_devices / tp_size;
|
|
self.create_tp_pp_plan(model, profiles, tp_size, pp_stages)
|
|
}
|
|
|
|
/// Pipeline-parallel first partitioning
|
|
fn partition_pp_first(
|
|
&self,
|
|
model: &ModelInfo,
|
|
profiles: &[LayerProfile],
|
|
) -> Result<ShardingPlan> {
|
|
let effective_memory = self.effective_memory_per_device();
|
|
|
|
// Maximize PP before using TP
|
|
let mut pp_stages = self.config.num_devices;
|
|
while pp_stages > 1 {
|
|
let memory_per_device = model.total_memory / pp_stages;
|
|
if memory_per_device <= effective_memory {
|
|
break;
|
|
}
|
|
pp_stages /= 2;
|
|
}
|
|
|
|
if pp_stages == 0 {
|
|
pp_stages = 1;
|
|
}
|
|
|
|
let tp_size = self.config.num_devices / pp_stages;
|
|
self.create_tp_pp_plan(model, profiles, tp_size, pp_stages)
|
|
}
|
|
|
|
/// Create pure pipeline parallel plan
|
|
fn create_pp_only_plan(
|
|
&self,
|
|
model: &ModelInfo,
|
|
profiles: &[LayerProfile],
|
|
) -> Result<ShardingPlan> {
|
|
let total_memory: usize = profiles.iter().map(|p| p.memory).sum();
|
|
let target_per_stage = total_memory / self.config.num_devices;
|
|
|
|
let mut plan = ShardingPlan::new(&model.name);
|
|
plan.tp_size = 1;
|
|
plan.num_pp_stages = self.config.num_devices;
|
|
plan.dp_size = 1;
|
|
|
|
let mut current_stage = 0;
|
|
let mut stage_memory = 0;
|
|
|
|
for profile in profiles {
|
|
if stage_memory + profile.memory > target_per_stage * 2
|
|
&& current_stage < self.config.num_devices - 1
|
|
{
|
|
current_stage += 1;
|
|
stage_memory = 0;
|
|
}
|
|
|
|
plan.add_partition(LayerPartition {
|
|
layer_name: profile.info.name.clone(),
|
|
strategy: ShardingStrategy::LayerWise,
|
|
device_ids: vec![current_stage],
|
|
tp_size: 1,
|
|
pp_stage: current_stage,
|
|
});
|
|
|
|
stage_memory += profile.memory;
|
|
}
|
|
|
|
plan.estimated_memory_per_device = target_per_stage;
|
|
Ok(plan)
|
|
}
|
|
|
|
/// Create combined TP + PP plan
|
|
fn create_tp_pp_plan(
|
|
&self,
|
|
model: &ModelInfo,
|
|
profiles: &[LayerProfile],
|
|
tp_size: usize,
|
|
pp_stages: usize,
|
|
) -> Result<ShardingPlan> {
|
|
let mut plan = ShardingPlan::new(&model.name);
|
|
plan.tp_size = tp_size;
|
|
plan.num_pp_stages = pp_stages;
|
|
plan.dp_size = self.config.num_devices / (tp_size * pp_stages);
|
|
|
|
let total_memory: usize = profiles.iter().map(|p| p.memory).sum();
|
|
let memory_per_stage = total_memory / pp_stages;
|
|
let target_per_stage = memory_per_stage / tp_size;
|
|
|
|
let mut current_stage = 0;
|
|
let mut stage_memory = 0;
|
|
|
|
for profile in profiles {
|
|
if stage_memory + profile.memory > target_per_stage * 2 && current_stage < pp_stages - 1
|
|
{
|
|
current_stage += 1;
|
|
stage_memory = 0;
|
|
}
|
|
|
|
let strategy = if profile.can_tp_shard && tp_size > 1 {
|
|
ShardingStrategy::TensorWise
|
|
} else if pp_stages > 1 {
|
|
ShardingStrategy::LayerWise
|
|
} else {
|
|
ShardingStrategy::Replicated
|
|
};
|
|
|
|
// Device IDs for this layer
|
|
let device_ids: Vec<usize> =
|
|
(0..tp_size).map(|t| current_stage * tp_size + t).collect();
|
|
|
|
plan.add_partition(LayerPartition {
|
|
layer_name: profile.info.name.clone(),
|
|
strategy,
|
|
device_ids,
|
|
tp_size,
|
|
pp_stage: current_stage,
|
|
});
|
|
|
|
stage_memory += profile.memory;
|
|
}
|
|
|
|
plan.estimated_memory_per_device = target_per_stage;
|
|
plan.communication_volume = profiles
|
|
.iter()
|
|
.filter(|p| p.can_tp_shard)
|
|
.map(|p| p.comm_cost)
|
|
.sum();
|
|
|
|
Ok(plan)
|
|
}
|
|
|
|
/// Score a sharding plan (higher is better)
|
|
fn score_plan(&self, plan: &ShardingPlan) -> f64 {
|
|
let effective_memory = self.effective_memory_per_device() as f64;
|
|
|
|
// Memory score (1.0 if fits well, decreasing as it gets tighter)
|
|
let memory_score = if plan.estimated_memory_per_device as f64 <= effective_memory {
|
|
1.0 - (plan.estimated_memory_per_device as f64 / effective_memory) * 0.5
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
// Compute balance score
|
|
let balance_score = if plan.compute_imbalance <= self.config.balance_threshold {
|
|
1.0 / plan.compute_imbalance as f64
|
|
} else {
|
|
0.5 / plan.compute_imbalance as f64
|
|
};
|
|
|
|
// Communication score (lower volume is better)
|
|
let max_comm = plan.estimated_memory_per_device as f64;
|
|
let comm_score = 1.0 - (plan.communication_volume as f64 / max_comm).min(1.0);
|
|
|
|
// Weighted combination
|
|
self.config.memory_weight as f64 * memory_score
|
|
+ self.config.compute_weight as f64 * balance_score
|
|
+ self.config.communication_weight as f64 * comm_score
|
|
}
|
|
}
|
|
|
|
/// Thread-safe shared auto partitioner
|
|
pub type SharedAutoPartitioner = Arc<AutoPartitioner>;
|
|
|
|
/// Create a shared auto partitioner
|
|
pub fn shared_auto_partitioner(config: AutoPartitionConfig) -> SharedAutoPartitioner {
|
|
Arc::new(AutoPartitioner::new(config))
|
|
}
|
|
|
|
// =============================================================================
|
|
// Tests
|
|
// =============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_auto_partition_config_default() {
|
|
let config = AutoPartitionConfig::default();
|
|
assert_eq!(config.strategy, AutoPartitionStrategy::Hybrid);
|
|
assert!((config.memory_margin - 0.1).abs() < 0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_auto_partition_config_memory_optimized() {
|
|
let config = AutoPartitionConfig::memory_optimized(8, 80 * 1024 * 1024 * 1024);
|
|
assert_eq!(config.strategy, AutoPartitionStrategy::MemoryOptimal);
|
|
assert_eq!(config.num_devices, 8);
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_info_creation() {
|
|
let layer = LayerInfo::new("fc1", 1000000, LayerType::Linear);
|
|
assert_eq!(layer.param_count, 1000000);
|
|
assert!(layer.is_tp_shardable);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_info_from_layers() {
|
|
let model = ModelInfo::from_layer_sizes(&[
|
|
("embed", 1_000_000),
|
|
("fc1", 500_000),
|
|
("fc2", 500_000),
|
|
]);
|
|
|
|
assert_eq!(model.layers.len(), 3);
|
|
assert_eq!(model.total_params, 500_000); // Divided by 4 for bytes->params
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharding_plan_creation() {
|
|
let mut plan = ShardingPlan::new("test_model");
|
|
plan.add_partition(LayerPartition {
|
|
layer_name: "layer1".to_string(),
|
|
strategy: ShardingStrategy::Replicated,
|
|
device_ids: vec![0],
|
|
tp_size: 1,
|
|
pp_stage: 0,
|
|
});
|
|
|
|
assert_eq!(plan.layer_partitions.len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_auto_partitioner_single_device() {
|
|
let config = AutoPartitionConfig {
|
|
num_devices: 1,
|
|
device_memory: 80 * 1024 * 1024 * 1024,
|
|
..Default::default()
|
|
};
|
|
let partitioner = AutoPartitioner::new(config);
|
|
|
|
let model = ModelInfo::from_layer_sizes(&[("embed", 1_000_000), ("fc1", 500_000)]);
|
|
|
|
let plan = partitioner.partition(&model).unwrap();
|
|
assert_eq!(plan.tp_size, 1);
|
|
assert_eq!(plan.num_pp_stages, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_auto_partitioner_large_model() {
|
|
let config = AutoPartitionConfig {
|
|
num_devices: 8,
|
|
device_memory: 10 * 1024 * 1024 * 1024, // 10 GB each
|
|
..Default::default()
|
|
};
|
|
let partitioner = AutoPartitioner::new(config);
|
|
|
|
// Create a model that needs distribution (but is tractable)
|
|
// With 10% margin, each device can hold ~9 GB
|
|
// Each layer should be ~1 GB of params = 4 GB memory (4 bytes * 4x for grad+optim)
|
|
// So 10 layers = 40 GB total, needs 8 devices
|
|
let layers: Vec<_> = (0..10)
|
|
.map(|i| (format!("layer{}", i), 1024 * 1024 * 1024usize)) // 1 GB params each
|
|
.collect();
|
|
let layer_refs: Vec<_> = layers.iter().map(|(n, s)| (n.as_str(), *s)).collect();
|
|
let model = ModelInfo::from_layer_sizes(&layer_refs);
|
|
|
|
let plan = partitioner.partition(&model).unwrap();
|
|
// Should use multiple devices
|
|
assert!(plan.tp_size > 1 || plan.num_pp_stages > 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_partition_strategy_variants() {
|
|
assert_eq!(
|
|
AutoPartitionStrategy::default(),
|
|
AutoPartitionStrategy::Hybrid
|
|
);
|
|
assert_ne!(
|
|
AutoPartitionStrategy::MemoryOptimal,
|
|
AutoPartitionStrategy::ComputeBalanced
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_block_index() {
|
|
assert_eq!(extract_block_index("transformer.12.attention"), Some(12));
|
|
assert_eq!(extract_block_index("block.5.mlp"), Some(5));
|
|
assert_eq!(extract_block_index("embedding"), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_type() {
|
|
assert!(matches!(LayerType::Linear, LayerType::Linear));
|
|
assert!(!matches!(LayerType::Linear, LayerType::Attention));
|
|
}
|
|
|
|
#[test]
|
|
fn test_effective_memory() {
|
|
let config = AutoPartitionConfig {
|
|
device_memory: 100,
|
|
memory_margin: 0.1,
|
|
..Default::default()
|
|
};
|
|
let partitioner = AutoPartitioner::new(config);
|
|
assert_eq!(partitioner.effective_memory_per_device(), 90);
|
|
}
|
|
}
|