从零开始用Python构建一个推理LLM(类似 o3 和 DeepSeek-R1)：详细的端到端指南

你可能已经看过很多关于用 Transformer 架构训练大模型的文章了。但现在，像 OpenAI 的 o3、谷歌的 Gemini 2.5 Pro，以及国内开源表现亮眼的 DeepSeek-R1 等新一代模型，已经不仅仅是在“给答案”，而是能像人一样，先思考、再回答。那么问题来了：我们能不能自己动手打造一个“会思考”的推理型 LLM？

从高层次来看，构建一个最简单的推理型 LLM 只需三步：

推理型 LLM 的训练流程，作者 Fareed Khan

1. 使用 Transformer 架构对 LLM 进行预训练；
2. 在预训练模型基础上进行监督微调（SFT）；
3. 在 SFT 模型基础上进行基于人类反馈的强化学习（RLHF），赋予模型推理能力。

本文中，我将逐步实现上述三个阶段，最终构建一个具备推理能力的 LLM，并检验其输出效果。

1、最终输出效果

以下是我们模型的输出对比（使用小样本数据 vs 使用完整训练数据）：

三个阶段的模型输出效果，作者 Fareed Khan

可以看到，使用更大规模数据集训练的模型表现明显更优，而使用小样本数据训练的模型效果较差，但对于学习和快速调试代码而言仍然十分有用。

在开始之前，请确保你对面向对象编程（OOP）和神经网络（NN）有基本了解，熟悉 PyTorch 也会对编程实现有很大帮助。

2、环境准备

本项目中我们将使用多个 Python 库，首先导入它们：

# 标准库
import os
import json
import math
import time
import random
import warnings
from typing import Optional, Tuple, List, Union, Iterator

# 数值计算库
import numpy as np

# PyTorch深度学习库
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from torch.utils.data import Dataset, DataLoader

# 用于自动混合精度和上下文管理
from contextlib import nullcontext

# 引入Hugging Face库，用于分词和模型处理
from transformers import AutoTokenizer, PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import (
    CausalLMOutputWithPast,  # 用于因果语言建模的输出类型
    BaseModelOutputWithPast  # 带缓存键值的基础模型输出类型
)
from transformers.activations import ACT2FN  # Transformer模型中使用的激活函数

# 引入Hugging Face的`tokenizers`库（高性能分词器库）
from tokenizers import Tokenizer as HFTokenizer  # 重命名以避免与transformers库的分词器冲突
from tokenizers import models as hf_models
from tokenizers import trainers as hf_trainers
from tokenizers import pre_tokenizers as hf_pre_tokenizers
from tokenizers import decoders as hf_decoders

导入了上述后续训练和预处理所需的模块后，我们接下来就可以开始训练分词器了。

3、三阶段的数据集

我们将进行三个阶段的训练，每个阶段的数据格式各不相同。

第一个数据集用于预训练阶段，即从零开始训练一个基于 Transformer 架构的大型语言模型，使其能够生成语法正确的文本。

下面定义一个示例预训练数据集：

# --- 预训练数据 ---
sample_pretrain_data = [
    {"text": "晴朗的蓝天下，阳光明媚。"},
    {"text": "猫咪喜欢追逐老鼠，也喜欢玩毛线球。"},
    {"text": "阅读书籍可以扩展你的知识和词汇量。"},
    {"text": "人工智能是一个快速发展的研究领域。"},
    {"text": "烤蛋糕需要面粉、糖、鸡蛋和黄油。"},
    {"text": "大型语言模型是在海量文本数据上训练出来的。"},
    {"text": "敏捷的棕色狐狸跳过了懒惰的狗。"}
]

在上述数据上训练好 Transformer 架构后，我们再使用监督微调（SFT）数据集，在预训练模型基础上进一步训练。下面定义 SFT 数据集：

# --- 监督微调（SFT）数据 ---
sample_sft_data = [
    {"conversations": [
        {"role": "user", "content": "你好，你怎么样？"},
        {"role": "assistant", "content": "我很好，谢谢！今天有什么可以帮你的吗？"}
    ]},
    {"conversations": [
        {"role": "user", "content": "法国的首都是哪里？"},
        {"role": "assistant", "content": "法国的首都是巴黎。"}
    ]},
    {"conversations": [
        {"role": "user", "content": "用简单的话解释一下什么是重力。"},
        {"role": "assistant", "content": "重力是一种将物体相互拉近的力，这就是为什么东西会掉到地上的原因！"}
    ]}
]

监督微调（SFT）通过提供问题及其正确答案，进一步提升大型语言模型的知识水平。

最后，在完成预训练和监督微调后，我们还会构建一个基于人类反馈的强化学习（RLHF）数据集，以赋予模型推理和决策能力。

# --- 推理数据示例 ---
sample_reasoning_data = [
    {"conversations": [
        {"role": "user", "content": "如果我有3个苹果，吃掉1个，还剩几个？"},
        {"role": "assistant", "content": "<think>用户一开始有3个苹果，吃掉了1个苹果，即从原有数量中减去1个。因此，3 - 1 = 2。</think><answer>你还剩下2个苹果。</answer>"}
    ]},
    {"conversations": [
        {"role": "user", "content": "什么是三原色？"},
        {"role": "assistant", "content": "<think>用户询问的是三原色，这些颜色无法通过其他颜色混合而成。在加色模型（如光学）中，标准的三原色为红、绿、蓝（RGB）；在减色模型（如颜料）中，通常认为是红、黄、蓝（RYB）或青、品红、黄（CMY）。</think><answer>通常所说的三原色是红色、黄色和蓝色。这些颜色可以混合产生其他各种颜色，但自身无法通过其他颜色混合而成。</answer>"}
    ]}
]

从上面的强化学习人类反馈（RLHF）数据集中可以看到，当用户提出问题时，回答中包含了``标签，这些标签用于引导大模型（LLM）在给出最终答案前先生成推理过程。

现在我们已经定义好了各个阶段的训练数据集，接下来就可以开始训练我们的分词器了。

4、训练 BPE 分词器

目前可用的分词器种类繁多，有些是预先在大型语料库上训练好的，比如tiktoken库中的分词器；也可以根据我们自己的数据集定制训练。一般来说，训练所用的英文语料库越大，分词器的词汇表规模也越大。

我们先用一个示例数据集来训练分词器，以便更清楚地了解代码的运行过程：

# 分词器训练数据
tokenizer_corpus = [
    "Hello world, this is a demonstration of building a thinking LLM.",
    "Language models learn from text data.",
    "Tokenization is a crucial first step.",
    "We will train a BPE tokenizer.",
    "Think before you answer.",
    "The answer is forty-two.",
    "<think>Let's consider the options.</think><answer>Option A seems best.</answer>",
    "<|im_start|>user\nWhat's up?<|im_end|>\n<|im_start|>assistant\nNot much!<|im_end|>"
]

# 将数据保存到临时文件中供分词器训练使用
tokenizer_corpus_file = os.path.join("YOUR_DIR_PATH", "tokenizer_corpus.txt")
with open(tokenizer_corpus_file, 'w', encoding='utf-8') as f:
    for line in tokenizer_corpus:
        f.write(line + "\n")

这里我只是定义了一些随机文本数据，你也可以使用莎士比亚数据集来进行训练。

下面我们创建一个简单的 BPE 分词器训练函数，用于在我们自己的语料库上训练分词器：

def train_demo_tokenizer(
    corpus_files: List[str],
    vocab_size: int,
    save_path: str,
    special_tokens: List[str]
):
    """
    使用 Hugging Face Tokenizers 库训练一个字节对编码（BPE）分词器。

    参数:
        corpus_files (List[str]): 包含训练语料的文件路径列表。
        vocab_size (int): 期望的词汇表大小。
        save_path (str): 训练好的分词器保存路径。
        special_tokens (List[str]): 要包含在分词器中的特殊标记列表。
    """
    # 初始化BPE分词器，并设置未知标记
    tokenizer_bpe = HFTokenizer(hf_models.BPE(unk_token="<unk>"))

    # 设置预分词器为字节级别（ByteLevel），以实现稳健的字节级处理
    tokenizer_bpe.pre_tokenizer = hf_pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True)

    # 设置解码器为字节级别（ByteLevel），确保字节级标记正确解码
    tokenizer_bpe.decoder = hf_decoders.ByteLevel()

# 创建BPE分词器训练器，指定词汇量、特殊标记和完整字节字母表
trainer = hf_trainers.BpeTrainer(
    vocab_size=vocab_size,
    special_tokens=special_tokens,
    show_progress=True,
    initial_alphabet=hf_pre_tokenizers.ByteLevel.alphabet()
)

# 确保语料文件列表格式统一
if isinstance(corpus_files, str):
    corpus_files = [corpus_files]

# 开始训练分词器
tokenizer_bpe.train(corpus_files, trainer=trainer)
print(f"分词器训练完成，词汇量为：{tokenizer_bpe.get_vocab_size()}")

# 将训练好的分词器保存为JSON文件
tokenizer_bpe.save(save_path)
print(f"分词器已保存至：{save_path}")

return tokenizer_bpe

我们定义了特殊标记``来处理未知词汇。此外，我们还需要定义一些特殊标记，用于表示文本的起始、结束以及填充（padding）。我们将这些标记放入一个列表中：

# 特殊标记列表
SPECIAL_TOKENS_LIST = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<pad>"]

现在，我们调用上述函数来训练分词器：

# 使用指定的语料文件、词汇量、保存路径和特殊标记训练并保存BPE分词器
trained_hf_tokenizer = train_demo_tokenizer(
    corpus_files=[tokenizer_corpus_file],
    vocab_size=1000,
    save_path="NOTEBOOK_TOKENIZER_PATH",
    special_tokens=SPECIAL_TOKENS_LIST
)

分词器训练完成后，我们用一个示例句子进行测试：

# 使用训练好的分词器进行编码和解码测试
test_sentence = "Hello <|im_start|> world <think>思考中</think><answer>答案</answer> <|im_end|>"

# 对句子进行编码
encoded = trained_hf_tokenizer.encode(test_sentence)

# 显示原始句子及其编码结果
print(f"原始句子：{test_sentence}")
print(f"编码后的ID序列：{encoded.ids}")
print(f"编码后的标记序列：{encoded.tokens}")

# 将编码后的ID序列解码回文本
decoded = trained_hf_tokenizer.decode(encoded.ids)
print(f"解码后的句子：{decoded}")

### 输出示例 ###
训练后的分词器词汇表（前10个及特殊标记）：
'ain': 328
'}': 96
'D': 39
'Ġfirst': 428
'Wh': 324
  ...

通过上述测试，我们可以直观看到训练好的 BPE 分词器如何将示例句子拆分成标记。接下来，我们将创建 Transformer 模型，用于第一阶段的训练。

5、Transformer 架构概览

我们先简单了解一下 Transformer 架构如何处理和理解文本。Transformer 的核心思想是将文本拆分成更小的单元（称为标记），然后预测序列中的下一个标记。一个 Transformer 模型通常由多个 Transformer 模块（block）堆叠而成，最后再通过一个输出层进行预测。

每个 Transformer 模块包含两个核心组件：

• 自注意力头（Self-Attention Heads）：负责确定输入文本中哪些部分最值得关注。例如，在处理一句话时，自注意力头可以识别词语之间的关系，比如代词与其所指代的名词之间的联系。
• 多层感知机（MLP）：这是一个简单的前馈神经网络，用于进一步处理自注意力头强调的信息。MLP 通常包含一个输入层（接收自注意力头的输出）、一个或多个隐藏层（增加处理复杂性）和一个输出层（将结果传递给下一个 Transformer 模块）。

换句话说，自注意力头决定了模型**“关注什么”，而MLP则决定了模型“如何处理这些关注的内容”**。通过堆叠多个Transformer模块，模型能够捕捉文本中更复杂的模式和关系，当然，这种能力并非绝对保证。

我们先来看一个更简单、更直观的架构图，以便后续实现代码时心中有数。

下面逐步梳理一下我们即将编写代码实现的架构流程：

1. 输入的词元（tokens）首先转换为词嵌入（embeddings），并加入位置信息。
2. 模型由 64 个结构完全相同的 Transformer 模块组成，依次对数据进行处理。
3. 每个 Transformer 模块首先通过多头注意力机制，捕捉词元之间的关联关系。
4. 接着，每个模块通过一个多层感知机（MLP）进一步处理数据，先扩展再压缩数据维度。
5. 每个步骤都使用残差连接（快捷连接），帮助信息更顺畅地流动。
6. 整个模型中广泛使用层归一化（Layer Normalization），以稳定训练过程。
7. 注意力机制负责计算每个词元之间的关注程度，决定哪些词元之间应当相互关注。
8. MLP将数据维度扩展至原来的4倍，经过ReLU激活函数后再压缩回原始维度。
9. 模型使用16个注意力头，以捕捉不同类型的词元关系。
10. 最终一层将处理后的数据映射为词汇表大小的预测结果。
11. 模型通过反复预测下一个最可能的词元，逐步生成文本。

6、多层感知机（MLP）

多层感知机（MLP）是 Transformer 结构中前馈网络的核心组件，负责引入非线性特性，学习嵌入表示中的复杂关系。在定义 MLP 模块时，一个关键参数是 n_embed，它定义了输入嵌入的维度。

MLP 通常包含一个隐藏的线性层，将输入维度扩展一定倍数（通常为 4 倍，我们这里也采用 4 倍），随后使用非线性激活函数（常用 ReLU）。这种结构使得网络能够学习更复杂的特征。最后，再通过一个投影线性层，将扩展后的表示压缩回原始嵌入维度。这一系列变换使 MLP 能够进一步精炼注意力机制学习到的表示。

MLP 图示，作者 Fareed Khan

# --- 多层感知机（MLP）类定义 ---

class MLP(nn.Module):
    """
    一个简单的多层感知机（MLP）模块，包含一个隐藏层。
    该模块用于Transformer模块中的前馈处理过程，
    先将输入嵌入维度扩展，经过ReLU激活后，再投影回原始维度。
    """
    def __init__(self, n_embed):
        super().__init__()
        self.hidden = nn.Linear(n_embed, 4 * n_embed)  # 线性层扩展嵌入维度
        self.relu = nn.ReLU()                          # ReLU激活函数
        self.proj = nn.Linear(4 * n_embed, n_embed)    # 线性层投影回原始维度

    def forward(self, x):
        """
        MLP的前向传播过程。
        参数:
            x (torch.Tensor): 输入张量，形状为(B, T, C)，
                              其中B为批次大小，T为序列长度，C为嵌入维度。
        返回:
            torch.Tensor: 输出张量，形状与输入相同。
        """
        x = self.forward_embedding(x)
        x = self.project_embedding(x)
        return x

    def forward_embedding(self, x):
        """
        经过隐藏线性层和ReLU激活函数。
        参数:
            x (torch.Tensor): 输入张量。
        返回:
            torch.Tensor: 经过隐藏层和ReLU后的输出。
        """
        x = self.relu(self.hidden(x))
        return x

    def project_embedding(self, x):
        """
        经过投影线性层。
        参数:
            x (torch.Tensor): 输入张量。
        返回:
            torch.Tensor: 经过投影层后的输出。
        """
        x = self.proj(x)
        return x

我们刚刚完成了多层感知器（MLP）部分的代码实现，其中__init__方法定义了一个隐藏层（线性层），用于将输入嵌入维度（n_embed）扩展到更高维度；同时定义了一个投影层，将扩展后的维度再压缩回原始维度。在隐藏层之后，我们使用了 ReLU 激活函数。

forward方法则定义了数据如何流经这些层次：输入数据首先经过隐藏层和 ReLU 激活（通过forward_embedding），再通过投影层（project_embedding）得到最终输出。

7、单头注意力机制（Single Head Attention）

注意力头（Attention Head）是我们模型的核心组件之一，其作用是帮助模型聚焦于输入序列中最相关的部分。

定义一个注意力头模块时，有几个关键参数需要关注：head_size（头部维度）、n_embed（输入嵌入维度）和context_length（上下文长度）。其中，head_size决定了键（key）、查询（query）和数值（value）投影后的维度大小，进而影响注意力机制的表达能力。

输入嵌入维度n_embed则决定了输入数据进入这些投影层时的维度大小。context_length用于生成因果掩码（causal mask），确保模型在计算注意力时只关注当前及之前的词元，而不会关注未来的词元。

在注意力头内部，我们使用无偏置的线性层（nn.Linear）分别初始化键、查询和值的投影。此外，我们还注册了一个大小为context_length × context_length的下三角矩阵（tril）作为缓冲区，用于实现因果掩码，防止注意力机制关注未来的词元。

单头注意力机制示意图，作者 Fareed Khan

# --- 注意力头类定义 ---

class Head(nn.Module):
    """
    单个注意力头模块。
    该模块负责计算注意力权重，并将其应用到数值（value）上。
    包含键（key）、查询（query）和值（value）的投影层，
    并使用因果掩码防止关注未来的词元。
    """
    def __init__(self, head_size, n_embed, context_length):
        super().__init__()
        self.key = nn.Linear(n_embed, head_size, bias=False)   # 键投影层
        self.query = nn.Linear(n_embed, head_size, bias=False) # 查询投影层
        self.value = nn.Linear(n_embed, head_size, bias=False) # 值投影层
        # 下三角矩阵用于因果掩码
        self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length)))

    def forward(self, x):
        """
        注意力头的前向传播过程。
        参数:
            x (torch.Tensor): 输入张量，形状为(B, T, C)。
        返回:
            torch.Tensor: 经过注意力机制处理后的输出张量。
        """
        B, T, C = x.shape
        k = self.key(x)     # (B, T, head_size)
        q = self.query(x)   # (B, T, head_size)
        scale_factor = 1 / math.sqrt(C)
        # 计算注意力权重：(B, T, head_size) @ (B, head_size, T) -> (B, T, T)
        attn_weights = q @ k.transpose(-2, -1) * scale_factor
        # 应用因果掩码，防止关注未来词元
        attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        attn_weights = F.softmax(attn_weights, dim=-1)
        v = self.value(x)   # (B, T, head_size)
        # 将注意力权重应用到数值上
        out = attn_weights @ v # (B, T, T) @ (B, T, head_size) -> (B, T, head_size)
        return out

我们定义的注意力头类（Head）的__init__方法中，分别初始化了键、查询和值的线性投影层，每个投影层都将输入嵌入维度（n_embed）映射到头部维度（head_size）。

此外，我们还基于context_length创建了一个下三角矩阵，用于实现因果掩码。

在forward方法中，我们首先通过查询和键的点积计算注意力权重，并进行缩放；然后应用因果掩码，确保模型不会关注未来的词元；接着使用 softmax 函数对权重进行归一化；最后将归一化后的注意力权重与数值相乘，得到最终的注意力输出。

8、多头注意力机制（Multi Head Attention）

为了捕捉输入序列中丰富多样的关系，我们将引入“多头注意力”这一概念。具体实现上，我们定义了一个MultiHeadAttention模块，它负责并行管理多个独立的注意力头（Attention Head）。

这里最关键的参数是n_head，它决定了并行运行的注意力头的数量。此外，输入嵌入维度（n_embed）和上下文长度（context_length）也是实例化每个注意力头所必需的参数。

每个注意力头独立处理输入数据，并将其投影到一个维度较低的子空间中，子空间的维度大小为n_embed // n_head。通过多个注意力头并行工作，模型能够同时关注输入序列的不同方面。

多头注意力示意图，作者 Fareed Khan

# --- 多头注意力模块定义 ---

class MultiHeadAttention(nn.Module):
    """
    多头注意力模块。
    该模块并行组合多个注意力头，并将各个头的输出拼接成最终输出。
    """
    def __init__(self, n_head, n_embed, context_length):
        super().__init__()
        self.heads = nn.ModuleList([
            Head(n_embed // n_head, n_embed, context_length) for _ in range(n_head)
        ])

    def forward(self, x):
        """
        多头注意力的前向传播过程。
        参数:
            x (torch.Tensor): 输入张量，形状为(B, T, C)。
        返回:
            torch.Tensor: 拼接所有注意力头输出后的张量。
        """
        # 沿最后一个维度（C维度）拼接各个头的输出
        x = torch.cat([h(x) for h in self.heads], dim=-1)
        return x

至此，我们定义了MultiHeadAttention类，它将多个注意力头组合起来。其中，__init__方法初始化了一个包含n_head个Head实例的列表，每个头的维度大小为n_embed // n_head。

forward方法则将输入x分别送入每个注意力头，并沿最后一个维度拼接各个头的输出，从而融合各个头所学习到的信息。

9、Transformer 模块

若要构建拥有数十亿参数规模的模型，我们必然需要更深的网络结构。因此，我们需要定义一个 Transformer 模块（Transformer Block），并将多个模块堆叠起来。

Transformer 模块的关键参数包括n_head、n_embed和context_length。每个模块由一个多头注意力层和一个前馈神经网络（MLP）组成，并在每个子层之前应用层归一化（Layer Normalization），在每个子层之后添加残差连接。

层归一化的参数由嵌入维度n_embed决定，有助于稳定训练过程。前面介绍过的多头注意力机制也使用了n_head、n_embed和context_length这几个参数。

模块中的前馈神经网络同样使用嵌入维度n_embed。这些组件协同工作，共同处理输入数据并学习复杂的模式。

Transformer模块示意图，作者 Fareed Khan

# --- Transformer模块定义 ---

```python
class Block(nn.Module):
    """
    单个Transformer模块。
    此模块由一个多头注意力层和一个多层感知机（MLP）组成，
    并配有层归一化和残差连接。
    """
    def __init__(self, n_head, n_embed, context_length):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embed)
        self.attn = MultiHeadAttention(n_head, n_embed, context_length)
        self.ln2 = nn.LayerNorm(n_embed)
        self.mlp = MLP(n_embed)

    def forward(self, x):
        """
        Transformer模块的前向传播。
        参数:
            x (torch.Tensor): 输入张量。
        返回:
            torch.Tensor: 经过Transformer模块处理后的输出张量。
        """
        # 多头注意力层，带残差连接
        x = x + self.attn(self.ln1(x))
        # MLP层，带残差连接
        x = x + self.mlp(self.ln2(x))
        return x

    def forward_embedding(self, x):
        """
        专注于嵌入和注意力部分的前向传播。
        参数:
            x (torch.Tensor): 输入张量。
        返回:
            tuple: 包含MLP嵌入后的输出和残差连接的中间结果。
        """
        res = x + self.attn(self.ln1(x))
        x = self.mlp.forward_embedding(self.ln2(res))
        return x, res

上述Block类实现了单个 Transformer 模块：

• 在__init__方法中，我们初始化了两个层归一化模块（ln1、ln2），一个多头注意力模块（MultiHeadAttention）和一个多层感知机模块（MLP）。这些模块的参数由n_head（注意力头数）、n_embed（嵌入维度）和context_length（上下文长度）决定。
• forward方法定义了模块的前向传播过程：输入首先经过层归一化，再通过多头注意力层，并与原输入进行残差连接；随后再次经过层归一化，进入 MLP 层，再次进行残差连接。
• forward_embedding方法则提供了一种替代的前向传播方式，专注于注意力和 MLP 嵌入阶段，返回 MLP 嵌入后的输出以及注意力层后的残差连接结果。

10、最终模型结构

到目前为止，我们实现了 Transformer 模型的若干小型组件。接下来，我们将整合词元嵌入（token embedding）、位置嵌入（position embedding）和多个 Transformer 模块，以实现序列到序列的任务。为此，我们需要定义几个关键参数：n_head（注意力头数）、n_embed（嵌入维度）、context_length（上下文长度）、vocab_size（词汇表大小）和N_BLOCKS（Transformer 模块数量）。

• vocab_size决定了词元嵌入层的规模，每个词元都会映射到一个维度为n_embed的稠密向量。
• context_length用于位置嵌入层，表示输入序列中每个词元的位置，同样具有维度n_embed。
• 注意力头数n_head和模块数量N_BLOCKS则决定了 Transformer 网络的深度和复杂程度。

这些参数共同定义了 Transformer 模型的架构与容量，下面我们就来实现它。

Transformer 类图示，作者 Fareed Khan

# --- Transformer模型类 ---

class Transformer(nn.Module):
    """
    Transformer模型的主类。
    该类将词元嵌入与位置嵌入相结合，随后通过一系列Transformer模块，最后经由线性层实现语言建模。
    """
    def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS):
        super().__init__()
        self.context_length = context_length
        self.N_BLOCKS = N_BLOCKS
        self.token_embed = nn.Embedding(vocab_size, n_embed)  # 词元嵌入层
        self.position_embed = nn.Embedding(context_length, n_embed)  # 位置嵌入层
        self.attn_blocks = nn.ModuleList([Block(n_head, n_embed, context_length) for _ in range(N_BLOCKS)])  # Transformer模块序列
        self.layer_norm = nn.LayerNorm(n_embed)  # 最终的层归一化
        self.lm_head = nn.Linear(n_embed, vocab_size)  # 语言建模的线性输出层
        self.register_buffer('pos_idxs', torch.arange(context_length))  # 预先计算的位置索引

    def _pre_attn_pass(self, idx):
        """
        将词元嵌入与位置嵌入相加。
        参数:
            idx (torch.Tensor): 输入的词元索引。
        返回:
            torch.Tensor: 词元嵌入与位置嵌入之和。
        """
        B, T = idx.shape
        tok_embedding = self.token_embed(idx)
        pos_embedding = self.position_embed(self.pos_idxs[:T])
        return tok_embedding + pos_embedding

    def forward(self, idx, targets=None):
        """
        Transformer模型的前向传播过程。
        参数:
            idx (torch.Tensor): 输入的词元索引序列。
            targets (torch.Tensor, 可选): 用于计算损失的目标词元索引序列，默认为None。
        返回:
            tuple: 模型输出的logits，以及若提供targets则返回损失值。
        """
        x = self._pre_attn_pass(idx)
        for block in self.attn_blocks:
            x = block(x)
        x = self.layer_norm(x)
        logits = self.lm_head(x)

        loss = None
        if targets isnotNone:
            B, T, C = logits.shape
            flat_logits = logits.view(B * T, C)
            targets = targets.view(B * T).long()
            loss = F.cross_entropy(flat_logits, targets)
        return logits, loss

    def forward_embedding(self, idx):
        """
        仅经过嵌入层和注意力模块的前向传播（不经过最终的线性层）。
        参数:
            idx (torch.Tensor): 输入的词元索引序列。
        返回:
            tuple: 注意力模块输出的结果及其残差。
        """
        x = self._pre_attn_pass(idx)
        residual = x
        for block in self.attn_blocks:
            x, residual = block.forward_embedding(x)
        return x, residual

    def generate(self, idx, max_new_tokens):
        """
        根据初始序列生成新的词元。
        参数:
            idx (torch.Tensor): 初始的词元索引序列。
            max_new_tokens (int): 需要生成的词元数量。
        返回:
            torch.Tensor: 扩展后的词元序列。
        """
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.context_length:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

我们的Transformer类的__init__方法初始化了词元嵌入层（token_embed）、位置嵌入层（position_embed）、一系列 Transformer 模块（attn_blocks）、最终的层归一化层（layer_norm）以及用于语言建模的线性输出层（lm_head）。

_pre_attn_pass方法负责将词元嵌入与位置嵌入相加。forward方法则将输入序列依次通过嵌入层、Transformer 模块序列，最后经过层归一化，得到输出 logits。

如果提供了目标词元序列，forward方法还会计算交叉熵损失。forward_embedding方法提供了一个中间的前向传播过程，仅经过嵌入层和注意力模块。generate方法实现了词元序列的生成功能。

11、Transformer 预训练

在开始训练之前，我们需要为 Transformer 模型设定一些超参数。为了便于演示，我们会将这些参数设定得较小，以便在普通计算机上也能顺利运行。

# --- 定义模型超参数 ---
CONTEXT_LENGTH = 64# 模型一次处理的最大标记数（上下文长度）
N_EMBED = 256        # 嵌入维度（每个标记向量的大小）
N_HEAD = 8           # 注意力头的数量
N_BLOCKS = 6         # Transformer模块的堆叠数量
VOCAB_SIZE = trained_hf_tokenizer.get_vocab_size()  # 从训练好的分词器获取词汇表大小
LEARNING_RATE = 3e-4# 优化器的学习率
BATCH_SIZE = 32      # 每次处理的序列数量（批次大小）
N_EPOCHS = 100       # 遍历整个数据集的次数（演示用）

# --- 设备配置 ---

DEVICE = 'cuda'if torch.cuda.is_available() else'cpu'

# --- 实例化Transformer模型 ---

model = Transformer(
    n_head=N_HEAD,
    n_embed=N_EMBED,
    context_length=CONTEXT_LENGTH,
    vocab_size=VOCAB_SIZE,
    N_BLOCKS=N_BLOCKS
).to(DEVICE)

# 输出模型参数数量
num_params = sum(p.numel() for p in model.parameters())
print(f"我们的Transformer模型共有{num_params/1e6:.2f}百万个参数。")

### 输出示例

使用设备：cuda（若无GPU则为cpu）
我们的Transformer模型共有2.15百万个参数。

我们搭建的这个模型大约有 200 万个参数。虽然与 GPT-3 这样的巨型模型相比微不足道，但作为理解 Transformer 机制的入门示例却非常合适。

接下来，我们会将数据输入模型，计算预测结果与真实值之间的差距（即“损失”），并使用优化器调整模型内部的参数（权重），从而逐步提升模型的表现。

# --- 优化器配置 ---
# AdamW是训练Transformer模型的常用优化器
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

# --- 混合精度训练（可选，但推荐GPU使用） ---
# autocast可自动启用混合精度训练，部分运算使用float16以加快训练速度并降低显存占用，同时保持其他运算的精度。
PTDTYPE = torch.float32  # 本文默认使用float32
autocast_ctx = nullcontext() if DEVICE == 'cpu'else torch.amp.autocast(device_type=DEVICE, dtype=PTDTYPE)
# GradScaler用于防止使用float16时梯度过小（下溢）的问题
scaler = torch.cuda.amp.GradScaler(enabled=(PTDTYPE != torch.float32 and DEVICE == 'cuda'))

# --- 预训练循环 ---
print(f"开始预训练，共计{N_EPOCHS}个周期...")
model.train()  # 将模型设置为训练模式

for epoch in range(N_EPOCHS):
    total_loss = 0
    batch_count = 0

    for xb, yb in pretrain_dataloader:
        # 将批次数据移动到指定设备
        xb, yb = xb.to(DEVICE), yb.to(DEVICE)

        # 清空上一轮迭代的梯度
        optimizer.zero_grad(set_to_none=True)

        with autocast_ctx:
            # 前向传播：获取模型的预测结果（logits）和损失值
            logits, loss = model(xb, targets=yb)

        if loss isnotNone:  # 确保损失值有效
            # 反向传播：计算损失相对于模型参数的梯度
            scaler.scale(loss).backward()
            # 优化器更新：调整模型参数
            scaler.step(optimizer)
            # 更新scaler状态
            scaler.update()

            total_loss += loss.item()
        batch_count += 1

        # 定期输出训练进度（例如每个周期输出两次）
        if batch_count % max(1, len(pretrain_dataloader)//2) == 0:
            print(f"周期 {epoch+1}/{N_EPOCHS}, 批次 {batch_count}/{len(pretrain_dataloader)}, 损失值: {loss.item():.4f}")

    avg_epoch_loss = total_loss / batch_count if batch_count > 0else float('nan')
    print(f"--- 第{epoch+1}个周期结束，平均预训练损失: {avg_epoch_loss:.4f} ---")

print("预训练完成。")

# --- 保存预训练模型 ---
PRETRAINED_MODEL_PATH = os.path.join(NOTEBOOK_OUT_DIR, "thinking_llm_pretrained.pth")
torch.save(model.state_dict(), PRETRAINED_MODEL_PATH)
print(f"预训练模型权重已保存至：{PRETRAINED_MODEL_PATH}")

当我们启动训练循环后，每个训练周期都会输出损失值：

### 输出示例 ###
开始进行100个周期的预训练...
周期 1/100，批次 3/7，损失：5.8732
周期 1/100，批次 6/7，损失：5.1234
--- 第1个周期结束，平均预训练损失：5.4321 ---
...（经过多个周期后）...
周期 100/100，批次 3/7，损失：1.5678
周期 100/100，批次 6/7，损失：1.4990

--- 第100个周期结束，平均预训练损失：1.5220 ---
预训练完成。
预训练模型权重已保存至：./out_notebook_scratch_blog/thinking_llm_pretrained.pth

此时，我们的模型已从预训练数据中初步掌握了一些基本的语言结构。如果现在调用模型的生成方法，它或许能生成一些看似连贯的文本（但由于数据量较小，生成的内容可能重复或毫无意义）。

12、预训练模型推理

我们快速测试一下，看看预训练后的模型是否学到了一些东西，即便只是微不足道的知识。

pt_model.eval()  # 设置为评估模式
test_prompt_str_pt = "Language models learn"
# 如果分词器不会自动添加BOS标记，则手动添加以与训练时保持一致
pt_test_input_ids = tokenizer(tokenizer.bos_token + test_prompt_str_pt, return_tensors="pt").input_ids.to(DEVICE)

with torch.no_grad(), autocast_ctx:
    generated_output_pt = pt_model.generate(
        pt_test_input_ids,
        max_new_tokens=15,
        do_sample=False,  # 此处使用贪婪策略进行测试
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id
    )
decoded_generated_pt = tokenizer.decode(generated_output_pt[0], skip_special_tokens=True)
print(f"提示词：'{test_prompt_str_pt}' -> 生成结果：'{decoded_generated_pt}'")

### 输出示例 ###
Language models learn learn learn learn learn

可以看到，模型明显出现了幻觉（重复生成相同词语），这在意料之中，因为我们只进行了少量周期的训练，且数据量较小。

我还使用了更大的数据集 Pile 数据集（约 5GB 子集） 来训练相同的模型架构，同样训练周期数下，生成结果如下：

# 200万参数LLM在小样本数据上的输出
Language models learn learn learn learn learn

# 200万参数LLM在PILE数据集上的输出
Language models were directly linked to

目前，我们已经拥有了一个初步预训练的语言模型。在进入下一步监督微调（SFT）之前，通常需要进行大量额外工作，比如检查并减少幻觉现象、过滤不雅用语，以及反复试验调整参数。

但本篇博文旨在简单演示“具备思考能力的语言模型”的基本实现原理，因此我们可以直接进入监督微调阶段（SFT）。下面我们就开始吧。

13、SFT 数据集加载器

预训练阶段通过大量通用文本数据，使模型掌握了语言的基本结构和规律。然而，为了让模型更好地完成特定任务（如回答问题或遵循指令），我们需要进行监督微调（SFT）。

在监督微调阶段，我们使用精心设计的数据集（例如我们之前的sample_sft_data），其中包含用户提示与理想助手回复的配对数据。

在微调聊天模型时，一个非常关键的技巧是损失掩码（loss masking）。训练时，模型输入序列同时包含用户的提示和助手的回复。

但我们只希望模型从助手的回复中学习，而不希望它从用户的提示中学习（因为在实际交互中，用户提示是模型的输入，而非模型需要生成的内容）。

损失掩码（Loss Masking）确保只有助理回复对应的词元参与损失计算，这样模型的学习就能专注于生成恰当且有帮助的回复，而在优化过程中忽略用户的提示内容。

首先，我们需要创建一个监督微调（SFT）数据集加载类，用于处理我们示例数据集的 SFT 训练：

class SFTDataset(Dataset):
    """
    用于监督微调（SFT）的简化数据集类，采用ChatML风格的格式化。
    仅助理的回复部分会参与损失计算。
    """
    def __init__(self, file_path, tokenizer, context_length):
        self.tokenizer = tokenizer
        self.context_length = context_length

        with open(file_path, 'r', encoding='utf-8') as f:
            self.conversations = [json.loads(line)['conversations'] for line in f]

        print(f"已加载 {len(self.conversations)} 条对话数据。")

    def __len__(self):
        return len(self.conversations)

    def __getitem__(self, idx):
        conversation = self.conversations[idx]

        tokens = [self.tokenizer.bos_token_id]
        labels = [-100]

        for turn in conversation:
            role, content = turn['role'], turn['content']

            # 编码当前对话轮次
            prefix = self.tokenizer.encode(f"<|im_start|>{role}\n", add_special_tokens=False).ids
            content_ids = self.tokenizer.encode(content, add_special_tokens=False).ids
            suffix = self.tokenizer.encode("<|im_end|>\n", add_special_tokens=False).ids

            # 拼接词元
            tokens.extend(prefix + content_ids + suffix)

            # 创建标签，仅助理回复部分参与损失计算
            if role == 'assistant':
                labels.extend([-100] * len(prefix))  # 前缀不参与损失
                labels.extend(content_ids)           # 助理内容参与损失
                labels.extend(suffix)                # 后缀也参与损失
            else:
                labels.extend([-100] * (len(prefix) + len(content_ids) + len(suffix)))  # 用户内容全部忽略

        # 确保序列长度为 context_length + 1，以便后续移位操作
        pad_id = self.tokenizer.pad_token_id
        tokens = tokens[:self.context_length + 1] + [pad_id] * max(0, self.context_length + 1 - len(tokens))
        labels = labels[:self.context_length + 1] + [-100] * max(0, self.context_length + 1 - len(labels))

        # 准备输入和目标序列（错位一个位置）
        input_ids = torch.tensor(tokens[:-1], dtype=torch.long)
        target_ids = torch.tensor(labels[1:], dtype=torch.long)

        return input_ids, target_ids

上述SFTDataset类的作用是构建输入序列和标签序列，其中只有助理回复的部分会参与训练时的损失计算。

也就是说，用户的输入虽然包含在模型的输入序列中，但在计算损失时会被忽略，从而使模型的学习专注于生成助理的回复。

在更复杂的场景下，可能需要更精细的词元对齐策略，但上述方法已足以覆盖基本的监督微调需求。

14、SFT 训练循环

接下来，我们编写监督微调的训练循环，使用之前训练好的基于 Transformer 架构的预训练模型：

# --- 监督微调（SFT）训练循环 ---
# 微调时通常使用较小的学习率
SFT_LEARNING_RATE = 1e-4
SFT_N_EPOCHS = 50# 监督微调通常在小型数据集上需要更多轮次

optimizer_sft = torch.optim.AdamW(model.parameters(), lr=SFT_LEARNING_RATE)
# 如果GPU设置不变，之前的scaler和autocast_ctx可以继续使用

print(f"开始监督微调，共计 {SFT_N_EPOCHS} 个epoch...")
model.train()  # 设置模型为训练模式

for epoch in range(SFT_N_EPOCHS):
    total_sft_loss = 0
    sft_batch_count = 0

    for xb_sft, yb_sft_labels in sft_dataloader:
        xb_sft, yb_sft_labels = xb_sft.to(DEVICE), yb_sft_labels.to(DEVICE)

        optimizer_sft.zero_grad(set_to_none=True)

        with autocast_ctx:
            # 前向传播：模型的forward方法接受输入序列`idx`和目标序列`targets`
            # 这里的`targets`即为`yb_sft_labels`，其中非助理部分已标记为-100
            # 模型内部会使用这些标签计算交叉熵损失
            logits_sft, loss_sft = model(xb_sft, targets=yb_sft_labels)

        if loss_sft isnotNone:
            scaler.scale(loss_sft).backward()
            scaler.step(optimizer_sft)
            scaler.update()
            total_sft_loss += loss_sft.item()
        sft_batch_count += 1

        if sft_batch_count % max(1, len(sft_dataloader)//1) == 0:  # 对于小型数据集，每个epoch记录一次
            print(f"SFT Epoch {epoch+1}/{SFT_N_EPOCHS}, Batch {sft_batch_count}/{len(sft_dataloader)}, Loss: {loss_sft.item():.4f}")
```bash

```bash

avg_epoch_sft_loss = total_sft_loss / sft_batch_count if sft_batch_count > 0 else float('nan')
print(f"--- 监督微调第 {epoch+1} 轮结束，平均损失：{avg_epoch_sft_loss:.4f} ---")

print("监督微调训练完成。")

上述循环非常简单，唯一的变化就是加入了监督微调损失（SFT Loss）的计算。当我们启动训练后，每个训练轮次都会打印出对应的损失值，输出示例如下：

### 输出示例 ###
开始进行50轮监督微调训练...
监督微调 第1/50轮, 批次 1/1, 损失: 2.1241
--- 监督微调第 1 轮结束，平均损失：2.1241 ---
...
监督微调 第50/50轮, 批次 1/1, 损失: 0.8721
--- 监督微调第 50 轮结束，平均损失：0.8721 ---
监督微调训练完成。

经过监督微调训练，我们原本只具备基础语言生成能力的预训练模型，现在已经初步掌握了语言问答的能力。比如当我们输入“你好”，它可能会给出“你好，有什么可以帮您的吗？”这样的较为连贯的回复。

15、SFT 模型推理

尽管我们用于监督微调的数据集规模很小，预训练模型本身也非常有限，但我们还是来看一下实际效果。以下展示了一个在 Pile 数据集上预训练、并在 OpenOrca 数据集上进行了 50 轮监督微调后的语言模型的输出情况：

# 准备一个简单的测试对话
chat_history = [{"role": "user", "content": "法国的首都是哪里？"}]
prompt = tokenizer.apply_chat_template(chat_history, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

# 生成回复
with torch.no_grad(), autocast_ctx_sft:
    outputs = sft_model_demo.generate(
        inputs.input_ids,
        max_new_tokens=200,
        do_sample=True,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id
    )

# 解码并记录回复
response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
logger(f"监督微调提示词: '{chat_history[0]['content']}' -> 模型生成回复: '{response}'")

我们提出了一个非常基础的问题，下面看看我们这个相对更好的模型给出的答案：

# 200万参数监督微调模型的输出（使用小规模监督微调数据）
unbiased asf gg link

# 200万参数监督微调模型的输出（使用ORCA监督微调数据）
paris is the bigger

可以看到，使用更大规模数据集训练的模型，已经开始给出更准确的回答了（尽管仍不完美）。

暂且假设我们的监督微调模型表现不错，接下来我们将进入下一阶段的训练——强化学习（RLHF），它能赋予语言模型更强的推理和思考能力。

16、RLHF 风格数据加载器

由于强化学习阶段的数据格式与监督微调阶段的数据格式在本质上都是对话形式，我们可以直接复用之前定义的SFTDataset结构。我们之前定义的sample_reasoning_data数据集中，助手的“思考过程”和“答案”部分已经包含在内容中。

SFTDataset会确保只使用助手的回复部分（包括特殊标记）来计算损失。

# --- 将推理数据保存到文件 ---
reasoning_file_path = os.path.join(NOTEBOOK_DATA_DIR, "reasoning_data.jsonl")
with open(reasoning_file_path, 'w', encoding='utf-8') as f:
    for item in sample_reasoning_data: # sample_reasoning_data 已在博客开头定义
        f.write(json.dumps(item) + '\n')
print(f"推理数据已保存至：{reasoning_file_path}")

# --- 实例化推理数据集与数据加载器 ---
# 我们复用之前定义的SFTDataset类，因为它能正确处理对话格式，
# 并且只对助手的回复部分计算损失。
# 其中，<think> 和 <answer> 标签都包含在助手的回复内容中。
reasoning_dataset = SFTDataset(
    file_path=reasoning_file_path,
    tokenizer=trained_hf_tokenizer,  # 我们训练好的分词器
    context_length=CONTEXT_LENGTH    # 前面定义的超参数
)

# 若显存受限，尤其是模型较大或上下文长度较长时，应减小批次大小
REASONING_BATCH_SIZE = max(1, BATCH_SIZE // 2)  # 例如：为预训练/SFT阶段批次大小的一半
reasoning_dataloader = DataLoader(reasoning_dataset, batch_size=REASONING_BATCH_SIZE, shuffle=True)

数据加载器会正确地准备序列，让模型学习生成完整的助手回复，包括关键的和标签。

17、带权重标签损失的训练循环

为了让模型更好地掌握和标签的结构，我们稍微修改一下损失函数的计算方式。

具体而言，我们会识别这些特殊标签对应的token，并在损失函数中赋予它们更高的权重，以明确告诉模型：准确生成这些结构性标签尤为重要。

# SFT_MODEL_PATH = os.path.join(NOTEBOOK_OUT_DIR, "thinking_llm_sft.pth") # 示例路径
if os.path.exists(SFT_MODEL_PATH):
    model.load_state_dict(torch.load(SFT_MODEL_PATH, map_location=DEVICE))
    print(f"已从 {SFT_MODEL_PATH} 加载SFT阶段训练好的模型权重。")
else:
    print(f"警告：未找到SFT模型路径 {SFT_MODEL_PATH}，推理训练将使用当前模型状态。")

# --- 推理训练阶段的超参数 ---
REASONING_LEARNING_RATE = 5e-5# 通常在最后微调阶段使用较小的学习率
REASONING_N_EPOCHS = 75         # 为掌握新结构，可能需要更多的训练轮次

optimizer_rsn = torch.optim.AdamW(model.parameters(), lr=REASONING_LEARNING_RATE)

# --- 定义特殊推理标签及其对应的token ID ---
# 这些标签应在训练分词器时已加入SPECIAL_TOKENS_LIST，
# 确保trained_hf_tokenizer将它们视为单个token。
think_token_id = trained_hf_tokenizer.token_to_id("<think>")
end_think_token_id = trained_hf_tokenizer.token_to_id("</think>")
answer_token_id = trained_hf_tokenizer.token_to_id("<answer>")
end_answer_token_id = trained_hf_tokenizer.token_to_id("</answer>")

定义完训练参数后，我们接下来定义 RLHF 训练阶段的训练循环。

# --- 推理训练循环 ---
print(f"开始推理训练，共计 {REASONING_N_EPOCHS} 个epoch...")
model.train()  # 将模型设为训练模式

for epoch in range(REASONING_N_EPOCHS):
    total_reasoning_loss = 0
    rsn_batch_count = 0

    for xb_rsn, yb_rsn_labels in reasoning_dataloader:
        xb_rsn, yb_rsn_labels = xb_rsn.to(DEVICE), yb_rsn_labels.to(DEVICE)

        optimizer_rsn.zero_grad(set_to_none=True)

        with autocast_ctx:
            # 获取模型输出的logits。我们模型的forward方法若传入targets会自动计算标准交叉熵损失。
            # 但我们需要加权损失，因此手动计算。
            logits_rsn, _ = model(xb_rsn)  # 仅传入输入索引，获取logits

            # 将logits和标签展平，以便逐token计算损失
            # logits_rsn维度：(B, T, 词汇表大小), yb_rsn_labels维度：(B, T)
            flat_logits_rsn = logits_rsn.view(-1, VOCAB_SIZE)  # (B*T, 词汇表大小)
            flat_labels_rsn = yb_rsn_labels.view(-1)           # (B*T)

            # 逐token计算原始损失（忽略标签为-100的token）
            raw_loss_per_token = loss_fct_no_reduction(flat_logits_rsn, flat_labels_rsn)

            # 为每个token创建权重张量：
            # 标签不为-100的token权重为1.0，否则为0.0
            token_weights = (flat_labels_rsn != -100).float()

            # 找出特殊推理标签在目标标签中的位置
            # is_special_tag_mask在flat_labels_rsn为特殊推理标签ID时为True
            is_special_tag_mask = torch.isin(flat_labels_rsn, special_reasoning_tag_ids)

# 对同时属于有效目标标签的特殊标记赋予更高的权重
# （即标签值不为-100且属于我们定义的特殊标记）
target_special_tags_mask = is_special_tag_mask & (flat_labels_rsn != -100)
token_weights[target_special_tags_mask] *= REASONING_TAG_LOSS_WEIGHT

# 计算最终的加权损失
# 具体为：每个token的损失乘以对应的权重后求和，
# 再除以实际目标token的数量（即标签不为-100的token数）
# 这样做的目的是使损失的量级与未加权的监督微调（SFT）大致相当，
# 同时又能突出强调特殊标记的重要性。
num_actual_target_tokens = (flat_labels_rsn != -100).float().sum().clamp(min=1)
weighted_loss_rsn = (raw_loss_per_token * token_weights).sum() / num_actual_target_tokens

# 反向传播与优化器更新，需检查损失是否有效
ifnot torch.isnan(weighted_loss_rsn) andnot torch.isinf(weighted_loss_rsn):
    scaler.scale(weighted_loss_rsn).backward()
    scaler.step(optimizer_rsn)
    scaler.update()
    total_reasoning_loss += weighted_loss_rsn.item()
else:
    # 如果某个批次中所有标签均为-100，可能导致实际目标token数为零，从而损失为NaN或Inf
    print(f"警告：推理训练第{epoch+1}轮，第{rsn_batch_count+1}批次出现NaN或Inf损失，本批次更新被跳过。")

rsn_batch_count += 1

# 记录训练进度
if rsn_batch_count % max(1, len(reasoning_dataloader)//1) == 0: # 如果批次数很少，每轮至少记录一次
    current_loss_item = weighted_loss_rsn.item() ifnot (torch.isnan(weighted_loss_rsn) or torch.isinf(weighted_loss_rsn)) else float('nan')
    print(f"推理训练第{epoch+1}/{REASONING_N_EPOCHS}轮，第{rsn_batch_count}/{len(reasoning_dataloader)}批次，加权损失：{current_loss_item:.4f}")

avg_epoch_rsn_loss = total_reasoning_loss / rsn_batch_count if rsn_batch_count > 0else float('nan')
print(f"--- 推理训练第{epoch+1}轮结束，平均推理损失：{avg_epoch_rsn_loss:.4f} ---")

print("推理训练完成。")

在上述循环中，我们专门针对推理任务对模型进行多轮训练。每个批次中，模型都会生成预测（logits），并与真实标签进行比较。

值得注意的是，我们在计算损失时，特意对某些特殊的推理标记赋予了更高的权重，以促使模型更加关注这些关键标记的学习。

在得到加权损失后，模型通过反向传播和优化器更新自身参数。同时，我们也加入了安全检查机制，以跳过可能出现异常损失（如 NaN 或 Inf）的批次。

训练过程中，我们会持续记录进度，并在每轮结束时报告平均损失。

### 输出示例 ###
已成功从 ./blog_output_dir/thinking_llm_sft.pth 加载监督微调（SFT）训练后的模型权重
用于加权的特殊推理标记ID为：[ID_think, ID_endthink, ID_answer, ID_endanswer]（实际token ID）
开始进行75轮推理训练...
推理训练第1/75轮，第1/1批次，加权损失：1.9503
--- 推理训练第1轮结束，平均推理损失：1.9503 ---
...（若干轮训练后）...
推理训练第75/75轮，第1/1批次，加权损失：0.3881
--- 推理训练第75轮结束，平均推理损失：0.3881 ---

训练启动后会打印损失信息。训练完成后，我们训练好的强化学习（RLHF）模型将被保存下来。

我们期望训练好的 RLHF 模型在生成回答时，会先生成``标记后的推理过程内容，再根据推理过程给出最终的回答。

18、推理型 LLM 的最终推理

如前所述，监督微调（SFT）或预训练只有在数据量更大、训练轮次更多、参数规模更大的情况下才会有效。强化学习人类反馈（RLHF）也是如此：我发现，当使用更大的RLHF数据集训练时，模型生成的回复质量明显更高。虽然仍不完美，但至少更接近我们对模型的最低期望。

使用RLHF模型进行推理时，代码与之前略有不同，因为RLHF模型的输出包含了特定的标签，而此前的监督微调或预训练模型输出的只是纯文本。

# 加载与已保存模型匹配的最终推理模型配置
final_model_config = DemoLLMConfig(
    vocab_size=DEMO_VOCAB_SIZE_FINAL,
    hidden_size=DEMO_HIDDEN_SIZE,
    intermediate_size=DEMO_INTERMEDIATE_SIZE,
    num_hidden_layers=DEMO_NUM_LAYERS,
    num_attention_heads=DEMO_NUM_ATTENTION_HEADS,
    num_key_value_heads=DEMO_NUM_KV_HEADS,
    max_position_embeddings=DEMO_MAX_SEQ_LEN,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
)

# 初始化模型并加载训练好的权重
final_thinking_llm = DemoLLMForCausalLM(final_model_config).to(DEVICE)
final_thinking_llm.load_state_dict(torch.load(final_reasoning_model_path, map_location=DEVICE))
final_thinking_llm.eval()

# 使用分词器的聊天模板准备模型输入提示
user_query = "如果我有3个苹果，吃掉1个，还剩几个？"
chat_history = [{"role": "user", "content": user_query}]
prompt_text = tokenizer.apply_chat_template(chat_history, tokenize=False, add_generation_prompt=True)

# 对提示进行分词并将输入ID移至设备
input_ids = tokenizer(prompt_text, return_tensors="pt").input_ids.to(DEVICE)

# 使用采样参数生成模型回复
with torch.no_grad(), autocast_ctx_rsn:
    generated_ids = final_thinking_llm.generate(
        input_ids,
        max_new_tokens=DEMO_MAX_SEQ_LEN - 10,
        do_sample=True,
        temperature=0.7,
        top_k=10,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
    )

# 解码生成的token，跳过特殊token
assistant_response = tokenizer.decode(generated_ids[0][input_ids.shape[1]:], skip_special_tokens=True)

# 辅助函数：提取标签内的内容
def extract_tag(text, tag):
    start_tag = f"<{tag}>"
    end_tag = f"</{tag}>"
    start_idx = text.find(start_tag)
    end_idx = text.find(end_tag)
    if start_idx != -1and end_idx != -1:
        return text[start_idx + len(start_tag):end_idx].strip()
    returnNone

# 提取<think>和<answer>部分（若存在）
think = extract_tag(assistant_response, "think") or"未找到"
answer = extract_tag(assistant_response, "answer") or assistant_response

# 输出解析后的思考过程和答案
print(f"<think>: {think}")
print(f"<answer>: {answer}")

我们之前已经看到，在极小的数据样本上训练模型并不能产生好的结果。

不过，我也尝试在更优质的数据集上对RLHF模型进行了微调，比如 HFHub RL 数据集。与在微小数据集上训练的模型相比，结果明显提升。请看下面的对比：

# 用户提问
如果我有3个苹果，吃掉1个，还剩几个？

# 200万参数RLHF模型输出（在小样本RLHF数据上训练）
<think> ugh far igu when

# 200万参数监督微调模型输出（在HFHub RLHF数据上训练）
<think> 用户问苹果问题还剩几个？</think>2

可以看到，在小样本数据上训练的模型输出毫无意义的乱码，而在更大数据集上训练的模型虽然仍存在一些小问题，但已经能够识别并使用思考标签了。

更重要的是，它给出了正确的答案。尽管思考标签中的内容还有些小瑕疵，但只要继续扩大数据规模，这些问题一定能得到改善。

那么，如何系统的去学习大模型LLM？

如果你也想系统学习AI大模型技术，想通过这项技能真正达到升职加薪，就业或是副业的目的，但是不知道该如何开始学习*_，因为网上的资料太多太杂乱了，如果不能系统的学习就相当于是白学。
为了帮助大家打破壁垒，快速了解大模型核心技术原理，学习相关大模型技术。从原理出发真正入局大模型。在这里我和MoPaaS魔泊云联合梳理打造了系统大模型学习脉络，这份 LLM大模型资料 分享出来：包括LLM大模型书籍、640套大模型行业报告、LLM大模型学习视频、LLM大模型学习路线、开源大模型学习教程等, 😝有需要的小伙伴，可以 扫描下方二维码免费领取🆓**⬇️⬇️⬇️

【大模型全套视频教程】

教程从当下的市场现状和趋势出发，分析各个岗位人才需求，带你充分了解自身情况，get 到适合自己的 AI 大模型入门学习路线。

从基础的 prompt 工程入手，逐步深入到 Agents，其中更是详细介绍了 LLM 最重要的编程框架 LangChain。最后把微调与预训练进行了对比介绍与分析。

同时课程详细介绍了AI大模型技能图谱知识树，规划属于你自己的大模型学习路线，并且专门提前收集了大家对大模型常见的疑问，集中解答所有疑惑！

深耕 AI 领域技术专家带你快速入门大模型

跟着行业技术专家免费学习的机会非常难得，相信跟着学习下来能够对大模型有更加深刻的认知和理解，也能真正利用起大模型，从而“弯道超车”，实现职业跃迁！

【精选AI大模型权威PDF书籍/教程】

精心筛选的经典与前沿并重的电子书和教程合集，包含《深度学习》等一百多本书籍和讲义精要等材料。绝对是深入理解理论、夯实基础的不二之选。

【AI 大模型面试题 】

除了 AI 入门课程，我还给大家准备了非常全面的**「AI 大模型面试题」，**包括字节、腾讯等一线大厂的 AI 岗面经分享、LLMs、Transformer、RAG 面试真题等，帮你在面试大模型工作中更快一步。

【大厂 AI 岗位面经分享（92份）】

【AI 大模型面试真题（102 道）】

【LLMs 面试真题（97 道）】

【640套 AI 大模型行业研究报告】

【AI大模型完整版学习路线图（2025版）】

明确学习方向，2025年 AI 要学什么，这一张图就够了！

👇👇点击下方卡片链接免费领取全部内容👇👇

抓住AI浪潮，重塑职业未来！

科技行业正处于深刻变革之中。英特尔等巨头近期进行结构性调整，缩减部分传统岗位，同时AI相关技术岗位（尤其是大模型方向）需求激增，已成为不争的事实。具备相关技能的人才在就业市场上正变得炙手可热。

行业趋势洞察：

• 转型加速： 传统IT岗位面临转型压力，拥抱AI技术成为关键。
• 人才争夺战： 拥有3-5年经验、扎实AI技术功底和真实项目经验的工程师，在头部大厂及明星AI企业中的薪资竞争力显著提升（部分核心岗位可达较高水平）。
• 门槛提高： “具备AI项目实操经验”正迅速成为简历筛选的重要标准，预计未来1-2年将成为普遍门槛。

与其观望，不如行动！

面对变革，主动学习、提升技能才是应对之道。掌握AI大模型核心原理、主流应用技术与项目实战经验，是抓住时代机遇、实现职业跃迁的关键一步。

01 为什么分享这份学习资料？

当前，我国在AI大模型领域的高质量人才供给仍显不足，行业亟需更多有志于此的专业力量加入。

因此，我们决定将这份精心整理的AI大模型学习资料，无偿分享给每一位真心渴望进入这个领域、愿意投入学习的伙伴！

我们希望能为你的学习之路提供一份助力。如果在学习过程中遇到技术问题，也欢迎交流探讨，我们乐于分享所知。

*02 这份资料的价值在哪里？*

专业背书，系统构建：

• 本资料由我与MoPaaS魔泊云的鲁为民博士共同整理。鲁博士拥有清华大学学士和美国加州理工学院博士学位，在人工智能领域造诣深厚：

• • 在IEEE Transactions等顶级学术期刊及国际会议发表论文超过50篇。
  • 拥有多项中美发明专利。
  • 荣获吴文俊人工智能科学技术奖（中国人工智能领域重要奖项）。

• 目前，我有幸与鲁博士共同进行人工智能相关研究。

内容实用，循序渐进：

• 资料体系化覆盖了从基础概念入门到核心技术进阶的知识点。

• 包含丰富的视频教程与实战项目案例，强调动手实践能力。

• 无论你是初探AI领域的新手，还是已有一定技术基础希望深入大模型的学习者，这份资料都能为你提供系统性的学习路径和宝贵的实践参考，助力你提升技术能力，向大模型相关岗位转型发展。

  

抓住机遇，开启你的AI学习之旅！