GPT-NeoX 是 EleutherAI 用于在 GPU 上训练大规模语言模型的库。其当前框架基于 NVIDIA 的 Megatron 语言模型，并融合了 DeepSpeed 的技术以及一些新颖的优化。该项目的目标是为训练大规模自回归语言模型提供集中且易于获取的技术集合，加速大规模训练研究。
主要能力：
- 支持 ZeRO 和 3D 并行的分布式训练。
- 支持多种系统和硬件，可通过 Slurm、MPI 和 IBM Job Step Manager 启动，已在 AWS、CoreWeave、ORNL Summit 和 Frontier、LUMI、CINECA Leonardo 等平台上进行过大规模运行。
- 包含前沿架构创新，如旋转和 alibi 位置嵌入、并行前馈注意力层以及 flash attention。
- 为流行架构提供预定义配置，包括 Pythia、PaLM、Falcon 以及 LLaMA 1 和 2。
- 支持课程学习（Curriculum Learning）。
- 易于与开源生态连接，包括 Hugging Face 的 tokenizers 和 transformers 库，可通过 WandB / Comet / TensorBoard 监控实验，并借助 Language Model Evaluation Harness 进行评估。
版本信息： 2023 年 3 月 9 日之前，GPT-NeoX 依赖于基于旧版 DeepSpeed（0.3.15）的 DeeperSpeed。为了迁移到最新的上游 DeepSpeed 版本，同时允许用户访问旧版 GPT-NeoX 和 DeeperSpeed，项目引入两个版本化发布：
- Version 2.0 的 GPT-NeoX 和 DeeperSpeed 是基于最新 DeepSpeed 构建的最新版本，将持续维护。
- Version 1.0 的 GPT-NeoX 和 DeeperSpeed 保留了 GPT-NeoX-20B 和 Pythia Suite 训练时使用的旧稳定版本快照。
使用方式： 代码库主要针对 Python 3.8-3.10 和 PyTorch 1.8-2.0 开发和测试。依赖 DeeperSpeed（项目对 DeepSpeed 库的分支，包含额外修改），推荐使用 Anaconda 或虚拟机进行环境隔离。 支持 AMD GPU（MI100、MI250X），通过 JIT 融合内核编译，融合内核会按需构建和加载。 使用 Flash-Attention 需安装 ./requirements/requirements-flashattention.txt 中的额外依赖，或在预装该库的 PyTorch NGC 容器中运行，然后在配置中设置注意力类型。 使用 Transformer Engine (TE) 需安装 ./requirements/requirements-transformer-engine.txt 中的额外依赖。 支持多节点训练，可使用 pdsh（默认）、MPI（支持 mvapich、openmpi、mpich、impi）或 Slurm 等启动器。需要配置包含 IP 地址和 GPU 槽位数的 hostfile，或在环境变量 DLTS_HOSTFILE 中指定路径。 提供 Apptainer（原 Singularity）和 Docker 容器，位于 containers 目录。 所有功能通过 deepy.py（deepseed 启动器的封装）启动，提供三个主要入口： train.py 用于训练和微调； eval.py 用于使用语言模型评估 harness 评估训练好的模型； generate.py 用于从训练好的模型采样文本。 例如启动训练：python3 deepy.py train.py /path/to/configs/my_model.yml
参数通过 YAML 配置文件传递给 deepy.py，configs 目录中有多种示例 .yml 文件。
适用场景： 该库已被 40 多个组织（学术界、工业界、政府）用于训练 50 多个模型，例如 EleutherAI 的 Pythia、Stability AI 的 StableLM、Together.ai 的 RedPajama-INCITE、橡树岭国家实验室的 FORGE、SB Intuitions 的 Sarashina1、Tilde 的 TildeOpen-30B。Arc Institute 的 Evo 2 使用了由 GPT-NeoX 衍生的 Savanna 框架训练。 适合需要从头训练数十亿参数规模模型的场景。
限制： 如果用户不打算从头训练数十亿参数的模型，该库可能不适合。对于通用推理需求，官方推荐使用支持 GPT-NeoX 模型的 Hugging Face transformers 库。
其他特性： 通过 megablocks 库支持 Dropless Mixture of Experts (DMoE)，兼容现有的 Megatron 张量并行和 DeepSpeed 流水线并行设置。该实现利用现有的张量并行组来分片专家权重，并使用 Sinkhorn 路由避免负载均衡问题。
