AI Infra之模型显存管理分析

AI Infra 之模型显存管理分析

摘要

某客户部署DeepSeek-R1模型时，发现显存占用持续上升。经排查，显存增长源于PyTorch的缓存机制，而非KV Cache不足或NCCL问题。PyTorch通过Block和BlockPool管理显存，申请大块内存并预留缓存，导致显存占用增加但不会自动释放。建议调整max_split_size_mb减少碎片化，或手动调用torch.cuda.empty_cache()释放显存。