nanoGPT源码浅析（上）

Transformer

Transformer中抛弃了传统的CNN和RNN，整个网络结构完全是由Attention机制组成。更准确地讲，Transformer由且仅由self-Attenion和Feed Forward Neural Network组成。一个基于Transformer的可训练的神经网络可以通过堆叠Transformer的形式进行搭建

Transformer结构

Transformer的Attention机制，将序列中的任意两个位置之间的距离是缩小为一个常量，同时具备更好的并行性。

Transformer本质上是Encoder-Decoder结构

Encoder中分为Self-Attention（处理数据得到加权特征向量）和Feed Forward NN（前馈神经网络，Relu（非线性转换）+线性激活函数（线性变换），每个神经元都与其他神经元相连接）

Decoder的Self-Attention和Feed Forward中多出Encoder-Decoder Attention模块，Self-Attention用于当前文字和前文之间的关系，Encoder-Decoder Attention用于给出当前文字与输入的Encoder向量的关系

Attention

Self-Attention（核心）

允许模型在处理一个序列时，考虑序列中每个元素与其他所有元素的关系并计算关联度（权重）

计算逻辑：

1. 对于输入序列中的每个元素都会计算一个查询向量Q、一个键向量K和一个值向量V（通过学习得到的权重矩阵与输入元素的线性变换所得，如：Qi=WQ*xi）
2. 计算注意力分数score，即查询向量和键向量的点积（相似度计算，如某元素的Q分别与所有元素的V相乘)/缩放因子（键向量的维度平方根，确保梯度稳定，类似于归一化）
3. softmax函数处理某元素所有score，进而得到其他元素对该元素的注意力权重
4. 其他元素的值向量V乘以与该元素的注意力权重再求和，得出该元素的向量输出Z

注：WQ Wk WV 是三个可训练的矩阵，相当于对x线性变换，增强拟合能力

典型参数（base）：encoder 6层，decoder 6层，隐层512维，forward2048，多头8个，KV矩阵维度64（512/8）

Multi-Head Attention

相当于多个不同的self-attention的集成

为了增强拟合能力，定义多组不同的WQ/Wk/WV,分别生成不同的Q/K/V，最终输出不同的Z，拼接所有Z生成拼接矩阵,再乘以W0（输出权重矩阵）降维，

计算规则：当前为首层时，直接对输入编码生成向量X；当前为后续层时，直接使用上一层输出Z

Encoder-Decoder Attention

前文提到decoder中的Encoder-Decoder Attention模块，与Self-attention不同点在于Q为Decoder上一个输出，而K和V来自Encoder的输出。Decode是顺序操作，k节点仅能看到k-1以及之前结果，故称之为masked-attention

损失层

Decode后的向量经一层softmax全连接层之后得到每个单词概率的输出向量。通过CTC等损失函数（转化为字符分数，输出正确的概率的乘积再取负对数）训练模型

位置编码

Transformer模型并没有捕捉顺序序列的能力，无论句子的结构怎么打乱，Transformer都会得到类似的结果（词袋模型）。为了保留位置信息，编码向量时引入位置编码（长度为dmodel的特征向量），即PE(pos,2i）=sin（pos/100002i/model）

• NanoGPT源码解读

NanoGPT是minGPT模型的重写，基于Pytorch编写，两个核心文件为model.py和train.py，前者是模型的定义文件（可以选择从微调1.3B参数的GPT-2，也可以重新训练），后者是模型训练文件

“NanoGPT 是用于训练和微调中型尺度 GPT 最简单、最快的库”

——前特斯拉AI总监、NanoGPT作者Andrej Karpathy

源码解读-model.py

class LayerNorm(nn.Module):
    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)

该类继接受两个参数 ndim 和 bias，forward接受input同时使用layer_norm函数对输入作层归一化处理（没啥好说的）。

class CausalSelfAttention(nn.Module):
        def __init__(self, config):
            super().__init__()
            assert config.n_embd % config.n_head == 0
            self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
            self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
            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
            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")
                self.register_buffer("bias",torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size))

该类用于实现自回归注意力机制，也就是前面说的那个Self-Attention。首先断言config的n_embd能被n_head整除，以确保通过线性变换后得到的矩阵具有相同维度。随后通过线性变换将输入投影为K、Q和V）。nn.Linear定义线性层，权重矩阵W0大小为n_embd×3*n_embd（前文知识，不再赘述）

随后定义将输出映射回嵌入维度的输出投影层，同时创建注意力机制和残差的 dropout 操作。

类中保存头数、嵌入维度、残差比例等超参，并根据输入序列特点创建因果掩码（负无穷大三角矩阵，结合softmax函数特性，以确保注意力仅应用于输入序列左侧）

def forward(self, x):
        B, T, C = x.size()
        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)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        if self.flash:
            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:
            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
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y

该函数实现自回归注意力机制的前向传播过程

首先获取输入张量x形状信息，包括批次大小（B）、序列长度（T）和嵌入维度（C），随后根据注意力头数目和嵌入维度，通过线性变换将x进行投影切分和变换，得出查询向量、键向量和值向量，并作Score计算。

接下来这段代码是Transformer的实现：

      att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))

在得到score矩阵后显然不够，需要进一步处理

      att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))

这里通过因果掩码处理，将未来位置的注意力分数屏蔽为负无穷大

      att = F.softmax(att, dim=-1)

对处理后的score矩阵softmax操作，得出权重矩阵att

      att = self.attn_dropout(att)

使用dropout随机失活，以提升鲁棒性、稳定性，减少训练复杂度和过拟合风险 （计划后续对DROPOUT算法详细讲解）

y = att @ v

通过注意力权重矩阵 att与相乘，最后对其转置和形状变换输出重组，再做一次线性变换并应用残差的 dropout 操作，（减轻梯度爆炸和梯度消失问题，后续详细讲）得到最终的输出y

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


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

MLP类实现了多层感知机模块。

初始化后创建了三个线性变换层c_fc、c_proj和激活函数gelu。c_fc输入大小为n_embd，输出大小为 4 * n_embd；c_proj 的输入大小为 4 * config.n_embd，输出大小为 n_embd，用于将上层的输出映射回嵌入维度。GELU 激活函数用于引入非线性变换。

前向传播方法forward中x经过线性变换c_fc(x)后通过 GELU 激活函数做非线性变换。接下来，再做线性变换c_proj（上层输出维度映射），随机失活后输出处理后的x。

而后者作为Transformer的基本块，实现了LayerNorm层（归一化处理以减少内部协变量偏移问题）ln_1、注意力机制attn、LayerNorm 层ln_2和MLP模块的划分。

简单来讲，上面三组代码段结合在一起实现了Transformer前向传播过程。

CausalSelfAttention类实现了自注意力机制的计算，MLP 类实现了多层感知机（MLP）模块的功能，Block 类将自注意力机制和 MLP 模块组合

前向传播中，x先经自注意力机制和归一化，然后与原始输入张量相加，再经MLP模块和残差连接得到最终的输出张量。

而数据类GPTConfig包含的以下超参数更是不难理解：

超参数有block_size（块）、vocab_size（词表大小）、n_layer（基本块层数）、n_head（注意力头数）n_embd（特征维度）、Dropout(丢弃率)、bias（是否偏置）等

前文只是前菜

接下来是最重要的一部分，

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)
        self.transformer.wte.weight = self.lm_head.weight
        self.apply(self._init_weights)
        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))
        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))

该类定义了 GPT（Generative Pre-trained Transformer）模型结构。

初始化后首先通过断言确保vocab_size和block_size非空，随后将传入的配置参数config存储在模型的config属性中，使nn.ModuleDict定义包含多个子模块的字典transformer。子模块包括词嵌入层（Word Token Embedding，twe，用于映射输入词向量），位置嵌入层（Positional Embedding，wpe，即位置编码），Drop（不赘述），基本块列表h（由多个基本块组成的列表，通过循环创建了n_layer 个基本块，归一层（ln_f，LayerNorm）。

而lm_head 线性层用于将最终的输出特征映射到词维度。

值得一提的是，初始化过程中调用_init_weigh对所有权重进行初始化操作，且依据GPT-2相关论文对残差投影层权重进行了特殊初始化。

最后计算模型的参数数量并输出。

补充一下上述函数源码，不再赘述

参数统计

def get_num_params(self, non_embedding=True):
        n_params = sum(p.numel() for p in self.parameters())
        if non_embedding:
            n_params -= self.transformer.wpe.weight.numel()
        return n_params

权重初始化

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

再次出现forward，这里不同于前文。简单来说，上个前向传播位于Block类中，处理单个基本块的前向传播。第二个前向传播定义在GPT类中，处理整个GPT模型前向传播，包含了多个基本块和嵌入层计算

def forward(self, idx, targets=None):
        device = idx.device
        b, t = idx.size()
        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
        pos = torch.arange(0, t, dtype=torch.long, device=device)
        tok_emb = self.transformer.wte(idx)
        pos_emb = self.transformer.wpe(pos) 
        x = self.transformer.drop(tok_emb + pos_emb)
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)
        if targets is not None:
            logits = self.lm_head(x)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        else:
            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
            loss = None
        return logits, loss

这里首先获取输入张量idx等并作参数合法性检查，随后调用self.transformer.wte(idx) 获取词嵌入层的输出（简单理解token就是离散文本单元，这里将索引映射到连续的实值向量空间中），得token embeddings，形状为 (b, t, n_embd)，对应同批次量、token长度、维度。同时调用 self.transformer.wpe(pos) 获取位置嵌入层的输出，得到 position embeddings（位置嵌入），形状为 (t, n_embd)，将 tok_emb 和 pos_emb相加（广播机制，复制b次pos_emb），处理后得到完整输入张量x。

接下来循环遍历 self.transformer.h 中的每个基本块并传入x,再将输出归一化，结合目标值计算logits并使用交叉熵损失函数计算损失并输出。

@torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

该方法以形状为（b，t）的序列 idx（即人为给定的输入）开始，首先检查上下文是否过长，如果超过block_size则截取部分作为条件序列输入模型，获得序列中最后一个时间步的logits（全连接层输出）。通过除以参数temperature进行缩放控制生成文本多样性。随后对 logits 裁剪保留前 k 个最可能的选项。这样可以限制生成的文本在给定的候选词范围内，softmax函数将其换为概率分布，得到下一个时间步的索引 idx_next并追加到原序列，循环直到生成所需数量的新标记。

最终，生成完整序列idx。

至此大框架梳理完毕。

另外：后续代码优化，不影响理解模型，这里一笔带过

def crop_block_size(self, block_size):
        assert block_size <= self.config.block_size
        self.config.block_size = block_size
        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
        for block in self.transformer.h:
            if hasattr(block.attn, 'bias'):
                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]

该方法在更新新的块大小后通过切片操作裁剪位置嵌入层权重参数，再遍历基本块并将偏置项做对应处理，以调整模型的块大小并减小模型的规模

@classmethod
def from_pretrained(cls, model_type, override_args=None)
def configure_optimizers(self, weight_decay, learning_rate, betas, device_type)

前者是预训练模型加载类方法（GPT2)，可选可不选，故不做赘述;后者是配置优化器方法，只保留需要进行梯度更新的参数。因源代码过长，节省篇幅，有兴趣的读者可自行学习