GPT-NeoX: EleutherAI's library for training billion-parameter language models on GPUs

What GPT-NeoX is for, and who it is not for

GPT-NeoX is EleutherAI's library for training large-scale language models on GPUs. It is an implementation of model parallel autoregressive transformers, based on NVIDIA's Megatron Language Model and augmented with techniques from DeepSpeed. The README describes the goal as making the repository a centralized place to gather techniques for training large-scale autoregressive language models.

The README is unusually direct about the audience. It states that if you are not looking to train models with billions of parameters from scratch, this is likely the wrong library to use, and recommends Hugging Face transformers for generic inference needs. That single sentence rules out a large share of people who search for GPT-NeoX after seeing the name attached to a released model.

The distinction matters because GPT-NeoX is also the name of a model, GPT-NeoX-20B, and of an architecture family that appears in Hugging Face. The repository here is the training code, not a weights download or an inference runtime. If your goal is to run a GPT-NeoX checkpoint, this is not the tool. If your goal is to produce one, it is.

How the training stack is put together

The architecture is a layered one. At the bottom sits DeepSpeed, which supplies the ZeRO optimizer and sharding machinery. On top of that sits Megatron-style tensor and pipeline parallelism, which is what lets a single model span many GPUs. GPT-NeoX adds its own optimizations and a configuration layer over both.

The repository layout reflects this. The megatron/ directory holds the model and parallelism code, configs/ holds YAML configuration files, and the top level exposes deepy.py as the launcher, train.py for training, generate.py for generation, eval.py for evaluation, and prepare_data.py for dataset preparation. The repository also has post-training/ and eval_tasks/ directories.

Parallelism is configured rather than inferred. The README lists distributed training with ZeRO and 3D parallelism as a major feature, and separately notes predefined configurations for Pythia, PaLM, Falcon, and LLaMA 1 and 2. Those configurations are starting points; the tensor, pipeline and data parallel degrees still have to match your cluster's topology.

On the systems side, the README lists launching via Slurm, MPI, and the IBM Job Step Manager. It also lists runs on AWS, CoreWeave, ORNL Summit and Frontier, PNNL, Argonne Polaris and LUMI. Those are deployment targets the project reports having been run on, not benchmarks you can read a throughput number from.

Installing GPT-NeoX and running a first training job

The README organizes setup into host setup, Flash Attention, Transformer Engine, multi-node launching and a containerized setup, with dependencies pinned under requirements/. The containers/ directory at the top level is the containerized path. Host setup is pip install -r requirements/requirements.txt from the repository root, with optional WandB, TensorBoard and Comet requirement files, and the README targets Python 3.8-3.10 and PyTorch 1.8-2.0 while recommending environment isolation because the codebase relies on DeeperSpeed.

Once the environment exists, the top level holds deepy.py and train.py for running jobs. The repository's configs/ directory holds the YAML files those runs read from, including the predefined configurations for Pythia, PaLM, Falcon, and LLaMA 1 and 2.

Data is prepared separately. The repository ships prepare_data.py for this, and the README has a Datasets section covering preconfigured datasets and custom data.

Monitoring is configured, not assumed. The README lists WandB, Comet and TensorBoard as supported experiment trackers.

Version 1.0 versus 2.0, and why the split exists

Before March 2023, GPT-NeoX depended on DeeperSpeed, a fork based on DeepSpeed 0.3.15. To move to current upstream DeepSpeed without stranding existing users, the project split both libraries into versioned releases.

Version 2.0 of GPT-NeoX and DeeperSpeed is built on the latest DeepSpeed and is the line the project says will be maintained going forward. Version 1.0 preserves the snapshots that GPT-NeoX-20B and the Pythia Suite were trained on. That is a reproducibility decision, and it is the right one, but it means the answer to "which version should I use" depends on whether you are reproducing a published model or starting something new.

The practical consequence is that instructions, config keys and dependency pins found in older write-ups may belong to the 1.0 line and not apply to a current checkout. The v2.0 release is dated 2023-03-10 and v1.0 2023-03-09, so the two are effectively siblings rather than one superseding the other over time.

Where GPT-NeoX is the wrong choice

The clearest limitation is the one the project states itself: this is not an inference library. Generation exists through generate.py, but the README explicitly directs people with generic inference needs to Hugging Face transformers. Running a released checkpoint through GPT-NeoX means carrying a DeepSpeed and Megatron dependency tree for something a standard transformers install handles.

The hardware requirement is the second constraint. Multi-GPU, multi-node training is the design center, and features like pipeline parallelism and Mixture-of-Experts only make sense at that scale. On a single GPU, the parallelism machinery is overhead rather than benefit.

There is also a moving dependency surface. The README has setup sections for Flash Attention, Transformer Engine and Mixture of Experts, and adds AMD GPU (MI100, MI250X) support through JIT fused-kernel compilation. Each of those is a new combination of library versions to keep aligned.

Finally, the project is a research framework. The README frames it as a place to gather techniques for large-scale training. That framing implies interfaces can change between releases in ways a production serving stack would not tolerate.

How it differs from Megatron-DeepSpeed

The nearest alternative is Megatron-DeepSpeed, and the README addresses it directly: GPT-NeoX uses many of the same features and technologies but with what the project describes as substantially increased usability and novel optimizations.

The difference in approach is mainly at the edges. Megatron-DeepSpeed is the combination of two upstream projects; GPT-NeoX is a curated distribution of that combination with its own launcher (deepy.py), its own YAML configuration layer, predefined configs for Pythia, PaLM, Falcon and LLaMA 1 and 2, and its own data preparation script. It also names Slurm, MPI and the IBM Job Step Manager as launch options.

If you are already running Megatron-DeepSpeed and your launcher and config format work, switching buys you the predefined configs and the integration surface (WandB, Comet, TensorBoard, Hugging Face tokenizers, the LM Evaluation Harness) rather than a different parallelism implementation. If you are starting from nothing, the config layer is the more meaningful difference, because it removes a category of setup work.

For anything below the billion-parameter line, neither is the right comparison. Hugging Face transformers is, and the README says so.

Maintenance, licensing and upgrade cost

The repository is not archived, and the last push was on 2026-10-11. The most recent tagged release is v2.0 from 2023-03-10, so the code has moved since the last version tag. Anyone depending on a tagged artifact should note that gap between the release line and the main branch.

The project is licensed Apache-2.0, which permits commercial use and modification and includes a patent grant. The LICENSE file is at the repository root. That covers the training code; it says nothing about the licences of datasets or base checkpoints you train on, which are separate questions and not ones this repository answers.

Upgrade cost is dominated by the dependency graph rather than the GPT-NeoX code itself. DeepSpeed, Flash Attention, Transformer Engine and PyTorch versions are coupled. An upgrade therefore means moving the whole set together. The versioned 1.0 and 2.0 releases exist precisely so that teams reproducing older models can stay put, which is the intended escape hatch.

Editorial conclusion

Adopt GPT-NeoX if you are pretraining a multi-billion-parameter autoregressive model on a multi-GPU cluster and need ZeRO, 3D parallelism and launch support for Slurm, MPI or the IBM Job Step Manager. Do not adopt it for inference or fine-tuning a small model; the README points those users to Hugging Face transformers instead. Before committing, verify that your cluster's CUDA and PyTorch versions satisfy the pinned requirements under requirements/, and check whether the versioned release you need is v2.0 or the v1.0 snapshot that GPT-NeoX-20B and Pythia were trained on.

Frequently asked questions

Official sources

1. EleutherAI/gpt-neox on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases