在Google Colab部署Llama-2-7B-chat-hf时加载检查点分片停滞无报错

我在Google Colab（Python 3.10.12）中尝试通过Hugging Face使用7B参数的Llama 2 chat模型，已通过Meta获取访问令牌。直接使用Hugging Face提供的代码并填入令牌，代码如下：

!pip install transformers
 
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

token = "---从Hugging Face复制粘贴的令牌---"

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf", token=token)
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", token=token)

模型开始下载，但执行到“Loading checkpoint shards:”步骤时停滞，无任何错误提示：

• 切换高显存GPU运行时：7B模型至少需要13GB显存，默认Colab可能分配CPU或低显存GPU。通过菜单栏「运行时→更改运行时类型」选择GPU硬件加速器，优先选A100或T4，避免显存不足导致加载停滞。
• 启用4bit量化加载：安装依赖库后，通过量化参数减少显存占用，加速加载：
  先执行安装命令：

  !pip install bitsandbytes accelerate
  

  修改模型加载代码：

  model = AutoModelForCausalLM.from_pretrained(
      "meta-llama/Llama-2-7b-chat-hf",
      token=token,
      load_in_4bit=True,
      device_map="auto",
      torch_dtype=torch.float16
  )
  

• 重启Colab运行时：若存在进程缓存或后台异常，重启运行时后重新执行代码，可解决无响应问题。
• 验证网络稳定性：分片加载停滞可能是网络波动导致下载中断，尝试切换稳定网络，或检查Colab的网络连接状态。
• 本地缓存加载（Google Drive）：将模型分片提前下载到Google Drive，从本地路径加载避免重复下载：
  1. 挂载Drive：

     from google.colab import drive
     drive.mount('/content/drive')
     

  2. 下载模型到Drive：

     from huggingface_hub import snapshot_download
     snapshot_download(
         "meta-llama/Llama-2-7b-chat-hf",
         local_dir="/content/drive/MyDrive/llama-2-7b-chat",
         token=token
     )
     

  3. 从本地路径加载：

     model = AutoModelForCausalLM.from_pretrained(
         "/content/drive/MyDrive/llama-2-7b-chat",
         device_map="auto"
     )
     

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