本地部署大模型，如何优化显存占用？
前置条件与工具准备
在开始优化显存之前，先确认你已安装以下工具：
- Python 3.8 或更高版本
- CUDA 工具包（建议 11.7 以上）
- PyTorch（根据 CUDA 版本安装，例如
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118） - Transformers 库（
pip install transformers） - Hugging Face 的 Accelerate 库（
pip install accelerate）
另外，显卡驱动需支持 CUDA。
运行 nvidia-smi 可查看显存总量与当前使用量。
第一步：选择合适的模型与量化方式
不是所有大模型都吃显存相同。
先学会量化——将模型参数从 32 位浮点（FP32）压缩为 16 位（FP16）或 8 位（INT8），显存占用几乎减半。
操作步骤：
- 在 Hugging Face 模型页面查看支持的精度（如
TheBloke/Llama-2-7B-Chat-GGUF等）。 - 加载量化模型时指定参数：
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16, # 关键：使用 FP16
device_map="auto"
)避坑： 如果你的显卡不支持 FP16（老卡如 GTX 10 系），不要强行指定，否则显存反而炸掉。
可用 torch.float32 并配合后面技巧。
第二步：开启梯度检查点（Gradient Checkpointing）
这是训练和推理通用的优化技巧。
它通过计算时只保留少量中间激活值，用时间换空间，显存占用可降低 30%~50%。
配置方法：
model.gradient_checkpointing_enable()放置的位置：在加载模型之后、执行推理或训练之前。
注意，梯度检查点只对训练有明显的正向收益，推理时也能减少峰值显存，但会让推理速度略有下降。
验证： 使用 torch.cuda.memory_summary() 查看启用前后的显存峰值差异。
第三步：使用卸载技术降低瞬时占用
大模型推理时，有些层可以暂时放到 CPU 或硬盘上，需要时再搬到 GPU。Accelerate 库提供了 device_map 参数，可以帮助自动分配。
实操：
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="balanced" # 或 "auto"，也可以自定义
)device_map="balanced" 会平均分配到所有 GPU 和 CPU 内存；
如果想更激进，可设为 "sequential"，按层顺序分配。
避坑： CPU 卸载会造成推理速度大幅下降（尤其需要频繁交互时），建议仅用于模型加载或单次生成。
对实时对话场景，优先使用纯 GPU 方案。
第四步：调整批大小与序列长度
推理时，批大小（batch size）和最大序列长度直接影响显存。
对于零基础用户，最简单的方法是每次只生成一条回复（batch size=1），并且限制输入输出长度。
配置示例：
inputs = tokenizer("你的问题", return_tensors="pt", truncation=True, max_length=512)
outputs = model.generate(
inputs.input_ids,
max_new_tokens=256, # 控制生成长度
do_sample=False
)设置 max_length=512 并限制 max_new_tokens=256，可以避免因上下文过长而爆显存。
效果验证： 对比 max_new_tokens=1024 与 256 时的显存占用，使用 nvidia-smi 实时观察。
常见问题与避坑
Q1：显存一直卡在 0 或者爆显存怎么办？
检查是否安装了 CPU 版本 PyTorch，运行 python -c "import torch; print(torch.cuda.is_available())" 确认。若输出 False，重装 CUDA 版 PyTorch。
Q2：量化后模型输出变差？
INT8 量化会引入一定精度损失，若任务对回答质量要求高，建议只使用 FP16，并配合梯度检查点。
Q3：装了量化库后报 CUDA OOM？
可能是其他进程占用了显存。关掉无关程序，或重启 kernel。也可先用 torch.cuda.empty_cache() 清空缓存。
核心避坑提醒： 不要在训练和推理环境混用 torch.inference_mode() 和梯度检查点，否则会报错；
推理时请确保模型已切换为 eval() 模式。
效果验证
运行以下代码观察显存变化：
import torch
import time
model.eval()
with torch.no_grad():
sample_input = tokenizer("请用一句话总结：显存优化", return_tensors="pt").to("cuda")
torch.cuda.reset_peak_memory_stats()
start = time.time()
output = model.generate(**sample_input, max_new_tokens=50)
end = time.time()
peak_memory = torch.cuda.max_memory_allocated() / 1024**2
print(f"峰值显存：{peak_memory:.1f} MB")
print(f"推理耗时：{end-start:.2f} 秒")如果峰值显存明显减少（比如从 12GB 掉到 6GB），说明优化生效。
还可以对比未优化的版本（关闭梯度检查点、使用 FP32、无卸载）的占用。
写在最后
本地部署大模型时，显存优化是一套组合拳——量化、梯度检查点、卸载、限制序列长度，四个技巧配合使用效果最好。
建议先完整执行本文步骤，再根据自己的硬件（显存大小、通量需求）微调参数。
遇到报错时优先检查配置是否匹配，祝各位顺利跑起自己的大模型！
