在Google Colab运行Llama 2时加载模型报错，需设置offload_dir

• 先创建一个本地目录用于存放模型卸载的文件，Colab里直接用代码创建就行：

  import os
  os.makedirs("./model_offload", exist_ok=True)
  

• 修改你原来的加载代码，加上offload_dir参数，指定刚才创建的目录：

  model = AutoPeftModelForCausalLM.from_pretrained(
      output_dir, 
      device_map="auto", 
      torch_dtype=torch.bfloat16,
      offload_dir="./model_offload"
  )
  

简单说明 #

当你设置device_map="auto"时，Hugging Face的库会自动把模型拆分到GPU、CPU甚至磁盘上运行，但需要你指定一个目录来存放那些从GPU卸载到磁盘的模型部分，不然就会触发这个报错。Colab的磁盘空间足够用，不用纠结容量问题。如果你的模型很小能完全塞进GPU显存，也可以直接把device_map改成"cuda"，这样就不用加offload_dir参数了。

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