UserWarning: aten::layer_norm has no autograd kernel registered

[jonas-eschle]

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

[ZailiWang]

Hi, would you share how to reproduce the error you met?

[jonas-eschle]

Took me a while but I could recreate a reproducible example, I've added it to the original post

[ZailiWang]

I tried to set up the env on 3 servers with different CPUs, including a similar server (Intel(R) Xeon(R) Gold 6238M CPU) with cuda version torch==2.6.0, as well as ipex==2.6.0. However I don't get the warning message on any of them.

(test) zailiw@xxx:/mnt/zailiw/cases$ python ipex_801_gnn.py
Warning occurs here
Done with graph

(test) zailiw@xxx:/mnt/zailiw/cases$ pip list|grep torch
intel_extension_for_pytorch 2.6.0
torch                       2.6.0+cu126
torch-geometric             2.6.1
torchaudio                  2.6.0+cu126
torchvision                 0.21.0+cu126

May I know how you installed torch? Using pip install directly or from source code build?
And does this warning blocks your training process?

[jonas-eschle]

Thanks a lot, I've also tried on two different machines before. I am not sure what exactly is causing it, I've created a reproducible example here: https://colab.research.google.com/drive/1u_tw4anVRCbAcBoJTm22kWcg4gvyjdLS?usp=sharing

Maybe it's one of the other deps? But this still should not cause torch to emit a warning I think?

It doesn't seem to block anything, but then I don't know either if it's just ignoring this layer in the backpropagation

[ZailiWang]

I am aware that you installed but did not import IPEX in the example code. Would you try without ipex package (pip uninstall intel_extension_for_pytorch)? If the warning persists, you can raise the ticket to PyTorch issue. Thanks.

[jonas-eschle]

@ZailiWang the issue does not persist. I was maybe unclear about this. That's why I raised the issue here: it only happens if the intel extension installed. The exact same code with pure pytorch does not raise an error.

While IPEX is not imported, it still changes the execution? I find a factor of two speedup when installing the extension

[ZailiWang]

Oh, I see. So I should re-try the reproduction code with import ipex added.

[jonas-eschle]

I think it's rather because I've also installed some other packages. Maybe there is some weird interference going on? But it's only torch_geometric

[ZailiWang]

I checked that layer_norm backward is registered in PyTorch rel 2.6, and IPEX should not fuse/fold the backward op, so it's basically not understandable why this message can be echoed in your env. It's also hardly possible that torch-geometric could change PyTorch op registration status.

Maybe we can wait a few days for PyTorch 2.7 release and check if the message persists in PT2.7.