diff --git a/DGraph/distributed/commInfo.py b/DGraph/distributed/commInfo.py index 0d5bd7e..0abd259 100644 --- a/DGraph/distributed/commInfo.py +++ b/DGraph/distributed/commInfo.py @@ -47,46 +47,102 @@ 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] - 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, + 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. + + 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 (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) - num_global = partitioning.size(0) + num_dst_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 - 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. Local + # src vertices are already mapped to [0, num_local) via g2l_dst above. + 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: 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_src_local, num_src_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 @@ -169,25 +225,53 @@ 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, + src_local_vertices_global=src_local_verts, ) 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/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 3d592e7..8289b76 100644 --- a/experiments/GraphCast/data_utils/graphcast_graph.py +++ b/experiments/GraphCast/data_utils/graphcast_graph.py @@ -57,29 +57,93 @@ 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 + + # 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 +) -> 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: @@ -90,7 +154,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,56 +168,51 @@ 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, - neighbor_partitioning=grid_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( global_edge_list=mesh_edges, partitioning=mesh_part, - neighbor_partitioning=mesh_part, rank=rank, world_size=world_size, ) # --- 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, - neighbor_partitioning=mesh_part, + src_partitioning=mesh_part, rank=rank, world_size=world_size, ) @@ -220,6 +285,58 @@ 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( + 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. + + 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. + + All indices/placements here are in the *original* (pre rank-sort) mesh + and grid vertex numbering. + """ + 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 --- + m2g_ranks = mesh_vertex_rank_placement[mesh2grid_mesh_src_indices.long()] + 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 + grid_partitioning = votes.argmax(dim=1) + + 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.""" @@ -238,17 +355,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] @@ -278,30 +391,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. - return meshtogrid_edge_placement + 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. - def get_grid2mesh_graph(self, mesh_graph_dict: dict): + 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"] - 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) @@ -310,60 +433,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] - grid_vertex_rank_placement = torch.zeros_like(lat_lon_grid_flat) - 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 + continuous_grid_mapping, renumbered_grid, inverse_grid_renumbering = ( + self._renumber_by_rank(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) + # 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_graph( + 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) - - 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] + if mesh2grid_raw is None: + mesh2grid_raw = self._get_raw_mesh2grid_edges() + edge_features, src_mesh_indices, dst_grid_indices = mesh2grid_raw - 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 = { - "node_features": torch.tensor([]), "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 @@ -396,42 +517,89 @@ 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"] + 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_graph( - grid_vertex_rank_placement, renumbered_vertices, renumbered_grid + mesh2grid_graph = self.get_mesh2grid_edges( + inverse_vertex_renumbering, + inverse_grid_renumbering, + mesh2grid_raw=mesh2grid_raw, ) - return DistributedGraphCastGraph( - rank=self.rank, + topology = GraphCastTopology( + rank=self.local_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_rank_placement=mesh_graph["node_rank_placement"], - mesh_graph_edge_rank_placement=mesh_graph["edge_rank_placement"], + 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"], - 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"], - grid2mesh_graph_node_features=grid2mesh_graph["node_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"], ) + + 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 + + # 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] + + 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=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=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/data_utils/preprocess.py b/experiments/GraphCast/data_utils/preprocess.py index 886452f..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.ones(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/dataset.py b/experiments/GraphCast/dataset.py index 3e4eb08..5e929f5 100644 --- a/experiments/GraphCast/dataset.py +++ b/experiments/GraphCast/dataset.py @@ -16,9 +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.utils import padded_size -from torch.nn.functional import pad +from data_utils.graphcast_graph import ( + DistributedGraphCastGraphGenerator, + move_graphcast_graph_to_device, +) class SyntheticWeatherDataset(Dataset): @@ -94,6 +95,21 @@ 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. + # 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 + ) print(f"Generated static graph in {time.time() - start_time:.2f} seconds.") self.extra_args: Dict[str, Any] = kwargs @@ -204,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), @@ -261,20 +270,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/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..4b50a78 --- /dev/null +++ b/experiments/GraphCast/gen_mesh_partitions.py @@ -0,0 +1,98 @@ +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 _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", + 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 = _parse_int_list(mesh_levels) + sizes = _parse_int_list(world_sizes) + + 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..f6b3639 --- /dev/null +++ b/experiments/GraphCast/inference_benchmark.py @@ -0,0 +1,453 @@ +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 ( # 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 +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 + ) + # 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"] + + 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/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..1f61574 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,27 +119,25 @@ def __init__(self, cfg: Config, comm, *args, **kwargs) -> None: comm: Communicator object """ super().__init__(*args, **kwargs) + hidden_dim = cfg.model.hidden_dim + + self.exchanger = HaloExchange(comm) - 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( + 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.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.mesh_node_mlp = MeshNodeBlock( + input_node_dim=hidden_dim, + input_edge_dim=hidden_dim, + output_node_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 + input_dim=hidden_dim, output_dim=hidden_dim, hidden_dim=hidden_dim ) def forward( @@ -144,27 +145,39 @@ def forward( 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 [neighbor=grid/halo, central=mesh] (col1 is + # always the locally-owned aggregation target per CommunicationPattern's contract) + edge_index = comm_pattern.local_edge_list + 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"): + 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_grid, + dst_node_features=mesh_node_features, + 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=mesh_node_features, + 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 +192,74 @@ 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 [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 + 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( + 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 +276,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 +299,48 @@ 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 [neighbor=mesh/halo, central=grid] (col1 is + # always the locally-owned aggregation target per CommunicationPattern's contract) + edge_index = comm_pattern.local_edge_list + 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"): + # 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,13 +355,14 @@ 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) 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( @@ -340,26 +376,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 +398,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 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/preprocess.py b/experiments/GraphCast/preprocess.py deleted file mode 100644 index e860da0..0000000 --- a/experiments/GraphCast/preprocess.py +++ /dev/null @@ -1,101 +0,0 @@ -import metis -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): - metis_graph = metis.networkx_to_metis(G) - (edgecuts, node_rank_placement) = metis.part_graph(metis_graph, nparts=num_ranks) - # 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) 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_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() 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 diff --git a/experiments/OGB/GCN.py b/experiments/OGB/GCN.py index 2b3cab4..a15ecb4 100644 --- a/experiments/OGB/GCN.py +++ b/experiments/OGB/GCN.py @@ -17,14 +17,165 @@ 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 + + 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 +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. + """ + # 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) + 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), dtype=torch.float32, device=device) + + # Create the sparse COO tensor + adj_sparse = torch.sparse_coo_tensor( + indices, + values, + size=(num_local_nodes, num_local_nodes + num_halo_nodes), + device=device, + ) + + # 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 + + +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__() + 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, + 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, 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)) + self.halo_exchanger = halo_exchanger + + 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) + + with TimingReport(f"process-{counter}"): + 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__() 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..9f6a526 --- /dev/null +++ b/experiments/OGB/benchmark_OGB_end_to_end.py @@ -0,0 +1,158 @@ +# 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(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() + 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__)) + dist_dataset = DGraphOGBDataset( + dname=f"ogbn-{dataset}", + comm=comm, + root_dir=f"{cur_dir}/data", + ) + + benchmark_ogb_end_to_end( + dataset=dist_dataset, + comm=comm, + lr=0.01, + epochs=100, + log_prefix=f"ogbn_{dataset}-{world_size}", + hidden_dims=128, + num_classes=num_classes, + num_layers=3, + device="cuda", + rank=rank, + local_rank=local_rank, + cur_dir=cur_dir, + ) + + +if __name__ == "__main__": + fire.Fire(main) diff --git a/experiments/OGB/main.py b/experiments/OGB/main.py index 9ad97e4..381ae86 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,19 @@ 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: + + 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 +109,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: @@ -144,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) @@ -339,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/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() 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") 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..6211f94 --- /dev/null +++ b/experiments/cost_model_benchmarks/analysis/fit_primitives.py @@ -0,0 +1,343 @@ +"""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 +from scipy.optimize import curve_fit +from scipy.special import expit + + +# --------------------------------------------------------------------------- +# 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 + max(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) + + 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. 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) + + # 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 + ] + + 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 + 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": 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, + } + + +# --------------------------------------------------------------------------- +# 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_crossover.py b/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py new file mode 100644 index 0000000..c5c2922 --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/bench_crossover.py @@ -0,0 +1,557 @@ +"""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 1 --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 gc + +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="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) + 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") + + 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: + 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, + ) + + # 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, + 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, + 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, fn_single + 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() 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..fa56f5d --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/bench_end_to_end.py @@ -0,0 +1,235 @@ +"""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, +) +from benchmarks.graph_data_common import ( + gen_erdos_renyi, + gen_sbm, + partition_balanced, + partition_metis, + partition_random, + build_local_comm_pattern, +) +from benchmarks.nn_layer_common import GCNLayer, EdgeConditionedLayer + +from DGraph.distributed import HaloExchange, CommunicationPattern +from DGraph import Communicator + +# =========================================================================== +# 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 + ) + + comm = Communicator(backend="nccl") + halo_exchange = HaloExchange(comm=comm) + + # --- Timed forward + backward --- + def one_layer(): + # Forward halo exchange + recv_buf = halo_exchange(x_local, comm_pattern=pattern) + # 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/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..63a6af3 --- /dev/null +++ b/experiments/cost_model_benchmarks/benchmarks/nn_layer_common.py @@ -0,0 +1,42 @@ +import torch +import torch.nn as nn +import torch.distributed as dist + +# =========================================================================== +# 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/run_local_compute_tests.sh b/experiments/cost_model_benchmarks/run_local_compute_tests.sh new file mode 100644 index 0000000..726bfe0 --- /dev/null +++ b/experiments/cost_model_benchmarks/run_local_compute_tests.sh @@ -0,0 +1,12 @@ +python -m benchmarks.bench_compute --model edge --sweep edges --min 1000 --max 1000000 --steps 10 \ + --fixed-value 200000 --feature-dim 512 --warmup 5 --trials 20 --output data/compute_edge_eswp_test.json --seed 4 + +python -m benchmarks.bench_compute --model gcn --sweep edges --min 1000 --max 1000000 --steps 10 \ + --fixed-value 200000 --feature-dim 512 --warmup 5 --trials 20 --output data/compute_gcn_eswp_test.json --seed 4 + +python -m analysis.fit_primitives --compute-gcn data/compute_gcn_*swp_test.json \ + --compute-edge data/compute_edge_*swp_test.json --output data/fitted_primitives.json + +python -m visualization.plot_compute --gcn-vertex data/compute_gcn_vswp_test.json \ + --gcn-edge data/compute_gcn_eswp_test.json --edge-vertex data/compute_edge_vswp_test.json \ + --edge-edge data/compute_edge_eswp_test.json --primitives data/fitted_primitives.json --output figures/compute diff --git a/experiments/cost_model_benchmarks/visualization/__init__.py b/experiments/cost_model_benchmarks/visualization/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/experiments/cost_model_benchmarks/visualization/plot_ablations.py b/experiments/cost_model_benchmarks/visualization/plot_ablations.py new file mode 100644 index 0000000..a23eb4f --- /dev/null +++ b/experiments/cost_model_benchmarks/visualization/plot_ablations.py @@ -0,0 +1,160 @@ +"""Visualization — Ablation Studies. + +Two-panel figure: + (a) Hierarchical model (intra+inter separate) vs. flat model (single B, t_L) + on a topology sweep (varying SBM inter-block density). + (b) Full model vs. model without the T_buffer_copy term. + +Reads ``data/predictions.json`` (which has the full breakdown per entry). + +Usage:: + + python -m visualization.plot_ablations \\ + --predictions data/predictions.json \\ + --output figures/ablations +""" + +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, +}) + +COLORS = { + "full": "#1f77b4", + "flat": "#ff7f0e", + "no_buffer": "#d62728", + "measured": "#2ca02c", +} + + +def relative_error(pred, meas): + return abs(pred - meas) / meas if meas > 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_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() 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..f33b193 --- /dev/null +++ b/experiments/cost_model_benchmarks/visualization/plot_gather.py @@ -0,0 +1,221 @@ +"""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 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): + """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("--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" + 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), + ]: + 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, + 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(), + ) + + 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" + 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()