← 全部论文 / 2022

GPT-NeoX-20B: An Open-Source Autoregressive Language Model

2022 超大规模稠密 LLM 训练
下载 PDF

GPT-NeoX-20B: An Open-Source Autoregressive Language Model - 中文验证版

英文原始依据卡片：gpt_neox_2022.md

状态：已翻译。

元数据

计算设置

论文明确报告了集群。GPT-NeoX-20B 在十二台 Supermicro AS-4124GO-NART 服务器上训练，每台配备八块 NVIDIA A100-SXM4-40GB GPU 和两颗 AMD EPYC 7532 CPU，共 96 块 A100 GPU。每块 GPU 可以通过四个 ConnectX-6 HCA 之一使用 GPUDirect RDMA 访问 InfiniBand 结构。两台 NVIDIA MQM8700-HS2R 交换机构成 InfiniBand 骨干。软件栈是 PyTorch 1.10.0，使用 CUDA 11.1 和 NCCL 2.10.3 构建，加上基于 Megatron 和 DeepSpeed 的代码库。

模型有 20B 参数、44 层、隐藏大小 6,144 和 64 个注意力头。训练使用 2,048 token 的上下文和约 3.15M token 的 batch（描述为每步 1,538 个上下文），共 150,000 步，使用余弦学习率衰减。模型在 The Pile 上训练，论文描述为来自 22 个数据源的超过 825 GiB 原始文本。

瓶颈

瓶颈是在具有 40 GB A100 的相对较小的开放集群上承载一个密集 20B 模型。论文直接说明，此规模的权重和优化器状态无法放入单块 GPU，因此内存放置决定了并行方案。与 80 GB A100 的超大规模运行相比，GPT-NeoX 具有更少的每设备内存和更少的总 GPU，因此必须将通信最密集的操作保持在节点内，即 NVLink/NVSwitch 带宽可用的地方。

这个边界是具体的。使用混合精度 Adam，一个 20B 模型具有大约 20B x 16 字节的训练状态下限，即约 320 GB，不包括激活值和临时缓冲区。一个 BF16/FP16 推理副本已经是约 40 GB 的权重；在 44 层、64 个头、头维度 96 和上下文 2048 下，batch-1 的 KV 缓存增加约 2.2 GB。这使得单块 40GB GPU 推理在没有更低精度或卸载的情况下不现实，也解释了为什么论文将推理框架化为在更大内存的 A6000 或一对消费级 GPU 上更经济，而微调仍然是集群任务。

通信是第二个瓶颈。张量并行在层内产生频繁的集合通信；流水线并行在阶段边界之间移动激活值；数据并行同步副本。因此论文选择的拓扑并非任意的：张量并行大小 2 和流水线并行大小 4 使那些昂贵的模型并行通信在八 GPU 节点内发生，而数据并行通信跨越节点边界通过 InfiniBand。

第三个瓶颈是 token 效率。在固定的训练步数、上下文长度和 batch size 下，使用更少 token 来表示同一语料库的分词器有效地减少了序列计算，并使模型在每个 token 预算内看到更多文本。论文将分词器设计，尤其是针对代码和 LaTeX 密集数据的空白处理，视为与计算相关的选择。

方法适配

GPT-NeoX 将 Megatron/DeepSpeed 训练适配到其 12 节点硬件。AdamW 通过 ZeRO 扩展，使优化器状态跨 rank 分布。来自 Megatron 的张量并行拆分大型层计算，而流水线并行跨 GPU 拆分模型。作者报告张量并行大小 2 和流水线并行大小 4 对其硬件最有效，因为它们将最重的模型并行流量保持在节点内，并将跨节点通信主要保留给数据并行。

架构包含一个特定的通信优化：注意力和前馈层并行计算然后求和，而不是串行运行。论文解释说，使用算子分片时，每次残差加法可能在前向和后向传递中需要 all-reduce 工作；并行执行注意力和前馈使在单个 all-reduce 之前进行局部归约成为可能。作者引用 Mesh Transformer JAX 中此模式带来 15% 的吞吐量提升，并具有可比拟的早期损失曲线。

其他设计选择减少了复杂性或提高了效率。GPT-NeoX 在前 25% 的嵌入维度上使用旋转位置嵌入，全部使用密集层而非 GPT-3 风格的交替密集/稀疏层，以及一种旨在防止激活值随深度和宽度增长的初始化方式。分词器在 The Pile 上训练，具有 GPT-2 的 50,257 词表大小，一致处理前缀空格，并包含最多 24 长度的重复空格 token 以压缩空白密集的文档。

证据

训练栈实现并保持了每块 GPU 117 TFLOP/s。这是关键的系统证据：在 96 块 A100-40GB GPU 上的公共 Megatron/DeepSpeed 栈可以在训练密集 20B 模型时维持高利用率。

分词器证据是量化的。在 The Pile 验证集上，GPT-2 分词器产生 383,111,734 个 token，而 GPT-NeoX 分词器产生 342,887,807 个 token，比率约为 0.895。论文将其概括为总体上大约 10% 更少的 token，在 EuroParl、GitHub 和 PubMed Central 组件上超过 20% 的节省。即使排除空白 token，GPT-NeoX 分词器在 The Pile 验证集上仍然使用大约 5% 更少的 token。论文中的一个代码示例从 55 个 GPT-2 token 下降到 39 个 GPT-NeoX token。

评估证据显示了一个有用的公共模型，而不仅仅是一次训练演练。在零样本表格中，GPT-NeoX-20B 在 LAMBADA（0.720 对 0.683）、PIQA（0.779 对 0.752）、TriviaQA（0.259 对 0.170）、ARC-Easy（0.723 对 0.670）和 ARC-Challenge（0.380 对 0.340）等任务上优于 GPT-J-6B。与 FairSeq 模型相比，作者说 GPT-NeoX-20B 在 32 项评估中的 22 项上表现更好，在四项上表现较差，在六项上在误差范围内。对于五样本评估，GPT-NeoX-20B 从零样本到五样本提高了 0.0598，远高于 FairSeq 13B 模型的 0.0183。

历史影响

GPT-NeoX 使一个 20B 密集自回归模型及其训练栈成为公共基础设施制品。从历史上看，它位于 GPT-J 规模的开放模型与 BLOOM 和 OPT 等 100B+ 开放科学努力之间：大到需要严肃的模型并行，但小到可以在 12 节点 A100 集群上训练。

本文持久的计算贡献是其操作方案：Megatron/DeepSpeed、ZeRO 优化器内存减少、与节点拓扑匹配的 TP/PP/DP 放置、针对实际语料库优化的分词器，以及权重、代码和中间检查点的发布。它为前沿实验室之外的研究人员提供了密集 LLM 训练和评估的具体路径。

局限

该模型仍远低于 175B 类前沿规模，论文承认推理并非普遍可及：它说推理在一对 RTX 3090 Ti GPU 或单块 A6000 上最经济，而微调需要显著更多的计算。即使权重是公开的，训练访问仍然是集群问题。

还存在方法论局限。超参数是从 GPT-3 和较小实验插值得到的，因为在 20B 规模进行完整扫描是不可行的。The Pile 按原样使用，没有大规模去重，作者把去重留给未来工作。许多设计选择针对代码性能，但论文说他们低估了代码评估的难度和成本，没有评估该领域。最后，架构包含一个已披露的疏忽：在意图使用绑定层归一化的地方应用了两个独立的 LayerNorm，注意到时已来不及重新开始。

链接

  • 计算范式：history/compute_regimes/hyperscale_dense_llm_training/README.md
  • 来源 PDF 和抽取文本见上述元数据。
  • 队列状态：read_complete。