大模型微调之添加special_token改变词表大小并进行微调方法——以llama3为例

背景：在微调任务中，常常需要修改词表大小对大模型进行微调。但是修改词表大小后会出现embedding层和lm_head不一致的问题。

embedding层：将token映射一个feature，大小一般是(vocabulary_size,hidden_dim)
lm_head:用来预测一个feature是什么token。大小一般是(hidden_dim, vocabulary_size)

llama-3-8B-Instruct介绍

llama-3-8B-Instruct是llama3的一个专门在指令微调上进行预训练的一个模型。
加载模型和tokenizer

from transformers import AutoTokenizer,AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained('meta-llama/Meta-Llama-3-8B-Instruct')
tokenizer = AutoTokenizer.from_pretrained('meta-llama/Meta-Llama-3-8B-Instruct')
print(model.get_input_embeddings())
print(model.lm_head)
print(len(tokenizer))

输出结果为：
Embedding(128256, 4096)
Linear(in_features=4096, out_features=128256, bias=False)
128256
可以看到词表的大小和embeding以及lm_head的关系

接下来我们需要扩充词表。扩充词表有两种，请一定要注意！！！一种是添加special_token，比如pad_token，另一种是添加普通的token。这篇博客聚焦于add_special_tokens函数。

添加特殊token

以llama3为例，他是没有pad_token的，一般的模型会使用pad_token = eos_token的方法进行训练。但是这个方法存在很大的缺陷！
1、在transformers的默认对齐方法中，有一个方案是

在上面的代码中，将所有pad_token_id标记为-100（即在模型训练中不会计算他的损失）在这样的背景下，eos_token也会被标记为-100，于是模型无法正确的预测结束符号。
因此在这样的背景下，需要添加特殊的pad_token。

1、下面是参考代码，以添加一个pad_token为例子，添加多个special_token自行更改，不做详细的解析。首先看看embedding层和lm_head的矩阵形状大小

embedding = model.get_input_embeddings()
lm_head = model.lm_head
print(embedding.weight.data.shape)
print(lm_head.weight.data.shape)

torch.Size([128256, 4096])
torch.Size([128256, 4096])
特别注意lm_head的维度大小

创建新的embedding和lm_head

#我这里只添加一个新的special_token,采用的是均值初始化。
#均值初始化的意思就是说，用已有的编码方式去编码特殊token，然后将编码的embedding的平均值初始化特殊token的embedding。
#要理解下面的代码需要知道embedding层的工作原理，以及one-hot向量

pad_token = "<|pad_id|>"
token_ids = tokenizer.encode(pad_token,add_special_tokens=False)
add_token_tokens = embeddings(torch.tensor(token_ids))
parameter = add_token_tokens.mean(dim = 0)
tokenizer.add_special_tokens({"pad_token":'<pad_id>'})

new_embedding = torch.nn.Embedding(len(tokenizer),4096)
new_embedding.weight.data[:len(tokenizer) - 1, :] = embeddings.weight.data
new_embedding.weight.data[len(tokenizer) - 1] = parameter
model.set_input_embeddings(new_embedding)

new_lm_head = torch.nn.Linear(in_features = 4096, out_features = len(tokenizer),bias = False)
lm_head = model.lm_head
new_lm_head.weight.data[:len(tokenizer) - 1,:] = lm_head.weight.data
original = 0
for id in token_ids:
    original += lm_head.weight.data[id]
mean_para = original / len(token_ids)
new_lm_head.weight.data[len(tokenizer) - 1] = mean_para

model.lm_head = new_lm_head

tokenizer.save_pretrained('llama-3-8B-extent')
model.save_pretrained('llama-3-8B-extent',max_shard_size = '8GB')

通过上面的代码，我们在网络结构层面和tokenizer上完成了对词表的扩充。但是需要注意的是还需要对网络的配置文件进行修改。我们上面的llama-3-8B-extent文件夹

点击config.json文件夹，找到vocab-size,将词表长度从128256改成len(tokenizer).比如我这里的是128257（因为我只添加了一个pad_token。同时在进行了上面的词表扩充后，记得对embeddings进行微调（虽然我没有微调也取得了较好的结果）