transformer package

transformer 软件包提供了一个可定制和可配置的 transformer 模型架构实现。transformer 堆栈的每个组件，从整个层到单个线性层，都可以通过使用 “spec” 参数（参见此处）替换不同的 PyTorch 模块来进行定制。transformer 的配置（隐藏层大小、层数、注意力头数等）通过 TransformerConfig 对象提供。

这是 transformer 层的完整注意力部分，包括自注意力或交叉注意力，包括查询、键和值投影、“核心”注意力计算（例如点积注意力）和最终输出线性投影。

class core.transformer.attention.Attention(*args: Any, **kwargs: Any)

基类： megatron.core.transformer.module.MegatronModule, abc.ABC

注意力层抽象类。

此层仅包含 “self attn” 和 “cross attn” 特化所需的通用模块。

flash_decoding(sequence_len_offset: torch.Tensor, query_layer: torch.Tensor, key_layer: torch.Tensor, value_layer: torch.Tensor, inference_key_memory: torch.Tensor, inference_value_memory: torch.Tensor, rotary_cos: torch.Tensor, rotary_sin: torch.Tensor)

flash 解码内核将在一次执行中执行以下操作：1. 使用预计算的 cos 和 sin 张量计算 RoPE 嵌入 2. 更新 KV 缓存 3. 执行 flash 注意力操作

forward(hidden_states, attention_mask, key_value_states=None, inference_params=None, rotary_pos_emb=None, rotary_pos_cos=None, rotary_pos_sin=None, attention_bias=None, packed_seq_params=None, sequence_len_offset=None)

执行通过注意力模块的前向传递。

abstract get_query_key_value_tensors(hidden_states, key_value_states)

此方法需要根据派生类是 “self-attn” 还是 “cross-attn” 来实现。

class core.transformer.attention.CrossAttention(*args: Any, **kwargs: Any)

基类： core.transformer.attention.Attention

交叉注意力层类

交叉注意力层接受大小为 [s, b, h] 的输入和大小为 [s, b, h] 的上下文，并返回相同大小的输出。

get_query_key_value_tensors(hidden_states, key_value_states)

从 hidden_states 派生 query 张量，并从 key_value_states 派生 key/value 张量。

class core.transformer.attention.CrossAttentionSubmodules(linear_q: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, type]] = None, linear_kv: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, type]] = None, core_attention: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, type]] = None, linear_proj: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, type]] = None)

基类： object

用于指定交叉注意力子模块的配置类。

core_attention: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

linear_kv: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

linear_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

linear_q: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

class core.transformer.attention.SelfAttention(*args: Any, **kwargs: Any)

基类： core.transformer.attention.Attention

自注意力层类

自注意力层接受大小为 [s, b, h] 的输入，并返回相同大小的输出。

get_query_key_value_tensors(hidden_states, key_value_states=None)

从 hidden_states 派生 query、key 和 value 张量。

run_realtime_tests()

执行一致性检查。

此函数确保在实验期间，跨设备的张量是相同的。这通常不能保证，因为可能存在静默硬件故障（例如，加载检查点时内存损坏，数据传输过程中遇到网络流量损坏）。

（TODO）将来，应在训练运行期间检查更多张量，并每 X 次迭代检查一次。这留待未来工作。可能不需要张量的相等性；传输哈希值就足够了。

class core.transformer.attention.SelfAttentionSubmodules(linear_qkv: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, type]] = None, core_attention: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, type]] = None, linear_proj: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, type]] = None, q_layernorm: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, type]] = None, k_layernorm: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, type]] = None)

基类： object

用于指定自注意力子模块的配置类。

core_attention: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

k_layernorm: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

linear_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

linear_qkv: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

q_layernorm: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

这是点积注意力的纯 PyTorch 实现。当训练速度很重要时，通常使用更高效的实现，例如 FlashAttention 或 CUDNN 的 FusedAttention。

class core.transformer.dot_product_attention.DotProductAttention(*args: Any, **kwargs: Any)

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

应用选择性激活重计算的区域。此区域是内存密集型的，但计算密集程度较低，这使得激活检查点对于 LLM (20B+) 更有效。有关更多详细信息，请参阅 Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198。

我们使用以下符号

h：隐藏层大小 n：注意力头数 p：张量模型并行分区数 b：批次大小 s：序列长度

forward(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor, attn_mask_type: Optional[megatron.core.transformer.enums.AttnMaskType] = None, attention_bias: Optional[torch.Tensor] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None)

前向传播。

class core.transformer.enums.AttnBackend(value)

基类： enum.Enum

注意力后端

auto = 5

flash = 1

fused = 2

local = 4

unfused = 3

class core.transformer.enums.AttnMaskType(value)

基类： enum.Enum

注意力掩码类型

arbitrary = 5

causal = 2

no_mask = 3

padding = 1

padding_causal = 4

class core.transformer.enums.AttnType(value)

基类： enum.Enum

注意力类型

cross_attn = 2

self_attn = 1

class core.transformer.enums.ModelType(value)

基类： enum.Enum

模型类型

encoder_or_decoder 用于 bert、gpt 等，encoder_and_decoder 用于多模态、T5 等

encoder_and_decoder = 2

encoder_or_decoder = 1

这提供了一个直通模块，可以在 spec 中使用，以指示不应执行该操作。例如，当将 LayerNorm 与后续线性层一起使用时，可以将 IdentityOp 作为 LayerNorm 模块传入以使用。

class core.transformer.identity_op.IdentityFuncOp(*args: Any, **kwargs: Any)

基类： core.transformer.identity_op.IdentityOp

这是 IdentityFuncOp(…)(x) -> IdentityOp(x) -> x 的占位符。这种函数对于诸如 bias_dropout_fusion 之类的操作非常方便，这些操作本身会在运行时根据传递的参数返回一个函数

forward(*args, **kwargs)

class core.transformer.identity_op.IdentityOp(*args: Any, **kwargs: Any)

基类： torch.nn.Module

这是 IdentityOp(x) -> x 的占位符

forward(x, *args, **kwargs)

这是 transformer 层的完整 MLP 部分，具有输入投影、非线性性和输出投影。

class core.transformer.mlp.MLP(*args: Any, **kwargs: Any)

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

MLP 将接受具有 h 隐藏状态的输入，将其投影到 4*h 隐藏维度，执行非线性变换，并将状态投影回 h 隐藏维度。

返回一个输出和一个要添加到输出的偏差。如果 config.add_bias_linear 为 False，则返回的偏差为 None。

我们使用以下符号

h：隐藏层大小 p：张量模型并行分区数 b：批次大小 s：序列长度

forward(hidden_states)

执行通过 MLP 块的前向传递。

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

class core.transformer.mlp.MLPSubmodules(linear_fc1: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None, linear_fc2: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None)

基类： object

linear_fc1: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

linear_fc2: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = None

core.transformer.mlp.apply_swiglu_sharded_factory(original_sh_ten, sharded_offsets)

这为 transformer 中使用的所有模块提供了一个公共基类，其中包含一些通用功能。

Megatron 模块。

class core.transformer.module.Float16Module(*args: Any, **kwargs: Any)

基类： core.transformer.module.MegatronModule

Float 16 模块。

config

Transformer 配置

类型

TransformerConfig

fp16

指定模型是否在 fp16 模式下运行

类型

bool

bf16

指定模型是否在 bf16 模式下运行

类型

bool

参数

config (TransformerConfig) – 用于初始化模型的 transformer 配置

forward(*inputs, **kwargs)

load_state_dict(state_dict, strict=True)

set_input_tensor(input_tensor)

sharded_state_dict(prefix='', *args, **kwargs)

从被包装的模块检索 sharded_state_dict。

state_dict(destination=None, prefix='', keep_vars=False)

state_dict_for_save_checkpoint(prefix='', keep_vars=False)

从被包装的模块检索 state_dict。

class core.transformer.module.MegatronModule(*args: Any, **kwargs: Any)

基类： torch.nn.Module

所有模型继承的 Megatron 基础模块。

torch Module 的 Megatron 特定扩展，支持流水线并行

参数

config (TransformerConfig) – Transformer 配置

set_is_first_microbatch()

如果存在 is_first_microbatch 标志且 config.fp8==True，则设置该标志。设置此标志后，TE 模块将更新其 fp8 参数缓存。

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

用于分布式检查点分片状态字典的默认实现。

sharded_state_dict 的通用定义只是在所有子模块上递归调用 sharded_state_dict_default （如果可能，调用 sharded_state_dict 方法，否则调用默认实现）。

参数

• prefix (str) – 状态字典键的前缀

• sharded_offsets (Tuple[Tuple[int, int, int]], optional) – sup-modules 已应用的分片（例如 PP 相关）。传递给 ShardedTensor

• metadata (dict, optional) – 递归传递给 sharded_state_dict 方法的元数据

返回

状态字典键到 ShardedTensor 的映射字典

返回类型

dict

state_dict_for_save_checkpoint(prefix: str = '', keep_vars: bool = False)

覆盖用于保存检查点的状态字典。使用此函数覆盖用于保存检查点的状态字典。

参数

• prefix (str, optional) – _description_。默认为 ''。

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

返回

_description_

返回类型

_type_

core.transformer.module.conversion_helper(val, conversion)

core.transformer.module.float16_to_fp32(val)

core.transformer.module.fp32_to_float16(val, float16_convertor)

core.transformer.module.param_is_not_shared(param)

几个 transformer 层的块或堆栈。这些层可以全部相同，也可以每个都唯一。

class core.transformer.transformer_block.TransformerBlock(*args: Any, **kwargs: Any)

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

Transformer 类。

forward(hidden_states: torch.Tensor, attention_mask: torch.Tensor, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, rotary_pos_emb: Optional[torch.Tensor] = None, rotary_pos_cos: Optional[torch.Tensor] = None, rotary_pos_sin: Optional[torch.Tensor] = None, attention_bias: Optional[torch.Tensor] = None, inference_params: Optional[megatron.core.InferenceParams] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, sequence_len_offset: Optional[torch.Tensor] = None)

执行通过 transformer 块的前向传递。

此方法处理 transformer 的核心计算，包括自注意力、可选的交叉注意力和前馈操作。

参数

• hidden_states (Tensor) – 形状为 [s, b, h] 的输入张量，其中 s 是序列长度，b 是批次大小，h 是隐藏层大小。

• attention_mask (Tensor) – 形状为 [1, 1, s, s] 的布尔张量，用于掩蔽自注意力。

• context (Tensor, optional) – 用于交叉注意力的上下文张量。

• context_mask (Tensor, optional) – 用于交叉注意力上下文的掩码

• rotary_pos_emb (Tensor, optional) – 旋转位置嵌入。

• attention_bias (Tensor) – 用于 Q * K.T 的偏差张量，形状可广播到 [b, num_head, sq, skv]，例如 [1, 1, sq, skv]。用作应用 TE cuDNN 注意力掩码的替代方法。

• inference_params (InferenceParams, optional) – 用于推理时优化的参数。

• packed_seq_params (PackedSeqParams, optional) – 用于打包序列处理的参数。

返回

形状为 [s, b, h] 的输出隐藏状态张量，以及可选的更新上下文张量（如果使用交叉注意力）。

返回类型

Union[Tensor, Tuple[Tensor, Tensor]]

get_cuda_graph_optional_args(attention_mask: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor, rotary_pos_emb: torch.Tensor, attention_bias: torch.Tensor, inference_params: megatron.core.InferenceParams, packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams)

获取 CUDA 图的可选张量参数。

set_input_tensor(input_tensor: torch.Tensor)

设置要使用的输入张量，以替代 forward() 的输入。

当进行流水线并行时，来自上一阶段的输入来自通信，而不是来自输入，因此模型的 forward_step_func 不会拥有它。因此，内部代码使用此函数绕过 forward_step_func 提供的输入

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

sharded_state_dict(prefix: str = '', sharded_offsets: tuple = (), metadata: dict = None)

参数

• 为 transformer 块生成分片状态字典。

• prefix (str, optional) – 要添加到状态字典中所有键的前缀。默认为空字符串。

• sharded_offsets (tuple, optional) – 分片偏移的元组。

返回

metadata (dict, optional) – 用于分片的其他元数据。可以指定层是否是非同构的。默认为 None。

返回类型

包含模型分片状态的字典。

ShardedStateDict

基类： object

class core.transformer.transformer_block.TransformerBlockSubmodules(layer_specs: Optional[List[megatron.core.transformer.spec_utils.ModuleSpec]] = None, layer_norm: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, torch.nn.Module]] = None)

用于指定 transformer 块子模块的数据类。

参数

• layer_specs (List[ModuleSpec], 可选) – Transformer 模块内各层的模块规范列表。每个规范通常定义一个完整的 Transformer 层（例如，自注意力，前馈网络）。

• layer_norm (Optional[Union[ModuleSpec, torch.nn.Module]], 可选) – 要应用的层归一化的规范或实例。

layer_norm: Optional[Union[megatron.core.transformer.spec_utils.ModuleSpec, torch.nn.Module]] = None

layer_specs: List[megatron.core.transformer.spec_utils.ModuleSpec] = None

core.transformer.transformer_block.get_num_layers_to_build(config: megatron.core.transformer.transformer_config.TransformerConfig) → int

确定当前流水线阶段要构建的 Transformer 层数。 :param config: 包含 Transformer 模型参数的配置对象。 :type config: TransformerConfig

返回

当前流水线阶段要构建的层数。

返回类型

int (整数)

包含 Transformer 的所有配置选项。使用 dataclass 可以将所有参数集中在一个数据类中，而不是通过多层函数调用传递多个参数，从而减少代码冗余。

class core.transformer.transformer_config.MLATransformerConfig(tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, virtual_pipeline_model_parallel_size: Optional[int] = None, sequence_parallel: bool = False, context_parallel_size: int = 1, hierarchical_context_parallel_sizes: Optional[list[int]] = None, expert_model_parallel_size: int = 1, expert_tensor_parallel_size: Optional[int] = None, moe_extended_tp: bool = False, perform_initialization: bool = True, use_cpu_initialization: bool = False, fp16: bool = False, bf16: bool = False, params_dtype: torch.dtype = torch.float32, timers: Optional[Callable] = None, finalize_model_grads_func: Optional[Callable] = None, grad_scale_func: Optional[Callable] = None, no_sync_func: Optional[Callable] = None, grad_sync_func: Optional[Callable] = None, param_sync_func: Optional[Callable] = None, deterministic_mode: bool = False, enable_autocast: bool = False, autocast_dtype: Optional[torch.dtype] = None, num_microbatches_with_partial_activation_checkpoints: Optional[int] = None, gradient_accumulation_fusion: bool = False, async_tensor_model_parallel_allreduce: bool = False, use_te_rng_tracker: bool = False, tp_comm_overlap: bool = False, tp_comm_bulk_wgrad: bool = True, tp_comm_bulk_dgrad: bool = True, tp_comm_overlap_ag: bool = True, tp_comm_overlap_rs: bool = True, tp_comm_overlap_rs_dgrad: bool = False, tp_comm_split_ag: bool = True, tp_comm_atomic_ag: bool = False, tp_comm_split_rs: bool = True, tp_comm_atomic_rs: bool = False, cross_entropy_loss_fusion: bool = False, tp_comm_overlap_disable_qkv: bool = False, tp_comm_overlap_disable_fc1: bool = False, tp_comm_bootstrap_backend: str = 'nccl', pipeline_dtype: Optional[torch.dtype] = None, variable_seq_lengths: bool = False, overlap_p2p_comm: bool = False, batch_p2p_comm: bool = True, batch_p2p_sync: bool = True, use_ring_exchange_p2p: bool = False, deallocate_pipeline_outputs: bool = False, defer_embedding_wgrad_compute: bool = False, wgrad_deferral_limit: int = 0, pipeline_model_parallel_split_rank: Optional[int] = None, overlap_p2p_comm_warmup_flush: bool = False, microbatch_group_size_per_vp_stage: Optional[int] = None, cpu_offloading: bool = False, cpu_offloading_num_layers: int = 0, _cpu_offloading_context: Optional[ContextManager] = None, cpu_offloading_activations: bool = True, cpu_offloading_weights: bool = True, barrier_with_L1_time: bool = True, num_layers: int = 0, first_pipeline_num_layers: Optional[int] = None, last_pipeline_num_layers: Optional[int] = None, hidden_size: int = 0, num_attention_heads: int = 0, attention_backend: megatron.core.transformer.enums.AttnBackend = megatron.core.transformer.enums.AttnBackend.auto, softmax_scale: Optional[float] = None, num_query_groups: Optional[int] = None, ffn_hidden_size: Optional[int] = None, kv_channels: Optional[int] = None, hidden_dropout: float = 0.1, attention_dropout: float = 0.1, fp32_residual_connection: bool = False, apply_residual_connection_post_layernorm: bool = False, layernorm_epsilon: float = 1e-05, layernorm_zero_centered_gamma: bool = False, add_bias_linear: bool = True, add_qkv_bias: bool = False, gated_linear_unit: bool = False, activation_func: Callable = torch.nn.functional.gelu, activation_func_fp8_input_store: bool = False, num_moe_experts: Optional[int] = None, rotary_interleaved: bool = False, window_size: Optional[Tuple[int, int]] = None, normalization: str = 'RMSNorm', qk_layernorm: bool = False, test_mode: bool = False, calculate_per_token_loss: bool = False, multi_latent_attention: bool = True, init_method: Optional[Callable] = None, output_layer_init_method: Optional[Callable] = None, init_method_std: float = 0.02, apply_query_key_layer_scaling: bool = False, attention_softmax_in_fp32: bool = True, bias_activation_fusion: bool = False, masked_softmax_fusion: bool = False, persist_layer_norm: bool = False, memory_efficient_layer_norm: bool = False, bias_dropout_fusion: bool = False, apply_rope_fusion: bool = False, recompute_granularity: Optional[str] = None, recompute_method: Optional[str] = None, recompute_num_layers: Optional[int] = None, distribute_saved_activations: Optional[bool] = None, fp8: Optional[str] = None, fp8_margin: int = 0, fp8_interval: int = 1, fp8_amax_history_len: int = 1, fp8_amax_compute_algo: str = 'most_recent', fp8_wgrad: bool = True, fp8_dot_product_attention: bool = False, fp8_multi_head_attention: bool = False, tp_only_amax_red: bool = False, moe_shared_expert_intermediate_size: Optional[int] = None, moe_shared_expert_overlap: bool = False, moe_layer_freq: int = 1, moe_ffn_hidden_size: Optional[int] = None, moe_router_load_balancing_type: str = 'aux_loss', moe_router_topk: int = 2, moe_router_topk_limited_devices: Optional[int] = None, moe_router_pre_softmax: bool = False, moe_router_topk_scaling_factor: Optional[float] = None, moe_grouped_gemm: bool = False, moe_use_legacy_grouped_gemm: bool = False, moe_aux_loss_coeff: float = 0, moe_z_loss_coeff: Optional[float] = None, moe_input_jitter_eps: Optional[float] = None, moe_token_dropping: bool = False, moe_token_dispatcher_type: str = 'allgather', moe_per_layer_logging: bool = False, moe_expert_capacity_factor: Optional[float] = None, moe_pad_expert_input_to_capacity: bool = False, moe_token_drop_policy: str = 'probs', moe_layer_recompute: bool = False, cp_comm_type: Optional[Union[str, List[str]]] = None, clone_scatter_output_in_embedding: bool = True, disable_parameter_transpose_cache: bool = False, enable_cuda_graph: bool = False, cuda_graph_retain_backward_graph: bool = False, external_cuda_graph: bool = False, config_logger_dir: str = '', flash_decode: bool = False, inference_rng_tracker: bool = False, q_lora_rank: int = 512, kv_lora_rank: int = 512, qk_head_dim: int = 128, qk_pos_emb_head_dim: int = 64, v_head_dim: int = 128, rotary_base: float = 10000, rotary_scaling_factor: float = 40, max_position_embeddings: int = 163840, beta_fast: float = 32, beta_slow: float = 1, mscale: float = 0.707, mscale_all_dim: float = 0.707)

基类: core.transformer.transformer_config.TransformerConfig

megatron-core 多潜在注意力 (MLA) Transformer 的配置对象。

初始化函数为每个参数都提供了一个参数，包括 ModelParallelConfig 中的参数。包含 MLA 中融合的 YaRN RoPE 参数。

beta_fast: float = 32

YaRN RoPE 的快速 Beta 值。

beta_slow: float = 1

YaRN RoPE 的慢速 Beta 值。

kv_lora_rank: int = 512

Key 和 Value 张量的低秩表示的秩。

max_position_embeddings: int = 163840

原始模型的最大位置嵌入。

mscale: float = 0.707

多潜在注意力中 YaRN RoPE 的 Mscale 值。

mscale_all_dim: float = 0.707

多潜在注意力中 YaRN RoPE 的所有维度的 Mscale 值。

multi_latent_attention: bool = True

是否使用多潜在注意力。

normalization: str = 'RMSNorm'

MLA 模型的默认归一化层是 RMSNorm。

q_lora_rank: int = 512

Query 张量的低秩表示的秩。

qk_head_dim: int = 128

QK 投影中 head 的维度。q_head_dim = qk_head_dim + qk_pos_emb_head_dim

qk_pos_emb_head_dim: int = 64

QK 投影中位置嵌入的维度。

rotary_base: float = 10000

旋转嵌入的旋转基数。

rotary_scaling_factor: float = 40

旋转嵌入的旋转缩放因子。

v_head_dim: int = 128

V 投影中 head 的维度。

class core.transformer.transformer_config.TransformerConfig(tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, virtual_pipeline_model_parallel_size: Optional[int] = None, sequence_parallel: bool = False, context_parallel_size: int = 1, hierarchical_context_parallel_sizes: Optional[list[int]] = None, expert_model_parallel_size: int = 1, expert_tensor_parallel_size: Optional[int] = None, moe_extended_tp: bool = False, perform_initialization: bool = True, use_cpu_initialization: bool = False, fp16: bool = False, bf16: bool = False, params_dtype: torch.dtype = torch.float32, timers: Optional[Callable] = None, finalize_model_grads_func: Optional[Callable] = None, grad_scale_func: Optional[Callable] = None, no_sync_func: Optional[Callable] = None, grad_sync_func: Optional[Callable] = None, param_sync_func: Optional[Callable] = None, deterministic_mode: bool = False, enable_autocast: bool = False, autocast_dtype: Optional[torch.dtype] = None, num_microbatches_with_partial_activation_checkpoints: Optional[int] = None, gradient_accumulation_fusion: bool = False, async_tensor_model_parallel_allreduce: bool = False, use_te_rng_tracker: bool = False, tp_comm_overlap: bool = False, tp_comm_bulk_wgrad: bool = True, tp_comm_bulk_dgrad: bool = True, tp_comm_overlap_ag: bool = True, tp_comm_overlap_rs: bool = True, tp_comm_overlap_rs_dgrad: bool = False, tp_comm_split_ag: bool = True, tp_comm_atomic_ag: bool = False, tp_comm_split_rs: bool = True, tp_comm_atomic_rs: bool = False, cross_entropy_loss_fusion: bool = False, tp_comm_overlap_disable_qkv: bool = False, tp_comm_overlap_disable_fc1: bool = False, tp_comm_bootstrap_backend: str = 'nccl', pipeline_dtype: Optional[torch.dtype] = None, variable_seq_lengths: bool = False, overlap_p2p_comm: bool = False, batch_p2p_comm: bool = True, batch_p2p_sync: bool = True, use_ring_exchange_p2p: bool = False, deallocate_pipeline_outputs: bool = False, defer_embedding_wgrad_compute: bool = False, wgrad_deferral_limit: int = 0, pipeline_model_parallel_split_rank: Optional[int] = None, overlap_p2p_comm_warmup_flush: bool = False, microbatch_group_size_per_vp_stage: Optional[int] = None, cpu_offloading: bool = False, cpu_offloading_num_layers: int = 0, _cpu_offloading_context: Optional[ContextManager] = None, cpu_offloading_activations: bool = True, cpu_offloading_weights: bool = True, barrier_with_L1_time: bool = True, num_layers: int = 0, first_pipeline_num_layers: Optional[int] = None, last_pipeline_num_layers: Optional[int] = None, hidden_size: int = 0, num_attention_heads: int = 0, attention_backend: megatron.core.transformer.enums.AttnBackend = megatron.core.transformer.enums.AttnBackend.auto, softmax_scale: Optional[float] = None, num_query_groups: Optional[int] = None, ffn_hidden_size: Optional[int] = None, kv_channels: Optional[int] = None, hidden_dropout: float = 0.1, attention_dropout: float = 0.1, fp32_residual_connection: bool = False, apply_residual_connection_post_layernorm: bool = False, layernorm_epsilon: float = 1e-05, layernorm_zero_centered_gamma: bool = False, add_bias_linear: bool = True, add_qkv_bias: bool = False, gated_linear_unit: bool = False, activation_func: Callable = torch.nn.functional.gelu, activation_func_fp8_input_store: bool = False, num_moe_experts: Optional[int] = None, rotary_interleaved: bool = False, window_size: Optional[Tuple[int, int]] = None, normalization: bool = 'LayerNorm', qk_layernorm: bool = False, test_mode: bool = False, calculate_per_token_loss: bool = False, multi_latent_attention: bool = False, init_method: Optional[Callable] = None, output_layer_init_method: Optional[Callable] = None, init_method_std: float = 0.02, apply_query_key_layer_scaling: bool = False, attention_softmax_in_fp32: bool = True, bias_activation_fusion: bool = False, masked_softmax_fusion: bool = False, persist_layer_norm: bool = False, memory_efficient_layer_norm: bool = False, bias_dropout_fusion: bool = False, apply_rope_fusion: bool = False, recompute_granularity: Optional[str] = None, recompute_method: Optional[str] = None, recompute_num_layers: Optional[int] = None, distribute_saved_activations: Optional[bool] = None, fp8: Optional[str] = None, fp8_margin: int = 0, fp8_interval: int = 1, fp8_amax_history_len: int = 1, fp8_amax_compute_algo: str = 'most_recent', fp8_wgrad: bool = True, fp8_dot_product_attention: bool = False, fp8_multi_head_attention: bool = False, tp_only_amax_red: bool = False, moe_shared_expert_intermediate_size: Optional[int] = None, moe_shared_expert_overlap: bool = False, moe_layer_freq: int = 1, moe_ffn_hidden_size: Optional[int] = None, moe_router_load_balancing_type: str = 'aux_loss', moe_router_topk: int = 2, moe_router_topk_limited_devices: Optional[int] = None, moe_router_pre_softmax: bool = False, moe_router_topk_scaling_factor: Optional[float] = None, moe_grouped_gemm: bool = False, moe_use_legacy_grouped_gemm: bool = False, moe_aux_loss_coeff: float = 0, moe_z_loss_coeff: Optional[float] = None, moe_input_jitter_eps: Optional[float] = None, moe_token_dropping: bool = False, moe_token_dispatcher_type: str = 'allgather', moe_per_layer_logging: bool = False, moe_expert_capacity_factor: Optional[float] = None, moe_pad_expert_input_to_capacity: bool = False, moe_token_drop_policy: str = 'probs', moe_layer_recompute: bool = False, cp_comm_type: Optional[Union[str, List[str]]] = None, clone_scatter_output_in_embedding: bool = True, disable_parameter_transpose_cache: bool = False, enable_cuda_graph: bool = False, cuda_graph_retain_backward_graph: bool = False, external_cuda_graph: bool = False, config_logger_dir: str = '', flash_decode: bool = False, inference_rng_tracker: bool = False)

Bases: core.model_parallel_config.ModelParallelConfig

megatron-core 转换器的配置对象。

初始化函数为每个参数都设置了实参，包括 ModelParallelConfig 中的参数。

activation_func: Callable

用于 MLP 中非线性激活的激活函数。

activation_func_fp8_input_store: bool = False

在 FP8 中存储 MLP 激活函数的输入，以便进行反向传播以节省内存。存储的输入在反向传播计算之前被转换回原始精度。

add_bias_linear: bool = True

在所有线性层（QKV 投影、核心注意力之后以及 MLP 层中的两个）中包含偏置项。

add_qkv_bias: bool = False

仅为 QKV 投影添加偏置项。

apply_query_key_layer_scaling: bool = False

如果为 true，则将 Q * K^T 缩放 1 / 层数。这提高了使用 fp16 训练时的数值稳定性。

apply_residual_connection_post_layernorm: bool = False

如果为 True，则使用原始 BERT 残差连接顺序。

apply_rope_fusion: bool = False

如果为 True，则使用融合 RoPE 内核。

attention_backend: megatron.core.transformer.enums.AttnBackend

要运行的注意力后端。默认情况下，我们让 Transformer Engine 决定要运行的最佳后端（本地情况除外）。如果注意力后端是本地的，我们使用 mcore 中的本地 PyTorch 实现。用户可以通过更改此配置来指定确切的后端。

attention_dropout: float = 0.1

后注意力 dropout 概率。

attention_softmax_in_fp32: bool = True

如果为 True，则在 fp32 中运行注意力掩码和 softmax。如果 apply_query_key_layer_scaling 为 True，则应为 True。

bias_activation_fusion: bool = False

如果为 True，则在可能的情况下融合偏置加法和激活函数。

bias_dropout_fusion: bool = False

如果为 True，则使用偏置 dropout 融合。

calculate_per_token_loss: bool = False

是否基于全局批次中实际的非填充 token 数量计算交叉熵损失，而不是假设所有 token 都是非填充的默认行为。

clone_scatter_output_in_embedding: bool = True

设置为 True 时，克隆 embedding 层中 scatter_to_sequence_parallel_region 的输出，以方便输入的垃圾回收。

config_logger_dir: str = ''

当为非空时，将入口点配置转储到 config_logger_dir

cp_comm_type: Union[str, List[str]] = None

上下文并行性的 GPU 间通信类型。 str：所有层共享相同的通信类型。 List[str]：每层都有其单独的通信类型。 每层的 cp_comm_type 可以是“p2p”或“all_gather”或“a2a”或“a2a+p2p”。 “p2p”：在环形拓扑中使用 P2P 通信交换 KV 块。 P2P 是异步的，可以与注意力计算重叠。“all_gather”：All-gather 以在注意力之前获取 KV 的完整序列。 all-gather 不是异步的，不能重叠。“a2a”：类似于 DeepSpeed Ulysses，在 CP 组中分散注意力头，然后聚集以获取 QKV 的完整序列。“a2a+p2p”：上下文并行性的分层实现到注意力。它在低级 CP 组（例如，通过 NVLink）中使用 A2A 通信，并在高级 CP 组（例如，通过 IBLink）中使用 P2P 通信。

cuda_graph_retain_backward_graph: bool = False

设置为 true 时，cudagraph 反向传播将使用 ‘retain_grad=True’ 进行图捕获。 这可能会为某些并非完全 cudagraph 安全的模块启用 cudagraph。 有关更多详细信息，请参阅：https://pytorch.ac.cn/docs/stable/generated/torch.Tensor.backward.html。

disable_parameter_transpose_cache: bool = False

设置为 true 时，参数转置不会为后续迭代缓存。

distribute_saved_activations: bool = None

如果为 True，则在模型并行组中分发重新计算的激活。

enable_cuda_graph: bool = False

设置为 true 时，TransformerLayer 层将替换为 CUDA 图版本。

external_cuda_graph: bool = False

设置为 true 时，TransformerLayer 层将替换为用户提供的 CUDA 图。

ffn_hidden_size: int = None

Transformer 前馈网络隐藏层大小。如果未提供，则设置为 4*hidden_size。

first_pipeline_num_layers: int = None

第一个流水线阶段的 transformer 层数。 None 意味着跨 PP 秩的平均层划分。

flash_decode: bool = False

在推理期间使用优化的闪存解码内核。

fp32_residual_connection: bool = False

如果为 true，则将残差连接移动到 fp32。

fp8: str = None

如果设置，则通过 Transformer Engine 启用 FP8 精度。 有 2 个预定义的选项 (1) ‘e4m3’ 统一地为所有 FP8 张量使用 e4m3，(2) ‘hybrid’ 为所有 FP8 激活和权重张量使用 e4m3，为所有 FP8 输出激活梯度张量使用 e5m2。

fp8_amax_compute_algo: str = 'most_recent'

用于为缩放因子计算选择 amax 值的算法。 有 2 个预定义的选项：max 选择历史窗口中最大的 amax，而 most_recent 始终选择最近看到的值。

fp8_amax_history_len: int = 1

用于缩放因子计算的 amax 历史窗口的长度。

fp8_dot_product_attention: bool = False

设置为 True 时，使用点积注意力的 FP8 实现。

fp8_interval: int = 1

自 TransformerEngine v1.8.0 起已弃用。 此标志将被忽略。 控制重新计算缩放因子的频率。

fp8_margin: int = 0

用于缩放因子计算的边距。

fp8_multi_head_attention: bool = False

设置为 True 时，使用多头注意力的 FP8 实现。

fp8_wgrad: bool = True

设置为 False 时，覆盖 FP8 配置选项，并在更高精度下执行 wgrad 计算。

gated_linear_unit: bool = False

为 MLP 中的第一个线性层使用门控线性单元。

hidden_dropout: float = 0.1

Transformer 隐藏状态的 Dropout 概率。

hidden_size: int = 0

Transformer 隐藏层大小。

inference_rng_tracker: bool = False

是否应该为推理实例化单独的 RNG 跟踪器。

init_method: Callable = None

初始化权重的方法。 请注意，偏置始终设置为零。 应该是一个接受单个 Tensor 并对其进行初始化的函数。 如果为 None，则将设置为 megatron.core.utils.init_method_normal(init_method_std)，它是均值 = 0.0 和标准差 = init_method_std 的 torch nn init normal。

init_method_std: float = 0.02

默认初始化方法的零均值正态分布的标准差，如果提供了 init_method 和 output_layer_init_method，则不使用。

kv_channels: int = None

多头注意力中的投影权重维度。 如果未提供，则设置为 hidden_size // num_attention_heads。

last_pipeline_num_layers: int = None

最后一个流水线阶段的 transformer 层数。 None 意味着跨 PP 秩的平均层划分。

layernorm_epsilon: float = 1e-05

任何 LayerNorm 操作的 epsilon 值。

layernorm_zero_centered_gamma: bool = False

如果设置为 True，则调整 LayerNorm 以将 gamma 值居中在 0 附近。 这提高了数值稳定性。

masked_softmax_fusion: bool = False

如果为 True，则使用 softmax 融合。

memory_efficient_layer_norm: bool = False

如果为 True，并且使用本地层（而不是来自 TransformerEngine），则告诉 Apex 使用内存高效的融合 LayerNorm 内核。 如果不使用 LayerNorm，则忽略。

moe_aux_loss_coeff: float = 0

辅助损失的缩放系数。 建议起始值为 1e-2。

moe_expert_capacity_factor: float = None

每个专家的容量因子，None 表示不会丢弃任何 token。 默认为 None。

类型

moe_expert_capacity_factor (float)

moe_ffn_hidden_size: int = None

MoE Feed-Forward Network 隐藏层大小

moe_grouped_gemm: bool = False

当每个 rank 有多个专家时，将多个本地（可能很小）的 gemm 压缩到单个内核启动中，通过利用 CUTLASS 2.8 版本后引入的 Grouped GEMM 功能（https://github.com/fanshiqing/grouped_gemm）来提高利用率和性能。

moe_input_jitter_eps: float = None

通过应用具有指定 epsilon 值的抖动，向输入张量添加噪声。

moe_layer_freq: int = 1

MoE 层和 Dense 层之间的频率。接受以下两种形式： - 整数 N：表示 1:N 的比例，意味着每 N-1 个 Dense 层有一个专家层。 - 包含 Python 列表表达式的字符串，用于定义自定义模式，例如：“([1]*3+[0]*1)*3” 求值为 [1,1,1,0,1,1,1,0,1,1,1,0]，其中 1 表示专家层，0 表示 Dense 层。

moe_layer_recompute: bool = False

检查点 moe_layer 以节省激活内存。

类型

内存优化

moe_pad_expert_input_to_capacity: bool = False

如果为 True，则填充每个专家的输入以匹配专家容量长度，仅在设置 moe_expert_capacity_factor 后生效。 默认设置为 False。

类型

moe_pad_expert_input_to_capacity (bool)

moe_per_layer_logging: bool = False

启用 MoE 的逐层日志记录，目前支持辅助损失和 z 损失。

moe_router_load_balancing_type: str = 'aux_loss'

路由器的负载均衡策略。“aux_loss” 对应于 GShard 和 SwitchTransformer 中使用的负载均衡损失；“seq_aux_loss” 对应于 DeepSeekV2 中使用的损失，它计算每个单独样本的损失；“sinkhorn” 对应于 S-BASE 中使用的均衡算法，“none” 表示不进行负载均衡。 默认为 “aux_loss”。

moe_router_pre_softmax: bool = False

为 MoE 启用 pre-softmax 路由，这意味着 softmax 在 top-k 选择之前。 默认情况下，softmax 在 top-k 之后完成。

moe_router_topk: int = 2

每个 token 要路由到的专家数量。

moe_router_topk_limited_devices: int = None

路由期间每个 token 要考虑的专家并行 rank 的数量。 通过首先为每个 token 选择 N 个 rank，然后在这些设备上的专家中进行 top-k 选择，从而在专家并行 rank 的子集上执行 top-k 路由。 None 表示没有设备限制。

moe_router_topk_scaling_factor: float = None

top-k 选择中路由分数的缩放因子，仅在启用 moe_router_pre_softmax 时有效。 默认为 None，表示不进行缩放。

moe_shared_expert_intermediate_size: int = None

共享专家总 ffn 隐藏层大小。 如果有多个共享专家，则应等于 “num_shared_experts * ffn_size_of_each_shared_expert”。 None 表示没有共享专家。

moe_shared_expert_overlap: bool = False

启用共享专家计算和调度器通信之间的重叠。 如果不启用此项，则共享专家将在路由专家之后执行。

moe_token_dispatcher_type: str = 'allgather'

要使用的 token 调度器类型。 默认为 “allgather”。 选项为 “allgather” 和 “alltoall”。

moe_token_drop_policy: str = 'probs'

丢弃 token 的策略。 可以是 “probs” 或 “position”。 如果为 “probs”，则将丢弃概率最低的 token。 如果为 “position”，则将丢弃每个批次末尾的 token。

moe_token_dropping: bool = False

此功能涉及选择性地丢弃和填充每个专家的 token 以达到指定的容量，类似于 GShard、Switch-Transformer 和 DeepSpeed-MoE。 请注意，当前不支持此功能，因此应保持为 False。

moe_use_legacy_grouped_gemm: bool = False

使用旧版 GroupedMLP 而不是 TEGroupedMLP。 注意：旧版将很快被弃用。

moe_z_loss_coeff: float = None

z-loss 的缩放系数。 建议起始值为 1e-3。

multi_latent_attention: bool = False

是否使用多潜在注意力。

normalization: bool = 'LayerNorm'

用于归一化层的 norm，有效选项为 LayerNorm 和 RMSNorm。

num_attention_heads: int = 0

Transformer 注意力头的数量。

num_layers: int = 0

Transformer 模块中 Transformer 层的数量。

num_moe_experts: int = None

MoE 层要使用的专家数量。 设置后，它将 MLP 替换为 MoE 层。 设置为 None 则不使用 MoE。

num_query_groups: int = None

组查询注意力的查询组数。 如果为 None，则使用普通注意力。

output_layer_init_method: Callable = None

初始化注意力模块和 MLP 模块的输出层权重的方法。 如果为 None，则将设置为 megatron.core.utils.scaled_init_method_normal(init_method_std)，它是 torch nn init normal，均值=0.0，标准差=init_method_std / math.sqrt(2.0 * num_layers)。

persist_layer_norm: bool = False

如果为 True，则使用持久化融合层归一化内核。 此内核仅支持一组固定的隐藏层大小。

qk_layernorm: bool = False

是否对 query 和 key 嵌入应用 LayerNorm。

recompute_granularity: str = None

确定要使用的激活重计算类型。 Megatron-core 支持 “selective” 激活检查点，其中仅检查点内存密集型注意力部分。 这些内存密集型激活的计算密集程度也较低，这使得激活检查点对于 LLM (20B+) 更有效。 有关更多详细信息，请参阅 Reducing Activation Recomputation in Large Transformer Models (https://arxiv.org/abs/2205.05198)。 “full” 将检查点整个 Transformer 层。 如果为 None，则不执行重计算，并且保存所有激活。 如果设置，则必须为 “selective” 或 “full”。 “selective” 始终使用所有层。

recompute_method: str = None

确定将重计算哪些 Transformer 层。“uniform” 将均匀地划分 Transformer 模块中 Transformer 层的总数，并在指定的粒度下重计算每个划分块的输入激活。“block” 将仅重计算每个流水线阶段中一组 Transformer 层的输入激活。 流水线阶段中的其余层将不进行任何激活重计算。 如果为 None 且启用了重计算，则所有层都将进行重计算。 如果设置，则必须为 “uniform” 或 “block”。

recompute_num_layers: int = None

当 recompute_method 为 uniform 时，recompute_num_layers 是每个均匀划分的重计算单元中的 Transformer 层数。 当 recompute_method 为 block 时，recompute_num_layers 是每个流水线阶段中要重计算的 Transformer 层数。 对于 “selective” 激活检查点，必须为 None。

rotary_interleaved: bool = False

True 表示旋转偶数和奇数维度的对（RoFormer 风格），False 表示旋转前半部分和后半部分的对（LLaMa 风格）。 默认为 False。

softmax_scale: float = None

用于注意力缩放的 Softmax 缩放。

test_mode: bool = False

是否运行实时测试。

tp_only_amax_red: bool = False

设置为 True 时，仅在 TP 或 TP-CP 域中减少 FP8 AMAX

use_te_rng_tracker: bool = False

是否使用 TE 或 MCore 版本的 RNG 跟踪器。

window_size: Optional[Tuple[int, int]] = None

如果不为 None，则将使用滑动窗口注意力。 窗口大小由元组内的数字指定；-1 是特殊值，表示 “无限窗口大小”。

单个标准 Transformer 层，包括注意力和 MLP 模块。

class core.transformer.transformer_layer.BaseTransformerLayer

基类: abc.ABC

类似于 TransformerLayer 实现的通用父类。

一个虚拟类，由类似的 TransformerLayer`s 例如此文件中的 `TransformerLayer 以及可能旨在将 TransformerBlock 用作基础模块的其他 TransformerLayer 实现子类化。 主要目的是检查规范中提供的任何层（或模块）是否是此类的子类，以允许为 TransformerBlock 中的所有层扇出该规范。 有关更多详细信息，请参见 transformer_block.py 文件中的 _get_block_submodules 方法实现。

class core.transformer.transformer_layer.TransformerLayer(*args: Any, **kwargs: Any)

基类: megatron.core.transformer.module.MegatronModule, core.transformer.transformer_layer.BaseTransformerLayer

单个 Transformer 层。

Transformer 层接受大小为 [s, b, h] 的输入，并返回相同大小的输出。

forward(hidden_states, attention_mask=None, context=None, context_mask=None, rotary_pos_emb=None, rotary_pos_cos=None, rotary_pos_sin=None, attention_bias=None, inference_params=None, packed_seq_params=None, sequence_len_offset=None)

执行通过 Transformer 层的前向传递。

此方法实现 Transformer 层的核心计算，包括自注意力、交叉注意力（如果适用）和前馈运算。

参数

• hidden_states (Tensor) – 形状为 [s, b, h] 的输入张量，其中 s 是序列长度，b 是批次大小，h 是隐藏层大小。

• attention_mask (Tensor) – 自注意力的 Mask 张量。

• context (Tensor, optional) – 用于交叉注意力的上下文张量。

• context_mask (Tensor, optional) – 交叉注意力的 Mask 张量。

• rotary_pos_emb (Tensor, optional) – 旋转位置嵌入。

• attention_bias (Tensor, optional) – Q * K.T 的 Bias 张量。

• inference_params (object, optional) – 推理时优化的参数。

• packed_seq_params (object, optional) – Packed 序列处理的参数。

返回

一个包含以下内容的元组

output (Tensor): 形状为 [s, b, h] 的转换后的隐藏状态。 context (Tensor): 如果使用交叉注意力，则为更新后的上下文张量；否则为 None。

返回类型

Tuple[Tensor, Tensor]

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

为 Transformer 层生成分片状态字典。

参数

• prefix (str, optional) – 要添加到状态字典中所有键的前缀。

• prefix (str, optional) – 要添加到状态字典中所有键的前缀。默认为空字符串。

• metadata (Optional[dict], optional) – 用于分片的其他元数据。

返回

包含 Transformer 层的分片状态的字典。

返回类型

包含模型分片状态的字典。

class core.transformer.transformer_layer.TransformerLayerSubmodules(input_layernorm: typing.Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = megatron.core.transformer.identity_op.IdentityOp, self_attention: typing.Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = megatron.core.transformer.identity_op.IdentityOp, self_attn_bda: typing.Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = megatron.core.transformer.identity_op.IdentityFuncOp, pre_cross_attn_layernorm: typing.Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = megatron.core.transformer.identity_op.IdentityOp, cross_attention: typing.Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = megatron.core.transformer.identity_op.IdentityOp, cross_attn_bda: typing.Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = megatron.core.transformer.identity_op.IdentityFuncOp, pre_mlp_layernorm: typing.Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = megatron.core.transformer.identity_op.IdentityOp, mlp: typing.Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = megatron.core.transformer.identity_op.IdentityOp, mlp_bda: typing.Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = megatron.core.transformer.identity_op.IdentityFuncOp, sharded_state_dict_keys_map: typing.Dict[str, str] = <factory>)

基类： object

用于指定 Transformer 层子模块的配置类。

此类定义了 Transformer 层各种组件的结构和默认实现，从而可以灵活地自定义层的架构。

参数

• input_layernorm (Union[ModuleSpec, type]) – 输入层归一化的规范。

• self_attention (Union[ModuleSpec, type]) – 自注意力机制的规范。

• self_attn_bda (Union[ModuleSpec, type]) – 自注意力之后 bias-dropout-add 运算的规范。

• pre_cross_attn_layernorm (Union[ModuleSpec, type]) – 交叉注意力之前层归一化的规范。

• cross_attention (Union[ModuleSpec, type]) – 交叉注意力机制的规范。

• cross_attn_bda (Union[ModuleSpec, type]) – 交叉注意力之后 bias-dropout-add 运算的规范。

• pre_mlp_layernorm (Union[ModuleSpec, type]) – MLP 之前层归一化的规范。

• mlp (Union[ModuleSpec, type]) – Dense 层中 MLP 的规范。

• mlp_bda (Union[ModuleSpec, type]) – MLP 之后 bias-dropout-add 运算的规范。

• sharded_state_dict_keys_map (Dict[str, str]) – 用于在 sharded_state_dict 方法中应用的分片张量键的映射。

sharded_state_dict_keys_map: Dict[str, str]

Transformer 实现中使用的各种实用程序。

Transformer 层的实用程序。

core.transformer.utils.attention_mask_func(attention_scores, attention_mask)

core.transformer.utils.erf_gelu(x)

core.transformer.utils.gelu_impl(x)

OpenAI 的 gelu 实现。

core.transformer.utils.get_default_causal_mask(sq: int) → torch.Tensor

返回 softmax 输入的因果上三角 Mask。

core.transformer.utils.get_linear_layer(rows, columns, init_method, perform_initialization=True)

具有权重初始化的简单线性层。

core.transformer.utils.make_sharded_object_for_checkpoint(obj: Any, key: str, sharded_offsets: Iterable[Tuple[int, int, int]] = (), replica_id: Union[None, int, Tuple[int, ...]] = None, **kwargs)

用于实例化非分片 ShardedObject（在 TP 和 DP 组之间复制）的帮助程序。

参数

• obj (object) – 要分片的任何对象

• key (str) – 对象的唯一标识符

• sharded_offsets (Iterable[Tuple[int, int, int]]) – 通常添加到 ShardedTensor 的偏移量，将用作 ShardedObject 的全局偏移量

• replica_id (Union[None, int, Tuple[int, ...]]) – 副本 ID

core.transformer.utils.make_sharded_tensors_for_checkpoint(state_dict: megatron.core.dist_checkpointing.mapping.StateDict, prefix: str, tensor_parallel_layers_axis_map: Optional[Dict[str, int]] = None, sharded_offsets: Iterable[Tuple[int, int, int]] = (), extra_state_suffix: str = '_extra_state')

使用 ShardedTensor 或 ShardedObject 包装来自 Transformer 层的张量。

对于给定的 state_dict，包装： - 所有 _extra_states 与 ShardedObject - 所有在 tensor_parallel_layers_axis_map 中指定的张量与 TP 和 DP 分片 ShardedTensor - 其他值与 DP 分片 ShardedTensor

参数

• state_dict (StateDict) – 要转换的 state_dict

• prefix (str) – 附加到最终状态字典中键的前缀

• tensor_parallel_layers_axis_map (Dict[str, int], optional) – 将层名称映射到 TP 分片轴的字典

• sharded_offsets (Iterable[Tuple[int, int, int]], optional) – 已应用的分片（例如，PP 相关），传递给 ShardedTensor

• extra_state_suffix (str, default = '_extra_state') – 具有此后缀的层将使用 ShardedObject 而不是 ShardedTensor 包装。

core.transformer.utils.openai_gelu(x)

core.transformer.utils.sharded_state_dict_default(module: torch.nn.Module, prefix: str = '', sharded_offsets: Tuple[Tuple[int, int, int]] = (), metadata: Optional[dict] = None) → megatron.core.dist_checkpointing.mapping.ShardedStateDict

为非 MegatronModules 提供 sharded_state_dict 方法的实现。

尽可能尝试调用 module.sharded_state_dict，否则使用常规状态字典并假定张量在 TP 和 DP 之间复制。

keep_vars=True 传递给 module.state_dict，以便稍后可以对优化器状态进行分片。

参数

• module (torch.nn.Module) – 我们要获取分片状态字典的模块

• prefix (str) – 状态字典键的前缀

• sharded_offsets (Tuple[Tuple[int, int, int]], optional) – sup-modules 已应用的分片（例如 PP 相关）。传递给 ShardedTensor

• metadata (dict, optional) – 传递给模块 sharded_state_dict 方法的元数据

返回

状态字典键到 ShardedTensor 的映射字典

返回类型

dict