GPT 模型架构

class GPTModel(LanguageModule):
    def __init__(self, config, transformer_layer_spec, ...):
        super().__init__(config)

        # ---- 1. Embedding ----
        self.embedding = LanguageModelEmbedding(
            config=self.config,
            vocab_size=vocab_size,
            max_sequence_length=max_sequence_length,
            position_embedding_type=position_embedding_type,
        )

        # ---- 2. TransformerBlock (Decoder) ----
        self.decoder = TransformerBlock(
            config=self.config,
            spec=transformer_layer_spec,    # Layer Spec 决定内部结构
            pre_process=self.pre_process,   # PP first stage?
            post_process=self.post_process, # PP last stage?
        )

        # ---- 3. Output Layer ----
        self.output_layer = tensor_parallel.ColumnParallelLinear(
            config.hidden_size,
            self.vocab_size,
            config=config,
            init_method=config.init_method,
            bias=False,
            skip_bias_add=False,
            gather_output=not self.parallel_output,
        )

def get_gpt_layer_local_spec(
    num_experts=None, moe_grouped_gemm=False, ...
) -> ModuleSpec:
    """返回基于 Megatron 原生实现的 GPT Layer Spec."""
    return ModuleSpec(
        module=TransformerLayer,
        submodules=TransformerLayerSubmodules(
            # ---- Self Attention ----
            self_attention=ModuleSpec(
                module=SelfAttention,
                params={"attn_mask_type": AttnMaskType.causal},
                submodules=SelfAttentionSubmodules(
                    linear_qkv=ColumnParallelLinear,    # QKV 投影
                    core_attention=DotProductAttention,  # 注意力计算
                    linear_proj=RowParallelLinear,       # 输出投影
                ),
            ),
            self_attn_bda=get_bias_dropout_add,

            # ---- MLP ----
            pre_mlp_layernorm=FusedLayerNorm,
            mlp=ModuleSpec(
                module=MLP,
                submodules=MLPSubmodules(
                    linear_fc1=ColumnParallelLinear,
                    linear_fc2=RowParallelLinear,
                ),
            ),
            mlp_bda=get_bias_dropout_add,
        ),
    )

class TransformerBlock(MegatronModule):
    def __init__(self, config, spec, ...):
        super().__init__(config=config)

        # 计算本 stage 应构建多少层
        self.num_layers_per_pipeline_rank = get_num_layers_to_build(config)

        # 根据 spec 构建层列表
        self._build_layers(spec)

        # 最后一个 PP stage 的最后加一个 Final LayerNorm
        if self.post_process and self.post_layer_norm:
            self.final_layernorm = build_module(
                spec.submodules.final_layernorm,
                config=self.config,
                hidden_size=self.config.hidden_size,
            )

def forward(self, hidden_states, attention_mask, ...):
    # ============ Self Attention Block ============
    # 保存残差
    residual = hidden_states

    # Pre-LN: 先 LayerNorm
    hidden_states = self.input_layernorm(hidden_states)

    # Self Attention
    attention_output_with_bias = self.self_attention(
        hidden_states, attention_mask=attention_mask, ...
    )

    # BDA: bias + dropout + residual add
    # training 时用融合 kernel，inference 时用普通实现
    with self.bias_dropout_add_exec_handler():
        hidden_states = self.self_attn_bda(
            self.training, self.config.bias_dropout_fusion
        )(attention_output_with_bias, residual, self.hidden_dropout)

    # ============ MLP Block ============
    residual = hidden_states

    # Pre-LN
    hidden_states = self.pre_mlp_layernorm(hidden_states)

    # MLP
    mlp_output_with_bias = self.mlp(hidden_states)

    # BDA: bias + dropout + residual add
    with self.bias_dropout_add_exec_handler():
        hidden_states = self.mlp_bda(
            self.training, self.config.bias_dropout_fusion
        )(mlp_output_with_bias, residual, self.hidden_dropout)

    return hidden_states

def forward(self, hidden_states, attention_mask, ...):
    # ---- QKV Projection ----
    # 单个 ColumnParallelLinear 同时产生 Q, K, V
    mixed_qkv, _ = self.linear_qkv(hidden_states)    # [s, b, proj_size]

    # 拆分 Q, K, V
    # proj_size = (n_q_per_rank + 2 * n_kv_per_rank) * d_head
    new_tensor_shape = mixed_qkv.size()[:-1] + (
        self.num_query_groups_per_partition,             # n_kv per rank
        (self.num_attention_heads_per_partition           # n_q per rank
         // self.num_query_groups_per_partition + 2),     # +2 for K, V
        self.hidden_size_per_attention_head,              # d_head
    )
    mixed_qkv = mixed_qkv.view(*new_tensor_shape)

    # 沿 group 维度拆分
    query, key, value = torch.split(
        mixed_qkv,
        [self.num_attention_heads_per_partition
         // self.num_query_groups_per_partition,          # Q heads per group
         1,                                              # 1 K head per group
         1],                                             # 1 V head per group
        dim=3,
    )

    # ---- Core Attention ----
    core_attn_out = self.core_attention(query, key, value, attention_mask)

    # ---- Output Projection ----
    output, bias = self.linear_proj(core_attn_out)     # RowParallelLinear
    return output, bias

def forward(self, hidden_states):
    # fc1: [s, b, h] -> [s, b, 2 * ffn_hidden / TP]
    intermediate_parallel, bias_parallel = self.linear_fc1(hidden_states)

    # SwiGLU activation
    if self.config.gated_linear_unit:
        # 拆成 gate 和 value 两半
        intermediate_parallel = self.activation_func(intermediate_parallel)
        # activation_func 内部做：
        #   gate, value = chunk(intermediate, 2, dim=-1)
        #   return silu(gate) * value

    # fc2: [s, b, ffn_hidden / TP] -> [s, b, h]
    output, output_bias = self.linear_fc2(intermediate_parallel)

    return output, output_bias

def forward(self, input_ids, position_ids, attention_mask, labels=None, ...):
    # ---- Stage 1: Embedding (仅 PP first stage) ----
    if self.pre_process:
        decoder_input = self.embedding(input_ids, position_ids)
    else:
        decoder_input = None  # 从上一个 PP stage 接收激活值

    # ---- Stage 2: TransformerBlock ----
    hidden_states = self.decoder(
        hidden_states=decoder_input,
        attention_mask=attention_mask,
        ...
    )

    # ---- Stage 3: Output (仅 PP last stage) ----
    if self.post_process:
        # 线性投影: [s, b, h] -> [s, b, V]
        logits, _ = self.output_layer(hidden_states)

        if labels is not None:
            # 计算 cross-entropy loss
            loss = self.compute_language_model_loss(labels, logits)
        else:
            loss = None

        return loss
    else:
        return hidden_states  # 中间 stage，传给下一个 PP stage