//! # 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()); } }