Fix layer_normalize gradients

[arrufat]

Closes #2902

[arrufat]

It still doesn't fix the discrepancy between GPU and CPU, but fixes a bug in the implementation.

[davisking]

I'm looking at this and not sure what the code is supposed to be doing. Go through the contracts and make sure they are right. Like

    void layer_normalize (
        const double eps, 
        resizable_tensor& dest,
        resizable_tensor& means,
        resizable_tensor& invstds,
        const tensor& src, 
        const tensor& gamma,
        const tensor& beta 
    );   
    /*!  
        requires
            - eps > 0
            - src.num_samples() == gamma.size() == beta.size()
            - have_same_dimensions(gamma, beta) == true
            - beta.num_samples() ==beta.nr() ==gamma.nc() == 1

That's saying beta and gamma are the same shape and all the dimensions are 1 except k, which would have to be src.num_sample(). So by that what was in the code before this would make sense. But running some of these I see that gamma and beta don't have the shape that requires clause says they do. So there is inconsistency in how these are being interpreted. I.e. is gamma.size() == src.num_samples() or is it src.k() * src.nr() * src.nc()?

That's totally at the root of the problem here. Everything starts with having contracts that are right. Then put DLIB_ASSERT statements that check all the requires statements so you know for sure they are not being violated. That will chase down the problem. Although you've got to decide what the arguments are first. Not sure what you want them to be for this layer. I would think gamma.size() == src.num_samples() as that's probably the most typical variant of layer norm though.

[davisking]

Inside layer_normalize_gradient() for the cuda code, there are also local resiable_tensor variables (dvars and dmeans), those can't live there. It's real expensive to be creating and destroying tensors. They all need to live in a layer object so they aren't created and destroyed when the network runs. But rather are allocated once. But more than that kernels run asynchronously in cuda, so that _cuda_layer_normalize_gradient kernel launches but then those two variables are immediately freed. That's probably why the cuda version isn't working. Since it's running on dangling pointers.

Sorry I never really looked at this. You write such good PRs I just kinda skimmed this one and was like "yeah another Adria PR, going to be great and looks great 👍 :D " without really reading it all.

[arrufat]

Oh, right, what was I thinking. It looks like I got confused half-way through the code, where I should normalize each channel independently, but I ended up trying to normalize along k, nr, and nc. Definitely, k should not be included in the normalization.

Hopefully, I will fix that and the dangling pointer tonight if life allows.

[arrufat]

I was just checking this again:
https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html

It seems like, it does normalize along C, H, W. and there's one beta and one gamma for each normalized item.

import torch
N, C, H, W = 20, 5, 10, 10
x = torch.randn(N, C, H, W)
layer_norm = torch.nn.LayerNorm([C, H, W])
y = layer_norm(x)
sum(p.numel() for p in layer_norm.parameters() if p.requires_grad)  # 1000 = 2 * (5 * 10 * 10) 

So, maybe the issue is just the dangling pointers? I will make sure the contracts are correct and respected, though.

EDIT: after necrobumping ConvNeXt, each LayerNorm only has 2 * C learnable parameters (beta and gamma), so the implementation here is wrong. You're right about the dimensions of beta and gamma: they should only have k parameters each.

[davisking]

Yeah, might be just the dangling pointer and everything else is fine.

…

On Thu, Aug 29, 2024 at 8:21 AM Adrià ***@***.***> wrote: I was just checking this again: https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html image.png (view on web) <https://github.com/user-attachments/assets/8a757676-7a92-4649-9bf0-2e5f23325901> It seems like, it does normalize along C, H, W. and there's one gamma and one beta for each normalized item. import torchN, C, H, W = 20, 5, 10, 10x = torch.randn(N, C, H, W)layer_norm = torch.nn.LayerNorm([C, H, W])y = layer_norm(x)sum(p.numel() for p in layer_norm.parameters() if p.requires_grad) # 1000 = 2 * (5 * 10 * 10) So, maybe the issue is just the dangling pointers? I will make sure the contracts are correct and respected, though. — Reply to this email directly, view it on GitHub <#3001 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABPYFR4GTZW35ALFN7QVDHLZT4G45AVCNFSM6AAAAABNIK6NHOVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDGMJXGQ4TONJVGQ> . You are receiving this because you commented.Message ID: ***@***.***>

[arrufat]

I am now confident about the CPU implementation, however, the CUDA version still fails.
There's a mismatch between CPU and CUDA and in the beta and gamma gradients, which are trivial, but I can't spot the mistake.
It's getting late, I'll try again later. Feel free to check, it must be something really stupid.

[arrufat]

I honestly don't know what else to do.
The CPU version works correctly now, but not the CUDA version.

However, if you run the test_layer_normalize with CUDA enabled, you'll see that all the functionality is on par with the CPU version:

• normalized output tensor
• src_grad
• gamma_grad
• beta_grad
• dmeans
• dvars

All of them are within the tolerance error of 1e-5 or 1e-4. However, test_layer with layer_norm still fails.

[arrufat]

Ok, fixed a race condition, now test_layer complains like this:

Average parameter gradient error is somewhat large at: 0.00713434                                                                                                                                                                                                                                                               

EDIT: after running a clean build, it's working!

[davisking]

Nice. I'm away from my computer. I'll look in a bit. Seems like you got it 🥳

[arrufat]

It took an awful amount of time...

[davisking]

Yeah that kind of warp reduction loop is the best way I know to do it too.

[arrufat]

After the previous two commits, the network went to train from 320 img/s to 2450 img/s (close to the official CUDA/CUDNN batch norm at 2560 img/s)

[davisking]

After the previous two commits, the network went to train from 320 img/s to 2450 img/s (close to the official CUDA/CUDNN batch norm at 2560 img/s)

Yeah that's awesome. All the tests are passing for me too on a GPU machine. Passing for you too now? Anything else you want to change before I merge it? :D

[arrufat]

Nothing else to add, I think it's done now. FINALLY.

And yes, tests are passing now :D

[davisking]

Yeah nice, thanks for all the good work. Looks perfect :D