GPT代码分析

主要参考代码：NanoGPT代码 中的 model.py

部分参考代码：从零开始手搓一个LLM 中的 model.py

部分参考代码：Transformers库中的GPT2模型

核心代码逻辑

对应代码中的：model.py/GPT/forward

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) # (b, t, n_embd)
    pos_emb = self.transformer.wpe(pos) # (t, n_embd)
    x = self.transformer.drop(tok_emb + pos_emb)
    for block in self.transformer.h:
        x = block(x)
    x = self.transformer.ln_f(x) # (b, t, n_embd)

    # 处理模型的输出
    if targets is not None: 
        # 训练模式 : (b, t, n_embd) --> (b, t, n_vocab)
        logits = self.lm_head(x) 
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
    else: 
        # 推理模式 : (b, t, n_embd) --> (b, 1, n_embd)--> (b, 1, n_vocab)
        # 这里使用了一个额外的 []，这使得结果保持了原始的三维形状。
        logits = self.lm_head(x[:, [-1], :]) 
        loss = None
    return logits, loss 

上面的代码还没有涉及怎么将概率分布采样为embedding，接下来进行完整分析。

对应代码中的：model.py/GPT/generate

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:]

        # 调用了forward函数
        logits, _ = self(idx_cond) 
        logits = logits[:, -1, :] / temperature # logits：(b, n_vocab)

        # 将topk以外的部分值置为无穷小，概率置为零
        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)

        # 从多项分布中采样，得到下一个字的embedding，并拼接
        idx_next = torch.multinomial(probs, num_samples=1)
        idx = torch.cat((idx, idx_next), dim=1)
    return idx

模型具体结构

宏观结构定义

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,))

• 如何计算模型中所有可学习的参数？

# self.parameters() 返回模型中所有可学习的参数（即权重和偏置）
# p.numel() 计算张量中元素的总数量

n_params = sum(p.numel() for p in self.parameters())

• 如何初始化模型的所有权重?

# apply() 接受一个函数，并将这个函数应用到模型的每个子模块上。
# self._init_weights 是一个函数，所以初始化就是self.apply(self._init_weights)
# torch.nn.init.normal_将张量的值从正态分布中采样

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)

具体模块定义

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

• LN 层的具体实现

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)

• MLP 层的具体实现

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)

• CausalSelfAttention 层的具体实现

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

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

    # 这个拆成多个头，我只能说不要太聪明，最后维度变成 (B, nh, T, hs)
    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) 

    # causal self-attention; 
    # Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
    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)))
        # 下面的这个bias由于注册过缓冲区所以会实现因果注意力！！！牛
        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)
    # contiguous() 保证其存储在内存中的顺序与其在张量空间中的顺序是一致的。
    y = y.transpose(1, 2).contiguous().view(B, T, C) 

    # 等维度的输出映射，再加一个dropout
    y = self.resid_dropout(self.c_proj(y))
    return y

• 理解mask_attention是防止padding的部分被关注！而不是实现因果注意力！

参考链接：transformer中: self-attention部分是否需要进行mask？

• 如何实现因果注意力？

# scaled_dot_product_attention中自带实现，所以如果版本好就不用手动实现了
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 
    # torch.tril 创建了一个下三角矩阵，其左下角的元素为 1，其余元素为 0
    # .view 将这个下三角矩阵重塑为 4 维张量，以匹配后续注意力计算的要求
    # self.register_buffer() 将bias注册到缓冲区
    self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                .view(1, 1, config.block_size, config.block_size))

补充：缓冲区是可以被模型访问但不会被优化器更新的张量，通常用于存储模型的固定参数、常量或中间结果。这些参数通常是不需要梯度的，但需要在模型的前向传播过程中被访问到。

或者，直接给torch.nn.functional.scaled_dot_product_attention() 传参数is_causal = True

• 为什么要用contiguous()保证存储顺序是一致的？

当我们对张量进行一些操作（例如转置、reshape等）后，可能会改变张量的存储顺序，使得它不再满足“行优先”的规则。这会导致一些问题，例如在使用一些底层的优化库（如cuBLAS、cuDNN等）进行计算时，这些库通常会假设张量的存储顺序是“行优先”，如果不满足这个假设，可能会导致计算结果错误或性能下降。