最近在试着把7B的模型部署到本地，显卡是3090（24G），按理说应该够用，但一加载就说CUDA out of memory。我试了加载时传low_cpu_mem_usage=True，也试了把max_memory分配给GPU，但好像没多大用。我看网上说用bitsandbytes做4bit量化，但装了之后又报错说找不到CUDA的setup.py，心态有点崩。就想问问大家，实际部署的时候一般怎么配参数？是不是需要把加载的模型分片到多卡或者用offload？另外，除了量化，还有哪些比较实用的显存优化思路？希望能给个可操作的步骤，谢谢了。
楼主
22天前
部署LLaMA模型一直OOM，大佬们都是怎么优化显存的？
请 登录 后发表回复
全部回复
共 81 条
2楼
1天前
24G跑7B按理说真够，但你这个报错大概率是加载时峰值爆了，因为transformers默认会把整个权重一次性载入显存，加上激活值和KV cache就超了。我建议你试试load_in_4bit=True配合bnb_4bit_compute_dtype=torch.float16，同时把device_map设成"auto"，这样bitsandbytes会自动把部分层放到CPU，虽然慢点但能跑起来。至于你说的setup.py报错，多半是bitsandbytes版本和CUDA不匹配，直接pip install bitsandbytes --upgrade试试，或者去GitHub下预编译的wheel。还有个思路是换用llama.cpp的GGUF格式，量化到Q4_K_M大概6G多，3090跑起来很轻松，而且CPU offload也简单，不用折腾transformers那套。另外可以把max_new_tokens限制短一点，或者用vLLM这类推理框架，它们有PagedAttention能省不少显存，但7B用vLLM有点大材小用。最后实在不行就上offload，把embedding和lm_head放CPU，但要注意这会增加延迟。反正别死磕bitsandbytes，多试几条路。