问题
最近在尝试使用lm-evaluation-harness时，在使用llama模型时，发现最开始的速度奇慢无比。其中定位到一部分问题是使用AutoTokenizer创建tokenizer缓慢，大约需要5分钟左右。这个问题该如何解决呢?
本文可以直接在Colab中打开
探索
首先安装依赖库
|
1 2 |
!pip install transformers==4.30.2 sentencepiece |
再设置BASE_MODEL为llama-7b
|
1 2 |
BASE_MODEL = "sharpbai/llama-7b-hf" |
使用LlamaTokenizer验证，可以看到实际执行时间为1s
|
1 2 3 4 |
from transformers import LlamaTokenizer tokernizer_llama = LlamaTokenizer.from_pretrained(BASE_MODEL) |
使用AutoTokenizer，发现会有非常长的时间在执行
|
1 2 3 4 |
from transformers import AutoTokenizer tokernizer_auto = AutoTokenizer.from_pretrained(BASE_MODEL) |
在colab的状态栏中，发现如下的代码在持续执行
|
1 2 3 4 |
from_pretrained() -> from_pretrained() -> _from_pretrained() -> __init__() -> __init() -> convert_slow_tokenizer() -> converted() -> tokenizer() -> extract() |
使用colab CPU实例，整个时间持续了5分49秒
这里发现了一点线索，convert_slow_tokenizer这个调用在transformers中可以看到何时被调用，检索后发现在生成fast的tokenizer时会使用。
于是使用如下代码查看刚才两个tokenizer的is_fast的状态
|
1 2 3 |
print(f"tokernizer_llama.is_fast: {tokernizer_llama.is_fast}") print(f"tokernizer_auto.is_fast: {tokernizer_auto.is_fast}") |
输出为
|
1 2 3 |
tokernizer_llama.is_fast: False tokernizer_auto.is_fast: True |
这里可以确认，AutoTokenizer默认返回的tokenizer是fast类型的，但LlamaTokenizer返回的是slow类型的
再次手工创建LlamaTokenizerFast来试试
|
1 2 3 4 |
from transformers import LlamaTokenizerFast tokernizer_llama_fast = LlamaTokenizerFast.from_pretrained(BASE_MODEL) |
看起来确实是LlamaTokenizerFast带来的问题。默认情况下，它会将slow版转换为fast版，上面的代码执行消耗了6分23秒
经过搜索，看到一个github issue描述创建LlamaTokenizerFast缓慢的问题，其中提到了会有一个一次性转换，转换后将模型保存起来就可以避免再次转换。那问题是不是就在这里呢?
正巧搜到了一个huggingface上的fast tokenizer，那就试一试
|
1 2 3 4 5 6 7 |
FAST_TOKENIZER_MODEL = "heilerich/llama-tokenizer-fast" from transformers import LlamaTokenizerFast tokernizer_llama_fast2 = LlamaTokenizerFast.from_pretrained(FAST_TOKENIZER_MODEL) |
输出如下
|
1 2 3 4 5 |
Downloading (…)okenizer_config.json: 0%| | 0.00/727 [00:00<?, ?B/s] Downloading tokenizer.model: 0%| | 0.00/500k [00:00<?, ?B/s] Downloading (…)/main/tokenizer.json: 0%| | 0.00/1.84M [00:00<?, ?B/s] Downloading (…)cial_tokens_map.json: 0%| | 0.00/411 [00:00<?, ?B/s] |
可以看到，这个外部的tokenizer多了一个tokenizer.json文件，然后加载速度就是秒级的了。
方案
如此一来，需要将转换为fast的tokenizer进行save_pretrained就可以了。使用以下代码进行修复
首先登录huggingface帐号
|
1 2 3 4 |
# 登录 from huggingface_hub import notebook_login notebook_login() |
然后生成LlamaTokenizerFast，再push_to_hub即可
|
1 2 3 4 5 6 7 |
# 独立补丁，用于修补缺失fast tokenizer的repo BASE_MODEL = 'openlm-research/open_llama_13b' SAVE_MODEL_NAME = BASE_MODEL.split('/')[-1] from transformers import LlamaTokenizerFast tokenizer_fast = LlamaTokenizerFast.from_pretrained(BASE_MODEL) tokenizer_fast.push_to_hub(SAVE_MODEL_NAME) |
验证一下解决方案有效，实测确实解决了问题
|
1 2 |
tokenizer_fast_test = LlamaTokenizerFast.from_pretrained(f"sharpbai/{SAVE_MODEL_NAME}") |
接下来就是批量处理
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
repair_models = [ 'sharpbai/llama-7b-hf', 'sharpbai/llama-13b-hf', 'sharpbai/vicuna-7b-v1.3', 'sharpbai/vicuna-13b-v1.3', 'sharpbai/alpaca-7b-merged', 'sharpbai/alpaca-lora-7b-merged', 'sharpbai/chinese-llama-plus-lora-7b-merged', 'sharpbai/chinese-alpaca-plus-lora-7b-merged', 'sharpbai/open_llama_7b', 'sharpbai/baichuan-llama-7b', 'sharpbai/baichuan-vicuna-7b', ] for BASE_MODEL in repair_models: print(f"fixing {BASE_MODEL}") SAVE_MODEL_NAME = BASE_MODEL.split('/')[-1] from transformers import LlamaTokenizerFast tokenizer_fast = LlamaTokenizerFast.from_pretrained(BASE_MODEL) tokenizer_fast.push_to_hub(SAVE_MODEL_NAME) |
在colab上执行，修复总共花了1个多小时
不过后来尝试了下在安装了比较新的transformers版本后，转换时间就变成瞬间了，命令如下
|
1 2 |
pip install git+https://github.com/huggingface/transformers@6c57ce15587810968d64fb4b5700b63726397194 |
结论
在保存llama的tokenizer权重时，需要使用LlamaTokenizerFast也加载一遍，并push_to_hub，即可解决问题