开源大模型扩充中文词表

当前开源大模型正在如火如荼的进行，随着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中介绍了这一词表扩充的脚本实现，如下：

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、效果测试

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对应的参数，再重新更新模型层的参数。

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）裁剪实验

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）实验效果

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次迭代裁剪到目标大小：

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)