兄弟们,最近搞大模型部署,显存动不动就OOM,是不是很烦?别慌,今天来点干货,聊聊内存优化的几种路子。
首先,最基础的 **混合精度训练(FP16/BF16)** 必须安排上。把参数从FP32降到半精度,显存直接砍半,效果损失微乎其微。注意,损失缩放(Loss Scaling)得调好,不然梯度爆炸教你做人。
其次,**梯度检查点(Gradient Checkpointing)** 是个好兄弟。牺牲一点计算时间,换来中间激活值不存显存,要算的时候再重算。比如训练LLaMA-70B,不开这个直接GG,开了之后显存占用能降30%以上。
再来点狠的:**ZeRO优化**(DeepSpeed的ZeRO Stage 1-3)。分布式环境下,把优化器状态、梯度、参数分片到各卡,显存利用率飙升。单卡玩家也可以用ZeRO-Offload,把部分数据卸载到CPU内存或NVMe上,适合穷人玩法。
最后,**量化(Quantization)** 是部署时的核弹。INT8或INT4量化,模型体积缩小4倍,推理速度提升。比如用bitsandbytes库,4bit量化后,70B模型都能塞进24G显存。
你们在实际部署中,遇到过哪些内存坑?是batch size调不动,还是推理时内存泄漏?评论区交流下。 |