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

> 原文链接：https://robotics.hellobike.com/pitfalls/2026-07-08-openpi-822
> 频道：pitfalls | 更新日期：2026-07-08 | 标签：auto, github-issue, openpi
> 核心摘要：JAX多GPU微调π0.5出现Loss/grad_norm为NaN且不稳定，通过增加--fsdp-devices参数降低显存占用可解决。

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

::: v-pre
> 本文由 AI 根据公开内容整理，可能有误，请以原文为准：[Loss becomes NaN when training with multiple GPUs](https://github.com/Physical-Intelligence/openpi/issues/822)（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 等，版权归原作者所有。
:::

## 原始参考来源 (Sources)

- [GitHub Physical-Intelligence/openpi #822：Loss becomes NaN when training with multiple GPUs](https://github.com/Physical-Intelligence/openpi/issues/822)