Closed rasbt closed 7 months ago
This is the relevant gpt-neox reference, not Megatron: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L669-L723
Thanks! It's still post LN though right?
# pseudocode:
# x = x + attn(ln1(x)) + mlp(ln2(x))
mlp_output = self.mlp(self.post_attention_layernorm(hidden_states))
I mean, Pre LN is what I use in practice, which is what Lit-GPT also does. But maybe to reproduce some of these models more exactly, we probably want to allow Post LN?
You have both, pre and post: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L673-L674 https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L690-L691 https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L706
Hi guys 👋 I think that the naming in HF implementation introduced some level of confusion. Both normalization layers are applied before (aka Pre-LayerNorm), it's just the names are a bit confusing:
self.input_layernorm
self.post_attention_layernorm
If to strip down their code, then the ordering will be:
self.input_layernorm
self.attention
self.post_attention_layernorm
self.mlp
This is reflected in the provided pseudo-code:
# pseudocode:
# x = x + attn(ln1(x))
# x = x + mlp(ln2(x))
Ah yes, thanks, this looks like pre-layernorm. Otherwise it would be x = x + ln1(attn(x))
Closing?
Yeah that's probably fine since no one is using Post-LN anymore.
It's a nitpick, but I think Pythia uses Post-LayerNorm. However, in Lit-GPT, we may use Pre-LayerNorm:
https://github.com/Lightning-AI/lit-gpt/blob/0f021f3ad8cd8d6fe30b0ef721a7a0e2dac15898/lit_gpt/model.py#L136C1-L142C44
I.e.,
I may see this incorrectly though, because in the GPTNeoX paper it looks like they use PreLayerNorm:
Long story short, maybe Block in Lit-GPT should have a Norm argument like