当前开源大模型正在如火如荼的进行,随着LLAMA,BLOOM为代表的开源社区逐步完善,如何基于这两个模型更好地使用低成本、高性能的中文场景需求,目前已经出现了多种具有代表性的工作。

不过很现实的问题是,LLaMA词表中仅包含很少的中文字符,其对中文并不友好,BLOOM作为一个多语言模型,词表有过大,在训练过程中并不平民化

因此,为了解决这个问题,通过干预词表,或通过增加词表,或裁剪词表,并加以预训练这一范式,已经逐步成为一个主流的方式。

因此,为了增强对该范式的认识,本文主要从LLAMA扩充词表以增强中文能力、Bloom裁剪词表以降低训练成本这两个角度进行介绍,充分借鉴了相关开源项目的代码原理一些实验论述,供大家一起参考。

一、LLaMA扩充词表以增强中文能力

《 Efficient and Effective Text Encoding for Chinese Llama and Alpaca》这一文章介绍了在LLaMA上进行中文词表扩充,以增强中文能力的工作。

项目地址:Github

1、LLaMA为什么要扩充词表

为什么要扩充词表?直接在原版LLaMA上用中文预训练不行吗?

这个是个有趣的问题,该项目对此做了解释,该工作认为:

首先,原版LLaMA模型的词表大小是32K,LLaMA已经在公开可用的语料库中预训练了1T到1.4T个token,其中大多数数据为英语,因此LLaMA理解和生成中文的能力受到限制,对比一下多语言模型XLM-R的词表大小为250K,预训练中没有出现过或者出现得很少的语言学习得不充分。

其次,LLaMA词表中仅包含很少的中文字符,原始LLaMA tokenizer词汇表中只有不到一千个中文字符,虽然可以通过回退到字节来支持所有的中文字符,但这种回退策略会显著增加序列长度,并降低处理中文文本的效率,所以在切词时会把中文切地更碎,需要多个byte token才能拼成一个完整的汉字,进而导致信息密度降低。

比如,在扩展词表后的模型中,单个汉字倾向于被切成1个token,而在原版LLaMA中可能就需要2-3个才能组合成一个汉字,显著降低编解码的效率。

为了解决这些问题,该项目在中文语料库上对LLaMA模型进行预训练,以增强其基本的中文理解和生成能力。

2、增加中文字符的核心思想

在具体实现上,该项目提出了以下两个解决方案来扩展LLaMA tokenizer中的中文词汇:

1)训练中文词表并与原Tokenizer合并

在通用中文语料上训练了基于sentencepiece的20K中文词表并与原版LLaMA模型的32K词表进行合并,排除重复的token后,得到的最终中文LLaMA词表大小为49953

需要注意的是,在fine-tune阶段Alpaca比LLaMA多一个pad token,所以中文Alpaca的词表大小为49954

2)训练过程中调整训练矩阵

为了适应新的tokenizer,需要将词嵌入和语言模型头从V × H调整为V’× H的形状,其中V = 32,000代表原始词汇表的大小,而V’ = 49,953则是Chinese LLaMA tokenizer的词汇表大小。

2、具体实现

项目地址: Github中介绍了这一词表扩充的脚本实现,如下:

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
import os
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"]="python"
from transformers import LlamaTokenizer
from sentencepiece import sentencepiece_model_pb2 as sp_pb2_model
import sentencepiece as spm
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--llama_tokenizer_dir', default=None, type=str, required=True)
parser.add_argument('--chinese_sp_model_file', default='./chinese_sp.model', type=str)
args = parser.parse_args()
llama_tokenizer_dir = args.llama_tokenizer_dir
chinese_sp_model_file = args.chinese_sp_model_file

# 加载tokenizer
llama_tokenizer = LlamaTokenizer.from_pretrained(llama_tokenizer_dir)
chinese_sp_model = spm.SentencePieceProcessor()
chinese_sp_model.Load(chinese_sp_model_file)
llama_spm = sp_pb2_model.ModelProto()
llama_spm.ParseFromString(llama_tokenizer.sp_model.serialized_model_proto())
chinese_spm = sp_pb2_model.ModelProto()
chinese_spm.ParseFromString(chinese_sp_model.serialized_model_proto())

# 输出各模型tokens的长度
print(len(llama_tokenizer),len(chinese_sp_model))
print(llama_tokenizer.all_special_tokens)
print(llama_tokenizer.all_special_ids)
print(llama_tokenizer.special_tokens_map)

# 将中文添加到LLaMA tokenizer
llama_spm_tokens_set=set(p.piece for p in llama_spm.pieces)
print(len(llama_spm_tokens_set))
print(f"Before:{len(llama_spm_tokens_set)}")
for p in chinese_spm.pieces:
piece = p.piece
if piece not in llama_spm_tokens_set:
new_p = sp_pb2_model.ModelProto().SentencePiece()
new_p.piece = piece
new_p.score = 0
llama_spm.pieces.append(new_p)
print(f"New model pieces: {len(llama_spm.pieces)}")
## 保存合并的tokenizer
output_sp_dir = 'merged_tokenizer_sp'
output_hf_dir = 'merged_tokenizer_hf' # the path to save Chinese-LLaMA tokenizer
os.makedirs(output_sp_dir,exist_ok=True)
with open(output_sp_dir+'/chinese_llama.model', 'wb') as f:
f.write(llama_spm.SerializeToString())
tokenizer = LlamaTokenizer(vocab_file=output_sp_dir+'/chinese_llama.model')
tokenizer.save_pretrained(output_hf_dir)
print(f"Chinese-LLaMA tokenizer has been saved to {output_hf_dir}")

可以通过执行如下脚本完成训练

1
2
3
python merge_tokenizers.py \
--llama_tokenizer_dir llama_tokenizer_dir \
--chinese_sp_model_file chinese_sp_model_file

其中:

  • llama_tokenizer_dir: 指向存放原版LLaMA tokenizer的目录;

  • chinese_sp_model_file: 指向用sentencepiece训练的中文词表文件

3、效果测试

1
2
3
4
5
6
llama_tokenizer = LlamaTokenizer.from_pretrained(llama_tokenizer_dir)
chinese_llama_tokenizer = LlamaTokenizer.from_pretrained(output_hf_dir)
content='''人工智能是计算机科学、心理学、哲学等学科融合的交叉学科'''
print("original sentence",len(text),text)
print("orginal tokenizer",len(llama_tokenizer.tokenize(text)), llama_tokenizer.tokenize(text))
print("Chinese okenizer", len(chinese_llama_tokenizer.tokenize(text)), chinese_llama_tokenizer.tokenize(text))

通过实验可以发现,使用中文LLaMA分词器相对于原始LLaMA分词器生成的token数减少了一半左右。

如上表所示,通过对比原始LLaMA分词器和中文LLaMA分词器,使用中文LLaMA分词器相对于原始的编码长度有明显的减少,这表明该项目提出的方法在提高LLaMA模型的中文理解和生成能力方面是有效的。

类似的项目还有chinese-vicuna

二、Bloom裁剪词表以降低训练成本

前面我们说到LLaMA当前存在不包含中文的问题,与其相反的一条路就是找一个开源的可以支持中文的模型。

1、为什么要对BLOOM进行词表裁剪

其中,BLOOM是一个典型的代表,BLOOM的词表有25万,LLaMA词表大小是3万,对于一个中文文本,LLaMA toeknzie之后的序列是2048,而BLOOM tokenize之后的序列可能只有1536或者更短,效率更高;

此外,BLOOM对非拉丁语系的性能表现更优。文章一文使用LLaMA和BLOOM分别在中文的指令微调数据上进行预训练,然后使用类似Vicuna的评测方式进行中文能力评测,让GPT-4作为评委为两者的回答打分,发现7B版本的BLOOM以及13B版本LLaMA,效果表明指标要高一些。

不过,Bloom作为一个多语言大模型,问题也是存在的,那就是参数量太大,词表往往很大,不够平民化。

因为,Bloom具有25万词表,在训练模型时,词表权重将会消耗非常大的显存,降低训练速度,产生OOM的现象,并且在一些下游任务中,往往只需要使用到一两种语言,例如在中文场景中,一般只会用到中英文。

所以,在具体下游任务中使用这些模型的时候,尤其是不需要其它语言,只需要中文和英文,可以对其vocab进行裁剪,既可以大大减少参数量,也能够保留模型的性能。

2、词表裁剪工具LLMPruner

LLMPruner是一个大语言模型裁剪工具,通过对大语言模型的冗余词表进行裁剪,减少模型参数量,降低显存占用,提升训练速度,并且能够保留预训练中学习到的知识。

例如,项目,对Bloom进行词表裁剪,保留常用的中英文token,词表由250880将至46145,缩减为原来的18.39%。

1)裁剪思想

词表裁剪的实现原理较为简单,核心在于提取输出new_token对应的原来模型里面的token对应的参数,再重新更新模型层的参数。

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
71
72
73
74
75
76
77
import os.path
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from tqdm import tqdm
class VocabularyPruner(object):
# 检查模型裁剪后,生成结果是否一致
def check(self, old_model_name_or_path, new_model_name_or_path, text):
max_length = 20
# 使用老模型对文本编码
old_model = AutoModelForCausalLM.from_pretrained(old_model_name_or_path)
old_tokenizer = AutoTokenizer.from_pretrained(old_model_name_or_path)
old_input_ids = old_tokenizer(text, return_tensors='pt').input_ids
old_output = old_model.generate(old_input_ids, max_length=max_length)
old_output_text = old_tokenizer.batch_decode(old_output)
print('old_output:{}'.format(old_output_text))
# 使用新模型对文本编码
new_model = AutoModelForCausalLM.from_pretrained(new_model_name_or_path)
new_tokenizer = AutoTokenizer.from_pretrained(new_model_name_or_path)
new_input_ids = new_tokenizer(text, return_tensors='pt').input_ids
new_output = new_model.generate(new_input_ids, max_length=max_length)
new_output_text = new_tokenizer.batch_decode(new_output)
print('new_output:{}'.format(new_output_text))
if old_output_text == new_output_text:
print('output is same, succeed to prune.')
else:
print('output is not same, fail to prune.')

def prune(self, model_name_or_path, new_tokenizer_name_or_path, save_path, new_name_or_path=None):
# 创建输出目录
if not os.path.exists(save_path):
os.makedirs(save_path)
# 加载新词表。如果是中文,就是中文的词表
new_tokenizer = AutoTokenizer.from_pretrained(new_tokenizer_name_or_path)
# 加载原词表。一般为多语言模型的词表
old_tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
# 检查新词表是否为原词表的子集
old_vocab = old_tokenizer.vocab
new_vocab = new_tokenizer.vocab
for token in tqdm(new_vocab.keys()):
if token not in old_vocab:
raise Exception('{} not exist'.format(token))
print('new_tokenizer is subset of old_tokenizer')
# 获得新词表中每个token_id到原词表的token_id的映射
new2old_token_id = {}
for token, token_id in tqdm(new_vocab.items()):
old_token_id = old_vocab[token]
new2old_token_id[token_id] = old_token_id
# 加载多语言模型
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype='auto')
# 计算原模型的参数量
old_params = sum(p.numel() for p in model.parameters())
print("Total params of original model: %.2fM" % (old_params / 1e6))
# 对于新词表中的每个token,取出其对应的权重,复制到新模型中
vocab_size = len(new_tokenizer)
hidden_size = model.config.hidden_size
new_embeds = torch.nn.Embedding(vocab_size, hidden_size, dtype=model.dtype)
new_lm_head = torch.nn.Linear(in_features=hidden_size, out_features=vocab_size, bias=False, dtype=model.dtype)
# 更新词表权重
self.update_ebeddings(model, new2old_token_id, new_embeds, new_lm_head)
model.config.__dict__['vocab_size'] = vocab_size
if new_name_or_path is not None:
model.config.__dict__['_name_or_path'] = new_name_or_path
# 计算新模型的参数量
new_params = sum(p.numel() for p in model.parameters())
print("Total params of new model : %.2fM" % (new_params / 1e6))
print('词表缩小为原来的:{}%'.format(round(len(new_tokenizer) / len(old_tokenizer), 4)*100))
print('模型参数量缩小为原来的:{}%'.format(round(new_params / old_params, 4)*100))
model.save_pretrained(save_path)
new_tokenizer.save_pretrained(save_path)
## 定义BloomVocabularyPruner
class BloomVocabularyPruner(VocabularyPruner):
def update_ebeddings(self, model, new2old_token_id, new_embeds, new_lm_head):
for token_id, old_token_id in tqdm(new2old_token_id.items()):
new_embeds.weight.data[token_id] = model.transformer.word_embeddings.weight.data[old_token_id]
new_lm_head.weight.data[token_id] = model.lm_head.weight.data[old_token_id]
model.transformer.word_embeddings.weight = new_embeds.weight
model.lm_head.weight = new_lm_head.weight

2)裁剪实验

1
2
3
4
5
6
7
8
9
10
11
from pruners.vocabulary_pruner import BloomVocabularyPruner
## 定义裁剪模型路径
model_name_or_path = 'bigscience/bloom-560m'
## 定义裁剪后词表路径
new_tokenizer_name_or_path = 'YeungNLP/bloom-396m-zh'
save_path = 'path-to-save'
## 定义裁剪工具,开始裁剪
pruner = BloomVocabularyPruner()
pruner.prune(model_name_or_path, new_tokenizer_name_or_path, save_path)
## 检测裁剪模型与原模型差别
pruner.check(model_name_or_path, save_path, text='长风破浪会有时')

3)实验效果

1
2
3
4
5
6
7
8
9
10
11
100%|██████████| 46145/46145 [00:00<00:00, 1309531.65it/s]
new_tokenizer is subset of old_tokenizer
100%|██████████| 46145/46145 [00:00<00:00, 1120687.88it/s]
Total params of original model: 559.21M
100%|██████████| 46145/46145 [00:01<00:00, 41641.55it/s]
Total params of new model : 396.82M
词表缩小为原来的:18.41%
模型参数量缩小为原来的:70.96000000000001%
old_output:['长风破浪会有时,直挂云帆济沧海。 愿你,在人生的旅途中,能遇见最美的风景,遇见最美的自己。</s>']
new_output:['长风破浪会有时,直挂云帆济沧海。 愿你,在人生的旅途中,能遇见最美的风景,遇见最美的自己。</s>']
output is same, succeed to prune.

3、模型裁剪工具TextPruner

TextPruner是一个为预训练语言模型设计,基于PyTorch实现的模型裁剪工具包。它提供了针对预训练模型的结构化裁剪功能,通过识别并移除模型结构中不重要的结构与神经元,达到压缩模型大小、提升模型推理速度的目的。

地址:https://github.com/airaria/TextPruner/

论文地址:https://aclanthology.org/2022.acl-demo.4/

TextPruner提供了3种裁剪模式,分别为词表裁剪(Vocabulary Pruning),Transformer裁剪(Transformer Pruning)和流水线裁剪(Pipeline Pruning)。

其中:

1)词表裁剪其实现思想在于,通过移除词表中未在具体任务未出现的token,可以实现减小模型体积,提升MLM等任务训练速度的效果,要进行词表裁剪,用户提供一个文本文件或字符串列表(list of strings)。TextPruner将从model和tokenizer中移除未在文本文件或列表中出现过的token。

1
2
3
from textpruner import VocabularyPruner
pruner = VocabularyPruner(model, tokenizer)
pruner.prune(dataiter=texts)

其中:model和tokenizer是要裁剪的模型和对应的分词器;
texts是字符串列表(list of strings),一般为任务相关数据的文本,用以确定裁剪后的词表大小。TextPruner将从model和tokenizer中移除未在其中出现过的token。

2)Transformer裁剪裁剪每个transformer模块的大小。一些研究表明transformer中的注意力头(attention heads)并不是同等重要,移除不重要的注意力头并不会显著降低模型性能。
TextPruner找到并移除每个transformer中“不重要”的注意力头和全连接层神经元,从而在减小模型体积的同时把对模型性能的影响尽可能降到最低。
例如,裁剪一个12层预训练模型,每层的注意力头目标数为8,全连接层的目标维数为2048,通过4次迭代裁剪到目标大小:

1
2
3
4
5
6
7
8
from textpruner import TransformerPruner, TransformerPruningConfig
transformer_pruning_config = TransformerPruningConfig(
target_ffn_size=2048,
target_num_of_heads=8,
pruning_method='iterative',
n_iters=4)
pruner = TransformerPruner(model,transformer_pruning_config=transformer_pruning_config)
pruner.prune(dataloader=dataloader, save_model=True)

其中:
transformer_pruning_config设置了具体的裁剪参数。dataloader用于向pruner提供数据用于计算各个注意力头的神经元的重要性,从而决定裁剪顺序。

总结

为了增强通过干预词表,或通过增加词表,或裁剪词表,并加以预训练以解决现有开源模型应用于中文领域范式的认识,本文主要从LLaMA扩充词表以增强中文能力、Bloom裁剪词表以降低训练成本这两个角度进行介绍,充分借鉴了相关开源项目的代码原理一些实验论述。

不过,通过构建多语言词表,直接加词或者减少词表,治标不治本,把词表搞大,总归会有OOV的情况出现,把词表搞小,通用能力又会受影响,如何找到一个更有效的方法进行处理,是个有趣的方向。

最后,感谢相关开源工作和科研工作者的五四奉献,对此表示崇高敬意。

参考文献

1、https://github.com/yangjianxin1/LLMPruner
2、https://zhuanlan.zhihu.com/p/623273021
3、https://www.cnblogs.com/xiximayou/p/17340154.html
4、https://github.com/ymcui/Chinese-LLaMA-Alpaca
5、https://mp.weixin.qq.com/s/pikAI1jL13kNsG8o4wzdHg