一顆小胡椒
自從chatgpt發布后一直關注如何和網絡安全融合，最近一直在學習GPT相關內容，傳統的用文本向量相似度索引的方式局限太大。
所以想到自己訓練或微調，受限于算力，很多地方都不好驗證，于是想做一個本地能跑起來的GPT學習研究。
GPT的發展史
- 2018年 GPT1發布，作者用5G的書籍文本做無監督學習預訓練，架構為12層transformer，每層12個注意頭，模型參數約1.17億，使用生成的預訓練模型再到微調運用到各類nlp任務。
- 2019年GPT2發布，對gpt1架構改動了一下，使用40GB文本，包含web爬蟲數據，redit等數據，模型參數達到15億。
- 提出了不修改預訓練模型的情況下，使用0樣本或少樣本學習，完成任務
- 2020年GPT3發布
- 和GPT2的區別就是將數據和模型都擴大了100倍，暴力出奇跡。使用45T文本，參數量達到1750億，效果炸裂。
- in-context learning
- GPT3.5
- 訓練獎勵模型，讓模型更偏向人類的思考方式
- 使用監督微調的方式，提供大量人類對話的例子，讓機器模仿人類
- GPT3很強大，但機器只是試圖完成文章，并不是“助手”
- SFT supervised fine-tuning
- RLHF align human
- GPT4
- 多模態混合模型
- GPT-4每個head都有2200億參數，是一個8路的混合模型，總參數達到1.6萬億
- 在GPT4論文有一個有意思的點，因為每次訓練都相當于黑盒，訓練的代價又過于昂貴，擔心loss不下降，所以GPT4先訓練了一個參數低100倍的小模型，基于這個模型用機器學習預測了GPT4模型量的loss值。
GPT是為了解決廣泛的nlp的任務所以才會在數據集和模型參數上不斷加倍，如果只是對一個垂直領域數據做問答和推理，是否可以用一個小模型達到效果。
數據收集
小模型整個訓練都是在Google colab上完成，免費提供的顯存大小只有16G，實際可用在13～15之間，后面很多地方受限于顯存大小，所以有些地方實現會非常簡單，后面會說到。
在數據的收集上，先進行一遍無監督學習，選取了seebug paper和一個poc倉庫
- https://github.com/Threekiii/Awesome-POC
- https://paper.seebug.org/
這只是一個簡單的測試，跑通后后面自然可以增大數據集
每個文章按1024大小進行分割，保存到json文件中，最后數據大小有31M。
Ps：按塊分割會造成很多信息不完整，數據收集這塊還是需要清洗后效果會更好。
數據處理
模型只是對數字進行計算，所以需要將文本轉換為文本向量，這里簡單的做法是將訓練集中每個字提取出來生成一個字表，字表的索引號就是該文本的向量。
最后生成的大小有4214。
Ps：這是簡單的做法，GPT的做法是使用 BPE（Byte Pair Encoding）算法處理，最后詞表有5w大小，詞表和顯存占用是線性關系所以用這個簡單的方法跑了。
數據集加載類
對每篇文章mask最后一個字用作預測，計算loss用mask第一個字的文本，gpt架構的神奇之處在于此，它只是預測最后一個字，而預測的這個字是根據學習文本的概率計算的。
# 定義數據集
class MyDataSet(Data.Dataset):
def __init__(self, datas):
self.datas = datas
def __getitem__(self, item):
data_item = self.datas[item]
decoder_input = data_item[:-1]
decoder_output = data_item[1:]
return {"decoder_input": decoder_input,
"decoder_output": decoder_output}
def padding_batch(self, batch): #
for d in batch: # 對當前batch的每一個decoder_input和decoder_output數據填充"<pad>"，填充到和batch里面的有的最大長度為止
input_len = len(d["decoder_input"])
output_len = len(d["decoder_output"])
d["decoder_input"].extend([special_char_pad] * (max_pos - input_len))
d["decoder_output"].extend([special_char_pad] * (max_pos - output_len))
decoder_inputs = torch.tensor([d["decoder_input"] for d in batch], dtype=torch.long) # 轉type
decoder_outputs = torch.tensor([d["decoder_output"] for d in batch], dtype=torch.long)
return decoder_inputs, decoder_outputs # 形狀[b,decoder_input_maxlen], [b,decoder_output_maxlen] type為torch.long
def __len__(self):
return len(self.datas)
訓練超參數
這里簡單訓練一個6層 8個注意頭的模型。
max_pos = 1024 # 一段話最多字 d_model = 768 # Embedding Size d_ff = 2048 # FeedForward dimension d_k = d_v = 64 # dimension of K(=Q), V n_layers = 6 # number of Encoder of Decoder Layer n_heads = 8 # number of heads in Multi-Head Attention
參數量在36M，及3600萬的小模型，顯存占用在11G左右，再大顯存就不夠了
模型層
模型使用是GPT2的結構
GPT( (decoder): Decoder( (tgt_emb): Embedding(6110, 768) (pos_emb): Embedding(1024, 768) (layers): ModuleList( (0-5): 6 x DecoderLayer( (dec_self_attn): MultiHeadAttention( (W_Q): Linear(in_features=768, out_features=512, bias=False) (W_K): Linear(in_features=768, out_features=512, bias=False) (W_V): Linear(in_features=768, out_features=512, bias=False) (fc): Linear(in_features=512, out_features=768, bias=False) (layernorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True) ) (pos_ffn): PoswiseFeedForwardNet( (fc): Sequential( (0): Linear(in_features=768, out_features=2048, bias=False) (1): ReLU() (2): Linear(in_features=2048, out_features=768, bias=False) ) (layernorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True) ) ) ) ) (projection): Linear(in_features=768, out_features=6110, bias=True) )
模型代碼
# 把數據里面<pad>對應的字符給mask掉，讓后面Q和K相似度矩陣的softmax中這些pad都為0，就不會被后續的V考慮
def get_attn_pad_mask(seq_q, seq_k): # 形狀都是[b, tgt_len <300]
batch_size, len_q = seq_q.size() # len_q = len_k = tgt_len
batch_size, len_k = seq_k.size()
# eq(zero) is PAD token.就是把數據里面<pad>對應的字符給mask掉，讓后面Q和K的softmax不考慮這些<pad>
pad_attn_mask = seq_k.data.eq(0).unsqueeze(
1) # [b, 1, tgt_len], id為0(也就是<pad>的id)的位置為True，其他位置為False。后面會把Ture位置的mask掉
return pad_attn_mask.expand(batch_size, len_q, len_k) # [b, tgt_len, tgt_len]
def get_attn_subsequence_mask(seq): # seq: [b, tgt_len]
attn_shape = [seq.size(0), seq.size(1), seq.size(1)] # [b, tgt_len, tgt_len]
subsequence_mask = np.triu(np.ones(attn_shape), k=1) # Upper triangular matrix(上三角矩陣)
subsequence_mask = torch.from_numpy(subsequence_mask).byte()
subsequence_mask = subsequence_mask.to(device)
return subsequence_mask # [b, tgt_len, tgt_len] 上三角矩陣,下0上1,dtype=torch.uint8
class ScaledDotProductAttention(nn.Module): # 計算Q和K的相似度矩陣，然后乘V。對應筆記里的圖
def __init__(self):
super(ScaledDotProductAttention, self).__init__()
def forward(self, Q, K, V,
attn_mask): # 前三者形狀相同[b, n_heads, tgt_len, d_k=64]，attn_mask:[b, n_heads, tgt_len, tgt_len]
scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k) # Q和K的相似度矩陣scores : [b, n_heads, tgt_len, tgt_len]
scores.masked_fill_(attn_mask, -1e9) # Fills elements of self tensor with value where mask is True.
# 就是scores矩陣里面和attn_mask=1對應位置的元素全部替換成-1e9，使其在下一步的softmax中變為0
attn = nn.Softmax(dim=-1)(scores) # [b, n_heads, tgt_len, tgt_len]
context = torch.matmul(attn, V) # [b, n_heads, tgt_len, d_v]
return context, attn
class MultiHeadAttention(nn.Module): # 多頭注意力機制
def __init__(self):
super(MultiHeadAttention, self).__init__()
self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False) # d_model=768 , d_v = d_k = 64 , n_heads=8
self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)
self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)
self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)
self.layernorm = nn.LayerNorm(d_model)
def forward(self, input_Q, input_K, input_V,
attn_mask): # 前三者形狀相同，都是[b, tgt_len, d_model] , attn_mask: [b, tgt_len, tgt_len]
residual, batch_size = input_Q, input_Q.size(0) #
# [b, tgt_len, d_model] --> [b, tgt_len, d_k * n_heads] -split-> (b, tgt_len, n_heads, d_k) -trans-> (b, n_heads, tgt_len, d_k)
Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2) # Q: [b, n_heads, tgt_len, d_k=64]
K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2) # K: [b, n_heads, tgt_len, d_k=64]
V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1, 2) # V: [b, n_heads, tgt_len, d_v=64]
attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1,
1) # 添加n_heads維度并復制。attn_mask : [b, n_heads, tgt_len, tgt_len]
context, attn = ScaledDotProductAttention()(Q, K, V, attn_mask) # 參考圖解，context形狀[b, n_heads, tgt_len, d_v]
context = context.transpose(1, 2).reshape(batch_size, -1, n_heads * d_v) # context: [b, tgt_len, n_heads * d_v]
output = self.fc(context) # [batch_size, tgt_len, d_model]
return self.layernorm(output + residual), attn
class PoswiseFeedForwardNet(nn.Module): # [b,tgt_len,d_model] -> [b,tgt_len,d_model] 輸入和輸出形狀不變
def __init__(self):
super(PoswiseFeedForwardNet, self).__init__()
self.fc = nn.Sequential(
nn.Linear(d_model, d_ff, bias=False),
nn.ReLU(),
nn.Linear(d_ff, d_model, bias=False)
)
self.layernorm = nn.LayerNorm(d_model)
def forward(self, inputs):
'''
inputs: [batch_size, seq_len, d_model]
'''
residual = inputs
output = self.fc(inputs)
return self.layernorm(output + residual) # [batch_size, seq_len, d_model]
class DecoderLayer(nn.Module):
def __init__(self):
super(DecoderLayer, self).__init__()
self.dec_self_attn = MultiHeadAttention() # 多頭注意力
# self.dec_enc_attn = MultiHeadAttention()
self.pos_ffn = PoswiseFeedForwardNet()
def forward(self, dec_inputs,
dec_self_attn_mask): # dec_inputs: [b, tgt_len, d_model] dec_self_attn_mask: [b, tgt_len, tgt_len]
# dec_outputs: [b, tgt_len, d_model], dec_self_attn: [b, n_heads, tgt_len, tgt_len]
dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask)
dec_outputs = self.pos_ffn(dec_outputs) # [b, tgt_len, d_model]
return dec_outputs, dec_self_attn # [b, tgt_len, d_model] , [b, n_heads, tgt_len, tgt_len]
class Decoder(nn.Module):
def __init__(self):
super(Decoder, self).__init__()
self.tgt_emb = nn.Embedding(vocab_size,
d_model) # 以矩陣形式抽取一行，會比直接用mlp高效。因為mlp會多很多無用運算 emb矩陣形狀(vocab_size,768)
self.pos_emb = nn.Embedding(max_pos, d_model) # 可學習的位置編碼 emb矩陣形狀(300,768)
self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])
def forward(self, dec_inputs): # 輸入dec_inputs形狀[b,tgt_len]
seq_len = dec_inputs.size(1) # tgt_len ，表示batch內最大長度，不會超過300
pos = torch.arange(seq_len, dtype=torch.long, device=device) # 給位編碼準備的值，[0,1,2,3,...,seq_len-1]
pos = pos.unsqueeze(0).expand_as(dec_inputs) # [tgt_len] -> [b, tgt_len]
dec_outputs = self.tgt_emb(dec_inputs) + self.pos_emb(pos) # [b, tgt_len, d_model=768]
dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs) # [b, tgt_len, tgt_len] 把<pad>給mask掉
dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs) # [b, tgt_len, tgt_len] 上三角矩陣
dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequence_mask),
0) # [b, tgt_len, tgt_len] 矩陣大于0的全為1，否則為0
dec_self_attns = []
for layer in self.layers:
# dec_outputs: [b, tgt_len, d_model], dec_self_attn: [b, n_heads, tgt_len, tgt_len], dec_enc_attn: [b, h_heads, tgt_len, src_len]
dec_outputs, dec_self_attn = layer(dec_outputs, dec_self_attn_mask)
dec_self_attns.append(dec_self_attn)
return dec_outputs, dec_self_attns
class GPT(nn.Module):
def __init__(self):
super(GPT, self).__init__()
self.decoder = Decoder()
self.projection = nn.Linear(d_model, vocab_size) # 768->vocab_size,也就是把最后的隱藏層節點768投影到字典個數的節點上
def forward(self, dec_inputs): # 輸入dec_inputs形狀[b,tgt_len] tgt_len<=300 (tgt_len是batch內最大長度)
dec_outputs, dec_self_attns = self.decoder(
dec_inputs) # dec_outpus: [b, tgt_len, d_model=768], dec_self_attns: [n_layers, b, n_heads, tgt_len, tgt_len]
dec_logits = self.projection(dec_outputs) # dec_logits: [b, tgt_len, vocab_size]
return dec_logits.view(-1, dec_logits.size(-1)), dec_self_attns # 左邊那個輸出形狀[b *tgt_len,vocab_size]
@torch.no_grad()
def generate(self, sentence, max_new_tokens, temperature=1.0, top_k=None):
"""
Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
the sequence max_new_tokens times, feeding the predictions back into the model each time.
Most likely you'll want to make sure to be in model.eval() mode of operation for this.
"""
idx = torch.tensor(encoder(sentence), dtype=torch.long, device=device).unsqueeze(
0) # [n] -> [1,n] 轉type，并放入指定設備
for _ in range(max_new_tokens):
# forward the model to get the logits for the index in the sequence
dec_outputs, _ = self.decoder(idx)
logits = self.projection(dec_outputs) # [1, tgt_len, vocab_size]
# pluck the logits at the final step and scale by desired temperature
logits = logits[:, -1, :] / temperature
# optionally crop the logits to only the top k options
if top_k is not None:
vv, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < vv[:, [-1]]] = -float('Inf')
# apply softmax to convert logits to (normalized) probabilities
probs = F.softmax(logits, dim=-1)
# sample from the distribution
# idx_next = torch.multinomial(probs, num_samples=1)
idx_next = torch.max(probs, dim=-1, keepdim=True)[1]
# append sampled index to the running sequence and continue
if idx_next.item() == special_char_sep:
break
idx = torch.cat(
[idx.detach(), idx_next], -1)
yield vocab_data_reverse[idx_next.item()]
訓練
訓練時額外的超參數
batch_size = 16 # 一次訓練多少個文本 epochs = 20 # 訓練幾輪 lr = 1e-4 # 學習率
訓練代碼
# 模型的訓練
import glob
import math
import time
from torch import optim
from tqdm import tqdm
def epoch_time(start_time, end_time): # 把秒數表示為分鐘和秒
elapsed_time = end_time - start_time
elapsed_mins = int(elapsed_time / 60)
elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
return elapsed_mins, elapsed_secs
def train_step(model, data_loader, optimizer, criterion, clip=1, print_every=None): # 每一個eopch的訓練
model.train() # 訓練模式
if print_every == 0:
print_every = 1
print_loss_total = 0 # 每次打印都重置，統計一定batch數內(默認10)的loss，每10個batch打印一次
epoch_loss = 0 # epoch的總loss
for i, (dec_inputs, dec_outputs) in enumerate(
tqdm(data_loader)): # dec_inputs: [b, tgt_len] , dec_outputs: [b, tgt_len]
optimizer.zero_grad()
dec_inputs, dec_outputs = dec_inputs.to(device), dec_outputs.to(device)
# outputs: [batch_size * tgt_len, tgt_vocab_size] tgt_len<=30
# with torch.cuda.amp.autocast(): # 半精度訓練
outputs, dec_self_attns = model(dec_inputs)
loss = criterion(outputs, dec_outputs.view(
-1)) # outputs :(b * tgt_len, vocab_size),dec_outputs.view(-1) :(b * tgt_len) tgt_len<=300
print_loss_total += loss.item()
epoch_loss += loss.item()
loss.backward() # 梯度反向傳播
# 梯度裁剪，防止梯度爆炸。如果loss超過clip，將梯度值縮小為原來的(loss/clip)分之一
torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
optimizer.step() # 更新模型權重
if print_every and (i + 1) % print_every == 0:
print_loss_avg = print_loss_total / print_every
print_loss_total = 0
print('\tCurrent Loss: %.4f' % print_loss_avg)
return epoch_loss / len(data_loader)
def train(model, data_loader, lr):
criterion = nn.CrossEntropyLoss(ignore_index=0).to(device) # 損失函數
optimizer = optim.AdamW(model.parameters(), lr=lr) # 優化器
for epoch in range(epochs):
start_time = time.time()
train_loss = train_step(model, data_loader, optimizer, criterion, CLIP, print_every=100) # 訓練一個epoch
end_time = time.time()
torch.save(model.state_dict(), r'model/GPT-%d.pt' % epoch) # 保存模型權重
epoch_mins, epoch_secs = epoch_time(start_time, end_time) # 把秒數表示為分鐘和秒
print(f'Epoch: {epoch + 1:02} | Time: {epoch_mins}m {epoch_secs}s')
print(f'\tTrain Loss: {train_loss:.3f}')
def print_num_parameters(model):
# Find total parameters and trainable parameters
total_params = sum(p.numel() for p in model.parameters())
print("number of parameters: %.2fM" % (total_params / 1e6,))
total_trainable_params = sum(
p.numel() for p in model.parameters() if p.requires_grad)
print("train of parameters: %.2fM" % (total_trainable_params / 1e6))
def split_array(array, num):
length = len(array)
chunk_size = math.ceil(length / num)
result = []
for i in range(0, chunk_size):
result.append(array[:num])
array = array[num:]
return result
def get_dataset_from_mk(folder):
dataset = []
for filename in glob.glob(folder + "*.md"):
with open(filename) as f:
data = f.read()
array_ = split_array(encoder(data), max_pos)
dataset.extend(array_)
return dataset
def get_dataset_from_json(filename):
with open(filename) as f:
data = json.load(f)
dataset = []
for item in data:
dataset.append(encoder(item))
return dataset
if __name__ == '__main__':
batch_size = 16
epochs = 10
shuffle = True
lr = 1e-4
filename = "data.json"
dataset = get_dataset_from_json(filename)
data_set = MyDataSet(dataset)
data_loader = Data.DataLoader(data_set,
batch_size=batch_size,
collate_fn=data_set.padding_batch,
shuffle=shuffle) # 對每個batch單獨調用collate_fn處理，因為batch內的句子長短不一，不能直接用torch的默認方法
model = GPT().to(device)
print_num_parameters(model)
train(model, data_loader, lr)
可以看到模型的參數大概是36M，在colab 使用T4訓練，一個Epoch在37分鐘。在訓練了三輪后停了，loss在1.8。
資源消耗
測試推理
epoch跑了三輪后
epoch跑了18輪后，loss降低到了0.9，在本地測試，隨機跑6個結果
語句不那么通順，還是有點胡言亂語。這和很多因素有關，有可能作為訓練的文本太少，機器還沒喂飽。不過它的優點是能在家用CPU推理，速度很快。
總結
這次實驗只是跑通了代碼，對效果還沒有調優，針對垂直領域的語言模型，應該還是需要在大的通用語言文本上訓練一次，理解世界，然后再針對垂直領域的數據繼續finetuning。如果這個猜想能成立的話，它就能在家用CPU上推理，成本大幅降低。
實驗到這里還引出了一個問題，不同參數的模型，訓練文本的極限是多少，小參數的模型能達到多少水平，這還需要一些算力做實驗，請見下篇文章。
D1Net
公安部網安局
GoUpSec
安全圈
D1Net
中國信息安全
奇安信集團
黑白之道
CNCERT國家工程研究中心
安全內參
信息安全與通信保密雜志社
信息安全與通信保密雜志社
