[Bug Report] RMSNormPre in Transformer_lens is maybe different from Llama source code?

[wangyifei0047]

In LlamaModeling.py, the LlamaRMSNorm function outputs the weights * scaled hidden_states like below

RMSNormPre definition in Transformer_lens: it seems that this function just outputs the scaled hidden_states

The way RMSNormPre by which Transformer_Block uses
it seems that in the forward process in Transformer_Block, the weights of LlamaRMSNorm still not be added.

I want to hook the values after applying RMSNorm on each residual stream, so I try to find the parameters in RMSNorm and find something weird.

• [yes] I have checked that there is no similar issue in the repo (required)

[4gatepylon]

Have you tried comparing intermediate values using hooks? It may be the case that they folded into the weights of a subsequent layer.