【NanoGPT 学习 01】model.py 代码详解

NanoGPT model.py 详解

前言

笔记介绍

这篇文章是自己在学习 NanoGPT 源码的时候的一些思考。

NanoGPT 是 Andrej Karpathy 大佬的开源项目。大佬是 OpenAI 的工程师，参与过 ChatGPT 项目的开发。

而 NanoGPT 正是他开源的一个入门级别的 LLM 项目，方便初学者学习 Transformer 架构。

以下笔记是对 model.py 模型架构文件的思考记录。

与 llama 的对比

之前看 transformers源码阅读 的时候，整理过 llama 的架构图

‍

可以依次作为参考，分析一下 NanoGPT 和 LLama 有什么相同与不同之处

‍

LayerNorm 层

代码实现

class LayerNorm(nn.Module):
    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """

    def __init__(self, ndim, bias):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(ndim))
        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None

    def forward(self, input):
        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)

​weight​和bias​变量扮演着极其重要的角色。这段代码实现了一个Layer Normalization（层标准化）的变体，其中包括可选的偏置项。我们来详细探讨这些组件的作用及其对Transformer模型的意义。

weight​ 和 bias​ 变量的作用

• **weight** （权重） ：这个参数是层标准化中的一个缩放因子，用于对标准化之后的数据进行重新缩放。在这段代码中，weight被初始化为与特征维度（ndim）大小相同的全1向量。这表明在开始时，并不对数据的缩放做任何改变，但随着训练的进行，这个缩放因子是可学习和调整的。
• **bias** （偏置） ：偏置参数用于在标准化之后对数据进行重新定位。在代码中，当bias参数设为True时，bias被初始化为与特征维度大小相同的全0向量。如果bias设为False，则不会对数据应用偏置项。这个参数也是可学习的，允许在标准化之后对数据进行进一步的调整。

传入到forward​处理函数中的效果

在forward​方法中，执行了层标准化操作，F.layer_norm​函数需要几个关键参数：输入input​、标准化所使用的维度（self.weight.shape​）、可学习的缩放（weight​）和偏置（bias​）参数，以及数值稳定性参数eps​。通过这些参数，层标准化函数对输入数据进行了两个主要步骤的处理：

1. 标准化：首先，它会根据输入数据的均值和方差对每个特征进行标准化，以确保结果具有0均值和单位方差。
2. 重新缩放和重定位 ：然后，通过应用缩放（weight）和偏置（bias）参数，对数据进行了重新缩放和重定位，以便调整数据分布。

对Transformer模型的运作的重要之处

在Transformer模型中，层标准化非常关键，原因包括：

• 促进稳定性：通过在各个子层（例如自注意力机制和前馈神经网络）之后应用层标准化，能够促进训练过程中的数值稳定性，避免梯度消失或梯度爆炸问题。
• 提升训练效率：层标准化有助于标准化输入数据的分布，这可以减少训练所需的时间，因为标准化的数据通常更容易通过梯度下降等优化算法进行学习。
• 增强模型泛化能力：通过重新缩放和重定位，层标准化使模型能够在训练过程中灵活调整数据传播，从而有助于提高模型对不同数据分布的适应性和泛化能力。

‍

Attention 注意力层

代码实现

class CausalSelfAttention(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        # regularization
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout
        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        if not self.flash:
            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
            # causal mask to ensure that attention is only applied to the left in the input sequence
            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                        .view(1, 1, config.block_size, config.block_size))

    def forward(self, x):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        if self.flash:
            # efficient attention using Flash Attention CUDA kernels
            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
        else:
            # manual implementation of attention
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y

‍

我们分 2 部分来看：先看 init 初始化函数

init 初始化分析

在 init 之中，主要初始化了：

• 2 个 nn.Linear 线性层：c_attn、c_proj

  • 1. self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
       这一行代码定义了一个线性变换层，其输入特征维度为config.n_embd，输出特征维度为3 * config.n_embd。这个输出维度是因为它同时处理key、query和value，每个都有config.n_embd维度。bias=config.bias表示是否添加偏置项。
    2. self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
       这行代码定义了另一个线性变换层，用于将自注意力机制的输出维度从config.n_embd变换回config.n_embd。这实际上是为了之后的残差连接准备。

• 2 个 Dropout 层，用于在自注意力机制和残差连接后进行正则化，以避免过拟合

  • attn_dropout、resid_dropout

• 计算注意力的不同方法

  • 对于 pytorch >= 2.0 的版本，支持 Flash Attention。该算法支持 CUDA 加速
  • 如果不支持的话，则只能自己通过"矩阵计算"来获取结果

‍

‍

自己实现的"矩阵乘法"获取当前 token 的关联权重，这个在我之前的笔记中有讲述过：github.com/HildaM/Nano...

大致是通过矩阵计算，获取当前 char 与前面出现过的字符之间的关联权重，具体简化的代码如下：

# version 4: self-attention!
torch.manual_seed(1337)
B,T,C = 4,8,32 # batch, time, channels
x = torch.randn(B,T,C)

# let's see a single Head perform self-attention
head_size = 16
key = nn.Linear(C, head_size, bias=False)   # bias=False，设置固定权重，方便复现结果
query = nn.Linear(C, head_size, bias=False)
value = nn.Linear(C, head_size, bias=False)
k = key(x)   # (B, T, 16)
q = query(x) # (B, T, 16)

# 计算出'注意力权重'
# 在矩阵乘法中，第一个矩阵的最后一个维度和第二个矩阵的倒数第二个维度必须相等
wei =  q @ k.transpose(-2, -1) #  (B, T, 16) @ (B, 16, T) ---> (B, T, T)


tril = torch.tril(torch.ones(T, T))
# wei = torch.zeros((T,T))
wei = wei.masked_fill(tril == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)


v = value(x) 
out = wei @ v
#out = wei @ x

out.shape

‍

forward 方法分析

def forward(self, x):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        if self.flash:
            # efficient attention using Flash Attention CUDA kernels
            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
        else:
            # manual implementation of attention
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y

forward 方法中定义了这一层具体的执行逻辑，在 pytorch 神经网络项目构建中：init 负责初始化变量、forward 方法负责具体执行逻辑。这是一个"标准"。

‍

在这个方法中主要定义了 Attention 的处理逻辑：

1. ​B, T, C = x.size()​：提取输入张量x​的维度：

   1. 批次大小（B, batch size）
   2. 序列长度（T, sequence length）
   3. 嵌入维度（C, embedding dimensionality，等同于n_embd）。

2. Q、K、V 计算

   1. ​q, k, v = self.c_attn(x).split(self.n_embd, dim=2)​

      1. 将输入x通过self.c_attn线性层来计算key（k）、query（q）和value（v）向量。这个线性层的输出是三倍于输入维度的，因为它同时为每个元素的key、query和value生成向量。随后，使用.split(self.n_embd, dim=2)将这三个向量分开。

   2. view 维度展开，transpose 维度置换

      1. 通过.view(B, T, self.n_head, C // self.n_head)和.transpose(1, 2)，将key、query和value的维度重排，以将注意力头（attention heads）置于批次维之后，同时将嵌入维度分为头数。
      2. 这样，每个头可以单独对序列的一部分进行操作。

3. Attention 算法处理

   1. 分为 Flash Attention 和普通矩阵计算 2 种处理方法

4. ​y = y.transpose(1, 2).contiguous().view(B, T, C)​

   1. 将注意力头合并回原始的嵌入维度中，为输出投影做准备。

5. Droupout 防止过拟合处理

   1. y = self.resid_dropout(self.c_proj(y))
   2. 将合并后的输出通过另一个线性层self.c_proj进行变换，然后应用残差dropout正则化，以避免过拟合。

‍

为什么要将 (B, T, C) 处理为 (B, n_head, T, c // n_head)

我们有以下变量定义：

• B: 批次大小（Batch size）
• T: 序列长度（Sequence length）
• C: 输入的总特征维度
• n_head: 头的数量
• C // n_head: 每个头处理的特征子集的维度

在多头自注意力中，原始的查询（query）、键（key）和值（value）张量都是形状为(B, T, C)​的三维张量。然后，为了实现多头处理，我们会将每个张量重塑为(B, T, n_head, C // n_head)​，这使我们可以在最后两个维度上拥有独立的头处理不同的特征子集。

为什么要交换第一和第二维度？

• 在重塑后，我们得到形状为(B, T, n_head, C // n_head)的张量。然而，为了在并行处理每个头的注意力计算时有效地组织数据，我们会希望将表示头数的维度（n_head）放到更前面的位置。这样，注意力计算可以在头维度上更自然地并行进行。
• 因此，执行.transpose(1, 2)操作后，张量的形状变为(B, n_head, T, C // n_head)。通过这种方式，每个头现在都可以独立处理其对应的序列部分，并且所有的这些处理可以并行执行。
• 在这种新的维度安排下，每个头都能够独立地对整个序列进行操作，而不是仅处理序列中的片段。这对于执行注意力机制至关重要，因为每个头需要访问整个序列来计算注意力分数。

这一步是实现多头注意力机制的关键步骤，它使模型能够利用多个头平行捕获信息的不同方面，增加了模型的表现力和灵活性。每个头捕获的信息在后续步骤中会被合并，形成最终的输出，这一输出可以捕捉到输入数据的更加丰富和细腻的特征。

一个实际例子讲解这样做的好处

假设我们在处理一个自然语言处理任务，其中：

• 批次大小（B）= 2，意味着我们同时处理2句话。
• 序列长度（T）= 3，每句话由3个单词组成。
• 输入的总特征维度（C）= 4，每个单词被表示为一个4维的向量。
• 头的数量（n_head）= 2，我们希望用两个不同的头来分别捕获不同的信息。

‍

原始的输入数据形状为(B, T, C)​，即(2, 3, 4)​。代表我们有2句话，每句话3个单词，每个单词用4维向量表示。

为了应用多头注意力，我们将原始数据重塑为(B, T, n_head, C // n_head)​，于是变成(2, 3, 2, 2)​的形状。这代表我们将每个4维的单词向量拆分为两个部分，分别让两个头处理。

在没有交换维度之前，对于每个头来说，它们是按照序列（单词）和特征维度来组织的。但我们的目标是让每个头能够独立并行地对整个输入序列进行操作，而不是在序列的维度上进行操作。为了达到这个目的，我们需要将头数的维度提前。

经过.transpose(1, 2)​操作后，形状变为(B, n_head, T, C // n_head)​，即(2, 2, 3, 2)​。这样，对于每个头来说，它现在可以独立地并且平行地处理两句话中的所有3个单词。

‍

假如我们有两句话：

1. "I like apples"
2. "You eat oranges"

在不使用维度交换的情况下，每个头在处理单词时要分别对不同句子中的相同位置的单词运行，这不是真正并行且效率较低。但在维度交换后，每个头可以同时、独立地处理两句话中的所有单词，比如头1可以同时处理"I", "like", "apples"和"You", "eat", "oranges"的第一部分，而头2处理它们的第二部分，这使得处理更加高效和平行化。

通过这种重新组织数据的方式，每个头都可以更加高效地并行处理整个输入数据，增强了模型的学习和表示能力，使每个头都能够专注于捕捉输入数据的不同特征和维度，从而提高模型整体的性能。

‍

‍

y.transpose(1, 2).contiguous().view(B, T, C)​ 合并头维度详解

在处理矩阵的时候，需要将原本 (B, nh, T, hs)​ 的矩阵形状，通过合并 (nh、hs）2 个维度的数据，变回原来的 (B, T, C)​ 形状。

这是为了"回到原始的嵌入空间"。

‍

"合并头维度"回到原始的嵌入空间是通过对多头注意力机制的输出进行重整和线性变换实现的。这一步是多头注意力机制的一个关键环节，它允许模型从多个表示子空间（每个头关注输入的不同方面）中合成信息。

多头注意力机制的设计初衷是让模型能够并行地从不同的表示子空间学习信息。每个"头"实际上是同一组参数的不同实例，它们各自从输入数据中捕获不同方面的特征。例如，一个头可能专注于序列中的语法结构，而另一个头可能捕获词语之间的语义关系。

头维度的分离和合并过程

1. 分离头维度：在输入数据上应用线性变换以生成query、key和value表示，然后将这些表示分割成多个头。每个头处理的是输入的一部分表示，这样可以并行处理多个子空间中的信息。
2. 合并头维度 ：在应用了注意力机制之后，我们得到了形状为(B, self.n_head, T, C // self.n_head)的张量------即，每个头生成的表示现在需要被合并回一个单一的高维空间以便于后续的处理或预测。

代码实例：

import torch

# 假设输入张量shape为：(B, n_head, T, C)，其中B=2, n_head=4, T=3, C=5
# 初始化一个随机张量来模拟多头注意力的输出
inputs = torch.rand(2, 4, 3, 5)

# 显示原始张量的形状
print("Original shape:", inputs.shape)  # 将输出 (2, 4, 3, 5)

# 第一步：合并头维度。我们要将形状转变为(B, T, n_head*C)，即将头维度和特征（嵌入）维度合并
# 这里我们首先对维度进行交换，然后使用view方法调整形状
merged = inputs.transpose(1, 2).reshape(2, 3, 4*5)  # 维度变化：先交换维度，然后reshape

# 显示处理后张量的形状
print("Merged shape:", merged.shape)  # 将输出 (2, 3, 20)

如何实现回到原始的嵌入空间

过程中如何具体实现的：

• 重整 ：首先，通过转置和调整形状操作，将多个头的输出合并。这里，我们先将头维度和序列长度的维度交换（transpose），然后使用view将头维度重新合并入嵌入维度中。这样子就从(B, self.n_head, T, C // self.n_head)变成了(B, T, C)的形状。此处的C是原始的嵌入维度，这意味着所有头的信息现在被压缩回了一个单一的高维表示中。
• 线性变换：在合并完成后，一般还会通过一个额外的线性层（输出层）来确保输出张量与原始的嵌入空间对齐，并可能进行进一步的转换。这一步是可选的，根据模型的设计而定。

‍

合并维度的考量

1. 信息综合利用 ：通过将n_head和C合并，实际上是在把各个头获得的不同"视角"或"解读"进行综合。每个头针对输入的特定方面给出了加权重要性，合并这些维度可以让这些学习到的特征被充分利用而不是被隔离。
2. 接口一致性 ：多头自注意力的设计需要与模型的其他部分（如前馈网络）兼容。通过把n_head和C合并，输出的维度被重新调整为与原来的嵌入维度相同或兼容的维度。这样做是为了确保数据在通过模型的不同部分时不需要复杂的转换，保持了数据流的顺畅和模型结构的简洁。
3. 模型性能优化：分头处理和后续的维度合并使得模型能够在较低的计算成本下捕捉到高维度的交互，相比于直接在高维空间进行单一的注意力操作，这种方式在提升表现的同时保持了较高的效率。

‍

结合论文背景

在《Attention is All You Need》中，多头自注意力机制的提出主要是为了增强模型的注意力机制，使模型具备同时从多个方面理解输入信息的能力。每个"头"能够捕捉到数据的不同特性或模式，而通过合并n_head​与C​的操作则是为了整合这些不同维度的信息，使之能够共同作用于后续的模型结构中。

此外，合并n_head​与C​还考虑到了模型的灵活性和扩展性。这种设计使得研究者和开发者能够根据需要调整头的数量和每个头的维度，而不需对模型的其他部分做出大的修改，有助于在不同的任务和数据集上进行调优和实验。

‍

未完待续......

下一篇：【NanoGPT 学习 02】model.py MLP 层和 Block 详解

‍

‍