具身开发者社区

外观

踩坑GitHub Physical-Intelligence/openpi3 个月前

JAX多卡微调π0.5时Loss与梯度变NaN，单卡正常 ​

JAX多GPU微调π0.5出现Loss/grad_norm为NaN且不稳定，通过增加--fsdp-devices参数降低显存占用可解决。

AI 整理
  • 具身模型
  • 端侧部署

本文由 AI 根据公开内容整理，可能有误，请以原文为准：Loss becomes NaN when training with multiple GPUs（GitHub Physical-Intelligence/openpi #822）

状态：有绕过办法

问题 ​

使用JAX在多GPU环境下对LIBERO或自定义数据集微调π0.5（含LoRA）时，Step 0即出现loss=nan, grad_norm=nan, param_norm=nan，或训练数十至数千步后突发NaN；相同代码与命令下结果非确定，有时正常有时NaN。切换至单GPU训练则数值稳定正常。

环境 ​

4x A6000多卡环境；JAX框架；π0.5模型；LIBERO或自定义数据集；LoRA微调；FSDP分布式训练策略。

解决办法 ​

  1. 检查多卡训练时的GPU显存与资源占用情况，确认是否存在显存不足导致的计算异常。
  2. 在训练启动命令中增加--fsdp-devices参数的值，以调整FSDP设备分片策略并降低单卡显存峰值占用。
  3. 重新启动多卡训练任务，验证loss=nan现象是否消失及训练稳定性是否恢复。

讨论原作者：@Elycyx 等，版权归原作者所有。

来源

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

  1. GitHub Physical-Intelligence/openpi #822：Loss becomes NaN when training with multiple GPUs访问 2026-10-26