Where is the actual code for LayerNorm (torch.nn.functional.layer_norm)

username: zeyuyun1
cooked: <p>I am looking for the implementation for <code>torch.nn.functional.layer_norm</code>, it links me to this <a href="https://pytorch.org/docs/stable/nn.functional.html?highlight=torch%20layer_norm#torch.nn.functional.layer_norm" rel="noopener nofollow ugc">doc</a>, which then link me to <a href="https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#layer_norm" rel="noopener nofollow ugc">this one</a></p>
<p>But I can’t find where is <code>torch.layer_norm</code>.</p>
<p>According to the documentation, it seems like the math is following:</p>
<pre><code class="lang-auto">
x = torch.randn(50,20,100)
mean = x.sum(axis = 0)/(x.shape[0])
std = (((x - mean)**2).sum()/(x.shape[0])).sqrt()

LayerNorm = torch.nn.LayerNorm(x.shape, elementwise_affine = True)

torch_layernorm = LayerNorm(x)
My_LayerNorm = (x - mean)/std*LayerNorm.weight+LayerNorm.bias

print(My_LayerNorm)
print(torch_layernorm)
</code></pre>
<p>However, the my output and LayerNorm output is different…</p>

username: ptrblck
cooked: <p>You can find the (CPU) C++ implementation <a href="https://github.com/pytorch/pytorch/blob/392abde8e64b0d91b7d52aecee8dce9aff8d0b2f/aten/src/ATen/native/layer_norm.cpp">here</a>.</p>

username: Xinchengzelin
cooked: <p>Hi, <a class="mention" href="/u/ptrblck">@ptrblck</a> ,  could you tell me where I can find  <code>native_layer_norm</code> in the line <code>return std::get&lt;0&gt;(at::native_layer_norm(input, normalized_shape, weight, bias, eps));</code></p>

username: Ice-Citron
cooked: <p>Heres what i written in my class which when I use it to train a 124m parameter LLM model. The loss curve looks the exact same (refer to forward pass method below):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/2/2205321494a950447610a118e6f1aa70df418ca7.png" data-download-href="https://discuss.pytorch.org/uploads/default/2205321494a950447610a118e6f1aa70df418ca7" title="Screenshot 2024-08-05 at 2.55.24 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/2/2205321494a950447610a118e6f1aa70df418ca7_2_690x381.png" alt="Screenshot 2024-08-05 at 2.55.24 AM" data-base62-sha1="4QXmFgt4nlT3BhTmrK1C3TCfxGL" width="690" height="381" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/2/2205321494a950447610a118e6f1aa70df418ca7_2_690x381.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/2/2205321494a950447610a118e6f1aa70df418ca7_2_1035x571.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/2/2205321494a950447610a118e6f1aa70df418ca7.png 2x" data-dominant-color="252624"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-08-05 at 2.55.24 AM</span><span class="informations">1307×723 77.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<ul>
<li>for reference. the run with the layer norm class used below can be seen in the run “radiant-flower-16”. “absurd-glitter-13”, “crisp-flower-10” and “kind-energy-14” are previous runs using torch’s default layernorm class. And hence you can see their loss is very simillar, hence I think my implementation is a direct replica afaik (haven’t tested it beyond this).</li>
<li>P.S. ignore the ethereal-oath-9, that was some bugs with logging.</li>
</ul>
<pre><code class="lang-auto">class LayerNorm(nn.Module):

    __constants__ = ['normalized_shape', 'eps', 'elementwise_affine']
    normalized_shape: Tuple[int, ...]
    eps: float
    elementwise_affine: bool

    def __init__(self, normalized_shape: _shape_t, eps: float = 1e-5, elementwise_affine: bool = True, bias: bool = True, device=None, dtype=None) -&gt; None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        if isinstance(normalized_shape, numbers.Integral):
            # mypy error: incompatible types in assignment
            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        if self.elementwise_affine:
            self.weight = Parameter(torch.empty(self.normalized_shape, **factory_kwargs))
            if bias:
                self.bias = Parameter(torch.empty(self.normalized_shape, **factory_kwargs))
            else:
                self.register_parameter('bias', None)
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)

        self.reset_parameters()

    def reset_parameters(self) -&gt; None:
        if self.elementwise_affine:
            init.ones_(self.weight)
            if self.bias is not None:
                init.zeros_(self.bias)

    def forward(self, x: Tensor) -&gt; Tensor:
        """Forward pass of LayerNorm."""
        mean = x.mean(dim=-1, keepdim=True)
        mean_x2 = torch.square(x).mean(dim=-1, keepdim=True)
        var = mean_x2 - torch.square(mean)

        x_norm = (x - mean) / torch.sqrt(var + self.eps)

        if self.elementwise_affine:
            x_norm = self.weight * x_norm + self.bias

        return x_norm
        
    def extra_repr(self) -&gt; str:
        return '{normalized_shape}, eps={eps}, elementwise_affine={elementwise_affine}'.format(**self.__dict__)```</code></pre>

username: Ice-Citron
cooked: <p>actually nevermind. There’s a chance that my code isn’t the exact thing. And it might have worked by chance, as I tested again.</p>

username: yuhsianglin
cooked: <p>One would only need to change the calculation by</p>
<pre><code class="lang-auto">My_LayerNorm = (x - mean) / (std**2 + epsilon)**0.5 * LayerNorm.weight + LayerNorm.bias
</code></pre>
<p>according to the <a href="https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html" rel="noopener nofollow ugc">documentation</a>. It will match the output of LayerNorm</p>

username: rk_evolving
cooked: <p>if you are dealing with text sequence, the target of layernorm should be the last dimension. in your computing, you choosed the first dimension. this might be the reason why the results are different.</p>