register fused rmsnorm as pytorch custom op

[tianyu-l]

Stack from ghstack (oldest at bottom):

• -> register fused rmsnorm as pytorch custom op #296

[msaroufim]

n00b but is there any benefit in making a triton function a custom op? User defined triton functions should just work with compile

[tianyu-l]

@msaroufim Hmm I don't know much. Maybe it depends on the way the triton function is wrapped? In this PR there is an autograd.function wrapping the triton function. cc: @Chillee @lessw2020 for more context.

Beyond compile, making it a custom op can be helpful for other things, e.g. registering customized DTensor sharding propagation rules, allow PP tracing to work (although we no longer have it in torchtitan) etc.

[msaroufim]

should we perhaps just remove this custom kernel? I believe in pytorch nightlies rms norm should now work

[tianyu-l]

I believe in pytorch nightlies rms norm should now work

@msaroufim oh which rms norm are you referring to?