跳转至

DBRX 介绍:一款全新的、强大的开源 LLM 模型

转载自 databricks

dbrx blog header

今天,我们很高兴向大家介绍 DBRX,这是由 Databricks 创建的一种开放的通用 LLM。 在一系列标准基准测试中,DBRX 在已建立的开放 LLM 中树立了新的技术水平。 此外,它为开放社区和构建自己的 LLM 的企业提供了之前仅限于闭源模型 API 的功能; 根据我们的测量,它超越了 GPT-3.5,并且与 Gemini 1.0 Pro 竞争。 它是一种特别强大的代码模型,在编程方面超越了专门的 CodeLLaMA-70B 模型,除了作为通用 LLM 的优势之外。

这种技术水平的提高伴随着训练和推理性能的显著改进。由于其细粒度的专家混合 (MoE) 架构,DBRX 在开放模型中的效率方面取得了突破性进展。 推理速度比 LLaMA2-70B 快 2 倍,而 DBRX 的总参数量和活跃参数量都只有 Grok-1 的 40%。当在 Mosaic AI Model Serving 上托管时, DBRX 可以以每个用户每秒 150 个 tok 的速度生成文本。我们的客户将发现,相比于训练相同最终模型质量的密集模型, 训练 MoE 模型的 FLOP 效率提高了约 2 倍。 从头到尾,我们的 DBRX 整体配方(包括预训练数据、模型架构和优化策略)可以在几乎 4 倍的计算资源下达到与我们上一代 MPT 模型相同的质量。

general knowledge infographic

图 1:DBRX 在语言理解(MMLU)、编程(HumanEval)和数学(GSM8K)等方面超越了已建立的开源模型。

基础模型(DBRX Base) 和微调模型(DBRX Instruct)的权重可在 Hugging Face 上以开放许可证获得。 从今天开始,DBRX 可供 Databricks 的客户通过 API 使用,并且 Databricks 的客户可以从头开始预训练自己的 DBRX 类模型, 或者使用我们的检查点之一继续训练,使用的是与我们构建模型时相同的工具和科学方法。DBRX 已经被集成到我们的 GenAI 动力产品中, 在诸如 SQL 的应用中,早期推出的版本已经超越了 GPT-3.5 Turbo,并且正在与 GPT-4 Turbo 竞争。它在 RAG 任务中也是开放模型和 GPT-3.5 Turbo 中的领先模型。

训练混合专家模型很困难。我们不得不克服各种科学和性能挑战,以构建一个足够稳健的流水线,以便以有效的方式可重复地训练 DBRX 类模型。 现在我们已经做到了这一点,我们拥有了一个独一无二的训练堆栈,允许任何企业从头开始训练世界级的 MoE 基础模型。我们期待与我们的客户分享这种能力,并与社区分享我们的经验教训。

从 Hugging Face(DBRX BaseDBRX Instruct)下载 DBRX, 或在我们的 HF Space 中尝试 DBRX Instruct,或者在 github 上查看我们的模型仓库:databricks/dbrx

DBRX 是什么?

DBRX 是一种基于 transformer 的仅解码器大型语言模型(LLM),使用下一个令牌预测进行训练。它采用精细的混合专家(MoE)架构, 总共有 1320 亿参数,其中 360 亿参数在任何输入上都是活跃的。它在 12T 个文本和代码数据上进行了预训练。 与 Mixtral 和 Grok-1 等其他开放的 MoE 模型相比,DBRX 是细粒度的,意味着它使用更多数量的较小专家。 DBRX 有 16 个专家,并选择其中 4 个,而 Mixtral 和 Grok-1 有 8 个专家,并选择其中 2 个。 这提供了 65 倍更多的专家组合可能性,我们发现这可以提高模型质量。DBRX 使用旋转位置编码(RoPE)、门控线性单元(GLU)和分组查询注意力(GQA)。 它使用 GPT-4 的分词器,该分词器提供在 tiktoken 仓库中提供的。我们根据详尽的评估和扩展实验做出了这些选择。

DBRX 在经过精心策划的 12T 个令牌的数据上进行了预训练,最大上下文长度为 32k 个令牌。我们估计,与用于预训练 MPT 系列模型的数据相比, 这个数据至少每个令牌更好了 2 倍。使用完整的 Databricks 工具套件,包括 Apache Spark™ 和 Databricks 笔记本进行数据处理, Unity Catalog 进行数据管理和治理,以及 MLflow 进行实验跟踪,我们开发了这个新的数据集。我们在预训练中使用了课程学习, 在训练过程中改变数据混合的方式,我们发现这样可以显著提高模型质量。

在基准测试中与领先的开放模型相比的质量

表 1 展示了 DBRX Instruct 和领先的已建立的开放模型的质量。 DBRX Instruct 在复合基准测试、编程和数学基准测试以及 MMLU 方面是领先的。它在标准基准测试上超越了所有聊天或指导微调模型。

复合基准测试。 我们在两个复合基准测试上评估了 DBRX Instruct 和其他模型: Hugging Face Open LLM Leaderboard (包括 ARC-Challenge、HellaSwag、MMLU、TruthfulQA、WinoGrande 和 GSM8k 的平均分)和 Databricks Model Gauntlet (包括超过 30 个任务,涵盖世界知识、常识推理、语言理解、阅读理解、符号问题解决和编程六个类别)。

在我们评估的模型中,DBRX Instruct 在两个复合基准测试中得分最高: Hugging Face Open LLM Leaderboard(74.5%,次高模型为 Mixtral Instruct 的 72.7%)和 Databricks Gauntlet(66.8%,次高模型为 Mixtral Instruct 的 60.7%)。

编程和数学。 DBRX Instruct 在编程和数学方面表现特别出色。在 HumanEval(70.1%,Grok-1 的 63.2%,Mixtral Instruct 的 54.8%, LLaMA2-70B 变种中表现最佳的 32.2%)和 GSM8k(66.9%,Grok-1 的 62.9%,Mixtral Instruct 的 61.1%,LLaMA2-70B 变种中表现最佳的 54.1%)上, 它的得分高于我们评估的其他开放模型。尽管 Grok-1 的参数数量是 DBRX 的 2.4 倍,但 DBRX 在 HumanEval 上的表现优于 Grok-1, 即使 DBRX Instruct 是为通用用途设计的(Meta 在 CodeLLaMA 博客 中报告的 HumanEval 上的得分为 70.1%,超过了专门用于编程的 CodeLLaMA-70B Instruct 模型的 67.8%)。

MMLU。 DBRX Instruct 在 MMLU 上的得分高于我们考虑的所有其他模型,达到了 73.7%。

模型 DBRX Instruct Mixtral Instruct Mixtral Base LLaMA2-70B Chat LLaMA2-70B Base Grok-11
Open LLM Leaderboard2(下面 6 行的平均值) 74.5% 72.7% 68.4% 62.4% 67.9%
ARC-challenge 25-shot 68.9% 70.1% 66.4% 64.6% 67.3%
HellaSwag 10-shot 89.0% 87.6% 86.5% 85.9% 87.3%
MMLU 5-shot 73.7% 71.4% 71.9% 63.9% 69.8% 73.0%
Truthful QA 0-shot 66.9% 65.0% 46.8% 52.8% 44.9%
WinoGrande 5-shot 81.8% 81.1% 81.7% 80.5% 83.7%
GSM8k CoT 5-shot maj@13 66.9% 61.1% 57.6% 26.7% 54.1% 62.9% (8-shot)
Gauntlet v0.34(30+ 个多样任务的平均值) 66.8% 60.7% 56.8% 52.8% 56.4%
HumanEval5 0-Shot, pass@1(编程) 70.1% 54.8% 40.2% 32.2% 31.0% 63.2%

表 1. DBRX Instruct 和领先的开放模型的质量。有关收集数字的详细信息,请参阅脚注。粗体和下划线表示最高得分。

在基准测试中与领先的闭源模型相比的质量

表 2 展示了 DBRX Instruct 和领先的闭源模型的质量。根据每个模型创建者报告的得分,DBRX Instruct 超越了 GPT-3.5(如 GPT-4 论文中所述),并且与 Gemini 1.0 Pro 和 Mistral Medium 竞争。

在我们考虑的几乎所有基准测试中,DBRX Instruct 超越或与 GPT-3.5 相匹配。DBRX Instruct 在 MMLU 方面的表现优于 GPT-3.5 (MMLU 中的总体得分为 73.7%,GPT-3.5 为 70.0%),在 HellaSwag(89.0% vs. 85.5%)和 WinoGrande(81.8% vs. 81.6%)等常识推理方面的表现也优于 GPT-3.5。 DBRX Instruct 在编程和数学推理方面的表现尤为出色,例如在 HumanEval(70.1% vs. 48.1%)和 GSM8k(72.8% vs. 57.1%)方面。 DBRX Instruct 与 Gemini 1.0 Pro 和 Mistral Medium 竞争。 DBRX Instruct 在 Inflection Corrected MTBench、MMLU、HellaSwag 和 HumanEval 上的得分高于 Gemini 1.0 Pro,而 Gemini 1.0 Pro 在 GSM8k 上更强。 DBRX Instruct 和 Mistral Medium 在 HellaSwag 上的得分相似,而 Mistral Medium 在 Winogrande 和 MMLU 上更强, DBRX Instruct 在 HumanEval、GSM8k 和 Inflection Corrected MTBench 上更强。

模型 DBRX Instruct GPT-3.57 GPT-48 Claude 3 Haiku Claude 3 Sonnet Claude 3 Opus Gemini 1.0 Pro Gemini 1.5 Pro Mistral Medium Mistral Large
MT Bench (Inflection corrected, n=5) 8.39 ± 0.08 8.41 ± 0.04 8.54 ± 0.09 9.03 ± 0.06 8.23 ± 0.08 8.05 ± 0.12 8.90 ± 0.06
MMLU 5-shot 73.7% 70.0% 86.4% 75.2% 79.0% 86.8% 71.8% 81.9% 75.3% 81.2%
HellaSwag 10-shot 89.0% 85.5% 95.3% 85.9% 89.0% 95.4% 84.7% 92.5% 88.0% 89.2%
HumanEval 0-Shot pass@1 (Programming) 70.1% temp=0, N=1 48.1% 67.0% 75.9% 73.0% 84.9% 67.7% 71.9% 38.4% 45.1%
GSM8k CoT maj@1 72.8% (5-shot) 57.1% (5-shot) 92.0% (5-shot) 88.9% 92.3% 95.0% 86.5%(maj1@32) 91.7% (11-shot) 66.7% (5-shot) 81.0% (5-shot)
WinoGrande 5-shot 81.8% 81.6% 87.5% 88.0% 86.7%

表 2. DBRX Instruct 和领先的闭源模型的质量。除 Inflection Corrected MTBench(我们在模型端点上自行测量的数据)外, 其他数字均由这些模型的创建者在各自的白皮书中报告。有关详细信息,请参阅脚注。

在长上下文任务和 RAG 上的质量

DBRX Instruct 在训练过程中使用了长达 32K 个令牌的上下文窗口。表 3 对其在性能上进行了与 Mixtral Instruct 和 GPT-3.5 Turbo、GPT-4 Turbo API 的最新版本在一组长上下文基准测试(来自 Lost in the Middle 论文的 KV-Pairs 和 HotpotQAXL,这是 HotPotQA 的修改版本,将任务扩展到更长的序列长度)上的表现进行了比较。GPT-4 Turbo 通常是这些任务中最好的模型。然而,除一个例外外,DBRX Instruct 在所有上下文长度和序列的所有部分上的表现都优于 GPT-3.5 Turbo。总体而言,DBRX Instruct 和 Mixtral Instruct 的性能相似。

模型 DBRX Instruct Mixtral Instruct GPT-3.5 Turbo (API) GPT-4 Turbo (API)
答案在上下文的前三分之一 45.1% 41.3% 37.3%* 49.3%
答案在上下文的中间三分之一 45.3% 42.7% 37.3%* 49.0%
答案在上下文的后三分之一 48.0% 44.4% 37.0%* 50.9%
2K 上下文 59.1% 64.6% 36.3% 69.3%
4K 上下文 65.1% 59.9% 35.9% 63.5%
8K 上下文 59.5% 55.3% 45.0% 61.5%
16K 上下文 27.0% 20.1% 31.7% 26.0%
32K 上下文 19.9% 14.0% 28.5%

表 3. 模型在 KV-Pairs 和 HotpotQAXL 基准测试中的平均性能。粗体是最高得分。 下划线是除 GPT-4 Turbo 外的最高得分。GPT-3.5 Turbo 支持最大上下文长度为 16K, 因此我们无法在 32K 上评估它。GPT-3.5 Turbo 的开始、中间和结尾的平均值仅包括不超过 16K 的上下文。

使用 RAG(检索增强生成)是利用模型上下文的最受欢迎的方法之一。在 RAG 中,从数据库中检索与提示相关的内容,并将其与提示一起提供给模型,以便为模型提供比其原本拥有的更多信息。表 4 展示了 DBRX 在两个 RAG 基准测试(自然问题和 HotPotQA)上的质量,当模型还提供了从维基百科文章语料库中使用嵌入模型 bge-large-en-v1.5 检索的前 10 段内容。DBRX Instruct 与开放模型如 Mixtral Instruct 和 LLaMA2-70B Chat 以及当前版本的 GPT-3.5 Turbo 相竞争。

模型 DBRX Instruct Mixtral Instruct LLaMa2-70B Chat GPT 3.5 Turbo (API) GPT 4 Turbo (API)
自然问题 60.0% 59.1% 56.5% 57.7% 63.9%
HotPotQA 55.0% 54.2% 54.7% 53.0% 62.9%

表 4. 在给定使用 bge-large-en-v1.5 检索的维基百科语料库的前 10 段内容的情况下,模型的性能。 准确率是通过与模型的答案匹配来衡量的。粗体是最高得分。下划线是除 GPT-4 Turbo 外的最高得分。

训练效率

模型质量必须放在训练和使用的效率背景下。这在 Databricks 尤为重要,因为我们构建像 DBRX 这样的模型是为了为客户建立训练自己基础模型的流程。

我们发现训练混合专家模型在训练效率方面提供了显著的改进(表 5)。例如,训练一个较小的 DBRX 家族成员,称为 DBRX MoE-B(23.5B 总参数,6.6B 活跃参数),在 Databricks LLM Gauntlet 上达到 45.5% 的得分所需的 FLOP 数比 LLaMA2-13B 达到 43.8% 的得分所需的 FLOP 数少 1.7 倍。DBRX MoE-B 的活跃参数数量也只有 LLaMA2-13B 的一半。

从整体上看,我们的端到端 LLM 预训练流程在过去十个月中变得几乎更加高效。2023 年 5 月 5 日,我们发布了 MPT-7B,这是一个在 1T 令牌上训练的 7B 参数模型,达到了 Databricks LLM Gauntlet 30.9% 的得分。DBRX 家族的一个成员称为 DBRX MoE-A(7.7B 总参数,2.2B 活跃参数),在 Databricks Gauntlet 上达到了 30.5% 的得分,所需的 FLOP 数比 MPT-7B 达到 30.9% 的得分少了 3.7 倍。这种效率的提高是许多改进的结果,包括使用 MoE 架构、网络的其他架构变化、更好的优化策略、更好的分词以及非常重要的是更好的预训练数据。

独立而言,更好的预训练数据对模型质量产生了重大影响。我们使用 DBRX 预训练数据训练了一个使用 1T 令牌的 7B 模型(称为 DBRX Dense-A)。它在 Databricks Gauntlet 上达到了 39.0% 的得分,而 MPT-7B 的得分为 30.9%。我们估计,我们的新预训练数据每个令牌至少比训练 MPT-7B 的数据更好了 2 倍。换句话说,我们估计只需要一半的令牌数量就可以达到相同的模型质量。我们通过使用 500B 个令牌对 DBRX Dense-A 进行训练来确定这一点;它在 Databricks Gauntlet 上的表现优于 MPT-7B,达到了 32.1%。除了更好的数据质量之外,另一个对令牌效率的重要贡献者可能是 GPT-4 的分词器,它具有大型词汇表,并且被认为在令牌效率方面特别高效。这些关于改善数据质量的经验直接转化为我们的客户用于根据自己的数据训练基础模型的实践和工具。

模型 总参数 活跃参数 Gauntlet 得分 相对 FLOP
DBRX MoE-A 7.7B 2.2B 30.5% 1x
MPT-7B(1T 令牌) 6.7B 30.9% 3.7x
DBRX Dense-A(1T 令牌) 6.7B 39.0% 3.7x
DBRX Dense-A(500B 令牌) 6.7B 32.1% 1.85x
DBRX MoE-B 23.5B 6.6B 45.5% 1x
LLaMA2-13B 13.0B 43.8% 1.7x

表 5. 我们用来验证 DBRX MoE 架构和端到端训练流程的几篇测试文章的详细信息

推理效率

图 2 显示了使用 NVIDIA TensorRT-LLM 在我们优化的服务基础设施和 16 位精度下为 DBRX 和类似模型提供的端到端推理效率。 我们希望这个基准尽可能接近实际使用情况,包括多个用户同时访问同一个推理服务器。 我们每秒生成一个新用户,每个用户请求包含大约 2000 个令牌的提示,并且每个响应包含 256 个令牌。

一般来说,MoE 模型在推理方面比其总参数数量所暗示的更快。这是因为它们对每个输入使用相对较少的参数。 我们发现 DBRX 在这方面也不例外。DBRX 的推理吞吐量比一个 132B 的非 MoE 模型高 2 到 3 倍。

推理效率和模型质量通常是相互制约的:较大的模型通常达到更高的质量,但较小的模型在推理方面更有效。 使用 MoE 架构可以实现比密集模型通常实现的更好的模型质量和推理效率的权衡。例如,DBRX 在质量上优于 LLaMA2-70B, 并且由于活跃参数数量约为一半,DBRX 的推理吞吐量比 LLaMA2-70B 快 2 倍(图 2)。 Mixtral 是 MoE 模型实现的改进 Pareto 前沿的另一个点:它比 DBRX 更小,因此在质量方面较低, 但在推理吞吐量方面更高。Databricks 基础模型 API 的用户可以在我们优化的模型服务平台上每秒看到 DBRX 达到 150 个 tok 的速度,使用 8 位量化。

dbrx inference efficiency

图 2. 在我们优化的服务基础设施上使用 NVIDIA TensorRT-LLM 以 16 位精度进行各种模型配置的推理吞吐量。 模型在整个节点上以张量并行方式运行。输入提示包含大约 2000 个提示令牌,我们生成 256 个输出令牌。每秒生成一个新的用户。

我们如何构建 DBRX

DBRX 是在连接了 3072 个 NVIDIA H100 的 3.2Tbps Infiniband 上进行训练的。 构建 DBRX 的主要过程 - 包括预训练、训练后处理、评估、红队测试和改进 - 在三个月的时间里进行。 这是在 Databricks 进行了几个月的科学、数据集研究和扩展实验的基础上进行的, 更不用说 Databricks 在 LLM 开发方面的多年经验,包括 MPT 和 Dolly 项目以及我们与客户一起构建和投入生产的成千上万个模型。

为了构建 DBRX,我们利用了与我们的客户可用的相同的 Databricks 工具套件。我们使用 Unity Catalog 管理和治理我们的训练数据。我们使用新获得的 Lilac AI 探索这些数据。我们使用 Apache Spark™ 和 Databricks 笔记本来处理和清理这些数据。我们使用我们开源的训练库的优化版本训练 DBRX:MegaBlocks、LLM Foundry、Composer 和 Streaming。我们使用 Mosaic AI Training 服务在数千个 GPU 上管理大规模的模型训练和微调。我们使用 MLflow 记录我们的结果。我们通过 Mosaic AI Model Serving 和 Inference Tables 收集人类反馈,以改善质量和安全性。我们使用 Databricks Playground 手动实验模型。我们发现 Databricks 工具在各自的用途中是最好的,并且我们受益于它们都是统一产品体验的一部分。

在 Databricks 上开始使用 DBRX

如果您想立即开始使用 DBRX,您可以通过 Databricks Mosaic AI Foundation Model APIs 轻松使用。您可以通过我们的按需定价快速入门,并通过我们的 AI Playground 聊天界面查询模型。对于生产应用程序,我们提供了一个预配吞吐量选项,以提供性能保证、支持微调模型,以及额外的安全性和合规性。要私下托管 DBRX,您可以从 Databricks Marketplace 下载模型,并在 Model Serving 上部署模型。

结论

在 Databricks,我们相信每个企业都应该能够在新兴的 GenAI 世界中掌控自己的数据和命运。DBRX 是我们下一代 GenAI 产品的核心支柱,我们期待着我们的客户利用 DBRX 的能力以及我们用于构建它的工具所带来的激动人心的旅程。在过去的一年中,我们与客户一起训练了数千个 LLM。DBRX 只是 Databricks 构建的强大高效模型的一个例子,该模型适用于各种应用,从内部功能到我们的客户的雄心勃勃的用例。

对于任何新模型来说,DBRX 的旅程只是一个开始,最好的工作将由那些在其上构建的人完成:企业和开放社区。这也只是我们在 DBRX 上的工作的开始,您应该期待更多的成果。

贡献

DBRX 的开发由 Mosaic 团队领导,该团队此前构建了 MPT 模型系列,与 Databricks 的各个部门的几十位工程师、律师、采购和财务专家、项目经理、营销人员、设计师和其他贡献者合作。我们对我们的同事、朋友、家人和社区在过去几个月中的耐心和支持表示感谢。

在创建 DBRX 时,我们站在开放和学术界的巨人的肩膀上。通过公开提供 DBRX,我们希望回馈社区,并希望我们将来能够共同构建更伟大的技术。在此背景下,我们非常感谢 Trevor Gale 及其 MegaBlocks 项目的工作和合作(Trevor 的博士导师是 Databricks CTO Matei Zaharia)、PyTorch 团队和 FSDP 项目、NVIDIATensorRT-LLM 项目、vLLM 团队和项目、EleutherAI 和他们的 LLM evaluation 项目、Lilac AI 的 Daniel Smilkov 和 Nikhil Thorat,以及我们在 Allen Institute for Artificial Intelligence (AI2) 的朋友们的工作和合作。

参考