[BUG] RMSNorm Implementation

[0seba]

I think you should be dividing by the scale in the following line

zeta/zeta/nn/modules/rms_norm.py

Line 35 in 7dbb6a6

return normed * self.scale * self.gamma

This this the scale definition

https://github.com/kyegomez/zeta/blob/7dbb6a62f83413977a922d5fc6dec1b11f734bc3/zeta/nn/modules/rms_norm.py#L29C9-L29C31

self.scale = dim**-0.5

And RMSNorm formula

Edit:

Also, I think the normalization should be in the dim -1, not -2

zeta/zeta/nn/modules/rms_norm.py

Line 34 in 7dbb6a6

normed = F.normalize(x, dim=-2)

Upvote & Fund

• We're using Polar.sh so you can upvote and help fund this issue.
• We receive the funding once the issue is completed & confirmed by you.
• Thank you in advance for helping prioritize & fund our backlog.

[github-actions]

Hello there, thank you for opening an Issue ! 🙏🏻 The team was notified and they will get back to you asap.

[github-actions]

Stale issue message