An implementation of model parallel autoregressive transformers on GPUs, based on the Megatron and DeepSpeed libraries
An implementation of model parallel autoregressive transformers on GPUs, based on the Megatron and DeepSpeed libraries
This repository records EleutherAI's library for training large-scale language models on GPUs. Our current framework is based on NVIDIA's Megatron Language Model and has been augmented with techniques from DeepSpeed as well as some novel optimizations. We aim to make this repo a centralized and accessible place to gather techniques for training large-scale autoregressive language models, and accelerate research into large-scale training. This library is in widespread use in academic, industry, and government labs, including by researchers at Oak Ridge National Lab, CarperAI, Stability AI, Together.ai, Korea University, Carnegie Mellon University, and the University of Tokyo among others. Uniquely among similar libraries GPT-NeoX supports a wide variety of systems and hardwares, including launching via Slurm, MPI, and the IBM Job Step Manager, and has been run at scale on AWS, CoreWeave, ORNL Summit, ORNL Frontier, LUMI, and others.
If you are not looking to train models with billions of parameters from scratch, this is likely the wrong library to use. For generic inference needs, we recommend you use the Hugging Face transformers library instead which supports GPT-NeoX models.
GPT-NeoX leverages many of the same features and technologies as the popular Megatron-DeepSpeed library but with substantially increased usability and novel optimizations. Major features include:
[10/9/2024] We now support Transformer Engine integration
[9/9/2024] We now support preference learning via DPO, KTO, and reward modeling
[9/9/2024] We now support integration with Comet ML, a machine learning monitoring platform
[5/21/2024] We now support RWKV with pipeline parallelism!. See the PRs for RWKV and RWKV+pipeline
[3/21/2024] We now support Mixture-of-Experts (MoE)
[3/17/2024] We now support AMD MI250X GPUs
[3/15/2024] We now support Mamba with tensor parallelism! See the PR
[8/10/2023] We now support checkpointing with AWS S3! Activate with the s3_path config option (for more detail, see the PR)
[9/20/2023] As of https://github.com/EleutherAI/gpt-neox/pull/1035, we have deprecated Flash Attention 0.x and 1.x, and migrated support to Flash Attention 2.x. We don't believe this will cause problems, but if you have a specific use-case that requires old flash support using the latest GPT-NeoX, please raise an issue.
[8/10/2023] We have experimental support for LLaMA 2 and Flash Attention v2 supported in our math-lm project that will be upstreamed later this month.
[5/17/2023] After fixing some miscellaneous bugs we now fully support bf16.
[4/11/2023] We have upgraded our Flash Attention implementation to now support Alibi positional embeddings.
[3/9/2023] We have released GPT-NeoX 2.0.0, an upgraded version built on the latest DeepSpeed which will be regularly synced with going forward.
Prior to 3/9/2023, GPT-NeoX relied on DeeperSpeed, which was based on an old version of DeepSpeed (0.3.15). In order to migrate to the latest upstream DeepSpeed version while allowing users to access the old versions of GPT-NeoX and DeeperSpeed, we have introduced two versioned releases for both libraries:
This codebase has primarily developed and tested for Python 3.8-3.10, and PyTorch 1.8-2.0. This is not a strict requirement, and other versions and combinations of libraries may work.
To install the remaining basic dependencies, run:
pip install -r requirements/requirements.txt
pip install -r requirements/requirements-wandb.txt # optional, if logging using WandB
pip install -r requirements/requirements-tensorboard.txt # optional, if logging via tensorboard
pip install -r requirements/requirements-comet.txt # optional, if logging via Comet
from the repository root.
[!Warning] Our codebase relies on DeeperSpeed, our fork of the DeepSpeed library with some added changes. We strongly recommend using Anaconda, a virtual machine, or some other form of environment isolation before continuing. Failure to do so may cause other repositories that rely on DeepSpeed to break.
We now support AMD GPUs (MI100, MI250X) through JIT fused-kernel compilation. Fused kernels will be built and loaded as needed. To avoid waiting during job launching, you can also do the following for manual pre-build:
python
from megatron.fused_kernels import load
load()
This will automatically adapts building process over different GPU vendors (AMD, NVIDIA) without platform specific code changes. To further test fused kernels using pytest, use pytest tests/model/test_fused_kernels.py
To use Flash-Attention, install the additional dependencies in ./requirements/requirements-flashattention.txt or use a PyTorch NGC container with it pre-installed (note that functionality is not guaranteed using versions different from our requirements file). Then set the attention type in your configuration accordingly (see configs). This can provide significant speed-ups over regular attention on certain GPU architectures, including Ampere GPUs (such as A100s); see the repository for more details.
To use Transformer Engine (TE), install the additional dependencies in ./requirements/requirements-transformer-engine.txt or use a PyTorch NGC container with it pre-installed (note that functionality is not guaranteed using versions different from our requirements file). See this config for an example of using TE on a 1.3B model. This can provide significant speed-ups over regular attention on certain GPU architectures, including Ampere and Hopper GPUs; see the repository for more details.
TE provides very efficient kernels for both A100 and H100 GPUs. We've run some sample ablations on A100:
and H100:
NeoX and Deep(er)Speed support training on multiple different nodes and you have the option of using a variety of different launchers to orchestrate multi-node jobs.
In general there needs to be a "hostfile" somewhere accessible with the format:
node1_ip slots=8
node2_ip slots=8
where the first column contains the IP address for each node in your setup and the number of slots is the number of GPUs that node has access to. In your config you must pass in the path to the hostfile with "hostfile": "/path/to/hostfile". Alternatively the path to the hostfile can be in the environment variable DLTS_HOSTFILE.
pdsh is the default launcher, and if you're using pdsh then all you must do (besides ensuring that pdsh is installed in your environment) is set {"launcher": "pdsh"} in your config files.
If using MPI then you must specify the MPI library (DeepSpeed/GPT-NeoX currently supports mvapich, openmpi, mpich, and impi, though openmpi is the most commonly used and tested) as well as pass the deepspeed_mpi flag in your config file:
{
"launcher": "openmpi",
"deepspeed_mpi": true
}
With your environment properly set up and the correct configuration files you can use deepy.py like a normal python script and start (for example) a training job with:
`python3 deep
No open issues yet, or sync has not completed.