【NanoGPT 学习 03】GPT 代码架构分析

409 阅读17分钟

前言

在前面两篇文章中,我将 GPT 架构中主要的模块都一一进行剖析:

  1. 【NanoGPT 学习 01】model.py 代码详解
  2. 【NanoGPT 学习 02】model.py MLP 层和 Block 详解

现在我们将其串联起来,构造一个 GPT 架构。

GPT Model Class

在前面将全部组件类分析完后,终于到了将它们“组合成”一个 GPT 模型的阶段了。

GPT class 有长,有很多方法。但是只需要理清楚 init 和 forward 方法就大致能够理解这个类的作用了,在 pytorch 中这两个类才是核心,其他都是辅助方法。

在分析前,先展示一下 Transformer 论文中的架构图:

image.png

代码实现

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 parameters in the model.
        For non-embedding count (default), the position embeddings get subtracted.
        The token embeddings would too, except due to the parameter sharing these
        params are actually used as weights in the final layer, so we include them.
        """
        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)

    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) # shape (t)

        # forward the GPT model itself
        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (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)

        if targets is not None:
            # if we are given some desired targets also calculate the loss
            logits = self.lm_head(x)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        else:
            # inference-time mini-optimization: only forward the lm_head on the very last position
            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
            loss = None

        return logits, loss

    def crop_block_size(self, block_size):
        # model surgery to decrease the block size if necessary
        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
        # but want to use a smaller block size for some smaller, simpler model
        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):
        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
        override_args = override_args or {} # default to empty dict
        # only dropout can be overridden see more notes below
        assert all(k == 'dropout' for k in override_args)
        from transformers import GPT2LMHeadModel
        print("loading weights from pretrained gpt: %s" % model_type)

        # n_layer, n_head and n_embd are determined from model_type
        config_args = {
            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M params
            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
        }[model_type]
        print("forcing vocab_size=50257, block_size=1024, bias=True")
        config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
        config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
        config_args['bias'] = True # always True for GPT model checkpoints
        # we can override the dropout rate, if desired
        if 'dropout' in override_args:
            print(f"overriding dropout rate to {override_args['dropout']}")
            config_args['dropout'] = override_args['dropout']
        # create a from-scratch initialized minGPT model
        config = GPTConfig(**config_args)
        model = GPT(config)
        sd = model.state_dict()
        sd_keys = sd.keys()
        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param

        # init a huggingface/transformers model
        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
        sd_hf = model_hf.state_dict()

        # copy while ensuring all of the parameters are aligned and match in names and shapes
        sd_keys_hf = sd_hf.keys()
        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
        # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
        # this means that we have to transpose these weights when we import them
        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
        for k in sd_keys_hf:
            if any(k.endswith(w) for w in transposed):
                # special treatment for the Conv1D weights we need to transpose
                assert sd_hf[k].shape[::-1] == sd[k].shape
                with torch.no_grad():
                    sd[k].copy_(sd_hf[k].t())
            else:
                # vanilla copy over the other parameters
                assert sd_hf[k].shape == sd[k].shape
                with torch.no_grad():
                    sd[k].copy_(sd_hf[k])

        return model

    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
        # start with all of the candidate parameters
        param_dict = {pn: p for pn, p in self.named_parameters()}
        # filter out those that do not require grad
        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
        optim_groups = [
            {'params': decay_params, 'weight_decay': weight_decay},
            {'params': nodecay_params, 'weight_decay': 0.0}
        ]
        num_decay_params = sum(p.numel() for p in decay_params)
        num_nodecay_params = sum(p.numel() for p in nodecay_params)
        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
        # Create AdamW optimizer and use the fused version if it is available
        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
        use_fused = fused_available and device_type == 'cuda'
        extra_args = dict(fused=True) if use_fused else dict()
        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
        print(f"using fused AdamW: {use_fused}")

        return optimizer

    def estimate_mfu(self, fwdbwd_per_iter, dt):
        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
        # first estimate the number of flops we do per iteration.
        # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
        N = self.get_num_params()
        cfg = self.config
        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
        flops_per_token = 6*N + 12*L*H*Q*T
        flops_per_fwdbwd = flops_per_token * T
        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
        # express our flops throughput as ratio of A100 bfloat16 peak flops
        # Add a small numbers to avoid ZeroDivisionError。
        # This usually happens when the time interval between two consecutive calls to time.time() is so small that it is considered as 0 under floating point precision.
        dt = dt + 1e-10
        flops_achieved = flops_per_iter * (1.0/dt) # per second
        flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
        mfu = flops_achieved / flops_promised
        return mfu

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
        the sequence max_new_tokens times, feeding the predictions back into the model each time.
        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
        """
        for _ in range(max_new_tokens):
            # if the sequence context is growing too long we must crop it at block_size
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            # forward the model to get the logits for the index in the sequence
            logits, _ = self(idx_cond)
            # pluck the logits at the final step and scale by desired temperature
            logits = logits[:, -1, :] / temperature
            # optionally crop the logits to only the top k options
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            # apply softmax to convert logits to (normalized) probabilities
            probs = F.softmax(logits, dim=-1)
            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1)
            # append sampled index to the running sequence and continue
            idx = torch.cat((idx, idx_next), dim=1)

        return idx

代码简要分析

构造函数 __init__​:

  • 这个函数初始化 GPT 模型,其中 config​ 参数包含了模型的各种配置,如词汇量大小 (vocab_size​)、块大小 (block_size​)、嵌入维度 (n_embd​) 等。

  • 检查配置确保 vocab_size​ 和 block_size​ 都已被指定。

  • 使用 nn.ModuleDict​ 和 nn.ModuleList​ 来构建 transformer 结构:

    • **词嵌入 (wte**​ ) : 将输入的词索引转换为向量表示。
    • **位置嵌入 (wpe**​ ) : 由于 GPT 使用的是自注意力机制,它需要位置嵌入来理解词的顺序。
    • **Dropout (drop**​ ) : 用于防止过拟合。
    • **Transformer 块 (h**​ ) : GPT 模型的核心,由多个相同的块组成。每个块内包含自注意力层和前馈网络。
    • **层归一化 (ln_f**​ ) : 应用于输出层,确保训练稳定。
  • 语言模型头 (lm_head​): 线性层,用于将 transformer 的输出转换回词汇空间的得分。

  • 参数初始化 (self._init_weights​): 根据 GPT-2 论文的建议,对模型的参数进行初始化。

  • 权重绑定 (weight tying​): 跟据 Weight Tying Explained | Papers With Code 的建议,词嵌入和语言模型头的权重是共享的。

forward​ 方法:

  • 这是模型的前向传播函数,接收索引 (idx​) 和可选的目标 (targets​)。
  • 首先,它会检查输入序列的长度是否超出模型允许的最大块大小 (block_size​)。
  • 利用词嵌入和位置嵌入将索引转换为向量。
  • 应用 dropout。
  • 通过 transformer 块传递数据。
  • 应用最后一层的归一化。
  • 如果提供了目标,则计算模型的损失;否则,只返回最后一位置的输出的 logits。

辅助方法和训练支持:

  • **参数初始化 (**​ **_init_weights**​ ) : 根据特定规则初始化模型的权重。
  • **从预训练加载 (from_pretrained**​ ) : 加载预训练的 GPT 模型。
  • **配置优化器 (configure_optimizers**​ ) : 根据配置选择优化器,并正确设定参数组(考虑权重衰减)。
  • **生成文本 (generate**​ ) : 在给定某些初始文本后,进一步生成文本。
  • **模型参数计数 (get_num_params**​ ) : 计算模型的参数数量。
  • **减小块大小 (crop_block_size**​ ) : 对模型进行修改,以使用比原始配置更小的块大小。
  • **估计模型效率 (estimate_mfu**​ ) : 计算模型在特定硬件上的效率,使用的是理论浮点运算数。

整体而言,这个类提供了一个 GPT 模型的 PyTorch 实现,包含了模型构建、参数初始化、前向传播以及其他辅助功能,适用于语言建模任务和文本生成。

init 初始化方法

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

基本验证和配置设置

  • super().__init__()​:这段代码调用了父类 nn.Module​ 的构造函数,确保所有继承自 nn.Module​ 的初始化代码都能被执行。
  • assert​ 语句确保传入的配置 (config​) 中必须包含 vocab_size​(词汇表的大小)和 block_size​(处理的最大序列长度)两个关键参数。这是因为,在构建模型时,这两个参数是必不可少的。

构建模型部件

nn.ModuleDict

首先,self.transformer​ 是通过 nn.ModuleDict​ 创建的,这是一种包含(key, module)对的有序字典,其中的每个模块都可以是神经网络中的一部分。在 PyTorch 中使用 ModuleDict​ 可以方便地对这些部件进行组织和索引。在本例中,self.transformer​ 包含了构成 Transformer 模型的所有核心部件。

词嵌入(Word Embedding) - wte

  • nn.Embedding(config.vocab_size, config.n_embd)​: 这个部分创建了一个词嵌入层,它的作用是将词(通常通过词汇索引表示)转换为固定维度的向量。这里,config.vocab_size​ 表示词汇表的大小,config.n_embd​ 是嵌入向量的维度。通过这种方式,模型可以学会每个词的密集向量表示,这比传统的 one-hot 编码方式 更为高效和表达力强。

位置嵌入(Positional Embedding) - wpe

  • nn.Embedding(config.block_size, config.n_embd)​: 位置嵌入层也是采用嵌入的方式,但它用来表示序列中词的位置信息。在 Transformer 模型中,由于采用的是自注意力机制,这种位置信息对于模型能够理解词序非常关键。config.block_size​ 表示一个序列中最大的词数量(即模型能处理的最大序列长度),这里同样使用了嵌入向量来表示每个位置。

Dropout - drop

  • nn.Dropout(config.dropout)​: 这是一个正则化技术,旨在防止模型过拟合。它通过在训练时随机"丢弃"一部分神经网络连接(即随机将部分神经元的输出设置为0),使得模型学习到的特征更加鲁棒。

Transformer块 - h

  • nn.ModuleList([Block(config) for _ in range(config.n_layer)])​: 这部分构造了一个由多个Transformer块组成的列表。Transformer块是模型处理输入和学习表示的核心,每个块由自注意力层和前馈网络组成。config.n_layer​ 指定了块的数量,通过堆叠多个这样的块,模型可以学习到更加复杂和高级的特征表示。

层归一化(Layer Normalization)- ln_f

  • LayerNorm(config.n_embd, bias=config.bias)​: 层归一化是另一种常用的正则化技术,它在模型的每个小批量处理中,对输入的每一层的特征进行归一化,使得输出的均值为0且标准差为1。这有助于加速模型的收敛,同时提高模型的稳定性。

self.transformer.wte.weight = self.lm_head.weight​ 权重共享深入理解

当我们谈论“权重共享”(Weight Sharing)时,我们可以把它想象成一位多面手艺人,这位艺人不仅擅长画画,还擅长雕塑;而且,让人惊讶的是,他使用画画的技巧来加强他的雕塑技艺,反之亦然。在深度学习,特别是自然语言处理(NLP)领域中,"权重共享" 也采用了类似的思想来提升模型的学习效率和效果。

权重共享是什么?

在NLP中,一个模型要学习的基本技能之一就是处理和理解语言中的单词。为了让计算机能够理解这些单词,我们通常会使用所谓的“词嵌入”技术,将每个单词转化为一个数值型的向量。你可以把每个词的嵌入向量想象成该词的多维空间表示,每个维度捕捉了这个词的某些语义特征。

权重共享技术在NLP模型中非常典型的应用就是在模型的输入层(将单词转换为向量的部分)和输出层(根据处理的上下文预测下一个单词的部分)使用相同的词嵌入矩阵。这就意味着模型在"理解"(编码)输入的单词和在"表达"(解码)要预测的单词时,使用的是同一套单词的表示。

为什么需要权重共享?

以我们的多面手艺人的比喻来说,通过在画画和雕塑之间共享技巧,艺人不仅能够在两个领域都取得进步,还能减少学习新技巧所需的时间和精力。在NLP模型中,权重共享有以下几个好处:

  • 减少参数数量:通过复用相同的词嵌入矩阵,模型需要学习的参数数量减少了。这意味着,我们需要更少的数据来训练模型,同时也降低了过拟合的风险。
  • 加快学习速度:因为模型在处理输入和产生输出时使用的是相同的单词表示,它可以更快地学习到单词之间的关系和模式。
  • 提高模型效果:权重共享有助于模型更好地理解语言的结构,因为它强制模型在"理解"和"生成"语言时采用一致的方式来处理单词。这通常能够帮助模型在各种语言任务上表现得更好。

权重共享的实践

在实际的模型设计中,实现权重共享通常很直接。如果你在使用一个深度学习框架(比如PyTorch),你可以简单地将输出层的权重设为输入层词嵌入矩阵的引用。这样,当词嵌入矩阵更新时,输出层的权重也会相应更新,反之亦然。

权重初始化

        # 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))
  1. 统一初始化所有权重

    • 首先,代码使用self.apply(self._init_weights)​语句。这个调用将_init_weights​这个函数应用于GPT模型的每一个模块(module)和子模块(sub-module)。_init_weights​函数根据模块的类型对其参数进行初始化。
    • 这里的思路是对不同类型的模块采用不同的初始化策略,比如对nn.Linear​模块(全连接层),使用正态分布(均值0,标准差0.02)来初始化权重,而偏置(如果存在)则初始化为0。
    • 对于nn.Embedding​模块(嵌入层),采用相同正态分布来初始化权重。
    • 这种初始化方式有助于在训练开始时避免梯度消失或者梯度爆炸的问题,同时保证每层的激活值分布大致相同。
  2. 特殊初始化残差投影的权重

    • 紧接着,代码遍历所有命名参数,专门对那些以'c_proj.weight'​结尾的参数执行了一个特殊的初始化过程。这些参数属于残差连接后的投影层(即attention模块输出后的线性层)。代码使用了一个正态分布进行初始化,其均值为0,但标准差不是固定值,而是0.02 / math.sqrt(2 * config.n_layer)​。

      • 这个标准差的计算考虑了模型层数(n_layer​),意味着随着模型层数的增加,每层线性层(projection layer)的权重标准差会减小。
    • 这种“缩放初始化”(scaled initialization)方法是从GPT-2论文中借鉴而来的。它的目的是要平衡前向传播和反向传播中的梯度大小,尤其是在训练深层模型时,这样做可以帮助减轻梯度消失或爆炸的问题。简而言之,这种特殊的初始化方式有利于提高深层网络的训练稳定性和效率。

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) # shape (t)

        # forward the GPT model itself
        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (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)

        if targets is not None:
            # if we are given some desired targets also calculate the loss
            logits = self.lm_head(x)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        else:
            # inference-time mini-optimization: only forward the lm_head on the very last position
            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
            loss = None

        return logits, loss

代码逻辑

初始化和输入处理

  1. 获取设备和输入大小

    • device = idx.device​获取输入序列(索引)所在的设备(CPU或GPU),以确保所有新创建的张量也在同一设备上。
    • b, t = idx.size()​获取输入序列的批次大小(b​)和序列长度(t​)。
  2. 断言检查

    • 检查输入序列长度是否不大于模型配置中的block_size​。这是因为模型被设计为只能处理或“看到”限定长度的序列。
  3. 生成位置向量

    • pos = torch.arange(0, t, dtype=torch.long, device=device)​生成一个从0到t-1​的位置向量,这个向量的长度与输入序列长度相同。每个位置即将与对应的词嵌入一起被用于表示序列中的单词。

前向传播

  1. 词嵌入与位置嵌入

    • 使用词嵌入层(wte​)将索引转换为词嵌入,得到形状为(b, t, n_embd)​的张量tok_emb​。
    • 使用位置嵌入层(wpe​)为序列的每个位置生成嵌入,得到(t, n_embd)​形状的pos_emb​。
    • 将词嵌入和位置嵌入相加(tok_emb + pos_emb​),得到表示输入序列的向量,并通过dropout层进行正则化。
  2. 通过Transformer块

    • 迭代地将加权输入x​通过所有的Transformer块(self.transformer.h​)。每个块将自注意力机制和前馈神经网络应用于其输入,并返回输出,该输出再作为下一个块的输入。
  3. 最终层标准化

    • 经过所有Transformer块处理后的输出会通过最终的层归一化(ln_f​),以便进行最后的处理。

计算损失或生成推理输出

  1. 处理目标(如果存在)

    • 如果有目标(targets​)提供,那么用lm_head​计算最终的逻辑层输出logits​,接着计算交叉熵损失。这一步主要用于训练时评估模型性能。
    • 如果没有提供目标,意味着处于推理模式。此时,仅针对最后位置的输出计算logits​,以进行下一词预测。这是一个优化措施,因为在推理时,我们通常只关心序列的最新预测输出。

返回值

  1. 输出

    • 方法最终返回计算得到的logits​(可能是整个序列或仅最后一个词的),以及(如果在训练模式下)计算得到的损失loss​。

总结

在 NanoGPT model.py 的 GPT 类中,其实还有其他辅助方法,但是只需要将 init 和 forward 方法看懂,就能很清楚理解这个类的作用。其他方法都是辅助方法。

如果需要看完整代码的话,可以去这里进行查看:github.com/karpathy/na…