[Hubo1231]

Summary

Improve JAX training throughput and stability for large-scale Pi0.5 training, with a validated FP16 cuDNN fused-attention path and an
explicit-attention fallback.

This PR also adds Dobot Nova5 dataset/policy support, refreshes production training configs, and updates the Flexiv inference recipes
to use serial grippers.

Changes

• Reduce DataLoader IPC overhead by transferring CPU Torch tensors through shared memory and constructing JAX arrays from zero-copy
  NumPy views.
• Explicitly shut down DataLoader workers when training finishes.
• Compute gradient and parameter norms only on logging steps.
• Optionally batch all camera views into one SigLIP encoder call.
• Add cuDNN fused attention for training:
  • use FP16 compute with dynamic loss scaling and a custom VJP;
  • zero query gradients for fully masked attention rows;
  • preserve explicit attention as the default and inference fallback;
  • add consistent sharding constraints for cuDNN operands.
• Update the CUDA stack to PyTorch 2.11 / cuDNN 9.19 and add scripts/check_cuda_stack.py.
• Add Dobot Nova5 LeRobot transforms, data configuration, registry entry, tests, and an example training config.
• Refresh the H100/H200 training examples and consolidate the optimization documentation.
• Update the Flexiv diagonal-02 and forward-04 inference recipes to use serial grippers.

Validation

• 70 passed:
  • src/openpi/models/pi0_test.py
  • src/openpi/training/data_loader_test.py
  • src/openpi/policies/dobot_nova5_policy_test.py
  • src/openpi/training/config_yaml_test.py
• git diff --check passes.
• FP16 cuDNN attention was previously validated for 3,000 steps from step 0 on 8×H100 with fixed seed and strict batch ordering.
• At the historical BF16 divergence point, FP16 loss remained within 0.5% of the explicit-attention baseline, and the recorded
  gradient norm remained below 1.3× the baseline.
• The optimized H100 configuration reduced measured step time from 1.328 s/step to 1.193 s/step, approximately 10.2%.

Notes

• Explicit attention remains the default. Production cuDNN training should use:

  model:
    use_cudnn_attention: true
    cudnn_attention_dtype: float16
  

• BF16 cuDNN attention is not production-safe; it diverged from the explicit baseline after approximately 1,000 steps.

• The FP16 implementation calls private JAX 0.5.3 cuDNN APIs and must be revalidated when upgrading JAX, CUDA, or cuDNN.

• max_token_len: 128 was validated only for the referenced earbuds dataset and must be checked again for other datasets and prompts.

[Vertax42]

我在这个分支上直接补了 6 个 commit(bbccfae..7ac2a6a),都是 review 里发现的问题。逐条说明:

1. bbccfae — 恢复三个被删但仍被引用的 config

4606754 和 e13ac0f 删掉了 pi05_droid、pi05_base_xtac_umi_pick_up_cube_0807_h200、pack_6_cosmetic_bottles_lora,但没改引用它们的地方,而 get_config() 找不到名字是直接抛 ValueError:

• xtac_umi_policy_test.py:112,126 两个测试直接失败
• serve_policy.py:70 的 DEFAULT_CHECKPOINT[EnvMode.DROID] 失效(这是那张表里唯一还能解析的一项)
• README、两个 examples README、_FULL_REFERENCE.yaml 的链接全部悬空

已从 7c51dc5 原样恢复。earbuds_0520_a100 保持删除 —— 0904 的 h100/h200 取代了它,而且只有 changelog 提到它。

2. 1db0ec9 — NaN loss 不能被 nanmean 吞掉

范数只在 log 步算,其余步填 NaN 占位,所以 reduce 整体换成了 jnp.nanmean。但这同样作用在 loss 上:一个发散的步骤现在会被悄悄从平均里剔除,而不是显示成 loss=nan。考虑到 per-step finite guard 已经移除、grad_norm 每 log_interval 才采样一次,这正好是 FP16 路径最需要的信号。现在只有 grad_norm/param_norm 走 nanmean。

3. bc477a6 — cuDNN 私有符号改为惰性解析 + 校验

顶层 import 那个私有模块其实没必要 —— jax.nn 自己就 import 了它,真消失的话 import jax 先挂。真正会坏的是符号名:_dot_product_attention_bwd_rule 之类被重命名,会在训练中从被 trace 的 VJP 内部抛 AttributeError。改成首次使用时解析并校验用到的 7 个符号,升级 JAX 后一次性给出可操作报错(点名缺失符号 + use_cudnn_attention=false 退路)。显式 attention、纯 CPU、推理都不会走到这个函数。

4. 9cbea2c — check_cuda_stack.py 补上 FP16 路径,版本比较放宽

README 说这个脚本是生产训练前的必过闸门,但 section 3 只跑了 bf16 的 jax.nn.dot_product_attention —— 也就是这个 PR 自己标记为不安全的那条路径。文档要求启用的 use_cudnn_attention: true + cudnn_attention_dtype: float16 走的是 _cudnn_attention_in_dtype,完全没被验证过。新增 section 4,在同样的生产 shape 上对 FP16 custom VJP 求导并与同一个 fp32 参考对比,同时断言全空 mask 行的 dQ 必须是 0(不是 NaN、不是垃圾值)。

版本比较改成只比 major.minor:pyproject 锁的是区间(torch>=2.11,<2.12 等),而脚本里写死 2.11.0+cu128,任何上游 patch 发布都会让闸门误报。+cu128 后缀仍然单独检查 —— CPU wheel 装得上,而且会静默掉这个脚本存在的全部意义。

5. fef66f2 — cube-color-sorting config 的复制粘贴

头部注释整块从 earbuds 配置复制,启动命令还写着训 earbuds 那个配置,已改。

这里有一点需要你确认: max_token_len 的注释也是复制来的,声称 128 是"Scanned over this dataset (4037 states)",但那 4037 个状态是耳机盒数据集的,docs/training-optimization.md 里也明确写了 128 只在那个数据集验证过、不要直接套用。我没敢替你断言它扫过 cube-color-sorting-0821,改成了标注"沿用自 earbuds、未重新扫描"。如果你确实扫过,把实测数字填回去就行;tokenizer 溢出只是 warning + 截断,训练里不会暴露出来,所以这个值得确认一下。

6. 7ac2a6a — data_loader close() 的 torch 私有依赖

_shutdown_workers / _workers / _iterator 都没有公开等价物,在 docstring 里记下这是 torch>=2.11,<2.12 这个 pin 在保护的东西,以及 getattr 保护会退化成旧的 __del__ 回收而不是抛错。

---

验证

145 passed                policies + training + examples
149 passed, 3 failed      3 个失败是 policy_test / policy_rtc_test,在基线分支上同样失败
                          (本机 8GB 显存 OOM,与改动无关)
ruff check / ruff format  通过

xtac_umi_policy_test 的两个失败已消失。另外单独验证了:版本闸门对 patch 升级放行、对 minor 升级和 CPU wheel 拦截;删掉一个 cuDNN 私有符号会立刻报错;NaN loss 确实能穿过 reduce 到达日志。

没动但建议单独处理的

• torch 2.6→2.11 是个大跨度升级,也会影响 PyTorch 训练路径和 transformers==5.3.0 兼容层,但验证清单里没有任何 PyTorch 测试。
• Flexiv taccap_follower→serial 夹爪 recipe 属于机器人硬件配置,和训练吞吐无关(它能通过 examples/flexiv_recipe_test.py,57 passed),建议拆成独立 PR。

[Vertax42]

见 review 评论;FP16 路径待 GPU 机验证