Debugging Graph Neural Networks: A Systematic Approach with TGraphX

GNN bugs are unusually hard because the data, the model, and the training loop are all interacting in ways that are not obvious from a stack trace. A CUDA assertion deep in a scatter operation could be caused by an out-of-bounds edge index, a missing self-loop, a label mismatch, or a model output shape error. Without a systematic approach, debugging becomes guesswork.

This article presents a workflow that uses TGraphX's debugging utilities to isolate problems efficiently.

The four-layer model of GNN bugs

GNN problems show up at different layers:

1. Data layer — shape mismatches, out-of-bounds edge indices, label misalignment.
2. Model layer — wrong layer for the input shape, missing activation, mismatched dimensions.
3. Training-loop layer — gradient flow issues, optimizer misconfiguration, wrong loss.
4. Reproducibility layer — different results across runs, can't reproduce baseline.

The order matters: a layer-N bug often manifests as a layer-N+1 symptom. A shape mismatch in the data layer can look like a model layer bug because that is where it raises.

Step 1: Validate the data

The first debugging step is always to validate the data:

import tgraphx as tgx
        
        tgx.validate_graph(g, strict=True)
        tgx.assert_tensor_native(g, min_rank=3)   # if you expect tensor features
        tgx.check_graph_invariants(g)
        tgx.check_leakage(train_mask, val_mask, test_mask, strict=True)
        

These run in milliseconds and catch the most common data-layer bugs:

• edge_index out of bounds.
• Inconsistent feature/label dimensions.
• Wrong tensor dtype (e.g., float when long expected).
• Device mismatch (CPU vs GPU).
• Mask leakage (same node in train and test).

If any of these raises, the bug is in the data layer. Fix it before looking at the model.

Step 2: Inspect a batch

If validation passes, inspect what a typical batch looks like:

from tgraphx import NeighborLoader
        
        loader = NeighborLoader(g, num_neighbors=[10, 5], batch_size=4, seed=42)
        batch = next(iter(loader))
        print(tgx.batch_summary(batch))
        

batch_summary prints a structured summary: node count, edge count, feature shape, label shape, device. If anything is off, the bug is in the loader configuration or the data preprocessing.

Step 3: Try a forward pass

import torch
        
        model = build_model(...)
        model.eval()
        with torch.no_grad():
            out = model(batch.x, batch.edge_index)
        print(f"Output shape: {out.shape}")
        print(f"Output stats:  mean={out.mean():.4f}, std={out.std():.4f}")
        print(f"NaN check:     {torch.isnan(out).any()}")
        print(f"Inf check:     {torch.isinf(out).any()}")
        

If forward works, the model architecture is consistent with the data shape. If output is all NaN or has unreasonable magnitudes, the problem is initialization or numerical stability.

Step 4: Try a backward pass

model.train()
        out = model(batch.x, batch.edge_index)
        seed_out = batch.seed_logits(out) if hasattr(batch, 'seed_logits') else out
        loss = torch.nn.functional.cross_entropy(seed_out, batch.seed_y if hasattr(batch, 'seed_y') else batch.y)
        loss.backward()
        
        # Check gradients
        for name, param in model.named_parameters():
            if param.grad is None:
                print(f"  ⚠️  No grad for {name}")
            else:
                print(f"  {name}: grad norm = {param.grad.norm().item():.4f}")
        

This catches:

• Layers that are not in the gradient graph (no requires_grad=True somewhere).
• Vanishing gradients (norm near zero).
• Exploding gradients (very large norm).

Step 5: Use tgx.explain_error for unfamiliar errors

When you hit a PyTorch error you do not recognize:

try:
            # ... offending code ...
            pass
        except Exception as e:
            print(tgx.explain_error(e))
        

explain_error matches the exception against a table of common GNN error patterns and returns:

• A plain-English description of likely causes.
• Suggested debugging steps.
• Pointers to relevant documentation.

It does not understand every error — but it covers the most common patterns and saves time on the easy ones.

Step 6: Bisect by simplification

When the bug is real but you cannot localize it, simplify:

• Reduce the graph to a few dozen nodes.
• Reduce the model to one layer.
• Reduce the training to one batch.
• Reduce the loss to a known-correct formulation.

If the bug persists in the simplified setup, it is reproducible and you can iterate quickly. If it disappears, add complexity back one piece at a time until it returns. The piece you add when it returns is where the bug lives.

Step 7: Audit the run directory

For mysterious reproducibility failures:

print(tgx.audit_run_dir("runs/exp_001"))
        

This checks whether all expected artifacts are present. If reproducibility_report.json is missing, you cannot diagnose run-to-run variance from the artifacts alone.

If the report is present, compare across runs:

import json
        
        r1 = json.load(open("runs/exp_001/reproducibility_report.json"))
        r2 = json.load(open("runs/exp_002/reproducibility_report.json"))
        for key in r1:
            if r1[key] != r2[key]:
                print(f"  Differ on {key}: {r1[key]} vs {r2[key]}")
        

If everything matches but results differ, the issue is hardware-level non-determinism (e.g., GPU model differences). If something differs (seed, package version), there is the cause.

Common pitfalls and their signatures

"CUDA error: device-side assert triggered": Almost always an out-of-bounds index — in edge_index or in label indexing. Run with CUDA_LAUNCH_BLOCKING=1 and find the assert. Then validate_graph(g) to confirm.

Validation accuracy stuck at uniform-random level: Labels misaligned with the input order. Verify with a manual assert (g.node_labels[:5] == expected_labels[:5]).all().

Loss decreasing on train, not on val: Standard overfitting. Not a bug, a hyperparameter issue. Reduce model capacity, add regularization.

Loss constant (not changing): Gradients are zero. Likely the model is not in the gradient graph somewhere, or learning rate is too low. Check gradient norms.

Different results across runs with same seed: Reproducibility flags not set. Wrap in tgx.reproducible(seed=42, deterministic=True).

Memory error on a small graph: Sampler num_neighbors is too large for the model. Reduce neighbor count or batch size.

Debug logging during training

For long-running issues, add periodic checks:

for epoch in range(num_epochs):
            for batch in loader:
                # ... train step ...
                if epoch % 10 == 0 and batch_idx == 0:
                    print(f"Epoch {epoch}: loss={loss.item():.4f}, "
                          f"output_std={out.std().item():.4f}")
        

If loss spikes, output std collapses, or anything unusual happens, you have a localized symptom to investigate.

Reproducibility for debugging

When you find a bug, reduce it to a deterministic reproducer:

with tgx.reproducible(seed=42, deterministic=True):
            # minimal code that reproduces the bug
            pass
        

A deterministic reproducer is essential for asking for help, for filing a bug report, and for verifying that your fix actually works.

FAQ

Q: Where is the list of patterns explain_error knows about?
A: tgraphx/ux/errors.py in the source repository. You can also extend it for project-specific patterns.

Q: What about debugging gradient instability?
A: Use torch.autograd.detect_anomaly() during the suspect forward/backward. It will tell you which operation produced the first NaN.

Q: How do I debug a custom layer?
A: Run the layer on a known small input and check output shape and values. Compare against a hand-computed reference for a tiny case.

Q: Is there a built-in pdb-style debugger?
A: No. Use Python's standard pdb or ipdb. TGraphX does not interfere with normal debugging tools.

Q: What about debugging hetero GNNs?
A: Hetero GNN is Experimental. Validation and debugging utilities have less coverage. For production hetero work, PyG's HeteroData is more mature.