Spacy Transformers自定义NER模型训练CUDA内存不足问题求助

使用Spacy Transformers（基于en_core_web_trf）训练包含90k条数据记录的自定义NER模型时，训练耗时异常漫长且中途中断，抛出CUDA内存不足错误。

环境配置 #

• Model: 自定义NER模型
• Library: Spacy Transformers
• Hardware: AWS EC2 g4dn.8xlarge实例（32 vCPUs，GPU显存14.76 GiB）
• Software: Python 3.6.9，spaCy 3.6.1

错误信息 #

RuntimeError: CUDA out of memory. Tried to allocate 114.00 MiB (GPU 0; 14.76 GiB total capacity; 11.19 GiB already allocated; 78.75 MiB free; 12.40 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

已尝试排查步骤 #

• 检查GPU内存占用并终止占用进程
• 调整训练批次大小

1. 优化PyTorch内存分配策略 #

设置环境变量缓解内存碎片问题，可根据实际情况调整max_split_size_mb的值（如256、512或1024）：

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512

或在Python代码开头添加：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'

2. 限制模型显存占用 #

• 训练时通过--gpu-memory参数限制显存使用比例，例如使用80%的显存：

python -m spacy train config.cfg --gpu-memory 0.8

• 在配置文件的[training]段中，将gpu_allocator设为"pytorch"，同时合理设置max_epochs，避免长期训练导致内存累积。

3. 数据预处理优化 #

• 截断长文本：将输入文本截断到Transformer模型支持的最大长度（如512 tokens），减少单条数据的显存占用。
• 启用懒加载：避免一次性加载全部90k条数据，采用分批懒加载模式降低内存压力。

4. 模型轻量化调整 #

• 替换为更小的预训练模型：比如使用en_distilbert_base_cased_trf（DistilBERT版本）或en_core_web_md，大幅降低显存消耗（若精度可接受）。
• 禁用冗余组件：在配置文件中关闭训练不需要的辅助功能，减少内存占用。

5. 训练过程内存释放 #

• 手动清除显存：在每个epoch结束后添加代码释放未使用的显存：

import torch
torch.cuda.empty_cache()

• 启用梯度累积：设置accumulate_grad_batches参数，用多个小批次的梯度累积替代大批次训练，平衡显存占用与训练效果。

6. 修复环境兼容性问题 #

Python 3.6.9与spaCy 3.6.1的兼容性存在隐患，建议升级Python到3.8-3.11（spaCy官方推荐版本），避免版本不兼容导致的内存泄漏或异常。

内容的提问来源于stack exchange，提问作者iamhimanshu0