几行代码教你轻松完成超大模型推理:LLaMA-30B+TITAN RTX*4+accelerate

这篇具有很好参考价值的文章主要介绍了几行代码教你轻松完成超大模型推理:LLaMA-30B+TITAN RTX*4+accelerate。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

是不是苦于没有ChatGPT的API key或者免费的token而无法愉快地和它玩耍?想不想在有限的计算资源上部署大模型并调戏大模型??想不想解锁大模型的除了对话之外的其它功能???几行代码教你搞定如何在有限的计算资源下部署超大模型并实现推理。

准备

超大语言模型。OPT,GPT,LLaMA都行,只要是开源的都行。去Hugging Face找一款心仪的模型,总有适合你的。我用的LLaMA-30B,你需要从官网上准备好下面这一堆文件:
几行代码教你轻松完成超大模型推理:LLaMA-30B+TITAN RTX*4+accelerate,llama

相应的环境依赖。作为调包侠,基本的pytorch、transformers等等就不用说了,这次介绍本期主角**accelerate**!!!
GPUs。TITAN RTX×4,也就是4×24G的算力。

# 踩坑 不用踩坑!!!因为我都踩好了。

直接开始

首先,简单阐述一下accelerate的实现原理:将模型切分存放,GPU不行放在CPU上,CPU放不下放在内存里。然后用哪一部分加载哪一部分就行。详细看油管的视频介绍。而计算哪部分需要被加载到哪里,需要经过预先的计算。可是不加载模型怎么计算内存并分配?新版本pytorch的init_empty_weights帮你实现。具体原理还是看看别人讲的吧,我主要负责踩坑。
上代码。

from transformers import LlamaTokenizer, AutoConfig, AutoModelForCausalLM
import os
import torch
from accelerate import infer_auto_device_map, init_empty_weights, load_checkpoint_and_dispatch

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'

# model_name = 'llama-7b'
model_name = 'llama-30b'

# 首先使用虚拟内存加载模型,这个内存被认为是无限大
# "/{}/".format(model_name)是你自己模型存放的路径,这个不会改的也别来部署大模型了
with init_empty_weights():
    model = AutoModelForCausalLM.from_pretrained("/{}/".format(model_name),
                                                 trust_remote_code=True, torch_dtype=torch.float16)

# The model weights are not tied.
# Please use the `tie_weights` method before using the `infer_auto_device` function.
model.tie_weights()
# 人为设置不同的device所分配的内存, devide:0分配得少一些方便后续推理,实际操作中不写也行
# max_memory={0: "16GiB", 1: "23GiB", 2: "23GiB", 3:"23GiB", "cpu":"80Gib"}  # cpu辛苦你啦
# device_map = infer_auto_device_map(model, max_memory=max_memory)    # 自动推断device_map
#
# print(device_map)  # 可以看看模型不同层被分在了哪里

model = load_checkpoint_and_dispatch(model, checkpoint="/{}/".format(model_name),
                    device_map='auto', offload_folder="offload", no_split_module_classes=["LlamaDecoderLayer"],
                    offload_state_dict=True, dtype=torch.float16).half()
        # half()强制数据转化,看名字应该是一个内存优化的数据类型

print("(*^_^*) model load finished!!!! please play with {}".format(model_name))


while True:
    text_in = input('please input your question: ')
    # 注意,使用Autokenizer会很慢,不知道为什么
    tokenizer = LlamaTokenizer.from_pretrained("/{}/".format(model_name))

    inputs = tokenizer(text_in, return_tensors="pt").to(0)

    outputs = model.generate(inputs["input_ids"], max_new_tokens=20)
    text_out = tokenizer.decode(outputs[0].tolist(), skip_special_tokens=True)

    print(text_out)

注意事项

  1. 在使用LLaMA的时候,千万不要使用Autokenizer,这会导致tokenizer加载速度很慢很慢很慢。
  2. 加载模型的时候一定注意使用:no_split_module_classes。它表示模型的哪些部分不能被拆分到不同的GPUs上。具体来说,包含残差的部分一定要设置no split!!!

最后,30B的模型占用的显存如下,非常可观:
几行代码教你轻松完成超大模型推理:LLaMA-30B+TITAN RTX*4+accelerate,llama文章来源地址https://www.toymoban.com/news/detail-544129.html

到了这里,关于几行代码教你轻松完成超大模型推理:LLaMA-30B+TITAN RTX*4+accelerate的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包