报错解决-Unrecognized CachingAllocator option: 1

前言

在服务器上跑Bert模型的时候,出现了如下的报错:

分析:Traceback (most recent call last):
  File "d:\xay\model\bert-short.py", line 136, in <module>
    model.to(device)
  File "C:\py\Anaconda3\envs\py39\lib\site-packages\torch\nn\modules\module.py", line 1152, in to
    return self._apply(convert)
  File "C:\py\Anaconda3\envs\py39\lib\site-packages\torch\nn\modules\module.py", line 802, in _apply
    module._apply(fn)
  File "C:\py\Anaconda3\envs\py39\lib\site-packages\torch\nn\modules\module.py", line 802, in _apply
    module._apply(fn)
  File "C:\py\Anaconda3\envs\py39\lib\site-packages\torch\nn\modules\module.py", line 802, in _apply
    module._apply(fn)
  [Previous line repeated 1 more time]
  File "C:\py\Anaconda3\envs\py39\lib\site-packages\torch\nn\modules\module.py", line 825, in _apply
    param_applied = fn(param)
  File "C:\py\Anaconda3\envs\py39\lib\site-packages\torch\nn\modules\module.py", line 1150, in convert
    return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
  File "C:\py\Anaconda3\envs\py39\lib\site-packages\torch\cuda\__init__.py", line 302, in _lazy_init
    torch._C._cuda_init()
RuntimeError: Unrecognized CachingAllocator option: 1

尝试了很多办法都没有解决，例如更换pytorch版本、pytorchvision等,最终解决方法如下：

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:256"

插入此行代码可以解决问题，原因是使用PyTorch进行深度学习训练时，由于动态内存分配策略，显存中的内存块可能会变得非常碎片化。这导致即使有足够的总显存，也可能出现"CUDA: Out of Memory"错误，因为无法找到足够大的连续空间来分配新的内存请求。

max_split_size_mb 是 PYTORCH_CUDA_ALLOC_CONF 环境变量中的一个重要参数，它定义了CUDA内存分配器在分配内存时可以拆分的最大块大小（以MB为单位）。通过设置这个参数，可以减少显存碎片化的程度，因为它限制了分配器拆分大块连续显存的概率，从而减少了因碎片化导致的内存分配问题。

通过设置 max_split_size_mb，可以优化PyTorch的CUDA内存分配策略，减少显存碎片化，从而避免"CUDA: Out of Memory"错误。具体的设置值需要根据实际情况进行调整，以达到最佳效果。