1128 lines
34 KiB
Rust
1128 lines
34 KiB
Rust
//! Model Sharding Utilities
|
|
//!
|
|
//! This module provides automatic model partitioning across devices:
|
|
//! - Balanced memory and compute distribution
|
|
//! - Layer-wise and tensor-wise sharding strategies
|
|
//! - Automatic shard placement optimization
|
|
//! - Cross-shard communication management
|
|
//! - Support for various model architectures
|
|
//!
|
|
//! Model sharding enables training models larger than single GPU memory
|
|
//! by distributing parameters, gradients, and optimizer states.
|
|
|
|
use crate::error::Result;
|
|
use parking_lot::{Mutex, RwLock};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
|
|
// =============================================================================
|
|
// Configuration
|
|
// =============================================================================
|
|
|
|
/// Sharding strategy for model partitioning
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum ShardingStrategy {
|
|
/// Shard by layers (pipeline-style)
|
|
LayerWise,
|
|
/// Shard tensors across devices (tensor parallelism)
|
|
TensorWise,
|
|
/// Shard by rows (for embeddings, linear layers)
|
|
RowWise,
|
|
/// Shard by columns (for linear layers)
|
|
ColumnWise,
|
|
/// Replicate across all devices
|
|
Replicated,
|
|
/// Custom sharding with explicit placement
|
|
Custom,
|
|
}
|
|
|
|
/// Dimension along which to shard
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum ShardDimension {
|
|
/// No sharding (replicated)
|
|
None,
|
|
/// Shard along first dimension
|
|
Dim0,
|
|
/// Shard along second dimension
|
|
Dim1,
|
|
/// Shard along last dimension
|
|
DimLast,
|
|
/// Shard along specific dimension
|
|
Dim(usize),
|
|
}
|
|
|
|
/// Configuration for model sharding
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ModelShardingConfig {
|
|
/// Number of devices to shard across
|
|
pub num_devices: usize,
|
|
/// Default sharding strategy
|
|
pub default_strategy: ShardingStrategy,
|
|
/// Memory limit per device (bytes, 0 = unlimited)
|
|
pub memory_limit_per_device: usize,
|
|
/// Balance memory across devices
|
|
pub balance_memory: bool,
|
|
/// Balance compute across devices
|
|
pub balance_compute: bool,
|
|
/// Minimum shard size (bytes)
|
|
pub min_shard_size: usize,
|
|
/// Enable automatic resharding on imbalance
|
|
pub auto_reshard: bool,
|
|
/// Threshold for resharding (imbalance ratio)
|
|
pub reshard_threshold: f32,
|
|
}
|
|
|
|
impl Default for ModelShardingConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
num_devices: 1,
|
|
default_strategy: ShardingStrategy::TensorWise,
|
|
memory_limit_per_device: 0,
|
|
balance_memory: true,
|
|
balance_compute: true,
|
|
min_shard_size: 1024,
|
|
auto_reshard: false,
|
|
reshard_threshold: 0.2,
|
|
}
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Parameter Info
|
|
// =============================================================================
|
|
|
|
/// Information about a model parameter
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ParameterInfo {
|
|
/// Parameter name
|
|
pub name: String,
|
|
/// Shape of the parameter
|
|
pub shape: Vec<usize>,
|
|
/// Data type (e.g., "float32", "float16")
|
|
pub dtype: String,
|
|
/// Size in bytes
|
|
pub size_bytes: usize,
|
|
/// Estimated compute cost (FLOPs for forward pass)
|
|
pub compute_cost: u64,
|
|
/// Layer index (for layer-wise sharding)
|
|
pub layer_index: Option<usize>,
|
|
/// Parameter type (weight, bias, etc.)
|
|
pub param_type: ParameterType,
|
|
/// Whether parameter requires gradient
|
|
pub requires_grad: bool,
|
|
}
|
|
|
|
impl ParameterInfo {
|
|
/// Create a new parameter info
|
|
pub fn new(name: String, shape: Vec<usize>, dtype: String) -> Self {
|
|
let elem_size = match dtype.as_str() {
|
|
"float32" | "f32" => 4,
|
|
"float16" | "f16" | "half" => 2,
|
|
"bfloat16" | "bf16" => 2,
|
|
"float64" | "f64" => 8,
|
|
"int32" | "i32" => 4,
|
|
"int64" | "i64" => 8,
|
|
_ => 4,
|
|
};
|
|
let num_elements: usize = shape.iter().product();
|
|
let size_bytes = num_elements * elem_size;
|
|
|
|
Self {
|
|
name,
|
|
shape,
|
|
dtype,
|
|
size_bytes,
|
|
compute_cost: num_elements as u64 * 2, // Simple estimate
|
|
layer_index: None,
|
|
param_type: ParameterType::Weight,
|
|
requires_grad: true,
|
|
}
|
|
}
|
|
|
|
/// Get number of elements
|
|
pub fn num_elements(&self) -> usize {
|
|
self.shape.iter().product()
|
|
}
|
|
|
|
/// Check if parameter can be sharded along a dimension
|
|
pub fn can_shard_dim(&self, dim: usize) -> bool {
|
|
dim < self.shape.len() && self.shape[dim] > 1
|
|
}
|
|
}
|
|
|
|
/// Type of parameter
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum ParameterType {
|
|
/// Weight matrix
|
|
Weight,
|
|
/// Bias vector
|
|
Bias,
|
|
/// Embedding table
|
|
Embedding,
|
|
/// Normalization scale
|
|
Scale,
|
|
/// Normalization shift
|
|
Shift,
|
|
/// Other parameter
|
|
Other,
|
|
}
|
|
|
|
// =============================================================================
|
|
// Shard Specification
|
|
// =============================================================================
|
|
|
|
/// Specification for how a parameter is sharded
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ShardSpec {
|
|
/// Parameter name
|
|
pub param_name: String,
|
|
/// Sharding strategy
|
|
pub strategy: ShardingStrategy,
|
|
/// Dimension to shard along
|
|
pub shard_dim: ShardDimension,
|
|
/// Device placements for each shard
|
|
pub placements: Vec<DevicePlacement>,
|
|
/// Original shape
|
|
pub original_shape: Vec<usize>,
|
|
/// Shard shapes
|
|
pub shard_shapes: Vec<Vec<usize>>,
|
|
}
|
|
|
|
impl ShardSpec {
|
|
/// Create a replicated shard spec
|
|
pub fn replicated(param_name: String, shape: Vec<usize>, devices: Vec<i32>) -> Self {
|
|
let placements = devices.iter().map(|&d| DevicePlacement::new(d)).collect();
|
|
let shard_shapes = vec![shape.clone(); devices.len()];
|
|
|
|
Self {
|
|
param_name,
|
|
strategy: ShardingStrategy::Replicated,
|
|
shard_dim: ShardDimension::None,
|
|
placements,
|
|
original_shape: shape,
|
|
shard_shapes,
|
|
}
|
|
}
|
|
|
|
/// Create a tensor-wise sharded spec
|
|
pub fn tensor_sharded(
|
|
param_name: String,
|
|
shape: Vec<usize>,
|
|
shard_dim: usize,
|
|
devices: Vec<i32>,
|
|
) -> Self {
|
|
let num_shards = devices.len();
|
|
let placements = devices.iter().map(|&d| DevicePlacement::new(d)).collect();
|
|
|
|
// Calculate shard shapes
|
|
let mut shard_shapes = Vec::new();
|
|
let dim_size = shape.get(shard_dim).copied().unwrap_or(1);
|
|
let base_size = dim_size / num_shards;
|
|
let remainder = dim_size % num_shards;
|
|
|
|
for i in 0..num_shards {
|
|
let mut shard_shape = shape.clone();
|
|
let shard_size = base_size + usize::from(i < remainder);
|
|
if shard_dim < shard_shape.len() {
|
|
shard_shape[shard_dim] = shard_size;
|
|
}
|
|
shard_shapes.push(shard_shape);
|
|
}
|
|
|
|
Self {
|
|
param_name,
|
|
strategy: ShardingStrategy::TensorWise,
|
|
shard_dim: ShardDimension::Dim(shard_dim),
|
|
placements,
|
|
original_shape: shape,
|
|
shard_shapes,
|
|
}
|
|
}
|
|
|
|
/// Get number of shards
|
|
pub fn num_shards(&self) -> usize {
|
|
self.placements.len()
|
|
}
|
|
|
|
/// Get total memory for all shards
|
|
pub fn total_memory(&self) -> usize {
|
|
self.shard_shapes
|
|
.iter()
|
|
.map(|s| s.iter().product::<usize>() * 4) // Assume float32
|
|
.sum()
|
|
}
|
|
}
|
|
|
|
/// Device placement for a shard
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct DevicePlacement {
|
|
/// Device ID
|
|
pub device_id: i32,
|
|
/// Memory offset on device
|
|
pub memory_offset: usize,
|
|
/// Priority (lower = preferred)
|
|
pub priority: i32,
|
|
}
|
|
|
|
impl DevicePlacement {
|
|
/// Create a new device placement
|
|
pub fn new(device_id: i32) -> Self {
|
|
Self {
|
|
device_id,
|
|
memory_offset: 0,
|
|
priority: 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Sharding Plan
|
|
// =============================================================================
|
|
|
|
/// A complete sharding plan for a model
|
|
#[derive(Debug, Clone)]
|
|
pub struct ShardingPlan {
|
|
/// Shard specifications per parameter
|
|
pub specs: HashMap<String, ShardSpec>,
|
|
/// Device memory usage
|
|
pub device_memory: HashMap<i32, usize>,
|
|
/// Device compute load
|
|
pub device_compute: HashMap<i32, u64>,
|
|
/// Total model size
|
|
pub total_size: usize,
|
|
/// Total compute cost
|
|
pub total_compute: u64,
|
|
/// Communication cost estimate
|
|
pub comm_cost: u64,
|
|
}
|
|
|
|
impl ShardingPlan {
|
|
/// Create a new empty sharding plan
|
|
pub fn new() -> Self {
|
|
Self {
|
|
specs: HashMap::new(),
|
|
device_memory: HashMap::new(),
|
|
device_compute: HashMap::new(),
|
|
total_size: 0,
|
|
total_compute: 0,
|
|
comm_cost: 0,
|
|
}
|
|
}
|
|
|
|
/// Add a shard spec to the plan
|
|
pub fn add_spec(&mut self, spec: ShardSpec) {
|
|
let param_name = spec.param_name.clone();
|
|
|
|
// Update device memory usage
|
|
for (i, placement) in spec.placements.iter().enumerate() {
|
|
let shard_size = spec.shard_shapes[i].iter().product::<usize>() * 4;
|
|
*self.device_memory.entry(placement.device_id).or_insert(0) += shard_size;
|
|
}
|
|
|
|
self.specs.insert(param_name, spec);
|
|
}
|
|
|
|
/// Get memory balance ratio (0 = perfect, 1 = all on one device)
|
|
pub fn memory_balance(&self) -> f32 {
|
|
if self.device_memory.is_empty() {
|
|
return 0.0;
|
|
}
|
|
|
|
let values: Vec<_> = self.device_memory.values().copied().collect();
|
|
let max = *values.iter().max().unwrap_or(&0) as f32;
|
|
let min = *values.iter().min().unwrap_or(&0) as f32;
|
|
let avg = values.iter().sum::<usize>() as f32 / values.len() as f32;
|
|
|
|
if avg == 0.0 { 0.0 } else { (max - min) / avg }
|
|
}
|
|
|
|
/// Get compute balance ratio
|
|
pub fn compute_balance(&self) -> f32 {
|
|
if self.device_compute.is_empty() {
|
|
return 0.0;
|
|
}
|
|
|
|
let values: Vec<_> = self.device_compute.values().copied().collect();
|
|
let max = *values.iter().max().unwrap_or(&0) as f32;
|
|
let min = *values.iter().min().unwrap_or(&0) as f32;
|
|
let avg = values.iter().sum::<u64>() as f32 / values.len() as f32;
|
|
|
|
if avg == 0.0 { 0.0 } else { (max - min) / avg }
|
|
}
|
|
|
|
/// Get number of parameters
|
|
pub fn num_parameters(&self) -> usize {
|
|
self.specs.len()
|
|
}
|
|
|
|
/// Get spec for a parameter
|
|
pub fn get_spec(&self, param_name: &str) -> Option<&ShardSpec> {
|
|
self.specs.get(param_name)
|
|
}
|
|
}
|
|
|
|
impl Default for ShardingPlan {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Model Sharder
|
|
// =============================================================================
|
|
|
|
/// Automatic model sharding utility
|
|
pub struct ModelSharder {
|
|
/// Configuration
|
|
config: ModelShardingConfig,
|
|
/// Available devices
|
|
devices: Vec<i32>,
|
|
/// Device memory capacities
|
|
device_capacities: HashMap<i32, usize>,
|
|
/// Current sharding plan
|
|
plan: RwLock<ShardingPlan>,
|
|
/// Parameter registry
|
|
parameters: RwLock<HashMap<String, ParameterInfo>>,
|
|
/// Statistics
|
|
stats: RwLock<ShardingStats>,
|
|
}
|
|
|
|
impl ModelSharder {
|
|
/// Create a new model sharder
|
|
pub fn new(config: ModelShardingConfig, devices: Vec<i32>) -> Self {
|
|
let device_capacities = devices
|
|
.iter()
|
|
.map(|&d| (d, config.memory_limit_per_device))
|
|
.collect();
|
|
|
|
Self {
|
|
config,
|
|
devices,
|
|
device_capacities,
|
|
plan: RwLock::new(ShardingPlan::new()),
|
|
parameters: RwLock::new(HashMap::new()),
|
|
stats: RwLock::new(ShardingStats::default()),
|
|
}
|
|
}
|
|
|
|
/// Register a parameter for sharding
|
|
pub fn register_parameter(&self, info: ParameterInfo) {
|
|
let mut params = self.parameters.write();
|
|
params.insert(info.name.clone(), info);
|
|
}
|
|
|
|
/// Register multiple parameters
|
|
pub fn register_parameters(&self, infos: Vec<ParameterInfo>) {
|
|
let mut params = self.parameters.write();
|
|
for info in infos {
|
|
params.insert(info.name.clone(), info);
|
|
}
|
|
}
|
|
|
|
/// Generate a sharding plan for all registered parameters
|
|
pub fn generate_plan(&self) -> Result<ShardingPlan> {
|
|
let params = self.parameters.read();
|
|
let mut plan = ShardingPlan::new();
|
|
|
|
// Sort parameters by size for better load balancing
|
|
let mut sorted_params: Vec<_> = params.values().collect();
|
|
sorted_params.sort_by(|a, b| b.size_bytes.cmp(&a.size_bytes));
|
|
|
|
for param in sorted_params {
|
|
let spec = self.create_shard_spec(param)?;
|
|
plan.add_spec(spec);
|
|
}
|
|
|
|
// Calculate totals
|
|
plan.total_size = params.values().map(|p| p.size_bytes).sum();
|
|
plan.total_compute = params.values().map(|p| p.compute_cost).sum();
|
|
|
|
// Estimate communication cost
|
|
plan.comm_cost = self.estimate_comm_cost(&plan);
|
|
|
|
// Update stats
|
|
{
|
|
let mut stats = self.stats.write();
|
|
stats.plans_generated += 1;
|
|
}
|
|
|
|
*self.plan.write() = plan.clone();
|
|
Ok(plan)
|
|
}
|
|
|
|
/// Create shard spec for a parameter
|
|
fn create_shard_spec(&self, param: &ParameterInfo) -> Result<ShardSpec> {
|
|
match self.config.default_strategy {
|
|
ShardingStrategy::Replicated => Ok(ShardSpec::replicated(
|
|
param.name.clone(),
|
|
param.shape.clone(),
|
|
self.devices.clone(),
|
|
)),
|
|
ShardingStrategy::TensorWise | ShardingStrategy::RowWise => {
|
|
// Shard along first dimension
|
|
let shard_dim = 0;
|
|
if param.can_shard_dim(shard_dim) && param.shape[shard_dim] >= self.devices.len() {
|
|
Ok(ShardSpec::tensor_sharded(
|
|
param.name.clone(),
|
|
param.shape.clone(),
|
|
shard_dim,
|
|
self.devices.clone(),
|
|
))
|
|
} else {
|
|
// Fall back to replication for small tensors
|
|
Ok(ShardSpec::replicated(
|
|
param.name.clone(),
|
|
param.shape.clone(),
|
|
self.devices.clone(),
|
|
))
|
|
}
|
|
}
|
|
ShardingStrategy::ColumnWise => {
|
|
// Shard along last dimension
|
|
let shard_dim = param.shape.len().saturating_sub(1);
|
|
if param.can_shard_dim(shard_dim) && param.shape[shard_dim] >= self.devices.len() {
|
|
Ok(ShardSpec::tensor_sharded(
|
|
param.name.clone(),
|
|
param.shape.clone(),
|
|
shard_dim,
|
|
self.devices.clone(),
|
|
))
|
|
} else {
|
|
Ok(ShardSpec::replicated(
|
|
param.name.clone(),
|
|
param.shape.clone(),
|
|
self.devices.clone(),
|
|
))
|
|
}
|
|
}
|
|
ShardingStrategy::LayerWise => {
|
|
// Assign entire layer to one device based on layer index
|
|
let device_idx = param.layer_index.unwrap_or(0) % self.devices.len();
|
|
Ok(ShardSpec::replicated(
|
|
param.name.clone(),
|
|
param.shape.clone(),
|
|
vec![self.devices[device_idx]],
|
|
))
|
|
}
|
|
ShardingStrategy::Custom => {
|
|
// Use default tensor-wise for custom
|
|
self.create_shard_spec(&ParameterInfo {
|
|
param_type: param.param_type,
|
|
..param.clone()
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Estimate communication cost for a plan
|
|
fn estimate_comm_cost(&self, plan: &ShardingPlan) -> u64 {
|
|
let mut cost = 0u64;
|
|
|
|
for spec in plan.specs.values() {
|
|
match spec.strategy {
|
|
ShardingStrategy::TensorWise
|
|
| ShardingStrategy::RowWise
|
|
| ShardingStrategy::ColumnWise => {
|
|
// AllGather or ReduceScatter needed
|
|
let shard_size: usize = spec
|
|
.shard_shapes
|
|
.iter()
|
|
.map(|s| s.iter().product::<usize>())
|
|
.sum();
|
|
cost += shard_size as u64 * 2; // Forward + backward
|
|
}
|
|
ShardingStrategy::Replicated => {
|
|
// AllReduce needed for gradients
|
|
let size: usize = spec.original_shape.iter().product();
|
|
cost += size as u64;
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
cost
|
|
}
|
|
|
|
/// Get current sharding plan
|
|
pub fn current_plan(&self) -> ShardingPlan {
|
|
self.plan.read().clone()
|
|
}
|
|
|
|
/// Get parameter info
|
|
pub fn get_parameter(&self, name: &str) -> Option<ParameterInfo> {
|
|
self.parameters.read().get(name).cloned()
|
|
}
|
|
|
|
/// Get all parameter names
|
|
pub fn parameter_names(&self) -> Vec<String> {
|
|
self.parameters.read().keys().cloned().collect()
|
|
}
|
|
|
|
/// Get statistics
|
|
pub fn stats(&self) -> ShardingStats {
|
|
self.stats.read().clone()
|
|
}
|
|
|
|
/// Get number of devices
|
|
pub fn num_devices(&self) -> usize {
|
|
self.devices.len()
|
|
}
|
|
|
|
/// Get device list
|
|
pub fn devices(&self) -> Vec<i32> {
|
|
self.devices.clone()
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Sharding Statistics
|
|
// =============================================================================
|
|
|
|
/// Statistics for model sharding
|
|
#[derive(Debug, Default, Clone)]
|
|
pub struct ShardingStats {
|
|
/// Number of plans generated
|
|
pub plans_generated: usize,
|
|
/// Number of resharding operations
|
|
pub reshard_count: usize,
|
|
/// Total parameters sharded
|
|
pub parameters_sharded: usize,
|
|
/// Total bytes sharded
|
|
pub bytes_sharded: usize,
|
|
}
|
|
|
|
// =============================================================================
|
|
// Shard Manager
|
|
// =============================================================================
|
|
|
|
/// Manages sharded parameter storage and communication
|
|
pub struct ShardManager {
|
|
/// Sharder reference
|
|
sharder: Arc<ModelSharder>,
|
|
/// Shard data storage (device_id -> param_name -> data)
|
|
shard_data: RwLock<HashMap<i32, HashMap<String, Vec<u8>>>>,
|
|
/// Communication pending
|
|
pending_comms: Mutex<Vec<PendingComm>>,
|
|
/// Next comm ID
|
|
next_comm_id: AtomicU64,
|
|
}
|
|
|
|
/// A pending communication operation
|
|
#[derive(Debug, Clone)]
|
|
pub struct PendingComm {
|
|
/// Communication ID
|
|
pub id: u64,
|
|
/// Type of communication
|
|
pub comm_type: CommType,
|
|
/// Source devices
|
|
pub src_devices: Vec<i32>,
|
|
/// Destination devices
|
|
pub dst_devices: Vec<i32>,
|
|
/// Parameter name
|
|
pub param_name: String,
|
|
/// Size in bytes
|
|
pub size_bytes: usize,
|
|
}
|
|
|
|
/// Type of communication
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum CommType {
|
|
/// AllGather shards
|
|
AllGather,
|
|
/// ReduceScatter gradients
|
|
ReduceScatter,
|
|
/// AllReduce gradients
|
|
AllReduce,
|
|
/// Point-to-point send
|
|
Send,
|
|
/// Point-to-point receive
|
|
Recv,
|
|
}
|
|
|
|
impl ShardManager {
|
|
/// Create a new shard manager
|
|
pub fn new(sharder: Arc<ModelSharder>) -> Self {
|
|
Self {
|
|
sharder,
|
|
shard_data: RwLock::new(HashMap::new()),
|
|
pending_comms: Mutex::new(Vec::new()),
|
|
next_comm_id: AtomicU64::new(1),
|
|
}
|
|
}
|
|
|
|
/// Store shard data
|
|
pub fn store_shard(&self, device_id: i32, param_name: String, data: Vec<u8>) {
|
|
let mut storage = self.shard_data.write();
|
|
storage
|
|
.entry(device_id)
|
|
.or_default()
|
|
.insert(param_name, data);
|
|
}
|
|
|
|
/// Get shard data
|
|
pub fn get_shard(&self, device_id: i32, param_name: &str) -> Option<Vec<u8>> {
|
|
let storage = self.shard_data.read();
|
|
storage.get(&device_id)?.get(param_name).cloned()
|
|
}
|
|
|
|
/// Schedule an AllGather operation
|
|
pub fn schedule_all_gather(&self, param_name: String) -> u64 {
|
|
let plan = self.sharder.current_plan();
|
|
let spec = plan.get_spec(¶m_name);
|
|
|
|
let devices: Vec<i32> = spec
|
|
.map(|s| s.placements.iter().map(|p| p.device_id).collect())
|
|
.unwrap_or_default();
|
|
|
|
let size = spec.map_or(0, ShardSpec::total_memory);
|
|
|
|
let id = self.next_comm_id.fetch_add(1, Ordering::SeqCst);
|
|
let comm = PendingComm {
|
|
id,
|
|
comm_type: CommType::AllGather,
|
|
src_devices: devices.clone(),
|
|
dst_devices: devices,
|
|
param_name,
|
|
size_bytes: size,
|
|
};
|
|
|
|
self.pending_comms.lock().push(comm);
|
|
id
|
|
}
|
|
|
|
/// Schedule a ReduceScatter operation
|
|
pub fn schedule_reduce_scatter(&self, param_name: String) -> u64 {
|
|
let plan = self.sharder.current_plan();
|
|
let spec = plan.get_spec(¶m_name);
|
|
|
|
let devices: Vec<i32> = spec
|
|
.map(|s| s.placements.iter().map(|p| p.device_id).collect())
|
|
.unwrap_or_default();
|
|
|
|
let size = spec.map_or(0, ShardSpec::total_memory);
|
|
|
|
let id = self.next_comm_id.fetch_add(1, Ordering::SeqCst);
|
|
let comm = PendingComm {
|
|
id,
|
|
comm_type: CommType::ReduceScatter,
|
|
src_devices: devices.clone(),
|
|
dst_devices: devices,
|
|
param_name,
|
|
size_bytes: size,
|
|
};
|
|
|
|
self.pending_comms.lock().push(comm);
|
|
id
|
|
}
|
|
|
|
/// Get pending communication count
|
|
pub fn pending_count(&self) -> usize {
|
|
self.pending_comms.lock().len()
|
|
}
|
|
|
|
/// Clear pending communications
|
|
pub fn clear_pending(&self) {
|
|
self.pending_comms.lock().clear();
|
|
}
|
|
|
|
/// Get sharder reference
|
|
pub fn sharder(&self) -> Arc<ModelSharder> {
|
|
self.sharder.clone()
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Thread-Safe Wrappers
|
|
// =============================================================================
|
|
|
|
/// Thread-safe shared model sharder
|
|
pub type SharedModelSharder = Arc<ModelSharder>;
|
|
|
|
/// Create a shared model sharder
|
|
pub fn shared_model_sharder(config: ModelShardingConfig, devices: Vec<i32>) -> SharedModelSharder {
|
|
Arc::new(ModelSharder::new(config, devices))
|
|
}
|
|
|
|
/// Thread-safe shared shard manager
|
|
pub type SharedShardManager = Arc<ShardManager>;
|
|
|
|
/// Create a shared shard manager
|
|
pub fn shared_shard_manager(sharder: SharedModelSharder) -> SharedShardManager {
|
|
Arc::new(ShardManager::new(sharder))
|
|
}
|
|
|
|
// =============================================================================
|
|
// Tests
|
|
// =============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_model_sharding_config_default() {
|
|
let config = ModelShardingConfig::default();
|
|
assert_eq!(config.num_devices, 1);
|
|
assert_eq!(config.default_strategy, ShardingStrategy::TensorWise);
|
|
assert!(config.balance_memory);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parameter_info_creation() {
|
|
let param = ParameterInfo::new(
|
|
"layer1.weight".to_string(),
|
|
vec![1024, 512],
|
|
"float32".to_string(),
|
|
);
|
|
|
|
assert_eq!(param.num_elements(), 1024 * 512);
|
|
assert_eq!(param.size_bytes, 1024 * 512 * 4);
|
|
assert!(param.can_shard_dim(0));
|
|
assert!(param.can_shard_dim(1));
|
|
assert!(!param.can_shard_dim(2));
|
|
}
|
|
|
|
#[test]
|
|
fn test_parameter_info_dtypes() {
|
|
let f32_param = ParameterInfo::new("p".to_string(), vec![100], "float32".to_string());
|
|
assert_eq!(f32_param.size_bytes, 400);
|
|
|
|
let f16_param = ParameterInfo::new("p".to_string(), vec![100], "float16".to_string());
|
|
assert_eq!(f16_param.size_bytes, 200);
|
|
|
|
let f64_param = ParameterInfo::new("p".to_string(), vec![100], "float64".to_string());
|
|
assert_eq!(f64_param.size_bytes, 800);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shard_spec_replicated() {
|
|
let spec = ShardSpec::replicated("weight".to_string(), vec![1024, 512], vec![0, 1, 2, 3]);
|
|
|
|
assert_eq!(spec.num_shards(), 4);
|
|
assert_eq!(spec.strategy, ShardingStrategy::Replicated);
|
|
assert_eq!(spec.shard_shapes.len(), 4);
|
|
|
|
// All shards should have same shape
|
|
for shape in &spec.shard_shapes {
|
|
assert_eq!(shape, &vec![1024, 512]);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_shard_spec_tensor_sharded() {
|
|
let spec = ShardSpec::tensor_sharded(
|
|
"weight".to_string(),
|
|
vec![1024, 512],
|
|
0, // Shard along dim 0
|
|
vec![0, 1, 2, 3],
|
|
);
|
|
|
|
assert_eq!(spec.num_shards(), 4);
|
|
assert_eq!(spec.strategy, ShardingStrategy::TensorWise);
|
|
|
|
// Each shard should have 1/4 of dim 0
|
|
for shape in &spec.shard_shapes {
|
|
assert_eq!(shape[0], 256);
|
|
assert_eq!(shape[1], 512);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_shard_spec_uneven_split() {
|
|
let spec =
|
|
ShardSpec::tensor_sharded("weight".to_string(), vec![10, 512], 0, vec![0, 1, 2, 3]);
|
|
|
|
// 10 / 4 = 2 with remainder 2
|
|
// First 2 shards get 3, last 2 get 2
|
|
let sizes: Vec<_> = spec.shard_shapes.iter().map(|s| s[0]).collect();
|
|
assert_eq!(sizes, vec![3, 3, 2, 2]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharding_plan() {
|
|
let mut plan = ShardingPlan::new();
|
|
|
|
let spec1 =
|
|
ShardSpec::tensor_sharded("layer1.weight".to_string(), vec![1024, 512], 0, vec![0, 1]);
|
|
let spec2 = ShardSpec::replicated("layer1.bias".to_string(), vec![512], vec![0, 1]);
|
|
|
|
plan.add_spec(spec1);
|
|
plan.add_spec(spec2);
|
|
|
|
assert_eq!(plan.num_parameters(), 2);
|
|
assert!(plan.get_spec("layer1.weight").is_some());
|
|
assert!(plan.get_spec("layer1.bias").is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharding_plan_memory_balance() {
|
|
let mut plan = ShardingPlan::new();
|
|
|
|
// Add balanced shards
|
|
let spec = ShardSpec::tensor_sharded("weight".to_string(), vec![1024, 1024], 0, vec![0, 1]);
|
|
plan.add_spec(spec);
|
|
|
|
let balance = plan.memory_balance();
|
|
assert!(balance < 0.1); // Should be well balanced
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_sharder_creation() {
|
|
let config = ModelShardingConfig {
|
|
num_devices: 4,
|
|
..Default::default()
|
|
};
|
|
let sharder = ModelSharder::new(config, vec![0, 1, 2, 3]);
|
|
|
|
assert_eq!(sharder.num_devices(), 4);
|
|
assert_eq!(sharder.devices(), vec![0, 1, 2, 3]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_register_parameter() {
|
|
let config = ModelShardingConfig::default();
|
|
let sharder = ModelSharder::new(config, vec![0, 1]);
|
|
|
|
let param = ParameterInfo::new(
|
|
"layer1.weight".to_string(),
|
|
vec![1024, 512],
|
|
"float32".to_string(),
|
|
);
|
|
sharder.register_parameter(param);
|
|
|
|
assert!(sharder.get_parameter("layer1.weight").is_some());
|
|
assert_eq!(sharder.parameter_names().len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_plan_tensor_wise() {
|
|
let config = ModelShardingConfig {
|
|
num_devices: 2,
|
|
default_strategy: ShardingStrategy::TensorWise,
|
|
..Default::default()
|
|
};
|
|
let sharder = ModelSharder::new(config, vec![0, 1]);
|
|
|
|
sharder.register_parameter(ParameterInfo::new(
|
|
"weight".to_string(),
|
|
vec![1024, 512],
|
|
"float32".to_string(),
|
|
));
|
|
|
|
let plan = sharder.generate_plan().unwrap();
|
|
|
|
assert_eq!(plan.num_parameters(), 1);
|
|
let spec = plan.get_spec("weight").unwrap();
|
|
assert_eq!(spec.strategy, ShardingStrategy::TensorWise);
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_plan_replicated() {
|
|
let config = ModelShardingConfig {
|
|
num_devices: 2,
|
|
default_strategy: ShardingStrategy::Replicated,
|
|
..Default::default()
|
|
};
|
|
let sharder = ModelSharder::new(config, vec![0, 1]);
|
|
|
|
sharder.register_parameter(ParameterInfo::new(
|
|
"weight".to_string(),
|
|
vec![1024, 512],
|
|
"float32".to_string(),
|
|
));
|
|
|
|
let plan = sharder.generate_plan().unwrap();
|
|
let spec = plan.get_spec("weight").unwrap();
|
|
assert_eq!(spec.strategy, ShardingStrategy::Replicated);
|
|
}
|
|
|
|
#[test]
|
|
fn test_small_tensor_fallback_to_replicated() {
|
|
let config = ModelShardingConfig {
|
|
num_devices: 4,
|
|
default_strategy: ShardingStrategy::TensorWise,
|
|
..Default::default()
|
|
};
|
|
let sharder = ModelSharder::new(config, vec![0, 1, 2, 3]);
|
|
|
|
// Tensor too small to shard across 4 devices
|
|
sharder.register_parameter(ParameterInfo::new(
|
|
"small".to_string(),
|
|
vec![2, 512],
|
|
"float32".to_string(),
|
|
));
|
|
|
|
let plan = sharder.generate_plan().unwrap();
|
|
let spec = plan.get_spec("small").unwrap();
|
|
// Should fall back to replicated
|
|
assert_eq!(spec.strategy, ShardingStrategy::Replicated);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shard_manager_creation() {
|
|
let config = ModelShardingConfig::default();
|
|
let sharder = Arc::new(ModelSharder::new(config, vec![0, 1]));
|
|
let manager = ShardManager::new(sharder);
|
|
|
|
assert_eq!(manager.pending_count(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shard_manager_store_get() {
|
|
let config = ModelShardingConfig::default();
|
|
let sharder = Arc::new(ModelSharder::new(config, vec![0, 1]));
|
|
let manager = ShardManager::new(sharder);
|
|
|
|
let data = vec![1u8, 2, 3, 4];
|
|
manager.store_shard(0, "weight".to_string(), data.clone());
|
|
|
|
let retrieved = manager.get_shard(0, "weight").unwrap();
|
|
assert_eq!(retrieved, data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_schedule_all_gather() {
|
|
let config = ModelShardingConfig::default();
|
|
let sharder = Arc::new(ModelSharder::new(config, vec![0, 1]));
|
|
|
|
sharder.register_parameter(ParameterInfo::new(
|
|
"weight".to_string(),
|
|
vec![1024, 512],
|
|
"float32".to_string(),
|
|
));
|
|
sharder.generate_plan().unwrap();
|
|
|
|
let manager = ShardManager::new(sharder);
|
|
let id = manager.schedule_all_gather("weight".to_string());
|
|
|
|
assert!(id > 0);
|
|
assert_eq!(manager.pending_count(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_schedule_reduce_scatter() {
|
|
let config = ModelShardingConfig::default();
|
|
let sharder = Arc::new(ModelSharder::new(config, vec![0, 1]));
|
|
|
|
sharder.register_parameter(ParameterInfo::new(
|
|
"weight".to_string(),
|
|
vec![1024, 512],
|
|
"float32".to_string(),
|
|
));
|
|
sharder.generate_plan().unwrap();
|
|
|
|
let manager = ShardManager::new(sharder);
|
|
let id = manager.schedule_reduce_scatter("weight".to_string());
|
|
|
|
assert!(id > 0);
|
|
assert_eq!(manager.pending_count(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_clear_pending() {
|
|
let config = ModelShardingConfig::default();
|
|
let sharder = Arc::new(ModelSharder::new(config, vec![0, 1]));
|
|
let manager = ShardManager::new(sharder.clone());
|
|
|
|
sharder.register_parameter(ParameterInfo::new(
|
|
"weight".to_string(),
|
|
vec![1024, 512],
|
|
"float32".to_string(),
|
|
));
|
|
sharder.generate_plan().unwrap();
|
|
|
|
manager.schedule_all_gather("weight".to_string());
|
|
manager.schedule_reduce_scatter("weight".to_string());
|
|
assert_eq!(manager.pending_count(), 2);
|
|
|
|
manager.clear_pending();
|
|
assert_eq!(manager.pending_count(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_model_sharder() {
|
|
let config = ModelShardingConfig::default();
|
|
let sharder = shared_model_sharder(config, vec![0, 1, 2, 3]);
|
|
|
|
assert_eq!(sharder.num_devices(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shared_shard_manager() {
|
|
let config = ModelShardingConfig::default();
|
|
let sharder = shared_model_sharder(config, vec![0, 1]);
|
|
let manager = shared_shard_manager(sharder);
|
|
|
|
assert_eq!(manager.pending_count(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharding_stats() {
|
|
let config = ModelShardingConfig::default();
|
|
let sharder = ModelSharder::new(config, vec![0, 1]);
|
|
|
|
sharder.register_parameter(ParameterInfo::new(
|
|
"weight".to_string(),
|
|
vec![1024, 512],
|
|
"float32".to_string(),
|
|
));
|
|
sharder.generate_plan().unwrap();
|
|
|
|
let stats = sharder.stats();
|
|
assert_eq!(stats.plans_generated, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_wise_sharding() {
|
|
let config = ModelShardingConfig {
|
|
num_devices: 4,
|
|
default_strategy: ShardingStrategy::LayerWise,
|
|
..Default::default()
|
|
};
|
|
let sharder = ModelSharder::new(config, vec![0, 1, 2, 3]);
|
|
|
|
let mut param = ParameterInfo::new(
|
|
"layer2.weight".to_string(),
|
|
vec![1024, 512],
|
|
"float32".to_string(),
|
|
);
|
|
param.layer_index = Some(2);
|
|
sharder.register_parameter(param);
|
|
|
|
let plan = sharder.generate_plan().unwrap();
|
|
let spec = plan.get_spec("layer2.weight").unwrap();
|
|
|
|
// Layer 2 should be on device 2 (2 % 4 = 2)
|
|
assert_eq!(spec.placements.len(), 1);
|
|
assert_eq!(spec.placements[0].device_id, 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_column_wise_sharding() {
|
|
let config = ModelShardingConfig {
|
|
num_devices: 2,
|
|
default_strategy: ShardingStrategy::ColumnWise,
|
|
..Default::default()
|
|
};
|
|
let sharder = ModelSharder::new(config, vec![0, 1]);
|
|
|
|
sharder.register_parameter(ParameterInfo::new(
|
|
"weight".to_string(),
|
|
vec![1024, 512],
|
|
"float32".to_string(),
|
|
));
|
|
|
|
let plan = sharder.generate_plan().unwrap();
|
|
let spec = plan.get_spec("weight").unwrap();
|
|
|
|
// Should shard along last dim (1)
|
|
assert_eq!(spec.shard_dim, ShardDimension::Dim(1));
|
|
for shape in &spec.shard_shapes {
|
|
assert_eq!(shape[0], 1024);
|
|
assert_eq!(shape[1], 256); // 512 / 2
|
|
}
|
|
}
|
|
}
|