optirs-tpu 0.3.1

OptiRS TPU coordination and pod management
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88

// Coordination Configuration Module

use crate::pod_coordination::types::*;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum BatchParallelizationStrategy {
    #[default]
    DataParallel,
    ModelParallel,
    PipelineParallel,
    Hybrid,
    Adaptive,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum CommunicationPattern {
    #[default]
    AllReduce,
    AllGather,
    AllToAll,
    Ring,
    Broadcast,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum CoordinationStrategy {
    #[default]
    Centralized,
    Decentralized,
    Hierarchical,
    Adaptive,
    Mesh,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum GradientAggregationMethod {
    #[default]
    Average,
    Sum,
    WeightedAverage,
    LocalSGD,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum LoadBalancingStrategy {
    Static,
    #[default]
    Dynamic,
    Adaptive,
    LatencyAware,
    BandwidthAware,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum MemoryManagementStrategy {
    StaticPartitioning,
    #[default]
    DynamicPartitioning,
    Pooling,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PodCoordinationConfig {
    pub topology: crate::pod_coordination::PodTopology,
    pub num_devices: usize,
    pub coordination_strategy: CoordinationStrategy,
    pub communication_pattern: CommunicationPattern,
    pub synchronization_mode: SynchronizationMode,
    pub batch_strategy: BatchParallelizationStrategy,
    pub gradient_aggregation: GradientAggregationMethod,
    pub enable_fault_tolerance: bool,
    pub heartbeat_interval_ms: u64,
    pub operation_timeout_ms: u64,
    pub enable_performance_monitoring: bool,
    pub load_balancing_strategy: LoadBalancingStrategy,
    pub memory_management: MemoryManagementStrategy,
    pub adaptive_optimization: bool,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum SynchronizationMode {
    #[default]
    Synchronous,
    Asynchronous,
    BulkSynchronous,
    Adaptive,
}