关于Jittor的cudaDeviceSynchronize()问题

我有一个业务需求,需要在同一进程中利用多个深度学习框架依次训练一个模型结构队列中的模型,目前涉及:tensorflow、torch和jittor,物理设备为总内存约20G的双GPU,cuda11.2。实际运行中,程序进行一开始并不会出现任何问题,但一般在程序进行约1h后,jittor会率先报错,然后引起其他框架的报错,同时使得当前进程因等待显卡消息陷入S态,无法中断或杀死,目前仅能通过reboot服务器关闭进程,报错信息如截图所示:


具体问题分析细节将会跟帖说明。

首先,基本确定是jittor框架的问题,原因如下:

  1. tensorflow和pytorch同时同卡运行并不会出现上述错误
  2. 在发现上述问题并确认多次问题存在后,我将jittor隔离在gpu:0上运行,tensorflow和pytorch在gpu:1上运行,测试程序正常运行4h没有出现任何问题
  3. 程序发生错误后陷入了S态,这与jittor报错中出错函数cudaDeviceSynchronize()(即等待cuda设备同步消息)功能吻合
    其次,分析问题的可能原因:
  4. 怀疑是否可能是jittor在多个框架分摊单卡内存时会出现非法内存访问的bug,我查阅了jittor文档,发现jittor在描述如何占用GPU资源以及如何设置占用策略方面缺少相关信息(如tf存在贪婪和非贪婪gpu占用,tf和torch均可指定gpu设备,但jittor似乎暂无这种参数设置,也没有描述默认的占用规则,目前实践发现仅能通过环境变量更改可视GPU但这对该进程是全局生效的),因此也无法笃定,望解答
  5. 也可能是GPU内存溢出导致。因为程序运行一定时间后报错这种触发机制显得不可靠,似乎遇到某个模型恰好导致内存情况改变显得更为可信。但是一方面,jittor似乎对GPU的占用并不大,不是很能理解为何两个框架可运行的场景下加入了几百MB的jittor即会出现此类问题;另一方面,我尝试单独运行训练了jittor的一个过大的模型,模型在编译阶段似乎就会直接报错,而不会出现上述bug,且即使真的是因为内存溢出,可能相关地方需要考虑这种情况并至少阻止难以挽回的进程S态的出现。
    由于本人对于底层cuda源码等并不了解,且jittor相关资料较少,难以尝试从源码和原理角度进行分析,以上是本人通过现象进行的分析,希望能够对定位问题有所帮助,感谢jittor开发者的支持,望解答,谢谢~

谢谢您的反馈,这个原因主要是jittor使用的是统一内存而不是设备内存,在设备内存不足的时候会导致出现illegal address的错误,为了规避这一点,您可以强制jittor使用设备内存,设置环境变量use_cuda_managed_allocator=0即可,另外建议您更新到最新版本的jittor

这种难以挽回的s态,主要是因为卡在cuda驱动层面,jittor框架很难直接控制

您好,谢谢您的答复,已解决我的疑问,现已基本理解jittor总体上的内存管理及配置方案:
use_cuda_managed_allocator=1时,采用统一内存管理,内存溢出可能会出现难以控制的S态
use_cuda_managed_allocator=0时,采用设备管理,类似tf的非贪婪式,可通过环境变量指定设备,内存溢出不会出现上述情况