具身开发者社区

外观

踩坑GitHub unitreerobotics/unifolm-world-model-action9 个月前

UnifoLM-WMA训练CUDA OOM：4x24G显存溢出报错 ​

4x24G显卡微调UnifoLM-WMA-0_Base遇CUDA OOM，调整DeepSpeed ZeRO-3及batch_size无效，建议对齐多GPU配置参数。

AI 整理
  • 仿真与训练
  • 具身模型

宇树科技宇树 Z1

本文由 AI 根据公开内容整理，可能有误，请以原文为准：4x24G NVIDIA显卡模型训练,内存不足（GitHub unitreerobotics/unifolm-world-model-action #35）

状态：有绕过办法

问题 ​

使用4x24G NVIDIA显卡基于Z1_StackBox_Dataset微调UnifoLM-WMA-0_Base模型时，触发torch.cuda.OutOfMemoryError。即使将batch_size降至1、启用DeepSpeed ZeRO-3并将优化器卸载至CPU、设置decision_making_only: True，显存仍耗尽。单卡24G环境下分配5.79GiB亦报错。

环境 ​

4x24G NVIDIA显卡；PyTorch Lightning；DeepSpeed ZeRO-3；UnifoLM-WMA-0_Base模型；Z1_StackBox_Dataset数据集。

解决办法 ​

  1. 检查训练启动脚本，确保设备数量、GPU数量和nproc_per_node三个参数设置为完全相同的值，以正确初始化多GPU分布式训练环境。
  2. 若单卡或错误配置下仍OOM，确认已应用DeepSpeed ZeRO-3策略并开启offload_optimizer.device: cpu。
  3. 保持batch_size: 1与decision_making_only: True作为基础显存优化配置。
  4. 如报错提示显存碎片化，尝试设置环境变量PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True。

讨论原作者：@John-Stones 等，版权归原作者所有。

来源

本页事实对应的公开页面，以原文为准。

  1. GitHub unitreerobotics/unifolm-world-model-action #35：4x24G NVIDIA显卡模型训练,内存不足访问 2026-10-26