【LLM From Scratch】4-在无标签数据上进行预训练
wbfwonderful Lv5

image

评估文本生成模型

在训练之前,模型会生成随机的下一个词元的概率向量。模型训练的目标是确保与图中框出的目标词元 ID 对应的概率值被最大化。

image

训练大语言模型的目标是最大化正确词元的可能性,这涉及增大其相对于其他词元的概率。通过这种方式,可以确保大语言模型始终选择目标词元(实质上是句子中的下一个单词)作为它生成的下一个词元。

训练大语言模型

简单的代码框架如下:

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
def train_model_simple(model, train_loader, val_loader, optimizer, device, num_epochs,
eval_freq, eval_iter, start_context, tokenizer):
# Initialize lists to track losses and tokens seen
train_losses, val_losses, track_tokens_seen = [], [], []
tokens_seen, global_step = 0, -1
#初始化训练模型而且给了空的队列
# Main training loop
for epoch in range(num_epochs):#训练次数
model.train() # Set model to training mode
#转移到训练模块
for input_batch, target_batch in train_loader:
#从loader里面调出输入跟目标
optimizer.zero_grad() # Reset loss gradients from previous batch iteration
#清空所有函数的梯度
loss = calc_loss_batch(input_batch, target_batch, model, device)
#计算损失函数
loss.backward() # Calculate loss gradients
#反向传播优化
optimizer.step() # Update model weights using loss gradients
#更新权重
tokens_seen += input_batch.numel()
#加一下一共有多少
global_step += 1
#看一下一共训练了多少步
# Optional evaluation step
if global_step % eval_freq == 0:
#按照一定的步数进行记录
train_loss, val_loss = evaluate_model(
model, train_loader, val_loader, device, eval_iter)
#计算损失函数
train_losses.append(train_loss)
val_losses.append(val_loss)
track_tokens_seen.append(tokens_seen)
#加到list中
print(f"Ep {epoch+1} (Step {global_step:06d}): "
f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")

# Print a sample text after each epoch
generate_and_print_sample(
model, tokenizer, device, start_context
)

return train_losses, val_losses, track_tokens_seen


def evaluate_model(model, train_loader, val_loader, device, eval_iter):
#评价模块
model.eval()
#检验模式
with torch.no_grad():
#我认为的双保险,防止梯度更新
train_loss = calc_loss_loader(train_loader, model, device, num_batches=eval_iter)
val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)
model.train()
# 在评估结束后切换回训练模式,确保模型能继续用于训练。
return train_loss, val_loss


def generate_and_print_sample(model, tokenizer, device, start_context):
model.eval()
context_size = model.pos_emb.weight.shape[0]
encoded = text_to_token_ids(start_context, tokenizer).to(device)
with torch.no_grad():
token_ids = generate_text_simple(
model=model, idx=encoded,
max_new_tokens=50, context_size=context_size
)
decoded_text = token_ids_to_text(token_ids, tokenizer)
print(decoded_text.replace("\n", " ")) # Compact print format
model.train()

控制随机性的解码策略

上一章末尾定义的 generate_text_simple 函数,生成的词元是从词汇表的所有词元中选择概率分数最大的那一个。这意味着,即使在相同的起始上下文(Every effort moves you)中多次运行前面的 generate_text_simple 函数,大语言模型也将始终生成相同的输出。代码如下:

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
def generate_text_simple(model, idx, max_new_tokens, context_size):
# 预测单词的模块
# idx 是当前上下文中的(batch, n_tokens)索引数组
for _ in range(max_new_tokens):
# 每次生成一个单词后,重新将其加入序列中
# 如果当前上下文长度超过模型支持的最大上下文长度,则截取
# 例如,如果LLM只支持5个token,而上下文长度为10
# 那么只使用最后5个token作为上下文
idx_cond = idx[:, -context_size:]
# 如果idx的长度超过模型支持的上下文长度size,只保留最后size个token
# 避免溢出
# 获取预测结果
with torch.no_grad(): # 在推理阶段,不需要计算梯度,因为没有反向传播
# 这样可以减少存储开销
logits = model(idx_cond)
# 模型输出结果
# 只关注最后一个时间步的输出
# (batch, n_tokens, vocab_size) 变为 (batch, vocab_size)
logits = logits[:, -1, :]
# 关注最后一个时间步
# 使用softmax函数计算概率
probas = torch.softmax(logits, dim=-1) # (batch, vocab_size)
# 归一化
# 获取具有最高概率值的词汇索引
idx_next = torch.argmax(probas, dim=-1, keepdim=True) # (batch, 1)
# 获取概率最高的词汇索引
# 将采样的索引添加到序列中
idx = torch.cat((idx, idx_next), dim=1) # (batch, n_tokens+1)

return idx

温度缩放(Temperature)

前面的方法是使用 torch.argmax(也称为贪婪解码)来采样具有最高概率的词元作为下一个词元。为了生成更多样化的文本,可以用一个从概率分布(这里是大语言模型在每个词元生成步骤为每个词汇条目生成的概率分数)中采样的函数来取代 argmax。为了实现一个概率采样过程,现在可以用 PyTorch 中的 multinomial 函数替换 argmax

1
next_token_id = torch.multinomial(probas, num_samples=1).item()

torch.multinomial 是 PyTorch 中用于多项分布采样(Multinomial Sampling)的函数,常用于根据概率分布从离散集合中随机抽取索引。它的核心功能是:给定一组权重(或概率),按照这些权重的比例随机采样若干个样本索引。
这里设置了 num_samples=1,所以本质上是采样,而不是选最大

通过一个称为温度缩放的概念,可以进一步控制分布和选择过程。温度缩放指的是将 logits 除以一个大于 0 的数。所有 logits 都除以相同的数(temperature),但 softmax 是一个非线性函数,而非线性变换对数值缩放的敏感性,使得 “都除以相同的数字” 并不会让比例保持完全相同。

image

一个例子:从结果中可以看出,当温度设置为 0.1 时,概率分布变得更加陡峭,接近于 torch.argmax 的行为,因此最可能的 token 几乎总是被选中

image

Top-k 采样

较高的温度值会导致下一个词元的概率分布更均匀,从而产生更多样化的输出,因为它降低了模型重复选择最可能词元的可能性。这种方法允许探索概率较低但可能更具创造性和趣味性的生成路径。然而,这种方法的一个缺点是,它有时会导致语法不正确或完全无意义的输出。

通过与概率采样和温度缩放相结合,Top-k 采样可以改善文本生成结果。在 Top-k 采样中,可以将采样的词元限制在前 k 个最可能的词元上,并通过掩码概率分数的方式来排除其他词元。

Top-k 方法用负无穷值(-inf)替换所有未选择的 logits,因此在计算 softmax 值时,非前 k 词元的概率分数为 0,剩余的概率总和为 1。

image

新的文本生成函数

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
def generate(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):
#生成模块
# For-loop is the same as before: Get logits, and only focus on last time step
for _ in range(max_new_tokens):
idx_cond = idx[:, -context_size:]
with torch.no_grad():
logits = model(idx_cond)
logits = logits[:, -1, :]
#计算预测值,但是切最后一个
# New: Filter logits with top_k sampling
#top K采样
if top_k is not None:
# Keep only top_k values
top_logits, _ = torch.topk(logits, top_k)
min_val = top_logits[:, -1]
logits = torch.where(logits < min_val, torch.tensor(float("-inf")).to(logits.device), logits)

# New: Apply temperature scaling
#温度校正
if temperature > 0.0:
logits = logits / temperature

# Apply softmax to get probabilities
probs = torch.softmax(logits, dim=-1) # (batch_size, context_len)

# Sample from the distribution
idx_next = torch.multinomial(probs, num_samples=1) # (batch_size, 1)
#从概率分布中采样下一个 token

# Otherwise same as before: get idx of the vocab entry with the highest logits value
else:
idx_next = torch.argmax(logits, dim=-1, keepdim=True) # (batch_size, 1)
#如果未启用采样,选择概率最高的 token 作为下一个 token
if idx_next == eos_id: # Stop generating early if end-of-sequence token is encountered and eos_id is specified
break

# Same as before: append sampled index to the running sequence
idx = torch.cat((idx, idx_next), dim=1) # (batch_size, num_tokens+1)

return idx

可以理解为:

  • 温度校正是更加平滑,防止数据差之毫厘以谬以千里
  • topK 是防止臭鱼烂虾进入筛选范围提高质量

加载和保存模型

保存

1
torch.save(model.state_dict(), "model.pth")

加载

1
2
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(torch.load("model.pth", map_location=device))

像 AdamW 这样的自适应优化器可以为每个模型权重存储额外的参数。AdamW 可以使用历史数据动态地调整每个模型参数的学习率。如果没有它,那么优化器就会重置,模型可能学习效果不佳,甚至无法正确收敛,这意味着模型将失去生成连贯文本的能力。可以使用 torch.save 保存模型和优化器的 state_dict 内容:

1
2
3
4
5
6
7
torch.save({
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
"model_and_optimizer.pth"
)
#全家整整齐齐地保存

加载优化器和模型的参数:

1
2
3
4
5
6
7
8
9
checkpoint = torch.load("model_and_optimizer.pth", weights_only=True)
#保存检查点
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])

optimizer = torch.optim.AdamW(model.parameters(), lr=0.0005, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train();
#调整到训练模式

加载 GPT-2 权重

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
import numpy as np

def load_weights_into_gpt(gpt, params):
gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params['wpe'])
#位置权重
gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params['wte'])
#单词全中
for b in range(len(params["blocks"])):
#三个参数的输入
q_w, k_w, v_w = np.split(
(params["blocks"][b]["attn"]["c_attn"])["w"], 3, axis=-1)
gpt.trf_blocks[b].att.W_query.weight = assign(
gpt.trf_blocks[b].att.W_query.weight, q_w.T)
gpt.trf_blocks[b].att.W_key.weight = assign(
gpt.trf_blocks[b].att.W_key.weight, k_w.T)
gpt.trf_blocks[b].att.W_value.weight = assign(
gpt.trf_blocks[b].att.W_value.weight, v_w.T)

q_b, k_b, v_b = np.split(
(params["blocks"][b]["attn"]["c_attn"])["b"], 3, axis=-1)
gpt.trf_blocks[b].att.W_query.bias = assign(
gpt.trf_blocks[b].att.W_query.bias, q_b)
gpt.trf_blocks[b].att.W_key.bias = assign(
gpt.trf_blocks[b].att.W_key.bias, k_b)
gpt.trf_blocks[b].att.W_value.bias = assign(
gpt.trf_blocks[b].att.W_value.bias, v_b)

gpt.trf_blocks[b].att.out_proj.weight = assign(
gpt.trf_blocks[b].att.out_proj.weight,
params["blocks"][b]["attn"]["c_proj"]["w"].T)
gpt.trf_blocks[b].att.out_proj.bias = assign(
gpt.trf_blocks[b].att.out_proj.bias,
params["blocks"][b]["attn"]["c_proj"]["b"])

gpt.trf_blocks[b].ff.layers[0].weight = assign(
gpt.trf_blocks[b].ff.layers[0].weight,
params["blocks"][b]["mlp"]["c_fc"]["w"].T)
gpt.trf_blocks[b].ff.layers[0].bias = assign(
gpt.trf_blocks[b].ff.layers[0].bias,
params["blocks"][b]["mlp"]["c_fc"]["b"])
gpt.trf_blocks[b].ff.layers[2].weight = assign(
gpt.trf_blocks[b].ff.layers[2].weight,
params["blocks"][b]["mlp"]["c_proj"]["w"].T)
gpt.trf_blocks[b].ff.layers[2].bias = assign(
gpt.trf_blocks[b].ff.layers[2].bias,
params["blocks"][b]["mlp"]["c_proj"]["b"])

gpt.trf_blocks[b].norm1.scale = assign(
gpt.trf_blocks[b].norm1.scale,
params["blocks"][b]["ln_1"]["g"])
gpt.trf_blocks[b].norm1.shift = assign(
gpt.trf_blocks[b].norm1.shift,
params["blocks"][b]["ln_1"]["b"])
gpt.trf_blocks[b].norm2.scale = assign(
gpt.trf_blocks[b].norm2.scale,
params["blocks"][b]["ln_2"]["g"])
gpt.trf_blocks[b].norm2.shift = assign(
gpt.trf_blocks[b].norm2.shift,
params["blocks"][b]["ln_2"]["b"])

gpt.final_norm.scale = assign(gpt.final_norm.scale, params["g"])
gpt.final_norm.shift = assign(gpt.final_norm.shift, params["b"])
gpt.out_head.weight = assign(gpt.out_head.weight, params["wte"])

#主要目的是将预训练的模型参数加载到一个gpt中
load_weights_into_gpt(gpt, params)
gpt.to(device);

默认情况下,GPTModel 实例使用随机权重初始化以进行预训练。使用 OpenAI 的模型权重的最后一步是用加载到 params 字典中的权重覆盖这些随机权重。为此,首先需要定义一个小的 assign 工具函数,该函数会检查两个张量或数组(left 和right)是否具有相同的维度或形状,并将 right 张量返回为可训练的 PyTorch 参数。