archi829/bottleneck-oracle-graphs

Dataset Card — bottleneck-oracle-graphs

Overview

A synthetic graph dataset of PyTorch profiler execution traces converted into heterogeneous graphs, designed for GNN-based bottleneck prediction in transformer model inference. Each graph represents one forward pass of a transformer variant, with nodes labelled by whether they lie on the critical execution path.

Dataset Construction

Traces were generated by profiling three transformer configurations (tiny, small, medium) using torch.profiler with CPU activity recording. Each raw Chrome trace was parsed into a directed acyclic graph using NetworkX, with the critical path computed via longest-path algorithm. Graphs were then converted to PyTorch Geometric HeteroData format.

Node Types

• —compute — standard CPU ops (matmul, attention, layernorm, etc.)
• —network — communication ops (allreduce, comm patterns)

Node Features (4-dim)

• —duration_ms — raw op duration in milliseconds
• —compute_ratio — op duration / total trace duration
• —norm_duration — op duration / max op duration in trace
• —comm_compute_ratio — 0.0 for compute nodes, 1.0 for network nodes

Labels

• —Node-level: is_critical (1 if op is on the critical path, 0 otherwise)
• —Graph-level: total forward pass time in milliseconds

Edge Types

• —compute → depends_on → compute
• —compute → sends_to → network
• —network → feeds → compute

Stats

• —501 graphs total (167 per config × 3 configs)
• —~574 nodes per graph (tiny config baseline)
• —All graphs are directed, no bidirectional leakage

Configs

Intended Use

Training GNNs to predict execution bottlenecks and critical path nodes from transformer profiler traces, without requiring real distributed training runs.

Loading

import torch
data = torch.load("graph_medium_0.pt", weights_only=False)
print(data['compute'].x.shape)   # node features
print(data['compute'].y)          # critical path labels
print(data.y)                     # graph-level step time