models.t5 包

基类：megatron.core.transformer.module.MegatronModule

T5 的 Masked LM 头

参数

• config (TransformerConfig) – transformer 配置

• parallel_output (bool) – 是否输出 logits 被分布式处理。

• vocab_size (int) – 词汇表大小

• pre_process (bool) – 包含嵌入层

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

forward(hidden_states: torch.Tensor, word_embeddings_weight: torch.Tensor) → torch.Tensor

前向传递。

参数

• hidden_states (Tensor) – 来自解码器的输出隐藏状态

• word_embeddings_weight (Tensor) – 词嵌入权重

返回

logits 张量

返回类型

Tensor

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

T5 语言模型。

参数

• config (TransformerConfig) – transformer 配置

• encoder_config (TransformerConfig) – 编码器 transformer 配置

• transformer_encoder_layer_spec (ModuleSpec) – 编码器的 transformer 层自定义规范

• transformer_decoder_layer_spec (ModuleSpec) – 解码器的 transformer 层自定义规范

• vocab_size (int) – 词汇表大小

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

• pre_process (bool) – 包含嵌入层（用于流水线并行）

• post_process (bool) – 包含输出层（用于流水线并行）

• fp16_lm_cross_entropy (bool, optional) – 默认为 False

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

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

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

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

• seq_len_interpolation_factor (float) – 线性插值 RoPE 以适应更长序列的比例因子。该值必须是大于 1.0 的浮点数。默认为 None。

• add_encoder (bool) – 创建编码器（用于流水线并行）。当使用流水线时，编码器将仅在流水线 ranks 的子集上创建。

• add_decoder (bool) – 包含输出层（用于流水线并行）。与 add_encoder 相同，当使用此模型和流水线时，解码器将仅在流水线 ranks 的子集上创建。

forward(encoder_input_ids: torch.Tensor, decoder_input_ids: torch.Tensor, encoder_attn_mask: torch.Tensor, decoder_attn_mask: torch.Tensor, encoder_decoder_attn_mask: torch.Tensor, lm_labels: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, output_encoder_hidden_only: bool = False, inference_params: Optional[megatron.core.InferenceParams] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None) → torch.Tensor

前向传递。

参数

• encoder_input_ids (Tensor) – 编码器的输入 id

• decoder_input_ids (Tensor) – 解码器的输入 id

• encoder_attn_mask (Tensor) – 编码器的自注意力掩码

• decoder_attn_mask (Tensor) – 解码器的自注意力掩码

• encoder_decoder_attn_mask (Tensor) – 编码器和解码器之间的交叉注意力掩码

• lm_labels (Tensor) – 解码器输出的标签

• inference_params (InferenceParams) – 推理的相关参数

返回

loss 张量

返回类型

Tensor

set_input_tensor(input_tensor)

参见 megatron.model.transformer.set_input_tensor()

sharded_state_dict(prefix: str = '', sharded_offsets: Tuple[Tuple[int, int, int]] = (), metadata: Optional[dict] = None) → megatron.core.dist_checkpointing.mapping.ShardedStateDict

分片状态字典实现，处理编码器和解码器层的重复。

一些层（输出、嵌入）在编码器和解码器之间共享。此方法为它们设置 replica_id，以确保只有一个层实例具有 replica_id (0, 0, 0)。

参数

• prefix (str) – 模块名称前缀。

• sharded_offsets (tuple) – PP 相关偏移量，预计在此模块级别为空。

• metadata (Optional[Dict]) – 控制分片状态字典创建的元数据。

返回

T5Model 的分片状态字典

返回类型

ShardedStateDict

shared_embedding_or_output_weight() → torch.Tensor

共享输入嵌入和输出 logit 权重的函数。