core/nn/pythia_hf_loader library

Loads HuggingFace Pythia (gpt_neox architecture) weights from a safetensors file into a PythiaModel.

Supported checkpoints (all use the standard GPT-NeoX tokenizer, vocab 50304, context 2048):

  • EleutherAI/pythia-14m — 6 layers, 128 embed, 8 heads
  • EleutherAI/pythia-70m — 6 layers, 512 embed, 8 heads
  • EleutherAI/pythia-160m — 12 layers, 768 embed, 12 heads
  • EleutherAI/pythia-410m — 24 layers, 1024 embed, 16 heads

HF key conventions handled here:

  • gpt_neox.embed_in.weight
  • gpt_neox.layers.{i}.input_layernorm.{weight,bias}
  • gpt_neox.layers.{i}.post_attention_layernorm.{weight,bias}
  • gpt_neox.layers.{i}.attention.query_key_value.{weight,bias} — shape [3*D, D] / [3*D]; per-head Q|K|V interleaved as [num_heads, 3, head_dim, ...] (not [3, num_heads, ...] the way GPT-2 does it).
  • gpt_neox.layers.{i}.attention.dense.{weight,bias} — output projection, standard [D, D] PyTorch Linear.
  • gpt_neox.layers.{i}.mlp.dense_h_to_4h.{weight,bias} — [4D, D].
  • gpt_neox.layers.{i}.mlp.dense_4h_to_h.{weight,bias} — [D, 4D].
  • gpt_neox.final_layer_norm.{weight,bias}
  • embed_out.weight — untied [V, D].

Ignored (safe): per-layer attention.rotary_emb.inv_freq, attention.bias, attention.masked_bias — we don't need HF's cached rotary freqs (RopeCache recomputes them) and we build causal masks on the fly.