Skip to content

在原生 Ubuntu 环境下进行 1.7B 模型全量微调时发生 OOM,但在 WSL2 下却能正常运行 #199

Description

@freenewfree

问题描述:
在显存为 16GB 的原生 Ubuntu 系统中,进行 moss_tts_local 1.7B 模型全量微调(Full SFT)时持续遇到 CUDA Out of Memory (OOM) 问题。奇怪的是,同样的训练链路在 WSL2 环境下可以运行(虽然速度较慢)。这表明两种环境在显存管理机制上存在显著差异。

环境配置:

GPU: 16GB VRAM

操作系统: 原生 Ubuntu (训练失败) vs. WSL2 (训练成功)

框架: DeepSpeed ZeRO-3, Accelerate

任务: 全量微调 (Full SFT)

配置: accelerate_zero3_1.7b.yaml

“内存悖论” (WSL2 与原生 Linux 的区别):

原生 Ubuntu: 训练进程频繁被系统强制终止,报 Signal 15 (SIGTERM/OOM Killer)。即便使用了 DeepSpeed ZeRO-3,在前向/反向传播期间(激活值缓存 + 梯度累积),显存峰值触及了 16GB 的硬件极限,导致进程被操作系统强行杀掉。

WSL2: 训练流程可以顺利完成。我们推测这是因为 WSL2 的 WDDM 驱动机制允许通过向系统 RAM “借用”空间作为虚拟显存(Virtual VRAM),从而缓冲了显存溢出。相比之下,原生 Linux 对物理显存限制非常严格,一旦超限即触发强制终止。

已排查步骤:

确认 CUDA 环境及 nvcc 配置正确 (12.4)。

正确配置了 DeepSpeed ZeRO-3 (Full sharding)。

将 gradient-accumulation-steps 调整为 16,以降低单次迭代的显存峰值。

强制使用 sdpa 注意力实现。

确认所有环境变量(如 CUDA_HOME)配置无误。

诉求/反馈:
目前看来,1.7B 模型全量微调确实需要超过 16GB 的物理显存,这对消费级显卡环境造成了硬性门槛。由于目前 MOSS-TTS 的实现未原生支持 LoRA/PEFT,这使得 16GB 显存的用户很难进行模型微调。

优化建议: 项目组是否有计划进行显存优化以降低激活值的内存占用?

兼容性: 是否有推荐的方法(或配置)来集成 LoRA/PEFT?或者在不依赖类似 WSL2 虚拟内存机制的前提下,有什么进一步降低显存需求的方法?

报错日志核心逻辑总结
1.参数解析阶段(已修复):

现象: sft.py: error: unrecognized arguments: --gradient_checkpointing

原因: 手动启动脚本时,使用了错误的参数格式(下划线而非连字符)。

结论: 该问题通过更正参数格式及转向 run_train.sh 官方脚本已解决。

2.环境依赖阶段(已修复):

现象: deepspeed.ops.op_builder.builder.MissingCUDAException: CUDA_HOME does not exist

原因: deepspeed 深度绑定了 nvcc 编译器进行算子编译,而原生 Linux 环境缺少该路径,导致检查机制报错。

结论: 通过安装 nvcc 并建立软链接(/usr/local/cuda-12.4/bin/nvcc)已规避。

3.运行崩溃阶段(物理瓶颈,无法修复):

现象: torch.OutOfMemoryError 以及随后的 SignalException: Process got signal: 15 (SIGTERM)。

深度原因:

内存溢出(OOM): 模型加载至显存后,反向传播过程中产生的激活值(Activations)和梯度状态在 16GB 显存内无法容纳,即便使用 ZeRO-3 将参数分片,依然触及物理显存上限。

系统干预(OOM Killer): 当显存耗尽时,Linux 内核识别到该进程正在拖垮系统性能,遂通过 Signal 15 强行杀死了训练进程。

结论: 这是一个硬性的物理资源瓶颈。原生 Ubuntu 系统严格执行了 16GB 的物理上限,而 WSL2 通过虚拟内存机制(WDDM)将其转嫁到了系统 RAM 和 Pagefile 中,才得以勉强运行。

希望能获得一些指导,谢谢!

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions