TGraphX › Insights › Profiling a TGraphX Workflow with PyTorch Tools
← Back to Insights
Tutorial · · 9 min read

Profiling a TGraphX Workflow with PyTorch Tools

Target keyword: profile PyTorch graph neural network workflow

Profiling a TGraphX Workflow with PyTorch Tools

Before optimizing a graph learning workflow, you need to know where time
actually goes. Guessing is rarely productive — the bottleneck is often not
where intuition says it is. This tutorial covers profiling a TGraphX training
loop using PyTorch's built-in profiling tools, with specific attention to the
operations that matter most in graph learning.

The Profiling Mindset

Profiling answers: "Where does wall time go?" For a GNN training loop, the
candidates are:

  1. Data loading and graph construction — building Graph objects from raw data
  2. Device transfer — .to(device) calls
  3. Message passing aggregation — scatter operations in layers
  4. Linear projections — the nn.Linear layers within each GNN layer
  5. Loss computation and backward pass — autograd overhead
  6. Optimizer step — weight update

Without profiling, people typically over-optimize message passing (which is
often already fast due to vectorized scatter) and under-optimize data loading
(which is often the actual bottleneck).

Basic Timing with CUDA Synchronization

The simplest profiling approach is manual timing. But on GPU, timing without
synchronization gives misleading results because CUDA operations are
asynchronous — the Python timer may return before the GPU has finished.

python
import torch
        import time
        
        def timed_section(label, fn, device):
            if device == "cuda":
                torch.cuda.synchronize()
            t0 = time.perf_counter()
            result = fn()
            if device == "cuda":
                torch.cuda.synchronize()
            t1 = time.perf_counter()
            print(f"{label}: {(t1 - t0)*1000:.2f} ms")
            return result
        
        from tgraphx.performance import recommended_device
        
        device = recommended_device()
        
        # Time the forward pass
        timed_section(
            "forward pass",
            lambda: model(g),
            device,
        )

This approach is useful for coarse timing but misses the breakdown between
operations within the forward pass.

Using torch.profiler for Detailed Tracing

PyTorch's torch.profiler provides operator-level profiling with CUDA event
tracking:

python
import torch
        from torch.profiler import profile, record_function, ProfilerActivity
        
        activities = [ProfilerActivity.CPU]
        if device == "cuda":
            activities.append(ProfilerActivity.CUDA)
        
        with profile(
            activities=activities,
            record_shapes=True,
            with_stack=False,
        ) as prof:
            with record_function("full_training_step"):
        
                with record_function("data_to_device"):
                    g = raw_graph.to(device)
        
                with record_function("forward_pass"):
                    out = model(g)
        
                with record_function("loss"):
                    loss = criterion(out, g.node_labels)
        
                with record_function("backward"):
                    loss.backward()
        
                with record_function("optimizer_step"):
                    optimizer.step()
                    optimizer.zero_grad()
        
        # Print top operations by self CPU time
        print(prof.key_averages().table(
            sort_by="cpu_time_total",
            row_limit=20,
        ))

The record_function context managers tag operations with human-readable
labels that appear in the profiler output.

Exporting a Chrome Trace

For visual exploration of the timeline:

python
prof.export_chrome_trace("trace.json")
        # Open in Chrome at chrome://tracing, or use Perfetto UI

The Chrome trace view shows a timeline where each operation occupies a
horizontal block. For GNN workflows, you're looking for:

  • Long blocks in the scatter aggregation (indicating the message passing
    kernel is the bottleneck)
  • Long gaps between GPU operations (indicating CPU overhead or data transfer
    bottlenecks)
  • Wide data loading blocks (indicating graph construction is the bottleneck)

What Operations to Look For

Operation in profiler What it represents in TGraphX
aten::index_add_ Message aggregation in SAGE/GIN layers
aten::scatter_reduce_ Alternative aggregation in scatter layers
aten::mm / aten::addmm Linear projections in layer weights
aten::softmax Attention coefficient normalization (GAT)
aten::to Device transfer (.to(device))
DataLoader worker time Graph construction in workers
aten::cat GraphBatch concatenation during collation

Profiling Data Loading Separately

Data loading happens in background workers and is not captured by
torch.profiler in the main process. To profile the data loading pipeline
separately:

python
from tgraphx import GraphDataLoader
        import time
        
        loader = GraphDataLoader(
            dataset=my_dataset,
            batch_size=32,
            num_workers=4,
        )
        
        # Time how long it takes to consume one batch
        t0 = time.perf_counter()
        batch = next(iter(loader))
        t1 = time.perf_counter()
        print(f"First batch load time: {(t1-t0)*1000:.1f} ms")
        
        # Compare to construction in main process (num_workers=0)
        loader_sync = GraphDataLoader(
            dataset=my_dataset,
            batch_size=32,
            num_workers=0,
        )
        t0 = time.perf_counter()
        batch = next(iter(loader_sync))
        t1 = time.perf_counter()
        print(f"Synchronous batch load time: {(t1-t0)*1000:.1f} ms")

If the synchronous time is much greater than the GPU forward pass time, more
workers will help. If they're similar, workers may not improve throughput.

Using record_function in a GNN Module

Add record_function markers inside your model for fine-grained layer
profiling:

python
import torch.nn as nn
        from torch.profiler import record_function
        
        class ProfiledGNN(nn.Module):
            def __init__(self, layer1, layer2):
                super().__init__()
                self.layer1 = layer1
                self.layer2 = layer2
        
            def forward(self, g):
                with record_function("layer1_message_pass"):
                    x = self.layer1(g)
                with record_function("relu"):
                    x = torch.relu(x)
                with record_function("layer2_message_pass"):
                    x = self.layer2(g)
                return x

This gives you per-layer timing in the profiler output.

Memory Profiling with estimate_message_memory

Before running on GPU, use TGraphX's built-in estimator:

python
from tgraphx.performance import estimate_message_memory
        import torch
        
        # Estimate before construction, as a sanity check
        est = estimate_message_memory(
            num_edges=g.edge_index.shape[1],
            node_feature_shape=tuple(g.node_features.shape[1:]),
            dtype=torch.float32,
        )
        print(f"Estimated scatter buffer memory: {est}")

For runtime memory profiling, use PyTorch's CUDA memory APIs:

python
if device == "cuda":
            torch.cuda.reset_peak_memory_stats()
            out = model(g)
            peak_mb = torch.cuda.max_memory_allocated() / 1024**2
            print(f"Peak GPU memory: {peak_mb:.1f} MB")

Interpreting Results: Common Patterns

Pattern 1: Data loading dominates
If loading one batch takes longer than the GPU forward + backward, increase
num_workers. If already at 4+ workers, consider caching preprocessed graphs.

Pattern 2: Device transfer is expensive
If aten::to takes significant time, batch your .to(device) calls rather
than moving individual tensors. g.to(device) is already batched for the
Graph object.

Pattern 3: aten::cat appears large
This is GraphBatch concatenation during collation. It scales with the number
of graphs per batch multiplied by tensor sizes. If this dominates, consider
pre-batching your dataset or reducing the number of Graph objects per batch.

Pattern 4: aten::index_add_ is the bottleneck
The message passing aggregation itself dominates. This is expected for
large, dense graphs. Options: reduce graph density, reduce feature dimensions,
or accept the cost as inherent to the graph structure.

Reproducible Profiling

Always profile under controlled conditions:

python
from tgraphx.reproducibility import set_seed
        from tgraphx.performance import env_report
        
        set_seed(42, deterministic=True)
        info = env_report(include_hardware=True)
        print("Profiling environment:", info)
        
        # Warm up (PyTorch JIT and CUDA have startup overhead)
        for _ in range(3):
            _ = model(g)
        if device == "cuda":
            torch.cuda.synchronize()
        
        # Now profile
        with profile(...) as prof:
            for _ in range(10):  # profile multiple iterations
                out = model(g)

Running multiple iterations and averaging removes run-to-run noise from
profiling results.