models.gpt 包

这是流行的 GPT 模型的实现。它支持多种功能，如模型并行化（张量并行、流水线并行、数据并行）、混合专家模型、FP8、分布式优化器等。我们不断添加新功能。请密切关注，或者如果您希望添加某些功能，请提出问题。

class core.models.gpt.gpt_model.GPTModel(*args: Any, **kwargs: Any)

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

GPT Transformer 语言模型。

参数

• config (TransformerConfig) – Transformer 配置

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

• vocab_size (int) – 词汇表大小

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

• pre_process (bool, optional) – 包含嵌入层（用于流水线并行）。默认为 True。

• post_process (bool, optional) – 包含输出层（用于流水线并行）。默认为 True。

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

• parallel_output (bool, optional) – 不收集输出，使其在张量并行 ranks 之间保持拆分状态。默认为 True。

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

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

• rotary_percent (float, optional) – 用于旋转位置嵌入的旋转维度百分比。除非 position_embedding_type 为 ‘rope’，否则将被忽略。默认为 1.0。

• rotary_base (int, optional) – 旋转位置嵌入的基周期。除非 position_embedding_type 为 ‘rope’，否则将被忽略。默认为 10000。

• scatter_embedding_sequence_parallel (bool, optional) – 嵌入是否应分散在序列并行区域中。默认为 True。

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

forward(input_ids: torch.Tensor, position_ids: torch.Tensor, attention_mask: torch.Tensor, decoder_input: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, inference_params: Optional[megatron.core.InferenceParams] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, extra_block_kwargs: Optional[dict] = None, runtime_gather_output: Optional[bool] = None) → torch.Tensor

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

如果给定标签，则返回 Loss 值；否则返回最终隐藏单元。

参数

runtime_gather_output (bool) – 在运行时收集输出。默认值 None 表示将使用构造函数中的 parallel_output 参数。

set_input_tensor(input_tensor: torch.Tensor) → None

设置模型的输入张量。

请参阅 megatron.model.transformer.set_input_tensor()

参数

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

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

GPTModel 分片状态字典实现，用于向后兼容（移除额外状态）。

参数

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

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

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

返回

GPTModel 的分片状态字典

返回类型

ShardedStateDict