【HuggingFace Transformers】OpenAIGPTModel源码解析

说到ChatGPT，大家可能都使用过吧。2022年，ChatGPT的推出引发了广泛的关注和讨论。这款对话生成模型不仅具备了强大的语言理解和生成能力，还能进行非常自然的对话，给用户带来了全新的互动体验。然而，ChatGPT的成功背后离不开它的前身——GPT。

1. GPT 介绍

GPT（Generative Pre-trained Transformer）是由OpenAI开发的一种基于Transformer架构的大型语言模型。它由多个堆叠的自注意力解码器层（Transformer Blocks）组成，每一层包含多头自注意力机制和前馈神经网络，并配有残差连接和层归一化以稳定训练。GPT采用自回归方式生成文本，通过在大规模互联网数据上进行预训练，具备强大的自然语言理解和生成能力，能够完成对话生成、文本补全等多种任务。其结构如下：

2. OpenAIGPTModel类 源码解析

源码地址：transformers/src/transformers/models/openai/modeling_openai.py

# -*- coding: utf-8 -*-
# @time: 2024/9/3 20:39
from typing import Optional, Union, Tuple

import torch

from torch import nn
from transformers import add_start_docstrings, OpenAIGPTPreTrainedModel
from transformers.modeling_outputs import BaseModelOutput
from transformers.models.openai.modeling_openai import OPENAI_GPT_START_DOCSTRING, Block, OPENAI_GPT_INPUTS_DOCSTRING, _CHECKPOINT_FOR_DOC, _CONFIG_FOR_DOC
from transformers.utils import add_start_docstrings_to_model_forward, add_code_sample_docstrings


@add_start_docstrings(
    "The bare OpenAI GPT transformer model outputting raw hidden-states without any specific head on top.",
    OPENAI_GPT_START_DOCSTRING,
)
class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)

        self.tokens_embed = nn.Embedding(config.vocab_size, config.n_embd)  # 定义 token 嵌入层
        self.positions_embed = nn.Embedding(config.n_positions, config.n_embd)  # 定义 position 嵌入层
        self.drop = nn.Dropout(config.embd_pdrop)  # 定义 drop 层
        self.h = nn.ModuleList([Block(config.n_positions, config, scale=True) for _ in range(config.n_layer)]) # 定义多个 Block 层

        # 注册一个缓冲区用于存储position_ids，初始化为从 0 到 config.n_positions 的序列
        self.register_buffer("position_ids", torch.arange(config.n_positions), persistent=False)
        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.tokens_embed

    def set_input_embeddings(self, new_embeddings):
        self.tokens_embed = new_embeddings

    def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
        """
        # 剪掉模型多头注意力机制中的一些头，heads_to_prune 是一个字典，键为layer_num，值为需要剪枝的 heads 列表。
        for layer, heads in heads_to_prune.items():
            self.h[layer].attn.prune_heads(heads)

    @add_start_docstrings_to_model_forward(OPENAI_GPT_INPUTS_DOCSTRING)
    @add_code_sample_docstrings(
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=BaseModelOutput,
        config_class=_CONFIG_FOR_DOC,
    )
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutput]:
        # 根据 config 配置设定 output_attentions, output_hidden_states, return_dict 的值
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # 获取 input_ids 或者 inputs_embeds 以及 input_shape
        if input_ids is not None and inputs_embeds is not None:  # 当 input_ids 和 inputs_embeds 同时存在时，抛出错误
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:  # 如果存在 input_ids，将其形状调整为 (batch_size, sequence_length)
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
        elif inputs_embeds is not None:  # 如果存在 inputs_embeds，获取其形状
            input_shape = inputs_embeds.size()[:-1]
        else:  # 如果 input_ids 和 inputs_embeds 都不存在，抛出错误
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        # 如果没有传入 position_ids，则生成默认的 position_ids
        if position_ids is None:
            # Code is different from when we had a single embedding matrix from position and token embeddings
            position_ids = self.position_ids[None, : input_shape[-1]]

        # ------------------------------------- 1. 获取 attention_mask -----------------------------#
        # Attention mask.
        if attention_mask is not None:
            # We create a 3D attention mask from a 2D tensor mask.
            # Sizes are [batch_size, 1, 1, to_seq_length]
            # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
            # this attention mask is more simple than the triangular masking of causal attention
            # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)  # 将 2D 掩码扩展为 3D 掩码，适用于批量输入

            # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
            # masked positions, this operation will create a tensor which is 0.0 for
            # positions we want to attend and the dtype's smallest value for masked positions.
            # Since we are adding it to the raw scores before the softmax, this is
            # effectively the same as removing these entirely.
            # 将注意力掩码转换为与模型参数相同的数据类型，并进行数值变换，torch.finfo(self.dtype).min 返回数据类型的最小值。
            attention_mask = attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
            attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
        # ----------------------------------------------------------------------------------------#

        # ------------------------------------- 2. 获取 head_mask ---------------------------------#
        # Prepare head mask if needed
        head_mask = self.get_head_mask(head_mask, self.config.n_layer)
        # ---------------------------------------------------------- -----------------------------#

        # ------------------------------------- 3. 获取 hidden_states -----------------------------#
        # 如果 inputs_embeds 为 None，则使用 tokens_embed 对 input_ids 计算
        if inputs_embeds is None:
            inputs_embeds = self.tokens_embed(input_ids)
        # 计算 position_embeds
        position_embeds = self.positions_embed(position_ids)
        # 如果存在 token_type_ids，使用 tokens_embed 计算；否则 token_type_embeds 为 0
        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1))
            token_type_embeds = self.tokens_embed(token_type_ids)
        else:
            token_type_embeds = 0
        # 计算 hidden_states，即inputs_embeds、position_embeds 和 token_type_embeds 之和，并使用 dropout
        hidden_states = inputs_embeds + position_embeds + token_type_embeds
        hidden_states = self.drop(hidden_states)
        # -------------------------------------------------------------------------------------#

        # 获取输出形状，以及初始化输出结果 all_attentions 和 all_hidden_states
        output_shape = input_shape + (hidden_states.size(-1),)
        all_attentions = () if output_attentions else None
        all_hidden_states = () if output_hidden_states else None

        # -----------------------------------4. Block逐层计算处理（核心部分）--------------------#
        for i, block in enumerate(self.h):
            # 如果需要输出 hidden states，将当前 hidden_states 添加到 all_hidden_states
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            # 通过当前 Block 处理 hidden_states，得到新的 hidden_states 和 attentions
            outputs = block(hidden_states, attention_mask, head_mask[i], output_attentions=output_attentions)
            hidden_states = outputs[0]
            # 如果需要输出 attentions，将当前 attentions 添加到 all_attentions
            if output_attentions:
                all_attentions = all_attentions + (outputs[1],)
        # ---------------------------------------------------------------------------------#

        # 将 hidden_states 的形状调整为输出形状
        hidden_states = hidden_states.view(*output_shape)

        # 如果需要输出 hidden states，将最后的 hidden_states 添加到 all_hidden_states
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        # -----------------------------------5. 根据配置的输出方式输出结果-------------------------------#
        if not return_dict:
            return tuple(v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None)

        return BaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )