开源GPT？nanoGPT啃代码记实（三）核心BLOCK模块和GPT模块

开源GPT？nanoGPT啃代码记实（三）

项目github：https://link.zhihu.com/?target=https%3A//github.com/karpathy/nanoGPT

今天继续来啃nanoGPT的代码，这个专栏的代码解析讲究一个从0开始，以完全不懂的身份0基础讲解，同时附上扒代码时候的个人理解。

模型架构model.py

MLP模块和BLOCK模块

class MLP(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.gelu    = nn.GELU()
        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        x = self.dropout(x)
        return x

这个MLP模块继承自pytorch的nn.Module，是一个简单的多层感知机模块，主要做了几件事，输入数据x首先经过c_fc线性变换，然后应用GELU激活函数，接着通过c_proj进行线性投影，最后经过Dropout层得到最终的输出结果。
这里有两个值得注意的地方，一个是为啥要用gelu激活函数？第二个是为什么要把神经元映射到4倍长的张量后激活，再变回来？

先介绍基础知识 GELU激活函数：它模拟了线性整流单元（ReLU）的行为，同时保持了更加连续和平滑的输出，这有助于模型的学习和泛化能力。

其中 是标准正态分布的累积分布函数（CDF）。

由于Φ(x)不是初等函数，实际实现中常使用近似算法来计算GELU函数的值： GELU近似表达式: GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2 / π) * (x + 0.044715 * x^3)))

所以为什么要用GELU呢：

1. GELU函数在整个实数域上都是连续和光滑的，这意味着它在任何点都有非零导数，且其输入GELU函数的输出分布接近于高斯分布。相比之下，ReLU函数在负半轴的导数为零，一旦输入低于某个阈值（即零），对应的神经元就无法更新权重。

2. ReLU在正区间内保持线性增长，但在较大正值时容易出现饱和现象，而GELU的渐进行为更为平缓，减少了饱和效应。

在实际应用中，尤其是对于自然语言处理任务，如BERT模型，GELU作为激活函数已被证明能提升模型的表现，总之由于其更复杂的非线性特性有助于模型捕捉更丰富的特征模式，使得在更复杂的大模型中，GLEU更受到青睐。

另一个问题是，为什么吃饱了没事干要投射到4层神经元上呢？其实很容易理解，就是增大模型的容量，更多的参数，给训练时以更多的拟合空间：

1. 增大模型容量： 将特征维度暂时扩大至四倍，增加了模型的学习能力，允许模型在更大空间中表达更复杂的函数关系，引入更多的非线性转换。

2. 保持局部感知性： 虽然中间层的维度增大，但最终通过第二个线性层又投影回原来的维度，这样既增强了模型的全局学习能力，又能保持对局部信息的敏感性。

接下来我们看Block模块，可以理解为transformer中的decoder模块了，由多个Attention功能的模块组成。

class Block(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

前向传播（forward方法）：

1. 首先，对输入x进行第一次Layer Normalization（使用self.ln_1），然后将其传递给self.attn执行自回归的注意力运算，将注意力机制的结果与未经处理的原始输入x相加，这是Transformer中经典的残差连接（residual connection）结构，便于梯度传播和防止过拟合训练。

2. 接下来，对上一步骤得到的结果再次进行Layer Normalization（使用self.ln_2），然后将其传递给self.mlp模块执行多层感知机操作。

3. 最后，将MLP处理后的结果与之前经过注意力机制后的结果再次相加，这也是一个残差连接结构，然后返回最终的输出结果。

附上自己画的block结构

整体来看，这个Block模块综合运用了Layer Normalization、自回归的Causal Self-Attention机制和MLP结构，共同构成了Transformer模型中的一个基础计算单元。通过堆叠多个此类Block来构建深层网络。

Block模块之外，作者写了一个model中最重要的模块，GPT模块，实际上前面的模块在我们transformer中都有所涉猎，只有GPT模块是在大模型中特有的东西。

@dataclass
class GPTConfig:
    block_size: int = 1024
    vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0
    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster

class GPT(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.vocab_size is not None
        assert config.block_size is not None
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = LayerNorm(config.n_embd, bias=config.bias),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        # with weight tying when using torch.compile() some warnings get generated:
        # "UserWarning: functional_call was passed multiple values for tied weights.
        # This behavior is deprecated and will be an error in future versions"
        # not 100% sure what this is, so far seems to be harmless. TODO investigate
        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying

        # init all weights
        self.apply(self._init_weights)
        # apply special scaled init to the residual projections, per GPT-2 paper
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))

        # report number of parameters
        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))

    def get_num_params(self, non_embedding=True):
        """
        Return the number of paramet

首先作者给出了一个 GPTConfig 数据类（使用 @dataclass 装饰器），用于存储GPT模型的超参数，其中几个重要的如下：

• block_size: 模型处理的最大序列长度，默认为1024。

• vocab_size: 模型词汇表的大小，默认设置为50304，针对GPT-2进行了微调以适应64的倍数，以提高效率。

• n_layer: Transformer编码器的层数，默认为12层。

接下来就是最重要的GPT类了 GPT类还是继承自 PyTorch 的 nn.Module 类

init 方法首先验证了vocab_size 和 block_size 是否已设置，并保存了传入的配置对象 config，说明了这俩参数的重要性。

定义了一个 transformer 模块字典包含：

1. wte：词嵌入层，将词汇表中的单词索引映射到高维向量空间。

2. wpe：位置嵌入层，用于编码输入序列的位置信息。

3. drop：Dropout层，挺熟悉不解释。

4. h：一个包含多个 Block 对象的列表，这些 Block 对象按照配置文件中的 n_layer 参数数量堆叠起来构成Transformer的核心部分。 这个层就是将我们之前构建好的block层叠加起来的核心层了，其中封装着重要的多层Block。

5. ln_f：最后一层LayerNorm层，在模型的最后阶段应用。

定义了一个线性层 lm_head，用于从隐藏状态预测词汇表上的概率分布。 将词嵌入层和解码头权重绑定在一起，采用权重共享技术（Weight Tying）。
这里需要注意为什么要将这两者绑定： 首先self.transformer.wte.weight 是词嵌入层的权重矩阵，用于将输入的词索引映射到相应的向量表示。而self.lm_head.weight 是解码层的权重矩阵。在GPT以及其他一些Transformer-based模型中，常常采用“权重绑带”（Weight Tying）的技术，即将词嵌入层的权重与解码层的权重共享，即执行 self.transformer.wte.weight = self.lm_head.weight。 这样做基于以下理由：

• 减少参数数量：共享权重可以降低模型参数总数，特别是在词汇表很大的情况下，显著减小模型的规模。

• 正则化效果：权重绑带可以被视为一种隐含的正则化技术，它限制了解码器在学习新的映射时的自由度，鼓励模型在嵌入空间中学习到更通用的表示。

• 理论上的一致性：在语言模型中，词嵌入应该反映词在语境中的含义和用法，而解码器的任务正是基于先前的上下文预测下一个可能出现的词。让这两个层共享权重有助于模型在学习词嵌入时同时考虑到生成任务的需求。

在这之后，初始化所有权重，并对特定类型的层（这里是指 residual projections 层）应用特殊的初始化策略，源自GPT-2论文中的建议。 随后打印出模型参数的数量。最后使用get_num_params 方法用于获取模型参数的数量（可选择排除词嵌入层参数）。

好的，最主要的类介绍完了！下一篇准备介绍剩下的train.py时作者的训练技巧。