gpt_neox

Model type	gpt_neox
Class	GPTNeoXCausalLMModel
Task	text-generation
Source	models/gpt_neox.py

Description

GPT-NeoX / Pythia causal language model.

Uses dual LayerNorm parallel residual connections — unlike standard sequential (Llama-style) or single-norm parallel (GPT-J style), each layer applies two separate norms to the original hidden states, one for attention and one for the MLP branch.

Attribute names align with HF GPTNeoXForCausalLM:

• gpt_neox for the backbone

• embed_out for the LM head (not lm_head)

GPT-NeoX only supports MHA (no GQA), so num_key_value_heads is forced to equal num_attention_heads to prevent weight shape mismatches.

Replicates HuggingFace’s GPTNeoXForCausalLM.

Usage

mobius build --model <MODEL_ID> output_dir/

from mobius import build

model = build("<MODEL_ID>")