gpt_neox

gpt_neox

GPT-NeoX (GPTNeoXForCausalLM; Pythia, GPT-NeoX-20B).

Its tree is gpt_neox.{embed_in, emb_dropout, layers[i].{input_layernorm, post_attention_layernorm, attention, mlp}, final_layer_norm}. The head is lm_head on current transformers and was embed_out before; that key binds only where it resolves, so both spellings map to lm_head.

RENAME module-attribute

RENAME = {'gpt_neox.embed_in': 'embed_tokens', 'gpt_neox.layers': 'layers', 'gpt_neox.final_layer_norm': 'norm', 'attention': 'self_attn', 'embed_out': 'lm_head', 'dense': 'o_proj', 'act': 'act_fn', 'dense_4h_to_h': 'down_proj'}

ENVOYS module-attribute

ENVOYS = {GPTNeoXLayer: Layer, GPTNeoXAttention: Attention, GPTNeoXMLP: Mlp}