ferrum_engine/parallel/
config.rs

1//! Parallel Configuration
2//!
3//! Configuration types for multi-GPU parallelism.
4
5use ferrum_types::Device;
6use serde::{Deserialize, Serialize};
7
8/// Type of parallelism strategy
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
10pub enum ParallelismType {
11    /// No parallelism (single GPU)
12    #[default]
13    None,
14    /// Tensor parallelism: split tensors across GPUs
15    Tensor,
16    /// Pipeline parallelism: split layers across GPUs
17    Pipeline,
18    /// Data parallelism: replicate model, split batches
19    Data,
20    /// Hybrid: combination of tensor and pipeline parallelism
21    Hybrid,
22}
23
24/// Parallel execution configuration
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct ParallelConfig {
27    /// Type of parallelism to use
28    pub parallelism_type: ParallelismType,
29    /// Devices to use for parallel execution
30    pub devices: Vec<Device>,
31    /// Number of tensor parallel ranks
32    pub tensor_parallel_size: usize,
33    /// Number of pipeline parallel stages
34    pub pipeline_parallel_size: usize,
35    /// Whether to enable memory optimization
36    pub enable_memory_optimization: bool,
37    /// Communication backend (nccl, gloo, etc.)
38    pub communication_backend: String,
39    /// Maximum chunk size for all-reduce operations
40    pub max_chunk_size: usize,
41    /// Enable overlapping communication with computation
42    pub overlap_communication: bool,
43}
44
45impl Default for ParallelConfig {
46    fn default() -> Self {
47        Self {
48            parallelism_type: ParallelismType::None,
49            devices: vec![Device::CPU],
50            tensor_parallel_size: 1,
51            pipeline_parallel_size: 1,
52            enable_memory_optimization: true,
53            communication_backend: "cpu".to_string(),
54            max_chunk_size: 1024 * 1024, // 1MB
55            overlap_communication: true,
56        }
57    }
58}
59
60impl ParallelConfig {
61    /// Create config for single GPU
62    pub fn single_gpu(device: Device) -> Self {
63        Self {
64            parallelism_type: ParallelismType::None,
65            devices: vec![device],
66            ..Default::default()
67        }
68    }
69
70    /// Create config for tensor parallelism
71    pub fn tensor_parallel(devices: Vec<Device>) -> Self {
72        let size = devices.len();
73        Self {
74            parallelism_type: ParallelismType::Tensor,
75            devices,
76            tensor_parallel_size: size,
77            pipeline_parallel_size: 1,
78            ..Default::default()
79        }
80    }
81
82    /// Create config for pipeline parallelism
83    pub fn pipeline_parallel(devices: Vec<Device>) -> Self {
84        let size = devices.len();
85        Self {
86            parallelism_type: ParallelismType::Pipeline,
87            devices,
88            tensor_parallel_size: 1,
89            pipeline_parallel_size: size,
90            ..Default::default()
91        }
92    }
93
94    /// Create config for data parallelism
95    pub fn data_parallel(devices: Vec<Device>) -> Self {
96        Self {
97            parallelism_type: ParallelismType::Data,
98            devices,
99            tensor_parallel_size: 1,
100            pipeline_parallel_size: 1,
101            ..Default::default()
102        }
103    }
104
105    /// Create hybrid config (tensor + pipeline)
106    pub fn hybrid(devices: Vec<Device>, tp_size: usize, pp_size: usize) -> Self {
107        Self {
108            parallelism_type: ParallelismType::Hybrid,
109            devices,
110            tensor_parallel_size: tp_size,
111            pipeline_parallel_size: pp_size,
112            ..Default::default()
113        }
114    }
115
116    /// Get the world size (total number of ranks)
117    pub fn world_size(&self) -> usize {
118        self.tensor_parallel_size * self.pipeline_parallel_size
119    }
120
121    /// Check if parallelism is enabled
122    pub fn is_parallel(&self) -> bool {
123        self.parallelism_type != ParallelismType::None && self.world_size() > 1
124    }
125
126    /// Get device for a specific rank
127    pub fn device_for_rank(&self, rank: usize) -> Option<&Device> {
128        self.devices.get(rank)
129    }
130
131    /// Calculate tensor parallel rank from global rank
132    pub fn tp_rank(&self, global_rank: usize) -> usize {
133        global_rank % self.tensor_parallel_size
134    }
135
136    /// Calculate pipeline parallel rank from global rank
137    pub fn pp_rank(&self, global_rank: usize) -> usize {
138        global_rank / self.tensor_parallel_size
139    }
140}
141
142/// Layer distribution configuration for pipeline parallelism
143#[derive(Debug, Clone, Serialize, Deserialize)]
144pub struct LayerDistribution {
145    /// Layer assignments per pipeline stage
146    pub stage_layers: Vec<LayerRange>,
147    /// Memory requirements per stage (in bytes)
148    pub stage_memory: Vec<usize>,
149}
150
151/// Range of layers assigned to a pipeline stage
152#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
153pub struct LayerRange {
154    /// First layer (inclusive)
155    pub start: usize,
156    /// Last layer (exclusive)
157    pub end: usize,
158}
159
160impl LayerRange {
161    pub fn new(start: usize, end: usize) -> Self {
162        Self { start, end }
163    }
164
165    pub fn len(&self) -> usize {
166        self.end - self.start
167    }
168
169    pub fn is_empty(&self) -> bool {
170        self.start >= self.end
171    }
172
173    pub fn contains(&self, layer: usize) -> bool {
174        layer >= self.start && layer < self.end
175    }
176}
177
178impl LayerDistribution {
179    /// Create even distribution of layers across stages
180    pub fn even_distribution(num_layers: usize, num_stages: usize) -> Self {
181        let layers_per_stage = num_layers / num_stages;
182        let remainder = num_layers % num_stages;
183
184        let mut stage_layers = Vec::with_capacity(num_stages);
185        let mut start = 0;
186
187        for stage in 0..num_stages {
188            let extra = if stage < remainder { 1 } else { 0 };
189            let end = start + layers_per_stage + extra;
190            stage_layers.push(LayerRange::new(start, end));
191            start = end;
192        }
193
194        Self {
195            stage_layers,
196            stage_memory: vec![0; num_stages], // To be filled based on actual model
197        }
198    }
199
200    /// Get stage for a given layer
201    pub fn stage_for_layer(&self, layer: usize) -> Option<usize> {
202        self.stage_layers
203            .iter()
204            .position(|range| range.contains(layer))
205    }
206}
207
208// ============================================================================
209// Tests
210// ============================================================================
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215
216    #[test]
217    fn test_single_gpu_config() {
218        let config = ParallelConfig::single_gpu(Device::CPU);
219        assert_eq!(config.parallelism_type, ParallelismType::None);
220        assert_eq!(config.world_size(), 1);
221        assert!(!config.is_parallel());
222    }
223
224    #[test]
225    fn test_tensor_parallel_config() {
226        let config = ParallelConfig::tensor_parallel(vec![Device::CUDA(0), Device::CUDA(1)]);
227        assert_eq!(config.parallelism_type, ParallelismType::Tensor);
228        assert_eq!(config.tensor_parallel_size, 2);
229        assert_eq!(config.world_size(), 2);
230        assert!(config.is_parallel());
231    }
232
233    #[test]
234    fn test_hybrid_config() {
235        let config = ParallelConfig::hybrid(
236            vec![
237                Device::CUDA(0),
238                Device::CUDA(1),
239                Device::CUDA(2),
240                Device::CUDA(3),
241            ],
242            2, // tp_size
243            2, // pp_size
244        );
245        assert_eq!(config.world_size(), 4);
246        assert_eq!(config.tp_rank(0), 0);
247        assert_eq!(config.tp_rank(1), 1);
248        assert_eq!(config.tp_rank(2), 0);
249        assert_eq!(config.tp_rank(3), 1);
250        assert_eq!(config.pp_rank(0), 0);
251        assert_eq!(config.pp_rank(1), 0);
252        assert_eq!(config.pp_rank(2), 1);
253        assert_eq!(config.pp_rank(3), 1);
254    }
255
256    #[test]
257    fn test_layer_distribution() {
258        let dist = LayerDistribution::even_distribution(32, 4);
259        assert_eq!(dist.stage_layers.len(), 4);
260        assert_eq!(dist.stage_layers[0].start, 0);
261        assert_eq!(dist.stage_layers[0].end, 8);
262        assert_eq!(dist.stage_layers[3].start, 24);
263        assert_eq!(dist.stage_layers[3].end, 32);
264    }
265
266    #[test]
267    fn test_layer_distribution_uneven() {
268        let dist = LayerDistribution::even_distribution(33, 4);
269        // 33 layers / 4 stages = 8 per stage with 1 remainder
270        // First stage gets extra layer
271        assert_eq!(dist.stage_layers[0].len(), 9);
272        assert_eq!(dist.stage_layers[1].len(), 8);
273    }
274
275    #[test]
276    fn test_stage_for_layer() {
277        let dist = LayerDistribution::even_distribution(32, 4);
278        assert_eq!(dist.stage_for_layer(0), Some(0));
279        assert_eq!(dist.stage_for_layer(7), Some(0));
280        assert_eq!(dist.stage_for_layer(8), Some(1));
281        assert_eq!(dist.stage_for_layer(31), Some(3));
282        assert_eq!(dist.stage_for_layer(32), None);
283    }
284}
ferrum_engine/parallel/config.rs

ferrum_engine/parallel/
config.rs
