Python JAX显存未释放问题及子进程解决方案
发布时间：2026-10-16| 4人浏览过
本文针对Python中使用JAX进行GPU计算时出现的显存无法释放问题，分析了即使设置XLA_PYTHON_CLIENT_PREALLOCATE=false及执行del操作后显存仍被占用的现象。通过引入multiprocessing模块，将GPU计算封装在子进程中，利用进程结束自动释放显存的机制，有效解决了僵尸显存占用问题，同时保留了函数返回值供后续使用。
1.技术背景
笔者在执行Jax任务时遇到一个异常现象：尽管仅分配了极小的矩阵空间，但在多次任务执行后，显存仍会突然溢出。此前已按照Jax官方文档配置XLA_PYTHON_CLIENT_PREALLOCATE参数为false，即禁止显存预分配（默认会预分配90%的显存空间）。参考网络上的类似问题（如链接1、2、3、4），均指出某些操作后显存未被释放。本文提供具体的复现案例及处理思路，欢迎提出更优方案。
2.问题复现
在未执行任何GPU任务前，通过nvidia-smi命令查看显存状态如下：
Tue Dec 14 16:14:32 2021
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.42.01 Driver Version: 470.42.01 CUDA Version: 11.4 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Quadro RTX 4000 On | 00000000:03:00.0 On | N/A |
| 30% 43C P8 20W / 125W | 1260MiB / 7979MiB | 10% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
| 1 Quadro RTX 4000 On | 00000000:A6:00.0 Off | N/A |
| 30% 34C P8 7W / 125W | 10MiB / 7982MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| 0 N/A N/A 1673 G /usr/lib/xorg/Xorg 110MiB |
| 0 N/A N/A 3015 G /usr/lib/xorg/Xorg 661MiB |
| 0 N/A N/A 3251 G /usr/bin/gnome-shell 132MiB |
| 0 N/A N/A 1142734 G ...AAAAAAAAA= --shared-files 64MiB |
| 0 N/A N/A 1337710 G ...AAAAAAAAA= --shared-files 80MiB |
| 0 N/A N/A 1371509 G ...369783.log --shared-files 63MiB |
| 0 N/A N/A 1506625 G ...AAAAAAAAA= --shared-files 89MiB |
| 1 N/A N/A 1673 G /usr/lib/xorg/Xorg 4MiB |
| 1 N/A N/A 3015 G /usr/lib/xorg/Xorg 4MiB |
+-----------------------------------------------------------------------------+
随后启动一个ipython终端，执行以下Jax任务代码：
In [1]: import numpy as np
In [2]: import os
...: os.environ['CUDA_VISIBLE_DEVICES']='1'
...: os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
In [3]: from jax import numpy as jnp
In [4]: a = np.ones(1000000)
In [5]: b = jnp.array(a)任务执行后，再次查看nvidia-smi结果，状态如下：
Tue Dec 14 16:18:26 2021
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.42.01 Driver Version: 470.42.01 CUDA Version: 11.4 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Quadro RTX 4000 On | 00000000:03:00.0 On | N/A |
| 30% 42C P8 20W / 125W | 1238MiB / 7979MiB | 10% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
| 1 Quadro RTX 4000 On | 00000000:A6:00.0 Off | N/A |
| 30% 36C P0 35W / 125W | 114MiB / 7982MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| 0 N/A N/A 1673 G /usr/lib/xorg/Xorg 110MiB |
| 0 N/A N/A 3015 G /usr/lib/xorg/Xorg 661MiB |
| 0 N/A N/A 3251 G /usr/bin/gnome-shell 129MiB |
| 0 N/A N/A 1142734 G ...AAAAAAAAA= --shared-files 44MiB |
| 0 N/A N/A 1337710 G ...AAAAAAAAA= --shared-files 80MiB |
| 0 N/A N/A 1371509 G ...369783.log --shared-files 63MiB |
| 0 N/A N/A 1506625 G ...AAAAAAAAA= --shared-files 89MiB |
| 1 N/A N/A 1673 G /usr/lib/xorg/Xorg 4MiB |
| 1 N/A N/A 3015 G /usr/lib/xorg/Xorg 4MiB |
| 1 N/A N/A 1743467 C /usr/local/bin/python 101MiB |
+-----------------------------------------------------------------------------+此时结果符合预期，该Python进程占用了101MB显存。若在ipython中手动删除该对象，操作如下：
In [6]: del b
In [7]: b
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
<ipython-input-7-89e6c98d9288> in <module>
----> 1 b
NameError: name 'b' is not defined再次查看nvidia-smi结果，显示如下：
Tue Dec 14 16:21:12 2021
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.42.01 Driver Version: 470.42.01 CUDA Version: 11.4 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Quadro RTX 4000 On | 00000000:03:00.0 On | N/A |
| 30% 42C P5 21W / 125W | 1231MiB / 7979MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
| 1 Quadro RTX 4000 On | 00000000:A6:00.0 Off | N/A |
| 30% 35C P8 7W / 125W | 114MiB / 7982MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| 0 N/A N/A 1673 G /usr/lib/xorg/Xorg 110MiB |
| 0 N/A N/A 3015 G /usr/lib/xorg/Xorg 662MiB |
| 0 N/A N/A 3251 G /usr/bin/gnome-shell 111MiB |
| 0 N/A N/A 1142734 G ...AAAAAAAAA= --shared-files 55MiB |
| 0 N/A N/A 1337710 G ...AAAAAAAAA= --shared-files 80MiB |
| 0 N/A N/A 1371509 G ...369783.log --shared-files 63MiB |
| 0 N/A N/A 1506625 G ...AAAAAAAAA= --shared-files 89MiB |
| 1 N/A N/A 1673 G /usr/lib/xorg/Xorg 4MiB |
| 1 N/A N/A 3015 G /usr/lib/xorg/Xorg 4MiB |
| 1 N/A N/A 1743467 C /usr/local/bin/python 101MiB |
+-----------------------------------------------------------------------------+
结果显示，尽管Python程序中该对象已被删除（NameError），但显存中的数据并未释放。若不进行干预，这块显存将一直被占用，直到python进程（作为该进程的守护进程）结束。
3.解决思路
目前尚不清楚该问题的具体发生机制，某些特定场景下的僵尸进程问题与本文复现场景略有不同。鉴于Python进程结束时显存能成功释放，本文尝试通过进程控制来解决显存分配与清空问题。以下是一个基于多进程实现的案例：
import os
os.environ['CUDA_VISIBLE_DEVICES']='1'
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
import time
from multiprocessing import Pool
import numpy as np
from jax import numpy as jnp
a = np.ones(1000000)
def f(a):
b = jnp.array(a)
time.sleep(2)
print('Array b has been deleted!')
return True
with Pool(1) as p:
res = p.map(f, [(a,)])
print ('Is jax array deleted successfully?\t{}'.format(res))
time.sleep(6)
在该程序中，我们将涉及GPU矩阵转换与分配的相关任务封装在一个函数中，并通过multiprocessing启动子进程执行该任务。此方法无需手动执行del删除操作，既实现了对象的即时销毁，又通过进程控制机制确保显存被清空。若进程执行出现问题，可通过terminate操作直接杀死进程，防止显存占用堆积。程序执行结果如下：
Array b has been deleted!
Is jax array deleted successfully? [True]
从nvidia-smi监控可见，显存占用仅持续了2秒（即函数内部设置的sleep参数时间）。随后的6秒sleep期间，显存占用已被清空，达到了预期目的。关键在于，我们仍能从函数获取返回值，确保后续所需参数不被同步销毁。需注意，若不使用子进程直接执行res=f(a)，临时变量b最终仍会残留在显存中，这是一个潜在风险。
4.总结概要
在使用某些python的GPU模块或编写CUDA代码时，常遇到显存被无端占用的情况，即使执行cudaFree()或Python的del操作也无法释放。最终发现，通过额外开启子进程封装相关操作，利用进程存活控制来管理GPU显存占用，是一种有效的临时规避方案。
到此这篇关于python中显存回收问题解决方法的文章就介绍到这了,更多相关python显存回收内容请搜索本站以前的文章或继续浏览下面的相关文章希望大家以后多多支持本站！
免责声明：本文内容仅供参考，不代表本站观点。转载或引用请注明出处。
