从零开始训练大模型—以RoBERTa为例_roberta 预训练源码

---

0 前言

本文讲述了如何从零开始训练一个大模型，这个从零开始值是指从源码层面自己处理数据、搭建模型。

1 数据预处理

• 加载数据

• 对过长的文本进行切分，设置max_len=128，也就是先按照标点符号把长句子切分为短句子，然后将短句子进行组合，但是不能超过最大长度，超过的部分组合成新的句子，所有数据处理完得到 Dataset。
  

• 将Dataset传给DataLoader，在DataLoader里面的collate_fn进行操作，得到模型的输入。
  collate_fn函数：
  1）对输入的文本进行编码：input_ids = tokenizer.encode(list(text))
  2）计算需要掩码的token数量：n_pred = min(max_pred, max(1, int(len(input_ids) * 0.15)))
  3）对input_ids进行处理得到待掩码的token：cand_maked_pos = [i for i, token in enumerate(input_ids) if token != word2idx[‘[CLS]’] and token != word2idx[‘[SEP]’]]
  4）打乱待掩码的token：shuffle(cand_maked_pos)
  5）进行掩码：

  for pos in cand_maked_pos[:n_pred]：
  	masked_pos.append(pos) 
  	masked_tokens.append(input_ids[pos])
  	if random() < 0.8: 
  		input_ids[pos] = word2idx['[MASK]']
  	elif random() > 0.9:
  		index = randint(0, vocab_size - 1) 
  		while index == 0 or index == 101 or index == 102 or index == 103: 
  			index = randint(0, vocab_size - 1)
  	input_ids[pos] = index
  1
  2
  3
  4
  5
  6
  7
  8
  9
  10

  6）对input_ids进行补零操作：
  n_pad = maxlen - len(input_ids)
  input_ids.extend([0] * n_pad)
  7）对掩码的token和对应的位置向量进行补零操作：
  n_pad = max_pred - n_pred
  masked_tokens.extend([0] * n_pad)
  masked_pos.extend([0] * n_pad)
  8）将得到的模型输入input_ids、masked_tokens、 masked_pos转换为torch tensor格式：
  torch.tensor(input_ids, dtype=torch.long)
  torch.tensor(masked_tokens, dtype=torch.long)
  torch.tensor(masked_pos, dtype=torch.long)

2 模型构建

2.1 模型介绍

1. RoBERTa是基于BERT进行改进得到的， RoBERTa 相较于 BERT 最大的改进有三点：
   1）动态 Masking： BERT的masking是在预处理时进行的，导致这种Masking是静态的，每个epoch的masking结果一致。而RoBERTa中使用Dynamic Masking，只是在序列送入模型中的时候才去进行动态的masking，这样在更大的数据集上或者更多步数的训练上会表现更好
   2）取消 NSP (Next Sentence predict) 任务：为了探索NSP训练策略对模型结果的影响，论文中设置了4种训练方式进行对比，最后发现没有NSP任务模型的训练结果会更好，下游任务的效果也会更好
   3）扩大 batch_size：论文中通过实验，证明了更大的batch_size可以得到更好的结果
2. RoBERTa的结构为：Embedding层；EncoderLayer层；全连接层
   1）Embedding层

   class Embedding(nn.Module):
   	def __init__(self):
   		super(Embedding, self).__init__()
   		self.tok_embed = nn.Embedding(vocab_size, hidden_size)
   		self.pos_embed = nn.Embedding(maxlen, hidden_size)
   		self.norm = nn.LayerNorm(hidden_size)
   		self.dropout = nn.Dropout(hidden_dropout_prob)
   		
   	def forward(self, input_ids):
   		seq_len = input_ids.size(1)
   		pos = torch.arange(seq_len, dtype=torch.long)
   		pos = pos.unsqueeze(0).expand_as(input_ids) # [seq_len] -> [batch_size, seq_len] 
   		embedding = self.tok_embed(input_ids) + self.pos_embed(pos) # [batch_size, seq_len, hidden_size] 
   		embedding = self.norm(embedding)
   		embedding = self.dropout(embedding)
   		return embedding
   1
   2
   3
   4
   5
   6
   7
   8
   9
   10
   11
   12
   13
   14
   15
   16

   2）EncoderLayer层
   EncoderLayer层的主体其实就是Transformer中的编码层，核心内容为自注意力机制，代码如下图所示。这部分内容较为复杂，如果想进一步了解可以参考笔者之前的一篇文章“基于模型结构与模型源码两个层面理解Transformer”。

   def forward(self, Q, K, V, attn_mask):
   	scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)
   	scores.masked_fill_(attn_mask, -1e9)
   	attn = nn.Softmax(dim=-1)(scores)
   	context = torch.matmul(attn, V)
   	return context
   1
   2
   3
   4
   5
   6

   3）全连接层

   self.linear = nn.Linear(hidden_size, hidden_size)
   self.activ2 = gelu
   embed_weight = self.embedding.tok_embed.weight
   self.fc2 = nn.Linear(hidden_size, vocab_size)
   self.fc2.weight = embed_weight
   
   masked_pos = masked_pos[:, :, None].expand(-1, -1, d_model) # [batch_size, max_pred, hidden_size]
   h_masked = self.activ2(self.linear(h_masked)) # [batch_size, max_pred, hidden_size]
   logits_lm = self.fc2(h_masked) # [batch_size, max_pred, vocab_size]
   return logits_lm
   1
   2
   3
   4
   5
   6
   7
   8
   9
   10

3. RoBERTa模型

   class RoBERTa(nn.Module):
   	def __init__(self):
   		super(RoBERTa, self).__init__()
   		self.embedding = Embedding()
   		self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])
   		self.linear = nn.Linear(hidden_size, hidden_size)
   		self.activ2 = gelu
   		embed_weight = self.embedding.tok_embed.weight
   		self.fc2 = nn.Linear(hidden_size, vocab_size)
   		self.fc2.weight = embed_weight
   	def forward(self, input_ids, masked_pos):
   		output = self.embedding(input_ids) # [batch_size, seq_len, hidden_size]
           enc_self_attn_mask = get_attn_pad_mask(input_ids, input_ids) # [batch_size, maxlen, maxlen]
           for layer in self.layers:
           	output = layer(output, enc_self_attn_mask) # [batch_size, seq_len, hidden_size]
           masked_pos = masked_pos[:, :, None].expand(-1, -1, hidden_size）  # [batch_size, max_pred, hidden_size]
           h_masked = torch.gather(output, 1, masked_pos) # [batch_size, max_pred, hidden_size]
           h_masked = self.activ2(self.linear(h_masked)) # [batch_size, max_pred, hidden_size]
           logits_lm = self.fc2(h_masked) # [batch_size, max_pred, vocab_size]
           return logits_lm
   1
   2
   3
   4
   5
   6
   7
   8
   9
   10
   11
   12
   13
   14
   15
   16
   17
   18
   19
   20

3 模型训练及保存

model = RoBERTa()
criterion = nn.CrossEntropyLoss(ignore_index=0) 
optimizer = optim.Adam(model.parameters(), lr=0.0001)

for epoch in range(epochs):
	loss = 0
	pbar = tqdm.tqdm(loader, desc='Train', nrows=200, ncols=100)
	for input_ids, masked_tokens, masked_pos in pbar:
		logits_lm = model(input_ids, masked_pos) # [batch_size, max_pred, vocab_size]
		loss_lm = criterion(logits_lm.view(-1, vocab_size), masked_tokens.view(-1)) 
		loss += loss_lm
		optimizer.zero_grad()
		loss_lm.backward()
		optimizer.step()
	print('Epoch:', '%04d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss))
	
    output_path = './outputs/'
    save_path = os.path.join(output_path, 'checkpoint_RoBERTa-' + str(epoch+1))
    torch.save(model, save_path)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

4 模型测试

下图为模型测试和测试结果，从预测结果可以看出，预测的两段文本，分别掩码了13和2个token，但是两段文本都只预测对了一个token。这是因为当时只用了少量的训练数据，而且模型未得到充分训练。
模型推理过程：

model=torch.load(save_path) 
input_ids = input_ids.numpy().tolist()
masked_tokens = masked_tokens.numpy().tolist()
masked_pos = masked_pos.numpy().tolist()
print([idx2word[w] for w in input_ids[0] if idx2word[w] != ‘[PAD]’])

logits_lm = model(torch.LongTensor([input_ids[0]]),torch.LongTensor([masked_pos[0]]))# logits_lm ：[batch_size, max_pred, vocab_size]
logits_lm = logits_lm.data.max(2)[1][0].data.numpy()  # 预测出的掩码位置的token，长度为max_pred
print(‘masked tokens list : ’,[pos for pos in masked_tokens[0] if pos != 0]) print('predict masked tokens list : ',[pos for pos in logits_lm if pos != 0])
1
2
3
4
5
6
7
8
9

测试结果1：

测试结果2：

5 常见大模型简述

5.1 ALBERT模型简述

ALBERT 的结构和 BERT 基本一样，采用了 Transformer 以及 GELU 激活函数。具体的创新部分有三个：

• embedding 层参数因式分解
• 跨层参数共享
• 将 NSP 任务改为 SOP 任务

1. embedding 层参数因式分解
   原始的 BERT 模型以及各种依据 Transformer 的预训连语言模型都有一个共同特点，即 E=H，其中 E 指的是 Embedding Dimension，H 指的是 Hidden Dimension。这就会导致一个问题，当提升 Hidden Dimension 时，Embedding Dimension 也需要提升，最终会导致参数量呈平方级的增加。所以 ALBERT 的作者将 E 和 H 进行解绑，具体的操作就是在 Embedding 后面加入一个矩阵进行维度变换。E 的维度是不变的，如果 H 增大了，我们只需要在 E 后面进行一个升维操作即可。
   举个例子：
   V * H = 30000 * 768 = 23 040 000
   V * E + E * H = 30000 * 256 + 256 * 768 = 7 876 608
   当V为30000，H为768，E为256时，参数量从2300万降低到780万
2. 跨层参数共享
   传统 Transformer 的每一层参数都是独立的，包括各层的 self-attention、全连接。这样就导致层数增加时，参数量也会明显上升。之前有工作试过单独将 self-attention 或者全连接层进行共享，都取得了一些效果。ALBERT 作者尝试将所有层的参数进行共享，相当于只学习第一层的参数，并在剩下的所有层中重用该层的参数，而不是每个层都学习不同的参数。
   作者通过实验发现了使用参数共享可以有效的提升模型稳定性，实验结果如下图：
   
3. 将NSP 任务改为 SOP 任务
   BERT 引入了一个叫做下一个句子预测的二分类问题。这是专门为提高使用句子对，如 “自然语言推理” 的下游任务的性能而创建的。但是在 RoBERTa 这样的论文中已经阐明了 NSP 的无效性，并且发现它对下游任务的影响是不可靠的。
   因此，ALBERT 提出了另一个任务 —— 句子顺序预测，关键思想是（1）从同一个文档中取两个连续的句子作为一个正样本；（2）交换这两个句子的顺序，并使用它作为一个负样本。SOP 提高了下游多种任务（SQUAD，MNLI，SST-2，RACE）的表现。
   

5.2 ELECTRA模型简述

ELECTRA最主要的贡献是提出了新的预训练任务和框架，把生成式的Masked language model(MLM)预训练任务改成了判别式的Replaced token detection(RTD)任务，判断当前token是否被语言模型替换过。

具体而言：首先按照一定的比例对于原始输入序列X-ORI进行随机MASK操作得到新序列X-MASK；其次将X-MASK作为生成器模型（Generator）的输入，该生成器模型用于对序列中那些被MASK操作的tokens生成新的token（此时生成器是面向所有词表而言），以此来产生新的序列X-Generator；之后将X-Generator作为判别器模型的输入（Discriminator），该判别器模型用于判别序列中每一个token是否是原始token（和X-ORI进行对比而言）。

ELECTRA和BERT的区别：

5.3 ERNIE模型简述

ERNIE相比于BERT，做出了如下改进：

• mask策略。BERT只使用了字级别的随机masking，但是ERNIE使用了字、短语、实体三个级别的masking，旨在使模型学习到更多高级的语义
• 添加更多优质中文语料。加入了百度百科、百度新闻、百度贴吧等中文语料，使得在中文NLP任务上效果更好
• 对话语言模型。对Dialog的角色，进行了Dialog embedding，从而加强模型在Dialog上的效果

1. mask策略，使用了字、短语、实体三个级别的masking，模型可以学习到更大语义单元的知识
   
2. 添加更多优质中文语料。加入了百度百科、百度新闻、百度贴吧等中文语料，使得在中文NLP任务上效果更好。
3. 对话语言模型，对Dialog的角色，进行了Dialog embedding，从而加强模型在Dialog上的效果。
   
   DLM任务可帮助ERNIE学习对话中的隐式关系，这也增强了模型学习语义表示的能力。DLM任务的模型体系结构与MLM任务的模型体系结构兼容，因此可以通过MLM任务对其进行预训练。

总结

本文是对过去看过内容的一个复盘，只涉及到训练大模型的主要部分，部分细节无法逐一展现，需要源码的可以私我，如果疑问欢迎评论区交流。