Transformers 文档
BitNet
并获得增强的文档体验
开始使用
该模型于 2025-04-16 发布,并于 2025-04-28 添加到 Hugging Face Transformers。
BitNet
概述
该模型在 4 万亿个 token 的语料库上进行训练,证明了原生的 1 位 LLM 在性能上可以媲美同等规模的领先的开源全精度模型,同时在计算效率(内存、能源、延迟)方面具有显著优势。
➡️ 技术报告: BitNet b1.58 2B4T 技术报告
➡️ 官方推理代码: microsoft/BitNet (bitnet.cpp)
模型变体
Hugging Face 上提供有多个版本的模型权重。
microsoft/bitnet-b1.58-2B-4T:包含为高效推理优化的打包的 1.58 位权重。部署时使用此模型。microsoft/bitnet-b1.58-2B-4T-bf16:包含 BF16 格式的主权重。仅用于训练或微调。microsoft/bitnet-b1.58-2B-4T-gguf:包含 GGUF 格式的模型权重,与bitnet.cpp库兼容,可用于 CPU 推理。
模型详情
- 架构: 基于 Transformer,使用
BitLinear层(BitNet 框架)修改。- 使用 Rotary Position Embeddings (RoPE)。
- 在 FFN 层中使用平方 ReLU (ReLU²) 激活函数。
- 采用
subln归一化。 - 线性或归一化层中不包含偏置项。
- 量化: 原生 1.58 位权重和 8 位激活(W1.58A8)。
- 在正向传播过程中,使用 absmean 量化将权重量化为三元值 {-1, 0, +1}。
- 使用 absmax 量化(每个 token)将激活量化为 8 位整数。
- 至关重要的是,该模型是从头开始使用此量化方案进行训练的,而不是在训练后进行量化。
- 参数: ~20 亿
- 训练 Token: 4 万亿
- 上下文长度: 最大序列长度为4096 个 token。
- 建议: 对于需要非常长上下文(超出预训练长度或用于专门的长推理任务)的任务,建议在最终微调阶段之前执行中间的长序列适应/训练。
- 训练阶段
- 预训练: 使用两阶段学习率和权重衰减计划,在公开文本/代码和合成数学数据上进行大规模训练。
- 监督微调 (SFT): 使用总和损失聚合和特定的超参数调整,在指令遵循和对话数据集上进行微调。
- 直接偏好优化 (DPO): 使用偏好对与人类偏好对齐。
- Tokenizer: LLaMA 3 Tokenizer(词汇量大小: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,
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
< source >( vocab_size: int | None = 128256 hidden_size: int | None = 2560 intermediate_size: int | None = 6912 num_hidden_layers: int | None = 30 num_attention_heads: int | None = 20 num_key_value_heads: int | None = 5 hidden_act: str | None = 'relu2' max_position_embeddings: int | None = 2048 initializer_range: float | None = 0.02 rms_norm_eps: int | None = 1e-05 use_cache: bool | None = True pad_token_id: int | None = None bos_token_id: int | None = 128000 eos_token_id: int | None = 128001 tie_word_embeddings: bool | None = False attention_bias: bool | None = False attention_dropout: str | None = 0.0 rope_parameters: transformers.modeling_rope_utils.RopeParameters | dict[str, transformers.modeling_rope_utils.RopeParameters] | None = None **kwargs )
参数
- vocab_size (
int|None, 默认为 128256) — BitNet 模型词汇量大小。定义在调用 BitNetModel 时传递的inputs_ids可以表示的不同 token 的数量。 - hidden_size (
int|None, 默认为 2560) — 隐藏表示的维度。 - intermediate_size (
int|None, 默认为 6912) — MLP 表示的维度。 - num_hidden_layers (
int|None, 默认为 30) — Transformer 解码器中的隐藏层数量。 - num_attention_heads (
int|None, 默认为 20) — Transformer 解码器中每个注意力层的注意力头数量。 - num_key_value_heads (
int|None, 默认为 5) — 这是实现分组查询注意力所需的 key_value 头数量。如果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|None, 默认为"relu2") — 解码器中的非线性激活函数(函数或字符串)。 - max_position_embeddings (
int|None, 默认为 2048) — 该模型可能使用的最大序列长度。 - initializer_range (
float|None, 默认为 0.02) — 用于初始化所有权重矩阵的截断正态初始化器的标准差。 - rms_norm_eps (
float|None, 默认为 1e-05) — rms 归一化层使用的 epsilon 值。 - use_cache (
bool|None, 默认为True) — 是否应返回模型最后的键/值注意力(并非所有模型都使用)。仅当config.is_decoder=True时相关。 - pad_token_id (
int|None) — 填充 token ID。 - bos_token_id (
int|None, 默认为 128000) — 开始流 token ID。 - eos_token_id (
int|None, 默认为 128001) — 结束流 token ID。 - tie_word_embeddings (
bool|None, 默认为False) — 是否绑定词嵌入。 - attention_bias (
bool|None, 默认为False) — 在自注意力过程中是否使用查询、键、值和输出投影层的偏置。 - attention_dropout (
float|None, 默认为 0.0) — 注意力概率的 dropout 比率。 - rope_parameters (
RopeParameters|None) — 包含 RoPE 嵌入配置参数的字典。字典应包含rope_theta的值,如果希望使用具有更长max_position_embeddings的 RoPE,则还可选地包含用于缩放的参数。
这是用于存储 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.configBitNetModel
class transformers.BitNetModel
< source >( config: BitNetConfig )
参数
- config (BitNetConfig) — 模型配置类,包含模型的所有参数。使用配置文件初始化不会加载模型权重,只会加载配置。请查看 from_pretrained() 方法来加载模型权重。
一个不带任何顶部特定头的、输出原始隐藏状态的裸露的 Bitnet 模型。
此模型继承自 PreTrainedModel。查看其父类文档,了解库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、修剪头等)。
此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它,并参考 PyTorch 文档了解一般用法和行为的所有相关信息。
forward
< source >( input_ids: torch.LongTensor | None = None attention_mask: torch.Tensor | None = None position_ids: torch.LongTensor | None = None past_key_values: transformers.cache_utils.Cache | None = None inputs_embeds: torch.FloatTensor | None = None cache_position: torch.LongTensor | None = None use_cache: bool | None = None **kwargs: typing_extensions.Unpack[transformers.utils.generic.TransformersKwargs] ) → transformers.modeling_outputs.BaseModelOutputWithPast 或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensor,形状为(batch_size, sequence_length),可选) — Vocabulary 中的输入序列 token 的索引。默认情况下将忽略 padding。可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.Tensor,形状为(batch_size, sequence_length),可选) — 用于避免对 padding token 索引执行 attention 的掩码。掩码值选择在[0, 1]之间:- 1 表示未掩码的 token,
- 0 表示已掩码的 token。
- position_ids (
torch.LongTensor,形状为(batch_size, sequence_length),可选) — 位置嵌入中每个输入序列 token 的位置索引。选择范围为[0, config.n_positions - 1]。 - past_key_values (
~cache_utils.Cache,可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键值对),可用于加速顺序解码。这通常是当use_cache=True或config.use_cache=True时,模型在上一个解码阶段返回的past_key_values。只有 Cache 实例被允许作为输入,请参阅我们的 kv cache 指南。如果不传递
past_key_values,则默认初始化 DynamicCache。模型将输出与输入相同的缓存格式。
如果使用
past_key_values,用户需要只输入未处理的input_ids(其过去键值状态未传递给此模型),形状为(batch_size, unprocessed_length),而不是所有input_ids,形状为(batch_size, sequence_length)。 - inputs_embeds (
torch.FloatTensor,形状为(batch_size, sequence_length, hidden_size),可选) — 可选地,代替传递input_ids,您可以选择直接传递嵌入表示。如果您想比模型的内部嵌入查找矩阵更精确地控制如何将input_ids索引转换为相关的向量,这很有用。 - cache_position (
torch.LongTensor,形状为(sequence_length),可选) — 描绘输入序列 token 在序列中位置的索引。与position_ids不同,此张量不受 padding 的影响。它用于在正确的位置更新缓存并推断完整的序列长度。 - use_cache (
bool,可选) — 如果设置为True,则返回past_key_values键值状态,并可用于加速解码(请参阅past_key_values)。
返回
transformers.modeling_outputs.BaseModelOutputWithPast or 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, optional, 当传递use_cache=True或当config.use_cache=True时返回) — 它是 Cache 实例。更多详情,请参阅我们的 kv cache 指南。Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if
config.is_encoder_decoder=Truein the cross-attention blocks) that can be used (seepast_key_valuesinput) to speed up sequential decoding. -
hidden_states (
tuple(torch.FloatTensor), optional, 当传递output_hidden_states=True或当config.output_hidden_states=True时返回) —torch.FloatTensor的元组(一个用于嵌入层的输出,如果模型有嵌入层;+一个用于每个层的输出),形状为(batch_size, sequence_length, hidden_size)。模型在每个层输出的隐藏状态以及可选的初始嵌入输出。
-
attentions (
tuple(torch.FloatTensor), optional, 当传递output_attentions=True或当config.output_attentions=True时返回) —torch.FloatTensor的元组(每个层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。注意力 softmax 后的注意力权重,用于计算自注意力头中的加权平均值。
BitNetModel 的 forward 方法,覆盖了 __call__ 特殊方法。
虽然 forward pass 的实现需要在此函数中定义,但你应该在之后调用
Module实例而不是这个,因为前者负责运行预处理和后处理步骤,而后者会静默地忽略它们。
BitNetForCausalLM
class transformers.BitNetForCausalLM
< source >( config model_args: ~utils.generic.ModelArgs | None = None adapter_args: ~utils.generic.AdapterArgs | None = None lora_args: ~utils.generic.LoRAArgs | None = None tokenizer_args: ~utils.generic.TokenizerArgs | None = None dataset_args: ~utils.generic.DatasetArgs | None = None data_args: ~utils.generic.DataArgs | None = None training_args: ~utils.generic.TrainingArgs | None = None generation_args: ~utils.generic.GenerationArgs | None = None vision_tower_args: ~utils.generic.VisionTowerArgs | None = None qlora_args: ~utils.generic.QLoRAArgs | None = None vision_tower_template_args: ~utils.generic.VisionTowerTemplateArgs | None = None video_tower_args: ~utils.generic.VideoTowerArgs | None = None vision_config: ~utils.generic.VisionConfig | None = None video_config: ~utils.generic.VideoConfig | None = None load_dataset: bool | None = None load_data_collator: bool | None = None load_processor: bool | None = None load_lora_adapter: bool | None = None load_adapter: bool | None = None load_qlora_adapter: bool | None = None **kwargs: typing_extensions.Unpack[transformers.modeling_utils.PreTrainedModelKwargs] )
参数
- config (BitNetForCausalLM) — 模型配置类,包含模型的所有参数。使用配置文件初始化不会加载模型权重,只会加载配置。请查看 from_pretrained() 方法来加载模型权重。
用于因果语言建模的 Bitnet 模型。
此模型继承自 PreTrainedModel。查看其父类文档,了解库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、修剪头等)。
此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它,并参考 PyTorch 文档了解一般用法和行为的所有相关信息。
forward
< source >( input_ids: torch.LongTensor | None = None attention_mask: torch.Tensor | None = None position_ids: torch.LongTensor | None = None past_key_values: transformers.cache_utils.Cache | None = None inputs_embeds: torch.FloatTensor | None = None labels: torch.LongTensor | None = None use_cache: bool | None = None cache_position: torch.LongTensor | None = None logits_to_keep: int | torch.Tensor = 0 **kwargs: typing_extensions.Unpack[transformers.utils.generic.TransformersKwargs] ) → transformers.modeling_outputs.CausalLMOutputWithPast 或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensor,形状为(batch_size, sequence_length),可选) — Vocabulary 中的输入序列 token 的索引。默认情况下将忽略 padding。可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.Tensor,形状为(batch_size, sequence_length),可选) — 用于避免对 padding token 索引执行 attention 的掩码。掩码值选择在[0, 1]之间:- 1 表示未掩码的 token,
- 0 表示已掩码的 token。
- position_ids (
torch.LongTensor,形状为(batch_size, sequence_length),可选) — 位置嵌入中每个输入序列 token 的位置索引。选择范围为[0, config.n_positions - 1]。 - past_key_values (
~cache_utils.Cache,可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键值对),可用于加速顺序解码。这通常是当use_cache=True或config.use_cache=True时,模型在上一个解码阶段返回的past_key_values。只有 Cache 实例被允许作为输入,请参阅我们的 kv cache 指南。如果不传递
past_key_values,则默认初始化 DynamicCache。模型将输出与输入相同的缓存格式。
如果使用
past_key_values,用户需要只输入未处理的input_ids(其过去键值状态未传递给此模型),形状为(batch_size, unprocessed_length),而不是所有input_ids,形状为(batch_size, sequence_length)。 - 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的 token 将被忽略(掩码),损失仅为标签在[0, transformers., config.vocab_size]范围内的 token 计算。 - use_cache (
bool, optional) — 如果设置为True,则返回past_key_values键值状态,并可用于加速解码(参见past_key_values)。 - cache_position (
torch.LongTensorof shape(sequence_length), optional) — 描绘了输入序列token在序列中的位置的索引。与position_ids相反,此张量不受填充的影响。它用于在正确的位置更新缓存并推断完整的序列长度。 - logits_to_keep (
Union[int, torch.Tensor], optional, defaults to0) — 如果是int,则为最后logits_to_keep个token计算logits。如果为0,则计算所有input_ids的logits(特殊情况)。仅生成需要最后一个token的logits,并且仅为该token计算logits可以节省内存,这对于长序列或大词汇量大小来说非常重要。如果为torch.Tensor,则必须是1D的,对应于在序列长度维度上要保留的索引。当使用打包张量格式(批次和序列长度的单维度)时,这很有用。
返回
transformers.modeling_outputs.CausalLMOutputWithPast or 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, optional, 当传递use_cache=True或当config.use_cache=True时返回) — 它是 Cache 实例。更多详情,请参阅我们的 kv cache 指南。包含预计算的隐藏状态(自注意力块中的键和值),可用于(参见
past_key_values输入)加速顺序解码。 -
hidden_states (
tuple(torch.FloatTensor), optional, 当传递output_hidden_states=True或当config.output_hidden_states=True时返回) —torch.FloatTensor的元组(一个用于嵌入层的输出,如果模型有嵌入层;+一个用于每个层的输出),形状为(batch_size, sequence_length, hidden_size)。模型在每个层输出的隐藏状态以及可选的初始嵌入输出。
-
attentions (
tuple(torch.FloatTensor), optional, 当传递output_attentions=True或当config.output_attentions=True时返回) —torch.FloatTensor的元组(每个层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。注意力 softmax 后的注意力权重,用于计算自注意力头中的加权平均值。
BitNetForCausalLM的前向方法,重写了__call__特殊方法。
虽然 forward pass 的实现需要在此函数中定义,但你应该在之后调用
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?"