Unsloth缓存配置
Unsloth 微调框架是否支持模型缓存目录配置？如何自定义路径？
Unsloth 微调框架是否支持模型缓存目录配置？如何自定义路径？
问题分析
Unsloth 是近期备受关注的大语言模型高效微调框架，其核心优势在于通过优化的 CUDA 内核实现了 2x-5x 的训练速度提升和 70% 的显存节省。然而，在实际部署中，开发者常遇到模型缓存目录的问题。
默认情况下，Unsloth（以及底层的 HuggingFace Transformers）会将模型下载到用户目录下的 .cache/huggingface 文件夹。对于系统盘空间有限的用户（如服务器上只有 50GB 系统盘），这会导致磁盘空间不足。更糟糕的是，某些企业环境的用户目录可能位于网络存储上，网络延迟会严重影响模型加载速度。
另一个常见场景是多用户共享服务器。如果每个用户都下载一份模型到自己的缓存目录，会造成大量重复存储。合理的做法是将模型存储在共享目录（如 /data/models），所有用户共用。
问题还涉及模型版本的隔离。同一模型的不同版本（如 Llama-3-8b 和 Llama-3-8b-Instruct）可能需要独立管理。Unsloth 默认的缓存结构是扁平的，难以直观区分。
解决原理
Unsloth 本身不直接管理模型缓存，它依赖 HuggingFace 的 transformers 和 huggingface_hub 库处理模型下载和加载。因此，配置缓存路径需要从 HuggingFace 生态的层面入手。
HuggingFace 提供了多种配置缓存路径的方式：
方式一：环境变量
最直接的方法是设置环境变量 HF_HOME 或 TRANSFORMERS_CACHE。HF_HOME 是 HuggingFace 生态的全局缓存根目录，影响所有 HuggingFace 库。TRANSFORMERS_CACHE 只针对 transformers 库，优先级高于 HF_HOME。
环境变量的生效时机很关键：必须在导入任何 HuggingFace 库之前设置。一旦库被导入，缓存路径就确定了，后续修改无效。
方式二：代码中设置
可以在 Python 代码中直接修改 HfApi 的缓存配置，或在调用 from_pretrained 时指定 cache_dir 参数。这种方式更灵活，可以针对不同模型使用不同路径。
方式三：配置文件
HuggingFace CLI 提供了 huggingface-cli 命令行工具，可以通过 huggingface-cli config set cache_dir <path> 持久化配置。
对于 Unsloth，由于它使用 FastLanguageModel.from_pretrained() 加载模型，底层最终调用 transformers，所以上述方法都有效。但需要注意 Unsloth 的某些优化（如 4-bit 量化加载）可能有额外的缓存需求。
程序实现与说明
import os
import sys
from pathlib import Path
# ================== 环境变量配置（推荐方式） ==================
# 必须在导入 unsloth 或 transformers 之前执行
def setup_huggingface_cache(cache_root: str):
"""
配置 HuggingFace 生态的缓存目录
通过环境变量实现，必须在导入相关库之前调用
cache_root: 缓存根目录路径
"""
# 转换为绝对路径并规范化
cache_path = Path(cache_root).resolve()
cache_path.mkdir(parents=True, exist_ok=True)
# HF_HOME: 全局缓存根目录，影响所有 HF 库
os.environ['HF_HOME'] = str(cache_path)
# TRANSFORMERS_CACHE: 专门针对 transformers，优先级更高
os.environ['TRANSFORMERS_CACHE'] = str(cache_path / 'transformers')
# HF_HUB_CACHE: huggingface_hub 库的缓存
os.environ['HF_HUB_CACHE'] = str(cache_path / 'hub')
# HF_DATASETS_CACHE: datasets 库专用
os.environ['HF_DATASETS_CACHE'] = str(cache_path / 'datasets')
# 禁用遥测（可选，减少网络请求）
os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
print(f"[缓存配置] 根目录: {cache_path}")
print(f"[缓存配置] TRANSFORMERS_CACHE: {os.environ['TRANSFORMERS_CACHE']}")
# 在脚本最开始调用
CUSTOM_CACHE_DIR = r"D:\models\hf_cache"
setup_huggingface_cache(CUSTOM_CACHE_DIR)
# ================== 导入 Unsloth ==================
from unsloth import FastLanguageModel
import torch
# ================== 模型加载 ==================
def load_model_with_cache(model_name: str, cache_dir: str = None):
"""
加载 Unsloth 优化的模型
支持自定义缓存目录
model_name: 模型名称（如 "unsloth/llama-3-8b-bnb-4bit"）
cache_dir: 可选，覆盖全局缓存配置
"""
# 如果未指定，使用环境变量中的配置
if cache_dir is None:
cache_dir = os.environ.get('TRANSFORMERS_CACHE')
# 确保目录存在
Path(cache_dir).mkdir(parents=True, exist_ok=True)
# 加载模型
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=model_name,
max_seq_length=2048, # 根据任务调整
dtype=None, # 自动检测
load_in_4bit=True, # 4-bit 量化
cache_dir=cache_dir # 指定缓存路径
)
print(f"[模型加载] 名称: {model_name}")
print(f"[模型加载] 缓存: {cache_dir}")
print(f"[模型加载] 参数量: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B")
return model, tokenizer
# ================== LoRA 配置 ==================
def setup_lora(model, r: int = 16):
"""
配置 LoRA 适配器
model: 基座模型
r: LoRA 秩，影响参数量和性能
"""
model = FastLanguageModel.get_peft_model(
model,
r=r,
target_modules=[
"q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"
],
lora_alpha=r, # 通常设为与 r 相同
lora_dropout=0,
bias="none",
use_gradient_checkpointing="unsloth",
random_state=42,
use_rslora=False,
loftq_config=None
)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"[LoRA配置] 可训练参数: {trainable:,} ({trainable/total*100:.2f}%)")
return model
# ================== 完整微调流程 ==================
def finetune_with_custom_cache():
"""
完整的 Unsloth 微调示例
包含缓存配置、模型加载、训练
"""
# 配置已在脚本开头完成
# 加载模型
model, tokenizer = load_model_with_cache(
"unsloth/llama-3-8b-bnb-4bit"
)
# 配置 LoRA
model = setup_lora(model, r=16)
# 准备数据集
from datasets import load_dataset
# 使用 Alpaca 格式数据集
dataset = load_dataset("yahma/alpaca-cleaned", split="train")
# 定义 Prompt 模板
alpaca_prompt = """Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
{}
### Response:
{}"""
EOS_TOKEN = tokenizer.eos_token
def format_prompts(examples):
texts = []
for instruction, output in zip(examples["instruction"], examples["output"]):
text = alpaca_prompt.format(instruction, output) + EOS_TOKEN
texts.append(text)
return {"text": texts}
dataset = dataset.map(format_prompts, batched=True)
# 训练参数
from trl import SFTTrainer
from transformers import TrainingArguments
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset,
dataset_text_field="text",
max_seq_length=2048,
args=TrainingArguments(
output_dir="./outputs",
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
max_steps=100,
learning_rate=2e-4,
fp16=not torch.cuda.is_bf16_supported(),
bf16=torch.cuda.is_bf16_supported(),
logging_steps=10,
save_steps=50
)
)
# 开始训练
print("\n[训练开始]")
trainer.train()
# 保存模型
model.save_pretrained_merged(
"./merged_model",
tokenizer,
save_method="merged_16bit"
)
print("[训练完成] 模型已保存")
if __name__ == "__main__":
finetune_with_custom_cache()
关键代码行解析：
os.environ['HF_HOME'] = str(cache_path)：设置 HuggingFace 全局缓存根目录。这个变量影响所有 HF 生态库，包括 transformers、datasets、huggingface_hub。os.environ['TRANSFORMERS_CACHE']：专门针对 transformers 库的缓存路径。如果同时设置了 HF_HOME 和 TRANSFORMERS_CACHE，后者优先级更高。cache_dir=cache_dir：在from_pretrained()中直接指定缓存路径。这种方式可以覆盖全局配置，为不同模型使用不同路径。load_in_4bit=True：Unsloth 的核心优势之一。使用 4-bit 量化加载模型，显存占用降低约 75%。use_gradient_checkpointing="unsloth"：Unsloth 优化的梯度检查点实现，比标准的梯度检查点更快。
注意事项：
- 环境变量必须在导入库之前设置，否则配置无效。
- Windows 路径建议使用原始字符串
r"..."或正斜杠/。 - 多 GPU 环境下，每个进程应使用独立的缓存目录或加锁，避免并发写入冲突。
- 对于 GGUF 格式模型，缓存机制可能不同，需要额外处理。