CUDA OOM 排查实践：从显存爆掉到复现指标
1. 背景
目标是使用 Docker 在 GPU 服务器上跑通一套代码，并复现 README 中给出的结果。
项目入口类似：
inference/run_classifier_multiview.py给出的环境要求包括 Python 3.9、CUDA 12.8、PyTorch 2.8.0、transformers、scikit-learn、dpkt 等依赖，并且数据处理流程包括 PCAP 处理、bi-gram 编码、训练/验证/测试集划分，以及基于 Dirichlet 不确定性的 open-set 评估。前期我已经完成了 Docker 环境部署、数据路径确认和快速训练验证。后续真正开始复现指标时，遇到的核心问题就是：
RuntimeError: CUDA out of memory这篇文章主要复盘这次 CUDA OOM 的排查过程，以及我最终如何确定稳定训练配置并复现 README 指标。
2. OOM 是什么
OOM 是 Out of Memory，意思是内存不足。
在深度学习训练中，最常见的是 CUDA OOM，也就是 GPU 显存不足。它和普通服务器内存不一样。
比如服务器有一张 40GB 显存的 A100，看起来显存很大，但如果已经有其他进程占用了十几 GB，实际可用于训练的显存就会少很多。训练 Transformer 类模型时，还需要保存前向传播的中间激活、梯度、优化器状态等内容，显存开销会明显高于推理。
本次项目中，模型不是简单的传统机器学习模型，而是带有 Transformer/UER 结构、多视角输入和 Dirichlet 不确定性建模的深度学习分类模型。项目目录中包含 uer 框架、bert_base_config.json、pre-trained_model.bin 等文件，README 也说明了会微调预训练模型并计算 Dirichlet 不确定性。因此，CUDA OOM 并不意外。
3. 第一次 OOM：默认 batch_size 太大
恢复训练代码后，我第一次直接按默认配置运行。脚本能正常读取数据，也能进入训练阶段，但随后在模型 forward 阶段报错：
RuntimeError: CUDA out of memory.
Tried to allocate 384.00 MiB
GPU 0; 39.49 GiB total capacity;
22.81 GiB already allocated;
337.44 MiB free;
23.39 GiB reserved in total by PyTorch这个报错信息很关键，它说明了几件事：
1. GPU 总显存约 39.49 GiB；
2. PyTorch 已经占用了 22GB 以上；
3. 剩余可用显存只有几百 MiB；
4. 模型还想再申请 384 MiB，但已经不够了。工作日志里也记录了：默认 batch_size=1024 太大，加上 GPU 上已有其他进程占用显存，多视角 Transformer attention 计算显存开销较高，所以训练进入 forward 阶段后触发 CUDA OOM。
4. 第一反应：不要盲目 kill 进程
遇到 OOM 时，第一反应不能是直接杀进程。
在多人共用 GPU 服务器时，nvidia-smi 里看到的进程可能是别人正在运行的服务或者实验。直接 kill 可能影响别人工作。
我的处理方式是先查看 GPU 状态：
nvidia-smi重点看：
Memory-Usage
GPU-Util
Processes如果 GPU 上已有其他进程长期占用显存，就应该优先考虑：
1. 换另一张 GPU；
2. 降低 batch size；
3. 等待空闲时间；
4. 和负责人确认是否可以释放 GPU。本次实践中，我没有直接 kill 其他进程，而是选择通过调整 batch size 来解决。
5. 为什么 batch size 会影响显存
batch_size 表示一次送入模型训练的样本数量。
对于深度学习训练，显存大致会被这些部分占用：
模型参数
梯度
优化器状态
输入 batch
中间激活
attention 计算缓存
loss 相关中间变量其中，batch size 越大，输入数据和中间激活越大，显存占用越高。
本项目中还有一个额外因素：多视角输入。也就是说，每个流量样本并不是单一路径输入，而是会被拆成多个 view，再分别进入编码器。这会进一步增加显存消耗。
所以虽然 A100 40GB 很强，但默认 batch_size=1024 对这个模型来说仍然太大。
6. 逐步降低 batch size，而不是一次改到很小
我没有直接随便改一个 batch size，而是做了一个简单的 batch 探测实验。
6.1 batch_size=16：能跑通
先把 batch size 从默认值降到 16，并且只跑 1 个 epoch 做快速验证：
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 python run_classifier_multiview.py \
--seq_length 128 \
--labels_num 200 \
--threshold 1 \
--lambda_epochs 1 \
--epochs_num 1 \
--batch_size 16 \
--packet_num 4 \
--packet_len 128 \
--device_id 0 \
--dataset_name android \
--train_dataset ./processed_data/android_train.txt \
--valid_dataset ./processed_data/android_valid.txt \
--test_dataset ./processed_data/android_encrypted_test.txt \
--open_dataset ./processed_data/attack_encrypted_open.txt结果是能跑通。单轮训练耗时约 248 秒，并成功生成 1.2G 的 fine-tuned 模型。这个阶段的准确率很低，Closed Balanced Acc 为 0.0050，Open Balanced Acc 为 0.0000，但这很正常，因为这里只是 1 epoch 快速验证，不代表模型真实性能。
这一步的价值是：
证明环境、数据读取、训练循环、模型保存和推理流程都能跑通。6.2 batch_size=32：稳定可用
接着我把 batch size 提到 32：
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 python run_classifier_multiview.py \
--seq_length 128 \
--labels_num 200 \
--threshold 1 \
--lambda_epochs 1 \
--epochs_num 1 \
--batch_size 32 \
--packet_num 4 \
--packet_len 128 \
--device_id 0 \
--dataset_name android \
--train_dataset ./processed_data/android_train.txt \
--valid_dataset ./processed_data/android_valid.txt \
--test_dataset ./processed_data/android_encrypted_test.txt \
--open_dataset ./processed_data/attack_encrypted_open.txt \
2>&1 | tee batch32_probe.logbatch_size=32 可以正常跑完 1 epoch，说明当前 GPU 剩余显存下，32 是一个可用配置。
6.3 batch_size=64：再次 OOM
继续尝试 batch_size=64 时，程序再次 OOM：
RuntimeError: CUDA out of memory.
Tried to allocate 96.00 MiB
GPU 0; 39.49 GiB total capacity;
23.54 GiB already allocated;
1.44 MiB free;
23.72 GiB reserved in total by PyTorch这说明在当前 GPU 占用情况下，64 已经超出安全范围。因此，我最终确定：
batch_size=32 是当前最合适的训练配置。7. PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 的作用
我在训练命令前加了：
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128它主要用于缓解 PyTorch 显存碎片问题。
有时 GPU 总显存看起来还有剩余，但由于显存被分割成很多不连续的小块，PyTorch 找不到足够连续的显存块，也可能 OOM。这个参数可以调整 PyTorch 的显存分配策略，降低碎片化风险。
但是它不是万能的。
如果 batch size 本身太大，显存总量确实不够，那还是会 OOM。本次 batch_size=64 就是这种情况：即使加了这个参数，仍然不够。
8. 从“能跑”到“能复现指标”
确定 batch_size=32 可用后，我没有直接跑默认 50 epoch，而是先跑了一个小完整测试：
batch_size=32
epochs_num=5
lambda_epochs=205 epoch 训练耗时约 17 分钟，结果如下：
加密应用流量 Closed Balanced Acc = 0.9187
非加密应用流量 Closed Balanced Acc = 0.9557这说明模型已经学起来了，closed-set 分类基本接近 README 目标。但当时 open-set 异常检测还很低，说明仅 5 epoch 可能不够充分，尤其是 Dirichlet 不确定性建模还没有完全收敛。
随后我继续使用：
batch_size=32
epochs_num=20
lambda_epochs=20进行完整训练。训练命令如下：
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 python run_classifier_multiview.py \
--seq_length 128 \
--labels_num 200 \
--threshold 1 \
--lambda_epochs 20 \
--epochs_num 20 \
--batch_size 32 \
--packet_num 4 \
--packet_len 128 \
--device_id 0 \
--dataset_name android \
--train_dataset ./processed_data/android_train.txt \
--valid_dataset ./processed_data/android_valid.txt \
--test_dataset ./processed_data/android_encrypted_test.txt \
--open_dataset ./processed_data/attack_encrypted_open.txt \
2>&1 | tee train_bs32_e20_lambda20_retry.log最终 20 epoch 训练耗时约 1 小时 08 分钟，并成功完成测试。训练日志显示，batch_size=32、epochs_num=20、lambda_epochs=20 下，程序正常读取训练集、验证集、测试集，完成 20/20 轮训练，并在测试集上得到 Closed Balanced Acc = 0.9438。
9. 最终复现结果
在当前可运行 Docker 环境下，虽然没有完全使用 README 的 Python 3.9 + PyTorch 2.8.0 环境，但通过调整 batch size 和训练轮数，仍然成功复现并超过了 README 中的核心指标。10. 为什么 1 epoch 结果很低，而 20 epoch 结果很高
一开始 1 epoch 的结果是：
Closed Balanced Acc = 0.0050
Open Balanced Acc = 0.0000看起来几乎没效果。但这不是模型坏了，而是训练远远不充分。
对于 200 类分类任务，随机猜测的水平本来就很低。1 epoch 只是为了测试代码链路，不足以让模型学习到有效分类能力。
更关键的是，这个项目还使用 Dirichlet 不确定性做 open-set 检测。open-set 检测不仅要学会“这个样本属于哪一类”，还要学会“这个样本是不是未知类”。这种不确定性建模通常需要更充分训练。
所以实验过程呈现出很典型的趋势：
1 epoch：链路验证，指标很低
5 epoch：closed 分类明显提升，open 检测仍弱
20 epoch：closed 和 open 都达到 README 目标这也是为什么实验日志不能只看一次结果，而要结合训练配置解释。
11. 这次 OOM 排查的经验总结
11.1 先看显存，再改参数
遇到 CUDA OOM，第一步应该是：
nvidia-smi看清楚：
GPU 总显存
当前显存占用
剩余显存
是否有其他进程11.2 不要一上来就跑完整训练
我这次的顺序是：
batch=16, epoch=1：验证链路
batch=32, epoch=1：探测稳定 batch
batch=64, epoch=1：确认显存上限
batch=32, epoch=5：小完整测试
batch=32, epoch=20：正式复现这种方式比直接跑 50 epoch 更稳，因为可以尽早发现 OOM、训练无效、模型未保存等问题。
11.3 OOM 不是代码逻辑错误
这点很重要。OOM 通常是资源配置问题，不是算法一定错了。关键是：
为什么 OOM
怎么定位
改了什么参数
每次实验结果是什么
最终如何确定稳定配置本次 OOM 的根因可以总结为：
默认 batch_size=1024 过大；
模型是多视角 Transformer，显存开销高；
训练阶段需要保存梯度和中间激活；
GPU 上已有其他进程占用显存；
因此需要降低 batch size。11.4 batch size 变小会影响速度和训练行为
把 batch size 从 1024 降到 32 后，单个 epoch 的 step 数会明显增加。比如 batch_size=32 时，训练实例数显示为 632。这意味着训练更慢，但能稳定运行。所以实际工程中经常要在这几者之间平衡：
显存占用
训练速度
指标复现
GPU 共享情况11.5 记录日志非常重要
我所有关键命令都加了：
2>&1 | tee xxx.log这样做的好处是：
终端能实时看到输出；
日志文件能留存证据；
后续可以整理实验表；
出问题时可以回溯。这次最终也把训练日志、评估日志和模型信息统一打包归档，方便后续交付和复盘。
12. 怎么讲
“你遇到过深度学习训练的 OOM 吗”
遇到过。我在部署一个模型时，恢复训练代码后默认
batch_size=1024，训练进入 Transformer forward 阶段后触发 CUDA OOM。通过nvidia-smi发现 A100 40GB 显卡上已有其他进程占用显存，实际可用显存不足。我没有直接 kill 进程，而是通过逐步降低 batch size 做探测：
batch_size=16能跑通，32稳定可用，64仍然 OOM，因此确定当前环境下的安全配置为batch_size=32。随后用batch_size=32、epochs_num=20、lambda_epochs=20完成完整训练，最终复现了 README 的指标。这个过程中我也保留了训练日志和评估日志，并把可运行 Docker 环境固化成镜像，保证后续可以复现。
这个回答能体现三点：
1. 我知道 OOM 的原理；
2. 我会用工程方法定位问题；
3. 我最后不是只解决报错，而是完成了指标复现。13. 最终总结
这次 OOM 排查让我最大的感受是：深度学习工程不是简单地“跑一条命令”，而是要在模型、数据、GPU、Docker 环境、训练参数之间做系统性排查。
本次实践中，我从默认 batch size OOM 开始，逐步定位到显存不足的原因，并通过 batch size 探测确定了稳定配置。最终在 batch_size=32 下完成 20 epoch 训练，并复现 README 指标。
整个过程可以总结为：
看错误日志
查 GPU 状态
确认是否资源问题
降低 batch size
做 batch 探测
小规模训练验证趋势
完整训练复现指标
保存模型和日志
总结可复用经验更多推荐
9
0- 0
已为社区贡献10条内容
所有评论(0)