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]>
1105 lines
39 KiB
Rust
1105 lines
39 KiB
Rust
//! Parallelism strategies for distributed training
|
|
//!
|
|
//! This module implements various parallelism approaches including:
|
|
//! - Data Parallel (DP)
|
|
//! - Tensor Parallel (TP)
|
|
//! - Pipeline Parallel (PP)
|
|
//! - Fully Sharded Data Parallel (FSDP/ZeRO)
|
|
|
|
use crate::error::{DistributedError, Result};
|
|
use crate::group::ProcessGroup;
|
|
use rtx_tensor::Tensor;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
|
|
/// Data Parallel implementation
|
|
pub struct DataParallel {
|
|
/// Process group for data parallel communication
|
|
pub process_group: ProcessGroup,
|
|
/// Gradient accumulation steps
|
|
pub accumulation_steps: usize,
|
|
/// Current accumulation step
|
|
current_step: usize,
|
|
/// Whether gradients are currently being accumulated
|
|
accumulating: bool,
|
|
}
|
|
|
|
impl DataParallel {
|
|
/// Create new data parallel instance
|
|
pub fn new(process_group: ProcessGroup, accumulation_steps: usize) -> Self {
|
|
Self {
|
|
process_group,
|
|
accumulation_steps: accumulation_steps.max(1),
|
|
current_step: 0,
|
|
accumulating: false,
|
|
}
|
|
}
|
|
|
|
/// Start gradient accumulation
|
|
pub fn start_accumulation(&mut self) {
|
|
self.accumulating = true;
|
|
self.current_step = 0;
|
|
}
|
|
|
|
/// Accumulate gradients (called after backward pass)
|
|
pub async fn accumulate_gradients(&mut self, gradients: &mut [Tensor]) -> Result<bool> {
|
|
if !self.accumulating {
|
|
return Err(DistributedError::parallelism(
|
|
"DataParallel",
|
|
"not accumulating - call start_accumulation() first",
|
|
));
|
|
}
|
|
|
|
self.current_step += 1;
|
|
|
|
// Check if we should synchronize
|
|
let should_sync = self.current_step >= self.accumulation_steps;
|
|
|
|
if should_sync {
|
|
self.sync_gradients(gradients).await?;
|
|
self.accumulating = false;
|
|
self.current_step = 0;
|
|
}
|
|
|
|
Ok(should_sync)
|
|
}
|
|
|
|
/// Synchronize gradients across all processes
|
|
async fn sync_gradients(&self, gradients: &mut [Tensor]) -> Result<()> {
|
|
use crate::comm::ReduceOp;
|
|
|
|
for gradient in gradients.iter_mut() {
|
|
// AllReduce to sum gradients across all processes
|
|
self.process_group
|
|
.all_reduce(gradient, ReduceOp::Sum)
|
|
.await?;
|
|
|
|
// Average the gradients
|
|
let world_size = self.process_group.world_size() as f32;
|
|
*gradient = gradient.div_scalar(world_size)?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// Fully Sharded Data Parallel (FSDP) implementation
|
|
pub struct Fsdp {
|
|
/// Process group for FSDP communication
|
|
pub process_group: ProcessGroup,
|
|
/// FSDP configuration
|
|
pub config: FsdpConfig,
|
|
/// Sharded parameters
|
|
sharded_params: HashMap<String, ShardedParameter>,
|
|
/// Memory usage statistics
|
|
memory_stats: FsdpMemoryStats,
|
|
}
|
|
|
|
/// FSDP configuration
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct FsdpConfig {
|
|
/// Sharding strategy
|
|
pub sharding_strategy: ShardingStrategy,
|
|
/// Minimum parameter size to shard (in elements)
|
|
pub min_param_size: usize,
|
|
/// CPU offloading enabled
|
|
pub cpu_offload: bool,
|
|
/// Mixed precision configuration
|
|
pub mixed_precision: bool,
|
|
/// Flatten parameters for sharding
|
|
pub flatten_parameters: bool,
|
|
}
|
|
|
|
/// Sharding strategies for FSDP
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum ShardingStrategy {
|
|
/// Full sharding (ZeRO-3)
|
|
FullShard,
|
|
/// Shard gradients only (ZeRO-2)
|
|
ShardGradOp,
|
|
/// No sharding, replicate parameters
|
|
NoShard,
|
|
}
|
|
|
|
/// Sharded parameter representation
|
|
#[derive(Debug)]
|
|
pub struct ShardedParameter {
|
|
/// Parameter name/identifier
|
|
pub name: String,
|
|
/// Local shard of the parameter
|
|
pub local_shard: Tensor,
|
|
/// Full parameter shape
|
|
pub full_shape: Vec<usize>,
|
|
/// Shard metadata
|
|
pub shard_metadata: ShardMetadata,
|
|
}
|
|
|
|
/// Metadata for parameter sharding
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ShardMetadata {
|
|
/// Rank that owns this shard
|
|
pub owner_rank: i32,
|
|
/// Start index in flattened parameter
|
|
pub start_idx: usize,
|
|
/// End index in flattened parameter
|
|
pub end_idx: usize,
|
|
/// Original parameter offset
|
|
pub param_offset: usize,
|
|
}
|
|
|
|
/// Memory usage statistics for FSDP
|
|
#[derive(Debug, Default)]
|
|
pub struct FsdpMemoryStats {
|
|
/// Total parameter memory (MB)
|
|
pub total_param_memory_mb: f32,
|
|
/// Local shard memory (MB)
|
|
pub local_shard_memory_mb: f32,
|
|
/// Peak memory during all-gather (MB)
|
|
pub peak_allgather_memory_mb: f32,
|
|
/// Memory saved compared to non-sharded (MB)
|
|
pub memory_saved_mb: f32,
|
|
}
|
|
|
|
impl Default for FsdpConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
sharding_strategy: ShardingStrategy::FullShard,
|
|
min_param_size: 1000, // Only shard parameters with >1K elements
|
|
cpu_offload: false,
|
|
mixed_precision: true,
|
|
flatten_parameters: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Fsdp {
|
|
/// Create new FSDP instance with default config
|
|
pub fn new(process_group: ProcessGroup) -> Result<Self> {
|
|
let config = FsdpConfig::default();
|
|
Ok(Self {
|
|
process_group,
|
|
config,
|
|
sharded_params: HashMap::new(),
|
|
memory_stats: FsdpMemoryStats::default(),
|
|
})
|
|
}
|
|
|
|
/// Create new FSDP instance with custom config
|
|
pub fn new_with_config(process_group: ProcessGroup, config: FsdpConfig) -> Self {
|
|
Self {
|
|
process_group,
|
|
config,
|
|
sharded_params: HashMap::new(),
|
|
memory_stats: FsdpMemoryStats::default(),
|
|
}
|
|
}
|
|
|
|
/// Shard model parameters across the process group
|
|
pub fn shard_parameters(&mut self, parameters: &Tensor) -> Result<Tensor> {
|
|
let world_size = self.process_group.world_size();
|
|
let rank = self.process_group.rank();
|
|
|
|
if world_size == 1 {
|
|
// No sharding needed for single GPU
|
|
return Ok(parameters.clone());
|
|
}
|
|
|
|
// Calculate shard size
|
|
let total_elements = parameters.numel();
|
|
let elements_per_shard = total_elements.div_ceil(world_size); // Round up
|
|
let start_idx = rank * elements_per_shard;
|
|
let end_idx = (start_idx + elements_per_shard).min(total_elements);
|
|
let shard_size = end_idx - start_idx;
|
|
|
|
// Create local shard (simplified - in real implementation would slice the actual tensor)
|
|
let shard_shape = crate::TensorShape::new(vec![shard_size])?;
|
|
let local_shard = Tensor::zeros(shard_shape, &crate::Device::default())?;
|
|
|
|
// Update memory statistics
|
|
let total_memory_mb = (total_elements * 4) as f32 / (1024.0 * 1024.0); // f32 = 4 bytes
|
|
let shard_memory_mb = (shard_size * 4) as f32 / (1024.0 * 1024.0);
|
|
|
|
self.memory_stats.total_param_memory_mb = total_memory_mb;
|
|
self.memory_stats.local_shard_memory_mb = shard_memory_mb;
|
|
self.memory_stats.memory_saved_mb = total_memory_mb - shard_memory_mb;
|
|
|
|
tracing::debug!(
|
|
"FSDP sharded {:.2}MB -> {:.2}MB ({:.1}% reduction)",
|
|
total_memory_mb,
|
|
shard_memory_mb,
|
|
(self.memory_stats.memory_saved_mb / total_memory_mb) * 100.0
|
|
);
|
|
|
|
Ok(local_shard)
|
|
}
|
|
|
|
/// Synchronize gradients across all processes
|
|
pub async fn sync_gradients(&self, gradients: &mut Tensor) -> Result<()> {
|
|
use crate::comm::ReduceOp;
|
|
|
|
// Perform AllReduce to sum gradients across all processes
|
|
self.process_group
|
|
.all_reduce(gradients, ReduceOp::Sum)
|
|
.await?;
|
|
|
|
// Average the gradients
|
|
let world_size = self.process_group.world_size() as f32;
|
|
*gradients = gradients.div_scalar(world_size)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get current memory statistics
|
|
pub fn memory_stats(&self) -> &FsdpMemoryStats {
|
|
&self.memory_stats
|
|
}
|
|
|
|
/// Calculate memory reduction percentage
|
|
pub fn memory_reduction_percent(&self) -> f32 {
|
|
if self.memory_stats.total_param_memory_mb > 0.0 {
|
|
(self.memory_stats.memory_saved_mb / self.memory_stats.total_param_memory_mb) * 100.0
|
|
} else {
|
|
0.0
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Tensor Parallel implementation
|
|
pub struct TensorParallel {
|
|
/// Process group for tensor parallel communication
|
|
pub process_group: ProcessGroup,
|
|
}
|
|
|
|
impl TensorParallel {
|
|
/// Create new tensor parallel instance
|
|
pub fn new(process_group: ProcessGroup) -> Result<Self> {
|
|
Ok(Self { process_group })
|
|
}
|
|
|
|
/// Parallel matrix multiplication using [`ColParallelLinear`].
|
|
///
|
|
/// `a` is `[batch, in_features]` and `b` is the weight matrix
|
|
/// `[out_features, in_features]` in row-major order. The method shards
|
|
/// with `tp_size = 1` (this rank is the sole rank) so the result equals a
|
|
/// full, un-sharded matmul.
|
|
pub fn matmul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor> {
|
|
let a_dims = a.shape().dims();
|
|
let b_dims = b.shape().dims();
|
|
|
|
if a_dims.len() != 2 || b_dims.len() != 2 {
|
|
return Err(DistributedError::parallelism(
|
|
"TensorParallel::matmul",
|
|
"expected 2-D tensors",
|
|
));
|
|
}
|
|
|
|
let batch = a_dims[0];
|
|
let in_features = a_dims[1];
|
|
let out_features = b_dims[0];
|
|
|
|
if b_dims[1] != in_features {
|
|
return Err(DistributedError::parallelism(
|
|
"TensorParallel::matmul",
|
|
format!(
|
|
"dimension mismatch: a has in_features={in_features} but b has inner dim={}",
|
|
b_dims[1]
|
|
),
|
|
));
|
|
}
|
|
|
|
let a_data = a.to_vec().map_err(|e| {
|
|
DistributedError::parallelism("TensorParallel::matmul", format!("to_vec a: {e}"))
|
|
})?;
|
|
let b_data = b.to_vec().map_err(|e| {
|
|
DistributedError::parallelism("TensorParallel::matmul", format!("to_vec b: {e}"))
|
|
})?;
|
|
|
|
// Use a single-rank ColParallelLinear (tp_size = 1) for the full matmul.
|
|
let col = ColParallelLinear::new(&b_data, None, in_features, out_features, 1, 0)?;
|
|
let output = col.forward_cpu(&a_data, batch);
|
|
|
|
tracing::debug!(
|
|
"Tensor parallel matmul: {batch}x{in_features} * {out_features}x{in_features}"
|
|
);
|
|
|
|
Tensor::from_vec(output, &[batch, out_features], &crate::Device::default()).map_err(|e| {
|
|
DistributedError::parallelism("TensorParallel::matmul", format!("from_vec: {e}"))
|
|
})
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Column-parallel linear layer
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Column-parallel linear layer for tensor parallelism.
|
|
///
|
|
/// The full weight matrix `[out_features, in_features]` is partitioned along
|
|
/// the output (row) dimension. Each TP rank holds
|
|
/// `[out_features / tp_size, in_features]` rows.
|
|
///
|
|
/// `forward_cpu(x)` computes `x @ weight_shard.T` locally. Because each
|
|
/// rank produces a non-overlapping output slice, the sharded outputs can be
|
|
/// concatenated (AllGather) by the caller, or fed directly into a
|
|
/// [`RowParallelLinear`] which will AllReduce the partial sums.
|
|
///
|
|
/// # Invariants
|
|
/// * `weight_full.len() == out_features * in_features`
|
|
/// * `out_features % tp_size == 0`
|
|
/// * `0 <= rank < tp_size`
|
|
pub struct ColParallelLinear {
|
|
/// Weight shard: `[out_features / tp_size, in_features]`, row-major.
|
|
pub weight_shard: Vec<f32>,
|
|
/// Optional bias shard: `[out_features / tp_size]`.
|
|
pub bias_shard: Option<Vec<f32>>,
|
|
/// Inner dimension (columns of the weight matrix).
|
|
pub in_features: usize,
|
|
/// Rows owned by this rank (`out_features / tp_size`).
|
|
pub out_features_per_rank: usize,
|
|
/// Total tensor-parallel degree.
|
|
pub tp_size: usize,
|
|
/// This rank's index within the TP group (`0..tp_size`).
|
|
pub rank: usize,
|
|
}
|
|
|
|
impl ColParallelLinear {
|
|
/// Construct a column-parallel linear layer from the full weight matrix.
|
|
///
|
|
/// # Arguments
|
|
/// * `weight_full` — flat, row-major `[out_features, in_features]`.
|
|
/// * `bias_full` — optional flat `[out_features]` bias.
|
|
/// * `in_features` — inner dimension K.
|
|
/// * `out_features` — outer dimension N (must be divisible by `tp_size`).
|
|
/// * `tp_size` — number of tensor-parallel ranks.
|
|
/// * `rank` — this rank's index (`0..tp_size`).
|
|
///
|
|
/// # Errors
|
|
/// Returns [`DistributedError`] if `out_features` is not evenly divisible
|
|
/// by `tp_size`, if the weight slice length is wrong, or if `rank >= tp_size`.
|
|
pub fn new(
|
|
weight_full: &[f32],
|
|
bias_full: Option<&[f32]>,
|
|
in_features: usize,
|
|
out_features: usize,
|
|
tp_size: usize,
|
|
rank: usize,
|
|
) -> Result<Self> {
|
|
if tp_size == 0 {
|
|
return Err(DistributedError::parallelism(
|
|
"ColParallelLinear",
|
|
"tp_size must be >= 1",
|
|
));
|
|
}
|
|
if rank >= tp_size {
|
|
return Err(DistributedError::parallelism(
|
|
"ColParallelLinear",
|
|
format!("rank {rank} is out of range for tp_size {tp_size}"),
|
|
));
|
|
}
|
|
if out_features % tp_size != 0 {
|
|
return Err(DistributedError::parallelism(
|
|
"ColParallelLinear",
|
|
format!("out_features {out_features} must be divisible by tp_size {tp_size}"),
|
|
));
|
|
}
|
|
|
|
let expected = out_features * in_features;
|
|
if weight_full.len() != expected {
|
|
return Err(DistributedError::parallelism(
|
|
"ColParallelLinear",
|
|
format!(
|
|
"weight_full length {} does not match out_features*in_features={expected}",
|
|
weight_full.len()
|
|
),
|
|
));
|
|
}
|
|
|
|
let shard_rows = out_features / tp_size;
|
|
let row_start = rank * shard_rows;
|
|
let row_end = row_start + shard_rows;
|
|
|
|
// Extract rows [row_start, row_end) from the weight matrix.
|
|
let weight_shard = weight_full[row_start * in_features..row_end * in_features].to_vec();
|
|
|
|
let bias_shard = bias_full.map(|bias| bias[row_start..row_end].to_vec());
|
|
|
|
Ok(Self {
|
|
weight_shard,
|
|
bias_shard,
|
|
in_features,
|
|
out_features_per_rank: shard_rows,
|
|
tp_size,
|
|
rank,
|
|
})
|
|
}
|
|
|
|
/// CPU reference forward pass.
|
|
///
|
|
/// Computes `x [batch, in_features] @ weight_shard.T`
|
|
/// → `[batch, out_features_per_rank]`.
|
|
/// Adds `bias_shard` if present.
|
|
pub fn forward_cpu(&self, x: &[f32], batch: usize) -> Vec<f32> {
|
|
let n = self.out_features_per_rank;
|
|
let k = self.in_features;
|
|
let mut out = vec![0.0_f32; batch * n];
|
|
|
|
for b in 0..batch {
|
|
for j in 0..n {
|
|
let mut acc = 0.0_f32;
|
|
for i in 0..k {
|
|
acc += x[b * k + i] * self.weight_shard[j * k + i];
|
|
}
|
|
out[b * n + j] = acc + self.bias_shard.as_ref().map_or(0.0, |bias| bias[j]);
|
|
}
|
|
}
|
|
|
|
out
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Row-parallel linear layer
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Row-parallel linear layer for tensor parallelism.
|
|
///
|
|
/// The full weight matrix `[out_features, in_features]` is partitioned along
|
|
/// the input (column) dimension. Each TP rank holds
|
|
/// `[out_features, in_features / tp_size]` columns.
|
|
///
|
|
/// Each rank receives a matching input shard `x_shard [batch, in_features/tp_size]`
|
|
/// (typically the output of a [`ColParallelLinear`]) and computes
|
|
/// `x_shard @ weight_shard.T` locally to produce a partial sum
|
|
/// `[batch, out_features]`. A subsequent AllReduce over all TP ranks yields
|
|
/// the full output.
|
|
///
|
|
/// # Invariants
|
|
/// * `weight_full.len() == out_features * in_features`
|
|
/// * `in_features % tp_size == 0`
|
|
/// * `0 <= rank < tp_size`
|
|
pub struct RowParallelLinear {
|
|
/// Weight shard: `[out_features, in_features / tp_size]`, row-major.
|
|
pub weight_shard: Vec<f32>,
|
|
/// Bias applied after AllReduce — **only rank 0 adds it** to avoid
|
|
/// double-counting during the reduce. Shape: `[out_features]`.
|
|
pub bias: Option<Vec<f32>>,
|
|
/// Output feature count (unchanged by column-sharding).
|
|
pub out_features: usize,
|
|
/// Columns owned by this rank (`in_features / tp_size`).
|
|
pub in_features_per_rank: usize,
|
|
/// Total tensor-parallel degree.
|
|
pub tp_size: usize,
|
|
/// This rank's index within the TP group (`0..tp_size`).
|
|
pub rank: usize,
|
|
/// Process group used for the AllReduce collective.
|
|
pub process_group: ProcessGroup,
|
|
}
|
|
|
|
impl RowParallelLinear {
|
|
/// Construct a row-parallel linear layer from the full weight matrix.
|
|
///
|
|
/// # Arguments
|
|
/// * `weight_full` — flat, row-major `[out_features, in_features]`.
|
|
/// * `bias` — optional flat `[out_features]` bias (owned; only
|
|
/// rank 0 applies it after AllReduce).
|
|
/// * `out_features` — output dimension M.
|
|
/// * `in_features` — inner dimension K (must be divisible by `tp_size`).
|
|
/// * `tp_size` — number of tensor-parallel ranks.
|
|
/// * `rank` — this rank's index (`0..tp_size`).
|
|
/// * `process_group`— process group for the AllReduce collective.
|
|
///
|
|
/// # Errors
|
|
/// Returns [`DistributedError`] if `in_features` is not evenly divisible
|
|
/// by `tp_size`, if the weight slice length is wrong, or if `rank >= tp_size`.
|
|
pub fn new(
|
|
weight_full: &[f32],
|
|
bias: Option<Vec<f32>>,
|
|
out_features: usize,
|
|
in_features: usize,
|
|
tp_size: usize,
|
|
rank: usize,
|
|
process_group: ProcessGroup,
|
|
) -> Result<Self> {
|
|
if tp_size == 0 {
|
|
return Err(DistributedError::parallelism(
|
|
"RowParallelLinear",
|
|
"tp_size must be >= 1",
|
|
));
|
|
}
|
|
if rank >= tp_size {
|
|
return Err(DistributedError::parallelism(
|
|
"RowParallelLinear",
|
|
format!("rank {rank} is out of range for tp_size {tp_size}"),
|
|
));
|
|
}
|
|
if in_features % tp_size != 0 {
|
|
return Err(DistributedError::parallelism(
|
|
"RowParallelLinear",
|
|
format!("in_features {in_features} must be divisible by tp_size {tp_size}"),
|
|
));
|
|
}
|
|
|
|
let expected = out_features * in_features;
|
|
if weight_full.len() != expected {
|
|
return Err(DistributedError::parallelism(
|
|
"RowParallelLinear",
|
|
format!(
|
|
"weight_full length {} does not match out_features*in_features={expected}",
|
|
weight_full.len()
|
|
),
|
|
));
|
|
}
|
|
|
|
let shard_cols = in_features / tp_size;
|
|
let col_start = rank * shard_cols;
|
|
let col_end = col_start + shard_cols;
|
|
|
|
// Extract columns [col_start, col_end) from every row of the weight matrix.
|
|
// weight_full is row-major [out_features, in_features], so row i spans
|
|
// indices [i*in_features .. (i+1)*in_features].
|
|
let mut weight_shard = Vec::with_capacity(out_features * shard_cols);
|
|
for row in 0..out_features {
|
|
let row_base = row * in_features;
|
|
weight_shard.extend_from_slice(&weight_full[row_base + col_start..row_base + col_end]);
|
|
}
|
|
|
|
Ok(Self {
|
|
weight_shard,
|
|
bias,
|
|
out_features,
|
|
in_features_per_rank: shard_cols,
|
|
tp_size,
|
|
rank,
|
|
process_group,
|
|
})
|
|
}
|
|
|
|
/// CPU reference forward pass (no network I/O).
|
|
///
|
|
/// Computes the local partial sum
|
|
/// `x_shard [batch, in_features/tp_size] @ weight_shard.T`
|
|
/// → `[batch, out_features]`.
|
|
///
|
|
/// Bias is added **only on rank 0** so that, after an AllReduce sum, the
|
|
/// bias is incorporated exactly once in the final result.
|
|
///
|
|
/// In production this is followed by
|
|
/// `process_group.all_reduce(&mut partial, ReduceOp::Sum)`.
|
|
pub fn forward_cpu(&self, x_shard: &[f32], batch: usize) -> Vec<f32> {
|
|
let n = self.out_features;
|
|
let k = self.in_features_per_rank;
|
|
let mut partial = vec![0.0_f32; batch * n];
|
|
|
|
for b in 0..batch {
|
|
for j in 0..n {
|
|
let mut acc = 0.0_f32;
|
|
for i in 0..k {
|
|
acc += x_shard[b * k + i] * self.weight_shard[j * k + i];
|
|
}
|
|
partial[b * n + j] = acc;
|
|
}
|
|
}
|
|
|
|
// Add bias only on rank 0 to avoid double-counting across TP ranks.
|
|
if self.rank == 0 {
|
|
if let Some(bias) = &self.bias {
|
|
for b in 0..batch {
|
|
for j in 0..n {
|
|
partial[b * n + j] += bias[j];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
partial
|
|
}
|
|
|
|
/// Async forward pass with a real AllReduce collective.
|
|
///
|
|
/// 1. Computes the local partial sum via [`forward_cpu`][Self::forward_cpu].
|
|
/// 2. Wraps the result in a [`Tensor`] and issues an AllReduce Sum across
|
|
/// all TP ranks via `self.process_group`.
|
|
/// 3. Returns the reduced flat `[batch * out_features]` buffer.
|
|
///
|
|
/// In CPU simulation mode (`Backend::Cpu`, single process) the AllReduce
|
|
/// is a no-op that scales by `world_size`; the returned data is therefore
|
|
/// correct only when `tp_size == 1` or when multiple real processes run.
|
|
/// Use [`forward_cpu`][Self::forward_cpu] directly for unit tests that
|
|
/// manually sum partials.
|
|
pub async fn forward(&self, x_shard: &[f32], batch: usize) -> Result<Vec<f32>> {
|
|
use crate::comm::ReduceOp;
|
|
|
|
let local = self.forward_cpu(x_shard, batch);
|
|
|
|
// Wrap in a Tensor for the collective API.
|
|
let mut t = Tensor::from_vec(
|
|
local,
|
|
&[batch, self.out_features],
|
|
&crate::Device::default(),
|
|
)
|
|
.map_err(|e| {
|
|
DistributedError::parallelism("RowParallelLinear::forward", format!("from_vec: {e}"))
|
|
})?;
|
|
|
|
// AllReduce: sum partial results across all TP ranks.
|
|
self.process_group.all_reduce(&mut t, ReduceOp::Sum).await?;
|
|
|
|
t.to_vec().map_err(|e| {
|
|
DistributedError::parallelism(
|
|
"RowParallelLinear::forward",
|
|
format!("to_vec after allreduce: {e}"),
|
|
)
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Pipeline Parallel implementation
|
|
pub struct PipelineParallel {
|
|
/// Process group for pipeline parallel communication
|
|
pub process_group: ProcessGroup,
|
|
/// Number of pipeline stages
|
|
pub num_stages: usize,
|
|
}
|
|
|
|
impl PipelineParallel {
|
|
/// Create new pipeline parallel instance
|
|
pub fn new(process_group: ProcessGroup, num_stages: usize) -> Result<Self> {
|
|
Ok(Self {
|
|
process_group,
|
|
num_stages,
|
|
})
|
|
}
|
|
|
|
/// Forward pass through pipeline stages
|
|
pub fn forward(&self, inputs: &Tensor) -> Result<Tensor> {
|
|
// Simplified pipeline forward pass
|
|
// In real implementation, this would:
|
|
// 1. Split model into stages
|
|
// 2. Pass activations between stages
|
|
// 3. Handle micro-batching
|
|
|
|
let outputs = inputs.clone();
|
|
|
|
tracing::debug!(
|
|
"Pipeline forward: {} stages, input shape: {:?}",
|
|
self.num_stages,
|
|
inputs.shape().dims()
|
|
);
|
|
|
|
Ok(outputs)
|
|
}
|
|
|
|
/// Backward pass through pipeline stages
|
|
pub fn backward(&self, grad_outputs: &Tensor) -> Result<Tensor> {
|
|
// Simplified pipeline backward pass
|
|
let grad_inputs = grad_outputs.clone();
|
|
|
|
tracing::debug!(
|
|
"Pipeline backward: {} stages, grad_output shape: {:?}",
|
|
self.num_stages,
|
|
grad_outputs.shape().dims()
|
|
);
|
|
|
|
Ok(grad_inputs)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::group::ProcessGroup;
|
|
use crate::{Backend, BackendConfig, Device, Tensor, TensorShape};
|
|
|
|
#[tokio::test]
|
|
async fn test_data_parallel() {
|
|
let config = BackendConfig::cpu();
|
|
let pg = ProcessGroup::new_with_config(Backend::Cpu, 2, 0, config)
|
|
.await
|
|
.unwrap();
|
|
let mut dp = DataParallel::new(pg, 2);
|
|
|
|
dp.start_accumulation();
|
|
|
|
let mut gradients = vec![
|
|
Tensor::ones(&TensorShape::new(vec![10, 10]).unwrap(), &Device::default()).unwrap(),
|
|
Tensor::ones(&TensorShape::new(vec![5, 5]).unwrap(), &Device::default()).unwrap(),
|
|
];
|
|
|
|
let should_sync = dp.accumulate_gradients(&mut gradients).await.unwrap();
|
|
assert!(!should_sync); // First accumulation shouldn't sync
|
|
|
|
let should_sync = dp.accumulate_gradients(&mut gradients).await.unwrap();
|
|
assert!(should_sync); // Second accumulation should sync
|
|
}
|
|
|
|
#[test]
|
|
fn test_fsdp_config() {
|
|
let config = FsdpConfig::default();
|
|
assert_eq!(config.sharding_strategy, ShardingStrategy::FullShard);
|
|
assert_eq!(config.min_param_size, 1000);
|
|
assert!(config.mixed_precision);
|
|
assert!(config.flatten_parameters);
|
|
assert!(!config.cpu_offload);
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tensor-parallel unit tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tp_tests {
|
|
use super::*;
|
|
use crate::{Backend, BackendConfig, Device, Tensor, TensorShape};
|
|
|
|
// ------------------------------------------------------------------
|
|
// Helpers
|
|
// ------------------------------------------------------------------
|
|
|
|
/// Naive CPU matmul: A [batch, k] @ W.T [k, n] → [batch, n].
|
|
/// Weight W is stored row-major as [n, k].
|
|
fn naive_matmul(a: &[f32], w: &[f32], batch: usize, k: usize, n: usize) -> Vec<f32> {
|
|
let mut out = vec![0.0_f32; batch * n];
|
|
for b in 0..batch {
|
|
for j in 0..n {
|
|
let mut acc = 0.0_f32;
|
|
for i in 0..k {
|
|
acc += a[b * k + i] * w[j * k + i];
|
|
}
|
|
out[b * n + j] = acc;
|
|
}
|
|
}
|
|
out
|
|
}
|
|
|
|
fn assert_vec_approx(got: &[f32], expected: &[f32], tol: f32, label: &str) {
|
|
assert_eq!(
|
|
got.len(),
|
|
expected.len(),
|
|
"{label}: length mismatch: got {} expected {}",
|
|
got.len(),
|
|
expected.len()
|
|
);
|
|
for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
|
|
assert!(
|
|
(g - e).abs() <= tol,
|
|
"{label}: element[{i}] got={g} expected={e} diff={}",
|
|
(g - e).abs()
|
|
);
|
|
}
|
|
}
|
|
|
|
fn cpu_pg() -> ProcessGroup {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
rt.block_on(async {
|
|
let config = BackendConfig::cpu();
|
|
ProcessGroup::new_with_config(Backend::Cpu, 1, 0, config)
|
|
.await
|
|
.unwrap()
|
|
})
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// ColParallelLinear tests
|
|
// ------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_col_parallel_single_rank() {
|
|
// tp_size = 1 → weight_shard == weight_full; output matches naive matmul.
|
|
let batch = 2;
|
|
let in_f = 3;
|
|
let out_f = 4;
|
|
|
|
// Weight [out_f, in_f]
|
|
let w: Vec<f32> = (1..=(out_f * in_f) as i32).map(|x| x as f32).collect();
|
|
// Input [batch, in_f]
|
|
let x: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();
|
|
|
|
let col = ColParallelLinear::new(&w, None, in_f, out_f, 1, 0).unwrap();
|
|
assert_eq!(col.weight_shard, w, "single rank should hold full weight");
|
|
assert_eq!(col.out_features_per_rank, out_f);
|
|
|
|
let got = col.forward_cpu(&x, batch);
|
|
let expected = naive_matmul(&x, &w, batch, in_f, out_f);
|
|
assert_vec_approx(&got, &expected, 1e-5, "col_single_rank");
|
|
}
|
|
|
|
#[test]
|
|
fn test_col_parallel_two_rank_sharding() {
|
|
// out_features = 4, tp_size = 2:
|
|
// rank 0 → rows [0, 1] (out cols 0,1 of full result)
|
|
// rank 1 → rows [2, 3] (out cols 2,3 of full result)
|
|
// Concatenated outputs must equal the full matmul.
|
|
let batch = 2;
|
|
let in_f = 3;
|
|
let out_f = 4;
|
|
|
|
let w: Vec<f32> = (1..=(out_f * in_f) as i32).map(|x| x as f32).collect();
|
|
let x: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();
|
|
|
|
let col0 = ColParallelLinear::new(&w, None, in_f, out_f, 2, 0).unwrap();
|
|
let col1 = ColParallelLinear::new(&w, None, in_f, out_f, 2, 1).unwrap();
|
|
|
|
assert_eq!(col0.out_features_per_rank, 2);
|
|
assert_eq!(col1.out_features_per_rank, 2);
|
|
|
|
let out0 = col0.forward_cpu(&x, batch); // [batch, 2]
|
|
let out1 = col1.forward_cpu(&x, batch); // [batch, 2]
|
|
|
|
// Interleave: for each batch row, concatenate out0[b] then out1[b] → [batch, 4]
|
|
let mut combined = Vec::with_capacity(batch * out_f);
|
|
for b in 0..batch {
|
|
combined.extend_from_slice(&out0[b * 2..(b + 1) * 2]);
|
|
combined.extend_from_slice(&out1[b * 2..(b + 1) * 2]);
|
|
}
|
|
|
|
let expected = naive_matmul(&x, &w, batch, in_f, out_f);
|
|
assert_vec_approx(&combined, &expected, 1e-5, "col_two_rank_sharding");
|
|
}
|
|
|
|
#[test]
|
|
fn test_col_parallel_with_bias() {
|
|
// Bias is added to the output slice owned by this rank.
|
|
let batch = 2;
|
|
let in_f = 3;
|
|
let out_f = 4;
|
|
|
|
let w: Vec<f32> = vec![1.0; out_f * in_f];
|
|
let x: Vec<f32> = vec![1.0; batch * in_f]; // all-ones → each output is in_f
|
|
let bias: Vec<f32> = vec![10.0, 20.0, 30.0, 40.0]; // one per output feature
|
|
|
|
// tp_size = 1: rank 0 gets all rows and the full bias.
|
|
let col = ColParallelLinear::new(&w, Some(&bias), in_f, out_f, 1, 0).unwrap();
|
|
let got = col.forward_cpu(&x, batch);
|
|
|
|
// Each output element should be in_f * 1.0 + bias[j] = 3 + bias[j]
|
|
let expected: Vec<f32> = (0..batch)
|
|
.flat_map(|_| bias.iter().map(|b| in_f as f32 + b))
|
|
.collect();
|
|
assert_vec_approx(&got, &expected, 1e-5, "col_with_bias");
|
|
}
|
|
|
|
#[test]
|
|
fn test_col_parallel_bias_two_ranks() {
|
|
// With tp_size = 2: rank 0 gets bias[0..2], rank 1 gets bias[2..4].
|
|
let in_f = 2;
|
|
let out_f = 4;
|
|
let batch = 1;
|
|
|
|
let w: Vec<f32> = vec![1.0; out_f * in_f];
|
|
let x: Vec<f32> = vec![1.0; in_f]; // [1, in_f]
|
|
let bias: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0];
|
|
|
|
let col0 = ColParallelLinear::new(&w, Some(&bias), in_f, out_f, 2, 0).unwrap();
|
|
let col1 = ColParallelLinear::new(&w, Some(&bias), in_f, out_f, 2, 1).unwrap();
|
|
|
|
let out0 = col0.forward_cpu(&x, batch); // expects [in_f + 1, in_f + 2] = [3,4]
|
|
let out1 = col1.forward_cpu(&x, batch); // expects [in_f + 3, in_f + 4] = [5,6]
|
|
|
|
assert_vec_approx(&out0, &[3.0, 4.0], 1e-5, "bias_rank0");
|
|
assert_vec_approx(&out1, &[5.0, 6.0], 1e-5, "bias_rank1");
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// RowParallelLinear tests
|
|
// ------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_row_parallel_single_rank() {
|
|
// tp_size = 1 → forward_cpu equals naive matmul.
|
|
let batch = 2;
|
|
let in_f = 4;
|
|
let out_f = 3;
|
|
|
|
let w: Vec<f32> = (1..=(out_f * in_f) as i32).map(|x| x as f32).collect();
|
|
let x: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();
|
|
|
|
let pg = cpu_pg();
|
|
let row = RowParallelLinear::new(&w, None, out_f, in_f, 1, 0, pg).unwrap();
|
|
assert_eq!(row.in_features_per_rank, in_f);
|
|
|
|
let got = row.forward_cpu(&x, batch);
|
|
let expected = naive_matmul(&x, &w, batch, in_f, out_f);
|
|
assert_vec_approx(&got, &expected, 1e-5, "row_single_rank");
|
|
}
|
|
|
|
#[test]
|
|
fn test_row_parallel_two_rank_sharding() {
|
|
// in_features = 4, tp_size = 2:
|
|
// rank 0 → cols [0,1] of W → partial sum from x[0..2]
|
|
// rank 1 → cols [2,3] of W → partial sum from x[2..4]
|
|
// Adding partial sums → full matmul output.
|
|
let batch = 2;
|
|
let in_f = 4;
|
|
let out_f = 3;
|
|
|
|
let w: Vec<f32> = (1..=(out_f * in_f) as i32).map(|x| x as f32).collect();
|
|
let x: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();
|
|
|
|
let pg0 = cpu_pg();
|
|
let pg1 = cpu_pg();
|
|
|
|
let row0 = RowParallelLinear::new(&w, None, out_f, in_f, 2, 0, pg0).unwrap();
|
|
let row1 = RowParallelLinear::new(&w, None, out_f, in_f, 2, 1, pg1).unwrap();
|
|
|
|
// x_shard for each rank: columns [0..2] and [2..4] of x per batch row.
|
|
let shard_size = in_f / 2;
|
|
let mut x_shard0 = Vec::with_capacity(batch * shard_size);
|
|
let mut x_shard1 = Vec::with_capacity(batch * shard_size);
|
|
for b in 0..batch {
|
|
x_shard0.extend_from_slice(&x[b * in_f..b * in_f + shard_size]);
|
|
x_shard1.extend_from_slice(&x[b * in_f + shard_size..b * in_f + in_f]);
|
|
}
|
|
|
|
let partial0 = row0.forward_cpu(&x_shard0, batch);
|
|
let partial1 = row1.forward_cpu(&x_shard1, batch);
|
|
|
|
// AllReduce (sum): add the two partials element-wise.
|
|
let combined: Vec<f32> = partial0
|
|
.iter()
|
|
.zip(partial1.iter())
|
|
.map(|(a, b)| a + b)
|
|
.collect();
|
|
|
|
let expected = naive_matmul(&x, &w, batch, in_f, out_f);
|
|
assert_vec_approx(&combined, &expected, 1e-5, "row_two_rank_sharding");
|
|
}
|
|
|
|
#[test]
|
|
fn test_row_parallel_bias_rank0_only() {
|
|
// Bias is added only by rank 0; rank 1 must not add it.
|
|
let batch = 1;
|
|
let in_f = 2;
|
|
let out_f = 2;
|
|
|
|
// Identity-ish weight [out_f, in_f] = [[1,0],[0,1]]
|
|
let w: Vec<f32> = vec![1.0, 0.0, 0.0, 1.0];
|
|
let bias: Vec<f32> = vec![100.0, 200.0];
|
|
let x: Vec<f32> = vec![3.0, 5.0]; // [1, 2]
|
|
|
|
let pg0 = cpu_pg();
|
|
let pg1 = cpu_pg();
|
|
|
|
// rank 0 gets cols [0..1], rank 1 gets cols [1..2].
|
|
let row0 = RowParallelLinear::new(&w, Some(bias.clone()), out_f, in_f, 2, 0, pg0).unwrap();
|
|
let row1 = RowParallelLinear::new(&w, Some(bias.clone()), out_f, in_f, 2, 1, pg1).unwrap();
|
|
|
|
// x_shard0 = [3.0], x_shard1 = [5.0]
|
|
let partial0 = row0.forward_cpu(&[3.0], batch);
|
|
let partial1 = row1.forward_cpu(&[5.0], batch);
|
|
|
|
// partial0 = [3*1 + bias[0], 3*0 + bias[1]] = [103, 200]
|
|
// partial1 = [5*0, 5*1 ] = [ 0, 5]
|
|
assert_vec_approx(&partial0, &[103.0, 200.0], 1e-5, "partial_rank0");
|
|
assert_vec_approx(&partial1, &[0.0, 5.0], 1e-5, "partial_rank1");
|
|
|
|
// AllReduce (sum) → [103, 205] = expected [3+100, 5+200]
|
|
let summed: Vec<f32> = partial0
|
|
.iter()
|
|
.zip(partial1.iter())
|
|
.map(|(a, b)| a + b)
|
|
.collect();
|
|
assert_vec_approx(&summed, &[103.0, 205.0], 1e-5, "bias_rank0_only_final");
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// TensorParallel::matmul
|
|
// ------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_matmul_not_zeros() {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
rt.block_on(async {
|
|
let config = BackendConfig::cpu();
|
|
let pg = ProcessGroup::new_with_config(Backend::Cpu, 1, 0, config)
|
|
.await
|
|
.unwrap();
|
|
let tp = TensorParallel::new(pg).unwrap();
|
|
|
|
let batch = 2;
|
|
let in_f = 3;
|
|
let out_f = 4;
|
|
|
|
let a_data: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();
|
|
let b_data: Vec<f32> = (1..=(out_f * in_f) as i32).map(|x| x as f32).collect();
|
|
|
|
let a = Tensor::from_vec(a_data.clone(), &[batch, in_f], &Device::default()).unwrap();
|
|
let b = Tensor::from_vec(b_data.clone(), &[out_f, in_f], &Device::default()).unwrap();
|
|
|
|
let result = tp.matmul(&a, &b).unwrap();
|
|
let result_data = result.to_vec().unwrap();
|
|
|
|
// Verify output is non-zero.
|
|
let all_zero = result_data.iter().all(|&v| v == 0.0);
|
|
assert!(!all_zero, "matmul output must not be all zeros");
|
|
|
|
// Verify correctness against naive matmul.
|
|
let expected = naive_matmul(&a_data, &b_data, batch, in_f, out_f);
|
|
assert_eq!(result.dims(), &[batch, out_f]);
|
|
assert_vec_approx(&result_data, &expected, 1e-4, "tp_matmul_correctness");
|
|
});
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// ColParallel → RowParallel round-trip
|
|
// ------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_col_then_row_roundtrip() {
|
|
// Simulate the standard TP linear pipeline:
|
|
// x [batch, in_f] → ColParallel → x_shard [batch, in_f/tp]
|
|
// x_shard → RowParallel → partial [batch, out_f]
|
|
// AllReduce (sum of tp partials) → [batch, out_f]
|
|
//
|
|
// With tp_size = 2 this is a genuine two-stage sharded pipeline.
|
|
// We run both ranks in the same process (CPU sim) and manually add
|
|
// their partial outputs to mimic the AllReduce.
|
|
let batch = 3;
|
|
let in_f = 4;
|
|
let hidden = 6; // ColParallel output / RowParallel input
|
|
let out_f = 5;
|
|
|
|
// Two weight matrices simulating a two-layer MLP.
|
|
let w1: Vec<f32> = (1..=(hidden * in_f) as i32)
|
|
.map(|x| x as f32 * 0.1)
|
|
.collect();
|
|
let w2: Vec<f32> = (1..=(out_f * hidden) as i32)
|
|
.map(|x| x as f32 * 0.05)
|
|
.collect();
|
|
let x: Vec<f32> = (1..=(batch * in_f) as i32).map(|x| x as f32).collect();
|
|
|
|
let tp_size = 2;
|
|
|
|
// --- Stage 1: ColParallel on w1 (shard along hidden / output dim) ---
|
|
let col0 = ColParallelLinear::new(&w1, None, in_f, hidden, tp_size, 0).unwrap();
|
|
let col1 = ColParallelLinear::new(&w1, None, in_f, hidden, tp_size, 1).unwrap();
|
|
|
|
let h_shard0 = col0.forward_cpu(&x, batch); // [batch, hidden/2]
|
|
let h_shard1 = col1.forward_cpu(&x, batch); // [batch, hidden/2]
|
|
|
|
// --- Stage 2: RowParallel on w2 (shard along hidden / input dim) ---
|
|
// Each rank receives the corresponding column shard from Stage 1.
|
|
let pg0 = cpu_pg();
|
|
let pg1 = cpu_pg();
|
|
let row0 = RowParallelLinear::new(&w2, None, out_f, hidden, tp_size, 0, pg0).unwrap();
|
|
let row1 = RowParallelLinear::new(&w2, None, out_f, hidden, tp_size, 1, pg1).unwrap();
|
|
|
|
let partial0 = row0.forward_cpu(&h_shard0, batch);
|
|
let partial1 = row1.forward_cpu(&h_shard1, batch);
|
|
|
|
// AllReduce (sum) → [batch, out_f]
|
|
let combined: Vec<f32> = partial0
|
|
.iter()
|
|
.zip(partial1.iter())
|
|
.map(|(a, b)| a + b)
|
|
.collect();
|
|
|
|
// Reference: single-rank two-layer matmul.
|
|
let h_ref = naive_matmul(&x, &w1, batch, in_f, hidden);
|
|
let out_ref = naive_matmul(&h_ref, &w2, batch, hidden, out_f);
|
|
|
|
assert_eq!(combined.len(), batch * out_f);
|
|
assert_vec_approx(&combined, &out_ref, 1e-3, "col_then_row_roundtrip");
|
|
}
|
|
}
|