models.bert 包

用于训练 bert 和类 bert 仅编码器模型的实用包。它可选地带有一个二元头，可用于分类任务。

class core.models.bert.bert_model.BertModel(*args: Any, **kwargs: Any)

基类： megatron.core.models.common.language_module.language_module.LanguageModule

Transformer 语言模型。

参数

• config (TransformerConfig) – transformer 配置

• num_tokentypes (int) – 当 args.bert_binary_head 为 True 时设置为 2，否则设置为 0。默认为 0。

• transformer_layer_spec (ModuleSpec) – 指定用于 transformer 层的模块

• vocab_size (int) – 词汇表大小

• max_sequence_length (int) – 序列的最大长度。这用于位置嵌入

• pre_process (bool) – 包含嵌入层（与流水线并行一起使用）

• post_process (bool) – 包含输出层（与流水线并行一起使用）

• parallel_output (bool) – 不收集输出，保持它们在张量并行 ranks 中拆分

• share_embeddings_and_output_weights (bool) – 当为 True 时，输入嵌入和输出 logits 权重是共享的。默认为 False。

• position_embedding_type (string) – 位置嵌入类型。选项 ['learned_absolute', 'rope']。默认为 'learned_absolute'。

• rotary_percent (float) – 用于 rotary 位置嵌入的 rotary 维度的百分比。默认为 1.0 (100%)。除非 position_embedding_type 为 'rope'，否则忽略。

bert_extended_attention_mask(attention_mask: torch.Tensor) → torch.Tensor

创建扩展的注意力掩码

将维度为 [batch size, 1, seq len] 的注意力掩码转换为 [batch size, 1, seq len, seq len] 或 [batch size, 1, 1, seq_len]，并使其变为二进制

参数

attention_mask (Tensor) – 输入注意力掩码

返回

扩展的二进制注意力掩码

返回类型

Tensor

bert_position_ids(token_ids)

bert 模型的位置 id

forward(input_ids: torch.Tensor, attention_mask: torch.Tensor, tokentype_ids: Optional[torch.Tensor] = None, lm_labels: Optional[torch.Tensor] = None, inference_params=None)

BERT 模型的前向函数

BERT 模型的前向函数。此函数将输入张量传递通过嵌入层，然后是编码器，最后进入后处理层（可选）。

如果给定标签，它会返回损失值；否则返回最终的隐藏单元

set_input_tensor(input_tensor: torch.Tensor) → None

设置模型的输入张量。

参见 megatron.model.transformer.set_input_tensor()

参数

input_tensor (Tensor) – 设置模型的输入张量。

core.models.bert.bert_model.get_te_version()

包含以实现向后兼容性。