gpt_neox

gpt_neox

GPT-NeoX (Pythia) on vLLM (vllm.model_executor.models.gpt_neox).

Its tree is transformers' (gpt_neox.{embed_in, layers[i].{input_layernorm, attention, post_attention_layernorm, mlp}, final_layer_norm} plus embed_out). The block is called with the positions and the residual stream and returns the stream, adding both contributions itself, in parallel or in sequence as the config's use_parallel_residual says.

RENAME module-attribute

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

ENVOYS module-attribute

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