`hetero_neighbor_sample` is non-deterministic

[pbielak]

🐛 Describe the bug

Hi, I encountered the problem that the NeighborLoader produces different outputs (different subgraphs) for hetergeneous graphs despite setting a fixed seed. In particular:

• Each time, I launch a script, the generated batches (subgraphs) are different, which makes it hard to get consistent results (i.e., different graphs -> different model outputs -> different metrics)
• I encountered the problem while dealing with model inference on heterogeneous graphs, but the same problem will also apply to training scenarios
• I tried using only torch.manual_seed(...) as well as the seed_everything(...) function from PyG - in both cases the NeighborLoader outputs are non-deterministic

A minimal reproducer is given below:

import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.loader import NeighborLoader
from torch_geometric import seed_everything


seed_everything(42)

data = OGB_MAG(
    root='./data',
    preprocess='metapath2vec',
    transform=T.ToUndirected(merge=True),
)[0]

val_loader = NeighborLoader(
    data,
    num_neighbors=[2, 2],
    shuffle=False,
    input_nodes=('paper', data['paper'].val_mask),
    batch_size=2,
    num_workers=0,
)

bt = next(iter(val_loader))
print(bt["paper"].n_id, bt["paper"].n_id.shape)

Expected behaviour: The exact same neighbor nodes are sampled for the given anchor nodes (here: 332, 756).

Actual runs of the script:

$ python bug.py
tensor([   332,    756, 228723,  93868, 450423,  36291, 716566, 651899, 494053,
        137372, 437078, 441757, 282790, 394355,  89966, 590215, 588425, 629091,
        560313, 571727, 629037, 577729, 707931,  15342, 634495]) torch.Size([25])

$ python bug.py
tensor([   332,    756, 420915,  93868,  15342, 734005, 670966, 726226, 733772,
        732766, 733204, 735112,    181, 736230, 707931, 401042, 715045]) torch.Size([17])

$ python bug.py
tensor([   332,    756, 228723,  93868, 450423,  36291, 729748, 726226, 441757,
        734491,     30, 736234, 736170,      0, 736388,    121, 736027, 707931,
        401042, 715045]) torch.Size([20])

The NeighborLoader calls internally the hetero_neighbor_sample kernel from pyg-lib. While looking into the implementation I found that the problem lies in the at::parallel_for call, i.e., https://github.com/pyg-team/pyg-lib/blob/master/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp#L714

• The random number generator is created in L606: https://github.com/pyg-team/pyg-lib/blob/master/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp#L606
• The kernel parallelizes computations for each edge type https://github.com/pyg-team/pyg-lib/blob/master/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp#L715
• All threads are using the same number generator, but we have no control over the thread launch order, so the first time we call the whole script, threads might be launched as follows: T1, T2, T3, T4, ... but the next time we run the script it may be as follows: T2, T1, T4, T3. This causes the RNG calls to be in different order each time and eventually this will result in different subgraphs.
• When disabling this parallelization (e.g., by using torch.set_num_threads(1)) the NeighborLoader outputs are deterministic, but the kernel computation is much slower (on custom graphs the slowdown was about 100x)
• I also tried to create the RNGs in each thread separately, but the performance was still degraded

I'm looking for a solution that could provide both deterministic outputs (when using a seed in the main script) and a decent performance.

Environment

• pyg-lib version: master
• PyTorch version: 2.8.0
• OS: Linux
• Python version: 3.10
• CUDA/cuDNN version: N/A
• How you installed PyTorch and pyg-lib (conda, pip, source): PyTorch - pip, pyg-lib - built from master
• Any other relevant information: N/A

[pbielak]

Follow up:

@rusty1s suggested that this non-determinism should not occur when using num_workers>0, but unfortunately it did not solve the issue

I wrote a short script to explore different configurations:

• first I create a small synthetic graph with multiple node types and relations (a star/snowflake like multirelational graph)
• one node of type n0, N nodes of type n1, N nodes of type n2, ...., N nodes of type n{T-1}
• each node of type n1...n{T-1} is connected to the one node of type n0 via relation r
• expected behaviour: the NeighborLoader should always (i.e. between different runs) retrieve the exact same subgraph (same set of nodes and edges)

The script checks the following configs:

• torch.set_num_threads(...) -> default value or 1
• num_workers -> 0, 1, T
  where each case is repeated 10 times

Script:

# check_neighbor_loader.py
import pandas as pd
import torch
import torch_geometric
from torch_geometric import seed_everything
from torch_geometric.data import HeteroData
from torch_geometric.loader import NeighborLoader


def query_loader(g: HeteroData, num_workers: int) -> dict[str, tuple[int]]:
    # Re-apply seed
    seed_everything(42)

    # Run NeighborLoader
    loader = NeighborLoader(
        g,
        num_neighbors={et: [3] for et in g.edge_types},
        batch_size=1,
        input_nodes=('n0', None),
        replace=True,
        num_workers=num_workers,
    )

    batch = loader(torch.tensor([0]))

    return {
        nt: tuple(batch[nt].n_id.tolist())
        for nt in g.node_types
    }


def main():
    N = 5
    T = 4
    dim = 2
    repeat = 10

    # Seed
    seed_everything(42)

    # Sample hetero graph
    g = HeteroData()

    g['n0'].x = torch.ones(1, dim)

    for i in range(1, T):
        g[f'n{i}'].x = torch.ones(N, dim) * torch.arange(1, N + 1).unsqueeze(dim=-1)
        g[f'n{i}', 'r', 'n0'].edge_index = torch.tensor([
            [0, 1, 2, 3, 4],
            [0, 0, 0, 0, 0],
        ])

    # Experiment with NeighborLoader settings
    records = []

    for num_workers in (0, 1, T):
        for _ in range(repeat):
            records.append({
                "torch::num_threads": torch.get_num_threads(),
                "num_workers": num_workers,
                **query_loader(g, num_workers),
            })

    torch.set_num_threads(1)
    for num_workers in (0, 1, T):
        for _ in range(repeat):
            records.append({
                "torch::num_threads": torch.get_num_threads(),
                "num_workers": num_workers,
                **query_loader(g, num_workers),
            })

    df = pd.DataFrame.from_records(records)
    df = df.groupby(["torch::num_threads", "num_workers"]).agg(set)
    print(df)


if __name__ == "__main__":
    main()

Script output:

$ python check_neighbor_loader.py
                                    n0                      n1                      n2                      n3
torch::num_threads num_workers
1                  0            {(0,)}                {(4, 1)}                {(0, 2)}                  {(1,)}
                   1            {(0,)}                {(4, 1)}                {(0, 2)}                  {(1,)}
                   4            {(0,)}                {(4, 1)}                {(0, 2)}                  {(1,)}
12                 0            {(0,)}        {(0, 2), (4, 1)}          {(1,), (0, 2)}          {(1,), (4, 1)}
                   1            {(0,)}  {(1,), (0, 2), (4, 1)}  {(1,), (0, 2), (4, 1)}  {(1,), (0, 2), (4, 1)}
                   4            {(0,)}          {(1,), (0, 2)}          {(1,), (4, 1)}  {(1,), (0, 2), (4, 1)}

When num_threads is set to 1 in pytorch, the results are fully deterministic (it always chooses nodes 1 and 4 for node type n1, nodes 0 and 2 for node type n2 and node 1 for node type n3). However, for real-world cases the sampling time increases over 100x (as mentioned in the original bug description).

When using the default num_threads value, regardless of the value of num_workers the results are non-deterministic - i.e. different iterations yield different sets of sampled nodes.

@rusty1s Do you have any other ideas?