Databricks DBRX 教程：一步步上手指南

Databricks 最近推出了 DBRX，这是一款采用细粒度专家混合（MoE）架构构建的通用开源大语言模型（LLM）。它并非用单一神经网络处理所有任务，而是由多个专门的“专家”网络组成，每个网络针对不同任务或数据类型进行了优化。

该模型在速度和成本效益上优于 GPT-3.5、Llama 2 等传统 LLM。测试显示，DBRX 在语言理解任务中的得分为 73.7%，高于 LLaMA2（69.8%）。

本文将更详细地介绍其能力，以及如何开始使用 Databricks DBRX。

什么是 Databricks DBRX？

DBRX 采用基于 Transformer 的仅解码器架构，通过下一个词预测进行训练。它使用细粒度的专家混合（MoE）架构。这些“专家”指的是基于 LLM 的专用代理，通过领域知识和推理能力进行了增强。

DBRX 包含数量可观但规模较小的专家（共 16 个），并会针对任一输入选择其中的一个子集（4 个专家）。

简而言之，这种细粒度方法提供了多达 65 倍的专家组合可能性，相比专家更少、每次选择专家数也更少的 Mixtral、Grok-1 等开源 MoE 模型，能显著提升模型质量。

以下是 DBRX 的一些关键信息：

1. 参数规模：DBRX 总参数量为 1320 亿，其中对任一输入处于激活状态的参数为 360 亿。
2. 预训练数据：在精心筛选的 12 万亿个标记上进行预训练，具备 32,000 个标记的上下文长度；按标记计，至少比用于预训练 MPT 系列模型的数据强 2 倍。

DBRX 是如何训练的？

该模型在极其庞大的数据集上进行预训练，其有效性据估计是 Databricks 先前数据集的两倍。训练过程中使用了一系列 Databricks 工具，如用于数据处理的 Apache Spark 与 Databricks 笔记本，以及用于数据治理的 Unity Catalog。

在训练阶段引入了课程式学习，并调整了数据配比，从而显著提升了模型质量。这些对训练数据配比的策略性改动，优化了模型处理多样化输入的能力。

DBRX 预训练所用的一些关键技术包括：

• 旋转位置编码（RoPE）：一种在 Transformer 中通过高维空间的旋转变换来编码标记位置的方法，以增强序列理解。
• 门控线性单元（GLU）：通过门控机制控制信息流动的神经网络层，有助于捕捉复杂模式。
• 分组查询注意力（GQA）：在查询分组内计算注意力权重的机制，以降低计算成本并提升聚焦能力。
• 来自 tiktoken 仓库的 GPT-4 分词器：用于 GPT-4 的特定分词实现，将文本转换为可供模型处理的标记，并针对 GPT-4 的处理需求进行优化。

DBRX 与相似模型对比如何？

Databricks 称，DBRX 在效率与任务表现上优于多款顶级开源模型。

这里有一份关于 DBRX 与竞品对比的详细报告：

1）DBRX vs. LLaMA2-70B

• 通用知识（MMLU）：DBRX Instruct 比 LLaMA2-70B 高出 9.8%。
• 常识推理（HellaSwag 数据集）：DBRX Instruct 继续领先，高出 3.1%。
• Databricks gauntlet（基于不同任务评估预训练 LLM 能力的模型）：DBRX Instruct 以 14% 的优势显著领先 LLaMA2-70B。
• 编程与数学推理：DBRX Instruct 在编程上为 37.9%，在数学推理上为 40.2%。

2）DBRX Instruct vs. Mixtral Instruct

• 通用知识：DBRX Instruct 以 2.3% 的优势略胜一筹。
• 常识推理（HellaSwag 数据集）：DBRX Instruct 高出 1.4%。
• Databricks gauntlet：DBRX Instruct 以 6.1% 的优势明显领先。
• 编程与数学推理：DBRX Instruct 在编程上高出 15.3%，在数学推理上高出 5.8%。

3）DBRX vs Grok-1

• 通用知识：两者表现接近，DBRX Instruct 以 0.7% 略占上风。
• 编程与数学推理：DBRX Instruct 在编程上领先 6.9%，在数学推理上领先 4%。

4）DBRX vs Mixtral Base

• 通用知识：DBRX Instruct 以 1.8% 略有领先。
• 常识推理（HellaSwag 数据集）：DBRX 高出 2.5%。
• Databricks gauntlet：DBRX 领先 Mixtral Base 10%。
• 编程与数学推理：DBRX Instruct 以 29.9% 的优势在编程推理上表现出色。

对比表

下方我们将上述对比汇总为表格，并基于部分结果绘制了图表：

将 DBRX 的质量与其他开源 LLM 对比 - 来源

如何使用 DBRX：一步步上手指南

在访问 DBRX 之前，请确保您的系统至少具备320GB 内存。然后按以下步骤访问 DBRX：

访问 DBRX

• 安装 transformers 库

pip install "transformers>=4.40.0"技术支持

• 然后从 Hugging Face 申请具有只读权限的访问令牌（用于下载模型）。
• 获得访问权限后，使用以下代码导入并加载模型。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-base", token="hf_YOUR_TOKEN")
model = AutoModelForCausalLM.from_pretrained("databricks/dbrx-base", device_map="auto", torch_dtype=torch.bfloat16, token="hf_YOUR_TOKEN")

# Directing tensors to "cuda" (GPU) for faster computation as GPUs are better at handling parallel tasks.
input_text = "Databricks was founded in "
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

# DBRX accepts a context length of up to 32768 tokens. Here`max_new_tokens` specifies the maximum number of tokens to generate.
outputs = model.generate(**input_ids, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))技术支持

想学习如何使用 Transformers 和 Hugging Face？阅读我们关于使用 Transformers 和 Hugging Face 的教程。

使用 DBRX 的基础任务

与其他 LLM 一样，DBRX 可帮助完成基础内容创作与提示响应任务。您可以用 DBRX 做这些事：

• 文本补全：DBRX 能基于您的提示生成文本响应。
• 语言理解：可执行解析自然语言输入、摘要长文档，或在编程语言或脚本风格之间进行转换等任务。
• 查询优化：DBRX 还能优化数据库查询以缩短执行时间，帮助数据工程师在不牺牲性能的前提下专注于其他关键任务。

使用 DBRX 的编程任务

DBRX 也可执行进阶的编码相关任务，帮助数据从业者或开发者：

• 生成代码与解释：可为多种编程任务生成代码，并解释现有函数或为特定问题给出算法建议，以理解并优化代码。

DBRX 对简单指令的响应 - 来源

• 使用 DBRX 进行调试：提供代码更正建议，以提升开发效率并加速调试流程。
• 漏洞识别：还能识别代码中的问题并提出缓解方案。

如何微调 DBRX

您可以使用Github 的开源 LLM Foundry对 DBRX 进行微调。但微调需要将训练样本格式化为字典：

formatted_example = {'prompt': <prompt_text>, 'response': <response_text>}技术支持

Prompt（提示）：这是您给模型的起始问题或指令。

Response（响应）：这是模型被训练要生成的答案。

您可以使用三种不同的数据集来微调任意 LLM：

1. 来自 HuggingFace Hub 的数据集。
2. 本地存储在您设备或机器上的数据集。
3. 采用 StreamingDataset 的 .mds 格式数据集。

1）使用 Hugging Face Hub

如果您想使用来自 HuggingFace Hub 的数据集进行微调，且该

数据集已有预定义的预处理函数或已采用

“prompt”/“response” 格式，您只需将数据加载器指向该数据集即可。

train_loader:
    name: finetuning
    dataset:
        hf_name: tatsu-lab/alpaca
        split: train
        ...技术支持

如果未定义预处理函数，请使用 preprocessing_fn 为数据加载器指定自定义预处理函数。

train_loader:
    name: finetuning
    dataset:
        hf_name: mosaicml/doge-facts
        preprocessing_fn: my_data.formatting:dogefacts_prep_fn
        split: train
        ...
技术支持

2）使用本地数据集

如果您的设备上已有微调数据集，请在 YAML 配置 yamls/finetune/1b_local_data_sft.yaml 中定义本地 JSONL 文件。

train_loader:
    name: finetuning
    dataset:
        hf_name: json # assuming data files are json formatted
        hf_kwargs:
            data_dir: /path/to/data/dir/
        preprocessing_fn: my.import.path:my_preprocessing_fn
        split: train
        ...
技术支持

如果您的本地数据已为 “prompt”/“response” 格式，则可跳过 preprocessing_fn。

3）使用 StreamingDataset

使用

convert_finetuning_dataset.py 脚本将 HuggingFace 数据集转换为 MDS 格式。

将 HuggingFace 数据集转换为流式格式后，只需按如下方式调整 YAML 配置。

train_loader:
    name: finetuning
    dataset:
        remote: s3://my-bucket/my-copy-doge-facts
        local: /tmp/mds-cache/
        split: train
        ...
技术支持

如需查看全参数微调，请参阅以下 YAML 配置：

dbrx-full-ft.yaml

结语

Databricks DBRX 通过多个专门网络来提升模型速度与成本效益。这种细粒度方法使其在处理复杂任务时优于其他 LLM。

想进一步学习大语言模型及其微调方法？请查看以下资源：

• Working with Hugging Face
• 微调 Llama 2
• 微调 Google Gemma
• 微调 Phi-2 模型