707 lines
22 KiB
Rust
707 lines
22 KiB
Rust
//! Device Mesh abstraction for multi-dimensional device organization.
|
|
//!
|
|
//! A DeviceMesh organizes devices (GPUs) into a logical multi-dimensional grid,
|
|
//! enabling hybrid parallelism strategies. Each dimension of the mesh can be
|
|
//! used for a different parallelism strategy (data, tensor, pipeline, etc.).
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_distributed::device_mesh::DeviceMesh;
|
|
//!
|
|
//! // Create a 2D mesh: 2 nodes x 4 GPUs = 8 total
|
|
//! // Dimension 0 ("dp") for data parallelism
|
|
//! // Dimension 1 ("tp") for tensor parallelism
|
|
//! let mesh = DeviceMesh::new(
|
|
//! vec![2, 4],
|
|
//! vec!["dp".to_string(), "tp".to_string()],
|
|
//! )?;
|
|
//!
|
|
//! // Get the process group for tensor parallelism
|
|
//! let tp_group = mesh.get_group("tp")?;
|
|
//! ```
|
|
|
|
use crate::error::{DistributedError, Result};
|
|
use crate::group::ProcessGroup;
|
|
use parking_lot::RwLock;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
|
// =============================================================================
|
|
// Mesh Dimension
|
|
// =============================================================================
|
|
|
|
/// Named dimension of the device mesh.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct MeshDimension {
|
|
/// Name of this dimension (e.g., "dp", "tp", "pp")
|
|
pub name: String,
|
|
/// Size of this dimension (number of devices)
|
|
pub size: usize,
|
|
/// Index in the mesh shape
|
|
pub index: usize,
|
|
}
|
|
|
|
impl MeshDimension {
|
|
/// Create a new mesh dimension.
|
|
pub fn new(name: impl Into<String>, size: usize, index: usize) -> Self {
|
|
Self {
|
|
name: name.into(),
|
|
size,
|
|
index,
|
|
}
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Device Mesh
|
|
// =============================================================================
|
|
|
|
/// Multi-dimensional device mesh for organizing distributed computation.
|
|
///
|
|
/// The mesh provides a logical view of devices and enables creating sub-groups
|
|
/// for different parallelism dimensions.
|
|
#[derive(Debug)]
|
|
pub struct DeviceMesh {
|
|
/// Shape of the mesh (e.g., [2, 4] for 2x4 mesh)
|
|
shape: Vec<usize>,
|
|
|
|
/// Named dimensions
|
|
dimensions: Vec<MeshDimension>,
|
|
|
|
/// Name to dimension index mapping
|
|
name_to_dim: HashMap<String, usize>,
|
|
|
|
/// Total number of devices
|
|
world_size: usize,
|
|
|
|
/// Local rank in the global mesh
|
|
local_rank: AtomicUsize,
|
|
|
|
/// Device list (flattened)
|
|
devices: Vec<DeviceInfo>,
|
|
|
|
/// Process groups for each dimension (lazily created)
|
|
dimension_groups: RwLock<HashMap<usize, Arc<ProcessGroup>>>,
|
|
|
|
/// Optional parent mesh (for submesh creation)
|
|
parent: Option<Arc<DeviceMesh>>,
|
|
}
|
|
|
|
/// Information about a device in the mesh.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct DeviceInfo {
|
|
/// Global rank of this device
|
|
pub global_rank: usize,
|
|
/// Coordinate in the mesh
|
|
pub coordinate: Vec<usize>,
|
|
/// Device type (e.g., "cuda:0", "cpu")
|
|
pub device_type: String,
|
|
/// Node this device belongs to
|
|
pub node_id: usize,
|
|
/// Local device index on the node
|
|
pub local_device_id: usize,
|
|
}
|
|
|
|
impl DeviceMesh {
|
|
/// Create a new device mesh with the given shape and dimension names.
|
|
///
|
|
/// # Arguments
|
|
/// * `shape` - Shape of the mesh (e.g., [2, 4] for 2x4)
|
|
/// * `dim_names` - Names for each dimension (e.g., ["dp", "tp"])
|
|
///
|
|
/// # Returns
|
|
/// A new DeviceMesh
|
|
pub fn new(shape: Vec<usize>, dim_names: Vec<impl Into<String>>) -> Result<Self> {
|
|
if shape.len() != dim_names.len() {
|
|
return Err(DistributedError::configuration(format!(
|
|
"Shape length {} must match dim_names length {}",
|
|
shape.len(),
|
|
dim_names.len()
|
|
)));
|
|
}
|
|
|
|
let world_size: usize = shape.iter().product();
|
|
if world_size == 0 {
|
|
return Err(DistributedError::configuration(
|
|
"Mesh cannot have zero devices",
|
|
));
|
|
}
|
|
|
|
let dimensions: Vec<MeshDimension> = dim_names
|
|
.into_iter()
|
|
.enumerate()
|
|
.map(|(i, name)| MeshDimension::new(name, shape[i], i))
|
|
.collect();
|
|
|
|
let name_to_dim: HashMap<String, usize> = dimensions
|
|
.iter()
|
|
.map(|d| (d.name.clone(), d.index))
|
|
.collect();
|
|
|
|
// Create device info for all devices
|
|
let devices = Self::create_device_list(&shape, world_size);
|
|
|
|
Ok(Self {
|
|
shape,
|
|
dimensions,
|
|
name_to_dim,
|
|
world_size,
|
|
local_rank: AtomicUsize::new(0),
|
|
devices,
|
|
dimension_groups: RwLock::new(HashMap::new()),
|
|
parent: None,
|
|
})
|
|
}
|
|
|
|
/// Create a simple 1D mesh.
|
|
pub fn new_simple(size: usize, dim_name: impl Into<String>) -> Self {
|
|
Self::new(vec![size], vec![dim_name]).expect("Simple mesh creation should not fail")
|
|
}
|
|
|
|
/// Create device list from shape.
|
|
fn create_device_list(shape: &[usize], world_size: usize) -> Vec<DeviceInfo> {
|
|
let mut devices = Vec::with_capacity(world_size);
|
|
|
|
for global_rank in 0..world_size {
|
|
let coordinate = Self::rank_to_coordinate_static(global_rank, shape);
|
|
devices.push(DeviceInfo {
|
|
global_rank,
|
|
coordinate,
|
|
device_type: format!("cuda:{}", global_rank % 8), // Assume 8 GPUs per node
|
|
node_id: global_rank / 8,
|
|
local_device_id: global_rank % 8,
|
|
});
|
|
}
|
|
|
|
devices
|
|
}
|
|
|
|
/// Convert global rank to mesh coordinate.
|
|
fn rank_to_coordinate_static(rank: usize, shape: &[usize]) -> Vec<usize> {
|
|
let mut coord = vec![0; shape.len()];
|
|
let mut remaining = rank;
|
|
|
|
for i in (0..shape.len()).rev() {
|
|
coord[i] = remaining % shape[i];
|
|
remaining /= shape[i];
|
|
}
|
|
|
|
coord
|
|
}
|
|
|
|
/// Get the shape of the mesh.
|
|
pub fn shape(&self) -> &[usize] {
|
|
&self.shape
|
|
}
|
|
|
|
/// Get the number of dimensions.
|
|
pub fn ndim(&self) -> usize {
|
|
self.shape.len()
|
|
}
|
|
|
|
/// Get total number of devices.
|
|
pub fn world_size(&self) -> usize {
|
|
self.world_size
|
|
}
|
|
|
|
/// Get the local rank.
|
|
pub fn local_rank(&self) -> usize {
|
|
self.local_rank.load(Ordering::SeqCst)
|
|
}
|
|
|
|
/// Set the local rank.
|
|
pub fn set_local_rank(&self, rank: usize) {
|
|
self.local_rank.store(rank, Ordering::SeqCst);
|
|
}
|
|
|
|
/// Get dimension by name.
|
|
pub fn get_dimension(&self, name: &str) -> Option<&MeshDimension> {
|
|
self.name_to_dim
|
|
.get(name)
|
|
.and_then(|&idx| self.dimensions.get(idx))
|
|
}
|
|
|
|
/// Get dimension by index.
|
|
pub fn get_dimension_by_index(&self, index: usize) -> Option<&MeshDimension> {
|
|
self.dimensions.get(index)
|
|
}
|
|
|
|
/// Get the size of a dimension by name.
|
|
pub fn dim_size(&self, name: &str) -> Option<usize> {
|
|
self.get_dimension(name).map(|d| d.size)
|
|
}
|
|
|
|
/// Convert global rank to mesh coordinate.
|
|
pub fn rank_to_coordinate(&self, rank: usize) -> Vec<usize> {
|
|
Self::rank_to_coordinate_static(rank, &self.shape)
|
|
}
|
|
|
|
/// Convert mesh coordinate to global rank.
|
|
pub fn coordinate_to_rank(&self, coord: &[usize]) -> Result<usize> {
|
|
if coord.len() != self.shape.len() {
|
|
return Err(DistributedError::configuration(format!(
|
|
"Coordinate length {} must match mesh dimensions {}",
|
|
coord.len(),
|
|
self.shape.len()
|
|
)));
|
|
}
|
|
|
|
for (i, (&c, &s)) in coord.iter().zip(self.shape.iter()).enumerate() {
|
|
if c >= s {
|
|
return Err(DistributedError::configuration(format!(
|
|
"Coordinate[{}] = {} exceeds dimension size {}",
|
|
i, c, s
|
|
)));
|
|
}
|
|
}
|
|
|
|
let mut rank = 0;
|
|
let mut multiplier = 1;
|
|
|
|
for i in (0..self.shape.len()).rev() {
|
|
rank += coord[i] * multiplier;
|
|
multiplier *= self.shape[i];
|
|
}
|
|
|
|
Ok(rank)
|
|
}
|
|
|
|
/// Get the local coordinate of this process in the mesh.
|
|
pub fn local_coordinate(&self) -> Vec<usize> {
|
|
self.rank_to_coordinate(self.local_rank())
|
|
}
|
|
|
|
/// Get device info for a rank.
|
|
pub fn get_device(&self, rank: usize) -> Option<&DeviceInfo> {
|
|
self.devices.get(rank)
|
|
}
|
|
|
|
/// Get all devices.
|
|
pub fn devices(&self) -> &[DeviceInfo] {
|
|
&self.devices
|
|
}
|
|
|
|
// =========================================================================
|
|
// Process Group Operations
|
|
// =========================================================================
|
|
|
|
/// Get ranks along a mesh dimension for the current coordinate.
|
|
///
|
|
/// Returns all ranks that share the same coordinates on all other dimensions.
|
|
pub fn get_ranks_along_dim(&self, dim: usize) -> Result<Vec<usize>> {
|
|
if dim >= self.ndim() {
|
|
return Err(DistributedError::configuration(format!(
|
|
"Dimension {} out of range for {}D mesh",
|
|
dim,
|
|
self.ndim()
|
|
)));
|
|
}
|
|
|
|
let local_coord = self.local_coordinate();
|
|
let mut ranks = Vec::with_capacity(self.shape[dim]);
|
|
|
|
for i in 0..self.shape[dim] {
|
|
let mut coord = local_coord.clone();
|
|
coord[dim] = i;
|
|
ranks.push(self.coordinate_to_rank(&coord)?);
|
|
}
|
|
|
|
Ok(ranks)
|
|
}
|
|
|
|
/// Get ranks along a mesh dimension by name.
|
|
pub fn get_ranks_along_dim_by_name(&self, name: &str) -> Result<Vec<usize>> {
|
|
let dim = self.name_to_dim.get(name).ok_or_else(|| {
|
|
DistributedError::configuration(format!("Unknown dimension: {}", name))
|
|
})?;
|
|
self.get_ranks_along_dim(*dim)
|
|
}
|
|
|
|
/// Get the rank of this process within a dimension.
|
|
pub fn get_rank_in_dim(&self, dim: usize) -> usize {
|
|
self.local_coordinate().get(dim).copied().unwrap_or(0)
|
|
}
|
|
|
|
/// Get the rank of this process within a dimension by name.
|
|
pub fn get_rank_in_dim_by_name(&self, name: &str) -> Option<usize> {
|
|
self.name_to_dim
|
|
.get(name)
|
|
.map(|&dim| self.get_rank_in_dim(dim))
|
|
}
|
|
|
|
/// Get or create a process group for a dimension.
|
|
///
|
|
/// Creates a ProcessGroup containing all ranks along the specified dimension
|
|
/// that share the same coordinates on all other dimensions. The group uses
|
|
/// RNCCL backend when the feature is enabled.
|
|
pub fn get_process_group(&self, dim: usize) -> Result<Arc<ProcessGroup>> {
|
|
// Check cache first
|
|
{
|
|
let groups = self.dimension_groups.read();
|
|
if let Some(pg) = groups.get(&dim) {
|
|
return Ok(pg.clone());
|
|
}
|
|
}
|
|
|
|
// Create new process group
|
|
let ranks = self.get_ranks_along_dim(dim)?;
|
|
let local_rank = self.local_rank();
|
|
|
|
// Find our position in the dimension
|
|
let rank_in_dim = ranks.iter().position(|&r| r == local_rank).unwrap_or(0);
|
|
let dim_size = ranks.len();
|
|
|
|
// Create process group with RNCCL backend when available
|
|
#[cfg(feature = "rnccl")]
|
|
let pg = {
|
|
use crate::backend::Backend;
|
|
let world_info =
|
|
crate::group::WorldInfo::new(dim_size as i32, rank_in_dim as i32, Backend::Rnccl);
|
|
ProcessGroup::new(Backend::Rnccl, world_info)?
|
|
};
|
|
|
|
#[cfg(not(feature = "rnccl"))]
|
|
let pg = {
|
|
use crate::backend::Backend;
|
|
let world_info =
|
|
crate::group::WorldInfo::new(dim_size as i32, rank_in_dim as i32, Backend::Cpu);
|
|
ProcessGroup::new(Backend::Cpu, world_info)?
|
|
};
|
|
|
|
let pg = Arc::new(pg);
|
|
|
|
// Cache and return
|
|
{
|
|
let mut groups = self.dimension_groups.write();
|
|
groups.insert(dim, pg.clone());
|
|
}
|
|
|
|
Ok(pg)
|
|
}
|
|
|
|
/// Get or create a process group for a dimension by name.
|
|
pub fn get_process_group_by_name(&self, name: &str) -> Result<Arc<ProcessGroup>> {
|
|
let dim = self.name_to_dim.get(name).ok_or_else(|| {
|
|
DistributedError::configuration(format!("Unknown dimension: {}", name))
|
|
})?;
|
|
self.get_process_group(*dim)
|
|
}
|
|
|
|
// =========================================================================
|
|
// Submesh Operations
|
|
// =========================================================================
|
|
|
|
/// Create a submesh by slicing along specified dimensions.
|
|
///
|
|
/// # Arguments
|
|
/// * `dim_indices` - Which dimensions to include in the submesh
|
|
///
|
|
/// # Returns
|
|
/// A new DeviceMesh representing the submesh
|
|
pub fn submesh(&self, dim_names: &[&str]) -> Result<Arc<DeviceMesh>> {
|
|
let dim_indices: Result<Vec<usize>> = dim_names
|
|
.iter()
|
|
.map(|name| {
|
|
self.name_to_dim.get(*name).copied().ok_or_else(|| {
|
|
DistributedError::configuration(format!("Unknown dimension: {}", name))
|
|
})
|
|
})
|
|
.collect();
|
|
let dim_indices = dim_indices?;
|
|
|
|
let new_shape: Vec<usize> = dim_indices.iter().map(|&i| self.shape[i]).collect();
|
|
let new_dim_names: Vec<String> = dim_indices
|
|
.iter()
|
|
.map(|&i| self.dimensions[i].name.clone())
|
|
.collect();
|
|
|
|
let submesh = DeviceMesh {
|
|
shape: new_shape.clone(),
|
|
dimensions: new_dim_names
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, name)| MeshDimension::new(name.clone(), new_shape[i], i))
|
|
.collect(),
|
|
name_to_dim: new_dim_names
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, name)| (name.clone(), i))
|
|
.collect(),
|
|
world_size: new_shape.iter().product(),
|
|
local_rank: AtomicUsize::new(self.local_rank()),
|
|
devices: Vec::new(), // Submesh devices would be computed
|
|
dimension_groups: RwLock::new(HashMap::new()),
|
|
parent: Some(Arc::new(DeviceMesh {
|
|
shape: self.shape.clone(),
|
|
dimensions: self.dimensions.clone(),
|
|
name_to_dim: self.name_to_dim.clone(),
|
|
world_size: self.world_size,
|
|
local_rank: AtomicUsize::new(self.local_rank()),
|
|
devices: self.devices.clone(),
|
|
dimension_groups: RwLock::new(HashMap::new()),
|
|
parent: None,
|
|
})),
|
|
};
|
|
|
|
Ok(Arc::new(submesh))
|
|
}
|
|
|
|
/// Flatten the mesh to 1D.
|
|
pub fn flatten(&self) -> Arc<DeviceMesh> {
|
|
Arc::new(DeviceMesh::new_simple(self.world_size, "flat"))
|
|
}
|
|
|
|
// =========================================================================
|
|
// Collective Operations
|
|
// =========================================================================
|
|
|
|
/// Check if two ranks are neighbors along a dimension.
|
|
pub fn are_neighbors(&self, rank1: usize, rank2: usize, dim: usize) -> bool {
|
|
let coord1 = self.rank_to_coordinate(rank1);
|
|
let coord2 = self.rank_to_coordinate(rank2);
|
|
|
|
// Must differ only in the specified dimension
|
|
for (i, (&c1, &c2)) in coord1.iter().zip(coord2.iter()).enumerate() {
|
|
if i == dim {
|
|
if (c1 as isize - c2 as isize).abs() != 1 {
|
|
return false;
|
|
}
|
|
} else if c1 != c2 {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
true
|
|
}
|
|
|
|
/// Get neighbor ranks along a dimension.
|
|
pub fn get_neighbors(&self, rank: usize, dim: usize) -> (Option<usize>, Option<usize>) {
|
|
let coord = self.rank_to_coordinate(rank);
|
|
let pos = coord[dim];
|
|
|
|
let prev = if pos > 0 {
|
|
let mut prev_coord = coord.clone();
|
|
prev_coord[dim] = pos - 1;
|
|
self.coordinate_to_rank(&prev_coord).ok()
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let next = if pos + 1 < self.shape[dim] {
|
|
let mut next_coord = coord.clone();
|
|
next_coord[dim] = pos + 1;
|
|
self.coordinate_to_rank(&next_coord).ok()
|
|
} else {
|
|
None
|
|
};
|
|
|
|
(prev, next)
|
|
}
|
|
|
|
/// Get the distance between two ranks (number of hops).
|
|
pub fn distance(&self, rank1: usize, rank2: usize) -> usize {
|
|
let coord1 = self.rank_to_coordinate(rank1);
|
|
let coord2 = self.rank_to_coordinate(rank2);
|
|
|
|
coord1
|
|
.iter()
|
|
.zip(coord2.iter())
|
|
.map(|(a, b)| (*a as isize - *b as isize).unsigned_abs())
|
|
.sum()
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Mesh Builder
|
|
// =============================================================================
|
|
|
|
/// Builder for creating device meshes with fluent API.
|
|
#[derive(Default)]
|
|
pub struct DeviceMeshBuilder {
|
|
shape: Vec<usize>,
|
|
dim_names: Vec<String>,
|
|
local_rank: Option<usize>,
|
|
}
|
|
|
|
impl DeviceMeshBuilder {
|
|
/// Create a new builder.
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Add a dimension to the mesh.
|
|
pub fn add_dim(mut self, name: impl Into<String>, size: usize) -> Self {
|
|
self.dim_names.push(name.into());
|
|
self.shape.push(size);
|
|
self
|
|
}
|
|
|
|
/// Set the local rank.
|
|
pub fn with_local_rank(mut self, rank: usize) -> Self {
|
|
self.local_rank = Some(rank);
|
|
self
|
|
}
|
|
|
|
/// Build the device mesh.
|
|
pub fn build(self) -> Result<DeviceMesh> {
|
|
let mesh = DeviceMesh::new(self.shape, self.dim_names)?;
|
|
if let Some(rank) = self.local_rank {
|
|
mesh.set_local_rank(rank);
|
|
}
|
|
Ok(mesh)
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Common Mesh Patterns
|
|
// =============================================================================
|
|
|
|
impl DeviceMesh {
|
|
/// Create a 2D mesh for data + tensor parallelism.
|
|
pub fn data_tensor_parallel(dp_size: usize, tp_size: usize) -> Result<Self> {
|
|
DeviceMeshBuilder::new()
|
|
.add_dim("dp", dp_size)
|
|
.add_dim("tp", tp_size)
|
|
.build()
|
|
}
|
|
|
|
/// Create a 3D mesh for data + tensor + pipeline parallelism.
|
|
pub fn parallelism_3d(dp_size: usize, tp_size: usize, pp_size: usize) -> Result<Self> {
|
|
DeviceMeshBuilder::new()
|
|
.add_dim("dp", dp_size)
|
|
.add_dim("tp", tp_size)
|
|
.add_dim("pp", pp_size)
|
|
.build()
|
|
}
|
|
|
|
/// Create a mesh for expert parallelism (MoE).
|
|
pub fn expert_parallel(dp_size: usize, ep_size: usize, tp_size: usize) -> Result<Self> {
|
|
DeviceMeshBuilder::new()
|
|
.add_dim("dp", dp_size)
|
|
.add_dim("ep", ep_size)
|
|
.add_dim("tp", tp_size)
|
|
.build()
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Tests
|
|
// =============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_mesh_creation() {
|
|
let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();
|
|
assert_eq!(mesh.shape(), &[2, 4]);
|
|
assert_eq!(mesh.world_size(), 8);
|
|
assert_eq!(mesh.ndim(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_simple_mesh() {
|
|
let mesh = DeviceMesh::new_simple(4, "dp");
|
|
assert_eq!(mesh.shape(), &[4]);
|
|
assert_eq!(mesh.world_size(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_rank_to_coordinate() {
|
|
let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();
|
|
|
|
assert_eq!(mesh.rank_to_coordinate(0), vec![0, 0]);
|
|
assert_eq!(mesh.rank_to_coordinate(1), vec![0, 1]);
|
|
assert_eq!(mesh.rank_to_coordinate(4), vec![1, 0]);
|
|
assert_eq!(mesh.rank_to_coordinate(7), vec![1, 3]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_coordinate_to_rank() {
|
|
let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();
|
|
|
|
assert_eq!(mesh.coordinate_to_rank(&[0, 0]).unwrap(), 0);
|
|
assert_eq!(mesh.coordinate_to_rank(&[0, 1]).unwrap(), 1);
|
|
assert_eq!(mesh.coordinate_to_rank(&[1, 0]).unwrap(), 4);
|
|
assert_eq!(mesh.coordinate_to_rank(&[1, 3]).unwrap(), 7);
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_ranks_along_dim() {
|
|
let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();
|
|
mesh.set_local_rank(5); // Coordinate [1, 1]
|
|
|
|
// Along dp dimension (dim 0): ranks with tp=1
|
|
let dp_ranks = mesh.get_ranks_along_dim(0).unwrap();
|
|
assert_eq!(dp_ranks, vec![1, 5]);
|
|
|
|
// Along tp dimension (dim 1): ranks with dp=1
|
|
let tp_ranks = mesh.get_ranks_along_dim(1).unwrap();
|
|
assert_eq!(tp_ranks, vec![4, 5, 6, 7]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dim_size_by_name() {
|
|
let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();
|
|
assert_eq!(mesh.dim_size("dp"), Some(2));
|
|
assert_eq!(mesh.dim_size("tp"), Some(4));
|
|
assert_eq!(mesh.dim_size("pp"), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_neighbors() {
|
|
let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();
|
|
|
|
// Rank 5 is at [1, 1]
|
|
let (prev, next) = mesh.get_neighbors(5, 1); // along tp
|
|
assert_eq!(prev, Some(4)); // [1, 0]
|
|
assert_eq!(next, Some(6)); // [1, 2]
|
|
|
|
// Edge cases
|
|
let (prev, next) = mesh.get_neighbors(4, 1); // [1, 0]
|
|
assert_eq!(prev, None);
|
|
assert_eq!(next, Some(5));
|
|
|
|
let (prev, next) = mesh.get_neighbors(7, 1); // [1, 3]
|
|
assert_eq!(prev, Some(6));
|
|
assert_eq!(next, None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_distance() {
|
|
let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();
|
|
|
|
assert_eq!(mesh.distance(0, 0), 0);
|
|
assert_eq!(mesh.distance(0, 1), 1); // [0,0] to [0,1]
|
|
assert_eq!(mesh.distance(0, 4), 1); // [0,0] to [1,0]
|
|
assert_eq!(mesh.distance(0, 7), 4); // [0,0] to [1,3]
|
|
}
|
|
|
|
#[test]
|
|
fn test_builder() {
|
|
let mesh = DeviceMeshBuilder::new()
|
|
.add_dim("dp", 2)
|
|
.add_dim("tp", 4)
|
|
.with_local_rank(5)
|
|
.build()
|
|
.unwrap();
|
|
|
|
assert_eq!(mesh.shape(), &[2, 4]);
|
|
assert_eq!(mesh.local_rank(), 5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_3d_parallelism() {
|
|
let mesh = DeviceMesh::parallelism_3d(2, 2, 2).unwrap();
|
|
assert_eq!(mesh.shape(), &[2, 2, 2]);
|
|
assert_eq!(mesh.world_size(), 8);
|
|
assert!(mesh.get_dimension("dp").is_some());
|
|
assert!(mesh.get_dimension("tp").is_some());
|
|
assert!(mesh.get_dimension("pp").is_some());
|
|
}
|
|
}
|