书生大模型实战营闯关 - 8GB显存玩转书生大模型demo

好名让狗申请了

于 2024-07-31 15:21:46 发布

阅读量1.2k

点赞数 26

CC 4.0 BY-SA版权

分类专栏：自然语言处理深度学习书生大模型实战营文章标签： python 开发语言人工智能

本文链接：https://blog.csdn.net/chen_blog/article/details/140821784

创建开发机

创建一个使用10%GPU算力，cuda12.2系统的开发机，并启动。由于开发机的IO性能较差，开发机共享盘中已经创建好了本次实验所需要的conda环境

# 启动共享的conda环境
conda activate /root/share/pre_envs/icamp3_demo

部署cli模型

创建目录

创建一个目录用于存放代码

mkdir -p /root/demo
touch /root/demo/cli_demo.py

编写代码

代码的主要逻辑是使用transformers加载共享目录下的internlm2-chat-1_8b模型

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 模型地址
model_name_or_path = '/root/share/new_models/Shanghai_AI_Laboratory/internlm2-chat-1_8b'

# 创建tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True, device_map='cuda:0') # 运行在第0块显卡上

# 创建大语言模型
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True, device_map=