閒言碎語

我在剛開始接觸 huggingface （後簡稱 hf）的 transformers 庫時候感覺很冗雜，比如就模型而言，有 PretrainedModel, AutoModel，還有各種 ModelForClassification, ModelForCausalLM, AutoModelForPreTraining, AutoModelForCausalLM等等；不僅如此，還設計了多到讓人頭皮發麻的各種 ModelOutput，比如BaseModelOutput, BaseModelOutputWithPast, CausalLMOutput等等。擁有選擇困難症的我選擇退出，所以之前一直沒怎麼用過這個大名鼎鼎的庫。今天咬着牙還是決定看看源碼，把這些東西搞清楚。

1. 長話短說

今天看了一下源碼和官方文檔，梳理過後會發現其實也不是很複雜，簡單理解就兩條：

ModelOutput(transformers.utils.ModelOutput)是所有模型輸出的基類。簡單理解它就是一個字典，在模型的 forward函數裏把原本的輸出做了一下封裝而已，方便用戶能直觀地知道輸出是什麼。例如CausalLMOutput顧名思義就是用於像 GPT 這樣自迴歸模型的輸出。
PreTrainedModel (transformers.modeling_utils.PretrainedModel) 是所有模型的基類。所以你如果看到一個模型取名爲LlamaForCausalLM，那你就可以知道這個模型的輸出格式大概率就是自迴歸輸出，即前面提到的CausalLMOutput。爲什麼說大概率呢，因爲自迴歸輸出還有蠻多種的，趕時間的朋友看到這就可以切換到其他文章了，至此你應該也能瞭解 transformers 最核心的模塊了。感興趣的可以繼續往下看，下面做一個簡單的總結和介紹。

2. 短話長說

2.1 ModelOutput

前面已經介紹過了，ModelOutput是所有模型輸出的基類。下面是其源碼核心部分，一些具體實現代碼刪除了，不過不影響理解。

可以看到 ModelOutput 其實就是一個有序的字典（OrderedDict）。

class ModelOutput(OrderedDict):
    def __init_subclass__(cls) -> None:
        """
        這個方法允許對 ModelOutput 的子類進行定製，使得子類在被創建時能夠執行特定的操作或註冊到某個系統中。
        """
        ...

    def __init__(self, *args, **kwargs):
        """
        初始化 ModelOutput 類的實例。
        """
        super().__init__(*args, **kwargs)

    def __post_init__(self):
        """
        在初始化 ModelOutput 類的實例之後執行的操作，允許進一步對實例進行處理或設置屬性。子類需要用 dataclass 裝飾器
        """
        ...

基於 ModelOutput，hf 預先定義了 40 多種不同的 sub-class，這些類是 Hugging Face Transformers 庫中用於表示不同類型模型輸出的基礎類，每個類都提供了特定類型模型輸出的結構和信息，以便於在實際任務中對模型輸出進行處理和使用。每個 sub-class 都需要用裝飾器 @dataclass。我們以CausalLMOutputWithPast爲例看一下源碼是什麼樣的：


@dataclass
class CausalLMOutputWithPast(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None

爲了保持代碼規範，我們需要在模型的forward函數中對輸出結果進行封裝，示例如下：

class MyModel(PretrainedModel):
    def __init__(self):
        self.model = ...

    def forward(self, inputs, labels):
        output = self.model(**inputs)
        hidden_states = ...
        loss = loss_fn(outputs, labels)
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

這裏簡單介紹以下幾種，更多的可以查看官方文檔和源碼：

BaseModelOutput: 該類是許多基本模型輸出的基礎，包含模型的一般輸出，如 logits、hidden_states 等。
BaseModelOutputWithNoAttention: 在模型輸出中不包含注意力（attention）信息。
BaseModelOutputWithPast: 包含過去隱藏狀態的模型輸出，適用於能夠迭代生成文本的模型，例如語言模型。
BaseModelOutputWithCrossAttentions: 在模型輸出中包含交叉注意力（cross attentions）信息，通常用於特定任務中需要跨注意力的情況，比如機器翻譯。
BaseModelOutputWithPastAndCrossAttentions: 同時包含過去隱藏狀態和交叉注意力的模型輸出。
MoEModelOutput: 包含混合專家模型（Mixture of Experts）輸出的模型。
MoECausalLMOutputWithPast: 混合專家語言模型的輸出，包括過去隱藏狀態。
Seq2SeqModelOutput: 序列到序列模型輸出的基類，適用於需要生成序列的模型。
CausalLMOutput: 用於生成式語言模型輸出的基礎類，提供生成文本的基本信息。
CausalLMOutputWithPast: 生成式語言模型輸出的類，包含過去隱藏狀態，用於連續生成文本的模型。

2.2 PretraiedModel

PreTrainedModel 是 Hugging Face Transformers 庫中定義預訓練模型的基類。它繼承了 nn.Module，同時混合了幾個不同的 mixin 類，如 ModuleUtilsMixin、GenerationMixin、PushToHubMixin 和 PeftAdapterMixin。這個基類提供了創建和定義預訓練模型所需的核心功能和屬性。

以下是 PreTrainedModel 中的部分代碼：

class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMixin, PeftAdapterMixin):
    config_class = None
    base_model_prefix = ""
    main_input_name = "input_ids"
    _auto_class = None
    _no_split_modules = None
    _skip_keys_device_placement = None
    _keep_in_fp32_modules = None
    ...

    def __init__(self, config: PretrainedConfig, *inputs, **kwargs):
        super().__init__()
        ...

在這個基類中，我們可以看到一些重要的屬性和方法：

config_class：指向特定預訓練模型類的配置文件，用於定義模型的配置。
base_model_prefix：基本模型前綴，在模型的命名中使用，例如在加載預訓練模型的權重時使用。
main_input_name：指定模型的主要輸入名稱，通常是 input_ids。
_init_weights 方法：用於初始化模型權重的方法。

在這個基類中，大多數屬性都被定義爲 None 或空字符串，這些屬性在具體的預訓練模型類中會被重寫或填充。接下來我們將看到如何使用 PretrainedModel 類定義 llama 模型。

class LlamaPreTrainedModel(PreTrainedModel):
    config_class = LlamaConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["LlamaDecoderLayer"]
    _skip_keys_device_placement = "past_key_values"
    _supports_flash_attn_2 = True

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

在這個例子中，首先定義了 LlamaPreTrainedModel 類作爲 llama 模型的基類，它繼承自 PreTrainedModel。在這個基類中，我們指定了一些 llama 模型特有的屬性，比如配置類 LlamaConfig、模型前綴 model、支持梯度檢查點（gradient checkpointing）、跳過的模塊列表 _no_split_modules 等等。

然後，我們基於這個基類分別定義了 LlamaModel、LlamaForCausalLM 和 LlamaForSequenceClassification。這些模型的邏輯關係如下圖所示：

LlamaModel是 llama 模型的主體定義類，也就是我們最常見的普pytorch 定義模型的方法、默認的輸出格式爲BaseModelOutputWithPast；

class LlamaModel(LlamaPreTrainedModel):

    def __init__(self, config: LlamaConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        ...

    def forward(self, ...):
        ...
        return BaseModelOutputWithPast(...)

LlamaForCausalLM 適用於生成式語言模型的 llama 模型，可以看到 backbone 就是 LlamaModel，增加了lm_head作爲分類器，輸出長度爲詞彙表達大小，用來預測下一個單詞。輸出格式爲CausalLMOutputWithPast；

class LlamaForCausalLM(LlamaPreTrainedModel):
    # 適用於生成式語言模型的 Llama 模型定義

    def __init__(self, config):
        super().__init__(config)
        self.model = LlamaModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        ...

    def forward(self, ...):
        outputs = self.model(...)
        ... # 後處理 outputs，以滿足輸出格式要求
        return CausalLMOutputWithPast(...)

LlamaForSequenceClassification 適用於序列分類任務的 llama 模型，同樣把 LlamaModel作爲 backbone，不過增加了score作爲分類器，輸出長度爲 label 的數量，用來預測類別。輸出格式爲SequenceClassifierOutputWithPast

class LlamaForSequenceClassification(LlamaPreTrainedModel):
    # 適用於序列分類任務的 Llama 模型定義

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.model = LlamaModel(config)
        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
        ...

    def forward(self, ...):
        outputs = self.model(...)
        ... # 後處理 outputs，以滿足輸出格式要求
        return SequenceClassifierOutputWithPast(...)

每個子類根據特定的任務或應用場景進行了定製，以滿足不同任務的需求。另外可以看到 hf 定義的模型都是由傳入的 config參數定義的，所以不同模型對應不同的配置啦，這也是爲什麼我們經常能看到有像 BertConfig，GPTConfig這些預先定義好的類。例如我們可以很方便地通過指定的字符串或者文件獲取和修改不同的參數配置：

config = BertConfig.from_pretrained(
    "bert-base-uncased"
)  # Download configuration from huggingface.co and cache.
config = BertConfig.from_pretrained(
    "./test/saved_model/"
)  # E.g. config (or model) was saved using *save_pretrained('./test/saved_model/')*
config = BertConfig.from_pretrained("./test/saved_model/my_configuration.json")
config = BertConfig.from_pretrained("bert-base-uncased", output_attentions=True, foo=False)

hf 爲了造福懶人，提供了更加簡便的 API，即 Auto 系列 API。至於有多簡便，看看下面的 demo 就知道了：

from transformers import AutoConfig, AutoModel

# Download configuration from huggingface.co and cache.
config = AutoConfig.from_pretrained("bert-base-cased")
model = AutoModel.from_config(config)

LLM 學習筆記-transformers庫的 PreTrainedModel 和 ModelOutput 到底是什麼？

1. 長話短說

2. 短話長說

2.1 ModelOutput

2.2 PretraiedModel

Pytorch 如何使用 storage 實現參數 offload？

TACC 集羣使用筆記

圖解 vLLM 的推理調度策略

大模型推理框架 vLLM 源碼解析（二）：Block 模塊分配和管理

OpenAI 的視頻生成大模型Sora的核心技術詳解（一）：Diffusion模型原理和代碼詳解

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結