Hugging Face 的 LOGO

Transformers 文档

BitNet

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

BitNet

概述

该模型在 4 万亿个 token 的语料库上进行训练，证明了原生的 1-bit LLM 能够达到与同等规模领先的开源全精度模型相当的性能，同时在计算效率（内存、能耗、延迟）方面具有显著优势。

➡️ 技术报告： BitNet b1.58 2B4T 技术报告

➡️ 官方推理代码： microsoft/BitNet (bitnet.cpp)

模型变体

Hugging Face 上提供了多个版本的模型权重

模型详情

  • 架构： 基于 Transformer，使用 BitLinear 层（BitNet 框架）进行了修改。
    • 使用旋转位置嵌入 (RoPE)。
    • 在 FFN 层中使用平方 ReLU (ReLU²) 激活函数。
    • 采用 subln 归一化。
    • 线性层或归一化层中没有偏置项。
  • 量化： 原生 1.58 位权重和 8 位激活 (W1.58A8)。
    • 在前向传播过程中，使用 absmean 量化将权重 量化为三元值 {-1, 0, +1}。
    • 使用 absmax 量化（逐个 token）将激活量化为 8 位整数。
    • 至关重要的是，该模型是使用此量化方案从头开始训练的，而不是训练后量化。
  • 参数： 约 20 亿
  • 训练 Token 数： 4 万亿
  • 上下文长度： 最大序列长度为 4096 个 token。
    • 建议： 为了在需要很长上下文的任务（超出预训练长度或用于专门的长推理任务）上获得最佳性能，我们建议在最终微调阶段之前进行中间的长序列适应/训练。
  • 训练阶段
    1. 预训练： 使用两阶段学习率和权重衰减计划，对公共文本/代码和合成数学数据进行大规模训练。
    2. 监督微调 (SFT)： 使用总和损失聚合和特定的超参数调整，在指令遵循和对话数据上进行微调。
    3. 直接偏好优化 (DPO)： 使用偏好对与人类偏好对齐。
  • 分词器： LLaMA 3 分词器（词汇表大小：128,256）。

使用技巧

关于效率的非常重要的说明

在使用标准 `transformers` 库运行此模型时，请不要期望获得性能效率上的提升（在速度、延迟或能耗方面）。

`transformers` 内部当前的执行路径不包含利用 BitNet 架构优势所需的专门、高度优化的计算内核。通过 `transformers` 运行此模型可能会导致推理速度和能耗与此框架内的标准全精度模型相当，甚至可能更差，无论是在 CPU 还是 GPU 上。

虽然由于量化权重您可能会观察到内存使用量减少，但主要的计算效率优势无法通过这种标准的 `transformers` 使用路径获得。

要实现技术报告中展示的效率优势，您必须使用专用的 C++ 实现：bitnet.cpp。

环境要求

pip install transformers

示例

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "microsoft/bitnet-b1.58-2B-4T"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16
)

# Apply the chat template
messages = [
    {"role": "system", "content": "You are a helpful AI assistant."},
    {"role": "user", "content": "How are you?"},
]
chat_input = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)

# Generate response
chat_outputs = model.generate(chat_input, max_new_tokens=50)
response = tokenizer.decode(chat_outputs[0][chat_input.shape[-1]:], skip_special_tokens=True) # Decode only the response part
print("\nAssistant Response:", response)

BitNetConfig

class transformers.BitNetConfig

< >

( vocab_size = 128256 hidden_size = 2560 intermediate_size = 6912 num_hidden_layers = 30 num_attention_heads = 20 num_key_value_heads = 5 hidden_act = 'relu2' max_position_embeddings = 2048 initializer_range = 0.02 rms_norm_eps = 1e-05 use_cache = True pad_token_id = None bos_token_id = 128000 eos_token_id = 128001 tie_word_embeddings = False rope_theta = 500000.0 attention_bias = False attention_dropout = 0.0 **kwargs )

参数

  • vocab_size (int, 可选, 默认为 128256) — BitNet 模型的词汇表大小。定义了调用 BitNetModel 时传入的 inputs_ids 可以表示的不同 token 的数量。
  • hidden_size (int, 可选, 默认为 2560) — 隐藏表示的维度。
  • intermediate_size (int, 可选, 默认为 6912) — MLP 表示的维度。
  • num_hidden_layers (int, 可选, 默认为 30) — Transformer 解码器中的隐藏层数量。
  • num_attention_heads (int, 可选, 默认为 20) — Transformer 解码器中每个注意力层的注意力头数量。
  • num_key_value_heads (int, 可选, 默认为 5) — 这是实现分组查询注意力 (Grouped Query Attention) 时应使用的键值头 (key_value heads) 的数量。如果 num_key_value_heads=num_attention_heads，模型将使用多头注意力 (MHA)；如果 num_key_value_heads=1，模型将使用多查询注意力 (MQA)；否则将使用 GQA。当将一个多头检查点转换为 GQA 检查点时，每个分组的键和值头应通过对该组内所有原始头进行均值池化来构建。更多详情，请查看[这篇论文](https://huggingface.co/papers/2305.13245)。如果未指定，将默认为 num_attention_heads。
  • hidden_act (str 或 function, 可选, 默认为 "relu2") — 解码器中的非线性激活函数（函数或字符串）。
  • max_position_embeddings (int, 可选, 默认为 2048) — 该模型可能使用的最大序列长度。
  • initializer_range (float, 可选, 默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
  • rms_norm_eps (float, 可选, 默认为 1e-05) — rms 归一化层使用的 epsilon 值。
  • use_cache (bool, 可选, 默认为 True) — 模型是否应返回最后一个键/值注意力（并非所有模型都使用）。仅当 config.is_decoder=True 时相关。
  • pad_token_id (int, 可选) — 填充 token 的 id。
  • bos_token_id (int, 可选, 默认为 128000) — 序列开始 token 的 id。
  • eos_token_id (int, 可选, 默认为 128001) — 序列结束 token 的 id。
  • tie_word_embeddings (bool, 可选, 默认为 False) — 是否绑定词嵌入权重
  • rope_theta (float, optional, 默认为 500000.0) — RoPE 嵌入的基础周期。
  • attention_bias (bool, optional, 默认为 False) — 是否在自注意力机制的查询、键、值和输出投影层中使用偏置。
  • attention_dropout (float, optional, 默认为 0.0) — 注意力概率的 dropout 比率。

这是一个用于存储 BitNetModel 配置的配置类。它用于根据指定的参数实例化一个 BitNet 模型，定义模型架构。使用默认值实例化配置将产生与 BitNet b1.58 2B4T microsoft/bitnet-b1.58-2B-4T 相似的配置。

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

>>> from transformers import BitNetModel, BitNetConfig

>>> # Initializing a BitNet style configuration
>>> configuration = BitNetConfig()

>>> # Initializing a model from the BitNet style configuration
>>> model = BitNetModel(configuration)

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

BitNetModel

class transformers.BitNetModel

< >

( config: BitNetConfig )

参数

  • config (BitNetConfig) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型相关的权重，只会加载配置。请查看 from_pretrained() 方法来加载模型权重。

裸的 Bitnet 模型，输出原始的隐藏状态，顶部没有任何特定的头。

该模型继承自 PreTrainedModel。请查看超类的文档，了解该库为其所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、修剪头部等）。

该模型也是 PyTorch torch.nn.Module 的子类。可以像常规 PyTorch 模块一样使用它，并参考 PyTorch 文档了解所有与通用用法和行为相关的事项。

forward

< >

( 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.FloatTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None cache_position: typing.Optional[torch.LongTensor] = None **flash_attn_kwargs: typing_extensions.Unpack[transformers.modeling_flash_attention_utils.FlashAttentionKwargs] ) → transformers.modeling_outputs.BaseModelOutputWithPast 或 tuple(torch.FloatTensor)

参数

  • input_ids (torch.LongTensor，形状为 (batch_size, sequence_length)，可选) — 词汇表中输入序列标记的索引。默认情况下将忽略填充。

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

    什么是输入 ID？

  • attention_mask (torch.Tensor，形状为 (batch_size, sequence_length)，可选) — 用于避免对填充标记索引执行注意力的掩码。掩码值选自 [0, 1]：

    • 1 表示未被屏蔽的标记，
    • 0 表示被屏蔽的标记。

    什么是注意力掩码？

  • position_ids (torch.LongTensor，形状为 (batch_size, sequence_length)，可选) — 位置嵌入中每个输入序列标记的位置索引。选自范围 [0, config.n_positions - 1]。

    什么是位置 ID？

  • past_key_values (~cache_utils.Cache，可选) — 预先计算的隐藏状态（自注意力块和交叉注意力块中的键和值），可用于加速顺序解码。这通常包括在解码的先前阶段由模型返回的 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（那些没有为其提供过去键值状态的 ID），形状为 (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。
  • cache_position (torch.LongTensor，形状为 (sequence_length)，可选) — 描绘输入序列标记在序列中位置的索引。与 position_ids 相反，此张量不受填充影响。它用于在正确的位置更新缓存并推断完整的序列长度。

返回

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

一个 transformers.modeling_outputs.BaseModelOutputWithPast 或一个 `torch.FloatTensor` 元组（如果传递了 `return_dict=False` 或当 `config.return_dict=False` 时），根据配置（BitNetConfig）和输入，包含各种元素。

  • last_hidden_state (torch.FloatTensor, 形状为 (batch_size, sequence_length, hidden_size)) — 模型最后一层输出的隐藏状态序列。

    如果使用了 past_key_values，则只输出形状为 (batch_size, 1, hidden_size) 的序列的最后一个隐藏状态。

  • past_key_values (Cache，可选，当传递 use_cache=True 或 config.use_cache=True 时返回) — 这是一个 Cache 实例。有关更多详细信息，请参阅我们的 kv 缓存指南。

    包含预先计算的隐藏状态（自注意力块中的键和值，如果 `config.is_encoder_decoder=True`，则还包括交叉注意力块中的键和值），可用于（请参阅 `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)`。

    模型在每个层输出的隐藏状态以及可选的初始嵌入输出。

  • attentions (tuple(torch.FloatTensor)，可选，当传递 `output_attentions=True` 或 `config.output_attentions=True` 时返回) — `torch.FloatTensor` 的元组（每层一个），形状为 `(batch_size, num_heads, sequence_length, sequence_length)`。

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

BitNetModel 的 forward 方法，覆盖了 `__call__` 特殊方法。

虽然前向传递的配方需要在此函数内定义，但之后应调用 `Module` 实例而不是此函数，因为前者会处理运行前处理和后处理步骤，而后者会静默地忽略它们。

BitNetForCausalLM

class transformers.BitNetForCausalLM

< >

( config )

参数

  • config (BitNetForCausalLM) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型相关的权重，只会加载配置。请查看 from_pretrained() 方法来加载模型权重。

用于因果语言建模的 Bitnet 模型。

该模型继承自 PreTrainedModel。请查看超类的文档，了解该库为其所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、修剪头部等）。

该模型也是 PyTorch torch.nn.Module 的子类。可以像常规 PyTorch 模块一样使用它，并参考 PyTorch 文档了解所有与通用用法和行为相关的事项。

forward

< >

( 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.FloatTensor] = 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 cache_position: typing.Optional[torch.LongTensor] = None logits_to_keep: typing.Union[int, torch.Tensor] = 0 **kwargs: typing_extensions.Unpack[transformers.models.bitnet.modeling_bitnet.KwargsForCausalLM] ) → transformers.modeling_outputs.CausalLMOutputWithPast 或 tuple(torch.FloatTensor)

参数

  • input_ids (torch.LongTensor，形状为 (batch_size, sequence_length)，可选) — 词汇表中输入序列标记的索引。默认情况下将忽略填充。

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

    什么是输入 ID？

  • attention_mask (torch.Tensor，形状为 (batch_size, sequence_length)，可选) — 用于避免对填充标记索引执行注意力的掩码。掩码值选自 [0, 1]：

    • 1 表示未被屏蔽的标记，
    • 0 表示被屏蔽的标记。

    什么是注意力掩码？

  • position_ids (torch.LongTensor，形状为 (batch_size, sequence_length)，可选) — 位置嵌入中每个输入序列标记的位置索引。选自范围 [0, config.n_positions - 1]。

    什么是位置 ID？

  • past_key_values (~cache_utils.Cache，可选) — 预先计算的隐藏状态（自注意力块和交叉注意力块中的键和值），可用于加速顺序解码。这通常包括在解码的先前阶段由模型返回的 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（那些没有为其提供过去键值状态的 ID），形状为 (batch_size, 1)，而不是所有形状为 (batch_size, sequence_length) 的 input_ids。

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

返回

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

一个 transformers.modeling_outputs.CausalLMOutputWithPast 或一个 `torch.FloatTensor` 元组（如果传递了 `return_dict=False` 或当 `config.return_dict=False` 时），根据配置（BitNetConfig）和输入，包含各种元素。

  • loss (torch.FloatTensor 形状为 (1,)，可选，当提供 labels 时返回) — 语言建模损失（用于下一个 token 预测）。

  • logits (形状为 (batch_size, sequence_length, config.vocab_size) 的 torch.FloatTensor) — 语言建模头部的预测分数（SoftMax 之前的每个词汇标记的分数）。

  • past_key_values (Cache，可选，当传递 use_cache=True 或 config.use_cache=True 时返回) — 这是一个 Cache 实例。有关更多详细信息，请参阅我们的 kv 缓存指南。

    包含预计算的隐藏状态（自注意力块中的键和值），可用于（参见 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)`。

    模型在每个层输出的隐藏状态以及可选的初始嵌入输出。

  • attentions (tuple(torch.FloatTensor)，可选，当传递 `output_attentions=True` 或 `config.output_attentions=True` 时返回) — `torch.FloatTensor` 的元组（每层一个），形状为 `(batch_size, num_heads, sequence_length, sequence_length)`。

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

BitNetForCausalLM 的 forward 方法，覆盖了 `__call__` 特殊方法。

虽然前向传递的配方需要在此函数内定义，但之后应调用 `Module` 实例而不是此函数，因为前者会处理运行前处理和后处理步骤，而后者会静默地忽略它们。

示例

>>> from transformers import AutoTokenizer, BitNetForCausalLM

>>> model = BitNetForCausalLM.from_pretrained("microsoft/bitnet-b1.58-2B-4T")
>>> tokenizer = AutoTokenizer.from_pretrained("microsoft/bitnet-b1.58-2B-4T")

>>> prompt = f'<|begin_of_text|>User: Hey, are you conscious? Can you talk to me?<|eot_id|>Assistant: '
>>> inputs = tokenizer(prompt, return_tensors="pt")

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=100)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"User: Hey, are you conscious? Can you talk to me?Assistant: No, I'm not conscious. I'm an artificial intelligence designed to assist with information and tasks. How can I help you today?"
< > 在 GitHub 上更新

←BioGpt Blenderbot→