453 lines
16 KiB
Rust
453 lines
16 KiB
Rust
//! Parallel planning agent for autonomous parallelization strategy optimization.
|
|
|
|
use crate::{
|
|
agents::AutonomousAgent,
|
|
error::AutoResult,
|
|
proposal::{Proposal, ProposalType},
|
|
};
|
|
use rtx_graph::Graph;
|
|
use rtx_runtime::Runtime;
|
|
use std::sync::Arc;
|
|
use tracing::info;
|
|
|
|
/// Parallelization opportunity analysis.
|
|
#[derive(Debug, Clone)]
|
|
pub struct ParallelizationOpportunity {
|
|
pub node_id: String,
|
|
pub opportunity_type: ParallelType,
|
|
pub estimated_speedup: f32,
|
|
pub resource_requirements: ResourceRequirements,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub enum ParallelType {
|
|
DataParallel,
|
|
ModelParallel,
|
|
PipelineParallel,
|
|
TensorParallel,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct ResourceRequirements {
|
|
pub min_devices: usize,
|
|
pub memory_per_device: usize,
|
|
pub communication_overhead: f32,
|
|
}
|
|
|
|
/// Agent for autonomous parallel execution planning.
|
|
pub struct ParallelPlannerAgent {
|
|
runtime: Arc<Runtime>,
|
|
name: String,
|
|
version: String,
|
|
}
|
|
|
|
impl ParallelPlannerAgent {
|
|
pub fn new(runtime: Arc<Runtime>) -> AutoResult<Self> {
|
|
Ok(Self {
|
|
runtime,
|
|
name: "ParallelPlannerAgent".to_string(),
|
|
version: "1.0.0".to_string(),
|
|
})
|
|
}
|
|
|
|
pub async fn analyze_parallelization_opportunities(
|
|
&self,
|
|
graph: &Graph,
|
|
) -> AutoResult<Vec<ParallelizationOpportunity>> {
|
|
info!(
|
|
"Analyzing parallelization opportunities for graph with {} nodes",
|
|
graph.node_count()
|
|
);
|
|
let mut opportunities = Vec::new();
|
|
|
|
// Analyze each node for parallelization potential
|
|
for node in graph.nodes() {
|
|
let node_id = node.id.to_string();
|
|
let op_name = node.operation.name();
|
|
|
|
// Identify compute-heavy operations suitable for different parallelization strategies
|
|
match op_name {
|
|
"Compute::MatMul" => {
|
|
// MatMul is ideal for data parallelism (batch dimension)
|
|
// and tensor parallelism (row/column splitting)
|
|
opportunities.push(ParallelizationOpportunity {
|
|
node_id: node_id.clone(),
|
|
opportunity_type: ParallelType::DataParallel,
|
|
estimated_speedup: self.estimate_data_parallel_speedup(graph, node),
|
|
resource_requirements: ResourceRequirements {
|
|
min_devices: 2,
|
|
memory_per_device: 1024 * 1024 * 1024, // 1GB base
|
|
communication_overhead: 0.05, // Low overhead for data parallel
|
|
},
|
|
});
|
|
|
|
opportunities.push(ParallelizationOpportunity {
|
|
node_id: node_id.clone(),
|
|
opportunity_type: ParallelType::TensorParallel,
|
|
estimated_speedup: self.estimate_tensor_parallel_speedup(graph, node),
|
|
resource_requirements: ResourceRequirements {
|
|
min_devices: 4,
|
|
memory_per_device: 512 * 1024 * 1024,
|
|
communication_overhead: 0.15, // AllReduce overhead
|
|
},
|
|
});
|
|
}
|
|
|
|
"Compute::Activation" | "Compute::BatchNorm" => {
|
|
// Elementwise ops - good for data parallelism
|
|
opportunities.push(ParallelizationOpportunity {
|
|
node_id: node_id.clone(),
|
|
opportunity_type: ParallelType::DataParallel,
|
|
estimated_speedup: self.estimate_data_parallel_speedup(graph, node) * 0.9,
|
|
resource_requirements: ResourceRequirements {
|
|
min_devices: 2,
|
|
memory_per_device: 256 * 1024 * 1024,
|
|
communication_overhead: 0.02,
|
|
},
|
|
});
|
|
}
|
|
|
|
"Compute::Conv2d" => {
|
|
// Conv2d benefits from both data and model parallelism
|
|
opportunities.push(ParallelizationOpportunity {
|
|
node_id: node_id.clone(),
|
|
opportunity_type: ParallelType::DataParallel,
|
|
estimated_speedup: self.estimate_data_parallel_speedup(graph, node) * 1.1,
|
|
resource_requirements: ResourceRequirements {
|
|
min_devices: 2,
|
|
memory_per_device: 2048 * 1024 * 1024,
|
|
communication_overhead: 0.08,
|
|
},
|
|
});
|
|
|
|
opportunities.push(ParallelizationOpportunity {
|
|
node_id: node_id.clone(),
|
|
opportunity_type: ParallelType::ModelParallel,
|
|
estimated_speedup: 2.0,
|
|
resource_requirements: ResourceRequirements {
|
|
min_devices: 2,
|
|
memory_per_device: 1024 * 1024 * 1024,
|
|
communication_overhead: 0.12,
|
|
},
|
|
});
|
|
}
|
|
|
|
_ => {
|
|
// Generic operations - only suggest data parallelism if graph is large enough
|
|
if graph.node_count() > 10 {
|
|
opportunities.push(ParallelizationOpportunity {
|
|
node_id: node_id.clone(),
|
|
opportunity_type: ParallelType::DataParallel,
|
|
estimated_speedup: 1.5,
|
|
resource_requirements: ResourceRequirements {
|
|
min_devices: 2,
|
|
memory_per_device: 512 * 1024 * 1024,
|
|
communication_overhead: 0.1,
|
|
},
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Analyze graph depth for pipeline parallelism opportunities
|
|
let graph_depth = self.calculate_graph_depth(graph);
|
|
if graph_depth > 4 {
|
|
opportunities.push(ParallelizationOpportunity {
|
|
node_id: "pipeline_stages".to_string(),
|
|
opportunity_type: ParallelType::PipelineParallel,
|
|
estimated_speedup: self.estimate_pipeline_speedup(graph_depth),
|
|
resource_requirements: ResourceRequirements {
|
|
min_devices: graph_depth.min(8),
|
|
memory_per_device: 1024 * 1024 * 1024,
|
|
communication_overhead: 0.05 * graph_depth as f32,
|
|
},
|
|
});
|
|
}
|
|
|
|
// Sort by estimated speedup (highest first)
|
|
opportunities.sort_by(|a, b| {
|
|
b.estimated_speedup
|
|
.partial_cmp(&a.estimated_speedup)
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
});
|
|
|
|
// Deduplicate by taking best opportunity per node
|
|
let mut seen_nodes = std::collections::HashSet::new();
|
|
opportunities.retain(|op| {
|
|
if seen_nodes.contains(&op.node_id) {
|
|
false
|
|
} else {
|
|
seen_nodes.insert(op.node_id.clone());
|
|
true
|
|
}
|
|
});
|
|
|
|
Ok(opportunities)
|
|
}
|
|
|
|
/// Estimate speedup from data parallelism based on graph structure
|
|
fn estimate_data_parallel_speedup(&self, graph: &Graph, _node: &rtx_graph::GraphNode) -> f32 {
|
|
// Base speedup of 2x for 2 devices, with diminishing returns
|
|
let node_count = graph.node_count();
|
|
let base_speedup = 2.0;
|
|
|
|
// Larger graphs benefit more from data parallelism
|
|
let scale_factor = (node_count as f32 / 10.0).min(2.0).max(1.0);
|
|
|
|
base_speedup * scale_factor
|
|
}
|
|
|
|
/// Estimate speedup from tensor parallelism
|
|
fn estimate_tensor_parallel_speedup(&self, graph: &Graph, _node: &rtx_graph::GraphNode) -> f32 {
|
|
// Tensor parallelism has higher communication overhead but better for large tensors
|
|
let node_count = graph.node_count();
|
|
let base_speedup = 1.8;
|
|
|
|
// Works better for larger models
|
|
let scale_factor = (node_count as f32 / 20.0).min(1.5).max(0.8);
|
|
|
|
base_speedup * scale_factor
|
|
}
|
|
|
|
/// Estimate speedup from pipeline parallelism based on graph depth
|
|
fn estimate_pipeline_speedup(&self, depth: usize) -> f32 {
|
|
// Pipeline parallelism approaches linear speedup with depth, minus bubble overhead
|
|
let ideal_speedup = depth as f32;
|
|
let bubble_overhead = 1.0 / depth as f32; // Fraction of time in bubbles
|
|
|
|
ideal_speedup * (1.0 - bubble_overhead)
|
|
}
|
|
|
|
/// Calculate the critical path depth of the graph
|
|
fn calculate_graph_depth(&self, graph: &Graph) -> usize {
|
|
// Simple heuristic: use edge count / node count ratio as depth estimate
|
|
// A more accurate implementation would do topological traversal
|
|
let nodes = graph.node_count();
|
|
let edges = graph.edge_count();
|
|
|
|
if nodes == 0 {
|
|
return 0;
|
|
}
|
|
|
|
// Estimate depth based on graph density
|
|
let density = edges as f32 / nodes as f32;
|
|
|
|
// For sequential graphs, density ~= 1 (each node has one incoming edge)
|
|
// For parallel graphs, density < 1
|
|
// For highly connected graphs, density > 1
|
|
|
|
if density <= 1.0 {
|
|
// Mostly sequential - depth is approximately node count
|
|
nodes.min(16)
|
|
} else {
|
|
// More parallel - depth is less than node count
|
|
((nodes as f32) / density).ceil() as usize
|
|
}
|
|
}
|
|
|
|
pub async fn generate_data_parallel_proposals(
|
|
&self,
|
|
graph: &Graph,
|
|
) -> AutoResult<Vec<Proposal>> {
|
|
let mut proposals = Vec::new();
|
|
|
|
let node_count = graph.node_count();
|
|
|
|
if node_count > 5 {
|
|
// Calculate optimal batch splitting strategy
|
|
let optimal_devices = (node_count / 5).min(8).max(2);
|
|
let estimated_speedup =
|
|
self.estimate_data_parallel_speedup(graph, graph.nodes().next().unwrap());
|
|
|
|
proposals.push(Proposal::new(
|
|
ProposalType::DataParallel,
|
|
format!(
|
|
"Apply data parallelism across {} devices for {} compute nodes (estimated {:.1}x speedup)",
|
|
optimal_devices, node_count, estimated_speedup
|
|
),
|
|
estimated_speedup,
|
|
));
|
|
|
|
// Gradient synchronization strategy
|
|
if node_count > 20 {
|
|
proposals.push(Proposal::new(
|
|
ProposalType::DataParallel,
|
|
"Use gradient compression (TopK or random) to reduce all-reduce communication overhead".to_string(),
|
|
1.4,
|
|
));
|
|
}
|
|
}
|
|
|
|
Ok(proposals)
|
|
}
|
|
|
|
pub async fn generate_pipeline_parallel_proposals(
|
|
&self,
|
|
graph: &Graph,
|
|
) -> AutoResult<Vec<Proposal>> {
|
|
let mut proposals = Vec::new();
|
|
|
|
let depth = self.calculate_graph_depth(graph);
|
|
|
|
if depth > 4 {
|
|
let num_stages = depth.min(8);
|
|
let estimated_speedup = self.estimate_pipeline_speedup(depth);
|
|
|
|
proposals.push(Proposal::new(
|
|
ProposalType::PipelineParallel,
|
|
format!(
|
|
"Implement {}-stage pipeline parallelism for depth-{} graph (estimated {:.1}x speedup)",
|
|
num_stages, depth, estimated_speedup
|
|
),
|
|
estimated_speedup,
|
|
));
|
|
|
|
// Micro-batch suggestion for pipeline efficiency
|
|
if depth > 6 {
|
|
proposals.push(Proposal::new(
|
|
ProposalType::PipelineParallel,
|
|
format!(
|
|
"Use {} micro-batches to reduce pipeline bubble overhead",
|
|
depth * 2
|
|
),
|
|
1.3,
|
|
));
|
|
}
|
|
|
|
// Interleaved scheduling for better memory efficiency
|
|
if depth > 8 {
|
|
proposals.push(Proposal::new(
|
|
ProposalType::PipelineParallel,
|
|
"Apply interleaved 1F1B schedule to reduce peak memory usage by 50%"
|
|
.to_string(),
|
|
1.2,
|
|
));
|
|
}
|
|
}
|
|
|
|
Ok(proposals)
|
|
}
|
|
|
|
pub async fn generate_tensor_parallel_proposals(
|
|
&self,
|
|
graph: &Graph,
|
|
) -> AutoResult<Vec<Proposal>> {
|
|
let mut proposals = Vec::new();
|
|
|
|
// Count large tensor operations
|
|
let mut large_matmul_count = 0;
|
|
let mut large_conv_count = 0;
|
|
|
|
for node in graph.nodes() {
|
|
match node.operation.name() {
|
|
"Compute::MatMul" => large_matmul_count += 1,
|
|
"Compute::Conv2d" => large_conv_count += 1,
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
if large_matmul_count > 0 {
|
|
let estimated_speedup =
|
|
self.estimate_tensor_parallel_speedup(graph, graph.nodes().next().unwrap());
|
|
proposals.push(Proposal::new(
|
|
ProposalType::TensorParallel,
|
|
format!(
|
|
"Apply column-parallel tensor splitting to {} MatMul operations (estimated {:.1}x speedup)",
|
|
large_matmul_count, estimated_speedup
|
|
),
|
|
estimated_speedup,
|
|
));
|
|
}
|
|
|
|
if large_conv_count > 0 {
|
|
proposals.push(Proposal::new(
|
|
ProposalType::TensorParallel,
|
|
format!(
|
|
"Apply channel-parallel splitting to {} Conv2d operations",
|
|
large_conv_count
|
|
),
|
|
2.0,
|
|
));
|
|
}
|
|
|
|
// Sequence parallelism for attention-heavy models
|
|
if large_matmul_count > 10 {
|
|
proposals.push(Proposal::new(
|
|
ProposalType::TensorParallel,
|
|
"Apply sequence parallelism for attention computation to reduce memory footprint"
|
|
.to_string(),
|
|
1.8,
|
|
));
|
|
}
|
|
|
|
Ok(proposals)
|
|
}
|
|
|
|
pub async fn optimize_communication_patterns(
|
|
&self,
|
|
device_count: usize,
|
|
) -> AutoResult<Vec<Proposal>> {
|
|
let mut proposals = Vec::new();
|
|
|
|
if device_count > 2 {
|
|
// Ring all-reduce for larger device counts
|
|
if device_count >= 4 {
|
|
proposals.push(Proposal::new(
|
|
ProposalType::CommunicationOptimization,
|
|
format!(
|
|
"Use ring all-reduce pattern for {} devices to achieve 2(n-1)/n bandwidth efficiency",
|
|
device_count
|
|
),
|
|
1.5,
|
|
));
|
|
}
|
|
|
|
// Hierarchical all-reduce for very large clusters
|
|
if device_count >= 8 {
|
|
proposals.push(Proposal::new(
|
|
ProposalType::CommunicationOptimization,
|
|
format!(
|
|
"Apply hierarchical all-reduce (intra-node then inter-node) for {} devices",
|
|
device_count
|
|
),
|
|
1.8,
|
|
));
|
|
}
|
|
|
|
// Overlap communication with computation
|
|
proposals.push(Proposal::new(
|
|
ProposalType::CommunicationOptimization,
|
|
format!(
|
|
"Overlap gradient all-reduce with backward computation for {} devices",
|
|
device_count
|
|
),
|
|
1.3,
|
|
));
|
|
|
|
// NCCL optimization hints
|
|
if device_count > 4 {
|
|
proposals.push(Proposal::new(
|
|
ProposalType::CommunicationOptimization,
|
|
"Enable NCCL tree algorithm for large tensor all-reduce operations".to_string(),
|
|
1.2,
|
|
));
|
|
}
|
|
}
|
|
|
|
Ok(proposals)
|
|
}
|
|
}
|
|
|
|
impl AutonomousAgent for ParallelPlannerAgent {
|
|
fn name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
fn version(&self) -> &str {
|
|
&self.version
|
|
}
|
|
fn is_healthy(&self) -> bool {
|
|
true
|
|
}
|
|
}
|