neox-ckpt-pythia-2.8b

Pythia-2.8B GPT-NeoX Checkpoints

This repository contains the raw GPT-NeoX training checkpoints for Pythia-2.8B. These are the native checkpoint files produced during training, stored in DeepSpeed's checkpoint format.

If you want to perform inference, use the HuggingFace Transformers-compatible weights at EleutherAI/pythia-2.8b instead. This repository is intended for research that requires access to optimizer states or the original training format.

Contents

Each branch contains a full training checkpoint at a given step, including:

• layer_XX-model_00-model_states.pt — model weight shards (one per layer)
• mp_rank_00_model_states.pt — model state metadata
• zero_pp_rank_*_optim_states.pt — ZeRO optimizer states (Adam moments, etc.)
• 2-8B.yml — GPT-NeoX training configuration

Branches

20 log-spaced checkpoints are available as branches:

• step0 — initialization
• step{1,2,4,8,16,32,64,128,256,512} — log-spaced early checkpoints
• step{1000,2000,4000,8000,16000,32000,64000,128000} — log-spaced training checkpoints
• step143000 — final model

Branch step143000 corresponds to the final model.

Note: To keep storage requirements manageable, this repository provides a log-spaced subset of 20 checkpoints rather than all 154 training checkpoints. If you need linearly-spaced checkpoints (every 1,000 steps), the HuggingFace Transformers-compatible weights for all 154 checkpoints are available at EleutherAI/pythia-2.8b.

Converting to HuggingFace Format

To convert a checkpoint to HuggingFace Transformers format, use the conversion script from GPT-NeoX:

python tools/convert_neox_to_hf.py \
    --input_dir /path/to/neox/checkpoint \
    --config_file /path/to/config.yml \
    --output_dir /path/to/hf/output

Pre-converted weights for all checkpoints are available at EleutherAI/pythia-2.8b.

Training Details

Trained on the Pile.

All Pythia models were trained for 143,000 steps with a batch size of 2M tokens (2,097,152 tokens per step), seeing a total of 299,892,736,000 tokens. See the Pythia paper and GitHub repository for full training details.

<figure> </figure>

Citation

@article{biderman2023pythia,
  title={Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling},
  author={Biderman, Stella and Schoelkopf, Hailey and Anthony, Quentin Gregory and Bradley, Herbie and O'Brien, Kyle and Hallahan, Eric and Khan, Mohammad Aflah and Purohit, Shivanshu and Prashanth, USVSN Sai and Raff, Edward and others},
  journal={International Conference on Machine Learning},
  year={2023}
}