两种方式加载Llama3到MPS为何内存占用存在差异？

内存差异的原因 #

两种加载方式的核心区别在于模型权重的加载路径：

• 方式一：通过device_map="mps"参数，Hugging Face的Accelerate库会直接将模型权重加载到MPS设备，全程不会在CPU内存中留存完整的权重副本，因此进程的CPU内存占用仅为必要的运行时开销，加上MPS显存的16GB，活动监视器显示总占用约16GB。
• 方式二：先将完整的模型权重加载到CPU内存（8B参数的float16模型权重约16GB），再通过.to("mps")转移到MPS设备。此时CPU内存中会残留未被及时回收的权重张量副本——即使你删除了原模型变量，PyTorch的内存管理机制可能因存在隐藏引用（如加载过程中生成的临时张量、内部缓存）而无法立刻释放这部分内存，导致进程CPU内存额外占用约4GB，总内存显示为20GB。

torch.mps.current_allocated_memory()仅统计MPS设备的显存占用，因此两种方式返回值一致，无法反映CPU内存的残留情况。

清除残留内存的方法 #

1. 优先使用方式一加载：这是最彻底的解决方案，直接避免CPU内存残留，代码如下：

llm = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    torch_dtype=torch.float16,
    device_map="mps",
    token=True,
)

2. 若必须先加载到CPU再转移：手动清理CPU内存的步骤如下：

# 先加载到CPU
llm_cpu = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    torch_dtype=torch.float16,
    token=True,
)
# 转移到MPS
llm = llm_cpu.to("mps")
# 删除CPU模型的引用
del llm_cpu
# 强制触发垃圾回收
import gc
gc.collect()
# 清空PyTorch的CPU内存缓存
torch.empty_cache()

注意：如果仍有残留，可能是Accelerate或PyTorch内部的隐藏缓存导致，此时建议切换到方式一加载。

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