Repository: pyg-team/pytorch_geometric
Pull request #9541 (2024-07-25 00:20:55)

Adapt cuGraph examples to new API
Updates the cuGraph examples to use the new non-dask API.  Adds a new multinode cuGraph example.

alexbarghi-nv: @puririshi98 
puririshi98: LGTM, @rusty1s @akihironitta anything else needed to merge?

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0f564a8f4..42e55cd28 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added a `residual` option in `GATConv` and `GATv2Conv` ([#9515](https://github.com/pyg-team/pytorch_geometric/pull/9515))
 - Added the `PatchTransformerAggregation` layer ([#9487](https://github.com/pyg-team/pytorch_geometric/pull/9487))
 - Added the `nn.nlp.LLM` model ([#9462](https://github.com/pyg-team/pytorch_geometric/pull/9462))
 - Added an example of training GNNs for a graph-level regression task ([#9070](https://github.com/pyg-team/pytorch_geometric/pull/9070))
@@ -37,6 +38,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Adapt `cugraph` examples to its new API ([#9541](https://github.com/pyg-team/pytorch_geometric/pull/9541))
 - Allow optional but untyped tensors in `MessagePassing` ([#9494](https://github.com/pyg-team/pytorch_geometric/pull/9494))
 - Added support for modifying `filename` of the stored partitioned file in `ClusterLoader` ([#9448](https://github.com/pyg-team/pytorch_geometric/pull/9448))
 - Support other than two-dimensional inputs in `AttentionalAggregation` ([#9433](https://github.com/pyg-team/pytorch_geometric/pull/9433))
diff --git a/examples/multi_gpu/README.md b/examples/multi_gpu/README.md
index df1b6e211..191ca4a86 100644
--- a/examples/multi_gpu/README.md
+++ b/examples/multi_gpu/README.md
@@ -9,8 +9,9 @@
 | [`distributed_sampling_multinode.py`](./distributed_sampling_multinode.py)         | multi-node  | Example for training GNNs on a homogeneous graph with neighbor sampling on multiple nodes.                                                                                                                                                                                                      |
 | [`distributed_sampling_multinode.sbatch`](./distributed_sampling_multinode.sbatch) | multi-node  | Example for submitting a training job to a Slurm cluster using [`distributed_sampling_multi_node.py`](./distributed_sampling_multinode.py).                                                                                                                                                     |
 | [`papers100m_gcn.py`](./papers100m_gcn.py)                                         | single-node | Example for training GNNs on the `ogbn-papers100M` homogeneous graph w/ ~1.6B edges.                                                                                                                                                                                                            |
-| [`papers100m_gcn_cugraph.py`](./papers100m_gcn_cugraph.py%60)                      | single-node | Example for accelerating GNN training on `ogbn-papers100M` using [CuGraph](...).                                                                                                                                                                                                                |
+| [`papers100m_gcn_cugraph.py`](./papers100m_gcn_cugraph.py)                         | single-node | Example for training GNNs on `ogbn-papers100M` using [CuGraph](...).                                                                                                                                                                                                                            |
 | [`papers100m_gcn_multinode.py`](./papers100m_gcn_multinode.py)                     | multi-node  | Example for training GNNs on a homogeneous graph on multiple nodes.                                                                                                                                                                                                                             |
+| [`papers100m_gcn_cugraph_multinode.py`](./papers100m_gcn_cugraph_multinode.py)     | multi-node  | Example for training GNNs on a homogeneous graph on multiple nodes using [CuGraph](...).                                                                                                                                                                                                        |
 | [`pcqm4m_ogb.py`](./pcqm4m_ogb.py)                                                 | single-node | Example for training GNNs for a graph-level regression task.                                                                                                                                                                                                                                    |
 | [`mag240m_graphsage.py`](./mag240m_graphsage.py)                                   | single-node | Example for training GNNs on a large heterogeneous graph.                                                                                                                                                                                                                                       |
 | [`taobao.py`](./taobao.py)                                                         | single-node | Example for training link prediction GNNs on a heterogeneous graph.                                                                                                                                                                                                                             |
diff --git a/examples/multi_gpu/papers100m_gcn_cugraph.py b/examples/multi_gpu/papers100m_gcn_cugraph.py
index 0473da463..2e7267bcc 100644
--- a/examples/multi_gpu/papers100m_gcn_cugraph.py
+++ b/examples/multi_gpu/papers100m_gcn_cugraph.py
@@ -1,17 +1,22 @@
+"""Single-node, multi-GPU example."""
+
 import argparse
 import os
+import os.path as osp
 import tempfile
 import time
 
-import numpy as np
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 import torch.nn.functional as F
+from cugraph.gnn import (
+    cugraph_comms_create_unique_id,
+    cugraph_comms_init,
+    cugraph_comms_shutdown,
+)
 from ogb.nodeproppred import PygNodePropPredDataset
-from torch.distributed.algorithms.join import Join
 from torch.nn.parallel import DistributedDataParallel
-from torchmetrics import Accuracy
 
 import torch_geometric
 
@@ -24,55 +29,24 @@ os.environ['CUDF_SPILL'] = '1'
 os.environ['RAPIDS_NO_INITIALIZE'] = '1'
 
 
-def start_dask_cluster():
-    from cugraph.testing.mg_utils import enable_spilling
-    from dask_cuda import LocalCUDACluster
+def init_pytorch_worker(rank, world_size, cugraph_id):
+    import rmm
 
-    cluster = LocalCUDACluster(
-        protocol="tcp",
-        rmm_pool_size=None,
-        memory_limit=None,
-        rmm_async=True,
+    rmm.reinitialize(
+        devices=rank,
+        managed_memory=True,
+        pool_allocator=True,
     )
 
-    from dask.distributed import Client
-    client = Client(cluster)
-    client.wait_for_workers(n_workers=len(cluster.workers))
-    client.run(enable_spilling)
-
-    print("Dask Cluster Setup Complete")
-    return client, cluster
-
-
-def shutdown_dask_client(client):
-    from cugraph.dask.comms import comms as Comms
-    Comms.destroy()
-    client.close()
-
-
-def pyg_num_work(world_size):
-    num_work = None
-    if hasattr(os, "sched_getaffinity"):
-        try:
-            num_work = len(os.sched_getaffinity(0)) / (2 * world_size)
-        except Exception:
-            pass
-    if num_work is None:
-        num_work = os.cpu_count() / (2 * world_size)
-    return int(num_work)
-
-
-def init_pytorch_worker(rank, world_size):
-    import rmm
-    if rank > 0:
-        rmm.reinitialize(devices=rank)
-
     import cupy
+
     cupy.cuda.Device(rank).use()
     from rmm.allocators.cupy import rmm_cupy_allocator
+
     cupy.cuda.set_allocator(rmm_cupy_allocator)
 
     from cugraph.testing.mg_utils import enable_spilling
+
     enable_spilling()
 
     torch.cuda.set_device(rank)
@@ -81,20 +55,19 @@ def init_pytorch_worker(rank, world_size):
     os.environ['MASTER_PORT'] = '12355'
     dist.init_process_group('nccl', rank=rank, world_size=world_size)
 
+    cugraph_comms_init(rank=rank, world_size=world_size, uid=cugraph_id,
+                       device=rank)
 
-def run_train(rank, data, world_size, model, epochs, batch_size, fan_out,
-              split_idx, num_classes, wall_clock_start, tempdir=None,
-              num_layers=3):
+
+def run(rank, data, world_size, cugraph_id, model, epochs, batch_size, fan_out,
+        split_idx, num_classes, wall_clock_start, tempdir=None, num_layers=3):
 
     init_pytorch_worker(
         rank,
         world_size,
+        cugraph_id,
     )
 
-    if rank == 0:
-        client, cluster = start_dask_cluster()
-        from cugraph.dask.comms import comms as Comms
-        Comms.initialize(p2p=True)
     model = model.to(rank)
     model = DistributedDataParallel(model, device_ids=[rank])
     optimizer = torch.optim.Adam(model.parameters(), lr=0.01,
@@ -104,230 +77,198 @@ def run_train(rank, data, world_size, model, epochs, batch_size, fan_out,
         num_neighbors=[fan_out] * num_layers,
         batch_size=batch_size,
     )
-    # Set Up Neighbor Loading
-    import cugraph
-    from cugraph_pyg.data import CuGraphStore
-    from cugraph_pyg.loader import BulkSampleLoader
-
-    # define the edges of the Graph
-    G = {("N", "E", "N"): data.edge_index}
-    # define the number of nodes in Graph
-    N = {"N": data.num_nodes}
-    # initialize feature store
-    fs = cugraph.gnn.FeatureStore(backend="torch")
-    # store node features as x
-    fs.add_data(data.x, "N", "x")
-    # store node labels as y
-    fs.add_data(data.y, "N", "y")
-    dist.barrier()
-
-    if rank == 0:
-        print("Rank 0 creating its cugraph store and \
-            initializing distributed graph")
-        cugraph_store = CuGraphStore(fs, G, N, multi_gpu=True)
-        print("Distributed graph initialization complete.")
+    from cugraph_pyg.data import GraphStore, TensorDictFeatureStore
+    from cugraph_pyg.loader import NeighborLoader
+
+    graph_store = GraphStore(is_multi_gpu=True)
+    ixr = torch.tensor_split(data.edge_index, world_size, dim=1)[rank]
+    graph_store[dict(
+        edge_type=('node', 'rel', 'node'),
+        layout='coo',
+        is_sorted=False,
+        size=(data.num_nodes, data.num_nodes),
+    )] = ixr
+
+    feature_store = TensorDictFeatureStore()
+    feature_store['node', 'x'] = data.x
+    feature_store['node', 'y'] = data.y
 
-    if rank != 0:
-        print(f"Rank {rank} waiting for distributed graph initialization")
     dist.barrier()
 
-    if rank != 0:
-        print(f"Rank {rank} proceeding with store creation")
-        cugraph_store = CuGraphStore(fs, {
-            k: len(v)
-            for k, v in G.items()
-        }, N, multi_gpu=False)
-        print(f"Rank {rank} created store")
-    dist.barrier()
+    ix_train = torch.tensor_split(split_idx['train'], world_size)[rank].cuda()
+    train_path = osp.join(tempdir, f'train_{rank}')
+    os.mkdir(train_path)
+    train_loader = NeighborLoader(
+        (feature_store, graph_store),
+        input_nodes=ix_train,
+        directory=train_path,
+        shuffle=True,
+        drop_last=True,
+        **kwargs,
+    )
 
-    if rank == 0:
-        # Direct cuGraph to sample offline prior to the training loop
-        # Sampling will occur in parallel but will be initiated on rank 0
-        for epoch in range(epochs):
-            train_path = os.path.join(tempdir, f'samples_{epoch}')
-            os.mkdir(train_path)
-            BulkSampleLoader(cugraph_store, cugraph_store,
-                             input_nodes=split_idx['train'],
-                             directory=train_path, shuffle=True,
-                             drop_last=True, **kwargs)
-
-            print('validation', len(split_idx['valid']))
-            eval_path = os.path.join(tempdir, f'samples_eval_{epoch}')
-            BulkSampleLoader(cugraph_store, cugraph_store,
-                             input_nodes=split_idx['valid'],
-                             directory=eval_path, **kwargs)
-
-        print('test', len(split_idx['test']))
-        test_path = os.path.join(tempdir, 'samples_test')
-        BulkSampleLoader(cugraph_store, cugraph_store,
-                         input_nodes=split_idx['test'], directory=test_path,
-                         **kwargs)
+    ix_val = torch.tensor_split(split_idx['valid'], world_size)[rank].cuda()
+    val_path = osp.join(tempdir, f'val_{rank}')
+    os.mkdir(val_path)
+    val_loader = NeighborLoader(
+        (feature_store, graph_store),
+        input_nodes=ix_val,
+        directory=val_path,
+        drop_last=True,
+        **kwargs,
+    )
+
+    ix_test = torch.tensor_split(split_idx['test'], world_size)[rank].cuda()
+    test_path = osp.join(tempdir, f'test_{rank}')
+    os.mkdir(test_path)
+    test_loader = NeighborLoader(
+        (feature_store, graph_store),
+        input_nodes=ix_test,
+        directory=test_path,
+        drop_last=True,
+        local_seeds_per_call=80000,
+        **kwargs,
+    )
 
     dist.barrier()
 
     eval_steps = 1000
     warmup_steps = 20
-    acc = Accuracy(task="multiclass", num_classes=num_classes).to(rank)
     dist.barrier()
     torch.cuda.synchronize()
-    if rank == 0:
-        prep_time = round(time.perf_counter() - wall_clock_start, 2)
-        print("Total time before training begins (prep_time) =", prep_time,
-              "seconds")
-        print("Beginning training...")
-    for epoch in range(epochs):
-        train_path = os.path.join(tempdir, f'samples_{epoch}')
-
-        input_files = np.array_split(np.array(os.listdir(train_path)),
-                                     world_size)[rank]
-
-        train_loader = BulkSampleLoader(cugraph_store, cugraph_store,
-                                        directory=train_path,
-                                        input_files=input_files)
-        with Join([model], divide_by_initial_world_size=False):
-            for i, batch in enumerate(train_loader):
-                if i == warmup_steps:
-                    torch.cuda.synchronize()
-                    start = time.time()
-                batch = batch.to(rank)
-
-                batch = batch.to_homogeneous()
-                batch_size = batch.num_sampled_nodes[0]
 
-                batch.y = batch.y.to(torch.long)
-                optimizer.zero_grad()
-                out = model(batch.x, batch.edge_index)
-                loss = F.cross_entropy(out[:batch_size], batch.y[:batch_size])
-                loss.backward()
-                optimizer.step()
-                if rank == 0 and i % 10 == 0:
-                    print("Epoch: " + str(epoch) + ", Iteration: " + str(i) +
-                          ", Loss: " + str(loss))
+    if rank == 0:
+        prep_time = time.perf_counter() - wall_clock_start
+        print(f"Preparation time: {prep_time:.2f}s")
+
+    for epoch in range(1, epochs + 1):
+        for i, batch in enumerate(train_loader):
+            if i == warmup_steps:
+                torch.cuda.synchronize()
+                start = time.time()
+
+            batch = batch.to(rank)
+            batch_size = batch.batch_size
+
+            batch.y = batch.y.to(torch.long)
+            optimizer.zero_grad()
+            out = model(batch.x, batch.edge_index)
+            loss = F.cross_entropy(out[:batch_size], batch.y[:batch_size])
+            loss.backward()
+            optimizer.step()
+            if rank == 0 and i % 10 == 0:
+                print(f"Epoch: {epoch:02d}, Iteration: {i}, Loss: {loss:.4f}")
         nb = i + 1.0
-        dist.barrier()
-        torch.cuda.synchronize()
-        if rank == 0:
-            print("Average Training Iteration Time:",
-                  (time.time() - start) / (nb - warmup_steps), "s/iter")
-        eval_path = os.path.join(tempdir, f'samples_eval_{epoch}')
-
-        input_files = np.array(os.listdir(eval_path))
-
-        eval_loader = BulkSampleLoader(cugraph_store, cugraph_store,
-                                       directory=eval_path,
-                                       input_files=input_files)
-        with Join([model], divide_by_initial_world_size=False):
-            with torch.no_grad():
-                for i, batch in enumerate(eval_loader):
-                    if i >= eval_steps:
-                        break
-
-                    batch = batch.to(rank)
-                    batch = batch.to_homogeneous()
-                    batch_size = batch.num_sampled_nodes[0]
-
-                    batch.y = batch.y.to(torch.long)
-                    out = model.module(batch.x, batch.edge_index)
-                    acc_i = acc(  # noqa
-                        out[:batch_size].softmax(dim=-1), batch.y[:batch_size])
-            acc_sum = acc.compute()
-            if rank == 0:
-                print(f"Validation Accuracy: {acc_sum * 100.0:.4f}%", )
-        dist.barrier()
 
-    with Join([model], divide_by_initial_world_size=False):
-        test_path = os.path.join(tempdir, 'samples_test')
-
-        input_files = np.array(os.listdir(test_path))
+        if rank == 0:
+            print(f"Avg Training Iteration Time: "
+                  f"{(time.time() - start) / (nb - warmup_steps):.4f} s/iter")
 
-        test_loader = BulkSampleLoader(cugraph_store, cugraph_store,
-                                       directory=test_path,
-                                       input_files=input_files)
         with torch.no_grad():
-            for i, batch in enumerate(test_loader):
+            total_correct = total_examples = 0
+            for i, batch in enumerate(val_loader):
+                if i >= eval_steps:
+                    break
+
                 batch = batch.to(rank)
-                batch = batch.to_homogeneous()
-                batch_size = batch.num_sampled_nodes[0]
+                batch_size = batch.batch_size
 
                 batch.y = batch.y.to(torch.long)
-                out = model.module(batch.x, batch.edge_index)
-                acc_i = acc(  # noqa
-                    out[:batch_size].softmax(dim=-1), batch.y[:batch_size])
-            acc_sum = acc.compute()
+                out = model(batch.x, batch.edge_index)[:batch_size]
+
+                pred = out.argmax(dim=-1)
+                y = batch.y[:batch_size].view(-1).to(torch.long)
+
+                total_correct += int((pred == y).sum())
+                total_examples += y.size(0)
+
+            acc_val = total_correct / total_examples
             if rank == 0:
-                print(f"Test Accuracy: {acc_sum * 100.0:.4f}%", )
-    dist.barrier()
+                print(f"Validation Accuracy: {acc_val * 100:.2f}%", )
+
+        torch.cuda.synchronize()
+
+    with torch.no_grad():
+        total_correct = total_examples = 0
+        for i, batch in enumerate(test_loader):
+            batch = batch.to(rank)
+            batch_size = batch.batch_size
+
+            batch.y = batch.y.to(torch.long)
+            out = model(batch.x, batch.edge_index)[:batch_size]
+
+            pred = out.argmax(dim=-1)
+            y = batch.y[:batch_size].view(-1).to(torch.long)
+
+            total_correct += int((pred == y).sum())
+            total_examples += y.size(0)
+
+        acc_test = total_correct / total_examples
+        if rank == 0:
+            print(f"Test Accuracy: {acc_test * 100:.2f}%", )
 
-    import gc
-    del cugraph_store
-    gc.collect()
-    shutdown_dask_client(client)
-    dist.barrier()
     if rank == 0:
-        total_time = round(time.perf_counter() - wall_clock_start, 2)
-        print("Total Program Runtime (total_time) =", total_time, "seconds")
-        print("total_time - prep_time =", total_time - prep_time, "seconds")
+        total_time = time.perf_counter() - wall_clock_start
+        print(f"Total Training Runtime: {total_time - prep_time}s")
+        print(f"Total Program Runtime: {total_time}s")
 
+    cugraph_comms_shutdown()
+    dist.destroy_process_group()
 
-if __name__ == '__main__':
 
+if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--hidden_channels', type=int, default=256)
     parser.add_argument('--num_layers', type=int, default=2)
     parser.add_argument('--lr', type=float, default=0.001)
-    parser.add_argument('--epochs', type=int, default=20)
+    parser.add_argument('--epochs', type=int, default=4)
     parser.add_argument('--batch_size', type=int, default=1024)
     parser.add_argument('--fan_out', type=int, default=30)
+    parser.add_argument('--dataset', type=str, default='ogbn-papers100M')
+    parser.add_argument('--root', type=str, default='dataset')
+    parser.add_argument('--tempdir_root', type=str, default=None)
     parser.add_argument(
-        "--use_gat_conv",
-        action='store_true',
-        help="Whether or not to use GATConv. (Defaults to using GCNConv)",
-    )
-    parser.add_argument(
-        "--n_gat_conv_heads",
+        '--n_devices',
         type=int,
-        default=4,
-        help="If using GATConv, number of attention heads to use",
+        default=-1,
+        help='How many GPUs to use. Defaults to all available GPUs',
     )
-    parser.add_argument(
-        "--n_devices", type=int, default=-1,
-        help="1-8 to use that many GPUs. Defaults to all available GPUs")
 
     args = parser.parse_args()
     wall_clock_start = time.perf_counter()
 
-    dataset = PygNodePropPredDataset(name='ogbn-papers100M',
-                                     root='/datasets/ogb_datasets')
+    from rmm.allocators.torch import rmm_torch_allocator
+
+    torch.cuda.memory.change_current_allocator(rmm_torch_allocator)
+
+    dataset = PygNodePropPredDataset(name=args.dataset, root=args.root)
     split_idx = dataset.get_idx_split()
     data = dataset[0]
     data.y = data.y.reshape(-1)
-    if args.use_gat_conv:
-        model = torch_geometric.nn.models.GAT(dataset.num_features,
-                                              args.hidden_channels,
-                                              args.num_layers,
-                                              dataset.num_classes,
-                                              heads=args.n_gat_conv_heads)
-    else:
-        model = torch_geometric.nn.models.GCN(dataset.num_features,
-                                              args.hidden_channels,
-                                              args.num_layers,
-                                              dataset.num_classes)
 
-    print("Data =", data)
-    if args.n_devices == -1:
+    model = torch_geometric.nn.models.GCN(
+        dataset.num_features,
+        args.hidden_channels,
+        args.num_layers,
+        dataset.num_classes,
+    )
+
+    if args.n_devices < 1:
         world_size = torch.cuda.device_count()
     else:
         world_size = args.n_devices
-    print('Let\'s use', world_size, 'GPUs!')
-    with tempfile.TemporaryDirectory() as tempdir:
-        if world_size > 1:
-            mp.spawn(
-                run_train,
-                args=(data, world_size, model, args.epochs, args.batch_size,
-                      args.fan_out, split_idx, dataset.num_classes,
-                      wall_clock_start, tempdir, args.num_layers),
-                nprocs=world_size, join=True)
-        else:
-            run_train(0, data, world_size, model, args.epochs, args.batch_size,
-                      args.fan_out, split_idx, dataset.num_classes,
-                      wall_clock_start, tempdir, args.num_layers)
+    print(f"Using {world_size} many GPUs...")
+
+    # Create the uid needed for cuGraph comms
+    cugraph_id = cugraph_comms_create_unique_id()
+
+    with tempfile.TemporaryDirectory(dir=args.tempdir_root) as tempdir:
+        mp.spawn(
+            run,
+            args=(data, world_size, cugraph_id, model, args.epochs,
+                  args.batch_size, args.fan_out, split_idx,
+                  dataset.num_classes, wall_clock_start, tempdir,
+                  args.num_layers),
+            nprocs=world_size,
+            join=True,
+        )
diff --git a/examples/multi_gpu/papers100m_gcn_cugraph_multinode.py b/examples/multi_gpu/papers100m_gcn_cugraph_multinode.py
new file mode 100644
index 000000000..cbb55f2b6
--- /dev/null
+++ b/examples/multi_gpu/papers100m_gcn_cugraph_multinode.py
@@ -0,0 +1,389 @@
+# Multi-node, multi-GPU example with WholeGraph feature storage.
+# It is recommended that you download the dataset first before running.
+
+# To run, use sbatch
+# (i.e. sbatch -N2 -p <partition> -A <account> -J <job name>)
+# with the script shown below:
+#
+# head_node_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+#
+# (yes || true) | srun -l \
+#        --container-image <container image> \
+#        --container-mounts "$(pwd):/workspace","/raid:/raid" \
+#          torchrun \
+#          --nnodes 2 \
+#          --nproc-per-node 8 \
+#          --rdzv-backend c10d \
+#          --rdzv-id 62 \
+#          --rdzv-endpoint $head_node_addr:29505 \
+#          /workspace/papers100m_gcn_cugraph_multinode.py \
+#            --epochs 1 \
+#            --dataset ogbn-papers100M \
+#            --dataset_root /workspace/datasets \
+#            --tempdir_root /raid/scratch
+
+import argparse
+import json
+import os
+import os.path as osp
+import tempfile
+import time
+from datetime import timedelta
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from cugraph.gnn import (
+    cugraph_comms_create_unique_id,
+    cugraph_comms_init,
+    cugraph_comms_shutdown,
+)
+from ogb.nodeproppred import PygNodePropPredDataset
+from pylibwholegraph.torch.initialize import finalize as wm_finalize
+from pylibwholegraph.torch.initialize import init as wm_init
+from torch.nn.parallel import DistributedDataParallel
+
+import torch_geometric
+
+# Allow computation on objects that are larger than GPU memory
+# https://docs.rapids.ai/api/cudf/stable/developer_guide/library_design/#spilling-to-host-memory
+os.environ['CUDF_SPILL'] = '1'
+
+# Ensures that a CUDA context is not created on import of rapids.
+# Allows pytorch to create the context instead
+os.environ['RAPIDS_NO_INITIALIZE'] = '1'
+
+
+def init_pytorch_worker(global_rank, local_rank, world_size, cugraph_id):
+    import rmm
+
+    rmm.reinitialize(
+        devices=local_rank,
+        managed_memory=True,
+        pool_allocator=True,
+    )
+
+    import cupy
+
+    cupy.cuda.Device(local_rank).use()
+    from rmm.allocators.cupy import rmm_cupy_allocator
+
+    cupy.cuda.set_allocator(rmm_cupy_allocator)
+
+    from cugraph.testing.mg_utils import enable_spilling
+
+    enable_spilling()
+
+    torch.cuda.set_device(local_rank)
+
+    cugraph_comms_init(rank=global_rank, world_size=world_size, uid=cugraph_id,
+                       device=local_rank)
+
+    wm_init(global_rank, world_size, local_rank, torch.cuda.device_count())
+
+
+def partition_data(dataset, split_idx, edge_path, feature_path, label_path,
+                   meta_path):
+    data = dataset[0]
+
+    os.makedirs(edge_path, exist_ok=True)
+    for (r, e) in enumerate(data.edge_index.tensor_split(world_size, dim=1)):
+        rank_path = osp.join(edge_path, f'rank={r}.pt')
+        torch.save(
+            e.clone(),
+            rank_path,
+        )
+
+    os.makedirs(feature_path, exist_ok=True)
+    for (r, f) in enumerate(torch.tensor_split(data.x, world_size)):
+        rank_path = osp.join(feature_path, f'rank={r}_x.pt')
+        torch.save(
+            f.clone(),
+            rank_path,
+        )
+    for (r, f) in enumerate(torch.tensor_split(data.y, world_size)):
+        rank_path = osp.join(feature_path, f'rank={r}_y.pt')
+        torch.save(
+            f.clone(),
+            rank_path,
+        )
+
+    os.makedirs(label_path, exist_ok=True)
+    for (d, i) in split_idx.items():
+        i_parts = torch.tensor_split(i, world_size)
+        for r, i_part in enumerate(i_parts):
+            rank_path = osp.join(label_path, f'rank={r}')
+            os.makedirs(rank_path, exist_ok=True)
+            torch.save(i_part, osp.join(rank_path, f'{d}.pt'))
+
+    meta = dict(
+        num_classes=int(dataset.num_classes),
+        num_features=int(dataset.num_features),
+        num_nodes=int(data.num_nodes),
+    )
+    with open(meta_path, 'w') as f:
+        json.dump(meta, f)
+
+
+def load_partitioned_data(rank, edge_path, feature_path, label_path, meta_path,
+                          wg_mem_type):
+    from cugraph_pyg.data import GraphStore, WholeFeatureStore
+
+    graph_store = GraphStore(is_multi_gpu=True)
+    feature_store = WholeFeatureStore(memory_type=wg_mem_type)
+
+    with open(meta_path) as f:
+        meta = json.load(f)
+
+    split_idx = {}
+    for split in ['train', 'test', 'valid']:
+        path = osp.join(label_path, f'rank={rank}', f'{split}.pt')
+        split_idx[split] = torch.load(path)
+
+    path = osp.join(feature_path, f'rank={rank}_x.pt')
+    feature_store['node', 'x'] = torch.load(path)
+    path = osp.join(feature_path, f'rank={rank}_y.pt')
+    feature_store['node', 'y'] = torch.load(path)
+
+    eix = torch.load(osp.join(edge_path, f'rank={rank}.pt'))
+    graph_store[dict(
+        edge_type=('node', 'rel', 'node'),
+        layout='coo',
+        is_sorted=False,
+        size=(meta['num_nodes'], meta['num_nodes']),
+    )] = eix
+
+    return (feature_store, graph_store), split_idx, meta
+
+
+def run(global_rank, data, split_idx, world_size, device, model, epochs,
+        batch_size, fan_out, num_classes, wall_clock_start, tempdir=None,
+        num_layers=3):
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01,
+                                 weight_decay=0.0005)
+
+    kwargs = dict(
+        num_neighbors=[fan_out] * num_layers,
+        batch_size=batch_size,
+    )
+    from cugraph_pyg.loader import NeighborLoader
+
+    ix_train = split_idx['train'].cuda()
+    train_path = osp.join(tempdir, f'train_{global_rank}')
+    os.mkdir(train_path)
+    train_loader = NeighborLoader(
+        data,
+        input_nodes=ix_train,
+        directory=train_path,
+        shuffle=True,
+        drop_last=True,
+        **kwargs,
+    )
+
+    ix_val = split_idx['valid'].cuda()
+    val_path = osp.join(tempdir, f'val_{global_rank}')
+    os.mkdir(val_path)
+    val_loader = NeighborLoader(
+        data,
+        input_nodes=ix_val,
+        directory=val_path,
+        shuffle=True,
+        drop_last=True,
+        **kwargs,
+    )
+
+    ix_test = split_idx['test'].cuda()
+    test_path = osp.join(tempdir, f'test_{global_rank}')
+    os.mkdir(test_path)
+    test_loader = NeighborLoader(
+        data,
+        input_nodes=ix_test,
+        directory=test_path,
+        shuffle=True,
+        drop_last=True,
+        local_seeds_per_call=80000,
+        **kwargs,
+    )
+
+    dist.barrier()
+
+    eval_steps = 1000
+    warmup_steps = 20
+    dist.barrier()
+    torch.cuda.synchronize()
+
+    if global_rank == 0:
+        prep_time = time.perf_counter() - wall_clock_start
+        print(f"Preparation time: {prep_time:.2f}s")
+
+    for epoch in range(epochs):
+        for i, batch in enumerate(train_loader):
+            if i == warmup_steps:
+                torch.cuda.synchronize()
+                start = time.time()
+
+            batch = batch.to(device)
+            batch_size = batch.batch_size
+
+            batch.y = batch.y.view(-1).to(torch.long)
+            optimizer.zero_grad()
+            out = model(batch.x, batch.edge_index)
+            loss = F.cross_entropy(out[:batch_size], batch.y[:batch_size])
+            loss.backward()
+            optimizer.step()
+            if global_rank == 0 and i % 10 == 0:
+                print(f"Epoch: {epoch:02d}, Iteration: {i}, Loss: {loss:.4f}")
+        nb = i + 1.0
+
+        if global_rank == 0:
+            print(f"Avg Training Iteration Time: "
+                  f"{(time.time() - start) / (nb - warmup_steps):.4f} s/iter")
+
+        with torch.no_grad():
+            total_correct = total_examples = 0
+            for i, batch in enumerate(val_loader):
+                if i >= eval_steps:
+                    break
+
+                batch = batch.to(device)
+                batch_size = batch.batch_size
+
+                batch.y = batch.y.to(torch.long)
+                out = model(batch.x, batch.edge_index)[:batch_size]
+
+                pred = out.argmax(dim=-1)
+                y = batch.y[:batch_size].view(-1).to(torch.long)
+
+                total_correct += int((pred == y).sum())
+                total_examples += y.size(0)
+
+            acc_val = total_correct / total_examples
+            if global_rank == 0:
+                print(f"Validation Accuracy: {acc_val * 100:.2f}%", )
+
+        torch.cuda.synchronize()
+
+    with torch.no_grad():
+        total_correct = total_examples = 0
+        for i, batch in enumerate(test_loader):
+            batch = batch.to(device)
+            batch_size = batch.batch_size
+
+            batch.y = batch.y.to(torch.long)
+            out = model(batch.x, batch.edge_index)[:batch_size]
+
+            pred = out.argmax(dim=-1)
+            y = batch.y[:batch_size].view(-1).to(torch.long)
+
+            total_correct += int((pred == y).sum())
+            total_examples += y.size(0)
+
+        acc_test = total_correct / total_examples
+        if global_rank == 0:
+            print(f"Test Accuracy: {acc_test * 100:.2f}%", )
+
+    if global_rank == 0:
+        total_time = time.perf_counter() - wall_clock_start
+        print(f"Total Training Runtime: {total_time - prep_time}s")
+        print(f"Total Program Runtime: {total_time}s")
+
+    wm_finalize()
+    cugraph_comms_shutdown()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--hidden_channels', type=int, default=256)
+    parser.add_argument('--num_layers', type=int, default=2)
+    parser.add_argument('--lr', type=float, default=0.001)
+    parser.add_argument('--epochs', type=int, default=4)
+    parser.add_argument('--batch_size', type=int, default=1024)
+    parser.add_argument('--fan_out', type=int, default=30)
+    parser.add_argument('--dataset', type=str, default='ogbn-papers100M')
+    parser.add_argument('--root', type=str, default='dataset')
+    parser.add_argument('--tempdir_root', type=str, default=None)
+    parser.add_argument('--skip_partition', action='store_true')
+    parser.add_argument('--wg_mem_type', type=str, default='distributed')
+
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    wall_clock_start = time.perf_counter()
+
+    # Set a very high timeout so that PyTorch does not crash while
+    # partitioning the data.
+    dist.init_process_group('nccl', timeout=timedelta(minutes=60))
+    world_size = dist.get_world_size()
+    assert dist.is_initialized()
+
+    global_rank = dist.get_rank()
+    local_rank = int(os.environ['LOCAL_RANK'])
+    device = torch.device(local_rank)
+
+    print(
+        f'Global rank: {global_rank},',
+        f'Local Rank: {local_rank},',
+        f'World size: {world_size}',
+    )
+
+    # Create the uid needed for cuGraph comms
+    if global_rank == 0:
+        cugraph_id = [cugraph_comms_create_unique_id()]
+    else:
+        cugraph_id = [None]
+    dist.broadcast_object_list(cugraph_id, src=0, device=device)
+    cugraph_id = cugraph_id[0]
+
+    init_pytorch_worker(global_rank, local_rank, world_size, cugraph_id)
+
+    edge_path = osp.join(args.root, f'{args.dataset}_eix_part')
+    feature_path = osp.join(args.root, f'{args.dataset}_fea_part')
+    label_path = osp.join(args.root, f'{args.dataset}_label_part')
+    meta_path = osp.join(args.root, f'{args.dataset}_meta.json')
+
+    # We partition the data to avoid loading it in every worker, which will
+    # waste memory and can lead to an out of memory exception.
+    # cugraph_pyg.GraphStore and cugraph_pyg.WholeFeatureStore are always
+    # constructed from partitions of the edge index and features, respectively,
+    # so this works well.
+    if not args.skip_partition and global_rank == 0:
+        print("Partitioning the data into equal size parts per worker")
+        dataset = PygNodePropPredDataset(name=args.dataset, root=args.root)
+        split_idx = dataset.get_idx_split()
+
+        partition_data(
+            dataset,
+            split_idx,
+            meta_path=meta_path,
+            label_path=label_path,
+            feature_path=feature_path,
+            edge_path=edge_path,
+        )
+
+    dist.barrier()
+    print("Loading partitioned data")
+    data, split_idx, meta = load_partitioned_data(
+        rank=global_rank,
+        edge_path=edge_path,
+        feature_path=feature_path,
+        label_path=label_path,
+        meta_path=meta_path,
+        wg_mem_type=args.wg_mem_type,
+    )
+    dist.barrier()
+
+    model = torch_geometric.nn.models.GCN(
+        meta['num_features'],
+        args.hidden_channels,
+        args.num_layers,
+        meta['num_classes'],
+    ).to(device)
+    model = DistributedDataParallel(model, device_ids=[local_rank])
+
+    with tempfile.TemporaryDirectory(dir=args.tempdir_root) as tempdir:
+        run(global_rank, data, split_idx, world_size, device, model,
+            args.epochs, args.batch_size, args.fan_out, meta['num_classes'],
+            wall_clock_start, tempdir, args.num_layers)
diff --git a/examples/ogbn_papers_100m_cugraph.py b/examples/ogbn_papers_100m_cugraph.py
index c9df4e07e..113a07328 100644
--- a/examples/ogbn_papers_100m_cugraph.py
+++ b/examples/ogbn_papers_100m_cugraph.py
@@ -1,7 +1,7 @@
 import argparse
 import os
-import time
-from typing import Optional
+import os.path as osp
+import tempfile
 
 import cupy
 import rmm
@@ -16,148 +16,171 @@ rmm.reinitialize(devices=[0], pool_allocator=True, managed_memory=True)
 cupy.cuda.set_allocator(rmm_cupy_allocator)
 torch.cuda.memory.change_current_allocator(rmm_torch_allocator)
 
-import cugraph  # noqa
+import cugraph_pyg  # noqa
 import torch.nn.functional as F  # noqa
+# Enable cudf spilling to save gpu memory
 from cugraph.testing.mg_utils import enable_spilling  # noqa
-from cugraph_pyg.data import CuGraphStore  # noqa
-from cugraph_pyg.loader import CuGraphNeighborLoader  # noqa
+from cugraph_pyg.loader import NeighborLoader  # noqa
 
-import torch_geometric  # noqa
-from torch_geometric.loader import NeighborLoader  # noqa
+enable_spilling()
+
+from ogb.nodeproppred import PygNodePropPredDataset  # noqa
+from tqdm import tqdm  # noqa
+
+from torch_geometric.nn import SAGEConv  # noqa
+from torch_geometric.utils import to_undirected  # noqa
 
 parser = argparse.ArgumentParser()
-parser.add_argument('--hidden_channels', type=int, default=256)
-parser.add_argument('--num_layers', type=int, default=2)
-parser.add_argument('--lr', type=float, default=0.001)
-parser.add_argument('--epochs', type=int, default=20)
+parser.add_argument('--epochs', type=int, default=10)
+parser.add_argument('--num_layers', type=int, default=3)
 parser.add_argument('--batch_size', type=int, default=1024)
-parser.add_argument('--fan_out', type=int, default=30)
-parser.add_argument(
-    "--use_gat_conv",
-    action='store_true',
-    help="Wether or not to use GATConv. (Defaults to using GCNConv)",
-)
-parser.add_argument(
-    "--n_gat_conv_heads",
-    type=int,
-    default=4,
-    help="If using GATConv, number of attention heads to use",
-)
+parser.add_argument('--num_neighbors', type=int, default=10)
+parser.add_argument('--channels', type=int, default=256)
+parser.add_argument('--lr', type=float, default=0.003)
+parser.add_argument('--dropout', type=float, default=0.5)
+parser.add_argument('--num_workers', type=int, default=12)
+parser.add_argument('--root', type=str, default=None)
+parser.add_argument('--tempdir_root', type=str, default=None)
 args = parser.parse_args()
-wall_clock_start = time.perf_counter()
-
-
-def get_num_workers() -> int:
-    try:
-        return len(os.sched_getaffinity(0)) // 2
-    except Exception:
-        return os.cpu_count() // 2
-
 
-kwargs = dict(
-    num_neighbors=[args.fan_out] * args.num_layers,
-    batch_size=args.batch_size,
-)
-# Set Up Neighbor Loading
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+root = args.root
+if root is None:
+    root = osp.dirname(osp.realpath(__file__))
+    root = osp.join(root, '..', 'data', 'papers100')
 
-enable_spilling()
-
-from ogb.nodeproppred import PygNodePropPredDataset  # noqa
-
-dataset = PygNodePropPredDataset(name='ogbn-papers100M',
-                                 root='/datasets/ogb_datasets')
+dataset = PygNodePropPredDataset('ogbn-papers100M', root)
 split_idx = dataset.get_idx_split()
-data = dataset[0]
 
-G = {("N", "E", "N"): data.edge_index}
-N = {"N": data.num_nodes}
-fs = cugraph.gnn.FeatureStore(backend="torch")
-fs.add_data(data.x, "N", "x")
-fs.add_data(data.y, "N", "y")
-cugraph_store = CuGraphStore(fs, G, N)
-train_loader = CuGraphNeighborLoader(cugraph_store,
-                                     input_nodes=split_idx['train'],
-                                     shuffle=True, drop_last=True, **kwargs)
-val_loader = CuGraphNeighborLoader(cugraph_store,
-                                   input_nodes=split_idx['valid'], **kwargs)
-test_loader = CuGraphNeighborLoader(cugraph_store,
-                                    input_nodes=split_idx['test'], **kwargs)
-
-if args.use_gat_conv:
-    model = torch_geometric.nn.models.GAT(
-        dataset.num_features, args.hidden_channels, args.num_layers,
-        dataset.num_classes, heads=args.n_gat_conv_heads).to(device)
-else:
-    model = torch_geometric.nn.models.GCN(
-        dataset.num_features,
-        args.hidden_channels,
-        args.num_layers,
-        dataset.num_classes,
-    ).to(device)
-
-optimizer = torch.optim.Adam(model.parameters(), lr=args.lr,
-                             weight_decay=0.0005)
-
-warmup_steps = 20
-
-
-def train():
-    model.train()
-    for i, batch in enumerate(train_loader):
-        batch = batch.to_homogeneous()
-
-        if i == warmup_steps:
-            torch.cuda.synchronize()
-            start_avg_time = time.perf_counter()
-        batch = batch.to(device)
-        optimizer.zero_grad()
-        batch_size = batch.num_sampled_nodes[0]
-        out = model(batch.x, batch.edge_index)[:batch_size]
-        y = batch.y[:batch_size].view(-1).to(torch.long)
-        loss = F.cross_entropy(out, y)
-        loss.backward()
-        optimizer.step()
-
-        if i % 10 == 0:
-            print(f'Epoch: {epoch:02d}, Iteration: {i}, Loss: {loss:.4f}')
-    torch.cuda.synchronize()
-    print(f'Average Training Iteration Time (s/iter): \
-            {(time.perf_counter() - start_avg_time)/(i-warmup_steps):.6f}')
-
-
-@torch.no_grad()
-def test(loader: NeighborLoader, val_steps: Optional[int] = None):
-    model.eval()
-
-    total_correct = total_examples = 0
-    for i, batch in enumerate(loader):
-        if val_steps is not None and i >= val_steps:
-            break
-        batch = batch.to_homogeneous()
-        batch = batch.to(device)
-        batch_size = batch.num_sampled_nodes[0]
-        out = model(batch.x, batch.edge_index)[:batch_size]
-        pred = out.argmax(dim=-1)
-        y = batch.y[:batch_size].view(-1).to(torch.long)
-
-        total_correct += int((pred == y).sum())
-        total_examples += y.size(0)
-
-    return total_correct / total_examples
-
-
-torch.cuda.synchronize()
-prep_time = round(time.perf_counter() - wall_clock_start, 2)
-print("Total time before training begins (prep_time)=", prep_time, "seconds")
-print("Beginning training...")
-for epoch in range(1, 1 + args.epochs):
-    train()
-    val_acc = test(val_loader, val_steps=100)
-    print(f'Val Acc: ~{val_acc:.4f}')
-
-test_acc = test(test_loader)
-print(f'Test Acc: {test_acc:.4f}')
-total_time = round(time.perf_counter() - wall_clock_start, 2)
-print("Total Program Runtime (total_time) =", total_time, "seconds")
-print("total_time - prep_time =", total_time - prep_time, "seconds")
+data = dataset[0]
+data.edge_index = to_undirected(data.edge_index, reduce="mean")
+
+graph_store = cugraph_pyg.data.GraphStore()
+graph_store[dict(
+    edge_type=('node', 'rel', 'node'),
+    layout='coo',
+    is_sorted=False,
+    size=(data.num_nodes, data.num_nodes),
+)] = data.edge_index
+
+feature_store = cugraph_pyg.data.TensorDictFeatureStore()
+feature_store['node', 'x'] = data.x
+feature_store['node', 'y'] = data.y
+
+data = (feature_store, graph_store)
+
+
+class SAGE(torch.nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+
+        self.convs = torch.nn.ModuleList()
+        self.convs.append(SAGEConv(in_channels, args.channels))
+        for _ in range(args.num_layers - 2):
+            self.convs.append(SAGEConv(args.channels, args.channels))
+        self.convs.append(SAGEConv(args.channels, out_channels))
+
+    def forward(self, x, edge_index):
+        for i, conv in enumerate(self.convs):
+            x = conv(x, edge_index)
+            if i != args.num_layers - 1:
+                x = x.relu()
+                x = F.dropout(x, p=args.dropout, training=self.training)
+        return x
+
+
+model = SAGE(dataset.num_features, dataset.num_classes).cuda()
+optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+
+
+def create_loader(
+    data,
+    num_neighbors,
+    input_nodes,
+    replace,
+    batch_size,
+    samples_dir,
+    stage_name,
+    shuffle=False,
+):
+    directory = osp.join(samples_dir, stage_name)
+    os.mkdir(directory)
+    return NeighborLoader(
+        data,
+        num_neighbors=num_neighbors,
+        input_nodes=input_nodes,
+        replace=replace,
+        batch_size=batch_size,
+        directory=directory,
+        shuffle=shuffle,
+    )
+
+
+with tempfile.TemporaryDirectory(dir=args.tempdir_root) as samples_dir:
+    loader_kwargs = dict(
+        data=data,
+        num_neighbors=[args.num_neighbors] * args.num_layers,
+        replace=False,
+        batch_size=args.batch_size,
+        samples_dir=samples_dir,
+    )
+
+    train_loader = create_loader(
+        input_nodes=split_idx['train'],
+        stage_name='train',
+        shuffle=True,
+        **loader_kwargs,
+    )
+
+    val_loader = create_loader(
+        input_nodes=split_idx['valid'],
+        stage_name='val',
+        **loader_kwargs,
+    )
+
+    test_loader = create_loader(
+        input_nodes=split_idx['test'],
+        stage_name='test',
+        **loader_kwargs,
+    )
+
+    def train():
+        model.train()
+
+        total_loss = total_correct = total_examples = 0
+        for batch in tqdm(train_loader):
+            batch = batch.cuda()
+            optimizer.zero_grad()
+            out = model(batch.x, batch.edge_index)[:batch.batch_size]
+            y = batch.y[:batch.batch_size].view(-1).to(torch.long)
+            loss = F.cross_entropy(out, y)
+            loss.backward()
+            optimizer.step()
+
+            total_loss += float(loss) * y.size(0)
+            total_correct += int(out.argmax(dim=-1).eq(y).sum())
+            total_examples += y.size(0)
+
+        return total_loss / total_examples, total_correct / total_examples
+
+    @torch.no_grad()
+    def test(loader):
+        model.eval()
+
+        total_correct = total_examples = 0
+        for batch in tqdm(loader):
+            batch = batch.cuda()
+            out = model(batch.x, batch.edge_index)[:batch.batch_size]
+            y = batch.y[:batch.batch_size].view(-1).to(torch.long)
+
+            total_correct += int(out.argmax(dim=-1).eq(y).sum())
+            total_examples += y.size(0)
+
+        return total_correct / total_examples
+
+    for epoch in range(1, args.epochs + 1):
+        loss, train_acc = train()
+        print(f'Epoch {epoch:02d}, Loss: {loss:.4f}, Train: {train_acc:.4f}')
+        val_acc = test(val_loader)
+        print(f'Epoch {epoch:02d}, Val: {val_acc:.4f}')
+        test_acc = test(test_loader)
+        print(f'Epoch {epoch:02d}, Test: {test_acc:.4f}')
diff --git a/examples/ppi.py b/examples/ppi.py
index 170294f96..49fc3d2ae 100644
--- a/examples/ppi.py
+++ b/examples/ppi.py
@@ -21,18 +21,16 @@ test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False)
 class Net(torch.nn.Module):
     def __init__(self):
         super().__init__()
-        self.conv1 = GATConv(train_dataset.num_features, 256, heads=4)
-        self.lin1 = torch.nn.Linear(train_dataset.num_features, 4 * 256)
-        self.conv2 = GATConv(4 * 256, 256, heads=4)
-        self.lin2 = torch.nn.Linear(4 * 256, 4 * 256)
+        self.conv1 = GATConv(train_dataset.num_features, 256, heads=4,
+                             residual=True)
+        self.conv2 = GATConv(4 * 256, 256, heads=4, residual=True)
         self.conv3 = GATConv(4 * 256, train_dataset.num_classes, heads=6,
-                             concat=False)
-        self.lin3 = torch.nn.Linear(4 * 256, train_dataset.num_classes)
+                             concat=False, residual=True)
 
     def forward(self, x, edge_index):
-        x = F.elu(self.conv1(x, edge_index) + self.lin1(x))
-        x = F.elu(self.conv2(x, edge_index) + self.lin2(x))
-        x = self.conv3(x, edge_index) + self.lin3(x)
+        x = F.elu(self.conv1(x, edge_index))
+        x = F.elu(self.conv2(x, edge_index))
+        x = self.conv3(x, edge_index)
         return x
 
 
diff --git a/examples/triangles_sag_pool.py b/examples/triangles_sag_pool.py
index a3d3898f7..8a6831d5c 100644
--- a/examples/triangles_sag_pool.py
+++ b/examples/triangles_sag_pool.py
@@ -71,7 +71,7 @@ model = Net(dataset.num_features).to(device)
 optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
 
-def train(epoch):
+def train():
     model.train()
 
     total_loss = 0
@@ -101,7 +101,7 @@ def test(loader):
 
 
 for epoch in range(1, 301):
-    loss = train(epoch)
+    loss = train()
     train_correct, train_ratio = test(train_loader)
     val_correct, val_ratio = test(val_loader)
     test_correct, test_ratio = test(test_loader)
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index e3caaecaa..3a1abb5e7 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -1165,7 +1165,7 @@ def size_repr(key: Any, value: Any, indent: int = 0) -> str:
                f'[{value.num_rows}, {value.num_cols}])')
     elif isinstance(value, str):
         out = f"'{value}'"
-    elif isinstance(value, Sequence):
+    elif isinstance(value, (Sequence, set)):
         out = str([len(value)])
     elif isinstance(value, Mapping) and len(value) == 0:
         out = '{}'
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index 17ee648fe..3029d78ea 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -1928,7 +1928,7 @@ def _spmm(
     if transpose and not input.is_sorted_by_col:
         cls_name = input.__class__.__name__
         raise ValueError(f"'matmul(..., transpose=True)' requires "
-                         f"'{cls_name}' to be sorted by colums")
+                         f"'{cls_name}' to be sorted by columns")
 
     if (torch_geometric.typing.WITH_TORCH_SPARSE and not is_compiling()
             and other.is_cuda):  # pragma: no cover
diff --git a/torch_geometric/explain/explanation.py b/torch_geometric/explain/explanation.py
index 72ef79bde..2aca9b280 100644
--- a/torch_geometric/explain/explanation.py
+++ b/torch_geometric/explain/explanation.py
@@ -340,10 +340,10 @@ class HeteroExplanation(HeteroData, ExplanationMixin):
         """
         node_mask_dict = self.node_mask_dict
         for node_mask in node_mask_dict.values():
-            if node_mask.dim() != 2 or node_mask.size(1) <= 1:
+            if node_mask.dim() != 2:
                 raise ValueError(f"Cannot compute feature importance for "
                                  f"object-level 'node_mask' "
-                                 f"(got shape {node_mask_dict.size()})")
+                                 f"(got shape {node_mask.size()})")
 
         if feat_labels is None:
             feat_labels = {}
diff --git a/torch_geometric/inspector.py b/torch_geometric/inspector.py
index 8f6e7bd04..32b4421f7 100644
--- a/torch_geometric/inspector.py
+++ b/torch_geometric/inspector.py
@@ -305,7 +305,7 @@ class Inspector:
         according to its function signature from a data blob.
 
         Args:
-            func (callabel or str): The function.
+            func (callable or str): The function.
             kwargs (dict[str, Any]): The data blob which may serve as inputs.
         """
         out_dict: Dict[str, Any] = {}
@@ -346,7 +346,7 @@ class Inspector:
             type annotations are not found.
 
         Args:
-            func (callabel or str): The function.
+            func (callable or str): The function.
             exclude (list[int or str]): A list of parameters to exclude, either
                 given by their name or index. (default: :obj:`None`)
         """
diff --git a/torch_geometric/nn/conv/gat_conv.py b/torch_geometric/nn/conv/gat_conv.py
index 8c7a6f34b..720dfb098 100644
--- a/torch_geometric/nn/conv/gat_conv.py
+++ b/torch_geometric/nn/conv/gat_conv.py
@@ -107,6 +107,8 @@ class GATConv(MessagePassing):
             :obj:`"min"`, :obj:`"max"`, :obj:`"mul"`). (default: :obj:`"mean"`)
         bias (bool, optional): If set to :obj:`False`, the layer will not learn
             an additive bias. (default: :obj:`True`)
+        residual (bool, optional): If set to :obj:`True`, the layer will add
+            a learnable skip-connection. (default: :obj:`False`)
         **kwargs (optional): Additional arguments of
             :class:`torch_geometric.nn.conv.MessagePassing`.
 
@@ -137,6 +139,7 @@ class GATConv(MessagePassing):
         edge_dim: Optional[int] = None,
         fill_value: Union[float, Tensor, str] = 'mean',
         bias: bool = True,
+        residual: bool = False,
         **kwargs,
     ):
         kwargs.setdefault('aggr', 'add')
@@ -151,6 +154,7 @@ class GATConv(MessagePassing):
         self.add_self_loops = add_self_loops
         self.edge_dim = edge_dim
         self.fill_value = fill_value
+        self.residual = residual
 
         # In case we are operating in bipartite graphs, we apply separate
         # transformations 'lin_src' and 'lin_dst' to source and target nodes:
@@ -176,10 +180,22 @@ class GATConv(MessagePassing):
             self.lin_edge = None
             self.register_parameter('att_edge', None)
 
-        if bias and concat:
-            self.bias = Parameter(torch.empty(heads * out_channels))
-        elif bias and not concat:
-            self.bias = Parameter(torch.empty(out_channels))
+        # The number of output channels:
+        total_out_channels = out_channels * (heads if concat else 1)
+
+        if residual:
+            self.res = Linear(
+                in_channels
+                if isinstance(in_channels, int) else in_channels[1],
+                total_out_channels,
+                bias=False,
+                weight_initializer='glorot',
+            )
+        else:
+            self.register_parameter('res', None)
+
+        if bias:
+            self.bias = Parameter(torch.empty(total_out_channels))
         else:
             self.register_parameter('bias', None)
 
@@ -195,6 +211,8 @@ class GATConv(MessagePassing):
             self.lin_dst.reset_parameters()
         if self.lin_edge is not None:
             self.lin_edge.reset_parameters()
+        if self.res is not None:
+            self.res.reset_parameters()
         glorot(self.att_src)
         glorot(self.att_dst)
         glorot(self.att_edge)
@@ -270,11 +288,16 @@ class GATConv(MessagePassing):
 
         H, C = self.heads, self.out_channels
 
+        res: Optional[Tensor] = None
+
         # We first transform the input node features. If a tuple is passed, we
         # transform source and target node features via separate weights:
         if isinstance(x, Tensor):
             assert x.dim() == 2, "Static graphs not supported in 'GATConv'"
 
+            if self.res is not None:
+                res = self.res(x)
+
             if self.lin is not None:
                 x_src = x_dst = self.lin(x).view(-1, H, C)
             else:
@@ -288,6 +311,9 @@ class GATConv(MessagePassing):
             x_src, x_dst = x
             assert x_src.dim() == 2, "Static graphs not supported in 'GATConv'"
 
+            if x_dst is not None and self.res is not None:
+                res = self.res(x_dst)
+
             if self.lin is not None:
                 # If the module is initialized as non-bipartite, we expect that
                 # source and destination node features have the same shape and
@@ -344,6 +370,9 @@ class GATConv(MessagePassing):
         else:
             out = out.mean(dim=1)
 
+        if res is not None:
+            out = out + res
+
         if self.bias is not None:
             out = out + self.bias
 
diff --git a/torch_geometric/nn/conv/gatv2_conv.py b/torch_geometric/nn/conv/gatv2_conv.py
index 1ce8ff21b..f3b2f4937 100644
--- a/torch_geometric/nn/conv/gatv2_conv.py
+++ b/torch_geometric/nn/conv/gatv2_conv.py
@@ -110,6 +110,8 @@ class GATv2Conv(MessagePassing):
             will be applied to the source and the target node of every edge,
             *i.e.* :math:`\mathbf{\Theta}_{s} = \mathbf{\Theta}_{t}`.
             (default: :obj:`False`)
+        residual (bool, optional): If set to :obj:`True`, the layer will add
+            a learnable skip-connection. (default: :obj:`False`)
         **kwargs (optional): Additional arguments of
             :class:`torch_geometric.nn.conv.MessagePassing`.
 
@@ -141,6 +143,7 @@ class GATv2Conv(MessagePassing):
         fill_value: Union[float, Tensor, str] = 'mean',
         bias: bool = True,
         share_weights: bool = False,
+        residual: bool = False,
         **kwargs,
     ):
         super().__init__(node_dim=0, **kwargs)
@@ -154,6 +157,7 @@ class GATv2Conv(MessagePassing):
         self.add_self_loops = add_self_loops
         self.edge_dim = edge_dim
         self.fill_value = fill_value
+        self.residual = residual
         self.share_weights = share_weights
 
         if isinstance(in_channels, int):
@@ -181,10 +185,22 @@ class GATv2Conv(MessagePassing):
         else:
             self.lin_edge = None
 
-        if bias and concat:
-            self.bias = Parameter(torch.empty(heads * out_channels))
-        elif bias and not concat:
-            self.bias = Parameter(torch.empty(out_channels))
+        # The number of output channels:
+        total_out_channels = out_channels * (heads if concat else 1)
+
+        if residual:
+            self.res = Linear(
+                in_channels
+                if isinstance(in_channels, int) else in_channels[1],
+                total_out_channels,
+                bias=False,
+                weight_initializer='glorot',
+            )
+        else:
+            self.register_parameter('res', None)
+
+        if bias:
+            self.bias = Parameter(torch.empty(total_out_channels))
         else:
             self.register_parameter('bias', None)
 
@@ -196,6 +212,8 @@ class GATv2Conv(MessagePassing):
         self.lin_r.reset_parameters()
         if self.lin_edge is not None:
             self.lin_edge.reset_parameters()
+        if self.res is not None:
+            self.res.reset_parameters()
         glorot(self.att)
         zeros(self.bias)
 
@@ -255,10 +273,16 @@ class GATv2Conv(MessagePassing):
         """
         H, C = self.heads, self.out_channels
 
+        res: Optional[Tensor] = None
+
         x_l: OptTensor = None
         x_r: OptTensor = None
         if isinstance(x, Tensor):
             assert x.dim() == 2
+
+            if self.res is not None:
+                res = self.res(x)
+
             x_l = self.lin_l(x).view(-1, H, C)
             if self.share_weights:
                 x_r = x_l
@@ -267,6 +291,10 @@ class GATv2Conv(MessagePassing):
         else:
             x_l, x_r = x[0], x[1]
             assert x[0].dim() == 2
+
+            if x_r is not None and self.res is not None:
+                res = self.res(x_r)
+
             x_l = self.lin_l(x_l).view(-1, H, C)
             if x_r is not None:
                 x_r = self.lin_r(x_r).view(-1, H, C)
@@ -305,6 +333,9 @@ class GATv2Conv(MessagePassing):
         else:
             out = out.mean(dim=1)
 
+        if res is not None:
+            out = out + res
+
         if self.bias is not None:
             out = out + self.bias
 
diff --git a/torch_geometric/nn/conv/hetero_conv.py b/torch_geometric/nn/conv/hetero_conv.py
index 2975854f0..29c6150ae 100644
--- a/torch_geometric/nn/conv/hetero_conv.py
+++ b/torch_geometric/nn/conv/hetero_conv.py
@@ -102,7 +102,7 @@ class HeteroConv(torch.nn.Module):
                 individual edge type, either as a :class:`torch.Tensor` of
                 shape :obj:`[2, num_edges]` or a
                 :class:`torch_sparse.SparseTensor`.
-            *args_dict (optional): Additional forward arguments of invididual
+            *args_dict (optional): Additional forward arguments of individual
                 :class:`torch_geometric.nn.conv.MessagePassing` layers.
             **kwargs_dict (optional): Additional forward arguments of
                 individual :class:`torch_geometric.nn.conv.MessagePassing`
diff --git a/torch_geometric/nn/conv/mixhop_conv.py b/torch_geometric/nn/conv/mixhop_conv.py
index c8e0be466..bc51a8550 100644
--- a/torch_geometric/nn/conv/mixhop_conv.py
+++ b/torch_geometric/nn/conv/mixhop_conv.py
@@ -14,7 +14,7 @@ from torch_geometric.utils import spmm
 
 class MixHopConv(MessagePassing):
     r"""The Mix-Hop graph convolutional operator from the
-    `"MixHop: Higher-Order Graph Convolutional Architecturesvia Sparsified
+    `"MixHop: Higher-Order Graph Convolutional Architectures via Sparsified
     Neighborhood Mixing" <https://arxiv.org/abs/1905.00067>`_ paper.
 
     .. math::
diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
index 6225fde7d..564b13161 100644
--- a/torch_geometric/nn/dense/linear.py
+++ b/torch_geometric/nn/dense/linear.py
@@ -58,7 +58,7 @@ def reset_bias_(bias: Optional[Tensor], in_channels: int,
 
 
 class Linear(torch.nn.Module):
-    r"""Applies a linear tranformation to the incoming data.
+    r"""Applies a linear transformation to the incoming data.
 
     .. math::
         \mathbf{x}^{\prime} = \mathbf{x} \mathbf{W}^{\top} + \mathbf{b}
@@ -192,7 +192,7 @@ class Linear(torch.nn.Module):
 
 
 class HeteroLinear(torch.nn.Module):
-    r"""Applies separate linear tranformations to the incoming data according
+    r"""Applies separate linear transformations to the incoming data according
     to types.
 
     For type :math:`\kappa`, it computes
@@ -365,7 +365,8 @@ class HeteroLinear(torch.nn.Module):
 
 
 class HeteroDictLinear(torch.nn.Module):
-    r"""Applies separate linear tranformations to the incoming data dictionary.
+    r"""Applies separate linear transformations to the incoming data
+    dictionary.
 
     For key :math:`\kappa`, it computes
 
diff --git a/torch_geometric/nn/fx.py b/torch_geometric/nn/fx.py
index d1141e44c..6a266bdd5 100644
--- a/torch_geometric/nn/fx.py
+++ b/torch_geometric/nn/fx.py
@@ -18,8 +18,8 @@ class Transformer:
     :class:`~torch.nn.Module`.
     :class:`Transformer` works entirely symbolically.
 
-    Methods in the :class:`Transformer` class can be overriden to customize the
-    behavior of transformation.
+    Methods in the :class:`Transformer` class can be overridden to customize
+    the behavior of transformation.
 
     .. code-block:: none
 
@@ -283,7 +283,7 @@ def symbolic_trace(
             # TODO We currently only trace top-level modules.
             return not isinstance(module, torch.nn.Sequential)
 
-        # Note: This is a hack around the fact that `Aggregaton.__call__`
+        # Note: This is a hack around the fact that `Aggregation.__call__`
         # is not patched by the base implementation of `trace`.
         # see https://github.com/pyg-team/pytorch_geometric/pull/5021 for
         # details on the rationale
diff --git a/torch_geometric/nn/nlp/sentence_transformer.py b/torch_geometric/nn/nlp/sentence_transformer.py
index cfb26176f..b1ae8254c 100644
--- a/torch_geometric/nn/nlp/sentence_transformer.py
+++ b/torch_geometric/nn/nlp/sentence_transformer.py
@@ -27,6 +27,8 @@ class SentenceTransformer(torch.nn.Module):
 
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.model = AutoModel.from_pretrained(model_name)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
 
     def forward(self, input_ids: Tensor, attention_mask: Tensor) -> Tensor:
         out = self.model(input_ids=input_ids, attention_mask=attention_mask)
diff --git a/torch_geometric/nn/summary.py b/torch_geometric/nn/summary.py
index 8256bacc9..a8aa2e472 100644
--- a/torch_geometric/nn/summary.py
+++ b/torch_geometric/nn/summary.py
@@ -141,7 +141,7 @@ def get_shape(inputs: Any) -> str:
 def postprocess(info_list: List[dict]) -> List[dict]:
     for idx, info in enumerate(info_list):
         depth = info['depth']
-        if idx > 0:  # root module (0) is exclued
+        if idx > 0:  # root module (0) is excluded
             if depth == 1:
                 prefix = '├─'
             else:
diff --git a/torch_geometric/nn/to_hetero_with_bases_transformer.py b/torch_geometric/nn/to_hetero_with_bases_transformer.py
index b99ee9e9f..5d2526b34 100644
--- a/torch_geometric/nn/to_hetero_with_bases_transformer.py
+++ b/torch_geometric/nn/to_hetero_with_bases_transformer.py
@@ -468,7 +468,7 @@ def get_edge_type(
 ###############################################################################
 
 # These methods are used to group the individual type-wise components into a
-# unfied single representation.
+# unified single representation.
 
 
 def group_node_placeholder(input_dict: Dict[NodeType, Tensor],
diff --git a/torch_geometric/sampler/neighbor_sampler.py b/torch_geometric/sampler/neighbor_sampler.py
index 758f14819..0ce692206 100644
--- a/torch_geometric/sampler/neighbor_sampler.py
+++ b/torch_geometric/sampler/neighbor_sampler.py
@@ -631,7 +631,7 @@ def edge_sample(
             if edge_label_time is not None:
                 dst_time = edge_label_time.repeat(1 + neg_sampling.amount)
 
-    # Heterogeneus Neighborhood Sampling ######################################
+    # Heterogeneous Neighborhood Sampling #####################################
 
     if input_type is not None:
         seed_time_dict = None
@@ -724,7 +724,7 @@ def edge_sample(
                 src_time,
             )
 
-    # Homogeneus Neighborhood Sampling ########################################
+    # Homogeneous Neighborhood Sampling #######################################
 
     else:
 
diff --git a/torch_geometric/transforms/add_metapaths.py b/torch_geometric/transforms/add_metapaths.py
index 9aca179fc..366a09644 100644
--- a/torch_geometric/transforms/add_metapaths.py
+++ b/torch_geometric/transforms/add_metapaths.py
@@ -37,7 +37,7 @@ class AddMetaPaths(BaseTransform):
     :class:`~torch_geometric.data.HeteroData` object as edge type
     :obj:`(src_node_type, "metapath_*", dst_node_type)`, where
     :obj:`src_node_type` and :obj:`dst_node_type` denote :math:`\mathcal{V}_1`
-    and :math:`\mathcal{V}_{\ell}`, repectively.
+    and :math:`\mathcal{V}_{\ell}`, respectively.
 
     In addition, a :obj:`metapath_dict` object is added to the
     :class:`~torch_geometric.data.HeteroData` object which maps the
@@ -108,12 +108,12 @@ class AddMetaPaths(BaseTransform):
         **kwargs: bool,
     ) -> None:
         if 'drop_orig_edges' in kwargs:
-            warnings.warn("'drop_orig_edges' is dprecated. Use "
+            warnings.warn("'drop_orig_edges' is deprecated. Use "
                           "'drop_orig_edge_types' instead")
             drop_orig_edge_types = kwargs['drop_orig_edges']
 
         if 'drop_unconnected_nodes' in kwargs:
-            warnings.warn("'drop_unconnected_nodes' is dprecated. Use "
+            warnings.warn("'drop_unconnected_nodes' is deprecated. Use "
                           "'drop_unconnected_node_types' instead")
             drop_unconnected_node_types = kwargs['drop_unconnected_nodes']
 
@@ -231,7 +231,7 @@ class AddRandomMetaPaths(BaseTransform):
             will drop node types not connected by any edge type.
             (default: :obj:`False`)
         walks_per_node (int, List[int], optional): The number of random walks
-            for each starting node in a metapth. (default: :obj:`1`)
+            for each starting node in a metapath. (default: :obj:`1`)
         sample_ratio (float, optional): The ratio of source nodes to start
             random walks from. (default: :obj:`1.0`)
     """
diff --git a/torch_geometric/transforms/node_property_split.py b/torch_geometric/transforms/node_property_split.py
index 1850c7a00..ead02249b 100644
--- a/torch_geometric/transforms/node_property_split.py
+++ b/torch_geometric/transforms/node_property_split.py
@@ -53,7 +53,7 @@ class NodePropertySplit(BaseTransform):
 
         property_name = 'popularity'
         ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
-        tranaform = NodePropertySplit(property_name, ratios)
+        transform = NodePropertySplit(property_name, ratios)
 
         data = transform(data)
     """
diff --git a/torch_geometric/transforms/pad.py b/torch_geometric/transforms/pad.py
index 74e08b50e..90f5c6222 100644
--- a/torch_geometric/transforms/pad.py
+++ b/torch_geometric/transforms/pad.py
@@ -262,7 +262,7 @@ class Pad(BaseTransform):
     All the attributes of node types other than :obj:`v0` and :obj:`v1` are
     padded using a value of :obj:`1.0`.
     All the attributes of the :obj:`('v0', 'e0', 'v1')` edge type are padded
-    usin a value of :obj:`3.5`.
+    using a value of :obj:`3.5`.
     The :obj:`edge_attr` attributes of the :obj:`('v1', 'e0', 'v0')` edge type
     are padded using a value of :obj:`-1.5`, and any other attributes of this
     edge type are padded using a value of :obj:`5.5`.
diff --git a/torch_geometric/transforms/svd_feature_reduction.py b/torch_geometric/transforms/svd_feature_reduction.py
index 624bb5e74..3f65982cf 100644
--- a/torch_geometric/transforms/svd_feature_reduction.py
+++ b/torch_geometric/transforms/svd_feature_reduction.py
@@ -11,7 +11,7 @@ class SVDFeatureReduction(BaseTransform):
     Decomposition (SVD) (functional name: :obj:`svd_feature_reduction`).
 
     Args:
-        out_channels (int): The dimensionlity of node features after
+        out_channels (int): The dimensionality of node features after
             reduction.
     """
     def __init__(self, out_channels: int):
diff --git a/torch_geometric/utils/_negative_sampling.py b/torch_geometric/utils/_negative_sampling.py
index 3deb6bcec..bb93a9385 100644
--- a/torch_geometric/utils/_negative_sampling.py
+++ b/torch_geometric/utils/_negative_sampling.py
@@ -265,7 +265,7 @@ def structured_negative_sampling_feasible(
     :meth:`~torch_geometric.utils.structured_negative_sampling` is feasible
     on the graph given by :obj:`edge_index`.
     :meth:`~torch_geometric.utils.structured_negative_sampling` is infeasible
-    if atleast one node is connected to all other nodes.
+    if at least one node is connected to all other nodes.
 
     Args:
         edge_index (LongTensor): The edge indices.
diff --git a/torch_geometric/utils/_scatter.py b/torch_geometric/utils/_scatter.py
index b4c8f5182..3d02d16a0 100644
--- a/torch_geometric/utils/_scatter.py
+++ b/torch_geometric/utils/_scatter.py
@@ -295,7 +295,7 @@ def group_cat(
     r"""Concatenates the given sequence of tensors :obj:`tensors` in the given
     dimension :obj:`dim`.
     Different from :meth:`torch.cat`, values along the concatenating dimension
-    are grouped according to the indicies defined in the :obj:`index` tensors.
+    are grouped according to the indices defined in the :obj:`index` tensors.
     All tensors must have the same shape (except in the concatenating
     dimension).
 
diff --git a/torch_geometric/utils/_tree_decomposition.py b/torch_geometric/utils/_tree_decomposition.py
index 5f6a7cf0d..14018e569 100644
--- a/torch_geometric/utils/_tree_decomposition.py
+++ b/torch_geometric/utils/_tree_decomposition.py
@@ -64,7 +64,7 @@ def tree_decomposition(
             xs.append(1)
 
     # Generate `atom2cliques` mappings.
-    atom2cliques: List[List[int]] = [[] for i in range(mol.GetNumAtoms())]
+    atom2cliques: List[List[int]] = [[] for _ in range(mol.GetNumAtoms())]
     for c in range(len(cliques)):
         for atom in cliques[c]:
             atom2cliques[atom].append(c)
diff --git a/torch_geometric/utils/augmentation.py b/torch_geometric/utils/augmentation.py
index 63f61d7bc..354fdb3b2 100644
--- a/torch_geometric/utils/augmentation.py
+++ b/torch_geometric/utils/augmentation.py
@@ -12,7 +12,7 @@ def shuffle_node(
     training: bool = True,
 ) -> Tuple[Tensor, Tensor]:
     r"""Randomly shuffle the feature matrix :obj:`x` along the
-    first dimmension.
+    first dimension.
 
     The method returns (1) the shuffled :obj:`x`, (2) the permutation
     indicating the orders of original nodes after shuffling.
diff --git a/torch_geometric/utils/map.py b/torch_geometric/utils/map.py
index 16277e6a8..554bf6109 100644
--- a/torch_geometric/utils/map.py
+++ b/torch_geometric/utils/map.py
@@ -14,7 +14,7 @@ def map_index(
     inclusive: bool = False,
 ) -> Tuple[Tensor, Optional[Tensor]]:
     r"""Maps indices in :obj:`src` to the positional value of their
-    corresponding occurence in :obj:`index`.
+    corresponding occurrence in :obj:`index`.
     Indices must be strictly positive.
 
     Args:


Test patch:
diff --git a/test/data/lightning/test_datamodule.py b/test/data/lightning/test_datamodule.py
index 81d9b9866..799cdaa3a 100644
--- a/test/data/lightning/test_datamodule.py
+++ b/test/data/lightning/test_datamodule.py
@@ -276,11 +276,19 @@ class LinearHeteroNodeModule(LightningModule):
         return torch.optim.Adam(self.parameters(), lr=0.01)
 
 
+@pytest.fixture
+def preserve_context():
+    num_threads = torch.get_num_threads()
+    yield
+    torch.distributed.destroy_process_group()
+    torch.set_num_threads(num_threads)
+
+
 @onlyCUDA
 @onlyFullTest
 @onlyNeighborSampler
 @withPackage('pytorch_lightning>=2.0.0', 'torchmetrics>=0.11.0')
-def test_lightning_hetero_node_data(get_dataset):
+def test_lightning_hetero_node_data(preserve_context, get_dataset):
     import pytorch_lightning as pl
 
     data = get_dataset(name='hetero')[0]
diff --git a/test/nn/conv/test_gat_conv.py b/test/nn/conv/test_gat_conv.py
index 2f30fc11d..6549911ac 100644
--- a/test/nn/conv/test_gat_conv.py
+++ b/test/nn/conv/test_gat_conv.py
@@ -11,13 +11,14 @@ from torch_geometric.typing import Adj, Size, SparseTensor
 from torch_geometric.utils import to_torch_csc_tensor
 
 
-def test_gat_conv():
+@pytest.mark.parametrize('residual', [False, True])
+def test_gat_conv(residual):
     x1 = torch.randn(4, 8)
     x2 = torch.randn(2, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
     adj1 = to_torch_csc_tensor(edge_index, size=(4, 4))
 
-    conv = GATConv(8, 32, heads=2)
+    conv = GATConv(8, 32, heads=2, residual=residual)
     assert str(conv) == 'GATConv(8, 32, heads=2)'
     out = conv(x1, edge_index)
     assert out.size() == (4, 64)
@@ -113,7 +114,7 @@ def test_gat_conv():
     # Test bipartite message passing:
     adj1 = to_torch_csc_tensor(edge_index, size=(4, 2))
 
-    conv = GATConv((8, 16), 32, heads=2)
+    conv = GATConv((8, 16), 32, heads=2, residual=residual)
     assert str(conv) == 'GATConv((8, 16), 32, heads=2)'
 
     out1 = conv((x1, x2), edge_index)
diff --git a/test/nn/conv/test_gatv2_conv.py b/test/nn/conv/test_gatv2_conv.py
index 16d82833c..3bca6530e 100644
--- a/test/nn/conv/test_gatv2_conv.py
+++ b/test/nn/conv/test_gatv2_conv.py
@@ -1,5 +1,6 @@
 from typing import Tuple
 
+import pytest
 import torch
 from torch import Tensor
 
@@ -10,13 +11,14 @@ from torch_geometric.typing import Adj, SparseTensor
 from torch_geometric.utils import to_torch_csc_tensor
 
 
-def test_gatv2_conv():
+@pytest.mark.parametrize('residual', [False, True])
+def test_gatv2_conv(residual):
     x1 = torch.randn(4, 8)
     x2 = torch.randn(2, 8)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
     adj1 = to_torch_csc_tensor(edge_index, size=(4, 4))
 
-    conv = GATv2Conv(8, 32, heads=2)
+    conv = GATv2Conv(8, 32, heads=2, residual=residual)
     assert str(conv) == 'GATv2Conv(8, 32, heads=2)'
     out = conv(x1, edge_index)
     assert out.size() == (4, 64)
diff --git a/test/utils/test_geodesic.py b/test/utils/test_geodesic.py
index c1daf41e2..de31eb512 100644
--- a/test/utils/test_geodesic.py
+++ b/test/utils/test_geodesic.py
@@ -1,6 +1,5 @@
 from math import sqrt
 
-import pytest
 import torch
 
 from torch_geometric.testing import withPackage
@@ -8,7 +7,6 @@ from torch_geometric.utils import geodesic_distance
 
 
 @withPackage('gdist')
-@pytest.mark.skip(reason="No way of currently testing this")
 def test_geodesic_distance():
     pos = torch.tensor([
         [0.0, 0.0, 0.0],
@@ -29,12 +27,12 @@ def test_geodesic_distance():
     assert torch.allclose(out, geodesic_distance(pos, face, num_workers=-1))
 
     out = geodesic_distance(pos, face, norm=False)
-    expected = [
+    expected = torch.tensor([
         [0, 2, 2, 2 * sqrt(2)],
         [2, 0, 2 * sqrt(2), 2],
         [2, 2 * sqrt(2), 0, 2],
         [2 * sqrt(2), 2, 2, 0],
-    ]
+    ])
     assert torch.allclose(out, expected)
 
     src = torch.tensor([0, 0, 0, 0])