加载Meta-Llama-3-8B-Instruct时导入flan.v2.mixtures引发CUDA OOM问题问询

核心原因 #

导入flan.v2.mixtures时，模块内部会自动执行初始化操作——比如预加载部分数据集样本、初始化预处理相关张量，甚至有些代码会默认把这些数据放到GPU上，额外占用大量显存。你的模型加载本身需要约5.9GiB，两者叠加后超出了GPU的24GiB显存上限，触发了OutOfMemoryError。

另外也存在一种可能：该模块依赖的其他库在导入时会抢占GPU显存，和PyTorch的显存管理产生冲突，间接导致模型加载时显存不足。

解决办法 #

• 延迟导入模块：不要在脚本开头就导入flan.v2.mixtures，等模型完全加载完成、显存稳定后再执行导入操作，避免显存叠加占用。
• 排查模块初始化代码：查看flan.v2.mixtures的源码，找到导入时自动执行的代码（比如自动创建数据集实例、移动张量到GPU的逻辑），修改为惰性加载，或者强制指定设备为CPU。
• 清理显存再导入：在导入该模块前执行torch.cuda.empty_cache()，释放临时显存占用；也可以设置torch.cuda.set_per_process_memory_fraction(0.9)，限制当前进程的显存占用比例，给数据集模块留出空间。
• 规范数据集加载流程：后续加载Flan v2数据集时，确保数据在CPU上完成预处理，采用批量加载的方式送入GPU，不要一次性把整个数据集加载到显存中。

内容的提问来源于stack exchange，提问作者Bongni