Description
Describe the bug
When training a GNN using pytorch, and using the LayerNorm from pytorch, I get the warning
UserWarning: aten::layer_norm: an autograd kernel was not registered to the Autograd key(s) but we are trying to backprop through it. This may lead to silently incorrect behavior. This behavior is deprecated and will be removed in a future version of PyTorch. If your operator is differentiable, please ensure you have registered an autograd kernel to the correct Autograd key (e.g. DispatchKey::Autograd, DispatchKey::CompositeImplicitAutograd). If your operator is not differentiable, or to squash this warning and use the previous behavior, please register torch::CppFunction::makeFallthrough() to DispatchKey::Autograd. (Triggered internally at /pytorch/torch/csrc/autograd/autograd_not_implemented_fallback.cpp:62.)
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
Reproducable example
As using the functional layer_norm could be the issue, you can change it in the init. Gives warning for both.
from __future__ import annotations
import torch
import torch.nn.functional as F
from torch import nn
from torch_geometric.data import Data
from torch_geometric.nn import MessagePassing
class CustomGNNLayer(MessagePassing):
def __init__(self, in_channels, out_channels):
super(CustomGNNLayer, self).__init__(aggr="add")
self.lin = nn.Linear(in_channels, out_channels)
# Use functional layer_norm directly (potential cause of the error)
self.use_functional_norm = False
# Also keep a LayerNorm module for comparison
self.layer_norm = nn.LayerNorm(out_channels)
def forward(self, x, edge_index):
# Propagate through edges
x = self.lin(x)
x = self.propagate(edge_index, x=x)
# Apply normalization - this might be where the error happens
if self.use_functional_norm:
# Using F.layer_norm directly might cause issues with autograd
x = F.layer_norm(x, [x.size(-1)], None, None, 1e-5)
else:
# This should work fine
x = self.layer_norm(x)
return x
def message(self, x_j):
return x_j
def update(self, aggr_out):
return aggr_out
# Create a simple GNN model
class SimpleGNN(nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels):
super().__init__()
self.conv1 = CustomGNNLayer(in_channels, hidden_channels)
self.conv2 = CustomGNNLayer(hidden_channels, out_channels)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.conv1(x, edge_index)
x = F.relu(x)
x = self.conv2(x, edge_index)
return x
# Generate synthetic graph data
def generate_graph_data(num_nodes=10, num_features=16, num_edges=20):
x = torch.randn(num_nodes, num_features)
edge_index = torch.randint(0, num_nodes, (2, num_edges))
data = Data(x=x, edge_index=edge_index)
return data
graph = generate_graph_data()
model = SimpleGNN(in_channels=16, hidden_channels=32, out_channels=8)
y = model(graph)
loss = torch.mean((y - 5) ** 2)
print("Warning occurs here")
_ = loss.backward()
print("Done with graph")
Versions
Collecting environment information...
PyTorch version: 2.6.0+cu124
PyTorch CXX11 ABI: No
IPEX version: 2.6.0+cpu
IPEX commit: 4784d0d
Build type: Release
OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: 10.0.0-4ubuntu1
IGC version: N/A
CMake version: version 3.16.3
Libc version: glibc-2.31
Python version: 3.11.9 | Intel Corporation | (main, Sep 9 2024, 23:42:49) [GCC 14.1.0] (64-bit runtime)
Python platform: Linux-5.4.0-208-generic-x86_64-with-glibc2.31
Is XPU available: False
DPCPP runtime: N/A
MKL version: N/A
GPU models and configuration onboard:
N/A
GPU models and configuration detected:
N/A
Driver version:
- intel_opencl: N/A
- level_zero: N/A
CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Byte Order: Little Endian
Address sizes: 45 bits physical, 48 bits virtual
CPU(s): 20
On-line CPU(s) list: 0-19
Thread(s) per core: 1
Core(s) per socket: 1
Socket(s): 20
NUMA node(s): 1
Vendor ID: GenuineIntel
CPU family: 6
Model: 85
Model name: Intel(R) Xeon(R) Platinum 8260 CPU @ 2.40GHz
Stepping: 7
CPU MHz: 2394.374
BogoMIPS: 4788.74
Hypervisor vendor: VMware
Virtualization type: full
L1d cache: 640 KiB
L1i cache: 640 KiB
L2 cache: 20 MiB
L3 cache: 715 MiB
NUMA node0 CPU(s): 0-19
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit: KVM: Vulnerable
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed: Mitigation; Enhanced IBRS
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI SW loop, KVM SW loop
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon nopl xtopology tsc_reliable nonstop_tsc cpuid pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 invpcid avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat pku ospke avx512_vnni md_clear flush_l1d arch_capabilities
Versions of relevant libraries:
[conda] intel-extension-for-pytorch 2.6.0 pypi_0 pypi
[conda] numpy 2.2.4 pypi_0 pypi
[conda] nvidia-cuda-cupti-cu12 12.4.127 pypi_0 pypi
[conda] nvidia-nccl-cu12 2.21.5 pypi_0 pypi
[conda] oneccl-bind-pt 2.6.0+cpu pypi_0 pypi
[conda] python 3.11.9 h2324612_8_cpython https://software.repos.intel.com/python/conda
[conda] pytorch-lightning 2.5.0.post0 pypi_0 pypi
[conda] torch 2.6.0 pypi_0 pypi
[conda] torch-geometric 2.6.1 pypi_0 pypi
[conda] torchmetrics 1.6.3 pypi_0 pypi