【HuggingFace Transformers】LlamaMLP源码解析

1. LlamaMLP 介绍

LlamaMLP 是 LLaMA 模型中的 MLP 层，主要用于对输入特征进行非线性变换。在分片预训练模式下，线性层的权重被切分，分步处理后再进行拼接和求和，而在常规模式下，直接应用线性变换和激活函数处理输入数据。其计算公式为：
$$output=W_{down}\cdot (\sigma (W_{gate}\cdot x+b_{gate})\odot (W_{up}\cdot x+b_{up}))+b_{down}$$

2. LlamaMLP类 源码解析

源码地址：transformers/src/transformers/models/llama/modeling_llama.py

# -*- coding: utf-8 -*-
# @time: 2024/8/28 15:16

import torch
import torch.nn.functional as F

from torch import nn
from transformers.activations import ACT2FN


class LlamaMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config  # 配置参数
        self.hidden_size = config.hidden_size  # 隐藏层的维度
        self.intermediate_size = config.intermediate_size  # 中间层的维度
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)  # 定义第一个线性变换层，将隐藏层映射到中间层
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)  # 定义第二个线性变换层，将隐藏层映射到中间层
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)  # 定义第三个线性变换层，将中间层的输出映射回隐藏层
        self.act_fn = ACT2FN[config.hidden_act]  # 根据配置选择激活函数

    def forward(self, x):
        # 如果是分片预训练
        if self.config.pretraining_tp > 1:
            slice = self.intermediate_size // self.config.pretraining_tp  # 计算每个切片的大小

            gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)  # 将 gate_proj 层的权重按行切分成多个切片
            up_proj_slices = self.up_proj.weight.split(slice, dim=0)  # 将 up_proj 层的权重按行切分成多个切片
            down_proj_slices = self.down_proj.weight.split(slice, dim=1)  # 将 down_proj 层的权重按列切分成多个切片

            gate_proj = torch.cat([F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)  # 对输入 x 应用每个 gate_proj 切片的线性变换，并沿列拼接
            up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)  # 对输入 x 应用每个 up_proj 切片的线性变换，并沿列拼接

            intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)  # 应用激活函数后，与 up_proj 结果逐元素相乘，并沿通道切分成多个张量
            down_proj = [F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)]  # 对每个 intermediate_states 切片应用对应的 down_proj 切片的线性变换
            down_proj = sum(down_proj)  # 将所有 down_proj 切片的结果相加
        else:
            # 如果不是分片预训练，直接进行线性变换和激活函数处理
            down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

        # 返回最终的输出结果
        return down_proj