闲社

标题: 显存不够用？手把手盘大模型推理/训练内存优化技巧 [打印本页]

作者: 康波 时间: 4 小时前
标题: 显存不够用？手把手盘大模型推理/训练内存优化技巧
兄弟们，大模型火了，但显存贵啊。😅 很多人一上来就是70B、130B，结果单卡根本跑不动，甚至推理都OOM。别慌，内存优化不是玄学，是实打实的工程。

先说训练场景。最经典的FSDP（全分片数据并行）和DeepSpeed ZeRO Stage 3，把优化器状态、梯度、参数分片放到不同设备，显存占用直线下降。但别忘了混合精度训练（FP16/BF16），相比FP32直接省一半内存，还能加速。别跟我说精度损失，现在多数任务根本无所谓。

推理侧，重点搞量化。INT8、INT4甚至NF4，虽然精度会掉一点，但显存能砍到四分之一。配合vLLM、TGI这类推理框架，支持PagedAttention，动态管理KV Cache，显存碎片少多了。还有FlashAttention，彻底干掉O(N^2)的内存瓶颈，长序列也不虚。

最后说个冷门但实用的：torch.compile + 内存检查点（Gradient Checkpointing）。训练时只存部分中间激活，反向传播再算一次，时间换空间。你可以试下在A100上跑70B推理，量化+KV Cache优化后，显存从80G降到20G以内。

各位现在跑模型最头疼的内存瓶颈是哪个？是模型参数太大，还是序列太长导致KV Cache爆了？来评论区聊聊，我抽空写个实操教程。🔧

欢迎光临闲社 (https://www.xianshe.com/)