Transformers 文档

DBRX

Transformers

加入 Hugging Face 社区

并获得增强的文档体验

协作开发模型、数据集和 Spaces

通过加速推理获得更快的示例

切换文档主题

开始使用

DBRX

概述

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

DBRX 在 12T tokens 的精心策划的数据和最大上下文长度为 32K tokens 的数据上进行了预训练。我们估计，这种数据在 token 对 token 的质量上至少比我们用于预训练 MPT 模型系列的数据好 2 倍。这个新的数据集是使用 Databricks 全套工具开发的，包括用于数据处理的 Apache Spark™ 和 Databricks notebooks，以及用于数据管理和治理的 Unity Catalog。我们在预训练中使用了课程学习，在训练过程中改变数据组合，我们发现这种方法可以显着提高模型质量。

有关 DBRX Instruct 和 DBRX Base 的更多详细信息，请参阅我们的技术博客文章。

此模型由 eitan-turok 和 abhi-db 贡献。原始代码可以在这里找到，但这可能不是最新的。

使用示例

generate() 方法可用于使用 DBRX 生成文本。您可以使用标准注意力实现、flash-attention 和 PyTorch 缩放点积注意力来生成。最后两种注意力实现可以提高速度。

from transformers import DbrxForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-instruct", token="YOUR_HF_TOKEN")
model = DbrxForCausalLM.from_pretrained(
    "databricks/dbrx-instruct",
    device_map="auto",
    torch_dtype=torch.bfloat16,
    token="YOUR_HF_TOKEN",
    )

input_text = "What does it take to build a great LLM?"
messages = [{"role": "user", "content": input_text}]
input_ids = tokenizer.apply_chat_template(messages, return_dict=True, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=200)
print(tokenizer.decode(outputs[0]))

如果您安装了 flash-attention（pip install flash-attn），则可以更快地生成。（有关 flash-attention 的 HuggingFace 文档可以在这里找到。）

from transformers import DbrxForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-instruct", token="YOUR_HF_TOKEN")
model = DbrxForCausalLM.from_pretrained(
    "databricks/dbrx-instruct",
    device_map="auto",
    torch_dtype=torch.bfloat16,
    token="YOUR_HF_TOKEN",
    attn_implementation="flash_attention_2",
    )

input_text = "What does it take to build a great LLM?"
messages = [{"role": "user", "content": input_text}]
input_ids = tokenizer.apply_chat_template(messages, return_dict=True, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=200)
print(tokenizer.decode(outputs[0]))

您还可以使用 PyTorch 缩放点积注意力来更快地生成。（有关缩放点积注意力的 HuggingFace 文档可以在这里找到。）

from transformers import DbrxForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-instruct", token="YOUR_HF_TOKEN")
model = DbrxForCausalLM.from_pretrained(
    "databricks/dbrx-instruct",
    device_map="auto",
    torch_dtype=torch.bfloat16,
    token="YOUR_HF_TOKEN",
    attn_implementation="sdpa",
    )

input_text = "What does it take to build a great LLM?"
messages = [{"role": "user", "content": input_text}]
input_ids = tokenizer.apply_chat_template(messages, return_dict=True, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=200)
print(tokenizer.decode(outputs[0]))

DbrxConfig

class transformers.DbrxConfig

< source >

( d_model: int = 2048 n_heads: int = 16 n_layers: int = 24 max_seq_len: int = 2048 vocab_size: int = 32000 resid_pdrop: float = 0.0 emb_pdrop: float = 0.0 attn_config: typing.Optional[transformers.models.dbrx.configuration_dbrx.DbrxAttentionConfig] = None ffn_config: typing.Optional[transformers.models.dbrx.configuration_dbrx.DbrxFFNConfig] = None use_cache: bool = True initializer_range: float = 0.02 output_router_logits: bool = False **kwargs: typing.Any )

参数

d_model (int, 可选, 默认为 2048) — 嵌入和隐藏状态的维度。
n_heads (int, 可选, 默认为 16) — Transformer 编码器中每个注意力层的注意力头的数量。
n_layers (int, 可选, 默认为 24) — Transformer 编码器中隐藏层的数量。
max_seq_len (int, 可选, 默认为 2048) — 模型的最长序列长度。
vocab_size (int, 可选, 默认为 32000) — Dbrx 模型的词汇表大小。定义了在调用 DbrxModel 时传递的 inputs_ids 可以表示的不同 token 的最大数量。
resid_pdrop (float, 可选, 默认为 0.0) — 在与残差结合之前应用于注意力输出的 dropout 概率。
emb_pdrop (float, 可选, 默认为 0.0) — 嵌入层的 dropout 概率。
attn_config (dict, 可选) — 用于配置模型注意力模块的字典。
ffn_config (dict, 可选) — 用于配置模型 FFN 模块的字典。
use_cache (bool, 可选, 默认为 True) — 模型是否应返回上次的键/值注意力（并非所有模型都使用）。
initializer_range (float, 可选, 默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
output_router_logits (bool, 可选, 默认为 False) — 模型是否应返回路由器 logits。启用此选项还将允许模型输出辅助损失。有关更多详细信息，请参见此处。

这是用于存储 DbrxModel 配置的配置类。它用于根据指定的参数实例化 Dbrx 模型，从而定义模型架构。使用默认值实例化配置将产生与 databricks/dbrx-instruct 架构不同的配置。

配置对象继承自 PretrainedConfig，可用于控制模型输出。有关更多信息，请阅读 PretrainedConfig 的文档。

示例

>>> from transformers import DbrxConfig, DbrxModel

>>> # Initializing a Dbrx configuration
>>> configuration = DbrxConfig(n_layers=2, d_model=256, n_heads=8, vocab_size=128)

>>> # Initializing a model (with random weights) from the configuration
>>> model = DbrxModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

DbrxModel

class transformers.DbrxModel

< source >

( config: DbrxConfig )

参数

config (DbrxConfig) — 具有模型所有参数的模型配置类。使用配置文件初始化不会加载与模型关联的权重，而只会加载配置。查看 from_pretrained() 方法以加载模型权重。
config (DbrxConfig) — 具有模型所有参数的模型配置类。使用配置文件初始化不会加载与模型关联的权重，而只会加载配置。查看 from_pretrained() 方法以加载模型权重。

裸机 DBRX 模型输出原始隐藏状态，顶部没有任何特定的 head。此模型继承自 PreTrainedModel。查看超类文档，了解库为其所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝 head 等）。

此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch 模块，并参阅 PyTorch 文档，了解与常规用法和行为相关的所有事项。

Transformer 解码器由 config.num_hidden_layers 组成。每一层都是一个 DbrxBlock 层。

前向传播

< source >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None position_ids: typing.Optional[torch.LongTensor] = None past_key_values: typing.Optional[transformers.cache_utils.Cache] = None inputs_embeds: typing.Optional[torch.Tensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None output_router_logits: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None cache_position: typing.Optional[torch.LongTensor] = None **kwargs )

参数

input_ids (torch.LongTensor，形状为 (batch_size, sequence_length)) — 词汇表中输入序列 token 的索引。如果您提供 padding，默认情况下将被忽略。

索引可以使用 AutoTokenizer 获得。请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call() 以了解详细信息。

什么是输入 ID？
attention_mask (torch.Tensor，形状为 (batch_size, sequence_length), 可选) — 用于避免在 padding token 索引上执行注意力的掩码。掩码值在 [0, 1] 中选择：
- 1 表示 token 未被掩盖，
- 0 表示 token 被掩盖。
什么是注意力掩码？

索引可以使用 AutoTokenizer 获得。请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call() 以了解详细信息。

如果使用 past_key_values，则可以选择仅输入最后的 decoder_input_ids（请参阅 past_key_values）。

如果您想更改 padding 行为，则应阅读 modeling_opt._prepare_decoder_attention_mask 并根据您的需要进行修改。有关默认策略的更多信息，请参见论文中的图 1。
- 1 表示 head 未被掩盖，
- 0 表示 head 被掩盖。
position_ids (torch.LongTensor，形状为 (batch_size, sequence_length), 可选) — 位置嵌入中每个输入序列 token 的位置索引。在范围 [0, config.n_positions - 1] 中选择。

什么是位置 ID？
past_key_values (Cache 或 tuple(tuple(torch.FloatTensor)), 可选) — 预先计算的隐藏状态（自注意力模块和交叉注意力模块中的键和值），可用于加速顺序解码。这通常包含模型在先前解码阶段返回的 past_key_values，当 use_cache=True 或 config.use_cache=True 时。

允许两种格式：
- Cache 实例，请参阅我们的 kv 缓存指南；
- 长度为 config.n_layers 的 tuple(torch.FloatTensor) 元组，每个元组包含 2 个形状为 (batch_size, num_heads, sequence_length, embed_size_per_head) 的张量。这也称为传统缓存格式。
模型将输出与作为输入馈送的缓存格式相同的格式。如果未传递 past_key_values，则将返回传统缓存格式。

如果使用 past_key_values，则用户可以选择仅输入最后一次的 input_ids （那些没有将其过去键值状态提供给此模型的）形状为 (batch_size, 1) 而不是所有形状为 (batch_size, sequence_length) 的 input_ids。
inputs_embeds (torch.FloatTensor，形状为 (batch_size, sequence_length, hidden_size)，可选) — （可选）您可以选择直接传递嵌入表示，而不是传递 input_ids。如果您希望比模型的内部嵌入查找矩阵更精细地控制如何将 input_ids 索引转换为关联的向量，这将非常有用。
use_cache (bool，可选) — 如果设置为 True，则返回 past_key_values 键值状态，并可用于加速解码（请参阅 past_key_values）。
output_attentions (bool，可选) — 是否返回所有注意力层的注意力张量。有关更多详细信息，请参阅返回张量下的 attentions。
output_hidden_states (bool，可选) — 是否返回所有层的隐藏状态。有关更多详细信息，请参阅返回张量下的 hidden_states。
output_router_logits (bool，可选) — 是否返回所有路由器的 logits。它们对于计算路由器损失很有用，不应在推理期间返回。
return_dict (bool，可选) — 是否返回 ModelOutput 而不是纯元组。
cache_position (torch.LongTensor，形状为 (sequence_length)，可选) — 索引，描述输入序列 token 在序列中的位置。与 position_ids 相反，此张量不受填充的影响。它用于在正确的位置更新缓存并推断完整序列长度。

DbrxModel forward 方法覆盖了 __call__ 特殊方法。

尽管 forward 传递的配方需要在该函数内定义，但应该在之后调用 Module 实例而不是此函数，因为前者负责运行预处理和后处理步骤，而后者会默默地忽略它们。

DbrxForCausalLM

class transformers.DbrxForCausalLM

< source >

( config: DbrxConfig )

参数

config (DbrxConfig) — 带有模型所有参数的模型配置类。使用配置文件初始化不会加载与模型关联的权重，仅加载配置。查看 from_pretrained() 方法以加载模型权重。

DBRX 模型转换器，用于因果语言建模。此模型继承自 PreTrainedModel。查看超类文档，了解库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch 模块，并参阅 PyTorch 文档，了解与常规用法和行为相关的所有事项。

前向传播

< source >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None position_ids: typing.Optional[torch.LongTensor] = None past_key_values: typing.Optional[transformers.cache_utils.Cache] = None inputs_embeds: typing.Optional[torch.Tensor] = None labels: typing.Optional[torch.LongTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None output_router_logits: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None cache_position: typing.Optional[torch.LongTensor] = None logits_to_keep: typing.Union[int, torch.Tensor] = 0 **kwargs ) → transformers.modeling_outputs.MoeCausalLMOutputWithPast 或 tuple(torch.FloatTensor)

参数

input_ids (torch.LongTensor，形状为 (batch_size, sequence_length)) — 词汇表中输入序列 token 的索引。默认情况下，如果您提供填充，则会忽略填充。

索引可以使用 AutoTokenizer 获得。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什么是输入 ID？
attention_mask (torch.Tensor，形状为 (batch_size, sequence_length)，可选) — 掩码，以避免对填充 token 索引执行注意力机制。掩码值在 [0, 1] 中选择：
- 1 表示 token 未被掩码，
- 0 表示 token 被掩码。
什么是注意力掩码？

索引可以使用 AutoTokenizer 获得。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

如果使用 past_key_values，则可以选择仅输入最后的 decoder_input_ids （请参阅 past_key_values）。

如果您想更改填充行为，则应阅读 modeling_opt._prepare_decoder_attention_mask 并根据您的需要进行修改。有关默认策略的更多信息，请参见论文中的图 1。
- 1 表示头 未被掩码，
- 0 表示头 被掩码。
position_ids (torch.LongTensor，形状为 (batch_size, sequence_length)，可选) — 每个输入序列 token 在位置嵌入中的位置索引。在范围 [0, config.n_positions - 1] 中选择。

什么是位置 ID？
past_key_values (Cache 或 tuple(tuple(torch.FloatTensor)), 可选) — 预先计算的隐藏状态（自注意力模块和交叉注意力模块中的键和值），可用于加速顺序解码。这通常包含模型在先前解码阶段返回的 past_key_values，当 use_cache=True 或 config.use_cache=True 时。

允许两种格式：
- Cache 实例，请参阅我们的 kv 缓存指南；
- 长度为 config.n_layers 的 tuple(torch.FloatTensor) 元组，每个元组包含 2 个形状为 (batch_size, num_heads, sequence_length, embed_size_per_head) 的张量。这也称为传统缓存格式。
模型将输出与作为输入馈送的缓存格式相同的格式。如果未传递 past_key_values，则将返回传统缓存格式。

如果使用 past_key_values，则用户可以选择仅输入最后一次的 input_ids （那些没有将其过去键值状态提供给此模型的）形状为 (batch_size, 1) 而不是所有形状为 (batch_size, sequence_length) 的 input_ids。
inputs_embeds (torch.FloatTensor，形状为 (batch_size, sequence_length, hidden_size)，可选) — （可选）您可以选择直接传递嵌入表示，而不是传递 input_ids。如果您希望比模型的内部嵌入查找矩阵更精细地控制如何将 input_ids 索引转换为关联的向量，这将非常有用。
use_cache (bool，可选) — 如果设置为 True，则返回 past_key_values 键值状态，并可用于加速解码（请参阅 past_key_values）。
output_attentions (bool，可选) — 是否返回所有注意力层的注意力张量。有关更多详细信息，请参阅返回张量下的 attentions。
output_hidden_states (bool，可选) — 是否返回所有层的隐藏状态。有关更多详细信息，请参阅返回张量下的 hidden_states。
output_router_logits (bool，可选) — 是否返回所有路由器的 logits。它们对于计算路由器损失很有用，不应在推理期间返回。
return_dict (bool，可选) — 是否返回 ModelOutput 而不是纯元组。
cache_position (torch.LongTensor，形状为 (sequence_length)，可选) — 索引，描述输入序列 token 在序列中的位置。与 position_ids 相反，此张量不受填充的影响。它用于在正确的位置更新缓存并推断完整序列长度。
labels (torch.LongTensor，形状为 (batch_size, sequence_length)，可选) — 用于计算掩码语言建模损失的标签。索引应在 [0, ..., config.vocab_size] 或 -100 中（请参阅 input_ids 文档字符串）。索引设置为 -100 的 token 将被忽略（掩码），损失仅针对标签在 [0, ..., config.vocab_size] 中的 token 计算。
logits_to_keep (int 或 torch.Tensor，可选) — 如果是 int，则计算最后 logits_to_keep 个 token 的 logits。如果为 0，则计算所有 input_ids 的 logits（特殊情况）。生成只需要最后一个 token 的 logits，并且仅针对该 token 计算它们可以节省内存，这对于长序列或大型词汇表大小来说变得非常重要。如果是 torch.Tensor，则必须是 1D，对应于要在序列长度维度中保留的索引。这在使用打包张量格式（批次和序列长度的单个维度）时很有用。

transformers.modeling_outputs.MoeCausalLMOutputWithPast 或 tuple(torch.FloatTensor)

一个 transformers.modeling_outputs.MoeCausalLMOutputWithPast 或一个 torch.FloatTensor 的元组 (如果传递了 return_dict=False 或当 config.return_dict=False 时)，包含各种元素，具体取决于配置 (DbrxConfig) 和输入。

loss (torch.FloatTensor，形状为 (1,)，可选，当提供 labels 时返回) — 语言建模损失（用于下一个 token 预测）。
logits (torch.FloatTensor，形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测得分（SoftMax 之前的每个词汇 token 的得分）。
aux_loss (torch.FloatTensor, 可选，当提供 labels 时返回) — 用于稀疏模块的 aux_loss。
router_logits (tuple(torch.FloatTensor), 可选，当传递 output_router_probs=True 和 config.add_router_probs=True 或当 config.output_router_probs=True 时返回) — torch.FloatTensor 元组（每层一个），形状为 (batch_size, sequence_length, num_experts)。

原始路由器 logits（post-softmax），由 MoE 路由器计算得出，这些项用于计算 Mixture of Experts 模型的辅助损失。
past_key_values (tuple(tuple(torch.FloatTensor)), 可选，当传递 use_cache=True 或当 config.use_cache=True 时返回) — 长度为 config.n_layers 的 tuple(torch.FloatTensor) 元组，其中每个元组包含 2 个形状为 (batch_size, num_heads, sequence_length, embed_size_per_head) 的张量）。

包含预计算的 hidden-states (隐藏状态)（自注意力模块中的键和值），可以用于（参见 past_key_values 输入）加速顺序解码。
hidden_states (tuple(torch.FloatTensor), 可选，当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — torch.FloatTensor 元组（一个用于词嵌入的输出，如果模型有词嵌入层，+ 每个层的输出各一个），形状为 (batch_size, sequence_length, hidden_size)。

模型在每一层输出的 hidden-states (隐藏状态)，加上可选的初始 embedding (词嵌入) 输出。
attentions (tuple(torch.FloatTensor), 可选，当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — torch.FloatTensor 元组（每层一个），形状为 (batch_size, num_heads, sequence_length, sequence_length)。

注意力 softmax 后的注意力权重，用于计算自注意力头的加权平均值。

DbrxForCausalLM forward 方法，覆盖了 __call__ 特殊方法。

示例

>> from transformers import AutoTokenizer, DbrxForCausalLM

>> model = DbrxForCausalLM.from_pretrained("databricks/dbrx-instruct")
>> tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-instruct")

>> prompt = "Hey, are you conscious? Can you talk to me?"
>> inputs = tokenizer(prompt, return_tensors="pt")

>> # Generate
>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."

< > 在 GitHub 上更新

←CTRL DeBERTa→