国产GPU清除显存碎片

在国产 DCU 上，你之前尝试的 expandable_segments 优化方法不可用，是因为 expandable_segments 是 PyTorch 针对 NVIDIA CUDA 平台开发的专属特性，DCU 基于 ROCm 平台，其显存管理机制与 CUDA 存在差异，因此 ROCm 后端尚未实现该功能。

不过，针对 DCU 的显存碎片和 OOM 问题，业界已有一些经过验证的解决方案。

🛠️ 国产 DCU 显存优化方案

你可以按照以下优先级尝试：

1. 尝试正确的环境变量（高优先级）

虽然 expandable_segments 在 DCU 上不可用，但针对 ROCm 平台，有一个特定的环境变量可以尝试：

export PYTORCH_HIP_ALLOC_CONF=expandable_segments:True

• 这个变量的作用：它是 PYTORCH_CUDA_ALLOC_CONF 在 HIP (ROCm) 后端上的对应版本。在 DCU 上，这个变量被支持的可能性远大于 CUDA 版本。有用户反馈，设置此变量后，因显存碎片导致的 OOM 问题得到了解决。
• 使用方法：在运行训练脚本前，在终端执行上述 export 命令，或将其与训练命令放在同一行执行。

2. 调整 PyTorch 显存分配策略（中优先级）

如果上述方法无效，可以通过调整 PyTorch 的显存分配策略来缓解碎片问题。

• 禁用缓存分配器：设置 PYTORCH_NO_CUDA_MEMORY_CACHING=1 可以让 PyTorch 在每次需要时直接向驱动申请和释放显存，而不是缓存起来。这能避免缓存导致的内存碎片，但会略微降低运行速度。有报告称此方法能有效改善问题。
• 设置最大分割大小：通过 PYTORCH_HIP_ALLOC_CONF=max_split_size_mb:128 可以限制 PyTorch 缓存分配器单次分配的最大内存块。这能避免产生过大且无法利用的显存碎片，但设置过小也可能影响性能，需要根据你的模型大小进行实验。

3. 使用替代内存分配器（备选方案）

如果 PyTorch 自带的分配器效果不佳，可以考虑使用系统级的、专为高性能计算优化的内存分配器，如 jemalloc 或 tcmalloc。

• 原理：这类分配器在处理大量、频繁的内存分配和释放时，能更有效地管理内存，减少碎片化。
• 使用方法：
  1. 安装：apt-get install libjemalloc-dev (或 libtcmalloc-minimal4)
  2. 预加载：在运行训练脚本前，通过 LD_PRELOAD 环境变量预加载：

     LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so python -u SwinUNETR/main_eval.py

4. 基础但关键的训练策略调整

在尝试上述高级优化之前，请确保已经应用了以下基础但至关重要的策略：

• 减小 batch_size：这是最直接、最有效的方法。将 batch_size 减小到 2 甚至 1，可以显著降低单次前向传播所需的显存，是缓解 OOM 最根本的手段。
• 启用梯度检查点 (use_checkpoint)：在你的 SwinUNETR 初始化中，将 use_checkpoint=True。这能以计算时间换取显存空间，是训练大模型的常用技巧。
• 使用混合精度训练 (AMP)：在训练循环中使用 torch.cuda.amp 进行自动混合精度训练，可以将显存占用降低约一半。

📌 总结与建议

1. 首选尝试：设置 PYTORCH_HIP_ALLOC_CONF=expandable_segments:True。
2. 如果仍然 OOM：则说明碎片问题比较严重，建议按顺序尝试：
   • 降低 batch_size 或启用 use_checkpoint。
   • 测试 PYTORCH_NO_CUDA_MEMORY_CACHING=1。
   • 考虑使用 jemalloc 或 tcmalloc。

你可以先从最简单的入手，逐步尝试这些方案，应该能找到适合你 DCU 环境的显存优化方法。