From df9b080dec1b3daa5f5b29394d0001035720f540 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Fri, 3 Apr 2026 14:59:27 -0400 Subject: [PATCH 01/39] Updated graphcast implementation --- .../GraphCast/data_utils/graphcast_graph.py | 115 +++--- experiments/GraphCast/dataset.py | 6 - experiments/GraphCast/layers.py | 68 ++-- experiments/GraphCast/model.py | 335 ++++++++++-------- 4 files changed, 287 insertions(+), 237 deletions(-) diff --git a/experiments/GraphCast/data_utils/graphcast_graph.py b/experiments/GraphCast/data_utils/graphcast_graph.py index 3d592e7..1137559 100644 --- a/experiments/GraphCast/data_utils/graphcast_graph.py +++ b/experiments/GraphCast/data_utils/graphcast_graph.py @@ -57,29 +57,29 @@ class GraphCastTopology: @dataclass class DistributedGraphCastGraph: + # Distributed environment info rank: int world_size: int ranks_per_graph: int + + # Graph metadata mesh_level: int lat_lon_grid: Tensor + + # Mesh vertex features mesh_graph_node_features: Tensor mesh_graph_edge_features: Tensor - mesh_graph_node_rank_placement: Tensor - mesh_graph_edge_rank_placement: Tensor - mesh_graph_src_indices: Tensor - mesh_graph_dst_indices: Tensor - mesh_graph_src_rank_placement: Tensor - mesh_graph_dst_rank_placement: Tensor - grid_rank_placement: Tensor + + # Grid vertex features mesh2grid_graph_node_features: Tensor - mesh2grid_graph_edge_features: Tensor - mesh2grid_graph_edge_rank_placement: Tensor - mesh2grid_graph_src_indices: Tensor - mesh2grid_graph_dst_indices: Tensor grid2mesh_graph_node_features: Tensor + + # Mesh <--> Grid edge features + mesh2grid_graph_edge_features: Tensor grid2mesh_graph_edge_features: Tensor - grid2mesh_graph_src_indices: Tensor - grid2mesh_graph_dst_indices: Tensor + + # Distributed graph info + distributed_comm_patterns: GraphCastCommPatterns def build_graphcast_comm_patterns(graph: GraphCastTopology) -> GraphCastCommPatterns: @@ -110,7 +110,6 @@ def build_graphcast_comm_patterns(graph: GraphCastTopology) -> GraphCastCommPatt grid2mesh_cp = build_communication_pattern( global_edge_list=grid2mesh_edges, partitioning=mesh_part, - neighbor_partitioning=grid_part, rank=rank, world_size=world_size, ) @@ -129,7 +128,6 @@ def build_graphcast_comm_patterns(graph: GraphCastTopology) -> GraphCastCommPatt mesh_cp = build_communication_pattern( global_edge_list=mesh_edges, partitioning=mesh_part, - neighbor_partitioning=mesh_part, rank=rank, world_size=world_size, ) @@ -147,7 +145,6 @@ def build_graphcast_comm_patterns(graph: GraphCastTopology) -> GraphCastCommPatt mesh2grid_cp = build_communication_pattern( global_edge_list=mesh2grid_edges, partitioning=grid_part, - neighbor_partitioning=mesh_part, rank=rank, world_size=world_size, ) @@ -220,6 +217,47 @@ def get_mesh_graph_partition(mesh_level: int, world_size: int): mesh_vertex_rank_placement = torch.tensor(mesh_vertex_rank_placement) return mesh_vertex_rank_placement + @staticmethod + def get_grid_vertex_partition( + lat: int, + lon: int, + mesh_vertex_rank_placement: torch.Tensor, + grid2mesh_grid_src_indices: torch.Tensor, + grid2mesh_mesh_dst_indices: torch.Tensor, + mesh2grid_mesh_src_indices: torch.Tensor, + world_size: int, + ) -> torch.Tensor: + """Generate the partitioning of grid vertices to minimize cross-rank edges. + + For each grid vertex, counts how many of its connected mesh vertices + (via both grid2mesh and mesh2grid edges) live on each rank, then assigns + the grid vertex to the rank with the plurality of connections. + + mesh2grid grid destinations are implicit: grid vertex i owns edges + [3i, 3i+1, 3i+2] since create_mesh2grid_graph assigns exactly 3 + edges (one face's vertices) per grid vertex. + """ + num_grid = lat * lon + votes = torch.zeros(num_grid, world_size, dtype=torch.long) + + # --- grid2mesh contribution: grid vertex is src, mesh vertex is dst --- + g2m_ranks = mesh_vertex_rank_placement[grid2mesh_mesh_dst_indices.long()] + # Flatten (grid_vertex, rank) into a 1D index for scatter_add_ + g2m_flat_idx = grid2mesh_grid_src_indices.long() * world_size + g2m_ranks + votes.view(-1).scatter_add_(0, g2m_flat_idx, torch.ones_like(g2m_flat_idx)) + + # --- mesh2grid contribution: mesh vertex is src, grid vertex is dst --- + # Each grid vertex i has exactly 3 mesh2grid edges at positions [3i, 3i+1, 3i+2] + m2g_grid_dst = torch.arange(num_grid, dtype=torch.long).repeat_interleave(3) + m2g_ranks = mesh_vertex_rank_placement[mesh2grid_mesh_src_indices.long()] + m2g_flat_idx = m2g_grid_dst * world_size + m2g_ranks + votes.view(-1).scatter_add_(0, m2g_flat_idx, torch.ones_like(m2g_flat_idx)) + + # Assign each grid vertex to the rank with the most connections + grid_partitioning = votes.argmax(dim=1) + + return grid_partitioning + def get_mesh_graph(self, mesh_vertex_rank_placement: torch.Tensor): """Get the graph for the distributed graphcast graph.""" @@ -327,8 +365,6 @@ def get_grid2mesh_graph(self, mesh_graph_dict: dict): grid_vertex_rank_placement ) - # TODO: Consider we can have it to so grid2mesh edges don't require a - # backpropagation. (If encoder is after the gather / scatter) grid2mesh_graph_dict = { "node_features": torch.tensor([]), "edge_features": edge_features, @@ -340,7 +376,7 @@ def get_grid2mesh_graph(self, mesh_graph_dict: dict): } return grid2mesh_graph_dict - def get_mesh2grid_graph( + def get_mesh2grid_edges( self, grid_vertex_rank_placement, renumbered_vertices, @@ -359,7 +395,6 @@ def get_mesh2grid_graph( mesh2grid_edge_rank_placement = grid_vertex_rank_placement[dst_grid_indices] mesh2grid_graph_dict = { - "node_features": torch.tensor([]), "edge_features": edge_features, "src_indices": src_mesh_indices, "dst_indices": dst_grid_indices, @@ -401,10 +436,26 @@ def get_graphcast_graph( grid_vertex_rank_placement = grid2mesh_graph["grid_vertex_rank_placement"] renumbered_grid = grid2mesh_graph["renumbered_grid"] - mesh2grid_graph = self.get_mesh2grid_graph( + mesh2grid_graph = self.get_mesh2grid_edges( grid_vertex_rank_placement, renumbered_vertices, renumbered_grid ) + topology = GraphCastTopology( + rank=self.local_rank, + world_size=self.world_size, + ranks_per_graph=self.ranks_per_graph, + mesh_rank_placement=mesh_vertex_rank_placement, + grid_rank_placement=grid_vertex_rank_placement, + mesh_graph_src_indices=mesh_graph["src_indices"], + mesh_graph_dst_indices=mesh_graph["dst_indices"], + mesh2grid_graph_src_indices=mesh2grid_graph["src_indices"], + mesh2grid_graph_dst_indices=mesh2grid_graph["dst_indices"], + grid2mesh_graph_src_indices=grid2mesh_graph["src_indices"], + grid2mesh_graph_dst_indices=grid2mesh_graph["dst_indices"], + ) + + comm_patterns = build_graphcast_comm_patterns(topology) + return DistributedGraphCastGraph( rank=self.rank, world_size=self.world_size, @@ -413,25 +464,9 @@ def get_graphcast_graph( lat_lon_grid=self.lat_lon_grid, mesh_graph_node_features=mesh_graph["node_features"], mesh_graph_edge_features=mesh_graph["edge_features"], - mesh_graph_node_rank_placement=mesh_graph["node_rank_placement"], - mesh_graph_edge_rank_placement=mesh_graph["edge_rank_placement"], - mesh_graph_src_indices=mesh_graph["src_indices"], - mesh_graph_dst_indices=mesh_graph["dst_indices"], - mesh_graph_src_rank_placement=mesh_graph["src_rank_placement"], - mesh_graph_dst_rank_placement=mesh_graph["dst_rank_placement"], - grid_rank_placement=grid2mesh_graph["grid_vertex_rank_placement"], - mesh2grid_graph_node_features=mesh2grid_graph["node_features"], - mesh2grid_graph_edge_features=mesh2grid_graph["edge_features"], - mesh2grid_graph_edge_rank_placement=mesh2grid_graph[ - "mesh2grid_edge_rank_placement" - ], - mesh2grid_graph_src_indices=mesh2grid_graph["src_indices"], - mesh2grid_graph_dst_indices=mesh2grid_graph["dst_indices"], + mesh2grid_graph_node_features=torch.tensor([]), grid2mesh_graph_node_features=grid2mesh_graph["node_features"], + mesh2grid_graph_edge_features=mesh2grid_graph["edge_features"], grid2mesh_graph_edge_features=grid2mesh_graph["edge_features"], - grid2mesh_graph_edge_rank_placement=grid2mesh_graph[ - "grid2mesh_edge_rank_placement" - ], - grid2mesh_graph_src_indices=grid2mesh_graph["src_indices"], - grid2mesh_graph_dst_indices=grid2mesh_graph["dst_indices"], + distributed_comm_patterns=comm_patterns, ) diff --git a/experiments/GraphCast/dataset.py b/experiments/GraphCast/dataset.py index 3e4eb08..b15e700 100644 --- a/experiments/GraphCast/dataset.py +++ b/experiments/GraphCast/dataset.py @@ -261,20 +261,14 @@ def test_synthetic_weather_dataset(num_days, batch_size=1): print("Mesh label:\t", static_graph.mesh_level) print("Mesh Node features:\t", static_graph.mesh_graph_node_features.shape) print("Mesh Edge features:\t", static_graph.mesh_graph_edge_features.shape) - print("Mesh src indices:\t", static_graph.mesh_graph_src_indices.shape) - print("Mesh dst indices:\t", static_graph.mesh_graph_dst_indices.shape) print("=" * 80) print( "mesh2grid edge features:\t", static_graph.mesh2grid_graph_edge_features.shape ) - print("mesh2grid src indices:\t", static_graph.mesh2grid_graph_src_indices.shape) - print("mesh2grid dst indices:\t", static_graph.mesh2grid_graph_dst_indices.shape) print("=" * 80) print( "grid2mesh edge features:\t", static_graph.grid2mesh_graph_edge_features.shape ) - print("grid2mesh src indices:\t", static_graph.grid2mesh_graph_src_indices.shape) - print("grid2mesh dst indices:\t", static_graph.grid2mesh_graph_dst_indices.shape) print("=" * 80) diff --git a/experiments/GraphCast/layers.py b/experiments/GraphCast/layers.py index 524987b..027974b 100644 --- a/experiments/GraphCast/layers.py +++ b/experiments/GraphCast/layers.py @@ -11,14 +11,12 @@ # https://github.com/LBANN and https://github.com/LLNL/LBANN. # # SPDX-License-Identifier: (Apache-2.0) - -from typing import Tuple, Union -import numpy as np import torch import torch.nn as nn -from typing import Optional -from DGraph.Communicator import Communicator -from dist_utils import SingleProcessDummyCommunicator +from DGraph.utils.TimingReport import TimingReport + +""" +Local only layers for mesh processing. These layers do not perform any communication and can be used in both GraphCast and MeshGraphNet.""" class MeshGraphMLP(nn.Module): @@ -72,7 +70,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: The transformed tensor """ - return self._model(x) + with TimingReport("MeshGraphMLP/forward"): + return self._model(x) class MeshNodeBlock(nn.Module): @@ -83,7 +82,6 @@ def __init__( input_node_dim: int, input_edge_dim: int, output_node_dim: int, - comm: Union[Communicator, SingleProcessDummyCommunicator], hidden_dim: int = 512, num_hidden_layers: int = 1, aggregation_type: str = "sum", @@ -102,7 +100,6 @@ def __init__( super(MeshNodeBlock, self).__init__() assert aggregation_type in ["sum"], "Only sum aggregation is supported for now." self.aggregation_type = aggregation_type - self.comm = comm self.mesh_mlp = MeshGraphMLP( input_dim=input_node_dim + input_edge_dim, output_dim=output_node_dim, @@ -115,7 +112,6 @@ def forward( node_features: torch.Tensor, edge_features: torch.Tensor, src_indices: torch.Tensor, - rank_mapping: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Compute the node block @@ -129,17 +125,23 @@ def forward( Returns: The updated node features """ - # Sum all the edge features for each node num_local_nodes = node_features.shape[0] - # TODO: This can be optimized by a fused gather-scatter operation - S.Z - - aggregated_edge_features = self.comm.scatter( - edge_features, src_indices, rank_mapping, num_local_nodes - ) - # Concatenate the node and edge features - x = torch.cat([node_features, aggregated_edge_features], dim=-1) - # Apply the MLP - node_features_new = self.mesh_mlp(x) + node_features + with TimingReport("MeshNodeBlock/scatter_add"): + aggregated_edge_features = torch.zeros( + num_local_nodes, + edge_features.shape[-1], + device=edge_features.device, + dtype=edge_features.dtype, + ) + aggregated_edge_features.scatter_add_( + 0, + src_indices.unsqueeze(-1).expand(-1, edge_features.shape[-1]), + edge_features, + ) + + with TimingReport("MeshNodeBlock/mlp"): + x = torch.cat([node_features, aggregated_edge_features], dim=-1) + node_features_new = self.mesh_mlp(x) + node_features return node_features_new @@ -152,7 +154,6 @@ def __init__( input_dst_node_dim: int, input_edge_dim: int, output_edge_dim: int, - comm: Union[Communicator, SingleProcessDummyCommunicator], hidden_dim: int = 512, num_hidden_layers: int = 1, aggregation_type: str = "sum", @@ -162,7 +163,6 @@ def __init__( input_node_dim (int): The dimensionality of the input node features. input_edge_dim (int): The dimensionality of the input edge features. output_edge_dim (int): The dimensionality of the output edge features. - comm (CommunicatorBase): The communicator to use for distributed training. hidden_dim (int, optional): The dimensionality of the hidden layers. Defaults to 512. aggregation_type (str, optional): The type of aggregation to use. Defaults to "sum". """ @@ -171,7 +171,6 @@ def __init__( super(MeshEdgeBlock, self).__init__() assert aggregation_type in ["sum"], "Only sum aggregation is supported for now." self.aggregation_type = aggregation_type - self.comm = comm self.mesh_mlp = MeshGraphMLP( input_dim=input_src_node_dim + input_dst_node_dim + input_edge_dim, output_dim=output_edge_dim, @@ -186,8 +185,6 @@ def forward( edge_features: torch.Tensor, src_indices: torch.Tensor, dst_indices: torch.Tensor, - src_rank_mapping: Optional[torch.Tensor] = None, - dst_rank_mapping: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Compute the edge block @@ -201,16 +198,13 @@ def forward( Returns: The updated edge features """ - # Concatenate the source and destination node features with the edge features - src_node_features = self.comm.gather( - src_node_features, src_indices, src_rank_mapping - ) - dst_node_features = self.comm.gather( - dst_node_features, dst_indices, dst_rank_mapping - ) - concatenated_features = torch.cat( - [src_node_features, dst_node_features, edge_features], dim=-1 - ) - # Apply the MLP - edge_features_new = self.mesh_mlp(concatenated_features) + edge_features + with TimingReport("MeshEdgeBlock/gather"): + src_node_features = src_node_features[src_indices] + dst_node_features = dst_node_features[dst_indices] + + with TimingReport("MeshEdgeBlock/mlp"): + concatenated_features = torch.cat( + [src_node_features, dst_node_features, edge_features], dim=-1 + ) + edge_features_new = self.mesh_mlp(concatenated_features) + edge_features return edge_features_new diff --git a/experiments/GraphCast/model.py b/experiments/GraphCast/model.py index da6459b..9183ad8 100644 --- a/experiments/GraphCast/model.py +++ b/experiments/GraphCast/model.py @@ -14,11 +14,14 @@ import torch import torch.nn as nn -from typing import Optional, Tuple +from typing import Tuple from torch import Tensor from layers import MeshEdgeBlock, MeshGraphMLP, MeshNodeBlock from graphcast_config import Config from data_utils.graphcast_graph import DistributedGraphCastGraph +from DGraph.distributed import HaloExchange +from DGraph.distributed.commInfo import CommunicationPattern +from DGraph.utils.TimingReport import TimingReport class GraphCastEmbedder(nn.Module): @@ -116,55 +119,62 @@ def __init__(self, cfg: Config, comm, *args, **kwargs) -> None: comm: Communicator object """ super().__init__(*args, **kwargs) + hidden_dim = cfg.model.hidden_dim - edge_block_invars = ( - cfg.model.hidden_dim, - cfg.model.hidden_dim, - cfg.model.hidden_dim, - cfg.model.hidden_dim, - comm, - cfg.model.hidden_dim, - ) - self.edge_mlp = MeshEdgeBlock(*edge_block_invars) - - node_block_invars = ( - cfg.model.hidden_dim, - cfg.model.hidden_dim, - cfg.model.hidden_dim, - comm, - cfg.model.hidden_dim, + self.exchanger = HaloExchange(comm) + + self.edge_mlp = MeshEdgeBlock( + input_src_node_dim=hidden_dim, + input_dst_node_dim=hidden_dim, + input_edge_dim=hidden_dim, + output_edge_dim=hidden_dim, + hidden_dim=hidden_dim, ) - self.mesh_node_mlp = MeshNodeBlock(*node_block_invars) - self.grid_node_mlp = MeshGraphMLP( - input_dim=cfg.model.hidden_dim, output_dim=cfg.model.hidden_dim + self.mesh_node_mlp = MeshNodeBlock( + input_node_dim=hidden_dim, + input_edge_dim=hidden_dim, + output_node_dim=hidden_dim, + hidden_dim=hidden_dim, ) + self.grid_node_mlp = MeshGraphMLP(input_dim=hidden_dim, output_dim=hidden_dim) def forward( self, grid_node_features: Tensor, mesh_node_features: Tensor, grid2mesh_edge_features: Tensor, - grid2mesh_edge_indices_src: Tensor, - grid2mesh_edge_indices_dst: Tensor, + comm_pattern: CommunicationPattern, ) -> Tuple[Tensor, Tensor]: + # local_edge_list: [E, 2] with [central=mesh, neighbor=grid/halo] + edge_index = comm_pattern.local_edge_list + dst_indices = edge_index[:, 0] # mesh (central, aggregation target) + src_indices = edge_index[:, 1] # grid/halo (neighbor, message source) + num_local = comm_pattern.num_local_vertices + + with TimingReport("encoder/halo_exchange"): + halo_features = self.exchanger(mesh_node_features, comm_pattern) + augmented = torch.cat([mesh_node_features, halo_features], dim=0) + + with TimingReport("encoder/edge_block"): + e_feats = self.edge_mlp( + src_node_features=augmented, + dst_node_features=augmented, + edge_features=grid2mesh_edge_features, + src_indices=src_indices, + dst_indices=dst_indices, + ) - e_feats = self.edge_mlp( - src_node_features=grid_node_features, - dst_node_features=mesh_node_features, - edge_features=grid2mesh_edge_features, - src_indices=grid2mesh_edge_indices_src, - dst_indices=grid2mesh_edge_indices_dst, - ) + with TimingReport("encoder/node_block"): + n_feats = self.mesh_node_mlp( + node_features=augmented[:num_local], + edge_features=e_feats, + src_indices=dst_indices, + ) - n_feats = self.mesh_node_mlp( - node_features=mesh_node_features, - edge_features=e_feats, - src_indices=grid2mesh_edge_indices_dst, - ) + with TimingReport("encoder/grid_mlp"): + grid_node_features = grid_node_features + self.grid_node_mlp(grid_node_features) mesh_node_features = mesh_node_features + n_feats - grid_node_features = grid_node_features + self.grid_node_mlp(grid_node_features) - return grid_node_features, mesh_node_features @@ -179,54 +189,73 @@ def __init__(self, cfg: Config, comm, *args, **kwargs): comm: Communicator object """ super().__init__() + hidden_dim = cfg.model.hidden_dim processor_layers = cfg.model.processor_layers - node_block_invars = ( - cfg.model.hidden_dim, - cfg.model.hidden_dim, - cfg.model.hidden_dim, - comm, - cfg.model.hidden_dim, + + self.exchanger = HaloExchange(comm) + + self.edge_processors = nn.ModuleList( + [ + MeshEdgeBlock( + input_src_node_dim=hidden_dim, + input_dst_node_dim=hidden_dim, + input_edge_dim=hidden_dim, + output_edge_dim=hidden_dim, + hidden_dim=hidden_dim, + ) + for _ in range(processor_layers) + ] ) - edge_block_invars = ( - cfg.model.hidden_dim, - cfg.model.hidden_dim, - cfg.model.hidden_dim, - cfg.model.hidden_dim, - comm, - cfg.model.hidden_dim, + self.node_processors = nn.ModuleList( + [ + MeshNodeBlock( + input_node_dim=hidden_dim, + input_edge_dim=hidden_dim, + output_node_dim=hidden_dim, + hidden_dim=hidden_dim, + ) + for _ in range(processor_layers) + ] ) - edge_layers = [] - node_layers = [] - for _ in range(processor_layers): - edge_layers.append(MeshEdgeBlock(*edge_block_invars)) - for _ in range(processor_layers): - node_layers.append(MeshNodeBlock(*node_block_invars)) - - self.edge_processors = nn.ModuleList(edge_layers) - self.node_processors = nn.ModuleList(node_layers) def forward( self, embedded_mesh_features: Tensor, embedded_mesh2mesh_edge_features: Tensor, - mesh2mesh_edge_indices_src: Tensor, - mesh2mesh_edge_indices_dst: Tensor, + comm_pattern: CommunicationPattern, ) -> Tuple[Tensor, Tensor]: e_feats = embedded_mesh2mesh_edge_features n_feats = embedded_mesh_features - for edge_layer, node_layer in zip(self.edge_processors, self.node_processors): - e_feats = edge_layer( - n_feats, - n_feats, - e_feats, - mesh2mesh_edge_indices_src, - mesh2mesh_edge_indices_dst, - ) - n_feats = node_layer( - n_feats, - e_feats, - mesh2mesh_edge_indices_src, - ) + + # local_edge_list: [E, 2] with [central=mesh_dst, neighbor=mesh_src] + edge_index = comm_pattern.local_edge_list + dst_indices = edge_index[:, 0] # central (aggregation target) + src_indices = edge_index[:, 1] # neighbor (message source) + num_local = comm_pattern.num_local_vertices + + for i, (edge_layer, node_layer) in enumerate( + zip(self.edge_processors, self.node_processors) + ): + with TimingReport(f"processor/layer_{i}/halo_exchange"): + halo_features = self.exchanger(n_feats, comm_pattern) + augmented = torch.cat([n_feats, halo_features], dim=0) + + with TimingReport(f"processor/layer_{i}/edge_block"): + e_feats = edge_layer( + src_node_features=augmented, + dst_node_features=augmented, + edge_features=e_feats, + src_indices=src_indices, + dst_indices=dst_indices, + ) + + with TimingReport(f"processor/layer_{i}/node_block"): + n_feats = node_layer( + node_features=augmented[:num_local], + edge_features=e_feats, + src_indices=dst_indices, + ) + return n_feats, e_feats @@ -243,26 +272,22 @@ def __init__(self, cfg: Config, comm, *args, **kwargs): comm: Communicator object """ super().__init__() - edge_block_invars = ( - cfg.model.hidden_dim, - cfg.model.hidden_dim, - cfg.model.hidden_dim, - cfg.model.hidden_dim, - comm, - cfg.model.hidden_dim, + hidden_dim = cfg.model.hidden_dim + + self.exchanger = HaloExchange(comm) + + self.edge_mlp = MeshEdgeBlock( + input_src_node_dim=hidden_dim, + input_dst_node_dim=hidden_dim, + input_edge_dim=hidden_dim, + output_edge_dim=hidden_dim, + hidden_dim=hidden_dim, ) - self.comm = comm - self.edge_mlp = MeshEdgeBlock(*edge_block_invars) - dst_node_input_dim = cfg.model.hidden_dim - dst_node_output_dim = cfg.model.hidden_dim - m2g_edge_output_dim = cfg.model.hidden_dim self.node_mlp = MeshNodeBlock( - input_node_dim=dst_node_input_dim, - input_edge_dim=m2g_edge_output_dim, - output_node_dim=dst_node_output_dim, - hidden_dim=cfg.model.hidden_dim, - comm=comm, - num_hidden_layers=1, + input_node_dim=hidden_dim, + input_edge_dim=hidden_dim, + output_node_dim=hidden_dim, + hidden_dim=hidden_dim, ) def forward( @@ -270,42 +295,47 @@ def forward( mesh2grid_edge_features: Tensor, grid_node_features: Tensor, mesh_node_features: Tensor, - mesh2grid_edge_indices_src: Tensor, - mesh2grid_edge_indices_dst: Tensor, + comm_pattern: CommunicationPattern, ) -> Tensor: """ Args: mesh2grid_edge_features (Tensor): The edge features from the mesh to the grid grid_node_features (Tensor): The grid node features mesh_node_features (Tensor): The mesh node features - mesh2grid_edge_indices_src (Tensor): The source indices for the mesh2grid - bipartitate edges. These are the indices - of the mesh nodes that are connected to - the grid nodes. - mesh2grid_edge_indices_dst (Tensor): The destination indices for the mesh2grid - bipartitate edges. These are the indices of - the grid nodes that are connected to the - mesh nodes. + comm_pattern (CommunicationPattern): Precomputed communication pattern + for the mesh2grid bipartite graph (partitioned by grid vertex placement). Returns: (Tensor): The updated grid node features """ - e_feats = self.edge_mlp( - src_node_features=mesh_node_features, - dst_node_features=grid_node_features, - edge_features=mesh2grid_edge_features, - src_indices=mesh2grid_edge_indices_src, - dst_indices=mesh2grid_edge_indices_dst, - ) - n_feats = self.node_mlp( - node_features=grid_node_features, - edge_features=e_feats, - src_indices=mesh2grid_edge_indices_dst, - ) + # local_edge_list: [E, 2] with [central=grid, neighbor=mesh/halo] + edge_index = comm_pattern.local_edge_list + dst_indices = edge_index[:, 0] # grid (central, aggregation target) + src_indices = edge_index[:, 1] # mesh/halo (neighbor, message source) + num_local = comm_pattern.num_local_vertices + + with TimingReport("decoder/halo_exchange"): + # Mesh nodes are the neighbors (sources); grid nodes are the central (destination). + halo_mesh_features = self.exchanger(mesh_node_features, comm_pattern) + augmented_mesh = torch.cat([mesh_node_features, halo_mesh_features], dim=0) + + with TimingReport("decoder/edge_block"): + e_feats = self.edge_mlp( + src_node_features=augmented_mesh, # mesh features (local + halo) + dst_node_features=grid_node_features, # grid features (destination side) + edge_features=mesh2grid_edge_features, + src_indices=src_indices, + dst_indices=dst_indices, + ) - n_feats = grid_node_features + n_feats + with TimingReport("decoder/node_block"): + n_feats = self.node_mlp( + node_features=grid_node_features[:num_local], # local grid nodes being updated + edge_features=e_feats, + src_indices=dst_indices, + ) - return n_feats + return grid_node_features + n_feats class DGraphCast(nn.Module): @@ -320,8 +350,7 @@ def __init__(self, cfg: Config, comm, *args, **kwargs): super().__init__() self.hidden_dim = cfg.model.hidden_dim self.output_grid_dim = cfg.model.output_grid_dim - self.comm = comm - self.embedder = GraphCastEmbedder(cfg=cfg, comm=comm, *args, **kwargs) + self.embedder = GraphCastEmbedder(cfg=cfg, *args, **kwargs) self.encoder = GraphCastEncoder(cfg=cfg, comm=comm, *args, **kwargs) self.processor = GraphCastProcessor(cfg=cfg, comm=comm, *args, **kwargs) self.decoder = GraphCastDecoder(cfg=cfg, comm=comm, *args, **kwargs) @@ -340,26 +369,21 @@ def forward( Returns: (Tensor): The predicted output grid """ - input_grid_features = input_grid_features.squeeze(0) input_mesh_features = static_graph.mesh_graph_node_features mesh2mesh_edge_features = static_graph.mesh_graph_edge_features grid2mesh_edge_features = static_graph.grid2mesh_graph_edge_features mesh2grid_edge_features = static_graph.mesh2grid_graph_edge_features - mesh2mesh_edge_indices_src = static_graph.mesh_graph_src_indices - mesh2mesh_edge_indices_dst = static_graph.mesh_graph_dst_indices - mesh2grid_edge_indices_src = static_graph.mesh2grid_graph_src_indices - mesh2grid_edge_indices_dst = static_graph.mesh2grid_graph_dst_indices - grid2mesh_edge_indices_src = static_graph.grid2mesh_graph_src_indices - grid2mesh_edge_indices_dst = static_graph.grid2mesh_graph_dst_indices - - out = self.embedder( - input_grid_features, - input_mesh_features, - mesh2mesh_edge_features, - grid2mesh_edge_features, - mesh2grid_edge_features, - ) + comm_patterns = static_graph.distributed_comm_patterns + + with TimingReport("model/embed"): + out = self.embedder( + input_grid_features, + input_mesh_features, + mesh2mesh_edge_features, + grid2mesh_edge_features, + mesh2grid_edge_features, + ) ( embedded_grid_features, embedded_mesh_features, @@ -367,28 +391,31 @@ def forward( embedded_grid2mesh_edge_features, embedded_mesh2grid_edge_features, ) = out - encoded_grid_features, encoded_mesh_features = self.encoder( - embedded_grid_features, - embedded_mesh_features, - embedded_grid2mesh_edge_features, - grid2mesh_edge_indices_src, - grid2mesh_edge_indices_dst, - ) - out = self.processor( - encoded_mesh_features, - embedded_mesh2mesh_edge_features, - mesh2mesh_edge_indices_src, - mesh2mesh_edge_indices_dst, - ) - processed_mesh_node_features, _ = out - x = self.decoder( - embedded_mesh2grid_edge_features, - encoded_grid_features, - processed_mesh_node_features, - mesh2grid_edge_indices_src, - mesh2grid_edge_indices_dst, - ) - output = self.final_prediction(x) + with TimingReport("model/encode"): + encoded_grid_features, encoded_mesh_features = self.encoder( + embedded_grid_features, + embedded_mesh_features, + embedded_grid2mesh_edge_features, + comm_patterns.grid2mesh, + ) + + with TimingReport("model/process"): + processed_mesh_node_features, _ = self.processor( + encoded_mesh_features, + embedded_mesh2mesh_edge_features, + comm_patterns.mesh, + ) + + with TimingReport("model/decode"): + x = self.decoder( + embedded_mesh2grid_edge_features, + encoded_grid_features, + processed_mesh_node_features, + comm_patterns.mesh2grid, + ) + + with TimingReport("model/final_prediction"): + output = self.final_prediction(x) output = input_grid_features + output return output From ec9b6b8e2463222b9465266d3ba0a20cf29646df Mon Sep 17 00:00:00 2001 From: Shehtab Date: Mon, 13 Apr 2026 02:35:57 -0400 Subject: [PATCH 02/39] Add compute benchmark pdf --- .../cost_model_benchmarks/figures/compute.pdf | Bin 0 -> 19308 bytes 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 experiments/cost_model_benchmarks/figures/compute.pdf diff --git a/experiments/cost_model_benchmarks/figures/compute.pdf b/experiments/cost_model_benchmarks/figures/compute.pdf new file mode 100644 index 0000000000000000000000000000000000000000..f60f93169802441205acbaa7ec6f0d27bcad62cb GIT binary patch literal 19308 zcmdtK1yof{^f#;$%9T(`QNRlb2okrrmrF@^ch@DAZc!;gK)OLfDFp-pDN#VWQ)vMy zB?J*fN|g5;)aNnye~;g{zV)tmz1%g-oP8$t?7h$I-^`pj%*x`DoNz7_ggJi<`tT_P z4uwJOO)oaEO zMeb+d)>bmTY>sn*V!mBUc(_QZxtO?^Ls8!d$|f!@=1z7{9`F*vt!8CnW^HE)MSs8R zWRFubcY*2w&5BC`teAVaK)GdX0Tx8R$70`O8K}We?7;u*0LVAUy^}e>-8c5RRn494 zU7c`1e?b0$_7%*{tW8AiJ%AQrz>fz8L7KVl*;BYPs5{}^k+Ts=idI?nJ z48?xqRm9HD9=L*o>mbK598&bx+sVVd6EXpG4!Kz7ugAYvAy$^PjHX+{G};5BgGh| zIJT+RHp+xq%~z=!>K{YxnMIS-rHPEUbpxmu&&|d6T0M*ofj7mvCbnIBqrCH_VfWI< z&5s{@NVNR5)0{qXzr=LB)(wE|aW8J(_U^pn)C ziGeKBorDhNMv9Zhv!ShLsFL0goS9BoklKBocIoA$1NAB4sM*wwpki6as2h8Bo~zM% z(oN#cJ={hLJ>F8zcZ|NeKFzfJ*a39OrM*ojc5rva-7|m<{{dc&*}&S+jzMAln(Z^n z&(m+#2N)NKm=w&^HyD)LD16tOc48M{qgf1=YEkPR_4`FY~pHu>QUVu zcU9cBOFW2UsrARR8T^u`I}5N$BJSZ z>dH@OOg7Y!G7htf4|2XdCHkd*I{LAY^^9yPq?he|qE+IPNR3b|gTxS@#yzDmqCS`- z(cWSvCzJX180&JRMGEY^;`-=_YW5KqS$>u?EL{BNcq7`2>b^pbtl~3`=H^wa>F3o5 z_?qY_NGPPB6>Qc9j9igIS0AY35fH{@qI0?2F21<B1@spdO}2*(4); z@mg7GlQaXozBCV6amcu?2f};=X(}WXR;Wk$Z6m4|eb#Gslubu^J6`u3j*%(0{yS^O=-+W#nx1AqNcD;9)wpw$`(UKIiZKG;;bRyJvq9#Pb z8%>JyLqC68e3RaNNm;RAJF)gHrRrICUw6V96OJhR?Qk)rNyk3cWwzP%cl3|@wm$HO zMPaulN9VP+qUM+?wrpyS1kxUfdrCk>>xOfDP*z*5OJFp>KGJOcDeLf?&DSg8vON1AEHy7hVHS8NOs@&9}haKIoRTHcXUOXLO*oIN&E*H7wbo<$Z z&-fD$wGf)sB(H_5sV|w(rWd``TQ#hXE?%_Vl3XN#c)t_=rI;5&d;Nu1Rh~zoVVl)Q z9zIR`MlK)u{?A|kDMh&0#0S+OPX4@60A|wuguec2Fx1`r?&{ajlouTn>}NPIU28K_ zwgImNXUU#i-B!@Y8sA!qkc+s^Cjo0q?CcRdRP)Z>14Z3Ay3Ge(m#%w6s} zmXw9B(D}vKp=qNUAust{bE)LzQSuWtQ2J*BCmtlNkJH~;kNmO|wiL&l zZ`yV_V7I-SilvD|`^5|#tHspfqK_>~U9~kwNK!I0ZKW9Q%1M}4vskXZzp62L znCLv_#knX0g!BEG6i@Gk>x~90DmuIj$Im4dt{4?o$UQN#)46(ex@Pv`9+`o&*HuSLjFo0d<=z_FmWxduP;z8i7e&(a#Tx)1qs#H0JmVbC9HknpPf=U`URxWf#Whb+x()+W#ese1& zbe1up8__3z7bbamz5i8!ISZqGBpIKDcOwU~5Bk=ub5L@q6G1^|^5r%|LvrQ|TNX*k z=_jyTbdH#$pgKLN&@9&MD=ce{H6jVlf!XFPx5R@50%Gn2$%{Yuz?B|&`^Dai&=A+D z$5oD28%dbM<~Ke%Iy>->XPmy@_$B{xOhuIYndRAqs8a$_t!&Y(q&V`3FU#UVChT06 zho0&5jy~y(_*~iTQ9?UW>YTt_q3?tdTYVhHheF7seFD2+zUG{h2IAwmCDb$O}106JJ?hFzv2?v6JKQa?{{;pW{nfc*mm%>p{zo zh3ooHUA_qKS}g3=Pa+YXzWnyTrZj(m)eAI^*Huy#^9?>7BtA})aJm-36>DE0WTz;<@Z9yLv9!LzG@=~pl`Wx zxgmx7TC>w(UxUj|mo%&jR`rOeto^)6>NM8hZ<~*vdg;}jbjhtd>CszSdHPb#yY?Lx zUr&Cn@d&F)-EyqI;?h<(B=`PCiH61a)ewQk^X0N3!L0ZCZ8BdDp!wF0EmVfRpV;Ma zwqJg5?pYUS{gLNVp3In0lImmEmP#Fj#^pcDG|nXtd>EQ{lFOhj{1nratouAMt9>|! z`2AY1R=^EW#+K_>`;f+rwe+-!6}K>CyZ+gRh7}U*N4m8O)+n$(seMuLQjl^v<{Fa4 zqQ(>26OWuP@mMv}##&Lwo|J>!8|>)#SdcTng|2tnF1z2#jK|R3DN0_eplYi?o

Y z(N--%_QY?0S?jDDZ+Hj{XRp0Kxn5(v@`*a~5rJni za}jS{w1#4yc9#VwtB7y^!Xq#HIdrY_cKG_xjZrpMR4NK`LV8ZG-cr5Ap2)H&G5ryq z5{>=Y(ib|O1jn+tCmyH=u9l~*d<`$7^Jw?sk3XHI<0HIt#7ow7DNkg#{f+T!U)C^> z$B7+TdOpH@OkR?%w5z8nE7P0tl=%H|40hQ#fP79{F!S==)*HG3#vL!v!q936;CR_e z8(V~5@C!q4K$LO`*-Xp8UB4%7 zJYNW|{*s*9_Zo6e46TVg{a#2ofF?LLuS|G*EcbnoMGE4m6GcFYFacjdQFyRqp4{+) zH14qP`XYb&nViV2I+`g?WnaY`E#gcASL32fse$TiDDt8pj-0I31o$r=0+8sQXy6Nx zJIBAVBpi@N8NXX^0Cd?d7-VX#EUWql@23UsE*$gzr&}N2#n&5ulMNwzsN%u}Izs#x zxgSC=XVvHUB~arXS9M<3dh?(u| zxy_5cw>KAUFFc%%cg(k_iAgSlU+W@R8REE*ZPCW7j?^oRaTQv2KLY?34z_@|w}4C` ziWV=tOpGSIn&Olp)YoyJB)`UqRdv25;_`TID#FGr*PMRGR+Ov+-IDtDw)>U%SW)K& zUiFfh1WNsh*Ci)CDlBB!P?%6Lp3gZB!C9Nj`m1(b(W+K{jdUt>MO zxhvQ^>0#ezh1l9>JaV$>l_phdA_!k2cW+gutcVES*M`|L+)6pGY+xw7{&&X{`3#<) zVxOpV3O4Fj;;JY#Jsi#wZP219x4Sq!-9A~?eto#)i7ZiK>|1xOa_-^$DnsJ40(%=m zG3%yGL$%#!EeEN}JUH3R%>CR>A?FCoHPv*-Z*c5+R+c6T6eZt`;L=aCi75&jwNH~& zj6paSdkZ&|ES#!*zZL{O%ev6c?*lr%!l~+K4!D@Nny@{nu5SGj*5FizST5 zh~vG^A71pQk@(!CaO4uhR~;i#FXql<|K%XD8fLs3vDZHjzSH6=sqO3X()a z!R5mm)2SuLyn~~!KdftvPtbjNIl@%wM)-ATYG~KOJqMn|pe5(ntsC7o;fW!2_~q6* zEw$7I@N+6!fkc7Zh|9JkT(dbYWr;qLD_bF}MJxk2zh2Ja6CJNi4V+8X7vXh=r>q?D zgGi5`iXalId*dk_EO>uHo`v`PC^XIUV`0&N+4LpEKzwocmhn*S)0gUE+~s_d$z+g< z+*{4gO)c|xv@Ooux#&s!{C&S znHUQ~NlR&Q?N_-ktPEaN>b4t3$5f~6Ryd#biwk;?FfLDpb2wEI=EnD!O`oOBVla46 zzjJf6vdM&{O^A zJ#DZWCN0&kP|? z2E)C}Qrs4@SjuoW2d!2zp2VY4Y`(6s3*qVbdBXQ$&yj1>^zOWlQT#2 zZ!ek7_b@fhkC<%LGFY#Z(ep>dpwV%+q+u zx_V*pj!*#9n&|vwDo-EH=S3+OL;SI+4dyPMsn&PK0k($Eq8cHTA6pP^uy%|ChLb?5&8b>_~bl~zKW>oyPk`~34!DR7LopAF^NY-#_kaO*t0pRufo|ytSd=*i|JqX%5JG}ZVALwZ2rx0_>&0JW-|4X-C`xokyn9@5WWDYlhbDvbReOj&oEfj`YipXIL)D z=(`yz*;*T&cBRA+M3J83uR6NT_?Clu-QH^wmN|7HxygVBH~ z7nL3&zklm3oW5FEf0xc2%RPY4v|znM}146G4G{OmK;ZtbD?{HCi)w?)X7 z_X?TmD3|L_Fv;mpzi$+$HY`^j5)GVbkTfJp_FcI-n6h=Ic>2<)WLDfIbv-BrT;;rr z3g*nFG^PsD=auEBmA9$=@ud63Ekf(ctqPV-IE%To zLw$uX3`fPS!wmzno^Ld`+M=G7p-WivZc5$vUy~_0bLSAI&|y0XC0Yd>k}es z;&Yj!0daWte#NY=#s-FXRHD>0X0xyOPld=@P(Q!Jl1uU8yyJv;I=_N!Dp~2vSCavA zc%K&NmJZZB|_?;VJ3CKibMd+wjDG!dp+I$Jl-n^)-th0?elYpCA1ZWof!A+osDdX0vD;p0)#N!*JAu_=?zPkkfLun1$3v6+`? zMR!{{&7_RwsZYpXscxQ5c5;8#eq}VizB;gdhj&ID`;h%g6#8V6#FevklJrYOJ(&-! zI&|2_Frf+t6mbp7P49Ka+oo@?6zb0maMX2fHdH*PJ*PE3&)zcJ?mwa$^Y#VyQ(=7H zjzeikfajFdH5O?H*N0Gz?7{9)HP0_%!B-3GY{|%nP?^N2nz($3Kmk)KRoT-ra{B+2dRuq(dxedK)J%RrTJZ z@8{tn!nkN<-oKi--1_)vZelk-n?=z(7ga91Q-KWht!p_S-eeA65$)M3y2!jP zJ$wu4#QNT<`E7sebmZ5GFnqa)M^+3UCkyQcxs0pS_#{lt6rJ*|Lo^gOl__6VHdGUz z_+%p?#^xL~uxQlKef&zi4Yh~t4gr1UIEx&q=E$oY#gzhqk>1U7xK3=_!2%f1>Cf--n%3?US;2hpzCE z1s>4;0hr{k9u^SC{o`f_-&e6V-3#2<8&YkjJSNOs{)E{S$8>}nzJ(1M+luDNWbv|e z(%R)v3XrNVn3SRPO_(xt^XoXgh{naGXCWd`T7fX{C2hS*9(5+whwolKj$B}{X9yc=yq*U<_$&ATzr15`3zM| zlQyBCN~Ik9M&zPRI%|5>g}XOxgf0h(CtR?c46duY&`Y{hxokPD!C+-D{=gQOBPrYR z@;QluLiDk1aU#od)O+f$-aPVP~ z`1yke<>=vU@2Am8ijf}}+Fnx7C3z$qMFxP?yYT8n; zbYg#(@`hNbvwGHgTR7tFJZ}_hTFA=AD-JiT7wNm>=r9JVxh-~2KlIIvrK>AP&Nzwe z8H{hyOFWxrhMpy&9#k=XpYq1vsM^hizU%_Iv*;ZWE((UO1@`+V?kMM%WZqL%amW>V zFrRqmHZ*p%hH%8}=L zPx?v`!($6)m<#Sb@_im6x@tH5?q)sop_nSP+mkHFW$8XTih_?W1$oH?gA1>ffD@)H zxZD>!%7Xn|5Ak$hcGrWNW4RTr#~Lj-3!PF`4Q|W68{2-# zwmlZYXVHhxL0&m&Mxnsl@k~ZFlqMFjNPpjr@kYKtYbM{knL>9yQ_{#wrB_K?v49=J@XT35;p6Pf-w~g#2&M~T~Y^mmOnTi?&o$X0K`9U!K3b{^k^(vMl zGs(H;=(&I=_{j9bxEnHv%AB`x3>1pQuC%1OS@ln@m^Mu7_6dZI*?pF@FlX^6i7*o# zJM9w$eRt&XX^vPlngXHIjkiqFMxtvxlzUbs#P4G4AdyRkW@BFGWXan>H;fqP+GNg+ z)g3F%1ml1{xlh_2Mgc;Jxk+yCqeVTZ6!;lC=^LNrC$fy!a3^2rnoj_bSo+EwlPjI~as^mF> zI2J=H?W3}|G35y1>M1AAwlM-S2UhRVUiFmtm+sb4G8$}SS>#vuNG(|+bFTK`Qa6v| zJ~Ig%z~up;6p8wakwrrF8~%H)=HM$ml@R%{G!OK`ddczwU6Vw{HXa>n_s!aiDw&zi z55+3oU=+*0zl(A@&v<`JhKRr|E~4d%Tz=S6R1W3c(bbvab;w#P+XNP(mlp*5HY_Ea zBt!c{c=;HjbzflpreEBn*PYSF6i?eS6<2D#et&K%gx|HePHgE|16^ZbPk zA<;;6z)=^g!w_f{7#elNlp=sqjO*U9AG2iPJf|o3+MSJo2=U~_A%fGShyz+b0Mlar z0@5N870SR29d(%e7Bvgxtp(x!480~dxKfnslr(W;`{5E|pHk)&zaUGWM5yhFCMdOd zaZ)giI9_=;3Bn%XXte=FmrBW;Mso+hl1 zkp7|PJqF&JPwNV%jN6~OP^_rU3^b>Z6XkfwD#SlM?%7{j(_dpX`Rv_Wxq;wIT|K8# z$)3^Cj32<-0f_2v)(s96Rs0D}u2b?HM>~)7#XCS4)9T^~4^>){upZFtK|qR!_rFCV z5e3XTpvS~F)BKoxc4q}*WX<7~duWnwGqhVm&KG}Kru*a^bCg$MH`>;^%*q$sv6}p6 z2`D8lB(`bH7aw`Zuhy7!CHV3?ocZjNm;8iHXS)rsxn_=;p|0w#><;ohai3Jix7~DN zZ%C1vLx%I6X>BIkZ5vLrN+uV)C*2H&DzU0bH7!F&J{)BYyAk*BqC=|9MxPvv6#t7T zMdECT@H=AP<}S7MS@NZs8MirKWBaQorAJP*MlqbG^K(np-KU6!bMFKn)x>hf ztkeDX-P>XgpNu{|Hyy_&2m7|Q(;e$tR2$ny^JhGL+V(Z8Wh{Tl>*A7ziM5Vg2T?`p zGYes&F*WSyL%~buAbcAK^z|Tk1pf=O5mc@W!-o@h7zCWruJv5bwWJP_yT7+Cy}iNM zM=hbs&XW|)#_bx*tfMx;cRjw5^y;lQc4la1b_wp7BIeCr$m)&Yx0IXX1Mwwds01sl z*Ks6!-dVbra_Bn_>=aM$jWr;5NNip3Gz002^`O03#zW;$rxvo)#C1s&A&kVb3WR%) zoFeWThMW~@SQVp;iY9%_v#0%K&6%W7ErK=5&c0^aFeR51?<8&Xohl{+?P4zSlh!;U z3O%*DoIcMg48NW&WqL|1yb;{!ct>lnntZKsti0Ayzuan&xLh+&+M z{OVmWe~%9TA780m@vo9J4~LUO)w4zw2W_{h^CB|kY0*B zLm7dg^wuDyz*z+|zqFPhx+gJh^LQ(&dd05uhnG)ioHr;T#uXPlbr3Yc=V|h!dix26 z3yzk$t0DL9FagEZi*4`sb}Mc-4jVVY_Lug5-3wsD zes9lX0j1mb0$uj5Fh3w5PK33wj!@r>3Kf$#UP{DGF{>*T2%TrDcq)TEk*kcsSF(~P zMy6!RO~x(MmRU<>*QQX^Z(b!<`MG+hF>G9N5aDeiDP$K=G(n3J~VtBotRw);^f zp6Yjc-SectDV>78{9;1}e^ZcCGL7}tmz6Axfzg-;Y~WAz0h;==f9StZpF|{<91ngB z-?`e4Ra)O~Gryy(Zf7s|=%7@_!VwhhQ)-2FPf9DHg<{F{r@dF^Ib@^7>!_pYlpo-N z_%BGyOXH%Y%js`2-O(;FNK8^^FHxN;wQj0^KB7C&Djwz#wi~=asJ^STs2}@5;bvhp zeCM=!75~jpnJq(!FmeNBUJOUrXF^WGS^9TKclSjeYnR zZQ0%cJAKL9Op$kU83pcTwRN&UrXy15o!IeA7VkWAaXPEv_#my5Cpn}7DP!@jM!Gnu z{;C0Q$yZ`Ei?XY%Jx&?@xyM$DE#-Q9v5(TK<8`cDk}0aP{Ck=el@x2Vknb3|MClu4#UugHzmRMhTx(tQ2A_I{P$>+4|-=j+Y~sO10@g+%;?{qX>Zp8#0#?O}T& zoc%$m`(s{SNX=*aD-_jK@-X8WHg264ay9+Tud*9v_0DZF7mH3rEZ^x=ZQ7(Rq3q!- zl#KFFg7R$DO_@VCFK->C&qhL(=Y*S|nIcs2 z{LT+MTRR>Lc47(k&Z?YQ+t!lJp-BTdr671;?W83 zZWN2I->Quk5R!to`7EJ>Z8Hd1=IFDTI~ytaidv3TS=<@Wm^ZGc4`d6~I$n}PZ?mRH z5tl_*Fher}o9KY4nWj$_5%j+`HTQe}W`?J*K4@7b#-PFBeSDYuu|BJ>XQ|sKaa7Sh zQV)vBNvwHmgKE|U*mK&7L0pOfT{P~thdsK)Z^ zs=IXr2B?=Uxonrljgn(#-OU3UiP|2LFRR{?&Z*L57+6U{6jtz9;C@6d`uWrofLl?LZ=x_?Ip5$Pb%!j#%XnJXx}oej6-O!a+Hso@2E5}_8Nzq%DHvq z1vth0=vj#|5ntx$L~PSjm-W|My~X#w4u@i%^RX}J>FVHK2%TMyZRS8Mcxtb^d>m}K zc^tKpu{$?FOpY=Si*7zzFp9mEk))#OQ6UyX(gkl}_(+=8Ktm+EBS85)f|-IYw)HhB zNxHWYeP(D1dg932j7zpuM@gqk8VoCasbo)GiQYV%D?jhMuvfTRRx>noj&9kf`2gh{ z1SS8*TXU+&IN^yLd)9Mt-QVKDYh1#`z}dZ_W9>%?NkW8Obf@oy-i2~#RlSLNXw~o< z3HN(V5a@Mr$X0t+SRE>ZQ|3u`TlHM0A9hS6>WmpIBDTmU+YMzKn{S$e(^ysvnfISq zGAI;G5^2lVZ)B!8B|l)oWJ!7{3BgtBEQo>5NopsUMA~VdL6V1~0;CxN@{+T<@kPi( zA4!IxohUzNb4R$`(g@3eFFAAC8ucjfUg=%Z&R?62r8jWoVm-y3PrLpJjdx1LI)6=A zy7r45*(c}6F|R%3=DMk0&RpwiTdP+$Zb}}q6GjhAzpe_J`5am}H@CfuXJo7wbbt&G zz?yK(UpPvINEuSG{k^05Q#t;~+Bt#ec~7lYo)?H6b@knyg*pfb^g5pBaxt3^4ZoO7 zJ5yT%#ZC9|?zKG;YJ`i7XK-pKxV4OQk~`NU;77_O+eQR3YROvkkCB@toU5sYG8#l| zYkt&;S?{UfACGuiYGEXpQ5(E;tf}a*md}@`ExvN~+lDV^6WEz?6oG5Li;Svh{`3=N z>Pt>}w@b#8FFH`I@AiqkUnVGH5_o+;w-3OvynnI6$OAxZ9~iOtHKnPO)6YkaGC+o( zH_15QUXRBZs!j{P;X4_A-Rzyop`LKKaVUju`2YqFK%xlvU+s-zNr5Th(kAe_$@x%j52cI?U0%=FA?+6?g z8|uqZ5j#cS+n=p=p=bK5=);86{7zkPseE@3JFGm?eO~p6z_CVFA(JY%_s9O z<5ls0fhit3MB&(R#O$tYY{_h+eerq;q1MjV!*`|Bau?h4{QVQVu~iVz-czo->d#sUi8}uHcmJUV2=|8%;4b|EDAVw2KYbV z(_=5!D z7Pq#rFb7W7fk)}|px~abvx5l^IOPZQ&jCDgXJc;R@;&KfZE58KMFBQsb=ASj+zyHY zX5F6VPWDi4dpn?0F780e#mdRt9Kvk@9M6MtTUfgRbZ3BHG_W};21Nr;hy>6ND7O?8 z0|-L~3cf?207wN2zSg4xxy>Nl=D^83Agd)10WDht5s={ul-mZ% zZ42eL1KNhd;lOwUGzTaM;MEnl21>gFVgtygfUXR{Q+?c8)@CkNz&Stg!T9+EXkqWdcqia6uIV~kh?z#-TIc)^K40{fR@CJxf(`x1sh zxV3*I0T}_uH?46Zc9u5g0OQ-2$27; z^&kO6hysih1@3>bJ2Pz06-0t6@;*hJ>x z0V)BWIt*ah7+!z@!1DHE@ILUa9vV2FiUG=lbz*_x9?koW2FClHL1G|a3^DQ zfv~R+u(@vm!t-5c|5WuY0WkQj3gA^ouow@>6~H-I95jO;>iQ7_DueuBWWNR!2FL`w zPzkaR+QC2dgSkNBzA*tcf!F`k0x|`ILH6%`uK)`SG(Q4p2;Xb~w1Z#6zN`Ve0Bz!@ zux}5bP5cP^w({L3KzsN(fPC#=19gD5@gsmeN5gmk9Rj?9HgG_|@&Y;m?E-MjKZhSS z@goL>0_YvGAA{+i!oIG*-3Rmh6F|EFwF<`nzJ&Vo{8mpeX{f0aFlQ(B6Ly3!te2n*Mgz1oR~P`M}{9a9;aq z928)0KV=3y0pC*QfTF<^;4wg70X)R}8TQ==AlUs2OTg~GrGEIG@4nR<3T*c6XSo6e z#=!k*Y@n#`&Dlax-%1iSyk^D2NN5pWF%PXT!JzpG9G z0R^6oZ&UyceV3d9fQi0~PJsss3;P56 zVdd6pr$nbd&G@cK)bn5USydo!7Riub$QHS!AvgfzeyC8Ql_B{uQP#*!d4bUzCOLIF znrG&U1HMD@b1?@q3#h3H?h4S%k1Bwy`wdvz+lhl;aJeAPhXCHK zfT7?>7y<#~L2$xQtS}fW+fPHNhm*Ml1Oz4!07HNO10XUF;I}QH|Ionk82p3UeW$@O zzz0`;rvWGgFq2>Np+T7ZJBR6aayLr=dZ3^;>;dV9N9x4GjaP zT))#0AiVmuEMNz}^9(~{{*V Date: Mon, 13 Apr 2026 13:01:13 -0400 Subject: [PATCH 03/39] Add cost model benchmarking --- .../analysis/__init__.py | 0 .../analysis/compute_predictions.py | 160 ++++++ .../analysis/fit_overhead.py | 226 ++++++++ .../analysis/fit_primitives.py | 255 +++++++++ .../benchmarks/__init__.py | 0 .../benchmarks/bench_compute.py | 232 ++++++++ .../benchmarks/bench_concurrency.py | 227 ++++++++ .../benchmarks/bench_end_to_end.py | 504 ++++++++++++++++++ .../benchmarks/bench_gather.py | 181 +++++++ .../benchmarks/bench_pingpong.py | 161 ++++++ .../benchmarks/common.py | 176 ++++++ .../cost_model_benchmarks/figures/compute.pdf | Bin 19308 -> 0 bytes .../run_local_compute_tests.sh | 12 + .../visualization/__init__.py | 0 .../visualization/plot_ablations.py | 160 ++++++ .../visualization/plot_compute.py | 172 ++++++ .../visualization/plot_gather.py | 120 +++++ .../visualization/plot_pingpong.py | 147 +++++ .../visualization/plot_tipping_point.py | 141 +++++ .../visualization/plot_validation.py | 127 +++++ 20 files changed, 3001 insertions(+) create mode 100644 experiments/cost_model_benchmarks/analysis/__init__.py create mode 100644 experiments/cost_model_benchmarks/analysis/compute_predictions.py create mode 100644 experiments/cost_model_benchmarks/analysis/fit_overhead.py create mode 100644 experiments/cost_model_benchmarks/analysis/fit_primitives.py create mode 100644 experiments/cost_model_benchmarks/benchmarks/__init__.py create mode 100644 experiments/cost_model_benchmarks/benchmarks/bench_compute.py create mode 100644 experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py create mode 100644 experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py create mode 100644 experiments/cost_model_benchmarks/benchmarks/bench_gather.py create mode 100644 experiments/cost_model_benchmarks/benchmarks/bench_pingpong.py create mode 100644 experiments/cost_model_benchmarks/benchmarks/common.py delete mode 100644 experiments/cost_model_benchmarks/figures/compute.pdf create mode 100644 experiments/cost_model_benchmarks/run_local_compute_tests.sh create mode 100644 experiments/cost_model_benchmarks/visualization/__init__.py create mode 100644 experiments/cost_model_benchmarks/visualization/plot_ablations.py create mode 100644 experiments/cost_model_benchmarks/visualization/plot_compute.py create mode 100644 experiments/cost_model_benchmarks/visualization/plot_gather.py create mode 100644 experiments/cost_model_benchmarks/visualization/plot_pingpong.py create mode 100644 experiments/cost_model_benchmarks/visualization/plot_tipping_point.py create mode 100644 experiments/cost_model_benchmarks/visualization/plot_validation.py diff --git a/experiments/cost_model_benchmarks/analysis/__init__.py b/experiments/cost_model_benchmarks/analysis/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/experiments/cost_model_benchmarks/analysis/compute_predictions.py b/experiments/cost_model_benchmarks/analysis/compute_predictions.py new file mode 100644 index 0000000..3e69f4d --- /dev/null +++ b/experiments/cost_model_benchmarks/analysis/compute_predictions.py @@ -0,0 +1,160 @@ +"""Analysis — Apply Assembled Cost Model and Compare to Measurements. + +Reads: + - data/fitted_primitives.json + - data/fitted_overhead.json + - All end-to-end JSON run files + +For every run, computes the predicted T_layer: + + T_layer = T_comp + max(T_intra, T_inter) + T_buffer_copy + T_overhead + +and records the measured median, predicted value, absolute error, and relative +error (as a fraction). + +Also computes aggregate MAPE for the fit subset and the held-out subset +(using the same --fit-filter expression as fit_overhead.py). + +Outputs ``data/predictions.json``. + +Usage:: + + python -m analysis.compute_predictions \\ + --primitives data/fitted_primitives.json \\ + --overhead data/fitted_overhead.json \\ + --e2e-runs data/e2e_*.json \\ + --fit-filter "world_size <= 8" \\ + --output data/predictions.json +""" + +import argparse +import json +from pathlib import Path + +import numpy as np + +from analysis.fit_overhead import ( + apply_filter, + load_e2e_runs, + predict_layer_time, +) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def parse_args(): + p = argparse.ArgumentParser(description="Apply cost model and compute predictions") + p.add_argument("--primitives", type=str, required=True) + p.add_argument("--overhead", type=str, required=True) + p.add_argument("--e2e-runs", nargs="+", required=True, metavar="FILE") + p.add_argument("--fit-filter", type=str, default="world_size <= 8", + help="Same expression used when fitting overhead (determines train/test split)") + p.add_argument("--output", type=str, default="data/predictions.json") + return p.parse_args() + + +def main(): + args = parse_args() + + with open(args.primitives) as f: + primitives = json.load(f) + with open(args.overhead) as f: + overhead_data = json.load(f) + + T_overhead = overhead_data.get("overhead_seconds", 0.0) + + all_runs = load_e2e_runs(args.e2e_runs) + fit_runs, held_runs = apply_filter(all_runs, args.fit_filter) + + fit_set = set(id(r) for r in fit_runs) + + prediction_entries = [] + for r in all_runs: + T_model_base = predict_layer_time(r["config"], r["per_rank_stats"], primitives) + T_pred = T_model_base + T_overhead + T_meas = r["measured_median"] + abs_err = abs(T_meas - T_pred) + rel_err = abs_err / T_meas if T_meas > 0 else float("nan") + + # Decompose prediction for ablation figures + net = primitives.get("network", {}) + intra_bytes = r["per_rank_stats"].get("c_intra_bytes", 0) + inter_bytes = r["per_rank_stats"].get("c_inter_bytes", 0) + + def net_time(nbytes, mode): + params = net.get(mode, None) + if params is None or nbytes == 0: + return 0.0 + return params.get("latency_seconds", 0.0) + nbytes / params.get("bandwidth_bytes_per_sec", 1e10) + + T_intra = net_time(intra_bytes, "intra") + T_inter = net_time(inter_bytes, "inter") + T_comm = max(T_intra, T_inter) + + F = r["config"]["feature_dim"] + send_bytes = r["per_rank_stats"].get("send_total", 0) * F * 4 + gath_params = primitives.get("gather", {}).get("clustered", {}).get("gather", None) + T_buffer_copy = 0.0 + if gath_params and send_bytes > 0: + B_g = gath_params.get("bandwidth_bytes_per_sec", 1e12) + T_buffer_copy = gath_params.get("intercept_seconds", 0.0) + send_bytes / B_g + + entry = { + "source_file": r["source_file"], + "config": r["config"], + "partition_stats": r["per_rank_stats"], + "measured_median_seconds": T_meas, + "predicted_seconds": T_pred, + "absolute_error_seconds": abs_err, + "relative_error": rel_err, + "in_fit_set": id(r) in fit_set, + "breakdown": { + "T_comp_seconds": T_model_base - T_comm - T_buffer_copy, + "T_comm_seconds": T_comm, + "T_buffer_copy_seconds": T_buffer_copy, + "T_overhead_seconds": T_overhead, + }, + } + prediction_entries.append(entry) + + # Aggregate MAPE + def mape(entries): + errs = [e["relative_error"] for e in entries + if not np.isnan(e["relative_error"])] + return float(np.mean(errs)) if errs else float("nan") + + fit_entries = [e for e in prediction_entries if e["in_fit_set"]] + held_entries = [e for e in prediction_entries if not e["in_fit_set"]] + + mape_fit = mape(fit_entries) + mape_held = mape(held_entries) + mape_total = mape(prediction_entries) + + print(f"[predictions] Fit MAPE={mape_fit*100:.2f}% " + f"Held-out MAPE={mape_held*100:.2f}% " + f"Total MAPE={mape_total*100:.2f}%") + + result = { + "fit_filter": args.fit_filter, + "T_overhead_seconds": T_overhead, + "aggregate": { + "mape_fit_set": mape_fit, + "mape_held_out": mape_held, + "mape_all": mape_total, + "num_fit": len(fit_entries), + "num_held_out": len(held_entries), + }, + "predictions": prediction_entries, + } + + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w") as f: + json.dump(result, f, indent=2) + print(f"[predictions] Written to {out_path}") + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/analysis/fit_overhead.py b/experiments/cost_model_benchmarks/analysis/fit_overhead.py new file mode 100644 index 0000000..426f440 --- /dev/null +++ b/experiments/cost_model_benchmarks/analysis/fit_overhead.py @@ -0,0 +1,226 @@ +"""Analysis — Fit Library Overhead Bias (T_overhead). + +Reads ``fitted_primitives.json`` and the small-K subset of end-to-end runs. +For each run it computes the model-predicted T_layer (without overhead), then +fits a single scalar T_overhead that minimises MAPE: + + MAPE = mean( |T_measured - (T_model + T_overhead)| / T_measured ) + +The subset used for fitting is controlled by ``--fit-filter``, which is +evaluated as a Python expression where each run's config fields are available +as local variables (e.g. ``"world_size <= 8"``). + +Outputs ``data/fitted_overhead.json``. + +Usage:: + + python -m analysis.fit_overhead \\ + --primitives data/fitted_primitives.json \\ + --e2e-runs data/e2e_*.json \\ + --fit-filter "world_size <= 8" \\ + --output data/fitted_overhead.json +""" + +import argparse +import json +from pathlib import Path + +import numpy as np + + +# --------------------------------------------------------------------------- +# Cost model (without overhead) +# --------------------------------------------------------------------------- + +def predict_layer_time(run_config: dict, per_rank_stats: dict, + primitives: dict) -> float: + """Predict T_layer for one rank using the assembled primitive model. + + T_layer = T_comp + max(T_intra, T_inter) + T_buffer_copy + + Parameters + ---------- + run_config : dict + Config block from the end-to-end JSON (feature_dim, model, etc.) + per_rank_stats : dict + Stats for rank 0 from per_rank_stats list. + primitives : dict + Loaded fitted_primitives.json. + """ + F = run_config["feature_dim"] + model_type = run_config.get("model", "gcn") + + n_local = per_rank_stats.get("n_local", 0) + n_halo = per_rank_stats.get("n_halo", 0) + n_total = n_local + n_halo + + # A rough edge count estimate: use avg_degree * n_local as a proxy + avg_degree = run_config.get("avg_degree", 20.0) + n_edges_local = int(n_local * avg_degree) + + # T_comp + comp_params = primitives.get("compute", {}).get(model_type, {}).get("forward", None) + if comp_params: + T_comp = (comp_params["coeff_V"] * n_total + + comp_params["coeff_E"] * n_edges_local + + comp_params["intercept"]) + T_comp = max(T_comp, 0.0) + else: + T_comp = 0.0 + + # T_intra and T_inter + intra_bytes = per_rank_stats.get("c_intra_bytes", 0) + inter_bytes = per_rank_stats.get("c_inter_bytes", 0) + + net = primitives.get("network", {}) + + def net_time(nbytes: int, mode: str) -> float: + params = net.get(mode, None) + if params is None or nbytes == 0: + return 0.0 + B = params.get("bandwidth_bytes_per_sec", 1e10) + t_L = params.get("latency_seconds", 0.0) + return t_L + nbytes / B + + T_intra = net_time(intra_bytes, "intra") + T_inter = net_time(inter_bytes, "inter") + T_comm = max(T_intra, T_inter) + + # T_buffer_copy (gather of send buffer) + send_bytes = per_rank_stats.get("send_total", 0) * F * 4 + gath_params = primitives.get("gather", {}).get("clustered", {}).get("gather", None) + if gath_params and send_bytes > 0: + B_g = gath_params.get("bandwidth_bytes_per_sec", 1e12) + T_buffer_copy = gath_params.get("intercept_seconds", 0.0) + send_bytes / B_g + else: + T_buffer_copy = 0.0 + + return T_comp + T_comm + T_buffer_copy + + +# --------------------------------------------------------------------------- +# Load helpers +# --------------------------------------------------------------------------- + +def load_e2e_runs(paths: list) -> list: + runs = [] + for p in paths: + with open(p) as f: + data = json.load(f) + config = data.get("config", {}) + for meas in data.get("measurements", []): + per_rank_stats = meas.get("per_rank_stats", [{}]) + rank0_stats = per_rank_stats[0] if per_rank_stats else {} + trials = meas.get("rank0_trials_seconds", []) + if not trials: + continue + runs.append({ + "config": config, + "per_rank_stats": rank0_stats, + "measured_median": float(np.median(trials)), + "source_file": str(p), + }) + return runs + + +def apply_filter(runs: list, filter_expr: str) -> tuple: + """Split runs into fit and held-out sets using filter_expr.""" + if not filter_expr: + return runs, [] + fit_runs, held_runs = [], [] + for r in runs: + env = dict(r["config"]) + env.update(r["per_rank_stats"]) + try: + if eval(filter_expr, {"__builtins__": {}}, env): + fit_runs.append(r) + else: + held_runs.append(r) + except Exception as e: + print(f"[fit_overhead] Warning: filter eval failed for run ({e}), including in fit set") + fit_runs.append(r) + return fit_runs, held_runs + + +# --------------------------------------------------------------------------- +# Scalar overhead fitting +# --------------------------------------------------------------------------- + +def fit_overhead_scalar(fit_runs: list, primitives: dict) -> tuple: + """Fit T_overhead to minimise MAPE on fit_runs. Returns (overhead, mape_in_sample).""" + if not fit_runs: + return 0.0, float("nan") + + residuals = [] + for r in fit_runs: + T_model = predict_layer_time(r["config"], r["per_rank_stats"], primitives) + residuals.append(r["measured_median"] - T_model) + + # Optimal scalar overhead that minimises sum of |err - overhead| / T_meas + # is the weighted median; for uniform weights it's just the median of residuals. + overhead = float(np.median(residuals)) + + mape = float(np.mean([ + abs(r["measured_median"] - (predict_layer_time(r["config"], r["per_rank_stats"], primitives) + overhead)) + / r["measured_median"] + for r in fit_runs + ])) + return overhead, mape + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def parse_args(): + p = argparse.ArgumentParser(description="Fit T_overhead from end-to-end runs") + p.add_argument("--primitives", type=str, required=True, + help="Path to fitted_primitives.json") + p.add_argument("--e2e-runs", nargs="+", required=True, metavar="FILE") + p.add_argument("--fit-filter", type=str, default="world_size <= 8", + help="Python expression evaluated per run; True → fit set") + p.add_argument("--output", type=str, default="data/fitted_overhead.json") + return p.parse_args() + + +def main(): + args = parse_args() + + with open(args.primitives) as f: + primitives = json.load(f) + + all_runs = load_e2e_runs(args.e2e_runs) + print(f"[fit_overhead] Loaded {len(all_runs)} run(s)") + + fit_runs, held_runs = apply_filter(all_runs, args.fit_filter) + print(f"[fit_overhead] Fit set: {len(fit_runs)} Held-out: {len(held_runs)}") + + overhead, mape_in = fit_overhead_scalar(fit_runs, primitives) + print(f"[fit_overhead] T_overhead = {overhead*1e3:.3f} ms in-sample MAPE = {mape_in*100:.2f}%") + + result = { + "overhead_seconds": overhead, + "fit_filter": args.fit_filter, + "num_fit_points": len(fit_runs), + "num_held_out": len(held_runs), + "in_sample_mape": mape_in, + "fit_subset_runs": [ + { + "source_file": r["source_file"], + "world_size": r["config"].get("world_size"), + "feature_dim": r["config"].get("feature_dim"), + "measured_median_seconds": r["measured_median"], + } + for r in fit_runs + ], + } + + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w") as f: + json.dump(result, f, indent=2) + print(f"[fit_overhead] Written to {out_path}") + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/analysis/fit_primitives.py b/experiments/cost_model_benchmarks/analysis/fit_primitives.py new file mode 100644 index 0000000..02c9233 --- /dev/null +++ b/experiments/cost_model_benchmarks/analysis/fit_primitives.py @@ -0,0 +1,255 @@ +"""Analysis — Fit Primitive Cost-Model Parameters. + +Reads JSON outputs from benchmarks 1.1, 1.3, and 1.4, fits the following +parameters by linear regression on **medians** of per-trial times: + +* Network: T = t_L + bytes / B → fit (t_L, B) for intra and inter +* Compute: T = coeff_V * |V| + coeff_E * |E| + intercept + (separate fits for GCN and edge-conditioned models) +* Gather: T = intercept + bytes / B_gather + (separate fits for contiguous, clustered, random distributions) + +Writes ``data/fitted_primitives.json``. + +Usage:: + + python -m analysis.fit_primitives \\ + --pingpong-intra data/pingpong_intra_*.json \\ + --pingpong-inter data/pingpong_inter_*.json \\ + --compute-gcn data/compute_gcn_*.json \\ + --compute-edge data/compute_edge_*.json \\ + --gather-contiguous data/gather_contiguous_*.json \\ + --gather-clustered data/gather_clustered_*.json \\ + --gather-random data/gather_random_*.json \\ + --output data/fitted_primitives.json +""" + +import argparse +import json +from pathlib import Path + +import numpy as np +from scipy import stats as sp_stats + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def load_json_files(paths: list) -> list: + records = [] + for p in paths: + with open(p) as f: + records.append(json.load(f)) + return records + + +def median_of_trials(trials: list) -> float: + return float(np.median(trials)) + + +def r_squared(y_true: np.ndarray, y_pred: np.ndarray) -> float: + ss_res = np.sum((y_true - y_pred) ** 2) + ss_tot = np.sum((y_true - np.mean(y_true)) ** 2) + return float(1.0 - ss_res / ss_tot) if ss_tot > 0 else 1.0 + + +def linear_fit(x: np.ndarray, y: np.ndarray): + """Fit y = slope * x + intercept via scipy linregress. Returns dict.""" + result = sp_stats.linregress(x, y) + y_pred = result.slope * x + result.intercept + r2 = r_squared(y, y_pred) + return { + "slope": float(result.slope), + "intercept": float(result.intercept), + "r_squared": r2, + } + + +# --------------------------------------------------------------------------- +# Network fit: T = t_L + bytes / B +# --------------------------------------------------------------------------- + +def fit_network(records: list) -> dict: + """Fit (t_L, B) from ping-pong records (one mode per call).""" + bytes_arr = [] + time_arr = [] + for rec in records: + for meas in rec["measurements"]: + nbytes = meas["params"]["message_bytes"] + t_med = median_of_trials(meas["trials_seconds"]) + bytes_arr.append(nbytes) + time_arr.append(t_med) + + bytes_arr = np.array(bytes_arr, dtype=float) + time_arr = np.array(time_arr, dtype=float) + + # T = t_L + bytes / B → T = intercept + slope * bytes + # so slope = 1/B, intercept = t_L + fit = linear_fit(bytes_arr, time_arr) + bandwidth = 1.0 / fit["slope"] if fit["slope"] > 0 else float("nan") + latency = fit["intercept"] + return { + "bandwidth_bytes_per_sec": bandwidth, + "latency_seconds": latency, + "r_squared": fit["r_squared"], + "_raw_slope": fit["slope"], + "_raw_intercept": fit["intercept"], + "_num_points": len(bytes_arr), + } + + +# --------------------------------------------------------------------------- +# Compute fit: T = coeff_V * |V| + coeff_E * |E| + intercept +# --------------------------------------------------------------------------- + +def fit_compute(records: list, timing_key: str = "forward_trials_seconds") -> dict: + """Fit compute cost as a function of |V| and |E|. + + Uses multiple linear regression: T = a * |V| + b * |E| + c + """ + V_arr, E_arr, T_arr = [], [], [] + for rec in records: + for meas in rec["measurements"]: + V_arr.append(meas["params"]["num_vertices"]) + E_arr.append(meas["params"]["num_edges"]) + T_arr.append(median_of_trials(meas[timing_key])) + + V_arr = np.array(V_arr, dtype=float) + E_arr = np.array(E_arr, dtype=float) + T_arr = np.array(T_arr, dtype=float) + + # Design matrix: [V, E, 1] + A = np.column_stack([V_arr, E_arr, np.ones_like(V_arr)]) + result, _, _, _ = np.linalg.lstsq(A, T_arr, rcond=None) + coeff_V, coeff_E, intercept = result + T_pred = A @ result + r2 = r_squared(T_arr, T_pred) + + return { + "coeff_V": float(coeff_V), + "coeff_E": float(coeff_E), + "intercept": float(intercept), + "r_squared": r2, + "_num_points": len(T_arr), + } + + +# --------------------------------------------------------------------------- +# Gather fit: T = intercept + bytes / B_gather +# --------------------------------------------------------------------------- + +def fit_gather(records: list, timing_key: str = "gather_trials_seconds") -> dict: + k_arr, T_arr, F_arr = [], [], [] + for rec in records: + F = rec["config"]["feature_dim"] + for meas in rec["measurements"]: + k = meas["params"]["k"] + t_med = median_of_trials(meas[timing_key]) + k_arr.append(k) + T_arr.append(t_med) + F_arr.append(F) + + k_arr = np.array(k_arr, dtype=float) + F_arr = np.array(F_arr, dtype=float) + T_arr = np.array(T_arr, dtype=float) + + bytes_arr = k_arr * F_arr * 4.0 # float32 + fit = linear_fit(bytes_arr, T_arr) + bandwidth = 1.0 / fit["slope"] if fit["slope"] > 0 else float("nan") + return { + "bandwidth_bytes_per_sec": bandwidth, + "intercept_seconds": fit["intercept"], + "r_squared": fit["r_squared"], + "_num_points": len(T_arr), + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def parse_args(): + p = argparse.ArgumentParser(description="Fit cost-model primitive parameters") + p.add_argument("--pingpong-intra", nargs="+", default=[], metavar="FILE") + p.add_argument("--pingpong-inter", nargs="+", default=[], metavar="FILE") + p.add_argument("--compute-gcn", nargs="+", default=[], metavar="FILE") + p.add_argument("--compute-edge", nargs="+", default=[], metavar="FILE") + p.add_argument("--gather-contiguous", nargs="+", default=[], metavar="FILE") + p.add_argument("--gather-clustered", nargs="+", default=[], metavar="FILE") + p.add_argument("--gather-random", nargs="+", default=[], metavar="FILE") + p.add_argument("--output", type=str, default="data/fitted_primitives.json") + return p.parse_args() + + +def main(): + args = parse_args() + result = {} + + # Network + net = {} + if args.pingpong_intra: + recs = load_json_files(args.pingpong_intra) + net["intra"] = fit_network(recs) + print(f"[network/intra] B={net['intra']['bandwidth_bytes_per_sec']/1e9:.2f} GB/s " + f"t_L={net['intra']['latency_seconds']*1e6:.2f} µs " + f"R²={net['intra']['r_squared']:.4f}") + if args.pingpong_inter: + recs = load_json_files(args.pingpong_inter) + net["inter"] = fit_network(recs) + print(f"[network/inter] B={net['inter']['bandwidth_bytes_per_sec']/1e9:.2f} GB/s " + f"t_L={net['inter']['latency_seconds']*1e6:.2f} µs " + f"R²={net['inter']['r_squared']:.4f}") + result["network"] = net + + # Compute + comp = {} + if args.compute_gcn: + recs = load_json_files(args.compute_gcn) + comp["gcn"] = { + "forward": fit_compute(recs, "forward_trials_seconds"), + "backward": fit_compute(recs, "backward_trials_seconds"), + } + print(f"[compute/gcn] coeff_V={comp['gcn']['forward']['coeff_V']:.3e} " + f"coeff_E={comp['gcn']['forward']['coeff_E']:.3e} " + f"R²={comp['gcn']['forward']['r_squared']:.4f}") + if args.compute_edge: + recs = load_json_files(args.compute_edge) + comp["edge"] = { + "forward": fit_compute(recs, "forward_trials_seconds"), + "backward": fit_compute(recs, "backward_trials_seconds"), + } + print(f"[compute/edge] coeff_V={comp['edge']['forward']['coeff_V']:.3e} " + f"coeff_E={comp['edge']['forward']['coeff_E']:.3e} " + f"R²={comp['edge']['forward']['r_squared']:.4f}") + result["compute"] = comp + + # Gather + gath = {} + for dist_name, files_attr in [ + ("contiguous", "gather_contiguous"), + ("clustered", "gather_clustered"), + ("random", "gather_random"), + ]: + files = getattr(args, files_attr.replace("-", "_")) + if files: + recs = load_json_files(files) + gath[dist_name] = { + "gather": fit_gather(recs, "gather_trials_seconds"), + "scatter_add": fit_gather(recs, "scatter_add_trials_seconds"), + } + print(f"[gather/{dist_name}] " + f"B_gather={gath[dist_name]['gather']['bandwidth_bytes_per_sec']/1e9:.2f} GB/s " + f"R²={gath[dist_name]['gather']['r_squared']:.4f}") + result["gather"] = gath + + # Write + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w") as f: + json.dump(result, f, indent=2) + print(f"[fit_primitives] Written to {out_path}") + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/benchmarks/__init__.py b/experiments/cost_model_benchmarks/benchmarks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_compute.py b/experiments/cost_model_benchmarks/benchmarks/bench_compute.py new file mode 100644 index 0000000..89f5249 --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/bench_compute.py @@ -0,0 +1,232 @@ +"""Benchmark 1.3 — GNN Layer Compute Primitive. + +Single-GPU benchmark. Fits f_comp(|Ṽ|, |Ẽ|) for two message-function +variants: + +* ``gcn`` — GCN-like: φ(h_b) = W h_b (source-only linear transform) +* ``edge`` — Edge-conditioned: φ(h_b, h_a, e_ba) = MLP([h_b, h_a, e_ba]) + with a 2-layer MLP (hidden dim = feature_dim) + +Two sweep modes (controlled by ``--sweep``): + +* ``vertices`` — vary |V| with |E| fixed at ``--fixed-value`` +* ``edges`` — vary |E| with |V| fixed at ``--fixed-value`` + +Usage:: + + python -m benchmarks.bench_compute \\ + --model edge --sweep vertices \\ + --min 1000 --max 100000 --steps 15 \\ + --fixed-value 500000 --feature-dim 128 \\ + --warmup 10 --trials 50 \\ + --output data/compute_edge_vswp.json --seed 42 +""" + +import argparse + +import numpy as np +import torch +import torch.nn as nn + +from benchmarks.common import ( + collect_metadata, + cuda_timed, + seed_everything, + write_result, +) + + +# --------------------------------------------------------------------------- +# Synthetic graph generation +# --------------------------------------------------------------------------- + + +def erdos_renyi_edges( + num_vertices: int, num_edges: int, device: torch.device +) -> torch.Tensor: + """Return an edge index tensor of shape [2, num_edges] (random, with replacement).""" + src = torch.randint(0, num_vertices, (num_edges,), device=device) + dst = torch.randint(0, num_vertices, (num_edges,), device=device) + return torch.stack([src, dst], dim=0) + + +# --------------------------------------------------------------------------- +# GNN layers +# --------------------------------------------------------------------------- + + +class GCNLayer(nn.Module): + """GCN-like: aggregate neighbour source features with a linear transform.""" + + def __init__(self, feature_dim: int): + super().__init__() + self.linear = nn.Linear(feature_dim, feature_dim, bias=False) + + def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor: + # x: [V, F], edge_index: [2, E] + src, dst = edge_index[0], edge_index[1] + # Message: transform source features + msg = self.linear(x[src]) # [E, F] + # Aggregate: scatter-add to destination + out = torch.zeros_like(x) + out.scatter_add_(0, dst.unsqueeze(1).expand_as(msg), msg) + return out + + +class EdgeConditionedLayer(nn.Module): + """Edge-conditioned: φ(h_b, h_a, e_ba) = MLP([h_b, h_a, e_ba]).""" + + def __init__(self, feature_dim: int): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(3 * feature_dim, feature_dim), + nn.ReLU(), + nn.Linear(feature_dim, feature_dim), + ) + + def forward( + self, x: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor + ) -> torch.Tensor: + # x: [V, F], edge_index: [2, E], edge_attr: [E, F] + src, dst = edge_index[0], edge_index[1] + msg_input = torch.cat([x[src], x[dst], edge_attr], dim=-1) # [E, 3F] + msg = self.mlp(msg_input) # [E, F] + out = torch.zeros_like(x) + out.scatter_add_(0, dst.unsqueeze(1).expand_as(msg), msg) + return out + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def parse_args(): + p = argparse.ArgumentParser(description="GNN compute primitive benchmark") + p.add_argument("--model", choices=["gcn", "edge"], required=True) + p.add_argument("--sweep", choices=["vertices", "edges"], required=True) + p.add_argument("--min", type=int, default=1_000, dest="sweep_min") + p.add_argument("--max", type=int, default=1_000_000, dest="sweep_max") + p.add_argument("--steps", type=int, default=15) + p.add_argument( + "--fixed-value", + type=int, + default=500_000, + help="Fixed |E| when sweeping vertices, or fixed |V| when sweeping edges", + ) + p.add_argument("--feature-dim", type=int, default=128) + p.add_argument("--warmup", type=int, default=10) + p.add_argument("--trials", type=int, default=50) + p.add_argument("--output", type=str, required=True) + p.add_argument("--seed", type=int, default=42) + return p.parse_args() + + +def main(): + args = parse_args() + seed_everything(args.seed) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + F = args.feature_dim + + # Build model + if args.model == "gcn": + model = GCNLayer(F).to(device) + else: + model = EdgeConditionedLayer(F).to(device) + + # Sweep points + sweep_vals = np.unique( + np.round( + np.logspace( + np.log10(args.sweep_min), + np.log10(args.sweep_max), + num=args.steps, + ) + ).astype(int) + ).tolist() + + measurements = [] + for val in sweep_vals: + if args.sweep == "vertices": + num_v, num_e = val, args.fixed_value + else: + num_v, num_e = args.fixed_value, val + + # Synthetic data + x = torch.randn(num_v, F, device=device, requires_grad=True) + edge_index = erdos_renyi_edges(num_v, num_e, device) + edge_attr = ( + torch.randn(num_e, F, device=device) if args.model == "edge" else None + ) + + # Forward timing + def fwd(): + if args.model == "gcn": + model(x, edge_index) + else: + model(x, edge_index, edge_attr) + + fwd_times = cuda_timed(fwd, warmup=args.warmup, trials=args.trials) + + # Backward timing (run fwd first to get a graph) + if args.model == "gcn": + out = model(x, edge_index) + else: + out = model(x, edge_index, edge_attr) + loss_ref = out.sum() + + def bwd(): + if x.grad is not None: + x.grad.zero_() + if args.model == "gcn": + out_inner = model(x, edge_index) + else: + out_inner = model(x, edge_index, edge_attr) + out_inner.sum().backward() + + bwd_times = cuda_timed(bwd, warmup=args.warmup, trials=args.trials) + + measurements.append( + { + "params": { + "num_vertices": num_v, + "num_edges": num_e, + "sweep_var": args.sweep, + "sweep_value": val, + "model": args.model, + "feature_dim": F, + }, + "forward_trials_seconds": fwd_times, + "backward_trials_seconds": bwd_times, + } + ) + med_fwd = sorted(fwd_times)[len(fwd_times) // 2] + med_bwd = sorted(bwd_times)[len(bwd_times) // 2] + print( + f"[compute/{args.model}] |V|={num_v:>8} |E|={num_e:>9} " + f"fwd {1e3*med_fwd:.2f} ms bwd {1e3*med_bwd:.2f} ms" + ) + + payload = { + "benchmark": "compute", + "metadata": collect_metadata(), + "config": { + "model": args.model, + "sweep": args.sweep, + "sweep_min": args.sweep_min, + "sweep_max": args.sweep_max, + "steps": args.steps, + "fixed_value": args.fixed_value, + "feature_dim": F, + "warmup": args.warmup, + "trials": args.trials, + "seed": args.seed, + }, + "measurements": measurements, + } + write_result(args.output, payload) + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py b/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py new file mode 100644 index 0000000..8d28cfd --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py @@ -0,0 +1,227 @@ +"""Benchmark 1.2 — Intra/Inter Concurrency Check. + +Verifies that T_both / max(T_intra, T_inter) ≈ 1, i.e. that NVLink and +InfiniBand transfers overlap when issued simultaneously. + +Requires exactly 4 ranks across 2 nodes: + Node A: rank 0 (A0), rank 1 (A1) + Node B: rank 2 (B0), rank 3 (B1) + +Three conditions at a fixed message size: + 1. intra-only — A0↔A1 and B0↔B1 (no cross-node traffic) + 2. inter-only — A0↔B0 and A1↔B1 (no intra-node traffic) + 3. concurrent — all four exchanges in the same window (separate streams) + +Each rank logs its own wall time per trial. Rank 0 collects and writes JSON. + +Usage:: + + srun -N 2 --ntasks-per-node 2 python -m benchmarks.bench_concurrency \\ + --message-bytes 16777216 --warmup 20 --trials 100 \\ + --output data/concurrency.json --seed 42 +""" + +import argparse +import os + +import torch +import torch.distributed as dist + +from benchmarks.common import ( + collect_metadata, + seed_everything, + setup_distributed, + write_result, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _exchange(tensor_send: torch.Tensor, tensor_recv: torch.Tensor, + peer: int, stream: torch.cuda.Stream) -> None: + """Non-blocking send+recv on *stream* with *peer*.""" + with torch.cuda.stream(stream): + send_op = dist.P2POp(dist.isend, tensor_send, peer) + recv_op = dist.P2POp(dist.irecv, tensor_recv, peer) + reqs = dist.batch_isend_irecv([send_op, recv_op]) + for req in reqs: + req.wait() + + +def timed_window(rank: int, + send_buf: torch.Tensor, + recv_buf: torch.Tensor, + peer: int, + stream: torch.cuda.Stream, + warmup: int, + trials: int) -> list: + """Time a single exchange window (one send+recv pair).""" + for _ in range(warmup): + _exchange(send_buf, recv_buf, peer, stream) + torch.cuda.synchronize() + dist.barrier() + + times = [] + for _ in range(trials): + dist.barrier() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + _exchange(send_buf, recv_buf, peer, stream) + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end) / 1_000.0) + return times + + +def timed_concurrent(rank: int, + intra_send: torch.Tensor, intra_recv: torch.Tensor, intra_peer: int, + inter_send: torch.Tensor, inter_recv: torch.Tensor, inter_peer: int, + intra_stream: torch.cuda.Stream, + inter_stream: torch.cuda.Stream, + warmup: int, + trials: int) -> list: + """Time intra and inter exchanges issued concurrently on separate streams.""" + for _ in range(warmup): + _exchange(intra_send, intra_recv, intra_peer, intra_stream) + _exchange(inter_send, inter_recv, inter_peer, inter_stream) + torch.cuda.synchronize() + dist.barrier() + + times = [] + for _ in range(trials): + dist.barrier() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + _exchange(intra_send, intra_recv, intra_peer, intra_stream) + _exchange(inter_send, inter_recv, inter_peer, inter_stream) + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end) / 1_000.0) + return times + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def parse_args(): + p = argparse.ArgumentParser(description="Intra/inter concurrency benchmark") + p.add_argument("--message-bytes", type=int, default=16_777_216) # 16 MiB + p.add_argument("--warmup", type=int, default=20) + p.add_argument("--trials", type=int, default=100) + p.add_argument("--output", type=str, required=True) + p.add_argument("--seed", type=int, default=42) + return p.parse_args() + + +def main(): + args = parse_args() + rank, world_size, local_rank = setup_distributed() + + if world_size != 4: + raise ValueError( + f"bench_concurrency requires exactly 4 ranks (got {world_size}).\n" + "Layout: rank 0,1 on node A; rank 2,3 on node B." + ) + + seed_everything(args.seed) + device = torch.device(f"cuda:{local_rank}") + + num_elems = max(1, args.message_bytes // 4) + send_buf = torch.randn(num_elems, dtype=torch.float32, device=device) + recv_buf = torch.zeros(num_elems, dtype=torch.float32, device=device) + + # Intra-node peers: 0↔1, 2↔3 + # Inter-node peers: 0↔2, 1↔3 + intra_peer = {0: 1, 1: 0, 2: 3, 3: 2}[rank] + inter_peer = {0: 2, 1: 3, 2: 0, 3: 1}[rank] + + intra_stream = torch.cuda.Stream(device=device) + inter_stream = torch.cuda.Stream(device=device) + + intra_send = send_buf.clone() + intra_recv = torch.zeros_like(recv_buf) + inter_send = send_buf.clone() + inter_recv = torch.zeros_like(recv_buf) + + # --- Condition 1: intra-only --- + times_intra = timed_window( + rank, intra_send, intra_recv, intra_peer, intra_stream, + args.warmup, args.trials + ) + dist.barrier() + + # --- Condition 2: inter-only --- + times_inter = timed_window( + rank, inter_send, inter_recv, inter_peer, inter_stream, + args.warmup, args.trials + ) + dist.barrier() + + # --- Condition 3: concurrent --- + times_concurrent = timed_concurrent( + rank, + intra_send, intra_recv, intra_peer, + inter_send, inter_recv, inter_peer, + intra_stream, inter_stream, + args.warmup, args.trials, + ) + dist.barrier() + + # Gather per-rank times to rank 0 + def gather_times(times_local): + obj = [None] * world_size + dist.all_gather_object(obj, times_local) + return obj + + intra_all = gather_times(times_intra) + inter_all = gather_times(times_inter) + conc_all = gather_times(times_concurrent) + + if rank == 0: + measurements = [ + { + "params": {"condition": "intra_only", "message_bytes": num_elems * 4}, + "per_rank_trials_seconds": intra_all, + }, + { + "params": {"condition": "inter_only", "message_bytes": num_elems * 4}, + "per_rank_trials_seconds": inter_all, + }, + { + "params": {"condition": "concurrent", "message_bytes": num_elems * 4}, + "per_rank_trials_seconds": conc_all, + }, + ] + payload = { + "benchmark": "concurrency", + "metadata": collect_metadata(), + "config": { + "message_bytes": args.message_bytes, + "warmup": args.warmup, + "trials": args.trials, + "world_size": world_size, + "seed": args.seed, + "rank_layout": "rank 0,1 on node A; rank 2,3 on node B", + }, + "measurements": measurements, + } + write_result(args.output, payload) + print( + f"[concurrency] intra median = " + f"{1e3*sorted(intra_all[0])[len(intra_all[0])//2]:.2f} ms | " + f"inter median = " + f"{1e3*sorted(inter_all[0])[len(inter_all[0])//2]:.2f} ms | " + f"concurrent median = " + f"{1e3*sorted(conc_all[0])[len(conc_all[0])//2]:.2f} ms" + ) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py new file mode 100644 index 0000000..7da40eb --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py @@ -0,0 +1,504 @@ +"""Benchmark 2.1 — End-to-End Halo Exchange. + +Measures full GNN layer wall time (forward + backward) across a sweep of +configurations on the full multi-node setup. Intended to be run as a SLURM +array job with one invocation per (K, F, graph) combination. + +This module contains a self-contained minimal halo-exchange implementation +(no dependency on the DGraph production library) so the benchmark remains +isolated and portable. + +Synthetic graphs: + * ``erdos_renyi`` — Erdős-Rényi with ``--avg-degree`` expected degree + * ``sbm`` — Stochastic Block Model; ``--sbm-inter-density`` controls + the fraction of inter-block edges (topology ablation) + +Partitioners: + * ``random`` — assign each vertex to a uniformly random rank + * ``balanced`` — contiguous vertex blocks of equal size + * ``metis`` — balanced k-way via pymetis (skipped if not installed) + +The benchmark logs, for every run: + * world_size K, feature dim F, graph type, partitioner + * per-rank partition statistics: + intra_halo_size — halo vertices on the same node + inter_halo_size — halo vertices on different nodes + c_intra, c_inter — communication volumes (bytes) + * per-trial layer times from rank 0 (and per-rank times for completeness) + +Usage:: + + torchrun --nnodes 2 --nproc_per_node 4 \\ + -m benchmarks.bench_end_to_end \\ + --graph erdos_renyi --num-vertices 100000 --avg-degree 20 \\ + --feature-dim 128 --model gcn --partitioner balanced \\ + --warmup 10 --trials 50 \\ + --output data/e2e_K8_F128_er_bal.json --seed 42 +""" + +import argparse +import os + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn + +from benchmarks.common import ( + collect_metadata, + cuda_timed, + seed_everything, + setup_distributed, + write_result, +) + + +# =========================================================================== +# Synthetic graph generators +# =========================================================================== + +def gen_erdos_renyi(num_vertices: int, avg_degree: float, + rng: np.random.Generator) -> np.ndarray: + """Return edge array of shape [E, 2] (src, dst) for an Erdős-Rényi digraph.""" + num_edges = int(num_vertices * avg_degree) + src = rng.integers(0, num_vertices, size=num_edges) + dst = rng.integers(0, num_vertices, size=num_edges) + return np.stack([src, dst], axis=1) + + +def gen_sbm(num_vertices: int, avg_degree: float, inter_density: float, + rng: np.random.Generator) -> np.ndarray: + """Return edges for a Stochastic Block Model graph. + + Vertices are split into blocks of equal size (one per rank for convenience, + though the actual partitioning is a separate step). The ratio of + intra-block to inter-block edges is controlled by *inter_density*. + """ + world_size = dist.get_world_size() if dist.is_initialized() else 4 + block_size = num_vertices // world_size + edges = [] + target_edges = int(num_vertices * avg_degree) + + intra_edges = int(target_edges * (1.0 - inter_density)) + inter_edges = target_edges - intra_edges + + # Intra-block edges + for b in range(world_size): + start = b * block_size + end = start + block_size + n = intra_edges // world_size + s = rng.integers(start, end, size=n) + d = rng.integers(start, end, size=n) + edges.append(np.stack([s, d], axis=1)) + + # Inter-block edges + s = rng.integers(0, num_vertices, size=inter_edges) + d = rng.integers(0, num_vertices, size=inter_edges) + # Force cross-block by offsetting dst block + d_block = (s // block_size + 1 + rng.integers(0, world_size - 1, + size=inter_edges)) % world_size + d = d_block * block_size + rng.integers(0, block_size, size=inter_edges) + d = np.clip(d, 0, num_vertices - 1) + edges.append(np.stack([s, d], axis=1)) + + return np.concatenate(edges, axis=0) + + +# =========================================================================== +# Partitioners +# =========================================================================== + +def partition_random(num_vertices: int, world_size: int, + rng: np.random.Generator) -> np.ndarray: + return rng.integers(0, world_size, size=num_vertices).astype(np.int64) + + +def partition_balanced(num_vertices: int, world_size: int) -> np.ndarray: + return np.floor(np.arange(num_vertices) * world_size / num_vertices).astype(np.int64) + + +def partition_metis(num_vertices: int, world_size: int, + edges: np.ndarray) -> np.ndarray: + try: + import pymetis + except ImportError: + raise RuntimeError( + "pymetis is not installed. Install it with: pip install pymetis\n" + "Or use --partitioner random or --partitioner balanced." + ) + # Build adjacency list for pymetis + adj = [[] for _ in range(num_vertices)] + for s, d in edges: + adj[s].append(int(d)) + adj[d].append(int(s)) + _, membership = pymetis.part_graph(world_size, adjacency=adj) + return np.array(membership, dtype=np.int64) + + +# =========================================================================== +# Minimal halo-exchange infrastructure +# =========================================================================== + +def build_local_comm_pattern(edges: np.ndarray, assignment: np.ndarray, + rank: int, world_size: int): + """Compute the local communication pattern for this rank. + + Returns a dict with: + local_vertices — np.ndarray of vertex IDs owned by this rank + local_edge_index — torch.Tensor [2, E_local] with local vertex IDs + remapped so that 0..n_local-1 are owned vertices + and n_local..n_local+n_halo-1 are halo vertices + send_counts — list[int] of length world_size: vertices to send + recv_counts — list[int] of length world_size: vertices to recv + send_idx — local indices (into local_vertices) to send per rank + halo_global_ids — global vertex IDs of halo vertices, in recv order + intra_halo_size — halo vertices from same node (ranks sharing node) + inter_halo_size — halo vertices from remote nodes + ranks_per_node — int (derived from LOCAL_RANK / RANK relationship) + """ + local_mask = assignment == rank + local_vertices = np.where(local_mask)[0] + n_local = len(local_vertices) + + # Global -> local index map + g2l = {int(v): i for i, v in enumerate(local_vertices)} + + # Find edges where dst is local + local_dst_mask = np.isin(edges[:, 1], local_vertices) + local_edges = edges[local_dst_mask] + + # Halo: src vertices not owned by this rank + halo_src_mask = ~np.isin(local_edges[:, 0], local_vertices) + halo_global = np.unique(local_edges[halo_src_mask, 0]) + + # Group halo vertices by owning rank + halo_owners = assignment[halo_global] + recv_by_rank = [] + halo_order = [] + for r in range(world_size): + verts = halo_global[halo_owners == r] + recv_by_rank.append(verts) + halo_order.extend(verts.tolist()) + halo_order = np.array(halo_order, dtype=np.int64) + + # Global halo id -> local halo index + halo_g2l = {int(v): n_local + i for i, v in enumerate(halo_order)} + all_g2l = {**g2l, **halo_g2l} + + # Find which local vertices other ranks need (send pattern) + # We exchange recv_counts via all_to_all to learn send_counts + recv_counts = [len(rv) for rv in recv_by_rank] + + # Build send: for each rank r, which of our local vertices does r need? + # We do a global exchange of halo_global per rank + all_recv = [None] * world_size + dist.all_gather_object(all_recv, halo_order.tolist()) + + send_idx_by_rank = [] + for r in range(world_size): + needed = np.array(all_recv[r], dtype=np.int64) + owned_mask = assignment[needed] == rank if len(needed) > 0 else np.array([], dtype=bool) + owned = needed[owned_mask] if len(needed) > 0 else np.array([], dtype=np.int64) + # Map to local indices + local_idxs = np.array([g2l[int(v)] for v in owned], dtype=np.int64) + send_idx_by_rank.append(local_idxs) + + send_counts = [len(s) for s in send_idx_by_rank] + + # Remap edges to local indices + valid_edge_mask = np.array([ + (int(s) in all_g2l) and (int(d) in all_g2l) + for s, d in local_edges + ]) + local_edges_valid = local_edges[valid_edge_mask] + if len(local_edges_valid) > 0: + remapped_src = np.array([all_g2l[int(s)] for s in local_edges_valid[:, 0]]) + remapped_dst = np.array([all_g2l[int(d)] for d in local_edges_valid[:, 1]]) + edge_index = torch.tensor( + np.stack([remapped_src, remapped_dst], axis=0), dtype=torch.long + ) + else: + edge_index = torch.zeros((2, 0), dtype=torch.long) + + # Compute intra / inter halo sizes + ranks_per_node = int(os.environ.get("LOCAL_WORLD_SIZE", + os.environ.get("SLURM_NTASKS_PER_NODE", "4"))) + my_node = rank // ranks_per_node + intra_halo_size = 0 + inter_halo_size = 0 + for r, verts in enumerate(recv_by_rank): + peer_node = r // ranks_per_node + if peer_node == my_node: + intra_halo_size += len(verts) + else: + inter_halo_size += len(verts) + + return { + "local_vertices": local_vertices, + "n_local": n_local, + "n_halo": len(halo_order), + "edge_index": edge_index, + "send_counts": send_counts, + "recv_counts": recv_counts, + "send_idx_by_rank": send_idx_by_rank, + "halo_order": halo_order, + "intra_halo_size": intra_halo_size, + "inter_halo_size": inter_halo_size, + "ranks_per_node": ranks_per_node, + } + + +class MinimalHaloExchange(torch.autograd.Function): + """Forward: gather boundary features → all_to_all → populate recv buffer. + Backward: reverse the transfer to accumulate gradients. + """ + + @staticmethod + def forward(ctx, x_local, send_idx_flat, send_counts, recv_counts, world_size): + # Gather send buffer + send_buf = x_local[send_idx_flat] # [total_send, F] + + # Split by destination rank + send_list = list(send_buf.split(send_counts, dim=0)) + recv_list = [torch.zeros(rc, x_local.shape[1], + dtype=x_local.dtype, device=x_local.device) + for rc in recv_counts] + + dist.all_to_all(recv_list, send_list) + + recv_buf = torch.cat(recv_list, dim=0) if sum(recv_counts) > 0 else \ + torch.zeros(0, x_local.shape[1], device=x_local.device) + + ctx.save_for_backward(send_idx_flat) + ctx.send_counts = send_counts + ctx.recv_counts = recv_counts + ctx.world_size = world_size + ctx.n_local = x_local.shape[0] + ctx.feature_dim = x_local.shape[1] + ctx.device = x_local.device + + return recv_buf + + @staticmethod + def backward(ctx, grad_recv): + send_idx_flat, = ctx.saved_tensors + send_counts = ctx.send_counts + recv_counts = ctx.recv_counts + world_size = ctx.world_size + n_local = ctx.n_local + F = ctx.feature_dim + device = ctx.device + + # Reverse: recv_counts become send_counts and vice versa + grad_recv_list = list(grad_recv.split(recv_counts, dim=0)) \ + if grad_recv.shape[0] > 0 else [torch.zeros(0, F, device=device)] * world_size + grad_send_list = [torch.zeros(sc, F, device=device) for sc in send_counts] + + dist.all_to_all(grad_send_list, grad_recv_list) + + grad_send = torch.cat(grad_send_list, dim=0) + + # Scatter-add back to local vertices + grad_x_local = torch.zeros(n_local, F, device=device, dtype=grad_recv.dtype) + grad_x_local.scatter_add_( + 0, + send_idx_flat.unsqueeze(1).expand_as(grad_send), + grad_send, + ) + return grad_x_local, None, None, None, None + + +# =========================================================================== +# GNN layers (same as bench_compute.py for consistency) +# =========================================================================== + +class GCNLayer(nn.Module): + def __init__(self, feature_dim: int): + super().__init__() + self.linear = nn.Linear(feature_dim, feature_dim, bias=False) + + def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor: + src, dst = edge_index[0], edge_index[1] + n_local = x.shape[0] + # Only update local vertices (dst < n_local guard not needed since + # edge_index already restricts to local dst) + msg = self.linear(x[src]) + out = torch.zeros_like(x) + out.scatter_add_(0, dst.unsqueeze(1).expand_as(msg), msg) + return out + + +class EdgeConditionedLayer(nn.Module): + def __init__(self, feature_dim: int): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(3 * feature_dim, feature_dim), + nn.ReLU(), + nn.Linear(feature_dim, feature_dim), + ) + + def forward(self, x: torch.Tensor, edge_index: torch.Tensor, + edge_attr: torch.Tensor) -> torch.Tensor: + src, dst = edge_index[0], edge_index[1] + msg = self.mlp(torch.cat([x[src], x[dst], edge_attr], dim=-1)) + out = torch.zeros_like(x) + out.scatter_add_(0, dst.unsqueeze(1).expand_as(msg), msg) + return out + + +# =========================================================================== +# Main +# =========================================================================== + +def parse_args(): + p = argparse.ArgumentParser(description="End-to-end halo exchange benchmark") + p.add_argument("--graph", choices=["erdos_renyi", "sbm"], default="erdos_renyi") + p.add_argument("--num-vertices", type=int, default=100_000) + p.add_argument("--avg-degree", type=float, default=20.0) + p.add_argument("--sbm-inter-density", type=float, default=0.1, + help="Fraction of inter-block edges for SBM graphs") + p.add_argument("--feature-dim", type=int, default=128) + p.add_argument("--model", choices=["gcn", "edge"], default="gcn") + p.add_argument("--partitioner", choices=["random", "balanced", "metis"], + default="balanced") + p.add_argument("--warmup", type=int, default=10) + p.add_argument("--trials", type=int, default=50) + p.add_argument("--output", type=str, required=True) + p.add_argument("--seed", type=int, default=42) + return p.parse_args() + + +def main(): + args = parse_args() + rank, world_size, local_rank = setup_distributed() + seed_everything(args.seed + rank) # per-rank seed for graph generation + rng = np.random.default_rng(args.seed) # shared seed for graph topology + device = torch.device(f"cuda:{local_rank}") + F = args.feature_dim + + # --- Generate graph on all ranks (same seed → identical graph) --- + if args.graph == "erdos_renyi": + edges = gen_erdos_renyi(args.num_vertices, args.avg_degree, rng) + else: + edges = gen_sbm(args.num_vertices, args.avg_degree, + args.sbm_inter_density, rng) + + # --- Partition --- + rng_part = np.random.default_rng(args.seed + 1) + if args.partitioner == "random": + assignment = partition_random(args.num_vertices, world_size, rng_part) + elif args.partitioner == "balanced": + assignment = partition_balanced(args.num_vertices, world_size) + else: + assignment = partition_metis(args.num_vertices, world_size, edges) + + # --- Build local comm pattern --- + pattern = build_local_comm_pattern(edges, assignment, rank, world_size) + n_local = pattern["n_local"] + n_halo = pattern["n_halo"] + edge_index = pattern["edge_index"].to(device) + + send_counts = pattern["send_counts"] + recv_counts = pattern["recv_counts"] + send_idx_flat = torch.cat([ + torch.tensor(s, dtype=torch.long) for s in pattern["send_idx_by_rank"] + ]).to(device) if sum(send_counts) > 0 else torch.zeros(0, dtype=torch.long, device=device) + + # --- Model --- + if args.model == "gcn": + layer = GCNLayer(F).to(device) + else: + layer = EdgeConditionedLayer(F).to(device) + layer.train() + + # --- Synthetic local node features --- + x_local = torch.randn(n_local, F, device=device, requires_grad=True) + edge_attr = torch.randn(edge_index.shape[1], F, device=device) \ + if args.model == "edge" else None + + # --- Timed forward + backward --- + def one_layer(): + # Forward halo exchange + recv_buf = MinimalHaloExchange.apply( + x_local, send_idx_flat, send_counts, recv_counts, world_size + ) + # Augment: local + halo + x_aug = torch.cat([x_local, recv_buf], dim=0) + # Message passing + if args.model == "gcn": + out = layer(x_aug, edge_index) + else: + out = layer(x_aug, edge_index, edge_attr) + # Backward + loss = out.sum() + loss.backward() + if x_local.grad is not None: + x_local.grad.zero_() + + # Barrier before timing + dist.barrier() + times_local = cuda_timed(one_layer, warmup=args.warmup, trials=args.trials) + dist.barrier() + + # Gather per-rank times and stats to rank 0 + stats_local = { + "rank": rank, + "n_local": n_local, + "n_halo": n_halo, + "intra_halo_size": pattern["intra_halo_size"], + "inter_halo_size": pattern["inter_halo_size"], + "c_intra_bytes": pattern["intra_halo_size"] * F * 4, + "c_inter_bytes": pattern["inter_halo_size"] * F * 4, + "send_total": sum(send_counts), + "recv_total": sum(recv_counts), + "trials_seconds": times_local, + } + + all_stats = [None] * world_size + dist.all_gather_object(all_stats, stats_local) + + if rank == 0: + med = sorted(times_local)[len(times_local) // 2] + print( + f"[e2e] K={world_size} F={F} {args.graph}/{args.partitioner}/{args.model} " + f"n_local={n_local} n_halo={n_halo} " + f"median {1e3*med:.2f} ms" + ) + payload = { + "benchmark": "end_to_end", + "metadata": collect_metadata(), + "config": { + "graph": args.graph, + "num_vertices": args.num_vertices, + "avg_degree": args.avg_degree, + "sbm_inter_density": args.sbm_inter_density, + "feature_dim": F, + "model": args.model, + "partitioner": args.partitioner, + "world_size": world_size, + "ranks_per_node": pattern["ranks_per_node"], + "warmup": args.warmup, + "trials": args.trials, + "seed": args.seed, + }, + "measurements": [ + { + "params": { + "world_size": world_size, + "feature_dim": F, + "graph": args.graph, + "partitioner": args.partitioner, + "model": args.model, + }, + "rank0_trials_seconds": times_local, + "per_rank_stats": all_stats, + } + ], + } + write_result(args.output, payload) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_gather.py b/experiments/cost_model_benchmarks/benchmarks/bench_gather.py new file mode 100644 index 0000000..d62c71b --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/bench_gather.py @@ -0,0 +1,181 @@ +"""Benchmark 1.4 — Buffer-Copy / Gather-Scatter Bandwidth. + +Single-GPU benchmark. Measures the effective bandwidth of ``x[idx]`` (gather) +and the corresponding ``scatter_add_`` (backward) for three index distributions: + +* ``contiguous`` — contiguous block starting at a random offset +* ``clustered`` — *c* cluster centres, each with a contiguous block of size + ``--cluster-size``; simulates a well-partitioned graph halo +* ``random`` — uniformly random indices (worst-case for cache) + +Sweeps *k* (number of gathered rows) from ``--min-k`` to ``--max-k``. + +Usage:: + + python -m benchmarks.bench_gather \\ + --distribution clustered \\ + --min-k 1000 --max-k 10000000 --steps 20 \\ + --N 20000000 --feature-dim 128 --cluster-size 64 \\ + --warmup 10 --trials 50 \\ + --output data/gather_clustered.json --seed 42 +""" + +import argparse + +import numpy as np +import torch + +from benchmarks.common import ( + collect_metadata, + cuda_timed, + seed_everything, + write_result, +) + + +# --------------------------------------------------------------------------- +# Index generators +# --------------------------------------------------------------------------- + +def contiguous_idx(N: int, k: int, device: torch.device) -> torch.Tensor: + start = torch.randint(0, max(1, N - k), (1,)).item() + return torch.arange(start, start + k, device=device) + + +def clustered_idx(N: int, k: int, cluster_size: int, + device: torch.device, rng: np.random.Generator) -> torch.Tensor: + """Draw cluster centres, then take a contiguous block around each.""" + num_clusters = max(1, k // cluster_size) + centres = rng.integers(0, N, size=num_clusters) + idx_parts = [] + for c in centres: + start = int(np.clip(c, 0, N - cluster_size)) + idx_parts.append(torch.arange(start, start + cluster_size, device=device)) + idx = torch.cat(idx_parts)[:k] + return idx + + +def random_idx(N: int, k: int, device: torch.device) -> torch.Tensor: + return torch.randperm(N, device=device)[:k] + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def parse_args(): + p = argparse.ArgumentParser(description="Gather/scatter bandwidth benchmark") + p.add_argument("--distribution", choices=["contiguous", "clustered", "random"], + required=True) + p.add_argument("--min-k", type=int, default=1_000) + p.add_argument("--max-k", type=int, default=10_000_000) + p.add_argument("--steps", type=int, default=20) + p.add_argument("--N", type=int, default=20_000_000, + help="Total number of rows in the source tensor x") + p.add_argument("--feature-dim", type=int, default=128) + p.add_argument("--cluster-size", type=int, default=64, + help="Rows per cluster (only used with --distribution clustered)") + p.add_argument("--warmup", type=int, default=10) + p.add_argument("--trials", type=int, default=50) + p.add_argument("--output", type=str, required=True) + p.add_argument("--seed", type=int, default=42) + return p.parse_args() + + +def main(): + args = parse_args() + seed_everything(args.seed) + rng = np.random.default_rng(args.seed) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + F = args.feature_dim + + # Pre-allocate the source tensor and gradient buffer once + x = torch.randn(args.N, F, device=device) + grad_y_template = torch.ones(1, F, device=device) # resized per k + + # Sweep k values + k_values = np.unique( + np.round( + np.logspace( + np.log10(args.min_k), + np.log10(args.max_k), + num=args.steps, + ) + ).astype(int) + ).tolist() + k_values = [min(int(k), args.N) for k in k_values] + + measurements = [] + for k in k_values: + # Build index tensor + if args.distribution == "contiguous": + idx = contiguous_idx(args.N, k, device) + elif args.distribution == "clustered": + idx = clustered_idx(args.N, k, args.cluster_size, device, rng) + else: + idx = random_idx(args.N, k, device) + + # Expand idx for scatter_add_: shape [k, F] + idx_expanded = idx.unsqueeze(1).expand(-1, F) + grad_y = torch.ones(k, F, device=device) + grad_x = torch.zeros_like(x) + + # --- Forward gather --- + def gather_fn(): + _ = x[idx] + + gather_times = cuda_timed(gather_fn, warmup=args.warmup, trials=args.trials) + + # --- Backward scatter-add --- + def scatter_fn(): + grad_x.zero_() + grad_x.scatter_add_(0, idx_expanded, grad_y) + + scatter_times = cuda_timed(scatter_fn, warmup=args.warmup, trials=args.trials) + + measurements.append({ + "params": { + "k": k, + "N": args.N, + "feature_dim": F, + "distribution": args.distribution, + "cluster_size": args.cluster_size if args.distribution == "clustered" else None, + }, + "gather_trials_seconds": gather_times, + "scatter_add_trials_seconds": scatter_times, + }) + + med_g = sorted(gather_times)[len(gather_times) // 2] + med_s = sorted(scatter_times)[len(scatter_times) // 2] + bytes_moved = k * F * 4 + bw_g = bytes_moved / med_g / 1e9 + bw_s = bytes_moved / med_s / 1e9 + print( + f"[gather/{args.distribution}] k={k:>9} " + f"gather {1e3*med_g:.2f} ms ({bw_g:.1f} GB/s) " + f"scatter {1e3*med_s:.2f} ms ({bw_s:.1f} GB/s)" + ) + + payload = { + "benchmark": "gather", + "metadata": collect_metadata(), + "config": { + "distribution": args.distribution, + "min_k": args.min_k, + "max_k": args.max_k, + "steps": args.steps, + "N": args.N, + "feature_dim": F, + "cluster_size": args.cluster_size, + "warmup": args.warmup, + "trials": args.trials, + "seed": args.seed, + }, + "measurements": measurements, + } + write_result(args.output, payload) + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_pingpong.py b/experiments/cost_model_benchmarks/benchmarks/bench_pingpong.py new file mode 100644 index 0000000..c844311 --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/bench_pingpong.py @@ -0,0 +1,161 @@ +"""Benchmark 1.1 — Network Bandwidth / Latency (Ping-Pong). + +Measures one-way transfer time across a sweep of message sizes using a +two-rank ping-pong pattern. Run once with both ranks on the same node +(intra-node, NVLink) and once with ranks on different nodes (inter-node, +InfiniBand). The SLURM script controls placement; this script only records +a --mode label. + +Usage (via torchrun / srun):: + + torchrun --nnodes 1 --nproc_per_node 2 \\ + -m benchmarks.bench_pingpong \\ + --mode intra --min-bytes 64 --max-bytes 67108864 --steps 21 \\ + --warmup 20 --trials 100 --output data/pingpong_intra.json --seed 42 +""" + +import argparse +import os + +import torch +import torch.distributed as dist + +from benchmarks.common import ( + collect_metadata, + seed_everything, + setup_distributed, + write_result, +) + + +# --------------------------------------------------------------------------- +# Ping-pong timing +# --------------------------------------------------------------------------- + +def pingpong_timed(rank: int, tensor: torch.Tensor, warmup: int, trials: int) -> list: + """Perform a ping-pong between rank 0 and rank 1. + + Returns per-trial *one-way* transfer times in seconds (rank 0 only). + Rank 1 returns an empty list. + """ + # Warmup + for _ in range(warmup): + if rank == 0: + dist.send(tensor, dst=1) + dist.recv(tensor, src=1) + else: + dist.recv(tensor, src=0) + dist.send(tensor, dst=0) + torch.cuda.synchronize() + dist.barrier() + + times = [] + for _ in range(trials): + dist.barrier() + start_evt = torch.cuda.Event(enable_timing=True) + end_evt = torch.cuda.Event(enable_timing=True) + + if rank == 0: + start_evt.record() + dist.send(tensor, dst=1) + dist.recv(tensor, src=1) + end_evt.record() + torch.cuda.synchronize() + # Round-trip / 2 = one-way + times.append(start_evt.elapsed_time(end_evt) / 2.0 / 1_000.0) + else: + dist.recv(tensor, src=0) + dist.send(tensor, dst=0) + + return times + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def parse_args(): + p = argparse.ArgumentParser(description="Ping-pong bandwidth/latency benchmark") + p.add_argument("--min-bytes", type=int, default=64) + p.add_argument("--max-bytes", type=int, default=67_108_864) # 64 MiB + p.add_argument("--steps", type=int, default=21, + help="Number of logarithmically-spaced message sizes") + p.add_argument("--warmup", type=int, default=20) + p.add_argument("--trials", type=int, default=100) + p.add_argument("--mode", choices=["intra", "inter"], default="inter", + help="Label only — actual placement is controlled by SLURM") + p.add_argument("--output", type=str, required=True) + p.add_argument("--seed", type=int, default=42) + return p.parse_args() + + +def main(): + args = parse_args() + rank, world_size, local_rank = setup_distributed() + + if world_size != 2: + raise ValueError(f"bench_pingpong requires exactly 2 ranks, got {world_size}") + + seed_everything(args.seed) + device = torch.device(f"cuda:{local_rank}") + + # Build logarithmically-spaced byte sizes (powers-of-2 friendly) + import numpy as np + byte_sizes = np.unique( + np.round( + np.logspace( + np.log2(args.min_bytes), + np.log2(args.max_bytes), + num=args.steps, + base=2, + ) + ).astype(int) + ).tolist() + + measurements = [] + for nbytes in byte_sizes: + # Float32 elements + num_elems = max(1, nbytes // 4) + tensor = torch.zeros(num_elems, dtype=torch.float32, device=device) + + times = pingpong_timed(rank, tensor, args.warmup, args.trials) + + if rank == 0: + measurements.append({ + "params": { + "message_bytes": num_elems * 4, + "num_elements": num_elems, + "mode": args.mode, + }, + "trials_seconds": times, + }) + print( + f"[pingpong] {num_elems * 4:>10} bytes | " + f"median {1e3 * float(sorted(times)[len(times)//2]):.3f} ms" + ) + + dist.barrier() + + if rank == 0: + payload = { + "benchmark": "pingpong", + "metadata": collect_metadata(), + "config": { + "min_bytes": args.min_bytes, + "max_bytes": args.max_bytes, + "steps": args.steps, + "warmup": args.warmup, + "trials": args.trials, + "mode": args.mode, + "world_size": world_size, + "seed": args.seed, + }, + "measurements": measurements, + } + write_result(args.output, payload) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/benchmarks/common.py b/experiments/cost_model_benchmarks/benchmarks/common.py new file mode 100644 index 0000000..27e8069 --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/common.py @@ -0,0 +1,176 @@ +"""Shared utilities for cost-model benchmarks: timing, logging, metadata.""" + +import json +import os +import random +import socket +import subprocess +import time +from pathlib import Path +from typing import Callable + +import numpy as np +import torch +import torch.distributed as dist + + +# --------------------------------------------------------------------------- +# Timing +# --------------------------------------------------------------------------- + + +def cuda_timed(fn: Callable, warmup: int = 10, trials: int = 50) -> list: + """Run *fn* with CUDA-event timing. Returns per-trial wall times in seconds. + + The function is invoked with no arguments. Callers should capture any + needed state via closure. Warmup iterations are discarded. + """ + # Warmup + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + times = [] + for _ in range(trials): + start_evt = torch.cuda.Event(enable_timing=True) + end_evt = torch.cuda.Event(enable_timing=True) + start_evt.record() + fn() + end_evt.record() + torch.cuda.synchronize() + # elapsed_time returns milliseconds + times.append(start_evt.elapsed_time(end_evt) / 1_000.0) + return times + + +# --------------------------------------------------------------------------- +# Metadata collection +# --------------------------------------------------------------------------- + + +def collect_metadata() -> dict: + """Return a dict of reproducibility metadata for the current run.""" + meta: dict = { + "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), + "hostname": socket.gethostname(), + } + + # GPU info + if torch.cuda.is_available(): + gpus = [] + for i in range(torch.cuda.device_count()): + props = torch.cuda.get_device_properties(i) + gpu_entry = { + "index": i, + "name": props.name, + "compute_capability": f"{props.major}.{props.minor}", + "total_memory_bytes": props.total_memory, + } + # UUID available in newer PyTorch builds + if hasattr(props, "uuid"): + gpu_entry["uuid"] = str(props.uuid) + gpus.append(gpu_entry) + meta["gpus"] = gpus + meta["cuda_version"] = torch.version.cuda + else: + meta["gpus"] = [] + meta["cuda_version"] = None + + meta["pytorch_version"] = torch.__version__ + + # NCCL version (tuple -> string) + try: + nccl_ver = torch.cuda.nccl.version() + meta["nccl_version"] = ".".join(str(x) for x in nccl_ver) + except Exception: + meta["nccl_version"] = "unknown" + + # SLURM environment variables + slurm_keys = [ + "SLURM_JOB_ID", + "SLURM_NODELIST", + "SLURM_NNODES", + "SLURM_NTASKS", + "SLURM_PROCID", + "SLURM_LOCALID", + "SLURM_ARRAY_JOB_ID", + "SLURM_ARRAY_TASK_ID", + ] + meta["slurm"] = {k: os.environ.get(k) for k in slurm_keys} + + # Git commit hash of the benchmark code + try: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + capture_output=True, + text=True, + check=True, + ) + meta["git_commit"] = result.stdout.strip() + except Exception: + meta["git_commit"] = "unknown" + + return meta + + +# --------------------------------------------------------------------------- +# JSON output +# --------------------------------------------------------------------------- + + +def write_result(path: str, payload: dict) -> None: + """Write *payload* as a JSON file at *path*, creating parents as needed. + + Expected schema:: + + { + "benchmark": "", + "metadata": { ... }, + "config": { ... }, + "measurements": [ + {"params": {...}, "trials_seconds": [t1, t2, ...]}, + ... + ] + } + """ + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + with open(p, "w") as fh: + json.dump(payload, fh, indent=2) + print(f"[write_result] Saved {p} ({p.stat().st_size} bytes)") + + +# --------------------------------------------------------------------------- +# Distributed setup +# --------------------------------------------------------------------------- + + +def setup_distributed() -> tuple[int, int, int]: + """Initialize torch.distributed with the env-var init method (NCCL). + + Expects MASTER_ADDR, MASTER_PORT, WORLD_SIZE, RANK, and LOCAL_RANK to be + set in the environment (standard for torchrun / SLURM + srun). + + Returns: + (rank, world_size, local_rank) + """ + if not dist.is_initialized(): + dist.init_process_group(backend="nccl", init_method="env://") + rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + return rank, world_size, local_rank + + +# --------------------------------------------------------------------------- +# Seeding +# --------------------------------------------------------------------------- + + +def seed_everything(seed: int) -> None: + """Set Python, NumPy, and PyTorch (CPU + CUDA) random seeds.""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) diff --git a/experiments/cost_model_benchmarks/figures/compute.pdf b/experiments/cost_model_benchmarks/figures/compute.pdf deleted file mode 100644 index f60f93169802441205acbaa7ec6f0d27bcad62cb..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 19308 zcmdtK1yof{^f#;$%9T(`QNRlb2okrrmrF@^ch@DAZc!;gK)OLfDFp-pDN#VWQ)vMy zB?J*fN|g5;)aNnye~;g{zV)tmz1%g-oP8$t?7h$I-^`pj%*x`DoNz7_ggJi<`tT_P z4uwJOO)oaEO zMeb+d)>bmTY>sn*V!mBUc(_QZxtO?^Ls8!d$|f!@=1z7{9`F*vt!8CnW^HE)MSs8R zWRFubcY*2w&5BC`teAVaK)GdX0Tx8R$70`O8K}We?7;u*0LVAUy^}e>-8c5RRn494 zU7c`1e?b0$_7%*{tW8AiJ%AQrz>fz8L7KVl*;BYPs5{}^k+Ts=idI?nJ z48?xqRm9HD9=L*o>mbK598&bx+sVVd6EXpG4!Kz7ugAYvAy$^PjHX+{G};5BgGh| zIJT+RHp+xq%~z=!>K{YxnMIS-rHPEUbpxmu&&|d6T0M*ofj7mvCbnIBqrCH_VfWI< z&5s{@NVNR5)0{qXzr=LB)(wE|aW8J(_U^pn)C ziGeKBorDhNMv9Zhv!ShLsFL0goS9BoklKBocIoA$1NAB4sM*wwpki6as2h8Bo~zM% z(oN#cJ={hLJ>F8zcZ|NeKFzfJ*a39OrM*ojc5rva-7|m<{{dc&*}&S+jzMAln(Z^n z&(m+#2N)NKm=w&^HyD)LD16tOc48M{qgf1=YEkPR_4`FY~pHu>QUVu zcU9cBOFW2UsrARR8T^u`I}5N$BJSZ z>dH@OOg7Y!G7htf4|2XdCHkd*I{LAY^^9yPq?he|qE+IPNR3b|gTxS@#yzDmqCS`- z(cWSvCzJX180&JRMGEY^;`-=_YW5KqS$>u?EL{BNcq7`2>b^pbtl~3`=H^wa>F3o5 z_?qY_NGPPB6>Qc9j9igIS0AY35fH{@qI0?2F21<B1@spdO}2*(4); z@mg7GlQaXozBCV6amcu?2f};=X(}WXR;Wk$Z6m4|eb#Gslubu^J6`u3j*%(0{yS^O=-+W#nx1AqNcD;9)wpw$`(UKIiZKG;;bRyJvq9#Pb z8%>JyLqC68e3RaNNm;RAJF)gHrRrICUw6V96OJhR?Qk)rNyk3cWwzP%cl3|@wm$HO zMPaulN9VP+qUM+?wrpyS1kxUfdrCk>>xOfDP*z*5OJFp>KGJOcDeLf?&DSg8vON1AEHy7hVHS8NOs@&9}haKIoRTHcXUOXLO*oIN&E*H7wbo<$Z z&-fD$wGf)sB(H_5sV|w(rWd``TQ#hXE?%_Vl3XN#c)t_=rI;5&d;Nu1Rh~zoVVl)Q z9zIR`MlK)u{?A|kDMh&0#0S+OPX4@60A|wuguec2Fx1`r?&{ajlouTn>}NPIU28K_ zwgImNXUU#i-B!@Y8sA!qkc+s^Cjo0q?CcRdRP)Z>14Z3Ay3Ge(m#%w6s} zmXw9B(D}vKp=qNUAust{bE)LzQSuWtQ2J*BCmtlNkJH~;kNmO|wiL&l zZ`yV_V7I-SilvD|`^5|#tHspfqK_>~U9~kwNK!I0ZKW9Q%1M}4vskXZzp62L znCLv_#knX0g!BEG6i@Gk>x~90DmuIj$Im4dt{4?o$UQN#)46(ex@Pv`9+`o&*HuSLjFo0d<=z_FmWxduP;z8i7e&(a#Tx)1qs#H0JmVbC9HknpPf=U`URxWf#Whb+x()+W#ese1& zbe1up8__3z7bbamz5i8!ISZqGBpIKDcOwU~5Bk=ub5L@q6G1^|^5r%|LvrQ|TNX*k z=_jyTbdH#$pgKLN&@9&MD=ce{H6jVlf!XFPx5R@50%Gn2$%{Yuz?B|&`^Dai&=A+D z$5oD28%dbM<~Ke%Iy>->XPmy@_$B{xOhuIYndRAqs8a$_t!&Y(q&V`3FU#UVChT06 zho0&5jy~y(_*~iTQ9?UW>YTt_q3?tdTYVhHheF7seFD2+zUG{h2IAwmCDb$O}106JJ?hFzv2?v6JKQa?{{;pW{nfc*mm%>p{zo zh3ooHUA_qKS}g3=Pa+YXzWnyTrZj(m)eAI^*Huy#^9?>7BtA})aJm-36>DE0WTz;<@Z9yLv9!LzG@=~pl`Wx zxgmx7TC>w(UxUj|mo%&jR`rOeto^)6>NM8hZ<~*vdg;}jbjhtd>CszSdHPb#yY?Lx zUr&Cn@d&F)-EyqI;?h<(B=`PCiH61a)ewQk^X0N3!L0ZCZ8BdDp!wF0EmVfRpV;Ma zwqJg5?pYUS{gLNVp3In0lImmEmP#Fj#^pcDG|nXtd>EQ{lFOhj{1nratouAMt9>|! z`2AY1R=^EW#+K_>`;f+rwe+-!6}K>CyZ+gRh7}U*N4m8O)+n$(seMuLQjl^v<{Fa4 zqQ(>26OWuP@mMv}##&Lwo|J>!8|>)#SdcTng|2tnF1z2#jK|R3DN0_eplYi?o

Y z(N--%_QY?0S?jDDZ+Hj{XRp0Kxn5(v@`*a~5rJni za}jS{w1#4yc9#VwtB7y^!Xq#HIdrY_cKG_xjZrpMR4NK`LV8ZG-cr5Ap2)H&G5ryq z5{>=Y(ib|O1jn+tCmyH=u9l~*d<`$7^Jw?sk3XHI<0HIt#7ow7DNkg#{f+T!U)C^> z$B7+TdOpH@OkR?%w5z8nE7P0tl=%H|40hQ#fP79{F!S==)*HG3#vL!v!q936;CR_e z8(V~5@C!q4K$LO`*-Xp8UB4%7 zJYNW|{*s*9_Zo6e46TVg{a#2ofF?LLuS|G*EcbnoMGE4m6GcFYFacjdQFyRqp4{+) zH14qP`XYb&nViV2I+`g?WnaY`E#gcASL32fse$TiDDt8pj-0I31o$r=0+8sQXy6Nx zJIBAVBpi@N8NXX^0Cd?d7-VX#EUWql@23UsE*$gzr&}N2#n&5ulMNwzsN%u}Izs#x zxgSC=XVvHUB~arXS9M<3dh?(u| zxy_5cw>KAUFFc%%cg(k_iAgSlU+W@R8REE*ZPCW7j?^oRaTQv2KLY?34z_@|w}4C` ziWV=tOpGSIn&Olp)YoyJB)`UqRdv25;_`TID#FGr*PMRGR+Ov+-IDtDw)>U%SW)K& zUiFfh1WNsh*Ci)CDlBB!P?%6Lp3gZB!C9Nj`m1(b(W+K{jdUt>MO zxhvQ^>0#ezh1l9>JaV$>l_phdA_!k2cW+gutcVES*M`|L+)6pGY+xw7{&&X{`3#<) zVxOpV3O4Fj;;JY#Jsi#wZP219x4Sq!-9A~?eto#)i7ZiK>|1xOa_-^$DnsJ40(%=m zG3%yGL$%#!EeEN}JUH3R%>CR>A?FCoHPv*-Z*c5+R+c6T6eZt`;L=aCi75&jwNH~& zj6paSdkZ&|ES#!*zZL{O%ev6c?*lr%!l~+K4!D@Nny@{nu5SGj*5FizST5 zh~vG^A71pQk@(!CaO4uhR~;i#FXql<|K%XD8fLs3vDZHjzSH6=sqO3X()a z!R5mm)2SuLyn~~!KdftvPtbjNIl@%wM)-ATYG~KOJqMn|pe5(ntsC7o;fW!2_~q6* zEw$7I@N+6!fkc7Zh|9JkT(dbYWr;qLD_bF}MJxk2zh2Ja6CJNi4V+8X7vXh=r>q?D zgGi5`iXalId*dk_EO>uHo`v`PC^XIUV`0&N+4LpEKzwocmhn*S)0gUE+~s_d$z+g< z+*{4gO)c|xv@Ooux#&s!{C&S znHUQ~NlR&Q?N_-ktPEaN>b4t3$5f~6Ryd#biwk;?FfLDpb2wEI=EnD!O`oOBVla46 zzjJf6vdM&{O^A zJ#DZWCN0&kP|? z2E)C}Qrs4@SjuoW2d!2zp2VY4Y`(6s3*qVbdBXQ$&yj1>^zOWlQT#2 zZ!ek7_b@fhkC<%LGFY#Z(ep>dpwV%+q+u zx_V*pj!*#9n&|vwDo-EH=S3+OL;SI+4dyPMsn&PK0k($Eq8cHTA6pP^uy%|ChLb?5&8b>_~bl~zKW>oyPk`~34!DR7LopAF^NY-#_kaO*t0pRufo|ytSd=*i|JqX%5JG}ZVALwZ2rx0_>&0JW-|4X-C`xokyn9@5WWDYlhbDvbReOj&oEfj`YipXIL)D z=(`yz*;*T&cBRA+M3J83uR6NT_?Clu-QH^wmN|7HxygVBH~ z7nL3&zklm3oW5FEf0xc2%RPY4v|znM}146G4G{OmK;ZtbD?{HCi)w?)X7 z_X?TmD3|L_Fv;mpzi$+$HY`^j5)GVbkTfJp_FcI-n6h=Ic>2<)WLDfIbv-BrT;;rr z3g*nFG^PsD=auEBmA9$=@ud63Ekf(ctqPV-IE%To zLw$uX3`fPS!wmzno^Ld`+M=G7p-WivZc5$vUy~_0bLSAI&|y0XC0Yd>k}es z;&Yj!0daWte#NY=#s-FXRHD>0X0xyOPld=@P(Q!Jl1uU8yyJv;I=_N!Dp~2vSCavA zc%K&NmJZZB|_?;VJ3CKibMd+wjDG!dp+I$Jl-n^)-th0?elYpCA1ZWof!A+osDdX0vD;p0)#N!*JAu_=?zPkkfLun1$3v6+`? zMR!{{&7_RwsZYpXscxQ5c5;8#eq}VizB;gdhj&ID`;h%g6#8V6#FevklJrYOJ(&-! zI&|2_Frf+t6mbp7P49Ka+oo@?6zb0maMX2fHdH*PJ*PE3&)zcJ?mwa$^Y#VyQ(=7H zjzeikfajFdH5O?H*N0Gz?7{9)HP0_%!B-3GY{|%nP?^N2nz($3Kmk)KRoT-ra{B+2dRuq(dxedK)J%RrTJZ z@8{tn!nkN<-oKi--1_)vZelk-n?=z(7ga91Q-KWht!p_S-eeA65$)M3y2!jP zJ$wu4#QNT<`E7sebmZ5GFnqa)M^+3UCkyQcxs0pS_#{lt6rJ*|Lo^gOl__6VHdGUz z_+%p?#^xL~uxQlKef&zi4Yh~t4gr1UIEx&q=E$oY#gzhqk>1U7xK3=_!2%f1>Cf--n%3?US;2hpzCE z1s>4;0hr{k9u^SC{o`f_-&e6V-3#2<8&YkjJSNOs{)E{S$8>}nzJ(1M+luDNWbv|e z(%R)v3XrNVn3SRPO_(xt^XoXgh{naGXCWd`T7fX{C2hS*9(5+whwolKj$B}{X9yc=yq*U<_$&ATzr15`3zM| zlQyBCN~Ik9M&zPRI%|5>g}XOxgf0h(CtR?c46duY&`Y{hxokPD!C+-D{=gQOBPrYR z@;QluLiDk1aU#od)O+f$-aPVP~ z`1yke<>=vU@2Am8ijf}}+Fnx7C3z$qMFxP?yYT8n; zbYg#(@`hNbvwGHgTR7tFJZ}_hTFA=AD-JiT7wNm>=r9JVxh-~2KlIIvrK>AP&Nzwe z8H{hyOFWxrhMpy&9#k=XpYq1vsM^hizU%_Iv*;ZWE((UO1@`+V?kMM%WZqL%amW>V zFrRqmHZ*p%hH%8}=L zPx?v`!($6)m<#Sb@_im6x@tH5?q)sop_nSP+mkHFW$8XTih_?W1$oH?gA1>ffD@)H zxZD>!%7Xn|5Ak$hcGrWNW4RTr#~Lj-3!PF`4Q|W68{2-# zwmlZYXVHhxL0&m&Mxnsl@k~ZFlqMFjNPpjr@kYKtYbM{knL>9yQ_{#wrB_K?v49=J@XT35;p6Pf-w~g#2&M~T~Y^mmOnTi?&o$X0K`9U!K3b{^k^(vMl zGs(H;=(&I=_{j9bxEnHv%AB`x3>1pQuC%1OS@ln@m^Mu7_6dZI*?pF@FlX^6i7*o# zJM9w$eRt&XX^vPlngXHIjkiqFMxtvxlzUbs#P4G4AdyRkW@BFGWXan>H;fqP+GNg+ z)g3F%1ml1{xlh_2Mgc;Jxk+yCqeVTZ6!;lC=^LNrC$fy!a3^2rnoj_bSo+EwlPjI~as^mF> zI2J=H?W3}|G35y1>M1AAwlM-S2UhRVUiFmtm+sb4G8$}SS>#vuNG(|+bFTK`Qa6v| zJ~Ig%z~up;6p8wakwrrF8~%H)=HM$ml@R%{G!OK`ddczwU6Vw{HXa>n_s!aiDw&zi z55+3oU=+*0zl(A@&v<`JhKRr|E~4d%Tz=S6R1W3c(bbvab;w#P+XNP(mlp*5HY_Ea zBt!c{c=;HjbzflpreEBn*PYSF6i?eS6<2D#et&K%gx|HePHgE|16^ZbPk zA<;;6z)=^g!w_f{7#elNlp=sqjO*U9AG2iPJf|o3+MSJo2=U~_A%fGShyz+b0Mlar z0@5N870SR29d(%e7Bvgxtp(x!480~dxKfnslr(W;`{5E|pHk)&zaUGWM5yhFCMdOd zaZ)giI9_=;3Bn%XXte=FmrBW;Mso+hl1 zkp7|PJqF&JPwNV%jN6~OP^_rU3^b>Z6XkfwD#SlM?%7{j(_dpX`Rv_Wxq;wIT|K8# z$)3^Cj32<-0f_2v)(s96Rs0D}u2b?HM>~)7#XCS4)9T^~4^>){upZFtK|qR!_rFCV z5e3XTpvS~F)BKoxc4q}*WX<7~duWnwGqhVm&KG}Kru*a^bCg$MH`>;^%*q$sv6}p6 z2`D8lB(`bH7aw`Zuhy7!CHV3?ocZjNm;8iHXS)rsxn_=;p|0w#><;ohai3Jix7~DN zZ%C1vLx%I6X>BIkZ5vLrN+uV)C*2H&DzU0bH7!F&J{)BYyAk*BqC=|9MxPvv6#t7T zMdECT@H=AP<}S7MS@NZs8MirKWBaQorAJP*MlqbG^K(np-KU6!bMFKn)x>hf ztkeDX-P>XgpNu{|Hyy_&2m7|Q(;e$tR2$ny^JhGL+V(Z8Wh{Tl>*A7ziM5Vg2T?`p zGYes&F*WSyL%~buAbcAK^z|Tk1pf=O5mc@W!-o@h7zCWruJv5bwWJP_yT7+Cy}iNM zM=hbs&XW|)#_bx*tfMx;cRjw5^y;lQc4la1b_wp7BIeCr$m)&Yx0IXX1Mwwds01sl z*Ks6!-dVbra_Bn_>=aM$jWr;5NNip3Gz002^`O03#zW;$rxvo)#C1s&A&kVb3WR%) zoFeWThMW~@SQVp;iY9%_v#0%K&6%W7ErK=5&c0^aFeR51?<8&Xohl{+?P4zSlh!;U z3O%*DoIcMg48NW&WqL|1yb;{!ct>lnntZKsti0Ayzuan&xLh+&+M z{OVmWe~%9TA780m@vo9J4~LUO)w4zw2W_{h^CB|kY0*B zLm7dg^wuDyz*z+|zqFPhx+gJh^LQ(&dd05uhnG)ioHr;T#uXPlbr3Yc=V|h!dix26 z3yzk$t0DL9FagEZi*4`sb}Mc-4jVVY_Lug5-3wsD zes9lX0j1mb0$uj5Fh3w5PK33wj!@r>3Kf$#UP{DGF{>*T2%TrDcq)TEk*kcsSF(~P zMy6!RO~x(MmRU<>*QQX^Z(b!<`MG+hF>G9N5aDeiDP$K=G(n3J~VtBotRw);^f zp6Yjc-SectDV>78{9;1}e^ZcCGL7}tmz6Axfzg-;Y~WAz0h;==f9StZpF|{<91ngB z-?`e4Ra)O~Gryy(Zf7s|=%7@_!VwhhQ)-2FPf9DHg<{F{r@dF^Ib@^7>!_pYlpo-N z_%BGyOXH%Y%js`2-O(;FNK8^^FHxN;wQj0^KB7C&Djwz#wi~=asJ^STs2}@5;bvhp zeCM=!75~jpnJq(!FmeNBUJOUrXF^WGS^9TKclSjeYnR zZQ0%cJAKL9Op$kU83pcTwRN&UrXy15o!IeA7VkWAaXPEv_#my5Cpn}7DP!@jM!Gnu z{;C0Q$yZ`Ei?XY%Jx&?@xyM$DE#-Q9v5(TK<8`cDk}0aP{Ck=el@x2Vknb3|MClu4#UugHzmRMhTx(tQ2A_I{P$>+4|-=j+Y~sO10@g+%;?{qX>Zp8#0#?O}T& zoc%$m`(s{SNX=*aD-_jK@-X8WHg264ay9+Tud*9v_0DZF7mH3rEZ^x=ZQ7(Rq3q!- zl#KFFg7R$DO_@VCFK->C&qhL(=Y*S|nIcs2 z{LT+MTRR>Lc47(k&Z?YQ+t!lJp-BTdr671;?W83 zZWN2I->Quk5R!to`7EJ>Z8Hd1=IFDTI~ytaidv3TS=<@Wm^ZGc4`d6~I$n}PZ?mRH z5tl_*Fher}o9KY4nWj$_5%j+`HTQe}W`?J*K4@7b#-PFBeSDYuu|BJ>XQ|sKaa7Sh zQV)vBNvwHmgKE|U*mK&7L0pOfT{P~thdsK)Z^ zs=IXr2B?=Uxonrljgn(#-OU3UiP|2LFRR{?&Z*L57+6U{6jtz9;C@6d`uWrofLl?LZ=x_?Ip5$Pb%!j#%XnJXx}oej6-O!a+Hso@2E5}_8Nzq%DHvq z1vth0=vj#|5ntx$L~PSjm-W|My~X#w4u@i%^RX}J>FVHK2%TMyZRS8Mcxtb^d>m}K zc^tKpu{$?FOpY=Si*7zzFp9mEk))#OQ6UyX(gkl}_(+=8Ktm+EBS85)f|-IYw)HhB zNxHWYeP(D1dg932j7zpuM@gqk8VoCasbo)GiQYV%D?jhMuvfTRRx>noj&9kf`2gh{ z1SS8*TXU+&IN^yLd)9Mt-QVKDYh1#`z}dZ_W9>%?NkW8Obf@oy-i2~#RlSLNXw~o< z3HN(V5a@Mr$X0t+SRE>ZQ|3u`TlHM0A9hS6>WmpIBDTmU+YMzKn{S$e(^ysvnfISq zGAI;G5^2lVZ)B!8B|l)oWJ!7{3BgtBEQo>5NopsUMA~VdL6V1~0;CxN@{+T<@kPi( zA4!IxohUzNb4R$`(g@3eFFAAC8ucjfUg=%Z&R?62r8jWoVm-y3PrLpJjdx1LI)6=A zy7r45*(c}6F|R%3=DMk0&RpwiTdP+$Zb}}q6GjhAzpe_J`5am}H@CfuXJo7wbbt&G zz?yK(UpPvINEuSG{k^05Q#t;~+Bt#ec~7lYo)?H6b@knyg*pfb^g5pBaxt3^4ZoO7 zJ5yT%#ZC9|?zKG;YJ`i7XK-pKxV4OQk~`NU;77_O+eQR3YROvkkCB@toU5sYG8#l| zYkt&;S?{UfACGuiYGEXpQ5(E;tf}a*md}@`ExvN~+lDV^6WEz?6oG5Li;Svh{`3=N z>Pt>}w@b#8FFH`I@AiqkUnVGH5_o+;w-3OvynnI6$OAxZ9~iOtHKnPO)6YkaGC+o( zH_15QUXRBZs!j{P;X4_A-Rzyop`LKKaVUju`2YqFK%xlvU+s-zNr5Th(kAe_$@x%j52cI?U0%=FA?+6?g z8|uqZ5j#cS+n=p=p=bK5=);86{7zkPseE@3JFGm?eO~p6z_CVFA(JY%_s9O z<5ls0fhit3MB&(R#O$tYY{_h+eerq;q1MjV!*`|Bau?h4{QVQVu~iVz-czo->d#sUi8}uHcmJUV2=|8%;4b|EDAVw2KYbV z(_=5!D z7Pq#rFb7W7fk)}|px~abvx5l^IOPZQ&jCDgXJc;R@;&KfZE58KMFBQsb=ASj+zyHY zX5F6VPWDi4dpn?0F780e#mdRt9Kvk@9M6MtTUfgRbZ3BHG_W};21Nr;hy>6ND7O?8 z0|-L~3cf?207wN2zSg4xxy>Nl=D^83Agd)10WDht5s={ul-mZ% zZ42eL1KNhd;lOwUGzTaM;MEnl21>gFVgtygfUXR{Q+?c8)@CkNz&Stg!T9+EXkqWdcqia6uIV~kh?z#-TIc)^K40{fR@CJxf(`x1sh zxV3*I0T}_uH?46Zc9u5g0OQ-2$27; z^&kO6hysih1@3>bJ2Pz06-0t6@;*hJ>x z0V)BWIt*ah7+!z@!1DHE@ILUa9vV2FiUG=lbz*_x9?koW2FClHL1G|a3^DQ zfv~R+u(@vm!t-5c|5WuY0WkQj3gA^ouow@>6~H-I95jO;>iQ7_DueuBWWNR!2FL`w zPzkaR+QC2dgSkNBzA*tcf!F`k0x|`ILH6%`uK)`SG(Q4p2;Xb~w1Z#6zN`Ve0Bz!@ zux}5bP5cP^w({L3KzsN(fPC#=19gD5@gsmeN5gmk9Rj?9HgG_|@&Y;m?E-MjKZhSS z@goL>0_YvGAA{+i!oIG*-3Rmh6F|EFwF<`nzJ&Vo{8mpeX{f0aFlQ(B6Ly3!te2n*Mgz1oR~P`M}{9a9;aq z928)0KV=3y0pC*QfTF<^;4wg70X)R}8TQ==AlUs2OTg~GrGEIG@4nR<3T*c6XSo6e z#=!k*Y@n#`&Dlax-%1iSyk^D2NN5pWF%PXT!JzpG9G z0R^6oZ&UyceV3d9fQi0~PJsss3;P56 zVdd6pr$nbd&G@cK)bn5USydo!7Riub$QHS!AvgfzeyC8Ql_B{uQP#*!d4bUzCOLIF znrG&U1HMD@b1?@q3#h3H?h4S%k1Bwy`wdvz+lhl;aJeAPhXCHK zfT7?>7y<#~L2$xQtS}fW+fPHNhm*Ml1Oz4!07HNO10XUF;I}QH|Ionk82p3UeW$@O zzz0`;rvWGgFq2>Np+T7ZJBR6aayLr=dZ3^;>;dV9N9x4GjaP zT))#0AiVmuEMNz}^9(~{{*V 0 else float("nan") + + +def parse_args(): + p = argparse.ArgumentParser(description="Plot ablation studies") + p.add_argument("--predictions", type=str, required=True) + p.add_argument("--output", type=str, default="figures/ablations") + return p.parse_args() + + +def main(): + args = parse_args() + + with open(args.predictions) as f: + data = json.load(f) + + entries = data["predictions"] + T_overhead = data.get("T_overhead_seconds", 0.0) + + # --- Panel (a): topology sweep (SBM inter-density) --- + # Filter to SBM entries, group by inter_density + sbm_entries = [e for e in entries + if e["config"].get("graph", "") == "sbm"] + + density_groups = {} + for e in sbm_entries: + d = e["config"].get("sbm_inter_density", 0.0) + density_groups.setdefault(d, []).append(e) + + densities = sorted(density_groups.keys()) + mape_full = [] + mape_flat = [] + + for d in densities: + grp = density_groups[d] + full_errs, flat_errs = [], [] + for e in grp: + T_meas = e["measured_median_seconds"] + # Full hierarchical prediction is already in predictions.json + T_pred_full = e["predicted_seconds"] + full_errs.append(relative_error(T_pred_full, T_meas)) + + # Flat model: use a single network term = T_intra + T_inter (not max) + # Approximate: flat model can't overlap, so T_comm = T_intra + T_inter + bd = e.get("breakdown", {}) + T_comm_hier = bd.get("T_comm_seconds", 0.0) + # Flat approximation: assume both intra and inter are sequential + c_intra = e["partition_stats"].get("c_intra_bytes", 0) + c_inter = e["partition_stats"].get("c_inter_bytes", 0) + # Without knowing the individual bandwidths, use ratio heuristic: + # flat ≈ 2 * max (conservative estimate) + T_comm_flat = T_comm_hier * 2.0 + T_pred_flat = (T_pred_full - T_comm_hier + T_comm_flat) + flat_errs.append(relative_error(T_pred_flat, T_meas)) + + mape_full.append(np.mean(full_errs) * 100 if full_errs else float("nan")) + mape_flat.append(np.mean(flat_errs) * 100 if flat_errs else float("nan")) + + # --- Panel (b): with vs. without T_buffer_copy --- + meas_all = np.array([e["measured_median_seconds"] * 1e3 for e in entries]) + pred_full = np.array([e["predicted_seconds"] * 1e3 for e in entries]) + pred_nobuf = np.array([ + (e["predicted_seconds"] - e.get("breakdown", {}).get("T_buffer_copy_seconds", 0.0)) * 1e3 + for e in entries + ]) + + fig, (ax_a, ax_b) = plt.subplots(1, 2, figsize=(9, 3.8)) + + # --- Panel (a) --- + if densities: + ax_a.plot(densities, mape_full, "o-", color=COLORS["full"], + markersize=5, linewidth=1.2, label="Hierarchical (intra+inter)") + ax_a.plot(densities, mape_flat, "s--", color=COLORS["flat"], + markersize=5, linewidth=1.2, label="Flat (single-tier)") + ax_a.set_xlabel("SBM inter-block edge density") + ax_a.set_ylabel("MAPE (%)") + ax_a.set_title("(a) Hierarchical vs. Flat Model\nover Topology Sweep", fontsize=9) + ax_a.legend() + ax_a.grid(True, linestyle=":", linewidth=0.4) + else: + ax_a.text(0.5, 0.5, "No SBM data found.\nRun bench_end_to_end with --graph sbm.", + ha="center", va="center", transform=ax_a.transAxes, fontsize=8) + ax_a.set_title("(a) Hierarchical vs. Flat Model", fontsize=9) + + # --- Panel (b) --- + lo = min(meas_all.min(), pred_full.min(), pred_nobuf.min()) * 0.9 + hi = max(meas_all.max(), pred_full.max(), pred_nobuf.max()) * 1.1 + ax_b.plot([lo, hi], [lo, hi], "k--", linewidth=0.8, label="Ideal") + ax_b.scatter(meas_all, pred_full, s=18, color=COLORS["full"], + alpha=0.8, label="Full model", zorder=3) + ax_b.scatter(meas_all, pred_nobuf, s=18, color=COLORS["no_buffer"], + alpha=0.6, marker="^", label="Without $T_{\\mathrm{buf}}$", zorder=3) + ax_b.set_xlim(lo, hi) + ax_b.set_ylim(lo, hi) + ax_b.set_xlabel("Measured $T_{\\mathrm{layer}}$ (ms)") + ax_b.set_ylabel("Predicted $T_{\\mathrm{layer}}$ (ms)") + ax_b.set_title("(b) Ablation: With vs. Without\n$T_{\\mathrm{buffer-copy}}$ Term", fontsize=9) + ax_b.legend() + ax_b.grid(True, linestyle=":", linewidth=0.4) + + fig.tight_layout() + + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(str(out) + ".pdf", bbox_inches="tight") + fig.savefig(str(out) + ".png", bbox_inches="tight", dpi=300) + print(f"[plot_ablations] Saved {out}.pdf and {out}.png") + print( + "Caption: Ablation studies. " + "(a) MAPE of the hierarchical cost model (separate intra/inter tiers) vs. " + "a flat model (single bandwidth parameter) across the SBM topology sweep. " + "The hierarchical model degrades more gracefully as inter-block density increases. " + "(b) Predicted vs. measured scatter with (blue circles) and without (red triangles) " + "the $T_{\\mathrm{buffer-copy}}$ term, demonstrating that omitting this term " + "causes systematic under-prediction for large halo sizes." + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/visualization/plot_compute.py b/experiments/cost_model_benchmarks/visualization/plot_compute.py new file mode 100644 index 0000000..20ef9ce --- /dev/null +++ b/experiments/cost_model_benchmarks/visualization/plot_compute.py @@ -0,0 +1,172 @@ +"""Visualization — GNN Compute Primitive Runtime. + +Two-panel figure: GCN-like (left) vs. edge-conditioned (right). +Each panel shows forward runtime vs. the swept variable (vertices or edges) +with the fitted linear model overlaid. + +Usage:: + + python -m visualization.plot_compute \\ + --gcn-vertex data/compute_gcn_vswp.json \\ + --gcn-edge data/compute_gcn_eswp.json \\ + --edge-vertex data/compute_edge_vswp.json \\ + --edge-edge data/compute_edge_eswp.json \\ + --primitives data/fitted_primitives.json \\ + --output figures/compute +""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +COLORS = {"gcn": "#2ca02c", "edge": "#ff7f0e"} +plt.rcParams.update( + { + "font.size": 9, + "axes.labelsize": 9, + "legend.fontsize": 8, + "xtick.labelsize": 8, + "ytick.labelsize": 8, + "figure.dpi": 300, + "text.usetex": False, + } +) + + +def load_compute_file(path: str, timing_key: str = "forward_trials_seconds"): + """Returns lists of (sweep_value, median, q25, q75).""" + with open(path) as f: + data = json.load(f) + sweep = data["config"]["sweep"] + rows = [] + for meas in data["measurements"]: + trials = np.array(meas[timing_key]) + rows.append( + ( + meas["params"]["sweep_value"], + meas["params"]["num_vertices"], + meas["params"]["num_edges"], + float(np.median(trials)), + float(np.percentile(trials, 25)), + float(np.percentile(trials, 75)), + ) + ) + rows.sort(key=lambda r: r[0]) + return sweep, rows + + +def fitted_compute(sweep_vals, fixed_val, sweep, model_type, primitives): + params = primitives.get("compute", {}).get(model_type, {}).get("forward", None) + if params is None: + return None + a, b, c = params["coeff_V"], params["coeff_E"], params["intercept"] + if sweep == "vertices": + V_arr = np.array(sweep_vals, dtype=float) + E_arr = np.full_like(V_arr, fixed_val) + else: + E_arr = np.array(sweep_vals, dtype=float) + V_arr = np.full_like(E_arr, fixed_val) + return a * V_arr + b * E_arr + c + + +def plot_one_panel(ax, rows, sweep, fixed_val, model_type, primitives, color, title): + xvals = [r[0] for r in rows] + meds = np.array([r[3] for r in rows]) * 1e3 + lo = np.array([r[3] - r[4] for r in rows]) * 1e3 + hi = np.array([r[5] - r[3] for r in rows]) * 1e3 + + ax.errorbar( + xvals, + meds, + yerr=[lo, hi], + fmt="o", + markersize=4, + color=color, + capsize=2, + linewidth=0.8, + elinewidth=0.8, + label="Measured (IQR)", + ) + + fit = fitted_compute(xvals, fixed_val, sweep, model_type, primitives) + if fit is not None: + ax.plot(xvals, fit * 1e3, "--", color=color, linewidth=1.2, label="Fit") + + xlabel = "|V| (vertices)" if sweep == "vertices" else "|E| (edges)" + ax.set_xlabel(xlabel) + ax.set_ylabel("Forward time (ms)") + ax.set_title(title, fontsize=9) + ax.set_xscale("log") + ax.set_yscale("log") + + ax.legend() + ax.grid(True, which="both", linestyle=":", linewidth=0.4) + + +def parse_args(): + p = argparse.ArgumentParser(description="Plot GNN compute primitive results") + p.add_argument("--gcn-vertex", type=str, default=None) + p.add_argument("--gcn-edge", type=str, default=None) + p.add_argument("--edge-vertex", type=str, default=None) + p.add_argument("--edge-edge", type=str, default=None) + p.add_argument("--primitives", type=str, default=None) + p.add_argument("--output", type=str, default="figures/compute") + return p.parse_args() + + +def main(): + args = parse_args() + primitives = {} + if args.primitives: + with open(args.primitives) as f: + primitives = json.load(f) + + fig, axes = plt.subplots(1, 2, figsize=(7, 3)) + + panel_map = [ + ("gcn", "edges", args.gcn_edge, axes[0], "GCN-like"), + ("edge", "edges", args.edge_edge, axes[1], "Edge-conditioned"), + ] + + for model_type, sweep_label, path, ax, title in panel_map: + if path is None: + ax.set_visible(False) + continue + sweep, rows = load_compute_file(path) + fixed_val = rows[0][2] if sweep == "vertices" else rows[0][1] # E or V fixed + plot_one_panel( + ax, + rows, + sweep, + fixed_val, + model_type, + primitives, + COLORS[model_type], + title, + ) + + # fig.suptitle("GNN Compute Primitive: Forward Runtime vs. Graph Size", fontsize=10) + fig.tight_layout() + + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(str(out) + ".pdf", bbox_inches="tight") + fig.savefig(str(out) + ".png", bbox_inches="tight", dpi=300) + print(f"[plot_compute] Saved {out}.pdf and {out}.png") + print( + "Caption: Forward runtime of a single GNN layer vs. subgraph size " + "(vertex sweep top row, edge sweep bottom row) for GCN-like (left) " + "and edge-conditioned (right) message functions. " + "Dashed lines: fitted model $T_{\\mathrm{comp}} = a|V| + b|E| + c$. " + "Error bars span IQR over 50+ trials." + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/visualization/plot_gather.py b/experiments/cost_model_benchmarks/visualization/plot_gather.py new file mode 100644 index 0000000..cddfdc4 --- /dev/null +++ b/experiments/cost_model_benchmarks/visualization/plot_gather.py @@ -0,0 +1,120 @@ +"""Visualization — Gather / Scatter-Add Bandwidth. + +Single plot with three curves (contiguous, clustered, random) showing +gather (or scatter-add) runtime vs. k (number of rows gathered). + +Usage:: + + python -m visualization.plot_gather \\ + --contiguous data/gather_contiguous_*.json \\ + --clustered data/gather_clustered_*.json \\ + --random data/gather_random_*.json \\ + --operation gather \\ + --output figures/gather +""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +COLORS = { + "contiguous": "#1f77b4", + "clustered": "#2ca02c", + "random": "#d62728", +} +plt.rcParams.update({ + "font.size": 9, "axes.labelsize": 9, "legend.fontsize": 8, + "xtick.labelsize": 8, "ytick.labelsize": 8, + "figure.dpi": 300, "text.usetex": False, +}) + + +def load_gather_file(paths: list, timing_key: str): + """Merge multiple JSON files, return sorted (k, median, q25, q75) arrays.""" + rows = [] + for p in paths: + with open(p) as f: + data = json.load(f) + for meas in data["measurements"]: + trials = np.array(meas[timing_key]) + rows.append(( + meas["params"]["k"], + float(np.median(trials)), + float(np.percentile(trials, 25)), + float(np.percentile(trials, 75)), + )) + rows.sort(key=lambda r: r[0]) + k_arr = np.array([r[0] for r in rows]) + med_arr = np.array([r[1] for r in rows]) + q25_arr = np.array([r[2] for r in rows]) + q75_arr = np.array([r[3] for r in rows]) + return k_arr, med_arr, q25_arr, q75_arr + + +def parse_args(): + p = argparse.ArgumentParser(description="Plot gather/scatter-add bandwidth") + p.add_argument("--contiguous", nargs="+", default=[], metavar="FILE") + p.add_argument("--clustered", nargs="+", default=[], metavar="FILE") + p.add_argument("--random", nargs="+", default=[], metavar="FILE") + p.add_argument("--operation", choices=["gather", "scatter_add"], default="gather") + p.add_argument("--output", type=str, default="figures/gather") + return p.parse_args() + + +def main(): + args = parse_args() + timing_key = ( + "gather_trials_seconds" if args.operation == "gather" + else "scatter_add_trials_seconds" + ) + + fig, ax = plt.subplots(figsize=(5, 3.5)) + + for dist_name, files in [ + ("contiguous", args.contiguous), + ("clustered", args.clustered), + ("random", args.random), + ]: + if not files: + continue + k, med, q25, q75 = load_gather_file(files, timing_key) + color = COLORS[dist_name] + ax.errorbar( + k * 1e-6, med * 1e3, + yerr=[(med - q25) * 1e3, (q75 - med) * 1e3], + fmt="o-", markersize=3, color=color, linewidth=0.9, + capsize=2, elinewidth=0.8, label=dist_name.capitalize(), + ) + + ax.set_xscale("log") + ax.set_yscale("log") + op_label = "Gather $x[\\mathrm{idx}]$" if args.operation == "gather" \ + else "Scatter-add (backward)" + ax.set_xlabel("k (millions of rows gathered)") + ax.set_ylabel(f"{op_label} time (ms)") + ax.set_title(f"Buffer-Copy Bandwidth: {op_label}", fontsize=9) + ax.legend() + ax.grid(True, which="both", linestyle=":", linewidth=0.4) + fig.tight_layout() + + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(str(out) + ".pdf", bbox_inches="tight") + fig.savefig(str(out) + ".png", bbox_inches="tight", dpi=300) + print(f"[plot_gather] Saved {out}.pdf and {out}.png") + print( + f"Caption: {op_label} time vs. gather size $k$ for three index " + "distributions: contiguous (best case, cache-friendly), clustered " + "(METIS-partitioned halo pattern), and random (worst case). " + "Error bars span IQR. The gap between contiguous/clustered and random " + "quantifies the cache-miss penalty relevant to poorly-partitioned graphs." + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/visualization/plot_pingpong.py b/experiments/cost_model_benchmarks/visualization/plot_pingpong.py new file mode 100644 index 0000000..b7625a4 --- /dev/null +++ b/experiments/cost_model_benchmarks/visualization/plot_pingpong.py @@ -0,0 +1,147 @@ +"""Visualization — Ping-Pong Bandwidth / Latency. + +Produces a log-log plot of one-way transfer time vs. message size for intra- +and inter-node measurements, with fitted lines overlaid and a residuals inset. + +Usage:: + + python -m visualization.plot_pingpong \\ + --intra data/pingpong_intra_*.json \\ + --inter data/pingpong_inter_*.json \\ + --primitives data/fitted_primitives.json \\ + --output figures/pingpong +""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.ticker import LogLocator + +# --------------------------------------------------------------------------- +# Shared style +# --------------------------------------------------------------------------- +COLORS = {"intra": "#1f77b4", "inter": "#d62728"} +plt.rcParams.update({ + "font.size": 9, + "axes.labelsize": 9, + "legend.fontsize": 8, + "xtick.labelsize": 8, + "ytick.labelsize": 8, + "figure.dpi": 300, + "text.usetex": False, +}) + + +def load_measurements(files: list) -> tuple: + """Return (byte_sizes, medians, q25, q75) from a list of JSON files.""" + byte_sizes, medians, q25s, q75s = [], [], [], [] + for p in files: + with open(p) as f: + data = json.load(f) + for meas in data["measurements"]: + trials = np.array(meas["trials_seconds"]) + byte_sizes.append(meas["params"]["message_bytes"]) + medians.append(float(np.median(trials))) + q25s.append(float(np.percentile(trials, 25))) + q75s.append(float(np.percentile(trials, 75))) + order = np.argsort(byte_sizes) + return ( + np.array(byte_sizes)[order], + np.array(medians)[order], + np.array(q25s)[order], + np.array(q75s)[order], + ) + + +def fitted_line(bytes_arr: np.ndarray, primitives: dict, mode: str) -> np.ndarray: + net = primitives.get("network", {}).get(mode, None) + if net is None: + return None + t_L = net["latency_seconds"] + B = net["bandwidth_bytes_per_sec"] + return t_L + bytes_arr / B + + +def parse_args(): + p = argparse.ArgumentParser(description="Plot ping-pong results") + p.add_argument("--intra", nargs="+", default=[], metavar="FILE") + p.add_argument("--inter", nargs="+", default=[], metavar="FILE") + p.add_argument("--primitives", type=str, default=None) + p.add_argument("--output", type=str, default="figures/pingpong") + return p.parse_args() + + +def main(): + args = parse_args() + + primitives = {} + if args.primitives: + with open(args.primitives) as f: + primitives = json.load(f) + + fig, axes = plt.subplots(1, 2, figsize=(7, 3.2)) + ax_main, ax_res = axes + + for mode, files in [("intra", args.intra), ("inter", args.inter)]: + if not files: + continue + xb, med, q25, q75 = load_measurements(files) + color = COLORS[mode] + label = "NVLink (intra)" if mode == "intra" else "InfiniBand (inter)" + + yerr_lo = med - q25 + yerr_hi = q75 - med + ax_main.errorbar( + xb * 1e-6, med * 1e3, + yerr=[yerr_lo * 1e3, yerr_hi * 1e3], + fmt="o", markersize=3, color=color, label=label, + capsize=2, linewidth=0.8, elinewidth=0.8, + ) + + fit = fitted_line(xb, primitives, mode) + if fit is not None: + ax_main.plot(xb * 1e-6, fit * 1e3, "--", color=color, + linewidth=1.2, label=f"{label} fit") + # Residuals + residuals = (med - fit) / med * 100 # percent + ax_res.plot(xb * 1e-6, residuals, "o-", markersize=3, color=color, + linewidth=0.8, label=label) + + ax_main.set_xscale("log") + ax_main.set_yscale("log") + ax_main.set_xlabel("Message size (MB)") + ax_main.set_ylabel("One-way transfer time (ms)") + ax_main.legend(loc="upper left") + ax_main.grid(True, which="both", linestyle=":", linewidth=0.4) + + ax_res.axhline(0, color="k", linewidth=0.6, linestyle="--") + ax_res.set_xscale("log") + ax_res.set_xlabel("Message size (MB)") + ax_res.set_ylabel("Residual (%)") + ax_res.legend() + ax_res.grid(True, which="both", linestyle=":", linewidth=0.4) + + fig.suptitle("Network Ping-Pong: Transfer Time vs. Message Size", fontsize=10) + fig.tight_layout() + + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(str(out) + ".pdf", bbox_inches="tight") + fig.savefig(str(out) + ".png", bbox_inches="tight", dpi=300) + print(f"[plot_pingpong] Saved {out}.pdf and {out}.png") + print( + "Caption: Log-log plot of one-way network transfer time vs. message size " + "for intra-node (NVLink) and inter-node (InfiniBand) communication. " + "Points show medians; error bars span the 25th--75th percentile (IQR). " + "Dashed lines show linear-latency fits $T = t_L + s/B$. " + "Right panel shows residuals (\\%) from the fit." + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py b/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py new file mode 100644 index 0000000..735c96d --- /dev/null +++ b/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py @@ -0,0 +1,141 @@ +"""Visualization — T_global vs. K (Tipping Point). + +Shows total training throughput or per-layer time as a function of the number +of GPUs K for a fixed graph. Overlays the cost-model prediction (solid) and +measured values (dashed with markers). Annotates K* — the point beyond which +adding more GPUs yields diminishing returns. + +T_global(K) is computed as: + + T_global(K) = T_layer(K) * num_layers * num_epochs + +For the tipping-point annotation K* is the largest K where the speedup +relative to K=1 is still within 10% of linear. + +Usage:: + + python -m visualization.plot_tipping_point \\ + --predictions data/predictions.json \\ + --num-layers 3 \\ + --num-epochs 100 \\ + --graph erdos_renyi \\ + --output figures/tipping_point +""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +plt.rcParams.update({ + "font.size": 9, "axes.labelsize": 9, "legend.fontsize": 8, + "xtick.labelsize": 8, "ytick.labelsize": 8, + "figure.dpi": 300, "text.usetex": False, +}) + + +def parse_args(): + p = argparse.ArgumentParser(description="Plot T_global vs. K tipping point") + p.add_argument("--predictions", type=str, required=True) + p.add_argument("--num-layers", type=int, default=3) + p.add_argument("--num-epochs", type=int, default=100) + p.add_argument("--graph", type=str, default=None, + help="Filter to this graph type (optional)") + p.add_argument("--output", type=str, default="figures/tipping_point") + return p.parse_args() + + +def main(): + args = parse_args() + + with open(args.predictions) as f: + data = json.load(f) + + entries = data["predictions"] + + # Filter by graph type if requested + if args.graph: + entries = [e for e in entries + if e["config"].get("graph", "") == args.graph] + + # Group by world_size + ws_to_meas = {} + ws_to_pred = {} + for e in entries: + K = e["config"].get("world_size", 1) + ws_to_meas.setdefault(K, []).append(e["measured_median_seconds"]) + ws_to_pred.setdefault(K, []).append(e["predicted_seconds"]) + + if not ws_to_meas: + print("[plot_tipping_point] No data found. Exiting.") + return + + K_vals = sorted(ws_to_meas.keys()) + T_meas = np.array([np.median(ws_to_meas[K]) for K in K_vals]) * args.num_layers * args.num_epochs + T_pred = np.array([np.median(ws_to_pred[K]) for K in K_vals]) * args.num_layers * args.num_epochs + K_arr = np.array(K_vals, dtype=float) + + # Ideal linear scaling from K=1 + T_single = T_meas[0] # K=1 reference + T_ideal = T_single / K_arr + + # Identify K*: largest K where speedup is ≥ 90% of ideal + speedup = T_single / T_meas + ideal_speedup = K_arr + efficiency = speedup / ideal_speedup + kstar_idx = np.where(efficiency >= 0.9)[0] + K_star = K_arr[kstar_idx[-1]] if len(kstar_idx) else K_arr[0] + + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 3.5)) + + # --- Left panel: T_global vs K --- + ax1.plot(K_arr, T_ideal, "k:", linewidth=1, label="Ideal linear scaling") + ax1.plot(K_arr, T_pred, "-", color="#1f77b4", linewidth=1.5, label="Predicted") + ax1.plot(K_arr, T_meas, "o--", color="#d62728", markersize=5, linewidth=1.2, + label="Measured") + ax1.axvline(K_star, color="gray", linestyle="--", linewidth=0.8) + ax1.text(K_star * 1.05, ax1.get_ylim()[1] * 0.95, + f"$K^* = {int(K_star)}$", fontsize=8, color="gray", va="top") + ax1.set_xlabel("Number of GPUs ($K$)") + ax1.set_ylabel(f"$T_{{\\mathrm{{global}}}}$ (s)\n" + f"({args.num_layers} layers × {args.num_epochs} epochs)") + ax1.set_title("Training Time vs. GPU Count", fontsize=9) + ax1.legend() + ax1.grid(True, linestyle=":", linewidth=0.4) + + # --- Right panel: scaling efficiency --- + ax2.plot(K_arr, efficiency * 100, "o-", color="#2ca02c", markersize=5, linewidth=1.2, + label="Scaling efficiency") + ax2.axhline(90, color="gray", linestyle="--", linewidth=0.8, label="90% threshold") + ax2.axvline(K_star, color="gray", linestyle="--", linewidth=0.8) + ax2.set_xlabel("Number of GPUs ($K$)") + ax2.set_ylabel("Scaling efficiency (%)") + ax2.set_title("Strong Scaling Efficiency", fontsize=9) + ax2.set_ylim(0, 110) + ax2.legend() + ax2.grid(True, linestyle=":", linewidth=0.4) + + fig.suptitle("Tipping-Point Analysis: When Does Adding GPUs Stop Helping?", fontsize=10) + fig.tight_layout() + + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(str(out) + ".pdf", bbox_inches="tight") + fig.savefig(str(out) + ".png", bbox_inches="tight", dpi=300) + print(f"[plot_tipping_point] Saved {out}.pdf and {out}.png K*={int(K_star)}") + print( + f"Caption: (Left) Total training time $T_{{\\mathrm{{global}}}}$ vs. GPU count $K$ " + f"for {args.num_layers}-layer GNN trained for {args.num_epochs} epochs. " + "Solid blue: cost-model prediction. Red dashed: measured. " + "Dotted black: ideal linear speedup. " + f"$K^* = {int(K_star)}$ marks the last GPU count with $\\geq 90\\%$ scaling efficiency. " + "(Right) Scaling efficiency $= \\text{{speedup}} / K$." + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/visualization/plot_validation.py b/experiments/cost_model_benchmarks/visualization/plot_validation.py new file mode 100644 index 0000000..65a6497 --- /dev/null +++ b/experiments/cost_model_benchmarks/visualization/plot_validation.py @@ -0,0 +1,127 @@ +"""Visualization — Predicted vs. Measured Scatter (Headline Figure). + +Reads ``data/predictions.json`` and produces a predicted-vs-measured scatter +plot. Points are colored by graph type (or fit/held-out split). + +Usage:: + + python -m visualization.plot_validation \\ + --predictions data/predictions.json \\ + --color-by split \\ + --output figures/validation +""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +plt.rcParams.update({ + "font.size": 9, "axes.labelsize": 9, "legend.fontsize": 8, + "xtick.labelsize": 8, "ytick.labelsize": 8, + "figure.dpi": 300, "text.usetex": False, +}) + + +def parse_args(): + p = argparse.ArgumentParser(description="Plot predicted vs. measured scatter") + p.add_argument("--predictions", type=str, required=True) + p.add_argument("--color-by", choices=["split", "graph", "world_size"], + default="split") + p.add_argument("--output", type=str, default="figures/validation") + return p.parse_args() + + +def main(): + args = parse_args() + + with open(args.predictions) as f: + data = json.load(f) + + entries = data["predictions"] + mape_fit = data["aggregate"]["mape_fit_set"] + mape_held = data["aggregate"]["mape_held_out"] + + meas = np.array([e["measured_median_seconds"] * 1e3 for e in entries]) + pred = np.array([e["predicted_seconds"] * 1e3 for e in entries]) + + # Color groups + if args.color_by == "split": + groups = { + "Fit set": [i for i, e in enumerate(entries) if e["in_fit_set"]], + "Held-out": [i for i, e in enumerate(entries) if not e["in_fit_set"]], + } + palette = {"Fit set": "#1f77b4", "Held-out": "#d62728"} + elif args.color_by == "graph": + graph_types = sorted(set(e["config"].get("graph", "unknown") for e in entries)) + palette = dict(zip(graph_types, ["#1f77b4", "#ff7f0e", "#2ca02c", "#9467bd"])) + groups = {g: [i for i, e in enumerate(entries) + if e["config"].get("graph", "unknown") == g] + for g in graph_types} + else: # world_size + ws_vals = sorted(set(e["config"].get("world_size", 1) for e in entries)) + colors = plt.cm.viridis(np.linspace(0, 1, len(ws_vals))) + palette = {f"K={w}": c for w, c in zip(ws_vals, colors)} + groups = {f"K={w}": [i for i, e in enumerate(entries) + if e["config"].get("world_size", 1) == w] + for w in ws_vals} + + fig, ax = plt.subplots(figsize=(4.5, 4.5)) + + all_vals = np.concatenate([meas, pred]) + lo, hi = all_vals.min() * 0.9, all_vals.max() * 1.1 + ax.plot([lo, hi], [lo, hi], "k--", linewidth=0.8, label="Ideal (y=x)") + + # 10% error bands + ax.fill_between([lo, hi], [lo * 0.9, hi * 0.9], [lo * 1.1, hi * 1.1], + alpha=0.08, color="gray") + + for label, idxs in groups.items(): + if not idxs: + continue + color = palette[label] + ax.scatter(meas[idxs], pred[idxs], s=22, color=color, + alpha=0.85, label=label, zorder=3) + + ax.set_xlim(lo, hi) + ax.set_ylim(lo, hi) + ax.set_xlabel("Measured $T_{\\mathrm{layer}}$ (ms)") + ax.set_ylabel("Predicted $T_{\\mathrm{layer}}$ (ms)") + ax.set_title("Cost Model Validation: Predicted vs. Measured", fontsize=9) + + # Annotate MAPE + mape_text = ( + f"Fit MAPE = {mape_fit*100:.1f}%\n" + f"Held-out MAPE = {mape_held*100:.1f}%" + ) + ax.text(0.04, 0.96, mape_text, transform=ax.transAxes, + verticalalignment="top", fontsize=8, + bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8)) + + ax.legend(loc="lower right", fontsize=8) + ax.grid(True, linestyle=":", linewidth=0.4) + fig.tight_layout() + + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(str(out) + ".pdf", bbox_inches="tight") + fig.savefig(str(out) + ".png", bbox_inches="tight", dpi=300) + print(f"[plot_validation] Saved {out}.pdf and {out}.png") + print( + "Caption: Predicted vs. measured layer time $T_{\\mathrm{layer}}$ for all " + "benchmarked configurations. Dashed diagonal: perfect prediction. " + "Shaded band: $\\pm10\\%$ error region. " + f"In-sample MAPE = {mape_fit*100:.1f}\\%, " + f"held-out MAPE = {mape_held*100:.1f}\\%. " + "Colors distinguish " + + ("fit-set vs. held-out configurations." if args.color_by == "split" + else f"configurations by {args.color_by}.") + ) + + +if __name__ == "__main__": + main() From 060db3ec265d0c7f030756366b9684f254245b4f Mon Sep 17 00:00:00 2001 From: Shehtab Date: Mon, 13 Apr 2026 14:44:41 -0400 Subject: [PATCH 04/39] Update the gather fit to use non-linear fit --- .../analysis/fit_primitives.py | 96 ++++++++++++++----- 1 file changed, 71 insertions(+), 25 deletions(-) diff --git a/experiments/cost_model_benchmarks/analysis/fit_primitives.py b/experiments/cost_model_benchmarks/analysis/fit_primitives.py index 02c9233..a18ec67 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_primitives.py +++ b/experiments/cost_model_benchmarks/analysis/fit_primitives.py @@ -30,12 +30,14 @@ import numpy as np from scipy import stats as sp_stats +from scipy.optimize import curve_fit # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- + def load_json_files(paths: list) -> list: records = [] for p in paths: @@ -70,6 +72,7 @@ def linear_fit(x: np.ndarray, y: np.ndarray): # Network fit: T = t_L + bytes / B # --------------------------------------------------------------------------- + def fit_network(records: list) -> dict: """Fit (t_L, B) from ping-pong records (one mode per call).""" bytes_arr = [] @@ -103,6 +106,7 @@ def fit_network(records: list) -> dict: # Compute fit: T = coeff_V * |V| + coeff_E * |E| + intercept # --------------------------------------------------------------------------- + def fit_compute(records: list, timing_key: str = "forward_trials_seconds") -> dict: """Fit compute cost as a function of |V| and |E|. @@ -136,16 +140,19 @@ def fit_compute(records: list, timing_key: str = "forward_trials_seconds") -> di # --------------------------------------------------------------------------- -# Gather fit: T = intercept + bytes / B_gather +# Gather fit: T = intercept + max(overhead, overhead + bytes / B_gather) # --------------------------------------------------------------------------- + def fit_gather(records: list, timing_key: str = "gather_trials_seconds") -> dict: k_arr, T_arr, F_arr = [], [], [] for rec in records: F = rec["config"]["feature_dim"] for meas in rec["measurements"]: k = meas["params"]["k"] + t_med = median_of_trials(meas[timing_key]) + k_arr.append(k) T_arr.append(t_med) F_arr.append(F) @@ -153,14 +160,42 @@ def fit_gather(records: list, timing_key: str = "gather_trials_seconds") -> dict k_arr = np.array(k_arr, dtype=float) F_arr = np.array(F_arr, dtype=float) T_arr = np.array(T_arr, dtype=float) - bytes_arr = k_arr * F_arr * 4.0 # float32 - fit = linear_fit(bytes_arr, T_arr) - bandwidth = 1.0 / fit["slope"] if fit["slope"] > 0 else float("nan") + + # 1. Define the piecewise model for curve_fit + def time_model(b, overhead, inv_bandwidth): + # Time is bounded by a constant overhead until bandwidth saturation is reached + return np.maximum(overhead, b * inv_bandwidth) + + # 2. Provide sensible initial guesses (p0) to help the optimizer + min_T = np.min(T_arr) + slope_guess = (np.max(T_arr) - min_T) / (np.max(bytes_arr) + 1e-9) + p0 = [min_T, slope_guess] + + # 3. Fit the curve + try: + # Bounds ensure overhead and time-per-byte are strictly positive + popt, _ = curve_fit(time_model, bytes_arr, T_arr, p0=p0, bounds=(0, np.inf)) + launch_overhead = popt[0] + inv_bandwidth = popt[1] + except Exception: + raise RuntimeError("Curve fit failed for gather primitive") + + bandwidth = 1.0 / inv_bandwidth if inv_bandwidth > 0 else float("nan") + + # 4. Calculate R-squared manually (curve_fit doesn't return it) + if not np.isnan(launch_overhead): + T_pred = time_model(bytes_arr, launch_overhead, inv_bandwidth) + ss_res = np.sum((T_arr - T_pred) ** 2) + ss_tot = np.sum((T_arr - np.mean(T_arr)) ** 2) + r_squared = 1 - (ss_res / ss_tot) if ss_tot > 0 else float("nan") + else: + r_squared = float("nan") + return { "bandwidth_bytes_per_sec": bandwidth, - "intercept_seconds": fit["intercept"], - "r_squared": fit["r_squared"], + "intercept_seconds": launch_overhead, + "r_squared": r_squared, "_num_points": len(T_arr), } @@ -169,6 +204,7 @@ def fit_gather(records: list, timing_key: str = "gather_trials_seconds") -> dict # Main # --------------------------------------------------------------------------- + def parse_args(): p = argparse.ArgumentParser(description="Fit cost-model primitive parameters") p.add_argument("--pingpong-intra", nargs="+", default=[], metavar="FILE") @@ -191,15 +227,19 @@ def main(): if args.pingpong_intra: recs = load_json_files(args.pingpong_intra) net["intra"] = fit_network(recs) - print(f"[network/intra] B={net['intra']['bandwidth_bytes_per_sec']/1e9:.2f} GB/s " - f"t_L={net['intra']['latency_seconds']*1e6:.2f} µs " - f"R²={net['intra']['r_squared']:.4f}") + print( + f"[network/intra] B={net['intra']['bandwidth_bytes_per_sec']/1e9:.2f} GB/s " + f"t_L={net['intra']['latency_seconds']*1e6:.2f} µs " + f"R²={net['intra']['r_squared']:.4f}" + ) if args.pingpong_inter: recs = load_json_files(args.pingpong_inter) net["inter"] = fit_network(recs) - print(f"[network/inter] B={net['inter']['bandwidth_bytes_per_sec']/1e9:.2f} GB/s " - f"t_L={net['inter']['latency_seconds']*1e6:.2f} µs " - f"R²={net['inter']['r_squared']:.4f}") + print( + f"[network/inter] B={net['inter']['bandwidth_bytes_per_sec']/1e9:.2f} GB/s " + f"t_L={net['inter']['latency_seconds']*1e6:.2f} µs " + f"R²={net['inter']['r_squared']:.4f}" + ) result["network"] = net # Compute @@ -210,37 +250,43 @@ def main(): "forward": fit_compute(recs, "forward_trials_seconds"), "backward": fit_compute(recs, "backward_trials_seconds"), } - print(f"[compute/gcn] coeff_V={comp['gcn']['forward']['coeff_V']:.3e} " - f"coeff_E={comp['gcn']['forward']['coeff_E']:.3e} " - f"R²={comp['gcn']['forward']['r_squared']:.4f}") + print( + f"[compute/gcn] coeff_V={comp['gcn']['forward']['coeff_V']:.3e} " + f"coeff_E={comp['gcn']['forward']['coeff_E']:.3e} " + f"R²={comp['gcn']['forward']['r_squared']:.4f}" + ) if args.compute_edge: recs = load_json_files(args.compute_edge) comp["edge"] = { "forward": fit_compute(recs, "forward_trials_seconds"), "backward": fit_compute(recs, "backward_trials_seconds"), } - print(f"[compute/edge] coeff_V={comp['edge']['forward']['coeff_V']:.3e} " - f"coeff_E={comp['edge']['forward']['coeff_E']:.3e} " - f"R²={comp['edge']['forward']['r_squared']:.4f}") + print( + f"[compute/edge] coeff_V={comp['edge']['forward']['coeff_V']:.3e} " + f"coeff_E={comp['edge']['forward']['coeff_E']:.3e} " + f"R²={comp['edge']['forward']['r_squared']:.4f}" + ) result["compute"] = comp # Gather gath = {} for dist_name, files_attr in [ ("contiguous", "gather_contiguous"), - ("clustered", "gather_clustered"), - ("random", "gather_random"), + ("clustered", "gather_clustered"), + ("random", "gather_random"), ]: files = getattr(args, files_attr.replace("-", "_")) if files: recs = load_json_files(files) gath[dist_name] = { - "gather": fit_gather(recs, "gather_trials_seconds"), - "scatter_add": fit_gather(recs, "scatter_add_trials_seconds"), + "gather": fit_gather(recs, "gather_trials_seconds"), + "scatter_add": fit_gather(recs, "scatter_add_trials_seconds"), } - print(f"[gather/{dist_name}] " - f"B_gather={gath[dist_name]['gather']['bandwidth_bytes_per_sec']/1e9:.2f} GB/s " - f"R²={gath[dist_name]['gather']['r_squared']:.4f}") + print( + f"[gather/{dist_name}] " + f"B_gather={gath[dist_name]['gather']['bandwidth_bytes_per_sec']/1e9:.2f} GB/s " + f"R²={gath[dist_name]['gather']['r_squared']:.4f}" + ) result["gather"] = gath # Write From ecddcd60be3a04bcae0606817bf9391db580a1b2 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Fri, 17 Apr 2026 13:24:13 -0400 Subject: [PATCH 05/39] Decompose common benchmark funcs --- .../analysis/fit_primitives.py | 96 ++++-- .../benchmarks/bench_end_to_end.py | 300 +++--------------- .../benchmarks/graph_data_common.py | 210 ++++++++++++ .../benchmarks/nn_layer_common.py | 41 +++ .../visualization/plot_gather.py | 151 +++++++-- 5 files changed, 494 insertions(+), 304 deletions(-) create mode 100644 experiments/cost_model_benchmarks/benchmarks/graph_data_common.py create mode 100644 experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py diff --git a/experiments/cost_model_benchmarks/analysis/fit_primitives.py b/experiments/cost_model_benchmarks/analysis/fit_primitives.py index a18ec67..6211f94 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_primitives.py +++ b/experiments/cost_model_benchmarks/analysis/fit_primitives.py @@ -31,6 +31,7 @@ import numpy as np from scipy import stats as sp_stats from scipy.optimize import curve_fit +from scipy.special import expit # --------------------------------------------------------------------------- @@ -140,7 +141,7 @@ def fit_compute(records: list, timing_key: str = "forward_trials_seconds") -> di # --------------------------------------------------------------------------- -# Gather fit: T = intercept + max(overhead, overhead + bytes / B_gather) +# Gather fit: T = intercept + max(overhead, bytes / B_gather) # --------------------------------------------------------------------------- @@ -150,9 +151,7 @@ def fit_gather(records: list, timing_key: str = "gather_trials_seconds") -> dict F = rec["config"]["feature_dim"] for meas in rec["measurements"]: k = meas["params"]["k"] - t_med = median_of_trials(meas[timing_key]) - k_arr.append(k) T_arr.append(t_med) F_arr.append(F) @@ -160,43 +159,86 @@ def fit_gather(records: list, timing_key: str = "gather_trials_seconds") -> dict k_arr = np.array(k_arr, dtype=float) F_arr = np.array(F_arr, dtype=float) T_arr = np.array(T_arr, dtype=float) + bytes_arr = k_arr * F_arr * 4.0 # float32 - # 1. Define the piecewise model for curve_fit - def time_model(b, overhead, inv_bandwidth): - # Time is bounded by a constant overhead until bandwidth saturation is reached - return np.maximum(overhead, b * inv_bandwidth) + # 1. The Piecewise Linear Model (No Logarithms) + def time_model(b, overhead, inv_bw_L2, inv_bw_HBM, L2_thresh, HBM_thresh): + # 1. Bucket the bytes into their respective physical regimes + # Bytes processed exclusively at L2 speeds + bytes_L2 = np.clip(b, 0, L2_thresh) + # Bytes processed exclusively at HBM speeds + bytes_HBM = np.maximum(0, b - HBM_thresh) + + # 2. Apply the specific bandwidth (slope) to each bucket + t_mem = (bytes_L2 * inv_bw_L2) + (bytes_HBM * inv_bw_HBM) - # 2. Provide sensible initial guesses (p0) to help the optimizer - min_T = np.min(T_arr) - slope_guess = (np.max(T_arr) - min_T) / (np.max(bytes_arr) + 1e-9) - p0 = [min_T, slope_guess] + # 3. Floor the total time by the kernel launch overhead + return np.maximum(overhead, t_mem) + + # 2. Strategic initial guesses + min_T, max_T = float(np.min(T_arr)), float(np.max(T_arr)) + max_b = float(np.max(bytes_arr)) + + # The asymptotic slope is the difference in max/min time over max bytes + inv_bw_HBM_guess = (max_T - min_T) / (max_b + 1e-9) + + p0 = [ + min_T, # overhead + inv_bw_HBM_guess * 0.3, # inv_bw_L2 + inv_bw_HBM_guess, # inv_bw_HBM + max_b * 0.00001, # L2_thresh + max_b * 0.001, # HBM_thresh + ] - # 3. Fit the curve try: - # Bounds ensure overhead and time-per-byte are strictly positive - popt, _ = curve_fit(time_model, bytes_arr, T_arr, p0=p0, bounds=(0, np.inf)) - launch_overhead = popt[0] - inv_bandwidth = popt[1] - except Exception: - raise RuntimeError("Curve fit failed for gather primitive") - - bandwidth = 1.0 / inv_bandwidth if inv_bandwidth > 0 else float("nan") - - # 4. Calculate R-squared manually (curve_fit doesn't return it) - if not np.isnan(launch_overhead): - T_pred = time_model(bytes_arr, launch_overhead, inv_bandwidth) + bounds = ([0, 0, 0, 0, 0], [np.inf, np.inf, np.inf, max_b, max_b]) + + # 3. The Magic Fix: sigma=T_arr weights the fit by relative error + popt, _ = curve_fit( + time_model, + bytes_arr, + T_arr, + p0=p0, + bounds=bounds, + method="trf", + sigma=T_arr, + absolute_sigma=False, + ) + overhead, inv_bw_L2, inv_bw_HBM, L2_thresh, HBM_thresh = popt + + except Exception as e: + print(f"Fit failed: {e}") + overhead = inv_bw_L2 = inv_bw_HBM = L2_thresh = np.nan + + bw_HBM = 1.0 / inv_bw_HBM if inv_bw_HBM > 0 else float("nan") + bw_L2 = 1.0 / inv_bw_L2 if inv_bw_L2 > 0 else float("nan") + + # 3. Calculate linear R-squared in linear space + if not np.isnan(overhead): + T_pred = time_model( + bytes_arr, overhead, inv_bw_L2, inv_bw_HBM, L2_thresh, HBM_thresh + ) ss_res = np.sum((T_arr - T_pred) ** 2) ss_tot = np.sum((T_arr - np.mean(T_arr)) ** 2) r_squared = 1 - (ss_res / ss_tot) if ss_tot > 0 else float("nan") else: r_squared = float("nan") + print( + f" Fitted gather: overhead={overhead*1e3:.3f} ms ", + f"BW_L2={bw_L2/1e9:.2f} GB/s BW_HBM={bw_HBM/1e9:.2f} GB/s " + f"L2_thresh={L2_thresh/1e6:.2f} MB HBM_thresh={HBM_thresh/1e6:.2f} MB " + f"R²={r_squared:.4f}", + ) + return { - "bandwidth_bytes_per_sec": bandwidth, - "intercept_seconds": launch_overhead, + "bandwidth_bytes_per_sec": bw_HBM, + "L2_bandwidth_bytes_per_sec": bw_L2, + "L2_inflection_bytes": L2_thresh, + "HBM_inflection_bytes": HBM_thresh, + "launch_overhead_seconds": overhead, "r_squared": r_squared, - "_num_points": len(T_arr), } diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py index 7da40eb..4cfe162 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py @@ -51,201 +51,14 @@ setup_distributed, write_result, ) - - -# =========================================================================== -# Synthetic graph generators -# =========================================================================== - -def gen_erdos_renyi(num_vertices: int, avg_degree: float, - rng: np.random.Generator) -> np.ndarray: - """Return edge array of shape [E, 2] (src, dst) for an Erdős-Rényi digraph.""" - num_edges = int(num_vertices * avg_degree) - src = rng.integers(0, num_vertices, size=num_edges) - dst = rng.integers(0, num_vertices, size=num_edges) - return np.stack([src, dst], axis=1) - - -def gen_sbm(num_vertices: int, avg_degree: float, inter_density: float, - rng: np.random.Generator) -> np.ndarray: - """Return edges for a Stochastic Block Model graph. - - Vertices are split into blocks of equal size (one per rank for convenience, - though the actual partitioning is a separate step). The ratio of - intra-block to inter-block edges is controlled by *inter_density*. - """ - world_size = dist.get_world_size() if dist.is_initialized() else 4 - block_size = num_vertices // world_size - edges = [] - target_edges = int(num_vertices * avg_degree) - - intra_edges = int(target_edges * (1.0 - inter_density)) - inter_edges = target_edges - intra_edges - - # Intra-block edges - for b in range(world_size): - start = b * block_size - end = start + block_size - n = intra_edges // world_size - s = rng.integers(start, end, size=n) - d = rng.integers(start, end, size=n) - edges.append(np.stack([s, d], axis=1)) - - # Inter-block edges - s = rng.integers(0, num_vertices, size=inter_edges) - d = rng.integers(0, num_vertices, size=inter_edges) - # Force cross-block by offsetting dst block - d_block = (s // block_size + 1 + rng.integers(0, world_size - 1, - size=inter_edges)) % world_size - d = d_block * block_size + rng.integers(0, block_size, size=inter_edges) - d = np.clip(d, 0, num_vertices - 1) - edges.append(np.stack([s, d], axis=1)) - - return np.concatenate(edges, axis=0) - - -# =========================================================================== -# Partitioners -# =========================================================================== - -def partition_random(num_vertices: int, world_size: int, - rng: np.random.Generator) -> np.ndarray: - return rng.integers(0, world_size, size=num_vertices).astype(np.int64) - - -def partition_balanced(num_vertices: int, world_size: int) -> np.ndarray: - return np.floor(np.arange(num_vertices) * world_size / num_vertices).astype(np.int64) - - -def partition_metis(num_vertices: int, world_size: int, - edges: np.ndarray) -> np.ndarray: - try: - import pymetis - except ImportError: - raise RuntimeError( - "pymetis is not installed. Install it with: pip install pymetis\n" - "Or use --partitioner random or --partitioner balanced." - ) - # Build adjacency list for pymetis - adj = [[] for _ in range(num_vertices)] - for s, d in edges: - adj[s].append(int(d)) - adj[d].append(int(s)) - _, membership = pymetis.part_graph(world_size, adjacency=adj) - return np.array(membership, dtype=np.int64) - - -# =========================================================================== -# Minimal halo-exchange infrastructure -# =========================================================================== - -def build_local_comm_pattern(edges: np.ndarray, assignment: np.ndarray, - rank: int, world_size: int): - """Compute the local communication pattern for this rank. - - Returns a dict with: - local_vertices — np.ndarray of vertex IDs owned by this rank - local_edge_index — torch.Tensor [2, E_local] with local vertex IDs - remapped so that 0..n_local-1 are owned vertices - and n_local..n_local+n_halo-1 are halo vertices - send_counts — list[int] of length world_size: vertices to send - recv_counts — list[int] of length world_size: vertices to recv - send_idx — local indices (into local_vertices) to send per rank - halo_global_ids — global vertex IDs of halo vertices, in recv order - intra_halo_size — halo vertices from same node (ranks sharing node) - inter_halo_size — halo vertices from remote nodes - ranks_per_node — int (derived from LOCAL_RANK / RANK relationship) - """ - local_mask = assignment == rank - local_vertices = np.where(local_mask)[0] - n_local = len(local_vertices) - - # Global -> local index map - g2l = {int(v): i for i, v in enumerate(local_vertices)} - - # Find edges where dst is local - local_dst_mask = np.isin(edges[:, 1], local_vertices) - local_edges = edges[local_dst_mask] - - # Halo: src vertices not owned by this rank - halo_src_mask = ~np.isin(local_edges[:, 0], local_vertices) - halo_global = np.unique(local_edges[halo_src_mask, 0]) - - # Group halo vertices by owning rank - halo_owners = assignment[halo_global] - recv_by_rank = [] - halo_order = [] - for r in range(world_size): - verts = halo_global[halo_owners == r] - recv_by_rank.append(verts) - halo_order.extend(verts.tolist()) - halo_order = np.array(halo_order, dtype=np.int64) - - # Global halo id -> local halo index - halo_g2l = {int(v): n_local + i for i, v in enumerate(halo_order)} - all_g2l = {**g2l, **halo_g2l} - - # Find which local vertices other ranks need (send pattern) - # We exchange recv_counts via all_to_all to learn send_counts - recv_counts = [len(rv) for rv in recv_by_rank] - - # Build send: for each rank r, which of our local vertices does r need? - # We do a global exchange of halo_global per rank - all_recv = [None] * world_size - dist.all_gather_object(all_recv, halo_order.tolist()) - - send_idx_by_rank = [] - for r in range(world_size): - needed = np.array(all_recv[r], dtype=np.int64) - owned_mask = assignment[needed] == rank if len(needed) > 0 else np.array([], dtype=bool) - owned = needed[owned_mask] if len(needed) > 0 else np.array([], dtype=np.int64) - # Map to local indices - local_idxs = np.array([g2l[int(v)] for v in owned], dtype=np.int64) - send_idx_by_rank.append(local_idxs) - - send_counts = [len(s) for s in send_idx_by_rank] - - # Remap edges to local indices - valid_edge_mask = np.array([ - (int(s) in all_g2l) and (int(d) in all_g2l) - for s, d in local_edges - ]) - local_edges_valid = local_edges[valid_edge_mask] - if len(local_edges_valid) > 0: - remapped_src = np.array([all_g2l[int(s)] for s in local_edges_valid[:, 0]]) - remapped_dst = np.array([all_g2l[int(d)] for d in local_edges_valid[:, 1]]) - edge_index = torch.tensor( - np.stack([remapped_src, remapped_dst], axis=0), dtype=torch.long - ) - else: - edge_index = torch.zeros((2, 0), dtype=torch.long) - - # Compute intra / inter halo sizes - ranks_per_node = int(os.environ.get("LOCAL_WORLD_SIZE", - os.environ.get("SLURM_NTASKS_PER_NODE", "4"))) - my_node = rank // ranks_per_node - intra_halo_size = 0 - inter_halo_size = 0 - for r, verts in enumerate(recv_by_rank): - peer_node = r // ranks_per_node - if peer_node == my_node: - intra_halo_size += len(verts) - else: - inter_halo_size += len(verts) - - return { - "local_vertices": local_vertices, - "n_local": n_local, - "n_halo": len(halo_order), - "edge_index": edge_index, - "send_counts": send_counts, - "recv_counts": recv_counts, - "send_idx_by_rank": send_idx_by_rank, - "halo_order": halo_order, - "intra_halo_size": intra_halo_size, - "inter_halo_size": inter_halo_size, - "ranks_per_node": ranks_per_node, - } +from benchmarks.graph_data_common import ( + gen_erdos_renyi, + gen_sbm, + partition_balanced, + partition_metis, + partition_random, +) +from benchmarks.nn_layer_common import GCNLayer, EdgeConditionedLayer class MinimalHaloExchange(torch.autograd.Function): @@ -260,14 +73,20 @@ def forward(ctx, x_local, send_idx_flat, send_counts, recv_counts, world_size): # Split by destination rank send_list = list(send_buf.split(send_counts, dim=0)) - recv_list = [torch.zeros(rc, x_local.shape[1], - dtype=x_local.dtype, device=x_local.device) - for rc in recv_counts] + recv_list = [ + torch.zeros( + rc, x_local.shape[1], dtype=x_local.dtype, device=x_local.device + ) + for rc in recv_counts + ] dist.all_to_all(recv_list, send_list) - recv_buf = torch.cat(recv_list, dim=0) if sum(recv_counts) > 0 else \ - torch.zeros(0, x_local.shape[1], device=x_local.device) + recv_buf = ( + torch.cat(recv_list, dim=0) + if sum(recv_counts) > 0 + else torch.zeros(0, x_local.shape[1], device=x_local.device) + ) ctx.save_for_backward(send_idx_flat) ctx.send_counts = send_counts @@ -281,7 +100,7 @@ def forward(ctx, x_local, send_idx_flat, send_counts, recv_counts, world_size): @staticmethod def backward(ctx, grad_recv): - send_idx_flat, = ctx.saved_tensors + (send_idx_flat,) = ctx.saved_tensors send_counts = ctx.send_counts recv_counts = ctx.recv_counts world_size = ctx.world_size @@ -290,8 +109,11 @@ def backward(ctx, grad_recv): device = ctx.device # Reverse: recv_counts become send_counts and vice versa - grad_recv_list = list(grad_recv.split(recv_counts, dim=0)) \ - if grad_recv.shape[0] > 0 else [torch.zeros(0, F, device=device)] * world_size + grad_recv_list = ( + list(grad_recv.split(recv_counts, dim=0)) + if grad_recv.shape[0] > 0 + else [torch.zeros(0, F, device=device)] * world_size + ) grad_send_list = [torch.zeros(sc, F, device=device) for sc in send_counts] dist.all_to_all(grad_send_list, grad_recv_list) @@ -308,59 +130,27 @@ def backward(ctx, grad_recv): return grad_x_local, None, None, None, None -# =========================================================================== -# GNN layers (same as bench_compute.py for consistency) -# =========================================================================== - -class GCNLayer(nn.Module): - def __init__(self, feature_dim: int): - super().__init__() - self.linear = nn.Linear(feature_dim, feature_dim, bias=False) - - def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor: - src, dst = edge_index[0], edge_index[1] - n_local = x.shape[0] - # Only update local vertices (dst < n_local guard not needed since - # edge_index already restricts to local dst) - msg = self.linear(x[src]) - out = torch.zeros_like(x) - out.scatter_add_(0, dst.unsqueeze(1).expand_as(msg), msg) - return out - - -class EdgeConditionedLayer(nn.Module): - def __init__(self, feature_dim: int): - super().__init__() - self.mlp = nn.Sequential( - nn.Linear(3 * feature_dim, feature_dim), - nn.ReLU(), - nn.Linear(feature_dim, feature_dim), - ) - - def forward(self, x: torch.Tensor, edge_index: torch.Tensor, - edge_attr: torch.Tensor) -> torch.Tensor: - src, dst = edge_index[0], edge_index[1] - msg = self.mlp(torch.cat([x[src], x[dst], edge_attr], dim=-1)) - out = torch.zeros_like(x) - out.scatter_add_(0, dst.unsqueeze(1).expand_as(msg), msg) - return out - - # =========================================================================== # Main # =========================================================================== + def parse_args(): p = argparse.ArgumentParser(description="End-to-end halo exchange benchmark") p.add_argument("--graph", choices=["erdos_renyi", "sbm"], default="erdos_renyi") p.add_argument("--num-vertices", type=int, default=100_000) p.add_argument("--avg-degree", type=float, default=20.0) - p.add_argument("--sbm-inter-density", type=float, default=0.1, - help="Fraction of inter-block edges for SBM graphs") + p.add_argument( + "--sbm-inter-density", + type=float, + default=0.1, + help="Fraction of inter-block edges for SBM graphs", + ) p.add_argument("--feature-dim", type=int, default=128) p.add_argument("--model", choices=["gcn", "edge"], default="gcn") - p.add_argument("--partitioner", choices=["random", "balanced", "metis"], - default="balanced") + p.add_argument( + "--partitioner", choices=["random", "balanced", "metis"], default="balanced" + ) p.add_argument("--warmup", type=int, default=10) p.add_argument("--trials", type=int, default=50) p.add_argument("--output", type=str, required=True) @@ -380,8 +170,7 @@ def main(): if args.graph == "erdos_renyi": edges = gen_erdos_renyi(args.num_vertices, args.avg_degree, rng) else: - edges = gen_sbm(args.num_vertices, args.avg_degree, - args.sbm_inter_density, rng) + edges = gen_sbm(args.num_vertices, args.avg_degree, args.sbm_inter_density, rng) # --- Partition --- rng_part = np.random.default_rng(args.seed + 1) @@ -400,9 +189,13 @@ def main(): send_counts = pattern["send_counts"] recv_counts = pattern["recv_counts"] - send_idx_flat = torch.cat([ - torch.tensor(s, dtype=torch.long) for s in pattern["send_idx_by_rank"] - ]).to(device) if sum(send_counts) > 0 else torch.zeros(0, dtype=torch.long, device=device) + send_idx_flat = ( + torch.cat( + [torch.tensor(s, dtype=torch.long) for s in pattern["send_idx_by_rank"]] + ).to(device) + if sum(send_counts) > 0 + else torch.zeros(0, dtype=torch.long, device=device) + ) # --- Model --- if args.model == "gcn": @@ -413,8 +206,11 @@ def main(): # --- Synthetic local node features --- x_local = torch.randn(n_local, F, device=device, requires_grad=True) - edge_attr = torch.randn(edge_index.shape[1], F, device=device) \ - if args.model == "edge" else None + edge_attr = ( + torch.randn(edge_index.shape[1], F, device=device) + if args.model == "edge" + else None + ) # --- Timed forward + backward --- def one_layer(): diff --git a/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py b/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py new file mode 100644 index 0000000..15358e5 --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py @@ -0,0 +1,210 @@ +import numpy as np +import torch.distributed as dist + + +# =========================================================================== +# Synthetic graph generators +# =========================================================================== + + +def gen_erdos_renyi( + num_vertices: int, avg_degree: float, rng: np.random.Generator +) -> np.ndarray: + """Return edge array of shape [E, 2] (src, dst) for an Erdős-Rényi digraph.""" + num_edges = int(num_vertices * avg_degree) + src = rng.integers(0, num_vertices, size=num_edges) + dst = rng.integers(0, num_vertices, size=num_edges) + return np.stack([src, dst], axis=1) + + +def gen_sbm( + num_vertices: int, avg_degree: float, inter_density: float, rng: np.random.Generator +) -> np.ndarray: + """Return edges for a Stochastic Block Model graph. + + Vertices are split into blocks of equal size (one per rank for convenience, + though the actual partitioning is a separate step). The ratio of + intra-block to inter-block edges is controlled by *inter_density*. + """ + world_size = dist.get_world_size() if dist.is_initialized() else 4 + block_size = num_vertices // world_size + edges = [] + target_edges = int(num_vertices * avg_degree) + + intra_edges = int(target_edges * (1.0 - inter_density)) + inter_edges = target_edges - intra_edges + + # Intra-block edges + for b in range(world_size): + start = b * block_size + end = start + block_size + n = intra_edges // world_size + s = rng.integers(start, end, size=n) + d = rng.integers(start, end, size=n) + edges.append(np.stack([s, d], axis=1)) + + # Inter-block edges + s = rng.integers(0, num_vertices, size=inter_edges) + d = rng.integers(0, num_vertices, size=inter_edges) + # Force cross-block by offsetting dst block + d_block = ( + s // block_size + 1 + rng.integers(0, world_size - 1, size=inter_edges) + ) % world_size + d = d_block * block_size + rng.integers(0, block_size, size=inter_edges) + d = np.clip(d, 0, num_vertices - 1) + edges.append(np.stack([s, d], axis=1)) + + return np.concatenate(edges, axis=0) + + +# =========================================================================== +# Partitioners +# =========================================================================== + + +def partition_random( + num_vertices: int, world_size: int, rng: np.random.Generator +) -> np.ndarray: + return rng.integers(0, world_size, size=num_vertices).astype(np.int64) + + +def partition_balanced(num_vertices: int, world_size: int) -> np.ndarray: + return np.floor(np.arange(num_vertices) * world_size / num_vertices).astype( + np.int64 + ) + + +def partition_metis( + num_vertices: int, world_size: int, edges: np.ndarray +) -> np.ndarray: + try: + import pymetis + except ImportError: + raise RuntimeError( + "pymetis is not installed. Install it with: pip install pymetis\n" + "Or use --partitioner random or --partitioner balanced." + ) + # Build adjacency list for pymetis + adj = [[] for _ in range(num_vertices)] + for s, d in edges: + adj[s].append(int(d)) + adj[d].append(int(s)) + _, membership = pymetis.part_graph(world_size, adjacency=adj) + return np.array(membership, dtype=np.int64) + + +# =========================================================================== +# Minimal halo-exchange infrastructure +# =========================================================================== + + +def build_local_comm_pattern( + edges: np.ndarray, assignment: np.ndarray, rank: int, world_size: int +): + """Compute the local communication pattern for this rank. + + Returns a CommunicationPattern object with: + local_vertices — np.ndarray of vertex IDs owned by this rank + local_edge_index — torch.Tensor [2, E_local] with local vertex IDs + remapped so that 0..n_local-1 are owned vertices + and n_local..n_local+n_halo-1 are halo vertices + send_counts — list[int] of length world_size: vertices to send + recv_counts — list[int] of length world_size: vertices to recv + send_idx — local indices (into local_vertices) to send per rank + halo_global_ids — global vertex IDs of halo vertices, in recv order + intra_halo_size — halo vertices from same node (ranks sharing node) + inter_halo_size — halo vertices from remote nodes + ranks_per_node — int (derived from LOCAL_RANK / RANK relationship) + """ + local_mask = assignment == rank + local_vertices = np.where(local_mask)[0] + n_local = len(local_vertices) + + # Global -> local index map + g2l = {int(v): i for i, v in enumerate(local_vertices)} + + # Find edges where dst is local + local_dst_mask = np.isin(edges[:, 1], local_vertices) + local_edges = edges[local_dst_mask] + + # Halo: src vertices not owned by this rank + halo_src_mask = ~np.isin(local_edges[:, 0], local_vertices) + halo_global = np.unique(local_edges[halo_src_mask, 0]) + + # Group halo vertices by owning rank + halo_owners = assignment[halo_global] + recv_by_rank = [] + halo_order = [] + for r in range(world_size): + verts = halo_global[halo_owners == r] + recv_by_rank.append(verts) + halo_order.extend(verts.tolist()) + halo_order = np.array(halo_order, dtype=np.int64) + + # Global halo id -> local halo index + halo_g2l = {int(v): n_local + i for i, v in enumerate(halo_order)} + all_g2l = {**g2l, **halo_g2l} + + # Find which local vertices other ranks need (send pattern) + # We exchange recv_counts via all_to_all to learn send_counts + recv_counts = [len(rv) for rv in recv_by_rank] + + # Build send: for each rank r, which of our local vertices does r need? + # We do a global exchange of halo_global per rank + all_recv = [None] * world_size + dist.all_gather_object(all_recv, halo_order.tolist()) + + send_idx_by_rank = [] + for r in range(world_size): + needed = np.array(all_recv[r], dtype=np.int64) + owned_mask = ( + assignment[needed] == rank if len(needed) > 0 else np.array([], dtype=bool) + ) + owned = needed[owned_mask] if len(needed) > 0 else np.array([], dtype=np.int64) + # Map to local indices + local_idxs = np.array([g2l[int(v)] for v in owned], dtype=np.int64) + send_idx_by_rank.append(local_idxs) + + send_counts = [len(s) for s in send_idx_by_rank] + + # Remap edges to local indices + valid_edge_mask = np.array( + [(int(s) in all_g2l) and (int(d) in all_g2l) for s, d in local_edges] + ) + local_edges_valid = local_edges[valid_edge_mask] + if len(local_edges_valid) > 0: + remapped_src = np.array([all_g2l[int(s)] for s in local_edges_valid[:, 0]]) + remapped_dst = np.array([all_g2l[int(d)] for d in local_edges_valid[:, 1]]) + edge_index = torch.tensor( + np.stack([remapped_src, remapped_dst], axis=0), dtype=torch.long + ) + else: + edge_index = torch.zeros((2, 0), dtype=torch.long) + + # Compute intra / inter halo sizes + ranks_per_node = int( + os.environ.get("LOCAL_WORLD_SIZE", os.environ.get("SLURM_NTASKS_PER_NODE", "4")) + ) + my_node = rank // ranks_per_node + intra_halo_size = 0 + inter_halo_size = 0 + for r, verts in enumerate(recv_by_rank): + peer_node = r // ranks_per_node + if peer_node == my_node: + intra_halo_size += len(verts) + else: + inter_halo_size += len(verts) + + return { + "local_vertices": local_vertices, + "n_local": n_local, + "n_halo": len(halo_order), + "edge_index": edge_index, + "send_counts": send_counts, + "recv_counts": recv_counts, + "send_idx_by_rank": send_idx_by_rank, + "halo_order": halo_order, + "intra_halo_size": intra_halo_size, + "inter_halo_size": inter_halo_size, + "ranks_per_node": ranks_per_node, + } diff --git a/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py b/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py new file mode 100644 index 0000000..893d268 --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py @@ -0,0 +1,41 @@ +import torch +import torch.nn as nn + +# =========================================================================== +# GNN layers (same as bench_compute.py for consistency) +# =========================================================================== + + +class GCNLayer(nn.Module): + def __init__(self, feature_dim: int): + super().__init__() + self.linear = nn.Linear(feature_dim, feature_dim, bias=False) + + def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor: + src, dst = edge_index[0], edge_index[1] + n_local = x.shape[0] + # Only update local vertices (dst < n_local guard not needed since + # edge_index already restricts to local dst) + msg = self.linear(x[src]) + out = torch.zeros_like(x) + out.scatter_add_(0, dst.unsqueeze(1).expand_as(msg), msg) + return out + + +class EdgeConditionedLayer(nn.Module): + def __init__(self, feature_dim: int): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(3 * feature_dim, feature_dim), + nn.ReLU(), + nn.Linear(feature_dim, feature_dim), + ) + + def forward( + self, x: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor + ) -> torch.Tensor: + src, dst = edge_index[0], edge_index[1] + msg = self.mlp(torch.cat([x[src], x[dst], edge_attr], dim=-1)) + out = torch.zeros_like(x) + out.scatter_add_(0, dst.unsqueeze(1).expand_as(msg), msg) + return out diff --git a/experiments/cost_model_benchmarks/visualization/plot_gather.py b/experiments/cost_model_benchmarks/visualization/plot_gather.py index cddfdc4..f33b193 100644 --- a/experiments/cost_model_benchmarks/visualization/plot_gather.py +++ b/experiments/cost_model_benchmarks/visualization/plot_gather.py @@ -19,19 +19,68 @@ import numpy as np import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt COLORS = { "contiguous": "#1f77b4", - "clustered": "#2ca02c", - "random": "#d62728", + "clustered": "#2ca02c", + "random": "#d62728", } -plt.rcParams.update({ - "font.size": 9, "axes.labelsize": 9, "legend.fontsize": 8, - "xtick.labelsize": 8, "ytick.labelsize": 8, - "figure.dpi": 300, "text.usetex": False, -}) +plt.rcParams.update( + { + "font.size": 9, + "axes.labelsize": 9, + "legend.fontsize": 8, + "xtick.labelsize": 8, + "ytick.labelsize": 8, + "figure.dpi": 300, + "text.usetex": False, + } +) + + +def fitted_gather( + min_val, + max_val, + feature_dim, + hbm_bandwidth_bytes_per_sec, + l2_bandwidth_bytes_per_sec, + launch_overhead_seconds, + L2_thresh, + HBM_thresh, +): + """Return a function that models gather time as a function of k.""" + + def time_model(b, overhead, inv_bw_L2, inv_bw_HBM, L2_thresh, HBM_thresh): + # 1. Bucket the bytes into their respective physical regimes + # Bytes processed exclusively at L2 speeds + bytes_L2 = np.clip(b, 0, L2_thresh) + # Bytes processed exclusively at HBM speeds + bytes_HBM = np.maximum(0, b - HBM_thresh) + + # 2. Apply the specific bandwidth (slope) to each bucket + t_mem = (bytes_L2 * inv_bw_L2) + (bytes_HBM * inv_bw_HBM) + + # 3. Floor the total time by the kernel launch overhead + return np.maximum(overhead, t_mem) + + x = np.linspace(min_val, max_val, 100) + inv_bw_L2 = 1.0 / l2_bandwidth_bytes_per_sec + inv_bw_HBM = 1.0 / hbm_bandwidth_bytes_per_sec + y = ( + time_model( + x * feature_dim * 4.0, + launch_overhead_seconds, + inv_bw_L2, + inv_bw_HBM, + L2_thresh, + HBM_thresh, + ) + * 1e3 + ) + return x, y def load_gather_file(paths: list, timing_key: str): @@ -42,14 +91,16 @@ def load_gather_file(paths: list, timing_key: str): data = json.load(f) for meas in data["measurements"]: trials = np.array(meas[timing_key]) - rows.append(( - meas["params"]["k"], - float(np.median(trials)), - float(np.percentile(trials, 25)), - float(np.percentile(trials, 75)), - )) + rows.append( + ( + meas["params"]["k"], + float(np.median(trials)), + float(np.percentile(trials, 25)), + float(np.percentile(trials, 75)), + ) + ) rows.sort(key=lambda r: r[0]) - k_arr = np.array([r[0] for r in rows]) + k_arr = np.array([r[0] for r in rows]) med_arr = np.array([r[1] for r in rows]) q25_arr = np.array([r[2] for r in rows]) q75_arr = np.array([r[3] for r in rows]) @@ -59,42 +110,92 @@ def load_gather_file(paths: list, timing_key: str): def parse_args(): p = argparse.ArgumentParser(description="Plot gather/scatter-add bandwidth") p.add_argument("--contiguous", nargs="+", default=[], metavar="FILE") - p.add_argument("--clustered", nargs="+", default=[], metavar="FILE") - p.add_argument("--random", nargs="+", default=[], metavar="FILE") - p.add_argument("--operation", choices=["gather", "scatter_add"], default="gather") - p.add_argument("--output", type=str, default="figures/gather") + p.add_argument("--clustered", nargs="+", default=[], metavar="FILE") + p.add_argument("--random", nargs="+", default=[], metavar="FILE") + p.add_argument("--operation", choices=["gather", "scatter_add"], default="gather") + p.add_argument("--fitted", type=str, default=None, metavar="FILE") + p.add_argument("--output", type=str, default="figures/gather") return p.parse_args() def main(): args = parse_args() timing_key = ( - "gather_trials_seconds" if args.operation == "gather" + "gather_trials_seconds" + if args.operation == "gather" else "scatter_add_trials_seconds" ) fig, ax = plt.subplots(figsize=(5, 3.5)) + if args.fitted: + with open(args.fitted) as f: + primitives = json.load(f) + min_k, max_k = 1e3, 1e9 for dist_name, files in [ ("contiguous", args.contiguous), - ("clustered", args.clustered), - ("random", args.random), + ("clustered", args.clustered), + ("random", args.random), ]: if not files: continue k, med, q25, q75 = load_gather_file(files, timing_key) + min_k = k[0] + max_k = k[-1] color = COLORS[dist_name] ax.errorbar( - k * 1e-6, med * 1e3, + k * 1e-6, + med * 1e3, yerr=[(med - q25) * 1e3, (q75 - med) * 1e3], - fmt="o-", markersize=3, color=color, linewidth=0.9, - capsize=2, elinewidth=0.8, label=dist_name.capitalize(), + fmt="o", + markersize=3, + color=color, + linewidth=0.9, + capsize=2, + elinewidth=0.8, + label=dist_name.capitalize(), ) + if args.fitted: + hbm_bw = primitives["gather"][dist_name]["gather"][ + "bandwidth_bytes_per_sec" + ] + l2_bw = primitives["gather"][dist_name]["gather"][ + "L2_bandwidth_bytes_per_sec" + ] + overhead = primitives["gather"][dist_name]["gather"][ + "launch_overhead_seconds" + ] + thresh = primitives["gather"][dist_name]["gather"]["L2_inflection_bytes"] + hbm_thresh = primitives["gather"][dist_name]["gather"][ + "HBM_inflection_bytes" + ] + x, y = fitted_gather( + min_k, + max_k, + feature_dim=512, + hbm_bandwidth_bytes_per_sec=hbm_bw, + l2_bandwidth_bytes_per_sec=l2_bw, + launch_overhead_seconds=overhead, + L2_thresh=thresh, + HBM_thresh=hbm_thresh, + ) + ax.plot( + x * 1e-6, + y, + "--", + color=color, + linewidth=1.2, + label=f"{dist_name.capitalize()}-Expected", + alpha=0.4, + ) ax.set_xscale("log") ax.set_yscale("log") - op_label = "Gather $x[\\mathrm{idx}]$" if args.operation == "gather" \ + op_label = ( + "Gather $x[\\mathrm{idx}]$" + if args.operation == "gather" else "Scatter-add (backward)" + ) ax.set_xlabel("k (millions of rows gathered)") ax.set_ylabel(f"{op_label} time (ms)") ax.set_title(f"Buffer-Copy Bandwidth: {op_label}", fontsize=9) From a89e2db662bd485677463b595d165679bad66a84 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Sat, 18 Apr 2026 11:14:44 -0400 Subject: [PATCH 06/39] Add crossover benchmark to measure communication compute tradeoff --- .../benchmarks/bench_crossover.py | 533 ++++++++++++++++++ 1 file changed, 533 insertions(+) create mode 100644 experiments/cost_model_benchmarks/benchmarks/bench_crossover.py diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py b/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py new file mode 100644 index 0000000..d6bc829 --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py @@ -0,0 +1,533 @@ +"""Benchmark 2.2 — Multi-GPU Crossover Point. + +Sweeps graph size to identify the point at which distributing computation +across K GPUs overcomes the overhead of halo-exchange communication. + +For each graph size N in the sweep: + + * **Single-GPU baseline**: rank 0 runs forward+backward on the *complete* + graph on one GPU. Measures raw compute cost T_comp(N). + * **Multi-GPU distributed**: all K ranks execute partitioned + forward+backward with halo exchange. Measures + T_comp(N/K) + T_comm. + +The *crossover* N* is the smallest graph size where +T_single(N) > T_multi(K, N), i.e. where distributing first becomes +beneficial. + +When ``--no-dist`` is set the script runs the single-GPU sweep only +(useful for characterising baseline compute without a multi-GPU +allocation). + +Synthetic graphs: + * ``erdos_renyi`` — Erdős-Rényi with ``--avg-degree`` expected degree + * ``sbm`` — Stochastic Block Model; ``--sbm-inter-density`` + controls the fraction of inter-block edges + +Partitioners: + * ``random`` — assign each vertex to a uniformly random rank + * ``balanced`` — contiguous vertex blocks of equal size + * ``metis`` — balanced k-way via pymetis (skipped if not installed) + +Usage — distributed (torchrun):: + + torchrun --nnodes 2 --nproc_per_node 4 \\ + -m benchmarks.benchmark_crossover \\ + --graph erdos_renyi \\ + --graph-sizes 10000,100000,1000000,10000000 \\ + --avg-degree 20 --feature-dim 128 --model gcn \\ + --partitioner balanced --warmup 10 --trials 50 \\ + --output data/crossover_K8_F128_er_bal.json --seed 42 + +Usage — single-GPU baseline only:: + + python -m benchmarks.benchmark_crossover --no-dist \\ + --graph erdos_renyi \\ + --graph-sizes 10000,100000,1000000,10000000 \\ + --avg-degree 20 --feature-dim 128 --model gcn \\ + --warmup 10 --trials 50 \\ + --output data/crossover_single_F128.json --seed 42 +""" + +import argparse +import os + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn + +from benchmarks.common import ( + collect_metadata, + cuda_timed, + seed_everything, + setup_distributed, + write_result, +) +from benchmarks.graph_data_common import ( + gen_erdos_renyi, + gen_sbm, + partition_balanced, + partition_metis, + partition_random, +) +from benchmarks.nn_layer_common import GCNLayer, EdgeConditionedLayer + +from DGraph.distributed import ( + HaloExchange, + CommunicationPattern, + build_communication_pattern, +) +from DGraph import Communicator + + +# --------------------------------------------------------------------------- +# Argument parsing +# --------------------------------------------------------------------------- + + +def parse_args(): + p = argparse.ArgumentParser(description="Multi-GPU crossover point benchmark") + p.add_argument("--graph", choices=["erdos_renyi", "sbm"], default="erdos_renyi") + p.add_argument( + "--graph-sizes", + type=str, + default="10000,100000,1000000,10000000", + help="Comma-separated list of num_vertices values to sweep", + ) + p.add_argument("--avg-degree", type=float, default=20.0) + p.add_argument( + "--sbm-inter-density", + type=float, + default=0.1, + help="Fraction of inter-block edges for SBM graphs", + ) + p.add_argument("--feature-dim", type=int, default=128) + p.add_argument("--model", choices=["gcn", "edge"], default="gcn") + p.add_argument( + "--partitioner", choices=["random", "balanced", "metis"], default="balanced" + ) + p.add_argument("--warmup", type=int, default=10) + p.add_argument("--trials", type=int, default=50) + p.add_argument("--output", type=str, required=True) + p.add_argument("--seed", type=int, default=42) + p.add_argument( + "--no-dist", + action="store_true", + help="Run single-GPU sweep only (no distributed setup required)", + ) + return p.parse_args() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _gen_graph(graph_type, num_vertices, avg_degree, sbm_inter_density, seed): + """Generate a synthetic graph reproducibly for a given graph size.""" + rng = np.random.default_rng(seed) + if graph_type == "erdos_renyi": + return gen_erdos_renyi(num_vertices, avg_degree, rng) + else: + return gen_sbm(num_vertices, avg_degree, sbm_inter_density, rng) + + +def _intra_inter_halo( + comm_pattern: CommunicationPattern, ranks_per_node: int +) -> tuple: + """Return (intra_halo_vertices, inter_halo_vertices) from recv_offset.""" + rank = comm_pattern.rank + my_node = rank // ranks_per_node + recv_counts = ( + comm_pattern.recv_offset[1:] - comm_pattern.recv_offset[:-1] + ).tolist() + intra = 0 + inter = 0 + for r, count in enumerate(recv_counts): + if r == rank: + continue + if (r // ranks_per_node) == my_node: + intra += int(count) + else: + inter += int(count) + return intra, inter + + +def _build_single_gpu_tensors(num_vertices, edges_np, F, model, device): + """Allocate full-graph tensors on *device*. May raise cuda.OutOfMemoryError.""" + # GCNLayer / EdgeConditionedLayer expect edge_index as [2, E] + edge_t = torch.from_numpy(edges_np.T.copy()).long().to(device) + x = torch.randn(num_vertices, F, device=device, requires_grad=True) + if model == "gcn": + layer = GCNLayer(F).to(device) + edge_attr = None + else: + layer = EdgeConditionedLayer(F).to(device) + edge_attr = torch.randn(edges_np.shape[0], F, device=device) + layer.train() + return x, edge_t, layer, edge_attr + + +def _single_gpu_fn(x, edge_t, layer, edge_attr): + """Return a zero-argument closure for single-GPU forward+backward.""" + if edge_attr is None: + def fn(): + out = layer(x, edge_t) + out.sum().backward() + if x.grad is not None: + x.grad.zero_() + else: + def fn(): + out = layer(x, edge_t, edge_attr) + out.sum().backward() + if x.grad is not None: + x.grad.zero_() + return fn + + +def _multi_gpu_fn(x_local, comm_pattern, layer, halo_exchange, edge_attr, model): + """Return a zero-argument closure for distributed forward+backward. + + ``comm_pattern.local_edge_list`` has shape ``[E, 2]``; we transpose it + once to ``[2, E]`` as required by GCNLayer / EdgeConditionedLayer. + """ + edge_index = comm_pattern.local_edge_list.T.contiguous() # [2, E_local] + + if model == "gcn": + def fn(): + recv_buf = halo_exchange(x_local, comm_pattern) + x_aug = torch.cat([x_local, recv_buf], dim=0) + out = layer(x_aug, edge_index) + out.sum().backward() + if x_local.grad is not None: + x_local.grad.zero_() + else: + def fn(): + recv_buf = halo_exchange(x_local, comm_pattern) + x_aug = torch.cat([x_local, recv_buf], dim=0) + out = layer(x_aug, edge_index, edge_attr) + out.sum().backward() + if x_local.grad is not None: + x_local.grad.zero_() + return fn + + +# --------------------------------------------------------------------------- +# Single-GPU-only path +# --------------------------------------------------------------------------- + + +def run_no_dist(args, graph_sizes, F): + """Sweep graph sizes on a single GPU and return the measurement list.""" + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + seed_everything(args.seed) + measurements = [] + + for num_vertices in graph_sizes: + # Unique seed per graph size so each topology is reproducible independently + graph_seed = args.seed + num_vertices + edges_np = _gen_graph( + args.graph, num_vertices, args.avg_degree, args.sbm_inter_density, graph_seed + ) + num_edges = int(edges_np.shape[0]) + + times_single = None + oom = False + try: + x, edge_t, layer, edge_attr = _build_single_gpu_tensors( + num_vertices, edges_np, F, args.model, device + ) + fn = _single_gpu_fn(x, edge_t, layer, edge_attr) + times_single = cuda_timed(fn, warmup=args.warmup, trials=args.trials) + # Free before next iteration + del x, edge_t, layer + if edge_attr is not None: + del edge_attr + torch.cuda.empty_cache() + except torch.cuda.OutOfMemoryError: + oom = True + print( + f"[crossover] OOM: N={num_vertices:,} exceeds single-GPU memory; " + "skipping and continuing" + ) + torch.cuda.empty_cache() + + med_str = ( + f"{1e3 * sorted(times_single)[len(times_single) // 2]:.2f} ms" + if times_single + else "OOM" + ) + print(f"[crossover] no-dist N={num_vertices:>12,} single={med_str}") + + measurements.append( + { + "params": { + "num_vertices": num_vertices, + "num_edges": num_edges, + "avg_degree": args.avg_degree, + "feature_dim": F, + "model": args.model, + "graph": args.graph, + }, + "single_gpu_trials_seconds": times_single, + "single_gpu_oom": oom, + "multi_gpu_trials_seconds_rank0": None, + "multi_gpu_trials_seconds_max": None, + "per_rank_stats": None, + "world_size": 1, + } + ) + + return measurements + + +# --------------------------------------------------------------------------- +# Distributed path +# --------------------------------------------------------------------------- + + +def run_distributed(args, graph_sizes, F, rank, world_size, local_rank): + """Run one crossover measurement per graph size using all K ranks.""" + device = torch.device(f"cuda:{local_rank}") + + comm = Communicator(backend="nccl") + halo_exchange = HaloExchange(comm=comm) + + ranks_per_node = int( + os.environ.get( + "LOCAL_WORLD_SIZE", os.environ.get("SLURM_NTASKS_PER_NODE", "4") + ) + ) + + measurements = [] + + for num_vertices in graph_sizes: + # All ranks generate *identical* graph topology (same seed per size) + graph_seed = args.seed + num_vertices + edges_np = _gen_graph( + args.graph, num_vertices, args.avg_degree, args.sbm_inter_density, graph_seed + ) + + # All ranks compute *identical* partition assignment + rng_part = np.random.default_rng(args.seed + 1 + num_vertices) + if args.partitioner == "random": + assignment_np = partition_random(num_vertices, world_size, rng_part) + elif args.partitioner == "balanced": + assignment_np = partition_balanced(num_vertices, world_size) + else: + assignment_np = partition_metis(num_vertices, world_size, edges_np) + + # Move to GPU for DGraph's build_communication_pattern (collective) + edges_t = torch.from_numpy(edges_np).long().to(device) # [E, 2] + partitioning_t = torch.from_numpy(assignment_np).long().to(device) # [V] + + # Collective: all ranks call this in sync (internally calls dist.all_gather) + comm_pattern = build_communication_pattern( + edges_t, partitioning_t, rank, world_size + ) + + n_local = comm_pattern.num_local_vertices + n_halo = comm_pattern.num_halo_vertices + n_local_edges = comm_pattern.local_edge_list.shape[0] + + x_local = torch.randn(n_local, F, device=device, requires_grad=True) + edge_attr_dist = ( + torch.randn(n_local_edges, F, device=device) + if args.model == "edge" + else None + ) + layer_dist = ( + GCNLayer(F).to(device) + if args.model == "gcn" + else EdgeConditionedLayer(F).to(device) + ) + layer_dist.train() + + fn_multi = _multi_gpu_fn( + x_local, comm_pattern, layer_dist, halo_exchange, edge_attr_dist, args.model + ) + + # ---- Time multi-GPU (all ranks participate) ---- + dist.barrier() + times_multi_local = cuda_timed(fn_multi, warmup=args.warmup, trials=args.trials) + dist.barrier() + + # Gather timings and partition stats from all ranks to rank 0 + all_times_multi = [None] * world_size + dist.all_gather_object(all_times_multi, times_multi_local) + + intra_halo, inter_halo = _intra_inter_halo(comm_pattern, ranks_per_node) + stats_local = { + "rank": rank, + "n_local": n_local, + "n_halo": n_halo, + "n_local_edges": n_local_edges, + "intra_halo_size": intra_halo, + "inter_halo_size": inter_halo, + "c_intra_bytes": intra_halo * F * 4, + "c_inter_bytes": inter_halo * F * 4, + "send_total": int(comm_pattern.send_offset[-1].item()), + "recv_total": int(comm_pattern.recv_offset[-1].item()), + "trials_seconds": times_multi_local, + } + all_stats = [None] * world_size + dist.all_gather_object(all_stats, stats_local) + + # Free GPU memory before the single-GPU run + del x_local, edges_t, partitioning_t, layer_dist + if edge_attr_dist is not None: + del edge_attr_dist + torch.cuda.empty_cache() + + # ---- Rank 0: time single-GPU baseline (non-collective) ---- + times_single = None + single_oom = False + if rank == 0: + edges_s = _gen_graph( + args.graph, + num_vertices, + args.avg_degree, + args.sbm_inter_density, + graph_seed, + ) + try: + x_s, edge_t_s, layer_s, edge_attr_s = _build_single_gpu_tensors( + num_vertices, edges_s, F, args.model, device + ) + fn_single = _single_gpu_fn(x_s, edge_t_s, layer_s, edge_attr_s) + times_single = cuda_timed( + fn_single, warmup=args.warmup, trials=args.trials + ) + del x_s, edge_t_s, layer_s + if edge_attr_s is not None: + del edge_attr_s + torch.cuda.empty_cache() + except torch.cuda.OutOfMemoryError: + single_oom = True + print( + f"[crossover] rank 0 OOM: N={num_vertices:,} exceeds single-GPU memory" + ) + torch.cuda.empty_cache() + + # Sync all ranks after rank 0 finishes its single-GPU benchmark + dist.barrier() + + if rank == 0: + n_trials = len(times_multi_local) + # Wall time = max latency across all ranks per trial + times_multi_max = [ + max(all_times_multi[r][i] for r in range(world_size)) + for i in range(n_trials) + ] + + med_multi = sorted(times_multi_max)[n_trials // 2] + if times_single: + med_single = sorted(times_single)[len(times_single) // 2] + speedup = med_single / med_multi if med_multi > 0 else float("nan") + else: + med_single = float("nan") + speedup = float("nan") + + print( + f"[crossover] K={world_size:>3} N={num_vertices:>12,} " + f"single={1e3*med_single:.2f}ms " + f"multi={1e3*med_multi:.2f}ms " + f"speedup={speedup:.2f}x" + ) + + measurements.append( + { + "params": { + "num_vertices": num_vertices, + "num_global_edges": int(edges_np.shape[0]), + "avg_degree": args.avg_degree, + "feature_dim": F, + "model": args.model, + "graph": args.graph, + "partitioner": args.partitioner, + "world_size": world_size, + "ranks_per_node": ranks_per_node, + }, + "single_gpu_trials_seconds": times_single, + "single_gpu_oom": single_oom, + "multi_gpu_trials_seconds_rank0": times_multi_local, + "multi_gpu_trials_seconds_max": times_multi_max, + "per_rank_stats": all_stats, + } + ) + + return measurements, ranks_per_node + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def main(): + args = parse_args() + graph_sizes = [int(s) for s in args.graph_sizes.split(",")] + F = args.feature_dim + + if args.no_dist: + measurements = run_no_dist(args, graph_sizes, F) + payload = { + "benchmark": "crossover", + "metadata": collect_metadata(), + "config": { + "graph": args.graph, + "graph_sizes": graph_sizes, + "avg_degree": args.avg_degree, + "sbm_inter_density": args.sbm_inter_density, + "feature_dim": F, + "model": args.model, + "partitioner": "none", + "world_size": 1, + "ranks_per_node": 1, + "warmup": args.warmup, + "trials": args.trials, + "seed": args.seed, + "mode": "single_gpu_only", + }, + "measurements": measurements, + } + write_result(args.output, payload) + return + + # ---- Distributed run ---- + rank, world_size, local_rank = setup_distributed() + seed_everything(args.seed + rank) + + measurements, ranks_per_node = run_distributed( + args, graph_sizes, F, rank, world_size, local_rank + ) + + if rank == 0: + payload = { + "benchmark": "crossover", + "metadata": collect_metadata(), + "config": { + "graph": args.graph, + "graph_sizes": graph_sizes, + "avg_degree": args.avg_degree, + "sbm_inter_density": args.sbm_inter_density, + "feature_dim": F, + "model": args.model, + "partitioner": args.partitioner, + "world_size": world_size, + "ranks_per_node": ranks_per_node, + "warmup": args.warmup, + "trials": args.trials, + "seed": args.seed, + "mode": "distributed", + }, + "measurements": measurements, + } + write_result(args.output, payload) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() From 62f02487f3063ce476640398f97208d6b273c5f9 Mon Sep 17 00:00:00 2001 From: M Shehtab Zaman Date: Fri, 24 Apr 2026 21:04:15 -0700 Subject: [PATCH 07/39] Add running crossover benchmark and visualization --- DGraph/distributed/commInfo.py | 15 +- .../benchmarks/bench_crossover.py | 54 ++++-- .../benchmarks/nn_layer_common.py | 1 + .../visualization/plot_crossover.py | 183 ++++++++++++++++++ .../visualization/plot_flop_crossover.py | 140 ++++++++++++++ 5 files changed, 370 insertions(+), 23 deletions(-) create mode 100644 experiments/cost_model_benchmarks/visualization/plot_crossover.py create mode 100644 experiments/cost_model_benchmarks/visualization/plot_flop_crossover.py diff --git a/DGraph/distributed/commInfo.py b/DGraph/distributed/commInfo.py index 0d5bd7e..edf36a3 100644 --- a/DGraph/distributed/commInfo.py +++ b/DGraph/distributed/commInfo.py @@ -47,23 +47,22 @@ def compute_halo_vertices( """ Computes halo vertices. Supports both homogeneous and bipartite/heterogeneous relations. """ - # Fallback for homogeneous graphs if dst_partitioning is None: dst_partitioning = src_partitioning src_rank = src_partitioning[edge_list[:, 0]] dst_rank = dst_partitioning[edge_list[:, 1]] - # Cross-rank mask: source is local, destination is remote - cross_mask = (src_rank == rank) & (dst_rank != rank) + # FIX: Cross-rank mask for pull model: source is remote, destination is local + cross_mask = (src_rank != rank) & (dst_rank == rank) - # Return unique destination vertex IDs from those edges - return torch.unique(edge_list[cross_mask, 1]) + # FIX: Return unique SOURCE vertex IDs from those edges + return torch.unique(edge_list[cross_mask, 0]) def compute_local_edge_list( global_edge_list: torch.Tensor, # [E, 2] - partitioning: torch.Tensor, # [V] + partitioning: torch.Tensor, # [V] (Acts as dst_partitioning) local_vertices_global: torch.Tensor, # [num_local] halo_vertices_global: torch.Tensor, # [num_halo] rank: int, @@ -72,8 +71,8 @@ def compute_local_edge_list( num_halo = halo_vertices_global.size(0) num_global = partitioning.size(0) - # Filter edges owned by this rank - local_edge_mask = partitioning[global_edge_list[:, 0]] == rank + # FIX: Filter edges where the DESTINATION is owned by this rank (Index 1) + local_edge_mask = partitioning[global_edge_list[:, 1]] == rank local_edges_global = global_edge_list[local_edge_mask] # Build inverse map: global_id -> local_idx via scatter diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py b/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py index d6bc829..c5c2922 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py @@ -31,7 +31,7 @@ Usage — distributed (torchrun):: - torchrun --nnodes 2 --nproc_per_node 4 \\ + torchrun --nnodes 1 --nproc_per_node 4 \\ -m benchmarks.benchmark_crossover \\ --graph erdos_renyi \\ --graph-sizes 10000,100000,1000000,10000000 \\ @@ -55,7 +55,7 @@ import numpy as np import torch import torch.distributed as dist -import torch.nn as nn +import gc from benchmarks.common import ( collect_metadata, @@ -92,7 +92,7 @@ def parse_args(): p.add_argument( "--graph-sizes", type=str, - default="10000,100000,1000000,10000000", + default="100,200,400,800,1000,2000,4000,8000,10000,16000,32000,100000,200000,400000", help="Comma-separated list of num_vertices values to sweep", ) p.add_argument("--avg-degree", type=float, default=20.0) @@ -133,9 +133,7 @@ def _gen_graph(graph_type, num_vertices, avg_degree, sbm_inter_density, seed): return gen_sbm(num_vertices, avg_degree, sbm_inter_density, rng) -def _intra_inter_halo( - comm_pattern: CommunicationPattern, ranks_per_node: int -) -> tuple: +def _intra_inter_halo(comm_pattern: CommunicationPattern, ranks_per_node: int) -> tuple: """Return (intra_halo_vertices, inter_halo_vertices) from recv_offset.""" rank = comm_pattern.rank my_node = rank // ranks_per_node @@ -172,17 +170,21 @@ def _build_single_gpu_tensors(num_vertices, edges_np, F, model, device): def _single_gpu_fn(x, edge_t, layer, edge_attr): """Return a zero-argument closure for single-GPU forward+backward.""" if edge_attr is None: + def fn(): out = layer(x, edge_t) out.sum().backward() if x.grad is not None: x.grad.zero_() + else: + def fn(): out = layer(x, edge_t, edge_attr) out.sum().backward() if x.grad is not None: x.grad.zero_() + return fn @@ -195,14 +197,18 @@ def _multi_gpu_fn(x_local, comm_pattern, layer, halo_exchange, edge_attr, model) edge_index = comm_pattern.local_edge_list.T.contiguous() # [2, E_local] if model == "gcn": + def fn(): recv_buf = halo_exchange(x_local, comm_pattern) x_aug = torch.cat([x_local, recv_buf], dim=0) + out = layer(x_aug, edge_index) out.sum().backward() if x_local.grad is not None: x_local.grad.zero_() + else: + def fn(): recv_buf = halo_exchange(x_local, comm_pattern) x_aug = torch.cat([x_local, recv_buf], dim=0) @@ -210,6 +216,7 @@ def fn(): out.sum().backward() if x_local.grad is not None: x_local.grad.zero_() + return fn @@ -228,7 +235,11 @@ def run_no_dist(args, graph_sizes, F): # Unique seed per graph size so each topology is reproducible independently graph_seed = args.seed + num_vertices edges_np = _gen_graph( - args.graph, num_vertices, args.avg_degree, args.sbm_inter_density, graph_seed + args.graph, + num_vertices, + args.avg_degree, + args.sbm_inter_density, + graph_seed, ) num_edges = int(edges_np.shape[0]) @@ -292,21 +303,23 @@ def run_distributed(args, graph_sizes, F, rank, world_size, local_rank): device = torch.device(f"cuda:{local_rank}") comm = Communicator(backend="nccl") - halo_exchange = HaloExchange(comm=comm) ranks_per_node = int( - os.environ.get( - "LOCAL_WORLD_SIZE", os.environ.get("SLURM_NTASKS_PER_NODE", "4") - ) + os.environ.get("LOCAL_WORLD_SIZE", os.environ.get("SLURM_NTASKS_PER_NODE", "4")) ) measurements = [] for num_vertices in graph_sizes: + halo_exchange = HaloExchange(comm=comm) # All ranks generate *identical* graph topology (same seed per size) graph_seed = args.seed + num_vertices edges_np = _gen_graph( - args.graph, num_vertices, args.avg_degree, args.sbm_inter_density, graph_seed + args.graph, + num_vertices, + args.avg_degree, + args.sbm_inter_density, + graph_seed, ) # All ranks compute *identical* partition assignment @@ -319,7 +332,7 @@ def run_distributed(args, graph_sizes, F, rank, world_size, local_rank): assignment_np = partition_metis(num_vertices, world_size, edges_np) # Move to GPU for DGraph's build_communication_pattern (collective) - edges_t = torch.from_numpy(edges_np).long().to(device) # [E, 2] + edges_t = torch.from_numpy(edges_np).long().to(device) # [E, 2] partitioning_t = torch.from_numpy(assignment_np).long().to(device) # [V] # Collective: all ranks call this in sync (internally calls dist.all_gather) @@ -375,15 +388,26 @@ def run_distributed(args, graph_sizes, F, rank, world_size, local_rank): dist.all_gather_object(all_stats, stats_local) # Free GPU memory before the single-GPU run - del x_local, edges_t, partitioning_t, layer_dist + del ( + x_local, + edges_t, + partitioning_t, + layer_dist, + halo_exchange, + comm_pattern, + fn_multi, + ) if edge_attr_dist is not None: del edge_attr_dist + gc.collect() torch.cuda.empty_cache() # ---- Rank 0: time single-GPU baseline (non-collective) ---- times_single = None single_oom = False if rank == 0: + gc.collect() + torch.cuda.empty_cache() edges_s = _gen_graph( args.graph, num_vertices, @@ -399,7 +423,7 @@ def run_distributed(args, graph_sizes, F, rank, world_size, local_rank): times_single = cuda_timed( fn_single, warmup=args.warmup, trials=args.trials ) - del x_s, edge_t_s, layer_s + del x_s, edge_t_s, layer_s, fn_single if edge_attr_s is not None: del edge_attr_s torch.cuda.empty_cache() diff --git a/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py b/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py index 893d268..63a6af3 100644 --- a/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py +++ b/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py @@ -1,5 +1,6 @@ import torch import torch.nn as nn +import torch.distributed as dist # =========================================================================== # GNN layers (same as bench_compute.py for consistency) diff --git a/experiments/cost_model_benchmarks/visualization/plot_crossover.py b/experiments/cost_model_benchmarks/visualization/plot_crossover.py new file mode 100644 index 0000000..cfa02cb --- /dev/null +++ b/experiments/cost_model_benchmarks/visualization/plot_crossover.py @@ -0,0 +1,183 @@ +import numpy as np +import matplotlib.pyplot as plt +import json +from pathlib import Path +import argparse + + +def parse_args(): + p = argparse.ArgumentParser(description="Plot Crossover Benchmark Results") + p.add_argument( + "--input", type=str, required=True, help="Path to benchmark JSON file" + ) + p.add_argument( + "--output", + type=str, + default="crossover_analysis.png", + help="Path to save the generated plot", + ) + return p.parse_args() + + +def plot_crossover_benchmark(payload, save_path="crossover_analysis.png"): + """ + Parses benchmark payload and plots Single GPU vs Multi-GPU execution times, + annotating the crossover point where distributed training becomes faster. + """ + # Sort measurements strictly by graph size (num_vertices) + measurements = sorted( + payload["measurements"], key=lambda x: x["params"]["num_vertices"] + ) + + vertices = [] + single_gpu_times = [] + multi_gpu_times = [] + single_gpu_oom_points = [] + + world_size = payload["config"]["world_size"] + model_name = payload["config"]["model"] + partitioner = payload["config"]["partitioner"] + feature_dim = payload["config"]["feature_dim"] + + for m in measurements: + v = m["params"]["num_global_edges"] + vertices.append(v) + + # Multi-GPU: Use the max time across ranks as it represents the true synchronous bottleneck + multi_time = np.median(m["multi_gpu_trials_seconds_max"]) + multi_gpu_times.append(multi_time) + + # Single-GPU: Handle OOM scenarios safely + if m.get("single_gpu_oom", False) or not m["single_gpu_trials_seconds"]: + single_gpu_times.append(np.nan) + single_gpu_oom_points.append(v) + else: + single_gpu_times.append(np.median(m["single_gpu_trials_seconds"])) + + vertices = np.array(vertices) + single_gpu_times = np.array(single_gpu_times) + multi_gpu_times = np.array(multi_gpu_times) + + # Initialize Plot + plt.figure(figsize=(10, 6), dpi=150) + plt.grid(True, which="both", ls="-", alpha=0.2) + + # Plot Valid Single GPU points + valid_mask = ~np.isnan(single_gpu_times) + plt.plot( + vertices[valid_mask], + single_gpu_times[valid_mask], + marker="o", + linestyle="-", + linewidth=2, + color="#1f77b4", + label="Single GPU", + ) + + # Plot Multi GPU points + plt.plot( + vertices, + multi_gpu_times, + marker="s", + linestyle="-", + linewidth=2, + color="#ff7f0e", + label=f"Distributed ({world_size} GPUs)", + ) + + # Annotate OOM boundaries + for oom_v in single_gpu_oom_points: + plt.axvline(x=oom_v, color="#d62728", linestyle="--", alpha=0.6) + plt.text( + oom_v * 1.02, + plt.ylim()[1] * 0.9, + "Single GPU OOM", + color="#d62728", + verticalalignment="top", + ) + + # Calculate and Annotate Crossover Point + crossover_found = False + for i in range(1, len(vertices)): + if valid_mask[i - 1] and valid_mask[i]: + diff_prev = multi_gpu_times[i - 1] - single_gpu_times[i - 1] + diff_curr = multi_gpu_times[i] - single_gpu_times[i] + + # A sign change indicates the lines crossed + if diff_prev * diff_curr < 0: + x1, x2 = vertices[i - 1], vertices[i] + y1_s, y2_s = single_gpu_times[i - 1], single_gpu_times[i] + y1_m, y2_m = multi_gpu_times[i - 1], multi_gpu_times[i] + + # Linear interpolation for precise intersection coordinates + m_s = (y2_s - y1_s) / (x2 - x1) + m_m = (y2_m - y1_m) / (x2 - x1) + + if m_s != m_m: + x_cross = x1 + (y1_m - y1_s) / (m_s - m_m) + y_cross = y1_s + m_s * (x_cross - x1) + + plt.plot( + x_cross, + y_cross, + marker="*", + color="#2ca02c", + markersize=15, + zorder=5, + ) + plt.annotate( + f"Crossover:\n~{int(x_cross):,} edges", + xy=(x_cross, y_cross), + xytext=(-20, 40), + textcoords="offset points", + fontsize=10, + fontweight="bold", + color="#2ca02c", + arrowprops=dict( + arrowstyle="->", + connectionstyle="arc3,rad=.2", + color="#2ca02c", + ), + ) + crossover_found = True + break + + # Formatting + plt.title( + f"GNN Distributed Scaling Crossover\nModel: {model_name} | Partitioner: {partitioner} | Feature Dim: {feature_dim}", + fontsize=14, + pad=15, + ) + plt.xlabel("Graph Size (Number of Edges)", fontsize=12) + plt.ylabel("Execution Time (Seconds)", fontsize=12) + plt.xscale( + "log" + ) # Using log scale for x-axis as graph sizes usually scale exponentially + plt.yscale("linear") + + plt.legend(loc="upper left", framealpha=0.9) + plt.tight_layout() + + # Output + plt.savefig(save_path) + print(f"Visualization saved to {save_path}") + if crossover_found: + print(f"Crossover point detected at approximately {int(x_cross):,} vertices.") + else: + print("No crossover point detected in the provided dataset.") + + +# Example usage assuming `payload` is already loaded in your environment: +# plot_crossover_benchmark(payload) + + +def main(): + args = parse_args() + with open(args.input, "r") as f: + payload = json.load(f) + + plot_crossover_benchmark(payload, save_path=args.output) + + +if __name__ == "__main__": + main() diff --git a/experiments/cost_model_benchmarks/visualization/plot_flop_crossover.py b/experiments/cost_model_benchmarks/visualization/plot_flop_crossover.py new file mode 100644 index 0000000..b2ea789 --- /dev/null +++ b/experiments/cost_model_benchmarks/visualization/plot_flop_crossover.py @@ -0,0 +1,140 @@ +import glob +import json +import re +import numpy as np +import matplotlib.pyplot as plt +from collections import defaultdict + + +def calculate_crossover(measurements): + """ + Calculates the exact crossover point (num_global_edges) where multi-GPU + execution becomes faster than single-GPU execution using linear interpolation. + """ + # Sort measurements strictly by graph size + measurements = sorted(measurements, key=lambda x: x["params"]["num_global_edges"]) + + vertices = [] + single_times = [] + multi_times = [] + + for m in measurements: + v = m["params"]["num_global_edges"] + + # Extract median times, ignore OOM or missing single GPU data + if m.get("single_gpu_oom", False) or not m.get("single_gpu_trials_seconds"): + continue + + s_time = np.median(m["single_gpu_trials_seconds"]) + m_time = np.median( + m["multi_gpu_trials_seconds_max"] + ) # Use max time for synchronous bottleneck + + vertices.append(v) + single_times.append(s_time) + multi_times.append(m_time) + + for i in range(1, len(vertices)): + diff_prev = multi_times[i - 1] - single_times[i - 1] + diff_curr = multi_times[i] - single_times[i] + + # Sign change indicates the lines crossed + if diff_prev * diff_curr < 0: + x1, x2 = vertices[i - 1], vertices[i] + y1_s, y2_s = single_times[i - 1], single_times[i] + y1_m, y2_m = multi_times[i - 1], multi_times[i] + + m_s = (y2_s - y1_s) / (x2 - x1) + m_m = (y2_m - y1_m) / (x2 - x1) + + if m_s != m_m: + x_cross = x1 + (y1_m - y1_s) / (m_s - m_m) + return x_cross + + return None # No crossover found in this dataset + + +def plot_crossover_dynamics( + file_pattern="results/crossover_*_world_F*.json", + save_path="results/crossover_vs_features.png", +): + """ + Parses benchmark files and plots the crossover point as a function of feature dimension, + with separate lines for each distributed world size. + """ + # Structure: data[world_size][feature_dim] = crossover_point + plot_data = defaultdict(dict) + + # Locate files + files = glob.glob(file_pattern) + if not files: + print(f"No files found matching pattern: {file_pattern}") + return + + # Regex to extract parameters from filename + pattern = re.compile(r"crossover_(\d+)_world_F(\d+)\.json") + + for filepath in files: + match = pattern.search(filepath) + if match: + world_size = int(match.group(1)) + feature_dim = int(match.group(2)) + + try: + with open(filepath, "r") as f: + payload = json.load(f) + + crossover_point = calculate_crossover(payload.get("measurements", [])) + if crossover_point is not None: + plot_data[world_size][feature_dim] = crossover_point + except Exception as e: + print(f"Error processing {filepath}: {e}") + + if not plot_data: + print("No valid crossover points could be calculated from the provided files.") + return + + # Initialize Plot + plt.figure(figsize=(10, 6), dpi=150) + plt.grid(True, which="both", ls="-", alpha=0.3) + + markers = ["o", "s", "^", "D", "v", "p"] + + # Plot lines per world size + for idx, (world_size, dim_data) in enumerate(sorted(plot_data.items())): + # Sort by feature dimension for sequential line plotting + sorted_dims = sorted(dim_data.keys()) + crossovers = [dim_data[d] for d in sorted_dims] + + plt.plot( + sorted_dims, + crossovers, + marker=markers[idx % len(markers)], + linestyle="-", + linewidth=2, + markersize=8, + label=f"{world_size} GPUs", + ) + + # Formatting + plt.title( + "GNN Communication Bottleneck:\nCrossover Threshold vs. Feature Dimension", + fontsize=14, + pad=15, + ) + plt.xlabel("Feature Dimension (F)", fontsize=12) + plt.ylabel("Crossover Point (Number of Edges)", fontsize=12) + + # Depending on the range of your feature dims, a log scale might be preferable for X + # plt.xscale("log", base=2) + plt.yscale("log") # Y-axis (vertices) usually scales exponentially + + plt.legend(title="World Size", loc="upper left", framealpha=0.9) + plt.tight_layout() + + plt.savefig(save_path) + print(f"Scaling dynamics visualization saved successfully to {save_path}") + + +if __name__ == "__main__": + plot_crossover_dynamics() From 1d8f90b82081d446809682ce4a8dca8b1d25b534 Mon Sep 17 00:00:00 2001 From: M Shehtab Zaman Date: Fri, 1 May 2026 14:31:13 -0700 Subject: [PATCH 08/39] Add GCN implementation --- experiments/OGB/GCN.py | 69 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/experiments/OGB/GCN.py b/experiments/OGB/GCN.py index 2b3cab4..c35879f 100644 --- a/experiments/OGB/GCN.py +++ b/experiments/OGB/GCN.py @@ -20,11 +20,80 @@ import torch.distributed as dist import sys +try: + from torch_geometric.nn.conv import GCNConv + + PYG_AVAILABLE = True +except ImportError: + print( + "torch_geometric not found, skipping import of GCNConv. Make sure to install torch_geometric if you want to use it." + ) + PYG_AVAILABLE = False + import torch import torch.nn as nn +class GCNLayer_impl(nn.Module): + def __init__(self, in_channels, out_channels): + super(GCNLayer_impl, self).__init__() + self.linear = nn.Linear(in_channels, out_channels) + self.act = nn.ReLU(inplace=True) + + def forward(self, x, edge_index): + source_vertices = edge_index[:, 0] + target_vertices = edge_index[:, 1] + x = self.linear(x) + + x_j = x[target_vertices, :] + out_channels = x_j.size(1) + out = torch.zeros(x.size(0), out_channels, dtype=x.dtype, device=x.device) + scatter_index = ( + source_vertices.unsqueeze(-1).expand(-1, out_channels).to(x.device) + ) + out = out.scatter_add(0, scatter_index, x_j) + out = self.act(out) + return out + + +class GCNLayer(nn.Module): + def __init__(self, in_channels, out_channels): + super(GCNLayer, self).__init__() + if PYG_AVAILABLE: + self.conv = GCNConv(in_channels, out_channels) + else: + self.conv = GCNLayer_impl(in_channels, out_channels) + + def forward(self, x, edge_index): + return self.conv(x, edge_index) + + +class GCNModel(nn.Module): + def __init__( + self, in_channels, hidden_channels, out_channels, num_layers, halo_exchanger + ): + super(GCNModel, self).__init__() + self.convs = nn.ModuleList() + self.convs.append(GCNLayer(in_channels, hidden_channels)) + for _ in range(num_layers - 2): + self.convs.append(GCNLayer(hidden_channels, hidden_channels)) + self.convs.append(GCNLayer(hidden_channels, out_channels)) + self.halo_exchanger = halo_exchanger + + def forward(self, x, comm_pattern): + edge_index = comm_pattern.local_edge_list + for conv in self.convs[:-1]: + boundary_features = self.halo_exchanger(x, comm_pattern) + x = torch.cat([x, boundary_features], dim=0) + x = conv(x, edge_index) + + boundary_features = self.halo_exchanger(x, comm_pattern) + x = torch.cat([x, boundary_features], dim=0) + x = self.convs[-1](x, edge_index) + return x + + class GraphConvLayer(nn.Module): def __init__(self, message_dim, out_channels): super(GraphConvLayer, self).__init__() From 0497ab32f5db0010156f13366964a84d8e75307a Mon Sep 17 00:00:00 2001 From: M Shehtab Zaman Date: Fri, 1 May 2026 16:03:13 -0700 Subject: [PATCH 09/39] Add sparse implementation of GCN --- experiments/OGB/GCN.py | 69 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 64 insertions(+), 5 deletions(-) diff --git a/experiments/OGB/GCN.py b/experiments/OGB/GCN.py index c35879f..6914104 100644 --- a/experiments/OGB/GCN.py +++ b/experiments/OGB/GCN.py @@ -35,6 +35,56 @@ import torch.nn as nn +def create_sparse_adj(edge_index, num_nodes, device="cpu"): + """ + Converts an edge_index of shape [num_edges, 2] into a PyTorch sparse tensor. + """ + # PyTorch sparse tensors expect indices in shape [2, num_edges] + # So we transpose the [num_edges, 2] tensor + + indices = edge_index.t().contiguous() + + # Adjacency values are 1 for unweighted graphs + values = torch.ones(edge_index.size(0), device=device) + # Create the sparse COO tensor + adj_sparse = torch.sparse_coo_tensor( + indices, values, size=(num_nodes, num_nodes), device=device + ) + + # # OPTIONAL BUT RECOMMENDED: Convert to CSR format for faster spmm operations + # # PyTorch's sparse.mm is significantly faster on CSR tensors than COO tensors. + if adj_sparse.is_coalesced() is False: + adj_sparse = adj_sparse.coalesce() + adj_sparse_csr = adj_sparse.to_sparse_csr() + + return adj_sparse_csr + + +class GCNLayer_impl_sparse(nn.Module): + def __init__(self, in_channels, out_channels): + super(GCNLayer_impl_sparse, self).__init__() + self.linear = nn.Linear(in_channels, out_channels) + self.act = nn.ReLU(inplace=True) + + def forward(self, x, adj_sparse): + """ + Args: + x: Node features tensor of shape [num_nodes, in_channels] + adj_sparse: Sparse adjacency tensor of shape [num_nodes, num_nodes] + """ + # 1. Feature transformation (Dense) + x = self.linear(x) + + # 2. Message Passing via Sparse Matrix Multiplication + # This single line replaces x_j, out.zeros(), and scatter_add() + out = torch.sparse.mm(adj_sparse, x) + + # 3. Activation + out = self.act(out) + + return out + + class GCNLayer_impl(nn.Module): def __init__(self, in_channels, out_channels): super(GCNLayer_impl, self).__init__() @@ -58,12 +108,15 @@ def forward(self, x, edge_index): class GCNLayer(nn.Module): - def __init__(self, in_channels, out_channels): + def __init__(self, in_channels, out_channels, use_sparse=False): super(GCNLayer, self).__init__() if PYG_AVAILABLE: self.conv = GCNConv(in_channels, out_channels) else: - self.conv = GCNLayer_impl(in_channels, out_channels) + if use_sparse: + self.conv = GCNLayer_impl_sparse(in_channels, out_channels) + else: + self.conv = GCNLayer_impl(in_channels, out_channels) def forward(self, x, edge_index): return self.conv(x, edge_index) @@ -71,14 +124,20 @@ def forward(self, x, edge_index): class GCNModel(nn.Module): def __init__( - self, in_channels, hidden_channels, out_channels, num_layers, halo_exchanger + self, + in_channels, + hidden_channels, + out_channels, + num_layers, + halo_exchanger, + is_sparse=False, ): super(GCNModel, self).__init__() self.convs = nn.ModuleList() self.convs.append(GCNLayer(in_channels, hidden_channels)) for _ in range(num_layers - 2): - self.convs.append(GCNLayer(hidden_channels, hidden_channels)) - self.convs.append(GCNLayer(hidden_channels, out_channels)) + self.convs.append(GCNLayer(hidden_channels, hidden_channels, is_sparse)) + self.convs.append(GCNLayer(hidden_channels, out_channels, is_sparse)) self.halo_exchanger = halo_exchanger def forward(self, x, comm_pattern): From 630f81e2428988cdc024637c0adce4d27478c561 Mon Sep 17 00:00:00 2001 From: M Shehtab Zaman Date: Wed, 6 May 2026 05:48:48 -0700 Subject: [PATCH 10/39] Add fixed end to end benchmark code for OGB with sparse GCN --- experiments/OGB/GCN.py | 24 +++- experiments/OGB/benchmark_OGB_end_to_end.py | 151 ++++++++++++++++++++ 2 files changed, 168 insertions(+), 7 deletions(-) create mode 100644 experiments/OGB/benchmark_OGB_end_to_end.py diff --git a/experiments/OGB/GCN.py b/experiments/OGB/GCN.py index 6914104..4c3b9ee 100644 --- a/experiments/OGB/GCN.py +++ b/experiments/OGB/GCN.py @@ -35,7 +35,7 @@ import torch.nn as nn -def create_sparse_adj(edge_index, num_nodes, device="cpu"): +def create_sparse_adj(edge_index, num_local_nodes, num_halo_nodes, device="cpu"): """ Converts an edge_index of shape [num_edges, 2] into a PyTorch sparse tensor. """ @@ -43,18 +43,28 @@ def create_sparse_adj(edge_index, num_nodes, device="cpu"): # So we transpose the [num_edges, 2] tensor indices = edge_index.t().contiguous() + indices = torch.stack( + [ + edge_index[:, 1], # Rows: Targets (strictly < num_local_nodes) + edge_index[:, 0], # Cols: Sources (< num_local_nodes + num_halo_nodes) + ] + ) # Adjacency values are 1 for unweighted graphs - values = torch.ones(edge_index.size(0), device=device) + values = torch.ones(edge_index.size(0), dtype=torch.float32, device=device) + # Create the sparse COO tensor adj_sparse = torch.sparse_coo_tensor( - indices, values, size=(num_nodes, num_nodes), device=device + indices, + values, + size=(num_local_nodes, num_local_nodes + num_halo_nodes), + device=device, ) - # # OPTIONAL BUT RECOMMENDED: Convert to CSR format for faster spmm operations - # # PyTorch's sparse.mm is significantly faster on CSR tensors than COO tensors. - if adj_sparse.is_coalesced() is False: + # Convert to CSR format for faster spmm operations + if not adj_sparse.is_coalesced(): adj_sparse = adj_sparse.coalesce() + adj_sparse_csr = adj_sparse.to_sparse_csr() return adj_sparse_csr @@ -134,7 +144,7 @@ def __init__( ): super(GCNModel, self).__init__() self.convs = nn.ModuleList() - self.convs.append(GCNLayer(in_channels, hidden_channels)) + self.convs.append(GCNLayer(in_channels, hidden_channels, is_sparse)) for _ in range(num_layers - 2): self.convs.append(GCNLayer(hidden_channels, hidden_channels, is_sparse)) self.convs.append(GCNLayer(hidden_channels, out_channels, is_sparse)) diff --git a/experiments/OGB/benchmark_OGB_end_to_end.py b/experiments/OGB/benchmark_OGB_end_to_end.py new file mode 100644 index 0000000..b9421e3 --- /dev/null +++ b/experiments/OGB/benchmark_OGB_end_to_end.py @@ -0,0 +1,151 @@ +# Copyright (c) 2014-2024, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LLNL/LBANN. +# +# SPDX-License-Identifier: (Apache-2.0) +""" +Distributed GCN benchmark on OGB node-property-prediction datasets. +""" + +import os +import json +from typing import Optional + +import fire +import numpy as np +import torch +import torch.distributed as dist +import torch.optim as optim +from torch.nn.parallel import DistributedDataParallel as DDP + +from DGraph.Communicator import Communicator +from DGraph.distributed import HaloExchange +from DGraph.utils.TimingReport import TimingReport + +from GCN import GCNModel, create_sparse_adj +from ogb_comm_dataset import DGraphOGBDataset + + +def benchmark_ogb_end_to_end( + dataset: DGraphOGBDataset, + comm: Communicator, + lr: float, + epochs: int, + log_prefix: str, + hidden_dims: int = 256, + num_classes: int = 40, + num_layers: int = 2, + device: str | torch.device = "cuda", + rank: int = 0, + local_rank: int = 0, + warm_up: int = 10, + cur_dir: Optional[str] = None, + convert_to_sparse_adj: bool = True, +): + + graph_dataset = dataset[0] + local_node_features, local_labels, comm_pattern = graph_dataset + + in_dim = local_node_features.shape[1] + local_masks = dataset.get_masks() + train_mask = local_masks["train_mask"].to(device) + val_mask = local_masks["val_mask"].to(device) + test_mask = local_masks["test_mask"].to(device) + + halo_exchanger = HaloExchange(comm) + if convert_to_sparse_adj: + comm_pattern.local_edge_list = create_sparse_adj( + comm_pattern.local_edge_list, + num_local_nodes=comm_pattern.num_local_vertices, + num_halo_nodes=comm_pattern.num_halo_vertices, + device=device, + ) + model = GCNModel( + in_channels=in_dim, + hidden_channels=hidden_dims, + out_channels=num_classes, + num_layers=num_layers, + halo_exchanger=halo_exchanger, + is_sparse=True, + ).to(device) + model = DDP(model, device_ids=[local_rank], output_device=local_rank) + optimizer = optim.Adam(model.parameters(), lr=lr) + criterion = torch.nn.CrossEntropyLoss() + + stream = torch.cuda.Stream() + TimingReport.init(communicator=comm) + + if rank == 0: + print( + f"Starting benchmark with {comm.get_world_size()} processes, local rank {local_rank}, global rank {rank}", + f" Local node features shape: {local_node_features.shape}, local labels shape: {local_labels.shape}", + ) + + for epoch in range(epochs): + with TimingReport( + "total-training-time", + ): + y = model(local_node_features.to(device), comm_pattern) + loss = y.sum() + # loss = criterion(y[train_mask], local_labels[train_mask]) + optimizer.zero_grad() + loss.backward() + optimizer.step() + + if rank == 0: + timers = TimingReport._timers + report = {name: np.mean(times[warm_up:]) for name, times in timers.items()} + print(report) + os.makedirs( + f"{cur_dir}/benchmark_results/benchmark_results", + exist_ok=True, + ) + with open( + f"{cur_dir}/benchmark_results/{log_prefix}_timing_report.json", + "w", + ) as f: + json.dump(report, f, indent=4) + + +def main(): + + comm = Communicator.init_process_group("nccl") + rank = comm.get_rank() + local_rank = rank % torch.cuda.device_count() + world_size = comm.get_world_size() + + torch.cuda.set_device(local_rank) + + cur_dir = os.path.dirname(os.path.abspath(__file__)) + dataset = DGraphOGBDataset( + dname="ogbn-products", + comm=comm, + root_dir=f"{cur_dir}/data", + ) + + benchmark_ogb_end_to_end( + dataset=dataset, + comm=comm, + lr=0.01, + epochs=100, + log_prefix=f"ogbn_arxiv-{world_size}", + hidden_dims=128, + num_classes=47, + num_layers=3, + device="cuda", + rank=rank, + local_rank=local_rank, + cur_dir=cur_dir, + ) + + +if __name__ == "__main__": + main() From b9e73e66a73be2d7a325f2e780cf5e5d1351cb6a Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 22 May 2026 10:00:42 -0400 Subject: [PATCH 11/39] Update OGB end to end with command line options --- experiments/OGB/benchmark_OGB_end_to_end.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/experiments/OGB/benchmark_OGB_end_to_end.py b/experiments/OGB/benchmark_OGB_end_to_end.py index b9421e3..48de151 100644 --- a/experiments/OGB/benchmark_OGB_end_to_end.py +++ b/experiments/OGB/benchmark_OGB_end_to_end.py @@ -115,7 +115,14 @@ def benchmark_ogb_end_to_end( json.dump(report, f, indent=4) -def main(): +def main(dataset: str = "arxiv"): + + assert dataset in ( + "arxiv", + "products", + ), f"Unsupported dataset '{dataset}'. Choose from: arxiv, products." + + num_classes = {"arxiv": 40, "products": 47}[dataset] comm = Communicator.init_process_group("nccl") rank = comm.get_rank() @@ -126,7 +133,7 @@ def main(): cur_dir = os.path.dirname(os.path.abspath(__file__)) dataset = DGraphOGBDataset( - dname="ogbn-products", + dname=f"ogbn-{dataset}", comm=comm, root_dir=f"{cur_dir}/data", ) @@ -136,9 +143,9 @@ def main(): comm=comm, lr=0.01, epochs=100, - log_prefix=f"ogbn_arxiv-{world_size}", + log_prefix=f"ogbn_{dataset}-{world_size}", hidden_dims=128, - num_classes=47, + num_classes=num_classes, num_layers=3, device="cuda", rank=rank, @@ -148,4 +155,4 @@ def main(): if __name__ == "__main__": - main() + fire.Fire(main) From 2e1f88101f902100972f0289ae24ddf3fcd8cc1e Mon Sep 17 00:00:00 2001 From: M Shehtab Zaman Date: Fri, 22 May 2026 09:11:37 -0700 Subject: [PATCH 12/39] Fix main with sparse GCN implementation for OGB --- experiments/OGB/GCN.py | 12 ++++++--- experiments/OGB/benchmark_OGB_end_to_end.py | 4 +-- experiments/OGB/main.py | 29 +++++++++++++++------ experiments/OGB/ogb_comm_dataset.py | 4 +++ 4 files changed, 35 insertions(+), 14 deletions(-) diff --git a/experiments/OGB/GCN.py b/experiments/OGB/GCN.py index 4c3b9ee..c04d277 100644 --- a/experiments/OGB/GCN.py +++ b/experiments/OGB/GCN.py @@ -17,8 +17,7 @@ from DGraph.distributed import HaloExchange, CommunicationPattern from DGraph.utils.TimingReport import TimingReport from DGraph import Communicator -import torch.distributed as dist -import sys +from typing import Union try: from torch_geometric.nn.conv import GCNConv @@ -35,7 +34,12 @@ import torch.nn as nn -def create_sparse_adj(edge_index, num_local_nodes, num_halo_nodes, device="cpu"): +def create_sparse_adj( + edge_index, + num_local_nodes, + num_halo_nodes, + device: Union[str, torch.device] = "cpu", +): """ Converts an edge_index of shape [num_edges, 2] into a PyTorch sparse tensor. """ @@ -173,7 +177,7 @@ def forward(self, x, edge_index, num_local_nodes, edge_features=None): source_vertices = edge_index[:, 0] target_vertices = edge_index[:, 1] - assert (source_vertices < num_local_nodes).all(), ( + assert (target_vertices < num_local_nodes).all(), ( f"Graph routing error: Found source_vertices >= num_local_nodes ({num_local_nodes}). " "Boundary nodes must only act as targets (x_j) in this aggregation scheme!" ) diff --git a/experiments/OGB/benchmark_OGB_end_to_end.py b/experiments/OGB/benchmark_OGB_end_to_end.py index 48de151..9f6a526 100644 --- a/experiments/OGB/benchmark_OGB_end_to_end.py +++ b/experiments/OGB/benchmark_OGB_end_to_end.py @@ -132,14 +132,14 @@ def main(dataset: str = "arxiv"): torch.cuda.set_device(local_rank) cur_dir = os.path.dirname(os.path.abspath(__file__)) - dataset = DGraphOGBDataset( + dist_dataset = DGraphOGBDataset( dname=f"ogbn-{dataset}", comm=comm, root_dir=f"{cur_dir}/data", ) benchmark_ogb_end_to_end( - dataset=dataset, + dataset=dist_dataset, comm=comm, lr=0.01, epochs=100, diff --git a/experiments/OGB/main.py b/experiments/OGB/main.py index 9ad97e4..c4f8492 100644 --- a/experiments/OGB/main.py +++ b/experiments/OGB/main.py @@ -14,6 +14,7 @@ """ Distributed GCN benchmark on OGB node-property-prediction datasets. """ + import os import json from typing import Optional @@ -28,8 +29,7 @@ from DGraph.Communicator import Communicator from DGraph.distributed import HaloExchange from DGraph.utils.TimingReport import TimingReport - -from GCN import CommAwareGCN as GCN +from GCN import GCNModel, create_sparse_adj from ogb_comm_dataset import DGraphOGBDataset from utils import ( calculate_accuracy, @@ -41,7 +41,6 @@ write_experiment_log, ) - # --------------------------------------------------------------------------- # Training / evaluation loop # --------------------------------------------------------------------------- @@ -77,9 +76,22 @@ def _run_experiment( """ # ---- Extract local data from the dataset -------------------------------- - comm_pattern = dataset.comm_pattern local_node_features, local_labels, comm_pattern = dataset[0] + + if not comm_pattern.local_edge_list.is_sparse_csr: + print( + comm_pattern.local_edge_list.is_sparse, + comm_pattern.local_edge_list.is_sparse_csr, + ) + comm_pattern.local_edge_list = create_sparse_adj( + comm_pattern.local_edge_list, + num_local_nodes=comm_pattern.num_local_vertices, + num_halo_nodes=comm_pattern.num_halo_vertices, + device=device, + ) + dist.barrier() + local_node_features = local_node_features.to(device) local_labels = local_labels.to(device) @@ -100,12 +112,13 @@ def _run_experiment( # ---- Model setup -------------------------------------------------------- halo_exchanger = HaloExchange(comm) - model = GCN( + model = GCNModel( in_channels=in_dim, - hidden_dims=hidden_dims, - num_classes=num_classes, + hidden_channels=hidden_dims, + out_channels=num_classes, + num_layers=3, halo_exchanger=halo_exchanger, - comm=comm, + is_sparse=True, ).to(device) if comm.get_world_size() > 1: diff --git a/experiments/OGB/ogb_comm_dataset.py b/experiments/OGB/ogb_comm_dataset.py index f739a49..b0f3ea2 100644 --- a/experiments/OGB/ogb_comm_dataset.py +++ b/experiments/OGB/ogb_comm_dataset.py @@ -164,6 +164,8 @@ def __getitem__( data, labels, comm_pattern = dataset[0] if rank == 0: + + breakpoint() import os print(comm_pattern.comm_map) @@ -172,3 +174,5 @@ def __getitem__( file_dir = os.path.dirname(file_path) print(f"Saving to {file_dir}/comm_map_{world_size}.pt") torch.save(comm_pattern.comm_map, f"{file_dir}/comm_map_{world_size}.pt") + + comm.barrier() From 01ca083a2575c87ee90ecf1c82cb699c56cb70cc Mon Sep 17 00:00:00 2001 From: M Shehtab Zaman Date: Fri, 22 May 2026 09:12:58 -0700 Subject: [PATCH 13/39] Remove extra prints --- experiments/OGB/main.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/experiments/OGB/main.py b/experiments/OGB/main.py index c4f8492..2e4853b 100644 --- a/experiments/OGB/main.py +++ b/experiments/OGB/main.py @@ -80,10 +80,7 @@ def _run_experiment( local_node_features, local_labels, comm_pattern = dataset[0] if not comm_pattern.local_edge_list.is_sparse_csr: - print( - comm_pattern.local_edge_list.is_sparse, - comm_pattern.local_edge_list.is_sparse_csr, - ) + comm_pattern.local_edge_list = create_sparse_adj( comm_pattern.local_edge_list, num_local_nodes=comm_pattern.num_local_vertices, From caa2ceeafc43464a7d9ee0506fc795db0e99593b Mon Sep 17 00:00:00 2001 From: M Shehtab Zaman Date: Fri, 10 Jul 2026 13:44:53 -0700 Subject: [PATCH 14/39] Completed OGB fix and experiments --- experiments/OGB/GCN.py | 23 ++++++++++++++++------- experiments/OGB/main.py | 13 +++++++------ experiments/OGB/utils.py | 2 +- 3 files changed, 24 insertions(+), 14 deletions(-) diff --git a/experiments/OGB/GCN.py b/experiments/OGB/GCN.py index c04d277..a15ecb4 100644 --- a/experiments/OGB/GCN.py +++ b/experiments/OGB/GCN.py @@ -124,7 +124,7 @@ def forward(self, x, edge_index): class GCNLayer(nn.Module): def __init__(self, in_channels, out_channels, use_sparse=False): super(GCNLayer, self).__init__() - if PYG_AVAILABLE: + if False: self.conv = GCNConv(in_channels, out_channels) else: if use_sparse: @@ -156,14 +156,23 @@ def __init__( def forward(self, x, comm_pattern): edge_index = comm_pattern.local_edge_list + counter = 1 for conv in self.convs[:-1]: + with TimingReport(f"feature-exchange-{counter}"): + boundary_features = self.halo_exchanger(x, comm_pattern) + + with TimingReport(f"process-{counter}"): + x = torch.cat([x, boundary_features], dim=0) + x = conv(x, edge_index) + + counter += 1 + + with TimingReport(f"feature-exchange-{counter}"): boundary_features = self.halo_exchanger(x, comm_pattern) - x = torch.cat([x, boundary_features], dim=0) - x = conv(x, edge_index) - boundary_features = self.halo_exchanger(x, comm_pattern) - x = torch.cat([x, boundary_features], dim=0) - x = self.convs[-1](x, edge_index) + with TimingReport(f"process-{counter}"): + x = torch.cat([x, boundary_features], dim=0) + x = self.convs[-1](x, edge_index) return x @@ -177,7 +186,7 @@ def forward(self, x, edge_index, num_local_nodes, edge_features=None): source_vertices = edge_index[:, 0] target_vertices = edge_index[:, 1] - assert (target_vertices < num_local_nodes).all(), ( + assert (source_vertices < num_local_nodes).all(), ( f"Graph routing error: Found source_vertices >= num_local_nodes ({num_local_nodes}). " "Boundary nodes must only act as targets (x_j) in this aggregation scheme!" ) diff --git a/experiments/OGB/main.py b/experiments/OGB/main.py index 2e4853b..381ae86 100644 --- a/experiments/OGB/main.py +++ b/experiments/OGB/main.py @@ -154,10 +154,11 @@ def _run_experiment( train_labels = local_labels[train_mask] loss = criterion(train_output, train_labels.reshape(-1)) - with TimingReport("backward"): - loss.backward() + comm.barrier() + # with TimingReport("backward"): + # loss.backward() - optimizer.step() + # optimizer.step() comm.barrier() end_time.record(stream) @@ -349,19 +350,19 @@ def main( visualize_trajectories( training_trajectories, "Training Loss", - f"{log_dir}/training_loss.png", + f"{log_dir}/{dataset}_world_{world_size}_training_loss.png", rank, ) visualize_trajectories( validation_trajectories, "Validation Loss", - f"{log_dir}/validation_loss.png", + f"{log_dir}/{dataset}_world_{world_size}_validation_loss.png", rank, ) visualize_trajectories( validation_accuracies, "Validation Accuracy", - f"{log_dir}/validation_accuracy.png", + f"{log_dir}/{dataset}_world_{world_size}_validation_accuracy.png", rank, ) diff --git a/experiments/OGB/utils.py b/experiments/OGB/utils.py index 30fe867..ab68517 100644 --- a/experiments/OGB/utils.py +++ b/experiments/OGB/utils.py @@ -17,7 +17,7 @@ def make_experiment_log(fname, rank): def write_experiment_log(log: str, fname: str, rank: int): if rank == 0: - with open(fname, "a") as f: + with open(fname, "w") as f: f.write(log + "\n") From fc37edc1709c7d78a3aa0a74cf267794e9b6c6b2 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 17 Jul 2026 09:48:05 -0400 Subject: [PATCH 15/39] Add GraphCast benchmarks --- DGraph/distributed/commInfo.py | 97 +++- .../GraphCast/data_utils/graphcast_graph.py | 105 +++- .../GraphCast/data_utils/preprocess.py | 2 +- experiments/GraphCast/dist_utils.py | 9 + experiments/GraphCast/gen_mesh_partitions.py | 83 ++++ experiments/GraphCast/inference_benchmark.py | 448 ++++++++++++++++++ experiments/GraphCast/model.py | 29 +- experiments/GraphCast/plot_scaling.py | 336 +++++++++++++ experiments/GraphCast/tests/conftest.py | 15 + .../GraphCast/tests/test_single_graph_data.py | 19 +- 10 files changed, 1082 insertions(+), 61 deletions(-) create mode 100644 experiments/GraphCast/gen_mesh_partitions.py create mode 100644 experiments/GraphCast/inference_benchmark.py create mode 100644 experiments/GraphCast/plot_scaling.py diff --git a/DGraph/distributed/commInfo.py b/DGraph/distributed/commInfo.py index edf36a3..1ff22b9 100644 --- a/DGraph/distributed/commInfo.py +++ b/DGraph/distributed/commInfo.py @@ -62,30 +62,62 @@ def compute_halo_vertices( def compute_local_edge_list( global_edge_list: torch.Tensor, # [E, 2] - partitioning: torch.Tensor, # [V] (Acts as dst_partitioning) - local_vertices_global: torch.Tensor, # [num_local] - halo_vertices_global: torch.Tensor, # [num_halo] + partitioning: torch.Tensor, # [V_dst] (Acts as dst_partitioning) + local_vertices_global: torch.Tensor, # [num_local], dst vertex space + halo_vertices_global: torch.Tensor, # [num_halo], src vertex space rank: int, + src_partitioning: Optional[torch.Tensor] = None, ) -> torch.Tensor: + """ + Remaps a global edge list [E, 2] (col0=src/neighbor, col1=dst/local) into + local numbering: dst ids -> [0, num_local), src halo ids -> + [num_local, num_local + num_halo). + + For homogeneous graphs (src_partitioning=None), src and dst share one vertex + ID space and a single remap table is used for both columns, identical to the + original behavior. For bipartite/heterogeneous graphs, src_partitioning's + vertex space may be sized differently than partitioning's (dst) space, so two + independent remap tables are built. + """ num_local = local_vertices_global.size(0) num_halo = halo_vertices_global.size(0) - num_global = partitioning.size(0) + num_dst_global = partitioning.size(0) # FIX: Filter edges where the DESTINATION is owned by this rank (Index 1) local_edge_mask = partitioning[global_edge_list[:, 1]] == rank local_edges_global = global_edge_list[local_edge_mask] - # Build inverse map: global_id -> local_idx via scatter - g2l = torch.empty(num_global, dtype=torch.long, device=global_edge_list.device) - g2l.scatter_(0, local_vertices_global, torch.arange(num_local, device=g2l.device)) - g2l.scatter_( - 0, - halo_vertices_global, - torch.arange(num_local, num_local + num_halo, device=g2l.device), + # dst remap table: dst global id -> local index [0, num_local) + g2l_dst = torch.empty( + num_dst_global, dtype=torch.long, device=global_edge_list.device + ) + g2l_dst.scatter_(0, local_vertices_global, torch.arange(num_local, device=g2l_dst.device)) + + if src_partitioning is None: + # Homogeneous: src and dst share one ID space and one remap table. + g2l_dst.scatter_( + 0, + halo_vertices_global, + torch.arange(num_local, num_local + num_halo, device=g2l_dst.device), + ) + g2l_src = g2l_dst + else: + # Heterogeneous/bipartite: halo vertices live in the (differently-sized) + # src vertex space, so they need their own remap table. + num_src_global = src_partitioning.size(0) + g2l_src = torch.empty( + num_src_global, dtype=torch.long, device=global_edge_list.device + ) + g2l_src.scatter_( + 0, + halo_vertices_global, + torch.arange(num_local, num_local + num_halo, device=g2l_src.device), + ) + + # Remap to local numbering: col0 via src table, col1 via dst table. + local_edge_list = torch.stack( + [g2l_src[local_edges_global[:, 0]], g2l_dst[local_edges_global[:, 1]]], dim=1 ) - - # Remap to local numbering - local_edge_list = g2l[local_edges_global] return local_edge_list @@ -168,25 +200,52 @@ def build_communication_pattern( partitioning: torch.Tensor, rank: int, world_size: int, + src_partitioning: Optional[torch.Tensor] = None, ) -> CommunicationPattern: """ Args: - global_edge_list (torch.Tensor)): A tensor of shape [E, 2] - partitioning (torch.Tensor): A tensor of shape [V] + global_edge_list (torch.Tensor)): A tensor of shape [E, 2], col0=src/neighbor, + col1=dst/local (dst is always the aggregation target and is guaranteed local) + partitioning (torch.Tensor): A tensor of shape [V_dst]; the dst/local-vertex + partitioning rank (int): Rank of this process world_size (int): Total number of processes + src_partitioning (Optional[torch.Tensor]): A tensor of shape [V_src], the + partitioning of the (possibly differently-sized) neighbor/src vertex + space, for bipartite/heterogeneous relations. Defaults to `partitioning` + for homogeneous graphs. Returns: CommunicationPattern """ + src_part = partitioning if src_partitioning is None else src_partitioning + local_verts = compute_local_vertices(partitioning, rank) - halo_verts = compute_halo_vertices(global_edge_list, partitioning, rank) + src_local_verts = ( + local_verts + if src_partitioning is None + else compute_local_vertices(src_partitioning, rank) + ) + + halo_verts = compute_halo_vertices( + global_edge_list, src_part, rank, dst_partitioning=partitioning + ) local_edges = compute_local_edge_list( - global_edge_list, partitioning, local_verts, halo_verts, rank + global_edge_list, + partitioning, + local_verts, + halo_verts, + rank, + src_partitioning=src_partitioning, ) send_idx, send_off = compute_boundary_vertices( - global_edge_list, partitioning, local_verts, rank, world_size + global_edge_list, + src_part, + src_local_verts, + rank, + world_size, + dst_partitioning=partitioning, ) comm = compute_comm_map(send_off, world_size) recv_off, recv_back_off = compute_recv_offsets(comm, rank) diff --git a/experiments/GraphCast/data_utils/graphcast_graph.py b/experiments/GraphCast/data_utils/graphcast_graph.py index 1137559..89d6def 100644 --- a/experiments/GraphCast/data_utils/graphcast_graph.py +++ b/experiments/GraphCast/data_utils/graphcast_graph.py @@ -90,7 +90,13 @@ def build_graphcast_comm_patterns(graph: GraphCastTopology) -> GraphCastCommPatt src = vertex originating the message (neighbor) dst = vertex aggregating the message (central) - We swap into [central, neighbor] ordering for the comm pattern edge list. + DGraph.distributed.commInfo's CommunicationPattern contract guarantees + col1 of local_edge_list is always the locally-owned aggregation target + (dst/central); col0 is the neighbor (src), which may be a halo vertex. + We therefore keep the edge list in [src (neighbor), dst (central)] = + [col0, col1] order, i.e. unchanged message-flow order, and pass the two + endpoints' partitionings as (dst partitioning, src partitioning) for the + bipartite grid2mesh/mesh2grid relations. """ rank = graph.rank world_size = graph.ranks_per_graph @@ -98,31 +104,29 @@ def build_graphcast_comm_patterns(graph: GraphCastTopology) -> GraphCastCommPatt grid_part = graph.grid_rank_placement # --- grid2mesh --- - # Message flow: grid → mesh. Central = mesh, neighbor = grid. - # Swap: message-flow (grid, mesh) → comm (mesh, grid) = (central, neighbor). + # Message flow: grid → mesh. Central/dst = mesh, neighbor/src = grid. grid2mesh_edges = torch.stack( [ - graph.grid2mesh_graph_dst_indices, # mesh (central, col 0) - graph.grid2mesh_graph_src_indices, - ], # grid (neighbor, col 1) + graph.grid2mesh_graph_src_indices, # grid (neighbor, col 0) + graph.grid2mesh_graph_dst_indices, + ], # mesh (central, col 1) dim=1, ) grid2mesh_cp = build_communication_pattern( global_edge_list=grid2mesh_edges, partitioning=mesh_part, + src_partitioning=grid_part, rank=rank, world_size=world_size, ) # --- mesh ↔ mesh --- - # Homogeneous, undirected. Central = mesh, neighbor = mesh. - # Message-flow src/dst are both mesh — swap is identity, but we keep - # the convention: col 0 = dst (central), col 1 = src (neighbor). + # Homogeneous, undirected. Central/dst = mesh, neighbor/src = mesh. mesh_edges = torch.stack( [ - graph.mesh_graph_dst_indices, # mesh (central, col 0) - graph.mesh_graph_src_indices, - ], # mesh (neighbor, col 1) + graph.mesh_graph_src_indices, # mesh (neighbor, col 0) + graph.mesh_graph_dst_indices, + ], # mesh (central, col 1) dim=1, ) mesh_cp = build_communication_pattern( @@ -133,18 +137,18 @@ def build_graphcast_comm_patterns(graph: GraphCastTopology) -> GraphCastCommPatt ) # --- mesh2grid --- - # Message flow: mesh → grid. Central = grid, neighbor = mesh. - # Swap: message-flow (mesh, grid) → comm (grid, mesh) = (central, neighbor). + # Message flow: mesh → grid. Central/dst = grid, neighbor/src = mesh. mesh2grid_edges = torch.stack( [ - graph.mesh2grid_graph_dst_indices, # grid (central, col 0) - graph.mesh2grid_graph_src_indices, - ], # mesh (neighbor, col 1) + graph.mesh2grid_graph_src_indices, # mesh (neighbor, col 0) + graph.mesh2grid_graph_dst_indices, + ], # grid (central, col 1) dim=1, ) mesh2grid_cp = build_communication_pattern( global_edge_list=mesh2grid_edges, partitioning=grid_part, + src_partitioning=mesh_part, rank=rank, world_size=world_size, ) @@ -333,6 +337,21 @@ def get_grid_placement( def get_grid2mesh_graph(self, mesh_graph_dict: dict): + # SUSPECTED BUG (unverified -- flagged for review, not fixed here): this is + # named "mesh_vertex_rank_placement" but is actually + # mesh_graph_dict["mesh_vertex_renumbering"] = the SORT PERMUTATION from + # get_mesh_graph (renumbered_nodes: sorted_position -> original_vertex_id), + # not a rank-lookup-by-original-id array. get_grid_placement below then + # does mesh_vertex_rank_placement[grid2mesh_mesh_dst_indices], i.e. indexes + # the permutation by ORIGINAL mesh vertex ids (dst_mesh_indices, from + # create_grid2mesh_graph, pre-renumbering) -- this looks like it computes + # "original id of whichever vertex sorts to position == this original id", + # not "rank that owns this original vertex". The value actually needed + # (rank of an original-id vertex) isn't exposed by mesh_graph_dict at all; + # it would require the inverse of renumbered_nodes (computed locally inside + # get_mesh_graph as reverse_renumbered_nodes but never returned), composed + # with node_rank_placement. Not fixed pending confirmation this analysis is + # correct (untestable without GPU access to run and inspect actual values). mesh_vertex_rank_placement = mesh_graph_dict["mesh_vertex_renumbering"] max_edge_len = max_edge_length( self.finest_mesh_vertices, self.finest_mesh_src, self.finest_mesh_dst @@ -356,13 +375,23 @@ def get_grid2mesh_graph(self, mesh_graph_dict: dict): contigous_edge_mapping, renumbered_edges = torch.sort(meshtogrid_edge_placement) src_grid_indices = src_grid_indices[renumbered_edges] - grid_vertex_rank_placement = torch.zeros_like(lat_lon_grid_flat) + # NOTE: must be a 1D tensor of length num_grid_vertices, NOT + # torch.zeros_like(lat_lon_grid_flat) (shape [num_grid_vertices, 2]) -- + # the latter made torch.sort() below sort each row's 2 (duplicate) entries + # against each other instead of sorting across all grid vertices, silently + # corrupting grid_vertex_rank_placement/renumbered_grid into non-1D, + # non-permutation tensors that broke every downstream consumer (comm-pattern + # construction via build_communication_pattern, get_mesh2grid_edges's + # renumbered_grid[dst_grid_indices] lookup) at any world_size, including 1. + grid_vertex_rank_placement = torch.zeros( + lat_lon_grid_flat.shape[0], dtype=meshtogrid_edge_placement.dtype + ) for i, rank in enumerate(meshtogrid_edge_placement): loc = src_grid_indices[i] grid_vertex_rank_placement[loc] = rank continuous_grid_mapping, renumbered_grid = torch.sort( - grid_vertex_rank_placement + grid_vertex_rank_placement, dim=0 ) grid2mesh_graph_dict = { @@ -456,17 +485,47 @@ def get_graphcast_graph( comm_patterns = build_graphcast_comm_patterns(topology) + # --- Slice node/edge feature tensors down to this rank's local partition --- + # HaloExchange (DGraph/distributed/haloExchange.py) expects x_local of shape + # [num_local, F]: only this rank's locally-owned vertices/edges, row-aligned + # with comm_patterns.*.local_edge_list's local numbering. mesh_graph/ + # grid2mesh_graph/mesh2grid_graph's raw feature tensors are the FULL + # (renumbered-by-rank, but unfiltered) global set, so they must be filtered + # here with the exact same masks build_communication_pattern used internally + # (vertex: placement==rank; edge: dst-owned-by-rank) so rows stay aligned. + rank = topology.rank + + mesh_node_local_mask = mesh_vertex_rank_placement == rank + local_mesh_node_features = mesh_graph["node_features"][mesh_node_local_mask] + + mesh_edge_local_mask = mesh_vertex_rank_placement[mesh_graph["dst_indices"]] == rank + local_mesh_edge_features = mesh_graph["edge_features"][mesh_edge_local_mask] + + grid2mesh_edge_local_mask = ( + mesh_vertex_rank_placement[grid2mesh_graph["dst_indices"]] == rank + ) + local_grid2mesh_edge_features = grid2mesh_graph["edge_features"][ + grid2mesh_edge_local_mask + ] + + mesh2grid_edge_local_mask = ( + grid_vertex_rank_placement[mesh2grid_graph["dst_indices"]] == rank + ) + local_mesh2grid_edge_features = mesh2grid_graph["edge_features"][ + mesh2grid_edge_local_mask + ] + return DistributedGraphCastGraph( rank=self.rank, world_size=self.world_size, ranks_per_graph=self.ranks_per_graph, mesh_level=self.mesh_level, lat_lon_grid=self.lat_lon_grid, - mesh_graph_node_features=mesh_graph["node_features"], - mesh_graph_edge_features=mesh_graph["edge_features"], + mesh_graph_node_features=local_mesh_node_features, + mesh_graph_edge_features=local_mesh_edge_features, mesh2grid_graph_node_features=torch.tensor([]), grid2mesh_graph_node_features=grid2mesh_graph["node_features"], - mesh2grid_graph_edge_features=mesh2grid_graph["edge_features"], - grid2mesh_graph_edge_features=grid2mesh_graph["edge_features"], + mesh2grid_graph_edge_features=local_mesh2grid_edge_features, + grid2mesh_graph_edge_features=local_grid2mesh_edge_features, distributed_comm_patterns=comm_patterns, ) diff --git a/experiments/GraphCast/data_utils/preprocess.py b/experiments/GraphCast/data_utils/preprocess.py index 886452f..ff86c22 100644 --- a/experiments/GraphCast/data_utils/preprocess.py +++ b/experiments/GraphCast/data_utils/preprocess.py @@ -17,7 +17,7 @@ def partition_graph(G: nx.Graph, num_ranks: int): except ImportError: raise ImportError("Please install metis to use this function.") if num_ranks == 1: - return np.ones(len(G.nodes), dtype=int) + return np.zeros(len(G.nodes), dtype=int) if num_ranks < 1: raise ValueError("Number of ranks must be greater than 0.") diff --git a/experiments/GraphCast/dist_utils.py b/experiments/GraphCast/dist_utils.py index 48a59cf..a8d4252 100644 --- a/experiments/GraphCast/dist_utils.py +++ b/experiments/GraphCast/dist_utils.py @@ -22,6 +22,15 @@ def get_rank(self): def get_world_size(self): return self._world_size + def init_process_group(self, backend: str, **kwargs): + pass + + def barrier(self) -> None: + pass + + def destroy(self) -> None: + pass + def scatter( self, tensor: torch.Tensor, src: torch.Tensor, rank_mappings, num_local_nodes ): diff --git a/experiments/GraphCast/gen_mesh_partitions.py b/experiments/GraphCast/gen_mesh_partitions.py new file mode 100644 index 0000000..155f2c3 --- /dev/null +++ b/experiments/GraphCast/gen_mesh_partitions.py @@ -0,0 +1,83 @@ +import os + +import torch +from fire import Fire + +from data_utils.graphcast_graph import DistributedGraphCastGraphGenerator +from data_utils.icosahedral_mesh import ( + get_hierarchy_of_triangular_meshes_for_sphere, + merge_meshes, +) + + +def num_mesh_vertices(mesh_level: int) -> int: + """Vertex count of the merged multi-mesh at this level. + + Mirrors DistributedGraphCastGraphGenerator.__init__'s own mesh construction + rather than a closed-form formula, so it stays correct if mesh generation + changes. Cross-checked against tests/test_single_graph_data.py: mesh_level=6 + -> 40962. + """ + meshes = get_hierarchy_of_triangular_meshes_for_sphere(splits=mesh_level) + merged = merge_meshes(meshes) + return len(merged.vertices) + + +def _partition_path(output_dir: str, mesh_level: int, world_size: int) -> str: + return os.path.join( + output_dir, f"mesh_vertex_rank_placement_L{mesh_level}_W{world_size}.pt" + ) + + +def generate_partition( + mesh_level: int, world_size: int, output_dir: str, force: bool = False +) -> str: + """Compute (or reuse if cached) the mesh vertex rank placement for + (mesh_level, world_size), saving it to output_dir. Returns the saved path.""" + os.makedirs(output_dir, exist_ok=True) + path = _partition_path(output_dir, mesh_level, world_size) + + if os.path.exists(path) and not force: + print(f"[skip] {path} already exists") + return path + + if world_size == 1: + placement = torch.zeros(num_mesh_vertices(mesh_level), dtype=torch.long) + else: + placement = DistributedGraphCastGraphGenerator.get_mesh_graph_partition( + mesh_level=mesh_level, world_size=world_size + ) + + torch.save(placement, path) + print( + f"[saved] {path} (mesh_level={mesh_level}, world_size={world_size}, " + f"num_vertices={placement.numel()})" + ) + return path + + +def main( + mesh_levels: str = "2,5,6", + world_sizes: str = "1,2,4", + output_dir: str = "partitions", + force: bool = False, +) -> None: + """Generate/cache mesh partitions for the cross product of mesh_levels x world_sizes. + + Args: + mesh_levels: comma-separated list of mesh levels, e.g. "2,5,6". + world_sizes: comma-separated list of world sizes, e.g. "1,2,4". + output_dir: directory to write mesh_vertex_rank_placement_L{L}_W{W}.pt files to. + force: recompute and overwrite even if a cached file already exists. + """ + levels = [int(x) for x in mesh_levels.split(",")] + sizes = [int(x) for x in world_sizes.split(",")] + + for mesh_level in levels: + for world_size in sizes: + generate_partition(mesh_level, world_size, output_dir, force=force) + + +if __name__ == "__main__": + # Usage: python gen_mesh_partitions.py --mesh_levels 2,5,6 --world_sizes 1,2,4 --output_dir partitions/ + Fire(main) diff --git a/experiments/GraphCast/inference_benchmark.py b/experiments/GraphCast/inference_benchmark.py new file mode 100644 index 0000000..b9548b5 --- /dev/null +++ b/experiments/GraphCast/inference_benchmark.py @@ -0,0 +1,448 @@ +import os +import sys + +import numpy as np +import torch +import torch.distributed as dist +from fire import Fire + +_GRAPHCAST_DIR = os.path.dirname(os.path.abspath(__file__)) +_EXPERIMENTS_DIR = os.path.dirname(_GRAPHCAST_DIR) +_COST_MODEL_DIR = os.path.join(_EXPERIMENTS_DIR, "cost_model_benchmarks") +sys.path.insert(0, _COST_MODEL_DIR) + +from benchmarks.common import ( # noqa: E402 + collect_metadata, + setup_distributed, + write_result, +) + +from DGraph.Communicator import Communicator # noqa: E402 +from DGraph.utils.TimingReport import TimingReport # noqa: E402 + +from data_utils.graphcast_graph import DistributedGraphCastGraphGenerator # noqa: E402 +from data_utils.utils import padded_size # noqa: E402 +from dataset import SyntheticWeatherDataset # noqa: E402 +from dist_utils import SingleProcessDummyCommunicator # noqa: E402 +from graphcast_config import Config # noqa: E402 +from model import DGraphCast # noqa: E402 + +GRID_LAT, GRID_LON = 721, 1440 +NUM_GRID_POINTS = GRID_LAT * GRID_LON + + +def build_comm(mode: str, backend: str, world_size: int): + """mode="single" -> SingleProcessDummyCommunicator (no real halo exchange). + mode="distributed" -> real DGraph Communicator (NCCL). Both branches assume + torch.distributed is already initialized (see setup_distributed() in main()).""" + if mode == "single": + return SingleProcessDummyCommunicator() + return Communicator.init_process_group(backend, ranks_per_graph=world_size) + + +def load_partition(partition_dir: str, mesh_level: int, world_size: int) -> torch.Tensor: + path = os.path.join( + partition_dir, f"mesh_vertex_rank_placement_L{mesh_level}_W{world_size}.pt" + ) + if not os.path.exists(path): + raise FileNotFoundError( + f"Missing partition file {path}. Generate it first with:\n" + f" python gen_mesh_partitions.py --mesh_levels {mesh_level} " + f"--world_sizes {world_size} --output_dir {partition_dir}" + ) + return torch.load(path) + + +def build_dataset( + cfg: Config, + mesh_level: int, + partition: torch.Tensor, + rank: int, + world_size: int, + device: torch.device, +) -> SyntheticWeatherDataset: + return SyntheticWeatherDataset( + channels=list(range(cfg.model.input_grid_dim)), + num_samples_per_year=2, + num_steps=1, + mesh_vertex_placement=partition, + mesh_level=mesh_level, + grid_size=(GRID_LAT, GRID_LON), + rank=rank, + world_size=world_size, + ranks_per_graph=world_size, + device=device, + ) + + +def summarize(trials_ms) -> dict: + if not trials_ms: + return {"mean": None, "std": None, "p50": None, "p95": None, "trials_ms": []} + arr = np.asarray(trials_ms, dtype=np.float64) + return { + "mean": float(arr.mean()), + "std": float(arr.std()), + "p50": float(np.percentile(arr, 50)), + "p95": float(np.percentile(arr, 95)), + "trials_ms": [float(x) for x in trials_ms], + } + + +def collect_phase_breakdown(measure_iters: int, processor_layers: int) -> dict: + """Reads DGraph.utils.TimingReport.TimingReport._timers, already populated by + model.py's existing TimingReport(...) blocks during run_forward_loop -- no new + instrumentation needed. Model control flow is deterministic (no data-dependent + branching), so the last `measure_iters` entries per name correspond exactly to + the measured (post-warmup) iterations.""" + names = [ + "model/embed", + "model/encode", + "model/process", + "model/decode", + "model/final_prediction", + "encoder/halo_exchange", + "encoder/edge_block", + "encoder/node_block", + "encoder/grid_mlp", + "decoder/halo_exchange", + "decoder/edge_block", + "decoder/node_block", + ] + for i in range(processor_layers): + names += [ + f"processor/layer_{i}/halo_exchange", + f"processor/layer_{i}/edge_block", + f"processor/layer_{i}/node_block", + ] + + breakdown = {} + for name in names: + values = TimingReport._timers.get(name, []) + values = values[-measure_iters:] if measure_iters > 0 else values + breakdown[name] = summarize(values) + return breakdown + + +def run_forward_loop( + model: torch.nn.Module, + in_data: torch.Tensor, + static_graph, + comm, + warmup_iters: int, + measure_iters: int, + device: torch.device, +) -> dict: + model.eval() + e2e_trials_ms = [] + oom = False + + try: + for i in range(warmup_iters + measure_iters): + if i == warmup_iters: + torch.cuda.reset_peak_memory_stats(device) + + comm.barrier() + start_evt = torch.cuda.Event(enable_timing=True) + end_evt = torch.cuda.Event(enable_timing=True) + start_evt.record() + with torch.no_grad(): + pred = model(in_data, static_graph) + end_evt.record() + torch.cuda.synchronize() + comm.barrier() + + if i >= warmup_iters: + e2e_trials_ms.append(start_evt.elapsed_time(end_evt)) + except RuntimeError as e: + if "out of memory" in str(e).lower(): + oom = True + else: + raise + + peak_memory_bytes = None if oom else torch.cuda.max_memory_allocated(device) + return { + "e2e_trials_ms": e2e_trials_ms, + "peak_memory_bytes": peak_memory_bytes, + "oom": oom, + } + + +def run_correctness_check( + cfg: Config, + mesh_level: int, + comm, + partition_dir: str, + rank: int, + world_size: int, + device: torch.device, + seed: int, + atol: float, + rtol: float, +) -> bool: + """Two-step correctness gate for world_size > 1 runs. + + Step 1 (structural): SyntheticWeatherDataset.__getitem__ shards grid input/ + output via a naive contiguous chunk, independent of + DistributedGraphCastGraphGenerator's connectivity-weighted grid_part voting + used to build the mesh2grid CommunicationPattern. If their per-rank local + grid vertex counts disagree, GraphCastDecoder silently operates on the + wrong slice of grid vertices. This step catches that mismatch and halts + before the (more expensive, more failure-prone) Step 2. + + Step 2 (value-level, only if Step 1 passes): every rank independently + builds an identical, full (world_size=1, SingleProcessDummyCommunicator) + reference forward pass -- redundant across ranks, but necessary because + DGraph.distributed.commInfo.compute_comm_map's dist.all_gather is a + collective on the REAL process group and must be called identically by + every rank to avoid deadlock. The reference and distributed runs use + different internal vertex renumberings (each DistributedGraphCastGraphGenerator + call re-derives its own grid renumbering for its own world_size), so both + are mapped back to a common "original grid index" frame before comparing. + """ + partition = load_partition(partition_dir, mesh_level, world_size) + + num_local_expected = padded_size(NUM_GRID_POINTS, world_size) // world_size + + dataset = build_dataset(cfg, mesh_level, partition, rank, world_size, device) + static_graph = dataset.get_static_graph() + num_local_actual = static_graph.distributed_comm_patterns.mesh2grid.num_local_vertices + + gathered = [None] * world_size + dist.all_gather_object(gathered, (rank, num_local_expected, num_local_actual)) + step1_ok = all(e == a for (_, e, a) in gathered) + + if rank == 0: + if step1_ok: + print("[correctness] STEP 1 PASS: dataset chunk size vs grid_part local " + "vertex count agree on every rank.") + else: + mismatches = [(r, e, a) for (r, e, a) in gathered if e != a] + print(f"[correctness] STEP 1 FAIL: per-rank (rank, expected, actual) " + f"mismatches: {mismatches}") + print("[correctness] Halting -- do not trust world_size>1 benchmark " + "numbers until this is resolved (see plan's known dataset.py vs " + "grid_part mismatch).") + dist.barrier() + if not step1_ok: + return False + + # --- Step 2: value-level check --- + lat_lon_grid = dataset.lat_lon_grid + + # Recover this run's grid renumbering (original grid index for each + # renumbered/rank-sorted slot) without re-triggering comm-pattern construction. + dist_generator = DistributedGraphCastGraphGenerator( + lat_lon_grid, mesh_level=mesh_level, ranks_per_graph=world_size, + rank=rank, world_size=world_size, + ) + dist_mesh_graph = dist_generator.get_mesh_graph(partition) + dist_grid2mesh = dist_generator.get_grid2mesh_graph(dist_mesh_graph) + dist_renumbered_grid = dist_grid2mesh["renumbered_grid"] # [N_grid] -> original grid idx + dist_grid_placement = dist_grid2mesh["grid_vertex_rank_placement"] # sorted by rank + + local_mask = dist_grid_placement == rank + local_original_grid_idx = dist_renumbered_grid[local_mask] + + # Every rank independently builds an identical world_size=1 reference (all + # ranks call the same dist collectives the same number of times -> no deadlock). + torch.manual_seed(seed) + ref_comm = SingleProcessDummyCommunicator() + ref_generator = DistributedGraphCastGraphGenerator( + lat_lon_grid, mesh_level=mesh_level, ranks_per_graph=1, rank=0, world_size=1 + ) + ref_placement = torch.zeros(ref_generator.mesh_vertices.shape[0], dtype=torch.long) + ref_static_graph = ref_generator.get_graphcast_graph( + mesh_vertex_rank_placement=ref_placement + ) + ref_mesh_graph = ref_generator.get_mesh_graph(ref_placement) + ref_grid2mesh = ref_generator.get_grid2mesh_graph(ref_mesh_graph) + ref_renumbered_grid = ref_grid2mesh["renumbered_grid"] + + torch.manual_seed(seed) + ref_model = DGraphCast(cfg, ref_comm).to(device).eval() + + canonical_input = torch.randn(NUM_GRID_POINTS, cfg.model.input_grid_dim) + ref_input = canonical_input[ref_renumbered_grid].unsqueeze(0).to(device) + + with torch.no_grad(): + ref_pred = ref_model(ref_input, ref_static_graph) # [N_grid, C], ref-renumbered order + + ref_pred_original_order = torch.zeros_like(ref_pred) + ref_pred_original_order[ref_renumbered_grid.to(device)] = ref_pred + + expected_local = ref_pred_original_order[local_original_grid_idx.to(device)] + + torch.manual_seed(seed) + dist_model = DGraphCast(cfg, comm).to(device).eval() + dist_input = canonical_input[local_original_grid_idx].unsqueeze(0).to(device) + + with torch.no_grad(): + actual_local = dist_model(dist_input, static_graph) + + close = torch.allclose(actual_local, expected_local, atol=atol, rtol=rtol) + max_abs_diff = (actual_local - expected_local).abs().max().item() + denom = expected_local.abs().clamp_min(1e-8) + max_rel_diff = ((actual_local - expected_local).abs() / denom).max().item() + + gathered2 = [None] * world_size + dist.all_gather_object(gathered2, (rank, bool(close), max_abs_diff, max_rel_diff)) + dist.barrier() + + if rank == 0: + all_close = all(c for (_, c, _, _) in gathered2) + if all_close: + print(f"[correctness] STEP 2 PASS: all ranks match reference within " + f"atol={atol}, rtol={rtol}. Per-rank (rank, close, max_abs_diff, " + f"max_rel_diff): {gathered2}") + else: + print(f"[correctness] STEP 2 FAIL. Per-rank (rank, close, max_abs_diff, " + f"max_rel_diff): {gathered2}") + return all_close + return close + + +def main( + mode: str = "distributed", + mesh_level: int = 6, + backend: str = "nccl", + batch_size: int = 1, + warmup_iters: int = 10, + measure_iters: int = 50, + partition_dir: str = "partitions", + output_dir: str = "benchmark_results", + seed: int = 42, + correctness: bool = False, + correctness_atol: float = 1e-4, + correctness_rtol: float = 1e-3, +) -> None: + assert mode in ("single", "distributed"), "mode must be 'single' or 'distributed'" + assert batch_size == 1, "only batch_size=1 is supported (matches GraphCast's one-step inference)" + assert not correctness or mode == "distributed", ( + "--correctness checks cross-rank partitioning; run it with --mode distributed " + "(and --nproc_per_node >= 2), not --mode single" + ) + + rank, world_size, local_rank = setup_distributed() + device = torch.device(f"cuda:{local_rank}") + + comm = build_comm(mode, backend, world_size) + + cfg = Config() + cfg.model.mesh_level = mesh_level + + # model.py's DGraphCast.forward unconditionally uses TimingReport(...) context + # managers, which raise if TimingReport.init() hasn't been called -- required + # before ANY forward pass, in both the correctness-check and benchmark paths. + TimingReport.init(comm) + + if correctness: + if world_size < 2: + if rank == 0: + print("[correctness] world_size < 2: nothing to check (single-GPU " + "baseline has no cross-rank partitioning to verify).") + return + ok = run_correctness_check( + cfg, mesh_level, comm, partition_dir, rank, world_size, device, + seed, correctness_atol, correctness_rtol, + ) + if rank == 0: + print(f"[correctness] Overall: {'PASS' if ok else 'FAIL'}") + return + + partition = load_partition(partition_dir, mesh_level, world_size) + dataset = build_dataset(cfg, mesh_level, partition, rank, world_size, device) + static_graph = dataset.get_static_graph() + + torch.manual_seed(seed) + model = DGraphCast(cfg, comm).to(device).eval() + + sample = dataset[0] + in_data = sample["invar"].unsqueeze(0).to(device) + + result = run_forward_loop( + model, in_data, static_graph, comm, warmup_iters, measure_iters, device + ) + + # Use a collective to agree on OOM across ranks: if only one rank OOMs (e.g. a + # slightly imbalanced METIS partition), letting ranks diverge onto different + # code paths here would deadlock the next collective call (all_gather_object + # below) on whichever ranks didn't OOM and are still waiting for it. + oom_flags = [None] * world_size + dist.all_gather_object(oom_flags, result["oom"]) + any_oom = any(oom_flags) + + if any_oom: + if rank == 0: + print(f"[inference_benchmark] OOM at mode={mode} mesh_level={mesh_level} " + f"world_size={world_size} (per-rank oom={oom_flags}) -- infeasible " + f"at this scale on this hardware, not a benchmark failure.") + payload = { + "benchmark": "graphcast_inference", + "metadata": collect_metadata(), + "config": { + "mode": mode, "backend": backend, "world_size": world_size, + "mesh_level": mesh_level, "hidden_dim": cfg.model.hidden_dim, + "processor_layers": cfg.model.processor_layers, + "batch_size": batch_size, "warmup_iters": warmup_iters, + "measure_iters": measure_iters, "seed": seed, + }, + "measurements": [{ + "params": {"world_size": world_size, "mesh_level": mesh_level}, + "oom": True, + "oom_per_rank": oom_flags, + }], + } + write_result( + os.path.join(output_dir, f"graphcast_infer_{mode}_W{world_size}_L{mesh_level}.json"), + payload, + ) + return + + latency_summary = summarize(result["e2e_trials_ms"]) + phase_breakdown = collect_phase_breakdown(measure_iters, cfg.model.processor_layers) + + peak_memory_list = [None] * world_size + dist.all_gather_object(peak_memory_list, result["peak_memory_bytes"]) + + trials_list = [None] * world_size + dist.all_gather_object(trials_list, result["e2e_trials_ms"]) + + if rank == 0: + p50_s = (latency_summary["p50"] or 0.0) / 1000.0 + throughput = (NUM_GRID_POINTS * batch_size / p50_s) if p50_s > 0 else None + + payload = { + "benchmark": "graphcast_inference", + "metadata": collect_metadata(), + "config": { + "mode": mode, "backend": backend, "world_size": world_size, + "mesh_level": mesh_level, "hidden_dim": cfg.model.hidden_dim, + "processor_layers": cfg.model.processor_layers, + "batch_size": batch_size, "warmup_iters": warmup_iters, + "measure_iters": measure_iters, "seed": seed, + }, + "measurements": [{ + "params": {"world_size": world_size, "mesh_level": mesh_level}, + "end_to_end_latency_ms": latency_summary, + "throughput_grid_points_per_sec": throughput, + "peak_memory_bytes": { + "per_rank": peak_memory_list, + "max": max(m for m in peak_memory_list if m is not None), + }, + "phase_breakdown_ms": phase_breakdown, + "per_rank_end_to_end_trials_ms": trials_list, + }], + } + write_result( + os.path.join(output_dir, f"graphcast_infer_{mode}_W{world_size}_L{mesh_level}.json"), + payload, + ) + print( + f"[inference_benchmark] mode={mode} world_size={world_size} " + f"mesh_level={mesh_level}: p50={latency_summary['p50']:.2f}ms " + f"throughput={throughput:.1f} grid-points/sec" + ) + + +if __name__ == "__main__": + Fire(main) diff --git a/experiments/GraphCast/model.py b/experiments/GraphCast/model.py index 9183ad8..ff3224e 100644 --- a/experiments/GraphCast/model.py +++ b/experiments/GraphCast/model.py @@ -136,7 +136,9 @@ def __init__(self, cfg: Config, comm, *args, **kwargs) -> None: output_node_dim=hidden_dim, hidden_dim=hidden_dim, ) - self.grid_node_mlp = MeshGraphMLP(input_dim=hidden_dim, output_dim=hidden_dim) + self.grid_node_mlp = MeshGraphMLP( + input_dim=hidden_dim, output_dim=hidden_dim, hidden_dim=hidden_dim + ) def forward( self, @@ -145,10 +147,11 @@ def forward( grid2mesh_edge_features: Tensor, comm_pattern: CommunicationPattern, ) -> Tuple[Tensor, Tensor]: - # local_edge_list: [E, 2] with [central=mesh, neighbor=grid/halo] + # local_edge_list: [E, 2] with [neighbor=grid/halo, central=mesh] (col1 is + # always the locally-owned aggregation target per CommunicationPattern's contract) edge_index = comm_pattern.local_edge_list - dst_indices = edge_index[:, 0] # mesh (central, aggregation target) - src_indices = edge_index[:, 1] # grid/halo (neighbor, message source) + src_indices = edge_index[:, 0] # grid/halo (neighbor, message source) + dst_indices = edge_index[:, 1] # mesh (central, aggregation target) num_local = comm_pattern.num_local_vertices with TimingReport("encoder/halo_exchange"): @@ -227,10 +230,11 @@ def forward( e_feats = embedded_mesh2mesh_edge_features n_feats = embedded_mesh_features - # local_edge_list: [E, 2] with [central=mesh_dst, neighbor=mesh_src] + # local_edge_list: [E, 2] with [neighbor=mesh_src, central=mesh_dst] (col1 is + # always the locally-owned aggregation target per CommunicationPattern's contract) edge_index = comm_pattern.local_edge_list - dst_indices = edge_index[:, 0] # central (aggregation target) - src_indices = edge_index[:, 1] # neighbor (message source) + src_indices = edge_index[:, 0] # neighbor (message source) + dst_indices = edge_index[:, 1] # central (aggregation target) num_local = comm_pattern.num_local_vertices for i, (edge_layer, node_layer) in enumerate( @@ -308,10 +312,11 @@ def forward( Returns: (Tensor): The updated grid node features """ - # local_edge_list: [E, 2] with [central=grid, neighbor=mesh/halo] + # local_edge_list: [E, 2] with [neighbor=mesh/halo, central=grid] (col1 is + # always the locally-owned aggregation target per CommunicationPattern's contract) edge_index = comm_pattern.local_edge_list - dst_indices = edge_index[:, 0] # grid (central, aggregation target) - src_indices = edge_index[:, 1] # mesh/halo (neighbor, message source) + src_indices = edge_index[:, 0] # mesh/halo (neighbor, message source) + dst_indices = edge_index[:, 1] # grid (central, aggregation target) num_local = comm_pattern.num_local_vertices with TimingReport("decoder/halo_exchange"): @@ -355,7 +360,9 @@ def __init__(self, cfg: Config, comm, *args, **kwargs): self.processor = GraphCastProcessor(cfg=cfg, comm=comm, *args, **kwargs) self.decoder = GraphCastDecoder(cfg=cfg, comm=comm, *args, **kwargs) self.final_prediction = MeshGraphMLP( - input_dim=self.hidden_dim, output_dim=self.output_grid_dim + input_dim=self.hidden_dim, + output_dim=self.output_grid_dim, + hidden_dim=self.hidden_dim, ) def forward( diff --git a/experiments/GraphCast/plot_scaling.py b/experiments/GraphCast/plot_scaling.py new file mode 100644 index 0000000..f1be288 --- /dev/null +++ b/experiments/GraphCast/plot_scaling.py @@ -0,0 +1,336 @@ +import glob +import json +import os + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 +from fire import Fire # noqa: E402 + +# Fixed categorical color/marker assignment by world_size, matching the tab10-based +# convention already used in experiments/cost_model_benchmarks/visualization/plot_crossover.py. +# World_size=1 (single-GPU baseline) is always drawn the same way so it reads as the +# reference line across every plot. +WORLD_SIZE_STYLE = { + 1: {"color": "#1f77b4", "marker": "o", "label": "1 GPU (baseline)"}, + 2: {"color": "#ff7f0e", "marker": "s", "label": "2 GPUs"}, + 4: {"color": "#2ca02c", "marker": "^", "label": "4 GPUs"}, + 8: {"color": "#9467bd", "marker": "D", "label": "8 GPUs"}, +} +OOM_COLOR = "#d62728" + + +def load_results(paths) -> list: + """Load and flatten measurements across multiple result JSON files, tagging + each with its run config.""" + records = [] + for path in paths: + with open(path) as f: + payload = json.load(f) + config = payload["config"] + for m in payload["measurements"]: + records.append({"config": config, "measurement": m}) + return records + + +def _style(world_size: int) -> dict: + return WORLD_SIZE_STYLE.get( + world_size, + {"color": "#7f7f7f", "marker": "x", "label": f"{world_size} GPUs"}, + ) + + +def _filter(records, mesh_level: int) -> list: + return sorted( + [r for r in records if r["config"]["mesh_level"] == mesh_level], + key=lambda r: r["config"]["world_size"], + ) + + +def plot_latency_vs_world_size(records, mesh_level: int, output_path: str) -> None: + rows = _filter(records, mesh_level) + if not rows: + return + + plt.figure(figsize=(8, 6), dpi=150) + plt.grid(True, which="both", ls="-", alpha=0.2) + + world_sizes, medians, p95s, oom_ws = [], [], [], [] + for r in rows: + m = r["measurement"] + ws = r["config"]["world_size"] + if m.get("oom"): + oom_ws.append(ws) + continue + world_sizes.append(ws) + medians.append(m["end_to_end_latency_ms"]["p50"]) + p95s.append(m["end_to_end_latency_ms"]["p95"]) + + if world_sizes: + colors = [_style(ws)["color"] for ws in world_sizes] + plt.errorbar( + world_sizes, + medians, + yerr=[np.array(p95s) - np.array(medians), np.zeros(len(medians))], + fmt="none", + ecolor="#7f7f7f", + alpha=0.5, + capsize=4, + ) + for ws, med in zip(world_sizes, medians): + style = _style(ws) + plt.scatter([ws], [med], color=style["color"], marker=style["marker"], + s=100, zorder=3, label=style["label"]) + plt.plot(world_sizes, medians, color="#7f7f7f", linestyle="-", linewidth=1, + alpha=0.6, zorder=2) + + for ws in oom_ws: + plt.axvline(x=ws, color=OOM_COLOR, linestyle="--", alpha=0.6) + plt.text(ws, plt.ylim()[1] * 0.9 if plt.ylim()[1] else 1, "OOM", + color=OOM_COLOR, rotation=90, va="top", ha="right") + + plt.title(f"GraphCast Inference Latency vs World Size (mesh_level={mesh_level})") + plt.xlabel("World size (# GPUs)") + plt.ylabel("End-to-end latency, p50 (ms)") + all_ws = sorted(set(world_sizes) | set(oom_ws)) + plt.xticks(all_ws, [str(w) for w in all_ws]) + plt.legend(loc="best", framealpha=0.9) + plt.tight_layout() + plt.savefig(output_path) + plt.close() + print(f"[plot_scaling] saved {output_path}") + + +def plot_memory_vs_world_size(records, mesh_level: int, output_path: str) -> None: + rows = _filter(records, mesh_level) + if not rows: + return + + plt.figure(figsize=(8, 6), dpi=150) + plt.grid(True, which="both", ls="-", alpha=0.2) + + world_sizes, peak_gb, oom_ws = [], [], [] + for r in rows: + m = r["measurement"] + ws = r["config"]["world_size"] + if m.get("oom"): + oom_ws.append(ws) + continue + world_sizes.append(ws) + peak_gb.append(m["peak_memory_bytes"]["max"] / (1024 ** 3)) + + for ws, mem in zip(world_sizes, peak_gb): + style = _style(ws) + plt.scatter([ws], [mem], color=style["color"], marker=style["marker"], + s=100, zorder=3, label=style["label"]) + if world_sizes: + plt.plot(world_sizes, peak_gb, color="#7f7f7f", linestyle="-", linewidth=1, + alpha=0.6, zorder=2) + + for ws in oom_ws: + plt.axvline(x=ws, color=OOM_COLOR, linestyle="--", alpha=0.6) + + plt.title(f"GraphCast Peak GPU Memory vs World Size (mesh_level={mesh_level})") + plt.xlabel("World size (# GPUs)") + plt.ylabel("Peak memory across ranks (GB)") + all_ws = sorted(set(world_sizes) | set(oom_ws)) + plt.xticks(all_ws, [str(w) for w in all_ws]) + plt.legend(loc="best", framealpha=0.9) + plt.tight_layout() + plt.savefig(output_path) + plt.close() + print(f"[plot_scaling] saved {output_path}") + + +PHASE_GROUPS = { + "embed": ["model/embed"], + "encoder": ["encoder/halo_exchange", "encoder/edge_block", "encoder/node_block", "encoder/grid_mlp"], + "processor": None, # filled dynamically: all processor/layer_*/* keys + "decoder": ["decoder/halo_exchange", "decoder/edge_block", "decoder/node_block"], + "final_prediction": ["model/final_prediction"], +} +PHASE_COLORS = { + "embed": "#1f77b4", + "encoder": "#ff7f0e", + "processor": "#2ca02c", + "decoder": "#9467bd", + "final_prediction": "#8c564b", +} + + +def plot_phase_breakdown_stacked(records, mesh_level: int, output_path: str) -> None: + rows = _filter(records, mesh_level) + rows = [r for r in rows if not r["measurement"].get("oom")] + if not rows: + return + + plt.figure(figsize=(9, 6), dpi=150) + world_sizes = [r["config"]["world_size"] for r in rows] + x = np.arange(len(world_sizes)) + + bottom = np.zeros(len(rows)) + for group_name, fixed_keys in PHASE_GROUPS.items(): + heights = [] + for r in rows: + breakdown = r["measurement"]["phase_breakdown_ms"] + if fixed_keys is not None: + keys = fixed_keys + else: + keys = [k for k in breakdown if k.startswith("processor/layer_")] + total = sum((breakdown.get(k) or {}).get("mean") or 0.0 for k in keys) + heights.append(total) + heights = np.array(heights) + plt.bar(x, heights, bottom=bottom, color=PHASE_COLORS[group_name], + label=group_name, width=0.6, edgecolor="white", linewidth=0.5) + bottom += heights + + plt.title(f"GraphCast Per-Phase Time Breakdown (mesh_level={mesh_level})") + plt.xlabel("World size (# GPUs)") + plt.ylabel("Mean time per phase (ms)") + plt.xticks(x, [str(w) for w in world_sizes]) + plt.legend(loc="best", framealpha=0.9) + plt.tight_layout() + plt.savefig(output_path) + plt.close() + print(f"[plot_scaling] saved {output_path}") + + +def plot_throughput_vs_world_size(records, mesh_level: int, output_path: str) -> None: + rows = _filter(records, mesh_level) + rows = [r for r in rows if not r["measurement"].get("oom")] + if not rows: + return + + plt.figure(figsize=(8, 6), dpi=150) + plt.grid(True, which="both", ls="-", alpha=0.2) + + world_sizes = [r["config"]["world_size"] for r in rows] + throughput = [r["measurement"]["throughput_grid_points_per_sec"] for r in rows] + + for ws, tp in zip(world_sizes, throughput): + style = _style(ws) + plt.scatter([ws], [tp], color=style["color"], marker=style["marker"], + s=100, zorder=3, label=style["label"]) + plt.plot(world_sizes, throughput, color="#7f7f7f", linestyle="-", linewidth=1, + alpha=0.6, zorder=2, label="Actual") + + if throughput and world_sizes[0] == 1: + ideal = [throughput[0] * ws for ws in world_sizes] + plt.plot(world_sizes, ideal, color="#7f7f7f", linestyle="--", linewidth=1, + alpha=0.8, label="Ideal linear scaling") + + plt.title(f"GraphCast Inference Throughput vs World Size (mesh_level={mesh_level})") + plt.xlabel("World size (# GPUs)") + plt.ylabel("Throughput (grid points / sec)") + plt.xticks(world_sizes, [str(w) for w in world_sizes]) + plt.legend(loc="best", framealpha=0.9) + plt.tight_layout() + plt.savefig(output_path) + plt.close() + print(f"[plot_scaling] saved {output_path}") + + +def plot_crossover_vs_mesh_level(records, output_path: str) -> None: + """bench_crossover.py-style plot: latency vs mesh_level, one line per world_size + (including the world_size=1 baseline), with sign-change interpolation marking + where each world_size first beats the baseline.""" + world_sizes = sorted({r["config"]["world_size"] for r in records}) + if 1 not in world_sizes or len(world_sizes) < 2: + return + + plt.figure(figsize=(9, 6), dpi=150) + plt.grid(True, which="both", ls="-", alpha=0.2) + + baseline_by_level = {} + for r in records: + if r["config"]["world_size"] == 1 and not r["measurement"].get("oom"): + baseline_by_level[r["config"]["mesh_level"]] = r["measurement"]["end_to_end_latency_ms"]["p50"] + + for ws in world_sizes: + rows = sorted( + [r for r in records if r["config"]["world_size"] == ws and not r["measurement"].get("oom")], + key=lambda r: r["config"]["mesh_level"], + ) + if not rows: + continue + levels = [r["config"]["mesh_level"] for r in rows] + latencies = [r["measurement"]["end_to_end_latency_ms"]["p50"] for r in rows] + style = _style(ws) + plt.plot(levels, latencies, color=style["color"], marker=style["marker"], + linewidth=2, label=style["label"]) + + if ws != 1: + for i in range(1, len(levels)): + l0, l1 = levels[i - 1], levels[i] + if l0 not in baseline_by_level or l1 not in baseline_by_level: + continue + diff_prev = latencies[i - 1] - baseline_by_level[l0] + diff_curr = latencies[i] - baseline_by_level[l1] + if diff_prev * diff_curr < 0: + m_ws = (latencies[i] - latencies[i - 1]) / (l1 - l0) + m_base = (baseline_by_level[l1] - baseline_by_level[l0]) / (l1 - l0) + if m_ws != m_base: + x_cross = l0 + (baseline_by_level[l0] - latencies[i - 1]) / (m_ws - m_base) + y_cross = latencies[i - 1] + m_ws * (x_cross - l0) + plt.plot(x_cross, y_cross, marker="*", color=style["color"], + markersize=15, zorder=5) + plt.annotate( + f"{ws} GPUs beat baseline\n~mesh_level {x_cross:.1f}", + xy=(x_cross, y_cross), xytext=(10, 20), + textcoords="offset points", fontsize=9, fontweight="bold", + color=style["color"], + arrowprops=dict(arrowstyle="->", color=style["color"]), + ) + break + + plt.title("GraphCast Distributed vs Single-GPU Crossover") + plt.xlabel("mesh_level") + plt.ylabel("End-to-end latency, p50 (ms)") + plt.yscale("log") + plt.legend(loc="best", framealpha=0.9) + plt.tight_layout() + plt.savefig(output_path) + plt.close() + print(f"[plot_scaling] saved {output_path}") + + +def main(inputs: str, output_dir: str = "figures") -> None: + """inputs: comma-separated list and/or glob patterns of result JSON paths.""" + os.makedirs(output_dir, exist_ok=True) + + paths = [] + for part in inputs.split(","): + part = part.strip() + matched = glob.glob(part) + paths.extend(matched if matched else [part]) + paths = sorted(set(paths)) + if not paths: + print(f"[plot_scaling] no files matched inputs={inputs!r}") + return + + records = load_results(paths) + mesh_levels = sorted({r["config"]["mesh_level"] for r in records}) + + for mesh_level in mesh_levels: + plot_latency_vs_world_size( + records, mesh_level, os.path.join(output_dir, f"latency_vs_world_size_L{mesh_level}.png") + ) + plot_memory_vs_world_size( + records, mesh_level, os.path.join(output_dir, f"memory_vs_world_size_L{mesh_level}.png") + ) + plot_phase_breakdown_stacked( + records, mesh_level, os.path.join(output_dir, f"phase_breakdown_L{mesh_level}.png") + ) + plot_throughput_vs_world_size( + records, mesh_level, os.path.join(output_dir, f"throughput_vs_world_size_L{mesh_level}.png") + ) + + plot_crossover_vs_mesh_level(records, os.path.join(output_dir, "crossover_vs_mesh_level.png")) + + +if __name__ == "__main__": + # Usage: + # python plot_scaling.py --inputs "benchmark_results/*.json" --output_dir figures/ + Fire(main) diff --git a/experiments/GraphCast/tests/conftest.py b/experiments/GraphCast/tests/conftest.py index 06e40b1..44050e6 100644 --- a/experiments/GraphCast/tests/conftest.py +++ b/experiments/GraphCast/tests/conftest.py @@ -11,6 +11,10 @@ def setup_data(): prev_dir = os.path.dirname(cur_dir) sys.path.append(prev_dir) from dataset import SyntheticWeatherDataset + from data_utils.icosahedral_mesh import ( + get_hierarchy_of_triangular_meshes_for_sphere, + merge_meshes, + ) latlon_res = (721, 1440) num_samples_per_year_train = 2 @@ -21,16 +25,27 @@ def setup_data(): start_year = 1980 use_time_of_year_index = True channels_list = [i for i in range(num_channels_climate)] + mesh_level = 6 cos_zenith_args = { "dt": dt, "start_year": start_year, } batch_size = 1 + + # world_size=1 placement: every mesh vertex is owned by rank 0. Computed by + # mirroring DistributedGraphCastGraphGenerator's own mesh construction rather + # than a hardcoded vertex count, so it stays correct if mesh_level changes. + _meshes = get_hierarchy_of_triangular_meshes_for_sphere(splits=mesh_level) + _merged_mesh = merge_meshes(_meshes) + mesh_vertex_placement = torch.zeros(len(_merged_mesh.vertices), dtype=torch.long) + test_dataset = SyntheticWeatherDataset( channels=channels_list, num_samples_per_year=num_samples_per_year_train, num_steps=1, + mesh_vertex_placement=mesh_vertex_placement, + mesh_level=mesh_level, grid_size=latlon_res, cos_zenith_args=cos_zenith_args, batch_size=batch_size, diff --git a/experiments/GraphCast/tests/test_single_graph_data.py b/experiments/GraphCast/tests/test_single_graph_data.py index 8b330d7..ead622a 100644 --- a/experiments/GraphCast/tests/test_single_graph_data.py +++ b/experiments/GraphCast/tests/test_single_graph_data.py @@ -17,17 +17,22 @@ def test_static_graph_data(setup_data): static_graph = _dataset.get_static_graph() - # These values are obtained from the original paper + # These values are obtained from the original paper. static_graph exposes + # aggregated node/edge features (not raw src/dst index tensors); per-edge-type + # connectivity is instead available via distributed_comm_patterns' local_edge_list, + # which at world_size=1 contains every edge (no cross-rank partitioning). assert static_graph.mesh_level == 6 assert static_graph.mesh_graph_node_features.shape == torch.Size([40962, 3]) assert static_graph.mesh_graph_edge_features.shape == torch.Size([655320, 4]) - assert static_graph.mesh_graph_src_indices.shape == torch.Size([655320]) - assert static_graph.mesh_graph_dst_indices.shape == torch.Size([655320]) + + comm_patterns = static_graph.distributed_comm_patterns + assert comm_patterns.mesh.local_edge_list.shape == torch.Size([655320, 2]) + assert comm_patterns.mesh.num_local_vertices == 40962 assert static_graph.grid2mesh_graph_edge_features.shape == torch.Size([1618824, 4]) - assert static_graph.grid2mesh_graph_src_indices.shape == torch.Size([1618824]) - assert static_graph.grid2mesh_graph_dst_indices.shape == torch.Size([1618824]) + assert comm_patterns.grid2mesh.local_edge_list.shape == torch.Size([1618824, 2]) + assert comm_patterns.grid2mesh.num_local_vertices == 40962 assert static_graph.mesh2grid_graph_edge_features.shape == torch.Size([3114720, 4]) - assert static_graph.mesh2grid_graph_src_indices.shape == torch.Size([3114720]) - assert static_graph.mesh2grid_graph_dst_indices.shape == torch.Size([3114720]) + assert comm_patterns.mesh2grid.local_edge_list.shape == torch.Size([3114720, 2]) + assert comm_patterns.mesh2grid.num_local_vertices == 721 * 1440 From ee4a4b54109e59301ef59a648769899c94ec71e2 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 17 Jul 2026 09:48:55 -0400 Subject: [PATCH 16/39] Add distributed comm pattern test for graphcast --- .../tests/test_distributed_comm_patterns.py | 90 +++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 experiments/GraphCast/tests/test_distributed_comm_patterns.py diff --git a/experiments/GraphCast/tests/test_distributed_comm_patterns.py b/experiments/GraphCast/tests/test_distributed_comm_patterns.py new file mode 100644 index 0000000..d4cb320 --- /dev/null +++ b/experiments/GraphCast/tests/test_distributed_comm_patterns.py @@ -0,0 +1,90 @@ +""" +Distributed regression test for GraphCast's comm-pattern construction. + +This test must be launched with >1 process so it actually exercises cross-rank +edges (world_size=1 has no halo vertices and would trivially pass even with a +broken bipartite partitioning). Run with, e.g.: + + torchrun --standalone --nproc_per_node 2 -m pytest \ + experiments/GraphCast/tests/test_distributed_comm_patterns.py + +It asserts the core invariant CommunicationPattern guarantees (see +DGraph/distributed/commInfo.py): for every edge type's local_edge_list, +column 1 (the aggregation target / central vertex) is always a locally-owned +vertex, i.e. `local_edge_list[:, 1] < num_local_vertices`. This is exactly the +invariant that GraphCast's data_utils/graphcast_graph.py violated before it was +fixed to match commInfo.py's (col0=neighbor/src, col1=central/dst) convention, +and would have caught that regression immediately. +""" + +import os +import sys + +import pytest +import torch + +_CUR_DIR = os.path.dirname(__file__) +_GRAPHCAST_DIR = os.path.dirname(_CUR_DIR) +sys.path.append(_GRAPHCAST_DIR) + +from data_utils.graphcast_graph import DistributedGraphCastGraphGenerator # noqa: E402 + + +def _requires_multi_rank(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required to build GraphCast comm patterns (commInfo.py's " + "compute_comm_map hardcodes .cuda() collectives).") + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + pytest.skip( + "torch.distributed process group is not initialized. Launch this test " + "via torchrun --nproc_per_node 2 -m pytest ... (see module docstring)." + ) + if torch.distributed.get_world_size() < 2: + pytest.skip( + "This test needs world_size >= 2 to exercise cross-rank edges; " + "run via torchrun --nproc_per_node 2." + ) + + +def test_local_edge_list_dst_column_is_always_local(): + _requires_multi_rank() + + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + torch.cuda.set_device(rank) + + mesh_level = 2 # small mesh for a fast test (162 mesh vertices) + + mesh_vertex_placement = DistributedGraphCastGraphGenerator.get_mesh_graph_partition( + mesh_level=mesh_level, world_size=world_size + ) + + latitudes = torch.linspace(-90, 90, steps=21) + longitudes = torch.linspace(-180, 180, steps=41)[1:] + lat_lon_grid = torch.stack( + torch.meshgrid(latitudes, longitudes, indexing="ij"), dim=-1 + ) + + graph = DistributedGraphCastGraphGenerator( + lat_lon_grid, + mesh_level=mesh_level, + ranks_per_graph=world_size, + rank=rank, + world_size=world_size, + ).get_graphcast_graph(mesh_vertex_rank_placement=mesh_vertex_placement) + + comm_patterns = graph.distributed_comm_patterns + for name, cp in ( + ("grid2mesh", comm_patterns.grid2mesh), + ("mesh", comm_patterns.mesh), + ("mesh2grid", comm_patterns.mesh2grid), + ): + dst_col = cp.local_edge_list[:, 1] + assert torch.all(dst_col < cp.num_local_vertices), ( + f"{name}: local_edge_list[:, 1] must always index a locally-owned " + f"vertex (< num_local_vertices={cp.num_local_vertices}), got max " + f"{dst_col.max().item() if dst_col.numel() else 'N/A'}" + ) + assert cp.num_local_vertices > 0, f"{name}: no local vertices on rank {rank}" + + torch.distributed.barrier() From 804b1e784cb15e577064a3a573baf38c8cbb9e7c Mon Sep 17 00:00:00 2001 From: Shehtab Date: Fri, 17 Jul 2026 14:30:18 -0400 Subject: [PATCH 17/39] Updated to fix the allocation issue on graphcast --- .../GraphCast/data_utils/graphcast_graph.py | 171 +++++++++--------- experiments/GraphCast/preprocess.py | 15 +- .../benchmarks/bench_end_to_end.py | 79 +------- 3 files changed, 105 insertions(+), 160 deletions(-) diff --git a/experiments/GraphCast/data_utils/graphcast_graph.py b/experiments/GraphCast/data_utils/graphcast_graph.py index 89d6def..29dd683 100644 --- a/experiments/GraphCast/data_utils/graphcast_graph.py +++ b/experiments/GraphCast/data_utils/graphcast_graph.py @@ -223,12 +223,12 @@ def get_mesh_graph_partition(mesh_level: int, world_size: int): @staticmethod def get_grid_vertex_partition( - lat: int, - lon: int, + num_grid: int, mesh_vertex_rank_placement: torch.Tensor, grid2mesh_grid_src_indices: torch.Tensor, grid2mesh_mesh_dst_indices: torch.Tensor, mesh2grid_mesh_src_indices: torch.Tensor, + mesh2grid_grid_dst_indices: torch.Tensor, world_size: int, ) -> torch.Tensor: """Generate the partitioning of grid vertices to minimize cross-rank edges. @@ -237,11 +237,9 @@ def get_grid_vertex_partition( (via both grid2mesh and mesh2grid edges) live on each rank, then assigns the grid vertex to the rank with the plurality of connections. - mesh2grid grid destinations are implicit: grid vertex i owns edges - [3i, 3i+1, 3i+2] since create_mesh2grid_graph assigns exactly 3 - edges (one face's vertices) per grid vertex. + All indices/placements here are in the *original* (pre rank-sort) mesh + and grid vertex numbering. """ - num_grid = lat * lon votes = torch.zeros(num_grid, world_size, dtype=torch.long) # --- grid2mesh contribution: grid vertex is src, mesh vertex is dst --- @@ -251,10 +249,8 @@ def get_grid_vertex_partition( votes.view(-1).scatter_add_(0, g2m_flat_idx, torch.ones_like(g2m_flat_idx)) # --- mesh2grid contribution: mesh vertex is src, grid vertex is dst --- - # Each grid vertex i has exactly 3 mesh2grid edges at positions [3i, 3i+1, 3i+2] - m2g_grid_dst = torch.arange(num_grid, dtype=torch.long).repeat_interleave(3) m2g_ranks = mesh_vertex_rank_placement[mesh2grid_mesh_src_indices.long()] - m2g_flat_idx = m2g_grid_dst * world_size + m2g_ranks + m2g_flat_idx = mesh2grid_grid_dst_indices.long() * world_size + m2g_ranks votes.view(-1).scatter_add_(0, m2g_flat_idx, torch.ones_like(m2g_flat_idx)) # Assign each grid vertex to the rank with the most connections @@ -262,6 +258,21 @@ def get_grid_vertex_partition( return grid_partitioning + @staticmethod + def _renumber_by_rank(rank_placement: torch.Tensor): + """Sort vertices by rank so each rank owns a contiguous numbering range. + + Returns (sorted_ranks, forward_perm, inverse_perm): forward_perm maps + new (sorted) position -> original vertex id, and is used to reorder + per-vertex feature tensors. inverse_perm maps original vertex id -> + new position, and is used to remap edge-index tensors that still + reference original vertex ids into the new numbering. + """ + sorted_ranks, forward_perm = torch.sort(rank_placement) + inverse_perm = torch.empty_like(forward_perm) + inverse_perm[forward_perm] = torch.arange(forward_perm.numel()) + return sorted_ranks, forward_perm, inverse_perm + def get_mesh_graph(self, mesh_vertex_rank_placement: torch.Tensor): """Get the graph for the distributed graphcast graph.""" @@ -280,17 +291,13 @@ def get_mesh_graph(self, mesh_vertex_rank_placement: torch.Tensor): assert num_nodes == mesh_vertex_rank_placement.size(0) - contiguous_rank_mapping, renumbered_nodes = torch.sort( - mesh_vertex_rank_placement + contiguous_rank_mapping, renumbered_nodes, reverse_renumbered_nodes = ( + self._renumber_by_rank(mesh_vertex_rank_placement) ) # renumber the nodes node_features = node_features[renumbered_nodes] - reverse_renumbered_nodes = torch.zeros_like(renumbered_nodes) - - reverse_renumbered_nodes[renumbered_nodes] = torch.arange(num_nodes) - # renumber the edges new_src_indices = reverse_renumbered_nodes[src_indices] new_dst_indices = reverse_renumbered_nodes[dst_indices] @@ -320,45 +327,40 @@ def get_mesh_graph(self, mesh_vertex_rank_placement: torch.Tensor): "edge_rank_placement": contigous_edge_mapping, "src_rank_placement": src_indices_rank_placement, "dst_rank_placement": dst_indices_rank_placement, + "original_rank_placement": mesh_vertex_rank_placement, "mesh_vertex_renumbering": renumbered_nodes, "renumbered_vertices": renumbered_nodes, + "inverse_vertex_renumbering": reverse_renumbered_nodes, } return mesh_graph_dict - def get_grid_placement( - self, mesh_vertex_rank_placement, grid2mesh_mesh_dst_indices - ): - meshtogrid_edge_placement = mesh_vertex_rank_placement[ - grid2mesh_mesh_dst_indices - ] + def _get_raw_mesh2grid_edges(self): + """mesh2grid edge_features/src(mesh)/dst(grid), all in original vertex ids.""" + lat_lon_grid_flat = self.lat_lon_grid.permute(2, 0, 1).view(2, -1).permute(1, 0) + return create_mesh2grid_graph( + lat_lon_grid_flat, self.mesh_vertices, self.mesh_faces + ) + + def get_grid2mesh_graph(self, mesh_graph_dict: dict, mesh2grid_raw: tuple = None): + """Get the grid2mesh bipartite graph, renumbered into per-rank-contiguous order. + + Each grid vertex is assigned to the rank holding the plurality of its + connected mesh vertices, counting both grid2mesh and mesh2grid edges + (see get_grid_vertex_partition) -- not just the grid2mesh direction. + + mesh2grid_raw lets callers that already computed the mesh2grid edges + (get_graphcast_graph) pass them in so the mesh2grid nearest-neighbor + search isn't repeated; if omitted (e.g. external callers that only need + the grid2mesh graph) it is computed here. + """ + inverse_vertex_renumbering = mesh_graph_dict["inverse_vertex_renumbering"] + original_mesh_rank_placement = mesh_graph_dict["original_rank_placement"] - return meshtogrid_edge_placement - - def get_grid2mesh_graph(self, mesh_graph_dict: dict): - - # SUSPECTED BUG (unverified -- flagged for review, not fixed here): this is - # named "mesh_vertex_rank_placement" but is actually - # mesh_graph_dict["mesh_vertex_renumbering"] = the SORT PERMUTATION from - # get_mesh_graph (renumbered_nodes: sorted_position -> original_vertex_id), - # not a rank-lookup-by-original-id array. get_grid_placement below then - # does mesh_vertex_rank_placement[grid2mesh_mesh_dst_indices], i.e. indexes - # the permutation by ORIGINAL mesh vertex ids (dst_mesh_indices, from - # create_grid2mesh_graph, pre-renumbering) -- this looks like it computes - # "original id of whichever vertex sorts to position == this original id", - # not "rank that owns this original vertex". The value actually needed - # (rank of an original-id vertex) isn't exposed by mesh_graph_dict at all; - # it would require the inverse of renumbered_nodes (computed locally inside - # get_mesh_graph as reverse_renumbered_nodes but never returned), composed - # with node_rank_placement. Not fixed pending confirmation this analysis is - # correct (untestable without GPU access to run and inspect actual values). - mesh_vertex_rank_placement = mesh_graph_dict["mesh_vertex_renumbering"] max_edge_len = max_edge_length( self.finest_mesh_vertices, self.finest_mesh_src, self.finest_mesh_dst ) - renumbered_vertices = mesh_graph_dict["node_rank_placement"] - # create the grid2mesh bipartite graph lat_lon_grid_flat = self.lat_lon_grid.permute(2, 0, 1).view(2, -1).permute(1, 0) @@ -367,67 +369,58 @@ def get_grid2mesh_graph(self, mesh_graph_dict: dict): ) edge_features, src_grid_indices, dst_mesh_indices = g2m_graph - meshtogrid_edge_placement = self.get_grid_placement( - mesh_vertex_rank_placement, dst_mesh_indices + if mesh2grid_raw is None: + mesh2grid_raw = self._get_raw_mesh2grid_edges() + _, m2g_src_mesh_indices, m2g_dst_grid_indices = mesh2grid_raw + + num_grid = lat_lon_grid_flat.shape[0] + grid_vertex_rank_placement = self.get_grid_vertex_partition( + num_grid=num_grid, + mesh_vertex_rank_placement=original_mesh_rank_placement, + grid2mesh_grid_src_indices=src_grid_indices, + grid2mesh_mesh_dst_indices=dst_mesh_indices, + mesh2grid_mesh_src_indices=m2g_src_mesh_indices, + mesh2grid_grid_dst_indices=m2g_dst_grid_indices, + world_size=self.world_size, ) - dst_mesh_indices = renumbered_vertices[dst_mesh_indices] - - contigous_edge_mapping, renumbered_edges = torch.sort(meshtogrid_edge_placement) - - src_grid_indices = src_grid_indices[renumbered_edges] - # NOTE: must be a 1D tensor of length num_grid_vertices, NOT - # torch.zeros_like(lat_lon_grid_flat) (shape [num_grid_vertices, 2]) -- - # the latter made torch.sort() below sort each row's 2 (duplicate) entries - # against each other instead of sorting across all grid vertices, silently - # corrupting grid_vertex_rank_placement/renumbered_grid into non-1D, - # non-permutation tensors that broke every downstream consumer (comm-pattern - # construction via build_communication_pattern, get_mesh2grid_edges's - # renumbered_grid[dst_grid_indices] lookup) at any world_size, including 1. - grid_vertex_rank_placement = torch.zeros( - lat_lon_grid_flat.shape[0], dtype=meshtogrid_edge_placement.dtype + continuous_grid_mapping, renumbered_grid, inverse_grid_renumbering = ( + self._renumber_by_rank(grid_vertex_rank_placement) ) - for i, rank in enumerate(meshtogrid_edge_placement): - loc = src_grid_indices[i] - grid_vertex_rank_placement[loc] = rank - continuous_grid_mapping, renumbered_grid = torch.sort( - grid_vertex_rank_placement, dim=0 - ) + # Remap edge endpoints from original ids into the new, rank-contiguous + # numbering: grid side via the grid renumbering, mesh side via the + # mesh renumbering computed in get_mesh_graph. + src_grid_indices = inverse_grid_renumbering[src_grid_indices.long()] + dst_mesh_indices = inverse_vertex_renumbering[dst_mesh_indices.long()] grid2mesh_graph_dict = { "node_features": torch.tensor([]), "edge_features": edge_features, "src_indices": src_grid_indices, "dst_indices": dst_mesh_indices, - "grid2mesh_edge_rank_placement": contigous_edge_mapping, "grid_vertex_rank_placement": continuous_grid_mapping, "renumbered_grid": renumbered_grid, + "inverse_grid_renumbering": inverse_grid_renumbering, } return grid2mesh_graph_dict def get_mesh2grid_edges( self, - grid_vertex_rank_placement, - renumbered_vertices, - renumbered_grid, + inverse_vertex_renumbering, + inverse_grid_renumbering, + mesh2grid_raw: tuple = None, ): - lat_lon_grid_flat = self.lat_lon_grid.permute(2, 0, 1).view(2, -1).permute(1, 0) + if mesh2grid_raw is None: + mesh2grid_raw = self._get_raw_mesh2grid_edges() + edge_features, src_mesh_indices, dst_grid_indices = mesh2grid_raw - m2g_graph = create_mesh2grid_graph( - lat_lon_grid_flat, self.mesh_vertices, self.mesh_faces - ) - - edge_features, src_mesh_indices, dst_grid_indices = m2g_graph - src_mesh_indices = renumbered_vertices[src_mesh_indices] - dst_grid_indices = renumbered_grid[dst_grid_indices] - - mesh2grid_edge_rank_placement = grid_vertex_rank_placement[dst_grid_indices] + src_mesh_indices = inverse_vertex_renumbering[src_mesh_indices.long()] + dst_grid_indices = inverse_grid_renumbering[dst_grid_indices.long()] mesh2grid_graph_dict = { "edge_features": edge_features, "src_indices": src_mesh_indices, "dst_indices": dst_grid_indices, - "mesh2grid_edge_rank_placement": mesh2grid_edge_rank_placement, } return mesh2grid_graph_dict @@ -460,13 +453,21 @@ def get_graphcast_graph( mesh_graph = self.get_mesh_graph(mesh_vertex_rank_placement) mesh_vertex_rank_placement = mesh_graph["node_rank_placement"] - renumbered_vertices = mesh_graph["renumbered_vertices"] - grid2mesh_graph = self.get_grid2mesh_graph(mesh_graph) + inverse_vertex_renumbering = mesh_graph["inverse_vertex_renumbering"] + + # Computed once and shared between get_grid2mesh_graph (needs it for the + # grid rank vote) and get_mesh2grid_edges (needs it for the edges + # themselves), so the mesh2grid nearest-neighbor search isn't repeated. + mesh2grid_raw = self._get_raw_mesh2grid_edges() + + grid2mesh_graph = self.get_grid2mesh_graph(mesh_graph, mesh2grid_raw=mesh2grid_raw) grid_vertex_rank_placement = grid2mesh_graph["grid_vertex_rank_placement"] - renumbered_grid = grid2mesh_graph["renumbered_grid"] + inverse_grid_renumbering = grid2mesh_graph["inverse_grid_renumbering"] mesh2grid_graph = self.get_mesh2grid_edges( - grid_vertex_rank_placement, renumbered_vertices, renumbered_grid + inverse_vertex_renumbering, + inverse_grid_renumbering, + mesh2grid_raw=mesh2grid_raw, ) topology = GraphCastTopology( diff --git a/experiments/GraphCast/preprocess.py b/experiments/GraphCast/preprocess.py index e860da0..93fc342 100644 --- a/experiments/GraphCast/preprocess.py +++ b/experiments/GraphCast/preprocess.py @@ -1,4 +1,3 @@ -import metis import torch import numpy as np import pickle @@ -15,8 +14,18 @@ def partition_graph(G: nx.Graph, num_ranks: int): - metis_graph = metis.networkx_to_metis(G) - (edgecuts, node_rank_placement) = metis.part_graph(metis_graph, nparts=num_ranks) + if num_ranks == 1: + return np.zeros(G.number_of_nodes(), dtype=int) + + try: + import pymetis + except ImportError: + raise ImportError( + "Please install pymetis to use this function (`pip install pymetis`)." + ) + + adjacency = [sorted(G.neighbors(node)) for node in range(G.number_of_nodes())] + (edgecuts, node_rank_placement) = pymetis.part_graph(num_ranks, adjacency=adjacency) # Node_rank_placement is of shape (num_nodes, ), where each element is the # rank of the node in the partitioning. diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py index 4cfe162..fa56f5d 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py @@ -57,78 +57,12 @@ partition_balanced, partition_metis, partition_random, + build_local_comm_pattern, ) from benchmarks.nn_layer_common import GCNLayer, EdgeConditionedLayer - -class MinimalHaloExchange(torch.autograd.Function): - """Forward: gather boundary features → all_to_all → populate recv buffer. - Backward: reverse the transfer to accumulate gradients. - """ - - @staticmethod - def forward(ctx, x_local, send_idx_flat, send_counts, recv_counts, world_size): - # Gather send buffer - send_buf = x_local[send_idx_flat] # [total_send, F] - - # Split by destination rank - send_list = list(send_buf.split(send_counts, dim=0)) - recv_list = [ - torch.zeros( - rc, x_local.shape[1], dtype=x_local.dtype, device=x_local.device - ) - for rc in recv_counts - ] - - dist.all_to_all(recv_list, send_list) - - recv_buf = ( - torch.cat(recv_list, dim=0) - if sum(recv_counts) > 0 - else torch.zeros(0, x_local.shape[1], device=x_local.device) - ) - - ctx.save_for_backward(send_idx_flat) - ctx.send_counts = send_counts - ctx.recv_counts = recv_counts - ctx.world_size = world_size - ctx.n_local = x_local.shape[0] - ctx.feature_dim = x_local.shape[1] - ctx.device = x_local.device - - return recv_buf - - @staticmethod - def backward(ctx, grad_recv): - (send_idx_flat,) = ctx.saved_tensors - send_counts = ctx.send_counts - recv_counts = ctx.recv_counts - world_size = ctx.world_size - n_local = ctx.n_local - F = ctx.feature_dim - device = ctx.device - - # Reverse: recv_counts become send_counts and vice versa - grad_recv_list = ( - list(grad_recv.split(recv_counts, dim=0)) - if grad_recv.shape[0] > 0 - else [torch.zeros(0, F, device=device)] * world_size - ) - grad_send_list = [torch.zeros(sc, F, device=device) for sc in send_counts] - - dist.all_to_all(grad_send_list, grad_recv_list) - - grad_send = torch.cat(grad_send_list, dim=0) - - # Scatter-add back to local vertices - grad_x_local = torch.zeros(n_local, F, device=device, dtype=grad_recv.dtype) - grad_x_local.scatter_add_( - 0, - send_idx_flat.unsqueeze(1).expand_as(grad_send), - grad_send, - ) - return grad_x_local, None, None, None, None - +from DGraph.distributed import HaloExchange, CommunicationPattern +from DGraph import Communicator # =========================================================================== # Main @@ -212,12 +146,13 @@ def main(): else None ) + comm = Communicator(backend="nccl") + halo_exchange = HaloExchange(comm=comm) + # --- Timed forward + backward --- def one_layer(): # Forward halo exchange - recv_buf = MinimalHaloExchange.apply( - x_local, send_idx_flat, send_counts, recv_counts, world_size - ) + recv_buf = halo_exchange(x_local, comm_pattern=pattern) # Augment: local + halo x_aug = torch.cat([x_local, recv_buf], dim=0) # Message passing From d22c19418b5d473b63600435bd2cec97673f2adc Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 17 Jul 2026 14:33:53 -0400 Subject: [PATCH 18/39] Fixing the incorrect encoder tensor --- experiments/GraphCast/model.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/experiments/GraphCast/model.py b/experiments/GraphCast/model.py index ff3224e..1f61574 100644 --- a/experiments/GraphCast/model.py +++ b/experiments/GraphCast/model.py @@ -155,13 +155,13 @@ def forward( num_local = comm_pattern.num_local_vertices with TimingReport("encoder/halo_exchange"): - halo_features = self.exchanger(mesh_node_features, comm_pattern) - augmented = torch.cat([mesh_node_features, halo_features], dim=0) + halo_features = self.exchanger(grid_node_features, comm_pattern) + augmented_grid = torch.cat([grid_node_features, halo_features], dim=0) with TimingReport("encoder/edge_block"): e_feats = self.edge_mlp( - src_node_features=augmented, - dst_node_features=augmented, + src_node_features=augmented_grid, + dst_node_features=mesh_node_features, edge_features=grid2mesh_edge_features, src_indices=src_indices, dst_indices=dst_indices, @@ -169,7 +169,7 @@ def forward( with TimingReport("encoder/node_block"): n_feats = self.mesh_node_mlp( - node_features=augmented[:num_local], + node_features=mesh_node_features, edge_features=e_feats, src_indices=dst_indices, ) From 145ec6322143653a612427bff6f47f4cbeed7e9f Mon Sep 17 00:00:00 2001 From: Shehtab Date: Fri, 17 Jul 2026 15:02:15 -0400 Subject: [PATCH 19/39] Bipartite fix with running single GPU run working --- DGraph/distributed/commInfo.py | 46 ++++++++++++---- DGraph/utils.py | 34 ------------ .../GraphCast/data_utils/graphcast_graph.py | 55 +++++++++++++++++++ experiments/GraphCast/dataset.py | 11 +++- experiments/GraphCast/inference_benchmark.py | 7 ++- 5 files changed, 107 insertions(+), 46 deletions(-) delete mode 100644 DGraph/utils.py diff --git a/DGraph/distributed/commInfo.py b/DGraph/distributed/commInfo.py index 1ff22b9..0abd259 100644 --- a/DGraph/distributed/commInfo.py +++ b/DGraph/distributed/commInfo.py @@ -67,17 +67,24 @@ def compute_local_edge_list( halo_vertices_global: torch.Tensor, # [num_halo], src vertex space rank: int, src_partitioning: Optional[torch.Tensor] = None, + src_local_vertices_global: Optional[torch.Tensor] = None, # [num_src_local], src space ) -> torch.Tensor: """ Remaps a global edge list [E, 2] (col0=src/neighbor, col1=dst/local) into - local numbering: dst ids -> [0, num_local), src halo ids -> - [num_local, num_local + num_halo). + local numbering. + + The dst column (col1) maps to ``[0, num_local)``. The src column (col0) maps + into the augmented src feature buffer ``[src_local ; halo]`` that the halo + exchange produces: locally-owned src vertices -> ``[0, num_src_local)`` and + halo (remote) src vertices -> ``[num_src_local, num_src_local + num_halo)``. For homogeneous graphs (src_partitioning=None), src and dst share one vertex - ID space and a single remap table is used for both columns, identical to the - original behavior. For bipartite/heterogeneous graphs, src_partitioning's - vertex space may be sized differently than partitioning's (dst) space, so two - independent remap tables are built. + ID space (num_src_local == num_local) and a single remap table is used for + both columns. For bipartite/heterogeneous graphs, src_partitioning's vertex + space may be sized differently than partitioning's (dst) space, so a separate + src remap table is built; it must cover BOTH local and halo src vertices, + since a bipartite relation can have on-rank src vertices (e.g. every edge is + intra-rank at world_size=1, so num_halo=0 and every src is local). """ num_local = local_vertices_global.size(0) num_halo = halo_vertices_global.size(0) @@ -94,7 +101,8 @@ def compute_local_edge_list( g2l_dst.scatter_(0, local_vertices_global, torch.arange(num_local, device=g2l_dst.device)) if src_partitioning is None: - # Homogeneous: src and dst share one ID space and one remap table. + # Homogeneous: src and dst share one ID space and one remap table. Local + # src vertices are already mapped to [0, num_local) via g2l_dst above. g2l_dst.scatter_( 0, halo_vertices_global, @@ -102,16 +110,33 @@ def compute_local_edge_list( ) g2l_src = g2l_dst else: - # Heterogeneous/bipartite: halo vertices live in the (differently-sized) - # src vertex space, so they need their own remap table. + # Heterogeneous/bipartite: src vertices live in the (differently-sized) + # src vertex space and need their own remap table covering both the + # locally-owned src vertices and the halo (remote) src vertices, matching + # the augmented src feature buffer layout [src_local ; halo]. + if src_local_vertices_global is None: + raise ValueError( + "src_local_vertices_global is required for bipartite/heterogeneous " + "edge lists (src_partitioning is not None)." + ) num_src_global = src_partitioning.size(0) + num_src_local = src_local_vertices_global.size(0) g2l_src = torch.empty( num_src_global, dtype=torch.long, device=global_edge_list.device ) + # local src -> [0, num_src_local) + g2l_src.scatter_( + 0, + src_local_vertices_global, + torch.arange(num_src_local, device=g2l_src.device), + ) + # halo src -> [num_src_local, num_src_local + num_halo) g2l_src.scatter_( 0, halo_vertices_global, - torch.arange(num_local, num_local + num_halo, device=g2l_src.device), + torch.arange( + num_src_local, num_src_local + num_halo, device=g2l_src.device + ), ) # Remap to local numbering: col0 via src table, col1 via dst table. @@ -238,6 +263,7 @@ def build_communication_pattern( halo_verts, rank, src_partitioning=src_partitioning, + src_local_vertices_global=src_local_verts, ) send_idx, send_off = compute_boundary_vertices( global_edge_list, diff --git a/DGraph/utils.py b/DGraph/utils.py deleted file mode 100644 index a61c55c..0000000 --- a/DGraph/utils.py +++ /dev/null @@ -1,34 +0,0 @@ -# Copyright (c) 2014-2024, Lawrence Livermore National Security, LLC. -# Produced at the Lawrence Livermore National Laboratory. -# Written by the LBANN Research Team (B. Van Essen, et al.) listed in -# the CONTRIBUTORS file. See the top-level LICENSE file for details. -# -# LLNL-CODE-697807. -# All rights reserved. -# -# This file is part of LBANN: Livermore Big Artificial Neural Network -# Toolkit. For details, see http://software.llnl.gov/LBANN or -# https://github.com/LBANN and https://github.com/LLNL/LBANN. -# -# SPDX-License-Identifier: (Apache-2.0) -import torch.distributed as dist - - -def largest_split(global_size, world_size): - return (global_size + world_size) // world_size - - -def split_per_rank(global_size, current_rank, world_size): - _split = largest_split(global_size, world_size) - if current_rank != world_size - 1: - return _split - else: - return global_size - (current_rank * _split) - - -def try_barrier(): - """Attempt a barrier but ignore any exceptions""" - try: - dist.barrier() - except: - pass diff --git a/experiments/GraphCast/data_utils/graphcast_graph.py b/experiments/GraphCast/data_utils/graphcast_graph.py index 29dd683..c58c5cb 100644 --- a/experiments/GraphCast/data_utils/graphcast_graph.py +++ b/experiments/GraphCast/data_utils/graphcast_graph.py @@ -82,6 +82,61 @@ class DistributedGraphCastGraph: distributed_comm_patterns: GraphCastCommPatterns +def _move_communication_pattern_to_device( + cp: CommunicationPattern, device +) -> CommunicationPattern: + """In-place move of every tensor field of a CommunicationPattern to device. + + build_communication_pattern leaves its outputs on a mix of devices (e.g. + compute_comm_map runs .cuda() collectives, while the index tensors stay on + CPU). The halo exchange indexes local features and runs NCCL all-to-all with + these tensors, so they must all sit on the same device as the node/edge + features. + """ + for field_name in ( + "local_edge_list", + "send_local_idx", + "send_offset", + "recv_offset", + "comm_map", + "put_forward_remote_offset", + "put_backward_remote_offset", + ): + value = getattr(cp, field_name) + if isinstance(value, torch.Tensor): + setattr(cp, field_name, value.to(device)) + return cp + + +def move_graphcast_graph_to_device( + graph: "DistributedGraphCastGraph", device +) -> "DistributedGraphCastGraph": + """In-place move of a DistributedGraphCastGraph (features + comm patterns) to device. + + The graph is constructed on CPU (nearest-neighbor search, METIS, etc.); this + lifts all of its tensors onto the compute device so the model can consume it + directly. Returns the same object for convenience. + """ + for field_name in ( + "lat_lon_grid", + "mesh_graph_node_features", + "mesh_graph_edge_features", + "mesh2grid_graph_node_features", + "grid2mesh_graph_node_features", + "mesh2grid_graph_edge_features", + "grid2mesh_graph_edge_features", + ): + value = getattr(graph, field_name) + if isinstance(value, torch.Tensor): + setattr(graph, field_name, value.to(device)) + + cps = graph.distributed_comm_patterns + _move_communication_pattern_to_device(cps.grid2mesh, device) + _move_communication_pattern_to_device(cps.mesh, device) + _move_communication_pattern_to_device(cps.mesh2grid, device) + return graph + + def build_graphcast_comm_patterns(graph: GraphCastTopology) -> GraphCastCommPatterns: """ Build CommunicationPatterns for all three GraphCast edge types. diff --git a/experiments/GraphCast/dataset.py b/experiments/GraphCast/dataset.py index b15e700..65046d8 100644 --- a/experiments/GraphCast/dataset.py +++ b/experiments/GraphCast/dataset.py @@ -16,7 +16,10 @@ import time from typing import Any, Dict, List, Optional, Tuple, Union from torch.utils.data import Dataset -from data_utils.graphcast_graph import DistributedGraphCastGraphGenerator +from data_utils.graphcast_graph import ( + DistributedGraphCastGraphGenerator, + move_graphcast_graph_to_device, +) from data_utils.utils import padded_size from torch.nn.functional import pad @@ -94,6 +97,12 @@ def __init__( rank=self.rank, world_size=self.world_size, ).get_graphcast_graph(mesh_vertex_rank_placement=mesh_vertex_placement) + # The graph is constructed on CPU; lift its features and comm-pattern + # index tensors onto the compute device so the model (which lives on + # self.device) can consume it without device-mismatch errors. + self.graph_cast_graph = move_graphcast_graph_to_device( + self.graph_cast_graph, self.device + ) print(f"Generated static graph in {time.time() - start_time:.2f} seconds.") self.extra_args: Dict[str, Any] = kwargs diff --git a/experiments/GraphCast/inference_benchmark.py b/experiments/GraphCast/inference_benchmark.py index b9548b5..f6b3639 100644 --- a/experiments/GraphCast/inference_benchmark.py +++ b/experiments/GraphCast/inference_benchmark.py @@ -20,7 +20,10 @@ from DGraph.Communicator import Communicator # noqa: E402 from DGraph.utils.TimingReport import TimingReport # noqa: E402 -from data_utils.graphcast_graph import DistributedGraphCastGraphGenerator # noqa: E402 +from data_utils.graphcast_graph import ( # noqa: E402 + DistributedGraphCastGraphGenerator, + move_graphcast_graph_to_device, +) from data_utils.utils import padded_size # noqa: E402 from dataset import SyntheticWeatherDataset # noqa: E402 from dist_utils import SingleProcessDummyCommunicator # noqa: E402 @@ -254,6 +257,8 @@ def run_correctness_check( ref_static_graph = ref_generator.get_graphcast_graph( mesh_vertex_rank_placement=ref_placement ) + # Built on CPU; move onto the compute device to match ref_model / ref_input. + ref_static_graph = move_graphcast_graph_to_device(ref_static_graph, device) ref_mesh_graph = ref_generator.get_mesh_graph(ref_placement) ref_grid2mesh = ref_generator.get_grid2mesh_graph(ref_mesh_graph) ref_renumbered_grid = ref_grid2mesh["renumbered_grid"] From af437216de6fdc16e7ede80b885a082012589cfb Mon Sep 17 00:00:00 2001 From: Shehtab Date: Fri, 17 Jul 2026 16:40:45 -0400 Subject: [PATCH 20/39] Removed stub preprocess code used for fast testing --- .../GraphCast/data_utils/preprocess.py | 21 ++-- experiments/GraphCast/gen_mesh_partitions.py | 19 ++- experiments/GraphCast/preprocess.py | 110 ------------------ 3 files changed, 30 insertions(+), 120 deletions(-) delete mode 100644 experiments/GraphCast/preprocess.py diff --git a/experiments/GraphCast/data_utils/preprocess.py b/experiments/GraphCast/data_utils/preprocess.py index ff86c22..d466d08 100644 --- a/experiments/GraphCast/data_utils/preprocess.py +++ b/experiments/GraphCast/data_utils/preprocess.py @@ -12,17 +12,22 @@ def graphcast_graph_to_nxgraph(mesh_graph): def partition_graph(G: nx.Graph, num_ranks: int): - try: - import metis - except ImportError: - raise ImportError("Please install metis to use this function.") - if num_ranks == 1: - return np.zeros(len(G.nodes), dtype=int) if num_ranks < 1: raise ValueError("Number of ranks must be greater than 0.") + if num_ranks == 1: + return np.zeros(G.number_of_nodes(), dtype=int) + + try: + import pymetis + except ImportError: + raise ImportError( + "Please install pymetis to use this function (`pip install pymetis`)." + ) - metis_graph = metis.networkx_to_metis(G) - (edgecuts, node_rank_placement) = metis.part_graph(metis_graph, nparts=num_ranks) + # pymetis takes an adjacency list indexed by contiguous node id 0..N-1 + # (the mesh graph's vertices are numbered this way). + adjacency = [sorted(G.neighbors(node)) for node in range(G.number_of_nodes())] + (edgecuts, node_rank_placement) = pymetis.part_graph(num_ranks, adjacency=adjacency) # Node_rank_placement is of shape (num_nodes, ), where each element is the # rank of the node in the partitioning. diff --git a/experiments/GraphCast/gen_mesh_partitions.py b/experiments/GraphCast/gen_mesh_partitions.py index 155f2c3..4b50a78 100644 --- a/experiments/GraphCast/gen_mesh_partitions.py +++ b/experiments/GraphCast/gen_mesh_partitions.py @@ -56,6 +56,21 @@ def generate_partition( return path +def _parse_int_list(value) -> list: + """Parse a CLI arg into a list of ints, tolerating Fire's auto-conversion. + + Fire turns "4,6" into a tuple (4, 6) and "1" into the int 1, so the value + can arrive as a str, int, or tuple/list. Normalize all of these to + list[int] so `--mesh_levels="4,6"`, `--mesh_levels=4`, and + `--mesh_levels 4 6` all work. + """ + if isinstance(value, str): + return [int(x) for x in value.split(",") if x != ""] + if isinstance(value, (list, tuple)): + return [int(x) for x in value] + return [int(value)] + + def main( mesh_levels: str = "2,5,6", world_sizes: str = "1,2,4", @@ -70,8 +85,8 @@ def main( output_dir: directory to write mesh_vertex_rank_placement_L{L}_W{W}.pt files to. force: recompute and overwrite even if a cached file already exists. """ - levels = [int(x) for x in mesh_levels.split(",")] - sizes = [int(x) for x in world_sizes.split(",")] + levels = _parse_int_list(mesh_levels) + sizes = _parse_int_list(world_sizes) for mesh_level in levels: for world_size in sizes: diff --git a/experiments/GraphCast/preprocess.py b/experiments/GraphCast/preprocess.py deleted file mode 100644 index 93fc342..0000000 --- a/experiments/GraphCast/preprocess.py +++ /dev/null @@ -1,110 +0,0 @@ -import torch -import numpy as np -import pickle -import graphcast_config -import networkx as nx -import numpy as np -from data_utils.utils import create_graph -from data_utils.icosahedral_mesh import ( - faces_to_edges, - get_hierarchy_of_triangular_meshes_for_sphere, - merge_meshes, -) -from fire import Fire - - -def partition_graph(G: nx.Graph, num_ranks: int): - if num_ranks == 1: - return np.zeros(G.number_of_nodes(), dtype=int) - - try: - import pymetis - except ImportError: - raise ImportError( - "Please install pymetis to use this function (`pip install pymetis`)." - ) - - adjacency = [sorted(G.neighbors(node)) for node in range(G.number_of_nodes())] - (edgecuts, node_rank_placement) = pymetis.part_graph(num_ranks, adjacency=adjacency) - # Node_rank_placement is of shape (num_nodes, ), where each element is the - # rank of the node in the partitioning. - - node_rank_placement = np.array(node_rank_placement) - return node_rank_placement - - -def node_renumbering(node_rank_placement, num_parts): - """The nodes are renumbered based on the rank mappings so the node features and - numbers are contiguous.""" - rearranged_indices = [] - for rank in range(num_parts): - mask = node_rank_placement == rank - indices = np.where(mask)[0] - rearranged_indices.append(indices) - renumbered_nodes = np.concatenate(rearranged_indices) - return renumbered_nodes - - -def edge_renumbering(edge_indices, renumbered_nodes): - """""" - src_indices = edge_indices[:, 0] - dst_indices = edge_indices[:, 1] - src_indices = renumbered_nodes[src_indices] - dst_indices = renumbered_nodes[dst_indices] - renumbered_edges = np.stack((src_indices, dst_indices)) - return renumbered_edges - - -def create_graphcast_meshgraph(mesh_level=6): - _meshes = get_hierarchy_of_triangular_meshes_for_sphere(splits=mesh_level) - finest_mesh = _meshes[-1] # get the last one in the list of meshes - mesh = merge_meshes(_meshes) - mesh_src, mesh_dst = faces_to_edges(mesh.faces) - mesh_vertices = np.array(mesh.vertices) - mesh_src = torch.tensor(mesh_src, dtype=torch.int32) - mesh_dst = torch.tensor(mesh_dst, dtype=torch.int32) - mesh_pos = torch.tensor(mesh_vertices, dtype=torch.float32) - mesh_graph = create_graph(mesh_src, mesh_dst, mesh_pos, to_bidirected=True) - return mesh_graph - - -def graphcast_graph_to_nxgraph(mesh_graph): - G = nx.Graph() - src_indices = mesh_graph[2].numpy() - dst_indices = mesh_graph[3].numpy() - edge_indices = np.stack((src_indices, dst_indices)).T - G.add_edges_from(edge_indices) - return G - - -def save_nx_meshgraph(nx_graph, filename): - with open(filename, "wb") as f: - pickle.dump(nx_graph, f) - - -def metis_partition_graph(nx_graph, num_ranks): - mesh_vertex_rank_placement = partition_graph(nx_graph, num_ranks) - - # Renumber the mesh graph vertices and edges - renumbered_nodes = node_renumbering(mesh_vertex_rank_placement, num_ranks) - renumbered_edges = edge_renumbering(nx_graph.edges, renumbered_nodes) - - # Generate rank placement for the grid nodes - - # Recalculate thee Grid2Mesh and Mesh2Grid edges using - # the new node numbering - - -def main(load_graph="", mesh_level=6, num_ranks=1): - if load_graph: - with open(load_graph, "rb") as f: - G = pickle.load(f) - else: - mesh_graph = create_graphcast_meshgraph(mesh_level=mesh_level) - G = graphcast_graph_to_nxgraph(mesh_graph) - - # if num_ranks > 1: - - -if __name__ == "__main__": - Fire(main) From fca2d701507e1628893dbcfb5fb8481956e393b7 Mon Sep 17 00:00:00 2001 From: Shehtab Date: Fri, 17 Jul 2026 17:15:37 -0400 Subject: [PATCH 21/39] Temporary fix on the grid out of bounds error --- .../GraphCast/data_utils/graphcast_graph.py | 18 ++++++++++ experiments/GraphCast/dataset.py | 34 +++++++++---------- 2 files changed, 35 insertions(+), 17 deletions(-) diff --git a/experiments/GraphCast/data_utils/graphcast_graph.py b/experiments/GraphCast/data_utils/graphcast_graph.py index c58c5cb..8289b76 100644 --- a/experiments/GraphCast/data_utils/graphcast_graph.py +++ b/experiments/GraphCast/data_utils/graphcast_graph.py @@ -81,6 +81,15 @@ class DistributedGraphCastGraph: # Distributed graph info distributed_comm_patterns: GraphCastCommPatterns + # Original grid-vertex ids owned by this rank, in the comm-pattern local + # order. This is the frame the grid2mesh/mesh2grid CommunicationPatterns use: + # local grid vertex i corresponds to original grid id + # local_grid_original_indices[i]. The dataset MUST select/reorder its grid + # input & output with this tensor so grid_node_features rows align with + # send_local_idx / local_edge_list. Kept on CPU (indexes CPU inputs in the + # dataset's __getitem__ before the device move). + local_grid_original_indices: Tensor + def _move_communication_pattern_to_device( cp: CommunicationPattern, device @@ -518,6 +527,11 @@ def get_graphcast_graph( grid2mesh_graph = self.get_grid2mesh_graph(mesh_graph, mesh2grid_raw=mesh2grid_raw) grid_vertex_rank_placement = grid2mesh_graph["grid_vertex_rank_placement"] inverse_grid_renumbering = grid2mesh_graph["inverse_grid_renumbering"] + # renumbered_grid maps rank-sorted slot -> original grid id. Selecting the + # slots this rank owns gives the original grid ids in comm-pattern local + # order, which the dataset uses to shard/reorder grid features (see the + # local_grid_original_indices field docstring). + renumbered_grid = grid2mesh_graph["renumbered_grid"] mesh2grid_graph = self.get_mesh2grid_edges( inverse_vertex_renumbering, @@ -551,6 +565,9 @@ def get_graphcast_graph( # (vertex: placement==rank; edge: dst-owned-by-rank) so rows stay aligned. rank = topology.rank + # Original grid ids owned by this rank, in comm-pattern local order. + local_grid_original_indices = renumbered_grid[grid_vertex_rank_placement == rank] + mesh_node_local_mask = mesh_vertex_rank_placement == rank local_mesh_node_features = mesh_graph["node_features"][mesh_node_local_mask] @@ -584,4 +601,5 @@ def get_graphcast_graph( mesh2grid_graph_edge_features=local_mesh2grid_edge_features, grid2mesh_graph_edge_features=local_grid2mesh_edge_features, distributed_comm_patterns=comm_patterns, + local_grid_original_indices=local_grid_original_indices, ) diff --git a/experiments/GraphCast/dataset.py b/experiments/GraphCast/dataset.py index 65046d8..5e929f5 100644 --- a/experiments/GraphCast/dataset.py +++ b/experiments/GraphCast/dataset.py @@ -20,8 +20,6 @@ DistributedGraphCastGraphGenerator, move_graphcast_graph_to_device, ) -from data_utils.utils import padded_size -from torch.nn.functional import pad class SyntheticWeatherDataset(Dataset): @@ -100,6 +98,15 @@ def __init__( # The graph is constructed on CPU; lift its features and comm-pattern # index tensors onto the compute device so the model (which lives on # self.device) can consume it without device-mismatch errors. + # Original grid ids this rank owns, in the comm-pattern local order. Grid + # input/output MUST be sharded/reordered with this (NOT a contiguous + # chunk): the grid2mesh/mesh2grid CommunicationPatterns number grid + # vertices by grid_part's connectivity-weighted vote, so a naive chunk + # both mis-sizes the buffer (send_local_idx out-of-bounds) and misorders + # its rows. Kept on CPU to index the CPU inputs in __getitem__. + self.local_grid_original_indices = ( + self.graph_cast_graph.local_grid_original_indices.detach().cpu().long() + ) self.graph_cast_graph = move_graphcast_graph_to_device( self.graph_cast_graph, self.device ) @@ -213,21 +220,14 @@ def __getitem__(self, idx: int): ) if self.world_size > 1: - # Get oartitioned inputs instead of the full graph - num_grid_nodes = in_var.shape[0] - padded_num_grid_nodes = padded_size(num_grid_nodes, self.ranks_per_graph) - - num_nodes_per_rank = padded_num_grid_nodes // self.ranks_per_graph - in_var = pad(in_var, (padded_num_grid_nodes - num_grid_nodes, 0), value=-0) - out_var = pad( - out_var, (padded_num_grid_nodes - num_grid_nodes, 0), value=-0 - ) - - start_index = self.rank * num_nodes_per_rank - end_index = start_index + num_nodes_per_rank - - in_var = in_var[start_index:end_index] - out_var = out_var[start_index:end_index] + # Select this rank's grid vertices in the comm-pattern's local order. + # local_grid_original_indices is a permutation-slice into the full + # grid (original grid ids owned by this rank, rank-sorted), so it both + # picks the right vertices and orders their rows to match + # send_local_idx / local_edge_list. No padding needed: the union of + # all ranks' indices covers every grid node exactly once. + in_var = in_var[self.local_grid_original_indices] + out_var = out_var[self.local_grid_original_indices] return { "invar": in_var.to(self.device), From 39af504f0ca9a092baecfa5fda8eacf30539f523 Mon Sep 17 00:00:00 2001 From: M Shehtab Zaman Date: Fri, 31 Jul 2026 13:44:12 -0700 Subject: [PATCH 22/39] Update inference benchmark for graphcast --- experiments/GraphCast/inference_benchmark.py | 162 +++++++++++++------ 1 file changed, 110 insertions(+), 52 deletions(-) diff --git a/experiments/GraphCast/inference_benchmark.py b/experiments/GraphCast/inference_benchmark.py index f6b3639..4c690ae 100644 --- a/experiments/GraphCast/inference_benchmark.py +++ b/experiments/GraphCast/inference_benchmark.py @@ -43,7 +43,9 @@ def build_comm(mode: str, backend: str, world_size: int): return Communicator.init_process_group(backend, ranks_per_graph=world_size) -def load_partition(partition_dir: str, mesh_level: int, world_size: int) -> torch.Tensor: +def load_partition( + partition_dir: str, mesh_level: int, world_size: int +) -> torch.Tensor: path = os.path.join( partition_dir, f"mesh_vertex_rank_placement_L{mesh_level}_W{world_size}.pt" ) @@ -208,7 +210,9 @@ def run_correctness_check( dataset = build_dataset(cfg, mesh_level, partition, rank, world_size, device) static_graph = dataset.get_static_graph() - num_local_actual = static_graph.distributed_comm_patterns.mesh2grid.num_local_vertices + num_local_actual = ( + static_graph.distributed_comm_patterns.mesh2grid.num_local_vertices + ) gathered = [None] * world_size dist.all_gather_object(gathered, (rank, num_local_expected, num_local_actual)) @@ -216,15 +220,21 @@ def run_correctness_check( if rank == 0: if step1_ok: - print("[correctness] STEP 1 PASS: dataset chunk size vs grid_part local " - "vertex count agree on every rank.") + print( + "[correctness] STEP 1 PASS: dataset chunk size vs grid_part local " + "vertex count agree on every rank." + ) else: mismatches = [(r, e, a) for (r, e, a) in gathered if e != a] - print(f"[correctness] STEP 1 FAIL: per-rank (rank, expected, actual) " - f"mismatches: {mismatches}") - print("[correctness] Halting -- do not trust world_size>1 benchmark " - "numbers until this is resolved (see plan's known dataset.py vs " - "grid_part mismatch).") + print( + f"[correctness] STEP 1 FAIL: per-rank (rank, expected, actual) " + f"mismatches: {mismatches}" + ) + print( + "[correctness] Halting -- do not trust world_size>1 benchmark " + "numbers until this is resolved (see plan's known dataset.py vs " + "grid_part mismatch)." + ) dist.barrier() if not step1_ok: return False @@ -235,12 +245,17 @@ def run_correctness_check( # Recover this run's grid renumbering (original grid index for each # renumbered/rank-sorted slot) without re-triggering comm-pattern construction. dist_generator = DistributedGraphCastGraphGenerator( - lat_lon_grid, mesh_level=mesh_level, ranks_per_graph=world_size, - rank=rank, world_size=world_size, + lat_lon_grid, + mesh_level=mesh_level, + ranks_per_graph=world_size, + rank=rank, + world_size=world_size, ) dist_mesh_graph = dist_generator.get_mesh_graph(partition) dist_grid2mesh = dist_generator.get_grid2mesh_graph(dist_mesh_graph) - dist_renumbered_grid = dist_grid2mesh["renumbered_grid"] # [N_grid] -> original grid idx + dist_renumbered_grid = dist_grid2mesh[ + "renumbered_grid" + ] # [N_grid] -> original grid idx dist_grid_placement = dist_grid2mesh["grid_vertex_rank_placement"] # sorted by rank local_mask = dist_grid_placement == rank @@ -270,7 +285,9 @@ def run_correctness_check( ref_input = canonical_input[ref_renumbered_grid].unsqueeze(0).to(device) with torch.no_grad(): - ref_pred = ref_model(ref_input, ref_static_graph) # [N_grid, C], ref-renumbered order + ref_pred = ref_model( + ref_input, ref_static_graph + ) # [N_grid, C], ref-renumbered order ref_pred_original_order = torch.zeros_like(ref_pred) ref_pred_original_order[ref_renumbered_grid.to(device)] = ref_pred @@ -296,12 +313,16 @@ def run_correctness_check( if rank == 0: all_close = all(c for (_, c, _, _) in gathered2) if all_close: - print(f"[correctness] STEP 2 PASS: all ranks match reference within " - f"atol={atol}, rtol={rtol}. Per-rank (rank, close, max_abs_diff, " - f"max_rel_diff): {gathered2}") + print( + f"[correctness] STEP 2 PASS: all ranks match reference within " + f"atol={atol}, rtol={rtol}. Per-rank (rank, close, max_abs_diff, " + f"max_rel_diff): {gathered2}" + ) else: - print(f"[correctness] STEP 2 FAIL. Per-rank (rank, close, max_abs_diff, " - f"max_rel_diff): {gathered2}") + print( + f"[correctness] STEP 2 FAIL. Per-rank (rank, close, max_abs_diff, " + f"max_rel_diff): {gathered2}" + ) return all_close return close @@ -321,7 +342,9 @@ def main( correctness_rtol: float = 1e-3, ) -> None: assert mode in ("single", "distributed"), "mode must be 'single' or 'distributed'" - assert batch_size == 1, "only batch_size=1 is supported (matches GraphCast's one-step inference)" + assert ( + batch_size == 1 + ), "only batch_size=1 is supported (matches GraphCast's one-step inference)" assert not correctness or mode == "distributed", ( "--correctness checks cross-rank partitioning; run it with --mode distributed " "(and --nproc_per_node >= 2), not --mode single" @@ -343,12 +366,22 @@ def main( if correctness: if world_size < 2: if rank == 0: - print("[correctness] world_size < 2: nothing to check (single-GPU " - "baseline has no cross-rank partitioning to verify).") + print( + "[correctness] world_size < 2: nothing to check (single-GPU " + "baseline has no cross-rank partitioning to verify)." + ) return ok = run_correctness_check( - cfg, mesh_level, comm, partition_dir, rank, world_size, device, - seed, correctness_atol, correctness_rtol, + cfg, + mesh_level, + comm, + partition_dir, + rank, + world_size, + device, + seed, + correctness_atol, + correctness_rtol, ) if rank == 0: print(f"[correctness] Overall: {'PASS' if ok else 'FAIL'}") @@ -378,27 +411,39 @@ def main( if any_oom: if rank == 0: - print(f"[inference_benchmark] OOM at mode={mode} mesh_level={mesh_level} " - f"world_size={world_size} (per-rank oom={oom_flags}) -- infeasible " - f"at this scale on this hardware, not a benchmark failure.") + print( + f"[inference_benchmark] OOM at mode={mode} mesh_level={mesh_level} " + f"world_size={world_size} (per-rank oom={oom_flags}) -- infeasible " + f"at this scale on this hardware, not a benchmark failure." + ) payload = { "benchmark": "graphcast_inference", "metadata": collect_metadata(), "config": { - "mode": mode, "backend": backend, "world_size": world_size, - "mesh_level": mesh_level, "hidden_dim": cfg.model.hidden_dim, + "mode": mode, + "backend": backend, + "world_size": world_size, + "mesh_level": mesh_level, + "hidden_dim": cfg.model.hidden_dim, "processor_layers": cfg.model.processor_layers, - "batch_size": batch_size, "warmup_iters": warmup_iters, - "measure_iters": measure_iters, "seed": seed, + "batch_size": batch_size, + "warmup_iters": warmup_iters, + "measure_iters": measure_iters, + "seed": seed, }, - "measurements": [{ - "params": {"world_size": world_size, "mesh_level": mesh_level}, - "oom": True, - "oom_per_rank": oom_flags, - }], + "measurements": [ + { + "params": {"world_size": world_size, "mesh_level": mesh_level}, + "oom": True, + "oom_per_rank": oom_flags, + } + ], } write_result( - os.path.join(output_dir, f"graphcast_infer_{mode}_W{world_size}_L{mesh_level}.json"), + os.path.join( + output_dir, + f"graphcast_infer_{mode}_W{world_size}_L{mesh_level}.json", + ), payload, ) return @@ -420,26 +465,35 @@ def main( "benchmark": "graphcast_inference", "metadata": collect_metadata(), "config": { - "mode": mode, "backend": backend, "world_size": world_size, - "mesh_level": mesh_level, "hidden_dim": cfg.model.hidden_dim, + "mode": mode, + "backend": backend, + "world_size": world_size, + "mesh_level": mesh_level, + "hidden_dim": cfg.model.hidden_dim, "processor_layers": cfg.model.processor_layers, - "batch_size": batch_size, "warmup_iters": warmup_iters, - "measure_iters": measure_iters, "seed": seed, + "batch_size": batch_size, + "warmup_iters": warmup_iters, + "measure_iters": measure_iters, + "seed": seed, }, - "measurements": [{ - "params": {"world_size": world_size, "mesh_level": mesh_level}, - "end_to_end_latency_ms": latency_summary, - "throughput_grid_points_per_sec": throughput, - "peak_memory_bytes": { - "per_rank": peak_memory_list, - "max": max(m for m in peak_memory_list if m is not None), - }, - "phase_breakdown_ms": phase_breakdown, - "per_rank_end_to_end_trials_ms": trials_list, - }], + "measurements": [ + { + "params": {"world_size": world_size, "mesh_level": mesh_level}, + "end_to_end_latency_ms": latency_summary, + "throughput_grid_points_per_sec": throughput, + "peak_memory_bytes": { + "per_rank": peak_memory_list, + "max": max(m for m in peak_memory_list if m is not None), + }, + "phase_breakdown_ms": phase_breakdown, + "per_rank_end_to_end_trials_ms": trials_list, + } + ], } write_result( - os.path.join(output_dir, f"graphcast_infer_{mode}_W{world_size}_L{mesh_level}.json"), + os.path.join( + output_dir, f"graphcast_infer_{mode}_W{world_size}_L{mesh_level}.json" + ), payload, ) print( @@ -448,6 +502,10 @@ def main( f"throughput={throughput:.1f} grid-points/sec" ) + if mode == "distributed": + dist.barrier() + dist.destroy_process_group() + if __name__ == "__main__": Fire(main) From a8a3a71690997ce4b1bceece8685c9f31f408f0b Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 31 Jul 2026 16:44:53 -0400 Subject: [PATCH 23/39] Updated cost model code for paper --- .../analysis/compute_predictions.py | 33 +----- .../analysis/fit_overhead.py | 101 ++++++++++++++---- .../analysis/fit_primitives.py | 55 +++++++++- .../benchmarks/bench_end_to_end.py | 1 + .../benchmarks/graph_data_common.py | 18 +++- .../visualization/plot_ablations.py | 38 +++++-- .../visualization/plot_tipping_point.py | 8 ++ 7 files changed, 195 insertions(+), 59 deletions(-) diff --git a/experiments/cost_model_benchmarks/analysis/compute_predictions.py b/experiments/cost_model_benchmarks/analysis/compute_predictions.py index 3e69f4d..5070fba 100644 --- a/experiments/cost_model_benchmarks/analysis/compute_predictions.py +++ b/experiments/cost_model_benchmarks/analysis/compute_predictions.py @@ -72,35 +72,12 @@ def main(): prediction_entries = [] for r in all_runs: - T_model_base = predict_layer_time(r["config"], r["per_rank_stats"], primitives) - T_pred = T_model_base + T_overhead + breakdown = predict_layer_time(r["config"], r["per_rank_stats"], primitives) + T_pred = breakdown["T_total"] + T_overhead T_meas = r["measured_median"] abs_err = abs(T_meas - T_pred) rel_err = abs_err / T_meas if T_meas > 0 else float("nan") - # Decompose prediction for ablation figures - net = primitives.get("network", {}) - intra_bytes = r["per_rank_stats"].get("c_intra_bytes", 0) - inter_bytes = r["per_rank_stats"].get("c_inter_bytes", 0) - - def net_time(nbytes, mode): - params = net.get(mode, None) - if params is None or nbytes == 0: - return 0.0 - return params.get("latency_seconds", 0.0) + nbytes / params.get("bandwidth_bytes_per_sec", 1e10) - - T_intra = net_time(intra_bytes, "intra") - T_inter = net_time(inter_bytes, "inter") - T_comm = max(T_intra, T_inter) - - F = r["config"]["feature_dim"] - send_bytes = r["per_rank_stats"].get("send_total", 0) * F * 4 - gath_params = primitives.get("gather", {}).get("clustered", {}).get("gather", None) - T_buffer_copy = 0.0 - if gath_params and send_bytes > 0: - B_g = gath_params.get("bandwidth_bytes_per_sec", 1e12) - T_buffer_copy = gath_params.get("intercept_seconds", 0.0) + send_bytes / B_g - entry = { "source_file": r["source_file"], "config": r["config"], @@ -111,9 +88,9 @@ def net_time(nbytes, mode): "relative_error": rel_err, "in_fit_set": id(r) in fit_set, "breakdown": { - "T_comp_seconds": T_model_base - T_comm - T_buffer_copy, - "T_comm_seconds": T_comm, - "T_buffer_copy_seconds": T_buffer_copy, + "T_comp_seconds": breakdown["T_comp"], + "T_comm_seconds": breakdown["T_comm"], + "T_buffer_copy_seconds": breakdown["T_buffer_copy"], "T_overhead_seconds": T_overhead, }, } diff --git a/experiments/cost_model_benchmarks/analysis/fit_overhead.py b/experiments/cost_model_benchmarks/analysis/fit_overhead.py index 426f440..25dbae1 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_overhead.py +++ b/experiments/cost_model_benchmarks/analysis/fit_overhead.py @@ -32,8 +32,32 @@ # Cost model (without overhead) # --------------------------------------------------------------------------- +def _gather_piecewise_time(nbytes: float, params: dict) -> float: + """Evaluate the fitted piecewise L2-cache/HBM-bandwidth gather model for + a single byte count. Mirrors ``time_model`` in + ``analysis/fit_primitives.py::fit_gather`` exactly — this is the single + place both ``predict_layer_time`` and ``compute_predictions.py`` should + call, rather than re-deriving a simplified linear approximation. + """ + overhead = params.get("launch_overhead_seconds", 0.0) + bw_HBM = params.get("bandwidth_bytes_per_sec") + bw_L2 = params.get("L2_bandwidth_bytes_per_sec") + if not bw_HBM or not bw_L2: + return overhead + + inv_bw_L2 = 1.0 / bw_L2 + inv_bw_HBM = 1.0 / bw_HBM + L2_thresh = params.get("L2_inflection_bytes", 0.0) + HBM_thresh = params.get("HBM_inflection_bytes", 0.0) + + bytes_L2 = min(max(nbytes, 0.0), L2_thresh) + bytes_HBM = max(0.0, nbytes - HBM_thresh) + t_mem = bytes_L2 * inv_bw_L2 + bytes_HBM * inv_bw_HBM + return max(overhead, t_mem) + + def predict_layer_time(run_config: dict, per_rank_stats: dict, - primitives: dict) -> float: + primitives: dict) -> dict: """Predict T_layer for one rank using the assembled primitive model. T_layer = T_comp + max(T_intra, T_inter) + T_buffer_copy @@ -46,6 +70,13 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, Stats for rank 0 from per_rank_stats list. primitives : dict Loaded fitted_primitives.json. + + Returns + ------- + dict with keys T_comp, T_intra, T_inter, T_comm, T_buffer_copy, T_total + (T_total = T_comp + T_comm + T_buffer_copy, i.e. without T_overhead). + Every component is floored at 0.0, since noisy fitted intercepts can + otherwise extrapolate to physically-meaningless negative times. """ F = run_config["feature_dim"] model_type = run_config.get("model", "gcn") @@ -54,9 +85,16 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, n_halo = per_rank_stats.get("n_halo", 0) n_total = n_local + n_halo - # A rough edge count estimate: use avg_degree * n_local as a proxy - avg_degree = run_config.get("avg_degree", 20.0) - n_edges_local = int(n_local * avg_degree) + # Prefer the real local edge count recorded by bench_end_to_end.py + # (edge_index.shape[1] in graph_data_common.build_local_comm_pattern); + # fall back to the avg_degree*n_local proxy only for older run data that + # predates recording it. The proxy ignores partitioner-dependent edge-cut + # effects (random/balanced/metis produce very different local edge + # densities for the same avg_degree). + n_edges_local = per_rank_stats.get("n_edges_local") + if n_edges_local is None: + avg_degree = run_config.get("avg_degree", 20.0) + n_edges_local = int(n_local * avg_degree) # T_comp comp_params = primitives.get("compute", {}).get(model_type, {}).get("forward", None) @@ -64,9 +102,9 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, T_comp = (comp_params["coeff_V"] * n_total + comp_params["coeff_E"] * n_edges_local + comp_params["intercept"]) - T_comp = max(T_comp, 0.0) else: T_comp = 0.0 + T_comp = max(T_comp, 0.0) # T_intra and T_inter intra_bytes = per_rank_stats.get("c_intra_bytes", 0) @@ -82,20 +120,28 @@ def net_time(nbytes: int, mode: str) -> float: t_L = params.get("latency_seconds", 0.0) return t_L + nbytes / B - T_intra = net_time(intra_bytes, "intra") - T_inter = net_time(inter_bytes, "inter") + T_intra = max(net_time(intra_bytes, "intra"), 0.0) + T_inter = max(net_time(inter_bytes, "inter"), 0.0) T_comm = max(T_intra, T_inter) - # T_buffer_copy (gather of send buffer) + # T_buffer_copy (gather of send buffer) — uses the fitted piecewise + # L2/HBM model, not a simplified single-slope linear stand-in. send_bytes = per_rank_stats.get("send_total", 0) * F * 4 gath_params = primitives.get("gather", {}).get("clustered", {}).get("gather", None) if gath_params and send_bytes > 0: - B_g = gath_params.get("bandwidth_bytes_per_sec", 1e12) - T_buffer_copy = gath_params.get("intercept_seconds", 0.0) + send_bytes / B_g + T_buffer_copy = _gather_piecewise_time(send_bytes, gath_params) else: T_buffer_copy = 0.0 - - return T_comp + T_comm + T_buffer_copy + T_buffer_copy = max(T_buffer_copy, 0.0) + + return { + "T_comp": T_comp, + "T_intra": T_intra, + "T_inter": T_inter, + "T_comm": T_comm, + "T_buffer_copy": T_buffer_copy, + "T_total": T_comp + T_comm + T_buffer_copy, + } # --------------------------------------------------------------------------- @@ -146,22 +192,37 @@ def apply_filter(runs: list, filter_expr: str) -> tuple: # Scalar overhead fitting # --------------------------------------------------------------------------- +def weighted_median(values: np.ndarray, weights: np.ndarray) -> float: + """Return m minimising sum(weights * |values - m|).""" + order = np.argsort(values) + v = values[order] + w = weights[order] + cum = np.cumsum(w) + cutoff = cum[-1] / 2.0 + idx = int(np.searchsorted(cum, cutoff)) + idx = min(idx, len(v) - 1) + return float(v[idx]) + + def fit_overhead_scalar(fit_runs: list, primitives: dict) -> tuple: """Fit T_overhead to minimise MAPE on fit_runs. Returns (overhead, mape_in_sample).""" if not fit_runs: return 0.0, float("nan") - residuals = [] - for r in fit_runs: - T_model = predict_layer_time(r["config"], r["per_rank_stats"], primitives) - residuals.append(r["measured_median"] - T_model) + residuals = np.array([ + r["measured_median"] - predict_layer_time(r["config"], r["per_rank_stats"], primitives)["T_total"] + for r in fit_runs + ]) + T_meas = np.array([r["measured_median"] for r in fit_runs]) - # Optimal scalar overhead that minimises sum of |err - overhead| / T_meas - # is the weighted median; for uniform weights it's just the median of residuals. - overhead = float(np.median(residuals)) + # minimize sum(|residual - overhead| / T_meas) over the scalar overhead + # == minimize sum(weight * |residual - overhead|) with weight = 1/T_meas, + # whose minimizer is the weighted median (NOT the plain median, which + # only minimizes the unweighted sum |residual - overhead|). + overhead = weighted_median(residuals, 1.0 / T_meas) mape = float(np.mean([ - abs(r["measured_median"] - (predict_layer_time(r["config"], r["per_rank_stats"], primitives) + overhead)) + abs(r["measured_median"] - (predict_layer_time(r["config"], r["per_rank_stats"], primitives)["T_total"] + overhead)) / r["measured_median"] for r in fit_runs ])) diff --git a/experiments/cost_model_benchmarks/analysis/fit_primitives.py b/experiments/cost_model_benchmarks/analysis/fit_primitives.py index 6211f94..0949a91 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_primitives.py +++ b/experiments/cost_model_benchmarks/analysis/fit_primitives.py @@ -69,6 +69,31 @@ def linear_fit(x: np.ndarray, y: np.ndarray): } +def weighted_linear_fit(x: np.ndarray, y: np.ndarray, sigma: np.ndarray = None): + """Fit y = slope * x + intercept via weighted least squares. + + ``sigma`` follows the same convention as ``scipy.optimize.curve_fit``: + residuals are weighted by ``1/sigma``, i.e. this minimises + ``sum(((y - pred) / sigma) ** 2)``. Passing ``sigma=y`` (as ``fit_gather`` + already does via ``curve_fit``) weights by *relative* error so that + points spanning many orders of magnitude don't dominate the fit. + """ + if sigma is None: + sigma = np.ones_like(y) + w = 1.0 / sigma + A = np.column_stack([x, np.ones_like(x)]) * w[:, None] + b = y * w + result, _, _, _ = np.linalg.lstsq(A, b, rcond=None) + slope, intercept = result + y_pred = slope * x + intercept + r2 = r_squared(y, y_pred) + return { + "slope": float(slope), + "intercept": float(intercept), + "r_squared": r2, + } + + # --------------------------------------------------------------------------- # Network fit: T = t_L + bytes / B # --------------------------------------------------------------------------- @@ -90,7 +115,11 @@ def fit_network(records: list) -> dict: # T = t_L + bytes / B → T = intercept + slope * bytes # so slope = 1/B, intercept = t_L - fit = linear_fit(bytes_arr, time_arr) + # Weighted by relative error (sigma=time_arr) for consistency with + # fit_gather's curve_fit(sigma=T_arr) — message sizes span 64B-64MiB, so + # an unweighted fit would be dominated by the largest sizes and leave the + # latency intercept (small messages) poorly constrained. + fit = weighted_linear_fit(bytes_arr, time_arr, sigma=time_arr) bandwidth = 1.0 / fit["slope"] if fit["slope"] > 0 else float("nan") latency = fit["intercept"] return { @@ -124,9 +153,17 @@ def fit_compute(records: list, timing_key: str = "forward_trials_seconds") -> di E_arr = np.array(E_arr, dtype=float) T_arr = np.array(T_arr, dtype=float) - # Design matrix: [V, E, 1] + # Design matrix: [V, E, 1], weighted by relative error (1/T) for + # consistency with fit_gather's curve_fit(sigma=T_arr) — |V|/|E| sweeps + # are log-spaced over several orders of magnitude, so an unweighted fit + # would be dominated by the largest graphs and leave the small-graph + # intercept (which matters most for the crossover/tipping-point + # analysis) poorly constrained. A = np.column_stack([V_arr, E_arr, np.ones_like(V_arr)]) - result, _, _, _ = np.linalg.lstsq(A, T_arr, rcond=None) + w = 1.0 / T_arr + A_w = A * w[:, None] + b_w = T_arr * w + result, _, _, _ = np.linalg.lstsq(A_w, b_w, rcond=None) coeff_V, coeff_E, intercept = result T_pred = A @ result r2 = r_squared(T_arr, T_pred) @@ -282,6 +319,18 @@ def main(): f"t_L={net['inter']['latency_seconds']*1e6:.2f} µs " f"R²={net['inter']['r_squared']:.4f}" ) + if args.pingpong_intra and args.pingpong_inter: + # Flat (single-tier) network fit: pools intra + inter data into one + # (B, t_L) pair, ignoring the intra/inter distinction. Used as a + # genuine baseline for the hierarchical-vs-flat ablation, rather than + # an ad hoc multiplier on the hierarchical prediction. + recs = load_json_files(args.pingpong_intra) + load_json_files(args.pingpong_inter) + net["flat"] = fit_network(recs) + print( + f"[network/flat] B={net['flat']['bandwidth_bytes_per_sec']/1e9:.2f} GB/s " + f"t_L={net['flat']['latency_seconds']*1e6:.2f} µs " + f"R²={net['flat']['r_squared']:.4f}" + ) result["network"] = net # Compute diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py index fa56f5d..9f0bb8b 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py @@ -180,6 +180,7 @@ def one_layer(): "inter_halo_size": pattern["inter_halo_size"], "c_intra_bytes": pattern["intra_halo_size"] * F * 4, "c_inter_bytes": pattern["inter_halo_size"] * F * 4, + "n_edges_local": edge_index.shape[1], "send_total": sum(send_counts), "recv_total": sum(recv_counts), "trials_seconds": times_local, diff --git a/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py b/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py index 15358e5..68db2e4 100644 --- a/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py +++ b/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py @@ -1,4 +1,7 @@ +import os + import numpy as np +import torch import torch.distributed as dist @@ -182,9 +185,20 @@ def build_local_comm_pattern( edge_index = torch.zeros((2, 0), dtype=torch.long) # Compute intra / inter halo sizes - ranks_per_node = int( - os.environ.get("LOCAL_WORLD_SIZE", os.environ.get("SLURM_NTASKS_PER_NODE", "4")) + ranks_per_node_str = os.environ.get( + "LOCAL_WORLD_SIZE", os.environ.get("SLURM_NTASKS_PER_NODE") ) + if ranks_per_node_str is None: + raise RuntimeError( + "Neither LOCAL_WORLD_SIZE nor SLURM_NTASKS_PER_NODE is set — cannot " + "determine GPUs-per-node. This value directly determines the " + "intra/inter halo split (c_intra_bytes vs c_inter_bytes) that the " + "hierarchical cost model relies on, so silently guessing would " + "bias every downstream fit. Launch with torchrun (which sets " + "LOCAL_WORLD_SIZE) or under srun/SLURM (which sets " + "SLURM_NTASKS_PER_NODE)." + ) + ranks_per_node = int(ranks_per_node_str) my_node = rank // ranks_per_node intra_halo_size = 0 inter_halo_size = 0 diff --git a/experiments/cost_model_benchmarks/visualization/plot_ablations.py b/experiments/cost_model_benchmarks/visualization/plot_ablations.py index a23eb4f..a3f1d5d 100644 --- a/experiments/cost_model_benchmarks/visualization/plot_ablations.py +++ b/experiments/cost_model_benchmarks/visualization/plot_ablations.py @@ -11,6 +11,7 @@ python -m visualization.plot_ablations \\ --predictions data/predictions.json \\ + --primitives data/fitted_primitives.json \\ --output figures/ablations """ @@ -44,6 +45,12 @@ def relative_error(pred, meas): def parse_args(): p = argparse.ArgumentParser(description="Plot ablation studies") p.add_argument("--predictions", type=str, required=True) + p.add_argument( + "--primitives", type=str, required=True, + help="fitted_primitives.json; needs the network.flat entry produced " + "by fit_primitives.py when both --pingpong-intra and " + "--pingpong-inter are supplied", + ) p.add_argument("--output", type=str, default="figures/ablations") return p.parse_args() @@ -53,10 +60,31 @@ def main(): with open(args.predictions) as f: data = json.load(f) + with open(args.primitives) as f: + primitives = json.load(f) entries = data["predictions"] T_overhead = data.get("T_overhead_seconds", 0.0) + flat_net = primitives.get("network", {}).get("flat") + if flat_net is None: + raise RuntimeError( + "fitted_primitives.json has no 'network.flat' entry. Re-run " + "analysis/fit_primitives.py with both --pingpong-intra and " + "--pingpong-inter so it can fit a genuine single-tier baseline " + "for this ablation, instead of the previous heuristic." + ) + + def flat_comm_time(c_intra_bytes: float, c_inter_bytes: float) -> float: + """Flat (single-tier) model: no intra/inter distinction, no overlap — + all halo bytes go over one fitted (B, t_L) pair, sequentially.""" + total_bytes = c_intra_bytes + c_inter_bytes + if total_bytes <= 0: + return 0.0 + B = flat_net.get("bandwidth_bytes_per_sec", 1e10) + t_L = flat_net.get("latency_seconds", 0.0) + return max(t_L + total_bytes / B, 0.0) + # --- Panel (a): topology sweep (SBM inter-density) --- # Filter to SBM entries, group by inter_density sbm_entries = [e for e in entries @@ -80,16 +108,14 @@ def main(): T_pred_full = e["predicted_seconds"] full_errs.append(relative_error(T_pred_full, T_meas)) - # Flat model: use a single network term = T_intra + T_inter (not max) - # Approximate: flat model can't overlap, so T_comm = T_intra + T_inter + # Flat model: swap the hierarchical (intra/inter, overlapped via + # max) comm term for one predicted by a genuinely fit single-tier + # (B, t_L) model over total halo bytes — not a heuristic ratio. bd = e.get("breakdown", {}) T_comm_hier = bd.get("T_comm_seconds", 0.0) - # Flat approximation: assume both intra and inter are sequential c_intra = e["partition_stats"].get("c_intra_bytes", 0) c_inter = e["partition_stats"].get("c_inter_bytes", 0) - # Without knowing the individual bandwidths, use ratio heuristic: - # flat ≈ 2 * max (conservative estimate) - T_comm_flat = T_comm_hier * 2.0 + T_comm_flat = flat_comm_time(c_intra, c_inter) T_pred_flat = (T_pred_full - T_comm_hier + T_comm_flat) flat_errs.append(relative_error(T_pred_flat, T_meas)) diff --git a/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py b/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py index 735c96d..f5cc0cb 100644 --- a/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py +++ b/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py @@ -75,6 +75,14 @@ def main(): return K_vals = sorted(ws_to_meas.keys()) + if K_vals[0] != 1: + raise ValueError( + f"No world_size=1 run found (smallest K in data is {K_vals[0]}). " + "Ideal-scaling and scaling-efficiency numbers are only meaningful " + "relative to a true single-GPU baseline — run bench_end_to_end.py " + "(or bench_crossover.py's --no-dist mode) at world_size=1 for this " + "graph/config before plotting the tipping point." + ) T_meas = np.array([np.median(ws_to_meas[K]) for K in K_vals]) * args.num_layers * args.num_epochs T_pred = np.array([np.median(ws_to_pred[K]) for K in K_vals]) * args.num_layers * args.num_epochs K_arr = np.array(K_vals, dtype=float) From 13ff6754c52eb31b22a46ac45eb149bd60ff4085 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 31 Jul 2026 17:20:42 -0400 Subject: [PATCH 24/39] Fix stale communication pattern code --- .../analysis/fit_overhead.py | 10 +- .../benchmarks/bench_crossover.py | 28 +--- .../benchmarks/bench_end_to_end.py | 53 +++---- .../benchmarks/graph_data_common.py | 147 +++++------------- 4 files changed, 75 insertions(+), 163 deletions(-) diff --git a/experiments/cost_model_benchmarks/analysis/fit_overhead.py b/experiments/cost_model_benchmarks/analysis/fit_overhead.py index 25dbae1..f623be1 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_overhead.py +++ b/experiments/cost_model_benchmarks/analysis/fit_overhead.py @@ -86,11 +86,11 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, n_total = n_local + n_halo # Prefer the real local edge count recorded by bench_end_to_end.py - # (edge_index.shape[1] in graph_data_common.build_local_comm_pattern); - # fall back to the avg_degree*n_local proxy only for older run data that - # predates recording it. The proxy ignores partitioner-dependent edge-cut - # effects (random/balanced/metis produce very different local edge - # densities for the same avg_degree). + # (edge_index.shape[1], from DGraph.distributed.build_communication_pattern's + # comm_pattern.local_edge_list); fall back to the avg_degree*n_local proxy + # only for older run data that predates recording it. The proxy ignores + # partitioner-dependent edge-cut effects (random/balanced/metis produce + # very different local edge densities for the same avg_degree). n_edges_local = per_rank_stats.get("n_edges_local") if n_edges_local is None: avg_degree = run_config.get("avg_degree", 20.0) diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py b/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py index c5c2922..04e278b 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py @@ -50,7 +50,6 @@ """ import argparse -import os import numpy as np import torch @@ -70,6 +69,8 @@ partition_balanced, partition_metis, partition_random, + get_ranks_per_node, + intra_inter_halo, ) from benchmarks.nn_layer_common import GCNLayer, EdgeConditionedLayer @@ -133,25 +134,6 @@ def _gen_graph(graph_type, num_vertices, avg_degree, sbm_inter_density, seed): return gen_sbm(num_vertices, avg_degree, sbm_inter_density, rng) -def _intra_inter_halo(comm_pattern: CommunicationPattern, ranks_per_node: int) -> tuple: - """Return (intra_halo_vertices, inter_halo_vertices) from recv_offset.""" - rank = comm_pattern.rank - my_node = rank // ranks_per_node - recv_counts = ( - comm_pattern.recv_offset[1:] - comm_pattern.recv_offset[:-1] - ).tolist() - intra = 0 - inter = 0 - for r, count in enumerate(recv_counts): - if r == rank: - continue - if (r // ranks_per_node) == my_node: - intra += int(count) - else: - inter += int(count) - return intra, inter - - def _build_single_gpu_tensors(num_vertices, edges_np, F, model, device): """Allocate full-graph tensors on *device*. May raise cuda.OutOfMemoryError.""" # GCNLayer / EdgeConditionedLayer expect edge_index as [2, E] @@ -304,9 +286,7 @@ def run_distributed(args, graph_sizes, F, rank, world_size, local_rank): comm = Communicator(backend="nccl") - ranks_per_node = int( - os.environ.get("LOCAL_WORLD_SIZE", os.environ.get("SLURM_NTASKS_PER_NODE", "4")) - ) + ranks_per_node = get_ranks_per_node() measurements = [] @@ -370,7 +350,7 @@ def run_distributed(args, graph_sizes, F, rank, world_size, local_rank): all_times_multi = [None] * world_size dist.all_gather_object(all_times_multi, times_multi_local) - intra_halo, inter_halo = _intra_inter_halo(comm_pattern, ranks_per_node) + intra_halo, inter_halo = intra_inter_halo(comm_pattern, ranks_per_node) stats_local = { "rank": rank, "n_local": n_local, diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py index 9f0bb8b..47d5375 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py @@ -4,9 +4,9 @@ configurations on the full multi-node setup. Intended to be run as a SLURM array job with one invocation per (K, F, graph) combination. -This module contains a self-contained minimal halo-exchange implementation -(no dependency on the DGraph production library) so the benchmark remains -isolated and portable. +Uses DGraph's production ``build_communication_pattern``/``HaloExchange`` +(``DGraph.distributed``) for graph partitioning/halo-exchange, matching +bench_crossover.py, so the benchmark measures the real communication path. Synthetic graphs: * ``erdos_renyi`` — Erdős-Rényi with ``--avg-degree`` expected degree @@ -57,11 +57,12 @@ partition_balanced, partition_metis, partition_random, - build_local_comm_pattern, + get_ranks_per_node, + intra_inter_halo, ) from benchmarks.nn_layer_common import GCNLayer, EdgeConditionedLayer -from DGraph.distributed import HaloExchange, CommunicationPattern +from DGraph.distributed import HaloExchange, CommunicationPattern, build_communication_pattern from DGraph import Communicator # =========================================================================== @@ -115,21 +116,16 @@ def main(): else: assignment = partition_metis(args.num_vertices, world_size, edges) - # --- Build local comm pattern --- - pattern = build_local_comm_pattern(edges, assignment, rank, world_size) - n_local = pattern["n_local"] - n_halo = pattern["n_halo"] - edge_index = pattern["edge_index"].to(device) - - send_counts = pattern["send_counts"] - recv_counts = pattern["recv_counts"] - send_idx_flat = ( - torch.cat( - [torch.tensor(s, dtype=torch.long) for s in pattern["send_idx_by_rank"]] - ).to(device) - if sum(send_counts) > 0 - else torch.zeros(0, dtype=torch.long, device=device) - ) + # --- Build local comm pattern (collective: internally calls dist.all_gather) --- + edges_t = torch.from_numpy(edges).long().to(device) # [E, 2] + assignment_t = torch.from_numpy(assignment).long().to(device) # [V] + comm_pattern = build_communication_pattern(edges_t, assignment_t, rank, world_size) + + n_local = comm_pattern.num_local_vertices + n_halo = comm_pattern.num_halo_vertices + edge_index = comm_pattern.local_edge_list.T.contiguous() # [2, E_local] + + ranks_per_node = get_ranks_per_node() # --- Model --- if args.model == "gcn": @@ -152,7 +148,7 @@ def main(): # --- Timed forward + backward --- def one_layer(): # Forward halo exchange - recv_buf = halo_exchange(x_local, comm_pattern=pattern) + recv_buf = halo_exchange(x_local, comm_pattern) # Augment: local + halo x_aug = torch.cat([x_local, recv_buf], dim=0) # Message passing @@ -172,17 +168,18 @@ def one_layer(): dist.barrier() # Gather per-rank times and stats to rank 0 + intra_halo, inter_halo = intra_inter_halo(comm_pattern, ranks_per_node) stats_local = { "rank": rank, "n_local": n_local, "n_halo": n_halo, - "intra_halo_size": pattern["intra_halo_size"], - "inter_halo_size": pattern["inter_halo_size"], - "c_intra_bytes": pattern["intra_halo_size"] * F * 4, - "c_inter_bytes": pattern["inter_halo_size"] * F * 4, + "intra_halo_size": intra_halo, + "inter_halo_size": inter_halo, + "c_intra_bytes": intra_halo * F * 4, + "c_inter_bytes": inter_halo * F * 4, "n_edges_local": edge_index.shape[1], - "send_total": sum(send_counts), - "recv_total": sum(recv_counts), + "send_total": int(comm_pattern.send_offset[-1].item()), + "recv_total": int(comm_pattern.recv_offset[-1].item()), "trials_seconds": times_local, } @@ -208,7 +205,7 @@ def one_layer(): "model": args.model, "partitioner": args.partitioner, "world_size": world_size, - "ranks_per_node": pattern["ranks_per_node"], + "ranks_per_node": ranks_per_node, "warmup": args.warmup, "trials": args.trials, "seed": args.seed, diff --git a/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py b/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py index 68db2e4..18a9216 100644 --- a/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py +++ b/experiments/cost_model_benchmarks/benchmarks/graph_data_common.py @@ -97,94 +97,29 @@ def partition_metis( # =========================================================================== -# Minimal halo-exchange infrastructure +# Communication-pattern derived stats +# +# The comm pattern itself (local vertex/edge remapping, send/recv CSR +# indexing, per-rank comm map) is built exclusively via +# ``DGraph.distributed.build_communication_pattern`` — see bench_crossover.py +# and bench_end_to_end.py. This module previously hand-rolled its own +# ``build_local_comm_pattern`` with an independent (and inconsistent) halo/ +# send/recv layout that was never a ``CommunicationPattern`` and was not +# ordered per-rank the way ``recv_offset``'s CSR layout requires, which made +# it unusable as input to ``DGraph.distributed.HaloExchange``. Only the +# small derived-stats helpers below remain; they consume a real +# ``CommunicationPattern``. # =========================================================================== -def build_local_comm_pattern( - edges: np.ndarray, assignment: np.ndarray, rank: int, world_size: int -): - """Compute the local communication pattern for this rank. - - Returns a CommunicationPattern object with: - local_vertices — np.ndarray of vertex IDs owned by this rank - local_edge_index — torch.Tensor [2, E_local] with local vertex IDs - remapped so that 0..n_local-1 are owned vertices - and n_local..n_local+n_halo-1 are halo vertices - send_counts — list[int] of length world_size: vertices to send - recv_counts — list[int] of length world_size: vertices to recv - send_idx — local indices (into local_vertices) to send per rank - halo_global_ids — global vertex IDs of halo vertices, in recv order - intra_halo_size — halo vertices from same node (ranks sharing node) - inter_halo_size — halo vertices from remote nodes - ranks_per_node — int (derived from LOCAL_RANK / RANK relationship) - """ - local_mask = assignment == rank - local_vertices = np.where(local_mask)[0] - n_local = len(local_vertices) - - # Global -> local index map - g2l = {int(v): i for i, v in enumerate(local_vertices)} - - # Find edges where dst is local - local_dst_mask = np.isin(edges[:, 1], local_vertices) - local_edges = edges[local_dst_mask] - - # Halo: src vertices not owned by this rank - halo_src_mask = ~np.isin(local_edges[:, 0], local_vertices) - halo_global = np.unique(local_edges[halo_src_mask, 0]) - - # Group halo vertices by owning rank - halo_owners = assignment[halo_global] - recv_by_rank = [] - halo_order = [] - for r in range(world_size): - verts = halo_global[halo_owners == r] - recv_by_rank.append(verts) - halo_order.extend(verts.tolist()) - halo_order = np.array(halo_order, dtype=np.int64) - - # Global halo id -> local halo index - halo_g2l = {int(v): n_local + i for i, v in enumerate(halo_order)} - all_g2l = {**g2l, **halo_g2l} - - # Find which local vertices other ranks need (send pattern) - # We exchange recv_counts via all_to_all to learn send_counts - recv_counts = [len(rv) for rv in recv_by_rank] - - # Build send: for each rank r, which of our local vertices does r need? - # We do a global exchange of halo_global per rank - all_recv = [None] * world_size - dist.all_gather_object(all_recv, halo_order.tolist()) - - send_idx_by_rank = [] - for r in range(world_size): - needed = np.array(all_recv[r], dtype=np.int64) - owned_mask = ( - assignment[needed] == rank if len(needed) > 0 else np.array([], dtype=bool) - ) - owned = needed[owned_mask] if len(needed) > 0 else np.array([], dtype=np.int64) - # Map to local indices - local_idxs = np.array([g2l[int(v)] for v in owned], dtype=np.int64) - send_idx_by_rank.append(local_idxs) - - send_counts = [len(s) for s in send_idx_by_rank] - - # Remap edges to local indices - valid_edge_mask = np.array( - [(int(s) in all_g2l) and (int(d) in all_g2l) for s, d in local_edges] - ) - local_edges_valid = local_edges[valid_edge_mask] - if len(local_edges_valid) > 0: - remapped_src = np.array([all_g2l[int(s)] for s in local_edges_valid[:, 0]]) - remapped_dst = np.array([all_g2l[int(d)] for d in local_edges_valid[:, 1]]) - edge_index = torch.tensor( - np.stack([remapped_src, remapped_dst], axis=0), dtype=torch.long - ) - else: - edge_index = torch.zeros((2, 0), dtype=torch.long) +def get_ranks_per_node() -> int: + """Return the number of ranks co-located on this node. - # Compute intra / inter halo sizes + Required to split a comm pattern's halo traffic into intra-/inter-node + volumes (``c_intra_bytes`` vs ``c_inter_bytes``) for the hierarchical + cost model. Refuses to guess: silently defaulting this would bias every + downstream fit. + """ ranks_per_node_str = os.environ.get( "LOCAL_WORLD_SIZE", os.environ.get("SLURM_NTASKS_PER_NODE") ) @@ -198,27 +133,27 @@ def build_local_comm_pattern( "LOCAL_WORLD_SIZE) or under srun/SLURM (which sets " "SLURM_NTASKS_PER_NODE)." ) - ranks_per_node = int(ranks_per_node_str) + return int(ranks_per_node_str) + + +def intra_inter_halo(comm_pattern, ranks_per_node: int) -> tuple: + """Split a ``CommunicationPattern``'s halo receive counts into + (intra_node, inter_node) vertex totals, using ``recv_offset`` — the + authoritative per-source-rank CSR layout ``build_communication_pattern`` + produces — rather than re-deriving halo ownership by hand. + """ + rank = comm_pattern.rank my_node = rank // ranks_per_node - intra_halo_size = 0 - inter_halo_size = 0 - for r, verts in enumerate(recv_by_rank): - peer_node = r // ranks_per_node - if peer_node == my_node: - intra_halo_size += len(verts) + recv_counts = ( + comm_pattern.recv_offset[1:] - comm_pattern.recv_offset[:-1] + ).tolist() + intra = 0 + inter = 0 + for r, count in enumerate(recv_counts): + if r == rank: + continue + if (r // ranks_per_node) == my_node: + intra += int(count) else: - inter_halo_size += len(verts) - - return { - "local_vertices": local_vertices, - "n_local": n_local, - "n_halo": len(halo_order), - "edge_index": edge_index, - "send_counts": send_counts, - "recv_counts": recv_counts, - "send_idx_by_rank": send_idx_by_rank, - "halo_order": halo_order, - "intra_halo_size": intra_halo_size, - "inter_halo_size": inter_halo_size, - "ranks_per_node": ranks_per_node, - } + inter += int(count) + return intra, inter From 6181ab0819d72f32650e102f734b6f78e83562a5 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 31 Jul 2026 17:29:06 -0400 Subject: [PATCH 25/39] Remove inter-node requirements --- .../analysis/fit_overhead.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/experiments/cost_model_benchmarks/analysis/fit_overhead.py b/experiments/cost_model_benchmarks/analysis/fit_overhead.py index f623be1..e43243e 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_overhead.py +++ b/experiments/cost_model_benchmarks/analysis/fit_overhead.py @@ -113,9 +113,23 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, net = primitives.get("network", {}) def net_time(nbytes: int, mode: str) -> float: - params = net.get(mode, None) - if params is None or nbytes == 0: + if nbytes == 0: return 0.0 + params = net.get(mode, None) + if params is None: + # A missing network fit is fine when this mode never carries + # traffic (e.g. single-node runs never have inter-node bytes, + # so "inter" is legitimately absent). But if a run actually + # has nonzero bytes for this mode, silently treating them as + # zero-cost would bias T_comm and every downstream prediction/ + # overhead fit without any signal that it happened. + raise ValueError( + f"per_rank_stats has {nbytes} bytes of {mode}-node " + f"communication but fitted_primitives.json has no '{mode}' " + f"network fit. Re-run fit_primitives.py with " + f"--pingpong-{mode} data, or confirm this run should never " + f"have {mode}-node traffic in the first place." + ) B = params.get("bandwidth_bytes_per_sec", 1e10) t_L = params.get("latency_seconds", 0.0) return t_L + nbytes / B From 3ad6957c72467103dc0eaa477e264ea9078a9463 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 31 Jul 2026 17:34:13 -0400 Subject: [PATCH 26/39] Add experiment script --- .../cost_model_benchmarks/run_experiments.sh | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100755 experiments/cost_model_benchmarks/run_experiments.sh diff --git a/experiments/cost_model_benchmarks/run_experiments.sh b/experiments/cost_model_benchmarks/run_experiments.sh new file mode 100755 index 0000000..aeb7914 --- /dev/null +++ b/experiments/cost_model_benchmarks/run_experiments.sh @@ -0,0 +1,148 @@ +#!/usr/bin/env bash +# Step-by-step driver for the cost-model benchmark pipeline (see README.md's +# "Pipeline" section) on a single node, sweeping the distributed +# (torchrun-launched) benchmarks over GPU_COUNTS. +# +# Single-node only: no inter-node ping-pong/concurrency/crossover runs are +# included, since this script never launches more than one node. +# +# Usage: +# ./run_experiments.sh +# +# Override any config variable via the environment, e.g.: +# GPU_COUNTS="2 4 8" MODEL=edge ./run_experiments.sh +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")" + +# --- Config ------------------------------------------------------------ +GPU_COUNTS=(${GPU_COUNTS:-2 4}) # world sizes to sweep via torchrun --nproc_per_node +SEED=${SEED:-42} +FEATURE_DIM=${FEATURE_DIM:-128} +MODEL=${MODEL:-gcn} # gcn | edge +GRAPH=${GRAPH:-erdos_renyi} # erdos_renyi | sbm +PARTITIONER=${PARTITIONER:-balanced} # random | balanced | metis +WARMUP=${WARMUP:-10} +TRIALS=${TRIALS:-50} +CROSSOVER_SIZES=${CROSSOVER_SIZES:-1000,10000,100000,1000000} +E2E_VERTEX_SIZES=(${E2E_VERTEX_SIZES:-10000 100000 1000000}) +DATA_DIR=${DATA_DIR:-data} +FIG_DIR=${FIG_DIR:-figures} + +mkdir -p "$DATA_DIR" "$FIG_DIR" + +step() { echo; echo "=== $* ==="; } + +# ------------------------------------------------------------------------- +# 1. Single-GPU primitive microbenchmarks (1.3 compute, 1.4 gather) — no +# torchrun, GPU-count independent. +# ------------------------------------------------------------------------- +step "1.3 compute -- ${MODEL}, vertex sweep" +python -m benchmarks.bench_compute --model "${MODEL}" --sweep vertices \ + --min 1000 --max 1000000 --steps 10 --fixed-value 200000 \ + --feature-dim "${FEATURE_DIM}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/compute_${MODEL}_vswp.json" --seed "${SEED}" + +step "1.3 compute -- ${MODEL}, edge sweep" +python -m benchmarks.bench_compute --model "${MODEL}" --sweep edges \ + --min 1000 --max 1000000 --steps 10 --fixed-value 200000 \ + --feature-dim "${FEATURE_DIM}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/compute_${MODEL}_eswp.json" --seed "${SEED}" + +step "1.4 gather -- contiguous / clustered / random" +for dist_name in contiguous clustered random; do + python -m benchmarks.bench_gather --distribution "${dist_name}" \ + --min-k 1000 --max-k 10000000 --steps 20 --N 20000000 \ + --feature-dim "${FEATURE_DIM}" --cluster-size 64 \ + --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/gather_${dist_name}.json" --seed "${SEED}" +done + +# ------------------------------------------------------------------------- +# 2. Intra-node ping-pong (1.1) -- fixed 2-rank NVLink/PCIe pair test, +# independent of GPU_COUNTS. +# ------------------------------------------------------------------------- +step "1.1 pingpong -- intra-node (2 ranks)" +torchrun --nnodes 1 --nproc_per_node 2 -m benchmarks.bench_pingpong \ + --mode intra --min-bytes 64 --max-bytes 67108864 --steps 21 \ + --warmup 20 --trials 100 \ + --output "${DATA_DIR}/pingpong_intra.json" --seed "${SEED}" + +# ------------------------------------------------------------------------- +# 3. Distributed benchmarks (2.1 end-to-end, 2.2 crossover), swept over +# GPU_COUNTS via torchrun. +# ------------------------------------------------------------------------- +for K in "${GPU_COUNTS[@]}"; do + step "2.2 crossover -- K=${K} GPUs" + torchrun --nnodes 1 --nproc_per_node "${K}" -m benchmarks.bench_crossover \ + --graph "${GRAPH}" --graph-sizes "${CROSSOVER_SIZES}" \ + --avg-degree 20 --feature-dim "${FEATURE_DIM}" --model "${MODEL}" \ + --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/crossover_K${K}.json" --seed "${SEED}" + + for N in "${E2E_VERTEX_SIZES[@]}"; do + step "2.1 end-to-end -- K=${K} GPUs, N=${N}" + torchrun --nnodes 1 --nproc_per_node "${K}" -m benchmarks.bench_end_to_end \ + --graph "${GRAPH}" --num-vertices "${N}" --avg-degree 20 \ + --feature-dim "${FEATURE_DIM}" --model "${MODEL}" \ + --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/e2e_K${K}_N${N}.json" --seed "${SEED}" + done +done + +# ------------------------------------------------------------------------- +# 4. Fit primitives, fit overhead, apply the assembled model. +# ------------------------------------------------------------------------- +step "fit_primitives" +python -m analysis.fit_primitives \ + --pingpong-intra "${DATA_DIR}/pingpong_intra.json" \ + --compute-gcn "${DATA_DIR}"/compute_gcn_*.json \ + --compute-edge "${DATA_DIR}"/compute_edge_*.json \ + --gather-contiguous "${DATA_DIR}/gather_contiguous.json" \ + --gather-clustered "${DATA_DIR}/gather_clustered.json" \ + --gather-random "${DATA_DIR}/gather_random.json" \ + --output "${DATA_DIR}/fitted_primitives.json" + +step "fit_overhead" +python -m analysis.fit_overhead \ + --primitives "${DATA_DIR}/fitted_primitives.json" \ + --e2e-runs "${DATA_DIR}"/e2e_*.json \ + --fit-filter "world_size <= 8" \ + --output "${DATA_DIR}/fitted_overhead.json" + +step "compute_predictions" +python -m analysis.compute_predictions \ + --primitives "${DATA_DIR}/fitted_primitives.json" \ + --overhead "${DATA_DIR}/fitted_overhead.json" \ + --e2e-runs "${DATA_DIR}"/e2e_*.json \ + --fit-filter "world_size <= 8" \ + --output "${DATA_DIR}/predictions.json" + +# ------------------------------------------------------------------------- +# 5. Visualization. +# ------------------------------------------------------------------------- +step "visualization" +python -m visualization.plot_compute \ + --gcn-vertex "${DATA_DIR}/compute_gcn_vswp.json" --gcn-edge "${DATA_DIR}/compute_gcn_eswp.json" \ + --edge-vertex "${DATA_DIR}/compute_edge_vswp.json" --edge-edge "${DATA_DIR}/compute_edge_eswp.json" \ + --primitives "${DATA_DIR}/fitted_primitives.json" --output "${FIG_DIR}/compute" + +python -m visualization.plot_gather \ + --contiguous "${DATA_DIR}/gather_contiguous.json" --clustered "${DATA_DIR}/gather_clustered.json" \ + --random "${DATA_DIR}/gather_random.json" --fitted "${DATA_DIR}/fitted_primitives.json" \ + --output "${FIG_DIR}/gather" + +python -m visualization.plot_pingpong \ + --intra "${DATA_DIR}/pingpong_intra.json" --primitives "${DATA_DIR}/fitted_primitives.json" \ + --output "${FIG_DIR}/pingpong" + +for K in "${GPU_COUNTS[@]}"; do + python -m visualization.plot_crossover --input "${DATA_DIR}/crossover_K${K}.json" \ + --output "${FIG_DIR}/crossover_K${K}.png" +done + +python -m visualization.plot_validation --predictions "${DATA_DIR}/predictions.json" \ + --color-by world_size --output "${FIG_DIR}/validation" + +echo +echo "Done. Data in ${DATA_DIR}/, figures in ${FIG_DIR}/." From a5174cfc8885024c3e822a590be9f295cd03d66a Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 31 Jul 2026 17:42:44 -0400 Subject: [PATCH 27/39] Fix shell script with syntax error --- .../cost_model_benchmarks/run_experiments.sh | 34 +++++++++++++------ 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/experiments/cost_model_benchmarks/run_experiments.sh b/experiments/cost_model_benchmarks/run_experiments.sh index aeb7914..d41589e 100755 --- a/experiments/cost_model_benchmarks/run_experiments.sh +++ b/experiments/cost_model_benchmarks/run_experiments.sh @@ -15,6 +15,11 @@ set -euo pipefail cd "$(dirname "${BASH_SOURCE[0]}")" +# An unmatched glob (e.g. a partial re-run missing an expected output file) +# should surface as a clear argparse error, not a literal glob-pattern +# string silently passed through to a file open() call. +shopt -s nullglob + # --- Config ------------------------------------------------------------ GPU_COUNTS=(${GPU_COUNTS:-2 4}) # world sizes to sweep via torchrun --nproc_per_node SEED=${SEED:-42} @@ -36,18 +41,25 @@ step() { echo; echo "=== $* ==="; } # ------------------------------------------------------------------------- # 1. Single-GPU primitive microbenchmarks (1.3 compute, 1.4 gather) — no # torchrun, GPU-count independent. +# +# fit_primitives.py fits GCN and edge-conditioned compute costs +# independently (--compute-gcn / --compute-edge), regardless of which +# single MODEL the distributed crossover/end-to-end sweep below uses — +# so both model types are always benchmarked here. # ------------------------------------------------------------------------- -step "1.3 compute -- ${MODEL}, vertex sweep" -python -m benchmarks.bench_compute --model "${MODEL}" --sweep vertices \ - --min 1000 --max 1000000 --steps 10 --fixed-value 200000 \ - --feature-dim "${FEATURE_DIM}" --warmup "${WARMUP}" --trials "${TRIALS}" \ - --output "${DATA_DIR}/compute_${MODEL}_vswp.json" --seed "${SEED}" - -step "1.3 compute -- ${MODEL}, edge sweep" -python -m benchmarks.bench_compute --model "${MODEL}" --sweep edges \ - --min 1000 --max 1000000 --steps 10 --fixed-value 200000 \ - --feature-dim "${FEATURE_DIM}" --warmup "${WARMUP}" --trials "${TRIALS}" \ - --output "${DATA_DIR}/compute_${MODEL}_eswp.json" --seed "${SEED}" +for m in gcn edge; do + step "1.3 compute -- ${m}, vertex sweep" + python -m benchmarks.bench_compute --model "${m}" --sweep vertices \ + --min 1000 --max 1000000 --steps 10 --fixed-value 200000 \ + --feature-dim "${FEATURE_DIM}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/compute_${m}_vswp.json" --seed "${SEED}" + + step "1.3 compute -- ${m}, edge sweep" + python -m benchmarks.bench_compute --model "${m}" --sweep edges \ + --min 1000 --max 1000000 --steps 10 --fixed-value 200000 \ + --feature-dim "${FEATURE_DIM}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/compute_${m}_eswp.json" --seed "${SEED}" +done step "1.4 gather -- contiguous / clustered / random" for dist_name in contiguous clustered random; do From d96d01c5be70f10f37829163bf4e5dde8dd2ccaa Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 31 Jul 2026 17:52:51 -0400 Subject: [PATCH 28/39] Update primitives to handle NaNs --- .../cost_model_benchmarks/analysis/fit_primitives.py | 2 +- experiments/cost_model_benchmarks/run_experiments.sh | 9 +++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/experiments/cost_model_benchmarks/analysis/fit_primitives.py b/experiments/cost_model_benchmarks/analysis/fit_primitives.py index 0949a91..5922eb2 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_primitives.py +++ b/experiments/cost_model_benchmarks/analysis/fit_primitives.py @@ -246,7 +246,7 @@ def time_model(b, overhead, inv_bw_L2, inv_bw_HBM, L2_thresh, HBM_thresh): except Exception as e: print(f"Fit failed: {e}") - overhead = inv_bw_L2 = inv_bw_HBM = L2_thresh = np.nan + overhead = inv_bw_L2 = inv_bw_HBM = L2_thresh = HBM_thresh = np.nan bw_HBM = 1.0 / inv_bw_HBM if inv_bw_HBM > 0 else float("nan") bw_L2 = 1.0 / inv_bw_L2 if inv_bw_L2 > 0 else float("nan") diff --git a/experiments/cost_model_benchmarks/run_experiments.sh b/experiments/cost_model_benchmarks/run_experiments.sh index d41589e..0038831 100755 --- a/experiments/cost_model_benchmarks/run_experiments.sh +++ b/experiments/cost_model_benchmarks/run_experiments.sh @@ -106,10 +106,15 @@ done # 4. Fit primitives, fit overhead, apply the assembled model. # ------------------------------------------------------------------------- step "fit_primitives" +# NOTE: explicit filenames, not a wildcard glob. A glob like +# compute_gcn_*.json would also match unrelated stale files that happen to +# share the prefix (e.g. data/compute_gcn_eswp_test.json left over from +# run_local_compute_tests.sh, which uses a different --feature-dim) and +# silently pool incompatible data into one regression. python -m analysis.fit_primitives \ --pingpong-intra "${DATA_DIR}/pingpong_intra.json" \ - --compute-gcn "${DATA_DIR}"/compute_gcn_*.json \ - --compute-edge "${DATA_DIR}"/compute_edge_*.json \ + --compute-gcn "${DATA_DIR}/compute_gcn_vswp.json" "${DATA_DIR}/compute_gcn_eswp.json" \ + --compute-edge "${DATA_DIR}/compute_edge_vswp.json" "${DATA_DIR}/compute_edge_eswp.json" \ --gather-contiguous "${DATA_DIR}/gather_contiguous.json" \ --gather-clustered "${DATA_DIR}/gather_clustered.json" \ --gather-random "${DATA_DIR}/gather_random.json" \ From f0ed3cd254afe394f68415af57a061fc9ffeed83 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 31 Jul 2026 18:07:20 -0400 Subject: [PATCH 29/39] Remove unnecessary e2e file --- .../cost_model_benchmarks/run_experiments.sh | 32 +++++++++++++++---- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/experiments/cost_model_benchmarks/run_experiments.sh b/experiments/cost_model_benchmarks/run_experiments.sh index 0038831..c33c0ac 100755 --- a/experiments/cost_model_benchmarks/run_experiments.sh +++ b/experiments/cost_model_benchmarks/run_experiments.sh @@ -34,6 +34,14 @@ E2E_VERTEX_SIZES=(${E2E_VERTEX_SIZES:-10000 100000 1000000}) DATA_DIR=${DATA_DIR:-data} FIG_DIR=${FIG_DIR:-figures} +# fit_overhead.py / compute_predictions.py default to "world_size <= 8", +# which never holds anything out for a single-node run limited to <=8 GPUs +# (held-out MAPE would always report NaN over an empty set). Hold out the +# largest swept GPU count instead, so held-out MAPE reflects genuine +# extrapolation across the GPU_COUNTS actually being run. +MAX_GPU_COUNT=$(printf '%s\n' "${GPU_COUNTS[@]}" | sort -n | tail -1) +FIT_FILTER=${FIT_FILTER:-"world_size < ${MAX_GPU_COUNT}"} + mkdir -p "$DATA_DIR" "$FIG_DIR" step() { echo; echo "=== $* ==="; } @@ -102,6 +110,18 @@ for K in "${GPU_COUNTS[@]}"; do done done +# Explicit list of exactly the e2e files this script just produced — not a +# e2e_*.json wildcard glob, which would also match any stray file sharing +# that prefix (e.g. a leftover data/e2e_K8_F128_er_bal.json from manually +# running a docstring example) and silently pool it into the fit/held-out +# split as if it were part of this run's matrix. +E2E_FILES=() +for K in "${GPU_COUNTS[@]}"; do + for N in "${E2E_VERTEX_SIZES[@]}"; do + E2E_FILES+=("${DATA_DIR}/e2e_K${K}_N${N}.json") + done +done + # ------------------------------------------------------------------------- # 4. Fit primitives, fit overhead, apply the assembled model. # ------------------------------------------------------------------------- @@ -120,19 +140,19 @@ python -m analysis.fit_primitives \ --gather-random "${DATA_DIR}/gather_random.json" \ --output "${DATA_DIR}/fitted_primitives.json" -step "fit_overhead" +step "fit_overhead (fit-filter: ${FIT_FILTER})" python -m analysis.fit_overhead \ --primitives "${DATA_DIR}/fitted_primitives.json" \ - --e2e-runs "${DATA_DIR}"/e2e_*.json \ - --fit-filter "world_size <= 8" \ + --e2e-runs "${E2E_FILES[@]}" \ + --fit-filter "${FIT_FILTER}" \ --output "${DATA_DIR}/fitted_overhead.json" -step "compute_predictions" +step "compute_predictions (fit-filter: ${FIT_FILTER})" python -m analysis.compute_predictions \ --primitives "${DATA_DIR}/fitted_primitives.json" \ --overhead "${DATA_DIR}/fitted_overhead.json" \ - --e2e-runs "${DATA_DIR}"/e2e_*.json \ - --fit-filter "world_size <= 8" \ + --e2e-runs "${E2E_FILES[@]}" \ + --fit-filter "${FIT_FILTER}" \ --output "${DATA_DIR}/predictions.json" # ------------------------------------------------------------------------- From a3442be1f6166f3a122b60f34be7a8a9d6806942 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 31 Jul 2026 18:12:08 -0400 Subject: [PATCH 30/39] Fix bug in compute code that double counted forward pass --- .../analysis/fit_overhead.py | 15 +++++++++++++-- .../benchmarks/bench_compute.py | 7 +++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/experiments/cost_model_benchmarks/analysis/fit_overhead.py b/experiments/cost_model_benchmarks/analysis/fit_overhead.py index e43243e..23b81df 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_overhead.py +++ b/experiments/cost_model_benchmarks/analysis/fit_overhead.py @@ -96,8 +96,19 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, avg_degree = run_config.get("avg_degree", 20.0) n_edges_local = int(n_local * avg_degree) - # T_comp - comp_params = primitives.get("compute", {}).get(model_type, {}).get("forward", None) + # T_comp — the end-to-end benchmark times forward + backward + grad-zero + # per iteration (bench_end_to_end.py::one_layer), so T_comp must be the + # forward+backward cost, not forward alone. + # + # NOTE: the fit stored under the "backward" key is ALREADY that combined + # cost. bench_compute.py's bwd() closure re-runs the forward pass inside + # its own timed region (it has to, to rebuild the autograd graph each + # trial), so "backward_trials_seconds" measures grad-zero + forward + + # backward — exactly the op set one_layer() performs. The key is a + # misnomer; "forward" is the only fit that is forward-alone. + # Using "forward" here undercounted T_comp and cost ~30% MAPE; adding + # forward+backward together instead would double-count the forward pass. + comp_params = primitives.get("compute", {}).get(model_type, {}).get("backward", None) if comp_params: T_comp = (comp_params["coeff_V"] * n_total + comp_params["coeff_E"] * n_edges_local diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_compute.py b/experiments/cost_model_benchmarks/benchmarks/bench_compute.py index 89f5249..ea7ed6d 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_compute.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_compute.py @@ -176,6 +176,13 @@ def fwd(): out = model(x, edge_index, edge_attr) loss_ref = out.sum() + # NOTE: this closure times grad-zero + FORWARD + backward, not the + # backward pass alone — the forward must be re-run inside the timed + # region to rebuild the autograd graph consumed by each .backward() + # call. So "backward_trials_seconds" below is really the combined + # forward+backward cost (which is what the assembled cost model's + # T_comp needs, since bench_end_to_end.py times the same op set). + # Do not subtract or add "forward_trials_seconds" to it. def bwd(): if x.grad is not None: x.grad.zero_() From 90762f62e0f25e07a5c76d77174c0e556bd5ac52 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 31 Jul 2026 19:00:03 -0400 Subject: [PATCH 31/39] Multi-node experiemts + removing dead code --- experiments/OGB/GCN.py | 183 ++------------ experiments/OGB/benchmark_OGB_end_to_end.py | 1 - experiments/OGB/main.py | 1 - .../analysis/fit_overhead.py | 37 ++- .../analysis/fit_primitives.py | 81 +++---- .../benchmarks/bench_concurrency.py | 12 + .../benchmarks/bench_pingpong.py | 11 + .../benchmarks/common.py | 54 +++++ .../run_experiments_multinode.sh | 229 ++++++++++++++++++ .../visualization/plot_gather.py | 87 +++---- 10 files changed, 413 insertions(+), 283 deletions(-) create mode 100755 experiments/cost_model_benchmarks/run_experiments_multinode.sh diff --git a/experiments/OGB/GCN.py b/experiments/OGB/GCN.py index a15ecb4..210722d 100644 --- a/experiments/OGB/GCN.py +++ b/experiments/OGB/GCN.py @@ -11,28 +11,13 @@ # https://github.com/LBANN and https://github.com/LLNL/LBANN. # # SPDX-License-Identifier: (Apache-2.0) -import torch -import torch.nn as nn - -from DGraph.distributed import HaloExchange, CommunicationPattern -from DGraph.utils.TimingReport import TimingReport -from DGraph import Communicator from typing import Union -try: - from torch_geometric.nn.conv import GCNConv - - PYG_AVAILABLE = True -except ImportError: - print( - "torch_geometric not found, skipping import of GCNConv. Make sure to install torch_geometric if you want to use it." - ) - PYG_AVAILABLE = False - - import torch import torch.nn as nn +from DGraph.utils.TimingReport import TimingReport + def create_sparse_adj( edge_index, @@ -42,11 +27,13 @@ def create_sparse_adj( ): """ Converts an edge_index of shape [num_edges, 2] into a PyTorch sparse tensor. + + Follows ``DGraph.distributed.commInfo``'s edge-list convention: column 1 is + the destination, which is always a locally-owned vertex, and column 0 is the + source/neighbour, which may be a halo vertex. Hence the adjacency is + rectangular: [num_local_nodes, num_local_nodes + num_halo_nodes]. """ # PyTorch sparse tensors expect indices in shape [2, num_edges] - # So we transpose the [num_edges, 2] tensor - - indices = edge_index.t().contiguous() indices = torch.stack( [ edge_index[:, 1], # Rows: Targets (strictly < num_local_nodes) @@ -74,17 +61,26 @@ def create_sparse_adj( return adj_sparse_csr -class GCNLayer_impl_sparse(nn.Module): +class GCNLayer(nn.Module): + """GCN layer using a sparse adjacency matmul for message passing. + + Transform-then-aggregate: a dense [V, F] x [F, F] linear followed by an + SpMM against the rectangular local adjacency. The SpMM consumes the + augmented [num_local + num_halo, F] feature matrix and emits only the + num_local rows, so the tensor does not grow from layer to layer. + """ + def __init__(self, in_channels, out_channels): - super(GCNLayer_impl_sparse, self).__init__() + super(GCNLayer, self).__init__() self.linear = nn.Linear(in_channels, out_channels) self.act = nn.ReLU(inplace=True) def forward(self, x, adj_sparse): """ Args: - x: Node features tensor of shape [num_nodes, in_channels] - adj_sparse: Sparse adjacency tensor of shape [num_nodes, num_nodes] + x: Node features tensor of shape [num_local + num_halo, in_channels] + adj_sparse: Sparse adjacency of shape + [num_local, num_local + num_halo] """ # 1. Feature transformation (Dense) x = self.linear(x) @@ -99,43 +95,6 @@ def forward(self, x, adj_sparse): return out -class GCNLayer_impl(nn.Module): - def __init__(self, in_channels, out_channels): - super(GCNLayer_impl, self).__init__() - self.linear = nn.Linear(in_channels, out_channels) - self.act = nn.ReLU(inplace=True) - - def forward(self, x, edge_index): - source_vertices = edge_index[:, 0] - target_vertices = edge_index[:, 1] - x = self.linear(x) - - x_j = x[target_vertices, :] - out_channels = x_j.size(1) - out = torch.zeros(x.size(0), out_channels, dtype=x.dtype, device=x.device) - scatter_index = ( - source_vertices.unsqueeze(-1).expand(-1, out_channels).to(x.device) - ) - out = out.scatter_add(0, scatter_index, x_j) - out = self.act(out) - return out - - -class GCNLayer(nn.Module): - def __init__(self, in_channels, out_channels, use_sparse=False): - super(GCNLayer, self).__init__() - if False: - self.conv = GCNConv(in_channels, out_channels) - else: - if use_sparse: - self.conv = GCNLayer_impl_sparse(in_channels, out_channels) - else: - self.conv = GCNLayer_impl(in_channels, out_channels) - - def forward(self, x, edge_index): - return self.conv(x, edge_index) - - class GCNModel(nn.Module): def __init__( self, @@ -144,17 +103,18 @@ def __init__( out_channels, num_layers, halo_exchanger, - is_sparse=False, ): super(GCNModel, self).__init__() self.convs = nn.ModuleList() - self.convs.append(GCNLayer(in_channels, hidden_channels, is_sparse)) + self.convs.append(GCNLayer(in_channels, hidden_channels)) for _ in range(num_layers - 2): - self.convs.append(GCNLayer(hidden_channels, hidden_channels, is_sparse)) - self.convs.append(GCNLayer(hidden_channels, out_channels, is_sparse)) + self.convs.append(GCNLayer(hidden_channels, hidden_channels)) + self.convs.append(GCNLayer(hidden_channels, out_channels)) self.halo_exchanger = halo_exchanger def forward(self, x, comm_pattern): + # comm_pattern.local_edge_list must already have been converted to a + # sparse CSR adjacency via create_sparse_adj(). edge_index = comm_pattern.local_edge_list counter = 1 for conv in self.convs[:-1]: @@ -174,96 +134,3 @@ def forward(self, x, comm_pattern): x = torch.cat([x, boundary_features], dim=0) x = self.convs[-1](x, edge_index) return x - - -class GraphConvLayer(nn.Module): - def __init__(self, message_dim, out_channels): - super(GraphConvLayer, self).__init__() - self.conv = nn.Linear(message_dim, out_channels) - self.act = nn.ReLU(inplace=True) - - def forward(self, x, edge_index, num_local_nodes, edge_features=None): - source_vertices = edge_index[:, 0] - target_vertices = edge_index[:, 1] - - assert (source_vertices < num_local_nodes).all(), ( - f"Graph routing error: Found source_vertices >= num_local_nodes ({num_local_nodes}). " - "Boundary nodes must only act as targets (x_j) in this aggregation scheme!" - ) - - x_i = x[source_vertices, :] - x_j = x[target_vertices, :] - - if edge_features is not None: - x_ij = torch.cat([x_i, x_j, edge_features], dim=1) - else: - x_ij = torch.cat([x_i, x_j], dim=1) - - m_ij = self.conv(x_ij) - m_ij = self.act(m_ij) - - out_channels = m_ij.size(1) - - # 1. Allocate ONLY for the local nodes (the sources we are aggregating to) - out = torch.zeros(num_local_nodes, out_channels, dtype=x.dtype, device=x.device) - - # 2. Scatter messages back to the SOURCE vertices - scatter_index = ( - source_vertices.unsqueeze(-1).expand(-1, out_channels).to(x.device) - ) - - # 3. Perform the aggregation - out = out.scatter_add(0, scatter_index, m_ij) - - return out - - -class CommAwareGCN(nn.Module): - """ - Least interesting GNN model to test distributed training - but good enough for the purpose of testing. - """ - - def __init__( - self, - in_channels: int, - hidden_dims: int, - num_classes: int, - halo_exchanger: HaloExchange, - comm: Communicator, - ): - super(CommAwareGCN, self).__init__() - self.halo_exchanger = halo_exchanger - - self.conv1 = GraphConvLayer(2 * in_channels, hidden_dims) - - self.conv2 = GraphConvLayer(2 * hidden_dims, hidden_dims) - - self.fc = nn.Linear(hidden_dims, num_classes) - self.softmax = nn.Softmax(dim=1) - self.comm = comm - - def forward( - self, local_node_features: torch.Tensor, comm_pattern: CommunicationPattern - ): - - num_local_nodes = local_node_features.shape[0] - - with TimingReport("feature-exchange-1"): - boundary_features = self.halo_exchanger(local_node_features, comm_pattern) - - with TimingReport("process-1"): - x = torch.cat([local_node_features, boundary_features], dim=0) - x = self.conv1(x, comm_pattern.local_edge_list, num_local_nodes) - - with TimingReport("feature-exchange-2"): - boundary_features = self.halo_exchanger(x, comm_pattern) - - with TimingReport("process-2"): - x = torch.cat([x, boundary_features], dim=0) - x = self.conv2(x, comm_pattern.local_edge_list, num_local_nodes) - - with TimingReport("final-fc"): - x = self.fc(x) - # x = self.softmax(x) - return x diff --git a/experiments/OGB/benchmark_OGB_end_to_end.py b/experiments/OGB/benchmark_OGB_end_to_end.py index 9f6a526..c576f9a 100644 --- a/experiments/OGB/benchmark_OGB_end_to_end.py +++ b/experiments/OGB/benchmark_OGB_end_to_end.py @@ -74,7 +74,6 @@ def benchmark_ogb_end_to_end( out_channels=num_classes, num_layers=num_layers, halo_exchanger=halo_exchanger, - is_sparse=True, ).to(device) model = DDP(model, device_ids=[local_rank], output_device=local_rank) optimizer = optim.Adam(model.parameters(), lr=lr) diff --git a/experiments/OGB/main.py b/experiments/OGB/main.py index 381ae86..9731b5b 100644 --- a/experiments/OGB/main.py +++ b/experiments/OGB/main.py @@ -115,7 +115,6 @@ def _run_experiment( out_channels=num_classes, num_layers=3, halo_exchanger=halo_exchanger, - is_sparse=True, ).to(device) if comm.get_world_size() > 1: diff --git a/experiments/cost_model_benchmarks/analysis/fit_overhead.py b/experiments/cost_model_benchmarks/analysis/fit_overhead.py index 23b81df..de959b4 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_overhead.py +++ b/experiments/cost_model_benchmarks/analysis/fit_overhead.py @@ -32,28 +32,21 @@ # Cost model (without overhead) # --------------------------------------------------------------------------- -def _gather_piecewise_time(nbytes: float, params: dict) -> float: - """Evaluate the fitted piecewise L2-cache/HBM-bandwidth gather model for - a single byte count. Mirrors ``time_model`` in - ``analysis/fit_primitives.py::fit_gather`` exactly — this is the single - place both ``predict_layer_time`` and ``compute_predictions.py`` should - call, rather than re-deriving a simplified linear approximation. +def _gather_time(nbytes: float, params: dict) -> float: + """Evaluate the fitted Hockney gather model for a single byte count: + + T(bytes) = launch_overhead + bytes / B_gather + + Mirrors ``time_model`` in ``analysis/fit_primitives.py::fit_gather`` + exactly — this is the single place both ``predict_layer_time`` and + ``compute_predictions.py`` should call, rather than re-deriving an + approximation of it. """ overhead = params.get("launch_overhead_seconds", 0.0) - bw_HBM = params.get("bandwidth_bytes_per_sec") - bw_L2 = params.get("L2_bandwidth_bytes_per_sec") - if not bw_HBM or not bw_L2: + bw = params.get("bandwidth_bytes_per_sec") + if not bw or not np.isfinite(bw): return overhead - - inv_bw_L2 = 1.0 / bw_L2 - inv_bw_HBM = 1.0 / bw_HBM - L2_thresh = params.get("L2_inflection_bytes", 0.0) - HBM_thresh = params.get("HBM_inflection_bytes", 0.0) - - bytes_L2 = min(max(nbytes, 0.0), L2_thresh) - bytes_HBM = max(0.0, nbytes - HBM_thresh) - t_mem = bytes_L2 * inv_bw_L2 + bytes_HBM * inv_bw_HBM - return max(overhead, t_mem) + return overhead + max(nbytes, 0.0) / bw def predict_layer_time(run_config: dict, per_rank_stats: dict, @@ -149,12 +142,12 @@ def net_time(nbytes: int, mode: str) -> float: T_inter = max(net_time(inter_bytes, "inter"), 0.0) T_comm = max(T_intra, T_inter) - # T_buffer_copy (gather of send buffer) — uses the fitted piecewise - # L2/HBM model, not a simplified single-slope linear stand-in. + # T_buffer_copy (gather of send buffer) — uses the fitted Hockney gather + # model, evaluated by the same function fit_gather fits. send_bytes = per_rank_stats.get("send_total", 0) * F * 4 gath_params = primitives.get("gather", {}).get("clustered", {}).get("gather", None) if gath_params and send_bytes > 0: - T_buffer_copy = _gather_piecewise_time(send_bytes, gath_params) + T_buffer_copy = _gather_time(send_bytes, gath_params) else: T_buffer_copy = 0.0 T_buffer_copy = max(T_buffer_copy, 0.0) diff --git a/experiments/cost_model_benchmarks/analysis/fit_primitives.py b/experiments/cost_model_benchmarks/analysis/fit_primitives.py index 5922eb2..44a7478 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_primitives.py +++ b/experiments/cost_model_benchmarks/analysis/fit_primitives.py @@ -178,7 +178,7 @@ def fit_compute(records: list, timing_key: str = "forward_trials_seconds") -> di # --------------------------------------------------------------------------- -# Gather fit: T = intercept + max(overhead, bytes / B_gather) +# Gather fit: T = launch_overhead + bytes / B_gather (Hockney) # --------------------------------------------------------------------------- @@ -199,63 +199,57 @@ def fit_gather(records: list, timing_key: str = "gather_trials_seconds") -> dict bytes_arr = k_arr * F_arr * 4.0 # float32 - # 1. The Piecewise Linear Model (No Logarithms) - def time_model(b, overhead, inv_bw_L2, inv_bw_HBM, L2_thresh, HBM_thresh): - # 1. Bucket the bytes into their respective physical regimes - # Bytes processed exclusively at L2 speeds - bytes_L2 = np.clip(b, 0, L2_thresh) - # Bytes processed exclusively at HBM speeds - bytes_HBM = np.maximum(0, b - HBM_thresh) + # Hockney model: a fixed per-kernel launch/setup cost plus a linear + # bandwidth term. + # + # T(bytes) = launch_overhead + bytes / B_gather + # + # This replaces an earlier 5-parameter piecewise "L2-cache vs HBM + # bandwidth" model. That model was unidentifiable on this benchmark's + # data and fit garbage: because it capped the L2 term at L2_thresh but + # only started the HBM term at a *separate*, larger HBM_thresh, every + # byte count between the two thresholds had zero marginal cost — a flat + # plateau spanning most of the sweep. It routinely converged with + # B_L2 far SLOWER than B_HBM (physically backwards) and R² as low as + # 0.60. + # + # The measured curve shows no two-regime structure to fit in the first + # place: effective bandwidth rises smoothly to a single asymptote, and + # the smallest transfers in the sweep are already large enough that any + # cache effect is masked by the launch-overhead floor. Two parameters + # are what the data supports, and this is the same Hockney form + # fit_network uses for the interconnect (R² >= 0.999 here). + def time_model(b, overhead, inv_bw): + return overhead + b * inv_bw - # 2. Apply the specific bandwidth (slope) to each bucket - t_mem = (bytes_L2 * inv_bw_L2) + (bytes_HBM * inv_bw_HBM) - - # 3. Floor the total time by the kernel launch overhead - return np.maximum(overhead, t_mem) - - # 2. Strategic initial guesses min_T, max_T = float(np.min(T_arr)), float(np.max(T_arr)) max_b = float(np.max(bytes_arr)) - - # The asymptotic slope is the difference in max/min time over max bytes - inv_bw_HBM_guess = (max_T - min_T) / (max_b + 1e-9) - - p0 = [ - min_T, # overhead - inv_bw_HBM_guess * 0.3, # inv_bw_L2 - inv_bw_HBM_guess, # inv_bw_HBM - max_b * 0.00001, # L2_thresh - max_b * 0.001, # HBM_thresh - ] + p0 = [min_T, (max_T - min_T) / (max_b + 1e-9)] try: - bounds = ([0, 0, 0, 0, 0], [np.inf, np.inf, np.inf, max_b, max_b]) - - # 3. The Magic Fix: sigma=T_arr weights the fit by relative error + # sigma=T_arr weights residuals by relative error, consistent with + # fit_network/fit_compute — byte counts span ~4 orders of magnitude, + # so an unweighted fit would be dominated by the largest transfers. popt, _ = curve_fit( time_model, bytes_arr, T_arr, p0=p0, - bounds=bounds, + bounds=([0, 0], [np.inf, np.inf]), method="trf", sigma=T_arr, absolute_sigma=False, ) - overhead, inv_bw_L2, inv_bw_HBM, L2_thresh, HBM_thresh = popt + overhead, inv_bw = popt except Exception as e: print(f"Fit failed: {e}") - overhead = inv_bw_L2 = inv_bw_HBM = L2_thresh = HBM_thresh = np.nan + overhead = inv_bw = np.nan - bw_HBM = 1.0 / inv_bw_HBM if inv_bw_HBM > 0 else float("nan") - bw_L2 = 1.0 / inv_bw_L2 if inv_bw_L2 > 0 else float("nan") + bw = 1.0 / inv_bw if inv_bw > 0 else float("nan") - # 3. Calculate linear R-squared in linear space if not np.isnan(overhead): - T_pred = time_model( - bytes_arr, overhead, inv_bw_L2, inv_bw_HBM, L2_thresh, HBM_thresh - ) + T_pred = time_model(bytes_arr, overhead, inv_bw) ss_res = np.sum((T_arr - T_pred) ** 2) ss_tot = np.sum((T_arr - np.mean(T_arr)) ** 2) r_squared = 1 - (ss_res / ss_tot) if ss_tot > 0 else float("nan") @@ -263,17 +257,12 @@ def time_model(b, overhead, inv_bw_L2, inv_bw_HBM, L2_thresh, HBM_thresh): r_squared = float("nan") print( - f" Fitted gather: overhead={overhead*1e3:.3f} ms ", - f"BW_L2={bw_L2/1e9:.2f} GB/s BW_HBM={bw_HBM/1e9:.2f} GB/s " - f"L2_thresh={L2_thresh/1e6:.2f} MB HBM_thresh={HBM_thresh/1e6:.2f} MB " - f"R²={r_squared:.4f}", + f" Fitted gather: launch_overhead={overhead*1e6:.2f} us " + f"BW={bw/1e9:.2f} GB/s R²={r_squared:.4f}" ) return { - "bandwidth_bytes_per_sec": bw_HBM, - "L2_bandwidth_bytes_per_sec": bw_L2, - "L2_inflection_bytes": L2_thresh, - "HBM_inflection_bytes": HBM_thresh, + "bandwidth_bytes_per_sec": bw, "launch_overhead_seconds": overhead, "r_squared": r_squared, } diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py b/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py index 8d28cfd..28c63fe 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py @@ -28,6 +28,7 @@ import torch.distributed as dist from benchmarks.common import ( + assert_placement, collect_metadata, seed_everything, setup_distributed, @@ -128,6 +129,17 @@ def main(): "Layout: rank 0,1 on node A; rank 2,3 on node B." ) + # The peer maps below hardcode "ranks 0,1 on node A; ranks 2,3 on node B". + # Nothing in the launch command enforces that, and running all 4 ranks on + # one node would make the "inter" peers actually intra-node — silently + # turning the overlap measurement (which the whole + # T_comm = max(T_intra, T_inter) assumption rests on) into a comparison of + # intra-node against intra-node. + assert_placement( + [(0, 1, True), (2, 3, True), (0, 2, False)], + context="layout: ranks 0,1 on node A; ranks 2,3 on node B", + ) + seed_everything(args.seed) device = torch.device(f"cuda:{local_rank}") diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_pingpong.py b/experiments/cost_model_benchmarks/benchmarks/bench_pingpong.py index c844311..b5facc6 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_pingpong.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_pingpong.py @@ -21,6 +21,7 @@ import torch.distributed as dist from benchmarks.common import ( + assert_placement, collect_metadata, seed_everything, setup_distributed, @@ -96,6 +97,16 @@ def main(): if world_size != 2: raise ValueError(f"bench_pingpong requires exactly 2 ranks, got {world_size}") + # --mode is only a label written into the output JSON; it does not change + # which ranks talk to each other (always 0 <-> 1). Verify the launcher + # actually placed those two ranks the way the label claims, so an + # --nnodes/--nproc_per_node mistake can't silently produce intra-node + # timings filed as inter-node (or vice versa). + assert_placement( + [(0, 1, args.mode == "intra")], + context=f"--mode {args.mode}", + ) + seed_everything(args.seed) device = torch.device(f"cuda:{local_rank}") diff --git a/experiments/cost_model_benchmarks/benchmarks/common.py b/experiments/cost_model_benchmarks/benchmarks/common.py index 27e8069..0898b1c 100644 --- a/experiments/cost_model_benchmarks/benchmarks/common.py +++ b/experiments/cost_model_benchmarks/benchmarks/common.py @@ -163,6 +163,60 @@ def setup_distributed() -> tuple[int, int, int]: return rank, world_size, local_rank +# --------------------------------------------------------------------------- +# Rank placement +# --------------------------------------------------------------------------- + + +def gather_node_of_rank() -> list: + """Return a list mapping each rank -> an opaque node identifier. + + Derived from the actual hostname each rank is running on, so it reflects + real placement rather than an assumed ``rank // ranks_per_node`` layout. + Collective: every rank must call it. + """ + hostnames = [None] * dist.get_world_size() + dist.all_gather_object(hostnames, socket.gethostname()) + return hostnames + + +def assert_placement(expected_same_node: list, context: str = "") -> None: + """Verify the real rank->node placement matches what a benchmark assumes. + + ``expected_same_node`` is a list of (rank_a, rank_b, same) triples: *same* + is True if the benchmark requires the two ranks to be co-located, False if + it requires them to be on different nodes. + + Benchmarks that measure intra- vs inter-node traffic encode their layout + assumption implicitly (in a ``--mode`` label, or a hardcoded peer map). + Nothing about the launch command enforces it, so a wrong ``--nnodes`` / + ``--nproc_per_node`` combination would silently mislabel intra-node + numbers as inter-node (or vice versa) and quietly corrupt the + hierarchical network fit that ``T_comm = max(T_intra, T_inter)`` rests on. + Fail loudly instead. + """ + node_of = gather_node_of_rank() + problems = [] + for rank_a, rank_b, same in expected_same_node: + actually_same = node_of[rank_a] == node_of[rank_b] + if actually_same != same: + problems.append( + f" ranks {rank_a} and {rank_b} must be on " + f"{'the SAME node' if same else 'DIFFERENT nodes'}, but rank " + f"{rank_a} is on {node_of[rank_a]!r} and rank {rank_b} is on " + f"{node_of[rank_b]!r}" + ) + if problems: + raise RuntimeError( + f"Rank placement does not match what this benchmark requires" + f"{' (' + context + ')' if context else ''}:\n" + + "\n".join(problems) + + f"\n\nObserved rank -> host mapping: {list(enumerate(node_of))}\n" + "Fix the launcher's --nnodes/--nproc_per_node (torchrun) or " + "-N/--ntasks-per-node (srun) so the layout matches." + ) + + # --------------------------------------------------------------------------- # Seeding # --------------------------------------------------------------------------- diff --git a/experiments/cost_model_benchmarks/run_experiments_multinode.sh b/experiments/cost_model_benchmarks/run_experiments_multinode.sh new file mode 100755 index 0000000..265faf9 --- /dev/null +++ b/experiments/cost_model_benchmarks/run_experiments_multinode.sh @@ -0,0 +1,229 @@ +#!/usr/bin/env bash +#SBATCH --job-name=cost_model_multinode +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=02:00:00 +#SBATCH --output=cost_model_multinode_%j.out +# NOTE: add your cluster's --partition/--account/--gres here, e.g. +# #SBATCH --gres=gpu:4 (or --gpus-per-node=4, cluster-dependent) +# +# Multi-node extension of run_experiments.sh: adds the inter-node primitives +# and the K = NNODES*GPUS_PER_NODE end-to-end/crossover points that a +# single-node run cannot produce. +# +# PREREQUISITE: run ./run_experiments.sh on one node first. This script reuses +# its compute_*/gather_*/pingpong_intra.json outputs and its K=2,4 e2e runs; +# it only produces what genuinely requires >1 node. +# +# Run it ONCE — it fans out to the other nodes itself via srun. Either: +# +# sbatch ./run_experiments_multinode.sh +# +# or from an interactive allocation (salloc -N 2 gives one shell on the head +# node; srun distributes from there): +# +# salloc -N 2 --ntasks-per-node=1 --gres=gpu:4 -t 2:00:00 +# ./run_experiments_multinode.sh +# +# Each step runs `srun --ntasks-per-node=1`, placing exactly one torchrun +# launcher on each node, which then spawns nproc_per_node worker processes +# locally. torchrun uses *static* rendezvous (explicit --node_rank taken from +# SLURM_NODEID) rather than c10d dynamic rendezvous, so global ranks are +# deterministically block-assigned (global_rank = node_rank*nproc + local_rank). +# bench_concurrency hardcodes a "ranks 0,1 on node A; ranks 2,3 on node B" +# layout, and c10d assigns node ranks in join order — which would scramble +# that mapping nondeterministically between runs. +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")" +shopt -s nullglob + +# --- Allocation discovery ---------------------------------------------- +if ! command -v srun >/dev/null 2>&1; then + echo "[error] srun not found — this script must run inside a SLURM allocation" >&2 + echo " (sbatch ./run_experiments_multinode.sh, or salloc then run it)" >&2 + exit 1 +fi +if [[ -z "${SLURM_JOB_ID:-}" ]]; then + echo "[error] no SLURM allocation detected (SLURM_JOB_ID unset)." >&2 + echo " Run under sbatch, or inside 'salloc -N 2 ... '." >&2 + exit 1 +fi + +NNODES=${NNODES:-${SLURM_JOB_NUM_NODES:-${SLURM_NNODES:-2}}} +if [[ "${NNODES}" -lt 2 ]]; then + echo "[error] allocation has ${NNODES} node(s); this script needs >= 2." >&2 + echo " For single-node runs use ./run_experiments.sh instead." >&2 + exit 1 +fi + +# Head node hostname, used as the torchrun rendezvous master by every node. +MASTER_ADDR=${MASTER_ADDR:-$(scontrol show hostnames "${SLURM_JOB_NODELIST}" | head -n1)} + +# --- Config ------------------------------------------------------------ +GPUS_PER_NODE=${GPUS_PER_NODE:-4} +MASTER_PORT=${MASTER_PORT:-29500} +SEED=${SEED:-42} +FEATURE_DIM=${FEATURE_DIM:-128} +MODEL=${MODEL:-gcn} +GRAPH=${GRAPH:-erdos_renyi} +PARTITIONER=${PARTITIONER:-balanced} +WARMUP=${WARMUP:-10} +TRIALS=${TRIALS:-50} +CONCURRENCY_BYTES=${CONCURRENCY_BYTES:-16777216} # 16 MiB +CROSSOVER_SIZES=${CROSSOVER_SIZES:-1000,10000,100000,1000000} +E2E_VERTEX_SIZES=(${E2E_VERTEX_SIZES:-10000 100000 1000000}) +SINGLE_NODE_GPU_COUNTS=(${SINGLE_NODE_GPU_COUNTS:-2 4}) # must match run_experiments.sh +DATA_DIR=${DATA_DIR:-data} +FIG_DIR=${FIG_DIR:-figures} + +K=$((NNODES * GPUS_PER_NODE)) # total world size for the full-scale runs + +mkdir -p "$DATA_DIR" "$FIG_DIR" + +echo "Allocation: ${NNODES} nodes, ${GPUS_PER_NODE} GPU/node -> K=${K}" +echo "Rendezvous master: ${MASTER_ADDR}" + +step() { echo; echo "=== $* ==="; } + +# Each torchrun invocation gets its own port. Back-to-back runs reusing one +# port can fail with "address already in use" while the previous rendezvous +# socket is still in TIME_WAIT. Incremented by a plain assignment in this +# shell — NOT inside $(...), whose subshell assignment would be discarded, +# silently pinning every step to the same port. +_port=$MASTER_PORT + +# Launch one torchrun per node via srun. +# $1 = nproc_per_node, rest = the module + its args. +# --node_rank must expand per-node, so the torchrun command runs under a +# remote shell where $SLURM_NODEID is set; args are %q-quoted so anything +# containing spaces survives that extra round of shell parsing intact. +launch() { + local nproc=$1; shift + _port=$((_port + 1)) + local quoted_cmd + quoted_cmd=$(printf '%q ' "$@") + srun --nodes="${NNODES}" --ntasks-per-node=1 --cpu-bind=none \ + bash -c "torchrun \ + --nnodes ${NNODES} \ + --nproc_per_node ${nproc} \ + --node_rank \${SLURM_NODEID} \ + --master_addr ${MASTER_ADDR} \ + --master_port ${_port} \ + ${quoted_cmd}" +} + +# ======================================================================= +# Multi-node benchmarks. Every step below needs more than one node; the +# single-node primitives come from run_experiments.sh. +# ======================================================================= + +# 1.1 inter-node ping-pong: exactly 2 ranks, one per node, so rank 0 and +# rank 1 are guaranteed to be on different nodes (bench_pingpong asserts it). +step "1.1 pingpong -- inter-node (1 rank/node)" +launch 1 -m benchmarks.bench_pingpong \ + --mode inter --min-bytes 64 --max-bytes 67108864 --steps 21 \ + --warmup 20 --trials 100 \ + --output "${DATA_DIR}/pingpong_inter.json" --seed "${SEED}" + +# 1.2 concurrency: exactly 4 ranks laid out 2-per-node, giving each rank both +# an intra-node peer and an inter-node peer. This is the benchmark that +# justifies T_comm = max(T_intra, T_inter); it cannot run on one node. +if [[ "${NNODES}" -eq 2 ]]; then + step "1.2 concurrency -- intra/inter overlap (2 ranks/node)" + launch 2 -m benchmarks.bench_concurrency \ + --message-bytes "${CONCURRENCY_BYTES}" --warmup 20 --trials 100 \ + --output "${DATA_DIR}/concurrency.json" --seed "${SEED}" +else + echo "[skip] bench_concurrency requires exactly 4 ranks (NNODES=2, 2/node); NNODES=${NNODES}" +fi + +# 2.2 crossover and 2.1 end-to-end at full scale across all nodes. +step "2.2 crossover -- K=${K} GPUs (${NNODES} nodes)" +launch "${GPUS_PER_NODE}" -m benchmarks.bench_crossover \ + --graph "${GRAPH}" --graph-sizes "${CROSSOVER_SIZES}" \ + --avg-degree 20 --feature-dim "${FEATURE_DIM}" --model "${MODEL}" \ + --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/crossover_K${K}.json" --seed "${SEED}" + +for N in "${E2E_VERTEX_SIZES[@]}"; do + step "2.1 end-to-end -- K=${K} GPUs, N=${N}" + launch "${GPUS_PER_NODE}" -m benchmarks.bench_end_to_end \ + --graph "${GRAPH}" --num-vertices "${N}" --avg-degree 20 \ + --feature-dim "${FEATURE_DIM}" --model "${MODEL}" \ + --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/e2e_K${K}_N${N}.json" --seed "${SEED}" +done + +# ======================================================================= +# Refit and plot (CPU-only post-processing, runs on the head node). +# ======================================================================= + +# Explicit file list (single-node K values plus this run's K), never a +# e2e_*.json glob — a glob would also pick up stray files that were not part +# of this run matrix. +E2E_FILES=() +for k in "${SINGLE_NODE_GPU_COUNTS[@]}" "${K}"; do + for N in "${E2E_VERTEX_SIZES[@]}"; do + f="${DATA_DIR}/e2e_K${k}_N${N}.json" + if [[ -f "$f" ]]; then + E2E_FILES+=("$f") + elif [[ "$k" -eq "$K" ]]; then + echo "[warn] missing ${f} — this run should have produced it; check the K=${K} step above for errors" + else + echo "[warn] missing ${f} — run ./run_experiments.sh first to produce the K=${k} points" + fi + done +done + +if [[ ${#E2E_FILES[@]} -eq 0 ]]; then + echo "[error] no end-to-end run files found in ${DATA_DIR}/ — nothing to fit against." >&2 + exit 1 +fi + +# Hold out the largest world size, so held-out MAPE measures extrapolation +# into the multi-node regime specifically. +FIT_FILTER=${FIT_FILTER:-"world_size < ${K}"} + +step "fit_primitives (intra + inter network)" +python -m analysis.fit_primitives \ + --pingpong-intra "${DATA_DIR}/pingpong_intra.json" \ + --pingpong-inter "${DATA_DIR}/pingpong_inter.json" \ + --compute-gcn "${DATA_DIR}/compute_gcn_vswp.json" "${DATA_DIR}/compute_gcn_eswp.json" \ + --compute-edge "${DATA_DIR}/compute_edge_vswp.json" "${DATA_DIR}/compute_edge_eswp.json" \ + --gather-contiguous "${DATA_DIR}/gather_contiguous.json" \ + --gather-clustered "${DATA_DIR}/gather_clustered.json" \ + --gather-random "${DATA_DIR}/gather_random.json" \ + --output "${DATA_DIR}/fitted_primitives.json" + +step "fit_overhead (fit-filter: ${FIT_FILTER})" +python -m analysis.fit_overhead \ + --primitives "${DATA_DIR}/fitted_primitives.json" \ + --e2e-runs "${E2E_FILES[@]}" \ + --fit-filter "${FIT_FILTER}" \ + --output "${DATA_DIR}/fitted_overhead.json" + +step "compute_predictions (fit-filter: ${FIT_FILTER})" +python -m analysis.compute_predictions \ + --primitives "${DATA_DIR}/fitted_primitives.json" \ + --overhead "${DATA_DIR}/fitted_overhead.json" \ + --e2e-runs "${E2E_FILES[@]}" \ + --fit-filter "${FIT_FILTER}" \ + --output "${DATA_DIR}/predictions.json" + +step "visualization" +# Now that an inter-node fit exists, plot both network tiers. +python -m visualization.plot_pingpong \ + --intra "${DATA_DIR}/pingpong_intra.json" \ + --inter "${DATA_DIR}/pingpong_inter.json" \ + --primitives "${DATA_DIR}/fitted_primitives.json" \ + --output "${FIG_DIR}/pingpong" + +python -m visualization.plot_crossover --input "${DATA_DIR}/crossover_K${K}.json" \ + --output "${FIG_DIR}/crossover_K${K}.png" + +python -m visualization.plot_validation --predictions "${DATA_DIR}/predictions.json" \ + --color-by world_size --output "${FIG_DIR}/validation" + +echo +echo "Done. Data in ${DATA_DIR}/, figures in ${FIG_DIR}/." diff --git a/experiments/cost_model_benchmarks/visualization/plot_gather.py b/experiments/cost_model_benchmarks/visualization/plot_gather.py index f33b193..5e7c988 100644 --- a/experiments/cost_model_benchmarks/visualization/plot_gather.py +++ b/experiments/cost_model_benchmarks/visualization/plot_gather.py @@ -45,50 +45,36 @@ def fitted_gather( min_val, max_val, feature_dim, - hbm_bandwidth_bytes_per_sec, - l2_bandwidth_bytes_per_sec, + bandwidth_bytes_per_sec, launch_overhead_seconds, - L2_thresh, - HBM_thresh, ): - """Return a function that models gather time as a function of k.""" - - def time_model(b, overhead, inv_bw_L2, inv_bw_HBM, L2_thresh, HBM_thresh): - # 1. Bucket the bytes into their respective physical regimes - # Bytes processed exclusively at L2 speeds - bytes_L2 = np.clip(b, 0, L2_thresh) - # Bytes processed exclusively at HBM speeds - bytes_HBM = np.maximum(0, b - HBM_thresh) - - # 2. Apply the specific bandwidth (slope) to each bucket - t_mem = (bytes_L2 * inv_bw_L2) + (bytes_HBM * inv_bw_HBM) - - # 3. Floor the total time by the kernel launch overhead - return np.maximum(overhead, t_mem) - - x = np.linspace(min_val, max_val, 100) - inv_bw_L2 = 1.0 / l2_bandwidth_bytes_per_sec - inv_bw_HBM = 1.0 / hbm_bandwidth_bytes_per_sec - y = ( - time_model( - x * feature_dim * 4.0, - launch_overhead_seconds, - inv_bw_L2, - inv_bw_HBM, - L2_thresh, - HBM_thresh, - ) - * 1e3 - ) + """Sample the fitted Hockney gather curve over a range of k values. + + Mirrors ``time_model`` in ``analysis/fit_primitives.py::fit_gather``: + ``T(bytes) = launch_overhead + bytes / B_gather``. Returns (k, T_ms). + """ + # Log-spaced, since the k sweep itself is log-spaced and the plot's x + # axis is logarithmic — linspace would leave the low decades unsampled + # and render the curve as a straight chord across them. + x = np.logspace(np.log10(min_val), np.log10(max_val), 200) + y = (launch_overhead_seconds + x * feature_dim * 4.0 / bandwidth_bytes_per_sec) * 1e3 return x, y def load_gather_file(paths: list, timing_key: str): - """Merge multiple JSON files, return sorted (k, median, q25, q75) arrays.""" + """Merge multiple JSON files, return sorted (k, median, q25, q75) arrays + plus the feature_dim the runs were measured at. + + feature_dim must come from the data, not a hardcoded constant: it sets the + bytes-per-index (``k * feature_dim * 4``) used to evaluate the fitted + curve, so a wrong value tilts the overlay by exactly that ratio. + """ rows = [] + feature_dims = set() for p in paths: with open(p) as f: data = json.load(f) + feature_dims.add(data["config"]["feature_dim"]) for meas in data["measurements"]: trials = np.array(meas[timing_key]) rows.append( @@ -104,7 +90,13 @@ def load_gather_file(paths: list, timing_key: str): med_arr = np.array([r[1] for r in rows]) q25_arr = np.array([r[2] for r in rows]) q75_arr = np.array([r[3] for r in rows]) - return k_arr, med_arr, q25_arr, q75_arr + if len(feature_dims) > 1: + raise ValueError( + f"Merged gather files have mismatched feature_dim values " + f"{sorted(feature_dims)}; the bytes-per-index conversion is only " + "well-defined for one. Plot them separately." + ) + return k_arr, med_arr, q25_arr, q75_arr, feature_dims.pop() def parse_args(): @@ -139,7 +131,7 @@ def main(): ]: if not files: continue - k, med, q25, q75 = load_gather_file(files, timing_key) + k, med, q25, q75, feature_dim = load_gather_file(files, timing_key) min_k = k[0] max_k = k[-1] color = COLORS[dist_name] @@ -157,28 +149,13 @@ def main(): ) if args.fitted: - hbm_bw = primitives["gather"][dist_name]["gather"][ - "bandwidth_bytes_per_sec" - ] - l2_bw = primitives["gather"][dist_name]["gather"][ - "L2_bandwidth_bytes_per_sec" - ] - overhead = primitives["gather"][dist_name]["gather"][ - "launch_overhead_seconds" - ] - thresh = primitives["gather"][dist_name]["gather"]["L2_inflection_bytes"] - hbm_thresh = primitives["gather"][dist_name]["gather"][ - "HBM_inflection_bytes" - ] + fit = primitives["gather"][dist_name][args.operation] x, y = fitted_gather( min_k, max_k, - feature_dim=512, - hbm_bandwidth_bytes_per_sec=hbm_bw, - l2_bandwidth_bytes_per_sec=l2_bw, - launch_overhead_seconds=overhead, - L2_thresh=thresh, - HBM_thresh=hbm_thresh, + feature_dim=feature_dim, + bandwidth_bytes_per_sec=fit["bandwidth_bytes_per_sec"], + launch_overhead_seconds=fit["launch_overhead_seconds"], ) ax.plot( x * 1e-6, From 0ce080eb87e53516df6b0ee3ee61fddead526291 Mon Sep 17 00:00:00 2001 From: M Shehtab Zaman Date: Fri, 31 Jul 2026 20:22:43 -0700 Subject: [PATCH 32/39] Update multinode to use HPC-launcher --- .../run_experiments_multinode.sh | 196 +++++++++--------- 1 file changed, 101 insertions(+), 95 deletions(-) diff --git a/experiments/cost_model_benchmarks/run_experiments_multinode.sh b/experiments/cost_model_benchmarks/run_experiments_multinode.sh index 265faf9..80091ad 100755 --- a/experiments/cost_model_benchmarks/run_experiments_multinode.sh +++ b/experiments/cost_model_benchmarks/run_experiments_multinode.sh @@ -1,68 +1,53 @@ -#!/usr/bin/env bash -#SBATCH --job-name=cost_model_multinode -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=1 -#SBATCH --time=02:00:00 -#SBATCH --output=cost_model_multinode_%j.out -# NOTE: add your cluster's --partition/--account/--gres here, e.g. -# #SBATCH --gres=gpu:4 (or --gpus-per-node=4, cluster-dependent) -# -# Multi-node extension of run_experiments.sh: adds the inter-node primitives -# and the K = NNODES*GPUS_PER_NODE end-to-end/crossover points that a -# single-node run cannot produce. -# -# PREREQUISITE: run ./run_experiments.sh on one node first. This script reuses -# its compute_*/gather_*/pingpong_intra.json outputs and its K=2,4 e2e runs; -# it only produces what genuinely requires >1 node. -# -# Run it ONCE — it fans out to the other nodes itself via srun. Either: -# -# sbatch ./run_experiments_multinode.sh -# -# or from an interactive allocation (salloc -N 2 gives one shell on the head -# node; srun distributes from there): -# -# salloc -N 2 --ntasks-per-node=1 --gres=gpu:4 -t 2:00:00 -# ./run_experiments_multinode.sh -# -# Each step runs `srun --ntasks-per-node=1`, placing exactly one torchrun -# launcher on each node, which then spawns nproc_per_node worker processes -# locally. torchrun uses *static* rendezvous (explicit --node_rank taken from -# SLURM_NODEID) rather than c10d dynamic rendezvous, so global ranks are -# deterministically block-assigned (global_rank = node_rank*nproc + local_rank). -# bench_concurrency hardcodes a "ranks 0,1 on node A; ranks 2,3 on node B" -# layout, and c10d assigns node ranks in join order — which would scramble -# that mapping nondeterministically between runs. set -euo pipefail cd "$(dirname "${BASH_SOURCE[0]}")" shopt -s nullglob -# --- Allocation discovery ---------------------------------------------- -if ! command -v srun >/dev/null 2>&1; then - echo "[error] srun not found — this script must run inside a SLURM allocation" >&2 - echo " (sbatch ./run_experiments_multinode.sh, or salloc then run it)" >&2 - exit 1 -fi -if [[ -z "${SLURM_JOB_ID:-}" ]]; then - echo "[error] no SLURM allocation detected (SLURM_JOB_ID unset)." >&2 - echo " Run under sbatch, or inside 'salloc -N 2 ... '." >&2 +# ======================================================================= +# Multi-node benchmark driver -- uses hpc-launcher's `torchrun-hpc` +# (https://github.com/llnl/HPC-launcher) instead of srun+torchrun. +# +# torchrun-hpc builds and submits its own SLURM batch script per +# invocation, which changes the operating model versus the old +# srun-inside-an-allocation approach: +# - This script no longer needs to run inside a pre-existing +# salloc/sbatch allocation. Run it directly, e.g. from a login node. +# - Every launch() call below requests its own (nnodes, procs/node) +# topology and is scheduled independently. Fixed-topology steps +# (pingpong-inter, concurrency) are no longer constrained to fit +# inside whatever NNODES was chosen for the full-scale crossover/e2e +# runs -- see the notes at each call site below. +# - MASTER_ADDR/MASTER_PORT/node_rank are handled internally by +# torchrun-hpc; no rendezvous port bookkeeping is needed here. +# - torchrun-hpc changes into its own launch directory (timestamped, by +# default) before running the command, and *relative* path arguments +# resolve against that directory, not this script's cwd. DATA_DIR and +# FIG_DIR are therefore canonicalized to absolute paths below. +# +# Assumption: torchrun-hpc blocks until the submitted job finishes +# (mirroring srun's blocking behavior), so that the files each step +# produces are on disk before the next step (or the fit/plot stage) reads +# them. Each launch() call is followed by a check that its expected +# output file actually landed, so the script fails fast with a clear +# message if that assumption doesn't hold for your installed version. +# If in doubt, sanity-check with `torchrun-hpc --dry-run` first. +# ======================================================================= + +if ! command -v torchrun-hpc >/dev/null 2>&1; then + echo "[error] torchrun-hpc not found on PATH." >&2 + echo " Install hpc-launcher with: pip install hpc-launcher" >&2 exit 1 fi -NNODES=${NNODES:-${SLURM_JOB_NUM_NODES:-${SLURM_NNODES:-2}}} +# --- Config ------------------------------------------------------------ +NNODES=${NNODES:-2} if [[ "${NNODES}" -lt 2 ]]; then - echo "[error] allocation has ${NNODES} node(s); this script needs >= 2." >&2 + echo "[error] NNODES=${NNODES}; this script needs >= 2." >&2 echo " For single-node runs use ./run_experiments.sh instead." >&2 exit 1 fi -# Head node hostname, used as the torchrun rendezvous master by every node. -MASTER_ADDR=${MASTER_ADDR:-$(scontrol show hostnames "${SLURM_JOB_NODELIST}" | head -n1)} - -# --- Config ------------------------------------------------------------ GPUS_PER_NODE=${GPUS_PER_NODE:-4} -MASTER_PORT=${MASTER_PORT:-29500} SEED=${SEED:-42} FEATURE_DIM=${FEATURE_DIM:-128} MODEL=${MODEL:-gcn} @@ -77,40 +62,57 @@ SINGLE_NODE_GPU_COUNTS=(${SINGLE_NODE_GPU_COUNTS:-2 4}) # must match run_experi DATA_DIR=${DATA_DIR:-data} FIG_DIR=${FIG_DIR:-figures} +# torchrun-hpc scheduling options. All optional -- leave unset to let +# hpc-launcher auto-detect system defaults (queue, time limit, account). +JOB_NAME_PREFIX=${JOB_NAME_PREFIX:-gnn_bench} +TORCHRUN_HPC_QUEUE=${TORCHRUN_HPC_QUEUE:-} +TORCHRUN_HPC_TIME_LIMIT=${TORCHRUN_HPC_TIME_LIMIT:-} # minutes +TORCHRUN_HPC_ACCOUNT=${TORCHRUN_HPC_ACCOUNT:-} +TORCHRUN_HPC_EXCLUSIVE=${TORCHRUN_HPC_EXCLUSIVE:-1} # non-empty = pass --exclusive +TORCHRUN_HPC_EXTRA_ARGS=${TORCHRUN_HPC_EXTRA_ARGS:-} # e.g. "-r mpi --comm-backend NCCL" + K=$((NNODES * GPUS_PER_NODE)) # total world size for the full-scale runs mkdir -p "$DATA_DIR" "$FIG_DIR" +# Canonicalize to absolute paths -- see launch-directory note above. +DATA_DIR=$(cd "${DATA_DIR}" && pwd) +FIG_DIR=$(cd "${FIG_DIR}" && pwd) -echo "Allocation: ${NNODES} nodes, ${GPUS_PER_NODE} GPU/node -> K=${K}" -echo "Rendezvous master: ${MASTER_ADDR}" +echo "Target topology: ${NNODES} nodes x ${GPUS_PER_NODE} GPU/node -> K=${K}" +echo "Each step below submits (and waits on) its own torchrun-hpc job." step() { echo; echo "=== $* ==="; } -# Each torchrun invocation gets its own port. Back-to-back runs reusing one -# port can fail with "address already in use" while the previous rendezvous -# socket is still in TIME_WAIT. Incremented by a plain assignment in this -# shell — NOT inside $(...), whose subshell assignment would be discarded, -# silently pinning every step to the same port. -_port=$MASTER_PORT - -# Launch one torchrun per node via srun. -# $1 = nproc_per_node, rest = the module + its args. -# --node_rank must expand per-node, so the torchrun command runs under a -# remote shell where $SLURM_NODEID is set; args are %q-quoted so anything -# containing spaces survives that extra round of shell parsing intact. +# Fails loudly if a step's expected output file didn't land -- see the +# blocking-behavior assumption noted at the top of this file. +expect_output() { + if [[ ! -s "$1" ]]; then + echo "[error] expected output file '$1' was not produced by the last step." >&2 + echo " If torchrun-hpc returned before the job finished, this" >&2 + echo " script's blocking assumption doesn't hold -- see header." >&2 + exit 1 + fi +} + +# Runs one Python module under torchrun-hpc. torchrun-hpc builds and +# submits the SLURM batch script itself; no srun/torchrun/rendezvous +# plumbing is needed here. +# $1 = nnodes, $2 = procs/node, $3 = python module (dotted path), +# rest = module args. launch() { - local nproc=$1; shift - _port=$((_port + 1)) - local quoted_cmd - quoted_cmd=$(printf '%q ' "$@") - srun --nodes="${NNODES}" --ntasks-per-node=1 --cpu-bind=none \ - bash -c "torchrun \ - --nnodes ${NNODES} \ - --nproc_per_node ${nproc} \ - --node_rank \${SLURM_NODEID} \ - --master_addr ${MASTER_ADDR} \ - --master_port ${_port} \ - ${quoted_cmd}" + local nnodes=$1 nproc=$2 module=$3; shift 3 + local -a extra=() + [[ -n "${TORCHRUN_HPC_EXCLUSIVE}" ]] && extra+=(--exclusive) + [[ -n "${TORCHRUN_HPC_QUEUE}" ]] && extra+=(--queue "${TORCHRUN_HPC_QUEUE}") + [[ -n "${TORCHRUN_HPC_TIME_LIMIT}" ]] && extra+=(--time-limit "${TORCHRUN_HPC_TIME_LIMIT}") + [[ -n "${TORCHRUN_HPC_ACCOUNT}" ]] && extra+=(--account "${TORCHRUN_HPC_ACCOUNT}") + # Intentional word-splitting of a user-supplied flag string. + local -a passthrough=(${TORCHRUN_HPC_EXTRA_ARGS}) + + torchrun-hpc \ + --nodes "${nnodes}" \ + --procs-per-node "${nproc}" \ + "${module}" "$@" } # ======================================================================= @@ -118,45 +120,49 @@ launch() { # single-node primitives come from run_experiments.sh. # ======================================================================= -# 1.1 inter-node ping-pong: exactly 2 ranks, one per node, so rank 0 and -# rank 1 are guaranteed to be on different nodes (bench_pingpong asserts it). -step "1.1 pingpong -- inter-node (1 rank/node)" -launch 1 -m benchmarks.bench_pingpong \ +# 1.1 inter-node ping-pong: fixed at exactly 2 nodes, 1 rank/node, so +# rank 0 and rank 1 are guaranteed to be on different nodes (bench_pingpong +# asserts it). This no longer depends on the full-scale NNODES setting. +step "1.1 pingpong -- inter-node (2 nodes, 1 rank/node)" +launch 2 1 benchmarks/bench_pingpong.py \ --mode inter --min-bytes 64 --max-bytes 67108864 --steps 21 \ --warmup 20 --trials 100 \ --output "${DATA_DIR}/pingpong_inter.json" --seed "${SEED}" +expect_output "${DATA_DIR}/pingpong_inter.json" -# 1.2 concurrency: exactly 4 ranks laid out 2-per-node, giving each rank both +# 1.2 concurrency: fixed at exactly 4 ranks, 2/node, giving each rank both # an intra-node peer and an inter-node peer. This is the benchmark that -# justifies T_comm = max(T_intra, T_inter); it cannot run on one node. -if [[ "${NNODES}" -eq 2 ]]; then - step "1.2 concurrency -- intra/inter overlap (2 ranks/node)" - launch 2 -m benchmarks.bench_concurrency \ - --message-bytes "${CONCURRENCY_BYTES}" --warmup 20 --trials 100 \ - --output "${DATA_DIR}/concurrency.json" --seed "${SEED}" -else - echo "[skip] bench_concurrency requires exactly 4 ranks (NNODES=2, 2/node); NNODES=${NNODES}" -fi +# justifies T_comm = max(T_intra, T_inter). Since this now schedules its +# own independent 2-node job, it no longer needs to be gated behind +# NNODES == 2 -- it always runs, regardless of the full-scale topology. +step "1.2 concurrency -- intra/inter overlap (2 nodes, 2 ranks/node)" +launch 2 2 benchmarks/bench_concurrency.py \ + --message-bytes "${CONCURRENCY_BYTES}" --warmup 20 --trials 100 \ + --output "${DATA_DIR}/concurrency.json" --seed "${SEED}" +expect_output "${DATA_DIR}/concurrency.json" # 2.2 crossover and 2.1 end-to-end at full scale across all nodes. step "2.2 crossover -- K=${K} GPUs (${NNODES} nodes)" -launch "${GPUS_PER_NODE}" -m benchmarks.bench_crossover \ +launch "${NNODES}" "${GPUS_PER_NODE}" benchmarks/bench_crossover.py \ --graph "${GRAPH}" --graph-sizes "${CROSSOVER_SIZES}" \ --avg-degree 20 --feature-dim "${FEATURE_DIM}" --model "${MODEL}" \ --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ --output "${DATA_DIR}/crossover_K${K}.json" --seed "${SEED}" +expect_output "${DATA_DIR}/crossover_K${K}.json" for N in "${E2E_VERTEX_SIZES[@]}"; do step "2.1 end-to-end -- K=${K} GPUs, N=${N}" - launch "${GPUS_PER_NODE}" -m benchmarks.bench_end_to_end \ + launch "${NNODES}" "${GPUS_PER_NODE}" benchmarks/bench_end_to_end.py \ --graph "${GRAPH}" --num-vertices "${N}" --avg-degree 20 \ --feature-dim "${FEATURE_DIM}" --model "${MODEL}" \ --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ --output "${DATA_DIR}/e2e_K${K}_N${N}.json" --seed "${SEED}" + expect_output "${DATA_DIR}/e2e_K${K}_N${N}.json" done # ======================================================================= -# Refit and plot (CPU-only post-processing, runs on the head node). +# Refit and plot (CPU-only post-processing, runs locally -- no launcher +# involved). # ======================================================================= # Explicit file list (single-node K values plus this run's K), never a @@ -226,4 +232,4 @@ python -m visualization.plot_validation --predictions "${DATA_DIR}/predictions.j --color-by world_size --output "${FIG_DIR}/validation" echo -echo "Done. Data in ${DATA_DIR}/, figures in ${FIG_DIR}/." +echo "Done. Data in ${DATA_DIR}/, figures in ${FIG_DIR}/." \ No newline at end of file From 4366d943160c10ea0a3b7cd5bfcaf570aec4ec0c Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 31 Jul 2026 23:57:05 -0400 Subject: [PATCH 33/39] Updated compute pred and scatter-add plotting --- .../analysis/compute_predictions.py | 7 ++- .../analysis/fit_overhead.py | 51 ++++++++++++++++--- .../benchmarks/bench_gather.py | 14 ++++- .../cost_model_benchmarks/run_experiments.sh | 16 ++++-- 4 files changed, 75 insertions(+), 13 deletions(-) diff --git a/experiments/cost_model_benchmarks/analysis/compute_predictions.py b/experiments/cost_model_benchmarks/analysis/compute_predictions.py index 5070fba..650843a 100644 --- a/experiments/cost_model_benchmarks/analysis/compute_predictions.py +++ b/experiments/cost_model_benchmarks/analysis/compute_predictions.py @@ -64,6 +64,11 @@ def main(): overhead_data = json.load(f) T_overhead = overhead_data.get("overhead_seconds", 0.0) + # Read from fitted_overhead.json rather than taking a second CLI flag: + # T_overhead was fitted under a specific contention assumption, so passing + # a different nics_per_node here would silently pair an overhead constant + # with a network model it was never fitted against. + nics_per_node = overhead_data.get("nics_per_node", 1) all_runs = load_e2e_runs(args.e2e_runs) fit_runs, held_runs = apply_filter(all_runs, args.fit_filter) @@ -72,7 +77,7 @@ def main(): prediction_entries = [] for r in all_runs: - breakdown = predict_layer_time(r["config"], r["per_rank_stats"], primitives) + breakdown = predict_layer_time(r["config"], r["per_rank_stats"], primitives, nics_per_node) T_pred = breakdown["T_total"] + T_overhead T_meas = r["measured_median"] abs_err = abs(T_meas - T_pred) diff --git a/experiments/cost_model_benchmarks/analysis/fit_overhead.py b/experiments/cost_model_benchmarks/analysis/fit_overhead.py index de959b4..754ab84 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_overhead.py +++ b/experiments/cost_model_benchmarks/analysis/fit_overhead.py @@ -50,10 +50,10 @@ def _gather_time(nbytes: float, params: dict) -> float: def predict_layer_time(run_config: dict, per_rank_stats: dict, - primitives: dict) -> dict: + primitives: dict, nics_per_node: int = 1) -> dict: """Predict T_layer for one rank using the assembled primitive model. - T_layer = T_comp + max(T_intra, T_inter) + T_buffer_copy + T_layer = T_comp + (T_intra + T_inter) + T_buffer_copy Parameters ---------- @@ -63,6 +63,12 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, Stats for rank 0 from per_rank_stats list. primitives : dict Loaded fitted_primitives.json. + nics_per_node : int + Number of network interfaces per node. Ranks co-located on a node + share them, so the per-pair inter-node bandwidth measured by + bench_pingpong (1 rank/node, hence 1 rank/NIC) is not what a rank + sees once several ranks per node contend for the same rail. See + the contention note on ``net_time`` below. Returns ------- @@ -73,6 +79,7 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, """ F = run_config["feature_dim"] model_type = run_config.get("model", "gcn") + ranks_per_node = run_config.get("ranks_per_node", 1) n_local = per_rank_stats.get("n_local", 0) n_halo = per_rank_stats.get("n_halo", 0) @@ -116,6 +123,19 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, net = primitives.get("network", {}) + # Ranks co-located on a node share that node's NICs. bench_pingpong + # measures the inter-node rate with 1 rank per node -- i.e. one rank per + # NIC -- so its B_inter is a *per-NIC* capacity, not a per-rank one. With + # r ranks per node and n NICs, ceil(r/n) ranks contend for each rail and + # each sees B_inter / ceil(r/n). + # + # Not a fitted parameter: nics_per_node is a topology constant (read from + # `nvidia-smi topo -m`), exactly like ranks_per_node above. Ignoring it + # (equivalently, pretending every rank owns a NIC) under-predicts K=8 on + # 2 nodes x 4 GPU by 15-31%, because 4 ranks share 2 NICs and each rank + # gets half the ping-pong rate. + ranks_per_nic = max(1, -(-ranks_per_node // max(1, nics_per_node))) + def net_time(nbytes: int, mode: str) -> float: if nbytes == 0: return 0.0 @@ -136,11 +156,21 @@ def net_time(nbytes: int, mode: str) -> float: ) B = params.get("bandwidth_bytes_per_sec", 1e10) t_L = params.get("latency_seconds", 0.0) + # Intra-node traffic goes over NVLink, which is all-to-all: co-located + # ranks do not share a single rail the way they share NICs, so only + # the inter-node rate is divided by the contention factor. + if mode == "inter": + B = B / ranks_per_nic return t_L + nbytes / B T_intra = max(net_time(intra_bytes, "intra"), 0.0) T_inter = max(net_time(inter_bytes, "inter"), 0.0) - T_comm = max(T_intra, T_inter) + # Additive, not max(). bench_concurrency measures exactly this overlap on + # this machine and finds none: with 16 MiB exchanges issued on separate + # CUDA streams, T_concurrent / max(T_intra, T_inter) = 1.60 while + # T_concurrent / (T_intra + T_inter) = 0.97. NVLink and IB transfers + # serialize, so the optimistic max() under-predicts every multi-node run. + T_comm = T_intra + T_inter # T_buffer_copy (gather of send buffer) — uses the fitted Hockney gather # model, evaluated by the same function fit_gather fits. @@ -222,13 +252,14 @@ def weighted_median(values: np.ndarray, weights: np.ndarray) -> float: return float(v[idx]) -def fit_overhead_scalar(fit_runs: list, primitives: dict) -> tuple: +def fit_overhead_scalar(fit_runs: list, primitives: dict, + nics_per_node: int = 1) -> tuple: """Fit T_overhead to minimise MAPE on fit_runs. Returns (overhead, mape_in_sample).""" if not fit_runs: return 0.0, float("nan") residuals = np.array([ - r["measured_median"] - predict_layer_time(r["config"], r["per_rank_stats"], primitives)["T_total"] + r["measured_median"] - predict_layer_time(r["config"], r["per_rank_stats"], primitives, nics_per_node)["T_total"] for r in fit_runs ]) T_meas = np.array([r["measured_median"] for r in fit_runs]) @@ -240,7 +271,7 @@ def fit_overhead_scalar(fit_runs: list, primitives: dict) -> tuple: overhead = weighted_median(residuals, 1.0 / T_meas) mape = float(np.mean([ - abs(r["measured_median"] - (predict_layer_time(r["config"], r["per_rank_stats"], primitives)["T_total"] + overhead)) + abs(r["measured_median"] - (predict_layer_time(r["config"], r["per_rank_stats"], primitives, nics_per_node)["T_total"] + overhead)) / r["measured_median"] for r in fit_runs ])) @@ -258,6 +289,11 @@ def parse_args(): p.add_argument("--e2e-runs", nargs="+", required=True, metavar="FILE") p.add_argument("--fit-filter", type=str, default="world_size <= 8", help="Python expression evaluated per run; True → fit set") + p.add_argument("--nics-per-node", type=int, default=1, + help="Network interfaces per node (topology constant, from " + "`nvidia-smi topo -m`; not fitted). Co-located ranks " + "share them, so each rank sees B_inter divided by " + "ceil(ranks_per_node / nics_per_node).") p.add_argument("--output", type=str, default="data/fitted_overhead.json") return p.parse_args() @@ -274,12 +310,13 @@ def main(): fit_runs, held_runs = apply_filter(all_runs, args.fit_filter) print(f"[fit_overhead] Fit set: {len(fit_runs)} Held-out: {len(held_runs)}") - overhead, mape_in = fit_overhead_scalar(fit_runs, primitives) + overhead, mape_in = fit_overhead_scalar(fit_runs, primitives, args.nics_per_node) print(f"[fit_overhead] T_overhead = {overhead*1e3:.3f} ms in-sample MAPE = {mape_in*100:.2f}%") result = { "overhead_seconds": overhead, "fit_filter": args.fit_filter, + "nics_per_node": args.nics_per_node, "num_fit_points": len(fit_runs), "num_held_out": len(held_runs), "in_sample_mape": mape_in, diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_gather.py b/experiments/cost_model_benchmarks/benchmarks/bench_gather.py index d62c71b..dc0bb5e 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_gather.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_gather.py @@ -128,8 +128,20 @@ def gather_fn(): gather_times = cuda_timed(gather_fn, warmup=args.warmup, trials=args.trials) # --- Backward scatter-add --- + # grad_x.zero_() is deliberately NOT in the timed region. grad_x is + # [N, F] (10+ GB at the default N), so zeroing it is an O(N) HBM + # write costing ~3.1 ms -- while the scatter-add being measured is + # O(k) and costs ~13 us at small k. Timing them together put a + # k-independent 3.1 ms floor under every measurement (233x the real + # cost at k=1000), which fit_gather then absorbed into + # launch_overhead_seconds. That silently made the fitted scatter + # model a function of N, even though it is applied to O(k) byte + # counts (send_total * F * 4) in the cost model. + # + # Not re-zeroing means grad_x accumulates across trials. That is + # harmless here: scatter_add_ runtime does not depend on the + # accumulated values, and grad_y is all-ones so the sums stay small. def scatter_fn(): - grad_x.zero_() grad_x.scatter_add_(0, idx_expanded, grad_y) scatter_times = cuda_timed(scatter_fn, warmup=args.warmup, trials=args.trials) diff --git a/experiments/cost_model_benchmarks/run_experiments.sh b/experiments/cost_model_benchmarks/run_experiments.sh index c33c0ac..9c1a95d 100755 --- a/experiments/cost_model_benchmarks/run_experiments.sh +++ b/experiments/cost_model_benchmarks/run_experiments.sh @@ -164,10 +164,18 @@ python -m visualization.plot_compute \ --edge-vertex "${DATA_DIR}/compute_edge_vswp.json" --edge-edge "${DATA_DIR}/compute_edge_eswp.json" \ --primitives "${DATA_DIR}/fitted_primitives.json" --output "${FIG_DIR}/compute" -python -m visualization.plot_gather \ - --contiguous "${DATA_DIR}/gather_contiguous.json" --clustered "${DATA_DIR}/gather_clustered.json" \ - --random "${DATA_DIR}/gather_random.json" --fitted "${DATA_DIR}/fitted_primitives.json" \ - --output "${FIG_DIR}/gather" +# bench_gather records BOTH gather_trials_seconds and +# scatter_add_trials_seconds, and fit_primitives fits both, but plot_gather +# renders one --operation per invocation (default: gather). Without the +# second call the scatter-add fit is never looked at by a human -- which is +# how an O(N) zero-fill inside the timed region went unnoticed there while +# the (clean) gather curve looked fine. +for op in gather scatter_add; do + python -m visualization.plot_gather \ + --contiguous "${DATA_DIR}/gather_contiguous.json" --clustered "${DATA_DIR}/gather_clustered.json" \ + --random "${DATA_DIR}/gather_random.json" --fitted "${DATA_DIR}/fitted_primitives.json" \ + --operation "${op}" --output "${FIG_DIR}/${op}" +done python -m visualization.plot_pingpong \ --intra "${DATA_DIR}/pingpong_intra.json" --primitives "${DATA_DIR}/fitted_primitives.json" \ From 4d4f9dfdf48ab63c8fa8bebd52afa3f974083b02 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Sat, 1 Aug 2026 00:21:50 -0400 Subject: [PATCH 34/39] Add SpMM GCN layer and F-dependent T_comp; fix concurrency benchmark instrumentation --- .../analysis/fit_overhead.py | 9 +- .../analysis/fit_primitives.py | 119 +++++++---- .../benchmarks/bench_compute.py | 156 ++++++++++---- .../benchmarks/bench_concurrency.py | 198 +++++++++++------- .../benchmarks/nn_layer_common.py | 52 +++++ .../cost_model_benchmarks/run_experiments.sh | 84 ++++++-- 6 files changed, 442 insertions(+), 176 deletions(-) diff --git a/experiments/cost_model_benchmarks/analysis/fit_overhead.py b/experiments/cost_model_benchmarks/analysis/fit_overhead.py index 754ab84..a92edcc 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_overhead.py +++ b/experiments/cost_model_benchmarks/analysis/fit_overhead.py @@ -27,6 +27,8 @@ import numpy as np +from analysis import compute_forms + # --------------------------------------------------------------------------- # Cost model (without overhead) @@ -110,9 +112,10 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, # forward+backward together instead would double-count the forward pass. comp_params = primitives.get("compute", {}).get(model_type, {}).get("backward", None) if comp_params: - T_comp = (comp_params["coeff_V"] * n_total - + comp_params["coeff_E"] * n_edges_local - + comp_params["intercept"]) + # Which columns this evaluates depends on the form fit_compute chose + # (fixed-F legacy, edge-centric, or vertex-centric). compute_forms owns + # both sides so the fit and this evaluation cannot drift apart. + T_comp = compute_forms.evaluate(comp_params, n_total, n_edges_local, F) else: T_comp = 0.0 T_comp = max(T_comp, 0.0) diff --git a/experiments/cost_model_benchmarks/analysis/fit_primitives.py b/experiments/cost_model_benchmarks/analysis/fit_primitives.py index 44a7478..b0dfcf6 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_primitives.py +++ b/experiments/cost_model_benchmarks/analysis/fit_primitives.py @@ -33,6 +33,8 @@ from scipy.optimize import curve_fit from scipy.special import expit +from analysis import compute_forms + # --------------------------------------------------------------------------- # Helpers @@ -137,44 +139,60 @@ def fit_network(records: list) -> dict: # --------------------------------------------------------------------------- -def fit_compute(records: list, timing_key: str = "forward_trials_seconds") -> dict: - """Fit compute cost as a function of |V| and |E|. +def fit_compute(records: list, timing_key: str = "forward_trials_seconds", + model_type: str = "gcn") -> dict: + """Fit compute cost as a function of |V|, |E| and the feature dim F. - Uses multiple linear regression: T = a * |V| + b * |E| + c + The design matrix depends on the layer's structure and on whether F was + actually swept — see ``analysis/compute_forms.py``, which owns the form + definitions so that fitting and prediction cannot disagree. With single-F + data this degrades to the original ``T = a|V| + b|E| + c``. """ - V_arr, E_arr, T_arr = [], [], [] + V_arr, E_arr, F_arr, T_arr = [], [], [], [] for rec in records: + default_F = rec.get("config", {}).get("feature_dim") for meas in rec["measurements"]: - V_arr.append(meas["params"]["num_vertices"]) - E_arr.append(meas["params"]["num_edges"]) + params = meas["params"] + V_arr.append(params["num_vertices"]) + E_arr.append(params["num_edges"]) + F_arr.append(params.get("feature_dim", default_F)) T_arr.append(median_of_trials(meas[timing_key])) + if any(f is None for f in F_arr): + raise ValueError( + "Some compute measurements have no feature_dim, in params or in " + "the file's config block. F is a design-matrix input now, so it " + "cannot be inferred. Re-run bench_compute.py to regenerate." + ) + V_arr = np.array(V_arr, dtype=float) E_arr = np.array(E_arr, dtype=float) + F_arr = np.array(F_arr, dtype=float) T_arr = np.array(T_arr, dtype=float) - # Design matrix: [V, E, 1], weighted by relative error (1/T) for - # consistency with fit_gather's curve_fit(sigma=T_arr) — |V|/|E| sweeps - # are log-spaced over several orders of magnitude, so an unweighted fit - # would be dominated by the largest graphs and leave the small-graph - # intercept (which matters most for the crossover/tipping-point - # analysis) poorly constrained. - A = np.column_stack([V_arr, E_arr, np.ones_like(V_arr)]) + form = compute_forms.select_form(model_type, len(np.unique(F_arr))) + + # Weighted by relative error (1/T) for consistency with fit_gather's + # curve_fit(sigma=T_arr) — |V|/|E| sweeps are log-spaced over several + # orders of magnitude, so an unweighted fit would be dominated by the + # largest graphs and leave the small-graph intercept (which matters most + # for the crossover/tipping-point analysis) poorly constrained. + A = np.column_stack(compute_forms.design_columns(form, V_arr, E_arr, F_arr)) w = 1.0 / T_arr - A_w = A * w[:, None] - b_w = T_arr * w - result, _, _, _ = np.linalg.lstsq(A_w, b_w, rcond=None) - coeff_V, coeff_E, intercept = result + result, _, _, _ = np.linalg.lstsq(A * w[:, None], T_arr * w, rcond=None) T_pred = A @ result r2 = r_squared(T_arr, T_pred) - return { - "coeff_V": float(coeff_V), - "coeff_E": float(coeff_E), - "intercept": float(intercept), + out = { + "form": form, "r_squared": r2, "_num_points": len(T_arr), + "_distinct_feature_dims": sorted(float(f) for f in np.unique(F_arr)), } + out.update( + {name: float(v) for name, v in zip(compute_forms.FORM_COEFFS[form], result)} + ) + return out # --------------------------------------------------------------------------- @@ -279,6 +297,10 @@ def parse_args(): p.add_argument("--pingpong-inter", nargs="+", default=[], metavar="FILE") p.add_argument("--compute-gcn", nargs="+", default=[], metavar="FILE") p.add_argument("--compute-edge", nargs="+", default=[], metavar="FILE") + p.add_argument("--compute-gcn-spmm", nargs="+", default=[], metavar="FILE", + help="Sweeps for the transform-then-aggregate SpMM GCN " + "layer (--model gcn_spmm). Fit with the vertex-centric " + "form; see analysis/compute_forms.py.") p.add_argument("--gather-contiguous", nargs="+", default=[], metavar="FILE") p.add_argument("--gather-clustered", nargs="+", default=[], metavar="FILE") p.add_argument("--gather-random", nargs="+", default=[], metavar="FILE") @@ -324,28 +346,43 @@ def main(): # Compute comp = {} - if args.compute_gcn: - recs = load_json_files(args.compute_gcn) - comp["gcn"] = { - "forward": fit_compute(recs, "forward_trials_seconds"), - "backward": fit_compute(recs, "backward_trials_seconds"), - } - print( - f"[compute/gcn] coeff_V={comp['gcn']['forward']['coeff_V']:.3e} " - f"coeff_E={comp['gcn']['forward']['coeff_E']:.3e} " - f"R²={comp['gcn']['forward']['r_squared']:.4f}" - ) - if args.compute_edge: - recs = load_json_files(args.compute_edge) - comp["edge"] = { - "forward": fit_compute(recs, "forward_trials_seconds"), - "backward": fit_compute(recs, "backward_trials_seconds"), + for model_type, files in [ + ("gcn", args.compute_gcn), + ("edge", args.compute_edge), + ("gcn_spmm", args.compute_gcn_spmm), + ]: + if not files: + continue + recs = load_json_files(files) + comp[model_type] = { + "forward": fit_compute(recs, "forward_trials_seconds", model_type), + "backward": fit_compute(recs, "backward_trials_seconds", model_type), } - print( - f"[compute/edge] coeff_V={comp['edge']['forward']['coeff_V']:.3e} " - f"coeff_E={comp['edge']['forward']['coeff_E']:.3e} " - f"R²={comp['edge']['forward']['r_squared']:.4f}" + fwd = comp[model_type]["forward"] + coeff_str = " ".join( + f"{name}={fwd[name]:.3e}" + for name in compute_forms.FORM_COEFFS[fwd["form"]] ) + print(f"[compute/{model_type}] form={fwd['form']} {coeff_str} " + f"R²={fwd['r_squared']:.4f}") + if fwd["form"] == "legacy_VE": + print(f"[compute/{model_type}] NOTE: only one feature dim " + f"({fwd['_distinct_feature_dims']}) in the input, so the " + "F-dependent form is unidentifiable and T_comp was fit at " + "fixed F. Sweep --feature-dim to get F*.") + else: + avg_deg = None + if fwd["form"] == "vertex_centric": + # F* depends on |E|/|V| for this form; report at the degree the + # benchmark actually used rather than inventing one. + degs = { + m["params"]["num_edges"] / m["params"]["num_vertices"] + for r in recs for m in r["measurements"] + } + avg_deg = sum(degs) / len(degs) + f_star = compute_forms.crossover_feature_dim(fwd, avg_deg) + print(f"[compute/{model_type}] F* = {f_star:.1f}" + ( + f" (at mean degree {avg_deg:.1f})" if avg_deg else "")) result["compute"] = comp # Gather diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_compute.py b/experiments/cost_model_benchmarks/benchmarks/bench_compute.py index ea7ed6d..445a441 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_compute.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_compute.py @@ -3,15 +3,22 @@ Single-GPU benchmark. Fits f_comp(|Ṽ|, |Ẽ|) for two message-function variants: -* ``gcn`` — GCN-like: φ(h_b) = W h_b (source-only linear transform) -* ``edge`` — Edge-conditioned: φ(h_b, h_a, e_ba) = MLP([h_b, h_a, e_ba]) - with a 2-layer MLP (hidden dim = feature_dim) +* ``gcn`` — GCN-like, applied per edge: φ(h_b) = W h_b. Costs |E|F². +* ``edge`` — Edge-conditioned: φ(h_b, h_a, e_ba) = MLP([h_b, h_a, e_ba]) + with a 2-layer MLP (hidden dim = feature_dim). Costs |E|F². +* ``gcn_spmm`` — A real GCN: transform-then-aggregate, dense [V,F]x[F,F] + followed by a sparse matmul. Costs |V|F² + |E|F. Two sweep modes (controlled by ``--sweep``): * ``vertices`` — vary |V| with |E| fixed at ``--fixed-value`` * ``edges`` — vary |E| with |V| fixed at ``--fixed-value`` +Sweeping ``--feature-dim`` across several runs is what makes ``T_comp``'s F +dependence identifiable: at a single F the F² and F terms are collinear, so +compute-bound and bandwidth-bound layers are indistinguishable. See +``analysis/compute_forms.py``. + Usage:: python -m benchmarks.bench_compute \\ @@ -73,6 +80,39 @@ def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor: return out +class GCNSpMMLayer(nn.Module): + """A real GCN: transform-then-aggregate via sparse matmul. + + ``GCNLayer`` above transforms *per edge* (``|E| * F^2``); this transforms + the vertex matrix once (``|V| * F^2``) and aggregates with an SpMM + (``|E| * F``). At average degree 20 the per-edge form does ~20x the FLOPs, + so the two need separate fits — see analysis/compute_forms.py. + """ + + def __init__(self, feature_dim: int): + super().__init__() + self.linear = nn.Linear(feature_dim, feature_dim, bias=False) + + def forward(self, x: torch.Tensor, adj_sparse: torch.Tensor) -> torch.Tensor: + return torch.sparse.mm(adj_sparse, self.linear(x)) + + +def build_sparse_adj(edge_index: torch.Tensor, num_vertices: int, + device: torch.device, sparse_format: str) -> torch.Tensor: + """Square [V, V] aggregation matrix: row = dst, column = src. + + Built once per sweep point, outside the timed region — in the distributed + setting the adjacency comes prebuilt from the comm pattern, so including + its construction here would measure something the model never pays for. + """ + indices = torch.stack([edge_index[1], edge_index[0]]) + values = torch.ones(edge_index.shape[1], dtype=torch.float32, device=device) + adj = torch.sparse_coo_tensor( + indices, values, size=(num_vertices, num_vertices), device=device + ).coalesce() + return adj.to_sparse_csr() if sparse_format == "csr" else adj + + class EdgeConditionedLayer(nn.Module): """Edge-conditioned: φ(h_b, h_a, e_ba) = MLP([h_b, h_a, e_ba]).""" @@ -103,7 +143,11 @@ def forward( def parse_args(): p = argparse.ArgumentParser(description="GNN compute primitive benchmark") - p.add_argument("--model", choices=["gcn", "edge"], required=True) + p.add_argument("--model", choices=["gcn", "edge", "gcn_spmm"], required=True) + p.add_argument("--sparse-format", choices=["csr", "coo"], default="csr", + help="Sparse layout for --model gcn_spmm. CSR is faster; " + "switch to coo if torch.sparse.mm's backward rejects " + "CSR on this PyTorch build.") p.add_argument("--sweep", choices=["vertices", "edges"], required=True) p.add_argument("--min", type=int, default=1_000, dest="sweep_min") p.add_argument("--max", type=int, default=1_000_000, dest="sweep_max") @@ -132,6 +176,8 @@ def main(): # Build model if args.model == "gcn": model = GCNLayer(F).to(device) + elif args.model == "gcn_spmm": + model = GCNSpMMLayer(F).to(device) else: model = EdgeConditionedLayer(F).to(device) @@ -153,46 +199,55 @@ def main(): else: num_v, num_e = args.fixed_value, val - # Synthetic data - x = torch.randn(num_v, F, device=device, requires_grad=True) - edge_index = erdos_renyi_edges(num_v, num_e, device) - edge_attr = ( - torch.randn(num_e, F, device=device) if args.model == "edge" else None - ) - - # Forward timing - def fwd(): - if args.model == "gcn": - model(x, edge_index) - else: - model(x, edge_index, edge_attr) - - fwd_times = cuda_timed(fwd, warmup=args.warmup, trials=args.trials) + # A sweep over feature dims will eventually hit a cell that does not + # fit (edge-conditioned at F=512 with |E|=1e6 is the tightest). Skip + # that point and keep going: an unguarded OOM here loses every + # measurement already collected in this run. + x = edge_index = edge_attr = adj = None + try: + # Synthetic data + x = torch.randn(num_v, F, device=device, requires_grad=True) + edge_index = erdos_renyi_edges(num_v, num_e, device) + edge_attr = ( + torch.randn(num_e, F, device=device) if args.model == "edge" else None + ) + adj = ( + build_sparse_adj(edge_index, num_v, device, args.sparse_format) + if args.model == "gcn_spmm" + else None + ) - # Backward timing (run fwd first to get a graph) - if args.model == "gcn": - out = model(x, edge_index) - else: - out = model(x, edge_index, edge_attr) - loss_ref = out.sum() - - # NOTE: this closure times grad-zero + FORWARD + backward, not the - # backward pass alone — the forward must be re-run inside the timed - # region to rebuild the autograd graph consumed by each .backward() - # call. So "backward_trials_seconds" below is really the combined - # forward+backward cost (which is what the assembled cost model's - # T_comp needs, since bench_end_to_end.py times the same op set). - # Do not subtract or add "forward_trials_seconds" to it. - def bwd(): - if x.grad is not None: - x.grad.zero_() - if args.model == "gcn": - out_inner = model(x, edge_index) - else: - out_inner = model(x, edge_index, edge_attr) - out_inner.sum().backward() - - bwd_times = cuda_timed(bwd, warmup=args.warmup, trials=args.trials) + def call_model(): + if args.model == "gcn": + return model(x, edge_index) + if args.model == "gcn_spmm": + return model(x, adj) + return model(x, edge_index, edge_attr) + + # Forward timing + fwd_times = cuda_timed(call_model, warmup=args.warmup, trials=args.trials) + + # NOTE: this closure times grad-zero + FORWARD + backward, not the + # backward pass alone — the forward must be re-run inside the timed + # region to rebuild the autograd graph consumed by each .backward() + # call. So "backward_trials_seconds" below is really the combined + # forward+backward cost (which is what the assembled cost model's + # T_comp needs, since bench_end_to_end.py times the same op set). + # Do not subtract or add "forward_trials_seconds" to it. + def bwd(): + if x.grad is not None: + x.grad.zero_() + call_model().sum().backward() + + bwd_times = cuda_timed(bwd, warmup=args.warmup, trials=args.trials) + except torch.cuda.OutOfMemoryError: + print( + f"[compute/{args.model}] OOM: |V|={num_v:,} |E|={num_e:,} F={F} " + "does not fit; skipping this point and continuing" + ) + del x, edge_index, edge_attr, adj + torch.cuda.empty_cache() + continue measurements.append( { @@ -211,15 +266,28 @@ def bwd(): med_fwd = sorted(fwd_times)[len(fwd_times) // 2] med_bwd = sorted(bwd_times)[len(bwd_times) // 2] print( - f"[compute/{args.model}] |V|={num_v:>8} |E|={num_e:>9} " + f"[compute/{args.model}] |V|={num_v:>8} |E|={num_e:>9} F={F:>4} " f"fwd {1e3*med_fwd:.2f} ms bwd {1e3*med_bwd:.2f} ms" ) + del x, edge_index, edge_attr, adj + torch.cuda.empty_cache() + + if not measurements: + raise RuntimeError( + f"Every sweep point OOMed for model={args.model} F={F} " + f"(|V|/|E| from {args.sweep_min:,} to {args.sweep_max:,}, fixed " + f"{args.fixed_value:,}). Writing an empty file would silently " + "produce a fit with no data behind it. Lower --max, --fixed-value, " + "or --feature-dim." + ) + payload = { "benchmark": "compute", "metadata": collect_metadata(), "config": { "model": args.model, + "sparse_format": args.sparse_format if args.model == "gcn_spmm" else None, "sweep": args.sweep, "sweep_min": args.sweep_min, "sweep_max": args.sweep_max, diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py b/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py index 28c63fe..bf09905 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_concurrency.py @@ -1,16 +1,24 @@ """Benchmark 1.2 — Intra/Inter Concurrency Check. -Verifies that T_both / max(T_intra, T_inter) ≈ 1, i.e. that NVLink and -InfiniBand transfers overlap when issued simultaneously. +Measures T_both / max(T_intra, T_inter), i.e. whether NVLink and InfiniBand +transfers overlap when issued together. A ratio near 1.0 justifies +``T_comm = max(T_intra, T_inter)`` in the cost model; a ratio near +``(T_intra + T_inter) / max(...)`` means they serialize and the tiers must be +added instead. Requires exactly 4 ranks across 2 nodes: Node A: rank 0 (A0), rank 1 (A1) Node B: rank 2 (B0), rank 3 (B1) -Three conditions at a fixed message size: - 1. intra-only — A0↔A1 and B0↔B1 (no cross-node traffic) - 2. inter-only — A0↔B0 and A1↔B1 (no intra-node traffic) - 3. concurrent — all four exchanges in the same window (separate streams) +Four conditions at a fixed message size: + 1. intra-only — A0↔A1 and B0↔B1 (no cross-node traffic) + 2. inter-only — A0↔B0 and A1↔B1 (no intra-node traffic) + 3. concurrent — both, posted as ONE batch_isend_irecv group call, + matching how DGraph's halo exchange issues a single + all_to_all_single over every peer + 4. sequential_calls — both, posted as two group calls with a wait in + between; the anti-pattern, measured so the cost of + not batching is quantified rather than assumed Each rank logs its own wall time per trial. Rank 0 collects and writes JSON. @@ -40,54 +48,77 @@ # Helpers # --------------------------------------------------------------------------- -def _exchange(tensor_send: torch.Tensor, tensor_recv: torch.Tensor, - peer: int, stream: torch.cuda.Stream) -> None: - """Non-blocking send+recv on *stream* with *peer*.""" - with torch.cuda.stream(stream): - send_op = dist.P2POp(dist.isend, tensor_send, peer) - recv_op = dist.P2POp(dist.irecv, tensor_recv, peer) - reqs = dist.batch_isend_irecv([send_op, recv_op]) - for req in reqs: +# NOTE ON WHAT CHANGED AND WHY +# +# The original version of this benchmark issued the intra and inter exchanges +# as two separate ``batch_isend_irecv`` group calls, each followed immediately +# by ``req.wait()``, with the two calls placed on different CUDA streams. It +# reported T_concurrent / max(T_intra, T_inter) = 1.60 -- i.e. no overlap -- +# and the cost model was briefly changed to add the two tiers because of it. +# +# That measurement was an artefact of how the ops were posted, not a property +# of the hardware: +# +# 1. ``req.wait()`` sat inside the per-peer helper, so the intra exchange was +# joined onto the calling stream *before* the inter exchange was even +# enqueued. The host serialized the two before NCCL saw them. +# 2. Both went through the default process group. NCCL orders operations on a +# communicator, so two group calls become two kernel launches executed in +# issue order. Putting them on different CUDA streams does not help: the +# user-facing stream is not where NCCL performs the transfer. +# 3. Most importantly, it did not match the code being modelled. +# DGraph's halo exchange issues a single ``dist.all_to_all_single`` over +# every peer at once (NCCLBackendEngine.py), which NCCL schedules across +# channels and *can* overlap NVLink and IB portions of. +# +# The end-to-end data agrees with (3): at K=8 the model's over-prediction is +# exactly T_intra (ratio 0.98-1.01 at the two larger sizes), i.e. intra-node +# traffic is fully hidden under the inter-node transfer in the real exchange. +# +# So the "concurrent" condition below now posts **one** group call containing +# both peers' ops -- matching the real exchange. The old two-call pattern is +# kept as a separate "sequential_calls" condition, because the gap between the +# two is a directly useful number: it is what NCCL leaves on the table when +# transfers are not batched into a single collective, and therefore part of +# the motivation for a finer-grained backend (NVSHMEM) that can issue +# NVLink stores and IB puts from within one kernel. +# +# CUDA streams are gone entirely: they gave the appearance of controlling +# concurrency while having no effect on it. + +def _ops_for(peer: int, send_buf: torch.Tensor, recv_buf: torch.Tensor) -> list: + """Build the (isend, irecv) P2POp pair for one peer, unposted.""" + return [ + dist.P2POp(dist.isend, send_buf, peer), + dist.P2POp(dist.irecv, recv_buf, peer), + ] + + +def _post_one_group(ops: list) -> None: + """Post *all* ops as a single NCCL group call and wait for completion. + + One group call is the point: NCCL is free to run the enclosed transfers + concurrently across channels. Splitting them across several calls forbids + that, regardless of streams. + """ + for req in dist.batch_isend_irecv(ops): req.wait() -def timed_window(rank: int, - send_buf: torch.Tensor, - recv_buf: torch.Tensor, - peer: int, - stream: torch.cuda.Stream, - warmup: int, - trials: int) -> list: - """Time a single exchange window (one send+recv pair).""" - for _ in range(warmup): - _exchange(send_buf, recv_buf, peer, stream) - torch.cuda.synchronize() - dist.barrier() +def _post_separate_groups(op_groups: list) -> None: + """Post each op list as its own group call, waiting between them. - times = [] - for _ in range(trials): - dist.barrier() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - _exchange(send_buf, recv_buf, peer, stream) - end.record() - torch.cuda.synchronize() - times.append(start.elapsed_time(end) / 1_000.0) - return times + Deliberately the anti-pattern, retained for measurement. + """ + for ops in op_groups: + for req in dist.batch_isend_irecv(ops): + req.wait() -def timed_concurrent(rank: int, - intra_send: torch.Tensor, intra_recv: torch.Tensor, intra_peer: int, - inter_send: torch.Tensor, inter_recv: torch.Tensor, inter_peer: int, - intra_stream: torch.cuda.Stream, - inter_stream: torch.cuda.Stream, - warmup: int, - trials: int) -> list: - """Time intra and inter exchanges issued concurrently on separate streams.""" +def timed(post_fn, warmup: int, trials: int) -> list: + """Time a comm window built by *post_fn* (a zero-arg closure).""" for _ in range(warmup): - _exchange(intra_send, intra_recv, intra_peer, intra_stream) - _exchange(inter_send, inter_recv, inter_peer, inter_stream) + post_fn() torch.cuda.synchronize() dist.barrier() @@ -97,8 +128,7 @@ def timed_concurrent(rank: int, start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() - _exchange(intra_send, intra_recv, intra_peer, intra_stream) - _exchange(inter_send, inter_recv, inter_peer, inter_stream) + post_fn() end.record() torch.cuda.synchronize() times.append(start.elapsed_time(end) / 1_000.0) @@ -152,34 +182,36 @@ def main(): intra_peer = {0: 1, 1: 0, 2: 3, 3: 2}[rank] inter_peer = {0: 2, 1: 3, 2: 0, 3: 1}[rank] - intra_stream = torch.cuda.Stream(device=device) - inter_stream = torch.cuda.Stream(device=device) - intra_send = send_buf.clone() intra_recv = torch.zeros_like(recv_buf) inter_send = send_buf.clone() inter_recv = torch.zeros_like(recv_buf) + intra_ops = lambda: _ops_for(intra_peer, intra_send, intra_recv) + inter_ops = lambda: _ops_for(inter_peer, inter_send, inter_recv) + # --- Condition 1: intra-only --- - times_intra = timed_window( - rank, intra_send, intra_recv, intra_peer, intra_stream, - args.warmup, args.trials + times_intra = timed( + lambda: _post_one_group(intra_ops()), args.warmup, args.trials ) dist.barrier() # --- Condition 2: inter-only --- - times_inter = timed_window( - rank, inter_send, inter_recv, inter_peer, inter_stream, - args.warmup, args.trials + times_inter = timed( + lambda: _post_one_group(inter_ops()), args.warmup, args.trials ) dist.barrier() - # --- Condition 3: concurrent --- - times_concurrent = timed_concurrent( - rank, - intra_send, intra_recv, intra_peer, - inter_send, inter_recv, inter_peer, - intra_stream, inter_stream, + # --- Condition 3: concurrent, ONE group call (matches the real exchange) --- + times_concurrent = timed( + lambda: _post_one_group(intra_ops() + inter_ops()), + args.warmup, args.trials, + ) + dist.barrier() + + # --- Condition 4: the anti-pattern, two group calls with a wait between --- + times_sequential = timed( + lambda: _post_separate_groups([intra_ops(), inter_ops()]), args.warmup, args.trials, ) dist.barrier() @@ -193,6 +225,7 @@ def gather_times(times_local): intra_all = gather_times(times_intra) inter_all = gather_times(times_inter) conc_all = gather_times(times_concurrent) + seq_all = gather_times(times_sequential) if rank == 0: measurements = [ @@ -205,9 +238,15 @@ def gather_times(times_local): "per_rank_trials_seconds": inter_all, }, { - "params": {"condition": "concurrent", "message_bytes": num_elems * 4}, + "params": {"condition": "concurrent", "message_bytes": num_elems * 4, + "posting": "single batch_isend_irecv group call"}, "per_rank_trials_seconds": conc_all, }, + { + "params": {"condition": "sequential_calls", "message_bytes": num_elems * 4, + "posting": "two group calls, wait between"}, + "per_rank_trials_seconds": seq_all, + }, ] payload = { "benchmark": "concurrency", @@ -223,13 +262,30 @@ def gather_times(times_local): "measurements": measurements, } write_result(args.output, payload) + def med(times_all): + t = sorted(times_all[0]) + return t[len(t) // 2] + + t_intra, t_inter = med(intra_all), med(inter_all) + t_conc, t_seq = med(conc_all), med(seq_all) + print( + f"[concurrency] intra = {1e3*t_intra:.3f} ms | " + f"inter = {1e3*t_inter:.3f} ms | " + f"concurrent (1 group call) = {1e3*t_conc:.3f} ms | " + f"sequential (2 calls) = {1e3*t_seq:.3f} ms" + ) + # The overlap verdict the cost model depends on: ratio ~1.0 means the + # tiers overlap (T_comm = max(...)); ratio ~ (T_intra+T_inter)/max + # means they serialize (T_comm = T_intra + T_inter). + print( + f"[concurrency] T_concurrent / max(T_intra, T_inter) = " + f"{t_conc / max(t_intra, t_inter):.2f} " + f"(1.00 = perfect overlap, " + f"{(t_intra + t_inter) / max(t_intra, t_inter):.2f} = full serialization)" + ) print( - f"[concurrency] intra median = " - f"{1e3*sorted(intra_all[0])[len(intra_all[0])//2]:.2f} ms | " - f"inter median = " - f"{1e3*sorted(inter_all[0])[len(inter_all[0])//2]:.2f} ms | " - f"concurrent median = " - f"{1e3*sorted(conc_all[0])[len(conc_all[0])//2]:.2f} ms" + f"[concurrency] cost of NOT batching into one group call: " + f"{1e3*(t_seq - t_conc):+.3f} ms ({100*(t_seq/t_conc - 1):+.1f}%)" ) dist.destroy_process_group() diff --git a/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py b/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py index 63a6af3..8f34ff9 100644 --- a/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py +++ b/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py @@ -23,6 +23,58 @@ def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor: return out +def create_sparse_adj( + edge_index: torch.Tensor, + num_rows: int, + num_cols: int, + device: torch.device, + sparse_format: str = "csr", +) -> torch.Tensor: + """Build the sparse aggregation matrix from an ``[2, E]`` edge index. + + Row = destination (the locally-owned aggregation target, ``edge_index[1]``), + column = source (the neighbour, possibly a halo vertex, ``edge_index[0]``) — + the convention ``DGraph.distributed.commInfo`` establishes and that + ``experiments/OGB/GCN.py::create_sparse_adj`` follows. The matrix is + therefore rectangular in the distributed case: + ``[num_local, num_local + num_halo]``. + + ``sparse_format="coo"`` is a fallback: CSR is faster for SpMM, but + ``torch.sparse.mm``'s backward for CSR operands is narrower than for COO + across PyTorch versions. If a run dies in backward, switch to coo before + concluding anything about the layer. + """ + indices = torch.stack([edge_index[1], edge_index[0]]) + values = torch.ones(edge_index.shape[1], dtype=torch.float32, device=device) + adj = torch.sparse_coo_tensor( + indices, values, size=(num_rows, num_cols), device=device + ).coalesce() + return adj.to_sparse_csr() if sparse_format == "csr" else adj + + +class GCNSpMMLayer(nn.Module): + """A real GCN: transform-then-aggregate, via a sparse matmul. + + ``GCNLayer`` above applies the linear transform *per edge* + (``linear(x[src])``), costing ``|E| * F^2``. This layer transforms the + ``[V, F]`` vertex matrix once (``|V| * F^2``) and then aggregates with an + SpMM (``|E| * F``), which is what a GCN actually is. At average degree 20 + the per-edge form does ~20x the FLOPs, so the two are worth fitting + separately rather than treating one as a stand-in for the other. + + The SpMM consumes the augmented ``[n_local + n_halo, F]`` feature matrix + and emits only the ``n_local`` rows, so the tensor does not grow from + layer to layer. + """ + + def __init__(self, feature_dim: int): + super().__init__() + self.linear = nn.Linear(feature_dim, feature_dim, bias=False) + + def forward(self, x: torch.Tensor, adj_sparse: torch.Tensor) -> torch.Tensor: + return torch.sparse.mm(adj_sparse, self.linear(x)) + + class EdgeConditionedLayer(nn.Module): def __init__(self, feature_dim: int): super().__init__() diff --git a/experiments/cost_model_benchmarks/run_experiments.sh b/experiments/cost_model_benchmarks/run_experiments.sh index 9c1a95d..cb51495 100755 --- a/experiments/cost_model_benchmarks/run_experiments.sh +++ b/experiments/cost_model_benchmarks/run_experiments.sh @@ -24,7 +24,16 @@ shopt -s nullglob GPU_COUNTS=(${GPU_COUNTS:-2 4}) # world sizes to sweep via torchrun --nproc_per_node SEED=${SEED:-42} FEATURE_DIM=${FEATURE_DIM:-128} -MODEL=${MODEL:-gcn} # gcn | edge +MODEL=${MODEL:-gcn} # gcn | edge | gcn_spmm (distributed runs) +# Layer variants and feature dims for the single-GPU compute sweeps (1.3). +# Sweeping F is what makes T_comp's F dependence identifiable: at a single F +# the F^2 and F terms are collinear, so a compute-bound layer and a +# bandwidth-bound one fit the data equally well. See analysis/compute_forms.py. +COMPUTE_MODELS=(${COMPUTE_MODELS:-gcn edge gcn_spmm}) +COMPUTE_FEATURE_DIMS=(${COMPUTE_FEATURE_DIMS:-32 64 128 256 512}) +# Defaults to FEATURE_DIM so overriding one does not silently leave the +# figures pointing at a different feature dim than the distributed runs use. +REF_FEATURE_DIM=${REF_FEATURE_DIM:-$FEATURE_DIM} GRAPH=${GRAPH:-erdos_renyi} # erdos_renyi | sbm PARTITIONER=${PARTITIONER:-balanced} # random | balanced | metis WARMUP=${WARMUP:-10} @@ -42,6 +51,21 @@ FIG_DIR=${FIG_DIR:-figures} MAX_GPU_COUNT=$(printf '%s\n' "${GPU_COUNTS[@]}" | sort -n | tail -1) FIT_FILTER=${FIT_FILTER:-"world_size < ${MAX_GPU_COUNT}"} +# Fail before spending GPU hours rather than at the plotting step at the end: +# plot_compute reads the REF_FEATURE_DIM cell, which only exists if that F was +# actually swept, and the distributed runs use MODEL, which only gets a fit if +# it was among the compute models. +if [[ ! " ${COMPUTE_FEATURE_DIMS[*]} " =~ " ${REF_FEATURE_DIM} " ]]; then + echo "[error] REF_FEATURE_DIM=${REF_FEATURE_DIM} is not in COMPUTE_FEATURE_DIMS" \ + "(${COMPUTE_FEATURE_DIMS[*]}); the figures would have no data to plot." >&2 + exit 1 +fi +if [[ ! " ${COMPUTE_MODELS[*]} " =~ " ${MODEL} " ]]; then + echo "[error] MODEL=${MODEL} is not in COMPUTE_MODELS (${COMPUTE_MODELS[*]});" \ + "the distributed runs would have no T_comp fit to predict with." >&2 + exit 1 +fi + mkdir -p "$DATA_DIR" "$FIG_DIR" step() { echo; echo "=== $* ==="; } @@ -55,18 +79,39 @@ step() { echo; echo "=== $* ==="; } # single MODEL the distributed crossover/end-to-end sweep below uses — # so both model types are always benchmarked here. # ------------------------------------------------------------------------- -for m in gcn edge; do - step "1.3 compute -- ${m}, vertex sweep" - python -m benchmarks.bench_compute --model "${m}" --sweep vertices \ - --min 1000 --max 1000000 --steps 10 --fixed-value 200000 \ - --feature-dim "${FEATURE_DIM}" --warmup "${WARMUP}" --trials "${TRIALS}" \ - --output "${DATA_DIR}/compute_${m}_vswp.json" --seed "${SEED}" - - step "1.3 compute -- ${m}, edge sweep" - python -m benchmarks.bench_compute --model "${m}" --sweep edges \ - --min 1000 --max 1000000 --steps 10 --fixed-value 200000 \ - --feature-dim "${FEATURE_DIM}" --warmup "${WARMUP}" --trials "${TRIALS}" \ - --output "${DATA_DIR}/compute_${m}_eswp.json" --seed "${SEED}" +# Both sweep directions are required per (model, F) cell: fit_compute needs a +# vertex sweep (|E| fixed) and an edge sweep (|V| fixed) to identify the |V| +# and |E| coefficients independently. Filenames carry F so the cells stay +# separable and a stale file from a different F cannot be pooled in. +# +# FIT_COMPUTE_ARGS is accumulated here rather than reconstructed later: the +# outer loop is per-model, so each --compute-* flag is followed contiguously +# by exactly the files just generated for it, which is what nargs="+" wants. +# (Built as a flat array, not an associative one — bash 3.2, still the default +# on macOS, has no `declare -A`.) +FIT_COMPUTE_ARGS=() +for m in "${COMPUTE_MODELS[@]}"; do + case "$m" in + gcn) FIT_COMPUTE_ARGS+=(--compute-gcn) ;; + edge) FIT_COMPUTE_ARGS+=(--compute-edge) ;; + gcn_spmm) FIT_COMPUTE_ARGS+=(--compute-gcn-spmm) ;; + *) echo "[error] unknown compute model '${m}'" >&2; exit 1 ;; + esac + for f in "${COMPUTE_FEATURE_DIMS[@]}"; do + for sweep in vertices edges; do + case "$sweep" in + vertices) tag=vswp ;; + edges) tag=eswp ;; + esac + out="${DATA_DIR}/compute_${m}_F${f}_${tag}.json" + step "1.3 compute -- ${m}, F=${f}, ${sweep} sweep" + python -m benchmarks.bench_compute --model "${m}" --sweep "${sweep}" \ + --min 1000 --max 1000000 --steps 10 --fixed-value 200000 \ + --feature-dim "${f}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${out}" --seed "${SEED}" + FIT_COMPUTE_ARGS+=("${out}") + done + done done step "1.4 gather -- contiguous / clustered / random" @@ -133,8 +178,7 @@ step "fit_primitives" # silently pool incompatible data into one regression. python -m analysis.fit_primitives \ --pingpong-intra "${DATA_DIR}/pingpong_intra.json" \ - --compute-gcn "${DATA_DIR}/compute_gcn_vswp.json" "${DATA_DIR}/compute_gcn_eswp.json" \ - --compute-edge "${DATA_DIR}/compute_edge_vswp.json" "${DATA_DIR}/compute_edge_eswp.json" \ + "${FIT_COMPUTE_ARGS[@]}" \ --gather-contiguous "${DATA_DIR}/gather_contiguous.json" \ --gather-clustered "${DATA_DIR}/gather_clustered.json" \ --gather-random "${DATA_DIR}/gather_random.json" \ @@ -159,9 +203,15 @@ python -m analysis.compute_predictions \ # 5. Visualization. # ------------------------------------------------------------------------- step "visualization" +# plot_compute shows one (model, F) cell per panel, so it gets the reference +# feature dim. NOTE: there is no figure for the F sweep itself yet — the +# fitted F* values are printed by fit_primitives and stored in +# fitted_primitives.json, but nothing plots T vs F. python -m visualization.plot_compute \ - --gcn-vertex "${DATA_DIR}/compute_gcn_vswp.json" --gcn-edge "${DATA_DIR}/compute_gcn_eswp.json" \ - --edge-vertex "${DATA_DIR}/compute_edge_vswp.json" --edge-edge "${DATA_DIR}/compute_edge_eswp.json" \ + --gcn-vertex "${DATA_DIR}/compute_gcn_F${REF_FEATURE_DIM}_vswp.json" \ + --gcn-edge "${DATA_DIR}/compute_gcn_F${REF_FEATURE_DIM}_eswp.json" \ + --edge-vertex "${DATA_DIR}/compute_edge_F${REF_FEATURE_DIM}_vswp.json" \ + --edge-edge "${DATA_DIR}/compute_edge_F${REF_FEATURE_DIM}_eswp.json" \ --primitives "${DATA_DIR}/fitted_primitives.json" --output "${FIG_DIR}/compute" # bench_gather records BOTH gather_trials_seconds and From 1cac71bb323206ad50438da9d7a775dc46cc2de6 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Sat, 1 Aug 2026 00:25:40 -0400 Subject: [PATCH 35/39] Add updated compute forms and mutlinode runner --- .../analysis/compute_forms.py | 131 ++++++++++++++++++ .../run_experiments_multinode.sh | 51 ++++++- 2 files changed, 179 insertions(+), 3 deletions(-) create mode 100644 experiments/cost_model_benchmarks/analysis/compute_forms.py diff --git a/experiments/cost_model_benchmarks/analysis/compute_forms.py b/experiments/cost_model_benchmarks/analysis/compute_forms.py new file mode 100644 index 0000000..be5b580 --- /dev/null +++ b/experiments/cost_model_benchmarks/analysis/compute_forms.py @@ -0,0 +1,131 @@ +"""Design matrices for the T_comp fit — the single definition shared by +``fit_primitives.fit_compute`` (which fits the coefficients) and +``fit_overhead.predict_layer_time`` (which evaluates them). + +Keeping one definition matters: a mismatch between the columns used to fit and +the expression used to predict is silent — it produces plausible numbers that +are simply wrong, with no exception and no bad R^2 to notice. + +Forms +----- +``legacy_VE`` T = a|V| + b|E| + c + The original fixed-F form. Exact as a linearization at one feature + dimension, but says nothing about how cost scales with F. + +``edge_centric`` T = a|E|F^2 + b|E|F + c|V|F + d + For layers that apply the dense transform *per edge*: ``GCNLayer`` + (``linear(x[src])``) and ``EdgeConditionedLayer`` (an MLP over + ``[h_src, h_dst, e]``). The F^2 term is the GEMM, the |E|F terms are the + gather and scatter traffic, the |V|F term is the output/zero-init. + +``vertex_centric`` T = a|V|F^2 + b|E|F + c|V|F + d + For transform-then-aggregate layers: ``GCNSpMMLayer``. The GEMM is now + over vertices, and the SpMM contributes |E|F. + +Why F matters +------------- +At a single F, |E|F^2 and |E|F are perfectly collinear, so fixed-F data cannot +distinguish a compute-bound layer from a bandwidth-bound one — both hypotheses +fit equally well. Sweeping F separates them, and the fitted ratio gives the +feature dimension at which a layer crosses over: + + F* = b/a (edge_centric) the F where bandwidth and compute terms match + +``fit_compute`` therefore refuses to fit an F-dependent form to single-F data +and falls back to ``legacy_VE``; see ``select_form``. +""" + +import numpy as np + + +# Which form each --model value takes. +MODEL_FORM = { + "gcn": "edge_centric", + "edge": "edge_centric", + "gcn_spmm": "vertex_centric", +} + +# Coefficient names per form, in design-matrix column order. +FORM_COEFFS = { + "legacy_VE": ["coeff_V", "coeff_E", "intercept"], + "edge_centric": ["coeff_EF2", "coeff_EF", "coeff_VF", "intercept"], + "vertex_centric": ["coeff_VF2", "coeff_EF", "coeff_VF", "intercept"], +} + + +def design_columns(form: str, V, E, F): + """Return the design-matrix columns for *form* as a list of arrays. + + V, E, F may be scalars or arrays; they are broadcast together. + """ + V = np.asarray(V, dtype=float) + E = np.asarray(E, dtype=float) + F = np.asarray(F, dtype=float) + ones = np.ones_like(V * E * F) + + if form == "legacy_VE": + return [V * ones, E * ones, ones] + if form == "edge_centric": + return [E * F * F, E * F, V * F, ones] + if form == "vertex_centric": + return [V * F * F, E * F, V * F, ones] + raise ValueError( + f"Unknown compute form {form!r}; expected one of {sorted(FORM_COEFFS)}" + ) + + +def select_form(model_type: str, distinct_feature_dims: int) -> str: + """Pick the fit form for *model_type*, downgrading if F was not swept. + + An F-dependent form needs at least two distinct feature dimensions to be + identifiable: at a single F, |V|F^2 and |V|F differ only by a constant + factor, so the columns are collinear and the split between them is + arbitrary. lstsq would still return *an* answer — one that fits the + training data and extrapolates nonsensically in F. Fall back instead. + """ + if distinct_feature_dims < 2: + return "legacy_VE" + return MODEL_FORM.get(model_type, "edge_centric") + + +def evaluate(params: dict, V, E, F) -> float: + """Evaluate a fitted compute model. *params* is a fit_compute() result.""" + form = params.get("form", "legacy_VE") + coeffs = [params[name] for name in FORM_COEFFS[form]] + cols = design_columns(form, V, E, F) + return float(sum(c * col for c, col in zip(coeffs, cols))) + + +def crossover_feature_dim(params: dict, avg_degree: float = None) -> float: + """F* — the feature dim where the GEMM term overtakes the aggregation term. + + Below F*, the layer's cost is dominated by moving features (linear in F); + above it, by the dense transform (quadratic in F). + + The two forms differ in whether F* depends on the graph: + + - ``edge_centric``: compute is ``a|E|F^2``, aggregation ``b|E|F``. The |E| + cancels, so ``F* = b/a`` — a property of the layer alone. + - ``vertex_centric``: compute is ``a|V|F^2``, aggregation ``b|E|F``. These + have different geometry, so ``F* = (b/a) * (|E|/|V|)`` = ``(b/a) * + avg_degree`` — the crossover moves with graph density, which is the + whole point of transform-then-aggregate. *avg_degree* is required here. + + Returns NaN for the legacy form, which has no F dependence. + """ + form = params.get("form", "legacy_VE") + if form == "legacy_VE": + return float("nan") + a = params[FORM_COEFFS[form][0]] # the F^2 coefficient + b = params["coeff_EF"] + if not a > 0: + return float("nan") + if form == "edge_centric": + return float(b / a) + if avg_degree is None: + raise ValueError( + "avg_degree is required for the vertex_centric form: its crossover " + "F* = (b/a) * |E|/|V| depends on graph density, unlike the " + "edge_centric form where |E| cancels." + ) + return float((b / a) * avg_degree) diff --git a/experiments/cost_model_benchmarks/run_experiments_multinode.sh b/experiments/cost_model_benchmarks/run_experiments_multinode.sh index 80091ad..a845324 100755 --- a/experiments/cost_model_benchmarks/run_experiments_multinode.sh +++ b/experiments/cost_model_benchmarks/run_experiments_multinode.sh @@ -48,6 +48,12 @@ if [[ "${NNODES}" -lt 2 ]]; then fi GPUS_PER_NODE=${GPUS_PER_NODE:-4} +# Network interfaces per node -- a topology constant, from `nvidia-smi topo -m`. +# Co-located ranks share NICs, so each rank sees B_inter divided by +# ceil(GPUS_PER_NODE / NICS_PER_NODE). Getting this wrong does not error, it +# just silently mispredicts: on 2 nodes x 4 GPU with 2 NICs, leaving it at 1 +# under-predicts every multi-node run by 15-31%. +NICS_PER_NODE=${NICS_PER_NODE:-1} SEED=${SEED:-42} FEATURE_DIM=${FEATURE_DIM:-128} MODEL=${MODEL:-gcn} @@ -73,6 +79,16 @@ TORCHRUN_HPC_EXTRA_ARGS=${TORCHRUN_HPC_EXTRA_ARGS:-} # e.g. "-r mpi --comm-bac K=$((NNODES * GPUS_PER_NODE)) # total world size for the full-scale runs +# The default NICS_PER_NODE=1 is only correct when each rank owns a NIC. Warn +# loudly otherwise: this is the one setting whose misconfiguration produces a +# confidently wrong model rather than an error. +if [[ "${NICS_PER_NODE}" -eq 1 && "${GPUS_PER_NODE}" -gt 1 ]]; then + echo "[warn] NICS_PER_NODE=1 with GPUS_PER_NODE=${GPUS_PER_NODE}: the model" >&2 + echo " will assume every rank has a NIC to itself. If this node has" >&2 + echo " fewer NICs than GPUs, inter-node time will be under-predicted." >&2 + echo " Check 'nvidia-smi topo -m' and set NICS_PER_NODE accordingly." >&2 +fi + mkdir -p "$DATA_DIR" "$FIG_DIR" # Canonicalize to absolute paths -- see launch-directory note above. DATA_DIR=$(cd "${DATA_DIR}" && pwd) @@ -191,22 +207,51 @@ fi # into the multi-node regime specifically. FIT_FILTER=${FIT_FILTER:-"world_size < ${K}"} +# The compute sweeps are single-GPU and are produced by run_experiments.sh, +# which writes one file per (model, F, sweep-direction) cell. Collect whatever +# cells exist rather than naming them: this script does not know which +# COMPUTE_FEATURE_DIMS the other one was run with. +# +# The F marker in the glob is deliberate — it matches only the current naming +# scheme, so a stray file from an ad-hoc run (which historically did pool +# incompatible feature dims into one regression) cannot be swept in. +FIT_COMPUTE_ARGS=() +for m in gcn edge gcn_spmm; do + files=("${DATA_DIR}/compute_${m}_F"*_vswp.json "${DATA_DIR}/compute_${m}_F"*_eswp.json) + if [[ ${#files[@]} -eq 0 ]]; then + echo "[warn] no compute sweeps found for model '${m}' — skipping its fit." \ + "Run ./run_experiments.sh to produce them." >&2 + continue + fi + case "$m" in + gcn) FIT_COMPUTE_ARGS+=(--compute-gcn "${files[@]}") ;; + edge) FIT_COMPUTE_ARGS+=(--compute-edge "${files[@]}") ;; + gcn_spmm) FIT_COMPUTE_ARGS+=(--compute-gcn-spmm "${files[@]}") ;; + esac +done + +if [[ ${#FIT_COMPUTE_ARGS[@]} -eq 0 ]]; then + echo "[error] no compute sweep files in ${DATA_DIR}/ — T_comp cannot be fit," >&2 + echo " so every prediction would be missing its dominant term." >&2 + exit 1 +fi + step "fit_primitives (intra + inter network)" python -m analysis.fit_primitives \ --pingpong-intra "${DATA_DIR}/pingpong_intra.json" \ --pingpong-inter "${DATA_DIR}/pingpong_inter.json" \ - --compute-gcn "${DATA_DIR}/compute_gcn_vswp.json" "${DATA_DIR}/compute_gcn_eswp.json" \ - --compute-edge "${DATA_DIR}/compute_edge_vswp.json" "${DATA_DIR}/compute_edge_eswp.json" \ + "${FIT_COMPUTE_ARGS[@]}" \ --gather-contiguous "${DATA_DIR}/gather_contiguous.json" \ --gather-clustered "${DATA_DIR}/gather_clustered.json" \ --gather-random "${DATA_DIR}/gather_random.json" \ --output "${DATA_DIR}/fitted_primitives.json" -step "fit_overhead (fit-filter: ${FIT_FILTER})" +step "fit_overhead (fit-filter: ${FIT_FILTER}, nics-per-node: ${NICS_PER_NODE})" python -m analysis.fit_overhead \ --primitives "${DATA_DIR}/fitted_primitives.json" \ --e2e-runs "${E2E_FILES[@]}" \ --fit-filter "${FIT_FILTER}" \ + --nics-per-node "${NICS_PER_NODE}" \ --output "${DATA_DIR}/fitted_overhead.json" step "compute_predictions (fit-filter: ${FIT_FILTER})" From 2b946f043eea539f7e62c1ddc390c4ee85cf0886 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Sat, 1 Aug 2026 00:45:28 -0400 Subject: [PATCH 36/39] Update visualization code and sweep across multiple graph sizes and modea --- .../analysis/fit_overhead.py | 22 ++- .../cost_model_benchmarks/run_experiments.sh | 127 +++++++++++------- .../run_experiments_multinode.sh | 87 +++++++----- .../visualization/plot_ablations.py | 33 ++++- .../visualization/plot_compute.py | 100 +++++++++++--- .../visualization/plot_tipping_point.py | 42 ++++++ 6 files changed, 303 insertions(+), 108 deletions(-) diff --git a/experiments/cost_model_benchmarks/analysis/fit_overhead.py b/experiments/cost_model_benchmarks/analysis/fit_overhead.py index a92edcc..50baabc 100644 --- a/experiments/cost_model_benchmarks/analysis/fit_overhead.py +++ b/experiments/cost_model_benchmarks/analysis/fit_overhead.py @@ -111,13 +111,21 @@ def predict_layer_time(run_config: dict, per_rank_stats: dict, # Using "forward" here undercounted T_comp and cost ~30% MAPE; adding # forward+backward together instead would double-count the forward pass. comp_params = primitives.get("compute", {}).get(model_type, {}).get("backward", None) - if comp_params: - # Which columns this evaluates depends on the form fit_compute chose - # (fixed-F legacy, edge-centric, or vertex-centric). compute_forms owns - # both sides so the fit and this evaluation cannot drift apart. - T_comp = compute_forms.evaluate(comp_params, n_total, n_edges_local, F) - else: - T_comp = 0.0 + if comp_params is None: + # Never silently fall back to zero: T_comp is the dominant term in every + # run measured so far (44 ms of a 49 ms layer at K=2), so a missing fit + # would not look like an error, it would look like a fast prediction. + available = sorted(primitives.get("compute", {})) + raise ValueError( + f"No compute fit for model {model_type!r} in fitted_primitives.json " + f"(have: {available or 'none'}), but an end-to-end run uses it. " + f"Re-run fit_primitives.py with the matching --compute-* sweeps " + f"for {model_type!r}." + ) + # Which columns this evaluates depends on the form fit_compute chose + # (fixed-F legacy, edge-centric, or vertex-centric). compute_forms owns + # both sides so the fit and this evaluation cannot drift apart. + T_comp = compute_forms.evaluate(comp_params, n_total, n_edges_local, F) T_comp = max(T_comp, 0.0) # T_intra and T_inter diff --git a/experiments/cost_model_benchmarks/run_experiments.sh b/experiments/cost_model_benchmarks/run_experiments.sh index cb51495..5bac7df 100755 --- a/experiments/cost_model_benchmarks/run_experiments.sh +++ b/experiments/cost_model_benchmarks/run_experiments.sh @@ -10,7 +10,7 @@ # ./run_experiments.sh # # Override any config variable via the environment, e.g.: -# GPU_COUNTS="2 4 8" MODEL=edge ./run_experiments.sh +# GPU_COUNTS="2 4 8" E2E_MODELS="gcn edge" ./run_experiments.sh set -euo pipefail cd "$(dirname "${BASH_SOURCE[0]}")" @@ -23,8 +23,17 @@ shopt -s nullglob # --- Config ------------------------------------------------------------ GPU_COUNTS=(${GPU_COUNTS:-2 4}) # world sizes to sweep via torchrun --nproc_per_node SEED=${SEED:-42} -FEATURE_DIM=${FEATURE_DIM:-128} -MODEL=${MODEL:-gcn} # gcn | edge | gcn_spmm (distributed runs) +FEATURE_DIM=${FEATURE_DIM:-128} # F for the gather microbenchmark (1.4) +# Layer variants and feature dims for the distributed (crossover + e2e) runs. +# Defaults are the paper configuration, so a plain ./run_experiments.sh +# produces exactly what the paper needs with nothing to override. +# +# F is a single value here while COMPUTE_FEATURE_DIMS below sweeps five: the +# F-dependence result is a compute-microbenchmark finding (it yields F*), and +# every extra F multiplies the whole distributed grid. Add to +# E2E_FEATURE_DIMS if you want the assembled model validated at other F too. +E2E_MODELS=(${E2E_MODELS:-gcn edge gcn_spmm}) +E2E_FEATURE_DIMS=(${E2E_FEATURE_DIMS:-128}) # Layer variants and feature dims for the single-GPU compute sweeps (1.3). # Sweeping F is what makes T_comp's F dependence identifiable: at a single F # the F^2 and F terms are collinear, so a compute-bound layer and a @@ -53,18 +62,28 @@ FIT_FILTER=${FIT_FILTER:-"world_size < ${MAX_GPU_COUNT}"} # Fail before spending GPU hours rather than at the plotting step at the end: # plot_compute reads the REF_FEATURE_DIM cell, which only exists if that F was -# actually swept, and the distributed runs use MODEL, which only gets a fit if -# it was among the compute models. +# actually swept, and each model in E2E_MODELS only gets a T_comp fit if it was +# also benchmarked as a compute model. if [[ ! " ${COMPUTE_FEATURE_DIMS[*]} " =~ " ${REF_FEATURE_DIM} " ]]; then echo "[error] REF_FEATURE_DIM=${REF_FEATURE_DIM} is not in COMPUTE_FEATURE_DIMS" \ "(${COMPUTE_FEATURE_DIMS[*]}); the figures would have no data to plot." >&2 exit 1 fi -if [[ ! " ${COMPUTE_MODELS[*]} " =~ " ${MODEL} " ]]; then - echo "[error] MODEL=${MODEL} is not in COMPUTE_MODELS (${COMPUTE_MODELS[*]});" \ - "the distributed runs would have no T_comp fit to predict with." >&2 - exit 1 -fi +for m in "${E2E_MODELS[@]}"; do + if [[ ! " ${COMPUTE_MODELS[*]} " =~ " ${m} " ]]; then + echo "[error] E2E_MODELS contains '${m}', which is not in COMPUTE_MODELS" \ + "(${COMPUTE_MODELS[*]}); its distributed runs would have no T_comp" \ + "fit to predict with." >&2 + exit 1 + fi +done +for f in "${E2E_FEATURE_DIMS[@]}"; do + if [[ ! " ${COMPUTE_FEATURE_DIMS[*]} " =~ " ${f} " ]]; then + echo "[warn] E2E_FEATURE_DIMS contains F=${f}, which is not in" \ + "COMPUTE_FEATURE_DIMS (${COMPUTE_FEATURE_DIMS[*]}). T_comp will be" \ + "extrapolated in F rather than interpolated there." >&2 + fi +done mkdir -p "$DATA_DIR" "$FIG_DIR" @@ -74,10 +93,10 @@ step() { echo; echo "=== $* ==="; } # 1. Single-GPU primitive microbenchmarks (1.3 compute, 1.4 gather) — no # torchrun, GPU-count independent. # -# fit_primitives.py fits GCN and edge-conditioned compute costs -# independently (--compute-gcn / --compute-edge), regardless of which -# single MODEL the distributed crossover/end-to-end sweep below uses — -# so both model types are always benchmarked here. +# fit_primitives.py fits each layer variant's compute cost independently +# (--compute-gcn / --compute-edge / --compute-gcn-spmm), so every model in +# COMPUTE_MODELS is benchmarked here regardless of which subset the +# distributed crossover/end-to-end sweep below exercises. # ------------------------------------------------------------------------- # Both sweep directions are required per (model, F) cell: fit_compute needs a # vertex sweep (|E| fixed) and an edge sweep (|V| fixed) to identify the |V| @@ -137,33 +156,41 @@ torchrun --nnodes 1 --nproc_per_node 2 -m benchmarks.bench_pingpong \ # 3. Distributed benchmarks (2.1 end-to-end, 2.2 crossover), swept over # GPU_COUNTS via torchrun. # ------------------------------------------------------------------------- -for K in "${GPU_COUNTS[@]}"; do - step "2.2 crossover -- K=${K} GPUs" - torchrun --nnodes 1 --nproc_per_node "${K}" -m benchmarks.bench_crossover \ - --graph "${GRAPH}" --graph-sizes "${CROSSOVER_SIZES}" \ - --avg-degree 20 --feature-dim "${FEATURE_DIM}" --model "${MODEL}" \ - --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ - --output "${DATA_DIR}/crossover_K${K}.json" --seed "${SEED}" - - for N in "${E2E_VERTEX_SIZES[@]}"; do - step "2.1 end-to-end -- K=${K} GPUs, N=${N}" - torchrun --nnodes 1 --nproc_per_node "${K}" -m benchmarks.bench_end_to_end \ - --graph "${GRAPH}" --num-vertices "${N}" --avg-degree 20 \ - --feature-dim "${FEATURE_DIM}" --model "${MODEL}" \ - --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ - --output "${DATA_DIR}/e2e_K${K}_N${N}.json" --seed "${SEED}" - done -done - -# Explicit list of exactly the e2e files this script just produced — not a -# e2e_*.json wildcard glob, which would also match any stray file sharing -# that prefix (e.g. a leftover data/e2e_K8_F128_er_bal.json from manually -# running a docstring example) and silently pool it into the fit/held-out -# split as if it were part of this run's matrix. +# Output filenames carry every axis that varies (K, model, F, N). They used to +# be e2e_K${K}_N${N}.json, which silently overwrote one model's results with +# the next as soon as more than one model was swept — the two runs differ only +# in fields that were not in the name. +# +# E2E_FILES is accumulated here rather than rebuilt afterwards from the same +# nested loops, so the list of files fed to the fit cannot drift out of sync +# with the list actually produced. It is an explicit list, not an e2e_*.json +# glob, which would also match stray files from ad-hoc runs (e.g. a leftover +# data/e2e_K8_F128_er_bal.json) and pool them into the fit/held-out split. E2E_FILES=() +CROSSOVER_TAGS=() for K in "${GPU_COUNTS[@]}"; do - for N in "${E2E_VERTEX_SIZES[@]}"; do - E2E_FILES+=("${DATA_DIR}/e2e_K${K}_N${N}.json") + for m in "${E2E_MODELS[@]}"; do + for f in "${E2E_FEATURE_DIMS[@]}"; do + tag="K${K}_${m}_F${f}" + + step "2.2 crossover -- K=${K} GPUs, model=${m}, F=${f}" + torchrun --nnodes 1 --nproc_per_node "${K}" -m benchmarks.bench_crossover \ + --graph "${GRAPH}" --graph-sizes "${CROSSOVER_SIZES}" \ + --avg-degree 20 --feature-dim "${f}" --model "${m}" \ + --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/crossover_${tag}.json" --seed "${SEED}" + CROSSOVER_TAGS+=("${tag}") + + for N in "${E2E_VERTEX_SIZES[@]}"; do + step "2.1 end-to-end -- K=${K} GPUs, model=${m}, F=${f}, N=${N}" + torchrun --nnodes 1 --nproc_per_node "${K}" -m benchmarks.bench_end_to_end \ + --graph "${GRAPH}" --num-vertices "${N}" --avg-degree 20 \ + --feature-dim "${f}" --model "${m}" \ + --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/e2e_${tag}_N${N}.json" --seed "${SEED}" + E2E_FILES+=("${DATA_DIR}/e2e_${tag}_N${N}.json") + done + done done done @@ -207,11 +234,17 @@ step "visualization" # feature dim. NOTE: there is no figure for the F sweep itself yet — the # fitted F* values are printed by fit_primitives and stored in # fitted_primitives.json, but nothing plots T vs F. -python -m visualization.plot_compute \ - --gcn-vertex "${DATA_DIR}/compute_gcn_F${REF_FEATURE_DIM}_vswp.json" \ - --gcn-edge "${DATA_DIR}/compute_gcn_F${REF_FEATURE_DIM}_eswp.json" \ - --edge-vertex "${DATA_DIR}/compute_edge_F${REF_FEATURE_DIM}_vswp.json" \ - --edge-edge "${DATA_DIR}/compute_edge_F${REF_FEATURE_DIM}_eswp.json" \ +PLOT_COMPUTE_ARGS=() +for m in "${COMPUTE_MODELS[@]}"; do + case "$m" in + gcn) vflag=--gcn-vertex; eflag=--gcn-edge ;; + edge) vflag=--edge-vertex; eflag=--edge-edge ;; + gcn_spmm) vflag=--spmm-vertex; eflag=--spmm-edge ;; + esac + PLOT_COMPUTE_ARGS+=("${vflag}" "${DATA_DIR}/compute_${m}_F${REF_FEATURE_DIM}_vswp.json") + PLOT_COMPUTE_ARGS+=("${eflag}" "${DATA_DIR}/compute_${m}_F${REF_FEATURE_DIM}_eswp.json") +done +python -m visualization.plot_compute "${PLOT_COMPUTE_ARGS[@]}" \ --primitives "${DATA_DIR}/fitted_primitives.json" --output "${FIG_DIR}/compute" # bench_gather records BOTH gather_trials_seconds and @@ -231,9 +264,9 @@ python -m visualization.plot_pingpong \ --intra "${DATA_DIR}/pingpong_intra.json" --primitives "${DATA_DIR}/fitted_primitives.json" \ --output "${FIG_DIR}/pingpong" -for K in "${GPU_COUNTS[@]}"; do - python -m visualization.plot_crossover --input "${DATA_DIR}/crossover_K${K}.json" \ - --output "${FIG_DIR}/crossover_K${K}.png" +for tag in "${CROSSOVER_TAGS[@]}"; do + python -m visualization.plot_crossover --input "${DATA_DIR}/crossover_${tag}.json" \ + --output "${FIG_DIR}/crossover_${tag}.png" done python -m visualization.plot_validation --predictions "${DATA_DIR}/predictions.json" \ diff --git a/experiments/cost_model_benchmarks/run_experiments_multinode.sh b/experiments/cost_model_benchmarks/run_experiments_multinode.sh index a845324..068b2d6 100755 --- a/experiments/cost_model_benchmarks/run_experiments_multinode.sh +++ b/experiments/cost_model_benchmarks/run_experiments_multinode.sh @@ -55,8 +55,17 @@ GPUS_PER_NODE=${GPUS_PER_NODE:-4} # under-predicts every multi-node run by 15-31%. NICS_PER_NODE=${NICS_PER_NODE:-1} SEED=${SEED:-42} -FEATURE_DIM=${FEATURE_DIM:-128} -MODEL=${MODEL:-gcn} +# Layer variants and feature dims for the distributed runs. Defaults are the +# paper configuration; keep these in sync with run_experiments.sh, since this +# script pools that script's K=2/4 files with its own K=${K} files for the fit. +# +# Note each launch() is a separately scheduled job, so the job count is +# |E2E_MODELS| x |E2E_FEATURE_DIMS| x (1 + |E2E_VERTEX_SIZES|). Adding a +# second feature dim doubles the queue time, which is why F stays single here +# while the compute microbenchmarks (single-GPU, in run_experiments.sh) sweep +# five values. +E2E_MODELS=(${E2E_MODELS:-gcn edge gcn_spmm}) +E2E_FEATURE_DIMS=(${E2E_FEATURE_DIMS:-128}) GRAPH=${GRAPH:-erdos_renyi} PARTITIONER=${PARTITIONER:-balanced} WARMUP=${WARMUP:-10} @@ -157,23 +166,34 @@ launch 2 2 benchmarks/bench_concurrency.py \ --output "${DATA_DIR}/concurrency.json" --seed "${SEED}" expect_output "${DATA_DIR}/concurrency.json" -# 2.2 crossover and 2.1 end-to-end at full scale across all nodes. -step "2.2 crossover -- K=${K} GPUs (${NNODES} nodes)" -launch "${NNODES}" "${GPUS_PER_NODE}" benchmarks/bench_crossover.py \ - --graph "${GRAPH}" --graph-sizes "${CROSSOVER_SIZES}" \ - --avg-degree 20 --feature-dim "${FEATURE_DIM}" --model "${MODEL}" \ - --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ - --output "${DATA_DIR}/crossover_K${K}.json" --seed "${SEED}" -expect_output "${DATA_DIR}/crossover_K${K}.json" - -for N in "${E2E_VERTEX_SIZES[@]}"; do - step "2.1 end-to-end -- K=${K} GPUs, N=${N}" - launch "${NNODES}" "${GPUS_PER_NODE}" benchmarks/bench_end_to_end.py \ - --graph "${GRAPH}" --num-vertices "${N}" --avg-degree 20 \ - --feature-dim "${FEATURE_DIM}" --model "${MODEL}" \ - --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ - --output "${DATA_DIR}/e2e_K${K}_N${N}.json" --seed "${SEED}" - expect_output "${DATA_DIR}/e2e_K${K}_N${N}.json" +# 2.2 crossover and 2.1 end-to-end at full scale across all nodes, swept over +# E2E_MODELS x E2E_FEATURE_DIMS. Filenames carry K, model and F: they used to +# be crossover_K${K}.json / e2e_K${K}_N${N}.json, which silently overwrote one +# model's results with the next as soon as more than one model was swept. +CROSSOVER_TAGS=() +for m in "${E2E_MODELS[@]}"; do + for f in "${E2E_FEATURE_DIMS[@]}"; do + tag="K${K}_${m}_F${f}" + + step "2.2 crossover -- K=${K} GPUs (${NNODES} nodes), model=${m}, F=${f}" + launch "${NNODES}" "${GPUS_PER_NODE}" benchmarks/bench_crossover.py \ + --graph "${GRAPH}" --graph-sizes "${CROSSOVER_SIZES}" \ + --avg-degree 20 --feature-dim "${f}" --model "${m}" \ + --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/crossover_${tag}.json" --seed "${SEED}" + expect_output "${DATA_DIR}/crossover_${tag}.json" + CROSSOVER_TAGS+=("${tag}") + + for N in "${E2E_VERTEX_SIZES[@]}"; do + step "2.1 end-to-end -- K=${K} GPUs, model=${m}, F=${f}, N=${N}" + launch "${NNODES}" "${GPUS_PER_NODE}" benchmarks/bench_end_to_end.py \ + --graph "${GRAPH}" --num-vertices "${N}" --avg-degree 20 \ + --feature-dim "${f}" --model "${m}" \ + --partitioner "${PARTITIONER}" --warmup "${WARMUP}" --trials "${TRIALS}" \ + --output "${DATA_DIR}/e2e_${tag}_N${N}.json" --seed "${SEED}" + expect_output "${DATA_DIR}/e2e_${tag}_N${N}.json" + done + done done # ======================================================================= @@ -186,15 +206,20 @@ done # of this run matrix. E2E_FILES=() for k in "${SINGLE_NODE_GPU_COUNTS[@]}" "${K}"; do - for N in "${E2E_VERTEX_SIZES[@]}"; do - f="${DATA_DIR}/e2e_K${k}_N${N}.json" - if [[ -f "$f" ]]; then - E2E_FILES+=("$f") - elif [[ "$k" -eq "$K" ]]; then - echo "[warn] missing ${f} — this run should have produced it; check the K=${K} step above for errors" - else - echo "[warn] missing ${f} — run ./run_experiments.sh first to produce the K=${k} points" - fi + for m in "${E2E_MODELS[@]}"; do + for fd in "${E2E_FEATURE_DIMS[@]}"; do + for N in "${E2E_VERTEX_SIZES[@]}"; do + f="${DATA_DIR}/e2e_K${k}_${m}_F${fd}_N${N}.json" + if [[ -f "$f" ]]; then + E2E_FILES+=("$f") + elif [[ "$k" -eq "$K" ]]; then + echo "[warn] missing ${f} — this run should have produced it; check the K=${K} step above for errors" + else + echo "[warn] missing ${f} — run ./run_experiments.sh with matching" \ + "E2E_MODELS/E2E_FEATURE_DIMS to produce the K=${k} points" + fi + done + done done done @@ -270,8 +295,10 @@ python -m visualization.plot_pingpong \ --primitives "${DATA_DIR}/fitted_primitives.json" \ --output "${FIG_DIR}/pingpong" -python -m visualization.plot_crossover --input "${DATA_DIR}/crossover_K${K}.json" \ - --output "${FIG_DIR}/crossover_K${K}.png" +for tag in "${CROSSOVER_TAGS[@]}"; do + python -m visualization.plot_crossover --input "${DATA_DIR}/crossover_${tag}.json" \ + --output "${FIG_DIR}/crossover_${tag}.png" +done python -m visualization.plot_validation --predictions "${DATA_DIR}/predictions.json" \ --color-by world_size --output "${FIG_DIR}/validation" diff --git a/experiments/cost_model_benchmarks/visualization/plot_ablations.py b/experiments/cost_model_benchmarks/visualization/plot_ablations.py index a3f1d5d..0544499 100644 --- a/experiments/cost_model_benchmarks/visualization/plot_ablations.py +++ b/experiments/cost_model_benchmarks/visualization/plot_ablations.py @@ -76,8 +76,16 @@ def main(): ) def flat_comm_time(c_intra_bytes: float, c_inter_bytes: float) -> float: - """Flat (single-tier) model: no intra/inter distinction, no overlap — - all halo bytes go over one fitted (B, t_L) pair, sequentially.""" + """Flat (single-tier) model: all halo bytes over one fitted (B, t_L). + + NOTE: this baseline is naive in *two* ways relative to the full model — + it has no intra/inter tier split, and it applies no per-NIC contention + correction (B is not divided by ranks-per-NIC). So the gap this panel + shows is the combined value of the hierarchy *and* the contention + term, not the hierarchy alone. Do not describe it as isolating the + hierarchy; to separate them you would need a third curve with tiers + but no contention. + """ total_bytes = c_intra_bytes + c_inter_bytes if total_bytes <= 0: return 0.0 @@ -96,6 +104,18 @@ def flat_comm_time(c_intra_bytes: float, c_inter_bytes: float) -> float: density_groups.setdefault(d, []).append(e) densities = sorted(density_groups.keys()) + # Panel (a) needs SBM runs at several inter-densities. With ER-only data + # it renders empty — and the caption below would still assert that "the + # hierarchical model degrades more gracefully as inter-block density + # increases", a claim with nothing behind it. Fail instead of emitting a + # figure whose caption states a result the panel does not show. + if len(densities) < 2: + raise SystemExit( + f"[plot_ablations] Panel (a) needs SBM runs at 2+ inter-densities; " + f"found {len(densities)} ({densities or 'none'}) among " + f"{len(entries)} predictions. Re-run bench_end_to_end.py with " + f"--graph sbm and a swept --sbm-inter-density, or drop this figure." + ) mape_full = [] mape_flat = [] @@ -108,9 +128,12 @@ def flat_comm_time(c_intra_bytes: float, c_inter_bytes: float) -> float: T_pred_full = e["predicted_seconds"] full_errs.append(relative_error(T_pred_full, T_meas)) - # Flat model: swap the hierarchical (intra/inter, overlapped via - # max) comm term for one predicted by a genuinely fit single-tier - # (B, t_L) model over total halo bytes — not a heuristic ratio. + # Flat model: swap the hierarchical comm term for one predicted by + # a genuinely fit single-tier (B, t_L) model over total halo bytes + # — not a heuristic ratio. Reading T_comm_seconds from the stored + # breakdown keeps this independent of how the full model combines + # its tiers (currently additive; see README's open question on + # sum vs max), so this panel needs no change if that flips. bd = e.get("breakdown", {}) T_comm_hier = bd.get("T_comm_seconds", 0.0) c_intra = e["partition_stats"].get("c_intra_bytes", 0) diff --git a/experiments/cost_model_benchmarks/visualization/plot_compute.py b/experiments/cost_model_benchmarks/visualization/plot_compute.py index 20ef9ce..9a692f2 100644 --- a/experiments/cost_model_benchmarks/visualization/plot_compute.py +++ b/experiments/cost_model_benchmarks/visualization/plot_compute.py @@ -25,6 +25,8 @@ matplotlib.use("Agg") import matplotlib.pyplot as plt +from analysis import compute_forms + COLORS = {"gcn": "#2ca02c", "edge": "#ff7f0e"} plt.rcParams.update( { @@ -40,10 +42,16 @@ def load_compute_file(path: str, timing_key: str = "forward_trials_seconds"): - """Returns lists of (sweep_value, median, q25, q75).""" + """Returns (sweep, rows, feature_dim). + + feature_dim must come from the data: T_comp's fitted form is a function of + F now, so evaluating the fit for the overlay needs the F this file was + measured at. One file is always a single (model, F) cell. + """ with open(path) as f: data = json.load(f) sweep = data["config"]["sweep"] + feature_dim = data["config"]["feature_dim"] rows = [] for meas in data["measurements"]: trials = np.array(meas[timing_key]) @@ -58,24 +66,37 @@ def load_compute_file(path: str, timing_key: str = "forward_trials_seconds"): ) ) rows.sort(key=lambda r: r[0]) - return sweep, rows + return sweep, rows, feature_dim + +def fitted_compute(sweep_vals, fixed_val, sweep, model_type, primitives, feature_dim): + """Evaluate the fitted T_comp curve for the overlay. -def fitted_compute(sweep_vals, fixed_val, sweep, model_type, primitives): + Delegates to analysis.compute_forms rather than unpacking coefficients by + name: fit_compute chooses its design matrix per layer type and per whether + F was swept, so the coefficient *names* vary ("coeff_V"/"coeff_E" only + exist in the fixed-F legacy form; the F-dependent forms use + "coeff_EF2"/"coeff_EF"/"coeff_VF"). Reading them directly here is what + raised KeyError: 'coeff_V' once the F sweep made the form edge_centric. + """ params = primitives.get("compute", {}).get(model_type, {}).get("forward", None) if params is None: return None - a, b, c = params["coeff_V"], params["coeff_E"], params["intercept"] if sweep == "vertices": V_arr = np.array(sweep_vals, dtype=float) E_arr = np.full_like(V_arr, fixed_val) else: E_arr = np.array(sweep_vals, dtype=float) V_arr = np.full_like(E_arr, fixed_val) - return a * V_arr + b * E_arr + c + # evaluate() returns a scalar, so map it over the sweep. + return np.array([ + compute_forms.evaluate(params, v, e, feature_dim) + for v, e in zip(V_arr, E_arr) + ]) -def plot_one_panel(ax, rows, sweep, fixed_val, model_type, primitives, color, title): +def plot_one_panel(ax, rows, sweep, fixed_val, model_type, primitives, color, title, + feature_dim): xvals = [r[0] for r in rows] meds = np.array([r[3] for r in rows]) * 1e3 lo = np.array([r[3] - r[4] for r in rows]) * 1e3 @@ -94,7 +115,7 @@ def plot_one_panel(ax, rows, sweep, fixed_val, model_type, primitives, color, ti label="Measured (IQR)", ) - fit = fitted_compute(xvals, fixed_val, sweep, model_type, primitives) + fit = fitted_compute(xvals, fixed_val, sweep, model_type, primitives, feature_dim) if fit is not None: ax.plot(xvals, fit * 1e3, "--", color=color, linewidth=1.2, label="Fit") @@ -115,6 +136,10 @@ def parse_args(): p.add_argument("--gcn-edge", type=str, default=None) p.add_argument("--edge-vertex", type=str, default=None) p.add_argument("--edge-edge", type=str, default=None) + p.add_argument("--spmm-vertex", type=str, default=None, + help="Vertex sweep for the SpMM GCN layer (--model gcn_spmm)") + p.add_argument("--spmm-edge", type=str, default=None, + help="Edge sweep for the SpMM GCN layer (--model gcn_spmm)") p.add_argument("--primitives", type=str, default=None) p.add_argument("--output", type=str, default="figures/compute") return p.parse_args() @@ -127,18 +152,37 @@ def main(): with open(args.primitives) as f: primitives = json.load(f) - fig, axes = plt.subplots(1, 2, figsize=(7, 3)) - - panel_map = [ - ("gcn", "edges", args.gcn_edge, axes[0], "GCN-like"), - ("edge", "edges", args.edge_edge, axes[1], "Edge-conditioned"), + # Columns = layer variant, rows = sweep direction. The previous version + # built a 1x2 grid from the *edge* sweeps only, leaving --gcn-vertex and + # --edge-vertex parsed but never plotted, while the caption described a + # two-row layout that did not exist. Only columns with at least one file + # are drawn, so passing a subset still produces a sensible figure. + columns = [ + ("gcn", "GCN-like (per-edge)", args.gcn_vertex, args.gcn_edge), + ("edge", "Edge-conditioned", args.edge_vertex, args.edge_edge), + ("gcn_spmm", "GCN (SpMM)", args.spmm_vertex, args.spmm_edge), ] + columns = [c for c in columns if c[2] or c[3]] + if not columns: + raise SystemExit( + "[plot_compute] No input files given. Pass at least one of " + "--gcn-vertex/--gcn-edge/--edge-vertex/--edge-edge/" + "--spmm-vertex/--spmm-edge." + ) - for model_type, sweep_label, path, ax, title in panel_map: + ncols = len(columns) + fig, axes = plt.subplots(2, ncols, figsize=(3.5 * ncols, 6), squeeze=False) + + panel_map = [] + for col, (model_type, title, vpath, epath) in enumerate(columns): + panel_map.append((model_type, vpath, axes[0][col], f"{title} — |V| sweep")) + panel_map.append((model_type, epath, axes[1][col], f"{title} — |E| sweep")) + + for model_type, path, ax, title in panel_map: if path is None: ax.set_visible(False) continue - sweep, rows = load_compute_file(path) + sweep, rows, feature_dim = load_compute_file(path) fixed_val = rows[0][2] if sweep == "vertices" else rows[0][1] # E or V fixed plot_one_panel( ax, @@ -148,7 +192,8 @@ def main(): model_type, primitives, COLORS[model_type], - title, + f"{title} (F={feature_dim})", + feature_dim, ) # fig.suptitle("GNN Compute Primitive: Forward Runtime vs. Graph Size", fontsize=10) @@ -159,12 +204,29 @@ def main(): fig.savefig(str(out) + ".pdf", bbox_inches="tight") fig.savefig(str(out) + ".png", bbox_inches="tight", dpi=300) print(f"[plot_compute] Saved {out}.pdf and {out}.png") + forms = sorted({ + primitives.get("compute", {}).get(m, {}).get("forward", {}).get("form") + for m, _, _, _ in columns + } - {None}) + form_expr = { + "legacy_VE": "$a|V| + b|E| + c$", + "edge_centric": "$a|E|F^2 + b|E|F + c|V|F + d$", + "vertex_centric": "$a|V|F^2 + b|E|F + c|V|F + d$", + } + # Older fitted_primitives.json files predate the "form" key; evaluate() + # treats those as legacy_VE, so say so rather than emitting a bare phrase. + fitted_desc = ( + " / ".join(form_expr.get(f, f) for f in forms) + if forms + else form_expr["legacy_VE"] + ) print( "Caption: Forward runtime of a single GNN layer vs. subgraph size " - "(vertex sweep top row, edge sweep bottom row) for GCN-like (left) " - "and edge-conditioned (right) message functions. " - "Dashed lines: fitted model $T_{\\mathrm{comp}} = a|V| + b|E| + c$. " - "Error bars span IQR over 50+ trials." + "(vertex sweep top row, edge sweep bottom row) for " + + ", ".join(t for _, t, _, _ in columns) + + " message functions. Dashed lines: fitted model " + + fitted_desc + + ". Error bars span IQR over 50+ trials." ) diff --git a/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py b/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py index f5cc0cb..a6c1b02 100644 --- a/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py +++ b/experiments/cost_model_benchmarks/visualization/plot_tipping_point.py @@ -45,6 +45,12 @@ def parse_args(): p.add_argument("--num-epochs", type=int, default=100) p.add_argument("--graph", type=str, default=None, help="Filter to this graph type (optional)") + p.add_argument("--model", type=str, default=None, + help="Restrict to one layer type (gcn | edge | gcn_spmm). " + "Required when predictions.json pools several, since a " + "scaling curve is per-configuration.") + p.add_argument("--num-vertices", type=int, default=None, + help="Restrict to one graph size, for the same reason.") p.add_argument("--output", type=str, default="figures/tipping_point") return p.parse_args() @@ -62,18 +68,54 @@ def main(): entries = [e for e in entries if e["config"].get("graph", "") == args.graph] + if args.model: + entries = [e for e in entries if e["config"].get("model", "") == args.model] + if args.num_vertices: + entries = [e for e in entries + if e["config"].get("num_vertices", 0) == args.num_vertices] + # Group by world_size ws_to_meas = {} ws_to_pred = {} + ws_to_cfg = {} for e in entries: K = e["config"].get("world_size", 1) ws_to_meas.setdefault(K, []).append(e["measured_median_seconds"]) ws_to_pred.setdefault(K, []).append(e["predicted_seconds"]) + c = e["config"] + ws_to_cfg.setdefault(K, []).append( + (c.get("model"), c.get("num_vertices"), c.get("graph"), + c.get("feature_dim"), c.get("partitioner")) + ) if not ws_to_meas: print("[plot_tipping_point] No data found. Exiting.") return + # A scaling curve is only meaningful for ONE configuration measured at + # several K. Previously this grouped by world_size alone and took a median + # over whatever else was in predictions.json — which silently averaged + # across N (1 ms and 49 ms points at the same K), and, once the drivers + # began sweeping E2E_MODELS, across layer types too. Refuse instead: the + # median of gcn and gcn_spmm is not a runtime of anything. + ambiguous = {K: sorted(set(cfgs)) for K, cfgs in ws_to_cfg.items() + if len(set(cfgs)) > 1} + if ambiguous: + K0 = sorted(ambiguous)[0] + varying = [ + name for i, name in enumerate( + ["model", "num_vertices", "graph", "feature_dim", "partitioner"]) + if len({c[i] for c in ambiguous[K0]}) > 1 + ] + raise ValueError( + f"predictions.json holds {len(ambiguous[K0])} different " + f"configurations at world_size={K0}, differing in " + f"{', '.join(varying)}. A tipping-point curve needs a single " + f"configuration swept over K. Narrow it with --model / " + f"--num-vertices / --graph. Configurations at K={K0}: " + f"{ambiguous[K0]}" + ) + K_vals = sorted(ws_to_meas.keys()) if K_vals[0] != 1: raise ValueError( From adb3a8c8f7ef314b196a3c949544d833c31cca27 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Sat, 1 Aug 2026 01:00:30 -0400 Subject: [PATCH 37/39] Update crossover code to include gcn_cpmm --- .../benchmarks/bench_crossover.py | 90 ++++++++++++++----- 1 file changed, 68 insertions(+), 22 deletions(-) diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py b/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py index 04e278b..d63b488 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py @@ -72,7 +72,12 @@ get_ranks_per_node, intra_inter_halo, ) -from benchmarks.nn_layer_common import GCNLayer, EdgeConditionedLayer +from benchmarks.nn_layer_common import ( + GCNLayer, + EdgeConditionedLayer, + GCNSpMMLayer, + create_sparse_adj, +) from DGraph.distributed import ( HaloExchange, @@ -104,7 +109,7 @@ def parse_args(): help="Fraction of inter-block edges for SBM graphs", ) p.add_argument("--feature-dim", type=int, default=128) - p.add_argument("--model", choices=["gcn", "edge"], default="gcn") + p.add_argument("--model", choices=["gcn", "edge", "gcn_spmm"], default="gcn") p.add_argument( "--partitioner", choices=["random", "balanced", "metis"], default="balanced" ) @@ -135,26 +140,47 @@ def _gen_graph(graph_type, num_vertices, avg_degree, sbm_inter_density, seed): def _build_single_gpu_tensors(num_vertices, edges_np, F, model, device): - """Allocate full-graph tensors on *device*. May raise cuda.OutOfMemoryError.""" - # GCNLayer / EdgeConditionedLayer expect edge_index as [2, E] + """Allocate full-graph tensors on *device*. May raise cuda.OutOfMemoryError. + + Returns (x, aux, layer, edge_attr) where ``aux`` is whatever the layer's + second forward argument is: a ``[2, E]`` edge index for GCNLayer / + EdgeConditionedLayer, or a prebuilt sparse ``[V, V]`` adjacency for + GCNSpMMLayer (which does not take an edge index at all — see + nn_layer_common.GCNSpMMLayer.forward). Built once here, outside the timed + region, matching bench_compute.py's build_sparse_adj convention. + """ edge_t = torch.from_numpy(edges_np.T.copy()).long().to(device) x = torch.randn(num_vertices, F, device=device, requires_grad=True) if model == "gcn": layer = GCNLayer(F).to(device) edge_attr = None + aux = edge_t + elif model == "gcn_spmm": + layer = GCNSpMMLayer(F).to(device) + edge_attr = None + aux = create_sparse_adj(edge_t, num_vertices, num_vertices, device) else: layer = EdgeConditionedLayer(F).to(device) edge_attr = torch.randn(edges_np.shape[0], F, device=device) + aux = edge_t layer.train() - return x, edge_t, layer, edge_attr + return x, aux, layer, edge_attr -def _single_gpu_fn(x, edge_t, layer, edge_attr): +def _single_gpu_fn(x, aux, layer, edge_attr, model): """Return a zero-argument closure for single-GPU forward+backward.""" - if edge_attr is None: + if model == "gcn_spmm": + + def fn(): + out = layer(x, aux) + out.sum().backward() + if x.grad is not None: + x.grad.zero_() + + elif edge_attr is None: def fn(): - out = layer(x, edge_t) + out = layer(x, aux) out.sum().backward() if x.grad is not None: x.grad.zero_() @@ -162,7 +188,7 @@ def fn(): else: def fn(): - out = layer(x, edge_t, edge_attr) + out = layer(x, aux, edge_attr) out.sum().backward() if x.grad is not None: x.grad.zero_() @@ -174,11 +200,30 @@ def _multi_gpu_fn(x_local, comm_pattern, layer, halo_exchange, edge_attr, model) """Return a zero-argument closure for distributed forward+backward. ``comm_pattern.local_edge_list`` has shape ``[E, 2]``; we transpose it - once to ``[2, E]`` as required by GCNLayer / EdgeConditionedLayer. + once to ``[2, E]`` as required by GCNLayer / EdgeConditionedLayer. For + GCNSpMMLayer we instead build the rectangular + ``[n_local, n_local + n_halo]`` sparse adjacency once here (outside the + timed closure, same as the single-GPU path) since that layer takes the + aggregation matrix directly rather than an edge index. """ edge_index = comm_pattern.local_edge_list.T.contiguous() # [2, E_local] - if model == "gcn": + if model == "gcn_spmm": + n_local = comm_pattern.num_local_vertices + n_halo = comm_pattern.num_halo_vertices + adj = create_sparse_adj( + edge_index, n_local, n_local + n_halo, x_local.device + ) + + def fn(): + recv_buf = halo_exchange(x_local, comm_pattern) + x_aug = torch.cat([x_local, recv_buf], dim=0) + out = layer(x_aug, adj) + out.sum().backward() + if x_local.grad is not None: + x_local.grad.zero_() + + elif model == "gcn": def fn(): recv_buf = halo_exchange(x_local, comm_pattern) @@ -228,13 +273,13 @@ def run_no_dist(args, graph_sizes, F): times_single = None oom = False try: - x, edge_t, layer, edge_attr = _build_single_gpu_tensors( + x, aux, layer, edge_attr = _build_single_gpu_tensors( num_vertices, edges_np, F, args.model, device ) - fn = _single_gpu_fn(x, edge_t, layer, edge_attr) + fn = _single_gpu_fn(x, aux, layer, edge_attr, args.model) times_single = cuda_timed(fn, warmup=args.warmup, trials=args.trials) # Free before next iteration - del x, edge_t, layer + del x, aux, layer if edge_attr is not None: del edge_attr torch.cuda.empty_cache() @@ -330,11 +375,12 @@ def run_distributed(args, graph_sizes, F, rank, world_size, local_rank): if args.model == "edge" else None ) - layer_dist = ( - GCNLayer(F).to(device) - if args.model == "gcn" - else EdgeConditionedLayer(F).to(device) - ) + if args.model == "gcn": + layer_dist = GCNLayer(F).to(device) + elif args.model == "gcn_spmm": + layer_dist = GCNSpMMLayer(F).to(device) + else: + layer_dist = EdgeConditionedLayer(F).to(device) layer_dist.train() fn_multi = _multi_gpu_fn( @@ -396,14 +442,14 @@ def run_distributed(args, graph_sizes, F, rank, world_size, local_rank): graph_seed, ) try: - x_s, edge_t_s, layer_s, edge_attr_s = _build_single_gpu_tensors( + x_s, aux_s, layer_s, edge_attr_s = _build_single_gpu_tensors( num_vertices, edges_s, F, args.model, device ) - fn_single = _single_gpu_fn(x_s, edge_t_s, layer_s, edge_attr_s) + fn_single = _single_gpu_fn(x_s, aux_s, layer_s, edge_attr_s, args.model) times_single = cuda_timed( fn_single, warmup=args.warmup, trials=args.trials ) - del x_s, edge_t_s, layer_s, fn_single + del x_s, aux_s, layer_s, fn_single if edge_attr_s is not None: del edge_attr_s torch.cuda.empty_cache() From 888b27f9864f6ef729ad235db005e936f8eef8f7 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Sat, 1 Aug 2026 01:06:06 -0400 Subject: [PATCH 38/39] Updated end to end test to include gcn spmm --- .../benchmarks/bench_end_to_end.py | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py index 47d5375..4d3cbe9 100644 --- a/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py +++ b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py @@ -60,7 +60,12 @@ get_ranks_per_node, intra_inter_halo, ) -from benchmarks.nn_layer_common import GCNLayer, EdgeConditionedLayer +from benchmarks.nn_layer_common import ( + GCNLayer, + EdgeConditionedLayer, + GCNSpMMLayer, + create_sparse_adj, +) from DGraph.distributed import HaloExchange, CommunicationPattern, build_communication_pattern from DGraph import Communicator @@ -82,7 +87,7 @@ def parse_args(): help="Fraction of inter-block edges for SBM graphs", ) p.add_argument("--feature-dim", type=int, default=128) - p.add_argument("--model", choices=["gcn", "edge"], default="gcn") + p.add_argument("--model", choices=["gcn", "edge", "gcn_spmm"], default="gcn") p.add_argument( "--partitioner", choices=["random", "balanced", "metis"], default="balanced" ) @@ -130,6 +135,8 @@ def main(): # --- Model --- if args.model == "gcn": layer = GCNLayer(F).to(device) + elif args.model == "gcn_spmm": + layer = GCNSpMMLayer(F).to(device) else: layer = EdgeConditionedLayer(F).to(device) layer.train() @@ -141,6 +148,14 @@ def main(): if args.model == "edge" else None ) + # GCNSpMMLayer takes the prebuilt rectangular [n_local, n_local+n_halo] + # aggregation matrix instead of an edge index (see nn_layer_common) — + # built once here, outside the timed closure, matching bench_crossover.py. + adj = ( + create_sparse_adj(edge_index, n_local, n_local + n_halo, device) + if args.model == "gcn_spmm" + else None + ) comm = Communicator(backend="nccl") halo_exchange = HaloExchange(comm=comm) @@ -154,6 +169,8 @@ def one_layer(): # Message passing if args.model == "gcn": out = layer(x_aug, edge_index) + elif args.model == "gcn_spmm": + out = layer(x_aug, adj) else: out = layer(x_aug, edge_index, edge_attr) # Backward From 4f37fc392c37b28e99e1b2a39a3e80dcf41b9642 Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Sat, 1 Aug 2026 01:15:42 -0400 Subject: [PATCH 39/39] Fix plot comute --- experiments/cost_model_benchmarks/visualization/plot_compute.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experiments/cost_model_benchmarks/visualization/plot_compute.py b/experiments/cost_model_benchmarks/visualization/plot_compute.py index 9a692f2..8c06909 100644 --- a/experiments/cost_model_benchmarks/visualization/plot_compute.py +++ b/experiments/cost_model_benchmarks/visualization/plot_compute.py @@ -27,7 +27,7 @@ from analysis import compute_forms -COLORS = {"gcn": "#2ca02c", "edge": "#ff7f0e"} +COLORS = {"gcn": "#2ca02c", "edge": "#ff7f0e", "gcn_spmm":"#613cd1"} plt.rcParams.update( { "font.size": 9,