TGraphX Performance FAQ: Speed, GPU Use, Parallelism, and Limits

Practical questions about TGraphX performance, answered from the source code.
No benchmark claims are made — these answers describe what the library does
and does not implement, not how fast it runs on any particular hardware.

Does TGraphX run on GPU?

Yes. recommended_device() returns 'cuda' if PyTorch detects a CUDA-capable
GPU. Graph.to(device) moves all tensors — node features, edge index, edge
features, node labels — to the specified device in one call.

from tgraphx.performance import recommended_device
        device = recommended_device()
        g = my_graph.to(device)
        model = MyGNN().to(device)
        

GPU execution is inherited from PyTorch's CUDA support. TGraphX does not
implement custom CUDA kernels; it uses PyTorch's built-in tensor operations.

Does TGraphX use custom CUDA kernels?

No. TGraphX uses PyTorch's index_add_ and scatter_reduce_ for message
passing aggregation. These are PyTorch built-ins. There are no .cu files or
compiled C extensions in TGraphX.

This means installation is straightforward (no compilation step), and the
performance ceiling for message passing is PyTorch's scatter implementation
on your GPU.

Is message passing vectorized, or is there a Python loop per node?

TensorMessagePassingLayer uses batched scatter operations — no Python loop
per node. The aggregation over all edges is expressed as a single tensor
operation dispatched to the GPU (or CPU) by PyTorch.

Does TGraphX support AMP (Automatic Mixed Precision)?

The library does not explicitly wrap training loops with AMP. Users can apply
torch.cuda.amp.autocast() themselves. However, note that _scatter.py
upcasts float16 to float32 before attention softmax. This means:

If you use AMP, expect the attention path to run in float32 regardless of the
AMP context. This is a correctness choice, not a limitation.

Does TGraphX support multi-GPU or distributed training?

No. distributed.py provides rank-zero printing helpers (for printing only
from the primary process in multi-process setups), but the source is explicit:
TGraphX does not ship a distributed training framework.

For multi-GPU training, use nn.DataParallel or PyTorch's
DistributedDataParallel around your model. TGraphX Graph objects work
normally within each process.

How do I estimate memory usage before running on GPU?

Use estimate_message_memory from tgraphx.performance:

from tgraphx.performance import estimate_message_memory
        import torch
        
        est = estimate_message_memory(
            num_edges=200_000,
            node_feature_shape=(64,),
            dtype=torch.float32
        )
        print(est)  # memory estimate for scatter buffers
        

This function estimates scatter buffer memory. It does not account for model
weights, optimizer state, or activation memory. Use it as a sanity check, not
a precise memory budget.

How many DataLoader workers should I use?

GraphDataLoader wraps PyTorch's DataLoader and accepts num_workers.
The right value depends on your hardware and graph construction cost:

• Start with num_workers=2
• If GPU utilization (from nvidia-smi) is below 80%, try increasing workers
• If RAM usage spikes, reduce workers
• If graph construction is very fast (tensors already in memory), num_workers=0
  (no background processes) may be faster due to reduced overhead

Does TGraphX support graph sampling for large graphs?

No. TGraphX does not include built-in neighborhood sampling (NeighborSampler,
ClusterGCN, etc.). For graphs that don't fit in GPU memory, you would need to
implement sampling yourself and pass sampled subgraphs as Graph objects.

For large-scale graphs with sampling requirements, DGL or PyG are more
appropriate tools.

Can I use torch.compile with TGraphX?

PyTorch's torch.compile can be applied to your model as usual — TGraphX
does not prevent this. However, TGraphX does not add torch.compile wrappers
itself. The scatter operations in _scatter.py should be compatible with
torch.compile, but testing this in your specific configuration is advisable
since torch.compile compatibility can depend on PyTorch version.

How do I check my environment before running experiments?

from tgraphx.performance import env_report
        
        info = env_report(include_hardware=True)
        # Returns: python version, PyTorch version, TGraphX version,
        #          CUDA availability, GPU name if present
        print(info)
        

This is cheap and should be the first call in any experiment script. Store
the output with your results for reproducibility.

Does deterministic mode affect performance?

Yes. set_seed(seed, deterministic=True) enables PyTorch's deterministic
algorithm mode. Some operations do not have deterministic CUDA implementations,
so PyTorch either uses a slower deterministic version or raises an error
for those operations.

If you encounter a RuntimeError about non-deterministic operations, you can
either accept the non-determinism (use deterministic=False) or work around
the specific operation. For most GNN workloads, deterministic mode works but
may reduce throughput.

Are there built-in benchmarks?

TGraphX exports OGBNodeEvaluator and run_v13_benchmark_suite from its
benchmarks module (Beta v0.5.0+). These provide OGB-compatible evaluation
metrics and a standard benchmark suite for node classification. They are
evaluation tools, not throughput benchmarks — they measure model accuracy
on standardized tasks, not wall-clock speed.

Advanced Questions

Q: Can I use torch.compile with TGraphX layers?

torch.compile (PyTorch 2.x) works on standard nn.Module subclasses. Since TGraphX layers derive from nn.Module and use standard PyTorch ops, torch.compile should be compatible in principle. Test on your specific model — custom message() implementations may need fullgraph=False for dynamic shapes.

import torch
        from tgraphx import TensorGraphSAGELayer
        
        layer = TensorGraphSAGELayer(in_shape=(16, 8, 8), out_shape=(32, 4, 4))
        compiled = torch.compile(layer, fullgraph=False)
        

Always verify outputs match the non-compiled version before committing to compiled training.

Q: How do I estimate memory for a large batch?

Use the built-in estimator before training:

from tgraphx import estimate_message_memory
        import torch
        
        # For 50K edges, [C=16, H=8, W=8] features, float32
        mem_mb = estimate_message_memory(
            num_edges=50_000,
            node_feature_shape=(16, 8, 8),
            dtype=torch.float32,
        )
        print(f"Peak message buffer: {mem_mb:.1f} MB")
        # Add model parameters + activations + optimizer states
        

This is a formula-based estimate. Add ~3x overhead for activations and gradients in a training context.

Q: What is the recommended way to monitor GPU utilization?

Use env_report(include_sensors=True) if pynvml is installed:

from tgraphx import env_report
        
        info = env_report(include_hardware=True, include_sensors=True)
        print(f"GPU util: {info.get('gpu_util_pct')} %")
        print(f"GPU temp: {info.get('gpu_temp_c')} °C")
        

For continuous monitoring, nvidia-smi dmon or PyTorch's torch.cuda.memory_stats() are more appropriate.

Q: Is there a difference between running TGraphX on CPU vs GPU?

The code path is identical — all ops are standard PyTorch. GPU training benefits from CUDA parallelism for large batches and feature tensors. For small graphs (under a few thousand nodes), CPU overhead from CUDA kernel launch may dominate, making CPU faster. Always profile.