如何在本地加载以Safe Tensor格式存储的微调LLAMA2模型

第一步：安装必要依赖 #

在Notebook中执行以下命令安装所需依赖库：

!pip install transformers accelerate safetensors bitsandbytes

• transformers：提供LLAMA2模型的标准加载接口
• safetensors：专门用于读取Safe Tensor格式的权重文件
• accelerate：优化模型加载与运行的硬件适配效率
• bitsandbytes：可选，用于低显存环境下的模型量化加载

第二步：加载模型与Tokenizer #

将代码中的./your-finetuned-llama2-model替换为你本地模型文件所在的目录路径（目录内需包含safetensors权重文件、config.json配置文件等核心文件）：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 本地模型目录路径
model_path = "./your-finetuned-llama2-model"

# 加载Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
# LLAMA默认无pad token，将eos token设为pad token
tokenizer.pad_token = tokenizer.eos_token

# 显存充足时加载完整模型
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    local_files_only=True,
    use_safetensors=True,
    device_map="auto"
)

# 显存不足时启用4bit量化加载（需bitsandbytes库）
# model = AutoModelForCausalLM.from_pretrained(
#     model_path,
#     local_files_only=True,
#     use_safetensors=True,
#     device_map="auto",
#     load_in_4bit=True,
#     bnb_4bit_quant_type="nf4",
#     bnb_4bit_compute_dtype=torch.bfloat16
# )

第三步：测试模型生成 #

加载完成后，用以下代码验证模型功能：

prompt = "请输入你的测试提示词"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# 生成文本
outputs = model.generate(
    **inputs,
    max_new_tokens=100,
    temperature=0.7,
    do_sample=True
)

# 解码并输出结果
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

关键注意事项 #

• 确保模型目录文件完整：至少包含safetensors权重文件、config.json、tokenizer_config.json和vocab.json
• 需使用支持LLAMA2的transformers版本（建议v4.30及以上）
• 低显存设备优先启用4bit/8bit量化加载，避免显存溢出

内容的提问来源于stack exchange，提问作者Tej Kiran Boppana