关于使用50系显卡对应的torch导致内存爆炸的问题

[Ethan-Miao]

您好，我在使用alphazero模型训练我自己设计的桌面类游戏时遇到了内存持续增长（1gb每分钟）且最终挤爆磁盘的现象。我使用了非常小的replaybufffer（2e3）后问题没有改进。我使用了memray采集了replay buffer满后的数据，发现内存持续增长的原因源于is_available at torch/cuda/init.py。
大模型给出的可能原因是PyTorch CUDA caching allocator 在吃 host 侧 pinned memory / 映射内存。
我的训练环境是wsl2+5070ti+96g内存。（可能是wsl2的问题？）
我在lightzero0.2.0的提供的五子棋文件中做了同样的replaybuffer的限制，结果同样是内存爆炸。
由于50系显卡要求sm120，我测试了torch支持的所有版本，均没有改善。
我想问一下lightzero是否有试过使用50系显卡训练，用的torch版本以及环境是什么？
附上我是用memray采集的内存占用：

[puyuan1996]

您好 @Ethan-Miao，

我们没有在50系显卡上测试过，也没有在wsl2上测试过。目前我们主要在linux上测试开发，主要版本参考 https://github.com/opendilab/LightZero/blob/main/requirements.txt。

（下面的回答由AI辅助，仅供参考）

感谢您提供的详细报告和 memray 分析。鉴于您使用的是 RTX 50 系显卡 (SM120) 配合 WSL2，您目前的运行环境属于“前沿架构”，这种组合下的驱动程序与内核交互极易出现内存泄漏问题。

memray 将问题指向 torch/cuda/__init__.py 且内存呈线性增长 (1GB/min)，这强烈暗示了 锁页内存 (Pinned Memory) 泄漏。这通常是 PyTorch DataLoader、CUDA 驱动程序以及 WSL2 翻译层之间交互异常导致的。

以下是建议的排查和修复步骤：

1. 禁用 Pin Memory (最可能的解决方案)
LightZero (基于 DI-engine) 通常会在数据收集或 Replay Buffer 中默认开启 pin_memory 以加速 CPU 到 GPU 的传输。在 WSL2 配合新显卡驱动时，这经常导致宿主机内存无法释放。

• 操作： 检查您的 LightZero 配置文件，查找 pin_memory 或 cuda_in_collector 相关设置。
• 修改： 显式将 pin_memory 设置为 False。

  # 配置修改示例
  policy=dict(
      learn=dict(
          learner=dict(
              dataloader=dict(pin_memory=False) # 显式禁用
          ),
          # ...
      ),
      # ...
  )

2. 检查多进程 Workers
如果您在配置中设置了 num_workers > 0，在 WSL2 上，子进程反复初始化 CUDA 上下文（这解释了为什么报错在 is_available）可能会导致句柄泄漏。

• 测试： 将 num_workers 设置为 0（单进程模式）。如果内存不再爆炸，则确认是 WSL2 的多进程 CUDA 上下文管理问题。

3. 限制 WSL2 内存 (缓解措施)
WSL2 默认会吞噬几乎所有宿主机内存用于缓存。虽然这不能根治泄漏，但可以防止磁盘被挤爆。

• 在 Windows 用户目录下创建 %UserProfile%\.wslconfig 文件：

  [wsl2]
  memory=32GB  # 限制为一个安全值
  swap=8GB

• 执行 wsl --shutdown 重启生效。

4. 终极方案：原生 Linux
RTX 50 系的原生 Linux 驱动通常比 WSL2 的直通驱动更稳定。如果上述方法无效，建议尝试双系统安装 Ubuntu 进行训练。这是解决新硬件驱动层内存泄漏最彻底的方法。