kwaai-distributed 0.3.36

Distributed ML operations for KwaaiNet - MoE, parameter averaging (Hivemind patterns)
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

//! # kwaai-distributed
//!
//! Distributed ML operations for KwaaiNet, implementing Hivemind patterns.
//!
//! This crate provides:
//!
//! - **Mixture of Experts (MoE)**: Distributed model layers across network
//! - **Decentralized Averaging**: Parameter sync without master node
//! - **Fault Tolerance**: Graceful handling of node failures
//!
//! ## Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────┐
//! │                  kwaai-distributed                   │
//! ├─────────────────┬─────────────────┬─────────────────┤
//! │   MoE Layer     │   Averaging     │  Fault Tolerance │
//! │  (Expert Routing)│ (Gradient Sync) │  (Retry/Fallback)│
//! ├─────────────────┴─────────────────┴─────────────────┤
//! │                    kwaai-p2p                         │
//! │               (P2P Networking / DHT)                 │
//! └─────────────────────────────────────────────────────┘
//! ```

pub mod averaging;
pub mod coordinator;
pub mod error;
pub mod expert;
pub mod moe;

pub use averaging::{AveragingResult, DecentralizedAverager, ParameterAverager};
pub use coordinator::DistributedCoordinator;
pub use error::{DistributedError, DistributedResult};
pub use expert::{Expert, ExpertId, ExpertRegistry};
pub use moe::{ExpertRouter, MixtureOfExperts, Routing};

/// Configuration for distributed operations
#[derive(Debug, Clone)]
pub struct DistributedConfig {
    /// Enable MoE distributed layers
    pub enable_moe: bool,
    /// Enable parameter averaging
    pub enable_averaging: bool,
    /// Number of experts to route to (top-k)
    pub moe_top_k: usize,
    /// Target averaging group size
    pub averaging_group_size: usize,
    /// Timeout for remote operations (ms)
    pub timeout_ms: u64,
    /// Maximum retry attempts
    pub max_retries: usize,
}

impl Default for DistributedConfig {
    fn default() -> Self {
        Self {
            enable_moe: true,
            enable_averaging: true,
            moe_top_k: 2,
            averaging_group_size: 4,
            timeout_ms: 5000,
            max_retries: 3,
        }
    }
}