查看: 1183|回复: 0

LLMPruner:大语言模型裁剪工具

[复制链接]

42

主题

870

帖子

1732

积分

金牌飞友

Rank: 6Rank: 6

积分
1732
飞币
858
注册时间
2017-8-9
发表于 2023-7-28 20:35:02 | 显示全部楼层 |阅读模式
在本文中,笔者将分享一个大语言模型裁剪工具,通过裁剪词表,减少模型参数。

项目链接:

https://github.com/yangjianxin1/LLMPruner

模型链接:

https://huggingface.co/YeungNLP

最近,大语言模型(LLM, Large Language Model)犹如雨后春笋般,虽然其效果惊艳,但参数量巨大,让普通玩家望而却步。 如今的大语言模型大多为多语种大预言模型(Multilingual Large Language Model),如LLaMA、mT5、Bloom等,其词表规模巨大,占据非常大部分的模型参数,如Bloom具有25万词表。 在训练模型时,词表权重将会消耗非常大的显存,需要更多的训练资源,并且会降低训练速度,产生OOM的现象。

然而在许多下游任务中,我们往往只需要使用到一两种语言,例如在中文场景中,一般只会用到中英文。 所以我们可以对大语言模型的词表进行裁剪,只留下所需的部分词表,这样不仅能够充分保留模型的预训练知识,并且减少模型参数量,降低显存占用,提升训练速度,使用更少的显卡进行下游任务的finetune训练。

基于上述原因,笔者开发了LLMPruner项目,目前主要包含裁剪后的各种参数规模的Bloom模型。对Bloom进行词表裁剪,保留常用的中英文token,词表由250880将至46145,缩减为原来的18.39%。

裁剪后的模型如下表所示,读者可按需自取。笔者使用裁剪后的模型在下游任务进行调优,取得了不错的效果,后续也会对该工作进行分享。

LLMPruner:大语言模型裁剪工具w2.jpg

裁剪方法:

from pruners.vocabulary_pruner import BloomVocabularyPruner
# 需要进行裁剪的模型路径model_name_or_path = 'bigscience/bloom-560m'# 自己制作的词表的路径new_tokenizer_name_or_path = 'YeungNLP/bloom-560m-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='长风破浪会有时')

模型使用方法:

from transformers import BloomTokenizerFast, BloomForCausalLMtokenizer = BloomTokenizerFast.from_pretrained('YeungNLP/bloom-1b4-zh')model = BloomForCausalLM.from_pretrained('YeungNLP/bloom-1b4-zh')print(tokenizer.batch_decode(model.generate(tokenizer.encode('长风破浪会有时', return_tensors='pt'))))

进技术交流群请添加AINLP小助手微信(id: ainlp2)

请备注具体方向+所用到的相关技术点
LLMPruner:大语言模型裁剪工具w3.jpg

关于AINLP

AINLP 是一个有趣有AI的自然语言处理社区,专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享,主题包括文本摘要、智能问答、聊天机器人、机器翻译、自动生成、知识图谱、预训练模型、推荐系统、计算广告、招聘信息、求职经验分享等,欢迎关注!加技术交流群请添加AINLP小助手微信(id:ainlp2),备注工作/研究方向+加群目的。

阅读至此了,分享、点赞、在看三选一吧
您需要登录后才可以回帖 登录 | 加入联盟

本版积分规则

快速回复 返回顶部 返回列表