deepspeed 执行chatglm推理会出现显存成倍的情况，具体如下：
deepspeed --num_gpus 2 test.py 服务器环境中有2个gpu，执行推理的时候，会出现显存为2倍参数量。 也就是可以理解为，并deepspeed加载后，并没有完成模型并行，2块gpu都加载了同样参数量，存在显存翻倍的情况。 测试bloom模型并未出现这个情况，具体不知道是不是和模型底层有关，chatglm为什么使用deepspeed推理时候会出现这个问题？
import torch import os from transformers import AutoTokenizer, AutoModelForCausalLM,AutoModel import deepspeed from modeling_chatglm import ChatGLMForConditionalGeneration import pandas as pd model_id = "chatglm-6b" tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) model = AutoModel.from_pretrained(model_id, trust_remote_code=True,torch_dtype=torch.float16) model = deepspeed.init_inference( model=model, # Transformers models mp_size=2, # Number of GPU dtype=torch.float16, # dtype of the weights (fp16) replace_method="auto", # Lets DS autmatically identify the layer to replace replace_with_kernel_inject=True, # replace the model with the kernel injector ) model.eval()
批量生成文案prompt = "请用c++写一段代码，代码的功能是实现对数据库的查询，需要有注释" response, history = model.module.chat(tokenizer, prompt, history=[], max_length=4000) print(response) print("\n\n"+"="70+" 显存占用 "+"="70) print(os.system("nvidia-smi")
Environment- OS:
- Python: 3.8.10
- Transformers:4.28.0
- PyTorch: 2.0
- CUDA Support (`python -c "import torch; print(torch.cuda.is_available())"`) :