Support backward for RMSNorm

[samsja]

Support Backward in RMSNorm

hello,

It seems that we can't use RMSNorm with a backward pass (see code here )

I was wondering if you plan to support it in the future, or if you knew any workaround to have a fast implementation of RMSNorm.

Thanks in advance

[danthe3rd]

Hi, This is a fairly old code that we won't improve moving forward (and it will most likely become deprecated). I would encourage you to try using torch compile for that particular code, and it should be fast + generate a correct backward pass

[samsja]

Hi, This is a fairly old code that we won't improve moving forward (and it will most likely become deprecated). I would encourage you to try using torch compile for that particular code, and it should be fast + generate a correct backward pass

I indeed ended using torch compile for this. Thanks for the update