CUDA error: an illegal memory access was encountered (pytorch CUDA extension) on GPU and Segmentation fault on CPU

singleroc: <p>HI, I am trying to call the native functions. However, I met ‘CUDA error: an illegal memory access was encountered’ when I ran the CUDA version and it gave ‘Segmentation fault’ when I switched to the CPU version.</p>
<p>I tested the code on pytorch 1.5/1.6/1.7 with cuda 9.2, all pytorch version gave the same error.</p>
<p>What I was trying to do:<br>
First wrap the functions:</p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;
#include &lt;ATen/NativeFunctions.h&gt;
#include &lt;ATen/Config.h&gt;

std::tuple&lt;at::Tensor, at::Tensor, at::Tensor&gt; layer_norm_forward_cpu(
    const at::Tensor &amp; input,
    const at::Tensor &amp; weight,
    const at::Tensor &amp; bias,
    int64_t M, int64_t N, double eps) {
    return at::native::layer_norm_cpu(input, weight, bias, M, N, eps);
}

std::tuple&lt;at::Tensor, at::Tensor, at::Tensor&gt; backward_layer_norm_cpu(
    const at::Tensor &amp; grad_out,
    const at::Tensor &amp; input,
    const at::Tensor &amp; mean,
    const at::Tensor &amp; rstd,
    const at::Tensor &amp; weight,
    int64_t M, int64_t N, std::array&lt;bool,3&gt; output_mask) {
    return at::native::layer_norm_backward_cpu(grad_out, input, mean, rstd, weight, M, N, output_mask);
}

std::tuple&lt;at::Tensor, at::Tensor, at::Tensor&gt; layer_norm_forward_cuda(
    const at::Tensor &amp; input,
    const at::Tensor &amp; weight,
    const at::Tensor &amp; bias,
    int64_t M, int64_t N, double eps) {
    return at::native::layer_norm_cuda(input, weight, bias, M, N, eps);
}

std::tuple&lt;at::Tensor, at::Tensor, at::Tensor&gt; backward_layer_norm_cuda(
    const at::Tensor &amp; grad_out,
    const at::Tensor &amp; input,
    const at::Tensor &amp; mean,
    const at::Tensor &amp; rstd,
    const at::Tensor &amp; weight,
    int64_t M, int64_t N, std::array&lt;bool,3&gt; output_mask) {
    return at::native::layer_norm_backward_cuda(grad_out, input, mean, rstd, weight, M, N, output_mask);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("layer_norm_forward_cpu",  &amp;layer_norm_forward_cpu, "layer norm forward (cpu version)");
    m.def("layer_norm_backward_cpu", &amp;backward_layer_norm_cpu, "layer norm backward (cpu version)");
    m.def("layer_norm_forward_cuda", &amp;layer_norm_forward_cuda, "layer norm forward (cuda version)");
    m.def("layer_norm_backward_cuda",&amp;backward_layer_norm_cuda, "layer norm backward (cuda version)");
}
</code></pre>
<p>and then call the functions in</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
import pdb

import native

class layer_norm(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, normalized_shape, weight, bias, eps, training):
        N = 1
        if isinstance(normalized_shape, int):
            N = normalized_shape
        elif isinstance(normalized_shape, (list, tuple)):
            for i in normalized_shape:
                N *= i
        else:
            raise RuntimeError("unexpected type of normalized_shape".format(type(normalized_shape)))
        M = x.nelement() // N

        if x.is_cuda:

            y, mean, rstd = native.layer_norm_forward_cuda(x, weight, bias, M, N, eps)
        else:
            y, mean, rstd = native.layer_norm_forward_cpu(x, weight, bias, M, N, eps)

        if training:
            ctx.layer_norm_input = x
            ctx.layer_norm_parameters = (mean, rstd, weight, M, N)
        return y

    @staticmethod
    def backward(ctx, grad_output):
        x = ctx.layer_norm_input
        mean, rstd, weight, M, N = ctx.layer_norm_parameters

        output_mask = [True, True, True]
        if grad_output.is_cuda:
            grad_input, grad_weight, grad_bias = native.layer_norm_backward_cuda(grad_output, x, mean, rstd, weight, M, N, output_mask)
        else:
            grad_input, grad_weight, grad_bias = native.layer_norm_backward_cpu(grad_output, x, mean, rstd, weight, M, N, output_mask)
        ctx.layer_norm_input = None
        ctx.layer_norm_parameters = None
        return grad_input, None, grad_weight, grad_bias, None, None, None, None

class LayerNorm(nn.LayerNorm):
    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True):
        nn.LayerNorm.__init__(self, normalized_shape, eps=eps, elementwise_affine=elementwise_affine)

    def forward(self, x):
        y = layer_norm.apply(x, self.normalized_shape, self.weight, self.bias, self.eps, self.training)
        return y


if __name__ == "__main__":
    seed = 2809
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.deterministic=True #https://github.com/pytorch/pytorch/issues/8019

    model = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            LayerNorm([64,56,56])
            )
    print(model)

    #model = model.cuda()
    model.train()
    iteration = 10
    for i in range(iteration):
        print("index: ", i)
        x = torch.rand(512,64,56,56)
        x = x - 0.5
        #x = x.cuda()

        y = model(x)
        z = y.sum()
        z.backward()

</code></pre>
<p>I also uploaded all the code to <a href="https://github.com/irving-qin/nativefunctions" class="inline-onebox" rel="noopener nofollow ugc">GitHub - irving-qin/nativefunctions</a><br>
Just run <code>bash install.sh</code> for testing.</p>
<p>The forward process of the wrapped layer norm seemed to be normal. However, it throwed errors in the backward function.</p>
<p>Thanks so much if anyone could give me some tips.</p>

ptrblck: <p>Your current code cannot be build using a recent source build and fails with:</p>
<pre><code class="lang-auto">/workspace/src/nativefunctions/native.cpp:39:66: error: could not convert ‘mean’ from ‘const at::Tensor’ to ‘c10::IntArrayRef’ {aka ‘c10::ArrayRef&lt;long int&gt;’}
   39 |     return at::native::layer_norm_backward_cuda(grad_out, input, mean, rstd, weight, M, N, output_mask);
      |                                                                  ^~~~
      |                                                                  |
      |                                                                  const at::Tensor
</code></pre>
<p>EDIT: it seems the <code>IntArrayRef normalized_shape</code> input is missing as seen <a href="https://github.com/pytorch/pytorch/blob/c3d40fdf5658d8f4d387a39633c11bbb67b17224/aten/src/ATen/native/cuda/layer_norm_kernel.cu#L475-L483">here</a>.</p>

singleroc: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="121574">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>t code cannot be build using a recent source build and fails with:</p>
</blockquote>
</aside>
<p>Hi, <a class="mention" href="/u/ptrblck">@ptrblck</a> Thank so much for your attention on the question.</p>
<p>It seems the api was changed since pytorch 1.8.</p>
<p>In my eary trival, I didn’t build the code again the source code of the pytorch. Instead I build it with a pre-installed version.</p>
<p>I found the function definition in path something like<br>
<code> ~/.pyenv/versions/3.6.5/lib/python3.6/site-packages/torch/include/ATen/NativeFunctions.h</code><br>
The path prefix might be different on different machines.</p>
<p>If I ran<br>
<code>grep layer_norm ~/.pyenv/versions/3.6.5/lib/python3.6/site-packages/torch/include/ATen/NativeFunctions.h</code> on pytorch1.5/1.6/1.7  it gave</p>
<pre><code class="lang-auto">CAFFE2_API Tensor layer_norm(const Tensor &amp; input, IntArrayRef normalized_shape, const Tensor &amp; weight={}, const Tensor &amp; bias={}, double eps=1e-05, bool cudnn_enable=true);
CAFFE2_API std::tuple&lt;Tensor,Tensor,Tensor&gt; layer_norm_cpu(const Tensor &amp; input, const Tensor &amp; weight, const Tensor &amp; bias, int64_t M, int64_t N, double eps);
CAFFE2_API std::tuple&lt;Tensor,Tensor,Tensor&gt; layer_norm_cuda(const Tensor &amp; input, const Tensor &amp; weight, const Tensor &amp; bias, int64_t M, int64_t N, double eps);
CAFFE2_API std::tuple&lt;Tensor,Tensor,Tensor&gt; layer_norm_backward_cpu(const Tensor &amp; grad_out, const Tensor &amp; input, const Tensor &amp; mean, const Tensor &amp; rstd, const Tensor &amp; weight, int64_t M, int64_t N, std::array&lt;bool,3&gt; output_mask);
CAFFE2_API std::tuple&lt;Tensor,Tensor,Tensor&gt; layer_norm_backward_cuda(const Tensor &amp; grad_out, const Tensor &amp; input, const Tensor &amp; mean, const Tensor &amp; rstd, const Tensor &amp; weight, int64_t M, int64_t N, std::array&lt;bool,3&gt; output_mask);
</code></pre>
<p>I just tried the grep on pytorch 1.8, it gave</p>
<pre><code class="lang-auto">TORCH_API Tensor layer_norm(const Tensor &amp; input, IntArrayRef normalized_shape, const Tensor &amp; weight={}, const Tensor &amp; bias={}, double eps=1e-05, bool cudnn_enable=true);
TORCH_API std::tuple&lt;Tensor,Tensor,Tensor&gt; layer_norm_cpu(const Tensor &amp; input, IntArrayRef normalized_shape, const Tensor &amp; weight, const Tensor &amp; bias, double eps);
TORCH_API std::tuple&lt;Tensor,Tensor,Tensor&gt; layer_norm_cuda(const Tensor &amp; input, IntArrayRef normalized_shape, const Tensor &amp; weight, const Tensor &amp; bias, double eps);
TORCH_API std::tuple&lt;Tensor,Tensor,Tensor&gt; math_native_layer_norm(const Tensor &amp; input, IntArrayRef normalized_shape, const Tensor &amp; weight, const Tensor &amp; bias, double eps);
TORCH_API std::tuple&lt;Tensor,Tensor,Tensor&gt; layer_norm_backward_cpu(const Tensor &amp; grad_out, const Tensor &amp; input, IntArrayRef normalized_shape, const Tensor &amp; mean, const Tensor &amp; rstd, const Tensor &amp; weight, const Tensor &amp; bias, std::array&lt;bool,3&gt; output_mask);
TORCH_API std::tuple&lt;Tensor,Tensor,Tensor&gt; layer_norm_backward_cuda(const Tensor &amp; grad_out, const Tensor &amp; input, IntArrayRef normalized_shape, const Tensor &amp; mean, const Tensor &amp; rstd, const Tensor &amp; weight, const Tensor &amp; bias, std::array&lt;bool,3&gt; output_mask);
</code></pre>
<p>I have updated the code on github repo (run <code>git pull</code>, please ) to make it support pytorch 1.8. However the error was still there.</p>

singleroc: <p>Hi, I found out the reason for the issue. It is caused by the grad_output to be not contiguous. After adding <code>grad_output = grad_output.contiguous()</code>, the error is gone.</p>