问题出现得很突然,日志倒是留得挺全。DataLoader 开了 pin_memory 和八个 worker,训练跑两小时后进程 RSS 涨到二十 GB,机器开始用交换分区,一个 epoch 从十八分钟拖到四十分钟;torch 2.4,主机内存三十二 GB,数据集是几万张图片。
刚开始盯着「数据集在初始化时缓存了全部图片数组,worker 继承后翻倍」,折腾了半个多小时,排掉了。
排查方向其实有限,先把可能性排一排。锁页缓冲池随 batch 增大,且不会被换出。还有一种是worker 数量多,每个进程各复制一份 batch。
还有一种情况,没有指定锁页设备,分配走了默认策略。
先看进程 RSS 与锁页内存的分布
python -c "import psutil,os;p=psutil.Process(os.getpid());print(p.memory_info().rss/2**30, p.memory_info().locked/2**30)"
再往下把 worker 数降到四并限制预取批次
python train.py --num-workers 4 --prefetch-factor 2
随后确认数据集对象加载后的常驻大小
python -c "import resource,dataset;print(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024)"
压测两小时观察 RSS 是否还单调上涨
python watch_mem.py --pid $PID --duration 7200 --out rss.csv
别只看报错消失,验证要多做一步。两小时压测 RSS 稳定在十二 GB 以内没有单调增长,epoch 耗时回到十八分钟,交换分区不再被使用。
容易忽略的是,pin_memory 的收益在主机到显存的拷贝,代价是锁页内存不可换出,worker 一多就会把内存吃满。如果同一现象隔几天又出现,说明根因没找到,别当偶发处理。
