Python多进程隔离在解决OOM当中的运用

一场与GPU显存的持久战，记录我如何用多进程隔离技术解决大模型训练中的OOM问题

问题背景：当BERT遇上显存瓶颈

在大型语言模型的研究中，我遇到了一个令人头疼的问题：使用TRAKer库对BERT模型进行数据贡献度分析时，程序在处理到某个阶段总会因为显存不足(OOM)而崩溃。

想象一下这样的场景：你精心设计的实验，在运行了几个小时后突然崩溃，屏幕上出现那个令人沮丧的CUDA out of memory错误。更糟糕的是，这个问题不是一开始就出现，而是在处理了20多个数据窗口后才爆发，这意味着每次都要重头开始，进度永远卡在某个点无法突破。

我的硬件配置是NVIDIA A100 40GB，理论上应该足够应对这个任务。但现实是，即使用尽了各种显存优化技巧，问题依然存在。

尝试过的"常规武器"

在采用多进程方案前，我尝试了所有能找到的显存优化方法：

1. 精度优化：FP16半精度

# 将模型转换为半精度
self.model = self.model.half().eval().to(DEVICE)

这确实将显存占用减半，但代价是计算速度变慢，而且依然没有从根本上解决OOM问题。

2. 批大小调整

不断减小batch_size，从16降到8，再到4，甚至2。虽然单个窗口能跑了，但总会在某个时刻还是遇到显存瓶颈。

3. 显存清理大法

我写了详细的显存清理函数：

def clean_memory(objs: list):
    # 删除Python对象
    for obj in objs:
        del obj
    # 垃圾回收
    gc.collect()
    # CUDA缓存清理
    if ch.cuda.is_available():
        ch.cuda.empty_cache()
        ch.cuda.ipc_collect(