运行环境说明
Conda虚拟环境 | Python 3.10.13 | Pip 24.1.2 |
torch 2.1.1 | torchvision 0.16.1 | torchaudio 2.1.1 |
transformers 4.37.1 | accelerate | |
RTX 4060 8GB | CUDA 12.2 | Qwen-14B-Chat-Int4 |
确认下torch安装信息
import torch
print(f"torch version: {torch.__version__}, cuda available: {torch.cuda.is_available()}")
模型推理
from transformers import AutoModelForCausalLM, AutoTokenizer
# -------------- 输入token准备 --------------
prompt = '最近几天天气很热, 给大家一些降暑的建议'
messages = [ {'role': 'system', 'content': '你是一个生活助手'}, {'role': 'user', 'content': prompt} ]
tokenizer = AutoTokenizer.from_pretrained(
pretrained_model_name_or_path='D:\\usr\\models\\language\\qwen-14b-chat-int4\\qwen\\Qwen-14B-Chat-Int4\\',
trust_remote_code=True
)
text = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
model_in = tokenizer([text], return_tensors='pt').to('cuda:0')
print(f"model input: {model_in}")
# -------------- 模型准备 --------------
model = AutoModelForCausalLM.from_pretrained(
pretrained_model_name_or_path='D:\\usr\\models\\language\\qwen-14b-chat-int4\\qwen\\Qwen-14B-Chat-Int4\\',
torch_dtype='auto', device_map='auto',
trust_remote_code=True
)
# -------------- 模型推理 --------------
gen_ids = model.generate(model_in.input_ids, max_new_tokens=512)
print(f"gen_ids: {gen_ids}")
gen_ids = [ output_ids[len(input_ids):] for input_ids, output_ids in zip(model_in.input_ids, gen_ids) ]
# -------------- 模型原始输出转为文本结果 --------------
response = tokenizer.batch_decode(gen_ids, skip_special_tokens=True)[0]
print(f"response: {response}")
过程及输出结果如下所示:
从如上过程输出可以看到,大模型大概按如下几步做了处理:
- 输入转为token
- 通过已训练好的模型内部自注意力的推理,得到输出token概率
- 重复推理,得到最终结果
在如上代码可以看到,在加载模型的时候,torch_dtype参数的值是auto,即按照模型的默认精度加载的,我们也可以指定精度,如设置其值为: torch.float16等,另外device_map参数值是auto,则模型加载是会安装这个和顺序来尝试加载:GPU显存、CPU内存、磁盘。
当我们想指定模型可以占用的最大显存时,可以按如下配置:
过程中,可以直接print(model)来输出一下模型的信息[print(model.config)可以输出模型的配置信息]:
查看模型的配置时,看到精度是fp16的:
可以[list(model.parameters())[0].dtype]验证下:
前面输出模型信息时,看到有40层,也可以输出每一层看看里面的参数:
for name, param in model.named_parameters():
print(f'{name} / {param.numel()} / {param.requires_grad} / {param[0].dtype}')
计算下模型总参数量:
# -- 计算模型总参数量:
count = 0
for _, param in model.named_parameters():
count += param.numel()
print(f"total parameters: {count}")
差不多就是14B的量级了。