trustformers_core/parallel/
mod.rs

1//! Parallel execution support for TrustformeRS
2//!
3//! This module provides infrastructure for various parallelism strategies including:
4//! - Data parallelism
5//! - Model parallelism (tensor and pipeline)
6//! - Hybrid parallelism
7//! - NUMA-aware optimization
8
9pub mod model_parallel;
10pub mod parallel_layers;
11pub mod pipeline_parallel;
12pub mod tensor_parallel;
13
14pub mod mpi_communicator;
15
16#[cfg(feature = "nccl")]
17pub mod nccl_communicator;
18
19pub use model_parallel::{
20    CommunicationBackend, Communicator, DeviceMesh, DistributedTensor, ModelParallelConfig,
21    ModelParallelContext, ModelParallelStrategy, PipelineOp, PipelineSchedule,
22    PipelineScheduleType, TensorPartition,
23};
24
25pub use parallel_layers::{
26    ActivationType, ColumnParallelLinear, ParallelMLP, ParallelMultiHeadAttention,
27    RowParallelLinear,
28};
29
30pub use tensor_parallel::{
31    AsyncTensorParallel, InitMethod, TensorParallelInit, TensorParallelOps, TensorParallelShapes,
32};
33
34pub use pipeline_parallel::{
35    MicrobatchManager, PipelineExecutor, PipelineLayer, PipelineModel, PipelineOptimizer,
36    PipelineStage,
37};
38
39pub use mpi_communicator::{mpi_utils, MpiCommunicatorImpl};
40
41#[cfg(feature = "nccl")]
42pub use nccl_communicator::{create_nccl_communicator, NcclCommunicator};
43
44use crate::errors::{runtime_error, Result};
45use parking_lot::RwLock;
46use std::sync::Arc;
47
48/// Core parallelism strategy
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum ParallelismStrategy {
51    /// Data parallelism only
52    Data,
53    /// Model parallelism (tensor or pipeline)
54    Model,
55    /// Hybrid (data + model)
56    Hybrid,
57    /// No parallelism (single device)
58    None,
59}
60
61/// Parallel execution context
62#[derive(Clone)]
63pub struct ParallelContext {
64    strategy: ParallelismStrategy,
65    num_devices: usize,
66    device_id: usize,
67    numa_config: Option<NumaConfig>,
68}
69
70/// NUMA configuration for CPU optimization
71#[derive(Debug, Clone)]
72pub struct NumaConfig {
73    pub node_id: usize,
74    pub cpu_affinity: Vec<usize>,
75    pub memory_policy: MemoryPolicy,
76}
77
78#[derive(Debug, Clone, Copy)]
79pub enum MemoryPolicy {
80    /// Bind memory to local NUMA node
81    BindLocal,
82    /// Interleave memory across nodes
83    Interleave,
84    /// Prefer local but allow remote
85    PreferLocal,
86}
87
88impl ParallelContext {
89    pub fn new(strategy: ParallelismStrategy, num_devices: usize) -> Self {
90        Self {
91            strategy,
92            num_devices,
93            device_id: 0,
94            numa_config: None,
95        }
96    }
97
98    pub fn with_device_id(mut self, device_id: usize) -> Self {
99        self.device_id = device_id;
100        self
101    }
102
103    pub fn with_numa_config(mut self, numa_config: NumaConfig) -> Self {
104        self.numa_config = Some(numa_config);
105        self
106    }
107
108    pub fn strategy(&self) -> ParallelismStrategy {
109        self.strategy
110    }
111
112    pub fn num_devices(&self) -> usize {
113        self.num_devices
114    }
115
116    pub fn device_id(&self) -> usize {
117        self.device_id
118    }
119}
120
121/// Parallel operations trait
122pub trait ParallelOps {
123    /// Execute operation in parallel context
124    fn parallel_execute<F, T>(&self, f: F) -> Result<T>
125    where
126        F: FnOnce(&ParallelContext) -> Result<T>;
127
128    /// Map operation across parallel devices
129    fn parallel_map<F, T>(&self, items: Vec<T>, f: F) -> Result<Vec<T>>
130    where
131        F: Fn(T, &ParallelContext) -> Result<T> + Send + Sync,
132        T: Send;
133}
134
135/// Global parallel context
136static PARALLEL_CONTEXT: RwLock<Option<Arc<ParallelContext>>> = RwLock::new(None);
137
138/// Initialize global parallel context
139pub fn init_parallelism(context: ParallelContext) {
140    *PARALLEL_CONTEXT.write() = Some(Arc::new(context));
141}
142
143/// Get global parallel context
144pub fn parallel_context() -> Option<Arc<ParallelContext>> {
145    PARALLEL_CONTEXT.read().clone()
146}
147
148/// Execute function in parallel context
149pub fn parallel_execute<F, T>(f: F) -> Result<T>
150where
151    F: FnOnce(&ParallelContext) -> Result<T>,
152{
153    let context =
154        parallel_context().ok_or_else(|| runtime_error("Parallel context not initialized"))?;
155    f(&context)
156}
157
158/// Map function across items in parallel
159pub fn parallel_map<F, T>(items: Vec<T>, f: F) -> Result<Vec<T>>
160where
161    F: Fn(T, &ParallelContext) -> Result<T> + Send + Sync,
162    T: Send,
163{
164    let context =
165        parallel_context().ok_or_else(|| runtime_error("Parallel context not initialized"))?;
166
167    // Simple implementation - in practice would use thread pool
168    items.into_iter().map(|item| f(item, &context)).collect()
169}
170
171/// Parallel chunk mapping for large datasets
172pub fn parallel_chunk_map<F, T>(items: Vec<T>, chunk_size: usize, f: F) -> Result<Vec<T>>
173where
174    F: Fn(Vec<T>, &ParallelContext) -> Result<Vec<T>> + Send + Sync,
175    T: Send + Clone,
176{
177    let context =
178        parallel_context().ok_or_else(|| runtime_error("Parallel context not initialized"))?;
179
180    let mut chunks = Vec::new();
181    let mut i = 0;
182    while i < items.len() {
183        let end = (i + chunk_size).min(items.len());
184        chunks.push(items[i..end].to_vec());
185        i = end;
186    }
187
188    let results: Result<Vec<Vec<T>>> = chunks.into_iter().map(|chunk| f(chunk, &context)).collect();
189
190    results.map(|vecs| vecs.into_iter().flatten().collect())
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    // ── 1. ParallelismStrategy variants are distinct ──────────────────────────
198
199    #[test]
200    fn test_parallelism_strategy_variants_distinct() {
201        assert_ne!(ParallelismStrategy::Data, ParallelismStrategy::Model);
202        assert_ne!(ParallelismStrategy::Hybrid, ParallelismStrategy::None);
203        assert_eq!(ParallelismStrategy::Data, ParallelismStrategy::Data);
204    }
205
206    // ── 2. ParallelContext constructs with correct strategy ───────────────────
207
208    #[test]
209    fn test_parallel_context_strategy() {
210        let ctx = ParallelContext::new(ParallelismStrategy::Data, 4);
211        assert_eq!(ctx.strategy(), ParallelismStrategy::Data);
212    }
213
214    // ── 3. ParallelContext num_devices ────────────────────────────────────────
215
216    #[test]
217    fn test_parallel_context_num_devices() {
218        let ctx = ParallelContext::new(ParallelismStrategy::Model, 8);
219        assert_eq!(ctx.num_devices(), 8);
220    }
221
222    // ── 4. ParallelContext default device_id is 0 ────────────────────────────
223
224    #[test]
225    fn test_parallel_context_default_device_id() {
226        let ctx = ParallelContext::new(ParallelismStrategy::None, 1);
227        assert_eq!(ctx.device_id(), 0);
228    }
229
230    // ── 5. with_device_id builder ─────────────────────────────────────────────
231
232    #[test]
233    fn test_parallel_context_with_device_id() {
234        let ctx = ParallelContext::new(ParallelismStrategy::Hybrid, 4).with_device_id(3);
235        assert_eq!(ctx.device_id(), 3);
236    }
237
238    // ── 6. with_numa_config sets numa config ─────────────────────────────────
239
240    #[test]
241    fn test_parallel_context_with_numa_config() {
242        let numa = NumaConfig {
243            node_id: 1,
244            cpu_affinity: vec![0, 1, 2, 3],
245            memory_policy: MemoryPolicy::BindLocal,
246        };
247        let ctx = ParallelContext::new(ParallelismStrategy::None, 1).with_numa_config(numa);
248        assert!(ctx.numa_config.is_some(), "numa_config must be set");
249    }
250
251    // ── 7. MemoryPolicy variants can be cloned ────────────────────────────────
252
253    #[test]
254    fn test_memory_policy_clone() {
255        let p = MemoryPolicy::Interleave;
256        let q = p;
257        let _ = q;
258    }
259
260    // ── 8. NumaConfig node_id stored correctly ────────────────────────────────
261
262    #[test]
263    fn test_numa_config_node_id() {
264        let numa = NumaConfig {
265            node_id: 2,
266            cpu_affinity: vec![4, 5],
267            memory_policy: MemoryPolicy::PreferLocal,
268        };
269        assert_eq!(numa.node_id, 2);
270    }
271
272    // ── 9. init_parallelism + parallel_context round-trip ────────────────────
273
274    #[test]
275    fn test_init_and_get_parallel_context() {
276        let ctx = ParallelContext::new(ParallelismStrategy::Data, 2);
277        init_parallelism(ctx);
278        let retrieved = parallel_context();
279        assert!(
280            retrieved.is_some(),
281            "parallel_context must return Some after init"
282        );
283        let c = retrieved.unwrap_or_else(|| panic!("context is None"));
284        assert_eq!(c.strategy(), ParallelismStrategy::Data);
285    }
286
287    // ── 10. parallel_execute returns error when not initialized ───────────────
288    // NOTE: Since global state may be set from test 9, this tests the happy path.
289
290    #[test]
291    fn test_parallel_execute_runs_closure() {
292        init_parallelism(ParallelContext::new(ParallelismStrategy::Data, 1));
293        let result = parallel_execute(|ctx| {
294            assert_eq!(ctx.num_devices(), 1);
295            Ok(42u32)
296        });
297        assert_eq!(result.unwrap_or(0), 42, "parallel_execute must run closure");
298    }
299
300    // ── 11. ParallelismStrategy is Copy ───────────────────────────────────────
301
302    #[test]
303    fn test_parallelism_strategy_is_copy() {
304        let s = ParallelismStrategy::Hybrid;
305        let t = s; // copy
306        assert_eq!(s, t);
307    }
308
309    // ── 12. parallel_map with initialized context ─────────────────────────────
310
311    #[test]
312    fn test_parallel_map_doubles_values() {
313        init_parallelism(ParallelContext::new(ParallelismStrategy::None, 1));
314        let items = vec![1u32, 2, 3, 4];
315        let result = parallel_map(items, |item, _ctx| Ok(item * 2));
316        let values = result.unwrap_or_default();
317        assert_eq!(
318            values,
319            vec![2u32, 4, 6, 8],
320            "parallel_map must double values"
321        );
322    }
323
324    // ── 13. NumaConfig cpu_affinity stored correctly ──────────────────────────
325
326    #[test]
327    fn test_numa_config_cpu_affinity() {
328        let affinity = vec![0usize, 2, 4, 6];
329        let numa = NumaConfig {
330            node_id: 0,
331            cpu_affinity: affinity.clone(),
332            memory_policy: MemoryPolicy::Interleave,
333        };
334        assert_eq!(numa.cpu_affinity, affinity);
335    }
336
337    // ── 14. ParallelContext clone works ───────────────────────────────────────
338
339    #[test]
340    fn test_parallel_context_clone() {
341        let ctx = ParallelContext::new(ParallelismStrategy::Hybrid, 3);
342        let _cloned = ctx.clone();
343    }
344}
trustformers_core/parallel/mod.rs

trustformers_core/parallel/
mod.rs
