GATConv on bipartite mapping triggers CUDA assert when batch-size > 1

[mishooax]

Hi,

I’m trying to map a tensor from one graph G (with 65k nodes) to a “latent space” graph G’ with just 5882 nodes (i.e. do a GATconv on a bipartite G -> G’ mapping):

def forward(self, x: torch.Tensor) -> torch.Tensor:
   bs = x.shape[0] # bs = 2
   x = einops.rearrange(x, "b n f -> (b n) f") # x.shape = [2 * 65160, NF]
   
   # self.forward_mapper is just a GATConv with x=(x, context) where context.shape = [bs*5882, 1]
   # e2h_edge_index.shape = [2, 76466]
   # e2h_edge_attr.shape = [76466, 1]
   # e2h_edge_index.min = 0, e2h_edge_index.max = 64590
   x_latent = self.forward_mapper(
            x,
            # expand edge index the correct number of times while adding the proper number to the edge index
            # this is required to handle multiple graphs in a single batch, correct?
            edge_index=torch.cat(
                [self.e2h_edge_index + i * torch.max(self.e2h_edge_index) + i for i in range(bs)],
                dim=1,
            ),
            # copy edge attributes bs times
            edge_attr=einops.repeat(self.e2h_edge_attr, "e f -> (repeat e) f", repeat=bs),
            batch_size=bs,
        )

    # unpack x_latent, do some more processing, etc.

The call to forward_mapper triggers a CUDA assert:

/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/Indexing.cu:975: indexSelectLargeIndex: block: [120,0,0], thread: [8,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/Indexing.cu:975: indexSelectLargeIndex: block: [120,0,0], thread: [9,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/Indexing.cu:975: indexSelectLargeIndex: block: [120,0,0], thread: [10,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/Indexing.cu:975: indexSelectLargeIndex: block: [120,0,0], thread: [11,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/Indexing.cu:975: indexSelectLargeIndex: block: [120,0,0], thread: [12,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
...
  File "...lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "...", line 84, in forward
    out = self.conv(x=(x, context), edge_index=edge_index, edge_attr=edge_attr, size=(x.shape[0], context.shape[0]))
  File "...lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "...lib/python3.9/site-packages/torch_geometric/nn/conv/gat_conv.py", line 246, in forward
    out = self.propagate(edge_index, x=x, alpha=alpha, size=size)
  File "...lib/python3.9/site-packages/torch_geometric/nn/conv/message_passing.py", line 391, in propagate
    out = self.aggregate(out, **aggr_kwargs)
  File "...lib/python3.9/site-packages/torch_geometric/nn/conv/message_passing.py", line 514, in aggregate
    return self.aggr_module(inputs, index, ptr=ptr, dim_size=dim_size,
  File "...lib/python3.9/site-packages/torch_geometric/nn/aggr/base.py", line 109, in __call__
    elif index.numel() > 0 and dim_size <= int(index.max()):
RuntimeError: CUDA error: device-side assert triggered

CUDA kernel errors might be asynchronously reported at some other API call, so the stack trace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

Any idea what’s going wrong? Thanks! 🙏

PS: My pyg version info (CUDA 11.6 and pytorch 1.12.1)

torch-geometric           2.1.0.post1              pypi_0    pypi
torch-scatter             2.0.9                    pypi_0    pypi
torch-sparse              0.6.15                   pypi_0    pypi

[EdisonLeeeee]

This seems to be a size mismatch error. Could you switch to the CPU to see what it breaks?

[EdisonLeeeee]

Exactly. You need to increment them with edge_index.max(1, keepdim=True).values instead, such as

self.e2h_edge_index + i * self.e2h_edge_index.max(1, keepdim=True).values + i for i in range(bs)

[mishooax]

Thanks. One more q: my e2h_edge_index that encodes the bipartite mapping looks as follows:

tensor([[... 3246, 3597, 3247, 3596, 3248, 3242, ...],
        [  34,   34,   34,   34,   34,   35 ...]])

I've generated this from a scipy "nearest-neighbors" COO adjacency matrix. The snippet above indicates that node 3246 in G is connected to node 34 in G', for example. Node indices for G run from 0 to 65159 and those for G' from 0 to 5881. Is this in agreement with what pyg expects? Or should I increment the node indices for G' to start from 65160 i.e. |G| + 1?

[EdisonLeeeee]

should I increment the node indices for G' to start from 65160 i.e. |G| + 1?

Not necessary. GATConv is initially designed to accept bipartite graphs like that. Your e2h_edge_index looks correct to me.

[mishooax]

@EdisonLeeeee - Cheers, much appreciated. One final q: should

self.e2h_edge_index + i * self.e2h_edge_index.max(1, keepdim=True).values + i for i in range(bs)

read

self.e2h_edge_index + i * self.e2h_edge_index.max(1, keepdim=True).values + 1 for i in range(bs)

(i.e. increment by 1 instead of i)? To me an extra "offset" of 1 seems to be sufficient to "separate" the node indices of the batched graphs. Thanks!

[EdisonLeeeee]

Actually, the edge_index should be incremented by num_nodes (number of nodes), where num_nodes=edge_index.max() + 1. Sorry I didn't look carefully into the code - the code should be:

self.e2h_edge_index + i * (self.e2h_edge_index.max(1, keepdim=True).values + 1) for i in range(bs)