rustytorch/crates/training/rtx-distributed/src/lib.rs

//! # RustyTorch++ Distributed Training
//!
//! This crate provides distributed training capabilities for RustyTorch++, including:
//! - Process group management with NCCL/RCCL integration
//! - Communication primitives (AllReduce, Broadcast, AllGather, ReduceScatter)
//! - Hybrid parallelism support (Data, Tensor, Pipeline, Sequence parallelism)
//! - FSDP/ZeRO-style parameter sharding with ≥40% memory reduction
//! - Elastic recovery with write-ahead logging checkpoints
//! - Topology discovery and optimization
//!
//! ## Design Principles
//!
//! - **Safety**: All distributed operations are memory-safe with proper error handling
//! - **Performance**: Target ≥0.8x scaling efficiency (1→8 GPUs), ≥0.7x multi-node
//! - **Reliability**: Elastic recovery from node failures with minimal data loss
//! - **Flexibility**: Support for various parallelism strategies and hybrid approaches
//!
//! ## Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────┐
//! │                    rtx-distributed                          │
//! ├─────────────────┬───────────────────┬───────────────────────┤
//! │ ProcessGroup    │ Communication     │ Parallelism           │
//! │ - World mgmt    │ - AllReduce       │ - Data Parallel       │
//! │ - Rank assign   │ - Broadcast       │ - Tensor Parallel     │
//! │ - Group split   │ - AllGather       │ - Pipeline Parallel   │
//! │ - Topology      │ - ReduceScatter   │ - FSDP/ZeRO          │
//! └─────────────────┴───────────────────┴───────────────────────┘
//! ```
//!
//! ## Usage Example
//!
//! ```rust,ignore
//! use rtx_distributed::{ProcessGroup, Backend, DistributedTensor};
//!
//! // Initialize process group
//! let pg = ProcessGroup::new(Backend::Nccl, world_size, rank)?;
//!
//! // Data parallel gradient averaging
//! let mut gradients = get_model_gradients();
//! pg.all_reduce(&mut gradients, ReduceOp::Sum)?;
//! gradients.div_scalar_(world_size as f32);
//!
//! // FSDP parameter sharding
//! let sharded_params = pg.shard_parameters(&model_params)?;
//! ```

pub mod activation_checkpointing;
pub mod async_grad;
pub mod auto_partition;
pub mod backend;
pub mod backend_benchmark;
#[cfg(feature = "hpc-channels")]
pub mod channels;
pub mod collective_fusion;
pub mod comm;
pub mod comm_overlap;
pub mod context_parallel;
pub mod coordination;
pub mod dcp;
pub mod device_mesh;
pub mod distributed_context;
pub mod distributed_dataloader;
pub mod distributed_transformer_trainer;
pub mod dtensor;
pub mod elastic_enhancements;
pub mod elastic_training;
pub mod error;
pub mod fault_tolerance;
pub mod fsdp2;
pub mod gradient_compression;
pub mod group;
pub mod hardware_topology;
#[cfg(feature = "hpc-channels")]
pub mod hpc_bridge;
pub mod hybrid_parallel;
#[cfg(feature = "cuda")]
pub mod kernels;
pub mod memory_pool;
pub mod mixed_precision;
pub mod model_sharding;
pub mod multi_gpu_trainer;
pub mod multi_node;
#[cfg(feature = "nccl")]
pub mod nccl;
#[cfg(feature = "nccl")]
pub mod nccl_comm;
pub mod nvlink_p2p;
pub mod parallel;
pub mod pipeline_parallel;
pub mod profiling;
#[cfg(feature = "rccl")]
pub mod rccl;
pub mod rdma_transport;
pub mod recovery;
pub mod resharding;
#[cfg(feature = "rnccl")]
pub mod rnccl_backend;
pub mod scaling_benchmarks;
pub mod tcp_backend;
pub mod tensor_ext;
pub mod topology;
pub mod zero_copy_fsdp;

#[cfg(test)]
pub mod multi_gpu_tests;

#[cfg(test)]
pub mod integration_tests;

#[cfg(all(test, feature = "nccl"))]
pub mod nccl_tests;

#[cfg(test)]
mod distributed_context_tests;

pub use activation_checkpointing::{
    ActivationCheckpointManager, ActivationStorage, CheckpointContext, CheckpointPolicy,
    CheckpointStats, CheckpointedSegment, MemoryAwareCheckpointer,
    SharedActivationCheckpointManager, StoredActivation, estimate_memory_savings,
    optimal_checkpoint_interval, shared_activation_checkpoint_manager,
    shared_activation_checkpoint_manager_with_limit,
};
pub use async_grad::{
    AggregatorStats, AllReduceState, AsyncAllReduceHandle, AsyncGradAggregator, AsyncGradConfig,
    GradientBucket as AsyncGradientBucket, GradientEntry, GradientPipeline, PipelineMessage,
    SharedAsyncGradAggregator, shared_async_grad_aggregator,
};
pub use backend::{Backend, BackendConfig};
#[cfg(feature = "hpc-channels")]
pub use channels::{
    SharedTrainingChannelBridge, TrainingChannelBridge, TrainingControlCommand,
    TrainingMetricsEvent, TrainingProgressEvent, TrainingStartEvent, TrainingStatus,
    TrainingStatusEvent, shared_channel_bridge,
};
pub use collective_fusion::{
    CollectiveFusionManager, CollectiveType, FusionBuffer, FusionBufferStats, FusionConfig,
    FusionGroup, FusionManagerStats, FusionOp, FusionStrategy, SharedFusionManager,
    shared_fusion_manager,
};
pub use comm::{AllGatherOutput, AllReduceOp, CommunicationPrimitive, ReduceOp};
pub use comm_overlap::{
    AutoTunerStats, BucketAutoTuner, BucketManager, BucketStats, CommOperation, CommState,
    DoubleBuffer, GradientBucket, OverlapConfig, OverlapScheduler, SchedulerStats,
    SharedOverlapScheduler, shared_overlap_scheduler,
};
pub use context_parallel::{
    ContextParallelConfig, ContextParallelGroup, ContextParallelStats, KVCache, RingAttentionState,
    RingAttentionStats, SequenceShardInfo,
};
pub use coordination::{
    ConsensusProtocol, DistributedCoordinator, GlobalStateManager, GradientAggregator,
    HyperparameterServer, LRSchedule, ModelSynchronizer, SynchronizationBarrier, TrainingScheduler,
    WorkloadDistributor,
};
pub use dcp::{
    AsyncSaveHandle, CheckpointValidation, DcpCompressionType, DcpConfig, DcpManager,
    DcpShardMetadata, DcpStats, DistributedStateDict, LoadPlan, LoadResult, RngStates, SavePlan,
    SaveResult, SerializedTensor, StorageBackend, TensorShardInfo, TrainingState,
    create_state_dict, state_dict_to_tensors,
};
pub use device_mesh::{DeviceInfo, DeviceMesh, DeviceMeshBuilder, MeshDimension};
pub use distributed_context::DistributedContext;
pub use distributed_dataloader::{
    DataBatch, DataLoaderConfig, DataLoaderStats, DataSample, DataSource, DistributedDataLoader,
    DistributionStrategy, InMemoryDataSource, IndexDistributor, PrefetchBuffer, PrefetchStats,
    PrefetchStrategy, SharedDataLoader, shared_dataloader,
};
pub use distributed_transformer_trainer::{
    CombinedTrainingMetrics, DistributedTrainingConfig, DistributedTrainingState,
    DistributedTransformerModel, DistributedTransformerTrainer,
};
pub use dtensor::{
    DTensor, DType as DTensorDType, PartialReduceOp, Placement, TensorSpec, ones as dtensor_ones,
    zeros as dtensor_zeros,
};
pub use elastic_enhancements::{
    AgentRunResult, AgentWorkerState, BatchSizeAdapter, BatchSizeConfig, BatchSizeStrategy,
    ElasticAgent, ElasticAgentConfig, ElasticEvent, ElasticEventHandler, ElasticEventManager,
    ElasticEventType, LRAdaptationConfig, LRAdaptationStrategy, LRAdapter, LoggingEventHandler,
    LossScalingRecovery, LossScalingRecoveryConfig, LossScalingState, SharedBatchSizeAdapter,
    SharedElasticAgent, SharedElasticEventManager, SharedLRAdapter, SharedLossScalingRecovery,
    shared_batch_adapter, shared_elastic_agent, shared_event_manager, shared_loss_scaling_recovery,
    shared_lr_adapter,
};
pub use elastic_training::{
    ElasticCluster, ElasticConfig, ElasticStats, HealthMonitor, RedistributionStats, ScaleOpState,
    ScaleOpType, ScaleOperation, ScalingPolicy, ShardMove, SharedElasticCluster,
    StateRedistributor, WorkerId, WorkerInfo, WorkerState, shared_elastic_cluster,
};
pub use error::{DistributedError, Result};
pub use fault_tolerance::{
    CheckpointManager, FailureMode, FaultDetector, HeartbeatMonitor, RecoveryManager,
    RecoveryStrategy, RedundancyLevel, ReplicaManager, StateReplication,
};
pub use fsdp2::{
    BackwardPrefetch, Fsdp2Config, Fsdp2ConfigBuilder, Fsdp2MemoryStats, Fsdp2Module,
    Fsdp2ShardedParam, MixedPrecisionPolicy, fully_shard, get_sharded_state_dict,
    load_sharded_state_dict,
};
pub use gradient_compression::{
    CompressedGradient, CompressionConfig, CompressionStats, CompressionType, ErrorFeedback,
    GradientCompressor, SharedGradientCompressor, shared_compressor,
};
pub use group::{ProcessGroup, WorldInfo};
pub use hardware_topology::{
    GpuDeviceProperties, HardwareTopology, NetworkInterface, NetworkInterfaceType, NumaNode,
    NumaTopology, NvLinkConnection, NvLinkTopology, PcieDevice, PcieDeviceType, PcieTopology,
};
#[cfg(feature = "hpc-channels")]
pub use hpc_bridge::{
    CheckpointLoadEvent, CheckpointSaveEvent, DataBatchEvent, DataPrefetchEvent,
    FaultDetectedEvent, FaultRecoveredEvent, FaultType, GradientPullEvent, GradientPushEvent,
    GradientSyncEvent, HpcRecoveryStrategy, ModelPipelineEvent, ModelShardEvent, PerfComputeEvent,
    PerfMemoryEvent, PerfThroughputEvent, SharedTorchChannelBridge, TensorParallelEvent,
    TorchChannelBridge, WorkerBarrierEvent, WorkerHeartbeatEvent, WorkerRegisterEvent,
};
pub use hybrid_parallel::{
    CommOp, CoordinatorStats, HybridParallelConfig, HybridParallelCoordinator, MeshCoordinate,
    ParallelDimension, PendingComm as HybridPendingComm, ProcessGroupInfo, ProcessGroupMesh,
    SharedHybridParallelCoordinator, shared_hybrid_coordinator,
};
#[cfg(feature = "cuda")]
pub use kernels::{CudaReduceOp, FusedKernelLauncher};
pub use memory_pool::{
    AllocationStrategy, BlockId, BlockState, MemoryAllocation, MemoryBlock, MemoryPool,
    MemoryPoolConfig, MemoryPoolStats, MultiDevicePoolManager, SharedMemoryPool, SizeClassCache,
    shared_memory_pool,
};
pub use mixed_precision::{
    GradScaler, MasterWeightManager, MasterWeightStats, MixedPrecisionConfig,
    MixedPrecisionTrainer, PrecisionType, ScalerState, ScalerStats, SharedGradScaler,
    SharedMixedPrecisionTrainer, bf16_to_fp32, fp16_to_fp32, shared_grad_scaler,
    shared_mixed_precision_trainer, to_bf16, to_fp16,
};
pub use model_sharding::{
    CommType, DevicePlacement, ModelSharder, ModelShardingConfig, ParameterInfo, ParameterType,
    PendingComm, ShardDimension, ShardManager, ShardSpec, ShardingPlan, ShardingStats,
    ShardingStrategy, SharedModelSharder, SharedShardManager, shared_model_sharder,
    shared_shard_manager,
};
pub use multi_gpu_trainer::{
    FaultTolerance, LoadBalancer, MultiGpuTrainer, ScalingOptimizer, TrainingMetrics,
};
pub use multi_node::{
    AggregationPattern, BandwidthOptimizer, CrossNodeCommunicator, HealthChecker, InterconnectType,
    MultiNodeCluster, NetworkTopology, NodeConfig, NodeInfo, NodeRole, RendezvousProtocol,
};
pub use nvlink_p2p::{
    GpuP2PInfo, P2PCapability, P2PConfig, P2PConnection, P2PManager, P2PStats, P2PTopology,
    SharedP2PManager, TransferMethod, shared_p2p_manager,
};
pub use parallel::{DataParallel, Fsdp, PipelineParallel, TensorParallel};
pub use pipeline_parallel::{
    AccumulatedGradients,
    LinearStageModule,
    MicroBatch,
    NoOpStageModule,
    PipelineConfig,
    PipelineExecutor,
    PipelineOp,
    PipelineSchedule,
    PipelineScheduler,
    PipelineStage,
    PipelineStats,
    SharedPipelineExecutor,
    // Stage module trait for actual tensor computation
    StageModule,
    StageState,
    StageStats,
    StashedActivation,
    // Re-export Tensor for convenience
    Tensor as PipelineTensor,
    shared_pipeline_executor,
};
pub use profiling::{
    CommEvent, MemorySnapshot, NvtxColor, OperationStats, Profiler, ProfilerBackend,
    ProfilingConfig, ProfilingLevel, ProfilingRange, RangeGuard, SharedProfiler, shared_profiler,
};
pub use rdma_transport::{
    GpuDirectConfig, GpuDirectRdmaTransport, QpState, RdmaCompletionQueue, RdmaConfig,
    RdmaConnectionInfo, RdmaContext, RdmaDeviceInfo, RdmaGid, RdmaMemoryRegion, RdmaNodeType,
    RdmaOpType, RdmaPeerConnection, RdmaPortInfo, RdmaPortState, RdmaSge, RdmaTransport,
    RdmaWorkCompletion, RdmaWorkRequest, WcStatus,
};
pub use recovery::{Checkpoint, ElasticRecovery, WalEntry};
pub use resharding::{
    placements_compatible, redistribute_like, replicate_tensor, reshard_dtensor, shard_tensor,
};
pub use scaling_benchmarks::{BenchmarkConfig, ScalingBenchmarkResults, ScalingBenchmarkSuite};
pub use tensor_ext::{TensorExt, TensorShapeExt};
pub use topology::{TopologyInfo, TopologyOptimizer};
pub use zero_copy_fsdp::{
    FsdpStats, GradientShard, ParameterShard, ShardMetadata, SharedZeroCopyFsdp, ZeroCopyFsdp,
    ZeroCopyFsdpConfig, ZeroStage, shared_zero_copy_fsdp,
};

/// Re-export rtx-tensor types for convenience
pub use rtx_tensor::{Device, Shape as TensorShape, Tensor};

/// Version information
pub const VERSION: &str = env!("CARGO_PKG_VERSION");

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_version_exists() {
        assert!(!VERSION.is_empty());
    }
}