在GeForce 1080 8Gb机器运行Llama2遇CUDA内存不足求解决方案

• 模型量化：用bitsandbytes库做4位或8位量化，大幅降低显存占用。加载模型时添加参数：

  from transformers import AutoModelForCausalLM, AutoTokenizer
  
  model = AutoModelForCausalLM.from_pretrained(
      "meta-llama/Llama-2-7b-hf",
      load_in_4bit=True,
      device_map="auto"
  )
  tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
  

  4位量化可将7B模型显存占用控制在2-3GB，完全适配8GB显卡。

• 启用梯度检查点：开启梯度检查点以牺牲少量计算时间为代价减少显存占用，加载模型后执行：

  model.gradient_checkpointing_enable()
  

• 限制生成序列长度：生成文本时设置max_new_tokens为较小值（如50、100），避免过长文本占用额外显存：

  inputs = tokenizer("Hello world!", return_tensors="pt").to("cuda")
  outputs = model.generate(**inputs, max_new_tokens=50)
  

• 清理显存碎片：在代码中定期清理无用张量与缓存：

  import torch
  del unused_tensor  # 替换为实际无用变量名
  torch.cuda.empty_cache()
  

• 切换更小模型：若业务允许，改用Llama-2-3B模型，无需量化即可在8GB显卡上运行，显存占用更低。

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