从 0 开始手撸 LLM

从 0 开始编写LLM

大型语言模型（LLM）是一种以其实现通用语言理解和生成能力而闻名的语言模型。通过在计算密集型自监督和半监督训练过程中从文本文档中学习统计关系来获得这些能力。 LLM 是遵循 Transformer 架构的人工神经网络。

这是一个基于神经网络的模型训练，采用了一些新旧技术: tokenization, embedding, position encoding, feed-forward, normalization, softmax, linear transformation, multi-head attention。

下图阐述了整个模型，接下来以一个小文本来演练下

0、数据准备

安装如下工具包

pip install numpy requests torch tiktoken matplotlib pandas

导入对应的包

# 导入对应的包

import os

import requests

import pandas as pd

import matplotlib.pyplot as plt

import math

import tiktoken

import torch

import torch.nn as nn

这里采用sales_textbook.txt的数据作为训练数据和验证数据。把文本数据加载进入内存

with open('sales_textbook.txt', 'r', encoding='utf-8') as f:

   text = f.read()

一些全局设置

# 设置一些参数

batch_size = 4  # 批次

context_length = 16  # token 长度，最长 16 个单词

d_model = 64  # 维度

num_layers = 8  # Number of transformer blocks ？？？

num_heads = 4  # 多头 # 我们的代码中通过 d_model / num_heads = 来获取 head_size

1、Tokenization

# 使用 tiktoken 工具库进行处理，该工具是openAI 提供的又快又轻量级，基于原始单词

encoding = tiktoken.get_encoding("cl100k_base")

tokenized_text = encoding.encode(text) # 整个文本的单词数量 77,919

vocab_size = len(set(tokenized_text)) # 单词数：相当于哈希表里面的字典 3,771

max_token_value = max(tokenized_text) # 每个单词对应一个数值，这里找到最大值，100069

print(f"文本单词数量: {len(tokenized_text)}")

print(f"词汇表数量: {vocab_size}")

print(f"最大的单词 token 对应的编码值: {max_token_value}")

2、词向量Word Embedding

# 2、词向量

# 把数据转成张量形式

tokenized_text=torch.tensor(tokenized_text, dtype=torch.long)

# 首先吧数据集分成训练数据和验证数据（8:2）

# Split train and validation

split_idx = int(len(tokenized_text) * 0.8) # 取出 80%的位置

train_data = tokenized_text[:split_idx]  #训练数据

val_data = tokenized_text[split_idx:]    #验证数据

# 准备训练数据

data = train_data

#随机取4个数字，范围 0-(77919*0.8-16)，为了构建批量数据

idxs = torch.randint(low=0, high=len(data) - context_length, size=(batch_size,))

print(idxs)

# x 是一个 4*16 的结构，因为 idxs 有 4 个，每个里面都是 16 的长度

# y比 x 移后一位

x_batch = torch.stack([data[idx:idx + context_length] for idx in idxs])

y_batch = torch.stack([data[idx + 1:idx + context_length + 1] for idx in idxs])

print(x_batch.shape,y_batch.shape)

# 看看 x_batch y_batch是什么？

import pandas as pd

pd.DataFrame(x_batch[0].numpy())

#看下对应的是什么

encoding.decode([15749])

#看看整个句子是什么

encoding.decode(x_batch[0].numpy())

到此，我们的数据准备已经完成。开始进入模型领域。

3、进入把数字转成输入向量

# 把 16 个数字转换成 64 维度（列）

# 找到最大的 token 值对应的数字

max_token_value=tokenized_text.max().item()

max_token_value

# 需要构造一个 max_token_value*64 维度的矩阵

input_embedding_lookup_table = torch.nn.Embedding(max_token_value+1, d_model)

#看下初始化权重信息，这些数据是在不断更新的

# 把 x_batch 与权重绑定

x_batch_embedding = input_embedding_lookup_table(x_batch)

y_batch_embedding = input_embedding_lookup_table(y_batch)

# 打印看看

x_batch_embedding.shape

torch.Size([4, 16, 64])

4、位置信息

# 先构建一个全 0编码 的 16*64 的矩阵，先把形状搞出来，16*64 方便与x_batch_embedding的形状进行@运算

position_encoding_lookup_table = torch.zeros(context_length, d_model)

# torch.arange形成一维16个数，然后扩展第二位

position = torch.arange(0, context_length, dtype=torch.float).unsqueeze(1)  # unsqueeze函数主要是对数据维度进行扩充。

#print(torch.arange(0, context_length, dtype=torch.float))

#print(position)

# 计算正玄余玄

div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

position_encoding_lookup_table[:, 0::2] = torch.sin(position * div_term) #偶数位

position_encoding_lookup_table[:, 1::2] = torch.cos(position * div_term) #奇数位

position_encoding_lookup_table = position_encoding_lookup_table.unsqueeze(0).expand(batch_size, -1, -1) #add batch to the first dimension

print("Position Encoding Look-up Table: ", position_encoding_lookup_table.shape)

pd.DataFrame(position_encoding_lookup_table[0].numpy())

# 最终输入 需要输入与位置信息加和

x=x_batch_embedding+position_encoding_lookup_table

y=y_batch_embedding+position_encoding_lookup_table

# 看看此时形状

x.shape,y.shape

(torch.Size([4, 16, 64]), torch.Size([4, 16, 64]))

pd.DataFrame(x[0].detach().numpy())查看下数据

# 上面的表格16 行代表 16 个单词，64 列代表 64 维度，里面的数字是权重

# 目前拿到给 transform 模型所需要的全部输入

# 顺便看看x_batch_embedding

pd.DataFrame(x_batch_embedding[0].detach().numpy())

5、Transform 核心区域

包括Q(query)、K(key)、V(Value)，多头机制等,Wq、Wk、Wv 是一个 64*64 的矩阵，与我们输入的x 进行相乘，得到 Q、K、V

# 准备 Wq、Wk、Wv，也是权重矩阵

Wq=nn.Linear(d_model,d_model)  # 64*64

Wk=nn.Linear(d_model,d_model)

Wv=nn.Linear(d_model,d_model)

# 此处相乘是 x 的后两维与 Wq 等相乘   16*64 X   64*64 = 16*64

Q=Wq(x)  # Wq*x

K=Wk(x)

V=Wv(x)

#得到 4*16*64 的 QKV

Q.shape

torch.Size([4, 16, 64])

6、多头机制

# 多头机制，为了把64 个维度进行拆分多份，每一个头里面有一部分维度，这里是 4 个头，分别进行部分计算，然后合并

#先拆 Q，里面第一维度第二维度不变，分别是 4，,16，增加第三维度4，第四维度从之前的第三维 64 变成 16；这样就变成（4，,16，4，,16）；

Q=Q.reshape(batch_size,context_length,num_heads,d_model//num_heads)

K=K.reshape(batch_size,context_length,num_heads,d_model//num_heads)

V=V.reshape(batch_size,context_length,num_heads,d_model//num_heads)

# 在把二三维度转换，变成（4，,4，,16，,16）因为接下来需要做注意力机制，关注的是 token，而不是多头

Q = Q.transpose(1, 2) # [4, 4, 16, 16]

K = K.transpose(1, 2) # [4, 4, 16, 16]

V = V.transpose(1, 2) # [4, 4, 16, 16]

7、注意力机制

# 注意力机制

# K 转置

output=Q @ K.transpose(-2,-1)/math.sqrt(d_model//num_heads)

#做 softmax之前做 mask 部分，一句话比如中国人民，在第三个字的时候，他只知道中国人，不知道第四位是民这个字，所以需要处理下，把未来的设置成 0

#形成一个三角区域

mask=torch.triu(torch.ones((context_length,context_length)),diagonal=1).bool()

output=output.masked_fill(mask,float('-inf'))

pd.DataFrame(output[0][0].detach().numpy())

8、Softmax

# 继续 softmax,把概率数值转换成 0-1 之间的百分比

attention_score=torch.softmax(output,dim=-1) #[4, 4, 16, 16] [batch_size, num_heads, context_length, context_length]

attention_score.shape

torch.Size([4, 4, 16, 16])

9、Attention 积分

# attention_score @ V 的操作

A=attention_score@V

#torch.Size([4, 4, 16, 16])

print(A.shape)

#合并多头,之前把 64 拆成 4*16，改变了位置，这边需要换回来

A=A.permute(0,2,1,3).reshape(batch_size,context_length,d_model)

print(A.shape)

torch.Size([4, 4, 16, 16])torch.Size([4, 16, 64])

10、Wo 操作

# 定义 Wo

Wo=nn.Linear(d_model,d_model)

output=Wo(A)

11、残差链接

此时多头注意力已经结束，需要做个残差链接，也很简单，只需要和原始数据求和即可

# 残差链接

# 从多头出来和原始数据进行

output=output+x

12、进入层归一化

# 层归一化

layer_norm=nn.LayerNorm(d_model)

layer_norm_output=layer_norm(output)

#前馈网络：先把维度放大，然后做激活函数，然后再维度缩回

output=nn.Linear(d_model,d_model*4)(layer_norm_output )

output=nn.ReLU()(output)

output=nn.Linear(d_model*4,d_model)(output)

output=output+layer_norm_output

13、再进入层归一化

#再进行一层归一化

output=layer_norm(output)

14、线性变化

# 最终一个线性变化,此时需要把所有的 token 作为最终预测，会得到一个分值，选择最大的那一个

output=nn.Linear(d_model,max_token_value+1)(output)

output.shape

torch.Size([4, 16, 100070])

15、最后的 Softmax

probabilities = torch.softmax(output, dim=-1)

pd.DataFrame(probabilities[0].detach().cpu().numpy())

16、Test

predicted_index = torch.argmax(output[0,0]).item()

encoding.decode([predicted_index])

#看看句子

encoding.decode(x_batch[0].tolist())

' Jenny'

' professionals can overcome price objections effectively.\n2. Competition Objections: Another common objection'