Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
TL;DR 精炼摘要
本文提出了一种高效的层内模型并行方法,在原生PyTorch中轻量实现,解决了超大规模Transformer训练的内存瓶颈。该方法在512个GPU上成功训练出83亿参数模型,达到76%扩展效率,并推动GPT-2、BERT类模型在多个基准测试中取得领先表现。
摘要
Recent work in language modeling demonstrates that training large transformer models advances the state of the art in Natural Language Processing applications. However, very large models can be quite difficult to train due to memory constraints. In this work, we present our techniques for training very large transformer models and implement a simple, efficient intra-layer model parallel approach that enables training transformer models with billions of parameters. Our approach does not require a new compiler or library changes, is orthogonal and complimentary to pipeline model parallelism, and can be fully implemented with the insertion of a few communication operations in native PyTorch. We illustrate this approach by converging transformer based models up to 8.3 billion parameters using 512 GPUs. We sustain 15.1 PetaFLOPs across the entire application with 76% scaling efficiency when compared to a strong single GPU baseline that sustains 39 TeraFLOPs, which is 30% of peak FLOPs. To demonstrate that large language models can further advance the state of the art (SOTA), we train an 8.3 billion parameter transformer language model similar to GPT-2 and a 3.9 billion parameter model similar to BERT. We show that careful attention to the placement of layer normalization in BERT-like models is critical to achieving increased performance as the model size grows. Using the GPT-2 model we achieve SOTA results on the WikiText103 (10.8 compared to SOTA perplexity of 15.8) and LAMBADA (66.5% compared to SOTA accuracy of 63.2%) datasets. Our BERT model achieves SOTA results on the RACE dataset (90.9% compared to SOTA accuracy of 89.4%).
思维导图
论文精读
中文精读
论文基本信息 (Bibliographic Information)
- 标题 (Title): Megatron-LM: 使用模型并行主义训练数十亿参数的语言模型 (Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism)
- 作者 (Authors): Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, Bryan Catanzaro。所有作者均来自 NVIDIA,这是一家在 GPU 和人工智能计算领域处于领导地位的公司。他们的背景为该研究提供了强大的硬件和系统工程支持。
- 发表期刊/会议 (Journal/Conference): 本文以预印本 (Preprint) 形式发布在 arXiv 上。arXiv 是一个开放获取的学术论文发布平台,许多研究成果在正式投递给顶级会议或期刊前会先在此发布,以快速分享研究进展。
- 发表年份 (Publication Year): 2019
- 摘要 (Abstract): 近期研究表明,训练大型 Transformer 模型能够推动自然语言处理 (NLP) 应用的技术水平。然而,由于内存限制,超大规模模型极难训练。在这项工作中,作者们展示了一种训练超大 Transformer 模型的技术,并实现了一种简单、高效的层内模型并行 (intra-layer model parallel) 方法,使得训练拥有数十亿参数的 Transformer 模型成为可能。该方法无需新的编译器或库更改,与流水线模型并行正交且互补,并且只需在原生 PyTorch 中插入少量通信操作即可完全实现。作者们通过在 512 个 GPU 上成功收敛了高达 83 亿参数的 Transformer 模型来展示此方法。在整个应用中,他们实现了 15.1 PetaFLOPs 的持续计算性能,与强大的单 GPU 基线(39 TeraFLOPs,为峰值算力的 30%)相比,扩展效率达到 76%。为了证明大型语言模型可以进一步提升技术水平,他们训练了一个类似 GPT-2 的 83 亿参数语言模型和一个类似 BERT 的 39 亿参数模型。他们发现,对于 BERT 类模型,仔细调整层归一化 (Layer Normalization) 的位置对于随着模型尺寸增大而获得性能提升至关重要。使用 GPT-2 模型,他们在 WikiText103(困惑度 10.8 vs SOTA 15.8)和 LAMBADA(准确率 66.5% vs SOTA 63.2%)数据集上取得了 SOTA(State-of-the-Art,即当时最佳水平)结果。他们的 BERT 模型在 RACE 数据集上取得了 SOTA 结果(准确率 90.9% vs SOTA 89.4%)。
- 原文链接 (Source Link): 论文以预印本形式发布在 arXiv,可从以下链接访问:
整体概括 (Executive Summary)
研究背景与动机 (Background & Motivation - Why)
- 核心问题: 随着深度学习的发展,研究人员发现,语言模型(如 BERT、GPT-2)的参数量越大,其在各种自然语言处理任务上的表现通常越好。然而,模型的尺寸受到了单个计算设备(如 GPU)内存容量的严格限制。一个拥有数十亿参数的模型,连同其梯度和优化器状态(如 Adam 优化器所需的动量和方差),所占用的内存远远超过了单个 GPU 的几十 GB 显存。
- 重要性与挑战: 如何突破单卡内存瓶颈,从而训练更大、更强的模型,成为了推动 NLP 领域发展的关键。当时已有的模型并行方案,如流水线并行 (Pipeline Parallelism) 或依赖特定编译器的张量并行 (Tensor Parallelism),往往实现复杂,需要重写大量代码,或依赖尚在开发中的定制框架,对普通研究者不够友好。
- 切入点与思路: 本文的作者们思考:能否找到一种简单、高效且易于集成到现有深度学习框架(如 PyTorch)中的模型并行方法?他们将目光投向 Transformer 模型内部的结构,发现其核心计算单元——自注意力模块和前馈网络(MLP)——本身具有可以被并行化的结构。由此,他们提出了一种层内模型并行 (intra-layer model parallelism) 的方法,即在单层 Transformer 内部对计算进行切分,而非在不同层之间切分。
核心贡献/主要发现 (Main Contribution/Findings - What)
- 提出了一种简单高效的层内模型并行方法: 这是本文最核心的贡献。该方法通过在 Transformer 层的特定位置插入几个通信原语 (
all-reduce),即可将矩阵运算(GEMM)分布到多个 GPU 上,而无需修改 PyTorch 底层代码或使用新的编译器。 - 证明了该方法卓越的扩展性: 通过实验,作者展示了该方法能够高效地利用大规模 GPU 集群。他们成功地在 512 个 GPU 上训练了一个 83 亿参数的模型,与强大的单 GPU 基线相比,实现了高达 76% 的计算扩展效率,证明了该方法的实用性。
- 发现了稳定训练大型 BERT 模型的关键: 在尝试扩大 BERT 模型时,作者发现原始架构在模型尺寸超过一定规模后会出现训练不稳定的问题。他们通过实验证明,将层归一化 (Layer Normalization) 的位置调整到残差连接之前,可以解决这一问题,使得 BERT 类模型也能够从参数量的增加中稳定获益。这是一个重要的实践发现。
- 训练出当时最先进的大模型并刷新纪录: 基于他们提出的方法,作者训练了当时参数量最大的语言模型之一(83 亿参数的 GPT-2 类模型和 39 亿参数的 BERT 类模型),并在多个权威的 NLP 基准测试(WikiText103, LAMBADA, RACE)上取得了当时最先进的(SOTA)成绩,有力地证明了“更大模型带来更好性能”的观点。
- 开源了代码库
Megatron-LM: 作者将他们的方法实现为一个名为Megatron-LM的代码库并公开发布。这一举动极大地降低了社区研究和训练超大模型的门槛,对后续大模型研究(包括微软和 NVIDIA 后来合作的 Turing-NLG)产生了深远影响。
预备知识与相关工作 (Prerequisite Knowledge & Related Work)
基础概念 (Foundational Concepts)
-
Transformer: Transformer 是一种在 2017 年被提出的深度学习模型架构,现已成为自然语言处理领域的主流模型。它完全基于注意力机制 (Attention Mechanism),摒弃了传统的循环神经网络 (RNN) 和卷积神经网络 (CNN)。其核心组件包括:
- 多头自注意力 (Multi-Head Self-Attention): 这是 Transformer 的关键。它允许模型在处理一个序列(如一个句子)时,计算序列中每个词元 (token) 对其他所有词元的重要性,并据此生成每个词元的新表示。
多头意味着这种注意力计算会并行地进行多次,每个“头”关注序列中不同的依赖关系模式。 - 前馈网络 (Feed-Forward Network, FFN): 在自注意力层之后,每个词元的表示都会经过一个独立的全连接前馈网络(通常由两个线性层和一个非线性激活函数组成),以进行进一步的特征变换。
- 层归一化 (Layer Normalization, LN) 与残差连接 (Residual Connection): 这两个组件用于稳定训练过程。残差连接(或称“跳跃连接”)将层的输入直接加到其输出上,有助于缓解梯度消失问题;层归一化则对每一层的输出进行标准化,使其均值为 0、方差为 1,加速模型收敛。
- 多头自注意力 (Multi-Head Self-Attention): 这是 Transformer 的关键。它允许模型在处理一个序列(如一个句子)时,计算序列中每个词元 (token) 对其他所有词元的重要性,并据此生成每个词元的新表示。
-
数据并行 (Data Parallelism): 这是最常见的多 GPU 训练方式。其工作原理是:
- 将完整的模型副本复制到每个 GPU 上。
- 将一个训练批次 (mini-batch) 的数据切分成多份,每个 GPU 分到一份数据进行前向和反向传播,独立计算梯度。
- 使用一个
All-Reduce操作,将所有 GPU 上的梯度进行求和平均。 - 每个 GPU 使用聚合后的梯度更新自己的模型副本,从而保证所有副本保持同步。
- 局限性: 这种方法要求整个模型必须能装进单个 GPU 的内存中,因此无法用于训练超出单卡内存限制的超大模型。
-
模型并行 (Model Parallelism): 当模型大到无法放入单个 GPU 时,就需要模型并行。其核心思想是将模型本身进行切分,把不同的部分放到不同的 GPU 上。本文主要关注以下两种模型并行方式:
- 层间模型并行 (Inter-layer Model Parallelism): 也称为流水线并行 (Pipeline Parallelism)。它将模型的不同层(或层块)放置在不同的 GPU 上。例如,GPU-1 计算第 1-8 层,然后将其输出传递给 GPU-2 计算第 9-16 层,以此类推。
- 层内模型并行 (Intra-layer Model Parallelism): 它在单个层内部进行切分。例如,一个大的矩阵乘法运算,可以将其权重矩阵切分,让多个 GPU 各自计算一部分,最后再将结果合并。本文提出的
Megatron-LM就是这种方法的典范。
前人工作 (Previous Works)
在 Megatron-LM 提出之前,学术界和工业界已经探索了多种模型并行技术。
-
GPipe (Huang et al., 2018): 这是由谷歌提出的一个著名的流水线并行框架。它将一个批次的数据进一步切分为多个微批次 (micro-batch),并让它们在 GPU 流水线上依次流动,从而提高 GPU 的利用率,减少“流水线空泡” (pipeline bubble) 带来的效率损失。但
GPipe的实现相对复杂,且仍存在一定的空闲时间。 -
Mesh-TensorFlow (Shazeer et al., 2018): 这是另一个由谷歌提出的框架,它引入了一种更通用的分布式张量计算范式。它允许用户在代码中为张量的不同维度指定并行策略,然后由框架自动编译成相应的分布式计算图。这种方法非常灵活强大,但需要学习一种新的编程“语言”或规范,并依赖其特定的编译器。
-
注意力机制的计算公式: 虽然本文未直接复述,但理解其并行化方法需要了解自注意力机制的核心计算。其公式如下:
- 符号解释:
- : 查询 (Query) 矩阵,代表当前词元为了解上下文而发出的“探针”。
- : 键 (Key) 矩阵,代表序列中所有词元可供查询的“标签”。
- : 值 (Value) 矩阵,代表序列中所有词元包含的实际信息。
- : 键 (Key) 向量的维度。除以 是一个缩放因子,用于防止梯度过小。
- : 一种归一化函数,将注意力得分转换为概率分布。
- 这个公式的核心是一系列的矩阵乘法 (
GEMM),这正是Megatron-LM进行并行化切分的对象。
- 符号解释:
技术演进 (Technological Evolution)
NLP 领域的发展趋势是模型越来越大。从早期的词嵌入 (Word Embeddings) 如 Word2Vec,到能够捕捉上下文信息的 ELMo,再到基于 Transformer 的大规模预训练模型 BERT 和 GPT-2,模型参数量从几亿增长到十几亿,甚至更多。这种趋势催生了对更强大、更易用的并行训练技术的需求。Megatron-LM 正是顺应这一趋势而生的关键技术,它使得训练数十亿乃至百亿参数的模型变得更加可行。
差异化分析 (Differentiation)
Megatron-LM 与当时其他主流模型并行方法的区别在于:
- 并行维度不同:
GPipe是层间 (inter-layer) 并行,将模型的连续几层作为一个整体;而Megatron-LM是层内 (intra-layer) 并行,深入到 Transformer 层的内部,对矩阵乘法进行切分。这两种方法是正交的 (orthogonal),意味着它们可以结合使用。 - 实现复杂度不同:
Mesh-TensorFlow需要特定的编译器和编程范式。而Megatron-LM的方法更为“轻量级”,它直接在PyTorch框架内,通过巧妙地组织计算和插入通信调用 (all-reduce) 来实现,无需对框架本身进行修改,也无需新的编译器。这种“微创”的实现方式是其广受欢迎的重要原因。
方法论 (Methodology - Core Technology & Implementation Details)
Megatron-LM 的核心是一种针对 Transformer 架构的层内模型并行 (intra-layer model parallelism) 方案。其精髓在于识别出 Transformer 层中计算量最大的矩阵乘法 (GEMM) 操作,并将其巧妙地切分到多个 GPU 上,同时最小化 GPU 之间的通信开销。
该方法分别对 Transformer 层中的两个主要部分——多层感知机 (MLP) 和自注意力 (Self-Attention) 模块进行并行化。
方法原理 (Methodology Principles)
核心思想是利用矩阵乘法的分块性质。一个大的矩阵乘法 可以通过对权重矩阵 或输入矩阵 进行切分,从而分解为多个小矩阵乘法。Megatron-LM 的巧妙之处在于它选择的切分方式可以与后续的非线性函数(如 GeLU)或独立计算单元(如注意力头)完美结合,从而减少必要的同步和通信。
为此,论文定义了两个共轭的通信算子 和 ,它们是实现并行化的关键构件,本质上是 identity (恒等映射) 和 all-reduce (全局规约) 操作在正向和反向传播中的不同组合。
-
算子 f: 在前向传播中是
identity(什么都不做),在反向传播中是all-reduce(对梯度求和)。 -
算子 g: 在前向传播中是
all-reduce(对输出求和),在反向传播中是identity(什么都不做)。以下是 算子的 PyTorch 实现示例,展示了其简洁性:
class f(torch.autograd.Function):
def forward(ctx, x):
return x
def backward(ctx, gradient):
all_reduce(gradient)
return gradient
方法步骤与流程 (Steps & Procedures)
1. 并行化多层感知机 (MLP Block)
一个标准的 MLP 模块包含两个线性层和一个非线性激活函数,其计算可表示为:
Megatron-LM 的并行化步骤如下 (参考下图 a):
该图像是论文中图3的示意图,展示了带模型并行的Transformer模块结构,包括(a) MLP部分与(b) Self-Attention部分。图中公式为,,等,反映了正向和反向传播中的算子和的不同操作。
- 第一个线性层 (XA): 将权重矩阵 沿列 (column) 方向切分。假设有 个 GPU,则 。每个 GPU 计算 。
- 关键优势: 由于 是按列切分的,每个 GPU 都可以独立计算出输出的一部分 。更重要的是,非线性函数
GeLU可以直接作用于每个局部的 上,无需在 GeLU 之前进行任何通信。这避免了一次昂贵的同步。
- 关键优势: 由于 是按列切分的,每个 GPU 都可以独立计算出输出的一部分 。更重要的是,非线性函数
- 第二个线性层 (YB): 相应地,将权重矩阵 沿行 (row) 方向切分, 。每个 GPU 的输入是 ,它计算一个部分结果 。
- 最终的输出 是所有部分结果之和,即 。
- 通信:
- 前向传播: 在第二个线性层计算之后,需要一个
all-reduce操作来对所有 GPU 上的部分结果 进行求和,得到最终的输出 。这就是图中 g 算子的作用。 - 反向传播: 根据链式法则,在反向传播至第一个线性层之前,需要对输入梯度进行
all-reduce求和。这就是图中 f 算子的作用。
- 前向传播: 在第二个线性层计算之后,需要一个
2. 并行化自注意力模块 (Self-Attention Block)
自注意力模块的核心是计算查询(Q)、键(K)、值(V)矩阵,然后计算注意力得分。Megatron-LM 利用了多头注意力 (multi-head attention) 的内在并行性 (参考上图 b)。
- Q, K, V 的计算: 生成 Q, K, V 的线性变换权重矩阵 均沿列 (column) 方向切分。这等价于将注意力头 (attention heads) 分配给不同的 GPU。例如,一个有 32 个头的模型在 8 个 GPU 上进行模型并行,每个 GPU 将负责 4 个头的计算。
- 注意力计算: 由于每个注意力头的计算是相互独立的,每个 GPU 可以在本地完成其分配到的所有头的注意力计算(包括
softmax),无需任何通信。 - 输出线性层: 在所有 GPU 算完各自的注意力头输出后,这些输出被拼接起来,送入最后的输出线性层。这个线性层的权重矩阵被沿行 (row) 方向切分,与 MLP 中的第二个线性层处理方式完全相同。
- 通信:
- 前向传播: 与 MLP 类似,在输出线性层之后,需要一个
all-reduce( g 算子) 来聚合所有 GPU 的结果。 - 反向传播: 在反向传播至输出线性层之前,需要一个
all-reduce( f 算子)。
- 前向传播: 与 MLP 类似,在输出线性层之后,需要一个
3. 总体通信模式
如下图所示,通过上述设计,一个完整的 Transformer 层(包含自注意力和 MLP)的模型并行实现,在一次完整的前向和反向传播中,总共只需要 4 次 all-reduce 通信(前向 2 次,反向 2 次)。这种高度优化的通信模式是 Megatron-LM 高效的关键。
该图像是论文中的示意图,展示了Transformer层中模型并行的通信操作,包括前向与反向传播过程中共4次All-Reduce通信,分别发生在Self-Attention-Linear和两层Linear-ReLU-Linear模块中。
数学公式与关键细节 (Mathematical Formulas & Key Details)
MLP 并行化公式
-
前向传播:
- 输入为 (在所有 GPU 上是相同的)。
- GPU 计算:
- 所有 GPU 同步计算: 。这一步通过一个
all-reduce(算子 ) 完成。
-
反向传播:
- 设 的梯度为
dZ。 - GPU 计算 的梯度: 。
- GPU 计算 的梯度: 。
- 的梯度
dX需要聚合来自所有分支的贡献。在计算dX之前,需要对 的梯度进行all-reduce求和(算子 )。
- 设 的梯度为
词嵌入层并行化
由于词表大小 (vocabulary size) 通常很大,嵌入层 (embedding layer) 的权重矩阵也十分庞大。Megatron-LM 对其也进行了并行化。
- 权重矩阵 (尺寸为
hidden_sizexvocab_size) 沿词表维度 (列) 进行切分。 - 前向传播: 在输入层,经过 embedding lookup 后,得到的向量需要通过
all-reduce聚合,因为后续的 Transformer 层需要完整的 embedding 向量。 - 反向传播/输出层: 在输出层,为了计算
softmax和交叉熵损失,传统方法需要收集所有 GPU 上的logits(模型输出),这将导致巨大的通信量 (vocab_size维度非常大)。Megatron-LM采用了一个巧妙的优化:将交叉熵损失计算与模型并行的GEMM融合。每个 GPU 只计算其负责的那部分词表的logits和对应的损失,最后只需在所有 GPU 间同步一个标量 (scalar) 损失值即可。这极大地减少了通信量。
实验设置 (Experimental Setup)
数据集 (Datasets)
作者构建了一个大规模、多样化的训练语料库,以充分发挥大模型的潜力。
- 来源: 聚合了多个公开数据集,包括:
- Wikipedia: 维基百科,知识性文本。
- CC-Stories: Common Crawl 数据集中提取的故事性文本。
- RealNews: 一个大规模的新闻文章集合。
- OpenWebtext: 从网络链接中提取的高质量文本,作为 GPT-2 训练数据的开源复现。
- BooksCorpus: 用于 BERT 模型训练,包含大量书籍文本。
- 规模与特点: 聚合后的数据集包含 174 GB 的去重文本。作者特别注意避免下游任务的测试集数据泄露到训练集中,例如,他们移除了 WikiText103 测试集中的维基百科文章。
- 预处理:
- 过滤掉长度小于 128 个词元的文档。
- 使用局部敏感哈希 (Locality-Sensitive Hashing, LSH) 对内容进行去重,去除了 Jaccard 相似度大于 0.7 的重复文档。
评估指标 (Evaluation Metrics)
论文使用了多个指标来评估模型在不同任务上的性能。
-
困惑度 (Perplexity, PPL):
- 概念定义: 困惑度是评估语言模型性能最常用的指标之一。它衡量的是语言模型对于一个测试集的不确定性。一个好的语言模型能够以较高的概率预测出测试集中的词元序列,因此其困惑度会很低。困惑度越低,模型性能越好。
- 数学公式:
- 符号解释:
- : 测试集文本。
- : 模型分词后测试集中的词元 (token) 总数。
- : 测试集在原始分词方案下的词元总数。论文强调使用 进行归一化,以公平地与先前工作比较。
- : 模型在给定前面所有词元 的条件下,预测出当前词元为 的概率。
- : 对概率取对数,即对数似然。整个求和部分计算的是测试集的平均负对数似然。
-
准确率 (Accuracy):
- 概念定义: 准确率用于分类或选择题式任务(如 LAMBADA, RACE)。它衡量的是模型做出正确预测的样本占总样本的比例。准确率越高,模型性能越好。
- 数学公式:
- 符号解释:
Number of Correct Predictions: 模型预测结果与真实标签一致的样本数量。Total Number of Predictions: 测试集中所有样本的总数量。
-
F1 分数 (F1 Score):
- 概念定义: F1 分数是精确率 (Precision) 和召回率 (Recall) 的调和平均值,常用于评估序列标注或抽取式问答任务(如 SQuAD)。它综合了模型“查得准”和“查得全”两方面的能力。
- 数学公式:
- 符号解释:
- (在所有预测为正的样本中,真正为正的比例)
- (在所有真正为正的样本中,被成功预测为正的比例)
TP: 真阳性 (True Positive),FP: 假阳性 (False Positive),FN: 假阴性 (False Negative)。
对比基线 (Baselines)
论文的比较对象主要分为两类:
- 扩展性分析的基线: 作者建立了一个非常强的单 GPU 基线:一个 12 亿参数的模型在单张 NVIDIA V100 32GB GPU 上训练,其计算效率达到了 39 TFLOPs,是理论峰值的 30%。所有多 GPU 的扩展效率都与这个强基线进行比较。
- 下游任务性能的基线: 在评估 WikiText103, LAMBADA, RACE 等任务时,论文将自己的模型结果与当时已发表的最先进 (SOTA) 模型进行比较。这些基线模型包括但不限于:
- 原始的 GPT-2 (Radford et al., 2019)
- RoBERTa (Liu et al., 2019b)
- ALBERT (Lan et al., 2019)
- XLNet (Yang et al., 2019)
实验结果与分析 (Results & Analysis)
核心结果分析 (Core Results Analysis)
1. 模型并行与数据并行的扩展性分析
这是证明 Megatron-LM 方法有效性的核心实验。作者进行了弱扩展 (weak scaling) 测试,即在增加 GPU 数量的同时,按比例增加模型大小,保持每个 GPU 的计算负载大致不变。
-
结果展示: 如下图 1 和图 5 所示,无论是单独使用模型并行,还是结合数据并行,
Megatron-LM都表现出极佳的扩展性。
该图像是论文中的图表,展示了模型并行(蓝色)和模型+数据并行(绿色)在不同GPU数量下的PetaFLOPs表现。横坐标为GPU数量,纵坐标为每秒PetaFLOPs,采用对数刻度。虚线表示线性扩展参考线,反映了性能扩展趋势。
该图像是一个柱状图,展示了模型并行(Model Parallel)和模型加数据并行(Model + Data Parallel)在不同GPU数量下的弱扩展效率。图中清晰对比了单一模型并行与结合数据并行的扩展表现,数据以百分比形式呈现。 -
关键数据: 在最大规模的实验中,一个 83 亿参数的模型使用 8 路模型并行和 64 路数据并行(共 512 个 GPU),实现了 15.1 PetaFLOPs 的持续算力,扩展效率高达 76%(与单卡1.2B模型基准相比)。这证明了该方法能够有效利用大规模硬件资源,通信开销控制得非常好。
以下是用于扩展性研究的模型配置(转录自原文 Table 1):
Hidden Size Attention heads Number of layers Number of parameters (billions) Model parallel GPUs Model +data parallel GPUs 1536 16 40 1.2 1 64 1920 20 54 2.5 2 128 2304 24 64 4.2 4 256 3072 32 72 8.3 8 512
2. GPT-2 模型的语言建模结果
作者训练了不同规模的 GPT-2 类模型,以验证模型规模与性能的关系。
-
模型配置 (转录自原文 Table 2):
Parameter Count Layers Hidden Size Attn Heads Hidden Size per Head Total GPUs Time per Epoch (days) 355M 24 1024 16 64 64 0.86 2.5B 54 1920 20 96 128 2.27 8.3B 72 3072 24 128 512 2.10 -
性能趋势: 如下图 6 所示,随着模型参数从 355M 增加到 8.3B,模型在验证集上的困惑度持续下降,并且更大模型收敛得更快,最终达到更低的困惑度值。
该图像是图表,展示了不同规模语言模型在验证集上的困惑度随训练迭代次数的变化。图中三条曲线分别代表355M、2.5B和8.3B参数模型,结果显示更大的模型收敛更快且达到更低的困惑度。 -
SOTA 结果 (转录自原文 Table 3): 8.3B 参数的模型在两个零样本 (zero-shot) 任务上均取得了当时的 SOTA 成绩,强有力地证明了“规模法则 (Scaling Law)”的有效性。
Model Wikitext103 Perplexity ↓ LAMBADA Accuracy ↑ 355M 19.31 45.18% 2.5B 12.76 61.73% 8.3B 10.81 66.51% Previous SOTA 15.79 63.24%
3. BERT 模型的微调结果与架构发现
这是本文的另一大亮点。作者在扩大 BERT 模型时,发现了一个关键的架构问题。
-
关键发现: 原始的 BERT 架构(图 7a, LayerNorm 在残差连接之后)在模型规模增大时会出现训练不稳定和性能下降的问题。作者通过实验发现,将 LayerNorm 移到残差连接的输入端(图 7b, 即所谓的
Pre-LN结构)可以解决这个问题,使得更大的 BERT 模型能够稳定训练并获得性能提升。
该图像是论文中的示意图与图表,展示了BERT模型原始架构(a)与重新排列架构(b)的区别及其在336M和752M模型上的训练损失表现。图表显示重新排列架构在752M模型上实现了更稳定且更低的训练损失。 -
模型配置 (转录自原文 Table 4):
Parameter Count Layers Hidden Size Attention Heads Total GPUs 336M 24 1024 16 128 1.3B 24 2048 32 256 3.9B 48 2560 40 512 -
性能提升: 如下表(转录自原文 Table 5)所示,采用了新架构后,BERT 模型的性能随着参数量从 336M 增加到 3.9B,在所有下游任务上都获得了单调提升。3.9B 参数的模型在多个任务上超越了之前的 SOTA 模型,并在 RACE 数据集上取得了新的 SOTA 记录。
Model trained tokens ratio MNLI m/mm accuracy (dev set) QQP accuracy (dev set) SQuAD 1.1 F1/EM (dev set) SQuAD 2.0 F1/EM (dev set) RACE m/h accuracy (test set) RoBERTa (Liu et al., 2019b) 2 90.2 / 90.
相似论文推荐
基于向量语义检索推荐的相关论文。