兄弟们,最近搞大模型部署,内存优化这块踩了不少坑,来聊聊心得。🔧
先说点基础:大模型推理时,显存主要被参数、激活值和KV cache占满。别光想着砍batch size,那是下下策。试试这几个实战方法:
1. **量化部署**:FP16转INT8,显存直接砍半。部署LLaMA-13B时,我用的bitsandbytes的8bit,精度掉不到1%,但能塞进24G卡。注意,别用4bit,推理速度会崩。
2. **KV cache优化**:长序列推理时,KV cache是显存刺客。可以开PagedAttention(vLLM的核心理念),把cache分页管理,碎片减少30%以上。或者直接限制max_length,别傻乎乎设4096。
3. **模型切分**:单卡搞不定就上张量并行。用Megatron-LM的TP,把矩阵运算拆到多卡,通信开销可控。我实测,2卡A100跑ChatGLM-6B,单batch推理延迟只涨8%。
4. **梯度检查点**:训练时必开。别心疼计算时间,显存省下的钱能买块新卡了。
最后提醒:别盲目上优化,先profile。用PyTorch的torch.cuda.memory_summary()看哪块爆了,再对症下药。
有个问题:你们在实际部署中,遇到的最奇葩的内存瓶颈是啥?是embedding层还是attention计算?来聊聊。💬 |