Some of the most influential language models of the open-source era were trained by a volunteer collective with no vendor behind them. GPT-NeoX is EleutherAI’s library for training large-scale autoregressive language models on GPU clusters, built on NVIDIA’s Megatron design and DeepSpeed’s ZeRO machinery, and released under the Apache 2.0 license. Its track record is extraordinary for a research codebase: GPT-NeoX-20B, one of the first capable open weights of its size, the entire Pythia scaling suite from 70 million to 12 billion parameters, StableLM for Stability AI, RedPajama-INCITE for Together, and FORGE at Oak Ridge National Laboratory all came out of this tree. The project counts more than 40 organizations, more than 50 models, and more than 60 publications in its USERS.md ledger, and it has run at scale on Frontier, Summit, LUMI, Leonardo, AWS, and CoreWeave on both NVIDIA and AMD hardware.

The repository rewards reading precisely because it refuses to hide the machinery. One YAML file names your architecture, another names your cluster, and the code meets you at the intersection: a launcher that merges configs and spawns processes, a typed argument object that every subsystem reads, Megatron-style sharded layers, a pipe-parallel wrapper, a memory-mapped data path, and checkpoint tooling that uploads straight to blob storage. It is a complete answer to the question of how billion-parameter training is actually organized. As always in this series, what follows is an educational tour of published source code.

GPT-NeoX overview: deepy merges configs and launches training, NeoXArgs drives the pretrain loop, which builds sharded transformer models, feeds them mmap-backed data, and checkpoints into storage with JIT-compiled kernels alongside.

Reading the overview from left to right:

• the deepy launcher wrapper — merges your YAML files and hands them to the DeepSpeed launcher, one process per GPU.
• the train entry point — a few dozen lines that consume arguments and call pretrain.
• the NeoXArgs configuration system — every knob of the framework merged into one validated object.
• the training loop — pretrain, forward_step, train_step, and evaluate in a 1,679-line module.
• the sharded transformer — attention and feedforward blocks cut across tensor-parallel ranks.
• the parallel layers — column and row parallel linears with vocab-parallel embeddings.
• the data loaders — blended, sampled iterators over memory-mapped token files.
• the JIT fused kernels — softmax and rotary kernels compiled on demand for NVIDIA and AMD.

Why You Need This

The first reason is that this is the most readable large-scale trainer in the open-source ecosystem. Where production frameworks interleave abstraction layers until the actual training step is unfindable, GPT-NeoX keeps an almost textbook layout: configuration in one directory, model in another, model-parallel utilities in a third, and a single training module that ties them together. If you have ever wanted to see, in one file, how gradient accumulation, ZeRO steps, gradient noise scale tracking, and evaluation interleave across a real 20-billion-parameter run, the answer is a 1,679-line Python module you can read in an afternoon.

The second reason is architectural literacy. The Megatron patterns implemented here — column-parallel and row-parallel linears, vocab-parallel embeddings that shard the token table, activations gathered and scattered between ranks — became the de facto vocabulary of every serious training stack. The repository also demonstrates the operational side that papers never show: hostfiles and pdsh, MPI and Slurm launchers, no-SSH modes for clusters that forbid direct connections, Docker and Apptainer deployments for national-lab systems, WandB and Comet monitoring, Nsight profiling hooks, and a checkpoint uploader. These are the unglamorous skills that decide whether a training run finishes at all.

The third reason is the model zoo hiding in the model directory. Beyond the standard attention stack you will find a dropless mixture-of-experts implementation routed with Sinkhorn balancing, a Mamba state-space block, RWKV v6 layers with their own CUDA kernel, gMLP blocks, rotary and AliBi positional schemes, and a Transformer Engine swap-in for Hopper-class GPUs. Each is written against the same sharded-layer interface, which makes the directory one of the best comparative anatomy lessons in modern architecture design available anywhere.

The detailed view: entry points on the left, the configuration merge in the middle, the training loop with its diagnostics, the model zoo, model-parallel utilities, and the data and kernel layer at the bottom.

The flow starts with deepy.py, the only command most users ever run. It calls NeoXArgs.consume_deepy_args to merge an arbitrary number of YAML files — later files override earlier ones, which is how the project composes a model config with a cluster config — injects a WandB API key into worker environments if one is configured, and then calls DeepSpeed’s launcher to replicate the script across every GPU in the hostfile. train.py is deliberately tiny: consume the merged arguments, configure distributed settings, build the tokenizer, initialize TensorBoard or Comet if requested, and call pretrain. eval.py and generate.py reuse the same argument system, so evaluation through the lm-evaluation-harness and text sampling both ride the identical configuration path into the loop. The heart of the argument system is the neox_arguments package: arguments.py (1,353 lines) handles merging and validation, neox_args.py (1,325 lines) declares every field as dataclass entries with documentation, and deepspeed_args.py (264 lines) translates user intent into the zero_optimization dictionaries DeepSpeed expects.

pretrain() in megatron/training.py is where the run takes shape. It calls initialize_megatron (194 lines) for process groups and seeds, builds the model through get_model (line 901), which returns a GPT2ModelPipe when pipe-parallel sizes are nonzero — that class, in gpt2_model.py, subclasses the pipe-module base and assigns layers to stages — then builds the optimizer (line 990) and learning-rate scheduler (line 1139), and enters the train loop at line 1465. Each iteration passes through train_step (line 1318) or its pipe variant train_step_pipe (line 1423), with backward_step at line 1299 and forward_step at line 564 handling the gradients and loss. Woven through it are timers, logging (a 537-line module feeding WandB, Comet, and TensorBoard), gradient noise scale tracking (181 lines), and muP parameter scaling support (196 lines). Evaluation lives at line 1618, and checkpointing flows through megatron/checkpointing.py (423 lines): save_ds_checkpoint at line 181, argument re-validation at line 53, old-tag pruning at line 158, and an upload_checkpoint path at line 315 that pushes artifacts to blob storage so a cluster job never dies with its weights locked inside.

The model block is the conceptual center. transformer.py (1,084 lines) defines ParallelMLP at line 88, ParallelSelfAttention at line 268, and ParallelTransformerLayer at line 908, plus pipe-aware subclasses at the bottom of the file. Everything heavy inside those blocks delegates to mpu/layers.py (720 lines): VocabParallelEmbedding at line 98 shards the vocabulary table across ranks so no GPU holds all embedding rows, ColumnParallelLinear at line 383 and RowParallelLinear at line 589 cut weight matrices along input or output dimensions, and the weight-initialization helpers at the top handle the delicate case of initializing sharded matrices identically on every rank. mpu/mappings.py (364 lines) supplies the gather/scatter glue and mpu/initialize.py (289 lines) builds the process groups these layers query at construction time.

Around that core, the zoo rotates through alternatives: moe_mlp.py (370 lines) and router.py (232 lines) implement dropless experts on the tensor-parallel group with Sinkhorn routing; transformer_engine.py (594 lines) swaps the block internals for NVIDIA’s TE kernels on Ampere and Hopper; positional_embeddings.py (217 lines) offers rotary and AliBi as interchangeable attention biases; and the fused_kernels package JIT-compiles masked-softmax and rotary kernels from .h and .cu sources, adapting to AMD’s MI100 and MI250X toolchains without code changes. The data side mirrors this discipline: preprocessing writes a .bin and .idx pair whose format is implemented twice, once in indexed_dataset.py (483 lines) as a mmap reader and once in helpers.cpp (671 lines) as the C++ extension that wrote it, and gpt2_dataset.py (447 lines) slices those token streams into fixed-length training sequences. A pairwise dataset module for preference-style records and a small sampler family round out the directory, so post-training data flows through the same loaders as pretraining. When training is done, tools/ckpts/convert_neox_to_hf.py exports checkpoints to GPTNeoX, Llama, or Mistral architectures for inference, and the post-training directory carries DPO, KTO, REINFORCE, reward-model, and SFT recipes with Llama 3 configs included.

From Install to a 125-Million-Parameter Run

The README’s supported range is Python 3.8 through 3.10 and PyTorch 1.8 through 2.0, and the whole journey fits in three commands:

pip install -r requirements/requirements.txt
python prepare_data.py -d ./data
python ./deepy.py train.py -d configs 125M.yml local_setup.yml

prepare_data.py downloads and tokenizes a sample dataset into the .bin and .idx pair the loaders expect. The -d flag sets a config directory prefix, so the model definition in 125M.yml combines with your local hardware settings in local_setup.yml — the README explicitly recommends this split of model parameters from data-path parameters. One notable dependency warning: the code relies on DeeperSpeed, EleutherAI’s fork of DeepSpeed, and the project strongly recommends a virtual environment because the fork can conflict with stock DeepSpeed installs elsewhere on the machine. Fused kernels compile on first launch, or you can pre-build them from a Python shell. For multi-node work, a hostfile with one “hostname slots=8” line per machine is all the launcher needs, whether you arrive via pdsh, MPI, or Slurm.

Honest limits: this is a pretraining rig, not a serving stack, and the README itself redirects inference users to Hugging Face Transformers; the dependency set is heavy and version-sensitive, anchored to a DeepSpeed fork that wants environment isolation; the configuration surface, while superbly documented, runs to hundreds of fields that interact in ways only experience teaches; and the smallest interesting runs still want a multi-GPU machine, so the barrier to a first experiment is a real GPU, not a laptop. None of that changes the essentials: this is the codebase that democratized billion-parameter pretraining, and its blend of readable structure and hard-won scale lessons makes it one of the most valuable trees an ML engineer can walk through.

Enjoyed this post? Never miss out on future posts by following us