What Can Be Parallelized in a Tensor-Valued Graph Pipeline?

Not all parts of a graph learning pipeline benefit equally from parallelization.
Some operations are inherently sequential; others can be parallelized at the
operation level, the sample level, or the data loading level. This article maps
out what TGraphX enables at each level, what requires user implementation, and
what is explicitly not supported.

The Four Levels of Parallelism

Level 1: Op-level         — within a single forward pass, on one device
        Level 2: Sample-level     — across multiple graphs in a mini-batch
        Level 3: Data loading     — constructing graphs in background processes
        Level 4: Multi-device     — across multiple GPUs or machines
        

TGraphX directly addresses the first three. The fourth is the user's
responsibility.

Level 1: Operation-Level Parallelism

Message passing aggregation is the core operation. TGraphX's
TensorMessagePassingLayer uses index_add_ and scatter_reduce_ from
PyTorch. These operate on entire tensors; the GPU executes them as parallel
reduction operations across all edges simultaneously.

The critical property: no Python loop per node or per edge. The entire
message passing step is expressed as a small number of tensor operations, each
of which PyTorch dispatches to a CUDA kernel (when running on GPU).

# What is NOT happening (slow):
        for node_idx in range(num_nodes):
            messages = [edge_feat[e] for e in edges_into[node_idx]]
            agg[node_idx] = sum(messages)
        
        # What IS happening (fast):
        # index_add_(dim, edge_dst_indices, edge_messages)
        # — all edges processed in one CUDA kernel call
        

This op-level parallelism is automatic — you get it by using the library's
layers rather than implementing your own aggregation.

The one notable exception: _scatter.py upcasts float16 to float32 for
the attention softmax in GAT-style layers. This is a correctness requirement
(float16 softmax is numerically unstable), but it means the attention path
does not run at full float16 throughput even when the rest of your model does.

Level 2: Sample-Level Parallelism (GraphBatch)

For datasets of many small graphs — molecular graphs, program graphs, patch
graphs — the sample-level parallelism is critical. A single graph with 20 nodes
occupies very little of a GPU's compute resources. Batching 64 such graphs
together makes the forward pass substantially more efficient.

GraphBatch.from_graphs([g1, g2, ..., g64]) combines multiple Graph objects
into a single batched graph by:

1. Concatenating node feature tensors along the node dimension
2. Concatenating edge feature tensors along the edge dimension
3. Offsetting each graph's edge indices so node indices don't collide
4. Recording a batch vector that maps each node to its source graph

from tgraphx import GraphBatch
        
        # 64 small graphs, each with ~20 nodes
        batch = GraphBatch.from_graphs(graphs[:64])
        # batch.node_features: [N_total, C, H, W]
        # batch.edge_index: [2, E_total] with offsets applied
        

The entire batch is processed in one forward pass. This is the primary way to
achieve GPU utilization for small-graph datasets.

Level 3: Data Loading Parallelism (GraphDataLoader)

Constructing Graph objects from raw data (building edge indices, normalizing
features, loading from disk) can be a significant bottleneck. If graph
construction happens in the main process while the GPU waits, you're leaving
throughput on the table.

GraphDataLoader wraps PyTorch's DataLoader and accepts num_workers:

from tgraphx import GraphDataLoader
        
        loader = GraphDataLoader(
            dataset=my_graph_dataset,
            batch_size=64,
            num_workers=4,   # 4 background processes for graph construction
            shuffle=True,
        )
        
        for batch in loader:
            batch = batch.to(device)
            out = model(batch)
            # ...
        

With num_workers > 0, graph construction happens in background worker
processes (on CPU). While the GPU runs the forward pass on one batch, workers
are constructing the next batch. This pipeline parallelism hides construction
latency behind compute.

Setting num_workers correctly requires experimentation. Too few workers and
construction becomes a bottleneck; too many and you exhaust CPU memory or
introduce overhead from process coordination.

Level 4: Multi-Device — User Responsibility

TGraphX's distributed.py provides helpers for rank-zero logging (printing
only from the primary process in a multi-process setup). It does not provide
a distributed training framework.

From the source documentation: TGraphX does not ship a distributed training
framework.

If you need multi-GPU training, you have two paths:

DataParallel (single machine, simple):

import torch.nn as nn
        
        model = nn.DataParallel(model)  # wraps your model, handles GPU splits
        # Each GPU receives a portion of the batch
        # Gradients are automatically synchronized
        

DistributedDataParallel (single or multi-machine, recommended for performance):

import torch.distributed as dist
        import torch.nn.parallel as para
        
        dist.init_process_group(backend='nccl')
        model = para.DistributedDataParallel(model, device_ids=[local_rank])
        # TGraphX Graph objects are used within each rank as normal
        

In both cases, TGraphX is used as a local computation library within each
process. The Graph and GraphBatch objects don't need to be aware of
distribution — each worker creates and processes its own graphs.

Pipeline Diagram

  Data pipeline parallelism in TGraphX
          ─────────────────────────────────────────────────────────────────
        
          Worker 0 ──► [construct Graph 0..63]  ──────────────────────────►
          Worker 1 ──► [construct Graph 64..127] ────────────────────────►
          Worker 2 ──► [construct Graph 128..191] ──────────────────────►
          Worker 3 ──► [construct Graph 192..255] ────────────────────►
        
          Main proc ──► [batch 0..63] ──► [.to(device)] ──► [forward] ──► [backward]
                                                   ↑
                                           GPU is busy here
                                           workers building next batch above
          ─────────────────────────────────────────────────────────────────
        
          Within forward pass:
          GraphBatch.node_features [N_total, ...] ──► scatter aggregation
                                                      (index_add_, scatter_reduce_)
                                                      ── all edges in parallel on GPU ──►
        

Parallelism Summary Table

Practical Guidance

For most research experiments on a single GPU:
- Use GraphBatch to batch your graphs
- Use GraphDataLoader with num_workers=2 to 4 as a starting point
- Do not worry about multi-GPU until you've verified the single-GPU
implementation is correct and you have a genuine scale requirement

For datasets of large single graphs (one graph that doesn't fit in memory),
TGraphX's current tooling is not sufficient. You would need to implement
neighborhood sampling yourself and pass sampled subgraphs to the model.