LlamaIndex加载Vicuna模型GPU显存翻倍问题排查求助

我在测试LlamaIndex时使用Vicuna-7B/13B模型，发现模型加载到CPU内存时占用正常，但迁移至GPU后VRAM占用量直接翻倍，导致13B模型无法运行。但使用FastChat的CLI运行13B模型时一切正常，VRAM和内存占用均约25GB。

相关代码实现 #

# define prompt helper
# set maximum input size
max_input_size = 2048
# set number of output tokens
num_output = 256 
# set maximum chunk overlap
max_chunk_overlap = 20
prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap)

def model_size(model: torch.nn.Module):
    return sum(p.numel() for p in model.parameters()) 

def model_memory_size(model: torch.nn.Module, dtype: torch.dtype=torch.float16):
    # Get the number of elements for each parameter
    num_elements = sum(p.numel() for p in model.parameters())
    # Get the number of bytes for the dtype
    dtype_size = torch.tensor([], dtype=dtype).element_size()
    return num_elements * dtype_size / (1024 ** 2)  # return in MB

class CustomLLM(LLM):
    model_name = "vicuna-7b"
    model_path = "../../../SharedData/vicuna-7b/"
    kwargs = {"torch_dtype": torch.float16}
    tokenizer_vicuna = AutoTokenizer.from_pretrained(model_path, use_fast=False)
    model_vicuna = AutoModelForCausalLM.from_pretrained(
            model_path, low_cpu_mem_usage=True, **kwargs 
    )
    # device = "cuda" 
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 
    print(device)
    print(f"Model size: {model_size(model_vicuna)/1e6} million parameters")
    dtype_current = next(model_vicuna.parameters()).dtype
    print(f"Model memory size: {model_memory_size(model_vicuna,dtype_current)} MB")
    print("Press any key to continue...")
    input()
    model_vicuna.to(device)
    
    @torch.inference_mode()
    def generate_response(self, prompt: str, max_new_tokens=num_output, temperature=0.7, top_k=0, top_p=1.0):
        encoded_prompt = self.tokenizer_vicuna.encode(prompt, return_tensors='pt').to(self.device)
        max_length = len(encoded_prompt[0]) + max_new_tokens
        with torch.no_grad():
            output = self.model_vicuna.generate(encoded_prompt, 
                                                max_length=max_length,
                                                temperature=temperature, 
                                                top_k=top_k, 
                                                top_p=top_p, 
                                                do_sample=True)
        response = self.tokenizer_vicuna.decode(output[0], skip_special_tokens=True)
        return response

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        prompt_length = len(prompt)
        response = self.generate_response(prompt)
        # only return newly generated tokens
        return response[prompt_length:]

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        return {"name_of_model": self.model_name}

    @property
    def _llm_type(self) -> str:
        return "custom"

运行输出 #

cuda
Model size: 6738.415616 million parameters
Model memory size: 12852.5078125 MB

nvidia-smi结果 #

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.161.03   Driver Version: 470.161.03   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA RTX A6000    Off  | 00000000:17:00.0 Off |                  Off |
| 30%   39C    P2    69W / 300W |  26747MiB / 48682MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      2205      G   /usr/libexec/Xorg                   9MiB |
|    0   N/A  N/A      2527      G   /usr/bin/gnome-shell                5MiB |
|    0   N/A  N/A   2270925      C   python                          26728MiB |
+-----------------------------------------------------------------------------+

已排查情况及需求 #

已确认模型dtype为float16，GPU内存占用26747MiB，CPU内存占用约12852MB，13B模型直接触发CUDA内存不足。求进一步排查建议。

内容的提问来源于stack exchange，提问作者TZXia