1. Day 2 学习内容

今天结合 GPT-1 与 GPT-2 论文继续学习 Decoder-only Transformer。GPT-2 延续了自回归语言模型的预训练目标，并重点展示了模型在不进行特定任务微调时的零样本能力；在网络结构上采用 Decoder-only 架构，并将 LayerNorm 放在注意力层和 MLP 之前，即 Pre-Norm。

GPT-2 Small 的词表大小为 50,257，上下文长度为 1,024，隐藏维度为 768，由 12 层 Transformer Block 和 12 个注意力头组成，所有 Block 之后还有一次最终 LayerNorm。

下图展示了 GPT-2 的整体数据流：Token Embedding 与 Position Embedding 相加后，依次经过多层掩码自注意力和前馈网络，最后输出下一个 Token 的预测分数。

GPT-2 Decoder-only Transformer 架构

Transformer 的基础原理已经在《02 Transformer 算法》中整理，这里主要参考 GPT-2 和 Andrej Karpathy 的实现方式，完成适合本机 RTX 5060 8 GB 显存的轻量版本。

今天的代码对应总体规划中的多头因果注意力、Transformer Block 和完整 Tiny-GPT 主体结构。

2. 代码实现

当前模型配置为 4 层 Transformer Block、4 个注意力头、256 维 Embedding 和 256 个 Token 的上下文长度。代码依次实现配置类、多头因果注意力、MLP、Pre-Norm Block，以及完整 Tiny-GPT 的前向传播和交叉熵损失。

model.py：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
import torch
from torch import nn
import math
from dataclasses import dataclass
import torch.nn.functional as F
from pathlib import Path


@dataclass
class GPTConfig:
    block_size: int = 256
    vocab_size: int = 73
    n_layer: int = 4
    n_head: int = 4
    n_embd: int = 256
    dropout: float = 0.1


# 拿到处理好的文件
TINY_GPT_DIR = Path(__file__).parent
SRC_DIR = TINY_GPT_DIR.parent
ROOT_DIR = SRC_DIR.parent

FILE_PATH = ROOT_DIR / "data" / "processed" / "alice.txt"
OUT_PATH = ROOT_DIR / "tokenizer.json"

# 设备检测：CPU 或者 GPU
device = "cpu"
if torch.cuda.is_available():
    device = "cuda"


# 实现多头掩码注意力机制
class CausalMaskedAttention(nn.Module):
    def __init__(self, config):
        super().__init__()

        # 将 [B, T, C] 变成 [B, T, 3C]，拆分为 Q、K、V
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)

        # 在最后输出之前经过线性层拼接多头
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)

        # Dropout
        self.attn_dropout = nn.Dropout(config.dropout)
        self.out_dropout = nn.Dropout(config.dropout)

        # 在 forward 中使用
        self.n_embd = config.n_embd
        self.n_head = config.n_head

    def forward(self, x):
        # 读取 B、T、C
        B, T, C = x.size()

        # 定义 Q、K、V，便于拆分
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.n_embd, dim=-1)

        # 拆分为多头并转换：[B, T, C] -> [B, T, H, D] -> [B, H, T, D]
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        k = k.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)

        # 缩放点积注意力：[B, H, T, D] @ [B, H, D, T] = [B, H, T, T]
        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(
            C // self.n_head
        )

        # 制作下三角掩码，并扩展为 [1, 1, T, T] 与 scores 广播
        mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T)
        scores = scores.masked_fill(mask == 0, float("-inf"))

        # Softmax
        weights = torch.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)

        # 乘以 V
        output = torch.matmul(weights, v)

        # 转置回 [B, H, T, D] -> [B, T, H, D] -> [B, T, C]
        output = output.transpose(1, 2).contiguous().view(B, T, C)

        # 经过最后一个线性层
        y = self.c_proj(output)
        y = self.out_dropout(y)
        return y


# 实现 MLP
class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.network = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            nn.GELU(approximate="tanh"),
            nn.Linear(4 * config.n_embd, config.n_embd),
            nn.Dropout(config.dropout),
        )

    def forward(self, x):
        return self.network(x)


# 每一层 Block，也就是 Decoder
class Block(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalMaskedAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        # GPT-2 架构采用 Pre-Norm 形式
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


# Tiny-GPT 完整结构
class Tiny_GPT(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.transformer = nn.ModuleDict(
            dict(
                wte=nn.Embedding(config.vocab_size, config.n_embd),
                wpe=nn.Embedding(config.block_size, config.n_embd),
                h=nn.ModuleList(
                    [Block(config) for _ in range(config.n_layer)]
                ),
                ln_f=nn.LayerNorm(config.n_embd),
            )
        )
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size)

    def forward(self, x, target=None):
        B, T = x.size()

        # 进行字符编码
        token_embd = self.transformer.wte(x)

        # 进行位置编码
        pos = torch.arange(0, T, dtype=torch.long, device=x.device)
        pos_embd = self.transformer.wpe(pos)

        # 相加得到 Block 的输入 x
        x = token_embd + pos_embd

        # 进行 Block 中的运算，每一层执行一次
        for layer in self.transformer.h:
            x = layer(x)

        # 最终 LayerNorm
        x = self.transformer.ln_f(x)

        # LM Head
        logits = self.lm_head(x)

        # 如果 target 不是 None，就计算 loss
        loss = None
        if target is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                target.reshape(-1),
            )

        return logits, loss

3. 重点知识

3.1 Dropout 的位置

当前实现主要在两个位置使用 Dropout：

  1. 注意力权重：Softmax 得到 weights 后、与 V 相乘前执行，避免模型过度依赖少数 Token 之间的注意力连接。
  2. 子层输出：多头注意力输出投影和 MLP 输出之后执行，降低相邻子层之间形成高度耦合依赖的风险。

3.2 交叉熵损失的输入形状

F.cross_entropy 要求预测值最后一维表示类别，因此需要先将 Batch 和时间维度合并：

  • 预测分数：logits.reshape(-1, logits.size(-1))，形状为 [B × T, vocab_size]。
  • 正确答案：target.reshape(-1)，形状为 [B × T]。

这样，每一个时间位置都会作为一个独立的下一个 Token 分类样本参与损失计算。