Falcon-7b微调模型加载内核崩溃，求助分片保存Hugging Face LLM方法

我按照LLM微调指南完成了Falcon-7B在Hugging Face数据集上的微调，但在Jupyter Notebook加载微调模型时内核崩溃。我的GPU是NVIDIA Quadro P5000（16GB VRAM），推测是显存不足导致，想通过分片保存模型解决，但不知道怎么操作，附上相关代码寻求帮助。

微调后保存模型的代码 #

# 定义微调模型的保存目录
output_dir = "./fine_tuned_model"

# 使用save_model方法保存微调后的模型
trainer.save_model(output_dir)

# 可选：将模型上传到Hugging Face模型仓库分享给他人
# trainer.push_to_hub("omarfarooq908/falcon-7b-finetuned01")

加载模型时内核崩溃的代码 #

from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM

config = PeftConfig.from_pretrained("omarfarooq908/falcon-7b-finetuned01")
model = AutoModelForCausalLM.from_pretrained("ybelkada/falcon-7b-sharded-bf16")
model = PeftModel.from_pretrained(model, "omarfarooq908/falcon-7b-finetuned01")

解决方案：分片保存模型+显存优化加载策略 #

1. 微调时配置分片保存 #

修改保存逻辑，指定分片大小，让模型按分片存储，避免单文件过大：

output_dir = "./fine_tuned_model"

# 保存时启用分片，设置每个分片最大为8GB（适配16GB显存的GPU）
trainer.save_model(output_dir, max_shard_size="8GB")

# 推送到Hub时同步启用分片
trainer.push_to_hub("omarfarooq908/falcon-7b-finetuned01", max_shard_size="8GB")

如果是训练前初始化Trainer，也可以直接在TrainingArguments中配置：

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./fine_tuned_model",
    # 其他训练参数（如batch_size、学习率等）保持不变
    save_safetensors=True,  # 推荐使用safetensors格式，更安全且显存友好
    max_shard_size="8GB"    # 按8GB分片保存模型文件
)

trainer = Trainer(
    model=model,
    args=training_args,
    # 其他训练相关参数
)

2. 加载模型时优化显存占用 #

结合分片加载+4-bit量化，进一步降低显存消耗：

from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# 配置4-bit量化参数，大幅压缩模型显存占用
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

config = PeftConfig.from_pretrained("omarfarooq908/falcon-7b-finetuned01")
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    quantization_config=bnb_config,
    device_map="auto",  # 自动分配模型层到GPU/CPU，平衡显存占用
    low_cpu_mem_usage=True,  # 优化CPU内存占用，间接降低GPU显存压力
    load_in_4bit=True
)
model = PeftModel.from_pretrained(model, "omarfarooq908/falcon-7b-finetuned01")

3. 额外显存清理技巧 #

在Jupyter Notebook中加载模型前，先执行以下代码清理显存：

import torch
torch.cuda.empty_cache()

同时确保transformers、peft、bitsandbytes库为最新版本，避免旧版本的显存泄漏问题。

内容的提问来源于stack exchange，提问作者Muhammad Omar Farooq