从0开始搭建大语言模型：构造GPT模型

接上文：【教程】从0开始搭建大语言模型：实现Attention机制

GPT模型

GPT，全称是Generative Pretrained Transformer，它是大型深度神经网络架构，旨在一次生成一个单词(或token)。GPT的大致流程为：

GPT-2的最小版本也有1.24亿个参数，我们通过以下Python字典指定小型GPT-2模型的配置：

GPT_CONFIG_124M = {
	"vocab_size": 50257, # Vocabulary size
	"context_length": 1024, # Context length
	"emb_dim": 768, # Embedding dimension
	"n_heads": 12, # Number of attention heads
	"n_layers": 12, # Number of layers
	"drop_rate": 0.1, # Dropout rate
	"qkv_bias": False # Query-Key-Value bias
}

参数的意义为：

• vocab_size：一个包含50257个单词的词汇表，由BPE tokenizer使用
• context_length：表示模型可以处理的最大输入token数
• emb_dim：表示embedding大小，将每个token转换为768维向量
• n_heads：表示多头注意机制中注意头的数量
• n_layers：指定模型中Transformer块的数量
• drop_rate：表示 dropout机制的概率(0.1表示删除10%的神经元)，以防止过拟合
• qkv_bias：决定是否在多头attention中查询、键和值的线性层中包含一个偏置向量

要构建一个GPT模型，我们需要依次完成下图的模块：

下图展示了输入数据如何tokenized、embedding并提供给GPT模型的整体概述：

在LLM中，embedding的输入token维度通常与输出维度匹配。

下面将展示如何实现GPT的各个模块。

Layer Normalization

由于梯度消失或梯度爆炸等问题，训练具有许多层的深度神经网络有时可能具有挑战性。这些问题导致训练动态不稳定，并使网络难以有效调整其权重，这意味着学习过程很难找到一组神经网络参数(权重)，以最小化损失函数。换句话说，网络很难学习数据中的基本模式，以使其做出准确的预测或决策。

为了解决这个问题，可以采用一些归一化来防止梯度消失或梯度爆炸。在LLM中，我们使用Layer Normalization来达到这一点。

层归一化背后的主要思想是调整神经网络层的激活(输出)，使其均值为0，方差为1，这种调整加快了收敛到有效权重的速度，并确保了一致、可靠的训练。

层归一化的一个示例如下：

要简单地实现层归一化，代码为：

torch.manual_seed(123)
batch_example = torch.randn(2, 5) #A
layer = nn.Sequential(nn.Linear(5, 6), nn.ReLU())
out = layer(batch_example)
print(out)
mean = out.mean(dim=-1, keepdim=True)
var = out.var(dim=-1, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)
# 层归一化
out_norm = (out - mean) / torch.sqrt(var)
mean = out_norm.mean(dim=-1, keepdim=True)
var = out_norm.var(dim=-1, keepdim=True)
print("Normalized layer outputs:\n", out_norm)
print("Mean:\n", mean)
print("Variance:\n", var)

需要注意，代码中dim参数指定在张量中计算统计量(这里是均值或方差)时应该沿着的维度，-1表示最后一维。不同dim的数据计算展示如下：

Layer Normalization类的代码为：

class LayerNorm(nn.Module):
	def __init__(self, emb_dim):
		super().__init__()
		self.eps = 1e-5
		self.scale = nn.Parameter(torch.ones(emb_dim))
		self.shift = nn.Parameter(torch.zeros(emb_dim))
	def forward(self, x):
		mean = x.mean(dim=-1, keepdim=True)
		var = x.var(dim=-1, keepdim=True, unbiased=False)
		norm_x = (x - mean) / torch.sqrt(var + self.eps)
		return self.scale * norm_x + self.shift

需要注意，因为layer normalization是作用于最后一个维度，因此self.scale和self.shift的参数维度也是emb_dim。eps是为了防止分母为0。scale和shift是两个可训练参数(与输入相同维度)，LLM会在训练期间自动调整。这允许模型学习最适合其处理数据的适当缩放和平移。

和批归一化的比较：与批归一化对批维度进行归一化不同，层归一化对特征维度进行归一化。LLM通常需要大量的计算资源，而可用的硬件或特定的用例可以决定训练或推理期间的批大小。由于层归一化独立于批量大小对每个输入进行归一化，因此它在这些场景中提供了更多的灵活性和稳定性。

GELU激活函数

在LLM中，通常采用GELU和SwiGLU函数，而不是传统的ReLU，GELU和SwiGLU是更复杂和平滑的激活函数，分别包含高斯门控线性单元和sigmoid门控线性单元。本部分主要介绍GELU。

GELU的公式为：GELU(x)=x Φ(x)，其中Φ(x)是标准高斯分布的累积分布函数。但通常不使用这种方法计算，而是使用另外一种计算更高效的估计：

GELU(x) ≈ 0.5 ⋅ x ⋅ (1 + tanh[√((2/π)) ⋅ (x + 0.044715 ⋅ x^3])

用代码实现为：

class GELU(nn.Module):
	def __init__(self):
		super().__init__()
	def forward(self, x):
		return 0.5 * x * (1 + torch.tanh(torch.sqrt(torch.tensor(2.0 / torch.pi)) *(x + 0.044715 * torch.pow(x, 3))))

ReLU和GELU的图像比较如下：

从图中可以看出，ReLU是一个分段线性函数，如果输入为正，则直接输出;否则，它输出0。GELU是一个光滑的非线性函数，它近似于ReLU，但对于负值具有非零梯度，这种平滑属性可以让模型在训练过程中更好地优化。

ReLU在零处有一个尖角，这有时会使优化变得更加困难，特别是在深度非常深或具有复杂架构的网络中。对于负数，GELU允许一个较小的非零输出。这一特性意味着在训练过程中，接受负输入的神经元仍然可以对学习过程做出贡献，尽管程度小于正输入。

Feed Forward网络

有了GELU后，我们将它应用在Feed Forward Network（FFN）中，FFN模块是一个由两个线性层和一个GELU激活函数组成的小型神经网络，代码如下：

class FeedForward(nn.Module):
	def __init__(self, cfg):
		super().__init__()
		self.layers = nn.Sequential(
			nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
			GELU(),
			nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
			)
	def forward(self, x):
		return self.layers(x)

这个模块可以增强模型从数据中学习和泛化的能力。尽管该模块的输入和输出维度相同，但它通过第一个线性层在内部将embedding维度扩展到更高维的空间。这种扩展之后是非线性的GELU激活，然后通过第二个线性变换收缩回到原始维度。这样的设计允许探索更丰富的表征空间

增强shortcut连接

shortcut connections，也被称为skip connections或者residual connections，它是为了缓解梯度消失的问题。梯度消失问题是指梯度(在训练期间指导权重更新)随着在各层中反向传播而逐渐变小的问题，使其难以有效训练较早的层。

有shortcut connections和没有shortcut connections的结构比较如下：

从图中可以看出，shortcut connections通过跳过一个或多个层为梯度创建了一个替代的、更短的路径，以通过网络，这是通过将一层的输出添加到后面一层的输出来实现的。

shortcut connections是非常大的模型(如LLM)的核心构建块，当我们训练GPT模型时，它们将通过确保跨层的一致梯度流来帮助促进更有效的训练。

如果你想对shortcut connections的作用实验，可以通过如下代码：

class ExampleDeepNeuralNetwork(nn.Module):
	def __init__(self, layer_sizes, use_shortcut):
		super().__init__()
		self.use_shortcut = use_shortcut
		self.layers = nn.ModuleList([
			# Implement 5 layers
			nn.Sequential(nn.Linear(layer_sizes[0], layer_sizes[1]), GELU()),
			nn.Sequential(nn.Linear(layer_sizes[1], layer_sizes[2]), GELU()),
			nn.Sequential(nn.Linear(layer_sizes[2], layer_sizes[3]), GELU()),
			nn.Sequential(nn.Linear(layer_sizes[3], layer_sizes[4]), GELU()),
			nn.Sequential(nn.Linear(layer_sizes[4], layer_sizes[5]), GELU())
		])
	def forward(self, x):
		for layer in self.layers:
			# 计算输出
			layer_output = layer(x)
			# shortcut是否被应用
			if self.use_shortcut and x.shape == layer_output.shape:
				x = x + layer_output
			else:
				x = layer_output
		return x

layer_sizes = [3, 3, 3, 3, 3, 1]
sample_input = torch.tensor([[1., 0., -1.]])
torch.manual_seed(123) # specify random seed for the initial weights for reproducibility
model_without_shortcut = ExampleDeepNeuralNetwork(
	layer_sizes, use_shortcut=False
)
# 打印梯度
def print_gradients(model, x):
	# 前向过程
	output = model(x)
	target = torch.tensor([[0.]])
	# 计算损失
	loss = nn.MSELoss()
	loss = loss(output, target)
	# 后向过程计算梯度
	loss.backward()
	for name, param in model.named_parameters():
	if 'weight' in name:
		# 输出梯度的均值
		print(f"{name} has gradient mean of {param.grad.abs().mean().item()}")
# 没有shortcut连接的输出
print_gradients(model_without_shortcut, sample_input)

torch.manual_seed(123)
model_with_shortcut = ExampleDeepNeuralNetwork(
	layer_sizes, use_shortcut=True
)
# 有shortcut连接的输出
print_gradients(model_with_shortcut, sample_input)

构造Transformer Block

Transformer块的组成如下：

当transformer块处理输入序列时，序列中的每个元素(例如，单词或子单词token)由固定大小的向量(768维)表示。transformer块内的操作，包括多头注意力和前馈层，旨在以保留其维度的方式变换这些向量。

多头注意力块中的自注意力机制识别并分析输入序列中元素之间的关系，而前馈网络在每个位置单独修改数据。这种组合不仅能够更细致地理解和处理输入，而且还增强了模型处理复杂数据模式的整体能力。

Transformer块的代码为：

class TransformerBlock(nn.Module):
	def __init__(self, cfg):
		super().__init__()
		self.att = MultiHeadAttention(
			d_in=cfg["emb_dim"],
			d_out=cfg["emb_dim"],
			block_size=cfg["context_length"],
			num_heads=cfg["n_heads"],
			dropout=cfg["drop_rate"],
			qkv_bias=cfg["qkv_bias"])
		self.ff = FeedForward(cfg)
		self.norm1 = LayerNorm(cfg["emb_dim"])
		self.norm2 = LayerNorm(cfg["emb_dim"])
		self.drop_resid = nn.Dropout(cfg["drop_rate"])
	def forward(self, x):
		#A
		shortcut = x
		x = self.norm1(x)
		x = self.att(x)
		x = self.drop_resid(x)
		x = x + shortcut # short连接
		shortcut = x #B
		x = self.norm2(x)
		x = self.ff(x)
		x = self.drop_resid(x)
		x = x + shortcut #C
		return x

需要注意的是，在MultiHeadAttention和FeedForward之前应用层归一化(LayerNorm)，在它们之后应用dropout，以使模型规范化并防止过拟合。这种在之前应用LayerNorm的方式称为Pre-LayerNorm。在自注意力和前馈网络之后应用层归一化的方式称为Post-LayerNorm，这种可能导致不稳定的训练。

transformer块在其输出中保持输入尺寸，这表明transformer架构在处理数据序列时不会改变它们在整个网络中的形状。这种设计使其能够在广泛的序列到序列任务中有效应用，其中每个输出向量直接对应于一个输入向量，保持一对一的关系。

然而，输出是一个上下文向量，它封装了来自整个输入序列的信息。这意味着虽然序列的物理维度(长度和特征大小)在通过transformer块时保持不变，但每个输出向量的内容被重新编码，以整合整个输入序列的上下文信息。

构造GPT模型

GPT模型的整体结构为：

其中最后一个transformer块的输出在到达线性输出层之前经过最后一个层归一化步骤。线性输出层将transformer的输出映射到高维空间(在本例中，50,257维，对应于模型的词汇表大小)，以预测序列中的下一个token。

在之前代码的基础上，我们可以构造最终的GPT模型，代码为：

class GPTModel(nn.Module):
	def __init__(self, cfg):
		super().__init__()
		self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
		self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
		self.drop_emb = nn.Dropout(cfg["drop_rate"])
		self.trf_blocks = nn.Sequential(
			*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])
			self.final_norm = LayerNorm(cfg["emb_dim"])
			self.out_head = nn.Linear(
			cfg["emb_dim"], cfg["vocab_size"], bias=False
		)
	def forward(self, in_idx):
		batch_size, seq_len = in_idx.shape
		tok_embeds = self.tok_emb(in_idx)
		#A
		pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
		x = tok_embeds + pos_embeds
		x = self.drop_emb(x)
		x = self.trf_blocks(x)
		x = self.final_norm(x)
		logits = self.out_head(x)
		return logits

输出模型的参数量大小：

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
total_params = sum(p.numel() for p in model.parameters())
print(f"Total number of parameters: {total_params:,}")

最终输出的参数量会是163 million，跟124 million不符合。

这个原因是原始的GPT-2架构中使用了一个名为权重绑定的概念，这意味着原始的GPT-2架构正在重用来自token embedding层和输出层的权重。

移除掉输出层的参数量后，最终的参数量会和124 million一致：

total_params_gpt2 = total_params - sum(p.numel() for p in model.out_head.parameters())
print(f"Number of trainable parameters considering weight tying: {total_params_gpt2:,}")

使用GPT模型生成文本

LLM每次生成一个word，过程展示如下：

模型在每次迭代中预测一个后续token，并将其附加到输入上下文以进行下一轮预测。

在GPT生成文本过程中：

• 模型会输出一个矩阵，它表示可能的下一个单词的向量。
• 接着提取与下一个token对应的向量，并通过softmax函数转换为概率分布。
• 在包含结果概率分数的向量中，位于最高值的索引，它转换为token ID。
• 然后，这个token ID被解码回文本，生成序列中的下一个token。
• 最后，这个标记被添加到前面的输入中，形成一个新的输入序列用于后续的迭代。

具体过程如下：

该部分的代码为：

def generate_text_simple(model, idx, max_new_tokens, context_size): #A
	for _ in range(max_new_tokens):
		idx_cond = idx[:, -context_size:] #B
		with torch.no_grad():
			logits = model(idx_cond)
		logits = logits[:, -1, :] #C
		probas = torch.softmax(logits, dim=-1) #D
		idx_next = torch.argmax(probas, dim=-1, keepdim=True) #E
		idx = torch.cat((idx, idx_next), dim=1) #F
	return idx

该代码迭代生成指定数量的新token，裁剪当前上下文以适应模型的最大上下文大小，计算预测，然后根据最高概率预测选择下一个token。

我们使用softmax函数将logits转换为概率分布，并通过torch.argmax确定具有最大值的位置。实际上，softmax是多余的，因为logit的位置softmax的值也最大。这么做是为了说明将logits转换为概率的整个过程，这可以增加额外的直觉，例如模型生成最有可能的下一个token，这被称为greedy decoding。

调用GPT模型来预测下一个文本的代码为：

# 编码
start_context = "Hello, I am"
encoded = tokenizer.encode(start_context)
print("encoded:", encoded)
encoded_tensor = torch.tensor(encoded).unsqueeze(0) #A
print("encoded_tensor.shape:", encoded_tensor.shape)

# GPT模型
model.eval() #A
out = generate_text_simple(
model=model,
idx=encoded_tensor,
max_new_tokens=6,
context_size=GPT_CONFIG_124M["context_length"]
)
print("Output:", out)
print("Output length:", len(out[0]))
# 解码
decoded_text = tokenizer.decode(out.squeeze(0).tolist())
print(decoded_text)