语言模型训练与推理：从概念到代码

一、简介

尽管大型语言模型（LLMs）已经取得了诸多成就，但驱动所有这些模型的核心概念却十分简单——我们只需要准确预测下一个token即可！尽管有些人（合理地）认为最近关于 LLMs 的研究已经超越了这一基本理念，但token预测仍然支撑着所有因果语言模型的预训练、微调（取决于具体变体）和推理过程，使其成为任何 LLMs 从业者必须理解的基本且重要的概念。

“令人惊讶的是，所有这些进展背后的仍然是原始的自回归机制，该机制逐个、从左到右地做出token级别的决策。”

在本概述中，我们将深入并实际地探讨token预测的概念，以了解语言模型如何在训练和推理过程中使用它。首先，我们将从概念层面学习这些理念。然后，我们将通过一个实际的实现（使用 PyTorch）来展示语言模型的预训练和推理过程，使token预测的概念更加具体。

1.1 相关背景概念

在深入探讨本概述主题之前，我们需要理解一些基本概念。在本节中，我们将快速概述这些重要概念，并为每个概念提供进一步阅读的链接。

Transformer 架构。首先，我们需要对 Transformer 架构有一个基本理解，特别是仅解码器的变体。幸运的是，我们之前已经广泛讨论过这些概念：

更根本地说，我们还需要理解自注意力的概念及其在 Transformer 架构中的作用。更具体地说，大型因果语言模型——我们将在本概述中研究的类型——使用一种特定的自注意力变体，称为多头因果自注意力。

使用 PyTorch 训练神经网络。本概述中我们将查看的代码是用 PyTorch 编写的，并且严重依赖于分布式训练技术，例如分布式数据并行（DDP）训练。要了解 PyTorch 和分布式训练的基础知识，请查看以下文章：

在 PyTorch 的基本（以及分布式）神经网络训练之外，我们还将看到自动混合精度（AMP）训练的使用，它会在训练过程中选择性地调整神经网络内的精度——在完整精度（ float32 ）和半精度（ float16 或 bfloat16）之间——以提高效率。简单来说，我们在神经网络内进行大量的矩阵乘法，如果能够以较低精度运行其中一些乘法，训练就会更快。有关更详尽（且实用）的 AMP 概述，请参见此处。

深度学习基础。本概述还需要对神经网络有一个基本的了解，包括它们是如何被训练和使用的。为了获得这些知识，我强烈推荐 fast.ai 的《为程序员准备的深度学习实践课程》，该课程经常更新，并且（在我看来）是任何人都能得到的最佳深度学习实践入门课程。

二、理解下一个Token预测

我们现在将学习关于下一个词预测（也称为标准的语言建模目标）——所有因果语言模型背后的工作核心。在本节中，我们将首先涵盖与分词相关的一些基本概念，然后我们将概述语言模型的预训练和推理过程，以及它们与下一个词预测概念的关系。

2.1 词和词汇表

在尝试理解下一个词预测时，我们可能首先会问：什么是词？简单来说，词就是一个文本序列中的单词或子词。给定一个原始文本序列作为输入，使用语言模型的第一步是将这个原始文本分词，或者将其分解为一系列离散的词；下面是一个示例。

为了执行这项分词操作，我们依赖于一个分词器。分词器是在未标记的文本语料库上进行训练的，以学习一组固定大小、唯一的现有token。这组固定大小的token被称为我们的词汇表，词汇表包含了语言模型所知的所有token。通常，我们应该确保用于训练分词器的数据能够准确反映模型在训练和推理过程中将看到的数据类型。鉴于词汇表的大小是固定的，这确保了我们看到的外部token在大多数情况下都存在于语言模型的词汇表中。

分词技术。存在多种不同的分词技术；此处提供了概述。有关为 LLMs 训练和使用流行分词器的详细信息，请参阅这篇文章，该文详细介绍了字节对编码（BPE）分词器——LLMs 中最常用的分词器。另一种最近变得流行的分词技术是字节级 BPE（BBPE），它依赖于字节（而不是文本字符）作为分词的基本单元。

词嵌入。一旦我们对文本进行了分词，我们就在语言模型参数的一部分中存储的嵌入层中查找每个词的嵌入 2。之后，由我们的输入构建的文本词序列就变成了词嵌入向量序列；见下文。

为了构建实际传递给我们的仅解码器 Transformer 架构的输入，还需要进行最后一步——添加位置编码。位置编码与词嵌入的尺寸相同，并类似处理（即，它们作为语言模型的一部分进行存储，并与其他模型参数一同训练）。然而，我们不是将嵌入与每个独特的词关联，而是将嵌入与词化输入中可能存在的每个独特位置关联；下方有示意图。

我们将这些嵌入添加到相应位置的词嵌入中。这种加性位置嵌入是必要的，因为自注意力操作没有任何方法来表示每个词的位置。通过添加位置嵌入，我们允许 Transformer 中的自注意力层在训练过程中将每个词的位置作为一个相关特征。最近的研究探索了将位置信息注入自注意力的新方法，从而产生了像 RoPE这样的技术。

上下文窗口。语言模型通过特定大小的token序列进行预训练，这个大小被称为上下文窗口的大小或上下文长度。这个大小——通常在 1K 到 8K 个token的范围内（尽管有些模型要大得多！）——通常是根据硬件和内存限制选择的 3。由于我们只学习这种长度的输入的位置嵌入，因此上下文窗口限制了 LLM 可以处理的输入数据量。然而，最近开发了像 ALiBi这样的技术，以实现超出训练期间所见输入的推断。

2.2 语言模型预训练

语言模型在多个步骤中进行训练，如上图所示。第一步（也是最耗计算资源的步骤）是预训练，我们将在本概述中重点介绍预训练。在预训练期间，我们获取一个大型无标签文本语料库，并通过 i) 从数据集中采样一些文本，以及 ii) 训练模型以预测下一个token来训练模型。这是一个自监督目标，因为不需要标签。相反，真实的下一个token已经存在于语料库本身中——监督的来源是隐含的。这种训练目标被称为下一个token预测，或标准的语言建模目标。

预测下一个词。在我们获得词嵌入（包含位置嵌入）后，我们将这些向量输入到一个仅包含解码器的 Transformer 中，该 Transformer 为每个词嵌入生成相应的输出向量；见下文。

对于每个token，给定一个输出向量，我们可以通过 i) 取出该token的输出向量，ii) 使用这个向量来预测序列中下一个token。下方有图示说明。

如上图所示，通过将一个token的输出向量作为输入传递给线性层来预测下一个token，该线性层输出一个与词汇表大小相同的向量。应用 softmax 变换后，形成了一个覆盖token词汇表的概率分布，我们可以在推理时从这个分布中采样下一个token，或者在预训练时训练模型以最大化正确下一个token的概率。

跨序列预测token。在预训练期间，我们不仅预测单个下一个token，而是对序列中的每个token都进行下一个token预测，并聚合它们的损失。由于使用了因果自注意力机制，每个输出token向量只考虑当前token以及序列中当前token之前的token。因此，可以使用单个解码器 Transformer 的前向传递来跨整个序列进行下一个token预测，因为每个token都不知道它之后的token。

2.3 自回归推理过程

现在，我们了解如何预训练语言模型，但在进行推理时，下一个token的预测也同样被使用！下一个token的预测是训练和使用 LLMs 的所有方面的基础。从初始（可能为空）的输入序列或前缀开始，语言模型通过遵循自回归下一个token预测过程（见上文）的以下步骤生成文本：

预测下一个token
将预测的token添加到当前输入序列
重复

选择下一个token。在上一节中，我们已经看到了如何创建一个token的概率分布。但是，我们实际上是如何从这个分布中选择下一个token的呢？通常，我们只是从这个分布中采样下一个token。然而，存在许多采样策略，它们通过修改token的概率分布来对这个方法进行微调。确切的解码方法取决于应用，但我们需要熟悉的主要概念和策略概述如下：

温度(Temperature)
贪婪解码(Greedy Decoding)
核采样(Nucleus Sampling)
Top-K 采样(Top-K Sampling)

三、创建一个最小实现

现在我们已经理解了下一token预测的概念，我们需要将所学到的想法变得更加具体。在本节中，我们将考察一个用 PyTorch 编写的实现——使用 LLM 进行预训练和推理（使用下一token预测）。这个实现源自 Andrei Karpathy 的 NanoGPT，与 GPT-2的规格相匹配。除了 GitHub 上提供的 NanoGPT 实现（见上文链接）之外，还有一个非常棒的配套教程视频。

虽然与大多数现代 LLMs 相比这个模型较小，但它是一个很好的例子，展示了语言模型在代码中的样子。在这里，我们将研究 NanoGPT 的实现，并将其与我们之前讨论的下一个token预测联系起来。

3.1 Decoder-Only Transformer

首先，我们将详细说明我们语言模型架构的实现，该架构基于仅解码器 Transformer。首先，我们将概述该架构的组件，从模型的单个模块到完整的多层架构。然后，我们将研究该模型架构如何在预训练和推理过程中使用下一个token预测。

@dataclass
class GPTConfig:
    block_size: int = 1024     # context window size
    vocab_size: int = 50304    # total number of tokens in the vocabulary
    n_layer: int = 12          # total number of layers
    n_head: int = 12           # number of heads for each attention operation
    n_embd: int = 768          # size of the token embeddings
    dropout: float = 0.0       # whether to use dropout
    bias: bool = True          # whether to use bias in Linear layers and LayerNorm

模型配置。我们首先需要关注的是模型架构的配置；见上文。正如我们所见，配置只是一个 Python 中的数据类，用于指定我们架构的各种超参数。上文中显示的设置对应于 GPT-2 论文中探索的最小模型架构，如下表所示。

该模型仅包含 11.7 亿个参数，实际上与原始 GPT 出版物中使用的基座 Transformer 架构完全相同。

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()

        # layer norm layers
        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)

        # multi-headed, causal self-attention layer
        self.attn = CausalSelfAttention(config)

        # two-layer feed-forward neural network
        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

一个单一模块。接下来，我们可以查看在仅解码器 Transformer 架构中单个模块的实现；见上文。在这里，我们看到一个仅解码器 Transformer 模块有两个组成部分：

多头因果自注意力
前馈神经网络

对于大多数语言模型（包括 NanoGPT），前馈网络是一个两层模型，其中隐藏层比输入层稍宽。在每个层之前，块的输入都会进行归一化处理，并且在层之间添加了残差连接。下方是示意图。

模型定义。现在我们已经了解了decoder-only Transformer 模块的结构，可以查看 NanoGPT 的完整模型定义。此定义如下提供，我们看到模型类的构造函数。

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        # full model definition
        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),
        ))

        # final linear layer to perform next token prediction
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # Weight Tying improves the performance of language models by tying (sharing)
        # the weights of the embedding and softmax layers. This also massively reduces
        # the total number of parameters in the language model
        self.transformer.wte.weight = self.lm_head.weight

        # initialize all weights
        self.apply(self._init_weights)

        # apply special scaled init to the residual projections (see 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))

如上所示，LLM 包含两个不同的嵌入层——一个用于存储 token 嵌入，另一个用于存储位置嵌入。有 1024 个位置嵌入，对应于用于训练 NanoGPT 的上下文长度（即配置中的 block_size 设置）。该语言模型总共有 12 个 Transformer 模块。除了采用 GPT-2 的一些特殊技术外，模型的权重是正态初始化的。

在基本 Transformer 架构之外，LLM 在第一层/最后一层的前向传递中使用了额外的 dropout 和 LayerNorm 模块。此外，我们还有一个用于下一个 token 预测的线性分类头，它与 token 嵌入层共享权重。这种权重共享方法称为权重绑定(weight tying)，可以在大幅减少模型参数总数的同时提高性能。

3.2 实现下一个Token预测

现在我们已经了解了 LLM 模型架构的实现，我们将查看一个具有相同架构的预训练和推理实现。预训练（如上所示）和推理都依赖于下一个token预测策略，在本节中，我们将概述这两个过程中下一个token预测的实现。

前向传播。为了理解如何训练 NanoGPT，我们需要了解模型的正向传播。我们可以考虑两种不同类型的前向传播——一种用于训练，一种用于推理。NanoGPT 的前向传播代码（即，此方法是之前提供的 GPT 模型类的一部分）如下所示。首先，我们将考虑此前向传播在预训练过程中的使用方式，稍后我们再回到推理过程。

def forward(self, idx, targets=None):
    # check whether input is on the CPU or GPU
    device = idx.device

    # get all positional embeddings
    pos = torch.arange(0, t, dtype=torch.long, device=device)
    pos_emb = self.transformer.wpe(pos)

    # get all token embeddings
    tok_emb = self.transformer.wte(idx)

    # add token/positional embeddings and apply dropout
    x = self.transformer.drop(tok_emb + pos_emb)

    # pass token vector sequence through all transformer blocks
    for block in self.transformer.h:
        x = block(x)

    # apply final layer normalization operation
    x = self.transformer.ln_f(x)

    if targets is not None:
        # if we are given target next tokens, then compute the loss over the entire input sequence
        logits = self.lm_head(x)
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
    else:
        # otherwise just compute the next token for the final token in the input sequence
        logits = self.lm_head(x[:, [-1], :])
        loss = None

    return logits, loss

前向传播如我们所预期的那样运作。我们以两个张量作为输入：

输入tensor (idx):一个矩阵，其中每一行包含一个 token id 序列，表示用于预训练（或推理）的文本序列。
目标tensor (targets): 与输入张量类似，但每个条目包含输入张量中每个 token 的真实下一个 token id。

这些张量中的每一个存储一个完整的 mini-batch，其中包含多个文本序列，训练迭代在这些序列上并行化。在这里，我们将假设目标张量不是 None 。在预训练期间这总是成立的，而在推理期间我们没有目标，只是自由地生成下一个 token。

# get all positional embeddings
pos = torch.arange(0, t, dtype=torch.long, device=device)
pos_emb = self.transformer.wpe(pos)

# get all token embeddings
tok_emb = self.transformer.wte(idx)

# add token/positional embeddings and apply dropout
x = self.transformer.drop(tok_emb + pos_emb)

前向传播的第一步是构建一个与我们位置和词嵌入相对应的矩阵；见上文。 idx 张量包含可以直接用于在词嵌入矩阵中查找的词 id。我们必须手动构建索引值来查找位置嵌入。位置嵌入和词嵌入相加，通过一个 dropout 层，然后通过所有 transformer 块。然后，在用下一个词预测目标计算损失之前，执行一个最终的 LayerNorm 操作。

if targets is not None:
    # if we are given target next tokens, then compute the loss over the entire input sequence
    logits = self.lm_head(x)
    loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

下一个 token 预测过程会输出一个潜在下一个 token 的分布——使用线性 lm_head 模块，其中 transformer 的每个 token 的输出向量被用作输入，针对输入序列中的每个 token。然后，我们对这个结果应用交叉熵损失(CrossEntropy)，从而训练模型正确预测整个输入序列中每个位置的下一个 token。

@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
    for _ in range(max_new_tokens):
        # crop the sequence so that it doesn't exceed the context window
        idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
        # perform a forward pass with
        logits, _ = self(idx_cond)
        # scale the outputted logits by the 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 a probability distribution over tokens
        probs = F.softmax(logits, dim=-1)
        # sample the next token from the outputted distribution
        idx_next = torch.multinomial(probs, num_samples=1)
        # append the sampled next token to the input and continue!
        idx = torch.cat((idx, idx_next), dim=1)
    return idx

进行推理。在预训练之后，我们可以通过下一个token预测来生成文本。如前所述，使用语言模型生成文本是一个自回归过程，该过程会迭代地预测每个下一个token。为了预测一个token，NanoGPT 遵循以下步骤：

对当前输入序列进行前向传递
根据指定的温度对输出的 logits 进行缩放
[可选] 仅保留最可能的 k 个 token（即 Top-K 采样）
应用 softmax 函数
从结果分布中采样下一个token

3.3 NanoGPT 训练

尽管分布式训练是一个复杂的话题，我们无法在本概述中全面涵盖，但为了完整性，我们将介绍 NanoGPT 预训练过程的实际要点。我们通常将 LLM 训练分布到多个计算设备上（例如 GPU 或 TPU）。从高层次来看，分布式训练是可取的或必要的几个原因如下：

预训练计算成本高昂，我们希望加快速度。
模型的大小可能太大，无法存储在单个设备上。

上述第二种情况特别适用于当前一代的语言模型，这些模型非常大，通常无法存储在单个设备上。存在多种分布式训练技术可以处理这些情况并加快训练过程；总结如下。

分布式训练设置。完整的预训练实现提供在 NanoGPT 仓库中的 train.py 文件内。模型使用单个 GPU 或分布式数据并行（DDP）方法进行训练。此训练框架的设置如下所示。

# RANK = -1 means we train on a single GPU
ddp = int(os.environ.get('RANK', -1)) != -1

# setup for DDP
if ddp:
    # wait for all processes to start so they can communicate
    init_process_group(backend=backend)

    # rank, local rank, and total number of processes
    ddp_rank = int(os.environ['RANK'])
    ddp_local_rank = int(os.environ['LOCAL_RANK'])
    ddp_world_size = int(os.environ['WORLD_SIZE'])

    # identify GPU to be used by this process and whether this process is master
    device = f'cuda:{ddp_local_rank}'
    torch.cuda.set_device(device)
    master_process = ddp_rank == 0

    # each process gets a different seed to ensure data is shuffled differently
    seed_offset = ddp_rank

    # world_size number of processes will be training simultaneously, so we can scale
    # down the desired gradient accumulation iterations per process proportionally
    assert gradient_accumulation_steps % ddp_world_size == 0
    gradient_accumulation_steps //= ddp_world_size
else:
    # if not ddp, we are running on a single gpu, and one process
    master_process = True
    seed_offset = 0
    ddp_world_size = 1

# total number of tokens that will be processed for every training iteration
tokens_per_iter = gradient_accumulation_steps * ddp_world_size * batch_size * block_size

如我们所见，使用 DDP 进行训练需要我们同时运行多个训练进程，这些进程将相互通信。进程的数量等于我们用于训练的总 GPU 数量（无论是在同一台机器上还是在多个节点上）。通过 DDP，我们可以将这些 GPU 上的训练过程并行化。为了协调运行的多个进程，我们必须为每个进程指定一个 rank。例如，如果有四个进程在四个 GPU 上运行训练，这些进程将在[0, 3]范围内各自拥有一个唯一的 rank。在上述代码中，所有 rank 信息都存储在一个环境变量中，该变量可以被进程访问。

梯度累积。在 NanoGPT 实现中，你可能会几次看到梯度累积这个术语。通常，我们通过以下方式训练神经网络：

计算一小批数据的损失
将这个损失反向传播以导出梯度
根据这个梯度更新模型的权重

梯度累积移除了上述的最后一步。相反，梯度会在多个模拟单个较大小批量的“微批量”数据中进行累积（即通过取平均值的方式）。一旦我们累积了足够多数据的梯度，就会更新权重。当所需的小批量大小对于所使用的硬件来说过大时，这种过程非常有用。我们可以简单地计算几个较小批量的梯度，并使用梯度累积来模拟较大的批量。更多详情请参见此处。

如果我们有一个更大的模型呢？使用 DDP，模型的副本会被发送到每个设备，我们通过 i) 在每个设备上对随机采样的数据进行梯度计算，以及 ii) 在每个设备上同步梯度后获取聚合模型更新，来并行训练这些模型的副本。对于许多现代 LLMs，我们可能无法将整个模型存储在单个设备的内存中，因此我们需要不同的训练方法。与 DDP 兼容且最流行的分布式训练算法之一是全分片数据并行（FSDP）训练。与 DDP 相比，这种方法更常用于训练现代 LLMs。

data_dir = os.path.join('data', dataset)
train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
val_data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')

def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([torch.from_numpy((data[i:i+block_size])).astype(np.int64) for i in ix])
    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size])).astype(np.int64) for i in ix])
    if device_type == 'cuda':
        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
    else:
        x, y = x.to(device), y.to(device)
    return x, y

加载数据。我们可以用多种方法创建一个用于训练语言模型的数据加载器。上面代码中展示了一个（简化）的例子。在这里，数据存储在一个文件中，我们为训练和验证数据分别有单独的文件。在训练期间，我们通过简单地获取大小为上下文窗口的随机数据块来加载数据。我们可以选择将数据放到 GPU 上，但整个过程足够简单！

# learning rate decay scheduler (cosine with warmup)
def get_lr(it):
    # 1) linear warmup for warmup_iters steps
    if it < warmup_iters:
        return learning_rate * it / warmup_iters
    # 2) if it > lr_decay_iters, return min learning rate
    if it > lr_decay_iters:
        return min_lr
    # 3) in between, use cosine decay down to min learning rate
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    assert 0 <= decay_ratio <= 1
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)

学习率。在预训练语言模型时，我们需要考虑的一个主要超参数是学习率。通常，我们在预训练期间会采用一个学习率计划。上面展示了一个典型的语言模型预训练学习率计划的示例实现。在这里，该计划有一个较短的（线性）预热期，随后是一个（余弦）衰减期，该衰减期持续指定的迭代次数；见下文。

训练循环。现在我们已经完成了所有必要的设置，我们可以最终实现我们语言模型的实际（预）训练循环；见下文。

while True:

    # determine and set the learning rate for this iteration
    lr = get_lr(iter_num) if decay_lr else learning_rate
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # forward backward update, with optional gradient accumulation to simulate larger batch size
    # and using the GradScaler if data type is float16
    for micro_step in range(gradient_accumulation_steps):
        if ddp:
            # DDP only needs to sync gradients across devices at the last micro step
            model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
        with ctx:
            # perform forward pass and compute next token prediction loss
            logits, loss = model(X, Y)
            # scale the loss to account for gradient accumulation
            loss = loss / gradient_accumulation_steps
        # immediately async prefetch next batch while model is doing the forward pass on the GPU
        X, Y = get_batch('train')
        # perform backpropagation (with gradient scaling if training in fp16)
        scaler.scale(loss).backward()

    # clip the gradient
    if grad_clip != 0.0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    # update the model weights using a PyTorch optimizer
    # use scaler if training in fp16
    scaler.step(optimizer)
    scaler.update()

    # flush the gradients as soon as we can, no need for this memory anymore
    optimizer.zero_grad(set_to_none=True)

    # termination conditions
    if iter_num > max_iters:
        break

在上述实现中可能存在一些不熟悉的组件（例如梯度裁剪和损失缩放）。这些更改大多与自动混合精度（AMP）训练相关，而自动混合精度训练是 NanoGPT 支持（但非强制）的组件。除了这些新增的细节外，上述代码与我们对预训练过程的先前讨论相符，并使用了标准的 PyTorch 语法。

3.4 结语

阅读关于 LLMs 的论文既有趣又有益，但我们仅靠阅读是有限的。如果我们想构建任何具体的东西，最终必须实现这些想法。在这篇概述中，我们首先学习了下一token预测的概念及其在因果语言模型中的应用。然后，我们探索了在 PyTorch 中用 LLM 进行预训练和推理的下一token预测的具体实现。虽然与当前研究中探索的一些庞大语言模型相比，这个实现很简单，但它为我们提供了一个实用的基础，让我们对 LLMs 有了更具体的理解。

Reference

LLM

#LLM

语言模型训练与推理：从概念到代码

https://mztchaoqun.com.cn/posts/D105_LM_Concept_to_Code/

作者

mztchaoqun

发布于

2026年1月16日

许可协议

Hyper-Connections 上一篇

SAPO：Soft Adaptive Policy Optimization 下一篇