EleutherAI/gpt-neox

An implementation of model parallel autoregressive transformers on GPUs, based on the Megatron and DeepSpeed libraries

What is GPT‑NeoX?

GPT‑NeoX is an open‑source library for training very large autoregressive language models (think billions of parameters) on GPU clusters. It builds on NVIDIA’s Megatron‑LM code‑base, adds DeepSpeed‑style optimizations, and includes a lot of engineering to make massive‑scale training practical on many supercomputing environments (AWS, Summit, Frontier, LUMI, etc.).

Who would use it?

  • Research labs that want to train new LLMs from scratch or continue training existing ones.
  • Engineers needing a production‑ready training stack that can run on Slurm, MPI, or IBM Job Step Manager clusters.
  • Anyone who already has the compute budget for multi‑node, multi‑GPU training and wants a turnkey framework rather than building everything from low‑level PyTorch.

If you only need to run inference on an existing model, the README advises using Hugging Face transformers instead.

Core capabilities (as described in the README)

Feature What it means
Distributed training Supports ZeRO‑stage optimizations and 3‑D parallelism (data, tensor, pipeline) to spread a model across many GPUs.
Hardware flexibility Works with NVIDIA and AMD GPUs, can be launched via Slurm, MPI, pdsh, or IBM Job Step Manager, and has been run on cloud (AWS, CoreWeave) and many supercomputers (Summit, Frontier, Polaris, LUMI).
Modern model tricks Includes rotary & alibi positional embeddings, flash attention, parallel feed‑forward layers, and support for Mixture‑of‑Experts, RWKV, Mamba, and other emerging architectures.
Curriculum learning & preference learning Built‑in pipelines for curriculum schedules and recent fine‑tuning methods like DPO, KTO, and reward modeling.
Ecosystem integration Uses Hugging Face tokenizers/transformers, logs to WandB, Comet, TensorBoard, and can evaluate with the EleutherAI LM‑Evaluation‑Harness.
Export Trained checkpoints can be converted to the Hugging Face format for downstream inference.
Container support Docker and Apptainer images are provided for reproducible environments.

How to get started (quick‑start outline)

  1. Set up a Python environment (Python 3.8‑3.10, PyTorch 1.8‑2.0) and install the required packages via the requirements/*.txt files.
  2. Choose a launch method – e.g., Slurm, MPI, or the default pdsh. Create a hostfile describing the nodes/GPUs.
  3. Pick or write a config – the repo ships YAML configs for popular model sizes (1‑3 B, 20 B, etc.) and architectures (Pythia, PaLM, Falcon, LLaMA 1/2).
  4. Run training with the helper script:
    python3 deepy.py train.py /path/to/your/config.yml
    
  5. Monitor via WandB, Comet, or TensorBoard, and optionally checkpoint to S3.
  6. Export the final checkpoint to Hugging Face format for inference.

When it’s not the right tool

  • You only need inference on an existing model (use transformers).
  • You lack access to multi‑GPU clusters or the budget for billion‑parameter training.
  • You prefer a higher‑level, single‑node trainer (e.g., Lightning, FastChat) that abstracts away the low‑level launch details.

Where to find more

  • Adoption & publications section lists labs and papers that have used GPT‑NeoX.
  • Licensing & contributing details are in the repo’s administrative notes.
  • Containers are under containers/ for Docker or Apptainer.

Bottom line: GPT‑NeoX is a serious, production‑grade training framework for large language models, aimed at researchers and engineers who need to run massive distributed training jobs across a variety of HPC environments. It is not a toy demo or a simple inference library.

Related

  • Project
  • Project
  • Project
  • Project