552 lines
19 KiB
Rust
552 lines
19 KiB
Rust
//! Resharding operations for DTensor redistribution.
|
|
//!
|
|
//! This module provides efficient implementations of tensor resharding operations
|
|
//! needed for distributed training, including:
|
|
//! - AllGather: Shard -> Replicate
|
|
//! - ReduceScatter: Partial -> Shard
|
|
//! - AllToAll: Shard(dim_a) -> Shard(dim_b)
|
|
//! - AllReduce: Partial -> Replicate
|
|
//!
|
|
//! These operations are fundamental to FSDP2's per-parameter sharding approach.
|
|
|
|
use crate::comm::{CommunicationPrimitive, ReduceOp};
|
|
use crate::device_mesh::DeviceMesh;
|
|
use crate::dtensor::{DTensor, PartialReduceOp, Placement, TensorSpec};
|
|
use crate::error::{DistributedError, Result};
|
|
use crate::group::ProcessGroup;
|
|
use rtx_tensor::{Device, Tensor};
|
|
use std::sync::Arc;
|
|
|
|
// =============================================================================
|
|
// Resharding Operations
|
|
// =============================================================================
|
|
|
|
/// Reshard a DTensor to a new placement specification.
|
|
///
|
|
/// This function determines the required communication pattern and
|
|
/// executes the resharding operation.
|
|
pub async fn reshard_dtensor(
|
|
dtensor: &DTensor,
|
|
target_placements: &[Placement],
|
|
pg: &ProcessGroup,
|
|
) -> Result<DTensor> {
|
|
let current_placements = &dtensor.spec().placements;
|
|
|
|
// No resharding needed if placements match
|
|
if current_placements == target_placements {
|
|
return Ok(dtensor.clone());
|
|
}
|
|
|
|
// Handle each mesh dimension's resharding
|
|
let mut result = dtensor.clone();
|
|
|
|
for (mesh_dim, (src, dst)) in current_placements
|
|
.iter()
|
|
.zip(target_placements.iter())
|
|
.enumerate()
|
|
{
|
|
if src == dst {
|
|
continue;
|
|
}
|
|
|
|
result = match (src, dst) {
|
|
// Shard -> Replicate: AllGather
|
|
(Placement::Shard { tensor_dim }, Placement::Replicate) => {
|
|
all_gather_reshard(&result, mesh_dim, *tensor_dim, pg).await?
|
|
}
|
|
|
|
// Replicate -> Shard: Local slice (no communication)
|
|
(Placement::Replicate, Placement::Shard { tensor_dim }) => {
|
|
replicate_to_shard(&result, mesh_dim, *tensor_dim)?
|
|
}
|
|
|
|
// Partial -> Replicate: AllReduce
|
|
(Placement::Partial { reduce_op }, Placement::Replicate) => {
|
|
all_reduce_reshard(&result, mesh_dim, *reduce_op, pg).await?
|
|
}
|
|
|
|
// Partial -> Shard: ReduceScatter
|
|
(Placement::Partial { reduce_op }, Placement::Shard { tensor_dim }) => {
|
|
reduce_scatter_reshard(&result, mesh_dim, *tensor_dim, *reduce_op, pg).await?
|
|
}
|
|
|
|
// Shard -> Shard (different dims): AllToAll
|
|
(
|
|
Placement::Shard {
|
|
tensor_dim: src_dim,
|
|
},
|
|
Placement::Shard {
|
|
tensor_dim: dst_dim,
|
|
},
|
|
) if src_dim != dst_dim => {
|
|
all_to_all_reshard(&result, mesh_dim, *src_dim, *dst_dim, pg).await?
|
|
}
|
|
|
|
// Same placement type - no action needed
|
|
(src_p, dst_p) if std::mem::discriminant(src_p) == std::mem::discriminant(dst_p) => {
|
|
result
|
|
}
|
|
|
|
// Unsupported transition
|
|
_ => {
|
|
return Err(DistributedError::configuration(format!(
|
|
"Unsupported resharding: {:?} -> {:?}",
|
|
src, dst
|
|
)));
|
|
}
|
|
};
|
|
}
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
// =============================================================================
|
|
// AllGather: Shard -> Replicate
|
|
// =============================================================================
|
|
|
|
/// Perform AllGather to convert sharded tensor to replicated.
|
|
async fn all_gather_reshard(
|
|
dtensor: &DTensor,
|
|
_mesh_dim: usize,
|
|
tensor_dim: usize,
|
|
pg: &ProcessGroup,
|
|
) -> Result<DTensor> {
|
|
let local_shard = dtensor.local_shard().clone();
|
|
let gathered = pg.allgather(&local_shard).await?;
|
|
|
|
// Concatenate gathered tensors along the sharded dimension
|
|
let full_tensor = match gathered {
|
|
crate::comm::AllGatherOutput::Tensor(t) => t,
|
|
crate::comm::AllGatherOutput::TensorList(tensors) => {
|
|
concat_along_dim(&tensors, tensor_dim)?
|
|
}
|
|
};
|
|
|
|
// Create new DTensor with replicated placement
|
|
let new_spec = TensorSpec::new(dtensor.global_shape().to_vec())
|
|
.with_placement(0, Placement::Replicate)
|
|
.with_requires_grad(dtensor.requires_grad());
|
|
|
|
DTensor::from_local(full_tensor, new_spec, dtensor.mesh().clone())
|
|
}
|
|
|
|
/// Concatenate tensors along a specific dimension.
|
|
fn concat_along_dim(tensors: &[Tensor], dim: usize) -> Result<Tensor> {
|
|
if tensors.is_empty() {
|
|
return Err(DistributedError::tensor(
|
|
"Cannot concatenate empty tensor list",
|
|
));
|
|
}
|
|
|
|
let first_shape = tensors[0].shape().dims().to_vec();
|
|
if dim >= first_shape.len() {
|
|
return Err(DistributedError::tensor(format!(
|
|
"Concatenation dim {} out of range for {}D tensor",
|
|
dim,
|
|
first_shape.len()
|
|
)));
|
|
}
|
|
|
|
// Calculate output shape
|
|
let total_dim_size: usize = tensors
|
|
.iter()
|
|
.map(|t| t.shape().dims().get(dim).copied().unwrap_or(0))
|
|
.sum();
|
|
|
|
let mut output_shape = first_shape.clone();
|
|
output_shape[dim] = total_dim_size;
|
|
|
|
// Simple case: dim 0 concatenation
|
|
if dim == 0 {
|
|
let mut combined_data = Vec::new();
|
|
for tensor in tensors {
|
|
let data = tensor
|
|
.data()
|
|
.map_err(|e| DistributedError::tensor(e.to_string()))?;
|
|
combined_data.extend(data);
|
|
}
|
|
return Tensor::from_data(combined_data, output_shape, &Device::cpu())
|
|
.map_err(|e| DistributedError::tensor(e.to_string()));
|
|
}
|
|
|
|
// General case: need to interleave data
|
|
let output_numel: usize = output_shape.iter().product();
|
|
let mut output_data = vec![0.0f32; output_numel];
|
|
|
|
// Calculate strides for output
|
|
let mut output_strides: Vec<usize> = vec![1; output_shape.len()];
|
|
for i in (0..output_shape.len() - 1).rev() {
|
|
output_strides[i] = output_strides[i + 1] * output_shape[i + 1];
|
|
}
|
|
|
|
// Copy each tensor's data to the correct position
|
|
let mut offset_in_dim = 0;
|
|
for tensor in tensors {
|
|
let tensor_shape = tensor.shape().dims().to_vec();
|
|
let tensor_data = tensor
|
|
.data()
|
|
.map_err(|e| DistributedError::tensor(e.to_string()))?;
|
|
let dim_size = tensor_shape[dim];
|
|
|
|
// Calculate tensor strides
|
|
let mut tensor_strides: Vec<usize> = vec![1; tensor_shape.len()];
|
|
for i in (0..tensor_shape.len() - 1).rev() {
|
|
tensor_strides[i] = tensor_strides[i + 1] * tensor_shape[i + 1];
|
|
}
|
|
|
|
// Copy data
|
|
let tensor_numel: usize = tensor_shape.iter().product();
|
|
for flat_idx in 0..tensor_numel {
|
|
// Convert to coordinates
|
|
let mut coords = vec![0; tensor_shape.len()];
|
|
let mut remaining = flat_idx;
|
|
for d in 0..tensor_shape.len() {
|
|
coords[d] = remaining / tensor_strides[d];
|
|
remaining %= tensor_strides[d];
|
|
}
|
|
|
|
// Adjust coordinate for concatenation dimension
|
|
coords[dim] += offset_in_dim;
|
|
|
|
// Calculate output flat index
|
|
let output_flat_idx: usize = coords
|
|
.iter()
|
|
.zip(output_strides.iter())
|
|
.map(|(c, s)| c * s)
|
|
.sum();
|
|
|
|
output_data[output_flat_idx] = tensor_data[flat_idx];
|
|
}
|
|
|
|
offset_in_dim += dim_size;
|
|
}
|
|
|
|
Tensor::from_data(output_data, output_shape, &Device::cpu())
|
|
.map_err(|e| DistributedError::tensor(e.to_string()))
|
|
}
|
|
|
|
// =============================================================================
|
|
// Replicate -> Shard: Local Slice
|
|
// =============================================================================
|
|
|
|
/// Extract local shard from replicated tensor (no communication).
|
|
fn replicate_to_shard(dtensor: &DTensor, mesh_dim: usize, tensor_dim: usize) -> Result<DTensor> {
|
|
let mesh = dtensor.mesh();
|
|
let mesh_size = mesh.shape().get(mesh_dim).copied().unwrap_or(1);
|
|
let mesh_coord = dtensor.local_coord().get(mesh_dim).copied().unwrap_or(0);
|
|
|
|
let global_shape = dtensor.global_shape();
|
|
let global_dim_size = global_shape.get(tensor_dim).copied().unwrap_or(0);
|
|
|
|
let shard_size = (global_dim_size + mesh_size - 1) / mesh_size;
|
|
let start_idx = mesh_coord * shard_size;
|
|
let end_idx = ((mesh_coord + 1) * shard_size).min(global_dim_size);
|
|
|
|
// Extract the local shard
|
|
let local_shard = slice_tensor(&dtensor.local_shard(), tensor_dim, start_idx, end_idx)?;
|
|
|
|
// Create new spec with sharded placement
|
|
let new_spec = TensorSpec::new(global_shape.to_vec())
|
|
.with_placement(mesh_dim, Placement::Shard { tensor_dim })
|
|
.with_requires_grad(dtensor.requires_grad());
|
|
|
|
DTensor::from_local(local_shard, new_spec, mesh.clone())
|
|
}
|
|
|
|
/// Slice a tensor along a dimension.
|
|
fn slice_tensor(tensor: &Tensor, dim: usize, start: usize, end: usize) -> Result<Tensor> {
|
|
let shape = tensor.shape().dims().to_vec();
|
|
if dim >= shape.len() {
|
|
return Err(DistributedError::tensor(format!(
|
|
"Slice dim {} out of range for {}D tensor",
|
|
dim,
|
|
shape.len()
|
|
)));
|
|
}
|
|
|
|
let mut output_shape = shape.clone();
|
|
output_shape[dim] = end - start;
|
|
|
|
// Calculate strides
|
|
let mut strides: Vec<usize> = vec![1; shape.len()];
|
|
for i in (0..shape.len() - 1).rev() {
|
|
strides[i] = strides[i + 1] * shape[i + 1];
|
|
}
|
|
|
|
let data = tensor
|
|
.data()
|
|
.map_err(|e| DistributedError::tensor(e.to_string()))?;
|
|
let output_numel: usize = output_shape.iter().product();
|
|
let mut output_data = vec![0.0f32; output_numel];
|
|
|
|
// Copy sliced data
|
|
let tensor_numel: usize = shape.iter().product();
|
|
let mut out_idx = 0;
|
|
|
|
for flat_idx in 0..tensor_numel {
|
|
// Convert to coordinates
|
|
let mut coords = vec![0; shape.len()];
|
|
let mut remaining = flat_idx;
|
|
for d in 0..shape.len() {
|
|
coords[d] = remaining / strides[d];
|
|
remaining %= strides[d];
|
|
}
|
|
|
|
// Check if this element is in the slice
|
|
if coords[dim] >= start && coords[dim] < end {
|
|
if out_idx < output_data.len() {
|
|
output_data[out_idx] = data[flat_idx];
|
|
out_idx += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
Tensor::from_data(output_data, output_shape, &Device::cpu())
|
|
.map_err(|e| DistributedError::tensor(e.to_string()))
|
|
}
|
|
|
|
// =============================================================================
|
|
// AllReduce: Partial -> Replicate
|
|
// =============================================================================
|
|
|
|
/// Perform AllReduce to reduce partial results across devices.
|
|
async fn all_reduce_reshard(
|
|
dtensor: &DTensor,
|
|
_mesh_dim: usize,
|
|
reduce_op: PartialReduceOp,
|
|
pg: &ProcessGroup,
|
|
) -> Result<DTensor> {
|
|
let mut local_tensor = dtensor.local_shard().clone();
|
|
|
|
let comm_reduce_op = match reduce_op {
|
|
PartialReduceOp::Sum => ReduceOp::Sum,
|
|
PartialReduceOp::Mean => ReduceOp::Sum, // Sum then divide
|
|
PartialReduceOp::Max => ReduceOp::Max,
|
|
PartialReduceOp::Min => ReduceOp::Min,
|
|
};
|
|
|
|
pg.allreduce(&mut local_tensor, comm_reduce_op).await?;
|
|
|
|
// For mean, divide by world size
|
|
if reduce_op == PartialReduceOp::Mean {
|
|
let world_size = pg.world_size() as f32;
|
|
local_tensor = local_tensor
|
|
.div_scalar(world_size)
|
|
.map_err(|e| DistributedError::tensor(e.to_string()))?;
|
|
}
|
|
|
|
// Create new DTensor with replicated placement
|
|
let new_spec = TensorSpec::new(dtensor.global_shape().to_vec())
|
|
.with_placement(0, Placement::Replicate)
|
|
.with_requires_grad(dtensor.requires_grad());
|
|
|
|
DTensor::from_local(local_tensor, new_spec, dtensor.mesh().clone())
|
|
}
|
|
|
|
// =============================================================================
|
|
// ReduceScatter: Partial -> Shard
|
|
// =============================================================================
|
|
|
|
/// Perform ReduceScatter to reduce and shard in one operation.
|
|
async fn reduce_scatter_reshard(
|
|
dtensor: &DTensor,
|
|
mesh_dim: usize,
|
|
tensor_dim: usize,
|
|
reduce_op: PartialReduceOp,
|
|
pg: &ProcessGroup,
|
|
) -> Result<DTensor> {
|
|
let local_tensor = dtensor.local_shard().clone();
|
|
|
|
let comm_reduce_op = match reduce_op {
|
|
PartialReduceOp::Sum => ReduceOp::Sum,
|
|
PartialReduceOp::Mean => ReduceOp::Sum,
|
|
PartialReduceOp::Max => ReduceOp::Max,
|
|
PartialReduceOp::Min => ReduceOp::Min,
|
|
};
|
|
|
|
let scattered: Tensor = pg.reduce_scatter(&local_tensor, comm_reduce_op)?;
|
|
|
|
// For mean, divide by world size
|
|
let result = if reduce_op == PartialReduceOp::Mean {
|
|
let world_size = pg.world_size() as f32;
|
|
scattered
|
|
.div_scalar(world_size)
|
|
.map_err(|e| DistributedError::tensor(e.to_string()))?
|
|
} else {
|
|
scattered
|
|
};
|
|
|
|
// Create new DTensor with sharded placement
|
|
let new_spec = TensorSpec::new(dtensor.global_shape().to_vec())
|
|
.with_placement(mesh_dim, Placement::Shard { tensor_dim })
|
|
.with_requires_grad(dtensor.requires_grad());
|
|
|
|
DTensor::from_local(result, new_spec, dtensor.mesh().clone())
|
|
}
|
|
|
|
// =============================================================================
|
|
// AllToAll: Shard(dim_a) -> Shard(dim_b)
|
|
// =============================================================================
|
|
|
|
/// Perform AllToAll to change sharding dimension.
|
|
async fn all_to_all_reshard(
|
|
dtensor: &DTensor,
|
|
mesh_dim: usize,
|
|
_src_dim: usize,
|
|
dst_dim: usize,
|
|
pg: &ProcessGroup,
|
|
) -> Result<DTensor> {
|
|
// AllToAll is more complex - for now, we use AllGather + local slice
|
|
// A more efficient implementation would use actual AllToAll communication
|
|
|
|
// First, gather to replicated
|
|
let gathered_spec = TensorSpec::new(dtensor.global_shape().to_vec())
|
|
.with_placement(0, Placement::Replicate)
|
|
.with_requires_grad(dtensor.requires_grad());
|
|
|
|
let local_shard = dtensor.local_shard().clone();
|
|
let gathered = pg.allgather(&local_shard).await?;
|
|
|
|
let full_tensor = match gathered {
|
|
crate::comm::AllGatherOutput::Tensor(t) => t,
|
|
crate::comm::AllGatherOutput::TensorList(tensors) => concat_along_dim(&tensors, 0)?,
|
|
};
|
|
|
|
let gathered_dtensor = DTensor::from_local(full_tensor, gathered_spec, dtensor.mesh().clone())?;
|
|
|
|
// Then slice for new sharding dimension
|
|
replicate_to_shard(&gathered_dtensor, mesh_dim, dst_dim)
|
|
}
|
|
|
|
// =============================================================================
|
|
// Utility Functions
|
|
// =============================================================================
|
|
|
|
/// Check if two DTensors have compatible placements for element-wise operations.
|
|
pub fn placements_compatible(a: &DTensor, b: &DTensor) -> bool {
|
|
a.spec().placements == b.spec().placements
|
|
}
|
|
|
|
/// Redistribute DTensor to match another DTensor's placement.
|
|
pub async fn redistribute_like(
|
|
dtensor: &DTensor,
|
|
target: &DTensor,
|
|
pg: &ProcessGroup,
|
|
) -> Result<DTensor> {
|
|
reshard_dtensor(dtensor, &target.spec().placements, pg).await
|
|
}
|
|
|
|
/// Create a replicated DTensor from a local tensor.
|
|
pub fn replicate_tensor(tensor: Tensor, mesh: Arc<DeviceMesh>) -> Result<DTensor> {
|
|
let shape: Vec<usize> = tensor.shape().dims().to_vec();
|
|
let spec = TensorSpec::new(shape).with_placement(0, Placement::Replicate);
|
|
DTensor::from_local(tensor, spec, mesh)
|
|
}
|
|
|
|
/// Create a sharded DTensor from a local shard.
|
|
pub fn shard_tensor(
|
|
local_shard: Tensor,
|
|
global_shape: Vec<usize>,
|
|
mesh_dim: usize,
|
|
tensor_dim: usize,
|
|
mesh: Arc<DeviceMesh>,
|
|
) -> Result<DTensor> {
|
|
let spec =
|
|
TensorSpec::new(global_shape).with_placement(mesh_dim, Placement::Shard { tensor_dim });
|
|
DTensor::from_local(local_shard, spec, mesh)
|
|
}
|
|
|
|
// =============================================================================
|
|
// Tests
|
|
// =============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_concat_along_dim0() {
|
|
let t1 = Tensor::from_data(vec![1.0, 2.0], vec![2], &Device::cpu()).unwrap();
|
|
let t2 = Tensor::from_data(vec![3.0, 4.0], vec![2], &Device::cpu()).unwrap();
|
|
|
|
let result = concat_along_dim(&[t1, t2], 0).unwrap();
|
|
let data = result.data().unwrap();
|
|
|
|
assert_eq!(data, vec![1.0, 2.0, 3.0, 4.0]);
|
|
assert_eq!(result.shape().dims(), &[4]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_concat_along_dim1() {
|
|
let t1 = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2], &Device::cpu()).unwrap();
|
|
let t2 = Tensor::from_data(vec![5.0, 6.0, 7.0, 8.0], vec![2, 2], &Device::cpu()).unwrap();
|
|
|
|
let result = concat_along_dim(&[t1, t2], 1).unwrap();
|
|
let data = result.data().unwrap();
|
|
|
|
// [1, 2] concat [5, 6] along dim 1 = [1, 2, 5, 6]
|
|
// [3, 4] concat [7, 8] along dim 1 = [3, 4, 7, 8]
|
|
assert_eq!(data, vec![1.0, 2.0, 5.0, 6.0, 3.0, 4.0, 7.0, 8.0]);
|
|
assert_eq!(result.shape().dims(), &[2, 4]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_slice_tensor_dim0() {
|
|
let tensor = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], vec![4], &Device::cpu()).unwrap();
|
|
let sliced = slice_tensor(&tensor, 0, 1, 3).unwrap();
|
|
|
|
assert_eq!(sliced.data().unwrap(), vec![2.0, 3.0]);
|
|
assert_eq!(sliced.shape().dims(), &[2]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_slice_tensor_dim1() {
|
|
let tensor = Tensor::from_data(
|
|
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
|
|
vec![2, 3],
|
|
&Device::cpu(),
|
|
)
|
|
.unwrap();
|
|
|
|
let sliced = slice_tensor(&tensor, 1, 0, 2).unwrap();
|
|
|
|
assert_eq!(sliced.data().unwrap(), vec![1.0, 2.0, 4.0, 5.0]);
|
|
assert_eq!(sliced.shape().dims(), &[2, 2]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_replicate_tensor() {
|
|
let mesh = Arc::new(DeviceMesh::new_simple(2, "dp"));
|
|
let tensor = Tensor::from_data(vec![1.0, 2.0, 3.0], vec![3], &Device::cpu()).unwrap();
|
|
|
|
let dtensor = replicate_tensor(tensor, mesh).unwrap();
|
|
|
|
assert!(!dtensor.spec().is_sharded());
|
|
assert_eq!(dtensor.global_shape(), &[3]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shard_tensor() {
|
|
let mesh = Arc::new(DeviceMesh::new_simple(2, "dp"));
|
|
let local_shard = Tensor::from_data(vec![1.0, 2.0], vec![2], &Device::cpu()).unwrap();
|
|
|
|
let dtensor = shard_tensor(local_shard, vec![4], 0, 0, mesh).unwrap();
|
|
|
|
assert!(dtensor.spec().is_sharded());
|
|
assert_eq!(dtensor.global_shape(), &[4]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_placements_compatible() {
|
|
let mesh = Arc::new(DeviceMesh::new_simple(2, "dp"));
|
|
|
|
let t1 = Tensor::from_data(vec![1.0, 2.0], vec![2], &Device::cpu()).unwrap();
|
|
let t2 = Tensor::from_data(vec![3.0, 4.0], vec![2], &Device::cpu()).unwrap();
|
|
|
|
let d1 = replicate_tensor(t1, mesh.clone()).unwrap();
|
|
let d2 = replicate_tensor(t2, mesh).unwrap();
|
|
|
|
assert!(placements_compatible(&d1, &d2));
|
|
}
|
|
}
|