From 9aef52dba290413d3599ae23d454a5b02b3bdea0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Thu, 27 Aug 2026 15:22:30 +0000
Subject: [PATCH 01/28] Optimize JAX training data pipeline

---
 scripts/train.py                        |  31 +++-
 src/openpi/training/config.py           |   9 ++
 src/openpi/training/data_loader.py      | 182 +++++++++++++++++++-----
 src/openpi/training/data_loader_test.py |  72 ++++++++++
 4 files changed, 261 insertions(+), 33 deletions(-)

diff --git a/scripts/train.py b/scripts/train.py
index 09d48e3..da3bb3a 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -280,6 +280,15 @@ def main(config: _config.TrainConfig):
             train_state, info = ptrain_step(train_rng, train_state, batch)
         t_dispatch = time.monotonic() - t0
 
+        # JAX dispatch is asynchronous. Fine-grained H2D profiling must first finish the
+        # current train step; otherwise block_until_ready(batch) also measures time spent
+        # queued behind model compute on the same devices.
+        t_compute_sync = 0.0
+        if config.profile_data_pipeline:
+            t0 = time.monotonic()
+            jax.block_until_ready((train_state, info))
+            t_compute_sync = time.monotonic() - t0
+
         infos.append(info)
 
         t_log = 0.0
@@ -296,9 +305,12 @@ def main(config: _config.TrainConfig):
         t0 = time.monotonic()
         batch = next(data_iter)
         t_next_batch = time.monotonic() - t0
+        data_profile = data_loader.profile_stats()
 
         t_ckpt = 0.0
-        if (step % config.save_interval == 0 and step > start_step) or step == config.num_train_steps - 1:
+        if (step % config.save_interval == 0 and step > start_step) or (
+            config.save_final_checkpoint and step == config.num_train_steps - 1
+        ):
             t0 = time.monotonic()
             _checkpoints.save_state(checkpoint_manager, train_state, data_loader, step)
             t_ckpt = time.monotonic() - t0
@@ -323,7 +335,24 @@ def main(config: _config.TrainConfig):
                 f"dispatch={t_dispatch:.2f}s next_batch={t_next_batch:.2f}s "
                 f"log={t_log:.2f}s ckpt={t_ckpt:.2f}s"
             )
+        if (
+            config.profile_data_pipeline
+            and data_profile is not None
+            and step % config.profile_log_interval == 0
+        ):
+            pbar.write(
+                f"[DATA_PROFILE step={step}] "
+                f"dispatch={t_dispatch:.4f}s compute_sync={t_compute_sync:.4f}s "
+                f"main_queue_wait={data_profile['main_queue_wait_s']:.4f}s "
+                f"worker_getitem={data_profile['worker_getitem_s']:.4f}s "
+                f"worker_collate={data_profile['worker_collate_s']:.4f}s "
+                f"jax_array_construct={data_profile['jax_array_construct_s']:.4f}s "
+                f"h2d_wait={data_profile['h2d_wait_s']:.4f}s "
+                f"next_batch_total={t_next_batch:.4f}s"
+            )
 
+    logging.info("Shutting down data loader")
+    data_loader.close()
     logging.info("Waiting for checkpoint manager to finish")
     checkpoint_manager.wait_until_finished()
 
diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py
index d580cc5..1b9a11f 100644
--- a/src/openpi/training/config.py
+++ b/src/openpi/training/config.py
@@ -615,6 +615,12 @@ class TrainConfig:
     # workers finish them. The sampler still visits every index exactly once per pass, so this
     # only reorders batches that were already randomly shuffled.
     strict_batch_order: bool = False
+    # If true, collect fine-grained data-pipeline timings. This adds per-sample clocks in
+    # workers and synchronizes device transfers, so it is intended only for short profiling
+    # runs rather than production training.
+    profile_data_pipeline: bool = False
+    # Number of steps between fine-grained data-pipeline timing lines.
+    profile_log_interval: int = 1
     # Number of train steps (batches) to run.
     num_train_steps: int = 30_000
 
@@ -622,6 +628,9 @@ class TrainConfig:
     log_interval: int = 100
     # How often (in steps) to save checkpoints.
     save_interval: int = 1000
+    # Whether to save the last step even when it is not on save_interval. Useful to disable
+    # for short profiling runs restored from an existing checkpoint.
+    save_final_checkpoint: bool = True
     # If set, any existing checkpoints matching step % keep_period == 0 will not be deleted.
     keep_period: int | None = 5000
 
diff --git a/src/openpi/training/data_loader.py b/src/openpi/training/data_loader.py
index 4bd5ef2..ad7a01e 100644
--- a/src/openpi/training/data_loader.py
+++ b/src/openpi/training/data_loader.py
@@ -1,4 +1,5 @@
 from collections.abc import Iterator, Sequence
+import functools
 import logging
 import multiprocessing
 import os
@@ -89,6 +90,28 @@ def data_config(self) -> _config.DataConfig:
     def __iter__(self) -> Iterator[T_co]:
         raise NotImplementedError("Subclasses of DataLoader should implement __iter__.")
 
+    def profile_stats(self) -> dict[str, float] | None:
+        """Return timings for the most recently yielded batch, when profiling is enabled."""
+        return None
+
+    def close(self) -> None:
+        """Release worker processes and other loader resources."""
+
+
+class _TimedDataset:
+    """Attach worker-side __getitem__ latency to samples for profiling."""
+
+    def __init__(self, dataset: Dataset):
+        self._dataset = dataset
+
+    def __getitem__(self, index: SupportsIndex):
+        started = time.monotonic()
+        sample = self._dataset[index]
+        return sample, time.monotonic() - started
+
+    def __len__(self) -> int:
+        return len(self._dataset)
+
 
 class TransformedDataset(Dataset[T_co]):
     def __init__(self, dataset: Dataset, transforms: Sequence[_transforms.DataTransformFn]):
@@ -385,6 +408,7 @@ def create_data_loader(
         skip_norm_stats=skip_norm_stats,
         framework=framework,
         strict_batch_order=config.strict_batch_order,
+        profile_data_pipeline=config.profile_data_pipeline,
     )
 
 
@@ -402,6 +426,7 @@ def create_torch_data_loader(
     seed: int = 0,
     framework: str = "jax",
     strict_batch_order: bool = False,
+    profile_data_pipeline: bool = False,
 ) -> DataLoader[tuple[_model.Observation, _model.Actions]]:
     """Create a data loader for training.
 
@@ -454,6 +479,7 @@ def create_torch_data_loader(
         seed=seed,
         framework=framework,
         strict_batch_order=strict_batch_order,
+        profile_data_pipeline=profile_data_pipeline,
     )
 
     return DataLoaderImpl(data_config, data_loader)
@@ -551,6 +577,7 @@ def __init__(
         seed: int = 0,
         framework: str = "jax",
         strict_batch_order: bool = False,
+        profile_data_pipeline: bool = False,
     ):
         """Create a PyTorch data loader.
 
@@ -572,6 +599,9 @@ def __init__(
         if jax.process_count() > 1:
             raise NotImplementedError("Data loading with multiple processes is not supported.")
 
+        if profile_data_pipeline:
+            dataset = _TimedDataset(dataset)
+
         if len(dataset) < local_batch_size:
             raise ValueError(f"Local batch size ({local_batch_size}) is larger than the dataset size ({len(dataset)}).")
 
@@ -584,6 +614,9 @@ def __init__(
                 jax.sharding.PartitionSpec("B"),
             )
         self._num_batches = num_batches
+        self._profile_data_pipeline = profile_data_pipeline
+        self._last_profile_stats: dict[str, float] | None = None
+        self._active_iterator = None
 
         mp_context = None
         if num_workers > 0:
@@ -607,7 +640,7 @@ def __init__(
             prefetch_factor=16
             if num_workers > 0
             else None,  # stall-fix: 32 over-buffered for cold start, 4 was too low vs slow workers, 16 is the sweet spot
-            collate_fn=_collate_fn,
+            collate_fn=functools.partial(_collate_fn, profile=profile_data_pipeline),
             worker_init_fn=_worker_init_fn,
             drop_last=True,
             # Hand batches over as workers finish them. In strict order, the one worker that is
@@ -625,38 +658,114 @@ def torch_loader(self) -> torch.utils.data.DataLoader:
 
     def __iter__(self):
         num_items = 0
-        while True:
-            t_inner_iter = time.monotonic()
-            logging.info("[TorchDataLoader] calling iter(self._data_loader) (spawning workers) ...")
-            data_iter = iter(self._data_loader)
-            logging.info(f"[TorchDataLoader] iter() returned in {time.monotonic() - t_inner_iter:.1f}s")
+        try:
             while True:
-                if self._num_batches is not None and num_items >= self._num_batches:
-                    return
-                t_batch = time.monotonic()
-                try:
-                    batch = next(data_iter)
-                except StopIteration:
-                    break  # We've exhausted the dataset. Create a new iterator and start over.
-                dt = time.monotonic() - t_batch
-                if num_items < 3 or dt > 3.0:
-                    logging.info(f"[TorchDataLoader] batch #{num_items} next() took {dt:.2f}s")
-                num_items += 1
-                # For JAX, convert to sharded arrays; for PyTorch, return torch tensors
-                if self._sharding is not None:
-                    yield jax.tree.map(
-                        lambda x: jax.make_array_from_process_local_data(self._sharding, x),
-                        batch,
-                    )
-                else:
-                    yield jax.tree.map(torch.as_tensor, batch)
-
-
-def _collate_fn(items):
-    """Collate the batch elements into batched numpy arrays."""
-    # Make sure to convert to numpy arrays before stacking since some of the incoming elements
-    # may be JAX arrays.
-    return jax.tree.map(lambda *xs: np.stack([np.asarray(x) for x in xs], axis=0), *items)
+                t_inner_iter = time.monotonic()
+                logging.info("[TorchDataLoader] calling iter(self._data_loader) (spawning workers) ...")
+                data_iter = iter(self._data_loader)
+                self._active_iterator = data_iter
+                logging.info(f"[TorchDataLoader] iter() returned in {time.monotonic() - t_inner_iter:.1f}s")
+                while True:
+                    if self._num_batches is not None and num_items >= self._num_batches:
+                        return
+                    t_batch = time.monotonic()
+                    try:
+                        batch = next(data_iter)
+                    except StopIteration:
+                        break  # We've exhausted the dataset. Create a new iterator and start over.
+                    dt = time.monotonic() - t_batch
+                    worker_stats: dict[str, float] = {}
+                    if self._profile_data_pipeline:
+                        batch, worker_stats = batch
+                    if num_items < 3 or dt > 3.0:
+                        logging.info(f"[TorchDataLoader] batch #{num_items} next() took {dt:.2f}s")
+                    num_items += 1
+                    # For JAX, convert to sharded arrays; for PyTorch, return torch tensors
+                    t_array = time.monotonic()
+                    if self._sharding is not None:
+                        batch = jax.tree.map(
+                            lambda x: jax.make_array_from_process_local_data(self._sharding, _to_numpy_view(x)),
+                            batch,
+                        )
+                    else:
+                        batch = jax.tree.map(torch.as_tensor, batch)
+                    array_construct_s = time.monotonic() - t_array
+
+                    h2d_wait_s = 0.0
+                    if self._profile_data_pipeline and self._sharding is not None:
+                        t_h2d = time.monotonic()
+                        jax.block_until_ready(batch)
+                        h2d_wait_s = time.monotonic() - t_h2d
+
+                    if self._profile_data_pipeline:
+                        self._last_profile_stats = {
+                            "main_queue_wait_s": dt,
+                            **worker_stats,
+                            "jax_array_construct_s": array_construct_s,
+                            "h2d_wait_s": h2d_wait_s,
+                        }
+                    yield batch
+        finally:
+            self.close()
+
+    def profile_stats(self) -> dict[str, float] | None:
+        return self._last_profile_stats
+
+    def close(self) -> None:
+        """Shut down the active multiprocessing iterator before interpreter teardown."""
+        data_iter = self._active_iterator
+        if data_iter is None:
+            return
+
+        workers = tuple(getattr(data_iter, "_workers", ()))
+        try:
+            shutdown_workers = getattr(data_iter, "_shutdown_workers", None)
+            if shutdown_workers is not None:
+                shutdown_workers()
+        except RuntimeError as error:
+            # Native libraries used by a worker can abort while reacting to the shutdown
+            # signal. PyTorch still terminates all remaining workers in its finally block;
+            # handle the resulting SIGCHLD report here instead of leaving it for __del__.
+            logging.warning(f"DataLoader worker exited while shutting down: {error}")
+        finally:
+            for worker in workers:
+                worker.join(timeout=5.0)
+                if worker.is_alive():
+                    worker.kill()
+                    worker.join()
+            self._active_iterator = None
+            if getattr(self._data_loader, "_iterator", None) is data_iter:
+                self._data_loader._iterator = None
+
+
+def _to_numpy_view(value):
+    """Return a NumPy view of a CPU tensor without copying its shared storage."""
+    if isinstance(value, torch.Tensor):
+        return value.numpy()
+    return np.asarray(value)
+
+
+def _collate_fn(items, *, profile: bool = False):
+    """Collate batch elements into CPU tensors so workers use shared-memory IPC."""
+    getitem_s = 0.0
+    if profile:
+        samples, durations = zip(*items, strict=True)
+        items = samples
+        getitem_s = sum(durations)
+    started = time.monotonic()
+    # Some inputs are JAX arrays, so stack through NumPy first. Returning Torch tensors is
+    # important here: the DataLoader multiprocessing reducer transfers their storage through
+    # shared memory instead of pickling each ~900 MiB NumPy batch into the result queue.
+    batch = jax.tree.map(
+        lambda *xs: torch.from_numpy(np.stack([np.asarray(x) for x in xs], axis=0)),
+        *items,
+    )
+    if not profile:
+        return batch
+    return batch, {
+        "worker_getitem_s": getitem_s,
+        "worker_collate_s": time.monotonic() - started,
+    }
 
 
 def _worker_init_fn(worker_id: int) -> None:
@@ -738,3 +847,12 @@ def data_config(self) -> _config.DataConfig:
     def __iter__(self):
         for batch in self._data_loader:
             yield _model.Observation.from_dict(batch), batch["actions"]
+
+    def profile_stats(self) -> dict[str, float] | None:
+        profile_stats = getattr(self._data_loader, "profile_stats", None)
+        return profile_stats() if profile_stats is not None else None
+
+    def close(self) -> None:
+        close = getattr(self._data_loader, "close", None)
+        if close is not None:
+            close()
diff --git a/src/openpi/training/data_loader_test.py b/src/openpi/training/data_loader_test.py
index e34d259..7562494 100644
--- a/src/openpi/training/data_loader_test.py
+++ b/src/openpi/training/data_loader_test.py
@@ -54,6 +54,78 @@ def test_torch_data_loader_parallel():
         assert all(x.shape[0] == 4 for x in jax.tree.leaves(batch))
 
 
+def test_torch_worker_batches_use_shared_memory_and_numpy_views_are_zero_copy():
+    config = pi0_config.Pi0Config(action_dim=24, action_horizon=50, max_token_len=48)
+    dataset = _data_loader.FakeDataset(config, 4)
+    loader = _data_loader.TorchDataLoader(
+        dataset,
+        local_batch_size=2,
+        num_batches=1,
+        num_workers=2,
+        framework="pytorch",
+    )
+
+    data_iter = iter(loader.torch_loader)
+    batch = next(data_iter)
+    del data_iter
+
+    for tensor in jax.tree.leaves(batch):
+        assert isinstance(tensor, torch.Tensor)
+        assert tensor.device.type == "cpu"
+        assert tensor.is_shared()
+
+        numpy_view = _data_loader._to_numpy_view(tensor)
+        assert numpy_view.__array_interface__["data"][0] == tensor.data_ptr()
+
+
+def test_torch_data_loader_close_shuts_down_workers():
+    config = pi0_config.Pi0Config(action_dim=24, action_horizon=50, max_token_len=48)
+    dataset = _data_loader.FakeDataset(config, 4)
+    loader = _data_loader.TorchDataLoader(
+        dataset,
+        local_batch_size=2,
+        num_workers=2,
+        framework="pytorch",
+    )
+
+    data_iter = iter(loader)
+    _ = next(data_iter)
+    workers = tuple(loader._active_iterator._workers)
+    assert all(worker.is_alive() for worker in workers)
+
+    data_iter.close()
+
+    assert loader._active_iterator is None
+    assert loader.torch_loader._iterator is None
+    assert all(not worker.is_alive() for worker in workers)
+
+
+def test_torch_data_loader_profile_stats():
+    config = pi0_config.Pi0Config(action_dim=24, action_horizon=50, max_token_len=48)
+    batch_size = jax.device_count()
+    dataset = _data_loader.FakeDataset(config, batch_size * 2)
+
+    loader = _data_loader.TorchDataLoader(
+        dataset,
+        local_batch_size=batch_size,
+        num_batches=1,
+        num_workers=2,
+        profile_data_pipeline=True,
+    )
+    _ = next(iter(loader))
+    stats = loader.profile_stats()
+
+    assert stats is not None
+    assert set(stats) == {
+        "main_queue_wait_s",
+        "worker_getitem_s",
+        "worker_collate_s",
+        "jax_array_construct_s",
+        "h2d_wait_s",
+    }
+    assert all(value >= 0 for value in stats.values())
+
+
 def test_with_fake_dataset():
     config = _config.get_config("debug_pi05")
 

From 8ccf99c6ee5c87f3ec31b970663928e664c1b37b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Fri, 28 Aug 2026 01:48:48 +0800
Subject: [PATCH 02/28] Profile and optimize H200 training

---
 docs/2026-08-28-h200-training-optimization.md | 161 ++++++++++++++++++
 scripts/train.py                              |  69 ++++++--
 src/openpi/models/pi0.py                      |  20 ++-
 src/openpi/models/pi0_config.py               |   3 +
 src/openpi/models/pi0_test.py                 |  22 +++
 src/openpi/training/config.py                 |   7 +
 6 files changed, 268 insertions(+), 14 deletions(-)
 create mode 100644 docs/2026-08-28-h200-training-optimization.md

diff --git a/docs/2026-08-28-h200-training-optimization.md b/docs/2026-08-28-h200-training-optimization.md
new file mode 100644
index 0000000..a4e791c
--- /dev/null
+++ b/docs/2026-08-28-h200-training-optimization.md
@@ -0,0 +1,161 @@
+# H200 训练性能优化记录
+
+**日期**：2026-08-27 至 2026-08-28
+
+**机器**：8×H200
+
+**分支**：`optimize-training-duration`
+
+**模型配置**：`pi05_base_bi_flexiv_earbud_case_insertion_0826_h200`
+
+**起始 Checkpoint**：step 49000
+
+本文记录一次从数据管线拆分计时到 GPU trace 的训练性能优化过程。目标不是只看某个局部计时，而是在相同 Checkpoint、相同训练配置下逐项修改，并用稳态端到端 step 时间判断收益。
+
+更早的 `in_order=False`、无限 sampler 和跳过未使用触觉视频解码等工作见 [data-pipeline-stalls.md](data-pipeline-stalls.md)。本文从 multiprocessing queue、JAX 数组构造和模型计算继续向下分析。
+
+## 1. 测量口径
+
+所有短跑都从原始 step 49000 Checkpoint 只读恢复到独立实验目录，关闭 W&B、周期保存和最终 Checkpoint 保存，避免覆盖训练结果。每个性能变体运行 40 step，丢弃前 10 step 的编译和预热数据，统计 step 49011–49040 的 30 个稳态样本。
+
+JAX dispatch 是异步的，所以不能把 `dispatch` 单独当成 step 时间。本文使用以下和作为主要指标：
+
+```text
+measured_step = dispatch + compute_sync + next_batch_total
+```
+
+数据管线 profiling 进一步拆成：
+
+- `main_queue_wait`：主进程等待 DataLoader 返回 batch；
+- `worker_getitem`：worker 内所有样本 `__getitem__` 时间之和，不是单个 step 的墙钟时间；
+- `worker_collate`：worker 内 collate 时间；
+- `jax_array_construct`：主进程构造分片 JAX array 的时间；
+- `h2d_wait`：显式等待 Host→Device 完成的时间。
+
+profiling 会增加计时和同步，因此只用于短跑，不应在生产训练中默认开启。
+
+## 2. 第一步：共享内存 Tensor 和零拷贝 NumPy 视图
+
+### 问题
+
+worker 原先返回 NumPy batch。大 batch 通过 multiprocessing queue 时需要 pickle/复制，主进程还要再次把它包装成 JAX array。此前已经消除了严格顺序交付和 epoch 边界停顿，但这里仍有一段稳定的数据传输开销。
+
+### 实现
+
+collate 改为返回 CPU Torch Tensor。PyTorch multiprocessing reducer 会通过共享内存传递 tensor storage；主进程用 `tensor.numpy()` 得到同一块 CPU storage 的零拷贝视图，再交给 `jax.make_array_from_process_local_data`。
+
+同时增加显式 `DataLoader.close()`：训练退出时主动调用 multiprocessing iterator 的 shutdown，等待 worker，超时后终止残留进程，避免把 native library 的析构异常留到解释器退出阶段。
+
+### 结果
+
+- 修改前：`3.275728 s/step`
+- 共享 Tensor 后：`2.828845 s/step`
+- step 时间减少 `13.64%`，等价吞吐提升 `15.80%`
+
+共享 Tensor 后的稳态数据侧均值为：
+
+- `main_queue_wait = 0.005030 s`
+- `worker_collate = 0.081114 s`
+- `jax_array_construct = 0.009046 s`
+- `h2d_wait = 0.006171 s`
+- `next_batch_total = 0.046491 s`
+
+这说明主进程队列等待、JAX array 构造和 H2D 已经不是主要瓶颈。`worker_getitem` 约 130 秒是 512 个样本耗时之和，由 64 个 worker 并行完成，不能直接加到 step 墙钟时间上。
+
+## 3. 第二步：20-step JAX/XProf trace
+
+训练脚本增加可控的 XProf 开关：指定绝对起始 step、trace step 数和输出目录。在 trace 起止位置显式 `block_until_ready`，保证采集边界对应完整训练迭代。
+
+本次采集 step 49011–49030，共 20 个完整 step。按每张 GPU、每个 step 汇总 trace 中记录的 kernel 时间：
+
+- GEMM：约 `390.449 ms`
+- Fusion：约 `183.930 ms`
+- NCCL：约 `45.990 ms`
+- 其他 kernel：约 `2.562 ms`
+- Memcpy：约 `1.509 ms`
+- Memset：约 `0.465 ms`
+
+在这些已记录的 kernel 时间中，GEMM 与 Fusion 约占 `92.6%`，NCCL 约占 `7.4%`。因此当前优先级应是减少或融合模型计算，而不是继续优化 Host→Device 或 NCCL。
+
+trace 中 kernel 可能重叠，不能把 kernel duration 简单相加后当作 GPU 利用率；同样，也不能仅凭 host dispatch 区间判断真实计算时间。
+
+## 4. 第三步：只在日志 step 计算范数
+
+### 问题
+
+训练原先每个 step 都计算整个 trainable pytree 的 `grad_norm` 和大部分 kernel 参数的 `param_norm`，但这些指标只在 `log_interval` step 被记录。
+
+### 实现
+
+`train_step` 增加静态布尔参数 `compute_metrics`：
+
+- 日志 step 计算真实 `grad_norm` 和 `param_norm`；
+- 非日志 step 返回同结构的 NaN 标量，保持 JIT 输出 pytree 稳定；
+- host 聚合使用 `nanmean`，忽略非日志 step 的占位值；
+- 该布尔值作为静态位置参数传给 JIT，分别编译计算指标和跳过指标的两个变体。
+
+使用位置参数是必要的：当前 JAX 在显式设置 `in_shardings` 时不接受 pjit 关键字实参。
+
+### 结果
+
+- 共享 Tensor 基线：`2.828845 s/step`
+- 仅日志 step 计算范数：`2.814683 s/step`
+- step 时间减少 `0.50%`，等价吞吐提升 `0.50%`
+
+收益存在但很小，说明范数归约不是当前主要瓶颈。这个改动仍值得保留，因为它与日志语义一致，且不会牺牲已记录指标。
+
+## 5. 第四步：三路相机合并为一次 SigLIP 调用
+
+### 问题
+
+`Pi0.embed_prefix` 原先对 head、left wrist、right wrist 三路相机分别调用一次相同的 SigLIP image encoder。这会重复进入相同计算图，并重复触发参数使用和 kernel launch。
+
+### 实现
+
+新增可选配置 `batch_image_views`。开启后：
+
+1. 将三路图像堆叠为 `(B, V, H, W, C)`；
+2. reshape 为 `(B*V, H, W, C)`，只调用一次 image encoder；
+3. 将 token reshape 回 `(B, V, ...)`，按原相机名拆分。
+
+布局使用 sample-major interleave，同一样本的三路视图连续排列。这样在数据并行/FSDP shard 中，各设备仍拿到完整样本的所有视图，不需要跨设备重排。
+
+该开关默认关闭。生产训练可通过以下参数启用：
+
+```bash
+--model.batch-image-views
+```
+
+### 结果
+
+- 只做范数优化：`2.814683 s/step`
+- 再合并三路相机：`2.773743 s/step`
+- 相对上一阶段 step 时间减少 `1.45%`，等价吞吐提升 `1.48%`
+- 相对共享 Tensor 基线累计吞吐提升 `1.99%`
+- 相对最初 NumPy queue 版本累计 step 时间减少 `15.32%`，等价吞吐提升 `18.10%`
+
+相机合并后，`dispatch` 从约 `2.657 s` 降到 `2.236 s`，但 `compute_sync` 从约 `0.115 s` 增到 `0.492 s`。这是异步工作在 dispatch 与同步边界之间重新分布，真实收益必须看两者与 `next_batch_total` 的总和，不能宣称 dispatch 单项显示的约 16% 提升。
+
+## 6. 正确性和运行验证
+
+- 增加单元测试，对三路相机分别编码与合并编码的 token 做逐元素相等比较；
+- 模型与 YAML 相关回归测试：`23 passed`；
+- 前一阶段数据管线与显式 shutdown 回归：`37 passed`；
+- 范数优化和相机合并各完成一次 40-step、8×H200 实际训练；
+- 两次运行都成功恢复 step 49000 Checkpoint，没有 OOM，没有保存新 Checkpoint；
+- `compute_metrics=True/False` 两个 JIT 分支都在真实训练中执行；
+- 两次运行结束后显式 shutdown 成功，worker 无残留，8 张 GPU 显存全部释放；
+- 本地与训练服务器上的修改文件 SHA256 一致。
+
+## 7. 当前结论与下一步
+
+当前稳态数据交付只占约 `0.04–0.05 s/step`，H2D 只有数毫秒；XProf 又显示已记录 kernel 时间主要集中在 GEMM/Fusion。因此，继续微调 queue、collate 或 H2D 不太可能带来显著收益。
+
+下一轮建议按以下顺序推进：
+
+1. **单独优化冷启动。** 64 个 worker 拉起约 320 秒，首批数据约 119 秒，总首批约 450 秒。它不影响长训练稳态吞吐，但严重影响恢复训练、调参和短 profiling。应检查 worker 是否串行重复初始化 JAX、数据集元数据或视频 backend，并评估持久 worker 池、延迟初始化和元数据预缓存。
+2. **继续拆最大的 GEMM/Fusion。** 优先检查 trace 中最大的 SigLIP/Gemma GEMM 和 `loop_transpose_fusion_4`，验证是否存在多余 transpose、重复 materialization 或可合并的前向计算。
+3. **做更长的相机合并 A/B。** 当前 30 个稳态 step 已确认约 1.5% 增量收益；在默认开启前，建议再做至少 200–500 个稳态 step，确认不同 batch 和训练阶段下收益稳定。
+4. **降低 checkpoint 干扰。** 长训练若不需要每 1000 step 恢复点，可提高 `save_interval`；此前生产日志已经确认 checkpoint 会造成下一 step 的数据等待尖峰。
+
+最终判断：共享内存 Tensor 是本轮最大且已验证的稳态收益；日志 step 范数是低风险小收益；三路相机合并方向正确但收益约 1.5%，应在更长 A/B 后再默认开启。后续优化应从数据搬运转向模型计算图，同时把约 450 秒冷启动作为独立问题处理。
diff --git a/scripts/train.py b/scripts/train.py
index da3bb3a..5ab0b86 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -140,6 +140,7 @@ def train_step(
     rng: at.KeyArrayLike,
     state: training_utils.TrainState,
     batch: tuple[_model.Observation, _model.Actions],
+    compute_metrics: bool = True,
 ) -> tuple[training_utils.TrainState, dict[str, at.Array]]:
     model = nnx.merge(state.model_def, state.params)
     model.train()
@@ -175,19 +176,29 @@ def loss_fn(
             ),
         )
 
-    # Filter out params that aren't kernels.
-    kernel_params = nnx.state(
-        model,
-        nnx.All(
-            nnx.Param,
-            nnx.Not(nnx_utils.PathRegex(".*/(bias|scale|pos_embedding|input_embedding)")),
-            lambda _, x: x.value.ndim > 1,
-        ),
-    )
+    if compute_metrics:
+        # These whole-model reductions are useful for diagnostics but need not run on
+        # every step when they are only logged every config.log_interval steps.
+        kernel_params = nnx.state(
+            model,
+            nnx.All(
+                nnx.Param,
+                nnx.Not(nnx_utils.PathRegex(".*/(bias|scale|pos_embedding|input_embedding)")),
+                lambda _, x: x.value.ndim > 1,
+            ),
+        )
+        grad_norm = optax.global_norm(grads)
+        param_norm = optax.global_norm(kernel_params)
+    else:
+        # Keep a stable output pytree for both JIT variants. The host-side logging
+        # reduction ignores these placeholders with nanmean.
+        grad_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
+        param_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
+
     info = {
         "loss": loss,
-        "grad_norm": optax.global_norm(grads),
-        "param_norm": optax.global_norm(kernel_params),
+        "grad_norm": grad_norm,
+        "param_norm": param_norm,
     }
     return new_state, info
 
@@ -258,6 +269,7 @@ def main(config: _config.TrainConfig):
         in_shardings=(replicated_sharding, train_state_sharding, data_sharding),
         out_shardings=(train_state_sharding, replicated_sharding),
         donate_argnums=(1,),
+        static_argnums=(3,),
     )
 
     start_step = int(train_state.step)
@@ -269,15 +281,36 @@ def main(config: _config.TrainConfig):
     )
 
     infos = []
+    xprof_active = False
     # --- stall diagnostics ---
     stall_threshold_s = 3.0  # log a warning when any phase exceeds this
     t_prev_loop_end = time.monotonic()
     for step in pbar:
+        if config.xprof_trace_dir is not None and step == config.xprof_start_step:
+            # Exclude work queued before the requested range and make the trace boundaries
+            # correspond to complete training iterations.
+            jax.block_until_ready((train_state, batch))
+            logging.info(
+                f"Starting JAX/XProf trace at step {step} for {config.xprof_num_steps} steps: "
+                f"{config.xprof_trace_dir}"
+            )
+            jax.profiler.start_trace(
+                config.xprof_trace_dir,
+                create_perfetto_link=False,
+                create_perfetto_trace=True,
+            )
+            xprof_active = True
+
         t_loop_start = time.monotonic()
 
         t0 = time.monotonic()
         with sharding.set_mesh(mesh):
-            train_state, info = ptrain_step(train_rng, train_state, batch)
+            train_state, info = ptrain_step(
+                train_rng,
+                train_state,
+                batch,
+                step % config.log_interval == 0,
+            )
         t_dispatch = time.monotonic() - t0
 
         # JAX dispatch is asynchronous. Fine-grained H2D profiling must first finish the
@@ -295,7 +328,7 @@ def main(config: _config.TrainConfig):
         if step % config.log_interval == 0:
             t0 = time.monotonic()
             stacked_infos = common_utils.stack_forest(infos)
-            reduced_info = jax.device_get(jax.tree.map(jnp.mean, stacked_infos))
+            reduced_info = jax.device_get(jax.tree.map(jnp.nanmean, stacked_infos))
             info_str = ", ".join(f"{k}={v:.4f}" for k, v in reduced_info.items())
             pbar.write(f"Step {step}: {info_str}")
             wandb.log(reduced_info, step=step)
@@ -351,6 +384,16 @@ def main(config: _config.TrainConfig):
                 f"next_batch_total={t_next_batch:.4f}s"
             )
 
+        if xprof_active and step + 1 == config.xprof_start_step + config.xprof_num_steps:
+            jax.block_until_ready((train_state, info, batch))
+            jax.profiler.stop_trace()
+            xprof_active = False
+            logging.info(f"Finished JAX/XProf trace at step {step}")
+
+    if xprof_active:
+        jax.block_until_ready((train_state, batch))
+        jax.profiler.stop_trace()
+
     logging.info("Shutting down data loader")
     data_loader.close()
     logging.info("Waiting for checkpoint manager to finish")
diff --git a/src/openpi/models/pi0.py b/src/openpi/models/pi0.py
index b57c487..8c9099c 100644
--- a/src/openpi/models/pi0.py
+++ b/src/openpi/models/pi0.py
@@ -16,6 +16,17 @@
 logger = logging.getLogger("openpi")
 
 
+def _encode_image_views_as_batch(image_encoder, images):
+    """Encode camera views as a sample-major interleaved batch."""
+    names = tuple(images)
+    stacked = jnp.stack([images[name] for name in names], axis=1)
+    batch_size, num_views = stacked.shape[:2]
+    flat = stacked.reshape((batch_size * num_views, *stacked.shape[2:]))
+    flat_tokens, _ = image_encoder(flat, train=False)
+    tokens = flat_tokens.reshape((batch_size, num_views, *flat_tokens.shape[1:]))
+    return {name: tokens[:, index] for index, name in enumerate(names)}
+
+
 def make_attn_mask(input_mask, mask_ar):
     """Adapted from big_vision.
 
@@ -116,6 +127,7 @@ def __init__(
 
         self._enable_training_time_rtc = config.enable_training_time_rtc
         self._max_delay = config.max_delay
+        self._batch_image_views = config.batch_image_views
 
     @at.typecheck
     def embed_prefix(
@@ -125,8 +137,14 @@ def embed_prefix(
         ar_mask = []
         tokens = []
         # embed images
+        batched_image_tokens = (
+            _encode_image_views_as_batch(self.PaliGemma.img, obs.images) if self._batch_image_views else None
+        )
         for name in obs.images:
-            image_tokens, _ = self.PaliGemma.img(obs.images[name], train=False)
+            if batched_image_tokens is None:
+                image_tokens, _ = self.PaliGemma.img(obs.images[name], train=False)
+            else:
+                image_tokens = batched_image_tokens[name]
 
             tokens.append(image_tokens)
             input_mask.append(
diff --git a/src/openpi/models/pi0_config.py b/src/openpi/models/pi0_config.py
index d7fd65f..c2746be 100644
--- a/src/openpi/models/pi0_config.py
+++ b/src/openpi/models/pi0_config.py
@@ -30,6 +30,9 @@ class Pi0Config(_model.BaseModelConfig):
     pi05: bool = False
     # This config option is not used directly by the model, but it is read by the ModelTransformFactory.
     discrete_state_input: bool = None  # type: ignore
+    # Encode all camera views in one larger SigLIP batch. This is equivalent to
+    # one encoder call per view but avoids repeated launches and FSDP parameter use.
+    batch_image_views: bool = False
 
     # training-time RTC config
     enable_training_time_rtc: bool = False
diff --git a/src/openpi/models/pi0_test.py b/src/openpi/models/pi0_test.py
index d3d70e8..239c8a7 100644
--- a/src/openpi/models/pi0_test.py
+++ b/src/openpi/models/pi0_test.py
@@ -7,6 +7,28 @@
 import openpi.models.pi0_config as _pi0_config
 
 
+class _FakeImageEncoder:
+    def __call__(self, images, *, train):
+        assert train is False
+        return jnp.mean(images, axis=(1, 2))[:, None, :], None
+
+
+def test_encode_image_views_as_batch_matches_separate_calls():
+    images = {
+        "base": jnp.arange(2 * 4 * 4 * 3, dtype=jnp.float32).reshape(2, 4, 4, 3),
+        "left": jnp.full((2, 4, 4, 3), 2.0),
+        "right": jnp.full((2, 4, 4, 3), 3.0),
+    }
+    encoder = _FakeImageEncoder()
+
+    actual = _pi0._encode_image_views_as_batch(encoder, images)
+    expected = {name: encoder(image, train=False)[0] for name, image in images.items()}
+
+    assert actual.keys() == expected.keys()
+    for name in images:
+        assert jnp.array_equal(actual[name], expected[name])
+
+
 def _get_frozen_state(config: _pi0_config.Pi0Config) -> nnx.State:
     abstract_model = nnx.eval_shape(config.create, jax.random.key(0))
 
diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py
index 1b9a11f..3e316a8 100644
--- a/src/openpi/training/config.py
+++ b/src/openpi/training/config.py
@@ -621,6 +621,11 @@ class TrainConfig:
     profile_data_pipeline: bool = False
     # Number of steps between fine-grained data-pipeline timing lines.
     profile_log_interval: int = 1
+    # If set, write a JAX/XProf trace to this directory. The trace begins at
+    # xprof_start_step and covers xprof_num_steps complete training iterations.
+    xprof_trace_dir: str | None = None
+    xprof_start_step: int = 0
+    xprof_num_steps: int = 20
     # Number of train steps (batches) to run.
     num_train_steps: int = 30_000
 
@@ -671,6 +676,8 @@ def trainable_filter(self) -> nnx.filterlib.Filter:
     def __post_init__(self) -> None:
         if self.resume and self.overwrite:
             raise ValueError("Cannot resume and overwrite at the same time.")
+        if self.xprof_num_steps <= 0:
+            raise ValueError("xprof_num_steps must be positive.")
 
 
 # YAML lookup paths, in priority order. First match wins.

From d0aec27a79c5879ed7f6a619a587203e29310437 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Fri, 28 Aug 2026 11:31:38 +0800
Subject: [PATCH 03/28] Update H100 RTC training configs

---
 ...exiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml} | 8 +++++---
 ..._base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml} | 2 ++
 2 files changed, 7 insertions(+), 3 deletions(-)
 rename configs/_examples/{pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0520_a100.yaml => pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml} (89%)
 rename configs/_examples/{pi05_base_bi_flexiv_unscrew_bottle_cap_0824_h100.yaml => pi05_base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml} (96%)

diff --git a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0520_a100.yaml b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml
similarity index 89%
rename from configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0520_a100.yaml
rename to configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml
index 4d3abc7..5c0df75 100644
--- a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0520_a100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml
@@ -18,11 +18,13 @@ data:
   repo_id: Xense/earbud_case_insertion_teleop_0515
   base_config:
     prompt_from_task: true
+    tactile: false
   default_prompt: Pick up each earbud case from the left stands, insert the matching
     earbuds, close the lid, and place the case in the box.
 batch_size: 256
 num_workers: 64
-num_train_steps: 20000
-save_interval: 5000
-keep_period: 20000
+num_train_steps: 60000
+save_interval: 10000
+keep_period: 10000
 fsdp_devices: 8
+strict_batch_order: false
diff --git a/configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_0824_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml
similarity index 96%
rename from configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_0824_h100.yaml
rename to configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml
index 975bb9f..934919b 100644
--- a/configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_0824_h100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml
@@ -23,6 +23,8 @@ batch_size: 256
 num_workers: 64
 num_train_steps: 80000
 fsdp_devices: 8
+save_interval: 10000
+keep_period: 10000
 # Whether batches are delivered in strict sampler order. false (the default) hands each
 # batch over as soon as a worker finishes it, which removes the head-of-line stall every
 # num_workers steps, and enables the InfiniteSampler so the epoch boundary never tears the

From 476babe5acdbadeab7f3c9e1f498b504925f9cf4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Sat, 29 Aug 2026 15:10:36 +0800
Subject: [PATCH 04/28] Enable cuDNN fused attention training

---
 README.md                                     |  30 +++-
 docs/2026-08-28-h200-training-optimization.md | 136 ++++++++++++++++--
 scripts/train.py                              |   3 +
 src/openpi/models/gemma.py                    |  37 +++--
 src/openpi/models/pi0.py                      |   1 +
 src/openpi/models/pi0_config.py               |   3 +
 6 files changed, 188 insertions(+), 22 deletions(-)

diff --git a/README.md b/README.md
index bec1e46..f15e4d6 100644
--- a/README.md
+++ b/README.md
@@ -184,6 +184,8 @@ model:
   paligemma_variant: gemma_2b
   action_expert_variant: gemma_300m
   enable_training_time_rtc: true
+  # Required when launching with the cuDNN 9.14 setup shown below.
+  use_cudnn_attention: true
   max_delay: 10
 
 data:
@@ -208,10 +210,28 @@ Then use it exactly like any other config:
 
 ```bash
 python scripts/compute_norm_stats.py --config-name my_task
-XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 python scripts/train.py my_task --exp-name=run_0520 --overwrite
-python scripts/serve_policy.py policy:checkpoint --policy.config=my_task --policy.dir=checkpoints/my_task/run_0520/<step>
+LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
+  XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+  python scripts/train.py my_task \
+    --exp-name=my_exp \
+    --overwrite
+python scripts/serve_policy.py policy:checkpoint --policy.config=my_task --policy.dir=checkpoints/my_task/my_exp/<step>
 ```
 
+The training command above selects cuDNN from the currently active Conda
+environment without hard-coding a machine-specific path. It must be paired
+with the following model setting (already shown in the example config):
+
+```yaml
+model:
+  use_cudnn_attention: true
+```
+
+At startup, `scripts/train.py` logs `JAX cuDNN runtime version: <version>`;
+verify that it reports `91400`. Use `--overwrite` only for a new run that may
+replace an existing experiment directory; use `--resume` to preserve and
+continue an existing run.
+
 ⚠️ Before committing a YAML into `configs/_examples/`, scrub any machine-local
 absolute paths from `weight_loader.params_path` (e.g. `/home/<you>/...`). Use
 upstream URLs (`gs://openpi-assets/...`) or paths that every contributor can
@@ -246,7 +266,11 @@ python scripts/compute_norm_stats.py --config-name pi05_base_xtac_umi_pick_up_cu
 Now we can kick off training with the following command (the `--overwrite` flag is used to overwrite existing checkpoints if you rerun fine-tuning with the same config):
 
 ```bash
-XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 python scripts/train.py pi05_base_xtac_umi_pick_up_cube_0807_h200 --exp-name=my_experiment --overwrite
+LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
+  XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+  python scripts/train.py pi05_base_xtac_umi_pick_up_cube_0807_h200 \
+    --exp-name=my_experiment \
+    --overwrite
 ```
 
 The command will log training progress to the console and save checkpoints to the `checkpoints` directory. You can also monitor training progress on the Weights & Biases dashboard. For maximally using the GPU memory, set `XLA_PYTHON_CLIENT_MEM_FRACTION=0.9` before running training -- this enables JAX to use up to 90% of the GPU memory (vs. the default of 75%).
diff --git a/docs/2026-08-28-h200-training-optimization.md b/docs/2026-08-28-h200-training-optimization.md
index a4e791c..7b91d20 100644
--- a/docs/2026-08-28-h200-training-optimization.md
+++ b/docs/2026-08-28-h200-training-optimization.md
@@ -1,5 +1,25 @@
 # H200 训练性能优化记录
 
+> **跨机器生产启动命令**：先激活目标 Conda 环境，再用 `$CONDA_PREFIX/lib` 相对该环境选择其中的 cuDNN。环境变量只作用于本次训练进程及其子进程，不修改环境内已安装的包。
+
+```bash
+LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
+  XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+  python scripts/train.py my_task \
+    --exp-name=run_0520 \
+    --overwrite
+```
+
+`my_task` 的模型配置需要包含 `use_cudnn_attention: true`；如果 YAML 中没有设置，则在命令末尾增加 `--model.use-cudnn-attention`。`--overwrite` 会删除同名实验的原 Checkpoint 目录；恢复已有训练应改用 `--resume`。
+
+训练启动后会在初始化日志中直接打印：
+
+```text
+JAX cuDNN runtime version: 91400
+```
+
+该值来自 `cuda_versions.cudnn_get_version()`。如果不是 `91400`，应停止训练并检查目标机器的 `$CONDA_PREFIX/lib`，不能只根据 PyTorch 输出判断 JAX 运行时版本。
+
 **日期**：2026-08-27 至 2026-08-28
 
 **机器**：8×H200
@@ -8,7 +28,7 @@
 
 **模型配置**：`pi05_base_bi_flexiv_earbud_case_insertion_0826_h200`
 
-**起始 Checkpoint**：step 49000
+**起始 Checkpoint**：step 49000（前四项优化）；step 16000（cuDNN Attention A/B）
 
 本文记录一次从数据管线拆分计时到 GPU trace 的训练性能优化过程。目标不是只看某个局部计时，而是在相同 Checkpoint、相同训练配置下逐项修改，并用稳态端到端 step 时间判断收益。
 
@@ -16,7 +36,7 @@
 
 ## 1. 测量口径
 
-所有短跑都从原始 step 49000 Checkpoint 只读恢复到独立实验目录，关闭 W&B、周期保存和最终 Checkpoint 保存，避免覆盖训练结果。每个性能变体运行 40 step，丢弃前 10 step 的编译和预热数据，统计 step 49011–49040 的 30 个稳态样本。
+前四项优化都从原始 step 49000 Checkpoint 只读恢复；cuDNN Attention A/B 从当前生产 step 16000 Checkpoint 只读恢复。所有短跑都使用独立实验目录，关闭 W&B、周期保存和最终 Checkpoint 保存，避免覆盖训练结果。每个变体运行到起始 step + 40，丢弃起始 step 到起始 step + 10 的编译、日志和预热数据，统计后续 30 个稳态样本。例如 step 49011–49040 或 step 16011–16040。
 
 JAX dispatch 是异步的，所以不能把 `dispatch` 单独当成 step 时间。本文使用以下和作为主要指标：
 
@@ -136,7 +156,99 @@ trace 中 kernel 可能重叠，不能把 kernel duration 简单相加后当作
 
 相机合并后，`dispatch` 从约 `2.657 s` 降到 `2.236 s`，但 `compute_sync` 从约 `0.115 s` 增到 `0.492 s`。这是异步工作在 dispatch 与同步边界之间重新分布，真实收益必须看两者与 `next_batch_total` 的总和，不能宣称 dispatch 单项显示的约 16% 提升。
 
-## 6. 正确性和运行验证
+## 6. 第五步：cuDNN fused Attention
+
+### 问题
+
+Gemma Attention 原先显式执行 QK einsum、mask、softmax 和 PV einsum。XProf 已显示 GEMM/Fusion 占已记录 kernel 时间的 92.6%，因此训练 Attention 是比继续压缩数据搬运更有价值的候选。当前模型使用 BF16 GQA 和 block attention mask，必须验证 cuDNN 是否支持完整的真实训练形状，不能只测试一个小型无 mask MHA。
+
+### 实现
+
+新增可选配置 `use_cudnn_attention`。开启后，训练时调用：
+
+```python
+jax.nn.dot_product_attention(
+    q,
+    k,
+    v,
+    mask=attn_mask,
+    scale=1.0,
+    implementation="cudnn",
+)
+```
+
+`q` 在调用前已经按 head dimension 缩放，因此这里显式使用 `scale=1.0`，避免重复缩放。该路径只在 `kv_cache is None` 的训练路径启用；带 KV cache 的推理继续使用原实现。开关默认关闭，可通过 `--model.use-cudnn-attention` 启用。
+
+### cuDNN 版本和动态库选择
+
+当前环境同时存在两套可被动态链接器找到的 cuDNN：
+
+- `$CONDA_PREFIX/lib` 中的 cuDNN 9.14；
+- Python `site-packages/nvidia/cudnn/lib` 路径，以及元数据版本为 `9.10.2.21` 的 `nvidia-cudnn-cu12` wheel。
+
+cuDNN 不是由 `torch.backends.cudnn.version()` 统一替整个 Python 环境选择的。PyTorch 和 JAX 都要由 Linux 动态链接器根据当前进程的 `LD_LIBRARY_PATH`、wheel 的 RPATH/RUNPATH、已经加载的同 SONAME 库和导入顺序解析 `libcudnn.so.9` 及其组件。因此必须在实际启动训练的同一个 shell 中检查 JAX，而不能用 PyTorch 的版本输出推断 JAX。
+
+当前已激活 shell 的实际状态是：
+
+```text
+LD_LIBRARY_PATH=$CONDA_PREFIX/lib
+PyTorch cuDNN runtime = 91400
+JAX cuDNN runtime     = 91400
+JAX build cuDNN       = 90101
+```
+
+如果显式移除 `LD_LIBRARY_PATH`，PyTorch 和 JAX 都会解析到 9.10.2：
+
+```bash
+env -u LD_LIBRARY_PATH python -c \
+  'import jax; from jax._src.lib import cuda_versions; print(cuda_versions.cudnn_get_version())'
+# 91002
+
+env -u LD_LIBRARY_PATH python -c \
+  'import torch; print(torch.backends.cudnn.version())'
+# 91002
+```
+
+所以“PyTorch 是 9.14、JAX 默认是 9.10”不是同一个进程环境下两个框架固定选择不同版本；准确说法是：**当前 shell 有 `$CONDA_PREFIX/lib` 时两者都是 9.14，去掉该搜索路径时两者都是 9.10.2。** 之前观察到的 JAX 9.10.2 是在 `env -u LD_LIBRARY_PATH` 的诊断进程中得到的。
+
+强制 cuDNN Attention 时，9.10.2 对本模型测试形状报错：
+
+```text
+XlaRuntimeError: INTERNAL: No valid engine configs for Matmul_MUL...
+```
+
+使用环境中已经安装的 9.14 后，小型 MHA、真实 BF16 GQA 形状、block mask 前向和反向全部通过。因此不需要下载新包，但生产启动前必须保证 JAX 实际显示 `91400`：
+
+```bash
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib"
+
+python -c \
+  'from jax._src.lib import cuda_versions; print(cuda_versions.cudnn_get_version())'
+# 必须输出 91400
+
+python scripts/train.py pi05_base_bi_flexiv_earbud_case_insertion_0826_h200 \
+  --resume \
+  --model.use-cudnn-attention
+```
+
+### A/B 结果
+
+测试条件为 8×H200、全局 batch 512、`batch_image_views=true`、相同 step 16000 Checkpoint 和相同数据配置。两组均使用 cuDNN 9.14 运行时，唯一变量是是否启用 fused Attention。统计 step 16011–16040 的 30 个稳态样本，不包含数据 worker 启动、首批解码、Checkpoint 恢复和 JIT/cuDNN 编译。
+
+- 原显式 Attention：均值 `2.795537 s/step`，中位数 `2.797350 s/step`，标准差 `0.007647 s`；
+- cuDNN Attention：均值 `1.978580 s/step`，中位数 `1.979450 s/step`，标准差 `0.011704 s`；
+- step 时间减少 `29.22%`；
+- 等价吞吐从 `183.15 samples/s` 提升到 `258.77 samples/s`，提升 `41.29%`；
+- 从生产 step 16000 跑到 step 60000，按纯稳态 step 时间估算可节省约 `9.98 小时`。
+
+分项均值也显示收益来自模型计算，而不是数据波动：
+
+- 原实现：`dispatch=2.237247 s`，`compute_sync=0.484967 s`，`next_batch_total=0.073323 s`；
+- cuDNN：`dispatch=1.516877 s`，`compute_sync=0.394600 s`，`next_batch_total=0.067103 s`。
+
+数值对齐测试中，BF16 masked attention 前向最大绝对误差为 `0.00390625`、平均绝对误差为 `0.00026494`；Q/K/V 梯度最大绝对误差分别为 `1.91e-6`、`4.77e-7` 和 `2.38e-7`，属于正常 BF16 数值差异。
+
+## 7. 正确性和运行验证
 
 - 增加单元测试，对三路相机分别编码与合并编码的 token 做逐元素相等比较；
 - 模型与 YAML 相关回归测试：`23 passed`；
@@ -146,16 +258,22 @@ trace 中 kernel 可能重叠，不能把 kernel duration 简单相加后当作
 - `compute_metrics=True/False` 两个 JIT 分支都在真实训练中执行；
 - 两次运行结束后显式 shutdown 成功，worker 无残留，8 张 GPU 显存全部释放；
 - 本地与训练服务器上的修改文件 SHA256 一致。
+- cuDNN Attention 小型与真实形状的前向/反向 smoke test 通过；
+- cuDNN Attention 相关模型回归：`7 passed`；
+- 原 Attention 和 cuDNN Attention 均完成一次 step 16000–16040 的 8×H200 短跑，没有 OOM、cuDNN engine fallback 或训练异常，也没有保存新 Checkpoint；
+- cuDNN 短跑退出时个别 DataLoader worker 在清理阶段打印 `killed by signal: Aborted`，发生在训练完成和计时结束之后，主训练进程与 GPU 资源均正常退出；这不是 Attention 计算错误。
 
-## 7. 当前结论与下一步
+## 8. 当前结论与下一步
 
 当前稳态数据交付只占约 `0.04–0.05 s/step`，H2D 只有数毫秒；XProf 又显示已记录 kernel 时间主要集中在 GEMM/Fusion。因此，继续微调 queue、collate 或 H2D 不太可能带来显著收益。
 
-下一轮建议按以下顺序推进：
+cuDNN fused Attention 是目前单项收益最大的模型侧优化：在已经开启三路相机合并的配置上，进一步减少 29.22% 的稳态 step 时间。生产训练建议同时开启 `batch_image_views` 和 `use_cudnn_attention`，并在启动后从 JAX 侧确认 cuDNN runtime 为 91400。PyTorch 的版本输出只能验证 PyTorch 当前进程，不能代替该检查。
+
+下一轮稳态优化建议按以下顺序推进：
 
-1. **单独优化冷启动。** 64 个 worker 拉起约 320 秒，首批数据约 119 秒，总首批约 450 秒。它不影响长训练稳态吞吐，但严重影响恢复训练、调参和短 profiling。应检查 worker 是否串行重复初始化 JAX、数据集元数据或视频 backend，并评估持久 worker 池、延迟初始化和元数据预缓存。
-2. **继续拆最大的 GEMM/Fusion。** 优先检查 trace 中最大的 SigLIP/Gemma GEMM 和 `loop_transpose_fusion_4`，验证是否存在多余 transpose、重复 materialization 或可合并的前向计算。
-3. **做更长的相机合并 A/B。** 当前 30 个稳态 step 已确认约 1.5% 增量收益；在默认开启前，建议再做至少 200–500 个稳态 step，确认不同 batch 和训练阶段下收益稳定。
+1. **做更长的 cuDNN 生产验证。** 先跑 200–500 个稳态 step，监控 loss、grad norm、显存和是否出现 engine 错误；短跑结果已足以支持开启，但长跑用于排除训练阶段相关问题。
+2. **缩短有效 token 长度。** 全数据扫描得到 Pi0.5 tokenized state 最大长度 165；可单独 A/B `max_token_len=168` 与当前 200，减少 Gemma attention/MLP 的序列计算，同时保留 3 token 余量。
+3. **继续拆最大的 GEMM/Fusion。** 优先检查启用 cuDNN 后的新 trace，确认剩余最大的 SigLIP/Gemma GEMM 和 fusion，避免继续针对已经被 fused Attention 消除的旧热点优化。
 4. **降低 checkpoint 干扰。** 长训练若不需要每 1000 step 恢复点，可提高 `save_interval`；此前生产日志已经确认 checkpoint 会造成下一 step 的数据等待尖峰。
 
-最终判断：共享内存 Tensor 是本轮最大且已验证的稳态收益；日志 step 范数是低风险小收益；三路相机合并方向正确但收益约 1.5%，应在更长 A/B 后再默认开启。后续优化应从数据搬运转向模型计算图，同时把约 450 秒冷启动作为独立问题处理。
+最终判断：共享内存 Tensor 消除了主要数据交付开销；日志 step 范数和三路相机合并是低风险增量收益；cuDNN fused Attention 是当前最大且已验证的稳态收益。后续应基于启用 cuDNN 后的新计算图继续优化，并按当前需求不把约 10 分钟冷启动作为训练吞吐优化目标。
diff --git a/scripts/train.py b/scripts/train.py
index 5ab0b86..f65f422 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -12,6 +12,7 @@
 import jax
 import jax.experimental
 import jax.numpy as jnp
+from jax._src.lib import cuda_versions
 import numpy as np
 import optax
 import tqdm_loggable.auto as tqdm
@@ -206,6 +207,8 @@ def loss_fn(
 def main(config: _config.TrainConfig):
     init_logging()
     logging.info(f"Running on: {platform.node()}")
+    cudnn_runtime_version = cuda_versions.cudnn_get_version() if cuda_versions is not None else None
+    logging.info(f"JAX cuDNN runtime version: {cudnn_runtime_version}")
 
     if config.batch_size % jax.device_count() != 0:
         raise ValueError(
diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index 023b497..90cbf47 100644
--- a/src/openpi/models/gemma.py
+++ b/src/openpi/models/gemma.py
@@ -165,6 +165,7 @@ class Attention(nn.Module):
     """Attention module."""
 
     configs: Sequence[Config]
+    use_cudnn_attention: bool = False
 
     @nn.compact
     def __call__(self, xs, positions, attn_mask, kv_cache):
@@ -219,22 +220,35 @@ def __call__(self, xs, positions, attn_mask, kv_cache):
             k = jnp.concatenate([cache_k, k], axis=1)
             v = jnp.concatenate([cache_v, v], axis=1)
 
-        q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=self.configs[0].num_kv_heads)
-        logits = jnp.einsum("BTKGH,BSKH->BKGTS", q, k, preferred_element_type=jnp.float32)
-
         if attn_mask.shape != (q.shape[0], 1, q.shape[1], k.shape[1]):
             raise ValueError(
                 f"Attention mask with shape {attn_mask.shape} but shapes for q and k are: {q.shape} and {k.shape}"
             )
 
-        # big_neg = jnp.finfo(logits.dtype).min
-        big_neg = -2.3819763e38  # See gemma/modules.py
-        masked_logits = jnp.where(attn_mask[:, :, None, :, :], logits, big_neg)
+        if self.use_cudnn_attention and kv_cache is None:
+            # q is already scaled above, so disable dot_product_attention's
+            # default head-dimension scaling. Keep cached inference on the
+            # existing implementation; this switch targets training throughput.
+            encoded = jax.nn.dot_product_attention(
+                q,
+                k,
+                v,
+                mask=attn_mask,
+                scale=1.0,
+                implementation="cudnn",
+            )
+        else:
+            grouped_q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=self.configs[0].num_kv_heads)
+            logits = jnp.einsum("BTKGH,BSKH->BKGTS", grouped_q, k, preferred_element_type=jnp.float32)
+
+            # big_neg = jnp.finfo(logits.dtype).min
+            big_neg = -2.3819763e38  # See gemma/modules.py
+            masked_logits = jnp.where(attn_mask[:, :, None, :, :], logits, big_neg)
 
-        probs = jax.nn.softmax(masked_logits, axis=-1).astype(dtype)
+            probs = jax.nn.softmax(masked_logits, axis=-1).astype(dtype)
 
-        encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v)
-        encoded = einops.rearrange(encoded, "B T K G H -> B T (K G) H")
+            encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v)
+            encoded = einops.rearrange(encoded, "B T K G H -> B T (K G) H")
 
         out = []
         start = 0
@@ -291,6 +305,7 @@ class Block(nn.Module):
     """Transformer block."""
 
     configs: tuple[Config, ...]
+    use_cudnn_attention: bool = False
 
     dropout: float = 0.0
     dropout_bdims: tuple[int, ...] = ()
@@ -300,7 +315,7 @@ def __call__(self, xs, kv_cache, positions, attn_mask, adarms_cond, deterministi
         xs = sharding.activation_sharding_constraint(xs)
         drop = nn.Dropout(self.dropout, self.dropout_bdims) if self.dropout else lambda x, _: x
 
-        attn = Attention(configs=self.configs, name="attn")
+        attn = Attention(configs=self.configs, use_cudnn_attention=self.use_cudnn_attention, name="attn")
 
         pre_attn = []
         gates = []
@@ -356,6 +371,7 @@ class Module(nn.Module):
     dropout: float = 0.0
     dropout_bdims: tuple[int, ...] = ()  # Every float is dropped independently.
     adarms: bool = False
+    use_cudnn_attention: bool = False
 
     def setup(self):
         # all experts must have the same depth
@@ -386,6 +402,7 @@ def setup(self):
             length=self.configs[0].depth,
         )(
             configs=self.configs,
+            use_cudnn_attention=self.use_cudnn_attention,
             dropout=self.dropout,
             dropout_bdims=self.dropout_bdims,
         )
diff --git a/src/openpi/models/pi0.py b/src/openpi/models/pi0.py
index 8c9099c..7420c04 100644
--- a/src/openpi/models/pi0.py
+++ b/src/openpi/models/pi0.py
@@ -94,6 +94,7 @@ def __init__(
                 configs=[paligemma_config, action_expert_config],
                 embed_dtype=config.dtype,
                 adarms=config.pi05,
+                use_cudnn_attention=config.use_cudnn_attention,
             )
         )
         llm.lazy_init(
diff --git a/src/openpi/models/pi0_config.py b/src/openpi/models/pi0_config.py
index c2746be..d506af7 100644
--- a/src/openpi/models/pi0_config.py
+++ b/src/openpi/models/pi0_config.py
@@ -33,6 +33,9 @@ class Pi0Config(_model.BaseModelConfig):
     # Encode all camera views in one larger SigLIP batch. This is equivalent to
     # one encoder call per view but avoids repeated launches and FSDP parameter use.
     batch_image_views: bool = False
+    # Use cuDNN fused attention for training-time Gemma attention. This requires
+    # a cuDNN runtime that supports the model's BF16 GQA shapes and attention mask.
+    use_cudnn_attention: bool = False
 
     # training-time RTC config
     enable_training_time_rtc: bool = False

From f42992e187eeb45afe609be63ab55dc266457d84 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Sat, 29 Aug 2026 15:23:47 +0800
Subject: [PATCH 05/28] Document cuDNN production validation

---
 docs/2026-08-28-h200-training-optimization.md | 41 ++++++++++++++++---
 1 file changed, 36 insertions(+), 5 deletions(-)

diff --git a/docs/2026-08-28-h200-training-optimization.md b/docs/2026-08-28-h200-training-optimization.md
index 7b91d20..a3296f5 100644
--- a/docs/2026-08-28-h200-training-optimization.md
+++ b/docs/2026-08-28-h200-training-optimization.md
@@ -10,6 +10,8 @@ LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
     --overwrite
 ```
 
+> **2026-08-29 长跑验证警告**：下面的命令和开关保留为复现记录，但当前 cuDNN Attention 不应继续用于生产训练。生产长跑在 step 16100 首次记录到 NaN 梯度，step 16200 起 loss 持续为 NaN。短跑的数值对齐不足以覆盖该问题，详见“生产长跑验证”。
+
 `my_task` 的模型配置需要包含 `use_cudnn_attention: true`；如果 YAML 中没有设置，则在命令末尾增加 `--model.use-cudnn-attention`。`--overwrite` 会删除同名实验的原 Checkpoint 目录；恢复已有训练应改用 `--resume`。
 
 训练启动后会在初始化日志中直接打印：
@@ -248,6 +250,34 @@ python scripts/train.py pi05_base_bi_flexiv_earbud_case_insertion_0826_h200 \
 
 数值对齐测试中，BF16 masked attention 前向最大绝对误差为 `0.00390625`、平均绝对误差为 `0.00026494`；Q/K/V 梯度最大绝对误差分别为 `1.91e-6`、`4.77e-7` 和 `2.38e-7`，属于正常 BF16 数值差异。
 
+### 2026-08-29 生产长跑验证
+
+生产训练从原 step 16000 Checkpoint 恢复，使用 8×H200、全局 batch 512、96 个 DataLoader worker、`batch_image_views=true`、`use_cudnn_attention=true`、`XLA_PYTHON_CLIENT_MEM_FRACTION=0.9`。启动日志确认 JAX 实际加载 cuDNN runtime `91400`。截至检查时，tmux 训练进程仍在运行，进度条约为 step 55100/60000；日志中没有 OOM、`No valid engine configs` 或 XLA runtime exception。
+
+速度统计使用日志中已经完整刷新的 step 16150–51650 记录。普通非日志、非保存 step 共 356 个样本：
+
+- **当前一个普通稳态 step：均值 `1.910590 s`，中位数 `1.910000 s`；**
+- 标准差 `0.010744 s`，范围 `1.86–1.95 s`；
+- 分项均值为 `dispatch=1.867331 s`、`next_batch=0.042612 s`；
+- 按 batch 512 折算吞吐约 `267.98 samples/s`；
+- 每 100 step 的日志 step 均值 `2.360375 s`，其中日志约 `0.436000 s`；
+- 每 1000 step 的 Checkpoint step 均值 `4.656286 s`，其中保存约 `2.288286 s`；
+- 按上述频率摊销日志和保存开销后，长期约为 **`1.917238 s/step`**。
+
+因此，生产长跑的速度与 40-step A/B 中的 `1.978580 s/step` 一致并略快；从纯吞吐看，cuDNN fused Attention 的约 1.9 秒/step 收益稳定存在。
+
+但是训练数值已经失效：
+
+```text
+Step 16100: grad_norm=nan, loss=0.1185, param_norm=nan
+Step 16200: grad_norm=nan, loss=nan,    param_norm=nan
+Step 16300 及之后: grad_norm=nan, loss=nan, param_norm=nan
+```
+
+step 16100 是第一个写入指标的生产日志点，因此只能确定 NaN 在 step 16000–16100 之间产生，不能从当前日志定位到更精确的单步。已完整刷新的日志一直到 step 51600 都没有再次出现有限 loss。周期保存又已经清理原 step 16000；当前保留的最后一个切换前 Checkpoint 是 step 15000，而 step 20000–55000 均是在 NaN 出现后保存，不能作为有效训练结果使用。
+
+结论是：**1.91 秒/step 是真实、稳定的硬件吞吐，但来自数值无效的训练，不能作为启用 cuDNN Attention 的生产验收。** 在定位并修复 NaN 前，应关闭 `use_cudnn_attention`；恢复有效训练时只能从仍保留的切换前 Checkpoint 重新开始。后续诊断需要缩短日志间隔并逐 step 检查 loss、梯度和参数有限性，重点覆盖真实 block mask 下的完整优化器更新，而不只是独立 Attention 前向/反向。
+
 ## 7. 正确性和运行验证
 
 - 增加单元测试，对三路相机分别编码与合并编码的 token 做逐元素相等比较；
@@ -262,18 +292,19 @@ python scripts/train.py pi05_base_bi_flexiv_earbud_case_insertion_0826_h200 \
 - cuDNN Attention 相关模型回归：`7 passed`；
 - 原 Attention 和 cuDNN Attention 均完成一次 step 16000–16040 的 8×H200 短跑，没有 OOM、cuDNN engine fallback 或训练异常，也没有保存新 Checkpoint；
 - cuDNN 短跑退出时个别 DataLoader worker 在清理阶段打印 `killed by signal: Aborted`，发生在训练完成和计时结束之后，主训练进程与 GPU 资源均正常退出；这不是 Attention 计算错误。
+- 后续生产长跑推翻了短跑的数值有效性结论：step 16100 首次记录到 NaN 梯度，step 16200 起 loss 持续为 NaN；这说明现有 smoke test 和 40-step 吞吐短跑不足以作为生产正确性验证。
 
 ## 8. 当前结论与下一步
 
 当前稳态数据交付只占约 `0.04–0.05 s/step`，H2D 只有数毫秒；XProf 又显示已记录 kernel 时间主要集中在 GEMM/Fusion。因此，继续微调 queue、collate 或 H2D 不太可能带来显著收益。
 
-cuDNN fused Attention 是目前单项收益最大的模型侧优化：在已经开启三路相机合并的配置上，进一步减少 29.22% 的稳态 step 时间。生产训练建议同时开启 `batch_image_views` 和 `use_cudnn_attention`，并在启动后从 JAX 侧确认 cuDNN runtime 为 91400。PyTorch 的版本输出只能验证 PyTorch 当前进程，不能代替该检查。
+cuDNN fused Attention 的速度收益是目前最大的模型侧收益：生产长跑普通 step 为 `1.910590 s`，摊销日志和 Checkpoint 后约 `1.917238 s/step`。但是该长跑从 step 16100 起出现 NaN，故当前实现没有通过生产正确性验收。确认 JAX runtime 为 91400 只能证明动态库选择正确，不能证明训练数值正确；在修复 NaN 前不得将 `use_cudnn_attention` 用于有效训练。
 
 下一轮稳态优化建议按以下顺序推进：
 
-1. **做更长的 cuDNN 生产验证。** 先跑 200–500 个稳态 step，监控 loss、grad norm、显存和是否出现 engine 错误；短跑结果已足以支持开启，但长跑用于排除训练阶段相关问题。
-2. **缩短有效 token 长度。** 全数据扫描得到 Pi0.5 tokenized state 最大长度 165；可单独 A/B `max_token_len=168` 与当前 200，减少 Gemma attention/MLP 的序列计算，同时保留 3 token 余量。
-3. **继续拆最大的 GEMM/Fusion。** 优先检查启用 cuDNN 后的新 trace，确认剩余最大的 SigLIP/Gemma GEMM 和 fusion，避免继续针对已经被 fused Attention 消除的旧热点优化。
+1. **先定位 cuDNN NaN。** 从保留的 step 15000 Checkpoint 做逐 step 有限性检查，至少覆盖 200 step；每步同步并检查 loss、梯度、参数和优化器状态，比较显式 Attention 与 cuDNN 的首次分叉位置。修复前默认关闭该开关。
+2. **缩短有效 token 长度。** 全数据扫描得到 Pi0.5 tokenized state 最大长度 165；可在显式 Attention 路径单独 A/B `max_token_len=168` 与当前 200，减少 Gemma attention/MLP 的序列计算，同时保留 3 token 余量。
+3. **重新做端到端正确性验收。** 修复后必须先完成至少 500 个有限 loss/gradient 的训练 step，再讨论更长的吞吐收益和默认开启；独立 Attention smoke test 不再视为充分证据。
 4. **降低 checkpoint 干扰。** 长训练若不需要每 1000 step 恢复点，可提高 `save_interval`；此前生产日志已经确认 checkpoint 会造成下一 step 的数据等待尖峰。
 
-最终判断：共享内存 Tensor 消除了主要数据交付开销；日志 step 范数和三路相机合并是低风险增量收益；cuDNN fused Attention 是当前最大且已验证的稳态收益。后续应基于启用 cuDNN 后的新计算图继续优化，并按当前需求不把约 10 分钟冷启动作为训练吞吐优化目标。
+最终判断：共享内存 Tensor 消除了主要数据交付开销；日志 step 范数和三路相机合并是低风险增量收益；cuDNN fused Attention 具有最大的速度收益，但生产长跑数值失败，当前只能作为实验性、不安全的优化。下一步的最高优先级是恢复到切换前 Checkpoint 并定位 NaN，而不是继续运行或推广这条路径。

From c1b70b2bd054754679561b8e6613d88d89f419c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Sat, 29 Aug 2026 08:34:50 +0000
Subject: [PATCH 06/28] Optimize H200 training and fix cuDNN attention NaNs

---
 ...lexiv_earbud_case_insertion_0826_h200.yaml |  31 +++++
 docs/2026-08-28-h200-training-optimization.md |  17 ++-
 ...2026-08-29-cudnn-attention-nan-incident.md | 117 ++++++++++++++++++
 scripts/train.py                              |  28 ++++-
 scripts/train_test.py                         |  12 ++
 src/openpi/models/gemma.py                    |  15 ++-
 src/openpi/models/pi0_test.py                 |  22 ++++
 7 files changed, 233 insertions(+), 9 deletions(-)
 create mode 100644 configs/_examples/pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml
 create mode 100644 docs/2026-08-29-cudnn-attention-nan-incident.md

diff --git a/configs/_examples/pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml b/configs/_examples/pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml
new file mode 100644
index 0000000..0f6527d
--- /dev/null
+++ b/configs/_examples/pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml
@@ -0,0 +1,31 @@
+model:
+  type: Pi0Config
+  max_token_len: 200
+  pi05: true
+  discrete_state_input: true
+  enable_training_time_rtc: true
+  batch_image_views: true
+  use_cudnn_attention: true
+weight_loader:
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+ema_decay: null
+data:
+  type: LeRobotBiFlexivDataConfig
+  repo_id: Xense/earbud_case_insertion_teleop_0515_left_right
+  base_config:
+    prompt_from_task: true
+    # 这个数据集有 7 路视频流，其中 4 路触觉流被 repack 直接丢弃。
+    # false（默认）= 不解码它们，单样本 __getitem__ 快约 1.7 倍。
+    # 用触觉的模型（Pi0Tactile 等）必须设 true。
+    tactile: false
+  default_prompt: Pick up each earbud case from the left stands, insert the matching earbuds, close the lid, and place the case in the box.
+batch_size: 512
+num_workers: 96
+num_train_steps: 60000
+fsdp_devices: 8
+# 是否严格按 sampler 顺序交付 batch。false（默认）= 谁先算完谁先交付，
+# 消除每 num_workers 步一次的队头阻塞停顿；同时启用 InfiniteSampler，
+# 去掉 epoch 边界重建迭代器的停顿。详见 docs/data-pipeline-stalls.md。
+# 代价：batch 的交付顺序不再逐次可复现（每个 index 每轮仍恰好出现一次）。
+strict_batch_order: false
diff --git a/docs/2026-08-28-h200-training-optimization.md b/docs/2026-08-28-h200-training-optimization.md
index a3296f5..7c67203 100644
--- a/docs/2026-08-28-h200-training-optimization.md
+++ b/docs/2026-08-28-h200-training-optimization.md
@@ -1,5 +1,8 @@
 # H200 训练性能优化记录
 
+> [!WARNING]
+> 2026-08-29 生产长跑发现：cuDNN Attention 遇到 padding 产生的全空 mask 行时，前向有限但 Q 梯度为 NaN。修复、影响范围和 checkpoint 恢复要求见 [2026-08-29-cudnn-attention-nan-incident.md](2026-08-29-cudnn-attention-nan-incident.md)。未包含该修复的实现不得用于有效训练。
+
 > **跨机器生产启动命令**：先激活目标 Conda 环境，再用 `$CONDA_PREFIX/lib` 相对该环境选择其中的 cuDNN。环境变量只作用于本次训练进程及其子进程，不修改环境内已安装的包。
 
 ```bash
@@ -10,7 +13,7 @@ LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
     --overwrite
 ```
 
-> **2026-08-29 长跑验证警告**：下面的命令和开关保留为复现记录，但当前 cuDNN Attention 不应继续用于生产训练。生产长跑在 step 16100 首次记录到 NaN 梯度，step 16200 起 loss 持续为 NaN。短跑的数值对齐不足以覆盖该问题，详见“生产长跑验证”。
+> **2026-08-29 长跑事故（已修复）**：原实现从 step 16100 起记录到 NaN；受污染 checkpoint 已删除，只能从保留的 step 15000 恢复。本次修复为全空 mask 行补一个不会影响有效 token 的 dummy key，并在训练更新中增加有限性保护；详见事故文档。
 
 `my_task` 的模型配置需要包含 `use_cudnn_attention: true`；如果 YAML 中没有设置，则在命令末尾增加 `--model.use-cudnn-attention`。`--overwrite` 会删除同名实验的原 Checkpoint 目录；恢复已有训练应改用 `--resume`。
 
@@ -293,18 +296,20 @@ step 16100 是第一个写入指标的生产日志点，因此只能确定 NaN 
 - 原 Attention 和 cuDNN Attention 均完成一次 step 16000–16040 的 8×H200 短跑，没有 OOM、cuDNN engine fallback 或训练异常，也没有保存新 Checkpoint；
 - cuDNN 短跑退出时个别 DataLoader worker 在清理阶段打印 `killed by signal: Aborted`，发生在训练完成和计时结束之后，主训练进程与 GPU 资源均正常退出；这不是 Attention 计算错误。
 - 后续生产长跑推翻了短跑的数值有效性结论：step 16100 首次记录到 NaN 梯度，step 16200 起 loss 持续为 NaN；这说明现有 smoke test 和 40-step 吞吐短跑不足以作为生产正确性验证。
+- 最小 H200 复现确认全空 mask 行会让原 cuDNN 路径产生 16384 个 Q 梯度 NaN；修复后 Q/K/V 梯度全部有限，且有效 query 输出不变；
+- 从干净 step 15000 checkpoint 做 8×H200、batch 512 的 60-step 实际短跑，loss、梯度、参数和更新均保持有限，`update_is_finite=1`。
 
 ## 8. 当前结论与下一步
 
 当前稳态数据交付只占约 `0.04–0.05 s/step`，H2D 只有数毫秒；XProf 又显示已记录 kernel 时间主要集中在 GEMM/Fusion。因此，继续微调 queue、collate 或 H2D 不太可能带来显著收益。
 
-cuDNN fused Attention 的速度收益是目前最大的模型侧收益：生产长跑普通 step 为 `1.910590 s`，摊销日志和 Checkpoint 后约 `1.917238 s/step`。但是该长跑从 step 16100 起出现 NaN，故当前实现没有通过生产正确性验收。确认 JAX runtime 为 91400 只能证明动态库选择正确，不能证明训练数值正确；在修复 NaN 前不得将 `use_cudnn_attention` 用于有效训练。
+cuDNN fused Attention 的速度收益是目前最大的模型侧收益：生产长跑普通 step 为 `1.910590 s`，摊销日志和 Checkpoint 后约 `1.917238 s/step`。原实现因全空 mask 行导致 NaN，现已修复并增加训练更新有限性保护。确认 JAX runtime 为 91400 仍只能证明动态库选择正确；生产验收还需要持续观察恢复训练的 loss、梯度和参数有限性。
 
 下一轮稳态优化建议按以下顺序推进：
 
-1. **先定位 cuDNN NaN。** 从保留的 step 15000 Checkpoint 做逐 step 有限性检查，至少覆盖 200 step；每步同步并检查 loss、梯度、参数和优化器状态，比较显式 Attention 与 cuDNN 的首次分叉位置。修复前默认关闭该开关。
-2. **缩短有效 token 长度。** 全数据扫描得到 Pi0.5 tokenized state 最大长度 165；可在显式 Attention 路径单独 A/B `max_token_len=168` 与当前 200，减少 Gemma attention/MLP 的序列计算，同时保留 3 token 余量。
-3. **重新做端到端正确性验收。** 修复后必须先完成至少 500 个有限 loss/gradient 的训练 step，再讨论更长的吞吐收益和默认开启；独立 Attention smoke test 不再视为充分证据。
+1. **完成修复后的生产长跑验收。** 从干净 step 15000 Checkpoint 恢复，至少完成 500 个有限 loss/gradient 的训练 step；有限性保护一旦拒绝更新必须立即停止，不能保存候选状态。
+2. **缩短有效 token 长度。** 全数据扫描得到 Pi0.5 tokenized state 最大长度 165；可单独 A/B `max_token_len=168` 与当前 200，减少 Gemma attention/MLP 的序列计算，同时保留 3 token 余量。
+3. **继续拆最大的 GEMM/Fusion。** 在生产正确性验收后重新采集 trace，再确认剩余最大的 SigLIP/Gemma GEMM 和 fusion。
 4. **降低 checkpoint 干扰。** 长训练若不需要每 1000 step 恢复点，可提高 `save_interval`；此前生产日志已经确认 checkpoint 会造成下一 step 的数据等待尖峰。
 
-最终判断：共享内存 Tensor 消除了主要数据交付开销；日志 step 范数和三路相机合并是低风险增量收益；cuDNN fused Attention 具有最大的速度收益，但生产长跑数值失败，当前只能作为实验性、不安全的优化。下一步的最高优先级是恢复到切换前 Checkpoint 并定位 NaN，而不是继续运行或推广这条路径。
+最终判断：共享内存 Tensor 消除了主要数据交付开销；日志 step 范数和三路相机合并是低风险增量收益；cuDNN fused Attention 的全空 mask 梯度 NaN 已定位并修复，最小复现和 60-step 真实训练通过。当前最高优先级是让从 step 15000 恢复的训练完成更长的有限性验收。
diff --git a/docs/2026-08-29-cudnn-attention-nan-incident.md b/docs/2026-08-29-cudnn-attention-nan-incident.md
new file mode 100644
index 0000000..28f727e
--- /dev/null
+++ b/docs/2026-08-29-cudnn-attention-nan-incident.md
@@ -0,0 +1,117 @@
+# cuDNN Attention 全空 Mask 导致训练 NaN
+
+**日期**：2026-08-29
+
+**机器**：8×H200
+
+**配置**：`pi05_base_bi_flexiv_earbud_case_insertion_0826_h200`
+
+**影响范围**：从 clean step 16000 切换到 `use_cudnn_attention=true` 后的生产训练
+
+## 摘要
+
+生产训练恢复后，第一个记录窗口已经出现 NaN 梯度，随后 loss 和参数持续为 NaN。根因不是 cuDNN 动态库版本、数据加载、显存或普通的梯度爆炸，而是 cuDNN fused Attention 在 attention mask 存在整行全 `False` 时的反向数值行为：前向输出有限，但 Q 梯度为 NaN。
+
+Pi0.5 的 prompt 会补齐到 `max_token_len=200`。padding token 在 `make_attn_mask` 中既不是有效 query，也不是有效 key，因此必然生成全 `False` query 行。原显式 Attention 用有限的 `big_neg` 处理这些行，不会产生同样的 Q 梯度 NaN；直接切换到 cuDNN 实现暴露了这一差异。
+
+训练已停止。当前最后一个确定干净且仍保留的 checkpoint 是 step 15000；step 16000 已被保留策略删除，step 20000–55000 均已污染。
+
+## 现象与时间线
+
+切换前的原 Attention 训练指标正常：
+
+```text
+Step 16000: grad_norm=0.6721, loss=0.1429, param_norm=1806.9792
+Step 16100: grad_norm=1.1132, loss=0.1456, param_norm=1806.9926
+Step 16200: grad_norm=0.8198, loss=0.1402, param_norm=1807.0051
+Step 16300: grad_norm=0.9735, loss=0.1410, param_norm=1807.0171
+```
+
+从 step 16000 checkpoint 用 cuDNN Attention 恢复后：
+
+```text
+Step 16100: grad_norm=nan, loss=0.1185, param_norm=nan
+Step 16200: grad_norm=nan, loss=nan,    param_norm=nan
+Step 16300 及以后: grad_norm=nan, loss=nan, param_norm=nan
+```
+
+checkpoint 目录名记录的是循环 step，而保存的 `train_state.step` 已经加一。因此恢复后第一个实际更新的前向 loss 仍有限，但该步反向产生 NaN 并污染参数；后续前向全部变为 NaN。
+
+日志又对整个窗口的所有字段使用 `nanmean`。第一个窗口中唯一有限的首步 loss 会让聚合结果仍显示为 `0.1185`，从而掩盖其余 NaN。grad norm 和 param norm 只在日志 step 真正计算，所以当时已经直接显示 NaN。
+
+## 最小复现
+
+在 JAX 实际加载 cuDNN runtime `91400` 的同一环境中，使用 BF16 GQA 和含 16 个全空 query 行的 mask：
+
+```text
+forward output finite: true
+Q/K/V gradient finite: [false, true, true]
+Q gradient NaN count: 16384
+```
+
+为每个全空 query 行临时开放第 0 个 key 后：
+
+```text
+Q/K/V gradient finite: [true, true, true]
+Q/K/V gradient NaN count: [0, 0, 0]
+valid query rows unchanged: true
+```
+
+这解释了完整训练的全部现象：首步前向正常、首步梯度 NaN、更新后参数 NaN、随后 loss NaN。
+
+## 修复
+
+cuDNN 分支在调用 `jax.nn.dot_product_attention` 前构造 `cudnn_mask`：
+
+1. 用 `any(..., axis=-1)` 找出至少有一个可见 key 的 query 行；
+2. 只对全空行开放一个 dummy key；
+3. 所有非空行保持逐元素不变；
+4. padded key 对有效 query 仍然不可见，因此 dummy 行的输出不能影响有效 token 或训练 loss。
+
+训练循环同时增加第二层保护：
+
+- 每步检查 loss、梯度和 optimizer update 是否全部有限；
+- 非有限更新不会写入参数或 optimizer state；
+- 在下一个日志边界停止训练，并在 checkpoint 保存前退出；
+- 抛出明确的 `FloatingPointError`，避免静默生成污染 checkpoint。
+
+## 验证要求
+
+修复验收按以下顺序执行：
+
+1. mask 单元测试确认只修改全空 query 行；
+2. H200/cuDNN 最小反向测试确认 Q/K/V 梯度全部有限；
+3. 从 step 15000 复制到独立诊断实验目录，禁用 W&B 和 checkpoint 保存；
+4. 至少运行 200–500 step，逐日志窗口确认 `update_is_finite=1`、loss/grad norm/param norm 有限；
+5. 通过后才能清理污染 checkpoint 或恢复生产训练。
+
+## 2026-08-29 修复验证结果
+
+代码与单元回归：
+
+- `python -m compileall` 通过；
+- `git diff --check` 通过；
+- `src/openpi/models/pi0_test.py`：`8 passed`；
+- `scripts/train_test.py`：`2 passed`，包括两步训练和 checkpoint 恢复；
+- H200/cuDNN BF16 GQA 最小反向测试：Q/K/V 梯度全部有限，NaN 数量均为 0，有效 mask 行逐元素不变。
+
+真实训练使用 8×H200、全局 batch 512、cuDNN runtime 91400，从 clean step 15000 通过 `/tmp` 独立目录只读恢复。W&B、周期保存和最终保存均关闭。为缩短冷启动只启用 16 个 DataLoader worker，因此本次不统计吞吐。
+
+主动结束前完成了 step 15001–15060 共 60 个更新；每 10 step 的检查结果如下：
+
+```text
+Step 15010: grad_norm=0.7763, loss=0.1438, param_norm=1806.8403, update_is_finite=1.0000
+Step 15020: grad_norm=1.0680, loss=0.1666, param_norm=1806.8416, update_is_finite=1.0000
+Step 15030: grad_norm=0.8248, loss=0.1507, param_norm=1806.8430, update_is_finite=1.0000
+Step 15040: grad_norm=0.8843, loss=0.1576, param_norm=1806.8450, update_is_finite=1.0000
+Step 15050: grad_norm=0.7334, loss=0.1528, param_norm=1806.8467, update_is_finite=1.0000
+Step 15060: grad_norm=1.4659, loss=0.1655, param_norm=1806.8483, update_is_finite=1.0000
+```
+
+原缺陷会在恢复后的第一个更新产生 NaN，因此 60 个连续有限更新已经验证修复覆盖了已知故障。短跑结束后 8 张 GPU 显存均为 0 MiB，原 checkpoint 未被修改。恢复生产前仍建议完成 200–500 step 的长验证，以覆盖更多数据样本和训练阶段。
+
+## 恢复注意事项
+
+不能直接对原实验执行普通 `--resume`，因为它会选择最新的 step 55000 污染 checkpoint。恢复时必须显式以保留的 step 15000 为源，使用独立实验目录验证；确认修复后，再决定是否归档或移除 step 20000–55000。
+
+原 step 16000 已删除，无法无损续接；从 step 15000 恢复意味着重跑约 1000 个有效 step。
diff --git a/scripts/train.py b/scripts/train.py
index f65f422..d368094 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -49,6 +49,10 @@ def format(self, record):
     logger.handlers[0].setFormatter(formatter)
 
 
+def _all_finite(tree):
+    return jnp.all(jnp.stack([jnp.all(jnp.isfinite(x)) for x in jax.tree.leaves(tree)]))
+
+
 def init_wandb(config: _config.TrainConfig, *, resuming: bool, log_code: bool = False, enabled: bool = True):
     if not enabled:
         wandb.init(mode="disabled")
@@ -161,8 +165,17 @@ def loss_fn(
     loss, grads = nnx.value_and_grad(loss_fn, argnums=diff_state)(model, train_rng, observation, actions)
 
     params = state.params.filter(config.trainable_filter)
-    updates, new_opt_state = state.tx.update(grads, state.opt_state, params)
-    new_params = optax.apply_updates(params, updates)
+    updates, candidate_opt_state = state.tx.update(grads, state.opt_state, params)
+    candidate_params = optax.apply_updates(params, updates)
+
+    # Never commit a non-finite update. This reduction is intentionally done on
+    # every step: a single NaN update can poison every later checkpoint before a
+    # periodic logging step notices it.
+    update_is_finite = _all_finite((loss, grads, updates))
+    new_params = jax.tree.map(lambda old, new: jnp.where(update_is_finite, new, old), params, candidate_params)
+    new_opt_state = jax.tree.map(
+        lambda old, new: jnp.where(update_is_finite, new, old), state.opt_state, candidate_opt_state
+    )
 
     # Update the model in place and return the new full state.
     nnx.update(model, new_params)
@@ -200,6 +213,7 @@ def loss_fn(
         "loss": loss,
         "grad_norm": grad_norm,
         "param_norm": param_norm,
+        "update_is_finite": update_is_finite,
     }
     return new_state, info
 
@@ -284,6 +298,7 @@ def main(config: _config.TrainConfig):
     )
 
     infos = []
+    nonfinite_step = None
     xprof_active = False
     # --- stall diagnostics ---
     stall_threshold_s = 3.0  # log a warning when any phase exceeds this
@@ -337,6 +352,13 @@ def main(config: _config.TrainConfig):
             wandb.log(reduced_info, step=step)
             infos = []
             t_log = time.monotonic() - t0
+            if float(reduced_info["update_is_finite"]) < 1.0:
+                nonfinite_step = step
+                pbar.write(
+                    f"[NONFINITE step={step}] At least one update since the previous log was skipped; "
+                    "stopping training."
+                )
+                break
 
         t0 = time.monotonic()
         batch = next(data_iter)
@@ -401,6 +423,8 @@ def main(config: _config.TrainConfig):
     data_loader.close()
     logging.info("Waiting for checkpoint manager to finish")
     checkpoint_manager.wait_until_finished()
+    if nonfinite_step is not None:
+        raise FloatingPointError(f"Non-finite loss, gradients, or updates detected by step {nonfinite_step}")
 
 
 if __name__ == "__main__":
diff --git a/scripts/train_test.py b/scripts/train_test.py
index 6165930..daa627c 100644
--- a/scripts/train_test.py
+++ b/scripts/train_test.py
@@ -6,11 +6,23 @@
 
 os.environ["JAX_PLATFORMS"] = "cpu"
 
+import jax.numpy as jnp
+
 from openpi.training import config as _config
 
 from . import train
 
 
+def test_nonfinite_update_is_rejected():
+    loss = jnp.asarray(1.0)
+    grads = {"finite": jnp.ones((2,)), "bad": jnp.array([0.0, jnp.nan])}
+    updates = {"finite": jnp.ones((2,)), "bad": jnp.zeros((2,))}
+
+    update_is_finite = train._all_finite((loss, grads, updates))
+
+    assert not bool(update_is_finite)
+
+
 @pytest.mark.parametrize("config_name", ["debug_pi05"])
 def test_train(tmp_path: pathlib.Path, config_name: str):
     config = dataclasses.replace(
diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index 90cbf47..e2d5ca8 100644
--- a/src/openpi/models/gemma.py
+++ b/src/openpi/models/gemma.py
@@ -160,6 +160,13 @@ def decode(self, x):
         return jnp.dot(x, self.input_embedding_table.T)
 
 
+def _make_cudnn_attention_mask_safe(attn_mask):
+    """Give fully masked query rows one dummy key for finite cuDNN gradients."""
+    row_has_key = jnp.any(attn_mask, axis=-1, keepdims=True)
+    dummy_key = jnp.arange(attn_mask.shape[-1]) == 0
+    return jnp.logical_or(attn_mask, jnp.logical_and(~row_has_key, dummy_key))
+
+
 @at.typecheck
 class Attention(nn.Module):
     """Attention module."""
@@ -229,11 +236,17 @@ def __call__(self, xs, positions, attn_mask, kv_cache):
             # q is already scaled above, so disable dot_product_attention's
             # default head-dimension scaling. Keep cached inference on the
             # existing implementation; this switch targets training throughput.
+            # cuDNN 9.14 returns finite outputs but NaN q-gradients for query
+            # rows whose mask is entirely false. Padding tokens create exactly
+            # those rows. Give each empty row one dummy key; padded query outputs
+            # cannot affect valid tokens because padded keys remain masked for
+            # every valid query.
+            cudnn_mask = _make_cudnn_attention_mask_safe(attn_mask)
             encoded = jax.nn.dot_product_attention(
                 q,
                 k,
                 v,
-                mask=attn_mask,
+                mask=cudnn_mask,
                 scale=1.0,
                 implementation="cudnn",
             )
diff --git a/src/openpi/models/pi0_test.py b/src/openpi/models/pi0_test.py
index 239c8a7..31ac570 100644
--- a/src/openpi/models/pi0_test.py
+++ b/src/openpi/models/pi0_test.py
@@ -29,6 +29,28 @@ def test_encode_image_views_as_batch_matches_separate_calls():
         assert jnp.array_equal(actual[name], expected[name])
 
 
+def test_cudnn_attention_mask_fills_only_empty_query_rows():
+    mask = jnp.array(
+        [
+            [
+                [
+                    [True, False, False, False],
+                    [True, True, False, False],
+                    [False, False, False, False],
+                    [False, False, False, False],
+                ]
+            ]
+        ],
+        dtype=jnp.bool_,
+    )
+
+    safe_mask = _gemma._make_cudnn_attention_mask_safe(mask)
+
+    assert jnp.array_equal(safe_mask[..., :2, :], mask[..., :2, :])
+    assert jnp.all(jnp.any(safe_mask, axis=-1))
+    assert jnp.array_equal(safe_mask[..., 2:, :], jnp.array([[[[True, False, False, False]] * 2]]))
+
+
 def _get_frozen_state(config: _pi0_config.Pi0Config) -> nnx.State:
     abstract_model = nnx.eval_shape(config.create, jax.random.key(0))
 

From 6305750afcb81cb1b7742269dae65b8bae5bf45f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Sat, 29 Aug 2026 11:48:36 +0000
Subject: [PATCH 07/28] Remove per-step finite update guard

---
 docs/2026-08-28-h200-training-optimization.md |  6 ++--
 ...2026-08-29-cudnn-attention-nan-incident.md | 11 +++-----
 scripts/train.py                              | 28 ++-----------------
 scripts/train_test.py                         | 12 --------
 4 files changed, 9 insertions(+), 48 deletions(-)

diff --git a/docs/2026-08-28-h200-training-optimization.md b/docs/2026-08-28-h200-training-optimization.md
index 7c67203..c294a77 100644
--- a/docs/2026-08-28-h200-training-optimization.md
+++ b/docs/2026-08-28-h200-training-optimization.md
@@ -13,7 +13,7 @@ LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
     --overwrite
 ```
 
-> **2026-08-29 长跑事故（已修复）**：原实现从 step 16100 起记录到 NaN；受污染 checkpoint 已删除，只能从保留的 step 15000 恢复。本次修复为全空 mask 行补一个不会影响有效 token 的 dummy key，并在训练更新中增加有限性保护；详见事故文档。
+> **2026-08-29 长跑事故（已修复）**：原实现从 step 16100 起记录到 NaN；受污染 checkpoint 已删除，并从保留的 step 15000 恢复。本次修复为全空 mask 行补一个不会影响有效 token 的 dummy key。逐 step 有限性保护因约 22.7% 的吞吐开销已按生产要求移除，训练继续按原逻辑记录 loss；详见事故文档。
 
 `my_task` 的模型配置需要包含 `use_cudnn_attention: true`；如果 YAML 中没有设置，则在命令末尾增加 `--model.use-cudnn-attention`。`--overwrite` 会删除同名实验的原 Checkpoint 目录；恢复已有训练应改用 `--resume`。
 
@@ -303,11 +303,11 @@ step 16100 是第一个写入指标的生产日志点，因此只能确定 NaN 
 
 当前稳态数据交付只占约 `0.04–0.05 s/step`，H2D 只有数毫秒；XProf 又显示已记录 kernel 时间主要集中在 GEMM/Fusion。因此，继续微调 queue、collate 或 H2D 不太可能带来显著收益。
 
-cuDNN fused Attention 的速度收益是目前最大的模型侧收益：生产长跑普通 step 为 `1.910590 s`，摊销日志和 Checkpoint 后约 `1.917238 s/step`。原实现因全空 mask 行导致 NaN，现已修复并增加训练更新有限性保护。确认 JAX runtime 为 91400 仍只能证明动态库选择正确；生产验收还需要持续观察恢复训练的 loss、梯度和参数有限性。
+cuDNN fused Attention 的速度收益是目前最大的模型侧收益：生产长跑普通 step 为 `1.910590 s`，摊销日志和 Checkpoint 后约 `1.917238 s/step`。原实现因全空 mask 行导致 NaN，现已修复。曾加入的逐 step 有限性保护把普通 step 增至 `2.345000 s`，因此已按生产要求移除，保留原有 loss/grad norm/param norm 日志。确认 JAX runtime 为 91400 仍只能证明动态库选择正确；生产验收需要持续观察日志。
 
 下一轮稳态优化建议按以下顺序推进：
 
-1. **完成修复后的生产长跑验收。** 从干净 step 15000 Checkpoint 恢复，至少完成 500 个有限 loss/gradient 的训练 step；有限性保护一旦拒绝更新必须立即停止，不能保存候选状态。
+1. **完成修复后的生产长跑验收。** 从干净 checkpoint 恢复并持续观察 loss、grad norm 和 param norm；当前策略不会自动拒绝或停止非有限更新，发现异常时需人工停止并回退到前一个干净 checkpoint。
 2. **缩短有效 token 长度。** 全数据扫描得到 Pi0.5 tokenized state 最大长度 165；可单独 A/B `max_token_len=168` 与当前 200，减少 Gemma attention/MLP 的序列计算，同时保留 3 token 余量。
 3. **继续拆最大的 GEMM/Fusion。** 在生产正确性验收后重新采集 trace，再确认剩余最大的 SigLIP/Gemma GEMM 和 fusion。
 4. **降低 checkpoint 干扰。** 长训练若不需要每 1000 step 恢复点，可提高 `save_interval`；此前生产日志已经确认 checkpoint 会造成下一 step 的数据等待尖峰。
diff --git a/docs/2026-08-29-cudnn-attention-nan-incident.md b/docs/2026-08-29-cudnn-attention-nan-incident.md
index 28f727e..17dad03 100644
--- a/docs/2026-08-29-cudnn-attention-nan-incident.md
+++ b/docs/2026-08-29-cudnn-attention-nan-incident.md
@@ -68,12 +68,9 @@ cuDNN 分支在调用 `jax.nn.dot_product_attention` 前构造 `cudnn_mask`：
 3. 所有非空行保持逐元素不变；
 4. padded key 对有效 query 仍然不可见，因此 dummy 行的输出不能影响有效 token 或训练 loss。
 
-训练循环同时增加第二层保护：
+训练循环曾增加逐 step 有限性保护，但生产统计显示普通 step 从 `1.910590 s` 增至 `2.345000 s`，约 98.7% 的新增时间位于模型更新的 `dispatch` 阶段。原因是每步扫描完整梯度和 optimizer update，并对全部参数和 optimizer state 做逐元素回退选择。根据 2026-08-29 的生产运行要求，该保护已移除，恢复原更新路径；训练仍按原逻辑在日志边界输出 `loss`、`grad_norm` 和 `param_norm`，不再自动拒绝更新或因非有限值停止。
 
-- 每步检查 loss、梯度和 optimizer update 是否全部有限；
-- 非有限更新不会写入参数或 optimizer state；
-- 在下一个日志边界停止训练，并在 checkpoint 保存前退出；
-- 抛出明确的 `FloatingPointError`，避免静默生成污染 checkpoint。
+这是一项明确的吞吐优先取舍：日志只能事后暴露数值异常，不能保证异常更新不会进入参数或下一个 checkpoint。
 
 ## 验证要求
 
@@ -82,7 +79,7 @@ cuDNN 分支在调用 `jax.nn.dot_product_attention` 前构造 `cudnn_mask`：
 1. mask 单元测试确认只修改全空 query 行；
 2. H200/cuDNN 最小反向测试确认 Q/K/V 梯度全部有限；
 3. 从 step 15000 复制到独立诊断实验目录，禁用 W&B 和 checkpoint 保存；
-4. 至少运行 200–500 step，逐日志窗口确认 `update_is_finite=1`、loss/grad norm/param norm 有限；
+4. 至少运行 200–500 step，逐日志窗口确认 loss/grad norm/param norm 有限；
 5. 通过后才能清理污染 checkpoint 或恢复生产训练。
 
 ## 2026-08-29 修复验证结果
@@ -112,6 +109,6 @@ Step 15060: grad_norm=1.4659, loss=0.1655, param_norm=1806.8483, update_is_finit
 
 ## 恢复注意事项
 
-不能直接对原实验执行普通 `--resume`，因为它会选择最新的 step 55000 污染 checkpoint。恢复时必须显式以保留的 step 15000 为源，使用独立实验目录验证；确认修复后，再决定是否归档或移除 step 20000–55000。
+污染的 step 20000–55000 已删除，生产训练已从 step 15000 恢复并保存 step 18000。逐 step 保护在 step 18600 报告最近 100 步中有 8 次更新被拒绝，当时窗口聚合 loss 仍为有限值 `0.7085`，因此只看 loss 不足以证明每次更新都有限。移除保护后的生产训练从最后保留的 step 18000 继续。
 
 原 step 16000 已删除，无法无损续接；从 step 15000 恢复意味着重跑约 1000 个有效 step。
diff --git a/scripts/train.py b/scripts/train.py
index d368094..f65f422 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -49,10 +49,6 @@ def format(self, record):
     logger.handlers[0].setFormatter(formatter)
 
 
-def _all_finite(tree):
-    return jnp.all(jnp.stack([jnp.all(jnp.isfinite(x)) for x in jax.tree.leaves(tree)]))
-
-
 def init_wandb(config: _config.TrainConfig, *, resuming: bool, log_code: bool = False, enabled: bool = True):
     if not enabled:
         wandb.init(mode="disabled")
@@ -165,17 +161,8 @@ def loss_fn(
     loss, grads = nnx.value_and_grad(loss_fn, argnums=diff_state)(model, train_rng, observation, actions)
 
     params = state.params.filter(config.trainable_filter)
-    updates, candidate_opt_state = state.tx.update(grads, state.opt_state, params)
-    candidate_params = optax.apply_updates(params, updates)
-
-    # Never commit a non-finite update. This reduction is intentionally done on
-    # every step: a single NaN update can poison every later checkpoint before a
-    # periodic logging step notices it.
-    update_is_finite = _all_finite((loss, grads, updates))
-    new_params = jax.tree.map(lambda old, new: jnp.where(update_is_finite, new, old), params, candidate_params)
-    new_opt_state = jax.tree.map(
-        lambda old, new: jnp.where(update_is_finite, new, old), state.opt_state, candidate_opt_state
-    )
+    updates, new_opt_state = state.tx.update(grads, state.opt_state, params)
+    new_params = optax.apply_updates(params, updates)
 
     # Update the model in place and return the new full state.
     nnx.update(model, new_params)
@@ -213,7 +200,6 @@ def loss_fn(
         "loss": loss,
         "grad_norm": grad_norm,
         "param_norm": param_norm,
-        "update_is_finite": update_is_finite,
     }
     return new_state, info
 
@@ -298,7 +284,6 @@ def main(config: _config.TrainConfig):
     )
 
     infos = []
-    nonfinite_step = None
     xprof_active = False
     # --- stall diagnostics ---
     stall_threshold_s = 3.0  # log a warning when any phase exceeds this
@@ -352,13 +337,6 @@ def main(config: _config.TrainConfig):
             wandb.log(reduced_info, step=step)
             infos = []
             t_log = time.monotonic() - t0
-            if float(reduced_info["update_is_finite"]) < 1.0:
-                nonfinite_step = step
-                pbar.write(
-                    f"[NONFINITE step={step}] At least one update since the previous log was skipped; "
-                    "stopping training."
-                )
-                break
 
         t0 = time.monotonic()
         batch = next(data_iter)
@@ -423,8 +401,6 @@ def main(config: _config.TrainConfig):
     data_loader.close()
     logging.info("Waiting for checkpoint manager to finish")
     checkpoint_manager.wait_until_finished()
-    if nonfinite_step is not None:
-        raise FloatingPointError(f"Non-finite loss, gradients, or updates detected by step {nonfinite_step}")
 
 
 if __name__ == "__main__":
diff --git a/scripts/train_test.py b/scripts/train_test.py
index daa627c..6165930 100644
--- a/scripts/train_test.py
+++ b/scripts/train_test.py
@@ -6,23 +6,11 @@
 
 os.environ["JAX_PLATFORMS"] = "cpu"
 
-import jax.numpy as jnp
-
 from openpi.training import config as _config
 
 from . import train
 
 
-def test_nonfinite_update_is_rejected():
-    loss = jnp.asarray(1.0)
-    grads = {"finite": jnp.ones((2,)), "bad": jnp.array([0.0, jnp.nan])}
-    updates = {"finite": jnp.ones((2,)), "bad": jnp.zeros((2,))}
-
-    update_is_finite = train._all_finite((loss, grads, updates))
-
-    assert not bool(update_is_finite)
-
-
 @pytest.mark.parametrize("config_name", ["debug_pi05"])
 def test_train(tmp_path: pathlib.Path, config_name: str):
     config = dataclasses.replace(

From 51d9af30b9089f925ef3d01e5c0a9859535c689a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Sat, 29 Aug 2026 13:32:32 +0000
Subject: [PATCH 08/28] Preserve the fast cuDNN mask for invalid queries

---
 docs/2026-08-28-h200-training-optimization.md |  4 ++--
 ...2026-08-29-cudnn-attention-nan-incident.md | 19 +++++++++--------
 src/openpi/models/gemma.py                    | 21 ++++++++-----------
 src/openpi/models/pi0_test.py                 | 16 +++++++++-----
 4 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/docs/2026-08-28-h200-training-optimization.md b/docs/2026-08-28-h200-training-optimization.md
index c294a77..0621f7c 100644
--- a/docs/2026-08-28-h200-training-optimization.md
+++ b/docs/2026-08-28-h200-training-optimization.md
@@ -13,7 +13,7 @@ LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
     --overwrite
 ```
 
-> **2026-08-29 长跑事故（已修复）**：原实现从 step 16100 起记录到 NaN；受污染 checkpoint 已删除，并从保留的 step 15000 恢复。本次修复为全空 mask 行补一个不会影响有效 token 的 dummy key。逐 step 有限性保护因约 22.7% 的吞吐开销已按生产要求移除，训练继续按原逻辑记录 loss；详见事故文档。
+> **2026-08-29 长跑事故（已修复）**：原实现从 step 16100 起记录到 NaN；受污染 checkpoint 已删除，并从保留的 step 15000 恢复。最终修复保留原始 attention mask，只对全空 query 行停止 Q 梯度。逐 step 有限性保护已按生产要求移除，训练继续按原逻辑记录 loss；详见事故文档。
 
 `my_task` 的模型配置需要包含 `use_cudnn_attention: true`；如果 YAML 中没有设置，则在命令末尾增加 `--model.use-cudnn-attention`。`--overwrite` 会删除同名实验的原 Checkpoint 目录；恢复已有训练应改用 `--resume`。
 
@@ -303,7 +303,7 @@ step 16100 是第一个写入指标的生产日志点，因此只能确定 NaN 
 
 当前稳态数据交付只占约 `0.04–0.05 s/step`，H2D 只有数毫秒；XProf 又显示已记录 kernel 时间主要集中在 GEMM/Fusion。因此，继续微调 queue、collate 或 H2D 不太可能带来显著收益。
 
-cuDNN fused Attention 的速度收益是目前最大的模型侧收益：生产长跑普通 step 为 `1.910590 s`，摊销日志和 Checkpoint 后约 `1.917238 s/step`。原实现因全空 mask 行导致 NaN，现已修复。曾加入的逐 step 有限性保护把普通 step 增至 `2.345000 s`，因此已按生产要求移除，保留原有 loss/grad norm/param norm 日志。确认 JAX runtime 为 91400 仍只能证明动态库选择正确；生产验收需要持续观察日志。
+cuDNN fused Attention 的旧长跑普通 step 为 `1.910590 s`，但该区间从最早日志点起梯度、参数已经 NaN，不能作为有效训练的吞吐基线。修复后连续有限训练的普通 step 稳定约为 `2.31 s`；移除逐 step 有限性保护和把安全 mask 计算移出扫描层都没有改变该结果。最终实现保留原 mask，仅停止全空 query 的 Q 梯度，并保留原有 loss/grad norm/param norm 日志。确认 JAX runtime 为 91400 仍只能证明动态库选择正确；生产验收需要持续观察日志。
 
 下一轮稳态优化建议按以下顺序推进：
 
diff --git a/docs/2026-08-29-cudnn-attention-nan-incident.md b/docs/2026-08-29-cudnn-attention-nan-incident.md
index 17dad03..6f13b33 100644
--- a/docs/2026-08-29-cudnn-attention-nan-incident.md
+++ b/docs/2026-08-29-cudnn-attention-nan-incident.md
@@ -61,14 +61,9 @@ valid query rows unchanged: true
 
 ## 修复
 
-cuDNN 分支在调用 `jax.nn.dot_product_attention` 前构造 `cudnn_mask`：
+cuDNN 分支保留原始 attention mask，不再给全空行增加 dummy key。调用 `jax.nn.dot_product_attention` 前用 `any(..., axis=-1)` 找出有效 query；Q 的前向值逐元素不变，但全空 query 行通过 `stop_gradient` 从反向链路中移除。这样保持原 mask 语义和有效 query 输出，同时把无效 Q 梯度固定为 0。
 
-1. 用 `any(..., axis=-1)` 找出至少有一个可见 key 的 query 行；
-2. 只对全空行开放一个 dummy key；
-3. 所有非空行保持逐元素不变；
-4. padded key 对有效 query 仍然不可见，因此 dummy 行的输出不能影响有效 token 或训练 loss。
-
-训练循环曾增加逐 step 有限性保护，但生产统计显示普通 step 从 `1.910590 s` 增至 `2.345000 s`，约 98.7% 的新增时间位于模型更新的 `dispatch` 阶段。原因是每步扫描完整梯度和 optimizer update，并对全部参数和 optimizer state 做逐元素回退选择。根据 2026-08-29 的生产运行要求，该保护已移除，恢复原更新路径；训练仍按原逻辑在日志边界输出 `loss`、`grad_norm` 和 `param_norm`，不再自动拒绝更新或因非有限值停止。
+训练循环曾增加逐 step 有限性保护。根据 2026-08-29 的生产运行要求，该保护已移除，恢复原更新路径；训练仍按原逻辑在日志边界输出 `loss`、`grad_norm` 和 `param_norm`，不再自动拒绝更新或因非有限值停止。最初曾把 `1.91 -> 2.3 s/step` 归因于该保护，后续 A/B 推翻了这一判断：移除保护后连续有限训练仍约为 `2.31 s/step`。
 
 这是一项明确的吞吐优先取舍：日志只能事后暴露数值异常，不能保证异常更新不会进入参数或下一个 checkpoint。
 
@@ -76,7 +71,7 @@ cuDNN 分支在调用 `jax.nn.dot_product_attention` 前构造 `cudnn_mask`：
 
 修复验收按以下顺序执行：
 
-1. mask 单元测试确认只修改全空 query 行；
+1. 单元测试确认 Q 前向逐元素不变，且只停止全空 query 行的梯度；
 2. H200/cuDNN 最小反向测试确认 Q/K/V 梯度全部有限；
 3. 从 step 15000 复制到独立诊断实验目录，禁用 W&B 和 checkpoint 保存；
 4. 至少运行 200–500 step，逐日志窗口确认 loss/grad norm/param norm 有限；
@@ -90,7 +85,7 @@ cuDNN 分支在调用 `jax.nn.dot_product_attention` 前构造 `cudnn_mask`：
 - `git diff --check` 通过；
 - `src/openpi/models/pi0_test.py`：`8 passed`；
 - `scripts/train_test.py`：`2 passed`，包括两步训练和 checkpoint 恢复；
-- H200/cuDNN BF16 GQA 最小反向测试：Q/K/V 梯度全部有限，NaN 数量均为 0，有效 mask 行逐元素不变。
+- H200/cuDNN BF16 GQA 真实序列形状测试：原始反向产生 71680 个 Q 梯度 NaN；Q stop-gradient 后 Q/K/V 梯度全部有限，NaN 数量均为 0，前向值和有效 query 输出逐元素不变。
 
 真实训练使用 8×H200、全局 batch 512、cuDNN runtime 91400，从 clean step 15000 通过 `/tmp` 独立目录只读恢复。W&B、周期保存和最终保存均关闭。为缩短冷启动只启用 16 个 DataLoader worker，因此本次不统计吞吐。
 
@@ -107,6 +102,12 @@ Step 15060: grad_norm=1.4659, loss=0.1655, param_norm=1806.8483, update_is_finit
 
 原缺陷会在恢复后的第一个更新产生 NaN，因此 60 个连续有限更新已经验证修复覆盖了已知故障。短跑结束后 8 张 GPU 显存均为 0 MiB，原 checkpoint 未被修改。恢复生产前仍建议完成 200–500 step 的长验证，以覆盖更多数据样本和训练阶段。
 
+### 速度复核
+
+上次长跑报告的普通 step `1.910590 s` 来自数值已经失效的区间：第一个生产日志窗口已出现 NaN 梯度，随后参数和 loss 持续为 NaN，因此不能作为有效训练的吞吐基线。修复后连续有限训练的普通 step 稳定约为 `2.31 s`。
+
+16-worker 短跑曾出现 `dispatch=1.82–1.83 s`，但它发生在 80 秒左右的数据断供之后，是 GPU 空闲后的瞬时样本；同一运行的连续 step 仍为 `2.30 s`，并且每约 16 step 出现 30–85 秒 `next_batch` stall。生产继续使用 96 workers 以避免数据断供，约 8 分钟冷启动不计入稳态吞吐。
+
 ## 恢复注意事项
 
 污染的 step 20000–55000 已删除，生产训练已从 step 15000 恢复并保存 step 18000。逐 step 保护在 step 18600 报告最近 100 步中有 8 次更新被拒绝，当时窗口聚合 loss 仍为有限值 `0.7085`，因此只看 loss 不足以证明每次更新都有限。移除保护后的生产训练从最后保留的 step 18000 继续。
diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index e2d5ca8..7955a06 100644
--- a/src/openpi/models/gemma.py
+++ b/src/openpi/models/gemma.py
@@ -160,11 +160,10 @@ def decode(self, x):
         return jnp.dot(x, self.input_embedding_table.T)
 
 
-def _make_cudnn_attention_mask_safe(attn_mask):
-    """Give fully masked query rows one dummy key for finite cuDNN gradients."""
-    row_has_key = jnp.any(attn_mask, axis=-1, keepdims=True)
-    dummy_key = jnp.arange(attn_mask.shape[-1]) == 0
-    return jnp.logical_or(attn_mask, jnp.logical_and(~row_has_key, dummy_key))
+def _stop_gradient_for_fully_masked_queries(q, attn_mask):
+    """Preserve the raw cuDNN mask while zeroing gradients for invalid queries."""
+    query_has_key = jnp.any(attn_mask, axis=-1)[:, 0, :, None, None]
+    return jnp.where(query_has_key, q, jax.lax.stop_gradient(q))
 
 
 @at.typecheck
@@ -236,17 +235,15 @@ def __call__(self, xs, positions, attn_mask, kv_cache):
             # q is already scaled above, so disable dot_product_attention's
             # default head-dimension scaling. Keep cached inference on the
             # existing implementation; this switch targets training throughput.
-            # cuDNN 9.14 returns finite outputs but NaN q-gradients for query
-            # rows whose mask is entirely false. Padding tokens create exactly
-            # those rows. Give each empty row one dummy key; padded query outputs
-            # cannot affect valid tokens because padded keys remain masked for
-            # every valid query.
-            cudnn_mask = _make_cudnn_attention_mask_safe(attn_mask)
+            # cuDNN 9.14 produces NaN q-gradients for fully masked query rows at
+            # the production GQA shape. Preserve the original fast mask and cut
+            # only those invalid query rows out of the backward path.
+            q = _stop_gradient_for_fully_masked_queries(q, attn_mask)
             encoded = jax.nn.dot_product_attention(
                 q,
                 k,
                 v,
-                mask=cudnn_mask,
+                mask=attn_mask,
                 scale=1.0,
                 implementation="cudnn",
             )
diff --git a/src/openpi/models/pi0_test.py b/src/openpi/models/pi0_test.py
index 31ac570..89e0c8e 100644
--- a/src/openpi/models/pi0_test.py
+++ b/src/openpi/models/pi0_test.py
@@ -29,7 +29,7 @@ def test_encode_image_views_as_batch_matches_separate_calls():
         assert jnp.array_equal(actual[name], expected[name])
 
 
-def test_cudnn_attention_mask_fills_only_empty_query_rows():
+def test_cudnn_attention_stops_gradient_only_for_empty_query_rows():
     mask = jnp.array(
         [
             [
@@ -44,11 +44,17 @@ def test_cudnn_attention_mask_fills_only_empty_query_rows():
         dtype=jnp.bool_,
     )
 
-    safe_mask = _gemma._make_cudnn_attention_mask_safe(mask)
+    q = jnp.arange(16, dtype=jnp.float32).reshape(1, 4, 2, 2)
 
-    assert jnp.array_equal(safe_mask[..., :2, :], mask[..., :2, :])
-    assert jnp.all(jnp.any(safe_mask, axis=-1))
-    assert jnp.array_equal(safe_mask[..., 2:, :], jnp.array([[[[True, False, False, False]] * 2]]))
+    def apply_and_sum(q_value):
+        return jnp.sum(_gemma._stop_gradient_for_fully_masked_queries(q_value, mask))
+
+    stopped_q = _gemma._stop_gradient_for_fully_masked_queries(q, mask)
+    grad = jax.grad(apply_and_sum)(q)
+
+    assert jnp.array_equal(stopped_q, q)
+    assert jnp.array_equal(grad[:, :2], jnp.ones_like(grad[:, :2]))
+    assert jnp.array_equal(grad[:, 2:], jnp.zeros_like(grad[:, 2:]))
 
 
 def _get_frozen_state(config: _pi0_config.Pi0Config) -> nnx.State:

From ad08973515a8b246e4acf6db6f4d503d3d6a1e4e Mon Sep 17 00:00:00 2001
From: Hubo1231 <1195132365@qq.com>
Date: Mon, 31 Aug 2026 16:46:47 +0800
Subject: [PATCH 09/28] Stabilize cuDNN attention training diagnostics

---
 README.md                                     |  38 +-
 ...exiv_cube_color_sorting_rtc_0829_h100.yaml |  31 ++
 configs/diag_cudnn_b1_099.yaml                |  45 +++
 configs/diag_cudnn_last3.yaml                 |  39 ++
 configs/diag_cudnn_strict_order.yaml          |  37 ++
 configs/diag_explicit_strict_order.yaml       |  37 ++
 docs/2026-08-28-h200-training-optimization.md |  31 +-
 ...2026-08-29-cudnn-attention-nan-incident.md |   3 +
 docs/2026-08-30-cudnn-attention-divergence.md | 375 ++++++++++++++++++
 pyproject.toml                                |  12 +-
 scripts/check_cuda_stack.py                   | 199 ++++++++++
 scripts/grad_ab_direction_probe.py            | 257 ++++++++++++
 scripts/grad_ab_probe.py                      | 162 ++++++++
 scripts/hybrid_update_ab_probe.py             | 192 +++++++++
 scripts/optimizer_update_ab_probe.py          | 194 +++++++++
 scripts/train.py                              |  52 +++
 scripts/train_test.py                         |  30 ++
 src/openpi/models/gemma.py                    |  80 +++-
 src/openpi/models/pi0.py                      |   2 +
 src/openpi/models/pi0_config.py               |   8 +
 src/openpi/models/pi0_test.py                 |   2 +
 src/openpi/training/config.py                 |   4 +
 22 files changed, 1793 insertions(+), 37 deletions(-)
 create mode 100644 configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml
 create mode 100644 configs/diag_cudnn_b1_099.yaml
 create mode 100644 configs/diag_cudnn_last3.yaml
 create mode 100644 configs/diag_cudnn_strict_order.yaml
 create mode 100644 configs/diag_explicit_strict_order.yaml
 create mode 100644 docs/2026-08-30-cudnn-attention-divergence.md
 create mode 100644 scripts/check_cuda_stack.py
 create mode 100644 scripts/grad_ab_direction_probe.py
 create mode 100644 scripts/grad_ab_probe.py
 create mode 100644 scripts/hybrid_update_ab_probe.py
 create mode 100644 scripts/optimizer_update_ab_probe.py

diff --git a/README.md b/README.md
index f15e4d6..10f3075 100644
--- a/README.md
+++ b/README.md
@@ -63,6 +63,10 @@ GIT_LFS_SKIP_SMUDGE=1 pip install -e .
 # 2. Install the main openpi package
 cd ../..
 GIT_LFS_SKIP_SMUDGE=1 pip install -e .
+
+# 3. Verify that JAX and PyTorch load one coherent pip cu128 stack
+#    (requires an accessible NVIDIA GPU)
+python scripts/check_cuda_stack.py
 ```
 
 ## Model Checkpoints
@@ -184,8 +188,9 @@ model:
   paligemma_variant: gemma_2b
   action_expert_variant: gemma_300m
   enable_training_time_rtc: true
-  # Required when launching with the cuDNN 9.14 setup shown below.
-  use_cudnn_attention: true
+  # Experimental and currently unsafe for production; keep disabled until the
+  # step-0 2000-step convergence gate documented below has passed.
+  use_cudnn_attention: false
   max_delay: 10
 
 data:
@@ -210,27 +215,35 @@ Then use it exactly like any other config:
 
 ```bash
 python scripts/compute_norm_stats.py --config-name my_task
-LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
-  XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
   python scripts/train.py my_task \
     --exp-name=my_exp \
     --overwrite
 python scripts/serve_policy.py policy:checkpoint --policy.config=my_task --policy.dir=checkpoints/my_task/my_exp/<step>
 ```
 
-The training command above selects cuDNN from the currently active Conda
-environment without hard-coding a machine-specific path. It must be paired
-with the following model setting (already shown in the example config):
+Do **not** prepend `$CONDA_PREFIX/lib` to `LD_LIBRARY_PATH`. The supported
+environment uses the pip CUDA 12.8 stack installed by `lerobot-xense`:
+PyTorch 2.11 pins cuDNN 9.19 and JAX shares that same runtime. Before training,
+run `python scripts/check_cuda_stack.py`; it rejects mixed library sources and
+checks the real BF16 GQA forward/backward shape.
+
+`use_cudnn_attention` remains experimental because it has diverged from the
+explicit-attention baseline in a step-0 long run. Production configs must use:
 
 ```yaml
 model:
-  use_cudnn_attention: true
+  use_cudnn_attention: false
 ```
 
 At startup, `scripts/train.py` logs `JAX cuDNN runtime version: <version>`;
-verify that it reports `91400`. Use `--overwrite` only for a new run that may
-replace an existing experiment directory; use `--resume` to preserve and
-continue an existing run.
+the unified PyTorch 2.11 cu128 environment reports `91900`. This number alone
+does not detect mixed dispatcher/engine libraries, so the stack-check script is
+the required gate. Use `--overwrite` only for a new run that may replace an
+existing experiment directory; use `--resume` to preserve and continue an
+existing run. See
+[`docs/2026-08-30-cudnn-attention-divergence.md`](docs/2026-08-30-cudnn-attention-divergence.md)
+before attempting to re-enable fused attention.
 
 ⚠️ Before committing a YAML into `configs/_examples/`, scrub any machine-local
 absolute paths from `weight_loader.params_path` (e.g. `/home/<you>/...`). Use
@@ -266,8 +279,7 @@ python scripts/compute_norm_stats.py --config-name pi05_base_xtac_umi_pick_up_cu
 Now we can kick off training with the following command (the `--overwrite` flag is used to overwrite existing checkpoints if you rerun fine-tuning with the same config):
 
 ```bash
-LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
-  XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
   python scripts/train.py pi05_base_xtac_umi_pick_up_cube_0807_h200 \
     --exp-name=my_experiment \
     --overwrite
diff --git a/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml
new file mode 100644
index 0000000..c15b505
--- /dev/null
+++ b/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml
@@ -0,0 +1,31 @@
+model:
+  type: Pi0Config
+  max_token_len: 200
+  pi05: true
+  discrete_state_input: true
+  enable_training_time_rtc: true
+  batch_image_views: true
+  use_cudnn_attention: true
+weight_loader:
+  # Was /home/li/hubo/xense-openpi/checkpoints/20000/params — a continuation of a
+  # 20k-step run that only existed on one laptop, so nobody else could start this
+  # config. Points at the upstream base now; to continue from your own checkpoint:
+  #   uv run scripts/train.py <this-config> --exp-name=... \
+  #     --weight_loader.params_path=/path/to/checkpoint/params
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+ema_decay: null
+data:
+  type: LeRobotBiFlexivDataConfig
+  repo_id: Xense/cube-color-sorting-0821
+  base_config:
+    prompt_from_task: true
+    tactile: false
+  default_prompt: Put the green cubes in the right bin and the yellow cubes in the left bin.
+batch_size: 256
+num_workers: 64
+num_train_steps: 60000
+save_interval: 10000
+keep_period: 10000
+fsdp_devices: 8
+strict_batch_order: false
diff --git a/configs/diag_cudnn_b1_099.yaml b/configs/diag_cudnn_b1_099.yaml
new file mode 100644
index 0000000..f8108f3
--- /dev/null
+++ b/configs/diag_cudnn_b1_099.yaml
@@ -0,0 +1,45 @@
+model:
+  type: Pi0Config
+  max_token_len: 200
+  pi05: true
+  discrete_state_input: true
+  enable_training_time_rtc: true
+  batch_image_views: true
+  use_cudnn_attention: true
+
+weight_loader:
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+
+optimizer:
+  type: AdamW
+  b1: 0.99
+  b2: 0.95
+  eps: 1.0e-8
+  weight_decay: 1.0e-10
+  clip_gradient_norm: 1.0
+
+ema_decay: null
+
+data:
+  type: LeRobotBiFlexivDataConfig
+  repo_id: Xense/cube-color-sorting-0821
+  assets:
+    assets_dir: /root/localstorage/xense-openpi/assets/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100
+    asset_id: Xense/cube-color-sorting-0821
+  base_config:
+    prompt_from_task: true
+    tactile: false
+  default_prompt: Put the green cubes in the right bin and the yellow cubes in the left bin.
+
+batch_size: 256
+num_workers: 64
+num_train_steps: 3000
+log_interval: 100
+save_interval: 100000
+save_final_checkpoint: false
+keep_period: null
+fsdp_devices: 8
+strict_batch_order: true
+wandb_enabled: false
+overwrite: true
diff --git a/configs/diag_cudnn_last3.yaml b/configs/diag_cudnn_last3.yaml
new file mode 100644
index 0000000..c1025bb
--- /dev/null
+++ b/configs/diag_cudnn_last3.yaml
@@ -0,0 +1,39 @@
+model:
+  type: Pi0Config
+  max_token_len: 200
+  pi05: true
+  discrete_state_input: true
+  enable_training_time_rtc: true
+  batch_image_views: true
+  use_cudnn_attention: true
+  cudnn_attention_layer_start: 15
+  cudnn_attention_num_layers: 3
+
+weight_loader:
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+
+ema_decay: null
+
+data:
+  type: LeRobotBiFlexivDataConfig
+  repo_id: Xense/cube-color-sorting-0821
+  assets:
+    assets_dir: /root/localstorage/xense-openpi/assets/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100
+    asset_id: Xense/cube-color-sorting-0821
+  base_config:
+    prompt_from_task: true
+    tactile: false
+  default_prompt: Put the green cubes in the right bin and the yellow cubes in the left bin.
+
+batch_size: 256
+num_workers: 64
+num_train_steps: 3000
+log_interval: 100
+save_interval: 1000000
+save_final_checkpoint: false
+keep_period: 1000
+fsdp_devices: 8
+strict_batch_order: true
+wandb_enabled: false
+overwrite: true
diff --git a/configs/diag_cudnn_strict_order.yaml b/configs/diag_cudnn_strict_order.yaml
new file mode 100644
index 0000000..7c2b3d7
--- /dev/null
+++ b/configs/diag_cudnn_strict_order.yaml
@@ -0,0 +1,37 @@
+model:
+  type: Pi0Config
+  max_token_len: 200
+  pi05: true
+  discrete_state_input: true
+  enable_training_time_rtc: true
+  batch_image_views: true
+  use_cudnn_attention: true
+
+weight_loader:
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+
+ema_decay: null
+
+data:
+  type: LeRobotBiFlexivDataConfig
+  repo_id: Xense/cube-color-sorting-0821
+  assets:
+    assets_dir: /root/localstorage/xense-openpi/assets/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100
+    asset_id: Xense/cube-color-sorting-0821
+  base_config:
+    prompt_from_task: true
+    tactile: false
+  default_prompt: Put the green cubes in the right bin and the yellow cubes in the left bin.
+
+batch_size: 256
+num_workers: 64
+num_train_steps: 3000
+log_interval: 100
+save_interval: 1000
+save_final_checkpoint: false
+keep_period: 1000
+fsdp_devices: 8
+strict_batch_order: true
+wandb_enabled: false
+overwrite: true
diff --git a/configs/diag_explicit_strict_order.yaml b/configs/diag_explicit_strict_order.yaml
new file mode 100644
index 0000000..d3e70a0
--- /dev/null
+++ b/configs/diag_explicit_strict_order.yaml
@@ -0,0 +1,37 @@
+model:
+  type: Pi0Config
+  max_token_len: 200
+  pi05: true
+  discrete_state_input: true
+  enable_training_time_rtc: true
+  batch_image_views: true
+  use_cudnn_attention: false
+
+weight_loader:
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+
+ema_decay: null
+
+data:
+  type: LeRobotBiFlexivDataConfig
+  repo_id: Xense/cube-color-sorting-0821
+  assets:
+    assets_dir: /root/localstorage/xense-openpi/assets/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100
+    asset_id: Xense/cube-color-sorting-0821
+  base_config:
+    prompt_from_task: true
+    tactile: false
+  default_prompt: Put the green cubes in the right bin and the yellow cubes in the left bin.
+
+batch_size: 256
+num_workers: 64
+num_train_steps: 3000
+log_interval: 100
+save_interval: 100000
+save_final_checkpoint: false
+keep_period: null
+fsdp_devices: 8
+strict_batch_order: true
+wandb_enabled: false
+overwrite: true
diff --git a/docs/2026-08-28-h200-training-optimization.md b/docs/2026-08-28-h200-training-optimization.md
index 0621f7c..8dbb3d7 100644
--- a/docs/2026-08-28-h200-training-optimization.md
+++ b/docs/2026-08-28-h200-training-optimization.md
@@ -1,18 +1,32 @@
 # H200 训练性能优化记录
 
+> [!CAUTION]
+> 2026-08-30 更新：**`use_cudnn_attention` 目前不得用于生产训练。** 全空 mask 修复只消除了瞬时 NaN；带该修复的训练仍会在 step 1000–1200 之间偏离已验证基线并缓慢发散至 NaN。单变量 A/B、已排除的候选和新的验收要求见 [2026-08-30-cudnn-attention-divergence.md](2026-08-30-cudnn-attention-divergence.md)。
+>
+> 本文以下关于启动命令、cuDNN 版本确认和 A/B 吞吐的内容写于该问题定位之前，其中 `LD_LIBRARY_PATH` 与「确认输出 91400」两条**已经过时**，见下方修订。
+
 > [!WARNING]
 > 2026-08-29 生产长跑发现：cuDNN Attention 遇到 padding 产生的全空 mask 行时，前向有限但 Q 梯度为 NaN。修复、影响范围和 checkpoint 恢复要求见 [2026-08-29-cudnn-attention-nan-incident.md](2026-08-29-cudnn-attention-nan-incident.md)。未包含该修复的实现不得用于有效训练。
 
-> **跨机器生产启动命令**：先激活目标 Conda 环境，再用 `$CONDA_PREFIX/lib` 相对该环境选择其中的 cuDNN。环境变量只作用于本次训练进程及其子进程，不修改环境内已安装的包。
+> **跨机器生产启动命令**（2026-08-30 修订）：**不要再导出 `LD_LIBRARY_PATH="$CONDA_PREFIX/lib"`。** conda cuDNN 已移除，环境统一到 pip cu128 栈；继续导出该目录会让其中残留的 conda-forge `libcublas`/`libcudart`/`libnccl` 遮蔽 pip 的同名库，重演同一类版本错配。
 
 ```bash
-LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
-  XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
   python scripts/train.py my_task \
     --exp-name=run_0520 \
     --overwrite
 ```
 
+<details><summary>原命令（已过时，仅供追溯）</summary>
+
+```bash
+LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
+  XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+  python scripts/train.py my_task --exp-name=run_0520 --overwrite
+```
+
+</details>
+
 > **2026-08-29 长跑事故（已修复）**：原实现从 step 16100 起记录到 NaN；受污染 checkpoint 已删除，并从保留的 step 15000 恢复。最终修复保留原始 attention mask，只对全空 query 行停止 Q 梯度。逐 step 有限性保护已按生产要求移除，训练继续按原逻辑记录 loss；详见事故文档。
 
 `my_task` 的模型配置需要包含 `use_cudnn_attention: true`；如果 YAML 中没有设置，则在命令末尾增加 `--model.use-cudnn-attention`。`--overwrite` 会删除同名实验的原 Checkpoint 目录；恢复已有训练应改用 `--resume`。
@@ -23,7 +37,16 @@ LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
 JAX cuDNN runtime version: 91400
 ```
 
-该值来自 `cuda_versions.cudnn_get_version()`。如果不是 `91400`，应停止训练并检查目标机器的 `$CONDA_PREFIX/lib`，不能只根据 PyTorch 输出判断 JAX 运行时版本。
+该值来自 `cuda_versions.cudnn_get_version()`。
+
+> [!IMPORTANT]
+> 2026-08-30 修订：**这个数字不足以证明 cuDNN 栈是干净的。** 它由 cuDNN 的引擎子库返回，因此一个 9.10.2 的 dispatcher stub 驱动 9.14 引擎库时，它照样打印 `91400`——事故期间的实际情况正是如此。请改用：
+>
+> ```bash
+> python scripts/check_cuda_stack.py
+> ```
+>
+> 该脚本遍历 `/proc/self/maps`，在 `libcudnn*` 来自多个目录或文件名版本与运行时报告不一致时判 FAIL，并在真实训练形状上跑一次前向+反向与 fp32 参考解对比。详见 [2026-08-30-cudnn-attention-divergence.md](2026-08-30-cudnn-attention-divergence.md)。
 
 **日期**：2026-08-27 至 2026-08-28
 
diff --git a/docs/2026-08-29-cudnn-attention-nan-incident.md b/docs/2026-08-29-cudnn-attention-nan-incident.md
index 6f13b33..c958bd5 100644
--- a/docs/2026-08-29-cudnn-attention-nan-incident.md
+++ b/docs/2026-08-29-cudnn-attention-nan-incident.md
@@ -8,6 +8,9 @@
 
 **影响范围**：从 clean step 16000 切换到 `use_cudnn_attention=true` 后的生产训练
 
+> [!CAUTION]
+> 2026-08-30 更新：本文的修复**有效但不完整**。它消除了这里描述的瞬时失效（第一步反向即产生 NaN），但带该修复的训练在另一台机器上仍会缓慢发散：前 1000 步与基线逐点吻合，随后在 step 1000–1200 之间分叉，约 4700 步后变为 NaN。本文「验证要求」一节的 200–500 步标准不足以发现它。见 [2026-08-30-cudnn-attention-divergence.md](2026-08-30-cudnn-attention-divergence.md)。
+
 ## 摘要
 
 生产训练恢复后，第一个记录窗口已经出现 NaN 梯度，随后 loss 和参数持续为 NaN。根因不是 cuDNN 动态库版本、数据加载、显存或普通的梯度爆炸，而是 cuDNN fused Attention 在 attention mask 存在整行全 `False` 时的反向数值行为：前向输出有限，但 Q 梯度为 NaN。
diff --git a/docs/2026-08-30-cudnn-attention-divergence.md b/docs/2026-08-30-cudnn-attention-divergence.md
new file mode 100644
index 0000000..cb20e98
--- /dev/null
+++ b/docs/2026-08-30-cudnn-attention-divergence.md
@@ -0,0 +1,375 @@
+# 启用 flash attention 提速：全层训练发散与分层稳定方案
+
+**日期**：2026-08-30
+
+**目标**：用 cuDNN fused attention（`use_cudnn_attention=true`）替代显式 attention，缩短训练时间。
+
+**机器**：8×H100 80GB（`cs-76d36-a14d7-server`）
+
+**状态**：**根因机制已定位，并得到一个通过 3000 步验收的保守修复。** 全 18 层 cuDNN 的约 1.5% 原始梯度差异具有不同于人造乘性噪声的坐标结构，经过 Adam 预条件后变成约 3.35% 的参数更新差异，并使约 0.79% 的更新坐标翻转符号；误差在早层产生时还会穿过后续网络放大。仅让最后 3/18 层使用 cuDNN，把更新差异降到 0.847%、符号翻转降到 0.193%，并从 step 0 稳定跑满 3000 步。该方案解决了发散，但实测步时约 `1.5 s/step`，只保留了小部分吞吐收益，尚未达到全层 cuDNN 的速度目标。
+
+## 为什么值得继续追
+
+同机同任务的实测步时：
+
+| 实现 | 步时 | 来源 |
+|---|---|---|
+| 显式 attention | `1.645 s/step` | 0828 run，98714 s / 60000 步，含日志与 checkpoint 摊销 |
+| cuDNN + stop-gradient mask | 约 `1.05 s/step` | 0829 run 的 `[TIMING]` 普通步 |
+| cuDNN + dummy-key mask | 约 `1.38 s/step` | 2026-08-30 诊断 run |
+
+按 60000 步计算，`1.645 -> 1.05` 约省 **9.9 小时**。
+
+> [!IMPORTANT]
+> `1.05 s/step` 取自一个**已经在发散**的 run。数值失效不改变 kernel 的执行成本，所以该数字作为吞吐参考是可用的；但正式验收时必须在一个不发散的 cuDNN run 上重测。
+
+## 失效现象
+
+`use_cudnn_attention=true` 时，训练在前约 1000 步与已验证基线逐点吻合，随后偏离并单调发散，最终变为 NaN。
+
+```text
+step    0828 显式（基线）      0829 cuDNN 9.14        cuDNN 9.19
+ 200    0.9244 /  2.89        0.9333 /  3.58         0.9348 /  3.53
+1000    0.4220 /  2.77        0.4306 /  2.50         0.4319 /  3.16
+1100    0.4077 /  1.79        0.4784 /  3.16         0.4723 /  3.37
+1200    0.3742 /  2.12        0.6484 / 31.40         0.6600 / 35.49
+2500    0.2457 /  1.33        5.7692 / 675.08        4.7055 / 759.18
+5900    0.1763 /  1.36        nan                    （已停止）
+```
+
+**因果起点是 step 1000–1200，不是 NaN 出现的 step 5900。** 后者只是累积发散溢出的位置。这个区别把复现成本从 27 小时降到约 35 分钟，是本次调查最有操作价值的一条。
+
+### 与 2026-08-29 事故是两种不同的失效
+
+| | 2026-08-29（H200） | 2026-08-30（H100） |
+|---|---|---|
+| 现象 | **瞬时**：step 16300 正常 → 16400 全 NaN | **缓慢**：贴合 1000 步后发散 4700 步 |
+| 根因 | 全空 mask 行的 Q 梯度 NaN | 数值扰动 + 配方裕度不足 |
+| 修复 | `_stop_gradient_for_fully_masked_queries` | 未解决 |
+
+全空 mask 修复是有效的，它消除了第一种模式。当前问题带着该修复仍然发生。
+
+## 责任归属：单变量 A/B
+
+同机、同数据集、同超参、同 seed，唯一变量是 `use_cudnn_attention`：
+
+| | `cve5hiyt`（0828） | `orbnh1b7`（0829） |
+|---|---|---|
+| commit | `d0aec27` | `51d9af3` |
+| `use_cudnn_attention` | 字段不存在（显式路径） | `true` |
+| 结果 | step 59900 **loss 0.0589** | step 5900 **NaN** |
+
+`git diff d0aec27 HEAD -- src/ scripts/` 共 72 行增删，全部属于该开关本身。排除 LR、数据、配置。
+
+## 结论：cuDNN 内核是正确的
+
+### 完整模型梯度 A/B
+
+同一份 pi05_base 权重、同一个真实 batch、同一个 rng，唯一变量是开关（`scripts/grad_ab_probe.py`，cuDNN 9.19，batch 32，FSDP 8 卡）：
+
+```text
+batch 0   loss 显式=5.859448  cuDNN=5.851340   |grad| 104.31 vs 103.65   ratio=0.9937
+batch 1                                        |grad|  67.19 vs  67.45   ratio=1.0038
+batch 2                                        |grad|  70.71 vs  70.61   ratio=0.9986
+
+相对梯度差：1.75e-02 / 1.66e-02 / 2.12e-02
+按模块：PaliGemma 1.87e-02  action_in_proj 1.66e-02
+        time_mlp_in 1.46e-02  time_mlp_out 1.20e-02  action_out_proj 5.87e-03
+```
+
+**相对差约 2%，梯度范数比值 0.994–1.004，没有任何模块偏离整体水平。** 这是 bf16 误差经 18 层 Gemma 反向累积的正常量级——单层测量中 cuDNN 对 fp32 参考解的相对误差是 `5.4e-03`、显式 bf16 路径是 `3.2e-03`，两者同阶。cuDNN 给出的不是更错的答案，只是另一个同样合法的 bf16 近似。
+
+### 扰动的噪声/确定性分解
+
+在同一 batch 上重复调用 cuDNN 反向，把差异拆开：
+
+```text
+batch    总差异(cuDNN vs 显式)   随机(cuDNN vs cuDNN)   确定性分量
+  0          1.6993e-02             9.7433e-03          1.3922e-02
+  1          1.7036e-02             8.7706e-03          1.4605e-02
+  2          2.2268e-02             9.1202e-03          2.0315e-02
+```
+
+随机部分约 `0.9%`，来自 cuDNN 反向的 atomic 累加（同一输入两次反向 bitwise 不同）。确定性部分约 `1.5%`，是主要成分。
+
+注意「确定性」在这里的含义：**给定 batch 可复现**，不等于**跨 batch 方向一致**。换一个 batch，舍入差异的方向随之改变；梯度范数比值为 1.00，也看不出系统性缩放。后续多 batch 实验已直接否定固定方向假设，见[跨 batch 固定方向假设不成立](#跨-batch-固定方向假设不成立)。
+
+## 已排除的候选
+
+全部在本机、真实训练形状下实测：`T=1018`（768 图像 + 200 prompt + 50 action）、`num_heads=8`、`num_kv_heads=1`、`head_dim=256`、bf16、含全空行的真实 block mask。参考解为同一显式路径在 fp32 下的结果。
+
+- **cuDNN 版本**：9.14 与 9.19 的发散曲线逐点重合（step 1200 分别为 31.40 与 35.49）。
+- **「静态序列长度非 128 倍数导致 SDPA backward 错误」**（cuDNN 9.15.1 release note）：
+  ```text
+  T=1018 (%128=122)   dQ rel-L2: 显式 3.23e-03   cuDNN 3.21e-03   NaN=0
+  T=1024 (%128=0)     dQ rel-L2: 显式 3.23e-03   cuDNN 3.21e-03   NaN=0
+  ```
+  两者无差异，均无 NaN。**该假设在本机本形状上不成立。**
+- **库打包错配**：环境统一为单源 pip 栈后（见下），照样发散。
+- **logit 幅度**：std 从 1 扫到 80，cuDNN 最差 `7.2e-03`、显式 `3.5e-03`，同量级。
+- **head_dim=256**（Hopper 上 cuDNN ≥9.5 才放开的新路径）：与 head_dim=128 表现一致。
+- **cotangent 稠密度**：只在 action token 上非零 与 在全部有效 token 上稠密，结果相同。
+- **dBias 反向路径**：`should_export_dbias` 要求 `bias.shape[0]==1`，实际 batch 为 256，不进入该路径。
+- **分片**：`activation_sharding_constraint` 使用 `PartitionSpec(DATA_AXIS)`，只按 batch 分片，attention 拿到完整序列。
+- **两版 mask 修复**：c1b70b2 给全空行开 dummy key，51d9af3 保留原 mask 只对 q 做 stop_gradient。在长/短两种 padding 密度下，两者在有效行上给出**逐位相同**的 dQ/dK/dV。且事故记录显示 c1b70b2 版本同样发散：从 step 15000 的 loss `0.145` 涨到 step 18600 的 `0.7085`，逐 step 保护在那里报告最近 100 步有 8 次更新非有限。**两版修复都不是原因，也都不是解药。**
+- **降低 `peak_lr`**：`2.5e-5 -> 1.5e-5` 只把发散起点从 step 约 1150 推迟到约 1450：
+  ```text
+  step   cuDNN @1.5e-5
+  1200   0.3954 /  2.95
+  1400   0.4089 /  2.84
+  1500   0.5155 /  7.17
+  1700   0.9654 / 26.85
+  1900   1.1773 / 88.68
+  ```
+  **LR 是速率参数，不是开关。**
+- **提高 `optimizer.b2`**：`0.95 -> 0.99` 未延后发散起点。2026-08-30 在统一后的 cuDNN 9.19 单源环境中，从 step 0 跑满 3000 个更新；8×H100、batch 256，W&B、周期保存和最终保存均关闭。启动前 `scripts/check_cuda_stack.py` 判定 PASS。所有已输出的指标均为有限值，但曲线仍在 step 1100–1200 分叉：
+  ```text
+  step    loss       grad_norm    param_norm
+   200    0.9082        3.1334     1802.3866
+  1000    0.4321        3.7520     1802.5017
+  1100    0.5189        4.3141     1802.5402
+  1200    0.9354       44.6635     1802.5638
+  1300    1.6713       52.6050     1802.5797
+  1500    2.2757       30.3204     1802.6188
+  2000    3.7292      193.8313     1802.7924
+  2500    4.7271      422.8436     1803.0002
+  2700    6.4255     1521.7623     1803.0616
+  2900    4.5535      819.7691     1803.1382
+  ```
+  训练完整执行了 3000 个更新，最后一个正式指标窗口是 step 2900。step 1200 的显式基线为 `loss=0.3742 / grad_norm=2.12`，因此新 run 已超过 3× 梯度验收上限约 7 倍。step 2500 的 loss `4.7271` 又几乎复现了 `b2=0.95` cuDNN run 的 `4.7055`；较高 `b2` 只降低了该点的梯度峰值（`422.84` vs `759.18`），没有改变发散轨迹。普通非日志 step 的 dispatch 约 `1.29–1.39 s`，无数据断供或系统错误。实验结束后 8 张 GPU 均归零，实验目录为空，没有污染 checkpoint。**提高二阶矩窗口不是解法。**
+- **显式 attention + 1.5% 逐 step 独立随机梯度扰动**：没有发散。2026-08-31 在训练配置中加入默认关闭的 `gradient_noise_scale` 诊断开关；显式梯度为 `g`，先生成 `r = g * N(0, I)`，再令 `delta = 0.015 * ||g|| / ||r|| * r`，优化器实际接收 `g + delta`。模型 RNG 保持原样，扰动 RNG 独立并按 step 可复现，因此每一步都精确满足 `||delta|| / ||g|| = 0.015`，同时扰动能量按原梯度能量分布，而不是按参数量集中到最大张量。
+
+  8×H100、batch 256、显式 attention、原 `b2=0.95`，从 step 0 跑满 3000 个更新；W&B、周期保存和最终保存均关闭。正式运行前，CPU helper/配置回归 `58 passed`，生产 batch 的 2-step GPU 校准也连续报告 `gradient_noise_rel_norm=0.0150`。正式曲线：
+  ```text
+  step    loss       applied grad_norm    clean grad_norm    noise rel-L2
+   100    2.8529          3.9906              3.9903           0.0150
+   200    0.9262          4.0429              4.0433           0.0150
+   500    0.5704          2.0811              2.0806           0.0150
+   900    0.4335          2.8625              2.8631           0.0150
+  2200    0.2698          1.7088              1.7083           0.0150
+  2500    0.2476          1.8600              1.8613           0.0150
+  2700    0.2387          1.5861              1.5854           0.0150
+  2900    0.2411          1.6389              1.6389           0.0150
+  ```
+  训练循环耗时 `1:22:45`，另有首次数据冷启动约 6.5 分钟；普通 step 的 dispatch 约 `1.57–1.64 s`。最后一个正式指标窗口是 step 2900。step 2500 的 loss `0.2476` 与历史显式基线 `0.2457` 基本重合，而同期 cuDNN 9.19 已经是 `4.7055 / 759.18`。实验结束后 8 张 GPU 均归零，目录为空，没有污染 checkpoint。
+
+  **结论**：配方并非对任意 1.5% 梯度扰动都没有裕度。该结果排除的是「逐 step 独立、零均值、按梯度能量加权」这一类随机扰动；它尚不能排除真实 cuDNN 差分具有跨 batch 的固定投影、时间相关性或不同的参数内结构。
+
+### 两个 mask 变体：等价，但一个明显更快
+
+`_stop_gradient_for_fully_masked_queries`（51d9af3）与 `_make_cudnn_attention_mask_safe`（c1b70b2）数值等价，但 dummy-key 版本要在**每一层**内重建整个 `(B, 1, T, S)` 掩码，batch 256 下实测 `1.38 s/step`，stop-gradient 版本约 `1.05 s/step`。H200 文档记录的 `1.91 -> 2.31 s/step` 回退是同一件事。
+
+因此代码保留 stop-gradient 版本。51d9af3 的标题「Preserve the **fast** cuDNN mask」描述的正是这个性能取舍，它不是缺陷来源。
+
+## 2026-08-31 后续调查：根因机制与稳定折中
+
+### 跨 batch 固定方向假设不成立
+
+`scripts/grad_ab_direction_probe.py` 对 8 个真实 batch 分别计算显式梯度和 3 次 cuDNN 反向的均值，再比较 `delta_b = g_cudnn,b - g_explicit,b`：
+
+```text
+跨 batch delta 两两 cosine：-0.004043 ± 0.126301，范围 [-0.247516, +0.226915]，n=28
+||delta|| / ||g_explicit||：  1.556024e-02 ± 2.76e-03
+cuDNN 同 batch 随机分量：    8.752676e-03
+cos(delta, gradient)：       -0.00113 ± 0.225
+cos(delta, parameter)：      +0.000016
+```
+
+差分没有稳定的跨 batch 方向，也不沿当前梯度或参数方向累积。因此没有继续做「显式 attention + 固定方向人造噪声」；该实验不会复现真实差分结构。
+
+### batch 到达顺序不是原因
+
+`strict_batch_order=true` 的全 18 层 cuDNN run 仍在原位置分叉：
+
+```text
+step 1000   loss=0.4304   grad_norm=2.3201
+step 1100   loss=0.4593   grad_norm=2.1643
+step 1200   loss=0.6034   grad_norm=15.4116
+```
+
+显式基线 step 1200 为约 `0.3742 / 2.12`。固定 sampler 输出顺序没有改变失效，因此 DataLoader worker 完成时序不是根因。后续诊断均保留 `strict_batch_order=true`，确保曲线可逐点比较。
+
+### 关键机制：Adam 放大了坐标级差异
+
+`scripts/optimizer_update_ab_probe.py` 从上述严格顺序 run 的 step 1000 checkpoint 恢复真实 Adam 状态，对同一真实 batch、同一 rng 计算实际优化器更新。原始梯度 L2 差异仍只有约 1.5%，但进入 Adam 后：
+
+```text
+                         相对更新差异     更新 cosine     符号翻转坐标
+真实全层 cuDNN 差分          3.3521%         0.999440        0.7885%
+旧的 1.5% 人造乘性噪声       0.7313%         --              0.1459%
+```
+
+真实差分造成的 Adam 更新差异是稳定人造噪声的 **4.6 倍**，符号翻转比例是 **5.4 倍**。主导部分在 PaliGemma（更新差异 3.353%、符号翻转 0.789%）。这解释了为何「相同梯度相对 L2」的人造噪声稳定，而 cuDNN 会发散：旧噪声按 `|g|` 分配能量，几乎不扰动小梯度/小二阶矩坐标；真实舍入差分具有不同的参数内结构，会被 Adam 的逐坐标归一化放大。
+
+把 Adam `eps` 从 `1e-8` 提高到 `1e-6`/`1e-5`，真实更新相对差异只从 3.35% 降到 3.22%/3.08%，符号翻转仍约 0.79%，不是有效修复。
+
+### 优化器时间平滑只能延迟，不能修复
+
+将 `b1` 从 `0.9` 提高到 `0.99`（`b2` 保持 `0.95`）后，全层 cuDNN 在 step 1200 暂时健康，但随后同样发散：
+
+```text
+step 1200   loss=0.4027   grad_norm=2.1161
+step 1400   loss=0.4950   grad_norm=6.2470
+step 1500   loss=0.7486   grad_norm=17.0
+step 2000   loss=2.749    grad_norm=67
+step 2900   loss=8.844    grad_norm=562
+```
+
+它只把分叉推迟约 300 步，和降低 LR、提高 `b2` 的结果一致：改变累积速率，不改变误差来源。
+
+### 层位置比层数更重要
+
+为 `Pi0Config` 增加 `cudnn_attention_layer_start` 和 `cudnn_attention_num_layers`，可只在连续的 Gemma 层段使用 cuDNN。其余层仍走显式 attention；默认值保持历史上的全层行为。CPU 模型回归 `8 passed`。
+
+在 step 1000 的真实 Adam 状态上做更新空间 A/B：
+
+```text
+cuDNN 层段       相对更新差异     更新 cosine     符号翻转坐标
+全 18 层            3.3118%         0.999451        0.7769%
+最前 6 层           3.2061%         0.999486        0.7548%
+中间 6 层           2.2169%         0.999754        0.5141%
+最后 6 层           1.1633%         0.999932        0.2720%
+最前 3 层           3.3016%         0.999455        0.7720%
+中间 3 层           1.9894%         0.999802        0.4634%
+最后 3 层           0.8470%         0.999964        0.1926%
+稳定人造噪声参考     0.7313%         --              0.1459%
+```
+
+只融合最前 3 层几乎等同全 18 层，说明早层的微小扰动会经后续 15 层反向/前向传播放大；不能按融合层数线性估计风险。最后 3 层最接近已经跑稳的人造噪声参考，因此选作首个训练候选。
+
+### 最后 3 层方案跑满 3000 步
+
+配置：`configs/diag_cudnn_last3.yaml`，8×H100、batch 256、`strict_batch_order=true`、原始 Adam（`b1=0.9, b2=0.95`）、从 pi05_base step 0 开始、W&B 和 checkpoint 保存关闭。层索引 15–17 使用 cuDNN，其余 15 层显式。正式曲线：
+
+```text
+step    loss      grad_norm    param_norm
+   0    6.4029      65.5204     1802.3865
+ 100    2.8439       5.1838     1802.3851
+ 500    0.5719       2.4526     1802.4014
+ 900    0.4349       2.5621     1802.4742
+1000    0.4240       2.3789     1802.5061
+1100    0.4050       2.1847     1802.5416
+1200    0.3757       2.3869     1802.5775
+1400    0.3390       1.8067     1802.6504
+1600    0.3238       2.6505     1802.7225
+2000    0.2793       1.8254     1802.8666
+2500    0.2463       2.0116     1803.0460
+2900    0.2424       1.4769     1803.1906
+```
+
+训练完整执行 3000 个更新；最后一个正式指标窗口是 step 2900。原全层 cuDNN 在 step 1200 已到 `0.6034 / 15.41`，b1=0.99 延迟方案在 step 2000 已到 `2.749 / 67`；本方案的曲线与显式稳定基线一致，满足 3000 步验收。
+
+训练循环耗时 `1:18:02`，普通步约 `1.5 s/step`；另有 404 秒的 64-worker 首批冷启动。相对本任务显式路径约 `1.57–1.64 s/step` 只有约 4–9% 加速，明显低于全层 cuDNN 的收益。**因此这是数值稳定的保守折中，不是最终吞吐目标。** 按用户指示，本轮到此停止，不继续运行最后 6 层或更大范围的训练。
+
+## 2026-08-31 调查计划原稿（结果见上）
+
+以下保留当时的计划作为调查轨迹。第 1 项已执行并否定固定方向；因此第 2 项不再能复现真实差分，未执行。实际随后转向 Adam 更新空间和分层位置分析，并得到上面的最后 3 层方案。
+
+### 1. 测量真实 cuDNN 差分的跨 batch 方向一致性
+
+扩展 `scripts/grad_ab_probe.py`，对多个真实 batch 计算 `delta_b = g_cudnn,b - g_explicit,b`，并报告：
+
+- 不同 batch 的 `delta_b` 两两 cosine；
+- `delta_b` 在显式梯度、参数向量和前一 batch 差分上的投影；
+- 上述量按 PaliGemma、action projection 和 time MLP 分模块的结果；
+- 同一 batch 多次 cuDNN 反向取平均后再计算，尽量剥离约 0.9% 的 atomic 随机分量。
+
+若跨 batch cosine 或固定投影显著非零，说明真实差分存在可累积方向；若接近零，则需要检查人造噪声与 cuDNN 差分在参数内分布、尾部和时间相关性上的差别。
+
+### 2. 显式 attention + 时间相关的固定方向扰动
+
+复用 `gradient_noise_scale=0.015`，但让随机方向不再按 step 变化，每一步只根据当前 `||g||` 重新缩放同一参数坐标方向。若这一版本发散而逐 step 独立版本不发散，时间相关性就是关键杠杆；若仍稳定，则必须更直接地复现真实 cuDNN 差分结构。
+
+### 3. `optimizer.weight_decay` `1e-10 -> 1e-4` 量级
+
+当前几乎为零。加大给慢漂移一个回复力。改动比前两项大，收敛行为也会变，排在后面。
+
+### 4. 若以上都不成立：绕开 attention 数值的提速路径
+
+- `max_token_len` `200 -> 168`（全数据扫描得到的 tokenized state 最大长度是 165，留 3 token 余量）：减少约 3% 的序列长度，不触碰数值；
+- 重新采集 XProf trace，确认剩余最大的 SigLIP/Gemma GEMM 与 fusion；
+- 提高 `save_interval`，减少 checkpoint 对下一步的数据等待尖峰。
+
+这三项合计的收益远小于 flash attention 的 36%，但零风险。
+
+### 不要重复的弯路
+
+- **不要再查 cuDNN 版本或库打包。** 9.14 与 9.19 表现一致，单源栈已验证。
+- **不要再改 mask 处理。** 两个变体逐位等价，且都发散。
+- **不要只调 `peak_lr`。** 按累积机制，任何取值大概率都只是推迟。
+- **不要再单独提高 `optimizer.b2`。** `0.99` 从 step 0 跑满 3000 个更新，仍在 step 1100–1200 分叉，并在 step 2700 出现 `grad_norm=1521.76`。
+- **不要重复逐 step 独立的 1.5% 随机梯度扰动。** 显式 attention 已跑满 3000 步，step 2500 loss `0.2476`，没有复现发散；下一步应测试真实差分的方向结构或时间相关扰动。
+- **不要跨 run 比较 `grad_norm`。** 本仓库默认 `strict_batch_order: false`，各 run 的 batch 序列不同，早期 `grad_norm` 差异来自数据顺序而非实现差异。曾据此误判 cuDNN 存在系统性梯度放大，被受控 A/B（范数比 1.00）否定。
+
+## 验收要求
+
+> [!WARNING]
+> 本次曾在 step 1300 仍健康时判定「降 LR 已解决」，随后 step 1450 开始发散。**判据必须是跑满预定步数并看完整条曲线，不能是「越过某一步还正常」。**
+
+启用 `use_cudnn_attention` 的验收：
+
+1. 从 **step 0** 起跑满 **3000 步**（最后 3 层方案实测训练循环约 78 分钟，另有约 6.7 分钟首次数据冷启动）；
+2. 与同配置显式 attention 的基线曲线逐点对比 step 200/500/1000/1500/2000/3000 的 `loss`；
+3. 全程 `grad_norm` 不得超过基线同点的 3 倍；
+4. 通过后才能进入长跑，且长跑期间仍需按日志窗口盯 `loss` 与 `grad_norm`。
+
+40 步吞吐短跑、从已收敛 checkpoint 出发的短跑、以及只验证「梯度有限」的 smoke test，都**不是**有效的正确性验收。2026-08-29 的修复通过了全部这三类检查，仍然发散。
+
+建议增加发散护栏：`grad_norm` 已在日志 step 计算，当它连续多个日志窗口超过滚动中位数的固定倍数时中止训练。本次若有该护栏，会在 step 1300 停下而不是跑满 60000 步。
+
+## 环境
+
+### dispatcher 与 engine 版本错配（已修复）
+
+事故期间 `cuda_versions.cudnn_get_version()` 返回 `91400`，但进程实际加载的是 **9.10.2 的 dispatcher stub 驱动 9.14.0 的引擎库**：
+
+```text
+site-packages/nvidia/cudnn/lib/libcudnn.so.9          <- 125 KB，dispatcher stub，9.10.2 wheel
+$CONDA_PREFIX/lib/libcudnn_graph.so.9.14.0            <- 真正执行的引擎库
+$CONDA_PREFIX/lib/libcudnn_engines_precompiled.so.9.14.0
+$CONDA_PREFIX/lib/libcudnn_ops.so.9.14.0
+```
+
+`cudnnGetVersion()` 的返回值来自引擎库，因此它**无法**发现这种错配。上一份文档要求「必须确认输出 91400」的验证方法是不充分的。正确做法是枚举进程实际映射的库：
+
+```bash
+python scripts/check_cuda_stack.py
+```
+
+它遍历 `/proc/self/maps`，在 `libcudnn*` 来自多个目录、或文件名版本与运行时报告不一致时判 FAIL，然后在真实训练形状上跑一次前向+反向并与 fp32 参考解对比。**在混装环境下该脚本判 FAIL，而三项数值检查全部通过**——这正是只看数值 smoke test 会漏掉的部分。
+
+### 统一到 pip cu128 栈（已完成）
+
+目标：CUDA Toolkit 12.8 + pip cu128 PyTorch + pip cuDNN 9.19，移除 conda cuDNN。
+
+依赖解析上的关键事实：
+
+- `torch 2.10.0+cu128` 把 cuDNN 钉死在 `nvidia-cudnn-cu12==9.10.2.21`（`==`，非 `>=`）；
+- `torch 2.11.0+cu128` 钉 `nvidia-cudnn-cu12==9.19.0.56` 和 `cuda-toolkit==12.8.1`；
+- PyPI 默认的 `torch 2.11.0` 是 **CUDA 13 构建**，钉 `nvidia-cudnn-cu13`，必须走 cu128 索引；
+- `jax-cuda12-plugin` 只要求 `nvidia-cudnn-cu12>=9.1,<10.0`，不构成约束。
+
+**torch 版本决定了训练进程加载哪个 cuDNN**，尽管 torch 在本仓库只用于 DataLoader。这是 `pyproject.toml` 把 torch 下界提到 `>=2.11` 的原因。
+
+另外两处容易漏掉的：
+
+- `torchcodec` 的 `+cu128` 构建只存在于 pytorch 索引，默认 PyPI 上的同版本号没有 `+cu128` 后缀，需在 `[tool.uv.sources]` 中路由；
+- 本地 fork `lerobot-xense`（editable 安装）在自己的 `pyproject.toml` 中声明 `torch<2.11.0`、`torchcodec<0.11.0`、`torchvision<0.26.0`，**必须同步放宽**，否则解析失败。不能改用 PyPI 的 lerobot 0.6——那会丢掉 fork 中的 xense 相机与 flexiv 机器人代码。
+
+启动命令必须**去掉** `LD_LIBRARY_PATH="$CONDA_PREFIX/lib"`。移除 conda cuDNN 后该目录仍保留 conda-forge 的 `libcublas`/`libcudart`/`libnccl` 等，继续导出会让它们遮蔽 pip cu128 的同名库，等于换一个库重演同一类错配。
+
+## 处置
+
+- 0829 run 在 step 5900 之后的所有 checkpoint（10000–50000）均为 NaN，无保留价值；
+- `cve5hiyt`（0828，显式 attention）仍是已跑满 60000 步的生产结果，loss `0.0589`；
+- 最后 3 层 cuDNN 已通过 3000 步短程验收，但尚未跑 60000 步，不能把它等同于完整生产验收；
+- 需要绝对稳妥时继续使用 `use_cudnn_attention: false`；需要先获取小幅吞吐收益时，可用 `use_cudnn_attention: true, cudnn_attention_layer_start: 15, cudnn_attention_num_layers: 3` 进入受监控长跑；
+- 不要把 `use_cudnn_attention: true` 且 `cudnn_attention_num_layers: null`（全 18 层）用于生产训练。
+
+## 相关脚本
+
+- `scripts/check_cuda_stack.py` —— 验证进程实际加载的 cuDNN 栈是否单源，并在真实训练形状上做前向/反向数值检查；
+- `scripts/grad_ab_probe.py` —— 完整模型梯度 A/B，比较显式与 cuDNN 的逐参数梯度，并把差异拆成随机与确定性分量；
+- `scripts/grad_ab_direction_probe.py` —— 多真实 batch 的 cuDNN 差分方向、投影和随机分量；
+- `scripts/optimizer_update_ab_probe.py` —— 从真实 checkpoint 恢复 Adam 状态，比较梯度差分进入优化器后的更新误差与符号翻转；
+- `scripts/hybrid_update_ab_probe.py` —— 比较不同 cuDNN 层段的 Adam 更新空间误差。
diff --git a/pyproject.toml b/pyproject.toml
index 77091ee..fe708f3 100755
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,9 +31,14 @@ dependencies = [
     "safetensors>=0.4.3",
     # >=2.6 for `DataLoader(in_order=...)`, which the data loader uses to hand batches
     # over as workers finish them (see TrainConfig.strict_batch_order).
-    "torch>=2.6,<2.11.0",
-    "torchcodec>=0.2.1,<0.11.0; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')",
-    "torchvision>=0.21.0,<0.26.0",
+    # >=2.11 is what pins cuDNN: torch 2.11.0+cu128 requires
+    # nvidia-cudnn-cu12==9.19.0.56, while 2.10 pinned ==9.10.2.21. JAX only asks for
+    # >=9.1,<10, so torch is what decides which cuDNN the training process loads even
+    # though torch itself only runs the DataLoader here. Keep torch and the cuDNN you
+    # want to train against in sync; see docs/2026-08-30-cudnn-attention-divergence.md.
+    "torch>=2.11,<2.12",
+    "torchcodec>=0.11,<0.12; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')",
+    "torchvision>=0.26,<0.27",
     "tqdm-loggable>=0.2",
     "typing-extensions>=4.12.2",
     "tyro>=0.9.5",
@@ -84,6 +89,7 @@ explicit = true
 [tool.uv.sources]
 torch = { index = "pytorch-cu128" }
 torchvision = { index = "pytorch-cu128" }
+torchcodec = { index = "pytorch-cu128" }
 xense-client = { workspace = true }
 dlimp = { git = "https://github.com/kvablack/dlimp", rev = "ad72ce3a9b414db2185bc0b38461d4101a65477a" }
 
diff --git a/scripts/check_cuda_stack.py b/scripts/check_cuda_stack.py
new file mode 100644
index 0000000..c6d5e4e
--- /dev/null
+++ b/scripts/check_cuda_stack.py
@@ -0,0 +1,199 @@
+"""Verify the CUDA/cuDNN stack a JAX training process will actually load.
+
+Checking cuda_versions.cudnn_get_version() alone is NOT sufficient: the number it
+returns comes from the cuDNN sub-libraries, so a 9.10.2 dispatcher stub driving
+9.14 engine libraries still reports 91400. This walks /proc/self/maps instead and
+then exercises the real production attention shape.
+
+Run it in the SAME shell that launches training:
+    python scripts/check_cuda_stack.py
+"""
+
+import collections
+import importlib.metadata
+import os
+import re
+import sys
+
+FAIL = []
+EXPECTED = {
+    "torch": "2.11.0+cu128",
+    "torchvision": "0.26.0+cu128",
+    "torchcodec": "0.11.1+cu128",
+    "nvidia-cudnn-cu12": "9.19.0.56",
+    "cuda": "12.8",
+    "cudnn": 91900,
+}
+
+
+def loaded_libs(pattern):
+    out = collections.OrderedDict()
+    with open("/proc/self/maps") as f:
+        for line in f:
+            m = re.search(r"(/\S*" + pattern + r"\S*)", line)
+            if m:
+                out[m.group(1)] = None
+    return list(out)
+
+
+def section(title):
+    print(f"\n{'=' * 70}\n{title}\n{'=' * 70}")
+
+
+def report(name, pattern):
+    libs = loaded_libs(pattern)
+    if not libs:
+        print(f"  {name}: (none loaded)")
+        return
+    dirs = {os.path.dirname(p) for p in libs}
+    for p in libs:
+        print(f"    {p}")
+    if len(dirs) > 1:
+        FAIL.append(f"{name} is loaded from {len(dirs)} different directories: {sorted(dirs)}")
+        print(f"  !! {name}: MIXED — {len(dirs)} source directories")
+    else:
+        print(f"  OK {name}: single source {dirs.pop()}")
+
+
+def main():
+    import jax
+    import jax.numpy as jnp
+    import numpy as np
+    import torch
+    import torchcodec
+    import torchvision
+    from jax._src.lib import cuda_versions
+
+    section("1. reported versions")
+    rt, build = cuda_versions.cudnn_get_version(), cuda_versions.cudnn_build_version()
+    cudnn_package_version = importlib.metadata.version("nvidia-cudnn-cu12")
+    torch_rt = torch.backends.cudnn.version()
+    print(f"  jax                    : {jax.__version__}")
+    print(f"  torch                  : {torch.__version__}")
+    print(f"  torchvision            : {torchvision.__version__}")
+    print(f"  torchcodec             : {torchcodec.__version__}")
+    print(f"  nvidia-cudnn-cu12      : {cudnn_package_version}")
+    print(f"  torch CUDA             : {torch.version.cuda}")
+    print(f"  torch cuDNN runtime    : {torch_rt}")
+    print(f"  cuDNN runtime (reported): {rt}")
+    print(f"  cuDNN build            : {build}")
+    print(f"  LD_LIBRARY_PATH        : {os.environ.get('LD_LIBRARY_PATH', '(unset)')}")
+
+    actual = {
+        "torch": torch.__version__,
+        "torchvision": torchvision.__version__,
+        "torchcodec": torchcodec.__version__,
+        "nvidia-cudnn-cu12": cudnn_package_version,
+        "cuda": torch.version.cuda,
+        "cudnn": rt,
+    }
+    for name, expected in EXPECTED.items():
+        if actual[name] != expected:
+            FAIL.append(f"{name} must be {expected}, got {actual[name]}")
+    if torch_rt != rt:
+        FAIL.append(f"PyTorch reports cuDNN {torch_rt}, while JAX reports {rt}")
+
+    # Force the CUDA libraries to actually load before inspecting the memory map.
+    dev = jax.devices()[0]
+    jnp.ones((8, 8)).block_until_ready()
+    print(f"  device                 : {dev} (compute {dev.compute_capability})")
+
+    section("2. libraries actually mapped into this process")
+    report("libcudnn", "libcudnn")
+    report("libcublas", "libcublas")
+    report("libnccl", "libnccl")
+
+    # A version embedded in a filename that disagrees with the reported runtime is
+    # the exact failure that produced a 9.10.2 dispatcher over 9.14 engines.
+    vers = set()
+    for p in loaded_libs("libcudnn"):
+        m = re.search(r"\.so\.(\d+)\.(\d+)\.(\d+)", p)
+        if m:
+            vers.add(f"{m.group(1)}.{m.group(2)}.{m.group(3)}")
+    if len(vers) > 1:
+        FAIL.append(f"cuDNN filenames span multiple versions: {sorted(vers)}")
+        print(f"  !! cuDNN file versions: MIXED {sorted(vers)}")
+    elif vers:
+        v = vers.pop()
+        maj, minor, patch = (int(x) for x in v.split("."))
+        want = maj * 10000 + minor * 100 + patch
+        print(f"  cuDNN file version     : {v}")
+        if want != rt:
+            FAIL.append(f"cuDNN filename says {v} ({want}) but runtime reports {rt}")
+            print(f"  !! filename {want} != reported runtime {rt}")
+
+    section("3. production attention shape: forward + backward vs fp32 reference")
+    # pi0.5 training shape: 3x256 image + 200 prompt + 50 action tokens, GQA 8:1,
+    # head_dim 256, bf16, real block mask including fully-masked padding rows.
+    B, NH, NKV, HD = 2, 8, 1, 256
+    T = 768 + 200 + 50
+
+    im = np.ones((B, T), dtype=bool)
+    for b, plen in enumerate((120, 165)):
+        im[b, 768 + plen : 768 + 200] = False
+    ar = np.zeros((T,), dtype=bool)
+    ar[768 + 200] = True
+    im_j = jnp.asarray(im)
+    cs = jnp.broadcast_to(jnp.cumsum(jnp.asarray(ar), -1), im_j.shape)
+    mask = jnp.logical_and(cs[:, None, :] <= cs[:, :, None], im_j[:, None, :] * im_j[:, :, None])[:, None]
+    valid = np.asarray(jnp.any(mask, -1))[0, 0]
+    n_empty = int((~np.asarray(jnp.any(mask, -1))).sum())
+    print(f"  shape B={B} T={T} heads={NH} kv={NKV} head_dim={HD} bf16")
+    print(f"  fully-masked query rows in batch: {n_empty}")
+
+    rng = np.random.default_rng(0)
+    q = jnp.asarray(rng.normal(0, 10 / np.sqrt(HD), (B, T, NH, HD)), jnp.bfloat16)
+    k = jnp.asarray(rng.normal(0, 1, (B, T, NKV, HD)), jnp.bfloat16)
+    v = jnp.asarray(rng.normal(0, 1, (B, T, NKV, HD)), jnp.bfloat16)
+    cot = np.asarray(rng.normal(0, 1, (B, T, NH, HD)), np.float32)
+    cot[:, : T - 50] = 0.0  # only action tokens carry loss, as in Pi0
+    cot = jnp.asarray(cot)
+
+    def reference(q, k, v, m):  # explicit path, fp32 throughout
+        q, k, v = (x.astype(jnp.float32) for x in (q, k, v))
+        gq = q.reshape(B, T, NKV, NH // NKV, HD)
+        lg = jnp.einsum("BTKGH,BSKH->BKGTS", gq, k, preferred_element_type=jnp.float32)
+        lg = jnp.where(m[:, :, None], lg, -2.3819763e38)
+        return jnp.einsum("BKGTS,BSKH->BTKGH", jax.nn.softmax(lg, -1), v).reshape(B, T, NH, HD)
+
+    def cudnn(q, k, v, m):  # what training runs with use_cudnn_attention=true
+        has_key = jnp.any(m, -1)[:, 0, :, None, None]
+        q = jnp.where(has_key, q, jax.lax.stop_gradient(q))
+        return jax.nn.dot_product_attention(q, k, v, mask=m, scale=1.0, implementation="cudnn")
+
+    def loss(fn):
+        return lambda q, k, v, m: jnp.sum(fn(q, k, v, m).astype(jnp.float32) * cot)
+
+    try:
+        gr = [np.asarray(x.astype(jnp.float32), np.float64) for x in jax.grad(loss(reference), (0, 1, 2))(q, k, v, mask)]
+        gc = [np.asarray(x.astype(jnp.float32), np.float64) for x in jax.grad(loss(cudnn), (0, 1, 2))(q, k, v, mask)]
+    except Exception as exc:  # noqa: BLE001
+        FAIL.append(f"cuDNN attention raised {type(exc).__name__}: {exc}")
+        print(f"  !! cuDNN attention FAILED: {type(exc).__name__}: {str(exc)[:300]}")
+        gr = gc = None
+
+    if gc is not None:
+        for i, nm in enumerate(("dQ", "dK", "dV")):
+            a, b = gc[i][:, valid], gr[i][:, valid]
+            rel = np.linalg.norm(a - b) / max(np.linalg.norm(b), 1e-30)
+            nan = int(np.isnan(gc[i]).sum())
+            status = "OK" if (nan == 0 and rel < 2e-2) else "!!"
+            if nan or rel >= 2e-2:
+                FAIL.append(f"{nm}: rel-L2={rel:.3e} NaN={nan}")
+            print(f"  {status} {nm}: rel-L2 vs fp32 = {rel:.3e}   NaN count = {nan}   (bf16 noise is ~3e-3)")
+
+    section("verdict")
+    if FAIL:
+        print("  FAIL — do not start production training:")
+        for f in FAIL:
+            print(f"    - {f}")
+        return 1
+    print("  PASS — single-source cuDNN stack, production-shape backward finite and accurate.")
+    print("  NOTE: this gates the kernel only. It does NOT prove training converges;")
+    print("        the 2026-08-30 divergence passed every check of this kind and still")
+    print("        diverged between step 1000 and 1200. Watch the loss curve.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/scripts/grad_ab_direction_probe.py b/scripts/grad_ab_direction_probe.py
new file mode 100644
index 0000000..e6745cb
--- /dev/null
+++ b/scripts/grad_ab_direction_probe.py
@@ -0,0 +1,257 @@
+"""Measure whether cuDNN-vs-explicit gradient deltas align across real batches.
+
+For each batch, one explicit-attention gradient is compared with the average of
+three cuDNN-attention gradients. Averaging suppresses the known nondeterministic
+atomic-accumulation component before measuring cross-batch directionality.
+"""
+
+import collections
+import dataclasses
+import functools
+import logging
+import sys
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+from jax._src.lib import cuda_versions
+import numpy as np
+import optax
+
+import openpi.training.config as _config
+import openpi.training.data_loader as _data_loader
+import openpi.training.sharding as sharding
+
+sys.path.insert(0, "scripts")
+import train as train_mod
+
+
+CONFIG = "pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100"
+BATCH = 32
+WORKERS = 16
+N_BATCHES = 8
+CUDNN_REPEATS = 3
+
+
+def _grad_fn(model_def, trainable_filter, params, rng, obs, act):
+    model = nnx.merge(model_def, params)
+    model.train()
+
+    def loss_fn(m, r, o, a):
+        return jnp.mean(m.compute_loss(r, o, a, train=True))
+
+    return nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, trainable_filter))(model, rng, obs, act)
+
+
+def _flat(state):
+    pure = state.to_pure_dict()
+    out = []
+
+    def rec(d, prefix):
+        if isinstance(d, dict):
+            for key, value in d.items():
+                rec(value, f"{prefix}/{key}" if prefix else str(key))
+        else:
+            out.append((prefix, d))
+
+    rec(pure, "")
+    return out
+
+
+def _to_host(state_tree):
+    return {key: np.asarray(jax.device_get(value), np.float32) for key, value in _flat(state_tree)}
+
+
+def _dot(a, b):
+    # Accumulate each tensor in float64 without materializing a float64 copy of
+    # the whole 3B-parameter model at once.
+    return float(np.vdot(a.astype(np.float64), b.astype(np.float64)))
+
+
+def _module(path):
+    return path.split("/", 1)[0]
+
+
+def _cos(dot, norm_a_sq, norm_b_sq):
+    return dot / max(np.sqrt(norm_a_sq * norm_b_sq), 1e-300)
+
+
+def main():
+    logging.basicConfig(level=logging.WARNING)
+    print("cuDNN runtime:", cuda_versions.cudnn_get_version(), flush=True)
+    print(f"probe: batches={N_BATCHES}, cudnn_repeats={CUDNN_REPEATS}, local_batch={BATCH}", flush=True)
+
+    cfg = _config.get_config(CONFIG)
+    cfg = dataclasses.replace(cfg, batch_size=BATCH, num_workers=WORKERS, wandb_enabled=False)
+    cfg_off = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=False))
+    cfg_on = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=True))
+
+    mesh = sharding.make_mesh(cfg.fsdp_devices)
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+    rng = jax.random.key(cfg.seed)
+    train_rng, init_rng = jax.random.split(rng)
+
+    print("loading pi05_base weights ...", flush=True)
+    state, _ = train_mod.init_train_state(cfg_off, init_rng, mesh, resume=False)
+    host_params = _to_host(state.params)
+    param_norm_sq = sum(_dot(value, value) for value in host_params.values())
+    print("  trainable param norm:", np.sqrt(param_norm_sq), flush=True)
+
+    gdef_off = state.model_def
+    gdef_on = nnx.graphdef(cfg_on.model.create(init_rng))
+    f_off = jax.jit(functools.partial(_grad_fn, gdef_off, cfg.trainable_filter))
+    f_on = jax.jit(functools.partial(_grad_fn, gdef_on, cfg.trainable_filter))
+
+    print("building data loader ...", flush=True)
+    dl = _data_loader.create_data_loader(cfg_off, sharding=data_sharding, shuffle=True)
+    iterator = iter(dl)
+
+    deltas = []
+    delta_norm_sq = []
+    pairwise_dot = np.zeros((N_BATCHES, N_BATCHES), np.float64)
+    per_batch_summary = []
+    per_module_summary = collections.defaultdict(list)
+
+    for batch_idx in range(N_BATCHES):
+        obs, act = next(iterator)
+        step_rng = jax.random.fold_in(train_rng, batch_idx)
+
+        with sharding.set_mesh(mesh):
+            loss_off, g_off = f_off(state.params, step_rng, obs, act)
+            loss_off = float(loss_off)
+            host_off = _to_host(g_off)
+        del g_off
+
+        host_on_sum = None
+        cudnn_losses = []
+        repeat_diffs_sq = []
+        first_on = None
+        for repeat in range(CUDNN_REPEATS):
+            with sharding.set_mesh(mesh):
+                loss_on, g_on = f_on(state.params, step_rng, obs, act)
+                cudnn_losses.append(float(loss_on))
+                host_on = _to_host(g_on)
+            del g_on
+            if host_on_sum is None:
+                host_on_sum = {key: value.copy() for key, value in host_on.items()}
+                first_on = {key: value.copy() for key, value in host_on.items()}
+            else:
+                repeat_diffs_sq.append(sum(_dot(host_on[key] - first_on[key], host_on[key] - first_on[key]) for key in host_on))
+                for key in host_on_sum:
+                    host_on_sum[key] += host_on[key]
+            del host_on
+
+        host_on_avg = {key: value / CUDNN_REPEATS for key, value in host_on_sum.items()}
+        del host_on_sum
+        delta = {key: host_on_avg[key] - host_off[key] for key in host_off}
+
+        totals = collections.defaultdict(lambda: collections.defaultdict(float))
+        previous_dots = np.zeros(batch_idx, np.float64)
+        for path, grad in host_off.items():
+            cudnn_grad = host_on_avg[path]
+            diff = delta[path]
+            param = host_params[path]
+            group = _module(path)
+            for bucket in (totals["ALL"], totals[group]):
+                bucket["g2"] += _dot(grad, grad)
+                bucket["c2"] += _dot(cudnn_grad, cudnn_grad)
+                bucket["d2"] += _dot(diff, diff)
+                bucket["p2"] += _dot(param, param)
+                bucket["dg"] += _dot(diff, grad)
+                bucket["dp"] += _dot(diff, param)
+            for previous_idx, previous_delta in enumerate(deltas):
+                previous_dots[previous_idx] += _dot(diff, previous_delta[path])
+
+        current_d2 = totals["ALL"]["d2"]
+        pairwise_dot[batch_idx, batch_idx] = current_d2
+        for previous_idx, dot in enumerate(previous_dots):
+            pairwise_dot[batch_idx, previous_idx] = dot
+            pairwise_dot[previous_idx, batch_idx] = dot
+
+        deltas.append(delta)
+        delta_norm_sq.append(current_d2)
+        random_rel = np.mean(np.sqrt(repeat_diffs_sq)) / max(np.sqrt(totals["ALL"]["g2"]), 1e-300)
+        all_stats = totals["ALL"]
+        summary = {
+            "rel_delta": np.sqrt(all_stats["d2"] / all_stats["g2"]),
+            "cos_d_g": _cos(all_stats["dg"], all_stats["d2"], all_stats["g2"]),
+            "proj_d_g": all_stats["dg"] / max(all_stats["g2"], 1e-300),
+            "cos_d_p": _cos(all_stats["dp"], all_stats["d2"], all_stats["p2"]),
+            "proj_d_p": all_stats["dp"] / max(all_stats["p2"], 1e-300),
+            "random_rel": random_rel,
+        }
+        per_batch_summary.append(summary)
+
+        print(f"\n=== batch {batch_idx} ===", flush=True)
+        print(
+            f"loss explicit={loss_off:.6f} cudnn_mean={np.mean(cudnn_losses):.6f} "
+            f"loss_delta={np.mean(cudnn_losses) - loss_off:+.3e}",
+            flush=True,
+        )
+        print(
+            f"|g|={np.sqrt(all_stats['g2']):.6f} |delta_mean|/|g|={summary['rel_delta']:.4e} "
+            f"repeat_random/|g|={random_rel:.4e}",
+            flush=True,
+        )
+        print(
+            f"delta vs grad: cos={summary['cos_d_g']:+.5f} projection/|g|={summary['proj_d_g']:+.4e}; "
+            f"delta vs params: cos={summary['cos_d_p']:+.5f} projection/|p|={summary['proj_d_p']:+.4e}",
+            flush=True,
+        )
+        if batch_idx:
+            cosines = [
+                _cos(pairwise_dot[batch_idx, previous_idx], current_d2, delta_norm_sq[previous_idx])
+                for previous_idx in range(batch_idx)
+            ]
+            print("delta cosine vs prior batches: " + " ".join(f"{value:+.5f}" for value in cosines), flush=True)
+
+        print("module metrics (rel_delta, cos(delta,grad), cos(delta,param)):", flush=True)
+        for group in sorted(key for key in totals if key != "ALL"):
+            stats = totals[group]
+            values = (
+                np.sqrt(stats["d2"] / max(stats["g2"], 1e-300)),
+                _cos(stats["dg"], stats["d2"], stats["g2"]),
+                _cos(stats["dp"], stats["d2"], stats["p2"]),
+            )
+            per_module_summary[group].append(values)
+            print(f"  {group:20s} {values[0]:.4e} {values[1]:+.5f} {values[2]:+.5f}", flush=True)
+        del host_off, host_on_avg, first_on
+
+    cosine_matrix = np.eye(N_BATCHES, dtype=np.float64)
+    off_diagonal = []
+    for row in range(N_BATCHES):
+        for col in range(row):
+            value = _cos(pairwise_dot[row, col], delta_norm_sq[row], delta_norm_sq[col])
+            cosine_matrix[row, col] = cosine_matrix[col, row] = value
+            off_diagonal.append(value)
+
+    print("\n=== cross-batch delta cosine matrix ===", flush=True)
+    print("        " + " ".join(f"b{i:02d}" for i in range(N_BATCHES)), flush=True)
+    for idx, row in enumerate(cosine_matrix):
+        print(f"b{idx:02d} " + " ".join(f"{value:+.5f}" for value in row), flush=True)
+    print(
+        f"off-diagonal cosine: mean={np.mean(off_diagonal):+.6f} "
+        f"std={np.std(off_diagonal, ddof=1):.6f} min={np.min(off_diagonal):+.6f} "
+        f"max={np.max(off_diagonal):+.6f} n={len(off_diagonal)}",
+        flush=True,
+    )
+
+    print("\n=== aggregate projections across batches (mean +/- sample std) ===", flush=True)
+    for key in ("rel_delta", "random_rel", "cos_d_g", "proj_d_g", "cos_d_p", "proj_d_p"):
+        values = np.asarray([row[key] for row in per_batch_summary])
+        print(f"{key:12s} {values.mean():+.6e} +/- {values.std(ddof=1):.6e}", flush=True)
+
+    print("\n=== aggregate module metrics (mean rel_delta, cos(delta,grad), cos(delta,param)) ===", flush=True)
+    for group in sorted(per_module_summary):
+        values = np.asarray(per_module_summary[group])
+        means = values.mean(axis=0)
+        stds = values.std(axis=0, ddof=1)
+        print(
+            f"{group:20s} rel={means[0]:.4e}+/-{stds[0]:.2e} "
+            f"cos_g={means[1]:+.5f}+/-{stds[1]:.2e} cos_p={means[2]:+.5f}+/-{stds[2]:.2e}",
+            flush=True,
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/grad_ab_probe.py b/scripts/grad_ab_probe.py
new file mode 100644
index 0000000..1748a04
--- /dev/null
+++ b/scripts/grad_ab_probe.py
@@ -0,0 +1,162 @@
+"""Full-model gradient A/B: explicit attention vs cuDNN fused attention.
+
+Identical weights, identical batch, identical rng. The only difference is
+Pi0Config.use_cudnn_attention. Reports per-parameter relative gradient error so a
+systematic discrepancy can be localised to a module.
+"""
+
+import collections
+import dataclasses
+import functools
+import logging
+import sys
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+from jax._src.lib import cuda_versions
+import numpy as np
+import optax
+
+import openpi.training.config as _config
+import openpi.training.data_loader as _data_loader
+import openpi.training.sharding as sharding
+
+sys.path.insert(0, "scripts")
+import train as train_mod
+
+CONFIG = "pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100"
+BATCH = 32
+WORKERS = 16
+N_BATCHES = 3
+
+
+def _grad_fn(model_def, trainable_filter, params, rng, obs, act):
+    model = nnx.merge(model_def, params)
+    model.train()
+
+    def loss_fn(m, r, o, a):
+        return jnp.mean(m.compute_loss(r, o, a, train=True))
+
+    return nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, trainable_filter))(model, rng, obs, act)
+
+
+def _flat(state):
+    pure = state.to_pure_dict()
+    out = []
+
+    def rec(d, prefix):
+        if isinstance(d, dict):
+            for k, v in d.items():
+                rec(v, f"{prefix}/{k}" if prefix else str(k))
+        else:
+            out.append((prefix, d))
+
+    rec(pure, "")
+    return out
+
+
+def main():
+    logging.basicConfig(level=logging.WARNING)
+    print("cuDNN runtime:", cuda_versions.cudnn_get_version())
+
+    cfg = _config.get_config(CONFIG)
+    cfg = dataclasses.replace(cfg, batch_size=BATCH, num_workers=WORKERS, wandb_enabled=False)
+    cfg_off = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=False))
+    cfg_on = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=True))
+
+    mesh = sharding.make_mesh(cfg.fsdp_devices)
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+
+    rng = jax.random.key(cfg.seed)
+    train_rng, init_rng = jax.random.split(rng)
+
+    print("loading pi05_base weights ...")
+    state, _ = train_mod.init_train_state(cfg_off, init_rng, mesh, resume=False)
+    print("  param global norm:", float(optax.global_norm(state.params)))
+
+    # Same parameter values, two attention implementations.
+    gdef_off = state.model_def
+    gdef_on = nnx.graphdef(cfg_on.model.create(init_rng))
+
+    f_off = jax.jit(functools.partial(_grad_fn, gdef_off, cfg.trainable_filter))
+    f_on = jax.jit(functools.partial(_grad_fn, gdef_on, cfg.trainable_filter))
+
+    print("building data loader (this spawns workers) ...")
+    dl = _data_loader.create_data_loader(cfg_off, sharding=data_sharding, shuffle=True)
+    it = iter(dl)
+
+    agg = collections.defaultdict(lambda: [0.0, 0.0])
+
+    def to_host(state_tree):
+        """Pull one gradient set to numpy so the device copy can be freed."""
+        return {k: np.asarray(jax.device_get(v), np.float32) for k, v in _flat(state_tree)}
+
+    for b in range(N_BATCHES):
+        obs, act = next(it)
+        step_rng = jax.random.fold_in(train_rng, b)
+
+        # One gradient set at a time: 3B params in fp32 is ~12 GB, two live copies
+        # plus the tree arithmetic does not fit alongside the activations.
+        with sharding.set_mesh(mesh):
+            loss_off, g_off = f_off(state.params, step_rng, obs, act)
+            loss_off = float(loss_off)
+            h_off = to_host(g_off)
+        del g_off
+        with sharding.set_mesh(mesh):
+            loss_on, g_on = f_on(state.params, step_rng, obs, act)
+            loss_on = float(loss_on)
+            h_on = to_host(g_on)
+        del g_on
+        # Second identical cuDNN call. cuDNN's backward accumulates with atomics, so
+        # this measures the RANDOM part of the cuDNN/explicit gap. Whatever is left
+        # over is systematic, and only a systematic part can accumulate over steps.
+        with sharding.set_mesh(mesh):
+            _, g_on2 = f_on(state.params, step_rng, obs, act)
+            h_on2 = to_host(g_on2)
+        del g_on2
+
+        sq_off = sq_on = sq_d = 0.0
+        rows = []
+        for path, av in h_off.items():
+            av = av.astype(np.float64)
+            bv = h_on[path].astype(np.float64)
+            na = float(np.linalg.norm(av))
+            nb = float(np.linalg.norm(bv))
+            nd = float(np.linalg.norm(bv - av))
+            sq_off += na**2
+            sq_on += nb**2
+            sq_d += nd**2
+            rows.append((nd / max(na, 1e-30), path, na))
+            key = path.split("/")[0]
+            agg[key][0] += nd**2
+            agg[key][1] += na**2
+
+        sq_rand = 0.0
+        for path, av in h_on.items():
+            sq_rand += float(np.linalg.norm(h_on2[path].astype(np.float64) - av.astype(np.float64)) ** 2)
+        n_off, n_on, n_d = np.sqrt(sq_off), np.sqrt(sq_on), np.sqrt(sq_d)
+        n_rand = np.sqrt(sq_rand)
+        print(f"\n=== batch {b} ===")
+        print(f"  loss    explicit={loss_off:.6f}  cudnn={loss_on:.6f}  diff={loss_on - loss_off:+.3e}")
+        print(f"  |grad|  explicit={n_off:.6f}  cudnn={n_on:.6f}  ratio={n_on / max(n_off, 1e-30):.4f}")
+        rel_total = n_d / max(n_off, 1e-30)
+        rel_rand = n_rand / max(n_off, 1e-30)
+        # Total gap and the random part are orthogonal in expectation.
+        rel_sys = np.sqrt(max(rel_total**2 - rel_rand**2, 0.0))
+        print(f"  cuDNN vs explicit (total)   = {rel_total:.4e}")
+        print(f"  cuDNN vs cuDNN    (random)  = {rel_rand:.4e}")
+        print(f"  implied SYSTEMATIC component= {rel_sys:.4e}")
+        rows.sort(reverse=True)
+        print("  worst 10 parameters by relative gradient difference:")
+        for rel, path, na in rows[:10]:
+            print(f"    {rel:10.3e}  |g|={na:11.4e}  {path}")
+        del h_off, h_on, h_on2
+
+    print("\n=== aggregate by top-level module ===")
+    for k, (dn, an) in sorted(agg.items(), key=lambda kv: -(kv[1][0] / max(kv[1][1], 1e-30))):
+        print(f"  {np.sqrt(dn / max(an, 1e-30)):10.3e}  {k}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/hybrid_update_ab_probe.py b/scripts/hybrid_update_ab_probe.py
new file mode 100644
index 0000000..622f930
--- /dev/null
+++ b/scripts/hybrid_update_ab_probe.py
@@ -0,0 +1,192 @@
+"""Compare optimizer-space updates at a saved training state.
+
+Raw gradient L2 can hide errors in coordinates with a small Adam second moment.
+This probe restores the real step-1000 optimizer state, computes gradients for
+the same real batch/rng, and compares the actual AdamW updates produced by
+explicit attention, cuDNN attention, and the prior multiplicative-noise control.
+"""
+
+import dataclasses
+import functools
+import logging
+import sys
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import openpi.training.checkpoints as _checkpoints
+import openpi.training.config as _config
+import openpi.training.data_loader as _data_loader
+import openpi.training.sharding as sharding
+
+sys.path.insert(0, "scripts")
+import train as train_mod
+
+
+CONFIG = "diag_cudnn_strict_order"
+EXP_NAME = "diag_cudnn_strict_order"
+CHECKPOINT_STEP = 1000
+BATCH = 32
+WORKERS = 16
+
+
+def _loss_and_update(model_def, trainable_filter, tx, noise_scale, opt_state, params, rng, obs, act):
+    model = nnx.merge(model_def, params)
+    model.train()
+
+    def loss_fn(m, r, o, a):
+        return jnp.mean(m.compute_loss(r, o, a, train=True))
+
+    loss, grads = nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, trainable_filter))(
+        model, rng, obs, act
+    )
+    if noise_scale:
+        noise_rng = jax.random.fold_in(rng, 0x6E6F6973)
+        grads, _ = train_mod._add_relative_gradient_noise(grads, noise_rng, noise_scale)
+    trainable_params = params.filter(trainable_filter)
+    updates, _ = tx.update(grads, opt_state, trainable_params)
+    return loss, updates
+
+
+def _flat(tree):
+    pure = tree.to_pure_dict()
+    out = []
+
+    def rec(value, prefix):
+        if isinstance(value, dict):
+            for key, child in value.items():
+                rec(child, f"{prefix}/{key}" if prefix else str(key))
+        else:
+            out.append((prefix, value))
+
+    rec(pure, "")
+    return out
+
+
+def _to_host(tree):
+    return {path: np.asarray(jax.device_get(value), np.float32) for path, value in _flat(tree)}
+
+
+def _dot(a, b):
+    return float(np.vdot(a.astype(np.float64), b.astype(np.float64)))
+
+
+def _metrics(reference, candidate):
+    totals = {"ALL": [0.0, 0.0, 0.0, 0.0, 0, 0]}
+    for path, ref in reference.items():
+        cand = candidate[path]
+        delta = cand - ref
+        group = path.split("/", 1)[0]
+        totals.setdefault(group, [0.0, 0.0, 0.0, 0.0, 0, 0])
+        for bucket in (totals["ALL"], totals[group]):
+            bucket[0] += _dot(ref, ref)
+            bucket[1] += _dot(cand, cand)
+            bucket[2] += _dot(delta, delta)
+            bucket[3] += _dot(ref, cand)
+            nonzero = (ref != 0) | (cand != 0)
+            bucket[4] += int(np.count_nonzero(nonzero & (np.signbit(ref) != np.signbit(cand))))
+            bucket[5] += int(np.count_nonzero(nonzero))
+    return totals
+
+
+def _print_comparison(label, reference, candidate):
+    metrics = _metrics(reference, candidate)
+    ref2, cand2, delta2, dot, flips, count = metrics["ALL"]
+    print(f"\n=== {label} ===", flush=True)
+    print(
+        f"|update_ref|={np.sqrt(ref2):.8e} |update_candidate|={np.sqrt(cand2):.8e} "
+        f"rel_update_delta={np.sqrt(delta2 / ref2):.6e} "
+        f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.8f} "
+        f"sign_flip_fraction={flips / max(count, 1):.6e}",
+        flush=True,
+    )
+    for group in sorted(key for key in metrics if key != "ALL"):
+        ref2, cand2, delta2, dot, flips, count = metrics[group]
+        print(
+            f"  {group:20s} rel={np.sqrt(delta2 / max(ref2, 1e-300)):.6e} "
+            f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.8f} "
+            f"sign_flips={flips / max(count, 1):.6e}",
+            flush=True,
+        )
+
+
+def main():
+    logging.basicConfig(level=logging.WARNING)
+    cfg = _config.get_config(CONFIG)
+    cfg = dataclasses.replace(
+        cfg,
+        exp_name=EXP_NAME,
+        batch_size=BATCH,
+        num_workers=WORKERS,
+        resume=True,
+        overwrite=False,
+        wandb_enabled=False,
+    )
+    cfg_on = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=True))
+    cfg_off = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=False))
+
+    mesh = sharding.make_mesh(cfg.fsdp_devices)
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+    rng = jax.random.key(cfg.seed)
+    train_rng, init_rng = jax.random.split(rng)
+
+    manager, resuming = _checkpoints.initialize_checkpoint_dir(
+        cfg_on.checkpoint_dir,
+        keep_period=cfg_on.keep_period,
+        overwrite=False,
+        resume=True,
+    )
+    if not resuming or CHECKPOINT_STEP not in manager.all_steps():
+        raise RuntimeError(f"checkpoint {CHECKPOINT_STEP} unavailable; found {manager.all_steps()}")
+    state_shape, state_sharding = train_mod.init_train_state(cfg_on, init_rng, mesh, resume=True)
+    state = _checkpoints.restore_state(manager, state_shape, None, step=CHECKPOINT_STEP)
+    jax.block_until_ready(state)
+    print(f"restored checkpoint directory step={CHECKPOINT_STEP}; train_state.step={int(state.step)}", flush=True)
+
+    print("building one-batch loader ...", flush=True)
+    loader = _data_loader.create_data_loader(cfg_off, sharding=data_sharding, shuffle=True)
+    obs, act = next(iter(loader))
+    step_rng = jax.random.fold_in(train_rng, state.step)
+
+    gdef_off = nnx.graphdef(cfg_off.model.create(init_rng))
+    hybrid_ranges = [
+        ("full_18", 0, None),
+        ("first_6", 0, 6),
+        ("middle_6", 6, 6),
+        ("last_6", 12, 6),
+        ("first_3", 0, 3),
+        ("middle_3", 7, 3),
+        ("last_3", 15, 3),
+    ]
+    graphdefs = {}
+    for name, start, count in hybrid_ranges:
+        model_config = dataclasses.replace(
+            cfg_on.model,
+            use_cudnn_attention=True,
+            cudnn_attention_layer_start=start,
+            cudnn_attention_num_layers=count,
+        )
+        graphdefs[name] = nnx.graphdef(model_config.create(init_rng))
+
+    fn = jax.jit(functools.partial(_loss_and_update, gdef_off, cfg.trainable_filter, state.tx, 0.0))
+    with sharding.set_mesh(mesh):
+        loss, updates = fn(state.opt_state, state.params, step_rng, obs, act)
+        reference = _to_host(updates)
+    del updates, fn
+    print(f"computed explicit reference: loss={float(loss):.8f}", flush=True)
+
+    for name, _, _ in hybrid_ranges:
+        fn = jax.jit(functools.partial(_loss_and_update, graphdefs[name], cfg.trainable_filter, state.tx, 0.0))
+        with sharding.set_mesh(mesh):
+            loss, updates = fn(state.opt_state, state.params, step_rng, obs, act)
+            candidate = _to_host(updates)
+        del updates, fn
+        print(f"computed {name}: loss={float(loss):.8f}", flush=True)
+        _print_comparison(name, reference, candidate)
+        del candidate
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/optimizer_update_ab_probe.py b/scripts/optimizer_update_ab_probe.py
new file mode 100644
index 0000000..487e205
--- /dev/null
+++ b/scripts/optimizer_update_ab_probe.py
@@ -0,0 +1,194 @@
+"""Compare optimizer-space updates at a saved training state.
+
+Raw gradient L2 can hide errors in coordinates with a small Adam second moment.
+This probe restores the real step-1000 optimizer state, computes gradients for
+the same real batch/rng, and compares the actual AdamW updates produced by
+explicit attention, cuDNN attention, and the prior multiplicative-noise control.
+"""
+
+import dataclasses
+import functools
+import logging
+import sys
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import openpi.training.checkpoints as _checkpoints
+import openpi.training.config as _config
+import openpi.training.data_loader as _data_loader
+import openpi.training.sharding as sharding
+
+sys.path.insert(0, "scripts")
+import train as train_mod
+
+
+CONFIG = "diag_cudnn_strict_order"
+EXP_NAME = "diag_cudnn_strict_order"
+CHECKPOINT_STEP = 1000
+BATCH = 32
+WORKERS = 16
+
+
+def _loss_and_update(model_def, trainable_filter, tx, noise_scale, opt_state, params, rng, obs, act):
+    model = nnx.merge(model_def, params)
+    model.train()
+
+    def loss_fn(m, r, o, a):
+        return jnp.mean(m.compute_loss(r, o, a, train=True))
+
+    loss, grads = nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, trainable_filter))(
+        model, rng, obs, act
+    )
+    if noise_scale:
+        noise_rng = jax.random.fold_in(rng, 0x6E6F6973)
+        grads, _ = train_mod._add_relative_gradient_noise(grads, noise_rng, noise_scale)
+    trainable_params = params.filter(trainable_filter)
+    updates, _ = tx.update(grads, opt_state, trainable_params)
+    return loss, updates
+
+
+def _flat(tree):
+    pure = tree.to_pure_dict()
+    out = []
+
+    def rec(value, prefix):
+        if isinstance(value, dict):
+            for key, child in value.items():
+                rec(child, f"{prefix}/{key}" if prefix else str(key))
+        else:
+            out.append((prefix, value))
+
+    rec(pure, "")
+    return out
+
+
+def _to_host(tree):
+    return {path: np.asarray(jax.device_get(value), np.float32) for path, value in _flat(tree)}
+
+
+def _dot(a, b):
+    return float(np.vdot(a.astype(np.float64), b.astype(np.float64)))
+
+
+def _metrics(reference, candidate):
+    totals = {"ALL": [0.0, 0.0, 0.0, 0.0, 0, 0]}
+    for path, ref in reference.items():
+        cand = candidate[path]
+        delta = cand - ref
+        group = path.split("/", 1)[0]
+        totals.setdefault(group, [0.0, 0.0, 0.0, 0.0, 0, 0])
+        for bucket in (totals["ALL"], totals[group]):
+            bucket[0] += _dot(ref, ref)
+            bucket[1] += _dot(cand, cand)
+            bucket[2] += _dot(delta, delta)
+            bucket[3] += _dot(ref, cand)
+            nonzero = (ref != 0) | (cand != 0)
+            bucket[4] += int(np.count_nonzero(nonzero & (np.signbit(ref) != np.signbit(cand))))
+            bucket[5] += int(np.count_nonzero(nonzero))
+    return totals
+
+
+def _print_comparison(label, reference, candidate):
+    metrics = _metrics(reference, candidate)
+    ref2, cand2, delta2, dot, flips, count = metrics["ALL"]
+    print(f"\n=== {label} ===", flush=True)
+    print(
+        f"|update_ref|={np.sqrt(ref2):.8e} |update_candidate|={np.sqrt(cand2):.8e} "
+        f"rel_update_delta={np.sqrt(delta2 / ref2):.6e} "
+        f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.8f} "
+        f"sign_flip_fraction={flips / max(count, 1):.6e}",
+        flush=True,
+    )
+    for group in sorted(key for key in metrics if key != "ALL"):
+        ref2, cand2, delta2, dot, flips, count = metrics[group]
+        print(
+            f"  {group:20s} rel={np.sqrt(delta2 / max(ref2, 1e-300)):.6e} "
+            f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.8f} "
+            f"sign_flips={flips / max(count, 1):.6e}",
+            flush=True,
+        )
+
+
+def main():
+    logging.basicConfig(level=logging.WARNING)
+    cfg = _config.get_config(CONFIG)
+    cfg = dataclasses.replace(
+        cfg,
+        exp_name=EXP_NAME,
+        batch_size=BATCH,
+        num_workers=WORKERS,
+        resume=True,
+        overwrite=False,
+        wandb_enabled=False,
+    )
+    cfg_on = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=True))
+    cfg_off = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=False))
+
+    mesh = sharding.make_mesh(cfg.fsdp_devices)
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+    rng = jax.random.key(cfg.seed)
+    train_rng, init_rng = jax.random.split(rng)
+
+    manager, resuming = _checkpoints.initialize_checkpoint_dir(
+        cfg_on.checkpoint_dir,
+        keep_period=cfg_on.keep_period,
+        overwrite=False,
+        resume=True,
+    )
+    if not resuming or CHECKPOINT_STEP not in manager.all_steps():
+        raise RuntimeError(f"checkpoint {CHECKPOINT_STEP} unavailable; found {manager.all_steps()}")
+    state_shape, state_sharding = train_mod.init_train_state(cfg_on, init_rng, mesh, resume=True)
+    state = _checkpoints.restore_state(manager, state_shape, None, step=CHECKPOINT_STEP)
+    jax.block_until_ready(state)
+    print(f"restored checkpoint directory step={CHECKPOINT_STEP}; train_state.step={int(state.step)}", flush=True)
+
+    print("building one-batch loader ...", flush=True)
+    loader = _data_loader.create_data_loader(cfg_off, sharding=data_sharding, shuffle=True)
+    obs, act = next(iter(loader))
+    step_rng = jax.random.fold_in(train_rng, state.step)
+
+    gdef_on = state.model_def
+    gdef_off = nnx.graphdef(cfg_off.model.create(init_rng))
+
+    # The optimizer state structure is unchanged when only Adam epsilon changes.
+    eps_values = (1e-8, 1e-6, 1e-5)
+    tx_by_eps = {1e-8: state.tx}
+    for eps in eps_values[1:]:
+        cfg_eps = dataclasses.replace(cfg_on, optimizer=dataclasses.replace(cfg_on.optimizer, eps=eps))
+        shape_eps, _ = train_mod.init_train_state(cfg_eps, init_rng, mesh, resume=True)
+        tx_by_eps[eps] = shape_eps.tx
+
+    host_updates = {}
+    variants = [
+        ("explicit_eps1e-8", gdef_off, tx_by_eps[1e-8], 0.0),
+        ("cudnn_eps1e-8", gdef_on, tx_by_eps[1e-8], 0.0),
+        ("explicit_noise015_eps1e-8", gdef_off, tx_by_eps[1e-8], 0.015),
+        ("explicit_eps1e-6", gdef_off, tx_by_eps[1e-6], 0.0),
+        ("cudnn_eps1e-6", gdef_on, tx_by_eps[1e-6], 0.0),
+        ("explicit_eps1e-5", gdef_off, tx_by_eps[1e-5], 0.0),
+        ("cudnn_eps1e-5", gdef_on, tx_by_eps[1e-5], 0.0),
+    ]
+    for name, graphdef, tx, noise_scale in variants:
+        fn = jax.jit(functools.partial(_loss_and_update, graphdef, cfg.trainable_filter, tx, noise_scale))
+        with sharding.set_mesh(mesh):
+            loss, updates = fn(state.opt_state, state.params, step_rng, obs, act)
+            loss = float(loss)
+            host_updates[name] = _to_host(updates)
+        del updates, fn
+        print(f"computed {name}: loss={loss:.8f}", flush=True)
+
+    _print_comparison("real cuDNN delta, Adam eps=1e-8", host_updates["explicit_eps1e-8"], host_updates["cudnn_eps1e-8"])
+    _print_comparison(
+        "prior multiplicative 1.5% noise, Adam eps=1e-8",
+        host_updates["explicit_eps1e-8"],
+        host_updates["explicit_noise015_eps1e-8"],
+    )
+    _print_comparison("real cuDNN delta, Adam eps=1e-6", host_updates["explicit_eps1e-6"], host_updates["cudnn_eps1e-6"])
+    _print_comparison("real cuDNN delta, Adam eps=1e-5", host_updates["explicit_eps1e-5"], host_updates["cudnn_eps1e-5"])
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/train.py b/scripts/train.py
index f65f422..b7ddd38 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -135,6 +135,30 @@ def init(rng: at.KeyArrayLike, partial_params: at.Params | None = None) -> train
     return train_state, state_sharding
 
 
+def _add_relative_gradient_noise(grads, rng: at.KeyArrayLike, relative_scale: float):
+    """Adds reproducible random noise with an exact tree-wide relative L2 norm.
+
+    The raw noise is weighted by each gradient element before global
+    normalization. This keeps the perturbation energy distributed like the
+    gradient energy instead of concentrating it in the largest parameter leaf.
+    """
+    leaves, tree_def = jax.tree.flatten(grads)
+    keys = jax.random.split(rng, len(leaves))
+    raw_noise = jax.tree.unflatten(
+        tree_def,
+        [grad * jax.random.normal(key, grad.shape, dtype=grad.dtype) for grad, key in zip(leaves, keys)],
+    )
+
+    grad_norm = optax.global_norm(grads)
+    raw_noise_norm = optax.global_norm(raw_noise)
+    tiny = jnp.finfo(grad_norm.dtype).tiny
+    multiplier = jnp.asarray(relative_scale, grad_norm.dtype) * grad_norm / jnp.maximum(raw_noise_norm, tiny)
+    noise = jax.tree.map(lambda value: value * multiplier, raw_noise)
+    noisy_grads = jax.tree.map(lambda grad, delta: grad + delta, grads, noise)
+    actual_relative_norm = multiplier * raw_noise_norm / jnp.maximum(grad_norm, tiny)
+    return noisy_grads, actual_relative_norm
+
+
 @at.typecheck
 def train_step(
     config: _config.TrainConfig,
@@ -160,6 +184,18 @@ def loss_fn(
     diff_state = nnx.DiffState(0, config.trainable_filter)
     loss, grads = nnx.value_and_grad(loss_fn, argnums=diff_state)(model, train_rng, observation, actions)
 
+    clean_grads = grads
+    gradient_noise_rel_norm = None
+    if config.gradient_noise_scale > 0:
+        # Keep the model RNG identical to the clean explicit-attention path.
+        # fold_in gives the diagnostic noise an independent, step-reproducible stream.
+        noise_rng = jax.random.fold_in(train_rng, 0x6E6F6973)
+        grads, gradient_noise_rel_norm = _add_relative_gradient_noise(
+            grads,
+            noise_rng,
+            config.gradient_noise_scale,
+        )
+
     params = state.params.filter(config.trainable_filter)
     updates, new_opt_state = state.tx.update(grads, state.opt_state, params)
     new_params = optax.apply_updates(params, updates)
@@ -189,11 +225,13 @@ def loss_fn(
             ),
         )
         grad_norm = optax.global_norm(grads)
+        clean_grad_norm = optax.global_norm(clean_grads) if gradient_noise_rel_norm is not None else None
         param_norm = optax.global_norm(kernel_params)
     else:
         # Keep a stable output pytree for both JIT variants. The host-side logging
         # reduction ignores these placeholders with nanmean.
         grad_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
+        clean_grad_norm = jnp.asarray(jnp.nan, dtype=loss.dtype) if gradient_noise_rel_norm is not None else None
         param_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
 
     info = {
@@ -201,6 +239,9 @@ def loss_fn(
         "grad_norm": grad_norm,
         "param_norm": param_norm,
     }
+    if gradient_noise_rel_norm is not None:
+        info["clean_grad_norm"] = clean_grad_norm
+        info["gradient_noise_rel_norm"] = gradient_noise_rel_norm
     return new_state, info
 
 
@@ -210,6 +251,13 @@ def main(config: _config.TrainConfig):
     cudnn_runtime_version = cuda_versions.cudnn_get_version() if cuda_versions is not None else None
     logging.info(f"JAX cuDNN runtime version: {cudnn_runtime_version}")
 
+    if not 0.0 <= config.gradient_noise_scale < 1.0:
+        raise ValueError(f"gradient_noise_scale must be in [0, 1), got {config.gradient_noise_scale}")
+    attention_backend = "cuDNN" if getattr(config.model, "use_cudnn_attention", False) else "explicit"
+    logging.info(
+        f"Attention backend: {attention_backend}; gradient_noise_scale={config.gradient_noise_scale:.6f}"
+    )
+
     if config.batch_size % jax.device_count() != 0:
         raise ValueError(
             f"Batch size {config.batch_size} must be divisible by the number of devices {jax.device_count()}."
@@ -334,6 +382,10 @@ def main(config: _config.TrainConfig):
             reduced_info = jax.device_get(jax.tree.map(jnp.nanmean, stacked_infos))
             info_str = ", ".join(f"{k}={v:.4f}" for k, v in reduced_info.items())
             pbar.write(f"Step {step}: {info_str}")
+            # tqdm_loggable can drop pbar.write() lines when stdout is redirected.
+            # Emit the same diagnostics through the normal logger so unattended
+            # training runs always retain loss and gradient trends.
+            logging.info("Step %d: %s", step, info_str)
             wandb.log(reduced_info, step=step)
             infos = []
             t_log = time.monotonic() - t0
diff --git a/scripts/train_test.py b/scripts/train_test.py
index 6165930..e5cba2d 100644
--- a/scripts/train_test.py
+++ b/scripts/train_test.py
@@ -6,11 +6,41 @@
 
 os.environ["JAX_PLATFORMS"] = "cpu"
 
+import jax
+import jax.numpy as jnp
+import numpy as np
+
 from openpi.training import config as _config
 
 from . import train
 
 
+def test_add_relative_gradient_noise():
+    grads = {
+        "a": jnp.arange(1, 13, dtype=jnp.float32).reshape(3, 4),
+        "b": jnp.asarray([-3.0, 0.0, 5.0], dtype=jnp.float32),
+    }
+    scale = 0.015
+    key = jax.random.key(7)
+
+    noisy, reported_rel_norm = train._add_relative_gradient_noise(grads, key, scale)
+    noise = jax.tree.map(lambda perturbed, clean: perturbed - clean, noisy, grads)
+    measured_rel_norm = np.sqrt(sum(float(jnp.sum(x**2)) for x in jax.tree.leaves(noise))) / np.sqrt(
+        sum(float(jnp.sum(x**2)) for x in jax.tree.leaves(grads))
+    )
+
+    assert float(reported_rel_norm) == pytest.approx(scale, rel=1e-6)
+    assert measured_rel_norm == pytest.approx(scale, rel=1e-6)
+    same_noise, _ = train._add_relative_gradient_noise(grads, key, scale)
+    different_noise, _ = train._add_relative_gradient_noise(grads, jax.random.key(8), scale)
+    for actual, expected in zip(jax.tree.leaves(same_noise), jax.tree.leaves(noisy)):
+        np.testing.assert_array_equal(actual, expected)
+    assert any(
+        not np.array_equal(actual, expected)
+        for actual, expected in zip(jax.tree.leaves(different_noise), jax.tree.leaves(noisy))
+    )
+
+
 @pytest.mark.parametrize("config_name", ["debug_pi05"])
 def test_train(tmp_path: pathlib.Path, config_name: str):
     config = dataclasses.replace(
diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index 7955a06..5c1ff6a 100644
--- a/src/openpi/models/gemma.py
+++ b/src/openpi/models/gemma.py
@@ -161,7 +161,23 @@ def decode(self, x):
 
 
 def _stop_gradient_for_fully_masked_queries(q, attn_mask):
-    """Preserve the raw cuDNN mask while zeroing gradients for invalid queries."""
+    """Cut fully masked query rows out of the backward pass, keeping the mask intact.
+
+    Padding tokens are neither valid queries nor valid keys, so make_attn_mask
+    produces query rows that are entirely false, and cuDNN 9.14 returned NaN
+    q-gradients for them at the production shape. Zeroing the gradient of those rows
+    fixes that without touching attn_mask.
+
+    The alternative -- opening a dummy key for each empty row -- is equivalent
+    numerically (measured bit-identical q/k/v-gradients on valid rows at the real
+    training shape) but rebuilds the whole (B, 1, T, S) mask inside every layer,
+    which costs about 0.3 s/step at batch 256. Prefer this one.
+
+    Neither variant is the cure for the 2026-08-30 divergence, which is still open:
+    the kernel itself is correct, but cuDNN's answer differs from the explicit path
+    by about 2% per step and this recipe has no margin for that. Do not spend time on
+    the mask again. See docs/2026-08-30-cudnn-attention-divergence.md.
+    """
     query_has_key = jnp.any(attn_mask, axis=-1)[:, 0, :, None, None]
     return jnp.where(query_has_key, q, jax.lax.stop_gradient(q))
 
@@ -174,7 +190,7 @@ class Attention(nn.Module):
     use_cudnn_attention: bool = False
 
     @nn.compact
-    def __call__(self, xs, positions, attn_mask, kv_cache):
+    def __call__(self, xs, positions, attn_mask, kv_cache, use_cudnn_attention=None):
         # all experts must share the same head dim, num heads, and num kv heads for self-attention to work
         assert all(config.head_dim == self.configs[0].head_dim for config in self.configs)
         assert all(config.num_heads == self.configs[0].num_heads for config in self.configs)
@@ -231,34 +247,46 @@ def __call__(self, xs, positions, attn_mask, kv_cache):
                 f"Attention mask with shape {attn_mask.shape} but shapes for q and k are: {q.shape} and {k.shape}"
             )
 
-        if self.use_cudnn_attention and kv_cache is None:
+        def cudnn_attention(operands):
+            q, k, v, mask = operands
             # q is already scaled above, so disable dot_product_attention's
             # default head-dimension scaling. Keep cached inference on the
             # existing implementation; this switch targets training throughput.
-            # cuDNN 9.14 produces NaN q-gradients for fully masked query rows at
-            # the production GQA shape. Preserve the original fast mask and cut
-            # only those invalid query rows out of the backward path.
-            q = _stop_gradient_for_fully_masked_queries(q, attn_mask)
-            encoded = jax.nn.dot_product_attention(
+            # Fully masked query rows make cuDNN produce NaN q-gradients; drop
+            # those rows from the backward pass instead of rebuilding the mask.
+            q = _stop_gradient_for_fully_masked_queries(q, mask)
+            return jax.nn.dot_product_attention(
                 q,
                 k,
                 v,
-                mask=attn_mask,
+                mask=mask,
                 scale=1.0,
                 implementation="cudnn",
             )
-        else:
+
+        def explicit_attention(operands):
+            q, k, v, mask = operands
             grouped_q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=self.configs[0].num_kv_heads)
             logits = jnp.einsum("BTKGH,BSKH->BKGTS", grouped_q, k, preferred_element_type=jnp.float32)
 
             # big_neg = jnp.finfo(logits.dtype).min
             big_neg = -2.3819763e38  # See gemma/modules.py
-            masked_logits = jnp.where(attn_mask[:, :, None, :, :], logits, big_neg)
+            masked_logits = jnp.where(mask[:, :, None, :, :], logits, big_neg)
 
             probs = jax.nn.softmax(masked_logits, axis=-1).astype(dtype)
 
             encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v)
-            encoded = einops.rearrange(encoded, "B T K G H -> B T (K G) H")
+            return einops.rearrange(encoded, "B T K G H -> B T (K G) H")
+
+        if use_cudnn_attention is None:
+            use_cudnn_attention = self.use_cudnn_attention
+        operands = (q, k, v, attn_mask)
+        if kv_cache is not None:
+            encoded = explicit_attention(operands)
+        elif isinstance(use_cudnn_attention, bool):
+            encoded = cudnn_attention(operands) if use_cudnn_attention else explicit_attention(operands)
+        else:
+            encoded = jax.lax.cond(use_cudnn_attention, cudnn_attention, explicit_attention, operands)
 
         out = []
         start = 0
@@ -316,16 +344,24 @@ class Block(nn.Module):
 
     configs: tuple[Config, ...]
     use_cudnn_attention: bool = False
+    cudnn_attention_layer_start: int = 0
+    cudnn_attention_num_layers: int | None = None
 
     dropout: float = 0.0
     dropout_bdims: tuple[int, ...] = ()
 
     @nn.compact
-    def __call__(self, xs, kv_cache, positions, attn_mask, adarms_cond, deterministic=True):
+    def __call__(self, xs, kv_cache, positions, attn_mask, adarms_cond, layer_index=None, deterministic=True):
         xs = sharding.activation_sharding_constraint(xs)
         drop = nn.Dropout(self.dropout, self.dropout_bdims) if self.dropout else lambda x, _: x
 
         attn = Attention(configs=self.configs, use_cudnn_attention=self.use_cudnn_attention, name="attn")
+        layer_uses_cudnn = self.use_cudnn_attention
+        if layer_uses_cudnn and self.cudnn_attention_num_layers is not None:
+            layer_uses_cudnn = jnp.logical_and(
+                layer_index >= self.cudnn_attention_layer_start,
+                layer_index < self.cudnn_attention_layer_start + self.cudnn_attention_num_layers,
+            )
 
         pre_attn = []
         gates = []
@@ -338,7 +374,9 @@ def __call__(self, xs, kv_cache, positions, attn_mask, adarms_cond, deterministi
             pre_attn.append(x)
 
         pre_attn = sharding.activation_sharding_constraint(pre_attn)
-        post_attn, kv_cache = attn(pre_attn, positions, attn_mask, kv_cache)
+        post_attn, kv_cache = attn(
+            pre_attn, positions, attn_mask, kv_cache, use_cudnn_attention=layer_uses_cudnn
+        )
         post_attn = jax.tree.map(lambda x: drop(x, deterministic), post_attn)
         post_attn = sharding.activation_sharding_constraint(post_attn)
         xs = [_gated_residual(x, y, gate) for x, y, gate in zip(xs, post_attn, gates, strict=True)]
@@ -382,6 +420,8 @@ class Module(nn.Module):
     dropout_bdims: tuple[int, ...] = ()  # Every float is dropped independently.
     adarms: bool = False
     use_cudnn_attention: bool = False
+    cudnn_attention_layer_start: int = 0
+    cudnn_attention_num_layers: int | None = None
 
     def setup(self):
         # all experts must have the same depth
@@ -395,7 +435,7 @@ def setup(self):
         block_cls = nn.remat(
             Block,
             prevent_cse=False,
-            static_argnums=(5,),  # 0=self, 6=deterministic
+            static_argnums=(6,),  # 0=self, 7=deterministic
             policy=jax.checkpoint_policies.nothing_saveable,
         )
         self.layers = nn.scan(
@@ -407,12 +447,15 @@ def setup(self):
                 nn.broadcast,
                 nn.broadcast,
                 nn.broadcast,
+                0,
                 nn.broadcast,
-            ),  # 0=kv_cache, 1=positions, 2=mask, 3=adarms_cond, 4=deterministic
+            ),  # 0=kv_cache, 1=positions, 2=mask, 3=adarms_cond, 4=layer_index, 5=deterministic
             length=self.configs[0].depth,
         )(
             configs=self.configs,
             use_cudnn_attention=self.use_cudnn_attention,
+            cudnn_attention_layer_start=self.cudnn_attention_layer_start,
+            cudnn_attention_num_layers=self.cudnn_attention_num_layers,
             dropout=self.dropout,
             dropout_bdims=self.dropout_bdims,
         )
@@ -439,7 +482,10 @@ def __call__(
         if adarms_cond is None:
             adarms_cond = [None] * len(self.configs)
 
-        embedded, kv_cache = self.layers(embedded, kv_cache, positions, mask, adarms_cond, deterministic)
+        layer_indices = jnp.arange(self.configs[0].depth, dtype=jnp.int32)
+        embedded, kv_cache = self.layers(
+            embedded, kv_cache, positions, mask, adarms_cond, layer_indices, deterministic
+        )
 
         assert all(e.dtype == jnp.dtype(self.embed_dtype) for e in embedded if e is not None)
 
diff --git a/src/openpi/models/pi0.py b/src/openpi/models/pi0.py
index 7420c04..ebb0363 100644
--- a/src/openpi/models/pi0.py
+++ b/src/openpi/models/pi0.py
@@ -95,6 +95,8 @@ def __init__(
                 embed_dtype=config.dtype,
                 adarms=config.pi05,
                 use_cudnn_attention=config.use_cudnn_attention,
+                cudnn_attention_layer_start=config.cudnn_attention_layer_start,
+                cudnn_attention_num_layers=config.cudnn_attention_num_layers,
             )
         )
         llm.lazy_init(
diff --git a/src/openpi/models/pi0_config.py b/src/openpi/models/pi0_config.py
index d506af7..6cfdfcf 100644
--- a/src/openpi/models/pi0_config.py
+++ b/src/openpi/models/pi0_config.py
@@ -36,6 +36,10 @@ class Pi0Config(_model.BaseModelConfig):
     # Use cuDNN fused attention for training-time Gemma attention. This requires
     # a cuDNN runtime that supports the model's BF16 GQA shapes and attention mask.
     use_cudnn_attention: bool = False
+    # Diagnostic hybrid mode. None keeps the historical all-layer behavior;
+    # otherwise only this contiguous layer range uses cuDNN.
+    cudnn_attention_layer_start: int = 0
+    cudnn_attention_num_layers: int | None = None
 
     # training-time RTC config
     enable_training_time_rtc: bool = False
@@ -46,6 +50,10 @@ def __post_init__(self):
             object.__setattr__(self, "max_token_len", 200 if self.pi05 else 48)
         if self.discrete_state_input is None:
             object.__setattr__(self, "discrete_state_input", self.pi05)
+        if self.cudnn_attention_layer_start < 0:
+            raise ValueError("cudnn_attention_layer_start must be non-negative")
+        if self.cudnn_attention_num_layers is not None and self.cudnn_attention_num_layers < 0:
+            raise ValueError("cudnn_attention_num_layers must be non-negative or None")
 
     @property
     @override
diff --git a/src/openpi/models/pi0_test.py b/src/openpi/models/pi0_test.py
index 89e0c8e..ef4fbb0 100644
--- a/src/openpi/models/pi0_test.py
+++ b/src/openpi/models/pi0_test.py
@@ -30,6 +30,7 @@ def test_encode_image_views_as_batch_matches_separate_calls():
 
 
 def test_cudnn_attention_stops_gradient_only_for_empty_query_rows():
+    # Rows 0 and 1 have keys; rows 2 and 3 are the fully masked rows padding produces.
     mask = jnp.array(
         [
             [
@@ -52,6 +53,7 @@ def apply_and_sum(q_value):
     stopped_q = _gemma._stop_gradient_for_fully_masked_queries(q, mask)
     grad = jax.grad(apply_and_sum)(q)
 
+    # Forward is untouched, so valid attention outputs cannot change.
     assert jnp.array_equal(stopped_q, q)
     assert jnp.array_equal(grad[:, :2], jnp.ones_like(grad[:, :2]))
     assert jnp.array_equal(grad[:, 2:], jnp.zeros_like(grad[:, 2:]))
diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py
index 3e316a8..b87145f 100644
--- a/src/openpi/training/config.py
+++ b/src/openpi/training/config.py
@@ -590,6 +590,10 @@ class TrainConfig:
     optimizer: _optimizer.OptimizerConfig = dataclasses.field(default_factory=_optimizer.AdamW)
     ema_decay: float | None = 0.99
 
+    # Diagnostic-only relative L2 perturbation applied to gradients before the
+    # optimizer update. Zero preserves the normal training path exactly.
+    gradient_noise_scale: float = 0.0
+
     # Specifies which weights should be frozen.
     freeze_filter: tyro.conf.Suppress[Filter] = dataclasses.field(default_factory=nnx.Nothing)
 

From c9e77cc6d786a5dd371f7e9ca5419eb68744c293 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Mon, 31 Aug 2026 17:00:58 +0800
Subject: [PATCH 10/28] cudnn issues

---
 docs/2026-08-31-cudnn-sdpa-nan-summary.md | 28 +++++++++++++++++++++++
 1 file changed, 28 insertions(+)
 create mode 100644 docs/2026-08-31-cudnn-sdpa-nan-summary.md

diff --git a/docs/2026-08-31-cudnn-sdpa-nan-summary.md b/docs/2026-08-31-cudnn-sdpa-nan-summary.md
new file mode 100644
index 0000000..5a3d0db
--- /dev/null
+++ b/docs/2026-08-31-cudnn-sdpa-nan-summary.md
@@ -0,0 +1,28 @@
+# cuDNN fused Attention NaN 问题概述
+
+## 问题原因
+
+训练启用 JAX cuDNN fused Attention 后，前向 loss 起初正常，但反向传播先产生 NaN 梯度，随后参数和 loss 被持续污染。当前运行时为 cuDNN 9.14。
+
+NVIDIA 已确认，旧版 cuDNN SDPA backward 在静态序列长度不是 128 倍数时可能产生错误结果或 NaN，并在 cuDNN 9.15.1 中修复。当前模型的 Attention 序列长度约为：
+
+```text
+3 × 256 个图像 token + 200 个文本 token + 50 个 action token = 1018
+1018 % 128 = 122
+```
+
+这与 NVIDIA 报告的触发条件高度一致。相关 issue 还表明：即使前向输出正确，cuDNN backward 的 `dQ` 仍可能出现整行 NaN。
+
+参考：
+
+- [NVIDIA/cudnn-frontend #160](https://github.com/NVIDIA/cudnn-frontend/issues/160)
+- [cuDNN 9.15.1 Release Notes](https://docs.nvidia.com/deeplearning/cudnn/backend/v9.15.1/release-notes.html)
+
+## 尝试解决办法
+
+1. **暂时关闭 cuDNN Attention**：恢复显式 Attention 路径，保证生产训练数值有效。
+2. **启用原生 padding 路径**：向 JAX `dot_product_attention` 传入值为完整物理长度的 `query_seq_lengths` 和 `key_value_seq_lengths`，强制 cuDNN 使用 padding kernel；现有 block mask 仍负责实际 token 屏蔽。
+3. **升级 cuDNN runtime**：将实际加载版本从 9.14 升到至少 9.15.1，并检查启动日志确认版本。
+4. **长度对照实验**：把 `max_token_len` 临时改为 206，使总长度变为 1024，以验证“非 128 倍数”是否为直接触发条件。
+
+以上方案需要从有效 checkpoint 使用固定真实 batch 做逐 step 有限性检查。至少验证 500 个连续 step 的 loss、梯度、参数和优化器状态均为有限值后，才能重新用于生产训练。

From 895811c44ff8d94e9301c0f58d95c0587059de60 Mon Sep 17 00:00:00 2001
From: Hubo1231 <1195132365@qq.com>
Date: Mon, 31 Aug 2026 21:58:26 +0800
Subject: [PATCH 11/28] docs: consolidate cuDNN attention incident reports

---
 ...2026-08-29-cudnn-attention-nan-incident.md | 118 -----
 docs/2026-08-30-cudnn-attention-divergence.md | 375 --------------
 docs/2026-08-31-cudnn-sdpa-nan-summary.md     | 483 +++++++++++++++++-
 3 files changed, 471 insertions(+), 505 deletions(-)
 delete mode 100644 docs/2026-08-29-cudnn-attention-nan-incident.md
 delete mode 100644 docs/2026-08-30-cudnn-attention-divergence.md

diff --git a/docs/2026-08-29-cudnn-attention-nan-incident.md b/docs/2026-08-29-cudnn-attention-nan-incident.md
deleted file mode 100644
index c958bd5..0000000
--- a/docs/2026-08-29-cudnn-attention-nan-incident.md
+++ /dev/null
@@ -1,118 +0,0 @@
-# cuDNN Attention 全空 Mask 导致训练 NaN
-
-**日期**：2026-08-29
-
-**机器**：8×H200
-
-**配置**：`pi05_base_bi_flexiv_earbud_case_insertion_0826_h200`
-
-**影响范围**：从 clean step 16000 切换到 `use_cudnn_attention=true` 后的生产训练
-
-> [!CAUTION]
-> 2026-08-30 更新：本文的修复**有效但不完整**。它消除了这里描述的瞬时失效（第一步反向即产生 NaN），但带该修复的训练在另一台机器上仍会缓慢发散：前 1000 步与基线逐点吻合，随后在 step 1000–1200 之间分叉，约 4700 步后变为 NaN。本文「验证要求」一节的 200–500 步标准不足以发现它。见 [2026-08-30-cudnn-attention-divergence.md](2026-08-30-cudnn-attention-divergence.md)。
-
-## 摘要
-
-生产训练恢复后，第一个记录窗口已经出现 NaN 梯度，随后 loss 和参数持续为 NaN。根因不是 cuDNN 动态库版本、数据加载、显存或普通的梯度爆炸，而是 cuDNN fused Attention 在 attention mask 存在整行全 `False` 时的反向数值行为：前向输出有限，但 Q 梯度为 NaN。
-
-Pi0.5 的 prompt 会补齐到 `max_token_len=200`。padding token 在 `make_attn_mask` 中既不是有效 query，也不是有效 key，因此必然生成全 `False` query 行。原显式 Attention 用有限的 `big_neg` 处理这些行，不会产生同样的 Q 梯度 NaN；直接切换到 cuDNN 实现暴露了这一差异。
-
-训练已停止。当前最后一个确定干净且仍保留的 checkpoint 是 step 15000；step 16000 已被保留策略删除，step 20000–55000 均已污染。
-
-## 现象与时间线
-
-切换前的原 Attention 训练指标正常：
-
-```text
-Step 16000: grad_norm=0.6721, loss=0.1429, param_norm=1806.9792
-Step 16100: grad_norm=1.1132, loss=0.1456, param_norm=1806.9926
-Step 16200: grad_norm=0.8198, loss=0.1402, param_norm=1807.0051
-Step 16300: grad_norm=0.9735, loss=0.1410, param_norm=1807.0171
-```
-
-从 step 16000 checkpoint 用 cuDNN Attention 恢复后：
-
-```text
-Step 16100: grad_norm=nan, loss=0.1185, param_norm=nan
-Step 16200: grad_norm=nan, loss=nan,    param_norm=nan
-Step 16300 及以后: grad_norm=nan, loss=nan, param_norm=nan
-```
-
-checkpoint 目录名记录的是循环 step，而保存的 `train_state.step` 已经加一。因此恢复后第一个实际更新的前向 loss 仍有限，但该步反向产生 NaN 并污染参数；后续前向全部变为 NaN。
-
-日志又对整个窗口的所有字段使用 `nanmean`。第一个窗口中唯一有限的首步 loss 会让聚合结果仍显示为 `0.1185`，从而掩盖其余 NaN。grad norm 和 param norm 只在日志 step 真正计算，所以当时已经直接显示 NaN。
-
-## 最小复现
-
-在 JAX 实际加载 cuDNN runtime `91400` 的同一环境中，使用 BF16 GQA 和含 16 个全空 query 行的 mask：
-
-```text
-forward output finite: true
-Q/K/V gradient finite: [false, true, true]
-Q gradient NaN count: 16384
-```
-
-为每个全空 query 行临时开放第 0 个 key 后：
-
-```text
-Q/K/V gradient finite: [true, true, true]
-Q/K/V gradient NaN count: [0, 0, 0]
-valid query rows unchanged: true
-```
-
-这解释了完整训练的全部现象：首步前向正常、首步梯度 NaN、更新后参数 NaN、随后 loss NaN。
-
-## 修复
-
-cuDNN 分支保留原始 attention mask，不再给全空行增加 dummy key。调用 `jax.nn.dot_product_attention` 前用 `any(..., axis=-1)` 找出有效 query；Q 的前向值逐元素不变，但全空 query 行通过 `stop_gradient` 从反向链路中移除。这样保持原 mask 语义和有效 query 输出，同时把无效 Q 梯度固定为 0。
-
-训练循环曾增加逐 step 有限性保护。根据 2026-08-29 的生产运行要求，该保护已移除，恢复原更新路径；训练仍按原逻辑在日志边界输出 `loss`、`grad_norm` 和 `param_norm`，不再自动拒绝更新或因非有限值停止。最初曾把 `1.91 -> 2.3 s/step` 归因于该保护，后续 A/B 推翻了这一判断：移除保护后连续有限训练仍约为 `2.31 s/step`。
-
-这是一项明确的吞吐优先取舍：日志只能事后暴露数值异常，不能保证异常更新不会进入参数或下一个 checkpoint。
-
-## 验证要求
-
-修复验收按以下顺序执行：
-
-1. 单元测试确认 Q 前向逐元素不变，且只停止全空 query 行的梯度；
-2. H200/cuDNN 最小反向测试确认 Q/K/V 梯度全部有限；
-3. 从 step 15000 复制到独立诊断实验目录，禁用 W&B 和 checkpoint 保存；
-4. 至少运行 200–500 step，逐日志窗口确认 loss/grad norm/param norm 有限；
-5. 通过后才能清理污染 checkpoint 或恢复生产训练。
-
-## 2026-08-29 修复验证结果
-
-代码与单元回归：
-
-- `python -m compileall` 通过；
-- `git diff --check` 通过；
-- `src/openpi/models/pi0_test.py`：`8 passed`；
-- `scripts/train_test.py`：`2 passed`，包括两步训练和 checkpoint 恢复；
-- H200/cuDNN BF16 GQA 真实序列形状测试：原始反向产生 71680 个 Q 梯度 NaN；Q stop-gradient 后 Q/K/V 梯度全部有限，NaN 数量均为 0，前向值和有效 query 输出逐元素不变。
-
-真实训练使用 8×H200、全局 batch 512、cuDNN runtime 91400，从 clean step 15000 通过 `/tmp` 独立目录只读恢复。W&B、周期保存和最终保存均关闭。为缩短冷启动只启用 16 个 DataLoader worker，因此本次不统计吞吐。
-
-主动结束前完成了 step 15001–15060 共 60 个更新；每 10 step 的检查结果如下：
-
-```text
-Step 15010: grad_norm=0.7763, loss=0.1438, param_norm=1806.8403, update_is_finite=1.0000
-Step 15020: grad_norm=1.0680, loss=0.1666, param_norm=1806.8416, update_is_finite=1.0000
-Step 15030: grad_norm=0.8248, loss=0.1507, param_norm=1806.8430, update_is_finite=1.0000
-Step 15040: grad_norm=0.8843, loss=0.1576, param_norm=1806.8450, update_is_finite=1.0000
-Step 15050: grad_norm=0.7334, loss=0.1528, param_norm=1806.8467, update_is_finite=1.0000
-Step 15060: grad_norm=1.4659, loss=0.1655, param_norm=1806.8483, update_is_finite=1.0000
-```
-
-原缺陷会在恢复后的第一个更新产生 NaN，因此 60 个连续有限更新已经验证修复覆盖了已知故障。短跑结束后 8 张 GPU 显存均为 0 MiB，原 checkpoint 未被修改。恢复生产前仍建议完成 200–500 step 的长验证，以覆盖更多数据样本和训练阶段。
-
-### 速度复核
-
-上次长跑报告的普通 step `1.910590 s` 来自数值已经失效的区间：第一个生产日志窗口已出现 NaN 梯度，随后参数和 loss 持续为 NaN，因此不能作为有效训练的吞吐基线。修复后连续有限训练的普通 step 稳定约为 `2.31 s`。
-
-16-worker 短跑曾出现 `dispatch=1.82–1.83 s`，但它发生在 80 秒左右的数据断供之后，是 GPU 空闲后的瞬时样本；同一运行的连续 step 仍为 `2.30 s`，并且每约 16 step 出现 30–85 秒 `next_batch` stall。生产继续使用 96 workers 以避免数据断供，约 8 分钟冷启动不计入稳态吞吐。
-
-## 恢复注意事项
-
-污染的 step 20000–55000 已删除，生产训练已从 step 15000 恢复并保存 step 18000。逐 step 保护在 step 18600 报告最近 100 步中有 8 次更新被拒绝，当时窗口聚合 loss 仍为有限值 `0.7085`，因此只看 loss 不足以证明每次更新都有限。移除保护后的生产训练从最后保留的 step 18000 继续。
-
-原 step 16000 已删除，无法无损续接；从 step 15000 恢复意味着重跑约 1000 个有效 step。
diff --git a/docs/2026-08-30-cudnn-attention-divergence.md b/docs/2026-08-30-cudnn-attention-divergence.md
deleted file mode 100644
index cb20e98..0000000
--- a/docs/2026-08-30-cudnn-attention-divergence.md
+++ /dev/null
@@ -1,375 +0,0 @@
-# 启用 flash attention 提速：全层训练发散与分层稳定方案
-
-**日期**：2026-08-30
-
-**目标**：用 cuDNN fused attention（`use_cudnn_attention=true`）替代显式 attention，缩短训练时间。
-
-**机器**：8×H100 80GB（`cs-76d36-a14d7-server`）
-
-**状态**：**根因机制已定位，并得到一个通过 3000 步验收的保守修复。** 全 18 层 cuDNN 的约 1.5% 原始梯度差异具有不同于人造乘性噪声的坐标结构，经过 Adam 预条件后变成约 3.35% 的参数更新差异，并使约 0.79% 的更新坐标翻转符号；误差在早层产生时还会穿过后续网络放大。仅让最后 3/18 层使用 cuDNN，把更新差异降到 0.847%、符号翻转降到 0.193%，并从 step 0 稳定跑满 3000 步。该方案解决了发散，但实测步时约 `1.5 s/step`，只保留了小部分吞吐收益，尚未达到全层 cuDNN 的速度目标。
-
-## 为什么值得继续追
-
-同机同任务的实测步时：
-
-| 实现 | 步时 | 来源 |
-|---|---|---|
-| 显式 attention | `1.645 s/step` | 0828 run，98714 s / 60000 步，含日志与 checkpoint 摊销 |
-| cuDNN + stop-gradient mask | 约 `1.05 s/step` | 0829 run 的 `[TIMING]` 普通步 |
-| cuDNN + dummy-key mask | 约 `1.38 s/step` | 2026-08-30 诊断 run |
-
-按 60000 步计算，`1.645 -> 1.05` 约省 **9.9 小时**。
-
-> [!IMPORTANT]
-> `1.05 s/step` 取自一个**已经在发散**的 run。数值失效不改变 kernel 的执行成本，所以该数字作为吞吐参考是可用的；但正式验收时必须在一个不发散的 cuDNN run 上重测。
-
-## 失效现象
-
-`use_cudnn_attention=true` 时，训练在前约 1000 步与已验证基线逐点吻合，随后偏离并单调发散，最终变为 NaN。
-
-```text
-step    0828 显式（基线）      0829 cuDNN 9.14        cuDNN 9.19
- 200    0.9244 /  2.89        0.9333 /  3.58         0.9348 /  3.53
-1000    0.4220 /  2.77        0.4306 /  2.50         0.4319 /  3.16
-1100    0.4077 /  1.79        0.4784 /  3.16         0.4723 /  3.37
-1200    0.3742 /  2.12        0.6484 / 31.40         0.6600 / 35.49
-2500    0.2457 /  1.33        5.7692 / 675.08        4.7055 / 759.18
-5900    0.1763 /  1.36        nan                    （已停止）
-```
-
-**因果起点是 step 1000–1200，不是 NaN 出现的 step 5900。** 后者只是累积发散溢出的位置。这个区别把复现成本从 27 小时降到约 35 分钟，是本次调查最有操作价值的一条。
-
-### 与 2026-08-29 事故是两种不同的失效
-
-| | 2026-08-29（H200） | 2026-08-30（H100） |
-|---|---|---|
-| 现象 | **瞬时**：step 16300 正常 → 16400 全 NaN | **缓慢**：贴合 1000 步后发散 4700 步 |
-| 根因 | 全空 mask 行的 Q 梯度 NaN | 数值扰动 + 配方裕度不足 |
-| 修复 | `_stop_gradient_for_fully_masked_queries` | 未解决 |
-
-全空 mask 修复是有效的，它消除了第一种模式。当前问题带着该修复仍然发生。
-
-## 责任归属：单变量 A/B
-
-同机、同数据集、同超参、同 seed，唯一变量是 `use_cudnn_attention`：
-
-| | `cve5hiyt`（0828） | `orbnh1b7`（0829） |
-|---|---|---|
-| commit | `d0aec27` | `51d9af3` |
-| `use_cudnn_attention` | 字段不存在（显式路径） | `true` |
-| 结果 | step 59900 **loss 0.0589** | step 5900 **NaN** |
-
-`git diff d0aec27 HEAD -- src/ scripts/` 共 72 行增删，全部属于该开关本身。排除 LR、数据、配置。
-
-## 结论：cuDNN 内核是正确的
-
-### 完整模型梯度 A/B
-
-同一份 pi05_base 权重、同一个真实 batch、同一个 rng，唯一变量是开关（`scripts/grad_ab_probe.py`，cuDNN 9.19，batch 32，FSDP 8 卡）：
-
-```text
-batch 0   loss 显式=5.859448  cuDNN=5.851340   |grad| 104.31 vs 103.65   ratio=0.9937
-batch 1                                        |grad|  67.19 vs  67.45   ratio=1.0038
-batch 2                                        |grad|  70.71 vs  70.61   ratio=0.9986
-
-相对梯度差：1.75e-02 / 1.66e-02 / 2.12e-02
-按模块：PaliGemma 1.87e-02  action_in_proj 1.66e-02
-        time_mlp_in 1.46e-02  time_mlp_out 1.20e-02  action_out_proj 5.87e-03
-```
-
-**相对差约 2%，梯度范数比值 0.994–1.004，没有任何模块偏离整体水平。** 这是 bf16 误差经 18 层 Gemma 反向累积的正常量级——单层测量中 cuDNN 对 fp32 参考解的相对误差是 `5.4e-03`、显式 bf16 路径是 `3.2e-03`，两者同阶。cuDNN 给出的不是更错的答案，只是另一个同样合法的 bf16 近似。
-
-### 扰动的噪声/确定性分解
-
-在同一 batch 上重复调用 cuDNN 反向，把差异拆开：
-
-```text
-batch    总差异(cuDNN vs 显式)   随机(cuDNN vs cuDNN)   确定性分量
-  0          1.6993e-02             9.7433e-03          1.3922e-02
-  1          1.7036e-02             8.7706e-03          1.4605e-02
-  2          2.2268e-02             9.1202e-03          2.0315e-02
-```
-
-随机部分约 `0.9%`，来自 cuDNN 反向的 atomic 累加（同一输入两次反向 bitwise 不同）。确定性部分约 `1.5%`，是主要成分。
-
-注意「确定性」在这里的含义：**给定 batch 可复现**，不等于**跨 batch 方向一致**。换一个 batch，舍入差异的方向随之改变；梯度范数比值为 1.00，也看不出系统性缩放。后续多 batch 实验已直接否定固定方向假设，见[跨 batch 固定方向假设不成立](#跨-batch-固定方向假设不成立)。
-
-## 已排除的候选
-
-全部在本机、真实训练形状下实测：`T=1018`（768 图像 + 200 prompt + 50 action）、`num_heads=8`、`num_kv_heads=1`、`head_dim=256`、bf16、含全空行的真实 block mask。参考解为同一显式路径在 fp32 下的结果。
-
-- **cuDNN 版本**：9.14 与 9.19 的发散曲线逐点重合（step 1200 分别为 31.40 与 35.49）。
-- **「静态序列长度非 128 倍数导致 SDPA backward 错误」**（cuDNN 9.15.1 release note）：
-  ```text
-  T=1018 (%128=122)   dQ rel-L2: 显式 3.23e-03   cuDNN 3.21e-03   NaN=0
-  T=1024 (%128=0)     dQ rel-L2: 显式 3.23e-03   cuDNN 3.21e-03   NaN=0
-  ```
-  两者无差异，均无 NaN。**该假设在本机本形状上不成立。**
-- **库打包错配**：环境统一为单源 pip 栈后（见下），照样发散。
-- **logit 幅度**：std 从 1 扫到 80，cuDNN 最差 `7.2e-03`、显式 `3.5e-03`，同量级。
-- **head_dim=256**（Hopper 上 cuDNN ≥9.5 才放开的新路径）：与 head_dim=128 表现一致。
-- **cotangent 稠密度**：只在 action token 上非零 与 在全部有效 token 上稠密，结果相同。
-- **dBias 反向路径**：`should_export_dbias` 要求 `bias.shape[0]==1`，实际 batch 为 256，不进入该路径。
-- **分片**：`activation_sharding_constraint` 使用 `PartitionSpec(DATA_AXIS)`，只按 batch 分片，attention 拿到完整序列。
-- **两版 mask 修复**：c1b70b2 给全空行开 dummy key，51d9af3 保留原 mask 只对 q 做 stop_gradient。在长/短两种 padding 密度下，两者在有效行上给出**逐位相同**的 dQ/dK/dV。且事故记录显示 c1b70b2 版本同样发散：从 step 15000 的 loss `0.145` 涨到 step 18600 的 `0.7085`，逐 step 保护在那里报告最近 100 步有 8 次更新非有限。**两版修复都不是原因，也都不是解药。**
-- **降低 `peak_lr`**：`2.5e-5 -> 1.5e-5` 只把发散起点从 step 约 1150 推迟到约 1450：
-  ```text
-  step   cuDNN @1.5e-5
-  1200   0.3954 /  2.95
-  1400   0.4089 /  2.84
-  1500   0.5155 /  7.17
-  1700   0.9654 / 26.85
-  1900   1.1773 / 88.68
-  ```
-  **LR 是速率参数，不是开关。**
-- **提高 `optimizer.b2`**：`0.95 -> 0.99` 未延后发散起点。2026-08-30 在统一后的 cuDNN 9.19 单源环境中，从 step 0 跑满 3000 个更新；8×H100、batch 256，W&B、周期保存和最终保存均关闭。启动前 `scripts/check_cuda_stack.py` 判定 PASS。所有已输出的指标均为有限值，但曲线仍在 step 1100–1200 分叉：
-  ```text
-  step    loss       grad_norm    param_norm
-   200    0.9082        3.1334     1802.3866
-  1000    0.4321        3.7520     1802.5017
-  1100    0.5189        4.3141     1802.5402
-  1200    0.9354       44.6635     1802.5638
-  1300    1.6713       52.6050     1802.5797
-  1500    2.2757       30.3204     1802.6188
-  2000    3.7292      193.8313     1802.7924
-  2500    4.7271      422.8436     1803.0002
-  2700    6.4255     1521.7623     1803.0616
-  2900    4.5535      819.7691     1803.1382
-  ```
-  训练完整执行了 3000 个更新，最后一个正式指标窗口是 step 2900。step 1200 的显式基线为 `loss=0.3742 / grad_norm=2.12`，因此新 run 已超过 3× 梯度验收上限约 7 倍。step 2500 的 loss `4.7271` 又几乎复现了 `b2=0.95` cuDNN run 的 `4.7055`；较高 `b2` 只降低了该点的梯度峰值（`422.84` vs `759.18`），没有改变发散轨迹。普通非日志 step 的 dispatch 约 `1.29–1.39 s`，无数据断供或系统错误。实验结束后 8 张 GPU 均归零，实验目录为空，没有污染 checkpoint。**提高二阶矩窗口不是解法。**
-- **显式 attention + 1.5% 逐 step 独立随机梯度扰动**：没有发散。2026-08-31 在训练配置中加入默认关闭的 `gradient_noise_scale` 诊断开关；显式梯度为 `g`，先生成 `r = g * N(0, I)`，再令 `delta = 0.015 * ||g|| / ||r|| * r`，优化器实际接收 `g + delta`。模型 RNG 保持原样，扰动 RNG 独立并按 step 可复现，因此每一步都精确满足 `||delta|| / ||g|| = 0.015`，同时扰动能量按原梯度能量分布，而不是按参数量集中到最大张量。
-
-  8×H100、batch 256、显式 attention、原 `b2=0.95`，从 step 0 跑满 3000 个更新；W&B、周期保存和最终保存均关闭。正式运行前，CPU helper/配置回归 `58 passed`，生产 batch 的 2-step GPU 校准也连续报告 `gradient_noise_rel_norm=0.0150`。正式曲线：
-  ```text
-  step    loss       applied grad_norm    clean grad_norm    noise rel-L2
-   100    2.8529          3.9906              3.9903           0.0150
-   200    0.9262          4.0429              4.0433           0.0150
-   500    0.5704          2.0811              2.0806           0.0150
-   900    0.4335          2.8625              2.8631           0.0150
-  2200    0.2698          1.7088              1.7083           0.0150
-  2500    0.2476          1.8600              1.8613           0.0150
-  2700    0.2387          1.5861              1.5854           0.0150
-  2900    0.2411          1.6389              1.6389           0.0150
-  ```
-  训练循环耗时 `1:22:45`，另有首次数据冷启动约 6.5 分钟；普通 step 的 dispatch 约 `1.57–1.64 s`。最后一个正式指标窗口是 step 2900。step 2500 的 loss `0.2476` 与历史显式基线 `0.2457` 基本重合，而同期 cuDNN 9.19 已经是 `4.7055 / 759.18`。实验结束后 8 张 GPU 均归零，目录为空，没有污染 checkpoint。
-
-  **结论**：配方并非对任意 1.5% 梯度扰动都没有裕度。该结果排除的是「逐 step 独立、零均值、按梯度能量加权」这一类随机扰动；它尚不能排除真实 cuDNN 差分具有跨 batch 的固定投影、时间相关性或不同的参数内结构。
-
-### 两个 mask 变体：等价，但一个明显更快
-
-`_stop_gradient_for_fully_masked_queries`（51d9af3）与 `_make_cudnn_attention_mask_safe`（c1b70b2）数值等价，但 dummy-key 版本要在**每一层**内重建整个 `(B, 1, T, S)` 掩码，batch 256 下实测 `1.38 s/step`，stop-gradient 版本约 `1.05 s/step`。H200 文档记录的 `1.91 -> 2.31 s/step` 回退是同一件事。
-
-因此代码保留 stop-gradient 版本。51d9af3 的标题「Preserve the **fast** cuDNN mask」描述的正是这个性能取舍，它不是缺陷来源。
-
-## 2026-08-31 后续调查：根因机制与稳定折中
-
-### 跨 batch 固定方向假设不成立
-
-`scripts/grad_ab_direction_probe.py` 对 8 个真实 batch 分别计算显式梯度和 3 次 cuDNN 反向的均值，再比较 `delta_b = g_cudnn,b - g_explicit,b`：
-
-```text
-跨 batch delta 两两 cosine：-0.004043 ± 0.126301，范围 [-0.247516, +0.226915]，n=28
-||delta|| / ||g_explicit||：  1.556024e-02 ± 2.76e-03
-cuDNN 同 batch 随机分量：    8.752676e-03
-cos(delta, gradient)：       -0.00113 ± 0.225
-cos(delta, parameter)：      +0.000016
-```
-
-差分没有稳定的跨 batch 方向，也不沿当前梯度或参数方向累积。因此没有继续做「显式 attention + 固定方向人造噪声」；该实验不会复现真实差分结构。
-
-### batch 到达顺序不是原因
-
-`strict_batch_order=true` 的全 18 层 cuDNN run 仍在原位置分叉：
-
-```text
-step 1000   loss=0.4304   grad_norm=2.3201
-step 1100   loss=0.4593   grad_norm=2.1643
-step 1200   loss=0.6034   grad_norm=15.4116
-```
-
-显式基线 step 1200 为约 `0.3742 / 2.12`。固定 sampler 输出顺序没有改变失效，因此 DataLoader worker 完成时序不是根因。后续诊断均保留 `strict_batch_order=true`，确保曲线可逐点比较。
-
-### 关键机制：Adam 放大了坐标级差异
-
-`scripts/optimizer_update_ab_probe.py` 从上述严格顺序 run 的 step 1000 checkpoint 恢复真实 Adam 状态，对同一真实 batch、同一 rng 计算实际优化器更新。原始梯度 L2 差异仍只有约 1.5%，但进入 Adam 后：
-
-```text
-                         相对更新差异     更新 cosine     符号翻转坐标
-真实全层 cuDNN 差分          3.3521%         0.999440        0.7885%
-旧的 1.5% 人造乘性噪声       0.7313%         --              0.1459%
-```
-
-真实差分造成的 Adam 更新差异是稳定人造噪声的 **4.6 倍**，符号翻转比例是 **5.4 倍**。主导部分在 PaliGemma（更新差异 3.353%、符号翻转 0.789%）。这解释了为何「相同梯度相对 L2」的人造噪声稳定，而 cuDNN 会发散：旧噪声按 `|g|` 分配能量，几乎不扰动小梯度/小二阶矩坐标；真实舍入差分具有不同的参数内结构，会被 Adam 的逐坐标归一化放大。
-
-把 Adam `eps` 从 `1e-8` 提高到 `1e-6`/`1e-5`，真实更新相对差异只从 3.35% 降到 3.22%/3.08%，符号翻转仍约 0.79%，不是有效修复。
-
-### 优化器时间平滑只能延迟，不能修复
-
-将 `b1` 从 `0.9` 提高到 `0.99`（`b2` 保持 `0.95`）后，全层 cuDNN 在 step 1200 暂时健康，但随后同样发散：
-
-```text
-step 1200   loss=0.4027   grad_norm=2.1161
-step 1400   loss=0.4950   grad_norm=6.2470
-step 1500   loss=0.7486   grad_norm=17.0
-step 2000   loss=2.749    grad_norm=67
-step 2900   loss=8.844    grad_norm=562
-```
-
-它只把分叉推迟约 300 步，和降低 LR、提高 `b2` 的结果一致：改变累积速率，不改变误差来源。
-
-### 层位置比层数更重要
-
-为 `Pi0Config` 增加 `cudnn_attention_layer_start` 和 `cudnn_attention_num_layers`，可只在连续的 Gemma 层段使用 cuDNN。其余层仍走显式 attention；默认值保持历史上的全层行为。CPU 模型回归 `8 passed`。
-
-在 step 1000 的真实 Adam 状态上做更新空间 A/B：
-
-```text
-cuDNN 层段       相对更新差异     更新 cosine     符号翻转坐标
-全 18 层            3.3118%         0.999451        0.7769%
-最前 6 层           3.2061%         0.999486        0.7548%
-中间 6 层           2.2169%         0.999754        0.5141%
-最后 6 层           1.1633%         0.999932        0.2720%
-最前 3 层           3.3016%         0.999455        0.7720%
-中间 3 层           1.9894%         0.999802        0.4634%
-最后 3 层           0.8470%         0.999964        0.1926%
-稳定人造噪声参考     0.7313%         --              0.1459%
-```
-
-只融合最前 3 层几乎等同全 18 层，说明早层的微小扰动会经后续 15 层反向/前向传播放大；不能按融合层数线性估计风险。最后 3 层最接近已经跑稳的人造噪声参考，因此选作首个训练候选。
-
-### 最后 3 层方案跑满 3000 步
-
-配置：`configs/diag_cudnn_last3.yaml`，8×H100、batch 256、`strict_batch_order=true`、原始 Adam（`b1=0.9, b2=0.95`）、从 pi05_base step 0 开始、W&B 和 checkpoint 保存关闭。层索引 15–17 使用 cuDNN，其余 15 层显式。正式曲线：
-
-```text
-step    loss      grad_norm    param_norm
-   0    6.4029      65.5204     1802.3865
- 100    2.8439       5.1838     1802.3851
- 500    0.5719       2.4526     1802.4014
- 900    0.4349       2.5621     1802.4742
-1000    0.4240       2.3789     1802.5061
-1100    0.4050       2.1847     1802.5416
-1200    0.3757       2.3869     1802.5775
-1400    0.3390       1.8067     1802.6504
-1600    0.3238       2.6505     1802.7225
-2000    0.2793       1.8254     1802.8666
-2500    0.2463       2.0116     1803.0460
-2900    0.2424       1.4769     1803.1906
-```
-
-训练完整执行 3000 个更新；最后一个正式指标窗口是 step 2900。原全层 cuDNN 在 step 1200 已到 `0.6034 / 15.41`，b1=0.99 延迟方案在 step 2000 已到 `2.749 / 67`；本方案的曲线与显式稳定基线一致，满足 3000 步验收。
-
-训练循环耗时 `1:18:02`，普通步约 `1.5 s/step`；另有 404 秒的 64-worker 首批冷启动。相对本任务显式路径约 `1.57–1.64 s/step` 只有约 4–9% 加速，明显低于全层 cuDNN 的收益。**因此这是数值稳定的保守折中，不是最终吞吐目标。** 按用户指示，本轮到此停止，不继续运行最后 6 层或更大范围的训练。
-
-## 2026-08-31 调查计划原稿（结果见上）
-
-以下保留当时的计划作为调查轨迹。第 1 项已执行并否定固定方向；因此第 2 项不再能复现真实差分，未执行。实际随后转向 Adam 更新空间和分层位置分析，并得到上面的最后 3 层方案。
-
-### 1. 测量真实 cuDNN 差分的跨 batch 方向一致性
-
-扩展 `scripts/grad_ab_probe.py`，对多个真实 batch 计算 `delta_b = g_cudnn,b - g_explicit,b`，并报告：
-
-- 不同 batch 的 `delta_b` 两两 cosine；
-- `delta_b` 在显式梯度、参数向量和前一 batch 差分上的投影；
-- 上述量按 PaliGemma、action projection 和 time MLP 分模块的结果；
-- 同一 batch 多次 cuDNN 反向取平均后再计算，尽量剥离约 0.9% 的 atomic 随机分量。
-
-若跨 batch cosine 或固定投影显著非零，说明真实差分存在可累积方向；若接近零，则需要检查人造噪声与 cuDNN 差分在参数内分布、尾部和时间相关性上的差别。
-
-### 2. 显式 attention + 时间相关的固定方向扰动
-
-复用 `gradient_noise_scale=0.015`，但让随机方向不再按 step 变化，每一步只根据当前 `||g||` 重新缩放同一参数坐标方向。若这一版本发散而逐 step 独立版本不发散，时间相关性就是关键杠杆；若仍稳定，则必须更直接地复现真实 cuDNN 差分结构。
-
-### 3. `optimizer.weight_decay` `1e-10 -> 1e-4` 量级
-
-当前几乎为零。加大给慢漂移一个回复力。改动比前两项大，收敛行为也会变，排在后面。
-
-### 4. 若以上都不成立：绕开 attention 数值的提速路径
-
-- `max_token_len` `200 -> 168`（全数据扫描得到的 tokenized state 最大长度是 165，留 3 token 余量）：减少约 3% 的序列长度，不触碰数值；
-- 重新采集 XProf trace，确认剩余最大的 SigLIP/Gemma GEMM 与 fusion；
-- 提高 `save_interval`，减少 checkpoint 对下一步的数据等待尖峰。
-
-这三项合计的收益远小于 flash attention 的 36%，但零风险。
-
-### 不要重复的弯路
-
-- **不要再查 cuDNN 版本或库打包。** 9.14 与 9.19 表现一致，单源栈已验证。
-- **不要再改 mask 处理。** 两个变体逐位等价，且都发散。
-- **不要只调 `peak_lr`。** 按累积机制，任何取值大概率都只是推迟。
-- **不要再单独提高 `optimizer.b2`。** `0.99` 从 step 0 跑满 3000 个更新，仍在 step 1100–1200 分叉，并在 step 2700 出现 `grad_norm=1521.76`。
-- **不要重复逐 step 独立的 1.5% 随机梯度扰动。** 显式 attention 已跑满 3000 步，step 2500 loss `0.2476`，没有复现发散；下一步应测试真实差分的方向结构或时间相关扰动。
-- **不要跨 run 比较 `grad_norm`。** 本仓库默认 `strict_batch_order: false`，各 run 的 batch 序列不同，早期 `grad_norm` 差异来自数据顺序而非实现差异。曾据此误判 cuDNN 存在系统性梯度放大，被受控 A/B（范数比 1.00）否定。
-
-## 验收要求
-
-> [!WARNING]
-> 本次曾在 step 1300 仍健康时判定「降 LR 已解决」，随后 step 1450 开始发散。**判据必须是跑满预定步数并看完整条曲线，不能是「越过某一步还正常」。**
-
-启用 `use_cudnn_attention` 的验收：
-
-1. 从 **step 0** 起跑满 **3000 步**（最后 3 层方案实测训练循环约 78 分钟，另有约 6.7 分钟首次数据冷启动）；
-2. 与同配置显式 attention 的基线曲线逐点对比 step 200/500/1000/1500/2000/3000 的 `loss`；
-3. 全程 `grad_norm` 不得超过基线同点的 3 倍；
-4. 通过后才能进入长跑，且长跑期间仍需按日志窗口盯 `loss` 与 `grad_norm`。
-
-40 步吞吐短跑、从已收敛 checkpoint 出发的短跑、以及只验证「梯度有限」的 smoke test，都**不是**有效的正确性验收。2026-08-29 的修复通过了全部这三类检查，仍然发散。
-
-建议增加发散护栏：`grad_norm` 已在日志 step 计算，当它连续多个日志窗口超过滚动中位数的固定倍数时中止训练。本次若有该护栏，会在 step 1300 停下而不是跑满 60000 步。
-
-## 环境
-
-### dispatcher 与 engine 版本错配（已修复）
-
-事故期间 `cuda_versions.cudnn_get_version()` 返回 `91400`，但进程实际加载的是 **9.10.2 的 dispatcher stub 驱动 9.14.0 的引擎库**：
-
-```text
-site-packages/nvidia/cudnn/lib/libcudnn.so.9          <- 125 KB，dispatcher stub，9.10.2 wheel
-$CONDA_PREFIX/lib/libcudnn_graph.so.9.14.0            <- 真正执行的引擎库
-$CONDA_PREFIX/lib/libcudnn_engines_precompiled.so.9.14.0
-$CONDA_PREFIX/lib/libcudnn_ops.so.9.14.0
-```
-
-`cudnnGetVersion()` 的返回值来自引擎库，因此它**无法**发现这种错配。上一份文档要求「必须确认输出 91400」的验证方法是不充分的。正确做法是枚举进程实际映射的库：
-
-```bash
-python scripts/check_cuda_stack.py
-```
-
-它遍历 `/proc/self/maps`，在 `libcudnn*` 来自多个目录、或文件名版本与运行时报告不一致时判 FAIL，然后在真实训练形状上跑一次前向+反向并与 fp32 参考解对比。**在混装环境下该脚本判 FAIL，而三项数值检查全部通过**——这正是只看数值 smoke test 会漏掉的部分。
-
-### 统一到 pip cu128 栈（已完成）
-
-目标：CUDA Toolkit 12.8 + pip cu128 PyTorch + pip cuDNN 9.19，移除 conda cuDNN。
-
-依赖解析上的关键事实：
-
-- `torch 2.10.0+cu128` 把 cuDNN 钉死在 `nvidia-cudnn-cu12==9.10.2.21`（`==`，非 `>=`）；
-- `torch 2.11.0+cu128` 钉 `nvidia-cudnn-cu12==9.19.0.56` 和 `cuda-toolkit==12.8.1`；
-- PyPI 默认的 `torch 2.11.0` 是 **CUDA 13 构建**，钉 `nvidia-cudnn-cu13`，必须走 cu128 索引；
-- `jax-cuda12-plugin` 只要求 `nvidia-cudnn-cu12>=9.1,<10.0`，不构成约束。
-
-**torch 版本决定了训练进程加载哪个 cuDNN**，尽管 torch 在本仓库只用于 DataLoader。这是 `pyproject.toml` 把 torch 下界提到 `>=2.11` 的原因。
-
-另外两处容易漏掉的：
-
-- `torchcodec` 的 `+cu128` 构建只存在于 pytorch 索引，默认 PyPI 上的同版本号没有 `+cu128` 后缀，需在 `[tool.uv.sources]` 中路由；
-- 本地 fork `lerobot-xense`（editable 安装）在自己的 `pyproject.toml` 中声明 `torch<2.11.0`、`torchcodec<0.11.0`、`torchvision<0.26.0`，**必须同步放宽**，否则解析失败。不能改用 PyPI 的 lerobot 0.6——那会丢掉 fork 中的 xense 相机与 flexiv 机器人代码。
-
-启动命令必须**去掉** `LD_LIBRARY_PATH="$CONDA_PREFIX/lib"`。移除 conda cuDNN 后该目录仍保留 conda-forge 的 `libcublas`/`libcudart`/`libnccl` 等，继续导出会让它们遮蔽 pip cu128 的同名库，等于换一个库重演同一类错配。
-
-## 处置
-
-- 0829 run 在 step 5900 之后的所有 checkpoint（10000–50000）均为 NaN，无保留价值；
-- `cve5hiyt`（0828，显式 attention）仍是已跑满 60000 步的生产结果，loss `0.0589`；
-- 最后 3 层 cuDNN 已通过 3000 步短程验收，但尚未跑 60000 步，不能把它等同于完整生产验收；
-- 需要绝对稳妥时继续使用 `use_cudnn_attention: false`；需要先获取小幅吞吐收益时，可用 `use_cudnn_attention: true, cudnn_attention_layer_start: 15, cudnn_attention_num_layers: 3` 进入受监控长跑；
-- 不要把 `use_cudnn_attention: true` 且 `cudnn_attention_num_layers: null`（全 18 层）用于生产训练。
-
-## 相关脚本
-
-- `scripts/check_cuda_stack.py` —— 验证进程实际加载的 cuDNN 栈是否单源，并在真实训练形状上做前向/反向数值检查；
-- `scripts/grad_ab_probe.py` —— 完整模型梯度 A/B，比较显式与 cuDNN 的逐参数梯度，并把差异拆成随机与确定性分量；
-- `scripts/grad_ab_direction_probe.py` —— 多真实 batch 的 cuDNN 差分方向、投影和随机分量；
-- `scripts/optimizer_update_ab_probe.py` —— 从真实 checkpoint 恢复 Adam 状态，比较梯度差分进入优化器后的更新误差与符号翻转；
-- `scripts/hybrid_update_ab_probe.py` —— 比较不同 cuDNN 层段的 Adam 更新空间误差。
diff --git a/docs/2026-08-31-cudnn-sdpa-nan-summary.md b/docs/2026-08-31-cudnn-sdpa-nan-summary.md
index 5a3d0db..fc8fbb6 100644
--- a/docs/2026-08-31-cudnn-sdpa-nan-summary.md
+++ b/docs/2026-08-31-cudnn-sdpa-nan-summary.md
@@ -1,28 +1,487 @@
-# cuDNN fused Attention NaN 问题概述
+# cuDNN fused Attention NaN 与训练发散：统一事故报告
 
-## 问题原因
+**覆盖日期**：2026-08-29 至 2026-08-31
 
-训练启用 JAX cuDNN fused Attention 后，前向 loss 起初正常，但反向传播先产生 NaN 梯度，随后参数和 loss 被持续污染。当前运行时为 cuDNN 9.14。
+**涉及机器**：8×H200、8×H100 80GB
 
-NVIDIA 已确认，旧版 cuDNN SDPA backward 在静态序列长度不是 128 倍数时可能产生错误结果或 NaN，并在 cuDNN 9.15.1 中修复。当前模型的 Attention 序列长度约为：
+**涉及实现**：JAX `dot_product_attention(..., implementation="cudnn")`、BF16 GQA
+
+**最终状态**：全空 mask 的瞬时 NaN 已修复；全 18 层 cuDNN 的缓慢训练发散仍未解决。生产训练不得启用全 18 层 cuDNN。
+
+> [!CAUTION]
+> 本次调查包含两种不同的失效，不能混为一个 cuDNN bug：
+>
+> 1. **瞬时 NaN**：全空 attention mask query 行使旧环境中的 cuDNN backward 产生 NaN Q 梯度。`_stop_gradient_for_fully_masked_queries` 已修复这一问题。
+> 2. **缓慢发散**：修复全空行后，全 18 层 cuDNN 仍在 step 1000–1200 与显式 attention 基线分叉，随后梯度爆炸并最终变为 NaN。根因是 BF16 实现差异的坐标结构被 Adam 放大，而不是非 128 倍序列长度、cuDNN 版本或 mask 修复。
+
+> [!IMPORTANT]
+> 截至本文最后一次验证：
+>
+> - **稳妥生产方案**：`use_cudnn_attention: false`；
+> - **已通过 3000 步短程验收的折中方案**：只在 Gemma 最后 3 层启用 cuDNN；
+> - **禁止用于生产**：全 18 层 cuDNN，即 `use_cudnn_attention: true` 且 `cudnn_attention_num_layers: null`；
+> - 已经出现 NaN 的参数和 Adam 状态不可修复，必须从确定干净的 checkpoint 恢复。
+
+## 1. 背景与真实训练形状
+
+Pi0.5 训练时的 Gemma attention 序列由三部分组成：
 
 ```text
 3 × 256 个图像 token + 200 个文本 token + 50 个 action token = 1018
-1018 % 128 = 122
+num_heads=8, num_kv_heads=1, head_dim=256, dtype=bf16
+```
+
+prompt 被补齐到 `max_token_len=200`。padding token 既不是有效 query，也不是有效 key，因此 `make_attn_mask` 会生成整行全 `False` 的 query 行。训练使用 block mask；只有 action token 上的 cotangent 非零。
+
+显式 attention 路径执行 QK einsum、有限 `big_neg` mask、softmax 和 PV einsum。cuDNN 路径调用：
+
+```python
+jax.nn.dot_product_attention(
+    q,
+    k,
+    v,
+    mask=attn_mask,
+    scale=1.0,
+    implementation="cudnn",
+)
+```
+
+`q` 已在调用前按 head dimension 缩放，因此必须使用 `scale=1.0`，避免重复缩放。
+
+## 2. 第一类失效：全空 mask 行导致瞬时 NaN
+
+### 2.1 事故范围
+
+**日期**：2026-08-29
+
+**机器**：8×H200
+
+**配置**：`pi05_base_bi_flexiv_earbud_case_insertion_0826_h200`
+
+**触发方式**：从 clean step 16000 切换到 `use_cudnn_attention=true` 后恢复生产训练。
+
+切换前显式 attention 指标正常：
+
+```text
+Step 16000: grad_norm=0.6721, loss=0.1429, param_norm=1806.9792
+Step 16100: grad_norm=1.1132, loss=0.1456, param_norm=1806.9926
+Step 16200: grad_norm=0.8198, loss=0.1402, param_norm=1807.0051
+Step 16300: grad_norm=0.9735, loss=0.1410, param_norm=1807.0171
+```
+
+从 checkpoint 切到 cuDNN 后：
+
+```text
+Step 16100: grad_norm=nan, loss=0.1185, param_norm=nan
+Step 16200: grad_norm=nan, loss=nan,    param_norm=nan
+Step 16300 及以后: grad_norm=nan, loss=nan, param_norm=nan
+```
+
+checkpoint 目录名记录循环 step，而保存的 `train_state.step` 已经加一。恢复后的第一个前向 loss 仍有限，但同一步反向产生 NaN 并污染参数，之后所有前向均为 NaN。
+
+日志窗口对 loss 使用 `nanmean`，第一个窗口里唯一有限的首步 loss 使聚合结果仍显示 `0.1185`，掩盖了该窗口其余 NaN。因此“窗口 loss 有限”不能证明窗口内每个更新都有限。
+
+### 2.2 最小复现与根因
+
+在当时实际加载 cuDNN runtime `91400` 的环境中，使用 BF16 GQA 和含 16 个全空 query 行的 mask：
+
+```text
+forward output finite: true
+Q/K/V gradient finite: [false, true, true]
+Q gradient NaN count: 16384
+```
+
+为全空 query 行临时开放一个 dummy key 后：
+
+```text
+Q/K/V gradient finite: [true, true, true]
+Q/K/V gradient NaN count: [0, 0, 0]
+valid query rows unchanged: true
+```
+
+完整生产形状也复现了相同现象：原始 backward 产生 71680 个 Q 梯度 NaN，K/V 梯度有限。
+
+### 2.3 实际修复
+
+最终实现保留原始 attention mask，不为全空行打开 dummy key。调用 cuDNN 前找出至少存在一个有效 key 的 query 行，并只停止全空 query 行的 Q 梯度：
+
+```python
+query_has_key = jnp.any(attn_mask, axis=-1)[:, 0, :, None, None]
+q = jnp.where(query_has_key, q, jax.lax.stop_gradient(q))
+```
+
+该变换具有以下语义：
+
+- Q 的前向值逐元素不变；
+- 有效 query 的输出和梯度不变；
+- 无效 query 的 Q 梯度固定为 0；
+- 原始 block mask 完整保留。
+
+dummy-key 版本在有效行上与 stop-gradient 版本给出逐位相同的 dQ/dK/dV，但每层重建 `(B, 1, T, S)` mask，batch 256 下约为 `1.38 s/step`；stop-gradient 版本约为 `1.05 s/step`，因此保留后者。
+
+### 2.4 修复验证及其局限
+
+当时完成的验证包括：
+
+- `src/openpi/models/pi0_test.py`：`8 passed`；
+- `scripts/train_test.py`：`2 passed`；
+- H200/cuDNN BF16 GQA 真实形状 Q/K/V 梯度全部有限；
+- 从 clean step 15000 独立恢复，连续执行 step 15001–15060；
+- 每 10 步的 loss、grad norm、param norm 和 `update_is_finite` 均有限。
+
+这足以证明第一类“第一步反向立即 NaN”已被修复，但不足以证明训练长期收敛。后续事实表明，60 步、200–500 步、已收敛 checkpoint 短跑和单次有限性 smoke test 都会漏掉第二类失效。
+
+### 2.5 checkpoint 处置
+
+H200 事故中最后一个确定干净且仍保留的 checkpoint 是 step 15000。原 step 16000 已被保留策略删除，step 20000–55000 受到 NaN 污染并已删除。从 step 15000 恢复意味着重跑约 1000 个有效 step。
+
+已经包含 NaN 的参数、Adam 一阶矩或二阶矩不能通过更换 attention 实现恢复；不得继续使用污染 checkpoint。
+
+## 3. 第二类失效：全层 cuDNN 在约 1000 步后缓慢发散
+
+### 3.1 现象
+
+全空 mask 修复后，全 18 层 cuDNN 在前约 1000 步与显式 attention 基线贴合，随后在 step 1000–1200 分叉，梯度持续放大，最终变为 NaN。
+
+```text
+step    显式基线             cuDNN 9.14            cuDNN 9.19
+ 200    0.9244 /   2.89      0.9333 /   3.58       0.9348 /   3.53
+1000    0.4220 /   2.77      0.4306 /   2.50       0.4319 /   3.16
+1100    0.4077 /   1.79      0.4784 /   3.16       0.4723 /   3.37
+1200    0.3742 /   2.12      0.6484 /  31.40       0.6600 /  35.49
+2500    0.2457 /   1.33      5.7692 / 675.08       4.7055 / 759.18
+5900    0.1763 /   1.36      nan                    已停止
+```
+
+每个单元格为 `loss / grad_norm`。
+
+因果起点是 step 1000–1200，而不是最终出现 NaN 的 step 5900。NaN 只是长期发散后的溢出结果。
+
+### 3.2 单变量责任归属
+
+同机、同数据集、同超参、同 seed 的 A/B 中，唯一变量是 `use_cudnn_attention`：
+
+```text
+cve5hiyt（0828）  commit d0aec27  显式 attention  step 59900 loss=0.0589
+orbnh1b7（0829）  commit 51d9af3  全层 cuDNN     step 5900  loss=nan
+```
+
+`git diff d0aec27 HEAD -- src/ scripts/` 的 72 行增删全部属于 attention 开关本身，排除了学习率、数据集和训练配置变化。
+
+### 3.3 cuDNN 单步结果并非明显错误
+
+同一权重、真实 batch 和 rng 的完整模型梯度 A/B：
+
+```text
+batch 0  loss 显式=5.859448  cuDNN=5.851340  |grad| 104.31 vs 103.65  ratio=0.9937
+batch 1                                      |grad|  67.19 vs  67.45  ratio=1.0038
+batch 2                                      |grad|  70.71 vs  70.61  ratio=0.9986
+
+相对梯度差：1.75e-02 / 1.66e-02 / 2.12e-02
+```
+
+梯度范数比值接近 1，没有单个模块出现异常放大。单层真实形状中，cuDNN 对 fp32 参考的相对误差约 `5.4e-03`，显式 BF16 路径约 `3.2e-03`，两者同阶。
+
+同一 batch 重复 cuDNN backward 的随机差异约 `0.9%`，来自 atomic 累加；cuDNN 与显式路径之间的确定性分量约 `1.5%`。因此 cuDNN 给出的是另一个合法的 BF16 近似，而不是单步明显错误的梯度。
+
+## 4. 根因机制：Adam 放大坐标级差异
+
+### 4.1 差分没有固定跨 batch 方向
+
+对 8 个真实 batch 分别计算显式梯度和 3 次 cuDNN backward 均值：
+
+```text
+跨 batch delta 两两 cosine：-0.004043 ± 0.126301，范围 [-0.247516, +0.226915]
+||delta|| / ||g_explicit||： 1.556024e-02 ± 2.76e-03
+cuDNN 同 batch 随机分量：   8.752676e-03
+cos(delta, gradient)：      -0.00113 ± 0.225
+cos(delta, parameter)：     +0.000016
+```
+
+差分没有稳定的跨 batch 方向，也不沿当前梯度或参数方向累积。
+
+### 4.2 Adam 更新空间暴露了真正差异
+
+从 strict-order 全层 cuDNN run 的 step 1000 checkpoint 恢复真实 Adam 状态，对同一真实 batch 和 rng 比较实际参数更新：
+
+```text
+                         相对更新差异   更新 cosine   符号翻转坐标
+真实全层 cuDNN 差分          3.3521%      0.999440       0.7885%
+1.5% 人造乘性噪声            0.7313%      --             0.1459%
+```
+
+cuDNN 原始梯度 L2 差异只有约 1.5%，但其误差分布会扰动小梯度、小二阶矩坐标，经过 Adam 逐坐标归一化后变成约 3.35% 的更新差异，并使约 0.79% 的非零更新坐标翻转符号。
+
+相比之下，按 `|g|` 分配能量的 1.5% 独立随机噪声几乎不扰动这些敏感坐标。该人造噪声从 step 0 稳定跑满 3000 步，step 2900 为 `loss=0.2411 / grad_norm=1.6389`，证明训练配方并非对任意 1.5% 梯度噪声都没有裕度。
+
+### 4.3 早层误差会穿过后续网络放大
+
+在同一个 step-1000 Adam 状态上，仅让不同连续层段使用 cuDNN：
+
+```text
+cuDNN 层段       相对更新差异   更新 cosine   符号翻转坐标
+全 18 层            3.3118%       0.999451       0.7769%
+最前 6 层           3.2061%       0.999486       0.7548%
+中间 6 层           2.2169%       0.999754       0.5141%
+最后 6 层           1.1633%       0.999932       0.2720%
+最前 3 层           3.3016%       0.999455       0.7720%
+中间 3 层           1.9894%       0.999802       0.4634%
+最后 3 层           0.8470%       0.999964       0.1926%
+稳定噪声参考         0.7313%       --             0.1459%
+```
+
+只融合最前 3 层几乎等同全 18 层。层的位置比数量更重要：早层的微小数值差异会穿过后续 15 层前向和反向传播，显著放大。
+
+## 5. 已否定的修复假设
+
+### 5.1 升级 cuDNN 不能解决缓慢发散
+
+cuDNN 9.14 与 9.19 的发散曲线在 step 1000–1200 的同一位置分叉。当前统一环境实际加载 `nvidia-cudnn-cu12==9.19.0.56`，全层训练仍发散。
+
+### 5.2 非 128 倍序列长度不是本任务根因
+
+cuDNN 9.15.1 release note 提到旧版 SDPA backward 在静态序列长度不是 128 倍数时可能产生错误结果或 NaN，因此曾怀疑 `T=1018` 命中该缺陷。真实形状对照结果：
+
+```text
+T=1018 (%128=122)  dQ rel-L2: 显式 3.23e-03  cuDNN 3.21e-03  NaN=0
+T=1024 (%128=0)    dQ rel-L2: 显式 3.23e-03  cuDNN 3.21e-03  NaN=0
 ```
 
-这与 NVIDIA 报告的触发条件高度一致。相关 issue 还表明：即使前向输出正确，cuDNN backward 的 `dQ` 仍可能出现整行 NaN。
+把 `max_token_len` 从 200 临时改为 206 不能改变数值行为。该 NVIDIA 缺陷与本任务的缓慢发散不是同一问题。
 
 参考：
 
 - [NVIDIA/cudnn-frontend #160](https://github.com/NVIDIA/cudnn-frontend/issues/160)
 - [cuDNN 9.15.1 Release Notes](https://docs.nvidia.com/deeplearning/cudnn/backend/v9.15.1/release-notes.html)
 
-## 尝试解决办法
+### 5.3 原生 seq_lengths / padding kernel 不能解决发散
+
+2026-08-31 对旧 summary 建议的方案进行了最终验证：向 cuDNN 调用同时传入完整物理长度的 `query_seq_lengths` 和 `key_value_seq_lengths`，值均为 1018；实际 padding 语义仍由原 block mask 负责。
+
+编译结果确认 seq lengths 已作为两个 `s32[B]` operand 进入 `__cudnn$fmhaScaleBiasSoftmax` 和 backward custom call，但未切换到数值上不同的算法。
+
+真实单层形状、三次重复 backward：
+
+```text
+NaN count，无 seq_lengths： [0, 0, 0]
+NaN count，有 seq_lengths： [0, 0, 0]
+相对 fp32 梯度误差，无 seq_lengths：约 5.66924e-03
+相对 fp32 梯度误差，有 seq_lengths：约 5.66927e-03
+两路径相对梯度差：5.3e-06 至 3.1e-05
+```
+
+两路径之差与同一 cuDNN backward 的 atomic 随机波动同量级。
+
+在 step-1000 真实 Adam 状态上的全模型对照：
+
+```text
+显式 loss：                 0.31219399
+padding-cuDNN loss：        0.31204814
+Adam update 相对差：        3.310997%
+update cosine：             0.99945178
+符号翻转坐标：              0.77606%
+```
+
+这与已知会发散的全 18 层 cuDNN 指纹 `3.3118% / 0.7769%` 几乎完全一致。把 Adam `eps` 提高到 `1e-6` 或 `1e-5` 后，update 相对差仍分别为 `3.1655%` 和 `3.0275%`，符号翻转仍约 `0.776%`。
+
+最后又完成了 3000-step、8×H100、batch 256、`strict_batch_order=true` 的全 18 层 padding-cuDNN 验收。W&B、周期 checkpoint 和最终 checkpoint 均关闭：
+
+```text
+step    loss      grad_norm    param_norm
+   0    6.4023       65.6324    1802.3865
+ 500    0.5719        2.5464    1802.4017
+1000    0.4298        2.4472    1802.5095
+1100    0.4591        2.2270    1802.5511
+1200    0.5902       13.5020    1802.5952
+1300    0.9563      108.9634    1802.6366
+1500    2.2053       45.7198    1802.7318
+2000    3.7753      150.5048    1802.9735
+2400    4.7852     6469.9019    1803.1691
+2500    5.7166      609.2015    1803.2141
+2900    6.2039      603.4764    1803.3545
+```
+
+step 1200 的显式基线约为 `0.3742 / 2.12`，padding-cuDNN 梯度已达到约 6.4 倍，超过 3 倍验收上限。该方案跑满 3000 步时尚未直接变成 NaN，但已沿已知轨迹持续发散，不能作为解决办法。
+
+### 5.4 其他已排除项
+
+- **两种 mask 修复**：dummy key 与 Q stop-gradient 在有效行上数值等价，两者都会出现第二类发散；
+- **动态库混装**：统一到单源 pip cu128/cuDNN 9.19 后仍发散；
+- **logit 幅度**：std 1–80 扫描中 cuDNN 与显式误差同阶；
+- **head_dim=256**：与 head_dim=128 表现一致；
+- **cotangent 稠密度**：只在 action token 非零和所有有效 token 稠密时结果一致；
+- **dBias 路径**：实际 batch 256 不满足 `bias.shape[0]==1`，未进入该路径；
+- **FSDP 分片**：attention 只按 batch 分片，每卡拥有完整序列；
+- **batch 到达顺序**：`strict_batch_order=true` 仍在 step 1000–1200 分叉；
+- **降低 peak LR**：`2.5e-5 -> 1.5e-5` 只把分叉推迟约 300 步；
+- **提高 Adam b2**：`0.95 -> 0.99` 从 step 0 跑满 3000 步仍发散；
+- **提高 Adam b1**：`0.9 -> 0.99` 只把明显分叉推迟到约 step 1400；
+- **提高 Adam eps**：update 差异和符号翻转几乎不变；
+- **1.5% 独立随机梯度噪声**：显式 attention 稳定跑满 3000 步，没有复现发散。
+
+## 6. 已验证的稳定折中：只融合最后 3 层
+
+配置 `configs/diag_cudnn_last3.yaml` 使用：
+
+```yaml
+model:
+  use_cudnn_attention: true
+  cudnn_attention_layer_start: 15
+  cudnn_attention_num_layers: 3
+```
+
+8×H100、batch 256、strict-order、原始 Adam，从 pi05_base step 0 开始，完整执行 3000 个更新：
+
+```text
+step    loss      grad_norm    param_norm
+   0    6.4029      65.5204     1802.3865
+ 500    0.5719       2.4526     1802.4014
+1000    0.4240       2.3789     1802.5061
+1100    0.4050       2.1847     1802.5416
+1200    0.3757       2.3869     1802.5775
+1600    0.3238       2.6505     1802.7225
+2000    0.2793       1.8254     1802.8666
+2500    0.2463       2.0116     1803.0460
+2900    0.2424       1.4769     1803.1906
+```
+
+该曲线与显式基线一致，满足 3000 步验收。普通 step 约 `1.5 s/step`，相比显式路径约 `1.57–1.64 s/step` 只有约 4–9% 加速，明显低于全层 cuDNN 的约 36% 潜在收益。
+
+最后 3 层方案只是已通过短程验收的保守折中，尚未完成 60000 步生产验收，不能等同于完整生产结论。
+
+## 7. CUDA/cuDNN 环境事故与当前要求
+
+### 7.1 历史混装问题
+
+事故期间 `cuda_versions.cudnn_get_version()` 返回 `91400`，但进程实际加载了 9.10.2 dispatcher stub 和 9.14 engine 库：
+
+```text
+site-packages/nvidia/cudnn/lib/libcudnn.so.9          <- 9.10.2 dispatcher
+$CONDA_PREFIX/lib/libcudnn_graph.so.9.14.0            <- 9.14 engine
+$CONDA_PREFIX/lib/libcudnn_engines_precompiled.so.9.14.0
+$CONDA_PREFIX/lib/libcudnn_ops.so.9.14.0
+```
+
+`cudnnGetVersion()` 的返回值来自 engine，因此只看 `91400` 无法发现 dispatcher/engine 混装。
+
+### 7.2 当前统一环境
+
+环境已统一为：
+
+```text
+CUDA Toolkit          12.8
+PyTorch               2.11.0+cu128
+torchvision           0.26.0+cu128
+torchcodec             0.11.1+cu128
+nvidia-cudnn-cu12      9.19.0.56
+JAX                    0.5.3
+```
+
+启动训练时不要导出 `LD_LIBRARY_PATH="$CONDA_PREFIX/lib"`。该目录中残留的 conda-forge `libcublas`、`libcudart` 或 `libnccl` 可能遮蔽 pip cu128 库。
+
+生产启动前在相同 shell 中运行：
+
+```bash
+env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py
+```
+
+脚本会枚举 `/proc/self/maps` 中实际加载的 cuDNN/cuBLAS/NCCL 路径，并在真实 attention 形状上做前向、反向和 fp32 参考对比。
+
+> [!NOTE]
+> `check_cuda_stack.py` PASS 只证明环境单源且单步内核有限、误差在阈值内。第二类缓慢发散曾通过所有此类 smoke test，因此 PASS 不能替代 3000 步训练验收。
+
+## 8. 生产处置与恢复规则
+
+### 8.1 当前可选配置
+
+需要绝对稳妥时：
+
+```yaml
+model:
+  use_cudnn_attention: false
+```
+
+需要先获取小幅吞吐收益并接受受监控长跑时：
+
+```yaml
+model:
+  use_cudnn_attention: true
+  cudnn_attention_layer_start: 15
+  cudnn_attention_num_layers: 3
+```
+
+不得用于生产：
+
+```yaml
+model:
+  use_cudnn_attention: true
+  cudnn_attention_num_layers: null  # 全 18 层
+```
+
+### 8.2 checkpoint 恢复
+
+- 一旦参数、Adam `mu` 或 `nu` 出现 NaN，后续 checkpoint 均视为污染；
+- 不得通过切换 attention、降低 LR 或修改 optimizer 后继续污染 checkpoint；
+- 必须回退到最后一个确定干净的 checkpoint；
+- 恢复前检查 checkpoint 内参数、optimizer state 和 step；
+- 使用独立实验名验证，关闭 W&B、周期保存和最终保存，避免覆盖生产目录；
+- 只有通过完整验收后才能恢复生产保存。
+
+### 8.3 建议增加发散护栏
+
+当前训练只在日志 step 计算 `grad_norm`。建议在 host 端维护滚动中位数，并在连续多个日志窗口超过固定倍数时中止。按本次曲线，该护栏可在 step 1200–1300 停止，而不是等到 loss 变成 NaN 或污染数万个 step。
+
+护栏只能减少损失，不能让全层 cuDNN 变得收敛。
+
+## 9. 新 attention 实现的统一验收标准
+
+任何新的全层 fused/flash attention backend、custom VJP 或稳定化方案必须：
+
+1. 从 **step 0** 开始；
+2. 使用 `strict_batch_order=true`，固定 seed、数据集、batch 和超参；
+3. 跑满 **3000 个更新**，不得只跑 40、60 或 500 步；
+4. 与显式 attention 基线逐点比较 step 200/500/1000/1500/2000/2900；
+5. 全程 loss、梯度、参数和 optimizer state 有限；
+6. `grad_norm` 不得超过基线同点的 3 倍；
+7. 通过后再进入有 checkpoint 的受监控长跑；
+8. 生产长跑仍需持续观察 loss 和 grad norm。
+
+以下均不是有效正确性验收：
+
+- 单次前向/反向有限；
+- 40-step 吞吐短跑；
+- 从已经收敛的 checkpoint 短跑；
+- 200–500 个连续有限 step；
+- 只检查最终是否出现 NaN；
+- 跨不同 batch 顺序的 run 直接比较单点 grad norm。
+
+## 10. 不要重复的弯路
+
+- 不要再次把全空 mask 的瞬时 NaN 与全层缓慢发散当成同一问题；
+- 不要再把 cuDNN 9.15.1 或 `T=1024` 当作本任务的修复；
+- 不要再尝试完整物理长度的 `query_seq_lengths/key_value_seq_lengths`；
+- 不要再修改 dummy-key/stop-gradient mask 语义来解决第二类发散；
+- 不要只调低 `peak_lr` 或提高 Adam `b1`、`b2`、`eps`；
+- 不要用 1.5% 独立随机噪声代替真实 cuDNN 差分；
+- 不要只看 `cuda_versions.cudnn_get_version()` 判断动态库是否干净；
+- 不要把最后 3 层的 3000 步通过误写成全 18 层或 60000 步生产通过。
+
+## 11. 相关脚本与配置
+
+- `scripts/check_cuda_stack.py`：检查实际映射的 CUDA/cuDNN 库并做真实形状数值 smoke test；
+- `scripts/grad_ab_probe.py`：完整模型显式/cuDNN 梯度 A/B；
+- `scripts/grad_ab_direction_probe.py`：多 batch 差分方向和随机分量；
+- `scripts/optimizer_update_ab_probe.py`：从真实 Adam 状态比较参数更新误差；
+- `scripts/hybrid_update_ab_probe.py`：比较不同 cuDNN 层段的更新空间误差；
+- `configs/diag_cudnn_strict_order.yaml`：全层 cuDNN strict-order 诊断；
+- `configs/diag_cudnn_last3.yaml`：最后 3 层 cuDNN 稳定折中诊断。
+
+## 12. 文档整合说明
+
+本文整合并取代以下三份调查记录作为当前权威结论：
 
-1. **暂时关闭 cuDNN Attention**：恢复显式 Attention 路径，保证生产训练数值有效。
-2. **启用原生 padding 路径**：向 JAX `dot_product_attention` 传入值为完整物理长度的 `query_seq_lengths` 和 `key_value_seq_lengths`，强制 cuDNN 使用 padding kernel；现有 block mask 仍负责实际 token 屏蔽。
-3. **升级 cuDNN runtime**：将实际加载版本从 9.14 升到至少 9.15.1，并检查启动日志确认版本。
-4. **长度对照实验**：把 `max_token_len` 临时改为 206，使总长度变为 1024，以验证“非 128 倍数”是否为直接触发条件。
+- `2026-08-29-cudnn-attention-nan-incident.md`：全空 mask 的瞬时 NaN 事故；
+- `2026-08-30-cudnn-attention-divergence.md`：全层 cuDNN 缓慢发散、Adam 机制和最后 3 层方案；
+- 原 `2026-08-31-cudnn-sdpa-nan-summary.md`：关于非 128 倍长度、升级 cuDNN 和 padding kernel 的待验证假设。
 
-以上方案需要从有效 checkpoint 使用固定真实 batch 做逐 step 有限性检查。至少验证 500 个连续 step 的 loss、梯度、参数和优化器状态均为有限值后，才能重新用于生产训练。
+整合完成后，前两份独立文件已删除，历史内容仍可通过 Git 提交记录追溯。后续调查和生产决策以本文为准。

From 3b2795ee051a5ba2f49915ecbd4e8cb9a1354e7a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Wed, 2 Sep 2026 10:56:21 +0800
Subject: [PATCH 12/28] feat: add Nova5 LeRobot data config

---
 configs/_examples/_FULL_REFERENCE.yaml        |   8 ++
 ...bot_nova5_loreal_returns_sorting_0831.yaml |  27 +++++
 src/openpi/policies/dobot_nova5_policy.py     | 107 ++++++++++++++++++
 .../policies/dobot_nova5_policy_test.py       |  78 +++++++++++++
 src/openpi/training/config.py                 |  54 +++++++++
 src/openpi/training/registry.py               |   1 +
 6 files changed, 275 insertions(+)
 create mode 100644 configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0831.yaml
 create mode 100644 src/openpi/policies/dobot_nova5_policy.py
 create mode 100644 src/openpi/policies/dobot_nova5_policy_test.py

diff --git a/configs/_examples/_FULL_REFERENCE.yaml b/configs/_examples/_FULL_REFERENCE.yaml
index 6519dc1..b68ccfc 100644
--- a/configs/_examples/_FULL_REFERENCE.yaml
+++ b/configs/_examples/_FULL_REFERENCE.yaml
@@ -121,6 +121,14 @@ data:
   # action_sequence_keys: [action]
   # repack_transforms is built from use_head_camera and cannot be overridden here.
 
+  # ---- LeRobotDobotNova5DataConfig-specific (uncomment when used) ----
+  # Single right-arm Dobot Nova5 DH. 10D state/action order:
+  # [tcp.x, tcp.y, tcp.z, tcp.r1-r6, gripper.pos]. The dataset camera columns
+  # are observation.images.head and observation.images.wrist.
+  # use_delta_cartesian_actions: true  # delta-encode TCP dims 0-8; keep gripper dim 9 absolute
+  # default_prompt: "..."
+  # action_sequence_keys: [action]
+
   # ---- LeRobotAlohaDataConfig-specific (uncomment when used) ----
   # use_delta_joint_actions: true       # delta-encode joints, keep grippers absolute
   # default_prompt: "..."
diff --git a/configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0831.yaml b/configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0831.yaml
new file mode 100644
index 0000000..4c511cc
--- /dev/null
+++ b/configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0831.yaml
@@ -0,0 +1,27 @@
+# JAX-only pi0.5 training config for the single-arm Dobot Nova5 dataset.
+model:
+  type: Pi0Config
+  max_token_len: 200
+  pi05: true
+  discrete_state_input: true
+  enable_training_time_rtc: true
+  use_cudnn_attention: false
+weight_loader:
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+ema_decay: null
+data:
+  type: LeRobotDobotNova5DataConfig
+  repo_id: Xense/loreal_returns_sorting_0831
+  use_delta_cartesian_actions: true
+  base_config:
+    prompt_from_task: true
+    tactile: false
+  default_prompt: Pick up returned cosmetic from conveyor, place neatly on left platform.
+batch_size: 256
+num_workers: 64
+num_train_steps: 60000
+save_interval: 10000
+keep_period: 10000
+fsdp_devices: 8
+strict_batch_order: false
diff --git a/src/openpi/policies/dobot_nova5_policy.py b/src/openpi/policies/dobot_nova5_policy.py
new file mode 100644
index 0000000..7780f50
--- /dev/null
+++ b/src/openpi/policies/dobot_nova5_policy.py
@@ -0,0 +1,107 @@
+"""Policy transforms for the single-arm Dobot Nova5 DH robot.
+
+The LeRobot dataset and driver expose the same Cartesian 10D layout::
+
+    [tcp.x, tcp.y, tcp.z, tcp.r1, tcp.r2, tcp.r3,
+     tcp.r4, tcp.r5, tcp.r6, gripper.pos]
+
+``r1-r6`` are the first two columns of the rotation matrix, so they are already
+a continuous representation and do not need an additional conversion.
+
+The robot has a head camera and one camera named ``wrist`` on the right arm.
+OpenPI still presents three image slots to the model: the unused left-wrist
+slot is a black image with ``image_mask=False``.
+"""
+
+import dataclasses
+from typing import ClassVar
+
+import einops
+import numpy as np
+
+from openpi import transforms
+
+
+TCP = slice(0, 9)
+GRIPPER = slice(9, 10)
+STATE_DIM = 10
+
+
+def make_dobot_nova5_example() -> dict:
+    """Create a policy input with the same keys and shapes as the Nova5 rig."""
+    return {
+        "state": np.ones((STATE_DIM,), dtype=np.float32),
+        "images": {
+            "head": np.random.randint(256, size=(3, 224, 224), dtype=np.uint8),
+            "wrist": np.random.randint(256, size=(3, 224, 224), dtype=np.uint8),
+        },
+        "prompt": "do something",
+    }
+
+
+@dataclasses.dataclass(frozen=True)
+class DobotNova5Inputs(transforms.DataTransformFn):
+    """Convert Nova5 LeRobot observations to OpenPI model inputs.
+
+    Expected images are channel-first ``head`` and ``wrist`` frames. The state
+    and optional action chunk use the 10D layout documented above.
+    """
+
+    EXPECTED_CAMERAS: ClassVar[tuple[str, ...]] = ("head", "wrist")
+
+    def __call__(self, data: dict) -> dict:
+        data = _decode_dobot_nova5(data)
+
+        images = data["images"]
+        unexpected = set(images) - set(self.EXPECTED_CAMERAS)
+        if unexpected:
+            raise ValueError(f"Unexpected cameras {tuple(sorted(unexpected))}; got {tuple(images)}")
+        missing = set(self.EXPECTED_CAMERAS) - set(images)
+        if missing:
+            raise ValueError(f"Missing required Nova5 cameras {tuple(sorted(missing))}; got {tuple(images)}")
+
+        head = images["head"]
+        wrist = images["wrist"]
+        result = {
+            "image": {
+                "base_0_rgb": head,
+                "left_wrist_0_rgb": np.zeros_like(wrist),
+                "right_wrist_0_rgb": wrist,
+            },
+            "image_mask": {
+                "base_0_rgb": np.True_,
+                "left_wrist_0_rgb": np.False_,
+                "right_wrist_0_rgb": np.True_,
+            },
+            "state": data["state"],
+        }
+
+        if "actions" in data:
+            result["actions"] = np.asarray(data["actions"])
+        if "prompt" in data:
+            result["prompt"] = data["prompt"]
+        return result
+
+
+@dataclasses.dataclass(frozen=True)
+class DobotNova5Outputs(transforms.DataTransformFn):
+    """Trim the model's padded action vector back to the robot's native 10D."""
+
+    def __call__(self, data: dict) -> dict:
+        return {"actions": np.asarray(data["actions"][:, :STATE_DIM])}
+
+
+def _decode_dobot_nova5(data: dict) -> dict:
+    """Normalize LeRobot images from CHW to HWC and floats to uint8."""
+
+    def convert_image(image):
+        image = np.asarray(image)
+        if np.issubdtype(image.dtype, np.floating):
+            image = (255 * image).astype(np.uint8)
+        return einops.rearrange(image, "c h w -> h w c")
+
+    return {
+        **data,
+        "state": np.asarray(data["state"]),
+        "images": {name: convert_image(image) for name, image in data["images"].items()},
+    }
diff --git a/src/openpi/policies/dobot_nova5_policy_test.py b/src/openpi/policies/dobot_nova5_policy_test.py
new file mode 100644
index 0000000..53f082a
--- /dev/null
+++ b/src/openpi/policies/dobot_nova5_policy_test.py
@@ -0,0 +1,78 @@
+import pathlib
+
+import numpy as np
+import pytest
+
+import openpi.models.pi0_config as pi0_config
+from openpi.policies import dobot_nova5_policy
+from openpi.training import config as _config
+from openpi.training import registry
+import openpi.transforms as _transforms
+
+
+def _example(*, with_actions: bool = False) -> dict:
+    example = dobot_nova5_policy.make_dobot_nova5_example()
+    if with_actions:
+        example["actions"] = np.zeros((50, dobot_nova5_policy.STATE_DIM), dtype=np.float32)
+    return example
+
+
+def test_inputs_map_single_wrist_to_right_slot_and_mask_left_slot():
+    result = dobot_nova5_policy.DobotNova5Inputs()(_example(with_actions=True))
+
+    assert set(result["image"]) == {"base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb"}
+    assert result["image_mask"] == {
+        "base_0_rgb": np.True_,
+        "left_wrist_0_rgb": np.False_,
+        "right_wrist_0_rgb": np.True_,
+    }
+    assert not result["image"]["left_wrist_0_rgb"].any()
+    assert result["state"].shape == (10,)
+    assert result["actions"].shape == (50, 10)
+
+
+def test_inputs_require_exact_dataset_camera_contract():
+    example = _example()
+    del example["images"]["wrist"]
+    with pytest.raises(ValueError, match="Missing required Nova5 cameras"):
+        dobot_nova5_policy.DobotNova5Inputs()(example)
+
+
+def test_outputs_trim_model_actions_to_ten_dims():
+    padded = np.random.rand(50, 32)
+    result = dobot_nova5_policy.DobotNova5Outputs()({"actions": padded})
+
+    assert result["actions"].shape == (50, 10)
+    np.testing.assert_array_equal(result["actions"], padded[:, :10])
+
+
+def test_delta_mask_keeps_only_gripper_absolute():
+    mask = _transforms.make_bool_mask(9, -1)
+
+    assert len(mask) == dobot_nova5_policy.STATE_DIM
+    assert all(mask[dobot_nova5_policy.TCP])
+    assert not any(mask[dobot_nova5_policy.GRIPPER])
+
+
+def test_data_config_matches_lerobot_dataset_columns(tmp_path: pathlib.Path):
+    factory = _config.LeRobotDobotNova5DataConfig(repo_id="Xense/loreal_returns_sorting_0831")
+    data_config = factory.create(tmp_path, pi0_config.Pi0Config(pi05=True))
+
+    assert registry.resolve(registry.all_registries()["DATA_CONFIGS"], "LeRobotDobotNova5DataConfig") is type(factory)
+    assert data_config.repack_transforms.inputs[0].structure == {
+        "images": {
+            "head": "observation.images.head",
+            "wrist": "observation.images.wrist",
+        },
+        "state": "observation.state",
+        "actions": "action",
+        "prompt": "task",
+    }
+    assert [type(transform) for transform in data_config.data_transforms.inputs] == [
+        dobot_nova5_policy.DobotNova5Inputs,
+        _transforms.DeltaActions,
+    ]
+    assert [type(transform) for transform in data_config.data_transforms.outputs] == [
+        _transforms.AbsoluteActions,
+        dobot_nova5_policy.DobotNova5Outputs,
+    ]
diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py
index b87145f..70b30b9 100644
--- a/src/openpi/training/config.py
+++ b/src/openpi/training/config.py
@@ -22,6 +22,7 @@
 import openpi.models.tokenizer as _tokenizer
 import openpi.policies.aloha_policy as aloha_policy
 import openpi.policies.bi_flexiv_policy as bi_flexiv_policy
+import openpi.policies.dobot_nova5_policy as dobot_nova5_policy
 import openpi.policies.droid_policy as droid_policy
 import openpi.policies.xtac_umi_policy as xtac_umi_policy
 import openpi.shared.download as _download
@@ -488,6 +489,59 @@ def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig
         )
 
 
+@dataclasses.dataclass(frozen=True)
+class LeRobotDobotNova5DataConfig(DataConfigFactory):
+    """Data config for the single right-arm Dobot Nova5 DH in LeRobot format.
+
+    State/action layout (10D): ``tcp.{x,y,z,r1-r6}`` followed by
+    ``gripper.pos``. Cameras are ``head`` and ``wrist``. This matches both the
+    ``dobot_nova5_dh`` LeRobot driver and datasets such as
+    ``Xense/loreal_returns_sorting_0831``.
+    """
+
+    use_delta_cartesian_actions: bool = True
+    default_prompt: str | None = None
+    action_sequence_keys: Sequence[str] = ("action",)
+
+    @override
+    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
+        data_transforms = _transforms.Group(
+            inputs=[dobot_nova5_policy.DobotNova5Inputs()],
+            outputs=[dobot_nova5_policy.DobotNova5Outputs()],
+        )
+
+        if self.use_delta_cartesian_actions:
+            # Cartesian TCP dims are relative to the current observation; the
+            # trailing gripper command remains absolute.
+            delta_action_mask = _transforms.make_bool_mask(9, -1)
+            data_transforms = data_transforms.push(
+                inputs=[_transforms.DeltaActions(delta_action_mask)],
+                outputs=[_transforms.AbsoluteActions(delta_action_mask)],
+            )
+
+        return dataclasses.replace(
+            self.create_base_config(assets_dirs, model_config),
+            repack_transforms=_transforms.Group(
+                inputs=[
+                    _transforms.RepackTransform(
+                        {
+                            "images": {
+                                "head": "observation.images.head",
+                                "wrist": "observation.images.wrist",
+                            },
+                            "state": "observation.state",
+                            "actions": "action",
+                            "prompt": "task",
+                        }
+                    )
+                ]
+            ),
+            data_transforms=data_transforms,
+            model_transforms=ModelTransformFactory(default_prompt=self.default_prompt)(model_config),
+            action_sequence_keys=self.action_sequence_keys,
+        )
+
+
 @dataclasses.dataclass(frozen=True)
 class LeRobotXtacUmiDataConfig(DataConfigFactory):
     """
diff --git a/src/openpi/training/registry.py b/src/openpi/training/registry.py
index 02ffff4..13c3478 100644
--- a/src/openpi/training/registry.py
+++ b/src/openpi/training/registry.py
@@ -76,6 +76,7 @@ def _populate_data_configs() -> None:
             "LeRobotDROIDDataConfig": _config.LeRobotDROIDDataConfig,
             "RLDSDroidDataConfig": _config.RLDSDroidDataConfig,
             "LeRobotBiFlexivDataConfig": _config.LeRobotBiFlexivDataConfig,
+            "LeRobotDobotNova5DataConfig": _config.LeRobotDobotNova5DataConfig,
             "LeRobotXtacUmiDataConfig": _config.LeRobotXtacUmiDataConfig,
             # SimpleDataConfig deliberately omitted: it carries lambdas (data_transforms)
             # that cannot be serialized to YAML.

From 3c5aac68696ad83f46cce6cc54d8694d76ec7637 Mon Sep 17 00:00:00 2001
From: Vertax42 <1007644658@qq.com>
Date: Thu, 3 Sep 2026 14:01:59 +0800
Subject: [PATCH 13/28] add loreal train config yaml

---
 ...> pi05_base_dobot_nova5_loreal_returns_sorting_0902_h100.yaml} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename configs/_examples/{pi05_base_dobot_nova5_loreal_returns_sorting_0831.yaml => pi05_base_dobot_nova5_loreal_returns_sorting_0902_h100.yaml} (100%)

diff --git a/configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0831.yaml b/configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0902_h100.yaml
similarity index 100%
rename from configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0831.yaml
rename to configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0902_h100.yaml

From df9755fad17439161c7972b648a24af684399467 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Thu, 3 Sep 2026 18:04:01 +0800
Subject: [PATCH 14/28] Improve cuDNN attention precision diagnostics

---
 ...s_case_insertion_teleop_rtc_0828_h100.yaml |   2 +-
 scripts/attention_precision_update_probe.py   | 216 ++++++++++++++++++
 src/openpi/models/gemma.py                    |  95 +++++++-
 src/openpi/models/pi0.py                      |   2 +
 src/openpi/models/pi0_config.py               |   7 +
 5 files changed, 309 insertions(+), 13 deletions(-)
 create mode 100644 scripts/attention_precision_update_probe.py

diff --git a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml
index 5c0df75..d5eadb9 100644
--- a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml
@@ -15,7 +15,7 @@ weight_loader:
 ema_decay: null
 data:
   type: LeRobotBiFlexivDataConfig
-  repo_id: Xense/earbud_case_insertion_teleop_0515
+  repo_id: Xense/earbud_case_insertion_teleop_0515_left_right
   base_config:
     prompt_from_task: true
     tactile: false
diff --git a/scripts/attention_precision_update_probe.py b/scripts/attention_precision_update_probe.py
new file mode 100644
index 0000000..71de30f
--- /dev/null
+++ b/scripts/attention_precision_update_probe.py
@@ -0,0 +1,216 @@
+"""Compare attention implementations against an fp32 attention reference at a real Adam state.
+
+The 2026-08-31 report showed that the full-18-layer cuDNN path diverges after
+~step 1000 although its raw gradient is only ~1.5% away from the explicit bf16
+path. This probe restores the real step-1000 train state (params + Adam moments),
+computes the gradient and the actual AdamW update for the same real batch/rng under
+
+  explicit_fp32ref  explicit attention with fp32 q/k/v/probs   (numerical reference)
+  explicit_bf16     production explicit path                    (known stable)
+  cudnn_bf16        historical cuDNN path                       (known to diverge)
+  cudnn_fp16        cuDNN kernel in float16 + dynamic loss scaling (candidate)
+
+and reports, for every variant, raw-gradient and Adam-update distance to the
+fp32 reference and to the explicit bf16 path. A candidate is promising when its
+update-space fingerprint (relative delta / sign-flip fraction) is at or below the
+explicit bf16 path's own distance from the fp32 reference.
+"""
+
+import dataclasses
+import functools
+import logging
+import sys
+import time
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import openpi.training.checkpoints as _checkpoints
+import openpi.training.config as _config
+import openpi.training.data_loader as _data_loader
+import openpi.training.sharding as sharding
+
+sys.path.insert(0, "scripts")
+import train as train_mod
+
+CONFIG = "diag_cudnn_strict_order"
+EXP_NAME = "diag_cudnn_strict_order"
+CHECKPOINT_STEP = 1000
+BATCH = 32
+WORKERS = 16
+NUM_BATCHES = 2
+
+
+def _loss_grads_update(model_def, trainable_filter, tx, opt_state, params, rng, obs, act):
+    model = nnx.merge(model_def, params)
+    model.train()
+
+    def loss_fn(m, r, o, a):
+        return jnp.mean(m.compute_loss(r, o, a, train=True))
+
+    loss, grads = nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, trainable_filter))(model, rng, obs, act)
+    trainable_params = params.filter(trainable_filter)
+    updates, _ = tx.update(grads, opt_state, trainable_params)
+    return loss, grads, updates
+
+
+def _flat(tree):
+    pure = tree.to_pure_dict()
+    out = []
+
+    def rec(value, prefix):
+        if isinstance(value, dict):
+            for key, child in value.items():
+                rec(child, f"{prefix}/{key}" if prefix else str(key))
+        else:
+            out.append((prefix, value))
+
+    rec(pure, "")
+    return out
+
+
+def _to_host(tree):
+    return {path: np.asarray(jax.device_get(value), np.float32) for path, value in _flat(tree)}
+
+
+def _dot(a, b):
+    return float(np.vdot(a.astype(np.float64), b.astype(np.float64)))
+
+
+def _metrics(reference, candidate):
+    totals = {"ALL": [0.0, 0.0, 0.0, 0.0, 0, 0]}
+    for path, ref in reference.items():
+        cand = candidate[path]
+        delta = cand - ref
+        group = path.split("/", 1)[0]
+        totals.setdefault(group, [0.0, 0.0, 0.0, 0.0, 0, 0])
+        for bucket in (totals["ALL"], totals[group]):
+            bucket[0] += _dot(ref, ref)
+            bucket[1] += _dot(cand, cand)
+            bucket[2] += _dot(delta, delta)
+            bucket[3] += _dot(ref, cand)
+            nonzero = (ref != 0) | (cand != 0)
+            bucket[4] += int(np.count_nonzero(nonzero & (np.signbit(ref) != np.signbit(cand))))
+            bucket[5] += int(np.count_nonzero(nonzero))
+    return totals
+
+
+def _summary(reference, candidate):
+    ref2, cand2, delta2, dot, flips, count = _metrics(reference, candidate)["ALL"]
+    return (
+        np.sqrt(delta2 / max(ref2, 1e-300)),
+        dot / max(np.sqrt(ref2 * cand2), 1e-300),
+        flips / max(count, 1),
+    )
+
+
+def _print_comparison(label, reference, candidate, per_group=False):
+    metrics = _metrics(reference, candidate)
+    ref2, cand2, delta2, dot, flips, count = metrics["ALL"]
+    print(
+        f"  {label:44s} rel_delta={np.sqrt(delta2 / max(ref2, 1e-300)):.4e} "
+        f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.7f} sign_flips={flips / max(count, 1):.4e}",
+        flush=True,
+    )
+    if per_group:
+        for group in sorted(key for key in metrics if key != "ALL"):
+            ref2, cand2, delta2, dot, flips, count = metrics[group]
+            print(
+                f"      {group:22s} rel={np.sqrt(delta2 / max(ref2, 1e-300)):.4e} "
+                f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.7f} sign_flips={flips / max(count, 1):.4e}",
+                flush=True,
+            )
+
+
+def _nonfinite(host):
+    return sum(int(np.count_nonzero(~np.isfinite(v))) for v in host.values())
+
+
+def main():
+    logging.basicConfig(level=logging.WARNING)
+    cfg = _config.get_config(CONFIG)
+    cfg = dataclasses.replace(
+        cfg,
+        exp_name=EXP_NAME,
+        batch_size=BATCH,
+        num_workers=WORKERS,
+        resume=True,
+        overwrite=False,
+        wandb_enabled=False,
+    )
+
+    def variant(**kw):
+        return dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, **kw))
+
+    variant_cfgs = {
+        "explicit_fp32ref": variant(use_cudnn_attention=False, explicit_attention_fp32=True),
+        "explicit_bf16": variant(use_cudnn_attention=False),
+        "cudnn_bf16": variant(use_cudnn_attention=True, cudnn_attention_dtype="bfloat16"),
+        "cudnn_fp16": variant(use_cudnn_attention=True, cudnn_attention_dtype="float16"),
+    }
+
+    mesh = sharding.make_mesh(cfg.fsdp_devices)
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+    rng = jax.random.key(cfg.seed)
+    train_rng, init_rng = jax.random.split(rng)
+
+    cfg_ckpt = variant_cfgs["cudnn_bf16"]
+    manager, resuming = _checkpoints.initialize_checkpoint_dir(
+        cfg_ckpt.checkpoint_dir, keep_period=cfg_ckpt.keep_period, overwrite=False, resume=True
+    )
+    if not resuming or CHECKPOINT_STEP not in manager.all_steps():
+        raise RuntimeError(f"checkpoint {CHECKPOINT_STEP} unavailable; found {manager.all_steps()}")
+    state_shape, _ = train_mod.init_train_state(cfg_ckpt, init_rng, mesh, resume=True)
+    state = _checkpoints.restore_state(manager, state_shape, None, step=CHECKPOINT_STEP)
+    jax.block_until_ready(state)
+    print(f"restored checkpoint step={CHECKPOINT_STEP}; train_state.step={int(state.step)}", flush=True)
+
+    print("building loader ...", flush=True)
+    loader = _data_loader.create_data_loader(variant_cfgs["explicit_bf16"], sharding=data_sharding, shuffle=True)
+    it = iter(loader)
+    batches = [next(it) for _ in range(NUM_BATCHES)]
+
+    graphdefs = {}
+    for name, vcfg in variant_cfgs.items():
+        graphdefs[name] = nnx.graphdef(vcfg.model.create(init_rng))
+    print("graphdefs built", flush=True)
+
+    fns = {
+        name: jax.jit(functools.partial(_loss_grads_update, graphdefs[name], cfg.trainable_filter, state.tx))
+        for name in variant_cfgs
+    }
+
+    for bi, (obs, act) in enumerate(batches):
+        step_rng = jax.random.fold_in(train_rng, state.step + bi)
+        grads_host, updates_host, losses = {}, {}, {}
+        for name, fn in fns.items():
+            t0 = time.time()
+            with sharding.set_mesh(mesh):
+                loss, grads, updates = fn(state.opt_state, state.params, step_rng, obs, act)
+                losses[name] = float(loss)
+                grads_host[name] = _to_host(grads)
+                updates_host[name] = _to_host(updates)
+            del grads, updates
+            print(
+                f"batch {bi} {name:18s} loss={losses[name]:.8f} "
+                f"nonfinite grads={_nonfinite(grads_host[name])} updates={_nonfinite(updates_host[name])} "
+                f"({time.time() - t0:.1f}s)",
+                flush=True,
+            )
+
+        for ref_name in ("explicit_fp32ref", "explicit_bf16"):
+            print(f"\n=== batch {bi}: RAW GRADIENT vs {ref_name} ===", flush=True)
+            for name in variant_cfgs:
+                if name != ref_name:
+                    _print_comparison(name, grads_host[ref_name], grads_host[name])
+            print(f"\n=== batch {bi}: ADAM UPDATE vs {ref_name} ===", flush=True)
+            for name in variant_cfgs:
+                if name != ref_name:
+                    _print_comparison(name, updates_host[ref_name], updates_host[name], per_group=(ref_name == "explicit_fp32ref"))
+        print("", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index 5c1ff6a..3f99c3b 100644
--- a/src/openpi/models/gemma.py
+++ b/src/openpi/models/gemma.py
@@ -182,12 +182,74 @@ def _stop_gradient_for_fully_masked_queries(q, attn_mask):
     return jnp.where(query_has_key, q, jax.lax.stop_gradient(q))
 
 
+def _cudnn_attention_call(q, k, v, attn_mask):
+    """The plain cuDNN fused attention call used by training (q is pre-scaled)."""
+    q = _stop_gradient_for_fully_masked_queries(q, attn_mask)
+    return jax.nn.dot_product_attention(q, k, v, mask=attn_mask, scale=1.0, implementation="cudnn")
+
+
+def _cudnn_attention_in_dtype(q, k, v, attn_mask, compute_dtype):
+    """Run cuDNN fused attention with q/k/v cast to ``compute_dtype`` (e.g. float16).
+
+    Motivation: the flash-attention backward forms ``dS = P * (dP - rowsum(dO * O))``
+    from the *rounded* stored output ``O``. For peaked attention rows (attention
+    sinks) ``dP`` and ``rowsum(dO * O)`` nearly cancel, so the rounding error of
+    ``O`` is a large, sign-consistent fraction of ``dS`` -- a bias the explicit
+    path does not have because it forms ``sum(P * dP)`` from the same rounded
+    ``dP``. float16 keeps 3 more mantissa bits than bfloat16, shrinking that bias
+    about 8x, but has a much smaller exponent range, so the incoming cotangent is
+    rescaled to a power of two near 1 (dynamic loss scaling) before the fused
+    backward and unscaled in float32 afterwards. Forward values are unchanged up
+    to the output rounding; q/k/v are exactly representable in float16 unless
+    they exceed its range, which the probe script checks.
+    """
+    out_dtype = q.dtype
+
+    # attn_mask is an explicit argument (not a closure) so that the backward,
+    # which flax remat traces separately, does not capture a leaked tracer.
+    @jax.custom_vjp
+    def attention(q, k, v, attn_mask):
+        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
+        return _cudnn_attention_call(qc, kc, vc, attn_mask).astype(out_dtype)
+
+    def attention_fwd(q, k, v, attn_mask):
+        return attention(q, k, v, attn_mask), (q, k, v, attn_mask)
+
+    def attention_bwd(residuals, d_out):
+        q, k, v, attn_mask = residuals
+        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
+        d_out32 = d_out.astype(jnp.float32)
+        max_abs = jnp.max(jnp.abs(d_out32))
+        tiny = jnp.finfo(jnp.float32).tiny
+        # Power-of-two scale so max|d_out * scale| lands in [0.25, 0.5): exact
+        # in the mantissa, leaves ~17 binades of headroom for dP = dO @ V^T
+        # and the dQ/dK/dV outputs inside the float16 kernel.
+        exponent = jnp.floor(jnp.log2(jnp.maximum(max_abs, tiny)))
+        scale = jnp.where(max_abs > 0, jnp.exp2(-(exponent + 2.0)), 1.0)
+        _, vjp_fn = jax.vjp(lambda a, b, c: _cudnn_attention_call(a, b, c, attn_mask), qc, kc, vc)
+        dq, dk, dv = vjp_fn((d_out32 * scale).astype(compute_dtype))
+        inv_scale = 1.0 / scale
+
+        def unscale(g):
+            return (g.astype(jnp.float32) * inv_scale).astype(out_dtype)
+
+        return unscale(dq), unscale(dk), unscale(dv), None
+
+    attention.defvjp(attention_fwd, attention_bwd)
+    return attention(q, k, v, attn_mask)
+
+
 @at.typecheck
 class Attention(nn.Module):
     """Attention module."""
 
     configs: Sequence[Config]
     use_cudnn_attention: bool = False
+    # Compute dtype handed to the cuDNN kernel: "bfloat16" (historical) or "float16"
+    # (see _cudnn_attention_in_dtype). Ignored by the explicit path.
+    cudnn_attention_dtype: str = "bfloat16"
+    # Diagnostic only: run the explicit path with fp32 q/k/v/probs as a reference.
+    explicit_attention_fp32: bool = False
 
     @nn.compact
     def __call__(self, xs, positions, attn_mask, kv_cache, use_cudnn_attention=None):
@@ -254,18 +316,15 @@ def cudnn_attention(operands):
             # existing implementation; this switch targets training throughput.
             # Fully masked query rows make cuDNN produce NaN q-gradients; drop
             # those rows from the backward pass instead of rebuilding the mask.
-            q = _stop_gradient_for_fully_masked_queries(q, mask)
-            return jax.nn.dot_product_attention(
-                q,
-                k,
-                v,
-                mask=mask,
-                scale=1.0,
-                implementation="cudnn",
-            )
+            compute_dtype = jnp.dtype(self.cudnn_attention_dtype)
+            if compute_dtype == q.dtype:
+                return _cudnn_attention_call(q, k, v, mask)
+            return _cudnn_attention_in_dtype(q, k, v, mask, compute_dtype)
 
         def explicit_attention(operands):
             q, k, v, mask = operands
+            if self.explicit_attention_fp32:
+                q, k, v = (x.astype(jnp.float32) for x in (q, k, v))
             grouped_q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=self.configs[0].num_kv_heads)
             logits = jnp.einsum("BTKGH,BSKH->BKGTS", grouped_q, k, preferred_element_type=jnp.float32)
 
@@ -273,9 +332,9 @@ def explicit_attention(operands):
             big_neg = -2.3819763e38  # See gemma/modules.py
             masked_logits = jnp.where(mask[:, :, None, :, :], logits, big_neg)
 
-            probs = jax.nn.softmax(masked_logits, axis=-1).astype(dtype)
+            probs = jax.nn.softmax(masked_logits, axis=-1).astype(v.dtype)
 
-            encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v)
+            encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v).astype(dtype)
             return einops.rearrange(encoded, "B T K G H -> B T (K G) H")
 
         if use_cudnn_attention is None:
@@ -346,6 +405,8 @@ class Block(nn.Module):
     use_cudnn_attention: bool = False
     cudnn_attention_layer_start: int = 0
     cudnn_attention_num_layers: int | None = None
+    cudnn_attention_dtype: str = "bfloat16"
+    explicit_attention_fp32: bool = False
 
     dropout: float = 0.0
     dropout_bdims: tuple[int, ...] = ()
@@ -355,7 +416,13 @@ def __call__(self, xs, kv_cache, positions, attn_mask, adarms_cond, layer_index=
         xs = sharding.activation_sharding_constraint(xs)
         drop = nn.Dropout(self.dropout, self.dropout_bdims) if self.dropout else lambda x, _: x
 
-        attn = Attention(configs=self.configs, use_cudnn_attention=self.use_cudnn_attention, name="attn")
+        attn = Attention(
+            configs=self.configs,
+            use_cudnn_attention=self.use_cudnn_attention,
+            cudnn_attention_dtype=self.cudnn_attention_dtype,
+            explicit_attention_fp32=self.explicit_attention_fp32,
+            name="attn",
+        )
         layer_uses_cudnn = self.use_cudnn_attention
         if layer_uses_cudnn and self.cudnn_attention_num_layers is not None:
             layer_uses_cudnn = jnp.logical_and(
@@ -422,6 +489,8 @@ class Module(nn.Module):
     use_cudnn_attention: bool = False
     cudnn_attention_layer_start: int = 0
     cudnn_attention_num_layers: int | None = None
+    cudnn_attention_dtype: str = "bfloat16"
+    explicit_attention_fp32: bool = False
 
     def setup(self):
         # all experts must have the same depth
@@ -456,6 +525,8 @@ def setup(self):
             use_cudnn_attention=self.use_cudnn_attention,
             cudnn_attention_layer_start=self.cudnn_attention_layer_start,
             cudnn_attention_num_layers=self.cudnn_attention_num_layers,
+            cudnn_attention_dtype=self.cudnn_attention_dtype,
+            explicit_attention_fp32=self.explicit_attention_fp32,
             dropout=self.dropout,
             dropout_bdims=self.dropout_bdims,
         )
diff --git a/src/openpi/models/pi0.py b/src/openpi/models/pi0.py
index ebb0363..e60d9f0 100644
--- a/src/openpi/models/pi0.py
+++ b/src/openpi/models/pi0.py
@@ -97,6 +97,8 @@ def __init__(
                 use_cudnn_attention=config.use_cudnn_attention,
                 cudnn_attention_layer_start=config.cudnn_attention_layer_start,
                 cudnn_attention_num_layers=config.cudnn_attention_num_layers,
+                cudnn_attention_dtype=config.cudnn_attention_dtype,
+                explicit_attention_fp32=config.explicit_attention_fp32,
             )
         )
         llm.lazy_init(
diff --git a/src/openpi/models/pi0_config.py b/src/openpi/models/pi0_config.py
index 6cfdfcf..cb48db0 100644
--- a/src/openpi/models/pi0_config.py
+++ b/src/openpi/models/pi0_config.py
@@ -40,6 +40,11 @@ class Pi0Config(_model.BaseModelConfig):
     # otherwise only this contiguous layer range uses cuDNN.
     cudnn_attention_layer_start: int = 0
     cudnn_attention_num_layers: int | None = None
+    # Compute dtype for the cuDNN kernel: "bfloat16" or "float16". float16 uses a
+    # dynamically loss-scaled custom VJP (see gemma._cudnn_attention_in_dtype).
+    cudnn_attention_dtype: str = "bfloat16"
+    # Diagnostic only: explicit attention with fp32 q/k/v/probs as a numerical reference.
+    explicit_attention_fp32: bool = False
 
     # training-time RTC config
     enable_training_time_rtc: bool = False
@@ -54,6 +59,8 @@ def __post_init__(self):
             raise ValueError("cudnn_attention_layer_start must be non-negative")
         if self.cudnn_attention_num_layers is not None and self.cudnn_attention_num_layers < 0:
             raise ValueError("cudnn_attention_num_layers must be non-negative or None")
+        if self.cudnn_attention_dtype not in ("bfloat16", "float16"):
+            raise ValueError("cudnn_attention_dtype must be 'bfloat16' or 'float16'")
 
     @property
     @override

From 3ea2b1c7a0a0c1d42b215c128d63fc6aaa0fc064 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Fri, 4 Sep 2026 00:39:06 +0800
Subject: [PATCH 15/28] Optimize cuDNN FP16 attention VJP

---
 ...026-09-03-cudnn-dterm-rounding-analysis.md | 196 ++++++++
 scripts/attention_dterm_cpu_analysis.py       | 176 +++++++
 scripts/attention_dterm_probe.py              | 449 ++++++++++++++++++
 scripts/attention_precision_update_probe.py   |   9 +-
 scripts/cudnn_attention_microbench.py         | 115 +++++
 scripts/cudnn_fp16_vjp_equivalence.py         | 112 +++++
 src/openpi/models/gemma.py                    |  60 ++-
 7 files changed, 1101 insertions(+), 16 deletions(-)
 create mode 100644 docs/2026-09-03-cudnn-dterm-rounding-analysis.md
 create mode 100644 scripts/attention_dterm_cpu_analysis.py
 create mode 100644 scripts/attention_dterm_probe.py
 create mode 100644 scripts/cudnn_attention_microbench.py
 create mode 100644 scripts/cudnn_fp16_vjp_equivalence.py

diff --git a/docs/2026-09-03-cudnn-dterm-rounding-analysis.md b/docs/2026-09-03-cudnn-dterm-rounding-analysis.md
new file mode 100644
index 0000000..2ab8164
--- /dev/null
+++ b/docs/2026-09-03-cudnn-dterm-rounding-analysis.md
@@ -0,0 +1,196 @@
+# 全层 cuDNN 发散的根因：flash 反向 D 项用了 bf16 舍入后的 O
+
+**日期**：2026-09-03
+
+**机器**：8×H100 80GB（ebcloud，hb-openpi 环境：jax 0.5.3、nvidia-cudnn-cu12 9.19.0.56，`check_cuda_stack.py` PASS）
+
+**前置报告**：`docs/2026-08-31-cudnn-sdpa-nan-summary.md`。那份报告证明了全 18 层 cuDNN 的缓慢发散是"cuDNN 与显式路径的 BF16 差分被 Adam 逐坐标放大"，但没有回答差分从哪里来、为什么它不像随机噪声。本文用真实激活回答这两个问题，并给出一个不改内核的缓解手段及其验收结果。
+
+> [!IMPORTANT]
+> 结论摘要：
+>
+> 1. cuDNN fused attention 的反向按 flash-attention 方式算 `dS = P ⊙ (dP − D)`，其中 `D_i = rowsum(dO_i ⊙ O_i)` 用的是**前向存下的、已舍入到 bf16 的 O**。显式路径用同一份舍入后的 dP 算 `D_i = Σ_j P_ij dP_ij`，抵消是自洽的。两者的差是一个每行标量 `Δ_i = dO_i · (bf16(O_i) − O_i)`，进入梯度后方向为 `−Δ_i · Σ_j P_ij K_j`，破坏了 softmax 梯度"行和为零"的不变量。
+> 2. 该分量在真实激活上可被逐层预测：cuDNN bf16 的 dQ 误差与预测方向的 cos 在全部行上为 0.36–0.55，在 peaked 行（maxP>0.99）上为 0.66–0.99；显式 bf16 与 cuDNN fp16 的 cos 均 ≈ 0。
+> 3. peaked 行的真实 dQ 是 ε 级抵消量，所以那里 cuDNN bf16 的相对误差达 7%–140%（显式路径 0.2%–0.5%）。它们绝对量小，但方向沿该行 argmax 键、且跨 batch 落在同一低维子空间，正是 Adam 会放大的结构性误差。
+> 4. 把 cuDNN 内核改为 fp16（多 3 位尾数，配合 2 的幂动态缩放余切）后，peaked 行误差降约 8 倍，全局 dQ 误差 1.76e-3 低于显式 bf16 的 2.33e-3，结构性分量消失（cos ≈ 0）。
+> 5. "把 V 减去 sink 键的 V 再送内核"无效，因为 peaked 行大多不在看 BOS sink，而是在看 action/text/image 的其他键；BOS 的 |V| 本来就只有 1–5。
+
+## 1. 方法
+
+两个探针，都用 `configs/diag_cudnn_strict_order.yaml` 的模型与数据、pi05_base 权重、一个真实 batch（B=8，单卡）：
+
+- `scripts/attention_dterm_probe.py`（GPU）：monkeypatch `gemma.Block.__call__` 与 `gemma._cudnn_attention_call`，在一次真实训练步里用 `jax.debug.callback` 抓下 18 层的 q/k/v/mask 和反向余切 dO（缓存到 `dterm_capture.npz`，1.5 GB）。然后在单层上以 fp32 显式 attention 为真值，比较五个变体的 dQ/dK/dV：`explicit_bf16`（生产路径）、`cudnn_bf16`（历史路径）、`cudnn_fp16`（候选）、`cudnn_bf16_shift`/`cudnn_fp16_shift`（V 减 V_sink）。同时按上面的公式算 H1 预测误差，用 `jax.lax.reduce_precision` 模拟 bf16 舍入（`astype` 往返会被 XLA 的 `allow_excess_precision` 删掉，预测会恒为零）。
+- `scripts/attention_dterm_cpu_analysis.py`（CPU，float64）：读同一份缓存，对 peaked 行逐行重建精确反向，回答"哪些行 peaked、真实 dQ 有多小、H1 误差占多少"。
+
+## 2. 事实
+
+### 2.1 sink 在哪里
+
+18 层、8 个样本，注意力质量最大的键全部是位置 768 = 文本第一个 token（BOS），占有效注意力质量 8%–59%（第 14–15 层最高）。BOS 键的 |K| 8–13（其他键中位数 17–29），**|V| 只有 0.8–5（其他键中位数 10–58）**，第 0/16/17 层例外（9–17）。
+
+### 2.2 哪些行是 peaked（maxP > 0.99）
+
+peaked 行占有效行 0.1%–7%，且**大多不在看 BOS**：
+
+```text
+层   peaked 行   其中看 BOS   其余在看
+ 0     3978        1270       img 1108, txt 1600
+ 4      131           0       act 131
+ 8       64           0       act 48, txt 16
+ 9       98           0       txt 90, img 8
+14     1885        1802       txt 74, img 9
+15      323         168       img 107, txt 48
+17      561          47       txt 357, img 153（该层非 action 行 dO ≡ 0）
+```
+
+### 2.3 peaked 行上的误差解剖（CPU，float64）
+
+每行几何平均。`|dQ_true|` 是精确梯度范数，`errH1` 是 H1 预测误差，`shift` 表示 V 减 V_768：
+
+```text
+层  行类别                  行数   ε=1−maxP  |dQ_true|   |errH1|/|dQ_true|  cos(errH1, K_argmax)
+ 0  看 BOS, P>0.99          1270   1.0e-3    3.8e-06     0.33 (shift: 9e-4)  1.00
+ 0  看其他键, P>0.99        2708   5.3e-4    1.1e-06     0.64                1.00
+ 4  看其他键(act), P>0.99    131   8.2e-4    1.3e-06     1.09                1.00
+ 8  看其他键, P>0.99          64   4.7e-3    2.0e-05     0.80                1.00
+11  看其他键, P>0.99          13   4.4e-3    3.1e-05     1.04                1.00
+14  看 BOS, P>0.99          1802   6.0e-3    5.6e-08     0.020 (shift: 5e-4) 1.00
+14  看其他键, P>0.99          83   2.0e-3    8.7e-07     0.58                1.00
+16  看其他键, P>0.99           4   4.7e-3    2.3e-10     1.47                1.00
+ *  普通行 (P≤0.9)          3000   ~0.7      0.5–3e-04   1.2e-3 – 2.2e-3     0.80–0.88
+```
+
+- peaked 行的真实 dQ 比普通行小两个数量级（ε 级抵消），H1 误差不随 ε 缩小，所以相对误差达 O(1)。
+- 看 BOS 的行因为 |V_768| 小、O 本身小，H1 只有 2%；shift 能再压 40 倍但没必要。
+- 看其他键的 peaked 行 |O| ≈ |V_argmax| = 13–115，H1 误差 60%–150%，shift（减的是 V_768）对它们完全无效。
+- 普通行上 H1 分量约 1.5e-3 相对、方向沿 argmax 键（cos 0.8）：绝对量占主导，但它是结构性的。
+
+### 2.4 GPU 单层实测：误差与 H1 预测的一致性
+
+`dQ rel` 是全部行的相对 L2 误差，`peaked rel` 只算 maxP>0.99 的行，`cos` 是误差向量与 H1 预测的 cos：
+
+```text
+层   变体           dQ rel    cos(all)   peaked rel   cos(peaked)   H1 预测 peaked rel
+ 0   explicit_bf16  2.46e-3   -0.02      2.8e-3       -0.04
+     cudnn_bf16     3.58e-3   +0.49      3.5e-1       +0.83         0.41
+     cudnn_fp16     1.71e-3   -0.02      3.9e-2       +0.13
+ 4   explicit_bf16  2.39e-3   +0.02      2.8e-3       -0.04
+     cudnn_bf16     4.12e-3   +0.44      5.3e-2       +0.47         0.16
+     cudnn_fp16     1.77e-3   +0.03      1.6e-2       +0.67
+ 8   explicit_bf16  2.88e-3   +0.00      3.5e-3       -0.10
+     cudnn_bf16     4.28e-3   +0.51      5.1e-1       +0.96         0.63
+     cudnn_fp16     1.78e-3   +0.03      7.2e-2       +0.57
+10   explicit_bf16  2.27e-3   -0.01      2.4e-3       +0.56
+     cudnn_bf16     4.36e-3   +0.55      5.0e-1       +0.91         0.53
+     cudnn_fp16     1.76e-3   -0.01      8.3e-2       +0.31
+11   cudnn_bf16     4.70e-3   +0.51      1.21         +0.76         0.86
+     cudnn_fp16     1.77e-3   -0.01      1.3e-1       -0.70
+14   explicit_bf16  2.22e-3   +0.05      2.3e-3       -0.03
+     cudnn_bf16     2.86e-3   +0.12      4.3e-1       +0.63         0.49
+     cudnn_fp16     1.72e-3   +0.04      3.5e-2       +0.13
+16   cudnn_bf16     2.46e-3   +0.39      1.41         +1.00         2.2
+     cudnn_fp16     1.71e-3   +0.01      1.69         -0.92
+18 层几何平均：explicit_bf16 2.33e-3 / cudnn_bf16 3.72e-3 / cudnn_fp16 1.76e-3
+18 层平均 cos(all)：explicit 0.006 / cudnn_bf16 0.418 / cudnn_fp16 0.006
+```
+
+- cuDNN bf16 的误差在**每一层**都与 H1 预测显著同向；显式与 fp16 都不同向。
+- H1 预测的 peaked 行误差量级与实测吻合（0.41/0.35、0.63/0.51、0.53/0.50、0.86/1.21、0.49/0.43、2.2/1.41）。cos 不到 1 是因为 cuDNN 前向的 PV 矩阵乘也把 P 舍入到 bf16，其内部 fp32 O 与本文的 fp32 O 差约 2⁻⁹，两者的舍入决策只是部分相关。
+- 第二个候选机制 H2（dP 先舍入到 bf16 再减 D）的 cos 为 0 或负值，排除。
+- dV 没有 D 项，cuDNN bf16 的 dV 误差（2.5e-3）与显式（1.9e-3）同阶，是对照。
+
+### 2.5 为什么这解释了 8-31 报告的全部现象
+
+- **1.5% 随机乘性噪声无害、cuDNN 差分有害**：H1 分量按行沿 `Σ_j P_ij K_j` 方向、符号随舍入随机，投到参数上落在 `x ⊗ K_argmax` 这类低维子空间，不与 |g| 成比例；sink 机制处于平衡态时这些坐标的真实梯度和二阶矩都很小，Adam 归一化后每步都在那里走满 lr 的随机步，累积到约 1000 步后 sink 结构失稳。乘性噪声在小 |g| 坐标上也小，Adam 看不见。
+- **跨 batch delta 两两 cos ±0.25**：在 30 亿维空间里这不是白噪声，正是"同一子空间、随机符号"。
+- **只融合前 3 层 ≈ 全 18 层**：sink/peaked 结构在早层形成，早层 q/k 投影的系统性扰动穿过后面 15 层。
+- **改 mask、seq_lengths、cuDNN 版本、LR、b1/b2/eps 都没用**：它们都不改变 O 的舍入。
+
+## 3. 缓解手段：fp16 内核 + 动态缩放
+
+`src/openpi/models/gemma.py` 的 `_cudnn_attention_in_dtype`：q/k/v 转 fp16（多 3 位尾数，H1 误差降约 8 倍），custom VJP 里把余切按 2 的幂缩放到 max|dO·s| ∈ [0.25, 0.5) 再送 fp16 反向，fp32 反缩放。真实激活的范围 max|q| ≤ 3.5、|k| ≤ 30、|v| ≤ 84、|dO| 1e-4–1e-2，均在 fp16 范围内。开关：`model.cudnn_attention_dtype: float16`（默认 `bfloat16`，逐位不变）。
+
+单层效果见 2.4：peaked 行误差 0.35 → 0.039（层 0）、0.51 → 0.072（层 8），全局误差低于显式 bf16。
+
+实现有两版：
+
+- **v1（验收 run 用的）**：custom VJP 的反向里用 `jax.vjp` 重新走一遍 `jax.nn.dot_product_attention` 前向再取反向；配合 `nothing_saveable` remat，每层反向比 bf16 路径多一次 attention 前向。实测 8×H100 batch 256 为 1.37 s/step（显式 1.57–1.64，cuDNN bf16 1.05）。
+- **v2（当前代码）**：直接调用 jax 0.5.3 私有的 `jax._src.cudnn.fused_attention_stablehlo._dot_product_attention_fwd_rule / _bwd_rule`（就是 `jax.nn.dot_product_attention(implementation="cudnn")` 自己 custom_vjp 里接的那两个函数），把 softmax stats 和输出当残差保存，反向只跑一次 fused backward。mask→bias（`where(mask, 0, -2<<14)`）、BTNH 布局、静态参数与公共路径一致；全 mask 的 query 行在反向里显式置零 dQ，等价于 bf16 路径的 `_stop_gradient_for_fully_masked_queries`。
+  `scripts/cudnn_fp16_vjp_equivalence.py` 在真实激活（层 0/4/8/14/17）上比较 v2 与 v1：前向逐位相同，dK/dV 逐位相同，dQ 相对差 5e-6–1.2e-5，与同一内核自身的 run-to-run 差（1.6e-6–1.4e-5，原子累加）同阶；全 mask 行 dQ 为 0，无非有限值。步时见第 4 节末。
+
+## 4. 3000 步 strict-order 验收（`configs/diag_cudnn_fp16.yaml`）：通过
+
+与 `diag_cudnn_strict_order` 同数据、seed、batch 顺序、超参，全 18 层 cuDNN fp16（v1 实现），从 pi05_base step 0 起，8×H100，batch 256。2026-09-03 21:00 启动、22:09 结束，3000 步全程有限。显式基线取自 8-31 报告（"显式基线"列是显式 strict-order run；"last3"列是只融合后 3 层、已判定与显式一致的 run，点更全）。
+
+```text
+step    fp16 loss / grad_norm    显式基线           last3 基线         历史 cuDNN bf16
+   0    6.4030 /  65.59          6.4023 / 65.63     6.4029 / 65.52     6.4023 / 65.63
+ 200    0.9355 /   3.23          0.9244 /  2.89                        0.9348 /  3.53
+ 500    0.5717 /   2.57                             0.5719 /  2.45     0.5719 /  2.55
+1000    0.4231 /   2.21          0.4220 /  2.77     0.4240 /  2.38     0.4319 /  3.16
+1100    0.4071 /   2.30          0.4077 /  1.79     0.4050 /  2.18     0.4723 /  3.37
+1200    0.3759 /   2.36          0.3742 /  2.12     0.3757 /  2.39     0.6600 / 35.49
+1600    0.3209 /   2.77                             0.3238 /  2.65
+2000    0.2792 /   1.93                             0.2793 /  1.83     3.7753 / 150.5
+2500    0.2464 /   1.70          0.2457 /  1.33     0.2463 /  2.01     4.7055 / 759.2
+2900    0.2427 /   1.59                             0.2424 /  1.48     6.2039 / 603.5
+```
+
+- 全部 30 个 log 点有限，grad_norm 最大值 4.75（step 100），之后 1.4–2.8，没有一个点超过基线同点的 1.3 倍（验收上限 3 倍）。
+- step 1000–1500 历史分叉点处 loss 与显式基线差 <0.5%，与 last3 基线差 <0.3%；step 2500/2900 的 loss 0.2464/0.2427 对基线 0.2457–0.2463/0.2424。历史 cuDNN bf16 在同一点是 4.7/6.2。
+- 3000 步全程 param_norm 1802.39 → 1803.19，与 last3 基线的 1803.19 一致。
+- step 1000 的 checkpoint 保存在远端 `checkpoints/diag_cudnn_fp16/diag_cudnn_fp16/1000`，可作 `attention_precision_update_probe.py` 的锚点。
+
+结论：把 cuDNN 内核的计算精度从 bf16 换成 fp16 就消除了全层 cuDNN 的发散，验证了第 2 节的机制。
+
+### 4.1 步时
+
+8×H100，batch 256，`[TIMING]` 行的 total（不含 log 的步）以及相邻 100 步 log 的墙钟：
+
+```text
+路径                                      s/step       来源
+显式 attention                             1.57–1.64    8-31 报告
+cuDNN bf16（发散）                          1.05         8-31 报告（mask 修复时的短跑）
+cuDNN fp16 v1（反向重算前向）                1.34–1.37    本次 3000 步验收 run
+cuDNN fp16 v2（直接调 fwd/bwd 规则）         1.31–1.35    `diag_cudnn_fp16_speed` 400 步（loss 逐点与 v1 一致：step 300 0.7243 对 0.7245）
+cuDNN bf16，今日同 config 复测              1.32–1.35    `diag_cudnn_bf16_speed` 400 步（step 300 loss 0.7250）
+显式 attention，今日同 config 复测          1.60–1.63    `diag_explicit_speed` 400 步（step 300 loss 0.7244）
+```
+
+去掉反向重算只省了约 2%。单卡微基准（`scripts/cudnn_attention_microbench.py`，层 0 真实激活平铺到每卡 batch 32，T=S=1018，8 q 头 / 1 kv 头，含 `nothing_saveable` remat）说明原因：
+
+```text
+变体                fwd      fwd+bwd(remat)   ×18 层
+cudnn_bf16         1.35 ms   5.24 ms          0.094 s/step
+cudnn_fp16_plain   1.53 ms   5.43 ms          0.098 s/step   （fp16 内核、jax 自带 VJP、无缩放）
+cudnn_fp16_v2      1.54 ms   5.53 ms          0.099 s/step
+explicit_bf16      4.46 ms   8.49 ms          0.153 s/step
+```
+
+18 层 attention 的前向加反向合计只占每步约 0.1 s；fp16 内核与 bf16 内核同速，动态缩放和直接调规则的开销可忽略；显式路径也只多 0.06 s。所以 attention 内核本身解释不了 1.33 s 对 1.05 s 的差距，也解释不了显式路径 1.6 s 对 cuDNN 的差距——差距来自整图层面（XLA 在显式路径下的 fusion/remat 决策、显存压力），要用同 config 的 bf16 复测和整步 profile 来定位，见上表最后一行。
+
+### 4.2 更新空间探针：单步指标不区分好坏
+
+`scripts/attention_precision_update_probe.py diag_cudnn_fp16 diag_cudnn_fp16 1000`，锚点是 fp16 验收 run 的 step-1000 train state（参数 + Adam 矩），batch 32，同一真实 batch 和 rng，以 fp32 显式 attention 为真值：
+
+```text
+batch 0            原始梯度 vs fp32：rel / sign_flips     Adam 更新 vs fp32：rel / cos / sign_flips
+explicit_bf16      3.37e-2 / 1.53e-2                      3.13e-2 / 0.99951 / 7.47e-3
+cudnn_bf16         3.18e-2 / 1.43e-2                      2.93e-2 / 0.99957 / 6.97e-3
+cudnn_fp16         3.01e-2 / 1.38e-2                      2.81e-2 / 0.99961 / 6.73e-3
+```
+
+会发散的 cudnn_bf16 与两条稳定路径在这个指标上不可区分（甚至略好于 explicit_bf16）。这与第 2 节一致：D 项误差在范数上只占 ~1e-3，它的危害来自方向的结构性（同一低维子空间、跨 step 累积），单步的 rel/cos/sign_flips 看不见。判定 fused attention 后端好坏要看 2.4 节的 cos(err, H1) 与 peaked 行误差，最终靠 3000 步 strict-order run；不要再用单步梯度/更新距离做验收指标。
+
+## 5. 备选方案（fp16 已通过验收，仅作记录）
+
+1. **前 k 层显式 + 其余 fp16 cuDNN**：8-31 报告的分段实验表明前 3 层贡献了几乎全部更新差异。
+2. **解析修正**：Δ_i 可用一次 fp16 前向估计（`dO·(O_bf16 − O_fp16)`），dQ 修正需一次 V:=K 的前向，dK 修正需一次 dO:=Δ⊙Q 的反向，总成本抵消吞吐收益，不推荐。
+3. 不要再试：减 V_sink 的 shift（2.3 节）、改 mask、改 Adam 参数。
+
+## 6. 相关脚本
+
+- `scripts/attention_dterm_probe.py`：抓真实激活 + 单层五变体误差 + H1/H2 预测；
+- `scripts/attention_dterm_cpu_analysis.py`：peaked 行的 float64 解剖；
+- `scripts/cudnn_fp16_vjp_equivalence.py`：v2（直接调规则）与 v1（反向重算）fp16 VJP 在真实激活上的等价性检查；
+- `scripts/attention_precision_update_probe.py [config] [exp_name] [step]`：真实 Adam 状态上的更新空间对比，锚点用 fp16 验收 run 的 step-1000 checkpoint。
diff --git a/scripts/attention_dterm_cpu_analysis.py b/scripts/attention_dterm_cpu_analysis.py
new file mode 100644
index 0000000..19d55ec
--- /dev/null
+++ b/scripts/attention_dterm_cpu_analysis.py
@@ -0,0 +1,176 @@
+"""CPU-only follow-up to attention_dterm_probe.py: anatomy of the peaked attention rows.
+
+Loads the captured q/k/v/mask/d_out of every layer (dterm_capture.npz) and, in float64,
+reconstructs the exact single-layer backward for the peaked rows to answer:
+
+  * how big are K/V of the BOS sink key (position 768) relative to other keys?
+  * which rows are peaked (maxP > 0.99): BOS-attending rows or other keys?
+  * on BOS rows, how much cancellation does the true dQ have, and how large is the
+    D-term error from rounding the stored O to bf16 (H1) relative to it -- with and
+    without shifting V by V_sink?
+
+Pure numpy; run anywhere the npz is:  python scripts/attention_dterm_cpu_analysis.py
+"""
+
+import sys
+
+import ml_dtypes
+import numpy as np
+
+CAPTURE_FILE = sys.argv[1] if len(sys.argv) > 1 else "/root/localstorage/logs/dterm_capture.npz"
+BIG_NEG = -2.3819763e38
+SINK = 768
+MAX_ROWS = 3000
+rng = np.random.default_rng(0)
+
+
+def region(pos):
+    if pos < 768:
+        return f"img{pos // 256}"
+    if pos < 968:
+        return "txt"
+    return "act"
+
+
+def bf16_round(x):
+    return x.astype(ml_dtypes.bfloat16).astype(np.float64)
+
+
+def load():
+    data = np.load(CAPTURE_FILE)
+    layers = {}
+    for key in data.files:
+        layer, name = key.split("/")
+        arr = data[key]
+        if arr.dtype.kind == "V" and arr.dtype.itemsize == 2:  # bf16 saved by np.savez as raw bytes
+            arr = arr.view(ml_dtypes.bfloat16)
+        layers.setdefault(int(layer), {})[name] = arr
+    return layers
+
+
+def row_anatomy(q_rows, k, v, mask_rows, d_rows, shift=None):
+    """q_rows/d_rows: R x H ; k, v: S x H (single kv head) ; mask_rows: R x S bool. float64."""
+    s = q_rows @ k.T  # R S
+    s = np.where(mask_rows, s, -np.inf)
+    s -= s.max(axis=1, keepdims=True)
+    p = np.exp(s)
+    p /= p.sum(axis=1, keepdims=True)
+    v_eff = v if shift is None else v - shift[None, :]
+    o = p @ v_eff  # R H   (fp64 "true" output of the kernel-internal V)
+    dp = d_rows @ v_eff.T  # R S
+    d_true = np.sum(p * dp, axis=1)  # R
+    ds = p * (dp - d_true[:, None])
+    dq = ds @ k  # R H
+    # H1: kernel D from bf16-rounded O
+    delta = np.sum(d_rows * (bf16_round(o) - o), axis=1)  # R
+    pk = p @ k  # R H
+    err_h1 = -delta[:, None] * pk
+    return dict(p=p, o=o, dp=dp, d_true=d_true, ds=ds, dq=dq, delta=delta, pk=pk, err_h1=err_h1)
+
+
+def norms(x):
+    return np.sqrt(np.sum(x * x, axis=-1))
+
+
+def main():
+    layers = load()
+    print(f"loaded {len(layers)} layers from {CAPTURE_FILE}", flush=True)
+    for layer in sorted(layers):
+        e = layers[layer]
+        q = e["q"].astype(np.float64)  # B T N H
+        k = e["k"].astype(np.float64)[:, :, 0]  # B S H
+        v = e["v"].astype(np.float64)[:, :, 0]
+        mask = e["mask"][:, 0]  # B T S
+        d = e["d_out"].astype(np.float64)  # B T N H
+        B, T, N, H = q.shape
+        valid_q = mask.any(-1)  # B T
+        valid_k = mask.any(1)  # B S
+
+        kn, vn = norms(k), norms(v)  # B S
+        med_k = np.median(kn[valid_k])
+        med_v = np.median(vn[valid_k])
+        print(
+            f"\n=== layer {layer:2d} ===  |K_768|={kn[:, SINK].mean():6.2f} (median valid |K|={med_k:6.2f})"
+            f"   |V_768|={vn[:, SINK].mean():6.2f} (median valid |V|={med_v:6.2f})",
+            flush=True,
+        )
+
+        # per-row softmax stats for all valid rows (float32 logits are exact enough for argmax/maxP)
+        s = np.einsum("btnh,bsh->btns", q.astype(np.float32), k.astype(np.float32))
+        s = np.where(mask[:, :, None, :], s, BIG_NEG)
+        s -= s.max(-1, keepdims=True)
+        p = np.exp(s)
+        p /= p.sum(-1, keepdims=True)
+        max_p = p.max(-1)  # B T N
+        arg_p = p.argmax(-1)
+        del s, p
+        vrow = np.broadcast_to(valid_q[:, :, None], max_p.shape)
+        peaked = (max_p > 0.99) & vrow
+        bos_rows = (arg_p == SINK) & vrow
+        n_peaked = int(peaked.sum())
+        n_peaked_bos = int((peaked & bos_rows).sum())
+        hist = {}
+        for pos in arg_p[peaked & ~bos_rows]:
+            hist[region(int(pos))] = hist.get(region(int(pos)), 0) + 1
+        print(
+            f"  peaked rows (maxP>0.99): {n_peaked}  of which BOS(768): {n_peaked_bos}; other peaked keys by region: {hist}"
+            f"   BOS-argmax rows: {int(bos_rows.sum())} ({bos_rows.sum() / vrow.sum():.1%} of valid rows),"
+            f" P>0.9: {int((bos_rows & (max_p > 0.9)).sum())}, P>0.99: {n_peaked_bos}",
+            flush=True,
+        )
+
+        def analyse(sel, label, shift_by_sink):
+            idx = np.argwhere(sel)
+            if len(idx) == 0:
+                print(f"  {label}: no rows", flush=True)
+                return
+            if len(idx) > MAX_ROWS:
+                idx = idx[rng.choice(len(idx), MAX_ROWS, replace=False)]
+            out = {"plain": [], "shift": []}
+            for b in np.unique(idx[:, 0]):
+                rows = idx[idx[:, 0] == b]
+                t_idx, n_idx = rows[:, 1], rows[:, 2]
+                q_rows = q[b, t_idx, n_idx]
+                d_rows = d[b, t_idx, n_idx]
+                m_rows = mask[b, t_idx]
+                for key, shift in (("plain", None), ("shift", v[b, SINK] if shift_by_sink else None)):
+                    if key == "shift" and not shift_by_sink:
+                        continue
+                    r = row_anatomy(q_rows, k[b], v[b], m_rows, d_rows, shift=shift)
+                    eps = 1.0 - r["p"].max(1)
+                    dq_n = norms(r["dq"])
+                    naive = eps * norms(d_rows) * med_v * med_k  # first-order scale of a peaked row's dQ
+                    ds_sink = np.abs(np.take_along_axis(r["ds"], r["p"].argmax(1)[:, None], 1)[:, 0])
+                    h1_n = norms(r["err_h1"])
+                    k_arg = k[b][r["p"].argmax(1)]  # R H
+                    cos_k = np.sum(r["err_h1"] * k_arg, 1) / np.maximum(h1_n * norms(k_arg), 1e-300)
+                    out[key].append(
+                        np.stack(
+                            [eps, dq_n, naive, ds_sink * norms(k_arg), np.abs(r["d_true"]), np.abs(r["delta"]), h1_n, cos_k, norms(r["o"])],
+                            1,
+                        )
+                    )
+            for key, chunks in out.items():
+                if not chunks:
+                    continue
+                a = np.concatenate(chunks)
+                eps, dq_n, naive, dssink_k, d_true, delta, h1_n, cos_k, o_n = a.T
+                gm = lambda x: float(np.exp(np.mean(np.log(np.maximum(x, 1e-300)))))
+                print(
+                    f"  {label:34s} [{key:5s}] rows={len(a):5d}  eps={gm(eps):.2e}  |dQ_true|={gm(dq_n):.2e}"
+                    f"  naive eps|dO||V||K|={gm(naive):.2e}  |dS_sink||K|={gm(dssink_k):.2e}  |D|={gm(d_true):.2e}"
+                    f"  |delta_H1|={gm(delta):.2e}  |O|={gm(o_n):.2e}"
+                    f"  |errH1|/|dQ_true|: gm={gm(h1_n / np.maximum(dq_n, 1e-300)):.2e} med={np.median(h1_n / np.maximum(dq_n, 1e-300)):.2e}"
+                    f"  cos(errH1,K_argmax)={np.mean(np.abs(cos_k)):.2f}",
+                    flush=True,
+                )
+
+        analyse(bos_rows & (max_p > 0.99), "BOS rows P>0.99", True)
+        analyse(bos_rows & (max_p > 0.9) & (max_p <= 0.99), "BOS rows 0.9<P<=0.99", True)
+        analyse(bos_rows & (max_p <= 0.9), "BOS rows P<=0.9", True)
+        analyse(peaked & ~bos_rows, "other peaked rows P>0.99", False)
+        analyse(vrow & ~bos_rows & (max_p <= 0.9), "ordinary rows (argmax!=768,P<=0.9)", False)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/attention_dterm_probe.py b/scripts/attention_dterm_probe.py
new file mode 100644
index 0000000..f0b391d
--- /dev/null
+++ b/scripts/attention_dterm_probe.py
@@ -0,0 +1,449 @@
+"""Test the flash-attention D-term hypothesis on real Pi0.5 activations.
+
+Hypothesis (docs/2026-08-31-cudnn-sdpa-nan-summary.md, open question): the cuDNN
+fused backward forms ``dS = P * (dP - D)`` with ``D_i = rowsum(dO_i * O_i)`` taken
+from the *bf16-rounded* stored output ``O``. On attention-sink rows (P nearly
+one-hot) ``dP`` and ``D`` cancel almost exactly, so the rounding error of ``O`` is
+a large fraction of ``dS`` and breaks the softmax invariant ``sum_j dS_ij = 0``.
+The explicit path forms ``D = sum_j P_ij dP_ij`` from the same rounded ``dP`` and
+keeps the invariant.
+
+If that is the mechanism, the cuDNN gradient error must be *predictable*::
+
+    delta_i   = dO_i . (bf16(O_i) - O_i)               (per query row / head)
+    err dQ_i  = -delta_i * sum_j P_ij K_j
+    err dK_j  = -sum_i delta_i P_ij Q_i
+
+This script captures q/k/v/mask and the true backward cotangent dO of every Gemma
+layer from a real training step (real batch, pi05_base weights), then on each layer
+compares, against an fp32 explicit reference:
+
+    explicit_bf16      production path (known stable)
+    cudnn_bf16         historical cuDNN path (known to diverge)
+    cudnn_fp16         cuDNN kernel in float16 + dynamic scaling (candidate)
+    cudnn_bf16_shift   cuDNN with V shifted by the sink key's V (candidate; makes
+                       O of sink rows small so its rounding error shrinks)
+    cudnn_fp16_shift   both
+
+and reports rel error, the cosine between each variant's error and the predicted
+D-term error, and the residual after subtracting the prediction. It also prints
+where the attention sinks are (position histogram) and the fp16 range headroom.
+
+Run on ONE GPU:
+    CUDA_VISIBLE_DEVICES=7 XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 PYTHONPATH=$PWD/src \
+        python scripts/attention_dterm_probe.py
+"""
+
+import dataclasses
+import functools
+import logging
+import sys
+import time
+
+import einops
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import openpi.models.gemma as _gemma
+import openpi.training.config as _config
+import openpi.training.data_loader as _data_loader
+import openpi.training.sharding as sharding
+
+sys.path.insert(0, "scripts")
+import train as train_mod  # noqa: E402
+
+CONFIG = "diag_cudnn_strict_order"
+CAPTURE_FILE = "/root/localstorage/logs/dterm_capture.npz"
+BATCH = 8
+WORKERS = 8
+BIG_NEG = -2.3819763e38
+NUM_LAYERS = 18
+
+# --------------------------------------------------------------------------------------
+# Capture q/k/v/mask/d_out of every layer from a real train step.
+# --------------------------------------------------------------------------------------
+_CAPTURE: dict[int, dict[str, np.ndarray]] = {}
+_LAYER = None  # tracer of the layer index of the Block currently being traced
+
+_orig_block_call = _gemma.Block.__call__
+_orig_cudnn_call = _gemma._cudnn_attention_call
+
+
+@functools.wraps(_orig_block_call)
+def _block_call(self, xs, kv_cache, positions, attn_mask, adarms_cond, layer_index=None, deterministic=True):
+    global _LAYER
+    _LAYER = layer_index
+    return _orig_block_call(self, xs, kv_cache, positions, attn_mask, adarms_cond, layer_index, deterministic)
+
+
+def _save_fwd(layer, q, k, v, mask):
+    entry = _CAPTURE.setdefault(int(layer), {})
+    entry.update(q=np.asarray(q), k=np.asarray(k), v=np.asarray(v), mask=np.asarray(mask))
+
+
+def _save_bwd(layer, d_out):
+    _CAPTURE.setdefault(int(layer), {})["d_out"] = np.asarray(d_out)
+
+
+@jax.custom_vjp
+def _tap(q, k, v, mask, layer):
+    return _orig_cudnn_call(q, k, v, mask)
+
+
+def _tap_fwd(q, k, v, mask, layer):
+    jax.debug.callback(_save_fwd, layer, q, k, v, mask)
+    return _tap(q, k, v, mask, layer), (q, k, v, mask, layer)
+
+
+def _tap_bwd(res, d_out):
+    q, k, v, mask, layer = res
+    jax.debug.callback(_save_bwd, layer, d_out)
+    _, vjp_fn = jax.vjp(lambda a, b, c: _orig_cudnn_call(a, b, c, mask), q, k, v)
+    dq, dk, dv = vjp_fn(d_out)
+    return dq, dk, dv, None, None
+
+
+_tap.defvjp(_tap_fwd, _tap_bwd)
+
+
+def _tapped_cudnn_call(q, k, v, attn_mask):
+    return _tap(q, k, v, attn_mask, _LAYER)
+
+
+def capture_real_activations(cfg):
+    _gemma.Block.__call__ = _block_call
+    _gemma._cudnn_attention_call = _tapped_cudnn_call
+    try:
+        mesh = sharding.make_mesh(cfg.fsdp_devices)
+        data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+        rng = jax.random.key(cfg.seed)
+        train_rng, init_rng = jax.random.split(rng)
+        state, _ = train_mod.init_train_state(cfg, init_rng, mesh, resume=False)
+        params, model_def = state.params, state.model_def
+        del state  # free the Adam moments
+        jax.block_until_ready(params)
+        print("weights loaded", flush=True)
+
+        loader = _data_loader.create_data_loader(cfg, sharding=data_sharding, shuffle=False)
+        obs, act = next(iter(loader))
+        print("batch loaded", flush=True)
+
+        def loss_fn(model, r, o, a):
+            return jnp.mean(model.compute_loss(r, o, a, train=True))
+
+        @jax.jit
+        def run(params, r, o, a):
+            model = nnx.merge(model_def, params)
+            model.train()
+            return nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, cfg.trainable_filter))(model, r, o, a)
+
+        with sharding.set_mesh(mesh):
+            loss, grads = run(params, jax.random.fold_in(train_rng, 0), obs, act)
+            jax.block_until_ready(grads)
+        print(f"real train step: loss={float(loss):.6f}", flush=True)
+        del grads, params
+    finally:
+        _gemma.Block.__call__ = _orig_block_call
+        _gemma._cudnn_attention_call = _orig_cudnn_call
+    missing = [i for i in range(NUM_LAYERS) if i not in _CAPTURE or "d_out" not in _CAPTURE[i]]
+    if missing:
+        raise RuntimeError(f"capture incomplete, missing layers {missing}; got {sorted(_CAPTURE)}")
+
+
+# --------------------------------------------------------------------------------------
+# Single-layer attention variants (mirror gemma.Attention exactly).
+# --------------------------------------------------------------------------------------
+def explicit_attention(q, k, v, mask, compute_dtype, probs_dtype, out_dtype):
+    num_kv = k.shape[2]
+    q, k, v = (x.astype(compute_dtype) for x in (q, k, v))
+    grouped_q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=num_kv)
+    logits = jnp.einsum("BTKGH,BSKH->BKGTS", grouped_q, k, preferred_element_type=jnp.float32)
+    masked = jnp.where(mask[:, :, None, :, :], logits, BIG_NEG)
+    probs = jax.nn.softmax(masked, axis=-1).astype(probs_dtype)
+    enc = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v).astype(out_dtype)
+    return einops.rearrange(enc, "B T K G H -> B T (K G) H")
+
+
+def _shift_v(v, sink_idx):
+    """Return v - v[:, sink] (rounded once) and the per-(B,K) shift c (stop-gradient)."""
+    c = jax.lax.stop_gradient(jnp.take_along_axis(v, sink_idx[:, None, :, None], axis=1))  # B 1 K H
+    v_shift = (v.astype(jnp.float32) - c.astype(jnp.float32)).astype(v.dtype)
+    return v_shift, c
+
+
+def cudnn_shift_attention(q, k, v, mask, sink_idx, compute_dtype):
+    v_shift, c = _shift_v(v, sink_idx)
+    if jnp.dtype(compute_dtype) == q.dtype:
+        out = _orig_cudnn_call(q, k, v_shift, mask)
+    else:
+        out = _gemma._cudnn_attention_in_dtype(q, k, v_shift, mask, compute_dtype)
+    has_key = jnp.any(mask, axis=-1)[:, 0, :, None, None]  # B T 1 1
+    c_full = einops.repeat(c, "B 1 K H -> B 1 (K G) H", G=q.shape[2] // k.shape[2])
+    return (out.astype(jnp.float32) + jnp.where(has_key, c_full.astype(jnp.float32), 0.0)).astype(out.dtype)
+
+
+VARIANT_NAMES = ("explicit_bf16", "cudnn_bf16", "cudnn_fp16", "cudnn_bf16_shift", "cudnn_fp16_shift")
+
+
+def _variant_output(name, q, k, v, mask, sink_idx):
+    if name == "explicit_bf16":
+        return explicit_attention(q, k, v, mask, jnp.bfloat16, jnp.bfloat16, jnp.bfloat16)
+    if name == "cudnn_bf16":
+        return _orig_cudnn_call(q, k, v, mask)
+    if name == "cudnn_fp16":
+        return _gemma._cudnn_attention_in_dtype(q, k, v, mask, jnp.float16)
+    if name == "cudnn_bf16_shift":
+        return cudnn_shift_attention(q, k, v, mask, sink_idx, jnp.bfloat16)
+    if name == "cudnn_fp16_shift":
+        return cudnn_shift_attention(q, k, v, mask, sink_idx, jnp.float16)
+    raise ValueError(name)
+
+
+def grads_of(fn, q, k, v, mask, d_out):
+    d32 = d_out.astype(jnp.float32)
+
+    def objective(q_, k_, v_):
+        return jnp.sum(fn(q_, k_, v_, mask).astype(jnp.float32) * d32)
+
+    return jax.grad(objective, argnums=(0, 1, 2))(q, k, v)
+
+
+@functools.partial(jax.jit, static_argnums=0)
+def variant_grads(name, q, k, v, mask, d_out, sink_idx):
+    return grads_of(lambda a, b, c, m: _variant_output(name, a, b, c, m, sink_idx), q, k, v, mask, d_out)
+
+
+@jax.jit
+def reference_and_prediction(q, k, v, mask, d_out):
+    """fp32 explicit gradients + the predicted D-term error of a bf16 flash backward."""
+    num_kv = k.shape[2]
+    q32, k32, v32, d32 = (x.astype(jnp.float32) for x in (q, k, v, d_out))
+    ref = grads_of(
+        lambda a, b, c, m: explicit_attention(a, b, c, m, jnp.float32, jnp.float32, jnp.float32), q32, k32, v32, mask, d_out
+    )
+
+    grouped_q = einops.rearrange(q32, "B T (K G) H -> B T K G H", K=num_kv)
+    logits = jnp.einsum("BTKGH,BSKH->BKGTS", grouped_q, k32, preferred_element_type=jnp.float32)
+    masked = jnp.where(mask[:, :, None, :, :], logits, BIG_NEG)
+    probs = jax.nn.softmax(masked, axis=-1)  # B K G T S
+    o32 = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v32)
+    # reduce_precision emulates the bf16 rounding; a plain astype round-trip is deleted by XLA
+    # (xla_allow_excess_precision) and would make the prediction identically zero.
+    o16 = jax.lax.reduce_precision(o32, exponent_bits=8, mantissa_bits=7)
+    d_o = einops.rearrange(d32, "B T (K G) H -> B T K G H", K=num_kv)
+    # H1: D formed from the bf16-rounded stored output O.
+    delta = jnp.sum(d_o * (o16 - o32), axis=-1)  # B T K G
+    pk = jnp.einsum("BKGTS,BSKH->BTKGH", probs, k32)
+    pred_dq = einops.rearrange(-delta[..., None] * pk, "B T K G H -> B T (K G) H")
+    pred_dk = -jnp.einsum("BTKG,BKGTS,BTKGH->BSKH", delta, probs, grouped_q)
+    # H2: dP = dO V^T rounded to bf16 before dS = P (dP - D) with an exact D.
+    dp = jnp.einsum("BTKGH,BSKH->BKGTS", d_o, v32)
+    ddp = jax.lax.reduce_precision(dp, exponent_bits=8, mantissa_bits=7) - dp
+    pred2_dq = einops.rearrange(jnp.einsum("BKGTS,BSKH->BTKGH", probs * ddp, k32), "B T K G H -> B T (K G) H")
+
+    valid_q = jnp.any(mask, axis=-1)[:, 0, :]  # B T
+    max_p = jnp.max(probs, axis=-1)  # B K G T
+    arg_p = jnp.argmax(probs, axis=-1)  # B K G T
+    mass = jnp.einsum("BKGTS,BT->BKS", probs, valid_q.astype(jnp.float32))  # B K S
+    return ref, (pred_dq, pred_dk, pred2_dq), dict(valid_q=valid_q, max_p=max_p, arg_p=arg_p, mass=mass, delta=delta)
+
+
+# --------------------------------------------------------------------------------------
+# Metrics (float64 on host).
+# --------------------------------------------------------------------------------------
+def _f64(x):
+    return np.asarray(jax.device_get(x)).astype(np.float64)
+
+
+def _norm(x):
+    return float(np.sqrt(np.sum(x * x)))
+
+
+def _cos(a, b):
+    return float(np.sum(a * b) / max(_norm(a) * _norm(b), 1e-300))
+
+
+def _region(pos):
+    if pos < 768:
+        return f"img{pos // 256}[{pos % 256}]"
+    if pos < 968:
+        return f"txt[{pos - 768}]"
+    return f"act[{pos - 968}]"
+
+
+def analyse_layer(layer, entry):
+    q, k, v, mask, d_out = (jnp.asarray(entry[n]) for n in ("q", "k", "v", "mask", "d_out"))
+    ref, (pred_dq, pred_dk, pred2_dq), stats = reference_and_prediction(q, k, v, mask, d_out)
+    ref = [_f64(g) for g in ref]
+    pred = [_f64(pred_dq), _f64(pred_dk)]
+    pred2 = _f64(pred2_dq)
+    k_np = _f64(k)  # B S K H
+    valid_q = np.asarray(stats["valid_q"])
+    max_p = np.asarray(stats["max_p"])  # B K G T
+    arg_p = np.asarray(stats["arg_p"])
+    mass = np.asarray(stats["mass"])  # B K S
+
+    # ---- sink statistics ----
+    d_np = np.asarray(entry["d_out"]).astype(np.float32)
+    masked_rows_dout = float(np.abs(d_np[~valid_q]).max()) if (~valid_q).any() else 0.0
+    vq = np.broadcast_to(valid_q[:, None, None, :], max_p.shape)
+    mp = max_p[vq]
+    ap = arg_p[vq]
+    frac = {t: float(np.mean(mp > t)) for t in (0.9, 0.99, 0.999)}
+    # per (B, K) sink key = argmax of attention mass
+    sink_idx = np.argmax(mass, axis=-1)  # B K
+    sink_share = mass.max(axis=-1) / mass.sum(axis=-1)  # fraction of all valid attention mass on the sink key
+    pos_counts = np.bincount(ap, minlength=q.shape[1])
+    top = np.argsort(pos_counts)[::-1][:4]
+    top_desc = ", ".join(f"{_region(int(p))}={pos_counts[p] / ap.size:.2%}" for p in top)
+    print(
+        f"\n=== layer {layer:2d} ===  valid rows maxP>0.9/0.99/0.999: {frac[0.9]:.3f}/{frac[0.99]:.3f}/{frac[0.999]:.3f}"
+        f"   sink key per (b,kv): {sorted(set(sink_idx.ravel().tolist()))} share={sink_share.mean():.2%}"
+        f"   argmax top: {top_desc}"
+        f"   max|q|={float(jnp.abs(q).max()):.1f} |k|={float(jnp.abs(k).max()):.1f} |v|={float(jnp.abs(v).max()):.1f}"
+        f" |dO|={float(jnp.abs(d_out).max()):.2e} |dO| on masked rows={masked_rows_dout:.1e}",
+        flush=True,
+    )
+
+    # ---- gradient variants ----
+    sink_idx_dev = jnp.asarray(sink_idx)
+
+    sink_rows = np.transpose(max_p > 0.99, (0, 3, 1, 2)).reshape(max_p.shape[0], max_p.shape[3], -1)  # B T N
+    arg_rows = np.transpose(arg_p, (0, 3, 1, 2)).reshape(sink_rows.shape)  # B T N: each row's own peak key
+    valid_rows = np.broadcast_to(valid_q[:, :, None], sink_rows.shape)
+    print(
+        f"  {'variant':18s} {'dQ rel':>9s} {'cos(err,pred)':>13s} {'resid':>9s} {'sinkrow rel':>11s} {'other rel':>9s} |"
+        f" {'dK rel':>9s} {'cos':>7s} {'resid':>9s} | {'dV rel':>9s} | nonfinite",
+        flush=True,
+    )
+    results = {}
+    for name in VARIANT_NAMES:
+        t0 = time.time()
+        g = variant_grads(name, q, k, v, mask, d_out, sink_idx_dev)
+        g = [_f64(x) for x in g]
+        nonfinite = sum(int(np.count_nonzero(~np.isfinite(x))) for x in g)
+        err = [g[i] - ref[i] for i in range(3)]
+        row = {}
+        for i, nm in enumerate(("dq", "dk", "dv")):
+            row[f"{nm}_rel"] = _norm(err[i]) / _norm(ref[i])
+            if i < 2:
+                row[f"{nm}_cos"] = _cos(err[i], pred[i])
+                row[f"{nm}_resid"] = _norm(err[i] - pred[i]) / _norm(ref[i])
+        # dQ error split by sink rows / other valid rows
+        e_rows = np.sqrt(np.sum(err[0] ** 2, axis=-1))  # B T N
+        r_rows = np.sqrt(np.sum(ref[0] ** 2, axis=-1))
+        sel_s = sink_rows & valid_rows
+        sel_o = (~sink_rows) & valid_rows
+        row["dq_sink_rel"] = _norm(e_rows[sel_s]) / max(_norm(r_rows[sel_s]), 1e-300)
+        row["dq_other_rel"] = _norm(e_rows[sel_o]) / max(_norm(r_rows[sel_o]), 1e-300)
+            # --- peaked-row structure: is the error along the row's peak key (a softmax row-sum violation)? ---
+        if sel_s.any():
+            e_s = err[0][sel_s]  # R H
+            b_idx = np.nonzero(sel_s)[0]
+            ks = k_np[b_idx, arg_rows[sel_s], 0, :]  # R H: K of each peaked row's own argmax key (single kv head)
+            ks_unit = ks / np.linalg.norm(ks, axis=-1, keepdims=True)
+            along = np.sum(e_s * ks_unit, axis=-1)
+            row["sink_frac_along_ksink"] = float(np.sum(along**2) / max(np.sum(e_s**2), 1e-300))
+            row["sink_cos_h1"] = _cos(e_s, pred[0][sel_s])
+            row["sink_cos_h2"] = _cos(e_s, pred2[sel_s])
+            row["sink_cos_h12"] = _cos(e_s, pred[0][sel_s] + pred2[sel_s])
+            row["sink_resid_h1"] = _norm(e_s - pred[0][sel_s]) / max(_norm(r_rows[sel_s]), 1e-300)
+            row["sink_resid_h12"] = _norm(e_s - pred[0][sel_s] - pred2[sel_s]) / max(_norm(r_rows[sel_s]), 1e-300)
+        else:
+            for key in ("sink_frac_along_ksink", "sink_cos_h1", "sink_cos_h2", "sink_cos_h12", "sink_resid_h1", "sink_resid_h12"):
+                row[key] = float("nan")
+        results[name] = row
+        print(
+            f"  {name:18s} {row['dq_rel']:9.3e} {row['dq_cos']:13.4f} {row['dq_resid']:9.3e} {row['dq_sink_rel']:11.3e}"
+            f" {row['dq_other_rel']:9.3e} | {row['dk_rel']:9.3e} {row['dk_cos']:7.4f} {row['dk_resid']:9.3e} |"
+            f" {row['dv_rel']:9.3e} | {nonfinite} ({time.time() - t0:.1f}s)"
+            f"\n      sink rows: frac(err along K_peak)={row['sink_frac_along_ksink']:.3f}  cos(err,H1)={row['sink_cos_h1']:+.3f}"
+            f"  cos(err,H2)={row['sink_cos_h2']:+.3f}  cos(err,H1+H2)={row['sink_cos_h12']:+.3f}"
+            f"  resid H1={row['sink_resid_h1']:.3e}  resid H1+H2={row['sink_resid_h12']:.3e}",
+            flush=True,
+        )
+    pred_rel = [_norm(pred[i]) / _norm(ref[i]) for i in range(2)]
+    n_sink = int(sel_s.sum())
+    sink_pred = _norm(pred[0][sel_s]) / max(_norm(r_rows[sel_s]), 1e-300) if n_sink else float("nan")
+    sink_pred2 = _norm(pred2[sel_s]) / max(_norm(r_rows[sel_s]), 1e-300) if n_sink else float("nan")
+    print(
+        f"  predicted error norm (rel to fp32 grads): H1 dQ {pred_rel[0]:.3e} dK {pred_rel[1]:.3e};"
+        f" on {n_sink} sink rows: H1 {sink_pred:.3e}  H2 {sink_pred2:.3e}",
+        flush=True,
+    )
+    results["_pred_rel"] = pred_rel
+    results["_sink"] = dict(frac=frac, sink_idx=sink_idx.ravel().tolist(), share=float(sink_share.mean()))
+    return results
+
+
+def main():
+    logging.basicConfig(level=logging.WARNING)
+    cfg = _config.get_config(CONFIG)
+    cfg = dataclasses.replace(
+        cfg,
+        exp_name="attention_dterm_probe",
+        batch_size=BATCH,
+        num_workers=WORKERS,
+        fsdp_devices=1,
+        wandb_enabled=False,
+        model=dataclasses.replace(
+            cfg.model,
+            use_cudnn_attention=True,
+            cudnn_attention_layer_start=0,
+            cudnn_attention_num_layers=None,
+            cudnn_attention_dtype="bfloat16",
+            explicit_attention_fp32=False,
+        ),
+    )
+    assert jax.device_count() == 1, "run with CUDA_VISIBLE_DEVICES=<one gpu>"
+    t0 = time.time()
+    import os
+
+    if os.path.exists(CAPTURE_FILE):
+        data = np.load(CAPTURE_FILE)
+        import ml_dtypes
+
+        for key in data.files:
+            layer, name = key.split("/")
+            arr = data[key]
+            if arr.dtype.kind == "V" and arr.dtype.itemsize == 2:  # bf16 saved by np.savez as raw bytes
+                arr = arr.view(ml_dtypes.bfloat16)
+            _CAPTURE.setdefault(int(layer), {})[name] = arr
+        print(f"loaded capture from {CAPTURE_FILE}", flush=True)
+    else:
+        capture_real_activations(cfg)
+        np.savez(CAPTURE_FILE, **{f"{l}/{n}": a for l, e in _CAPTURE.items() for n, a in e.items()})
+        print(f"captured {len(_CAPTURE)} layers in {time.time() - t0:.0f}s -> {CAPTURE_FILE}", flush=True)
+    e = _CAPTURE[0]
+    print(f"shapes q{e['q'].shape} k{e['k'].shape} v{e['v'].shape} mask{e['mask'].shape} dO{e['d_out'].shape}", flush=True)
+
+    per_layer = {}
+    for layer in range(NUM_LAYERS):
+        per_layer[layer] = analyse_layer(layer, _CAPTURE[layer])
+
+    names = [n for n in per_layer[0] if not n.startswith("_")]
+    print("\n=== SUMMARY over 18 layers (geometric mean of rel errors, mean of cosines) ===", flush=True)
+    print(f"  {'variant':18s} {'dQ rel':>9s} {'cos(err,pred)':>13s} {'resid':>9s} {'sinkrow rel':>11s} | {'dK rel':>9s} {'cos':>7s} {'resid':>9s} | {'dV rel':>9s}")
+    for name in names:
+        gm = lambda key: float(np.exp(np.mean([np.log(max(per_layer[l][name][key], 1e-300)) for l in per_layer])))
+        mean = lambda key: float(np.mean([per_layer[l][name][key] for l in per_layer]))
+        print(
+            f"  {name:18s} {gm('dq_rel'):9.3e} {mean('dq_cos'):13.4f} {gm('dq_resid'):9.3e} {gm('dq_sink_rel'):11.3e} |"
+            f" {gm('dk_rel'):9.3e} {mean('dk_cos'):7.4f} {gm('dk_resid'):9.3e} | {gm('dv_rel'):9.3e}",
+            flush=True,
+        )
+    for name in names:
+        mean = lambda key: float(np.nanmean([per_layer[l][name][key] for l in per_layer]))
+        print(
+            f"  {name:18s} sink rows: frac along K_peak={mean('sink_frac_along_ksink'):.3f}"
+            f" cos H1={mean('sink_cos_h1'):+.3f} cos H2={mean('sink_cos_h2'):+.3f} cos H1+H2={mean('sink_cos_h12'):+.3f}"
+        )
+    print("  predicted D-term rel norm per layer (dQ):", " ".join(f"{per_layer[l]['_pred_rel'][0]:.1e}" for l in per_layer))
+    print("  sink share per layer:", " ".join(f"{per_layer[l]['_sink']['share']:.2f}" for l in per_layer))
+    print("  sink keys per layer:", [sorted(set(per_layer[l]["_sink"]["sink_idx"])) for l in per_layer])
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/attention_precision_update_probe.py b/scripts/attention_precision_update_probe.py
index 71de30f..b1a04ec 100644
--- a/scripts/attention_precision_update_probe.py
+++ b/scripts/attention_precision_update_probe.py
@@ -35,9 +35,12 @@
 sys.path.insert(0, "scripts")
 import train as train_mod
 
-CONFIG = "diag_cudnn_strict_order"
-EXP_NAME = "diag_cudnn_strict_order"
-CHECKPOINT_STEP = 1000
+# Usage: python scripts/attention_precision_update_probe.py [config] [exp_name] [checkpoint_step]
+# The checkpoint only supplies the parameters and Adam moments; every variant is
+# evaluated from that same state, so any clean step-1000 run works as the anchor.
+CONFIG = sys.argv[1] if len(sys.argv) > 1 else "diag_cudnn_strict_order"
+EXP_NAME = sys.argv[2] if len(sys.argv) > 2 else CONFIG
+CHECKPOINT_STEP = int(sys.argv[3]) if len(sys.argv) > 3 else 1000
 BATCH = 32
 WORKERS = 16
 NUM_BATCHES = 2
diff --git a/scripts/cudnn_attention_microbench.py b/scripts/cudnn_attention_microbench.py
new file mode 100644
index 0000000..ddb344f
--- /dev/null
+++ b/scripts/cudnn_attention_microbench.py
@@ -0,0 +1,115 @@
+"""Time the attention variants used by gemma.Attention at the training shape on one GPU.
+
+Answers "where does the float16 cuDNN path lose time against the bfloat16 cuDNN
+path?" by timing forward and forward+backward of each variant on real captured
+activations (dterm_capture.npz, B=8, T=S=1018, 8 q heads / 1 kv head, head_dim 256),
+optionally tiled along batch to the per-device training batch.
+
+    CUDA_VISIBLE_DEVICES=7 PYTHONPATH=$PWD/src python scripts/cudnn_attention_microbench.py \
+        /root/localstorage/logs/dterm_capture.npz [batch_tile]
+"""
+
+import sys
+import time
+
+import einops
+import jax
+import jax.numpy as jnp
+import ml_dtypes
+import numpy as np
+
+import openpi.models.gemma as _gemma
+
+CAPTURE_FILE = sys.argv[1] if len(sys.argv) > 1 else "/root/localstorage/logs/dterm_capture.npz"
+TILE = int(sys.argv[2]) if len(sys.argv) > 2 else 4  # 8 * 4 = 32 = per-device batch at global batch 256 on 8 GPUs
+LAYER = 0
+ITERS = 20
+
+
+def explicit_bf16(q, k, v, mask):
+    num_kv_heads = k.shape[2]
+    grouped_q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=num_kv_heads)
+    logits = jnp.einsum("BTKGH,BSKH->BKGTS", grouped_q, k, preferred_element_type=jnp.float32)
+    masked_logits = jnp.where(mask[:, :, None, :, :], logits, -2.3819763e38)
+    probs = jax.nn.softmax(masked_logits, axis=-1).astype(v.dtype)
+    encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v).astype(q.dtype)
+    return einops.rearrange(encoded, "B T K G H -> B T (K G) H")
+
+
+def cudnn_bf16(q, k, v, mask):
+    return _gemma._cudnn_attention_call(q, k, v, mask)
+
+
+def cudnn_fp16_plain(q, k, v, mask):
+    # float16 kernel through jax's own custom_vjp, no loss scaling: isolates kernel cost.
+    q16, k16, v16 = (x.astype(jnp.float16) for x in (q, k, v))
+    return _gemma._cudnn_attention_call(q16, k16, v16, mask).astype(q.dtype)
+
+
+def cudnn_fp16_v2(q, k, v, mask):
+    return _gemma._cudnn_attention_in_dtype(q, k, v, mask, jnp.float16)
+
+
+VARIANTS = {
+    "cudnn_bf16": cudnn_bf16,
+    "cudnn_fp16_plain": cudnn_fp16_plain,
+    "cudnn_fp16_v2": cudnn_fp16_v2,
+    "explicit_bf16": explicit_bf16,
+}
+
+
+def _timeit(fn, *args):
+    out = fn(*args)
+    jax.block_until_ready(out)
+    t0 = time.perf_counter()
+    for _ in range(ITERS):
+        out = fn(*args)
+    jax.block_until_ready(out)
+    return (time.perf_counter() - t0) / ITERS * 1e3
+
+
+def main():
+    data = np.load(CAPTURE_FILE)
+    e = {}
+    for name in ("q", "k", "v", "mask", "d_out"):
+        arr = data[f"{LAYER}/{name}"]
+        if arr.dtype.kind == "V" and arr.dtype.itemsize == 2:
+            arr = arr.view(ml_dtypes.bfloat16)
+        if TILE > 1:
+            arr = np.concatenate([arr] * TILE, axis=0)
+        e[name] = jnp.asarray(arr)
+    q, k, v, mask, d_out = e["q"], e["k"], e["v"], e["mask"], e["d_out"]
+    print(f"shape q={q.shape} k={k.shape} mask={mask.shape} dtype={q.dtype} device={jax.devices()[0]}", flush=True)
+    d32 = d_out.astype(jnp.float32)
+
+    for name, fn in VARIANTS.items():
+        fwd = jax.jit(fn)
+
+        def objective(q_, k_, v_, fn=fn):
+            return jnp.sum(fn(q_, k_, v_, mask).astype(jnp.float32) * d32)
+
+        fwd_bwd = jax.jit(jax.grad(objective, argnums=(0, 1, 2)))
+        # Same structure as training: flax remat(nothing_saveable) recomputes the forward in the backward.
+        fwd_bwd_remat = jax.jit(
+            jax.grad(
+                lambda q_, k_, v_, fn=fn: jnp.sum(
+                    jax.checkpoint(lambda a, b, c: fn(a, b, c, mask), policy=jax.checkpoint_policies.nothing_saveable)(
+                        q_, k_, v_
+                    ).astype(jnp.float32)
+                    * d32
+                ),
+                argnums=(0, 1, 2),
+            )
+        )
+        t_fwd = _timeit(fwd, q, k, v, mask)
+        t_fb = _timeit(fwd_bwd, q, k, v)
+        t_fbr = _timeit(fwd_bwd_remat, q, k, v)
+        print(
+            f"{name:18s} fwd {t_fwd:7.2f} ms | fwd+bwd {t_fb:7.2f} ms | fwd+bwd(remat) {t_fbr:7.2f} ms"
+            f" | x18 layers: fwd+bwd(remat) {t_fbr * 18 / 1e3:.3f} s/step",
+            flush=True,
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/cudnn_fp16_vjp_equivalence.py b/scripts/cudnn_fp16_vjp_equivalence.py
new file mode 100644
index 0000000..e75f20f
--- /dev/null
+++ b/scripts/cudnn_fp16_vjp_equivalence.py
@@ -0,0 +1,112 @@
+"""Check that the direct-rule float16 cuDNN VJP in gemma.py matches the recompute-based one.
+
+`gemma._cudnn_attention_in_dtype` now calls JAX's cuDNN fwd/bwd rules directly
+(keeping softmax stats + output as residuals). This script re-implements the
+previous version (jax.vjp over jax.nn.dot_product_attention inside the backward)
+and compares forward outputs and dQ/dK/dV on real captured activations. The two
+backwards run the same kernel on the same inputs, so their difference must be at
+the level of the kernel's own run-to-run (atomic accumulation) noise.
+
+    CUDA_VISIBLE_DEVICES=7 PYTHONPATH=$PWD/src python scripts/cudnn_fp16_vjp_equivalence.py \
+        /root/localstorage/logs/dterm_capture.npz
+"""
+
+import sys
+import time
+
+import jax
+import jax.numpy as jnp
+import ml_dtypes
+import numpy as np
+
+import openpi.models.gemma as _gemma
+
+CAPTURE_FILE = sys.argv[1] if len(sys.argv) > 1 else "/root/localstorage/logs/dterm_capture.npz"
+LAYERS = (0, 4, 8, 14, 17)
+
+
+def _reference_in_dtype(q, k, v, attn_mask, compute_dtype):
+    """The previous implementation: recompute the forward inside the backward via jax.vjp."""
+    out_dtype = q.dtype
+
+    @jax.custom_vjp
+    def attention(q, k, v, attn_mask):
+        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
+        return _gemma._cudnn_attention_call(qc, kc, vc, attn_mask).astype(out_dtype)
+
+    def attention_fwd(q, k, v, attn_mask):
+        return attention(q, k, v, attn_mask), (q, k, v, attn_mask)
+
+    def attention_bwd(residuals, d_out):
+        q, k, v, attn_mask = residuals
+        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
+        d_out32 = d_out.astype(jnp.float32)
+        max_abs = jnp.max(jnp.abs(d_out32))
+        tiny = jnp.finfo(jnp.float32).tiny
+        exponent = jnp.floor(jnp.log2(jnp.maximum(max_abs, tiny)))
+        scale = jnp.where(max_abs > 0, jnp.exp2(-(exponent + 2.0)), 1.0)
+        _, vjp_fn = jax.vjp(lambda a, b, c: _gemma._cudnn_attention_call(a, b, c, attn_mask), qc, kc, vc)
+        dq, dk, dv = vjp_fn((d_out32 * scale).astype(compute_dtype))
+        inv_scale = 1.0 / scale
+        return (*((g.astype(jnp.float32) * inv_scale).astype(out_dtype) for g in (dq, dk, dv)), None)
+
+    attention.defvjp(attention_fwd, attention_bwd)
+    return attention(q, k, v, attn_mask)
+
+
+def _fwd_and_grads(fn, q, k, v, mask, d_out):
+    d32 = d_out.astype(jnp.float32)
+
+    def objective(q_, k_, v_):
+        out = fn(q_, k_, v_, mask, jnp.float16)
+        return jnp.sum(out.astype(jnp.float32) * d32), out
+
+    (_, out), grads = jax.value_and_grad(objective, argnums=(0, 1, 2), has_aux=True)(q, k, v)
+    return out, grads
+
+
+new_fn = jax.jit(lambda *a: _fwd_and_grads(_gemma._cudnn_attention_in_dtype, *a))
+old_fn = jax.jit(lambda *a: _fwd_and_grads(_reference_in_dtype, *a))
+
+
+def _f64(x):
+    return np.asarray(jax.device_get(x)).astype(np.float64)
+
+
+def _rel(a, b):
+    return float(np.sqrt(np.sum((a - b) ** 2) / max(np.sum(b**2), 1e-300)))
+
+
+def main():
+    data = np.load(CAPTURE_FILE)
+    for layer in LAYERS:
+        e = {}
+        for name in ("q", "k", "v", "mask", "d_out"):
+            arr = data[f"{layer}/{name}"]
+            if arr.dtype.kind == "V" and arr.dtype.itemsize == 2:
+                arr = arr.view(ml_dtypes.bfloat16)
+            e[name] = jnp.asarray(arr)
+        args = (e["q"], e["k"], e["v"], e["mask"], e["d_out"])
+        t0 = time.time()
+        out_new, g_new = new_fn(*args)
+        _, g_new2 = new_fn(*args)
+        out_old, g_old = old_fn(*args)
+        _, g_old2 = old_fn(*args)
+        jax.block_until_ready((g_new, g_new2, g_old, g_old2))
+        valid_q = np.asarray(jnp.any(e["mask"], axis=-1)[:, 0, :])
+        dq_new = _f64(g_new[0])
+        masked_dq = float(np.abs(dq_new[~valid_q]).max()) if (~valid_q).any() else 0.0
+        nonfinite = sum(int(np.count_nonzero(~np.isfinite(_f64(x)))) for x in (*g_new, out_new))
+        print(
+            f"layer {layer:2d}: fwd new-vs-old rel={_rel(_f64(out_new), _f64(out_old)):.2e}"
+            f" (bit-identical: {bool(np.array_equal(np.asarray(out_new), np.asarray(out_old)))})"
+            f" | new-vs-old dQ {_rel(dq_new, _f64(g_old[0])):.2e} dK {_rel(_f64(g_new[1]), _f64(g_old[1])):.2e}"
+            f" dV {_rel(_f64(g_new[2]), _f64(g_old[2])):.2e}"
+            f" | new run-to-run dQ {_rel(dq_new, _f64(g_new2[0])):.2e} old run-to-run dQ {_rel(_f64(g_old[0]), _f64(g_old2[0])):.2e}"
+            f" | masked-row |dQ| max={masked_dq:.1e} nonfinite={nonfinite} ({time.time() - t0:.1f}s)",
+            flush=True,
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index 3f99c3b..6357201 100644
--- a/src/openpi/models/gemma.py
+++ b/src/openpi/models/gemma.py
@@ -32,6 +32,12 @@
 import einops
 import flax.linen as nn
 import jax
+
+# Private JAX API (pinned jax 0.5.3): the cuDNN fused-attention fwd/bwd rules that
+# jax.nn.dot_product_attention(implementation="cudnn") wires into its custom_vjp.
+# Used by _cudnn_attention_in_dtype to run the float16 kernel with loss scaling
+# without recomputing the forward inside the backward.
+from jax._src.cudnn import fused_attention_stablehlo as _cudnn_fa
 import jax.numpy as jnp
 
 import openpi.models.lora as lora
@@ -205,19 +211,28 @@ def _cudnn_attention_in_dtype(q, k, v, attn_mask, compute_dtype):
     """
     out_dtype = q.dtype
 
+    # The forward and backward call JAX's cuDNN fwd/bwd rules directly (the same
+    # functions jax.nn.dot_product_attention's own custom_vjp uses) instead of
+    # re-running the forward through jax.vjp inside the backward: the softmax
+    # stats and the output are kept as residuals, so with flax remat the backward
+    # costs one recomputed forward plus one fused backward, like the bf16 path.
     # attn_mask is an explicit argument (not a closure) so that the backward,
     # which flax remat traces separately, does not capture a leaked tracer.
     @jax.custom_vjp
     def attention(q, k, v, attn_mask):
-        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
-        return _cudnn_attention_call(qc, kc, vc, attn_mask).astype(out_dtype)
+        return attention_fwd(q, k, v, attn_mask)[0]
 
     def attention_fwd(q, k, v, attn_mask):
-        return attention(q, k, v, attn_mask), (q, k, v, attn_mask)
+        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
+        bias = jnp.where(attn_mask, jnp.asarray(0, compute_dtype), _cudnn_fa.get_large_negative_number(compute_dtype))
+        zeros = jnp.zeros(0, dtype=compute_dtype)
+        out, res = _cudnn_fa._dot_product_attention_fwd_rule(
+            qc, kc, vc, bias, zeros, zeros, zeros, zeros, *_cudnn_static_args(bias.shape, qc.shape)
+        )
+        return out.astype(out_dtype), (res, attn_mask)
 
     def attention_bwd(residuals, d_out):
-        q, k, v, attn_mask = residuals
-        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
+        res, attn_mask = residuals
         d_out32 = d_out.astype(jnp.float32)
         max_abs = jnp.max(jnp.abs(d_out32))
         tiny = jnp.finfo(jnp.float32).tiny
@@ -226,8 +241,14 @@ def attention_bwd(residuals, d_out):
         # and the dQ/dK/dV outputs inside the float16 kernel.
         exponent = jnp.floor(jnp.log2(jnp.maximum(max_abs, tiny)))
         scale = jnp.where(max_abs > 0, jnp.exp2(-(exponent + 2.0)), 1.0)
-        _, vjp_fn = jax.vjp(lambda a, b, c: _cudnn_attention_call(a, b, c, attn_mask), qc, kc, vc)
-        dq, dk, dv = vjp_fn((d_out32 * scale).astype(compute_dtype))
+        grads = _cudnn_fa._dot_product_attention_bwd_rule(
+            *_cudnn_static_args(res[3].shape, res[0].shape), res, (d_out32 * scale).astype(compute_dtype)
+        )
+        dq, dk, dv = grads[:3]
+        # Same semantics as _stop_gradient_for_fully_masked_queries: fully masked
+        # query rows get a zero (never NaN) q-gradient.
+        query_has_key = jnp.any(attn_mask, axis=-1)[:, 0, :, None, None]
+        dq = jnp.where(query_has_key, dq, jnp.zeros_like(dq))
         inv_scale = 1.0 / scale
 
         def unscale(g):
@@ -239,6 +260,23 @@ def unscale(g):
     return attention(q, k, v, attn_mask)
 
 
+def _cudnn_static_args(bias_shape, query_shape):
+    """Static parameters jax.nn.dot_product_attention(..., scale=1.0, implementation="cudnn") uses."""
+    layout = _cudnn_fa._normalize_layout("BTNH")
+    has_dbias = _cudnn_fa.should_export_dbias(bias_shape, query_shape, layout.value)
+    # (scale, seed, dropout_rate, variadic_args, mask_type, layout, sliding_window_length, cudnn_version)
+    return (
+        1.0,
+        42,
+        0.0,
+        (True, has_dbias),
+        _cudnn_fa.MaskType.NO_MASK,
+        layout.value,
+        None,
+        _cudnn_fa.check_cudnn_version(),
+    )
+
+
 @at.typecheck
 class Attention(nn.Module):
     """Attention module."""
@@ -441,9 +479,7 @@ def __call__(self, xs, kv_cache, positions, attn_mask, adarms_cond, layer_index=
             pre_attn.append(x)
 
         pre_attn = sharding.activation_sharding_constraint(pre_attn)
-        post_attn, kv_cache = attn(
-            pre_attn, positions, attn_mask, kv_cache, use_cudnn_attention=layer_uses_cudnn
-        )
+        post_attn, kv_cache = attn(pre_attn, positions, attn_mask, kv_cache, use_cudnn_attention=layer_uses_cudnn)
         post_attn = jax.tree.map(lambda x: drop(x, deterministic), post_attn)
         post_attn = sharding.activation_sharding_constraint(post_attn)
         xs = [_gated_residual(x, y, gate) for x, y, gate in zip(xs, post_attn, gates, strict=True)]
@@ -554,9 +590,7 @@ def __call__(
             adarms_cond = [None] * len(self.configs)
 
         layer_indices = jnp.arange(self.configs[0].depth, dtype=jnp.int32)
-        embedded, kv_cache = self.layers(
-            embedded, kv_cache, positions, mask, adarms_cond, layer_indices, deterministic
-        )
+        embedded, kv_cache = self.layers(embedded, kv_cache, positions, mask, adarms_cond, layer_indices, deterministic)
 
         assert all(e.dtype == jnp.dtype(self.embed_dtype) for e in embedded if e is not None)
 

From cf07bbeaf32e92da69df7431f3400c942953edbc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Fri, 4 Sep 2026 10:43:18 +0800
Subject: [PATCH 16/28] Speed up H100 training: token length 128, XLA
 collective flags, remat knobs

A/B on 8xH100 (docs/2026-09-04-h100-speed-ab.md): trimming max_token_len
200->128 (-6.2%) and enabling XLA collective combining / pipelining
(-3.5%) together cut the earbuds 0904 step time 1.328 -> 1.193 s
(-10.2%) with identical loss.

- Add configs/_examples/..._0904_h100.yaml with max_token_len 128 and the
  XLA_FLAGS launch line in its header.
- Add scripts/token_len_scan.py: runs the real transform chain over real
  dataset states to measure tokenized prompt length before setting
  max_token_len.
- Add Pi0Config.gemma_remat_policy / siglip_remat_policy (defaults
  unchanged). Neither "none" nor dots_with_no_batch_dims_saveable fits at
  batch 256 on 80 GB, but they are usable on larger GPUs or smaller
  per-device batches.
- Pin q/k/v/mask shardings before the cuDNN attention call; the custom
  partitioner otherwise fails under non-default remat policies.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 ...s_case_insertion_teleop_rtc_0904_h100.yaml |  47 +++++++
 docs/2026-09-04-h100-speed-ab.md              |  79 ++++++++++++
 scripts/token_len_scan.py                     | 115 ++++++++++++++++++
 src/openpi/models/gemma.py                    |  22 +++-
 src/openpi/models/pi0.py                      |   2 +
 src/openpi/models/pi0_config.py               |   8 ++
 src/openpi/models/siglip.py                   |  15 ++-
 7 files changed, 276 insertions(+), 12 deletions(-)
 create mode 100644 configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
 create mode 100644 docs/2026-09-04-h100-speed-ab.md
 create mode 100644 scripts/token_len_scan.py

diff --git a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
new file mode 100644
index 0000000..75aefd2
--- /dev/null
+++ b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
@@ -0,0 +1,47 @@
+# Launch (8xH100, hb-openpi env). XLA_FLAGS below measured -3.5% step time on 2026-09-04
+# (FSDP collective combining + pipelined all-gather/reduce-scatter); see
+# docs/2026-09-04-h100-speed-ab.md for the A/B table.
+#   XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true \
+#     --xla_gpu_all_gather_combine_threshold_bytes=1073741824 \
+#     --xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824 \
+#     --xla_gpu_all_reduce_combine_threshold_bytes=1073741824 \
+#     --xla_gpu_enable_pipelined_all_gather=true \
+#     --xla_gpu_enable_pipelined_reduce_scatter=true \
+#     --xla_gpu_enable_while_loop_double_buffering=true" \
+#   XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 python scripts/train.py <this-config> --exp-name=...
+model:
+  type: Pi0Config
+  # Scanned 2026-09-04 over this dataset (4037 states incl. per-dim extremes) with
+  # scripts/token_len_scan.py: the longest tokenized "Task: ..., State: ...;\nAction: " prompt
+  # is 116 tokens, so 128 leaves 12 spare. Down from 200: -6.2% step time, identical loss
+  # (padding tokens were masked anyway). Re-scan when the dataset or prompt changes.
+  max_token_len: 128
+  pi05: true
+  discrete_state_input: true
+  enable_training_time_rtc: true
+  batch_image_views: true
+  use_cudnn_attention: true
+  cudnn_attention_dtype: float16
+weight_loader:
+  # Was /home/li/hubo/xense-openpi/checkpoints/20000/params — a continuation of a
+  # 20k-step run that only existed on one laptop, so nobody else could start this
+  # config. Points at the upstream base now; to continue from your own checkpoint:
+  #   uv run scripts/train.py <this-config> --exp-name=... \
+  #     --weight_loader.params_path=/path/to/checkpoint/params
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+ema_decay: null
+data:
+  type: LeRobotBiFlexivDataConfig
+  repo_id: Xense/earbud_case_insertion_teleop_0515_left_right
+  base_config:
+    prompt_from_task: true
+    tactile: false
+  default_prompt: Pick up each earbud case from the left stands, insert the matching earbuds, close the lid, and place the case in the box.
+batch_size: 256
+num_workers: 64
+num_train_steps: 60000
+save_interval: 5000
+keep_period: 10000
+fsdp_devices: 8
+strict_batch_order: false
diff --git a/docs/2026-09-04-h100-speed-ab.md b/docs/2026-09-04-h100-speed-ab.md
new file mode 100644
index 0000000..80c3ddf
--- /dev/null
+++ b/docs/2026-09-04-h100-speed-ab.md
@@ -0,0 +1,79 @@
+# 2026-09-04 H100 训练提速 A/B（earbuds 0904 配置）
+
+**机器**：ebcloud 8×H100 80GB，hb-openpi 环境（jax 0.5.3，cuDNN 9.19）
+**配置**：`configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml`
+（pi05、cuDNN fp16 attention、batch_image_views、全局 batch 256、64 workers、fsdp 8）
+**数据**：`Xense/earbud_case_insertion_teleop_0515_left_right`，776 集 / 1,145,613 帧
+
+> 注：实验当天远端机器的 localstorage 随后被清空，本文与相关代码是事后在本机重建并提交的；
+> 数字来自当时的日志摘录。
+
+## 口径
+
+每个变体从 pi05_base 权重起跑 400 步，`strict_batch_order: true`，关 checkpoint 和 W&B，
+其余与生产配置相同（短跑用一份复制自生产 yaml 的临时 config，CLI 覆盖单个变量）。
+主指标是日志里 `Step 100` 到 `Step 300` 两条时间戳之差除以 200，即含日志开销的端到端墙钟 s/step。
+step 300 的 loss 用来确认数值没变。
+
+## 结果
+
+| 变体 | s/step | 相对基线 | loss@300 | 结论 |
+| --- | --- | --- | --- | --- |
+| E0 基线（原配置） | 1.328 | — | 0.7356 | 等数据 0.03 s，数据侧不是瓶颈 |
+| E1 `max_token_len` 200→128 | 1.246 | **−6.2%** | 0.7356 | 采纳 |
+| E2 关闭 remat（gemma+siglip） | OOM | — | — | XLA 试图分配 280 GiB，不可行 |
+| E2b gemma remat `dots_with_no_batch_dims_saveable` | 编译失败 | — | — | cuDNN 自定义分片器报 q/k/v 分片不一致 |
+| E2c 同上 + q/k/v 显式分片约束 | OOM | — | — | 编译过了，但单块残差 84.6 GB，仍放不下 |
+| E3 XLA 集合通信标志 | 1.281 | **−3.5%** | 0.7357 | 采纳 |
+| E4 `fsdp_devices` 8→4 | 1.335 | +0.5% | 0.7354 | 无收益，放弃 |
+| **E5 = E1 + E3** | **1.193** | **−10.2%** | 0.7356 | 生产采用；吞吐 192.8→214.6 samples/s |
+
+60000 步按稳态估算：22.1 h → 19.9 h，省约 2.2 h。
+
+## 各项说明
+
+### E1：token 长度按数据集实测收紧
+
+`scripts/token_len_scan.py` 走真实 transform 链（repack → BiFlexivInputs → Normalize →
+TokenizePrompt），对 4037 个状态（含每一维的最小/最大值样本）做 tokenize：min 110、p50 113、p99 115、
+**max 116**。原来 200 里有 84 个是纯 padding，但 attention/MLP 照算。128 留 12 个余量。
+换数据集或改 prompt 时必须重扫；tokenizer 超长会 warning 并截断，不会报错。
+
+```bash
+JAX_PLATFORMS=cpu uv run scripts/token_len_scan.py \
+  --config-name pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100
+```
+
+### E3：XLA 标志
+
+```
+--xla_gpu_enable_latency_hiding_scheduler=true
+--xla_gpu_all_gather_combine_threshold_bytes=1073741824
+--xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824
+--xla_gpu_all_reduce_combine_threshold_bytes=1073741824
+--xla_gpu_enable_pipelined_all_gather=true
+--xla_gpu_enable_pipelined_reduce_scatter=true
+--xla_gpu_enable_while_loop_double_buffering=true
+```
+
+FSDP 每步的参数 all-gather / 梯度 reduce-scatter 合并成大块并与计算流水。没有逐项拆分测哪一个起作用。
+通过 `XLA_FLAGS` 环境变量传入，已写在 0904 yaml 头部注释里。
+
+### E2：remat 在 80 GB 上没有空间
+
+`gemma.py` / `siglip.py` 每层用 `nothing_saveable` 全量重算。本次给 `Pi0Config` 加了
+`gemma_remat_policy` / `siglip_remat_policy`（默认不变），并在 cuDNN 调用前给 q/k/v/mask 加了统一的
+batch 分片约束（`gemma.py`，否则非默认 remat 策略下编译失败）。但 batch 256（每卡 32）时：
+完全不 remat 要 280 GiB 单块，只保存无 batch 维 matmul 也要 84.6 GB，都放不下。
+在 H200 141 GB 或每卡 batch 更小时这两个开关才可能兑现，理论上限约为一次前向的重算（15–25%）。
+
+### 其他观察
+
+- 数据侧：等数据 0.03–0.04 s/step，64 workers 足够；worker 启动一次约 330 s（spawn 上下文，64 进程各自
+  import jax/torch/lerobot 并反序列化 1.1M 行数据集），是每次启动的一次性开销。
+- `ResizeImages` 实际调用的是 `xense_client.image_tools`（PIL），worker 里不跑 JAX resize。
+- fsdp 4 相对 8 没有收益：参数 all-gather 的量只差 12.5%，而每卡持有的优化器状态翻倍。
+- 短跑结束后偶有 DataLoader worker 成为孤儿进程（ppid 1，每个约 3 GB RSS），一次残留了 39 个；
+  跑完批量实验后用 `ps -eo pid,ppid,cmd | grep spawn_main` 检查一下。
+- 该配置 `num_train_steps: 60000` 但 `lr_schedule.decay_steps` 仍是默认 30000，后半程一直在底线 LR；
+  这是训练质量问题，本次未动。
diff --git a/scripts/token_len_scan.py b/scripts/token_len_scan.py
new file mode 100644
index 0000000..59fd8da
--- /dev/null
+++ b/scripts/token_len_scan.py
@@ -0,0 +1,115 @@
+"""Measure the real tokenized prompt length of a training config's dataset.
+
+For Pi0.5 the language prefix is ``Task: <prompt>, State: <discretized state>;\\nAction: `` and
+everything up to ``max_token_len`` is padded (and masked) -- but attention/MLP still run over the
+padding. On 2026-09-04 trimming ``max_token_len`` 200 -> 128 gave -6.2% step time with identical
+loss (see docs/2026-09-04-h100-speed-ab.md). This script runs the *actual* transform chain
+(repack -> data transforms -> model transforms) on real states with dummy images and reports the
+token-length distribution so ``max_token_len`` can be set with a known margin.
+
+Note: the tokenizer only warns and truncates when the prompt exceeds ``max_token_len``.
+
+Example:
+    JAX_PLATFORMS=cpu uv run scripts/token_len_scan.py --config-name pi05_droid --samples 4000
+"""
+
+import dataclasses
+import glob
+import json
+import logging
+import os
+import pathlib
+
+import numpy as np
+import pyarrow.parquet as pq
+import tyro
+
+import openpi.training.config as _config
+import openpi.transforms as _transforms
+
+
+@dataclasses.dataclass
+class Args:
+    config_name: str
+    # Random states to sample; per-dimension min/max states are always added so the worst case is covered.
+    samples: int = 4000
+    # Override the LeRobot dataset root (default: $HF_LEROBOT_HOME/<repo_id>).
+    dataset_root: str | None = None
+    # Image shape handed to the transform chain (channels-first uint8, like LeRobot video frames).
+    image_shape: tuple[int, int, int] = (3, 480, 640)
+    seed: int = 0
+
+
+def _lerobot_home() -> pathlib.Path:
+    if "HF_LEROBOT_HOME" in os.environ:
+        return pathlib.Path(os.environ["HF_LEROBOT_HOME"]).expanduser()
+    return pathlib.Path(os.environ.get("HF_HOME", "~/.cache/huggingface")).expanduser() / "lerobot"
+
+
+def _load_tasks(root: pathlib.Path) -> list[str]:
+    jsonl = root / "meta" / "tasks.jsonl"
+    if jsonl.exists():
+        return sorted({json.loads(line)["task"] for line in jsonl.read_text().splitlines() if line.strip()})
+    parquet = root / "meta" / "tasks.parquet"
+    if parquet.exists():
+        table = pq.read_table(parquet).to_pandas().reset_index()
+        col = "task" if "task" in table.columns else table.columns[0]
+        return sorted(set(table[col].astype(str)))
+    raise FileNotFoundError(f"No tasks metadata under {root / 'meta'}")
+
+
+def main(args: Args) -> None:
+    cfg = _config.get_config(args.config_name)
+    data_config = cfg.data.create(cfg.assets_dirs, cfg.model)
+    if data_config.repo_id is None:
+        raise ValueError("config has no repo_id")
+    chain = _transforms.compose(
+        [
+            *data_config.repack_transforms.inputs,
+            *data_config.data_transforms.inputs,
+            *data_config.model_transforms.inputs,
+        ]
+    )
+    root = pathlib.Path(args.dataset_root) if args.dataset_root else _lerobot_home() / data_config.repo_id
+    tasks = _load_tasks(root)
+    logging.info("tasks: %s", tasks)
+
+    # Discover the raw column/key names from the repack transform so this works for any robot config.
+    repack = next(t for t in data_config.repack_transforms.inputs if isinstance(t, _transforms.RepackTransform))
+    structure = repack.structure
+    image_keys = list(structure["images"].values())
+    state_key, action_key, prompt_key = structure["state"], structure["actions"], structure["prompt"]
+
+    files = sorted(glob.glob(str(root / "data" / "*" / "*.parquet")))
+    states = np.concatenate(
+        [np.stack(pq.read_table(f, columns=[state_key]).column(0).to_pylist()) for f in files]
+    ).astype(np.float32)
+    logging.info("states: %s", states.shape)
+
+    rng = np.random.default_rng(args.seed)
+    idx = rng.choice(len(states), size=min(args.samples, len(states)), replace=False)
+    idx = np.unique(np.concatenate([idx, states.argmin(0), states.argmax(0)]))
+
+    dummy_img = np.zeros(args.image_shape, np.uint8)
+    action_dim = states.shape[1]
+    lengths = []
+    for i in idx:
+        for task in tasks:
+            sample = dict.fromkeys(image_keys, dummy_img)
+            sample[state_key] = states[i]
+            sample[action_key] = np.zeros((cfg.model.action_horizon, action_dim), np.float32)
+            sample[prompt_key] = task
+            out = chain(sample)
+            lengths.append(int(np.asarray(out["tokenized_prompt_mask"]).sum()))
+    lengths = np.asarray(lengths)
+    cap = cfg.model.max_token_len
+    print(
+        f"n={len(lengths)} min={lengths.min()} p50={np.percentile(lengths, 50):.0f} "
+        f"p99={np.percentile(lengths, 99):.0f} max={lengths.max()} "
+        f"(max_token_len={cap}; samples_at_cap={(lengths >= cap).sum()})"
+    )
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    main(tyro.cli(Args))
diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index 6357201..20a8bb1 100644
--- a/src/openpi/models/gemma.py
+++ b/src/openpi/models/gemma.py
@@ -191,6 +191,10 @@ def _stop_gradient_for_fully_masked_queries(q, attn_mask):
 def _cudnn_attention_call(q, k, v, attn_mask):
     """The plain cuDNN fused attention call used by training (q is pre-scaled)."""
     q = _stop_gradient_for_fully_masked_queries(q, attn_mask)
+    # The cuDNN custom partitioner requires q/k/v to carry identical shardings. Pin them to the
+    # batch axis explicitly; otherwise XLA sharding propagation can differ per operand (seen with
+    # non-default remat policies) and compilation fails with "should have same sharding".
+    q, k, v, attn_mask = sharding.activation_sharding_constraint((q, k, v, attn_mask))
     return jax.nn.dot_product_attention(q, k, v, mask=attn_mask, scale=1.0, implementation="cudnn")
 
 
@@ -225,6 +229,8 @@ def attention(q, k, v, attn_mask):
     def attention_fwd(q, k, v, attn_mask):
         qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
         bias = jnp.where(attn_mask, jnp.asarray(0, compute_dtype), _cudnn_fa.get_large_negative_number(compute_dtype))
+        # See _cudnn_attention_call: the custom partitioner needs identical q/k/v/bias shardings.
+        qc, kc, vc, bias = sharding.activation_sharding_constraint((qc, kc, vc, bias))
         zeros = jnp.zeros(0, dtype=compute_dtype)
         out, res = _cudnn_fa._dot_product_attention_fwd_rule(
             qc, kc, vc, bias, zeros, zeros, zeros, zeros, *_cudnn_static_args(bias.shape, qc.shape)
@@ -527,6 +533,7 @@ class Module(nn.Module):
     cudnn_attention_num_layers: int | None = None
     cudnn_attention_dtype: str = "bfloat16"
     explicit_attention_fp32: bool = False
+    remat_policy: str = "nothing_saveable"
 
     def setup(self):
         # all experts must have the same depth
@@ -537,12 +544,15 @@ def setup(self):
             embed_dim=self.configs[0].width,  # embedder for first expert only
             name="embedder",
         )
-        block_cls = nn.remat(
-            Block,
-            prevent_cse=False,
-            static_argnums=(6,),  # 0=self, 7=deterministic
-            policy=jax.checkpoint_policies.nothing_saveable,
-        )
+        if self.remat_policy == "none":
+            block_cls = Block
+        else:
+            block_cls = nn.remat(
+                Block,
+                prevent_cse=False,
+                static_argnums=(6,),  # 0=self, 7=deterministic
+                policy=getattr(jax.checkpoint_policies, self.remat_policy),
+            )
         self.layers = nn.scan(
             block_cls,
             variable_axes={"params": 0},
diff --git a/src/openpi/models/pi0.py b/src/openpi/models/pi0.py
index e60d9f0..284fc79 100644
--- a/src/openpi/models/pi0.py
+++ b/src/openpi/models/pi0.py
@@ -99,6 +99,7 @@ def __init__(
                 cudnn_attention_num_layers=config.cudnn_attention_num_layers,
                 cudnn_attention_dtype=config.cudnn_attention_dtype,
                 explicit_attention_fp32=config.explicit_attention_fp32,
+                remat_policy=config.gemma_remat_policy,
             )
         )
         llm.lazy_init(
@@ -112,6 +113,7 @@ def __init__(
                 variant="So400m/14",
                 pool_type="none",
                 scan=True,
+                remat_policy=config.siglip_remat_policy,
                 dtype_mm=config.dtype,
             )
         )
diff --git a/src/openpi/models/pi0_config.py b/src/openpi/models/pi0_config.py
index cb48db0..076d59a 100644
--- a/src/openpi/models/pi0_config.py
+++ b/src/openpi/models/pi0_config.py
@@ -45,6 +45,14 @@ class Pi0Config(_model.BaseModelConfig):
     cudnn_attention_dtype: str = "bfloat16"
     # Diagnostic only: explicit attention with fp32 q/k/v/probs as a numerical reference.
     explicit_attention_fp32: bool = False
+    # Activation rematerialization policy for the Gemma transformer blocks and the SigLIP
+    # encoder blocks. "nothing_saveable" (default) recomputes every block in the backward
+    # pass; "none" disables remat entirely (fastest, most memory); any other name is looked
+    # up in jax.checkpoint_policies (e.g. "dots_with_no_batch_dims_saveable").
+    # NOTE: on 8xH100 80GB with global batch 256 neither "none" (XLA wants a 280 GiB buffer)
+    # nor "dots_with_no_batch_dims_saveable" (84.6 GB) fits; see docs/2026-09-04-h100-speed-ab.md.
+    gemma_remat_policy: str = "nothing_saveable"
+    siglip_remat_policy: str = "nothing_saveable"
 
     # training-time RTC config
     enable_training_time_rtc: bool = False
diff --git a/src/openpi/models/siglip.py b/src/openpi/models/siglip.py
index fd45c29..5b02eca 100644
--- a/src/openpi/models/siglip.py
+++ b/src/openpi/models/siglip.py
@@ -124,12 +124,15 @@ def __call__(self, x, deterministic=True):
         out = {}
 
         if self.scan:
-            block = nn.remat(
-                Encoder1DBlock,
-                prevent_cse=False,
-                static_argnums=(2,),  # 0=self, 2=deterministic
-                policy=getattr(jax.checkpoint_policies, self.remat_policy, None),
-            )
+            if self.remat_policy == "none":
+                block = Encoder1DBlock
+            else:
+                block = nn.remat(
+                    Encoder1DBlock,
+                    prevent_cse=False,
+                    static_argnums=(2,),  # 0=self, 2=deterministic
+                    policy=getattr(jax.checkpoint_policies, self.remat_policy, None),
+                )
             x, scan_out = nn.scan(
                 block,
                 variable_axes={"params": 0},

From cf03854317e51142e74c64cc1557986c305f62a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Fri, 4 Sep 2026 12:30:29 +0800
Subject: [PATCH 17/28] docs: consolidate training optimization history

---
 README.md                                     |  13 +-
 configs/_examples/_FULL_REFERENCE.yaml        |   2 +-
 ...lexiv_earbud_case_insertion_0826_h200.yaml |   2 +-
 ...s_case_insertion_teleop_rtc_0904_h100.yaml |   2 +-
 ...base_bi_flexiv_pick_up_cube_0824_h100.yaml |   2 +-
 ...exiv_unscrew_bottle_cap_rtc_0824_h100.yaml |   2 +-
 docs/2026-08-28-h200-training-optimization.md | 338 --------
 docs/2026-08-31-cudnn-sdpa-nan-summary.md     | 487 -----------
 ...026-09-03-cudnn-dterm-rounding-analysis.md | 196 -----
 docs/2026-09-04-h100-speed-ab.md              |  79 --
 docs/data-pipeline-stalls.md                  | 326 -------
 docs/training-optimization.md                 | 797 ++++++++++++++++++
 pyproject.toml                                |   2 +-
 scripts/attention_dterm_probe.py              |   2 +-
 scripts/token_len_scan.py                     |   2 +-
 src/openpi/models/gemma.py                    |   8 +-
 src/openpi/models/pi0_config.py               |   2 +-
 17 files changed, 818 insertions(+), 1444 deletions(-)
 delete mode 100644 docs/2026-08-28-h200-training-optimization.md
 delete mode 100644 docs/2026-08-31-cudnn-sdpa-nan-summary.md
 delete mode 100644 docs/2026-09-03-cudnn-dterm-rounding-analysis.md
 delete mode 100644 docs/2026-09-04-h100-speed-ab.md
 delete mode 100644 docs/data-pipeline-stalls.md
 create mode 100644 docs/training-optimization.md

diff --git a/README.md b/README.md
index 10f3075..fe522ce 100644
--- a/README.md
+++ b/README.md
@@ -228,12 +228,15 @@ PyTorch 2.11 pins cuDNN 9.19 and JAX shares that same runtime. Before training,
 run `python scripts/check_cuda_stack.py`; it rejects mixed library sources and
 checks the real BF16 GQA forward/backward shape.
 
-`use_cudnn_attention` remains experimental because it has diverged from the
-explicit-attention baseline in a step-0 long run. Production configs must use:
+Full-layer cuDNN attention with the default `bfloat16` compute dtype is unsafe:
+it diverged from the explicit-attention baseline after roughly 1,000 steps.
+Either keep the explicit path (`use_cudnn_attention: false`) or use the
+3,000-step strict-order validated FP16 path:
 
 ```yaml
 model:
-  use_cudnn_attention: false
+  use_cudnn_attention: true
+  cudnn_attention_dtype: float16
 ```
 
 At startup, `scripts/train.py` logs `JAX cuDNN runtime version: <version>`;
@@ -242,8 +245,8 @@ does not detect mixed dispatcher/engine libraries, so the stack-check script is
 the required gate. Use `--overwrite` only for a new run that may replace an
 existing experiment directory; use `--resume` to preserve and continue an
 existing run. See
-[`docs/2026-08-30-cudnn-attention-divergence.md`](docs/2026-08-30-cudnn-attention-divergence.md)
-before attempting to re-enable fused attention.
+[`docs/training-optimization.md`](docs/training-optimization.md)
+for the numerical root cause, validation criteria, and fallback configuration.
 
 ⚠️ Before committing a YAML into `configs/_examples/`, scrub any machine-local
 absolute paths from `weight_loader.params_path` (e.g. `/home/<you>/...`). Use
diff --git a/configs/_examples/_FULL_REFERENCE.yaml b/configs/_examples/_FULL_REFERENCE.yaml
index b68ccfc..5763168 100644
--- a/configs/_examples/_FULL_REFERENCE.yaml
+++ b/configs/_examples/_FULL_REFERENCE.yaml
@@ -247,7 +247,7 @@ strict_batch_order: false     # false ⇒ hand batches over as workers finish th
                               #   and lets the loader iterate without epoch boundaries.
                               #   true ⇒ strict sampler order, reproducible batch sequence,
                               #   one slow worker blocks everything behind it.
-                              #   See docs/data-pipeline-stalls.md.
+                              #   See docs/training-optimization.md.
 num_train_steps: 30000        # total gradient steps
 
 log_interval: 100             # how often to push metrics to wandb
diff --git a/configs/_examples/pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml b/configs/_examples/pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml
index 0f6527d..b787476 100644
--- a/configs/_examples/pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml
@@ -26,6 +26,6 @@ num_train_steps: 60000
 fsdp_devices: 8
 # 是否严格按 sampler 顺序交付 batch。false（默认）= 谁先算完谁先交付，
 # 消除每 num_workers 步一次的队头阻塞停顿；同时启用 InfiniteSampler，
-# 去掉 epoch 边界重建迭代器的停顿。详见 docs/data-pipeline-stalls.md。
+# 去掉 epoch 边界重建迭代器的停顿。详见 docs/training-optimization.md。
 # 代价：batch 的交付顺序不再逐次可复现（每个 index 每轮仍恰好出现一次）。
 strict_batch_order: false
diff --git a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
index 75aefd2..bff7ef0 100644
--- a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
@@ -1,6 +1,6 @@
 # Launch (8xH100, hb-openpi env). XLA_FLAGS below measured -3.5% step time on 2026-09-04
 # (FSDP collective combining + pipelined all-gather/reduce-scatter); see
-# docs/2026-09-04-h100-speed-ab.md for the A/B table.
+# docs/training-optimization.md for the A/B table.
 #   XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true \
 #     --xla_gpu_all_gather_combine_threshold_bytes=1073741824 \
 #     --xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824 \
diff --git a/configs/_examples/pi05_base_bi_flexiv_pick_up_cube_0824_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_pick_up_cube_0824_h100.yaml
index 25f81a2..61a818e 100644
--- a/configs/_examples/pi05_base_bi_flexiv_pick_up_cube_0824_h100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_pick_up_cube_0824_h100.yaml
@@ -26,7 +26,7 @@ fsdp_devices: 8
 # Whether batches are delivered in strict sampler order. false (the default) hands each
 # batch over as soon as a worker finishes it, which removes the head-of-line stall every
 # num_workers steps, and enables the InfiniteSampler so the epoch boundary never tears the
-# prefetch pipeline down. See docs/data-pipeline-stalls.md.
+# prefetch pipeline down. See docs/training-optimization.md.
 # Cost: batch delivery order is no longer reproducible run to run (every index is still
 # visited exactly once per pass).
 strict_batch_order: false
diff --git a/configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml
index 934919b..b3951c5 100644
--- a/configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml
@@ -28,7 +28,7 @@ keep_period: 10000
 # Whether batches are delivered in strict sampler order. false (the default) hands each
 # batch over as soon as a worker finishes it, which removes the head-of-line stall every
 # num_workers steps, and enables the InfiniteSampler so the epoch boundary never tears the
-# prefetch pipeline down. See docs/data-pipeline-stalls.md.
+# prefetch pipeline down. See docs/training-optimization.md.
 # Cost: batch delivery order is no longer reproducible run to run (every index is still
 # visited exactly once per pass).
 strict_batch_order: false
diff --git a/docs/2026-08-28-h200-training-optimization.md b/docs/2026-08-28-h200-training-optimization.md
deleted file mode 100644
index 8dbb3d7..0000000
--- a/docs/2026-08-28-h200-training-optimization.md
+++ /dev/null
@@ -1,338 +0,0 @@
-# H200 训练性能优化记录
-
-> [!CAUTION]
-> 2026-08-30 更新：**`use_cudnn_attention` 目前不得用于生产训练。** 全空 mask 修复只消除了瞬时 NaN；带该修复的训练仍会在 step 1000–1200 之间偏离已验证基线并缓慢发散至 NaN。单变量 A/B、已排除的候选和新的验收要求见 [2026-08-30-cudnn-attention-divergence.md](2026-08-30-cudnn-attention-divergence.md)。
->
-> 本文以下关于启动命令、cuDNN 版本确认和 A/B 吞吐的内容写于该问题定位之前，其中 `LD_LIBRARY_PATH` 与「确认输出 91400」两条**已经过时**，见下方修订。
-
-> [!WARNING]
-> 2026-08-29 生产长跑发现：cuDNN Attention 遇到 padding 产生的全空 mask 行时，前向有限但 Q 梯度为 NaN。修复、影响范围和 checkpoint 恢复要求见 [2026-08-29-cudnn-attention-nan-incident.md](2026-08-29-cudnn-attention-nan-incident.md)。未包含该修复的实现不得用于有效训练。
-
-> **跨机器生产启动命令**（2026-08-30 修订）：**不要再导出 `LD_LIBRARY_PATH="$CONDA_PREFIX/lib"`。** conda cuDNN 已移除，环境统一到 pip cu128 栈；继续导出该目录会让其中残留的 conda-forge `libcublas`/`libcudart`/`libnccl` 遮蔽 pip 的同名库，重演同一类版本错配。
-
-```bash
-XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
-  python scripts/train.py my_task \
-    --exp-name=run_0520 \
-    --overwrite
-```
-
-<details><summary>原命令（已过时，仅供追溯）</summary>
-
-```bash
-LD_LIBRARY_PATH="$CONDA_PREFIX/lib" \
-  XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
-  python scripts/train.py my_task --exp-name=run_0520 --overwrite
-```
-
-</details>
-
-> **2026-08-29 长跑事故（已修复）**：原实现从 step 16100 起记录到 NaN；受污染 checkpoint 已删除，并从保留的 step 15000 恢复。最终修复保留原始 attention mask，只对全空 query 行停止 Q 梯度。逐 step 有限性保护已按生产要求移除，训练继续按原逻辑记录 loss；详见事故文档。
-
-`my_task` 的模型配置需要包含 `use_cudnn_attention: true`；如果 YAML 中没有设置，则在命令末尾增加 `--model.use-cudnn-attention`。`--overwrite` 会删除同名实验的原 Checkpoint 目录；恢复已有训练应改用 `--resume`。
-
-训练启动后会在初始化日志中直接打印：
-
-```text
-JAX cuDNN runtime version: 91400
-```
-
-该值来自 `cuda_versions.cudnn_get_version()`。
-
-> [!IMPORTANT]
-> 2026-08-30 修订：**这个数字不足以证明 cuDNN 栈是干净的。** 它由 cuDNN 的引擎子库返回，因此一个 9.10.2 的 dispatcher stub 驱动 9.14 引擎库时，它照样打印 `91400`——事故期间的实际情况正是如此。请改用：
->
-> ```bash
-> python scripts/check_cuda_stack.py
-> ```
->
-> 该脚本遍历 `/proc/self/maps`，在 `libcudnn*` 来自多个目录或文件名版本与运行时报告不一致时判 FAIL，并在真实训练形状上跑一次前向+反向与 fp32 参考解对比。详见 [2026-08-30-cudnn-attention-divergence.md](2026-08-30-cudnn-attention-divergence.md)。
-
-**日期**：2026-08-27 至 2026-08-28
-
-**机器**：8×H200
-
-**分支**：`optimize-training-duration`
-
-**模型配置**：`pi05_base_bi_flexiv_earbud_case_insertion_0826_h200`
-
-**起始 Checkpoint**：step 49000（前四项优化）；step 16000（cuDNN Attention A/B）
-
-本文记录一次从数据管线拆分计时到 GPU trace 的训练性能优化过程。目标不是只看某个局部计时，而是在相同 Checkpoint、相同训练配置下逐项修改，并用稳态端到端 step 时间判断收益。
-
-更早的 `in_order=False`、无限 sampler 和跳过未使用触觉视频解码等工作见 [data-pipeline-stalls.md](data-pipeline-stalls.md)。本文从 multiprocessing queue、JAX 数组构造和模型计算继续向下分析。
-
-## 1. 测量口径
-
-前四项优化都从原始 step 49000 Checkpoint 只读恢复；cuDNN Attention A/B 从当前生产 step 16000 Checkpoint 只读恢复。所有短跑都使用独立实验目录，关闭 W&B、周期保存和最终 Checkpoint 保存，避免覆盖训练结果。每个变体运行到起始 step + 40，丢弃起始 step 到起始 step + 10 的编译、日志和预热数据，统计后续 30 个稳态样本。例如 step 49011–49040 或 step 16011–16040。
-
-JAX dispatch 是异步的，所以不能把 `dispatch` 单独当成 step 时间。本文使用以下和作为主要指标：
-
-```text
-measured_step = dispatch + compute_sync + next_batch_total
-```
-
-数据管线 profiling 进一步拆成：
-
-- `main_queue_wait`：主进程等待 DataLoader 返回 batch；
-- `worker_getitem`：worker 内所有样本 `__getitem__` 时间之和，不是单个 step 的墙钟时间；
-- `worker_collate`：worker 内 collate 时间；
-- `jax_array_construct`：主进程构造分片 JAX array 的时间；
-- `h2d_wait`：显式等待 Host→Device 完成的时间。
-
-profiling 会增加计时和同步，因此只用于短跑，不应在生产训练中默认开启。
-
-## 2. 第一步：共享内存 Tensor 和零拷贝 NumPy 视图
-
-### 问题
-
-worker 原先返回 NumPy batch。大 batch 通过 multiprocessing queue 时需要 pickle/复制，主进程还要再次把它包装成 JAX array。此前已经消除了严格顺序交付和 epoch 边界停顿，但这里仍有一段稳定的数据传输开销。
-
-### 实现
-
-collate 改为返回 CPU Torch Tensor。PyTorch multiprocessing reducer 会通过共享内存传递 tensor storage；主进程用 `tensor.numpy()` 得到同一块 CPU storage 的零拷贝视图，再交给 `jax.make_array_from_process_local_data`。
-
-同时增加显式 `DataLoader.close()`：训练退出时主动调用 multiprocessing iterator 的 shutdown，等待 worker，超时后终止残留进程，避免把 native library 的析构异常留到解释器退出阶段。
-
-### 结果
-
-- 修改前：`3.275728 s/step`
-- 共享 Tensor 后：`2.828845 s/step`
-- step 时间减少 `13.64%`，等价吞吐提升 `15.80%`
-
-共享 Tensor 后的稳态数据侧均值为：
-
-- `main_queue_wait = 0.005030 s`
-- `worker_collate = 0.081114 s`
-- `jax_array_construct = 0.009046 s`
-- `h2d_wait = 0.006171 s`
-- `next_batch_total = 0.046491 s`
-
-这说明主进程队列等待、JAX array 构造和 H2D 已经不是主要瓶颈。`worker_getitem` 约 130 秒是 512 个样本耗时之和，由 64 个 worker 并行完成，不能直接加到 step 墙钟时间上。
-
-## 3. 第二步：20-step JAX/XProf trace
-
-训练脚本增加可控的 XProf 开关：指定绝对起始 step、trace step 数和输出目录。在 trace 起止位置显式 `block_until_ready`，保证采集边界对应完整训练迭代。
-
-本次采集 step 49011–49030，共 20 个完整 step。按每张 GPU、每个 step 汇总 trace 中记录的 kernel 时间：
-
-- GEMM：约 `390.449 ms`
-- Fusion：约 `183.930 ms`
-- NCCL：约 `45.990 ms`
-- 其他 kernel：约 `2.562 ms`
-- Memcpy：约 `1.509 ms`
-- Memset：约 `0.465 ms`
-
-在这些已记录的 kernel 时间中，GEMM 与 Fusion 约占 `92.6%`，NCCL 约占 `7.4%`。因此当前优先级应是减少或融合模型计算，而不是继续优化 Host→Device 或 NCCL。
-
-trace 中 kernel 可能重叠，不能把 kernel duration 简单相加后当作 GPU 利用率；同样，也不能仅凭 host dispatch 区间判断真实计算时间。
-
-## 4. 第三步：只在日志 step 计算范数
-
-### 问题
-
-训练原先每个 step 都计算整个 trainable pytree 的 `grad_norm` 和大部分 kernel 参数的 `param_norm`，但这些指标只在 `log_interval` step 被记录。
-
-### 实现
-
-`train_step` 增加静态布尔参数 `compute_metrics`：
-
-- 日志 step 计算真实 `grad_norm` 和 `param_norm`；
-- 非日志 step 返回同结构的 NaN 标量，保持 JIT 输出 pytree 稳定；
-- host 聚合使用 `nanmean`，忽略非日志 step 的占位值；
-- 该布尔值作为静态位置参数传给 JIT，分别编译计算指标和跳过指标的两个变体。
-
-使用位置参数是必要的：当前 JAX 在显式设置 `in_shardings` 时不接受 pjit 关键字实参。
-
-### 结果
-
-- 共享 Tensor 基线：`2.828845 s/step`
-- 仅日志 step 计算范数：`2.814683 s/step`
-- step 时间减少 `0.50%`，等价吞吐提升 `0.50%`
-
-收益存在但很小，说明范数归约不是当前主要瓶颈。这个改动仍值得保留，因为它与日志语义一致，且不会牺牲已记录指标。
-
-## 5. 第四步：三路相机合并为一次 SigLIP 调用
-
-### 问题
-
-`Pi0.embed_prefix` 原先对 head、left wrist、right wrist 三路相机分别调用一次相同的 SigLIP image encoder。这会重复进入相同计算图，并重复触发参数使用和 kernel launch。
-
-### 实现
-
-新增可选配置 `batch_image_views`。开启后：
-
-1. 将三路图像堆叠为 `(B, V, H, W, C)`；
-2. reshape 为 `(B*V, H, W, C)`，只调用一次 image encoder；
-3. 将 token reshape 回 `(B, V, ...)`，按原相机名拆分。
-
-布局使用 sample-major interleave，同一样本的三路视图连续排列。这样在数据并行/FSDP shard 中，各设备仍拿到完整样本的所有视图，不需要跨设备重排。
-
-该开关默认关闭。生产训练可通过以下参数启用：
-
-```bash
---model.batch-image-views
-```
-
-### 结果
-
-- 只做范数优化：`2.814683 s/step`
-- 再合并三路相机：`2.773743 s/step`
-- 相对上一阶段 step 时间减少 `1.45%`，等价吞吐提升 `1.48%`
-- 相对共享 Tensor 基线累计吞吐提升 `1.99%`
-- 相对最初 NumPy queue 版本累计 step 时间减少 `15.32%`，等价吞吐提升 `18.10%`
-
-相机合并后，`dispatch` 从约 `2.657 s` 降到 `2.236 s`，但 `compute_sync` 从约 `0.115 s` 增到 `0.492 s`。这是异步工作在 dispatch 与同步边界之间重新分布，真实收益必须看两者与 `next_batch_total` 的总和，不能宣称 dispatch 单项显示的约 16% 提升。
-
-## 6. 第五步：cuDNN fused Attention
-
-### 问题
-
-Gemma Attention 原先显式执行 QK einsum、mask、softmax 和 PV einsum。XProf 已显示 GEMM/Fusion 占已记录 kernel 时间的 92.6%，因此训练 Attention 是比继续压缩数据搬运更有价值的候选。当前模型使用 BF16 GQA 和 block attention mask，必须验证 cuDNN 是否支持完整的真实训练形状，不能只测试一个小型无 mask MHA。
-
-### 实现
-
-新增可选配置 `use_cudnn_attention`。开启后，训练时调用：
-
-```python
-jax.nn.dot_product_attention(
-    q,
-    k,
-    v,
-    mask=attn_mask,
-    scale=1.0,
-    implementation="cudnn",
-)
-```
-
-`q` 在调用前已经按 head dimension 缩放，因此这里显式使用 `scale=1.0`，避免重复缩放。该路径只在 `kv_cache is None` 的训练路径启用；带 KV cache 的推理继续使用原实现。开关默认关闭，可通过 `--model.use-cudnn-attention` 启用。
-
-### cuDNN 版本和动态库选择
-
-当前环境同时存在两套可被动态链接器找到的 cuDNN：
-
-- `$CONDA_PREFIX/lib` 中的 cuDNN 9.14；
-- Python `site-packages/nvidia/cudnn/lib` 路径，以及元数据版本为 `9.10.2.21` 的 `nvidia-cudnn-cu12` wheel。
-
-cuDNN 不是由 `torch.backends.cudnn.version()` 统一替整个 Python 环境选择的。PyTorch 和 JAX 都要由 Linux 动态链接器根据当前进程的 `LD_LIBRARY_PATH`、wheel 的 RPATH/RUNPATH、已经加载的同 SONAME 库和导入顺序解析 `libcudnn.so.9` 及其组件。因此必须在实际启动训练的同一个 shell 中检查 JAX，而不能用 PyTorch 的版本输出推断 JAX。
-
-当前已激活 shell 的实际状态是：
-
-```text
-LD_LIBRARY_PATH=$CONDA_PREFIX/lib
-PyTorch cuDNN runtime = 91400
-JAX cuDNN runtime     = 91400
-JAX build cuDNN       = 90101
-```
-
-如果显式移除 `LD_LIBRARY_PATH`，PyTorch 和 JAX 都会解析到 9.10.2：
-
-```bash
-env -u LD_LIBRARY_PATH python -c \
-  'import jax; from jax._src.lib import cuda_versions; print(cuda_versions.cudnn_get_version())'
-# 91002
-
-env -u LD_LIBRARY_PATH python -c \
-  'import torch; print(torch.backends.cudnn.version())'
-# 91002
-```
-
-所以“PyTorch 是 9.14、JAX 默认是 9.10”不是同一个进程环境下两个框架固定选择不同版本；准确说法是：**当前 shell 有 `$CONDA_PREFIX/lib` 时两者都是 9.14，去掉该搜索路径时两者都是 9.10.2。** 之前观察到的 JAX 9.10.2 是在 `env -u LD_LIBRARY_PATH` 的诊断进程中得到的。
-
-强制 cuDNN Attention 时，9.10.2 对本模型测试形状报错：
-
-```text
-XlaRuntimeError: INTERNAL: No valid engine configs for Matmul_MUL...
-```
-
-使用环境中已经安装的 9.14 后，小型 MHA、真实 BF16 GQA 形状、block mask 前向和反向全部通过。因此不需要下载新包，但生产启动前必须保证 JAX 实际显示 `91400`：
-
-```bash
-export LD_LIBRARY_PATH="$CONDA_PREFIX/lib"
-
-python -c \
-  'from jax._src.lib import cuda_versions; print(cuda_versions.cudnn_get_version())'
-# 必须输出 91400
-
-python scripts/train.py pi05_base_bi_flexiv_earbud_case_insertion_0826_h200 \
-  --resume \
-  --model.use-cudnn-attention
-```
-
-### A/B 结果
-
-测试条件为 8×H200、全局 batch 512、`batch_image_views=true`、相同 step 16000 Checkpoint 和相同数据配置。两组均使用 cuDNN 9.14 运行时，唯一变量是是否启用 fused Attention。统计 step 16011–16040 的 30 个稳态样本，不包含数据 worker 启动、首批解码、Checkpoint 恢复和 JIT/cuDNN 编译。
-
-- 原显式 Attention：均值 `2.795537 s/step`，中位数 `2.797350 s/step`，标准差 `0.007647 s`；
-- cuDNN Attention：均值 `1.978580 s/step`，中位数 `1.979450 s/step`，标准差 `0.011704 s`；
-- step 时间减少 `29.22%`；
-- 等价吞吐从 `183.15 samples/s` 提升到 `258.77 samples/s`，提升 `41.29%`；
-- 从生产 step 16000 跑到 step 60000，按纯稳态 step 时间估算可节省约 `9.98 小时`。
-
-分项均值也显示收益来自模型计算，而不是数据波动：
-
-- 原实现：`dispatch=2.237247 s`，`compute_sync=0.484967 s`，`next_batch_total=0.073323 s`；
-- cuDNN：`dispatch=1.516877 s`，`compute_sync=0.394600 s`，`next_batch_total=0.067103 s`。
-
-数值对齐测试中，BF16 masked attention 前向最大绝对误差为 `0.00390625`、平均绝对误差为 `0.00026494`；Q/K/V 梯度最大绝对误差分别为 `1.91e-6`、`4.77e-7` 和 `2.38e-7`，属于正常 BF16 数值差异。
-
-### 2026-08-29 生产长跑验证
-
-生产训练从原 step 16000 Checkpoint 恢复，使用 8×H200、全局 batch 512、96 个 DataLoader worker、`batch_image_views=true`、`use_cudnn_attention=true`、`XLA_PYTHON_CLIENT_MEM_FRACTION=0.9`。启动日志确认 JAX 实际加载 cuDNN runtime `91400`。截至检查时，tmux 训练进程仍在运行，进度条约为 step 55100/60000；日志中没有 OOM、`No valid engine configs` 或 XLA runtime exception。
-
-速度统计使用日志中已经完整刷新的 step 16150–51650 记录。普通非日志、非保存 step 共 356 个样本：
-
-- **当前一个普通稳态 step：均值 `1.910590 s`，中位数 `1.910000 s`；**
-- 标准差 `0.010744 s`，范围 `1.86–1.95 s`；
-- 分项均值为 `dispatch=1.867331 s`、`next_batch=0.042612 s`；
-- 按 batch 512 折算吞吐约 `267.98 samples/s`；
-- 每 100 step 的日志 step 均值 `2.360375 s`，其中日志约 `0.436000 s`；
-- 每 1000 step 的 Checkpoint step 均值 `4.656286 s`，其中保存约 `2.288286 s`；
-- 按上述频率摊销日志和保存开销后，长期约为 **`1.917238 s/step`**。
-
-因此，生产长跑的速度与 40-step A/B 中的 `1.978580 s/step` 一致并略快；从纯吞吐看，cuDNN fused Attention 的约 1.9 秒/step 收益稳定存在。
-
-但是训练数值已经失效：
-
-```text
-Step 16100: grad_norm=nan, loss=0.1185, param_norm=nan
-Step 16200: grad_norm=nan, loss=nan,    param_norm=nan
-Step 16300 及之后: grad_norm=nan, loss=nan, param_norm=nan
-```
-
-step 16100 是第一个写入指标的生产日志点，因此只能确定 NaN 在 step 16000–16100 之间产生，不能从当前日志定位到更精确的单步。已完整刷新的日志一直到 step 51600 都没有再次出现有限 loss。周期保存又已经清理原 step 16000；当前保留的最后一个切换前 Checkpoint 是 step 15000，而 step 20000–55000 均是在 NaN 出现后保存，不能作为有效训练结果使用。
-
-结论是：**1.91 秒/step 是真实、稳定的硬件吞吐，但来自数值无效的训练，不能作为启用 cuDNN Attention 的生产验收。** 在定位并修复 NaN 前，应关闭 `use_cudnn_attention`；恢复有效训练时只能从仍保留的切换前 Checkpoint 重新开始。后续诊断需要缩短日志间隔并逐 step 检查 loss、梯度和参数有限性，重点覆盖真实 block mask 下的完整优化器更新，而不只是独立 Attention 前向/反向。
-
-## 7. 正确性和运行验证
-
-- 增加单元测试，对三路相机分别编码与合并编码的 token 做逐元素相等比较；
-- 模型与 YAML 相关回归测试：`23 passed`；
-- 前一阶段数据管线与显式 shutdown 回归：`37 passed`；
-- 范数优化和相机合并各完成一次 40-step、8×H200 实际训练；
-- 两次运行都成功恢复 step 49000 Checkpoint，没有 OOM，没有保存新 Checkpoint；
-- `compute_metrics=True/False` 两个 JIT 分支都在真实训练中执行；
-- 两次运行结束后显式 shutdown 成功，worker 无残留，8 张 GPU 显存全部释放；
-- 本地与训练服务器上的修改文件 SHA256 一致。
-- cuDNN Attention 小型与真实形状的前向/反向 smoke test 通过；
-- cuDNN Attention 相关模型回归：`7 passed`；
-- 原 Attention 和 cuDNN Attention 均完成一次 step 16000–16040 的 8×H200 短跑，没有 OOM、cuDNN engine fallback 或训练异常，也没有保存新 Checkpoint；
-- cuDNN 短跑退出时个别 DataLoader worker 在清理阶段打印 `killed by signal: Aborted`，发生在训练完成和计时结束之后，主训练进程与 GPU 资源均正常退出；这不是 Attention 计算错误。
-- 后续生产长跑推翻了短跑的数值有效性结论：step 16100 首次记录到 NaN 梯度，step 16200 起 loss 持续为 NaN；这说明现有 smoke test 和 40-step 吞吐短跑不足以作为生产正确性验证。
-- 最小 H200 复现确认全空 mask 行会让原 cuDNN 路径产生 16384 个 Q 梯度 NaN；修复后 Q/K/V 梯度全部有限，且有效 query 输出不变；
-- 从干净 step 15000 checkpoint 做 8×H200、batch 512 的 60-step 实际短跑，loss、梯度、参数和更新均保持有限，`update_is_finite=1`。
-
-## 8. 当前结论与下一步
-
-当前稳态数据交付只占约 `0.04–0.05 s/step`，H2D 只有数毫秒；XProf 又显示已记录 kernel 时间主要集中在 GEMM/Fusion。因此，继续微调 queue、collate 或 H2D 不太可能带来显著收益。
-
-cuDNN fused Attention 的旧长跑普通 step 为 `1.910590 s`，但该区间从最早日志点起梯度、参数已经 NaN，不能作为有效训练的吞吐基线。修复后连续有限训练的普通 step 稳定约为 `2.31 s`；移除逐 step 有限性保护和把安全 mask 计算移出扫描层都没有改变该结果。最终实现保留原 mask，仅停止全空 query 的 Q 梯度，并保留原有 loss/grad norm/param norm 日志。确认 JAX runtime 为 91400 仍只能证明动态库选择正确；生产验收需要持续观察日志。
-
-下一轮稳态优化建议按以下顺序推进：
-
-1. **完成修复后的生产长跑验收。** 从干净 checkpoint 恢复并持续观察 loss、grad norm 和 param norm；当前策略不会自动拒绝或停止非有限更新，发现异常时需人工停止并回退到前一个干净 checkpoint。
-2. **缩短有效 token 长度。** 全数据扫描得到 Pi0.5 tokenized state 最大长度 165；可单独 A/B `max_token_len=168` 与当前 200，减少 Gemma attention/MLP 的序列计算，同时保留 3 token 余量。
-3. **继续拆最大的 GEMM/Fusion。** 在生产正确性验收后重新采集 trace，再确认剩余最大的 SigLIP/Gemma GEMM 和 fusion。
-4. **降低 checkpoint 干扰。** 长训练若不需要每 1000 step 恢复点，可提高 `save_interval`；此前生产日志已经确认 checkpoint 会造成下一 step 的数据等待尖峰。
-
-最终判断：共享内存 Tensor 消除了主要数据交付开销；日志 step 范数和三路相机合并是低风险增量收益；cuDNN fused Attention 的全空 mask 梯度 NaN 已定位并修复，最小复现和 60-step 真实训练通过。当前最高优先级是让从 step 15000 恢复的训练完成更长的有限性验收。
diff --git a/docs/2026-08-31-cudnn-sdpa-nan-summary.md b/docs/2026-08-31-cudnn-sdpa-nan-summary.md
deleted file mode 100644
index fc8fbb6..0000000
--- a/docs/2026-08-31-cudnn-sdpa-nan-summary.md
+++ /dev/null
@@ -1,487 +0,0 @@
-# cuDNN fused Attention NaN 与训练发散：统一事故报告
-
-**覆盖日期**：2026-08-29 至 2026-08-31
-
-**涉及机器**：8×H200、8×H100 80GB
-
-**涉及实现**：JAX `dot_product_attention(..., implementation="cudnn")`、BF16 GQA
-
-**最终状态**：全空 mask 的瞬时 NaN 已修复；全 18 层 cuDNN 的缓慢训练发散仍未解决。生产训练不得启用全 18 层 cuDNN。
-
-> [!CAUTION]
-> 本次调查包含两种不同的失效，不能混为一个 cuDNN bug：
->
-> 1. **瞬时 NaN**：全空 attention mask query 行使旧环境中的 cuDNN backward 产生 NaN Q 梯度。`_stop_gradient_for_fully_masked_queries` 已修复这一问题。
-> 2. **缓慢发散**：修复全空行后，全 18 层 cuDNN 仍在 step 1000–1200 与显式 attention 基线分叉，随后梯度爆炸并最终变为 NaN。根因是 BF16 实现差异的坐标结构被 Adam 放大，而不是非 128 倍序列长度、cuDNN 版本或 mask 修复。
-
-> [!IMPORTANT]
-> 截至本文最后一次验证：
->
-> - **稳妥生产方案**：`use_cudnn_attention: false`；
-> - **已通过 3000 步短程验收的折中方案**：只在 Gemma 最后 3 层启用 cuDNN；
-> - **禁止用于生产**：全 18 层 cuDNN，即 `use_cudnn_attention: true` 且 `cudnn_attention_num_layers: null`；
-> - 已经出现 NaN 的参数和 Adam 状态不可修复，必须从确定干净的 checkpoint 恢复。
-
-## 1. 背景与真实训练形状
-
-Pi0.5 训练时的 Gemma attention 序列由三部分组成：
-
-```text
-3 × 256 个图像 token + 200 个文本 token + 50 个 action token = 1018
-num_heads=8, num_kv_heads=1, head_dim=256, dtype=bf16
-```
-
-prompt 被补齐到 `max_token_len=200`。padding token 既不是有效 query，也不是有效 key，因此 `make_attn_mask` 会生成整行全 `False` 的 query 行。训练使用 block mask；只有 action token 上的 cotangent 非零。
-
-显式 attention 路径执行 QK einsum、有限 `big_neg` mask、softmax 和 PV einsum。cuDNN 路径调用：
-
-```python
-jax.nn.dot_product_attention(
-    q,
-    k,
-    v,
-    mask=attn_mask,
-    scale=1.0,
-    implementation="cudnn",
-)
-```
-
-`q` 已在调用前按 head dimension 缩放，因此必须使用 `scale=1.0`，避免重复缩放。
-
-## 2. 第一类失效：全空 mask 行导致瞬时 NaN
-
-### 2.1 事故范围
-
-**日期**：2026-08-29
-
-**机器**：8×H200
-
-**配置**：`pi05_base_bi_flexiv_earbud_case_insertion_0826_h200`
-
-**触发方式**：从 clean step 16000 切换到 `use_cudnn_attention=true` 后恢复生产训练。
-
-切换前显式 attention 指标正常：
-
-```text
-Step 16000: grad_norm=0.6721, loss=0.1429, param_norm=1806.9792
-Step 16100: grad_norm=1.1132, loss=0.1456, param_norm=1806.9926
-Step 16200: grad_norm=0.8198, loss=0.1402, param_norm=1807.0051
-Step 16300: grad_norm=0.9735, loss=0.1410, param_norm=1807.0171
-```
-
-从 checkpoint 切到 cuDNN 后：
-
-```text
-Step 16100: grad_norm=nan, loss=0.1185, param_norm=nan
-Step 16200: grad_norm=nan, loss=nan,    param_norm=nan
-Step 16300 及以后: grad_norm=nan, loss=nan, param_norm=nan
-```
-
-checkpoint 目录名记录循环 step，而保存的 `train_state.step` 已经加一。恢复后的第一个前向 loss 仍有限，但同一步反向产生 NaN 并污染参数，之后所有前向均为 NaN。
-
-日志窗口对 loss 使用 `nanmean`，第一个窗口里唯一有限的首步 loss 使聚合结果仍显示 `0.1185`，掩盖了该窗口其余 NaN。因此“窗口 loss 有限”不能证明窗口内每个更新都有限。
-
-### 2.2 最小复现与根因
-
-在当时实际加载 cuDNN runtime `91400` 的环境中，使用 BF16 GQA 和含 16 个全空 query 行的 mask：
-
-```text
-forward output finite: true
-Q/K/V gradient finite: [false, true, true]
-Q gradient NaN count: 16384
-```
-
-为全空 query 行临时开放一个 dummy key 后：
-
-```text
-Q/K/V gradient finite: [true, true, true]
-Q/K/V gradient NaN count: [0, 0, 0]
-valid query rows unchanged: true
-```
-
-完整生产形状也复现了相同现象：原始 backward 产生 71680 个 Q 梯度 NaN，K/V 梯度有限。
-
-### 2.3 实际修复
-
-最终实现保留原始 attention mask，不为全空行打开 dummy key。调用 cuDNN 前找出至少存在一个有效 key 的 query 行，并只停止全空 query 行的 Q 梯度：
-
-```python
-query_has_key = jnp.any(attn_mask, axis=-1)[:, 0, :, None, None]
-q = jnp.where(query_has_key, q, jax.lax.stop_gradient(q))
-```
-
-该变换具有以下语义：
-
-- Q 的前向值逐元素不变；
-- 有效 query 的输出和梯度不变；
-- 无效 query 的 Q 梯度固定为 0；
-- 原始 block mask 完整保留。
-
-dummy-key 版本在有效行上与 stop-gradient 版本给出逐位相同的 dQ/dK/dV，但每层重建 `(B, 1, T, S)` mask，batch 256 下约为 `1.38 s/step`；stop-gradient 版本约为 `1.05 s/step`，因此保留后者。
-
-### 2.4 修复验证及其局限
-
-当时完成的验证包括：
-
-- `src/openpi/models/pi0_test.py`：`8 passed`；
-- `scripts/train_test.py`：`2 passed`；
-- H200/cuDNN BF16 GQA 真实形状 Q/K/V 梯度全部有限；
-- 从 clean step 15000 独立恢复，连续执行 step 15001–15060；
-- 每 10 步的 loss、grad norm、param norm 和 `update_is_finite` 均有限。
-
-这足以证明第一类“第一步反向立即 NaN”已被修复，但不足以证明训练长期收敛。后续事实表明，60 步、200–500 步、已收敛 checkpoint 短跑和单次有限性 smoke test 都会漏掉第二类失效。
-
-### 2.5 checkpoint 处置
-
-H200 事故中最后一个确定干净且仍保留的 checkpoint 是 step 15000。原 step 16000 已被保留策略删除，step 20000–55000 受到 NaN 污染并已删除。从 step 15000 恢复意味着重跑约 1000 个有效 step。
-
-已经包含 NaN 的参数、Adam 一阶矩或二阶矩不能通过更换 attention 实现恢复；不得继续使用污染 checkpoint。
-
-## 3. 第二类失效：全层 cuDNN 在约 1000 步后缓慢发散
-
-### 3.1 现象
-
-全空 mask 修复后，全 18 层 cuDNN 在前约 1000 步与显式 attention 基线贴合，随后在 step 1000–1200 分叉，梯度持续放大，最终变为 NaN。
-
-```text
-step    显式基线             cuDNN 9.14            cuDNN 9.19
- 200    0.9244 /   2.89      0.9333 /   3.58       0.9348 /   3.53
-1000    0.4220 /   2.77      0.4306 /   2.50       0.4319 /   3.16
-1100    0.4077 /   1.79      0.4784 /   3.16       0.4723 /   3.37
-1200    0.3742 /   2.12      0.6484 /  31.40       0.6600 /  35.49
-2500    0.2457 /   1.33      5.7692 / 675.08       4.7055 / 759.18
-5900    0.1763 /   1.36      nan                    已停止
-```
-
-每个单元格为 `loss / grad_norm`。
-
-因果起点是 step 1000–1200，而不是最终出现 NaN 的 step 5900。NaN 只是长期发散后的溢出结果。
-
-### 3.2 单变量责任归属
-
-同机、同数据集、同超参、同 seed 的 A/B 中，唯一变量是 `use_cudnn_attention`：
-
-```text
-cve5hiyt（0828）  commit d0aec27  显式 attention  step 59900 loss=0.0589
-orbnh1b7（0829）  commit 51d9af3  全层 cuDNN     step 5900  loss=nan
-```
-
-`git diff d0aec27 HEAD -- src/ scripts/` 的 72 行增删全部属于 attention 开关本身，排除了学习率、数据集和训练配置变化。
-
-### 3.3 cuDNN 单步结果并非明显错误
-
-同一权重、真实 batch 和 rng 的完整模型梯度 A/B：
-
-```text
-batch 0  loss 显式=5.859448  cuDNN=5.851340  |grad| 104.31 vs 103.65  ratio=0.9937
-batch 1                                      |grad|  67.19 vs  67.45  ratio=1.0038
-batch 2                                      |grad|  70.71 vs  70.61  ratio=0.9986
-
-相对梯度差：1.75e-02 / 1.66e-02 / 2.12e-02
-```
-
-梯度范数比值接近 1，没有单个模块出现异常放大。单层真实形状中，cuDNN 对 fp32 参考的相对误差约 `5.4e-03`，显式 BF16 路径约 `3.2e-03`，两者同阶。
-
-同一 batch 重复 cuDNN backward 的随机差异约 `0.9%`，来自 atomic 累加；cuDNN 与显式路径之间的确定性分量约 `1.5%`。因此 cuDNN 给出的是另一个合法的 BF16 近似，而不是单步明显错误的梯度。
-
-## 4. 根因机制：Adam 放大坐标级差异
-
-### 4.1 差分没有固定跨 batch 方向
-
-对 8 个真实 batch 分别计算显式梯度和 3 次 cuDNN backward 均值：
-
-```text
-跨 batch delta 两两 cosine：-0.004043 ± 0.126301，范围 [-0.247516, +0.226915]
-||delta|| / ||g_explicit||： 1.556024e-02 ± 2.76e-03
-cuDNN 同 batch 随机分量：   8.752676e-03
-cos(delta, gradient)：      -0.00113 ± 0.225
-cos(delta, parameter)：     +0.000016
-```
-
-差分没有稳定的跨 batch 方向，也不沿当前梯度或参数方向累积。
-
-### 4.2 Adam 更新空间暴露了真正差异
-
-从 strict-order 全层 cuDNN run 的 step 1000 checkpoint 恢复真实 Adam 状态，对同一真实 batch 和 rng 比较实际参数更新：
-
-```text
-                         相对更新差异   更新 cosine   符号翻转坐标
-真实全层 cuDNN 差分          3.3521%      0.999440       0.7885%
-1.5% 人造乘性噪声            0.7313%      --             0.1459%
-```
-
-cuDNN 原始梯度 L2 差异只有约 1.5%，但其误差分布会扰动小梯度、小二阶矩坐标，经过 Adam 逐坐标归一化后变成约 3.35% 的更新差异，并使约 0.79% 的非零更新坐标翻转符号。
-
-相比之下，按 `|g|` 分配能量的 1.5% 独立随机噪声几乎不扰动这些敏感坐标。该人造噪声从 step 0 稳定跑满 3000 步，step 2900 为 `loss=0.2411 / grad_norm=1.6389`，证明训练配方并非对任意 1.5% 梯度噪声都没有裕度。
-
-### 4.3 早层误差会穿过后续网络放大
-
-在同一个 step-1000 Adam 状态上，仅让不同连续层段使用 cuDNN：
-
-```text
-cuDNN 层段       相对更新差异   更新 cosine   符号翻转坐标
-全 18 层            3.3118%       0.999451       0.7769%
-最前 6 层           3.2061%       0.999486       0.7548%
-中间 6 层           2.2169%       0.999754       0.5141%
-最后 6 层           1.1633%       0.999932       0.2720%
-最前 3 层           3.3016%       0.999455       0.7720%
-中间 3 层           1.9894%       0.999802       0.4634%
-最后 3 层           0.8470%       0.999964       0.1926%
-稳定噪声参考         0.7313%       --             0.1459%
-```
-
-只融合最前 3 层几乎等同全 18 层。层的位置比数量更重要：早层的微小数值差异会穿过后续 15 层前向和反向传播，显著放大。
-
-## 5. 已否定的修复假设
-
-### 5.1 升级 cuDNN 不能解决缓慢发散
-
-cuDNN 9.14 与 9.19 的发散曲线在 step 1000–1200 的同一位置分叉。当前统一环境实际加载 `nvidia-cudnn-cu12==9.19.0.56`，全层训练仍发散。
-
-### 5.2 非 128 倍序列长度不是本任务根因
-
-cuDNN 9.15.1 release note 提到旧版 SDPA backward 在静态序列长度不是 128 倍数时可能产生错误结果或 NaN，因此曾怀疑 `T=1018` 命中该缺陷。真实形状对照结果：
-
-```text
-T=1018 (%128=122)  dQ rel-L2: 显式 3.23e-03  cuDNN 3.21e-03  NaN=0
-T=1024 (%128=0)    dQ rel-L2: 显式 3.23e-03  cuDNN 3.21e-03  NaN=0
-```
-
-把 `max_token_len` 从 200 临时改为 206 不能改变数值行为。该 NVIDIA 缺陷与本任务的缓慢发散不是同一问题。
-
-参考：
-
-- [NVIDIA/cudnn-frontend #160](https://github.com/NVIDIA/cudnn-frontend/issues/160)
-- [cuDNN 9.15.1 Release Notes](https://docs.nvidia.com/deeplearning/cudnn/backend/v9.15.1/release-notes.html)
-
-### 5.3 原生 seq_lengths / padding kernel 不能解决发散
-
-2026-08-31 对旧 summary 建议的方案进行了最终验证：向 cuDNN 调用同时传入完整物理长度的 `query_seq_lengths` 和 `key_value_seq_lengths`，值均为 1018；实际 padding 语义仍由原 block mask 负责。
-
-编译结果确认 seq lengths 已作为两个 `s32[B]` operand 进入 `__cudnn$fmhaScaleBiasSoftmax` 和 backward custom call，但未切换到数值上不同的算法。
-
-真实单层形状、三次重复 backward：
-
-```text
-NaN count，无 seq_lengths： [0, 0, 0]
-NaN count，有 seq_lengths： [0, 0, 0]
-相对 fp32 梯度误差，无 seq_lengths：约 5.66924e-03
-相对 fp32 梯度误差，有 seq_lengths：约 5.66927e-03
-两路径相对梯度差：5.3e-06 至 3.1e-05
-```
-
-两路径之差与同一 cuDNN backward 的 atomic 随机波动同量级。
-
-在 step-1000 真实 Adam 状态上的全模型对照：
-
-```text
-显式 loss：                 0.31219399
-padding-cuDNN loss：        0.31204814
-Adam update 相对差：        3.310997%
-update cosine：             0.99945178
-符号翻转坐标：              0.77606%
-```
-
-这与已知会发散的全 18 层 cuDNN 指纹 `3.3118% / 0.7769%` 几乎完全一致。把 Adam `eps` 提高到 `1e-6` 或 `1e-5` 后，update 相对差仍分别为 `3.1655%` 和 `3.0275%`，符号翻转仍约 `0.776%`。
-
-最后又完成了 3000-step、8×H100、batch 256、`strict_batch_order=true` 的全 18 层 padding-cuDNN 验收。W&B、周期 checkpoint 和最终 checkpoint 均关闭：
-
-```text
-step    loss      grad_norm    param_norm
-   0    6.4023       65.6324    1802.3865
- 500    0.5719        2.5464    1802.4017
-1000    0.4298        2.4472    1802.5095
-1100    0.4591        2.2270    1802.5511
-1200    0.5902       13.5020    1802.5952
-1300    0.9563      108.9634    1802.6366
-1500    2.2053       45.7198    1802.7318
-2000    3.7753      150.5048    1802.9735
-2400    4.7852     6469.9019    1803.1691
-2500    5.7166      609.2015    1803.2141
-2900    6.2039      603.4764    1803.3545
-```
-
-step 1200 的显式基线约为 `0.3742 / 2.12`，padding-cuDNN 梯度已达到约 6.4 倍，超过 3 倍验收上限。该方案跑满 3000 步时尚未直接变成 NaN，但已沿已知轨迹持续发散，不能作为解决办法。
-
-### 5.4 其他已排除项
-
-- **两种 mask 修复**：dummy key 与 Q stop-gradient 在有效行上数值等价，两者都会出现第二类发散；
-- **动态库混装**：统一到单源 pip cu128/cuDNN 9.19 后仍发散；
-- **logit 幅度**：std 1–80 扫描中 cuDNN 与显式误差同阶；
-- **head_dim=256**：与 head_dim=128 表现一致；
-- **cotangent 稠密度**：只在 action token 非零和所有有效 token 稠密时结果一致；
-- **dBias 路径**：实际 batch 256 不满足 `bias.shape[0]==1`，未进入该路径；
-- **FSDP 分片**：attention 只按 batch 分片，每卡拥有完整序列；
-- **batch 到达顺序**：`strict_batch_order=true` 仍在 step 1000–1200 分叉；
-- **降低 peak LR**：`2.5e-5 -> 1.5e-5` 只把分叉推迟约 300 步；
-- **提高 Adam b2**：`0.95 -> 0.99` 从 step 0 跑满 3000 步仍发散；
-- **提高 Adam b1**：`0.9 -> 0.99` 只把明显分叉推迟到约 step 1400；
-- **提高 Adam eps**：update 差异和符号翻转几乎不变；
-- **1.5% 独立随机梯度噪声**：显式 attention 稳定跑满 3000 步，没有复现发散。
-
-## 6. 已验证的稳定折中：只融合最后 3 层
-
-配置 `configs/diag_cudnn_last3.yaml` 使用：
-
-```yaml
-model:
-  use_cudnn_attention: true
-  cudnn_attention_layer_start: 15
-  cudnn_attention_num_layers: 3
-```
-
-8×H100、batch 256、strict-order、原始 Adam，从 pi05_base step 0 开始，完整执行 3000 个更新：
-
-```text
-step    loss      grad_norm    param_norm
-   0    6.4029      65.5204     1802.3865
- 500    0.5719       2.4526     1802.4014
-1000    0.4240       2.3789     1802.5061
-1100    0.4050       2.1847     1802.5416
-1200    0.3757       2.3869     1802.5775
-1600    0.3238       2.6505     1802.7225
-2000    0.2793       1.8254     1802.8666
-2500    0.2463       2.0116     1803.0460
-2900    0.2424       1.4769     1803.1906
-```
-
-该曲线与显式基线一致，满足 3000 步验收。普通 step 约 `1.5 s/step`，相比显式路径约 `1.57–1.64 s/step` 只有约 4–9% 加速，明显低于全层 cuDNN 的约 36% 潜在收益。
-
-最后 3 层方案只是已通过短程验收的保守折中，尚未完成 60000 步生产验收，不能等同于完整生产结论。
-
-## 7. CUDA/cuDNN 环境事故与当前要求
-
-### 7.1 历史混装问题
-
-事故期间 `cuda_versions.cudnn_get_version()` 返回 `91400`，但进程实际加载了 9.10.2 dispatcher stub 和 9.14 engine 库：
-
-```text
-site-packages/nvidia/cudnn/lib/libcudnn.so.9          <- 9.10.2 dispatcher
-$CONDA_PREFIX/lib/libcudnn_graph.so.9.14.0            <- 9.14 engine
-$CONDA_PREFIX/lib/libcudnn_engines_precompiled.so.9.14.0
-$CONDA_PREFIX/lib/libcudnn_ops.so.9.14.0
-```
-
-`cudnnGetVersion()` 的返回值来自 engine，因此只看 `91400` 无法发现 dispatcher/engine 混装。
-
-### 7.2 当前统一环境
-
-环境已统一为：
-
-```text
-CUDA Toolkit          12.8
-PyTorch               2.11.0+cu128
-torchvision           0.26.0+cu128
-torchcodec             0.11.1+cu128
-nvidia-cudnn-cu12      9.19.0.56
-JAX                    0.5.3
-```
-
-启动训练时不要导出 `LD_LIBRARY_PATH="$CONDA_PREFIX/lib"`。该目录中残留的 conda-forge `libcublas`、`libcudart` 或 `libnccl` 可能遮蔽 pip cu128 库。
-
-生产启动前在相同 shell 中运行：
-
-```bash
-env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py
-```
-
-脚本会枚举 `/proc/self/maps` 中实际加载的 cuDNN/cuBLAS/NCCL 路径，并在真实 attention 形状上做前向、反向和 fp32 参考对比。
-
-> [!NOTE]
-> `check_cuda_stack.py` PASS 只证明环境单源且单步内核有限、误差在阈值内。第二类缓慢发散曾通过所有此类 smoke test，因此 PASS 不能替代 3000 步训练验收。
-
-## 8. 生产处置与恢复规则
-
-### 8.1 当前可选配置
-
-需要绝对稳妥时：
-
-```yaml
-model:
-  use_cudnn_attention: false
-```
-
-需要先获取小幅吞吐收益并接受受监控长跑时：
-
-```yaml
-model:
-  use_cudnn_attention: true
-  cudnn_attention_layer_start: 15
-  cudnn_attention_num_layers: 3
-```
-
-不得用于生产：
-
-```yaml
-model:
-  use_cudnn_attention: true
-  cudnn_attention_num_layers: null  # 全 18 层
-```
-
-### 8.2 checkpoint 恢复
-
-- 一旦参数、Adam `mu` 或 `nu` 出现 NaN，后续 checkpoint 均视为污染；
-- 不得通过切换 attention、降低 LR 或修改 optimizer 后继续污染 checkpoint；
-- 必须回退到最后一个确定干净的 checkpoint；
-- 恢复前检查 checkpoint 内参数、optimizer state 和 step；
-- 使用独立实验名验证，关闭 W&B、周期保存和最终保存，避免覆盖生产目录；
-- 只有通过完整验收后才能恢复生产保存。
-
-### 8.3 建议增加发散护栏
-
-当前训练只在日志 step 计算 `grad_norm`。建议在 host 端维护滚动中位数，并在连续多个日志窗口超过固定倍数时中止。按本次曲线，该护栏可在 step 1200–1300 停止，而不是等到 loss 变成 NaN 或污染数万个 step。
-
-护栏只能减少损失，不能让全层 cuDNN 变得收敛。
-
-## 9. 新 attention 实现的统一验收标准
-
-任何新的全层 fused/flash attention backend、custom VJP 或稳定化方案必须：
-
-1. 从 **step 0** 开始；
-2. 使用 `strict_batch_order=true`，固定 seed、数据集、batch 和超参；
-3. 跑满 **3000 个更新**，不得只跑 40、60 或 500 步；
-4. 与显式 attention 基线逐点比较 step 200/500/1000/1500/2000/2900；
-5. 全程 loss、梯度、参数和 optimizer state 有限；
-6. `grad_norm` 不得超过基线同点的 3 倍；
-7. 通过后再进入有 checkpoint 的受监控长跑；
-8. 生产长跑仍需持续观察 loss 和 grad norm。
-
-以下均不是有效正确性验收：
-
-- 单次前向/反向有限；
-- 40-step 吞吐短跑；
-- 从已经收敛的 checkpoint 短跑；
-- 200–500 个连续有限 step；
-- 只检查最终是否出现 NaN；
-- 跨不同 batch 顺序的 run 直接比较单点 grad norm。
-
-## 10. 不要重复的弯路
-
-- 不要再次把全空 mask 的瞬时 NaN 与全层缓慢发散当成同一问题；
-- 不要再把 cuDNN 9.15.1 或 `T=1024` 当作本任务的修复；
-- 不要再尝试完整物理长度的 `query_seq_lengths/key_value_seq_lengths`；
-- 不要再修改 dummy-key/stop-gradient mask 语义来解决第二类发散；
-- 不要只调低 `peak_lr` 或提高 Adam `b1`、`b2`、`eps`；
-- 不要用 1.5% 独立随机噪声代替真实 cuDNN 差分；
-- 不要只看 `cuda_versions.cudnn_get_version()` 判断动态库是否干净；
-- 不要把最后 3 层的 3000 步通过误写成全 18 层或 60000 步生产通过。
-
-## 11. 相关脚本与配置
-
-- `scripts/check_cuda_stack.py`：检查实际映射的 CUDA/cuDNN 库并做真实形状数值 smoke test；
-- `scripts/grad_ab_probe.py`：完整模型显式/cuDNN 梯度 A/B；
-- `scripts/grad_ab_direction_probe.py`：多 batch 差分方向和随机分量；
-- `scripts/optimizer_update_ab_probe.py`：从真实 Adam 状态比较参数更新误差；
-- `scripts/hybrid_update_ab_probe.py`：比较不同 cuDNN 层段的更新空间误差；
-- `configs/diag_cudnn_strict_order.yaml`：全层 cuDNN strict-order 诊断；
-- `configs/diag_cudnn_last3.yaml`：最后 3 层 cuDNN 稳定折中诊断。
-
-## 12. 文档整合说明
-
-本文整合并取代以下三份调查记录作为当前权威结论：
-
-- `2026-08-29-cudnn-attention-nan-incident.md`：全空 mask 的瞬时 NaN 事故；
-- `2026-08-30-cudnn-attention-divergence.md`：全层 cuDNN 缓慢发散、Adam 机制和最后 3 层方案；
-- 原 `2026-08-31-cudnn-sdpa-nan-summary.md`：关于非 128 倍长度、升级 cuDNN 和 padding kernel 的待验证假设。
-
-整合完成后，前两份独立文件已删除，历史内容仍可通过 Git 提交记录追溯。后续调查和生产决策以本文为准。
diff --git a/docs/2026-09-03-cudnn-dterm-rounding-analysis.md b/docs/2026-09-03-cudnn-dterm-rounding-analysis.md
deleted file mode 100644
index 2ab8164..0000000
--- a/docs/2026-09-03-cudnn-dterm-rounding-analysis.md
+++ /dev/null
@@ -1,196 +0,0 @@
-# 全层 cuDNN 发散的根因：flash 反向 D 项用了 bf16 舍入后的 O
-
-**日期**：2026-09-03
-
-**机器**：8×H100 80GB（ebcloud，hb-openpi 环境：jax 0.5.3、nvidia-cudnn-cu12 9.19.0.56，`check_cuda_stack.py` PASS）
-
-**前置报告**：`docs/2026-08-31-cudnn-sdpa-nan-summary.md`。那份报告证明了全 18 层 cuDNN 的缓慢发散是"cuDNN 与显式路径的 BF16 差分被 Adam 逐坐标放大"，但没有回答差分从哪里来、为什么它不像随机噪声。本文用真实激活回答这两个问题，并给出一个不改内核的缓解手段及其验收结果。
-
-> [!IMPORTANT]
-> 结论摘要：
->
-> 1. cuDNN fused attention 的反向按 flash-attention 方式算 `dS = P ⊙ (dP − D)`，其中 `D_i = rowsum(dO_i ⊙ O_i)` 用的是**前向存下的、已舍入到 bf16 的 O**。显式路径用同一份舍入后的 dP 算 `D_i = Σ_j P_ij dP_ij`，抵消是自洽的。两者的差是一个每行标量 `Δ_i = dO_i · (bf16(O_i) − O_i)`，进入梯度后方向为 `−Δ_i · Σ_j P_ij K_j`，破坏了 softmax 梯度"行和为零"的不变量。
-> 2. 该分量在真实激活上可被逐层预测：cuDNN bf16 的 dQ 误差与预测方向的 cos 在全部行上为 0.36–0.55，在 peaked 行（maxP>0.99）上为 0.66–0.99；显式 bf16 与 cuDNN fp16 的 cos 均 ≈ 0。
-> 3. peaked 行的真实 dQ 是 ε 级抵消量，所以那里 cuDNN bf16 的相对误差达 7%–140%（显式路径 0.2%–0.5%）。它们绝对量小，但方向沿该行 argmax 键、且跨 batch 落在同一低维子空间，正是 Adam 会放大的结构性误差。
-> 4. 把 cuDNN 内核改为 fp16（多 3 位尾数，配合 2 的幂动态缩放余切）后，peaked 行误差降约 8 倍，全局 dQ 误差 1.76e-3 低于显式 bf16 的 2.33e-3，结构性分量消失（cos ≈ 0）。
-> 5. "把 V 减去 sink 键的 V 再送内核"无效，因为 peaked 行大多不在看 BOS sink，而是在看 action/text/image 的其他键；BOS 的 |V| 本来就只有 1–5。
-
-## 1. 方法
-
-两个探针，都用 `configs/diag_cudnn_strict_order.yaml` 的模型与数据、pi05_base 权重、一个真实 batch（B=8，单卡）：
-
-- `scripts/attention_dterm_probe.py`（GPU）：monkeypatch `gemma.Block.__call__` 与 `gemma._cudnn_attention_call`，在一次真实训练步里用 `jax.debug.callback` 抓下 18 层的 q/k/v/mask 和反向余切 dO（缓存到 `dterm_capture.npz`，1.5 GB）。然后在单层上以 fp32 显式 attention 为真值，比较五个变体的 dQ/dK/dV：`explicit_bf16`（生产路径）、`cudnn_bf16`（历史路径）、`cudnn_fp16`（候选）、`cudnn_bf16_shift`/`cudnn_fp16_shift`（V 减 V_sink）。同时按上面的公式算 H1 预测误差，用 `jax.lax.reduce_precision` 模拟 bf16 舍入（`astype` 往返会被 XLA 的 `allow_excess_precision` 删掉，预测会恒为零）。
-- `scripts/attention_dterm_cpu_analysis.py`（CPU，float64）：读同一份缓存，对 peaked 行逐行重建精确反向，回答"哪些行 peaked、真实 dQ 有多小、H1 误差占多少"。
-
-## 2. 事实
-
-### 2.1 sink 在哪里
-
-18 层、8 个样本，注意力质量最大的键全部是位置 768 = 文本第一个 token（BOS），占有效注意力质量 8%–59%（第 14–15 层最高）。BOS 键的 |K| 8–13（其他键中位数 17–29），**|V| 只有 0.8–5（其他键中位数 10–58）**，第 0/16/17 层例外（9–17）。
-
-### 2.2 哪些行是 peaked（maxP > 0.99）
-
-peaked 行占有效行 0.1%–7%，且**大多不在看 BOS**：
-
-```text
-层   peaked 行   其中看 BOS   其余在看
- 0     3978        1270       img 1108, txt 1600
- 4      131           0       act 131
- 8       64           0       act 48, txt 16
- 9       98           0       txt 90, img 8
-14     1885        1802       txt 74, img 9
-15      323         168       img 107, txt 48
-17      561          47       txt 357, img 153（该层非 action 行 dO ≡ 0）
-```
-
-### 2.3 peaked 行上的误差解剖（CPU，float64）
-
-每行几何平均。`|dQ_true|` 是精确梯度范数，`errH1` 是 H1 预测误差，`shift` 表示 V 减 V_768：
-
-```text
-层  行类别                  行数   ε=1−maxP  |dQ_true|   |errH1|/|dQ_true|  cos(errH1, K_argmax)
- 0  看 BOS, P>0.99          1270   1.0e-3    3.8e-06     0.33 (shift: 9e-4)  1.00
- 0  看其他键, P>0.99        2708   5.3e-4    1.1e-06     0.64                1.00
- 4  看其他键(act), P>0.99    131   8.2e-4    1.3e-06     1.09                1.00
- 8  看其他键, P>0.99          64   4.7e-3    2.0e-05     0.80                1.00
-11  看其他键, P>0.99          13   4.4e-3    3.1e-05     1.04                1.00
-14  看 BOS, P>0.99          1802   6.0e-3    5.6e-08     0.020 (shift: 5e-4) 1.00
-14  看其他键, P>0.99          83   2.0e-3    8.7e-07     0.58                1.00
-16  看其他键, P>0.99           4   4.7e-3    2.3e-10     1.47                1.00
- *  普通行 (P≤0.9)          3000   ~0.7      0.5–3e-04   1.2e-3 – 2.2e-3     0.80–0.88
-```
-
-- peaked 行的真实 dQ 比普通行小两个数量级（ε 级抵消），H1 误差不随 ε 缩小，所以相对误差达 O(1)。
-- 看 BOS 的行因为 |V_768| 小、O 本身小，H1 只有 2%；shift 能再压 40 倍但没必要。
-- 看其他键的 peaked 行 |O| ≈ |V_argmax| = 13–115，H1 误差 60%–150%，shift（减的是 V_768）对它们完全无效。
-- 普通行上 H1 分量约 1.5e-3 相对、方向沿 argmax 键（cos 0.8）：绝对量占主导，但它是结构性的。
-
-### 2.4 GPU 单层实测：误差与 H1 预测的一致性
-
-`dQ rel` 是全部行的相对 L2 误差，`peaked rel` 只算 maxP>0.99 的行，`cos` 是误差向量与 H1 预测的 cos：
-
-```text
-层   变体           dQ rel    cos(all)   peaked rel   cos(peaked)   H1 预测 peaked rel
- 0   explicit_bf16  2.46e-3   -0.02      2.8e-3       -0.04
-     cudnn_bf16     3.58e-3   +0.49      3.5e-1       +0.83         0.41
-     cudnn_fp16     1.71e-3   -0.02      3.9e-2       +0.13
- 4   explicit_bf16  2.39e-3   +0.02      2.8e-3       -0.04
-     cudnn_bf16     4.12e-3   +0.44      5.3e-2       +0.47         0.16
-     cudnn_fp16     1.77e-3   +0.03      1.6e-2       +0.67
- 8   explicit_bf16  2.88e-3   +0.00      3.5e-3       -0.10
-     cudnn_bf16     4.28e-3   +0.51      5.1e-1       +0.96         0.63
-     cudnn_fp16     1.78e-3   +0.03      7.2e-2       +0.57
-10   explicit_bf16  2.27e-3   -0.01      2.4e-3       +0.56
-     cudnn_bf16     4.36e-3   +0.55      5.0e-1       +0.91         0.53
-     cudnn_fp16     1.76e-3   -0.01      8.3e-2       +0.31
-11   cudnn_bf16     4.70e-3   +0.51      1.21         +0.76         0.86
-     cudnn_fp16     1.77e-3   -0.01      1.3e-1       -0.70
-14   explicit_bf16  2.22e-3   +0.05      2.3e-3       -0.03
-     cudnn_bf16     2.86e-3   +0.12      4.3e-1       +0.63         0.49
-     cudnn_fp16     1.72e-3   +0.04      3.5e-2       +0.13
-16   cudnn_bf16     2.46e-3   +0.39      1.41         +1.00         2.2
-     cudnn_fp16     1.71e-3   +0.01      1.69         -0.92
-18 层几何平均：explicit_bf16 2.33e-3 / cudnn_bf16 3.72e-3 / cudnn_fp16 1.76e-3
-18 层平均 cos(all)：explicit 0.006 / cudnn_bf16 0.418 / cudnn_fp16 0.006
-```
-
-- cuDNN bf16 的误差在**每一层**都与 H1 预测显著同向；显式与 fp16 都不同向。
-- H1 预测的 peaked 行误差量级与实测吻合（0.41/0.35、0.63/0.51、0.53/0.50、0.86/1.21、0.49/0.43、2.2/1.41）。cos 不到 1 是因为 cuDNN 前向的 PV 矩阵乘也把 P 舍入到 bf16，其内部 fp32 O 与本文的 fp32 O 差约 2⁻⁹，两者的舍入决策只是部分相关。
-- 第二个候选机制 H2（dP 先舍入到 bf16 再减 D）的 cos 为 0 或负值，排除。
-- dV 没有 D 项，cuDNN bf16 的 dV 误差（2.5e-3）与显式（1.9e-3）同阶，是对照。
-
-### 2.5 为什么这解释了 8-31 报告的全部现象
-
-- **1.5% 随机乘性噪声无害、cuDNN 差分有害**：H1 分量按行沿 `Σ_j P_ij K_j` 方向、符号随舍入随机，投到参数上落在 `x ⊗ K_argmax` 这类低维子空间，不与 |g| 成比例；sink 机制处于平衡态时这些坐标的真实梯度和二阶矩都很小，Adam 归一化后每步都在那里走满 lr 的随机步，累积到约 1000 步后 sink 结构失稳。乘性噪声在小 |g| 坐标上也小，Adam 看不见。
-- **跨 batch delta 两两 cos ±0.25**：在 30 亿维空间里这不是白噪声，正是"同一子空间、随机符号"。
-- **只融合前 3 层 ≈ 全 18 层**：sink/peaked 结构在早层形成，早层 q/k 投影的系统性扰动穿过后面 15 层。
-- **改 mask、seq_lengths、cuDNN 版本、LR、b1/b2/eps 都没用**：它们都不改变 O 的舍入。
-
-## 3. 缓解手段：fp16 内核 + 动态缩放
-
-`src/openpi/models/gemma.py` 的 `_cudnn_attention_in_dtype`：q/k/v 转 fp16（多 3 位尾数，H1 误差降约 8 倍），custom VJP 里把余切按 2 的幂缩放到 max|dO·s| ∈ [0.25, 0.5) 再送 fp16 反向，fp32 反缩放。真实激活的范围 max|q| ≤ 3.5、|k| ≤ 30、|v| ≤ 84、|dO| 1e-4–1e-2，均在 fp16 范围内。开关：`model.cudnn_attention_dtype: float16`（默认 `bfloat16`，逐位不变）。
-
-单层效果见 2.4：peaked 行误差 0.35 → 0.039（层 0）、0.51 → 0.072（层 8），全局误差低于显式 bf16。
-
-实现有两版：
-
-- **v1（验收 run 用的）**：custom VJP 的反向里用 `jax.vjp` 重新走一遍 `jax.nn.dot_product_attention` 前向再取反向；配合 `nothing_saveable` remat，每层反向比 bf16 路径多一次 attention 前向。实测 8×H100 batch 256 为 1.37 s/step（显式 1.57–1.64，cuDNN bf16 1.05）。
-- **v2（当前代码）**：直接调用 jax 0.5.3 私有的 `jax._src.cudnn.fused_attention_stablehlo._dot_product_attention_fwd_rule / _bwd_rule`（就是 `jax.nn.dot_product_attention(implementation="cudnn")` 自己 custom_vjp 里接的那两个函数），把 softmax stats 和输出当残差保存，反向只跑一次 fused backward。mask→bias（`where(mask, 0, -2<<14)`）、BTNH 布局、静态参数与公共路径一致；全 mask 的 query 行在反向里显式置零 dQ，等价于 bf16 路径的 `_stop_gradient_for_fully_masked_queries`。
-  `scripts/cudnn_fp16_vjp_equivalence.py` 在真实激活（层 0/4/8/14/17）上比较 v2 与 v1：前向逐位相同，dK/dV 逐位相同，dQ 相对差 5e-6–1.2e-5，与同一内核自身的 run-to-run 差（1.6e-6–1.4e-5，原子累加）同阶；全 mask 行 dQ 为 0，无非有限值。步时见第 4 节末。
-
-## 4. 3000 步 strict-order 验收（`configs/diag_cudnn_fp16.yaml`）：通过
-
-与 `diag_cudnn_strict_order` 同数据、seed、batch 顺序、超参，全 18 层 cuDNN fp16（v1 实现），从 pi05_base step 0 起，8×H100，batch 256。2026-09-03 21:00 启动、22:09 结束，3000 步全程有限。显式基线取自 8-31 报告（"显式基线"列是显式 strict-order run；"last3"列是只融合后 3 层、已判定与显式一致的 run，点更全）。
-
-```text
-step    fp16 loss / grad_norm    显式基线           last3 基线         历史 cuDNN bf16
-   0    6.4030 /  65.59          6.4023 / 65.63     6.4029 / 65.52     6.4023 / 65.63
- 200    0.9355 /   3.23          0.9244 /  2.89                        0.9348 /  3.53
- 500    0.5717 /   2.57                             0.5719 /  2.45     0.5719 /  2.55
-1000    0.4231 /   2.21          0.4220 /  2.77     0.4240 /  2.38     0.4319 /  3.16
-1100    0.4071 /   2.30          0.4077 /  1.79     0.4050 /  2.18     0.4723 /  3.37
-1200    0.3759 /   2.36          0.3742 /  2.12     0.3757 /  2.39     0.6600 / 35.49
-1600    0.3209 /   2.77                             0.3238 /  2.65
-2000    0.2792 /   1.93                             0.2793 /  1.83     3.7753 / 150.5
-2500    0.2464 /   1.70          0.2457 /  1.33     0.2463 /  2.01     4.7055 / 759.2
-2900    0.2427 /   1.59                             0.2424 /  1.48     6.2039 / 603.5
-```
-
-- 全部 30 个 log 点有限，grad_norm 最大值 4.75（step 100），之后 1.4–2.8，没有一个点超过基线同点的 1.3 倍（验收上限 3 倍）。
-- step 1000–1500 历史分叉点处 loss 与显式基线差 <0.5%，与 last3 基线差 <0.3%；step 2500/2900 的 loss 0.2464/0.2427 对基线 0.2457–0.2463/0.2424。历史 cuDNN bf16 在同一点是 4.7/6.2。
-- 3000 步全程 param_norm 1802.39 → 1803.19，与 last3 基线的 1803.19 一致。
-- step 1000 的 checkpoint 保存在远端 `checkpoints/diag_cudnn_fp16/diag_cudnn_fp16/1000`，可作 `attention_precision_update_probe.py` 的锚点。
-
-结论：把 cuDNN 内核的计算精度从 bf16 换成 fp16 就消除了全层 cuDNN 的发散，验证了第 2 节的机制。
-
-### 4.1 步时
-
-8×H100，batch 256，`[TIMING]` 行的 total（不含 log 的步）以及相邻 100 步 log 的墙钟：
-
-```text
-路径                                      s/step       来源
-显式 attention                             1.57–1.64    8-31 报告
-cuDNN bf16（发散）                          1.05         8-31 报告（mask 修复时的短跑）
-cuDNN fp16 v1（反向重算前向）                1.34–1.37    本次 3000 步验收 run
-cuDNN fp16 v2（直接调 fwd/bwd 规则）         1.31–1.35    `diag_cudnn_fp16_speed` 400 步（loss 逐点与 v1 一致：step 300 0.7243 对 0.7245）
-cuDNN bf16，今日同 config 复测              1.32–1.35    `diag_cudnn_bf16_speed` 400 步（step 300 loss 0.7250）
-显式 attention，今日同 config 复测          1.60–1.63    `diag_explicit_speed` 400 步（step 300 loss 0.7244）
-```
-
-去掉反向重算只省了约 2%。单卡微基准（`scripts/cudnn_attention_microbench.py`，层 0 真实激活平铺到每卡 batch 32，T=S=1018，8 q 头 / 1 kv 头，含 `nothing_saveable` remat）说明原因：
-
-```text
-变体                fwd      fwd+bwd(remat)   ×18 层
-cudnn_bf16         1.35 ms   5.24 ms          0.094 s/step
-cudnn_fp16_plain   1.53 ms   5.43 ms          0.098 s/step   （fp16 内核、jax 自带 VJP、无缩放）
-cudnn_fp16_v2      1.54 ms   5.53 ms          0.099 s/step
-explicit_bf16      4.46 ms   8.49 ms          0.153 s/step
-```
-
-18 层 attention 的前向加反向合计只占每步约 0.1 s；fp16 内核与 bf16 内核同速，动态缩放和直接调规则的开销可忽略；显式路径也只多 0.06 s。所以 attention 内核本身解释不了 1.33 s 对 1.05 s 的差距，也解释不了显式路径 1.6 s 对 cuDNN 的差距——差距来自整图层面（XLA 在显式路径下的 fusion/remat 决策、显存压力），要用同 config 的 bf16 复测和整步 profile 来定位，见上表最后一行。
-
-### 4.2 更新空间探针：单步指标不区分好坏
-
-`scripts/attention_precision_update_probe.py diag_cudnn_fp16 diag_cudnn_fp16 1000`，锚点是 fp16 验收 run 的 step-1000 train state（参数 + Adam 矩），batch 32，同一真实 batch 和 rng，以 fp32 显式 attention 为真值：
-
-```text
-batch 0            原始梯度 vs fp32：rel / sign_flips     Adam 更新 vs fp32：rel / cos / sign_flips
-explicit_bf16      3.37e-2 / 1.53e-2                      3.13e-2 / 0.99951 / 7.47e-3
-cudnn_bf16         3.18e-2 / 1.43e-2                      2.93e-2 / 0.99957 / 6.97e-3
-cudnn_fp16         3.01e-2 / 1.38e-2                      2.81e-2 / 0.99961 / 6.73e-3
-```
-
-会发散的 cudnn_bf16 与两条稳定路径在这个指标上不可区分（甚至略好于 explicit_bf16）。这与第 2 节一致：D 项误差在范数上只占 ~1e-3，它的危害来自方向的结构性（同一低维子空间、跨 step 累积），单步的 rel/cos/sign_flips 看不见。判定 fused attention 后端好坏要看 2.4 节的 cos(err, H1) 与 peaked 行误差，最终靠 3000 步 strict-order run；不要再用单步梯度/更新距离做验收指标。
-
-## 5. 备选方案（fp16 已通过验收，仅作记录）
-
-1. **前 k 层显式 + 其余 fp16 cuDNN**：8-31 报告的分段实验表明前 3 层贡献了几乎全部更新差异。
-2. **解析修正**：Δ_i 可用一次 fp16 前向估计（`dO·(O_bf16 − O_fp16)`），dQ 修正需一次 V:=K 的前向，dK 修正需一次 dO:=Δ⊙Q 的反向，总成本抵消吞吐收益，不推荐。
-3. 不要再试：减 V_sink 的 shift（2.3 节）、改 mask、改 Adam 参数。
-
-## 6. 相关脚本
-
-- `scripts/attention_dterm_probe.py`：抓真实激活 + 单层五变体误差 + H1/H2 预测；
-- `scripts/attention_dterm_cpu_analysis.py`：peaked 行的 float64 解剖；
-- `scripts/cudnn_fp16_vjp_equivalence.py`：v2（直接调规则）与 v1（反向重算）fp16 VJP 在真实激活上的等价性检查；
-- `scripts/attention_precision_update_probe.py [config] [exp_name] [step]`：真实 Adam 状态上的更新空间对比，锚点用 fp16 验收 run 的 step-1000 checkpoint。
diff --git a/docs/2026-09-04-h100-speed-ab.md b/docs/2026-09-04-h100-speed-ab.md
deleted file mode 100644
index 80c3ddf..0000000
--- a/docs/2026-09-04-h100-speed-ab.md
+++ /dev/null
@@ -1,79 +0,0 @@
-# 2026-09-04 H100 训练提速 A/B（earbuds 0904 配置）
-
-**机器**：ebcloud 8×H100 80GB，hb-openpi 环境（jax 0.5.3，cuDNN 9.19）
-**配置**：`configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml`
-（pi05、cuDNN fp16 attention、batch_image_views、全局 batch 256、64 workers、fsdp 8）
-**数据**：`Xense/earbud_case_insertion_teleop_0515_left_right`，776 集 / 1,145,613 帧
-
-> 注：实验当天远端机器的 localstorage 随后被清空，本文与相关代码是事后在本机重建并提交的；
-> 数字来自当时的日志摘录。
-
-## 口径
-
-每个变体从 pi05_base 权重起跑 400 步，`strict_batch_order: true`，关 checkpoint 和 W&B，
-其余与生产配置相同（短跑用一份复制自生产 yaml 的临时 config，CLI 覆盖单个变量）。
-主指标是日志里 `Step 100` 到 `Step 300` 两条时间戳之差除以 200，即含日志开销的端到端墙钟 s/step。
-step 300 的 loss 用来确认数值没变。
-
-## 结果
-
-| 变体 | s/step | 相对基线 | loss@300 | 结论 |
-| --- | --- | --- | --- | --- |
-| E0 基线（原配置） | 1.328 | — | 0.7356 | 等数据 0.03 s，数据侧不是瓶颈 |
-| E1 `max_token_len` 200→128 | 1.246 | **−6.2%** | 0.7356 | 采纳 |
-| E2 关闭 remat（gemma+siglip） | OOM | — | — | XLA 试图分配 280 GiB，不可行 |
-| E2b gemma remat `dots_with_no_batch_dims_saveable` | 编译失败 | — | — | cuDNN 自定义分片器报 q/k/v 分片不一致 |
-| E2c 同上 + q/k/v 显式分片约束 | OOM | — | — | 编译过了，但单块残差 84.6 GB，仍放不下 |
-| E3 XLA 集合通信标志 | 1.281 | **−3.5%** | 0.7357 | 采纳 |
-| E4 `fsdp_devices` 8→4 | 1.335 | +0.5% | 0.7354 | 无收益，放弃 |
-| **E5 = E1 + E3** | **1.193** | **−10.2%** | 0.7356 | 生产采用；吞吐 192.8→214.6 samples/s |
-
-60000 步按稳态估算：22.1 h → 19.9 h，省约 2.2 h。
-
-## 各项说明
-
-### E1：token 长度按数据集实测收紧
-
-`scripts/token_len_scan.py` 走真实 transform 链（repack → BiFlexivInputs → Normalize →
-TokenizePrompt），对 4037 个状态（含每一维的最小/最大值样本）做 tokenize：min 110、p50 113、p99 115、
-**max 116**。原来 200 里有 84 个是纯 padding，但 attention/MLP 照算。128 留 12 个余量。
-换数据集或改 prompt 时必须重扫；tokenizer 超长会 warning 并截断，不会报错。
-
-```bash
-JAX_PLATFORMS=cpu uv run scripts/token_len_scan.py \
-  --config-name pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100
-```
-
-### E3：XLA 标志
-
-```
---xla_gpu_enable_latency_hiding_scheduler=true
---xla_gpu_all_gather_combine_threshold_bytes=1073741824
---xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824
---xla_gpu_all_reduce_combine_threshold_bytes=1073741824
---xla_gpu_enable_pipelined_all_gather=true
---xla_gpu_enable_pipelined_reduce_scatter=true
---xla_gpu_enable_while_loop_double_buffering=true
-```
-
-FSDP 每步的参数 all-gather / 梯度 reduce-scatter 合并成大块并与计算流水。没有逐项拆分测哪一个起作用。
-通过 `XLA_FLAGS` 环境变量传入，已写在 0904 yaml 头部注释里。
-
-### E2：remat 在 80 GB 上没有空间
-
-`gemma.py` / `siglip.py` 每层用 `nothing_saveable` 全量重算。本次给 `Pi0Config` 加了
-`gemma_remat_policy` / `siglip_remat_policy`（默认不变），并在 cuDNN 调用前给 q/k/v/mask 加了统一的
-batch 分片约束（`gemma.py`，否则非默认 remat 策略下编译失败）。但 batch 256（每卡 32）时：
-完全不 remat 要 280 GiB 单块，只保存无 batch 维 matmul 也要 84.6 GB，都放不下。
-在 H200 141 GB 或每卡 batch 更小时这两个开关才可能兑现，理论上限约为一次前向的重算（15–25%）。
-
-### 其他观察
-
-- 数据侧：等数据 0.03–0.04 s/step，64 workers 足够；worker 启动一次约 330 s（spawn 上下文，64 进程各自
-  import jax/torch/lerobot 并反序列化 1.1M 行数据集），是每次启动的一次性开销。
-- `ResizeImages` 实际调用的是 `xense_client.image_tools`（PIL），worker 里不跑 JAX resize。
-- fsdp 4 相对 8 没有收益：参数 all-gather 的量只差 12.5%，而每卡持有的优化器状态翻倍。
-- 短跑结束后偶有 DataLoader worker 成为孤儿进程（ppid 1，每个约 3 GB RSS），一次残留了 39 个；
-  跑完批量实验后用 `ps -eo pid,ppid,cmd | grep spawn_main` 检查一下。
-- 该配置 `num_train_steps: 60000` 但 `lr_schedule.decay_steps` 仍是默认 30000，后半程一直在底线 LR；
-  这是训练质量问题，本次未动。
diff --git a/docs/data-pipeline-stalls.md b/docs/data-pipeline-stalls.md
deleted file mode 100644
index e9d5f1f..0000000
--- a/docs/data-pipeline-stalls.md
+++ /dev/null
@@ -1,326 +0,0 @@
-# 数据管线优化：跳过触觉解码 + 消除加载器停顿
-
-**日期**：2026-08-19 实现与验证，2026-08-25 生产运行复核
-**机器**：ebcloud H100（8×H100 80GB / 192 逻辑核 / 2 TB 内存）
-**起因**：`parquet-vs-mp4` 实验记录 §8 里未执行的前两项建议
-
----
-
-## 0. 结论速览
-
-
-| 改动                                    | 目标                                | 实测结果                                                   |
-| --------------------------------------- | ----------------------------------- | ---------------------------------------------------------- |
-| **A. 跳过 4 路未用的触觉解码**          | §8 第一项，预估「约 1.8×」        | 数据侧**1.71×**，端到端 **0%**——预估的前提不成立        |
-| **B1. 批次乱序交付** (`in_order=False`) | 每`num_workers` 步一次的约 6 s 停顿 | 受控 A/B：停顿**20 次 → 0 次**，最大等待 7.83 s → 0.43 s |
-| **B2. 无限 sampler**                    | epoch 边界重建迭代器                | 边界停顿**63.0 s → 0 s**                                  |
-| **B 合计（受控 400 步 A/B）**           | —                                  | 770.8 s →**715.7 s = 1.08×**                             |
-| **B 合计（65,000 步生产运行）**         | —                                  | 数据停顿损失**6.84 h → 131 s**，约 **1.14×**             |
-| 剩余瓶颈                                | —                                  | 加载器吞吐天花板约 1.84 s/batch，**不是排序问题**          |
-
-一句话：**A 是对的优化但在H100上兑现不了；B 才是真正拿到的收益。**
-
----
-
-## 1. 改动 A：`tactile` 开关
-
-### 1.1 问题
-
-`LeRobotDataset.__getitem__` 无条件遍历 `meta.video_keys` 解码**全部** 7 路视频流，
-而 `LeRobotBiFlexivDataConfig` 的 repack 只取 `head` / `left_wrist` / `right_wrist`：
-
-```
-observation.images.head / left_wrist / right_wrist    640×480   ← 用
-observation.images.{left,right}_tactile_{left,right}  400×700   ← 解完直接丢
-```
-
-触觉那四路的后缀取决于录制时间：lerobot-xense `1146d034` 之前是 USB 枚举序
-（`_tactile_{0,1}`），之后是指位（`_tactile_{left,right}`）。两种数据集都在盘上，
-所以 `TACTILE_KEY_MARKER` 按子串匹配 `tactile`，两种都吃得下。
-
-H.264 随机寻址解码只有顺序解码的 1/200，所以这 4 路是纯浪费。
-（Xense 的 bi_flexiv 数据集全都是这个结构：optical-module-insertion、pick_up_cube、
-unscrew-the-bottle-cap、bottle-sorting、stack-cubes 均为 3 用 + 4 触觉。）
-
-### 1.2 实现
-
-- `DataConfig.tactile: bool = False`
-- `SelectiveVideoLeRobotDataset` 覆盖 `_get_query_timestamps`，把非白名单的流从查询里滤掉，
-  `_query_videos` 就永远不会被要求解它们
-- 启动时校验：若 repack 引用了被禁用的流，直接 `ValueError`，不会静默拿到黑图
-
-**没有改 lerobot，没有转换任何数据。**
-
-```yaml
-data:
-  type: LeRobotBiFlexivDataConfig
-  repo_id: Xense/<dataset>
-  base_config:
-    prompt_from_task: true
-    tactile: false        # 省略即为 false；用触觉的模型（Pi05Tactile 等）才设 true
-```
-
-### 1.3 结果
-
-
-| 指标                            | tactile=True | tactile=False       | 倍数       |
-| ------------------------------- | ------------ | ------------------- | ---------- |
-| 单样本`__getitem__`（H100）     | 262.4 ms     | **153.3 ms**        | **1.71×** |
-| 单样本`__getitem__`（RTX 5090） | 192.1 ms     | 103.2 ms            | 1.86×     |
-| 全 transform 链输出             | —           | 20 样本**逐位相同** | —         |
-
-省下的 109.1 ms/样本，与 §8 预估的「288 ms 解码里的 110 ms」几乎完全吻合。
-
-**但 400 步受控 A/B 显示端到端毫无变化：**
-
-
-| 指标               | tactile=True | tactile=False |
-| ------------------ | ------------ | ------------- |
-| 400 步总时长       | 768.4 s      | 770.8 s       |
-| GPU dispatch 合计  | 615.0 s      | 617.6 s       |
-| 等数据合计         | 139.6 s      | 139.5 s       |
-| 停顿步数（>0.5 s） | 20           | 20            |
-| 启动到首个 batch   | 413.7 s      | **376.4 s**   |
-
-唯一实测到的收益是启动快 37.3 s（worker spawn 301 s 两边一样，差的是「单个 worker
-独自造完一个 256 batch」的时间）。
-
-### 1.4 为什么 1.71× 兑现成 0
-
-§8 那句「约 1.8×」是 h200 箱 **batch 512** 的推算，当时 mp4 加载相对 GPU 的余量只有 11%。
-这台机 batch 256 + 192 核，加载器本来就跟得上 GPU，多出来的余量 GPU 用不上。
-
-更关键的是：第 2 节的诊断显示，剩下的等待是加载器的**吞吐天花板**，
-而这个天花板与解码量无关——这正是 tactile 开关一秒都省不下来的直接原因。
-
-### 1.5 保留它的理由
-
-逐位无损、零风险、每样本少烧约 40% 的 CPU、每个 epoch 边界省约 37 s。
-换更大 batch 或核更少的机器时它才会真正兑现。**但别指望它让当前配置变快。**
-
----
-
-## 2. 改动 B：消除Dataloader停顿
-
-### 2.1 诊断装置
-
-为了不必每次花 20 分钟跑训练，做了一个复现装置：真数据 + 真 transform 链 + 真 torch
-DataLoader，用 `sleep(1.55)` 替代 GPU 步，同时抓 torch 内部迭代器的两个私有状态：
-
-```
-outstanding = _send_idx - _rcvd_idx     已派给 worker 但还没回来的任务数
-buffered    = _task_info 里已完成的结果  算完了、但还轮不到它交付的 batch 数
-```
-
-**这两个数把「算不出来」和「算出来了但过不去」区分开——是整件事里最值钱的一步。**
-
-### 2.2 发现一：队头阻塞（head-of-line blocking）
-
-
-| batch | wait   | outstanding | buffered |
-| ----- | ------ | ----------- | -------- |
-| 64    | 3.02 s | 1024        | 10       |
-| 66    | 9.84 s | 1024        | **41**   |
-| 128   | 3.16 s | 1024        | 10       |
-| 130   | 9.66 s | 1024        | **41**   |
-
-`outstanding` **恒为 1024**（= `prefetch_factor 16 × 64 workers`，任务管线一直是满的），
-而停顿发生时 `buffered` 高达 41——**41 个 batch 已经算完躺在那儿，训练循环却在等**。
-
-原因：torch 严格按 sampler 顺序交付。任务 *j* 派给 worker `j mod 64`，主进程必须先拿到
-*j* 才能拿 *j+1*。这一轮里最慢的那个 worker，堵住排在它后面的所有 batch。每 64 步绕一圈，
-就必然撞上一次最慢的那个。`prefetch_factor` 调多深都没用——问题不是没算，是不让过。
-
-**修复**：`DataLoader(in_order=False)`（PyTorch ≥2.6 为这个场景加的官方参数）。
-谁先算完谁先交付。这里排序毫无意义——sampler 本来就是随机打散的。
-
-正确性已验证：`in_order=False` 下一个 pass 仍然**每个 index 恰好出现一次**，不重不漏，
-只是交付顺序变了。
-
-### 2.3 发现二：epoch 边界重建迭代器
-
-`TorchDataLoader.__iter__` 在 sampler 耗尽时 `break` 出去重新 `iter()`。预取管线被整个
-拆掉重建，训练循环要**等一个 worker 独自造完一整个 batch**。把 epoch 缩短到 80 个 batch
-强制触发，实测：
-
-```
-EPOCH BOUNDARY at batch 80: rebuilding iterator
-STALL batch 80: wait=62.96s outstanding=79 buffered=60
-```
-
-**63.0 秒** = 256 样本 × 246 ms/样本。
-
-**修复**：`InfiniteSampler`——永不停止、每轮重新打散，`StopIteration` 再也不会触发。
-代价为零：openpi 本来就不 checkpoint 数据加载器状态（`restore_state` 直接
-`del data_loader`），batch 序列从来就不可恢复。
-
-### 2.4 诊断装置上的结果
-
-150 个 batch，模拟 1.55 s GPU 步：
-
-
-| 变体                 | 总时长      | 冷启动 | 等数据 | 停顿数 | 最大等待   |
-| -------------------- | ----------- | ------ | ------ | ------ | ---------- |
-| 基线                 | 349.5 s     | 66.0 s | 49.9 s | 14     | 9.8 s      |
-| `in_order=False`     | 336.3 s     | 57.3 s | 45.7 s | **0**  | **0.42 s** |
-| 基线（epoch=80）     | 376.8 s     | 67.1 s | 76.2 s | 6      | **63.0 s** |
-| 两项都修（epoch=80） | **330.0 s** | 53.6 s | 43.1 s | **0**  | **0.39 s** |
-
-### 2.5 受控真实训练 A/B（400 步，同机同配置）
-
-
-| 指标                   | 修复前  | 修复后      | 倍数       |
-| ---------------------- | ------- | ----------- | ---------- |
-| **400 步总时长**       | 770.8 s | **715.7 s** | **1.08×** |
-| GPU dispatch 合计      | 617.6 s | 578.4 s     | 1.07×     |
-| 等数据合计             | 139.5 s | 122.7 s     | 1.14×     |
-| **停顿步数（>0.5 s）** | **20**  | **0**       | —         |
-| 最大`next_batch`       | 7.83 s  | **0.43 s**  | 18×       |
-| 中位`next_batch`       | 0.030 s | 0.290 s     | —         |
-
-修复前的 20 次停顿：步 63/64/69、127/128/133、191/192/197、255/256/261、319/320/325、
-383/384/389——每 64 步一簇。**修复后一次都没有。**
-
-> 中位 `next_batch` 从 0.030 s 涨到 0.290 s **不是退步**：同样的产能缺口，以前攒成每 64 步
-> 一次的 6 秒尖峰，现在均摊到每一步。总等待反而少了 16.8 s。
-
----
-
-## 3. 生产运行复核（2026-08-25）
-
-**配置**：`pi05_base_bi_flexiv_unscrew_bottle_cap_0824_h100`
-（`Xense/unscrew-the-bottle-cap-0821`，100 ep / 84,419 帧，batch 256，workers 64，fsdp 8）
-**运行**：64,951 步 / 30.60 h
-
-日志确认补丁生效：`skipping video decode for 4 of 7 stream(s)`，且调用栈行号是
-`data_loader.py:352`（旧代码 `:292`）。
-
-### 3.1 与另一次运行的对照（注意：不是受控 A/B）
-
-对照的 `优化前.log` 是**另一个任务、另一个数据集**
-（`earbuds_case_insertion_teleop_rtc_0715`），绝对时长不可比，只有归一化后的停顿指标有意义。
-
-
-|                                      | 优化前（earbuds，99,951 步 / 51.46 h） | 优化后（unscrew，64,951 步 / 30.60 h） |
-| ------------------------------------ | -------------------------------------- | -------------------------------------- |
-| **平均每步**                         | 1.853 s                                | **1.696 s**                            |
-| **数据停顿次数**（`next_batch`>3 s） | **3,345**                              | **36**                                 |
-| **数据停顿总损失**                   | **6.84 h**（占 13.3%）                 | **131 s**（占 0.12%）                  |
-| 最大`next_batch`                     | 10.32 s                                | 5.38 s                                 |
-
-归一化到每万步：停顿次数 **334.7 → 5.5**（少 61 倍），损失 **2464 s → 20 s**（少 123 倍）。
-
-### 3.2 这次运行省了多少
-
-数据集只有 84,419 帧，**一个 epoch 才 329 步**——65,000 步里有 **197 次 epoch 边界**。
-按修复前每次约 63 s 算，光边界就该吃掉 **3.45 h**。
-
-日志验证：落在 epoch 边界附近（±2 步）的数据停顿 **0 次**（无限 sampler 生效）；
-落在 `num_workers` 倍数上的也没有（乱序交付生效）。
-
-按优化前的停顿比率外推，这次本该损失约 **4.45 h**，实际只损失 **0.04 h**：
-
-
-|                  |               |
-| ---------------- | ------------- |
-| 实际用时         | 30.60 h       |
-| 不优化的估计用时 | 约 35.0 h     |
-| **提速**         | **约 1.14×** |
-
-### 3.3 残留的 36 次停顿：全是 checkpoint，不是加载器
-
-那 36 次 `next_batch > 3 s` 的步号是 1001、2001、3001…——**36/36 全部落在 checkpoint
-保存的下一步**。`save_interval=1000`，65,000 步存了 64 次，每次阻塞 7–11 s（合计 389 s），
-异步落盘线程又和 64 个 worker 抢 CPU，导致下一步取数慢 3–5 s。
-
-合计约 **520 s（0.14 h）**，很好摘：
-
-```yaml
-save_interval: 5000      # 从 1000 提到 5000，checkpoint 开销降到 1/5
-keep_period: 10000
-```
-
----
-
-## 4. 剩下的瓶颈：吞吐天花板，不是排序
-
-修完之后 `buffered` **恒为 0**——结果一到就被取走，中位等待 0.21–0.29 s，峰值 0.36 s。
-这说明加载器已经在满负荷跑，稳态产能约 **1.84 s/batch**，而 GPU 每 1.55 s 就要一个。
-残留的约 14% 等待就是这个差额。
-
-折算：64 workers 每秒产出 139 样本 = **每 worker 460 ms/样本**，而单进程实测只要 153 ms
-——并发到 64 个 worker 时每样本成本涨了 3 倍，在 192 核的机器上并没有线性扩展。
-
-**这个天花板不在解码上**：直接证据是 tactile 开/关对等待时间毫无影响。
-候选嫌疑（未验证，按可疑度排序）：
-
-1. 每个 worker 里跑 `jax.image.resize`（`resize_with_pad` 是 jax 实现，64 个 worker 进程各自初始化 JAX）
-2. `_collate_fn` 逐样本 `np.stack`
-3. numpy 数组经 multiprocessing queue 的 pickle 传输（每 batch 约 115 MB uint8）
-
-**下一步该测什么**：在 64 路并发下给 worker 内的 `__getitem__` 分段计时
-（解码 / transform / collate 各占多少），再决定优化哪一段。在此之前别再猜。
-
----
-
-## 5. 代码位置与开关
-
-改动全部在远程 `~/localstorage/xense-openpi`（**本机仓库保持未改**）：
-
-
-| 文件                                                          | 改动                                                                                             |
-| ------------------------------------------------------------- | ------------------------------------------------------------------------------------------------ |
-| `src/openpi/training/config.py`                               | `DataConfig.tactile`、`TrainConfig.strict_batch_order`                                           |
-| `src/openpi/training/data_loader.py`                          | `SelectiveVideoLeRobotDataset`、`_resolve_decode_video_keys`、`InfiniteSampler`、`in_order` 接线 |
-| `src/openpi/training/data_loader_test.py`                     | 10 个新测试（`pytest src/openpi/training/` 共 **68 passed**）                                    |
-| `configs/README.md`、`configs/_examples/_FULL_REFERENCE.yaml` | `tactile` 文档                                                                                   |
-
-```yaml
-data:
-  base_config:
-    tactile: false          # 默认；true = 解码触觉流
-
-strict_batch_order: false   # 默认；true = 恢复严格顺序交付
-```
-
-**两个开关的默认值就是优化后的行为**，已有 yaml 不改也在用；写出来只是让配置自解释。
-
-**`strict_batch_order` 的取舍**：默认 `false` 意味着 batch 的**交付顺序**不再逐次可复现。
-每个 index 每轮仍恰好出现一次，样本内容不受影响，但两次同 seed 的训练看到的 batch 顺序会不同。
-需要逐 batch 复现时设 `true`，代价是队头阻塞回来。
-（注意：因为 openpi 不保存数据加载器状态，**断点续训本来就已经打乱了顺序**。）
-
----
-
-## 6. 复现
-
-诊断装置与基准脚本在远程 `/root/tac_bench/`：
-
-```bash
-conda activate lerobot-xense && cd ~/localstorage/xense-openpi
-export PYTHONPATH=$PWD/src
-
-# 触觉解码 A/B + 逐位一致性
-JAX_PLATFORMS=cpu python /root/tac_bench/bench_decode.py --n 60
-
-# 停顿诊断（--in-order / --infinite / --epoch-batches 组合出 4 个变体）
-JAX_PLATFORMS=cpu python /root/tac_bench/diag_rotation.py \
-    --batches 150 --epoch-batches 80 --in-order false --infinite --label fixed
-
-# in_order=False 的覆盖性验证
-python /root/tac_bench/check_in_order.py
-```
-
----
-
-## 7. 给后来者的四条
-
-1. **先看停顿的周期再动手。** 周期 = `num_workers` 是队头阻塞，周期 = 一个 epoch 的步数是
-   边界重建，周期 = `save_interval` 是 checkpoint——都不是「解码太慢」。看错了会花一天优化
-   解码然后发现端到端纹丝不动。
-2. **`buffered` 这个数值千金难买。** 停顿时它是 0 说明真的算不出来，是 41 说明算出来了
-   过不去——两种情况的修法完全相反。
-3. **别把别的机器上的预估当结论。** §8 的「1.8×」在 h200 箱 batch 512 下是对的，
-   搬到 batch 256 + 192 核就一点不剩。
-4. **别拿两个不同任务的日志比绝对时长。** 要比就比归一化后的指标（每万步的停顿次数与损失），
-   或者干脆跑受控 A/B。
diff --git a/docs/training-optimization.md b/docs/training-optimization.md
new file mode 100644
index 0000000..5bb5101
--- /dev/null
+++ b/docs/training-optimization.md
@@ -0,0 +1,797 @@
+# openpi 训练提速与 cuDNN Attention 数值事故：完整记录
+
+**覆盖时间**：2026-08-19 → 2026-09-04
+**分支**：`optimize-training-duration`（性能主线 commit `ec72f34` → `cf07bbe`；其中
+`ec72f34` 已进入本分支基线 `main`，分支新增优化从 `9aef52d` 开始）
+**机器**：8×H200、8×H100 80GB（ebcloud）
+**模型**：Pi0.5（`pi05_base` 微调，BF16 GQA，18 层 Gemma）
+
+本文按时间线合并了原先分散的 5 份记录（数据管线、H200 优化、cuDNN NaN 事故、D 项根因分析、
+H100 A/B），并与 commit 对齐。每个阶段给出：动机 → 实现（关键代码）→ 实测 → 结论。
+被推翻的中间结论保留在「弯路」小节里，避免后来人重走。
+
+---
+
+## 0. 总账
+
+| 阶段 | 日期 | commit | 改动 | 端到端收益 |
+| --- | --- | --- | --- | --- |
+| 一 | 08-19 ~ 08-27 | `ec72f34` | 跳过触觉解码 / `in_order=False` / 无限 sampler | 消除周期性加载器停顿，生产 **≈1.14×** |
+| 二 | 08-27 ~ 08-28 | `9aef52d` `8ccf99c` | 共享内存 Tensor / XProf / 只在日志步算范数 / 三路相机合批 | step 时间 **−15.3%** |
+| 三 | 08-29 ~ 08-31 | `476babe`…`895811c` | cuDNN fused attention（收益 −29%，但**训练发散**） | 净收益 0，转为诊断 |
+| 四 | 09-03 ~ 09-04 | `df9755f` `3ea2b1c` | 定位根因 + fp16 内核（3000 步验收通过） | 显式 1.60 → fp16 **1.33 s/step** |
+| 五 | 09-04 | `cf07bbe` | `max_token_len` 200→128、XLA 集合通信标志 | 再 **−10.2%**，1.328 → 1.193 s/step |
+
+**当前生产配置**（`configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml`）：
+
+```yaml
+model:
+  max_token_len: 128          # 阶段五，按数据集实测收紧
+  batch_image_views: true     # 阶段二
+  use_cudnn_attention: true   # 阶段三
+  cudnn_attention_dtype: float16   # 阶段四：必须 float16，bfloat16 会发散
+data:
+  base_config:
+    tactile: false            # 阶段一
+strict_batch_order: false     # 阶段一
+```
+
+配合 `XLA_FLAGS`（见 §5.2）。上表的测试跨越不同机器、batch 和任务，**不能把各阶段百分比直接连乘成
+一个受控的总提速数字**；可复现的最终同配置结论是阶段五的 1.328 → 1.193 s/step。
+
+性能主线的 commit 对应关系如下。`3b2795e`、`3c5aac6` 只是夹在分支中的 Nova5 配置改动，
+与本文提速结论无关，故不纳入阶段收益。
+
+| 日期 | commit | 作用 |
+| --- | --- | --- |
+| 08-27 | `ec72f34` | 合入跳过触觉解码、乱序交付与无限 sampler；同时把 PyTorch 下限提高到支持 `in_order` 的 2.6 |
+| 08-27 | `9aef52d` | 共享内存 Tensor、零拷贝 NumPy 视图、显式关闭 DataLoader 与数据管线 profiling |
+| 08-28 | `8ccf99c` | XProf 开关、只在日志步计算范数、三路相机合批及其测试 |
+| 08-28 | `d0aec27` | 更新 H100 RTC 训练配置；没有单独的性能 A/B |
+| 08-29 | `476babe` `f42992e` | 引入 cuDNN fused attention，并记录最初的短跑/生产验证 |
+| 08-29 | `c1b70b2` `6305750` `51d9af3` | 处理全空 mask NaN；撤回逐步有限性保护；用 stop-gradient 保留快速原始 mask |
+| 08-31 | `ad08973` `c9e77cc` `895811c` | 建立 strict-order/分层诊断，定位长期发散，并第一次合并事故文档 |
+| 09-03 | `df9755f` | 增加 attention 精度与 Adam 更新空间诊断 |
+| 09-04 | `3ea2b1c` | 定位 D 项舍入机制，实现并验证 fp16 custom VJP |
+| 09-04 | `cf07bbe` | token 长度、XLA 集合通信和 remat A/B，形成 H100 最终配置 |
+
+---
+
+## 1. 阶段一：数据管线（2026-08-19 实现，08-25 生产复核，commit `ec72f34`）
+
+起点是「训练每隔几十步卡 6 秒」。三个改动，其中只有两个真正兑现。
+
+### 1.1 改动 A：跳过未使用的触觉视频解码（收益 0）
+
+`LeRobotDataset.__getitem__` 无条件解码 `meta.video_keys` 里的**全部** 7 路视频，而 repack 只取
+head / left_wrist / right_wrist——4 路触觉流解完直接丢。H.264 随机寻址解码只有顺序解码的 1/200，纯浪费。
+
+实现是覆盖 `_get_query_timestamps`，让 `_query_videos` 永远不被要求解它们
+（`src/openpi/training/data_loader.py:203`）：
+
+```python
+class SelectiveVideoLeRobotDataset(lerobot_dataset.LeRobotDataset):
+    decode_video_keys: frozenset[str] | None = None   # 必须是普通属性：dataset 要 pickle 给 spawn 的 worker
+
+    @override
+    def _get_query_timestamps(self, current_ts, query_indices=None):
+        query_timestamps = super()._get_query_timestamps(current_ts, query_indices)
+        if self.decode_video_keys is None:
+            return query_timestamps        # None = 不过滤，忘了设也退化成原行为
+        return {key: ts for key, ts in query_timestamps.items() if key in self.decode_video_keys}
+```
+
+白名单由 `_resolve_decode_video_keys` 按子串 `tactile` 匹配（两代命名 `_tactile_{0,1}` 和
+`_tactile_{left,right}` 都能吃下），并且**启动时校验**：如果 repack 引用了被禁用的流直接 `ValueError`，
+不会静默拿到黑图。没有改 lerobot，没有转换数据。
+
+**结果**：单样本 `__getitem__` 262.4 ms → 153.3 ms（**1.71×**），全 transform 链输出 20 个样本**逐位相同**。
+**但 400 步端到端 A/B 完全没变**（768.4 s vs 770.8 s，停顿 20 次 vs 20 次）。唯一实测收益是启动快 37.3 s。
+
+原因：这台机 batch 256 + 192 核，加载器本来就跟得上 GPU；省下的 CPU 余量 GPU 用不上。更关键的是
+§1.4 的天花板与解码量无关。**保留它**（零风险、每样本少烧 40% CPU、epoch 边界省 37 s），
+但别指望它让当前配置变快。
+
+### 1.2 改动 B1：队头阻塞 → `in_order=False`
+
+诊断装置是关键：真数据 + 真 transform + 真 DataLoader，用 `sleep(1.55)` 替代 GPU 步，同时抓 torch
+迭代器的两个私有状态：
+
+```
+outstanding = _send_idx - _rcvd_idx     已派给 worker 但还没回来的任务数
+buffered    = _task_info 里已完成的结果  算完了、但还轮不到它交付的 batch 数
+```
+
+**这两个数把「算不出来」和「算出来了但过不去」区分开——是整件事里最值钱的一步。**
+
+| batch | wait | outstanding | buffered |
+| --- | --- | --- | --- |
+| 64 | 3.02 s | 1024 | 10 |
+| 66 | 9.84 s | 1024 | **41** |
+| 130 | 9.66 s | 1024 | **41** |
+
+`outstanding` 恒为 1024（`prefetch_factor 16 × 64 workers`，管线一直是满的），停顿时 `buffered` 高达 41
+——**41 个 batch 已经算完躺在那儿，训练循环却在等**。torch 严格按 sampler 顺序交付，任务 *j* 派给
+worker `j mod 64`，这一轮最慢的那个 worker 堵住排在它后面的所有 batch，每 64 步撞一次。
+调深 `prefetch_factor` 没用——问题不是没算，是不让过。
+
+修复：`DataLoader(in_order=False)`（PyTorch ≥2.6 为这个场景加的官方参数），谁先算完谁先交付。
+sampler 本来就是随机打散的，排序毫无意义。已验证一个 pass 内每个 index 仍恰好出现一次，不重不漏。
+
+### 1.3 改动 B2：epoch 边界重建迭代器 → `InfiniteSampler`
+
+`TorchDataLoader.__iter__` 在 sampler 耗尽时 `break` 出去重新 `iter()`，预取管线被整个拆掉重建，
+训练循环要等**一个 worker 独自造完一整个 batch**：
+
+```
+STALL batch 80: wait=62.96s outstanding=79 buffered=60
+```
+
+**63.0 秒** = 256 样本 × 246 ms/样本。修复（`data_loader.py:529`）：
+
+```python
+class InfiniteSampler(torch.utils.data.Sampler[int]):
+    """永不停止、每轮重新打散，StopIteration 再也不会触发。"""
+    def __iter__(self):
+        epoch = 0
+        while True:
+            if self._shuffle:
+                generator = torch.Generator().manual_seed(self._seed + epoch)
+                yield from torch.randperm(self._num_samples, generator=generator).tolist()
+            else:
+                yield from range(self._num_samples)
+            epoch += 1
+```
+
+代价为零：openpi 本来就不 checkpoint 数据加载器状态（`restore_state` 直接 `del data_loader`），
+batch 序列从来就不可恢复。
+
+### 1.4 结果与残留
+
+受控 400 步真实训练 A/B（同机同配置，唯一变量是 B1+B2）：
+
+| 指标 | 修复前 | 修复后 | 倍数 |
+| --- | --- | --- | --- |
+| 400 步总时长 | 770.8 s | **715.7 s** | **1.08×** |
+| 停顿步数（>0.5 s） | **20** | **0** | — |
+| 最大 `next_batch` | 7.83 s | **0.43 s** | 18× |
+| 中位 `next_batch` | 0.030 s | 0.290 s | — |
+
+> 中位数从 0.030 涨到 0.290 s **不是退步**：同样的产能缺口，以前攒成每 64 步一次的 6 秒尖峰，
+> 现在均摊到每一步，总等待反而少 16.8 s。
+
+**生产复核**（08-25，`unscrew-the-bottle-cap-0821`，64,951 步 / 30.60 h）：数据停顿从归一化后的
+每万步 334.7 次 / 2464 s 降到 **5.5 次 / 20 s**。该数据集一个 epoch 才 329 步，65,000 步里有 197 次
+epoch 边界，按修复前每次 63 s 算光边界就该吃掉 3.45 h。实际停顿损失 0.04 h，**约 1.14×**。
+
+残留的 36 次 `next_batch > 3 s` **全部落在 checkpoint 保存的下一步**（步号 1001、2001、3001…）。
+`save_interval` 从 1000 提到 5000 即可，已写进后续所有生产 yaml。
+
+修完之后 `buffered` **恒为 0**，加载器满负荷跑，稳态产能约 1.84 s/batch。折算每 worker 460 ms/样本，
+而单进程只要 153 ms——并发到 64 个 worker 时每样本成本涨了 3 倍。这个天花板不在解码上
+（直接证据：tactile 开/关对等待时间毫无影响）。阶段五复测显示等数据只剩 0.03–0.04 s/step，
+数据侧已不是瓶颈，这条线到此为止。
+
+### 1.5 开关
+
+```yaml
+data:
+  base_config:
+    tactile: false          # 默认；true = 解码触觉流（Pi05Tactile 等模型才需要）
+strict_batch_order: false   # 默认；true = 恢复严格顺序交付，队头阻塞会回来
+```
+
+两个开关的默认值就是优化后的行为。`strict_batch_order: false` 意味着 batch 的**交付顺序**不再逐次可复现
+（每个 index 每轮仍恰好出现一次，样本内容不变）。诊断用的单变量 A/B 必须设 `true`——阶段三/四的所有
+3000 步验收都是这么跑的。
+
+---
+
+## 2. 阶段二：主机传输与计算侧（2026-08-27 ~ 08-28，commit `9aef52d` `8ccf99c`）
+
+**测量口径**（后续所有短跑沿用）：从同一 checkpoint 只读恢复，独立实验目录，关 W&B 和所有 checkpoint 保存；
+跑到起始 step + 40，丢弃前 10 步的编译/预热，统计后 30 个稳态样本。JAX dispatch 是异步的，
+所以主指标是三者之和，不能单看 `dispatch`：
+
+```
+measured_step = dispatch + compute_sync + next_batch_total
+```
+
+### 2.1 共享内存 Tensor（commit `9aef52d`，−13.6%）
+
+worker 原先返回 NumPy batch，过 multiprocessing queue 要 pickle/复制约 900 MiB，主进程再包装成 JAX array。
+collate 改为返回 **CPU Torch Tensor**——PyTorch 的 multiprocessing reducer 会通过共享内存传 storage：
+
+```python
+def _collate_fn(items, *, profile=False):
+    """Collate batch elements into CPU tensors so workers use shared-memory IPC."""
+    batch = jax.tree.map(
+        lambda *xs: torch.from_numpy(np.stack([np.asarray(x) for x in xs], axis=0)),
+        *items,
+    )
+```
+
+主进程再用 `tensor.numpy()` 拿同一块 storage 的**零拷贝视图**交给 `jax.make_array_from_process_local_data`。
+同时加了显式 `DataLoader.close()`，训练退出时主动 shutdown worker 并超时终止残留进程。
+
+`3.275728 → 2.828845 s/step`（**−13.64%**）。之后数据侧分项均值：`main_queue_wait 0.005 s`、
+`jax_array_construct 0.009 s`、`h2d_wait 0.006 s`、`next_batch_total 0.046 s`——队列、array 构造和 H2D
+都已经不是瓶颈。
+
+### 2.2 XProf trace 定方向（commit `8ccf99c`）
+
+训练脚本加了可控的 trace 开关（绝对起始 step + step 数 + 输出目录，起止显式 `block_until_ready`）。
+采集 step 49011–49030 后按每卡每步汇总 kernel 时间：
+
+```
+GEMM     390.4 ms      Fusion   183.9 ms      NCCL   46.0 ms
+其他       2.6 ms      Memcpy     1.5 ms      Memset  0.5 ms
+```
+
+GEMM + Fusion 占已记录 kernel 时间的 **92.6%**，NCCL 7.4%。**继续优化 H2D 或 NCCL 没有意义，
+下一步必须动模型计算。** 这一步直接决定了阶段三去碰 attention。
+
+（注意：trace 里 kernel 可能重叠，不能把 duration 相加当 GPU 利用率。）
+
+### 2.3 只在日志步计算范数（commit `8ccf99c`，−0.5%）
+
+`grad_norm` / `param_norm` 是整个 trainable pytree 的归约，但只在 `log_interval` 步被记录。
+`train_step` 加静态布尔 `compute_metrics`（`scripts/train.py:168`）：
+
+```python
+    if compute_metrics:
+        grad_norm = optax.global_norm(grads)
+        param_norm = optax.global_norm(kernel_params)
+    else:
+        # 保持两个 JIT 变体的输出 pytree 一致；host 端用 nanmean 忽略占位值
+        grad_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
+        param_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
+```
+
+必须用**位置参数**传给 JIT：当前 JAX 在显式设置 `in_shardings` 时不接受 pjit 关键字实参。
+收益很小但语义正确、零代价，保留。
+
+### 2.4 三路相机合并为一次 SigLIP 调用（commit `8ccf99c`，−1.45%）
+
+`Pi0.embed_prefix` 原先对 head / left_wrist / right_wrist 分别调用同一个 image encoder 三次
+（`src/openpi/models/pi0.py:19`）：
+
+```python
+def _encode_image_views_as_batch(image_encoder, images):
+    """Encode camera views as a sample-major interleaved batch."""
+    names = tuple(images)
+    stacked = jnp.stack([images[name] for name in names], axis=1)   # (B, V, H, W, C)
+    batch_size, num_views = stacked.shape[:2]
+    flat = stacked.reshape((batch_size * num_views, *stacked.shape[2:]))
+    flat_tokens, _ = image_encoder(flat, train=False)               # 只调一次
+    tokens = flat_tokens.reshape((batch_size, num_views, *flat_tokens.shape[1:]))
+    return {name: tokens[:, index] for index, name in enumerate(names)}
+```
+
+布局是 **sample-major interleave**：同一样本的三路视图连续排列，所以 FSDP/DP 分片时每张卡仍拿到
+完整样本的所有视图，不需要跨设备重排。开关 `model.batch_image_views`（默认关）。
+单元测试逐元素比较了分别编码与合并编码的 token。
+
+累计：`3.275728 → 2.773743 s/step`，**−15.32%**（等价吞吐 +18.10%）。
+
+随后 `d0aec27` 只更新了两份 H100 RTC 配置的名称和参数，没有改变上述实现，也没有独立 A/B，
+因此不把它计算为新的提速阶段。
+
+> 一个容易被误读的现象：相机合并后 `dispatch` 从 2.657 s 降到 2.236 s，但 `compute_sync` 从 0.115 s
+> 涨到 0.492 s。这是异步工作在 dispatch 与同步边界之间重新分布，**不能宣称 dispatch 单项的 16% 提升**。
+
+---
+
+## 3. 阶段三：cuDNN fused Attention 与两类失效（08-29 ~ 08-31）
+
+> **一句话**：吞吐收益是真的（−29%），但全 18 层 cuDNN bf16 会让训练在约 1000 步后缓慢发散至 NaN。
+> 这个阶段的产出是**两类失效的严格区分**和一份验收标准，真正的修复在阶段四。
+
+### 3.1 实现与收益（commit `476babe`）
+
+真实训练形状：`3×256 图像 token + 200 文本 token + 50 action token = 1018`，
+`num_heads=8, num_kv_heads=1, head_dim=256, bf16`，block mask。
+
+```python
+jax.nn.dot_product_attention(q, k, v, mask=attn_mask, scale=1.0, implementation="cudnn")
+```
+
+`q` 在调用前已按 head dimension 缩放，所以必须 `scale=1.0` 避免重复缩放。只在 `kv_cache is None`
+的训练路径启用，带 KV cache 的推理保持原实现。
+
+8×H200 / batch 512 / step 16000 恢复的 A/B：`2.795537 → 1.978580 s/step`（**−29.22%**），
+等价吞吐 183.15 → 258.77 samples/s。数值对齐测试当时也过了：BF16 前向最大绝对误差 0.0039，
+Q/K/V 梯度最大绝对误差 1.9e-6 / 4.8e-7 / 2.4e-7。
+
+**这些都是真的。问题在于它们全都不足以作为正确性验收。**
+
+### 3.2 第一类失效：全空 mask 行的瞬时 NaN（08-29，commit `c1b70b2` `51d9af3`）
+
+prompt 被补齐到 `max_token_len=200`，padding token 既不是有效 query 也不是有效 key，
+所以 `make_attn_mask` 会生成整行全 `False` 的 query 行。切到 cuDNN 后从 clean step 16000 恢复：
+
+```
+Step 16000（切换前，显式）: grad_norm=0.6721, loss=0.1429, param_norm=1806.9792
+Step 16100（切换后，cuDNN）: grad_norm=nan,    loss=0.1185, param_norm=nan
+Step 16200 及以后:          grad_norm=nan,    loss=nan,    param_norm=nan
+```
+
+最小复现（BF16 GQA + 16 个全空 query 行）：**前向有限，Q 梯度 16384 个 NaN，K/V 梯度有限**。
+真实生产形状是 71680 个 Q 梯度 NaN。
+
+> **一个隐蔽的坑**：日志窗口对 loss 用 `nanmean`，第一个窗口里唯一有限的首步 loss 让聚合结果显示
+> `0.1185`，掩盖了窗口其余全是 NaN。**「窗口 loss 有限」不能证明窗口内每个更新都有限。**
+
+最终修复保留原始 mask，只对全空行停 Q 梯度（`src/openpi/models/gemma.py` 中的
+`_stop_gradient_for_fully_masked_queries`）：
+
+```python
+def _stop_gradient_for_fully_masked_queries(q, attn_mask):
+    query_has_key = jnp.any(attn_mask, axis=-1)[:, 0, :, None, None]
+    return jnp.where(query_has_key, q, jax.lax.stop_gradient(q))
+```
+
+语义：Q 前向值逐元素不变；有效 query 的输出和梯度不变；无效 query 的 Q 梯度固定为 0；block mask 完整保留。
+
+另一个方案是给每个空行开一个 dummy key——在有效行上给出**逐位相同**的 dQ/dK/dV，但每层要重建
+`(B, 1, T, S)` mask，batch 256 下 1.38 vs 1.05 s/step，所以选了 stop-gradient。
+
+提交演进要注意：`c1b70b2` 最初还加入了逐 step 的有限性检查，并采用 dummy-key mask；
+`6305750` 恢复原训练更新语义、移除这项逐步保护，`51d9af3` 再把 dummy-key 换成上面的 stop-gradient，
+这才是后续诊断和当前代码使用的实现。
+
+checkpoint 处置：原 step 16000 已被保留策略删除，step 20000–55000 全部受污染并删除，
+只能从 step 15000 重来（重跑约 1000 个有效 step）。**含 NaN 的参数或 Adam mu/nu 不能通过更换
+attention 实现恢复。**
+
+### 3.3 第二类失效：全层 cuDNN 在约 1000 步后缓慢发散（08-30 ~ 08-31，commit `ad08973` `895811c`）
+
+修完全空 mask 后，全 18 层 cuDNN 前约 1000 步与显式基线贴合，随后分叉：
+
+```
+step    显式基线          cuDNN 9.14        cuDNN 9.19
+ 200    0.9244 /   2.89   0.9333 /   3.58   0.9348 /   3.53
+1000    0.4220 /   2.77   0.4306 /   2.50   0.4319 /   3.16
+1200    0.3742 /   2.12   0.6484 /  31.40   0.6600 /  35.49     ← 分叉点
+2500    0.2457 /   1.33   5.7692 / 675.08   4.7055 / 759.18
+5900    0.1763 /   1.36   nan               已停止
+```
+
+单元格为 `loss / grad_norm`。**因果起点是 step 1000–1200，不是最终 NaN 的 step 5900。**
+单变量归因：同机同数据同超参同 seed，`git diff d0aec27 HEAD -- src/ scripts/` 的 72 行增删全部属于
+attention 开关本身。
+
+**但 cuDNN 的单步结果并不明显错误**：同权重同 batch 的完整模型梯度 A/B，梯度范数比 0.9937–1.0038，
+相对梯度差 1.7e-2–2.1e-2；单层对 fp32 参考的相对误差 cuDNN 5.4e-3 vs 显式 bf16 3.2e-3，同阶。
+cuDNN 同 batch 重复 backward 的随机差异（原子累加）就有 0.9%。**它给的是另一个合法的 BF16 近似。**
+
+### 3.4 机制：Adam 逐坐标放大结构性差异
+
+从 step-1000 checkpoint 恢复真实 Adam 状态，对同一 batch 比较实际参数更新：
+
+```
+                        相对更新差异   更新 cosine   符号翻转坐标
+真实全层 cuDNN 差分         3.3521%      0.999440      0.7885%
+1.5% 人造乘性噪声           0.7313%      --            0.1459%
+```
+
+梯度 L2 只差 1.5%，但经 Adam 逐坐标归一化后变成 3.35% 的更新差异，0.79% 的坐标符号翻转。
+而按 `|g|` 分配能量的 1.5% 独立随机噪声几乎不扰动这些坐标，**从 step 0 稳定跑满 3000 步**
+（step 2900 `loss=0.2411 / grad_norm=1.6389`）——证明训练配方并非对任意 1.5% 梯度噪声都没裕度，
+问题在**误差的结构**而非幅度。
+
+分层实验显示位置比数量更重要：
+
+```
+cuDNN 层段      相对更新差异   符号翻转坐标
+全 18 层           3.3118%       0.7769%
+最前 3 层          3.3016%       0.7720%     ← 只融合前 3 层几乎等同全 18 层
+中间 3 层          1.9894%       0.4634%
+最后 3 层          0.8470%       0.1926%
+稳定噪声参考       0.7313%       0.1459%
+```
+
+早层的微小数值差异会穿过后续 15 层前向和反向放大。
+
+### 3.5 当时的折中：只融合最后 3 层
+
+`configs/diag_cudnn_last3.yaml`，8×H100 / batch 256 / strict-order，从 step 0 跑满 3000 步，
+曲线与显式基线一致（step 1200 `0.3757 / 2.39`，step 2900 `0.2424 / 1.48`）。
+但 `~1.5 s/step` 对显式的 1.57–1.64，只有 4–9% 加速，远低于全层的潜在 36%。
+
+层选择开关（`Pi0Config`，阶段四后已不需要，保留作诊断用）：
+
+```yaml
+model:
+  use_cudnn_attention: true
+  cudnn_attention_layer_start: 15
+  cudnn_attention_num_layers: 3
+```
+
+### 3.6 CUDA/cuDNN 环境事故
+
+事故期间 `cuda_versions.cudnn_get_version()` 返回 `91400`，但进程实际加载的是
+**9.10.2 的 dispatcher stub + 9.14 的 engine 库**——`cudnnGetVersion()` 的返回值来自 engine，
+所以**只看这个数字发现不了混装**。
+
+```
+site-packages/nvidia/cudnn/lib/libcudnn.so.9        <- 9.10.2 dispatcher
+$CONDA_PREFIX/lib/libcudnn_graph.so.9.14.0          <- 9.14 engine
+```
+
+当前环境已统一（CUDA 12.8 / PyTorch 2.11.0+cu128 / nvidia-cudnn-cu12 9.19.0.56 / JAX 0.5.3）。
+**不要再导出 `LD_LIBRARY_PATH="$CONDA_PREFIX/lib"`**，残留的 conda-forge `libcublas`/`libcudart`/`libnccl`
+会遮蔽 pip cu128 库。生产启动前在同一 shell 跑：
+
+```bash
+env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py
+```
+
+该脚本遍历 `/proc/self/maps`，在 `libcudnn*` 来自多个目录或版本不一致时判 FAIL，
+并在真实训练形状上做前向+反向与 fp32 参考对比。
+
+> **PASS 只证明环境单源、单步内核有限。第二类缓慢发散曾通过所有这类 smoke test。**
+
+---
+
+## 4. 阶段四：根因与 fp16 修复（09-03 ~ 09-04，commit `df9755f` `3ea2b1c`）
+
+阶段三证明了「BF16 差分被 Adam 放大」，但没回答**差分从哪来、为什么它不像随机噪声**。这一阶段用真实激活回答了。
+
+### 4.1 机制：flash 反向的 D 项用了 bf16 舍入后的 O
+
+cuDNN fused attention 的反向按 flash-attention 方式算 `dS = P ⊙ (dP − D)`，其中
+`D_i = rowsum(dO_i ⊙ O_i)` 用的是**前向存下的、已舍入到 bf16 的 O**。
+显式路径用同一份舍入后的 dP 算 `D_i = Σ_j P_ij dP_ij`，抵消是自洽的。
+
+两者之差是一个每行标量：
+
+```
+Δ_i = dO_i · (bf16(O_i) − O_i)
+```
+
+进入梯度后方向为 `−Δ_i · Σ_j P_ij K_j`，**破坏了 softmax 梯度「行和为零」的不变量**。
+
+对 **peaked 行**（`maxP > 0.99`，即注意力 sink）杀伤最大：那里 `dP` 与 `rowsum(dO·O)` 几乎完全抵消，
+真实 dQ 是 ε 级量，比普通行小两个数量级，而 H1 误差不随 ε 缩小——所以相对误差达 O(1)。
+
+### 4.2 证据
+
+两个探针：`scripts/attention_dterm_probe.py`（GPU，monkeypatch `gemma.Block.__call__`，
+用 `jax.debug.callback` 在一次真实训练步里抓下 18 层的 q/k/v/mask 和反向余切 dO，缓存 1.5 GB）；
+`scripts/attention_dterm_cpu_analysis.py`（CPU float64，对 peaked 行逐行重建精确反向）。
+
+> 实现细节：模拟 bf16 舍入必须用 `jax.lax.reduce_precision`，`astype` 往返会被 XLA 的
+> `allow_excess_precision` 删掉，预测会恒为零。
+
+单层实测（`cos` 是误差向量与 H1 预测方向的 cos）：
+
+```
+层   变体           dQ rel    cos(all)   peaked rel   cos(peaked)
+ 0   explicit_bf16  2.46e-3   -0.02      2.8e-3       -0.04
+     cudnn_bf16     3.58e-3   +0.49      3.5e-1       +0.83
+     cudnn_fp16     1.71e-3   -0.02      3.9e-2       +0.13
+ 8   explicit_bf16  2.88e-3   +0.00      3.5e-3       -0.10
+     cudnn_bf16     4.28e-3   +0.51      5.1e-1       +0.96
+     cudnn_fp16     1.78e-3   +0.03      7.2e-2       +0.57
+
+18 层几何平均：explicit_bf16 2.33e-3 / cudnn_bf16 3.72e-3 / cudnn_fp16 1.76e-3
+18 层平均 cos(all)：explicit 0.006 / cudnn_bf16 0.418 / cudnn_fp16 0.006
+```
+
+- cuDNN bf16 的误差在**每一层**都与 H1 预测显著同向；显式和 fp16 都不同向。
+- H1 预测的 peaked 行误差量级与实测吻合（0.41/0.35、0.63/0.51、0.86/1.21）。cos 不到 1 是因为
+  cuDNN 前向的 PV 乘也把 P 舍入到 bf16。
+- 候选机制 H2（dP 先舍入再减 D）的 cos 为 0 或负，排除。
+- dV 没有 D 项，cuDNN bf16 的 dV 误差（2.5e-3）与显式（1.9e-3）同阶——干净的对照。
+
+**这解释了阶段三的全部现象**：H1 分量按行沿 `Σ_j P_ij K_j` 方向、符号随舍入随机，投到参数上落在
+`x ⊗ K_argmax` 这类**低维子空间**，不与 `|g|` 成比例。sink 机制处于平衡态时这些坐标的真实梯度和
+二阶矩都很小，Adam 归一化后每步都在那里走满 lr 的随机步，约 1000 步后 sink 结构失稳。
+乘性噪声在小 `|g|` 坐标上也小，Adam 看不见。跨 batch delta 两两 cos ±0.25 在 30 亿维空间里不是白噪声，
+正是「同一子空间、随机符号」。只融合前 3 层≈全 18 层，因为 sink/peaked 结构在早层形成。
+改 mask、seq_lengths、cuDNN 版本、LR、b1/b2/eps 都没用，**因为它们都不改变 O 的舍入**。
+
+### 4.3 修复：fp16 内核 + 2 的幂动态缩放
+
+fp16 比 bf16 多 3 位尾数，H1 误差降约 8 倍；代价是指数范围小得多，所以余切要做动态 loss scaling。
+真实激活的范围 `max|q| ≤ 3.5`、`|k| ≤ 30`、`|v| ≤ 84`、`|dO| 1e-4–1e-2`，都在 fp16 范围内。
+
+关键实现（`src/openpi/models/gemma.py:201`，`_cudnn_attention_in_dtype`）：
+
+```python
+    def attention_fwd(q, k, v, attn_mask):
+        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
+        bias = jnp.where(attn_mask, jnp.asarray(0, compute_dtype),
+                         _cudnn_fa.get_large_negative_number(compute_dtype))
+        # cuDNN 的 custom partitioner 要求 q/k/v/bias 分片一致
+        qc, kc, vc, bias = sharding.activation_sharding_constraint((qc, kc, vc, bias))
+        out, res = _cudnn_fa._dot_product_attention_fwd_rule(...)
+        return out.astype(out_dtype), (res, attn_mask)
+
+    def attention_bwd(residuals, d_out):
+        res, attn_mask = residuals
+        d_out32 = d_out.astype(jnp.float32)
+        max_abs = jnp.max(jnp.abs(d_out32))
+        # 2 的幂缩放，使 max|d_out * scale| 落在 [0.25, 0.5)：尾数上精确，
+        # 给内核里的 dP = dO @ V^T 和 dQ/dK/dV 留约 17 个 binade 余量
+        exponent = jnp.floor(jnp.log2(jnp.maximum(max_abs, tiny)))
+        scale = jnp.where(max_abs > 0, jnp.exp2(-(exponent + 2.0)), 1.0)
+        grads = _cudnn_fa._dot_product_attention_bwd_rule(..., (d_out32 * scale).astype(compute_dtype))
+        dq, dk, dv = grads[:3]
+        # 与 _stop_gradient_for_fully_masked_queries 同语义：全 mask 行的 dQ 置零而非 NaN
+        query_has_key = jnp.any(attn_mask, axis=-1)[:, 0, :, None, None]
+        dq = jnp.where(query_has_key, dq, jnp.zeros_like(dq))
+        ...
+```
+
+开关：`model.cudnn_attention_dtype: float16`（默认 `bfloat16`，逐位不变）。
+
+**两版实现**：
+
+- **v1**（3000 步验收用的）：custom VJP 的反向里用 `jax.vjp` 重走一遍 `jax.nn.dot_product_attention`
+  前向再取反向。配合 `nothing_saveable` remat，每层反向比 bf16 路径多一次 attention 前向。1.37 s/step。
+- **v2**（当前代码，commit `3ea2b1c`）：直接调 jax 0.5.3 私有的
+  `jax._src.cudnn.fused_attention_stablehlo._dot_product_attention_fwd_rule / _bwd_rule`
+  ——就是 `jax.nn.dot_product_attention(implementation="cudnn")` 自己 custom_vjp 里接的那两个函数。
+  softmax stats 和输出当残差保存，反向只跑一次 fused backward。1.31–1.35 s/step。
+
+  `attn_mask` 必须作为**显式参数**而非闭包传入，否则 flax remat 单独 trace 反向时会捕获泄漏的 tracer。
+
+  `scripts/cudnn_fp16_vjp_equivalence.py` 在真实激活（层 0/4/8/14/17）上比较 v2 与 v1：
+  前向逐位相同，dK/dV 逐位相同，dQ 相对差 5e-6–1.2e-5，与同一内核自身的 run-to-run 差
+  （1.6e-6–1.4e-5，原子累加）同阶。
+
+### 4.4 3000 步 strict-order 验收：通过
+
+`configs/diag_cudnn_fp16.yaml`，与 `diag_cudnn_strict_order` 同数据/seed/batch 顺序/超参，
+全 18 层 cuDNN fp16，从 pi05_base step 0 起，8×H100 batch 256：
+
+```
+step    fp16 loss / grad_norm    显式基线           last3 基线        历史 cuDNN bf16
+   0    6.4030 /  65.59          6.4023 / 65.63     6.4029 / 65.52    6.4023 / 65.63
+1000    0.4231 /   2.21          0.4220 /  2.77     0.4240 /  2.38    0.4319 /  3.16
+1200    0.3759 /   2.36          0.3742 /  2.12     0.3757 /  2.39    0.6600 / 35.49   ← 历史分叉点
+2000    0.2792 /   1.93                             0.2793 /  1.83    3.7753 / 150.5
+2900    0.2427 /   1.59                             0.2424 /  1.48    6.2039 / 603.5
+```
+
+30 个 log 点全部有限，`grad_norm` 最大 4.75（step 100），之后 1.4–2.8，**没有一个点超过基线同点的
+1.3 倍**（验收上限 3 倍）。历史分叉点处 loss 与显式基线差 <0.5%。
+**把 cuDNN 内核精度从 bf16 换成 fp16 就消除了发散，验证了 §4.1 的机制。**
+
+### 4.5 步时：attention 内核本身不解释差距
+
+```
+路径                                   s/step      来源
+显式 attention                          1.57–1.64   8-31 报告 / 今日复测 1.60–1.63
+cuDNN bf16（发散）                       1.05        8-31 报告
+cuDNN fp16 v1（反向重算前向）             1.34–1.37   3000 步验收 run
+cuDNN fp16 v2（直接调 fwd/bwd 规则）      1.31–1.35   400 步（loss 逐点与 v1 一致）
+cuDNN bf16，今日同 config 复测            1.32–1.35   400 步
+```
+
+去掉反向重算只省约 2%。单卡微基准（`scripts/cudnn_attention_microbench.py`，层 0 真实激活，
+T=S=1018，含 `nothing_saveable` remat）说明原因：
+
+```
+变体                fwd      fwd+bwd(remat)   ×18 层
+cudnn_bf16         1.35 ms   5.24 ms          0.094 s/step
+cudnn_fp16_v2      1.54 ms   5.53 ms          0.099 s/step
+explicit_bf16      4.46 ms   8.49 ms          0.153 s/step
+```
+
+18 层 attention 前向+反向合计只占每步约 **0.1 s**；fp16 内核与 bf16 同速，动态缩放开销可忽略，
+显式路径也只多 0.06 s。**所以 attention 内核解释不了 1.33 vs 1.05 的差距，也解释不了显式的 1.6 s**
+——差距在整图层面（XLA 的 fusion/remat 决策、显存压力）。今日同 config 的 bf16 复测已是 1.32–1.35，
+与 fp16 持平，说明 8-31 那个 1.05 s 不是同一编译环境下的可比数字。
+
+### 4.6 一个必须记住的负面结论：单步指标不区分好坏
+
+`scripts/attention_precision_update_probe.py`，锚点是 fp16 验收 run 的 step-1000 train state
+（参数 + Adam 矩），以 fp32 显式 attention 为真值：
+
+```
+                   原始梯度 rel / sign_flips     Adam 更新 rel / cos / sign_flips
+explicit_bf16      3.37e-2 / 1.53e-2            3.13e-2 / 0.99951 / 7.47e-3
+cudnn_bf16         3.18e-2 / 1.43e-2            2.93e-2 / 0.99957 / 6.97e-3     ← 会发散的那个
+cudnn_fp16         3.01e-2 / 1.38e-2            2.81e-2 / 0.99961 / 6.73e-3
+```
+
+**会发散的 cudnn_bf16 在这个指标上与两条稳定路径不可区分，甚至略好于 explicit_bf16。**
+D 项误差在范数上只占 ~1e-3，危害来自方向的结构性（同一低维子空间、跨 step 累积）。
+判定 fused attention 后端好坏要看 §4.2 的 `cos(err, H1)` 与 peaked 行误差，最终靠 3000 步 strict-order run。
+**不要再用单步梯度/更新距离做验收指标。**
+
+---
+
+## 5. 阶段五：H100 端到端 A/B（09-04，commit `cf07bbe`）
+
+**配置**：`pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml`
+（pi05、cuDNN fp16、`batch_image_views`、全局 batch 256、64 workers、fsdp 8）
+**数据**：`Xense/earbud_case_insertion_teleop_0515_left_right`，776 集 / 1,145,613 帧
+**口径**：每个变体从 pi05_base 起跑 400 步，`strict_batch_order: true`，关 checkpoint 和 W&B；
+主指标是日志 `Step 100` 到 `Step 300` 的时间戳之差除以 200（含日志开销的端到端墙钟）；
+`loss@300` 用来确认数值没变。
+
+| 变体 | s/step | 相对基线 | loss@300 | 结论 |
+| --- | --- | --- | --- | --- |
+| E0 基线 | 1.328 | — | 0.7356 | 等数据 0.03 s，数据侧不是瓶颈 |
+| E1 `max_token_len` 200→128 | 1.246 | **−6.2%** | 0.7356 | 采纳 |
+| E2 关闭 remat | OOM | — | — | XLA 试图分配 280 GiB |
+| E2b gemma remat `dots_with_no_batch_dims_saveable` | 编译失败 | — | — | cuDNN 分片器报 q/k/v 分片不一致 |
+| E2c 同上 + q/k/v 显式分片约束 | OOM | — | — | 编译过了，单块残差 84.6 GB |
+| E3 XLA 集合通信标志 | 1.281 | **−3.5%** | 0.7357 | 采纳 |
+| E4 `fsdp_devices` 8→4 | 1.335 | +0.5% | 0.7354 | 无收益 |
+| **E5 = E1 + E3** | **1.193** | **−10.2%** | 0.7356 | 生产采用；192.8 → 214.6 samples/s |
+
+60000 步按稳态估算：22.1 h → 19.9 h，省约 2.2 h。
+
+### 5.1 E1：token 长度按数据集实测收紧
+
+`scripts/token_len_scan.py` 走真实 transform 链（repack → BiFlexivInputs → Normalize → TokenizePrompt），
+对 4037 个状态（含每一维的最小/最大值样本）做 tokenize：min 110、p50 113、p99 115、**max 116**。
+原来 200 里有 84 个是纯 padding，但 attention/MLP 照算。128 留 12 个余量。
+
+```bash
+JAX_PLATFORMS=cpu uv run scripts/token_len_scan.py \
+  --config-name pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100
+```
+
+> **换数据集或改 prompt 时必须重扫**：tokenizer 超长只会 warning 并截断，不会报错。
+
+### 5.2 E3：XLA 标志
+
+```
+--xla_gpu_enable_latency_hiding_scheduler=true
+--xla_gpu_all_gather_combine_threshold_bytes=1073741824
+--xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824
+--xla_gpu_all_reduce_combine_threshold_bytes=1073741824
+--xla_gpu_enable_pipelined_all_gather=true
+--xla_gpu_enable_pipelined_reduce_scatter=true
+--xla_gpu_enable_while_loop_double_buffering=true
+```
+
+FSDP 每步的参数 all-gather / 梯度 reduce-scatter 合并成大块并与计算流水。
+**没有逐项拆分测哪一个起作用。** 通过 `XLA_FLAGS` 环境变量传入，已写在 0904 yaml 头部注释里。
+
+### 5.3 E2：remat 在 80 GB 上没有空间
+
+`gemma.py` / `siglip.py` 每层用 `nothing_saveable` 全量重算。本次给 `Pi0Config` 加了
+`gemma_remat_policy` / `siglip_remat_policy`（默认不变），并在 cuDNN 调用前给 q/k/v/mask 加了统一的
+batch 分片约束——否则非默认 remat 策略下编译失败：
+
+```python
+def _cudnn_attention_call(q, k, v, attn_mask):
+    q = _stop_gradient_for_fully_masked_queries(q, attn_mask)
+    # cuDNN 的 custom partitioner 要求 q/k/v 分片一致。显式钉到 batch 轴，
+    # 否则 XLA 分片传播可能逐 operand 不同（非默认 remat 策略下出现过），编译报 "should have same sharding"。
+    q, k, v, attn_mask = sharding.activation_sharding_constraint((q, k, v, attn_mask))
+    return jax.nn.dot_product_attention(q, k, v, mask=attn_mask, scale=1.0, implementation="cudnn")
+```
+
+但 batch 256（每卡 32）时：完全不 remat 要 280 GiB 单块，只保存无 batch 维 matmul 也要 84.6 GB，都放不下。
+**在 H200 141 GB 或每卡 batch 更小时这两个开关才可能兑现**，理论上限约为一次前向的重算（15–25%）。
+
+### 5.4 其他观察
+
+- worker 启动一次约 330 s（spawn 上下文，64 进程各自 import jax/torch/lerobot 并反序列化 1.1M 行数据集），
+  是每次启动的一次性开销。
+- `ResizeImages` 实际调用的是 `xense_client.image_tools`（PIL），worker 里**不跑 JAX resize**
+  ——阶段一怀疑的头号嫌疑人被排除。
+- fsdp 4 相对 8 没有收益：参数 all-gather 的量只差 12.5%，而每卡持有的优化器状态翻倍。
+- 短跑结束后偶有 DataLoader worker 成为孤儿进程（ppid 1，每个约 3 GB RSS，一次残留 39 个）；
+  批量实验后用 `ps -eo pid,ppid,cmd | grep spawn_main` 检查。
+- 该配置 `num_train_steps: 60000` 但 `lr_schedule.decay_steps` 仍是默认 30000，后半程一直在底线 LR。
+  **这是训练质量问题，本次未动。**
+
+---
+
+## 6. 新 attention 后端的验收标准
+
+任何新的全层 fused/flash attention backend、custom VJP 或稳定化方案必须：
+
+1. 从 **step 0** 开始；
+2. `strict_batch_order: true`，固定 seed、数据集、batch 和超参；
+3. 跑满 **3000 个更新**，不得只跑 40、60 或 500 步；
+4. 与显式基线逐点比较 step 200/500/1000/1500/2000/2900；
+5. 全程 loss、梯度、参数和 optimizer state 有限；
+6. `grad_norm` 不得超过基线同点的 3 倍；
+7. 通过后再进入有 checkpoint 的受监控长跑。
+
+**以下均不是有效验收**：单次前向/反向有限；40-step 吞吐短跑；从已收敛 checkpoint 短跑；
+200–500 个连续有限 step；只检查最终是否 NaN；单步梯度/Adam 更新距离（§4.6）；
+跨不同 batch 顺序的 run 直接比较单点 grad norm。
+
+**建议增加发散护栏**：当前只在日志步算 `grad_norm`。在 host 端维护滚动中位数，
+连续多个窗口超过固定倍数时中止——按 §3.3 的曲线可在 step 1200–1300 停下，而不是污染数万个 step。
+
+**checkpoint 恢复规则**：一旦参数或 Adam `mu`/`nu` 出现 NaN，后续 checkpoint 全部视为污染，
+不得通过切换 attention、降 LR 或改 optimizer 继续；必须回退到最后一个确定干净的 checkpoint，
+用独立实验名验证（关 W&B 和所有保存）后才能恢复生产保存。
+
+---
+
+## 7. 不要重复的弯路
+
+**测量方法**
+
+1. **先看停顿的周期再动手。** 周期 = `num_workers` 是队头阻塞，周期 = 一个 epoch 的步数是边界重建，
+   周期 = `save_interval` 是 checkpoint——都不是「解码太慢」。
+2. **`buffered` 千金难买。** 停顿时它是 0 说明真的算不出来，是 41 说明算出来了过不去，两种修法完全相反。
+3. **别把别的机器上的预估当结论。** 「跳过触觉解码约 1.8×」在 h200 箱 batch 512 下是对的，
+   搬到 batch 256 + 192 核就一点不剩。
+4. **别拿两个不同任务的日志比绝对时长。** 要比就比归一化指标，或者跑受控 A/B。
+5. **别单看 `dispatch`。** JAX 异步，收益可能只是从 dispatch 挪到 compute_sync。
+6. **`nanmean` 会掩盖 NaN。** 窗口 loss 有限不证明窗口内每个更新有限。
+
+**cuDNN attention**
+
+7. 不要再把全空 mask 的瞬时 NaN 与全层缓慢发散当成同一问题。
+8. 不要再把 cuDNN 9.15.1 的非 128 倍序列长度缺陷或 `T=1024` 当作本任务的修复
+   （实测 T=1018 与 T=1024 的 dQ 误差同为 3.2e-3，NaN=0）。
+9. 不要再尝试传完整物理长度的 `query_seq_lengths`/`key_value_seq_lengths`
+   （编译确认进了 custom call，但数值指纹 `3.311% / 0.776%` 与已知发散路径完全一致，3000 步照样发散）。
+10. 不要再改 dummy-key / stop-gradient 的 mask 语义来解决发散——两者在有效行上数值等价。
+11. 不要只调 `peak_lr`（2.5e-5→1.5e-5 只推迟约 300 步）、Adam `b1`（0.9→0.99 推迟到约 1400 步）、
+    `b2`（0.95→0.99 照样发散）或 `eps`（更新差异几乎不变）。
+12. 不要用独立随机噪声代替真实 cuDNN 差分做实验——1.5% 乘性噪声能稳定跑满 3000 步。
+13. 不要试「V 减去 sink 键的 V」：peaked 行大多不在看 BOS，而是在看 action/text/image 的其他键。
+14. 不要只看 `cuda_versions.cudnn_get_version()` 判断动态库是否干净（它由 engine 返回，看不见 dispatcher 混装）。
+15. 用 `jax.lax.reduce_precision` 模拟低精度舍入，`astype` 往返会被 XLA 删掉。
+
+---
+
+## 8. 脚本与配置索引
+
+**诊断脚本**
+
+| 脚本 | 用途 |
+| --- | --- |
+| `scripts/check_cuda_stack.py` | 枚举 `/proc/self/maps` 里实际加载的 cuDNN/cuBLAS/NCCL，+ 真实形状数值 smoke test |
+| `scripts/token_len_scan.py` | 走真实 transform 链扫数据集的 tokenized prompt 长度分布 |
+| `scripts/attention_dterm_probe.py` | 抓真实激活 + 单层五变体误差 + H1/H2 预测（阶段四主力） |
+| `scripts/attention_dterm_cpu_analysis.py` | peaked 行的 float64 逐行解剖 |
+| `scripts/cudnn_fp16_vjp_equivalence.py` | fp16 VJP v2（直接调规则）与 v1（反向重算）的等价性 |
+| `scripts/cudnn_attention_microbench.py` | 单层 attention 前向/反向微基准 |
+| `scripts/attention_precision_update_probe.py` | 真实 Adam 状态上的更新空间对比（**注意 §4.6：该指标不区分好坏**） |
+| `scripts/grad_ab_probe.py` | 完整模型显式/cuDNN 梯度 A/B |
+| `scripts/grad_ab_direction_probe.py` | 多 batch 差分方向和随机分量 |
+| `scripts/hybrid_update_ab_probe.py` | 不同 cuDNN 层段的更新空间误差 |
+| `scripts/optimizer_update_ab_probe.py` | 从真实 Adam 状态比较参数更新误差 |
+
+**诊断配置**：`configs/diag_cudnn_strict_order.yaml`（全层 bf16，会发散）、
+`configs/diag_explicit_strict_order.yaml`（基线）、`configs/diag_cudnn_last3.yaml`（后 3 层折中）、
+`configs/diag_cudnn_fp16.yaml`（fp16 验收）、`configs/diag_cudnn_b1_099.yaml`。
+
+**代码位置**
+
+| 文件 | 相关改动 |
+| --- | --- |
+| `src/openpi/training/data_loader.py` | `SelectiveVideoLeRobotDataset`、`_resolve_decode_video_keys`、`InfiniteSampler`、`in_order` 接线、Torch Tensor collate、`close()` |
+| `src/openpi/training/config.py` | `DataConfig.tactile`、`TrainConfig.strict_batch_order` |
+| `src/openpi/models/pi0_config.py` | `batch_image_views`、`use_cudnn_attention`、`cudnn_attention_{layer_start,num_layers,dtype}`、`gemma_remat_policy`、`siglip_remat_policy` |
+| `src/openpi/models/pi0.py` | `_encode_image_views_as_batch` |
+| `src/openpi/models/gemma.py` | `_stop_gradient_for_fully_masked_queries`、`_cudnn_attention_call`、`_cudnn_attention_in_dtype`、`_cudnn_static_args`、逐层 cuDNN 开关 |
+| `scripts/train.py` | `compute_metrics` 静态分支、XProf trace 开关、数据管线 profiling |
+
+---
+
+## 9. 生产启动
+
+```bash
+env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py    # 先确认环境单源
+
+XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true \
+  --xla_gpu_all_gather_combine_threshold_bytes=1073741824 \
+  --xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824 \
+  --xla_gpu_all_reduce_combine_threshold_bytes=1073741824 \
+  --xla_gpu_enable_pipelined_all_gather=true \
+  --xla_gpu_enable_pipelined_reduce_scatter=true \
+  --xla_gpu_enable_while_loop_double_buffering=true" \
+XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+  python scripts/train.py <config> --exp-name=<name>
+```
+
+换数据集时记得：重扫 `max_token_len`（§5.1）、重算 norm stats、检查 `lr_schedule.decay_steps`
+是否匹配 `num_train_steps`（§5.4）。
diff --git a/pyproject.toml b/pyproject.toml
index fe708f3..a104f90 100755
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -35,7 +35,7 @@ dependencies = [
     # nvidia-cudnn-cu12==9.19.0.56, while 2.10 pinned ==9.10.2.21. JAX only asks for
     # >=9.1,<10, so torch is what decides which cuDNN the training process loads even
     # though torch itself only runs the DataLoader here. Keep torch and the cuDNN you
-    # want to train against in sync; see docs/2026-08-30-cudnn-attention-divergence.md.
+    # want to train against in sync; see docs/training-optimization.md.
     "torch>=2.11,<2.12",
     "torchcodec>=0.11,<0.12; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')",
     "torchvision>=0.26,<0.27",
diff --git a/scripts/attention_dterm_probe.py b/scripts/attention_dterm_probe.py
index f0b391d..2d1347f 100644
--- a/scripts/attention_dterm_probe.py
+++ b/scripts/attention_dterm_probe.py
@@ -1,6 +1,6 @@
 """Test the flash-attention D-term hypothesis on real Pi0.5 activations.
 
-Hypothesis (docs/2026-08-31-cudnn-sdpa-nan-summary.md, open question): the cuDNN
+Mechanism under test (confirmed in docs/training-optimization.md, section 4): the cuDNN
 fused backward forms ``dS = P * (dP - D)`` with ``D_i = rowsum(dO_i * O_i)`` taken
 from the *bf16-rounded* stored output ``O``. On attention-sink rows (P nearly
 one-hot) ``dP`` and ``D`` cancel almost exactly, so the rounding error of ``O`` is
diff --git a/scripts/token_len_scan.py b/scripts/token_len_scan.py
index 59fd8da..432c79c 100644
--- a/scripts/token_len_scan.py
+++ b/scripts/token_len_scan.py
@@ -3,7 +3,7 @@
 For Pi0.5 the language prefix is ``Task: <prompt>, State: <discretized state>;\\nAction: `` and
 everything up to ``max_token_len`` is padded (and masked) -- but attention/MLP still run over the
 padding. On 2026-09-04 trimming ``max_token_len`` 200 -> 128 gave -6.2% step time with identical
-loss (see docs/2026-09-04-h100-speed-ab.md). This script runs the *actual* transform chain
+loss (see docs/training-optimization.md). This script runs the *actual* transform chain
 (repack -> data transforms -> model transforms) on real states with dummy images and reports the
 token-length distribution so ``max_token_len`` can be set with a known margin.
 
diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index 20a8bb1..f50aa39 100644
--- a/src/openpi/models/gemma.py
+++ b/src/openpi/models/gemma.py
@@ -179,10 +179,10 @@ def _stop_gradient_for_fully_masked_queries(q, attn_mask):
     training shape) but rebuilds the whole (B, 1, T, S) mask inside every layer,
     which costs about 0.3 s/step at batch 256. Prefer this one.
 
-    Neither variant is the cure for the 2026-08-30 divergence, which is still open:
-    the kernel itself is correct, but cuDNN's answer differs from the explicit path
-    by about 2% per step and this recipe has no margin for that. Do not spend time on
-    the mask again. See docs/2026-08-30-cudnn-attention-divergence.md.
+    Neither mask variant cures the 2026-08-30 bfloat16 divergence: the kernel's
+    rounded backward residual creates a structured error on peaked attention rows.
+    Use the validated float16 custom VJP instead; do not spend time on the mask again.
+    See docs/training-optimization.md.
     """
     query_has_key = jnp.any(attn_mask, axis=-1)[:, 0, :, None, None]
     return jnp.where(query_has_key, q, jax.lax.stop_gradient(q))
diff --git a/src/openpi/models/pi0_config.py b/src/openpi/models/pi0_config.py
index 076d59a..95ebad7 100644
--- a/src/openpi/models/pi0_config.py
+++ b/src/openpi/models/pi0_config.py
@@ -50,7 +50,7 @@ class Pi0Config(_model.BaseModelConfig):
     # pass; "none" disables remat entirely (fastest, most memory); any other name is looked
     # up in jax.checkpoint_policies (e.g. "dots_with_no_batch_dims_saveable").
     # NOTE: on 8xH100 80GB with global batch 256 neither "none" (XLA wants a 280 GiB buffer)
-    # nor "dots_with_no_batch_dims_saveable" (84.6 GB) fits; see docs/2026-09-04-h100-speed-ab.md.
+    # nor "dots_with_no_batch_dims_saveable" (84.6 GB) fits; see docs/training-optimization.md.
     gemma_remat_policy: str = "nothing_saveable"
     siglip_remat_policy: str = "nothing_saveable"
 

From 460675421c65df6a197c25dff361cab7714e65e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Fri, 4 Sep 2026 13:50:17 +0800
Subject: [PATCH 18/28] configs: refresh optimized training examples

---
 README.md                                     | 18 +++++++++
 ...one_stand_lora_0422_merged_fixed_h100.yaml | 21 -----------
 ...s_case_insertion_teleop_rtc_0828_h100.yaml | 30 ---------------
 ...s_case_insertion_teleop_rtc_0904_h100.yaml | 37 +++++++++++++------
 ..._case_insertion_teleop_rtc_0904_h200.yaml} |  0
 ...i_flexiv_pack_6_cosmetic_bottles_lora.yaml | 24 ------------
 configs/_examples/pi05_droid.yaml             | 12 ------
 7 files changed, 44 insertions(+), 98 deletions(-)
 delete mode 100644 configs/_examples/pi05_base_bi_flexiv_assemble_box_with_phone_stand_lora_0422_merged_fixed_h100.yaml
 delete mode 100644 configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml
 rename configs/_examples/{pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml => pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h200.yaml} (100%)
 delete mode 100644 configs/_examples/pi05_base_bi_flexiv_pack_6_cosmetic_bottles_lora.yaml
 delete mode 100644 configs/_examples/pi05_droid.yaml

diff --git a/README.md b/README.md
index fe522ce..119a64a 100644
--- a/README.md
+++ b/README.md
@@ -222,6 +222,24 @@ XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
 python scripts/serve_policy.py policy:checkpoint --policy.config=my_task --policy.dir=checkpoints/my_task/my_exp/<step>
 ```
 
+For the optimized 8×H100 FSDP launch, set the measured XLA collective
+combining and pipelining flags before Python starts:
+
+```bash
+env -u LD_LIBRARY_PATH \
+  XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true \
+    --xla_gpu_all_gather_combine_threshold_bytes=1073741824 \
+    --xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824 \
+    --xla_gpu_all_reduce_combine_threshold_bytes=1073741824 \
+    --xla_gpu_enable_pipelined_all_gather=true \
+    --xla_gpu_enable_pipelined_reduce_scatter=true \
+    --xla_gpu_enable_while_loop_double_buffering=true" \
+  XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+  python scripts/train.py my_task \
+    --exp-name=my_exp \
+    --overwrite
+```
+
 Do **not** prepend `$CONDA_PREFIX/lib` to `LD_LIBRARY_PATH`. The supported
 environment uses the pip CUDA 12.8 stack installed by `lerobot-xense`:
 PyTorch 2.11 pins cuDNN 9.19 and JAX shares that same runtime. Before training,
diff --git a/configs/_examples/pi05_base_bi_flexiv_assemble_box_with_phone_stand_lora_0422_merged_fixed_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_assemble_box_with_phone_stand_lora_0422_merged_fixed_h100.yaml
deleted file mode 100644
index e39b6e6..0000000
--- a/configs/_examples/pi05_base_bi_flexiv_assemble_box_with_phone_stand_lora_0422_merged_fixed_h100.yaml
+++ /dev/null
@@ -1,21 +0,0 @@
-model:
-  type: Pi0Config
-  max_token_len: 200
-  pi05: true
-  discrete_state_input: true
-  enable_training_time_rtc: true
-weight_loader:
-  type: CheckpointWeightLoader
-  params_path: gs://openpi-assets/checkpoints/pi05_base/params
-ema_decay: null
-data:
-  type: LeRobotBiFlexivDataConfig
-  repo_id: Xense/assemble_box_with_phone_stand0410_merged_fixed
-  base_config:
-    prompt_from_task: true
-  default_prompt: Assemble the packaging by folding the flat box into shape, placing
-    the metal phone stand inside, and closing the box properly.
-batch_size: 256
-num_workers: 64
-num_train_steps: 80000
-fsdp_devices: 8
diff --git a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml
deleted file mode 100644
index d5eadb9..0000000
--- a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0828_h100.yaml
+++ /dev/null
@@ -1,30 +0,0 @@
-model:
-  type: Pi0Config
-  max_token_len: 200
-  pi05: true
-  discrete_state_input: true
-  enable_training_time_rtc: true
-weight_loader:
-  # Was /home/li/hubo/xense-openpi/checkpoints/20000/params — a continuation of a
-  # 20k-step run that only existed on one laptop, so nobody else could start this
-  # config. Points at the upstream base now; to continue from your own checkpoint:
-  #   uv run scripts/train.py <this-config> --exp-name=... \
-  #     --weight_loader.params_path=/path/to/checkpoint/params
-  type: CheckpointWeightLoader
-  params_path: gs://openpi-assets/checkpoints/pi05_base/params
-ema_decay: null
-data:
-  type: LeRobotBiFlexivDataConfig
-  repo_id: Xense/earbud_case_insertion_teleop_0515_left_right
-  base_config:
-    prompt_from_task: true
-    tactile: false
-  default_prompt: Pick up each earbud case from the left stands, insert the matching
-    earbuds, close the lid, and place the case in the box.
-batch_size: 256
-num_workers: 64
-num_train_steps: 60000
-save_interval: 10000
-keep_period: 10000
-fsdp_devices: 8
-strict_batch_order: false
diff --git a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
index bff7ef0..6ee2686 100644
--- a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
@@ -1,6 +1,9 @@
-# Launch (8xH100, hb-openpi env). XLA_FLAGS below measured -3.5% step time on 2026-09-04
-# (FSDP collective combining + pipelined all-gather/reduce-scatter); see
-# docs/training-optimization.md for the A/B table.
+# Production launch (8xH100, hb-openpi env). XLA_FLAGS must be set before Python/JAX
+# starts and therefore cannot be injected by a TrainConfig YAML. The flags below measured
+# -3.5% step time on 2026-09-04 (FSDP collective combining + pipelined collectives); see
+# docs/training-optimization.md for the complete A/B table and numerical validation.
+#   env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py
+#   env -u LD_LIBRARY_PATH \
 #   XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true \
 #     --xla_gpu_all_gather_combine_threshold_bytes=1073741824 \
 #     --xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824 \
@@ -8,7 +11,9 @@
 #     --xla_gpu_enable_pipelined_all_gather=true \
 #     --xla_gpu_enable_pipelined_reduce_scatter=true \
 #     --xla_gpu_enable_while_loop_double_buffering=true" \
-#   XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 python scripts/train.py <this-config> --exp-name=...
+#   XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+#   python scripts/train.py pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100 \
+#     --exp-name=<experiment-name>
 model:
   type: Pi0Config
   # Scanned 2026-09-04 over this dataset (4037 states incl. per-dim extremes) with
@@ -19,15 +24,20 @@ model:
   pi05: true
   discrete_state_input: true
   enable_training_time_rtc: true
+  # Encode all three visual views in one SigLIP call.
   batch_image_views: true
+  # Run all 18 Gemma layers through the validated FP16 cuDNN custom VJP. Leaving
+  # cudnn_attention_num_layers null means there is no partial-layer restriction.
   use_cudnn_attention: true
+  cudnn_attention_layer_start: 0
+  cudnn_attention_num_layers: null
   cudnn_attention_dtype: float16
+  # On 8xH100 80GB / global batch 256, weaker remat policies require either a
+  # 280 GiB allocation or an 84.6 GB residual. Full recomputation is the fastest
+  # policy that actually fits this shape.
+  gemma_remat_policy: nothing_saveable
+  siglip_remat_policy: nothing_saveable
 weight_loader:
-  # Was /home/li/hubo/xense-openpi/checkpoints/20000/params — a continuation of a
-  # 20k-step run that only existed on one laptop, so nobody else could start this
-  # config. Points at the upstream base now; to continue from your own checkpoint:
-  #   uv run scripts/train.py <this-config> --exp-name=... \
-  #     --weight_loader.params_path=/path/to/checkpoint/params
   type: CheckpointWeightLoader
   params_path: gs://openpi-assets/checkpoints/pi05_base/params
 ema_decay: null
@@ -40,8 +50,13 @@ data:
   default_prompt: Pick up each earbud case from the left stands, insert the matching earbuds, close the lid, and place the case in the box.
 batch_size: 256
 num_workers: 64
-num_train_steps: 60000
-save_interval: 5000
+num_train_steps: 80000
+# Whole-model grad/parameter norms are computed only on these logging steps.
+log_interval: 100
+# Fine-grained pipeline profiling adds timers and device synchronization; keep it off
+# for production now that shared-memory Tensor IPC is the default loader path.
+profile_data_pipeline: false
+save_interval: 10000
 keep_period: 10000
 fsdp_devices: 8
 strict_batch_order: false
diff --git a/configs/_examples/pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h200.yaml
similarity index 100%
rename from configs/_examples/pi05_base_bi_flexiv_earbud_case_insertion_0826_h200.yaml
rename to configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h200.yaml
diff --git a/configs/_examples/pi05_base_bi_flexiv_pack_6_cosmetic_bottles_lora.yaml b/configs/_examples/pi05_base_bi_flexiv_pack_6_cosmetic_bottles_lora.yaml
deleted file mode 100644
index d9dfebd..0000000
--- a/configs/_examples/pi05_base_bi_flexiv_pack_6_cosmetic_bottles_lora.yaml
+++ /dev/null
@@ -1,24 +0,0 @@
-model:
-  type: Pi0Config
-  max_token_len: 200
-  paligemma_variant: gemma_2b_lora
-  action_expert_variant: gemma_300m_lora
-  pi05: true
-  discrete_state_input: true
-  enable_training_time_rtc: true
-weight_loader:
-  type: CheckpointWeightLoader
-  params_path: gs://openpi-assets/checkpoints/pi05_base/params
-ema_decay: null
-data:
-  type: LeRobotBiFlexivDataConfig
-  repo_id: Xense/pack_6_cosmetic_bottles_into_carton
-  base_config:
-    prompt_from_task: true
-  default_prompt: Pick up six cosmetic bottles one by one and pack them into the carton
-    box. The box is narrow, so align each bottle carefully and insert it precisely.
-batch_size: 64
-num_train_steps: 20000
-# freeze_filter is intentionally absent. The loader sees the `lora` substring in the
-# model variants above and derives the flax filter via Pi0Config.get_freeze_filter();
-# see test_lora_freeze_filter_auto_derived in src/openpi/training/yaml_loader_test.py.
diff --git a/configs/_examples/pi05_droid.yaml b/configs/_examples/pi05_droid.yaml
deleted file mode 100644
index 5ad3728..0000000
--- a/configs/_examples/pi05_droid.yaml
+++ /dev/null
@@ -1,12 +0,0 @@
-model:
-  type: Pi0Config
-  action_horizon: 15
-  max_token_len: 200
-  pi05: true
-  discrete_state_input: true
-data:
-  type: DroidInferenceDataConfig
-  assets:
-    asset_id: droid
-  base_config:
-    prompt_from_task: true

From 72ace2a9061087489b34ad989a3712a46195ddd7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Sun, 6 Sep 2026 23:38:26 +0800
Subject: [PATCH 19/28] docs: tidy training optimization report

---
 docs/training-optimization.md | 133 +++++++++++++++++-----------------
 1 file changed, 68 insertions(+), 65 deletions(-)

diff --git a/docs/training-optimization.md b/docs/training-optimization.md
index 5bb5101..56f2005 100644
--- a/docs/training-optimization.md
+++ b/docs/training-optimization.md
@@ -14,13 +14,14 @@ H100 A/B），并与 commit 对齐。每个阶段给出：动机 → 实现（
 
 ## 0. 总账
 
-| 阶段 | 日期 | commit | 改动 | 端到端收益 |
-| --- | --- | --- | --- | --- |
-| 一 | 08-19 ~ 08-27 | `ec72f34` | 跳过触觉解码 / `in_order=False` / 无限 sampler | 消除周期性加载器停顿，生产 **≈1.14×** |
-| 二 | 08-27 ~ 08-28 | `9aef52d` `8ccf99c` | 共享内存 Tensor / XProf / 只在日志步算范数 / 三路相机合批 | step 时间 **−15.3%** |
-| 三 | 08-29 ~ 08-31 | `476babe`…`895811c` | cuDNN fused attention（收益 −29%，但**训练发散**） | 净收益 0，转为诊断 |
-| 四 | 09-03 ~ 09-04 | `df9755f` `3ea2b1c` | 定位根因 + fp16 内核（3000 步验收通过） | 显式 1.60 → fp16 **1.33 s/step** |
-| 五 | 09-04 | `cf07bbe` | `max_token_len` 200→128、XLA 集合通信标志 | 再 **−10.2%**，1.328 → 1.193 s/step |
+
+| 阶段 | 日期          | commit               | 改动                                                      | 端到端收益                             |
+| ---- | ------------- | -------------------- | --------------------------------------------------------- | -------------------------------------- |
+| 一   | 08-19 ~ 08-27 | `ec72f34`            | 跳过触觉解码 /`in_order=False` / 无限 sampler             | 消除周期性加载器停顿，生产**≈1.14×** |
+| 二   | 08-27 ~ 08-28 | `9aef52d` `8ccf99c`  | 共享内存 Tensor / XProf / 只在日志步算范数 / 三路相机合批 | step 时间**−15.3%**                   |
+| 三   | 08-29 ~ 08-31 | `476babe`…`895811c` | cuDNN fused attention（收益 −29%，但**训练发散**）       | 净收益 0，转为诊断                     |
+| 四   | 09-03 ~ 09-04 | `df9755f` `3ea2b1c`  | 定位根因 + fp16 内核（3000 步验收通过）                   | 显式 1.60 → fp16**1.33 s/step**       |
+| 五   | 09-04         | `cf07bbe`            | `max_token_len` 200→128、XLA 集合通信标志                | 再**−10.2%**，1.328 → 1.193 s/step   |
 
 **当前生产配置**（`configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml`）：
 
@@ -42,18 +43,19 @@ strict_batch_order: false     # 阶段一
 性能主线的 commit 对应关系如下。`3b2795e`、`3c5aac6` 只是夹在分支中的 Nova5 配置改动，
 与本文提速结论无关，故不纳入阶段收益。
 
-| 日期 | commit | 作用 |
-| --- | --- | --- |
-| 08-27 | `ec72f34` | 合入跳过触觉解码、乱序交付与无限 sampler；同时把 PyTorch 下限提高到支持 `in_order` 的 2.6 |
-| 08-27 | `9aef52d` | 共享内存 Tensor、零拷贝 NumPy 视图、显式关闭 DataLoader 与数据管线 profiling |
-| 08-28 | `8ccf99c` | XProf 开关、只在日志步计算范数、三路相机合批及其测试 |
-| 08-28 | `d0aec27` | 更新 H100 RTC 训练配置；没有单独的性能 A/B |
-| 08-29 | `476babe` `f42992e` | 引入 cuDNN fused attention，并记录最初的短跑/生产验证 |
-| 08-29 | `c1b70b2` `6305750` `51d9af3` | 处理全空 mask NaN；撤回逐步有限性保护；用 stop-gradient 保留快速原始 mask |
-| 08-31 | `ad08973` `c9e77cc` `895811c` | 建立 strict-order/分层诊断，定位长期发散，并第一次合并事故文档 |
-| 09-03 | `df9755f` | 增加 attention 精度与 Adam 更新空间诊断 |
-| 09-04 | `3ea2b1c` | 定位 D 项舍入机制，实现并验证 fp16 custom VJP |
-| 09-04 | `cf07bbe` | token 长度、XLA 集合通信和 remat A/B，形成 H100 最终配置 |
+
+| 日期  | commit                        | 作用                                                                                     |
+| ----- | ----------------------------- | ---------------------------------------------------------------------------------------- |
+| 08-27 | `ec72f34`                     | 合入跳过触觉解码、乱序交付与无限 sampler；同时把 PyTorch 下限提高到支持`in_order` 的 2.6 |
+| 08-27 | `9aef52d`                     | 共享内存 Tensor、零拷贝 NumPy 视图、显式关闭 DataLoader 与数据管线 profiling             |
+| 08-28 | `8ccf99c`                     | XProf 开关、只在日志步计算范数、三路相机合批及其测试                                     |
+| 08-28 | `d0aec27`                     | 更新 H100 RTC 训练配置；没有单独的性能 A/B                                               |
+| 08-29 | `476babe` `f42992e`           | 引入 cuDNN fused attention，并记录最初的短跑/生产验证                                    |
+| 08-29 | `c1b70b2` `6305750` `51d9af3` | 处理全空 mask NaN；撤回逐步有限性保护；用 stop-gradient 保留快速原始 mask                |
+| 08-31 | `ad08973` `c9e77cc` `895811c` | 建立 strict-order/分层诊断，定位长期发散，并第一次合并事故文档                           |
+| 09-03 | `df9755f`                     | 增加 attention 精度与 Adam 更新空间诊断                                                  |
+| 09-04 | `3ea2b1c`                     | 定位 D 项舍入机制，实现并验证 fp16 custom VJP                                            |
+| 09-04 | `cf07bbe`                     | token 长度、XLA 集合通信和 remat A/B，形成 H100 最终配置                                 |
 
 ---
 
@@ -104,11 +106,12 @@ buffered    = _task_info 里已完成的结果  算完了、但还轮不到它
 
 **这两个数把「算不出来」和「算出来了但过不去」区分开——是整件事里最值钱的一步。**
 
-| batch | wait | outstanding | buffered |
-| --- | --- | --- | --- |
-| 64 | 3.02 s | 1024 | 10 |
-| 66 | 9.84 s | 1024 | **41** |
-| 130 | 9.66 s | 1024 | **41** |
+
+| batch | wait   | outstanding | buffered |
+| ----- | ------ | ----------- | -------- |
+| 64    | 3.02 s | 1024        | 10       |
+| 66    | 9.84 s | 1024        | **41**   |
+| 130   | 9.66 s | 1024        | **41**   |
 
 `outstanding` 恒为 1024（`prefetch_factor 16 × 64 workers`，管线一直是满的），停顿时 `buffered` 高达 41
 ——**41 个 batch 已经算完躺在那儿，训练循环却在等**。torch 严格按 sampler 顺序交付，任务 *j* 派给
@@ -150,12 +153,13 @@ batch 序列从来就不可恢复。
 
 受控 400 步真实训练 A/B（同机同配置，唯一变量是 B1+B2）：
 
-| 指标 | 修复前 | 修复后 | 倍数 |
-| --- | --- | --- | --- |
-| 400 步总时长 | 770.8 s | **715.7 s** | **1.08×** |
-| 停顿步数（>0.5 s） | **20** | **0** | — |
-| 最大 `next_batch` | 7.83 s | **0.43 s** | 18× |
-| 中位 `next_batch` | 0.030 s | 0.290 s | — |
+
+| 指标               | 修复前  | 修复后      | 倍数       |
+| ------------------ | ------- | ----------- | ---------- |
+| 400 步总时长       | 770.8 s | **715.7 s** | **1.08×** |
+| 停顿步数（>0.5 s） | **20**  | **0**       | —         |
+| 最大`next_batch`   | 7.83 s  | **0.43 s**  | 18×       |
+| 中位`next_batch`   | 0.030 s | 0.290 s     | —         |
 
 > 中位数从 0.030 涨到 0.290 s **不是退步**：同样的产能缺口，以前攒成每 64 步一次的 6 秒尖峰，
 > 现在均摊到每一步，总等待反而少 16.8 s。
@@ -340,10 +344,6 @@ def _stop_gradient_for_fully_masked_queries(q, attn_mask):
 `6305750` 恢复原训练更新语义、移除这项逐步保护，`51d9af3` 再把 dummy-key 换成上面的 stop-gradient，
 这才是后续诊断和当前代码使用的实现。
 
-checkpoint 处置：原 step 16000 已被保留策略删除，step 20000–55000 全部受污染并删除，
-只能从 step 15000 重来（重跑约 1000 个有效 step）。**含 NaN 的参数或 Adam mu/nu 不能通过更换
-attention 实现恢复。**
-
 ### 3.3 第二类失效：全层 cuDNN 在约 1000 步后缓慢发散（08-30 ~ 08-31，commit `ad08973` `895811c`）
 
 修完全空 mask 后，全 18 层 cuDNN 前约 1000 步与显式基线贴合，随后分叉：
@@ -614,16 +614,17 @@ D 项误差在范数上只占 ~1e-3，危害来自方向的结构性（同一低
 主指标是日志 `Step 100` 到 `Step 300` 的时间戳之差除以 200（含日志开销的端到端墙钟）；
 `loss@300` 用来确认数值没变。
 
-| 变体 | s/step | 相对基线 | loss@300 | 结论 |
-| --- | --- | --- | --- | --- |
-| E0 基线 | 1.328 | — | 0.7356 | 等数据 0.03 s，数据侧不是瓶颈 |
-| E1 `max_token_len` 200→128 | 1.246 | **−6.2%** | 0.7356 | 采纳 |
-| E2 关闭 remat | OOM | — | — | XLA 试图分配 280 GiB |
-| E2b gemma remat `dots_with_no_batch_dims_saveable` | 编译失败 | — | — | cuDNN 分片器报 q/k/v 分片不一致 |
-| E2c 同上 + q/k/v 显式分片约束 | OOM | — | — | 编译过了，单块残差 84.6 GB |
-| E3 XLA 集合通信标志 | 1.281 | **−3.5%** | 0.7357 | 采纳 |
-| E4 `fsdp_devices` 8→4 | 1.335 | +0.5% | 0.7354 | 无收益 |
-| **E5 = E1 + E3** | **1.193** | **−10.2%** | 0.7356 | 生产采用；192.8 → 214.6 samples/s |
+
+| 变体                                              | s/step    | 相对基线    | loss@300 | 结论                               |
+| ------------------------------------------------- | --------- | ----------- | -------- | ---------------------------------- |
+| E0 基线                                           | 1.328     | —          | 0.7356   | 等数据 0.03 s，数据侧不是瓶颈      |
+| E1`max_token_len` 200→128                        | 1.246     | **−6.2%**  | 0.7356   | 采纳                               |
+| E2 关闭 remat                                     | OOM       | —          | —       | XLA 试图分配 280 GiB               |
+| E2b gemma remat`dots_with_no_batch_dims_saveable` | 编译失败  | —          | —       | cuDNN 分片器报 q/k/v 分片不一致    |
+| E2c 同上 + q/k/v 显式分片约束                     | OOM       | —          | —       | 编译过了，单块残差 84.6 GB         |
+| E3 XLA 集合通信标志                               | 1.281     | **−3.5%**  | 0.7357   | 采纳                               |
+| E4`fsdp_devices` 8→4                             | 1.335     | +0.5%       | 0.7354   | 无收益                             |
+| **E5 = E1 + E3**                                  | **1.193** | **−10.2%** | 0.7356   | 生产采用；192.8 → 214.6 samples/s |
 
 60000 步按稳态估算：22.1 h → 19.9 h，省约 2.2 h。
 
@@ -746,19 +747,20 @@ def _cudnn_attention_call(q, k, v, attn_mask):
 
 **诊断脚本**
 
-| 脚本 | 用途 |
-| --- | --- |
-| `scripts/check_cuda_stack.py` | 枚举 `/proc/self/maps` 里实际加载的 cuDNN/cuBLAS/NCCL，+ 真实形状数值 smoke test |
-| `scripts/token_len_scan.py` | 走真实 transform 链扫数据集的 tokenized prompt 长度分布 |
-| `scripts/attention_dterm_probe.py` | 抓真实激活 + 单层五变体误差 + H1/H2 预测（阶段四主力） |
-| `scripts/attention_dterm_cpu_analysis.py` | peaked 行的 float64 逐行解剖 |
-| `scripts/cudnn_fp16_vjp_equivalence.py` | fp16 VJP v2（直接调规则）与 v1（反向重算）的等价性 |
-| `scripts/cudnn_attention_microbench.py` | 单层 attention 前向/反向微基准 |
-| `scripts/attention_precision_update_probe.py` | 真实 Adam 状态上的更新空间对比（**注意 §4.6：该指标不区分好坏**） |
-| `scripts/grad_ab_probe.py` | 完整模型显式/cuDNN 梯度 A/B |
-| `scripts/grad_ab_direction_probe.py` | 多 batch 差分方向和随机分量 |
-| `scripts/hybrid_update_ab_probe.py` | 不同 cuDNN 层段的更新空间误差 |
-| `scripts/optimizer_update_ab_probe.py` | 从真实 Adam 状态比较参数更新误差 |
+
+| 脚本                                          | 用途                                                                            |
+| --------------------------------------------- | ------------------------------------------------------------------------------- |
+| `scripts/check_cuda_stack.py`                 | 枚举`/proc/self/maps` 里实际加载的 cuDNN/cuBLAS/NCCL，+ 真实形状数值 smoke test |
+| `scripts/token_len_scan.py`                   | 走真实 transform 链扫数据集的 tokenized prompt 长度分布                         |
+| `scripts/attention_dterm_probe.py`            | 抓真实激活 + 单层五变体误差 + H1/H2 预测（阶段四主力）                          |
+| `scripts/attention_dterm_cpu_analysis.py`     | peaked 行的 float64 逐行解剖                                                    |
+| `scripts/cudnn_fp16_vjp_equivalence.py`       | fp16 VJP v2（直接调规则）与 v1（反向重算）的等价性                              |
+| `scripts/cudnn_attention_microbench.py`       | 单层 attention 前向/反向微基准                                                  |
+| `scripts/attention_precision_update_probe.py` | 真实 Adam 状态上的更新空间对比（**注意 §4.6：该指标不区分好坏**）              |
+| `scripts/grad_ab_probe.py`                    | 完整模型显式/cuDNN 梯度 A/B                                                     |
+| `scripts/grad_ab_direction_probe.py`          | 多 batch 差分方向和随机分量                                                     |
+| `scripts/hybrid_update_ab_probe.py`           | 不同 cuDNN 层段的更新空间误差                                                   |
+| `scripts/optimizer_update_ab_probe.py`        | 从真实 Adam 状态比较参数更新误差                                                |
 
 **诊断配置**：`configs/diag_cudnn_strict_order.yaml`（全层 bf16，会发散）、
 `configs/diag_explicit_strict_order.yaml`（基线）、`configs/diag_cudnn_last3.yaml`（后 3 层折中）、
@@ -766,14 +768,15 @@ def _cudnn_attention_call(q, k, v, attn_mask):
 
 **代码位置**
 
-| 文件 | 相关改动 |
-| --- | --- |
-| `src/openpi/training/data_loader.py` | `SelectiveVideoLeRobotDataset`、`_resolve_decode_video_keys`、`InfiniteSampler`、`in_order` 接线、Torch Tensor collate、`close()` |
-| `src/openpi/training/config.py` | `DataConfig.tactile`、`TrainConfig.strict_batch_order` |
-| `src/openpi/models/pi0_config.py` | `batch_image_views`、`use_cudnn_attention`、`cudnn_attention_{layer_start,num_layers,dtype}`、`gemma_remat_policy`、`siglip_remat_policy` |
-| `src/openpi/models/pi0.py` | `_encode_image_views_as_batch` |
-| `src/openpi/models/gemma.py` | `_stop_gradient_for_fully_masked_queries`、`_cudnn_attention_call`、`_cudnn_attention_in_dtype`、`_cudnn_static_args`、逐层 cuDNN 开关 |
-| `scripts/train.py` | `compute_metrics` 静态分支、XProf trace 开关、数据管线 profiling |
+
+| 文件                                 | 相关改动                                                                                                                                  |
+| ------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------- |
+| `src/openpi/training/data_loader.py` | `SelectiveVideoLeRobotDataset`、`_resolve_decode_video_keys`、`InfiniteSampler`、`in_order` 接线、Torch Tensor collate、`close()`         |
+| `src/openpi/training/config.py`      | `DataConfig.tactile`、`TrainConfig.strict_batch_order`                                                                                    |
+| `src/openpi/models/pi0_config.py`    | `batch_image_views`、`use_cudnn_attention`、`cudnn_attention_{layer_start,num_layers,dtype}`、`gemma_remat_policy`、`siglip_remat_policy` |
+| `src/openpi/models/pi0.py`           | `_encode_image_views_as_batch`                                                                                                            |
+| `src/openpi/models/gemma.py`         | `_stop_gradient_for_fully_masked_queries`、`_cudnn_attention_call`、`_cudnn_attention_in_dtype`、`_cudnn_static_args`、逐层 cuDNN 开关    |
+| `scripts/train.py`                   | `compute_metrics` 静态分支、XProf trace 开关、数据管线 profiling                                                                          |
 
 ---
 

From 95785a7440d59ed33a7f9c2029c5d51ecb8a6b34 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Mon, 7 Sep 2026 17:24:22 +0800
Subject: [PATCH 20/28] policies: disable cuDNN attention for JAX inference

Pi0/Pi05 checkpoints trained with use_cudnn_attention now load with the flag
cleared. The cuDNN path also runs during lazy initialization and prefix prefill,
where no KV cache exists yet, and JAX 0.5.3 only permits head_dim=256 cuDNN
attention on Hopper -- so serving such a checkpoint on a non-Hopper GPU failed
at trace time.

The attention backend does not change the checkpoint parameter structure, and
the denoising steps already ran on the explicit path because they carry a KV
cache, so only the single prefill call changes.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 src/openpi/policies/policy_config.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/openpi/policies/policy_config.py b/src/openpi/policies/policy_config.py
index 6570df0..1f97ef1 100644
--- a/src/openpi/policies/policy_config.py
+++ b/src/openpi/policies/policy_config.py
@@ -1,3 +1,4 @@
+import dataclasses
 import logging
 import os
 import pathlib
@@ -5,6 +6,7 @@
 
 import jax.numpy as jnp
 
+from openpi.models import pi0_config
 import openpi.models.model as _model
 import openpi.policies.policy as _policy
 import openpi.shared.download as download
@@ -41,6 +43,8 @@ def create_trained_policy(
     Note:
         The function automatically detects whether the model is PyTorch-based by checking for the
         presence of "model.safensors" in the checkpoint directory.
+        JAX Pi0/Pi05 policies use explicit attention for inference, even when the training
+        config enables cuDNN attention, so checkpoints can be served on non-Hopper GPUs.
     """
     repack_transforms = repack_transforms or transforms.Group()
     checkpoint_dir = download.maybe_download(str(checkpoint_dir))
@@ -54,7 +58,15 @@ def create_trained_policy(
         model = train_config.model.load_pytorch(train_config, weight_path)
         model.paligemma_with_expert.to_bfloat16_for_selected_params("bfloat16")
     else:
-        model = train_config.model.load(_model.restore_params(checkpoint_dir / "params", dtype=jnp.bfloat16))
+        model_config = train_config.model
+        if isinstance(model_config, pi0_config.Pi0Config) and model_config.use_cudnn_attention:
+            # This training optimization also runs during lazy initialization and prefix
+            # prefill (before a KV cache exists). JAX 0.5.3 only permits head_dim=256
+            # cuDNN attention on Hopper, so disable it before load() traces the model.
+            # The attention backend does not change the checkpoint parameter structure.
+            logging.info("Using explicit attention for inference (disabling training-time cuDNN attention).")
+            model_config = dataclasses.replace(model_config, use_cudnn_attention=False)
+        model = model_config.load(_model.restore_params(checkpoint_dir / "params", dtype=jnp.bfloat16))
     data_config = train_config.data.create(train_config.assets_dirs, train_config.model)
     if norm_stats is None:
         # We are loading the norm stats from the checkpoint instead of the config assets dir to make sure

From e13ac0fcaaeb9a71d8c754a2f713fedf7ad4856c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Thu, 10 Sep 2026 17:28:01 +0800
Subject: [PATCH 21/28] Clean up training diagnostics and refresh optimized
 configs

---
 README.md                                     |   4 +-
 ...exiv_cube_color_sorting_rtc_0829_h100.yaml |  33 +-
 ...s_case_insertion_teleop_rtc_0904_h100.yaml |  17 +-
 ...se_bi_flexiv_optical_module_0611_h100.yaml |  21 +
 ..._base_xtac_umi_pick_up_cube_0807_h200.yaml |  42 -
 configs/diag_cudnn_b1_099.yaml                |  45 -
 configs/diag_cudnn_last3.yaml                 |  39 -
 configs/diag_cudnn_strict_order.yaml          |  37 -
 configs/diag_explicit_strict_order.yaml       |  37 -
 docs/training-optimization.md                 | 840 ++----------------
 scripts/attention_dterm_cpu_analysis.py       | 176 ----
 scripts/attention_dterm_probe.py              | 449 ----------
 scripts/attention_precision_update_probe.py   | 219 -----
 scripts/check_cuda_stack.py                   |  28 +-
 scripts/cudnn_attention_microbench.py         | 115 ---
 scripts/cudnn_fp16_vjp_equivalence.py         | 112 ---
 scripts/grad_ab_direction_probe.py            | 257 ------
 scripts/grad_ab_probe.py                      | 162 ----
 scripts/hybrid_update_ab_probe.py             | 192 ----
 scripts/optimizer_update_ab_probe.py          | 194 ----
 scripts/token_len_scan.py                     | 115 ---
 scripts/train.py                              | 106 +--
 scripts/train_test.py                         |  30 -
 src/openpi/models/gemma.py                    |  61 +-
 src/openpi/models/pi0.py                      |   5 -
 src/openpi/models/pi0_config.py               |  19 -
 src/openpi/models/siglip.py                   |  15 +-
 src/openpi/policies/dobot_nova5_policy.py     |   1 -
 src/openpi/training/config.py                 |  20 -
 src/openpi/training/data_loader.py            |  73 +-
 src/openpi/training/data_loader_test.py       |  26 -
 31 files changed, 165 insertions(+), 3325 deletions(-)
 create mode 100644 configs/_examples/pi05_base_bi_flexiv_optical_module_0611_h100.yaml
 delete mode 100644 configs/_examples/pi05_base_xtac_umi_pick_up_cube_0807_h200.yaml
 delete mode 100644 configs/diag_cudnn_b1_099.yaml
 delete mode 100644 configs/diag_cudnn_last3.yaml
 delete mode 100644 configs/diag_cudnn_strict_order.yaml
 delete mode 100644 configs/diag_explicit_strict_order.yaml
 delete mode 100644 scripts/attention_dterm_cpu_analysis.py
 delete mode 100644 scripts/attention_dterm_probe.py
 delete mode 100644 scripts/attention_precision_update_probe.py
 delete mode 100644 scripts/cudnn_attention_microbench.py
 delete mode 100644 scripts/cudnn_fp16_vjp_equivalence.py
 delete mode 100644 scripts/grad_ab_direction_probe.py
 delete mode 100644 scripts/grad_ab_probe.py
 delete mode 100644 scripts/hybrid_update_ab_probe.py
 delete mode 100644 scripts/optimizer_update_ab_probe.py
 delete mode 100644 scripts/token_len_scan.py

diff --git a/README.md b/README.md
index 119a64a..ec60af2 100644
--- a/README.md
+++ b/README.md
@@ -188,8 +188,8 @@ model:
   paligemma_variant: gemma_2b
   action_expert_variant: gemma_300m
   enable_training_time_rtc: true
-  # Experimental and currently unsafe for production; keep disabled until the
-  # step-0 2000-step convergence gate documented below has passed.
+  # Explicit attention is the default. For validated FP16 cuDNN training,
+  # see the configuration and validation notes below.
   use_cudnn_attention: false
   max_delay: 10
 
diff --git a/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml
index c15b505..e912566 100644
--- a/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml
@@ -1,17 +1,35 @@
+# Production launch (8xH100, hb-openpi env). XLA_FLAGS must be set before Python/JAX
+# starts and therefore cannot be injected by a TrainConfig YAML. The flags below measured
+# -3.5% step time on 2026-09-04 (FSDP collective combining + pipelined collectives); see
+# docs/training-optimization.md for the performance results and numerical validation.
+#   env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py
+#   env -u LD_LIBRARY_PATH \
+#   XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true \
+#     --xla_gpu_all_gather_combine_threshold_bytes=1073741824 \
+#     --xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824 \
+#     --xla_gpu_all_reduce_combine_threshold_bytes=1073741824 \
+#     --xla_gpu_enable_pipelined_all_gather=true \
+#     --xla_gpu_enable_pipelined_reduce_scatter=true \
+#     --xla_gpu_enable_while_loop_double_buffering=true" \
+#   XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+#   python scripts/train.py pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100 \
+#     --exp-name=<experiment-name>
 model:
   type: Pi0Config
-  max_token_len: 200
+  # Scanned 2026-09-04 over this dataset (4037 states incl. per-dim extremes) with
+  # The longest tokenized "Task: ..., State: ...;\nAction: " prompt
+  # is 116 tokens, so 128 leaves 12 spare. Down from 200: -6.2% step time, identical loss
+  # (padding tokens were masked anyway). Re-scan when the dataset or prompt changes.
+  max_token_len: 128
   pi05: true
   discrete_state_input: true
   enable_training_time_rtc: true
+  # Encode all three visual views in one SigLIP call.
   batch_image_views: true
+  # Run all 18 Gemma layers through the validated FP16 cuDNN custom VJP.
   use_cudnn_attention: true
+  cudnn_attention_dtype: float16
 weight_loader:
-  # Was /home/li/hubo/xense-openpi/checkpoints/20000/params — a continuation of a
-  # 20k-step run that only existed on one laptop, so nobody else could start this
-  # config. Points at the upstream base now; to continue from your own checkpoint:
-  #   uv run scripts/train.py <this-config> --exp-name=... \
-  #     --weight_loader.params_path=/path/to/checkpoint/params
   type: CheckpointWeightLoader
   params_path: gs://openpi-assets/checkpoints/pi05_base/params
 ema_decay: null
@@ -25,7 +43,10 @@ data:
 batch_size: 256
 num_workers: 64
 num_train_steps: 60000
+# Whole-model grad/parameter norms are computed only on these logging steps.
+log_interval: 100
 save_interval: 10000
 keep_period: 10000
 fsdp_devices: 8
 strict_batch_order: false
+project_name: xense-openpi-hubo
diff --git a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
index 6ee2686..c8eb631 100644
--- a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
@@ -1,7 +1,7 @@
 # Production launch (8xH100, hb-openpi env). XLA_FLAGS must be set before Python/JAX
 # starts and therefore cannot be injected by a TrainConfig YAML. The flags below measured
 # -3.5% step time on 2026-09-04 (FSDP collective combining + pipelined collectives); see
-# docs/training-optimization.md for the complete A/B table and numerical validation.
+# docs/training-optimization.md for the performance results and numerical validation.
 #   env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py
 #   env -u LD_LIBRARY_PATH \
 #   XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true \
@@ -17,7 +17,7 @@
 model:
   type: Pi0Config
   # Scanned 2026-09-04 over this dataset (4037 states incl. per-dim extremes) with
-  # scripts/token_len_scan.py: the longest tokenized "Task: ..., State: ...;\nAction: " prompt
+  # The longest tokenized "Task: ..., State: ...;\nAction: " prompt
   # is 116 tokens, so 128 leaves 12 spare. Down from 200: -6.2% step time, identical loss
   # (padding tokens were masked anyway). Re-scan when the dataset or prompt changes.
   max_token_len: 128
@@ -26,17 +26,9 @@ model:
   enable_training_time_rtc: true
   # Encode all three visual views in one SigLIP call.
   batch_image_views: true
-  # Run all 18 Gemma layers through the validated FP16 cuDNN custom VJP. Leaving
-  # cudnn_attention_num_layers null means there is no partial-layer restriction.
+  # Run all 18 Gemma layers through the validated FP16 cuDNN custom VJP.
   use_cudnn_attention: true
-  cudnn_attention_layer_start: 0
-  cudnn_attention_num_layers: null
   cudnn_attention_dtype: float16
-  # On 8xH100 80GB / global batch 256, weaker remat policies require either a
-  # 280 GiB allocation or an 84.6 GB residual. Full recomputation is the fastest
-  # policy that actually fits this shape.
-  gemma_remat_policy: nothing_saveable
-  siglip_remat_policy: nothing_saveable
 weight_loader:
   type: CheckpointWeightLoader
   params_path: gs://openpi-assets/checkpoints/pi05_base/params
@@ -53,9 +45,6 @@ num_workers: 64
 num_train_steps: 80000
 # Whole-model grad/parameter norms are computed only on these logging steps.
 log_interval: 100
-# Fine-grained pipeline profiling adds timers and device synchronization; keep it off
-# for production now that shared-memory Tensor IPC is the default loader path.
-profile_data_pipeline: false
 save_interval: 10000
 keep_period: 10000
 fsdp_devices: 8
diff --git a/configs/_examples/pi05_base_bi_flexiv_optical_module_0611_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_optical_module_0611_h100.yaml
new file mode 100644
index 0000000..cfe9b36
--- /dev/null
+++ b/configs/_examples/pi05_base_bi_flexiv_optical_module_0611_h100.yaml
@@ -0,0 +1,21 @@
+model:
+  type: Pi0Config
+  max_token_len: 200
+  pi05: true
+  discrete_state_input: true
+  enable_training_time_rtc: true
+weight_loader:
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+ema_decay: null
+data:
+  type: LeRobotBiFlexivDataConfig
+  repo_id: Xense/optical-module-insertion-0731
+  base_config:
+    prompt_from_task: true
+  default_prompt: Pick up the bag, open it, inspect its contents, close it, and place
+    it on the opposite side.
+batch_size: 256
+num_workers: 64
+num_train_steps: 60000
+fsdp_devices: 8
diff --git a/configs/_examples/pi05_base_xtac_umi_pick_up_cube_0807_h200.yaml b/configs/_examples/pi05_base_xtac_umi_pick_up_cube_0807_h200.yaml
deleted file mode 100644
index 3fd461d..0000000
--- a/configs/_examples/pi05_base_xtac_umi_pick_up_cube_0807_h200.yaml
+++ /dev/null
@@ -1,42 +0,0 @@
-# XTac-UMI bimanual dataset, recorded on the handheld bi_taccap_gripper rig.
-#
-# State/action are 20D in the rig's own per-side-grouped feature order:
-#   [left_tcp(0-8), left_gripper(9), right_tcp(10-18), right_gripper(19)]
-# No head camera: base_0_rgb is black with image_mask=False; the two wrist
-# cameras stay in their matching wrist slots. The dataset also carries four
-# tactile streams (left/right x tactile_left/tactile_right) that this config
-# does not consume - pi05 takes wrist RGB and the 20D vector only.
-#
-# Deploy target: bi_flexiv (dual Flexiv Rizon4 RT with taccap grippers), which
-# assembles the same per-side-grouped order — see
-# examples/xtac_umi_bi_flexiv_rizon4_rt and src/openpi/policies/xtac_umi_policy.py.
-model:
-  type: Pi0Config
-  max_token_len: 200
-  pi05: true
-  discrete_state_input: true
-  # enable_training_time_rtc: true
-weight_loader:
-  type: CheckpointWeightLoader
-  params_path: gs://openpi-assets/checkpoints/pi05_base/params
-ema_decay: null
-data:
-  type: LeRobotXtacUmiDataConfig
-  repo_id: TacVerse/taccap-g1-pick-up-cube-0807
-  use_delta_cartesian_actions: true
-  # use_head_camera: true   # only for a bench that actually recorded one
-  base_config:
-    prompt_from_task: true
-  # Only a fallback: prompt_from_task above means the dataset's own task string
-  # is what actually reaches the model.
-  default_prompt: Use the left arm to pick the red cube and place it into the blue
-    plastic box first, then use the right arm to grasp the green cube and put it into
-    the same blue plastic box after the first action is fully finished
-batch_size: 256
-num_workers: 64
-# 150 episodes / 43447 frames, i.e. ~290 frames (~9.7 s) per episode - roughly
-# half the frames of the 0710 sort-defective-parts set. 30k steps at batch 256 is
-# ~177 passes over the data, in the same range as the other single-task configs
-# here; raise it only if the loss is still moving.
-num_train_steps: 30000
-fsdp_devices: 8
diff --git a/configs/diag_cudnn_b1_099.yaml b/configs/diag_cudnn_b1_099.yaml
deleted file mode 100644
index f8108f3..0000000
--- a/configs/diag_cudnn_b1_099.yaml
+++ /dev/null
@@ -1,45 +0,0 @@
-model:
-  type: Pi0Config
-  max_token_len: 200
-  pi05: true
-  discrete_state_input: true
-  enable_training_time_rtc: true
-  batch_image_views: true
-  use_cudnn_attention: true
-
-weight_loader:
-  type: CheckpointWeightLoader
-  params_path: gs://openpi-assets/checkpoints/pi05_base/params
-
-optimizer:
-  type: AdamW
-  b1: 0.99
-  b2: 0.95
-  eps: 1.0e-8
-  weight_decay: 1.0e-10
-  clip_gradient_norm: 1.0
-
-ema_decay: null
-
-data:
-  type: LeRobotBiFlexivDataConfig
-  repo_id: Xense/cube-color-sorting-0821
-  assets:
-    assets_dir: /root/localstorage/xense-openpi/assets/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100
-    asset_id: Xense/cube-color-sorting-0821
-  base_config:
-    prompt_from_task: true
-    tactile: false
-  default_prompt: Put the green cubes in the right bin and the yellow cubes in the left bin.
-
-batch_size: 256
-num_workers: 64
-num_train_steps: 3000
-log_interval: 100
-save_interval: 100000
-save_final_checkpoint: false
-keep_period: null
-fsdp_devices: 8
-strict_batch_order: true
-wandb_enabled: false
-overwrite: true
diff --git a/configs/diag_cudnn_last3.yaml b/configs/diag_cudnn_last3.yaml
deleted file mode 100644
index c1025bb..0000000
--- a/configs/diag_cudnn_last3.yaml
+++ /dev/null
@@ -1,39 +0,0 @@
-model:
-  type: Pi0Config
-  max_token_len: 200
-  pi05: true
-  discrete_state_input: true
-  enable_training_time_rtc: true
-  batch_image_views: true
-  use_cudnn_attention: true
-  cudnn_attention_layer_start: 15
-  cudnn_attention_num_layers: 3
-
-weight_loader:
-  type: CheckpointWeightLoader
-  params_path: gs://openpi-assets/checkpoints/pi05_base/params
-
-ema_decay: null
-
-data:
-  type: LeRobotBiFlexivDataConfig
-  repo_id: Xense/cube-color-sorting-0821
-  assets:
-    assets_dir: /root/localstorage/xense-openpi/assets/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100
-    asset_id: Xense/cube-color-sorting-0821
-  base_config:
-    prompt_from_task: true
-    tactile: false
-  default_prompt: Put the green cubes in the right bin and the yellow cubes in the left bin.
-
-batch_size: 256
-num_workers: 64
-num_train_steps: 3000
-log_interval: 100
-save_interval: 1000000
-save_final_checkpoint: false
-keep_period: 1000
-fsdp_devices: 8
-strict_batch_order: true
-wandb_enabled: false
-overwrite: true
diff --git a/configs/diag_cudnn_strict_order.yaml b/configs/diag_cudnn_strict_order.yaml
deleted file mode 100644
index 7c2b3d7..0000000
--- a/configs/diag_cudnn_strict_order.yaml
+++ /dev/null
@@ -1,37 +0,0 @@
-model:
-  type: Pi0Config
-  max_token_len: 200
-  pi05: true
-  discrete_state_input: true
-  enable_training_time_rtc: true
-  batch_image_views: true
-  use_cudnn_attention: true
-
-weight_loader:
-  type: CheckpointWeightLoader
-  params_path: gs://openpi-assets/checkpoints/pi05_base/params
-
-ema_decay: null
-
-data:
-  type: LeRobotBiFlexivDataConfig
-  repo_id: Xense/cube-color-sorting-0821
-  assets:
-    assets_dir: /root/localstorage/xense-openpi/assets/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100
-    asset_id: Xense/cube-color-sorting-0821
-  base_config:
-    prompt_from_task: true
-    tactile: false
-  default_prompt: Put the green cubes in the right bin and the yellow cubes in the left bin.
-
-batch_size: 256
-num_workers: 64
-num_train_steps: 3000
-log_interval: 100
-save_interval: 1000
-save_final_checkpoint: false
-keep_period: 1000
-fsdp_devices: 8
-strict_batch_order: true
-wandb_enabled: false
-overwrite: true
diff --git a/configs/diag_explicit_strict_order.yaml b/configs/diag_explicit_strict_order.yaml
deleted file mode 100644
index d3e70a0..0000000
--- a/configs/diag_explicit_strict_order.yaml
+++ /dev/null
@@ -1,37 +0,0 @@
-model:
-  type: Pi0Config
-  max_token_len: 200
-  pi05: true
-  discrete_state_input: true
-  enable_training_time_rtc: true
-  batch_image_views: true
-  use_cudnn_attention: false
-
-weight_loader:
-  type: CheckpointWeightLoader
-  params_path: gs://openpi-assets/checkpoints/pi05_base/params
-
-ema_decay: null
-
-data:
-  type: LeRobotBiFlexivDataConfig
-  repo_id: Xense/cube-color-sorting-0821
-  assets:
-    assets_dir: /root/localstorage/xense-openpi/assets/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100
-    asset_id: Xense/cube-color-sorting-0821
-  base_config:
-    prompt_from_task: true
-    tactile: false
-  default_prompt: Put the green cubes in the right bin and the yellow cubes in the left bin.
-
-batch_size: 256
-num_workers: 64
-num_train_steps: 3000
-log_interval: 100
-save_interval: 100000
-save_final_checkpoint: false
-keep_period: null
-fsdp_devices: 8
-strict_batch_order: true
-wandb_enabled: false
-overwrite: true
diff --git a/docs/training-optimization.md b/docs/training-optimization.md
index 56f2005..93cab98 100644
--- a/docs/training-optimization.md
+++ b/docs/training-optimization.md
@@ -1,800 +1,104 @@
-# openpi 训练提速与 cuDNN Attention 数值事故：完整记录
+# 训练优化与 cuDNN attention
 
-**覆盖时间**：2026-08-19 → 2026-09-04
-**分支**：`optimize-training-duration`（性能主线 commit `ec72f34` → `cf07bbe`；其中
-`ec72f34` 已进入本分支基线 `main`，分支新增优化从 `9aef52d` 开始）
-**机器**：8×H200、8×H100 80GB（ebcloud）
-**模型**：Pi0.5（`pi05_base` 微调，BF16 GQA，18 层 Gemma）
+本文保留训练配置、实现约束和历史验证结论。详细实验脚本及排障过程可从本分支
+`72ace2a` 及更早的提交中查阅；下述性能数字来自历史实验，并非每台机器上的保证。
 
-本文按时间线合并了原先分散的 5 份记录（数据管线、H200 优化、cuDNN NaN 事故、D 项根因分析、
-H100 A/B），并与 commit 对齐。每个阶段给出：动机 → 实现（关键代码）→ 实测 → 结论。
-被推翻的中间结论保留在「弯路」小节里，避免后来人重走。
+## 生产配置
 
----
-
-## 0. 总账
-
-
-| 阶段 | 日期          | commit               | 改动                                                      | 端到端收益                             |
-| ---- | ------------- | -------------------- | --------------------------------------------------------- | -------------------------------------- |
-| 一   | 08-19 ~ 08-27 | `ec72f34`            | 跳过触觉解码 /`in_order=False` / 无限 sampler             | 消除周期性加载器停顿，生产**≈1.14×** |
-| 二   | 08-27 ~ 08-28 | `9aef52d` `8ccf99c`  | 共享内存 Tensor / XProf / 只在日志步算范数 / 三路相机合批 | step 时间**−15.3%**                   |
-| 三   | 08-29 ~ 08-31 | `476babe`…`895811c` | cuDNN fused attention（收益 −29%，但**训练发散**）       | 净收益 0，转为诊断                     |
-| 四   | 09-03 ~ 09-04 | `df9755f` `3ea2b1c`  | 定位根因 + fp16 内核（3000 步验收通过）                   | 显式 1.60 → fp16**1.33 s/step**       |
-| 五   | 09-04         | `cf07bbe`            | `max_token_len` 200→128、XLA 集合通信标志                | 再**−10.2%**，1.328 → 1.193 s/step   |
-
-**当前生产配置**（`configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml`）：
+Pi0.5 的优化示例见
+[`pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml`](../configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml)。
 
 ```yaml
 model:
-  max_token_len: 128          # 阶段五，按数据集实测收紧
-  batch_image_views: true     # 阶段二
-  use_cudnn_attention: true   # 阶段三
-  cudnn_attention_dtype: float16   # 阶段四：必须 float16，bfloat16 会发散
-data:
-  base_config:
-    tactile: false            # 阶段一
-strict_batch_order: false     # 阶段一
-```
-
-配合 `XLA_FLAGS`（见 §5.2）。上表的测试跨越不同机器、batch 和任务，**不能把各阶段百分比直接连乘成
-一个受控的总提速数字**；可复现的最终同配置结论是阶段五的 1.328 → 1.193 s/step。
-
-性能主线的 commit 对应关系如下。`3b2795e`、`3c5aac6` 只是夹在分支中的 Nova5 配置改动，
-与本文提速结论无关，故不纳入阶段收益。
-
-
-| 日期  | commit                        | 作用                                                                                     |
-| ----- | ----------------------------- | ---------------------------------------------------------------------------------------- |
-| 08-27 | `ec72f34`                     | 合入跳过触觉解码、乱序交付与无限 sampler；同时把 PyTorch 下限提高到支持`in_order` 的 2.6 |
-| 08-27 | `9aef52d`                     | 共享内存 Tensor、零拷贝 NumPy 视图、显式关闭 DataLoader 与数据管线 profiling             |
-| 08-28 | `8ccf99c`                     | XProf 开关、只在日志步计算范数、三路相机合批及其测试                                     |
-| 08-28 | `d0aec27`                     | 更新 H100 RTC 训练配置；没有单独的性能 A/B                                               |
-| 08-29 | `476babe` `f42992e`           | 引入 cuDNN fused attention，并记录最初的短跑/生产验证                                    |
-| 08-29 | `c1b70b2` `6305750` `51d9af3` | 处理全空 mask NaN；撤回逐步有限性保护；用 stop-gradient 保留快速原始 mask                |
-| 08-31 | `ad08973` `c9e77cc` `895811c` | 建立 strict-order/分层诊断，定位长期发散，并第一次合并事故文档                           |
-| 09-03 | `df9755f`                     | 增加 attention 精度与 Adam 更新空间诊断                                                  |
-| 09-04 | `3ea2b1c`                     | 定位 D 项舍入机制，实现并验证 fp16 custom VJP                                            |
-| 09-04 | `cf07bbe`                     | token 长度、XLA 集合通信和 remat A/B，形成 H100 最终配置                                 |
-
----
-
-## 1. 阶段一：数据管线（2026-08-19 实现，08-25 生产复核，commit `ec72f34`）
-
-起点是「训练每隔几十步卡 6 秒」。三个改动，其中只有两个真正兑现。
-
-### 1.1 改动 A：跳过未使用的触觉视频解码（收益 0）
-
-`LeRobotDataset.__getitem__` 无条件解码 `meta.video_keys` 里的**全部** 7 路视频，而 repack 只取
-head / left_wrist / right_wrist——4 路触觉流解完直接丢。H.264 随机寻址解码只有顺序解码的 1/200，纯浪费。
-
-实现是覆盖 `_get_query_timestamps`，让 `_query_videos` 永远不被要求解它们
-（`src/openpi/training/data_loader.py:203`）：
-
-```python
-class SelectiveVideoLeRobotDataset(lerobot_dataset.LeRobotDataset):
-    decode_video_keys: frozenset[str] | None = None   # 必须是普通属性：dataset 要 pickle 给 spawn 的 worker
-
-    @override
-    def _get_query_timestamps(self, current_ts, query_indices=None):
-        query_timestamps = super()._get_query_timestamps(current_ts, query_indices)
-        if self.decode_video_keys is None:
-            return query_timestamps        # None = 不过滤，忘了设也退化成原行为
-        return {key: ts for key, ts in query_timestamps.items() if key in self.decode_video_keys}
-```
-
-白名单由 `_resolve_decode_video_keys` 按子串 `tactile` 匹配（两代命名 `_tactile_{0,1}` 和
-`_tactile_{left,right}` 都能吃下），并且**启动时校验**：如果 repack 引用了被禁用的流直接 `ValueError`，
-不会静默拿到黑图。没有改 lerobot，没有转换数据。
-
-**结果**：单样本 `__getitem__` 262.4 ms → 153.3 ms（**1.71×**），全 transform 链输出 20 个样本**逐位相同**。
-**但 400 步端到端 A/B 完全没变**（768.4 s vs 770.8 s，停顿 20 次 vs 20 次）。唯一实测收益是启动快 37.3 s。
-
-原因：这台机 batch 256 + 192 核，加载器本来就跟得上 GPU；省下的 CPU 余量 GPU 用不上。更关键的是
-§1.4 的天花板与解码量无关。**保留它**（零风险、每样本少烧 40% CPU、epoch 边界省 37 s），
-但别指望它让当前配置变快。
-
-### 1.2 改动 B1：队头阻塞 → `in_order=False`
-
-诊断装置是关键：真数据 + 真 transform + 真 DataLoader，用 `sleep(1.55)` 替代 GPU 步，同时抓 torch
-迭代器的两个私有状态：
-
-```
-outstanding = _send_idx - _rcvd_idx     已派给 worker 但还没回来的任务数
-buffered    = _task_info 里已完成的结果  算完了、但还轮不到它交付的 batch 数
-```
-
-**这两个数把「算不出来」和「算出来了但过不去」区分开——是整件事里最值钱的一步。**
-
-
-| batch | wait   | outstanding | buffered |
-| ----- | ------ | ----------- | -------- |
-| 64    | 3.02 s | 1024        | 10       |
-| 66    | 9.84 s | 1024        | **41**   |
-| 130   | 9.66 s | 1024        | **41**   |
-
-`outstanding` 恒为 1024（`prefetch_factor 16 × 64 workers`，管线一直是满的），停顿时 `buffered` 高达 41
-——**41 个 batch 已经算完躺在那儿，训练循环却在等**。torch 严格按 sampler 顺序交付，任务 *j* 派给
-worker `j mod 64`，这一轮最慢的那个 worker 堵住排在它后面的所有 batch，每 64 步撞一次。
-调深 `prefetch_factor` 没用——问题不是没算，是不让过。
-
-修复：`DataLoader(in_order=False)`（PyTorch ≥2.6 为这个场景加的官方参数），谁先算完谁先交付。
-sampler 本来就是随机打散的，排序毫无意义。已验证一个 pass 内每个 index 仍恰好出现一次，不重不漏。
-
-### 1.3 改动 B2：epoch 边界重建迭代器 → `InfiniteSampler`
-
-`TorchDataLoader.__iter__` 在 sampler 耗尽时 `break` 出去重新 `iter()`，预取管线被整个拆掉重建，
-训练循环要等**一个 worker 独自造完一整个 batch**：
-
-```
-STALL batch 80: wait=62.96s outstanding=79 buffered=60
-```
-
-**63.0 秒** = 256 样本 × 246 ms/样本。修复（`data_loader.py:529`）：
-
-```python
-class InfiniteSampler(torch.utils.data.Sampler[int]):
-    """永不停止、每轮重新打散，StopIteration 再也不会触发。"""
-    def __iter__(self):
-        epoch = 0
-        while True:
-            if self._shuffle:
-                generator = torch.Generator().manual_seed(self._seed + epoch)
-                yield from torch.randperm(self._num_samples, generator=generator).tolist()
-            else:
-                yield from range(self._num_samples)
-            epoch += 1
-```
-
-代价为零：openpi 本来就不 checkpoint 数据加载器状态（`restore_state` 直接 `del data_loader`），
-batch 序列从来就不可恢复。
-
-### 1.4 结果与残留
-
-受控 400 步真实训练 A/B（同机同配置，唯一变量是 B1+B2）：
-
-
-| 指标               | 修复前  | 修复后      | 倍数       |
-| ------------------ | ------- | ----------- | ---------- |
-| 400 步总时长       | 770.8 s | **715.7 s** | **1.08×** |
-| 停顿步数（>0.5 s） | **20**  | **0**       | —         |
-| 最大`next_batch`   | 7.83 s  | **0.43 s**  | 18×       |
-| 中位`next_batch`   | 0.030 s | 0.290 s     | —         |
-
-> 中位数从 0.030 涨到 0.290 s **不是退步**：同样的产能缺口，以前攒成每 64 步一次的 6 秒尖峰，
-> 现在均摊到每一步，总等待反而少 16.8 s。
-
-**生产复核**（08-25，`unscrew-the-bottle-cap-0821`，64,951 步 / 30.60 h）：数据停顿从归一化后的
-每万步 334.7 次 / 2464 s 降到 **5.5 次 / 20 s**。该数据集一个 epoch 才 329 步，65,000 步里有 197 次
-epoch 边界，按修复前每次 63 s 算光边界就该吃掉 3.45 h。实际停顿损失 0.04 h，**约 1.14×**。
-
-残留的 36 次 `next_batch > 3 s` **全部落在 checkpoint 保存的下一步**（步号 1001、2001、3001…）。
-`save_interval` 从 1000 提到 5000 即可，已写进后续所有生产 yaml。
-
-修完之后 `buffered` **恒为 0**，加载器满负荷跑，稳态产能约 1.84 s/batch。折算每 worker 460 ms/样本，
-而单进程只要 153 ms——并发到 64 个 worker 时每样本成本涨了 3 倍。这个天花板不在解码上
-（直接证据：tactile 开/关对等待时间毫无影响）。阶段五复测显示等数据只剩 0.03–0.04 s/step，
-数据侧已不是瓶颈，这条线到此为止。
-
-### 1.5 开关
-
-```yaml
+  type: Pi0Config
+  pi05: true
+  max_token_len: 128
+  batch_image_views: true
+  use_cudnn_attention: true
+  cudnn_attention_dtype: float16
 data:
   base_config:
-    tactile: false          # 默认；true = 解码触觉流（Pi05Tactile 等模型才需要）
-strict_batch_order: false   # 默认；true = 恢复严格顺序交付，队头阻塞会回来
-```
-
-两个开关的默认值就是优化后的行为。`strict_batch_order: false` 意味着 batch 的**交付顺序**不再逐次可复现
-（每个 index 每轮仍恰好出现一次，样本内容不变）。诊断用的单变量 A/B 必须设 `true`——阶段三/四的所有
-3000 步验收都是这么跑的。
-
----
-
-## 2. 阶段二：主机传输与计算侧（2026-08-27 ~ 08-28，commit `9aef52d` `8ccf99c`）
-
-**测量口径**（后续所有短跑沿用）：从同一 checkpoint 只读恢复，独立实验目录，关 W&B 和所有 checkpoint 保存；
-跑到起始 step + 40，丢弃前 10 步的编译/预热，统计后 30 个稳态样本。JAX dispatch 是异步的，
-所以主指标是三者之和，不能单看 `dispatch`：
-
-```
-measured_step = dispatch + compute_sync + next_batch_total
-```
-
-### 2.1 共享内存 Tensor（commit `9aef52d`，−13.6%）
-
-worker 原先返回 NumPy batch，过 multiprocessing queue 要 pickle/复制约 900 MiB，主进程再包装成 JAX array。
-collate 改为返回 **CPU Torch Tensor**——PyTorch 的 multiprocessing reducer 会通过共享内存传 storage：
-
-```python
-def _collate_fn(items, *, profile=False):
-    """Collate batch elements into CPU tensors so workers use shared-memory IPC."""
-    batch = jax.tree.map(
-        lambda *xs: torch.from_numpy(np.stack([np.asarray(x) for x in xs], axis=0)),
-        *items,
-    )
-```
-
-主进程再用 `tensor.numpy()` 拿同一块 storage 的**零拷贝视图**交给 `jax.make_array_from_process_local_data`。
-同时加了显式 `DataLoader.close()`，训练退出时主动 shutdown worker 并超时终止残留进程。
-
-`3.275728 → 2.828845 s/step`（**−13.64%**）。之后数据侧分项均值：`main_queue_wait 0.005 s`、
-`jax_array_construct 0.009 s`、`h2d_wait 0.006 s`、`next_batch_total 0.046 s`——队列、array 构造和 H2D
-都已经不是瓶颈。
-
-### 2.2 XProf trace 定方向（commit `8ccf99c`）
-
-训练脚本加了可控的 trace 开关（绝对起始 step + step 数 + 输出目录，起止显式 `block_until_ready`）。
-采集 step 49011–49030 后按每卡每步汇总 kernel 时间：
-
-```
-GEMM     390.4 ms      Fusion   183.9 ms      NCCL   46.0 ms
-其他       2.6 ms      Memcpy     1.5 ms      Memset  0.5 ms
-```
-
-GEMM + Fusion 占已记录 kernel 时间的 **92.6%**，NCCL 7.4%。**继续优化 H2D 或 NCCL 没有意义，
-下一步必须动模型计算。** 这一步直接决定了阶段三去碰 attention。
-
-（注意：trace 里 kernel 可能重叠，不能把 duration 相加当 GPU 利用率。）
-
-### 2.3 只在日志步计算范数（commit `8ccf99c`，−0.5%）
-
-`grad_norm` / `param_norm` 是整个 trainable pytree 的归约，但只在 `log_interval` 步被记录。
-`train_step` 加静态布尔 `compute_metrics`（`scripts/train.py:168`）：
-
-```python
-    if compute_metrics:
-        grad_norm = optax.global_norm(grads)
-        param_norm = optax.global_norm(kernel_params)
-    else:
-        # 保持两个 JIT 变体的输出 pytree 一致；host 端用 nanmean 忽略占位值
-        grad_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
-        param_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
-```
-
-必须用**位置参数**传给 JIT：当前 JAX 在显式设置 `in_shardings` 时不接受 pjit 关键字实参。
-收益很小但语义正确、零代价，保留。
-
-### 2.4 三路相机合并为一次 SigLIP 调用（commit `8ccf99c`，−1.45%）
-
-`Pi0.embed_prefix` 原先对 head / left_wrist / right_wrist 分别调用同一个 image encoder 三次
-（`src/openpi/models/pi0.py:19`）：
-
-```python
-def _encode_image_views_as_batch(image_encoder, images):
-    """Encode camera views as a sample-major interleaved batch."""
-    names = tuple(images)
-    stacked = jnp.stack([images[name] for name in names], axis=1)   # (B, V, H, W, C)
-    batch_size, num_views = stacked.shape[:2]
-    flat = stacked.reshape((batch_size * num_views, *stacked.shape[2:]))
-    flat_tokens, _ = image_encoder(flat, train=False)               # 只调一次
-    tokens = flat_tokens.reshape((batch_size, num_views, *flat_tokens.shape[1:]))
-    return {name: tokens[:, index] for index, name in enumerate(names)}
-```
-
-布局是 **sample-major interleave**：同一样本的三路视图连续排列，所以 FSDP/DP 分片时每张卡仍拿到
-完整样本的所有视图，不需要跨设备重排。开关 `model.batch_image_views`（默认关）。
-单元测试逐元素比较了分别编码与合并编码的 token。
-
-累计：`3.275728 → 2.773743 s/step`，**−15.32%**（等价吞吐 +18.10%）。
-
-随后 `d0aec27` 只更新了两份 H100 RTC 配置的名称和参数，没有改变上述实现，也没有独立 A/B，
-因此不把它计算为新的提速阶段。
-
-> 一个容易被误读的现象：相机合并后 `dispatch` 从 2.657 s 降到 2.236 s，但 `compute_sync` 从 0.115 s
-> 涨到 0.492 s。这是异步工作在 dispatch 与同步边界之间重新分布，**不能宣称 dispatch 单项的 16% 提升**。
-
----
-
-## 3. 阶段三：cuDNN fused Attention 与两类失效（08-29 ~ 08-31）
-
-> **一句话**：吞吐收益是真的（−29%），但全 18 层 cuDNN bf16 会让训练在约 1000 步后缓慢发散至 NaN。
-> 这个阶段的产出是**两类失效的严格区分**和一份验收标准，真正的修复在阶段四。
-
-### 3.1 实现与收益（commit `476babe`）
-
-真实训练形状：`3×256 图像 token + 200 文本 token + 50 action token = 1018`，
-`num_heads=8, num_kv_heads=1, head_dim=256, bf16`，block mask。
-
-```python
-jax.nn.dot_product_attention(q, k, v, mask=attn_mask, scale=1.0, implementation="cudnn")
-```
-
-`q` 在调用前已按 head dimension 缩放，所以必须 `scale=1.0` 避免重复缩放。只在 `kv_cache is None`
-的训练路径启用，带 KV cache 的推理保持原实现。
-
-8×H200 / batch 512 / step 16000 恢复的 A/B：`2.795537 → 1.978580 s/step`（**−29.22%**），
-等价吞吐 183.15 → 258.77 samples/s。数值对齐测试当时也过了：BF16 前向最大绝对误差 0.0039，
-Q/K/V 梯度最大绝对误差 1.9e-6 / 4.8e-7 / 2.4e-7。
-
-**这些都是真的。问题在于它们全都不足以作为正确性验收。**
-
-### 3.2 第一类失效：全空 mask 行的瞬时 NaN（08-29，commit `c1b70b2` `51d9af3`）
-
-prompt 被补齐到 `max_token_len=200`，padding token 既不是有效 query 也不是有效 key，
-所以 `make_attn_mask` 会生成整行全 `False` 的 query 行。切到 cuDNN 后从 clean step 16000 恢复：
-
-```
-Step 16000（切换前，显式）: grad_norm=0.6721, loss=0.1429, param_norm=1806.9792
-Step 16100（切换后，cuDNN）: grad_norm=nan,    loss=0.1185, param_norm=nan
-Step 16200 及以后:          grad_norm=nan,    loss=nan,    param_norm=nan
+    tactile: false
+strict_batch_order: false
 ```
 
-最小复现（BF16 GQA + 16 个全空 query 行）：**前向有限，Q 梯度 16384 个 NaN，K/V 梯度有限**。
-真实生产形状是 71680 个 Q 梯度 NaN。
-
-> **一个隐蔽的坑**：日志窗口对 loss 用 `nanmean`，第一个窗口里唯一有限的首步 loss 让聚合结果显示
-> `0.1185`，掩盖了窗口其余全是 NaN。**「窗口 loss 有限」不能证明窗口内每个更新都有限。**
-
-最终修复保留原始 mask，只对全空行停 Q 梯度（`src/openpi/models/gemma.py` 中的
-`_stop_gradient_for_fully_masked_queries`）：
-
-```python
-def _stop_gradient_for_fully_masked_queries(q, attn_mask):
-    query_has_key = jnp.any(attn_mask, axis=-1)[:, 0, :, None, None]
-    return jnp.where(query_has_key, q, jax.lax.stop_gradient(q))
-```
-
-语义：Q 前向值逐元素不变；有效 query 的输出和梯度不变；无效 query 的 Q 梯度固定为 0；block mask 完整保留。
-
-另一个方案是给每个空行开一个 dummy key——在有效行上给出**逐位相同**的 dQ/dK/dV，但每层要重建
-`(B, 1, T, S)` mask，batch 256 下 1.38 vs 1.05 s/step，所以选了 stop-gradient。
-
-提交演进要注意：`c1b70b2` 最初还加入了逐 step 的有限性检查，并采用 dummy-key mask；
-`6305750` 恢复原训练更新语义、移除这项逐步保护，`51d9af3` 再把 dummy-key 换成上面的 stop-gradient，
-这才是后续诊断和当前代码使用的实现。
-
-### 3.3 第二类失效：全层 cuDNN 在约 1000 步后缓慢发散（08-30 ~ 08-31，commit `ad08973` `895811c`）
-
-修完全空 mask 后，全 18 层 cuDNN 前约 1000 步与显式基线贴合，随后分叉：
-
-```
-step    显式基线          cuDNN 9.14        cuDNN 9.19
- 200    0.9244 /   2.89   0.9333 /   3.58   0.9348 /   3.53
-1000    0.4220 /   2.77   0.4306 /   2.50   0.4319 /   3.16
-1200    0.3742 /   2.12   0.6484 /  31.40   0.6600 /  35.49     ← 分叉点
-2500    0.2457 /   1.33   5.7692 / 675.08   4.7055 / 759.18
-5900    0.1763 /   1.36   nan               已停止
-```
-
-单元格为 `loss / grad_norm`。**因果起点是 step 1000–1200，不是最终 NaN 的 step 5900。**
-单变量归因：同机同数据同超参同 seed，`git diff d0aec27 HEAD -- src/ scripts/` 的 72 行增删全部属于
-attention 开关本身。
-
-**但 cuDNN 的单步结果并不明显错误**：同权重同 batch 的完整模型梯度 A/B，梯度范数比 0.9937–1.0038，
-相对梯度差 1.7e-2–2.1e-2；单层对 fp32 参考的相对误差 cuDNN 5.4e-3 vs 显式 bf16 3.2e-3，同阶。
-cuDNN 同 batch 重复 backward 的随机差异（原子累加）就有 0.9%。**它给的是另一个合法的 BF16 近似。**
+`max_token_len: 128` 只针对上述耳机盒数据集做过验证：历史采样的 4037 个状态
+（含各维极值样本）最长为 116 tokens。更换数据集、prompt 或状态维度后需重新检查
+完整 transform 链的 token 长度；tokenizer 超长会警告并截断。不要直接套用 128。
+训练前也需重新计算 norm stats，并检查 `lr_schedule.decay_steps` 是否符合训练计划。
 
-### 3.4 机制：Adam 逐坐标放大结构性差异
+## 保留的优化
 
-从 step-1000 checkpoint 恢复真实 Adam 状态，对同一 batch 比较实际参数更新：
+- DataLoader worker 返回 CPU Torch Tensor，通过共享内存传输 batch；主进程以
+  `tensor.numpy()` 的零拷贝视图构造 JAX array，减少大批量 NumPy 数据的进程间复制。
+- 训练结束显式关闭 DataLoader，回收 worker。
+- 梯度和参数范数仅在 `log_interval` 步计算，减少全模型归约；loss 每步计算。
+- `batch_image_views` 将多路相机按样本交错合批，只调用一次 SigLIP，然后还原各视角。
+- cuDNN FP16 fused attention 使用动态 loss scaling 和 custom VJP，复用前向残差。
 
-```
-                        相对更新差异   更新 cosine   符号翻转坐标
-真实全层 cuDNN 差分         3.3521%      0.999440      0.7885%
-1.5% 人造乘性噪声           0.7313%      --            0.1459%
-```
+main 已包含的跳过未使用触觉视频、乱序 batch 交付和无限 sampler 继续保留。
+`strict_batch_order: false` 可避免队头阻塞，但交付顺序不保证跨运行一致；对比训练数值时设为 `true`。
+Gemma 与 SigLIP 保持默认的全量 rematerialization；历史 8×H100、batch 256 实验中，
+关闭 remat 或保存更多残差均超出显存容量。
 
-梯度 L2 只差 1.5%，但经 Adam 逐坐标归一化后变成 3.35% 的更新差异，0.79% 的坐标符号翻转。
-而按 `|g|` 分配能量的 1.5% 独立随机噪声几乎不扰动这些坐标，**从 step 0 稳定跑满 3000 步**
-（step 2900 `loss=0.2411 / grad_norm=1.6389`）——证明训练配方并非对任意 1.5% 梯度噪声都没裕度，
-问题在**误差的结构**而非幅度。
+## cuDNN 数值约束
 
-分层实验显示位置比数量更重要：
+全层 BF16 cuDNN attention 曾在约 1000 步后偏离显式基线并发散，不能依据短跑有限就用于生产。
+当前显式路径仍是默认回退方案：`use_cudnn_attention: false`。
+启用 cuDNN 时使用 `cudnn_attention_dtype: float16`。
 
-```
-cuDNN 层段      相对更新差异   符号翻转坐标
-全 18 层           3.3118%       0.7769%
-最前 3 层          3.3016%       0.7720%     ← 只融合前 3 层几乎等同全 18 层
-中间 3 层          1.9894%       0.4634%
-最后 3 层          0.8470%       0.1926%
-稳定噪声参考       0.7313%       0.1459%
-```
+排障中发现两个独立问题：
 
-早层的微小数值差异会穿过后续 15 层前向和反向放大。
+- 全空 attention mask 行可能产生 NaN query 梯度。实现将这些行的 query 梯度置零，保留原 mask。
+- peaked attention 行的 backward 对存储输出的舍入误差敏感。FP16 提供更多尾数位；
+  反向将余切按 2 的幂缩放后送入内核，再以 FP32 还原梯度，缓解 FP16 指数范围限制。
 
-### 3.5 当时的折中：只融合最后 3 层
+当前 VJP 直接调用固定 JAX 0.5.3 的私有 cuDNN forward/backward rule。
+升级 JAX 或 CUDA/cuDNN 时必须复核接口、分片约束和训练收敛性。
+q/k/v 与 mask/bias 显式使用一致的 batch 分片，避免 cuDNN partitioner 编译失败。
+通过 `policy_config.create_trained_policy` 加载 JAX 推理策略时会关闭 cuDNN attention；
+带 KV cache 的 attention 也继续使用显式实现。
 
-`configs/diag_cudnn_last3.yaml`，8×H100 / batch 256 / strict-order，从 step 0 跑满 3000 步，
-曲线与显式基线一致（step 1200 `0.3757 / 2.39`，step 2900 `0.2424 / 1.48`）。
-但 `~1.5 s/step` 对显式的 1.57–1.64，只有 4–9% 加速，远低于全层的潜在 36%。
+## 历史验证与性能
 
-层选择开关（`Pi0Config`，阶段四后已不需要，保留作诊断用）：
+2026-09-04，8×H100、全局 batch 256、固定 seed 和严格 batch 顺序，从 pi05_base step 0
+起跑的 FP16 VJP 初版通过 3000 步验收。历史 BF16 分叉点附近，FP16 loss 与显式基线差小于
+0.5%，记录的梯度范数未超过比较基线同点的 1.3 倍。
 
-```yaml
-model:
-  use_cudnn_attention: true
-  cudnn_attention_layer_start: 15
-  cudnn_attention_num_layers: 3
-```
+当前直接调用 forward/backward rule 的版本另做过真实激活等价性检查与 400 步对照：
+前向、dK/dV 逐位相同，dQ 差异与内核自身跨运行的原子累加差异同阶，loss 与初版逐点一致。
+这与初版的 3000 步训练验证是不同的验证范围。
 
-### 3.6 CUDA/cuDNN 环境事故
+同日 H100 端到端实验中，FP16 基线为 1.328 s/step；token 长度 200→128 后为
+1.246 s/step；单独加入下述 XLA 标志为 1.281 s/step；二者组合为 1.193 s/step，
+相对基线减少约 10.2%。口径为 400 步运行中 step 100 到 step 300 的日志墙钟差除以 200。
+这些实验使用同一配置和严格 batch 顺序，不能与其他机器、任务或 batch 的提速比例直接相乘。
 
-事故期间 `cuda_versions.cudnn_get_version()` 返回 `91400`，但进程实际加载的是
-**9.10.2 的 dispatcher stub + 9.14 的 engine 库**——`cudnnGetVersion()` 的返回值来自 engine，
-所以**只看这个数字发现不了混装**。
+新的 attention 后端或 VJP 改动，应从 step 0 固定数据、seed、batch 顺序及超参跑满
+3000 次更新，与显式基线对比 loss 和梯度趋势；确认参数与 optimizer state 有限，
+梯度范数不超过基线同点的 3 倍后，再进入有 checkpoint 的受监控长跑。
+单次前后向、吞吐短跑或单步梯度距离不能替代收敛验证。
 
-```
-site-packages/nvidia/cudnn/lib/libcudnn.so.9        <- 9.10.2 dispatcher
-$CONDA_PREFIX/lib/libcudnn_graph.so.9.14.0          <- 9.14 engine
-```
+## 环境与启动
 
-当前环境已统一（CUDA 12.8 / PyTorch 2.11.0+cu128 / nvidia-cudnn-cu12 9.19.0.56 / JAX 0.5.3）。
-**不要再导出 `LD_LIBRARY_PATH="$CONDA_PREFIX/lib"`**，残留的 conda-forge `libcublas`/`libcudart`/`libnccl`
-会遮蔽 pip cu128 库。生产启动前在同一 shell 跑：
+保持 pip CUDA 12.8 库来源一致，避免将 `$CONDA_PREFIX/lib` 加到 `LD_LIBRARY_PATH` 前面。
+当前依赖配置使用 PyTorch 2.11 对应的 cuDNN 9.19；训练启动时记录 JAX cuDNN runtime 版本。
+运行 `check_cuda_stack.py` 检查实际加载的库来源与 BF16 GQA 前后向；该检查不替代 FP16 收敛验证。
 
 ```bash
 env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py
-```
-
-该脚本遍历 `/proc/self/maps`，在 `libcudnn*` 来自多个目录或版本不一致时判 FAIL，
-并在真实训练形状上做前向+反向与 fp32 参考对比。
-
-> **PASS 只证明环境单源、单步内核有限。第二类缓慢发散曾通过所有这类 smoke test。**
-
----
-
-## 4. 阶段四：根因与 fp16 修复（09-03 ~ 09-04，commit `df9755f` `3ea2b1c`）
-
-阶段三证明了「BF16 差分被 Adam 放大」，但没回答**差分从哪来、为什么它不像随机噪声**。这一阶段用真实激活回答了。
-
-### 4.1 机制：flash 反向的 D 项用了 bf16 舍入后的 O
-
-cuDNN fused attention 的反向按 flash-attention 方式算 `dS = P ⊙ (dP − D)`，其中
-`D_i = rowsum(dO_i ⊙ O_i)` 用的是**前向存下的、已舍入到 bf16 的 O**。
-显式路径用同一份舍入后的 dP 算 `D_i = Σ_j P_ij dP_ij`，抵消是自洽的。
-
-两者之差是一个每行标量：
-
-```
-Δ_i = dO_i · (bf16(O_i) − O_i)
-```
-
-进入梯度后方向为 `−Δ_i · Σ_j P_ij K_j`，**破坏了 softmax 梯度「行和为零」的不变量**。
-
-对 **peaked 行**（`maxP > 0.99`，即注意力 sink）杀伤最大：那里 `dP` 与 `rowsum(dO·O)` 几乎完全抵消，
-真实 dQ 是 ε 级量，比普通行小两个数量级，而 H1 误差不随 ε 缩小——所以相对误差达 O(1)。
-
-### 4.2 证据
-
-两个探针：`scripts/attention_dterm_probe.py`（GPU，monkeypatch `gemma.Block.__call__`，
-用 `jax.debug.callback` 在一次真实训练步里抓下 18 层的 q/k/v/mask 和反向余切 dO，缓存 1.5 GB）；
-`scripts/attention_dterm_cpu_analysis.py`（CPU float64，对 peaked 行逐行重建精确反向）。
-
-> 实现细节：模拟 bf16 舍入必须用 `jax.lax.reduce_precision`，`astype` 往返会被 XLA 的
-> `allow_excess_precision` 删掉，预测会恒为零。
-
-单层实测（`cos` 是误差向量与 H1 预测方向的 cos）：
-
-```
-层   变体           dQ rel    cos(all)   peaked rel   cos(peaked)
- 0   explicit_bf16  2.46e-3   -0.02      2.8e-3       -0.04
-     cudnn_bf16     3.58e-3   +0.49      3.5e-1       +0.83
-     cudnn_fp16     1.71e-3   -0.02      3.9e-2       +0.13
- 8   explicit_bf16  2.88e-3   +0.00      3.5e-3       -0.10
-     cudnn_bf16     4.28e-3   +0.51      5.1e-1       +0.96
-     cudnn_fp16     1.78e-3   +0.03      7.2e-2       +0.57
-
-18 层几何平均：explicit_bf16 2.33e-3 / cudnn_bf16 3.72e-3 / cudnn_fp16 1.76e-3
-18 层平均 cos(all)：explicit 0.006 / cudnn_bf16 0.418 / cudnn_fp16 0.006
-```
-
-- cuDNN bf16 的误差在**每一层**都与 H1 预测显著同向；显式和 fp16 都不同向。
-- H1 预测的 peaked 行误差量级与实测吻合（0.41/0.35、0.63/0.51、0.86/1.21）。cos 不到 1 是因为
-  cuDNN 前向的 PV 乘也把 P 舍入到 bf16。
-- 候选机制 H2（dP 先舍入再减 D）的 cos 为 0 或负，排除。
-- dV 没有 D 项，cuDNN bf16 的 dV 误差（2.5e-3）与显式（1.9e-3）同阶——干净的对照。
-
-**这解释了阶段三的全部现象**：H1 分量按行沿 `Σ_j P_ij K_j` 方向、符号随舍入随机，投到参数上落在
-`x ⊗ K_argmax` 这类**低维子空间**，不与 `|g|` 成比例。sink 机制处于平衡态时这些坐标的真实梯度和
-二阶矩都很小，Adam 归一化后每步都在那里走满 lr 的随机步，约 1000 步后 sink 结构失稳。
-乘性噪声在小 `|g|` 坐标上也小，Adam 看不见。跨 batch delta 两两 cos ±0.25 在 30 亿维空间里不是白噪声，
-正是「同一子空间、随机符号」。只融合前 3 层≈全 18 层，因为 sink/peaked 结构在早层形成。
-改 mask、seq_lengths、cuDNN 版本、LR、b1/b2/eps 都没用，**因为它们都不改变 O 的舍入**。
-
-### 4.3 修复：fp16 内核 + 2 的幂动态缩放
-
-fp16 比 bf16 多 3 位尾数，H1 误差降约 8 倍；代价是指数范围小得多，所以余切要做动态 loss scaling。
-真实激活的范围 `max|q| ≤ 3.5`、`|k| ≤ 30`、`|v| ≤ 84`、`|dO| 1e-4–1e-2`，都在 fp16 范围内。
-
-关键实现（`src/openpi/models/gemma.py:201`，`_cudnn_attention_in_dtype`）：
-
-```python
-    def attention_fwd(q, k, v, attn_mask):
-        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
-        bias = jnp.where(attn_mask, jnp.asarray(0, compute_dtype),
-                         _cudnn_fa.get_large_negative_number(compute_dtype))
-        # cuDNN 的 custom partitioner 要求 q/k/v/bias 分片一致
-        qc, kc, vc, bias = sharding.activation_sharding_constraint((qc, kc, vc, bias))
-        out, res = _cudnn_fa._dot_product_attention_fwd_rule(...)
-        return out.astype(out_dtype), (res, attn_mask)
-
-    def attention_bwd(residuals, d_out):
-        res, attn_mask = residuals
-        d_out32 = d_out.astype(jnp.float32)
-        max_abs = jnp.max(jnp.abs(d_out32))
-        # 2 的幂缩放，使 max|d_out * scale| 落在 [0.25, 0.5)：尾数上精确，
-        # 给内核里的 dP = dO @ V^T 和 dQ/dK/dV 留约 17 个 binade 余量
-        exponent = jnp.floor(jnp.log2(jnp.maximum(max_abs, tiny)))
-        scale = jnp.where(max_abs > 0, jnp.exp2(-(exponent + 2.0)), 1.0)
-        grads = _cudnn_fa._dot_product_attention_bwd_rule(..., (d_out32 * scale).astype(compute_dtype))
-        dq, dk, dv = grads[:3]
-        # 与 _stop_gradient_for_fully_masked_queries 同语义：全 mask 行的 dQ 置零而非 NaN
-        query_has_key = jnp.any(attn_mask, axis=-1)[:, 0, :, None, None]
-        dq = jnp.where(query_has_key, dq, jnp.zeros_like(dq))
-        ...
-```
-
-开关：`model.cudnn_attention_dtype: float16`（默认 `bfloat16`，逐位不变）。
-
-**两版实现**：
-
-- **v1**（3000 步验收用的）：custom VJP 的反向里用 `jax.vjp` 重走一遍 `jax.nn.dot_product_attention`
-  前向再取反向。配合 `nothing_saveable` remat，每层反向比 bf16 路径多一次 attention 前向。1.37 s/step。
-- **v2**（当前代码，commit `3ea2b1c`）：直接调 jax 0.5.3 私有的
-  `jax._src.cudnn.fused_attention_stablehlo._dot_product_attention_fwd_rule / _bwd_rule`
-  ——就是 `jax.nn.dot_product_attention(implementation="cudnn")` 自己 custom_vjp 里接的那两个函数。
-  softmax stats 和输出当残差保存，反向只跑一次 fused backward。1.31–1.35 s/step。
-
-  `attn_mask` 必须作为**显式参数**而非闭包传入，否则 flax remat 单独 trace 反向时会捕获泄漏的 tracer。
-
-  `scripts/cudnn_fp16_vjp_equivalence.py` 在真实激活（层 0/4/8/14/17）上比较 v2 与 v1：
-  前向逐位相同，dK/dV 逐位相同，dQ 相对差 5e-6–1.2e-5，与同一内核自身的 run-to-run 差
-  （1.6e-6–1.4e-5，原子累加）同阶。
-
-### 4.4 3000 步 strict-order 验收：通过
-
-`configs/diag_cudnn_fp16.yaml`，与 `diag_cudnn_strict_order` 同数据/seed/batch 顺序/超参，
-全 18 层 cuDNN fp16，从 pi05_base step 0 起，8×H100 batch 256：
-
-```
-step    fp16 loss / grad_norm    显式基线           last3 基线        历史 cuDNN bf16
-   0    6.4030 /  65.59          6.4023 / 65.63     6.4029 / 65.52    6.4023 / 65.63
-1000    0.4231 /   2.21          0.4220 /  2.77     0.4240 /  2.38    0.4319 /  3.16
-1200    0.3759 /   2.36          0.3742 /  2.12     0.3757 /  2.39    0.6600 / 35.49   ← 历史分叉点
-2000    0.2792 /   1.93                             0.2793 /  1.83    3.7753 / 150.5
-2900    0.2427 /   1.59                             0.2424 /  1.48    6.2039 / 603.5
-```
-
-30 个 log 点全部有限，`grad_norm` 最大 4.75（step 100），之后 1.4–2.8，**没有一个点超过基线同点的
-1.3 倍**（验收上限 3 倍）。历史分叉点处 loss 与显式基线差 <0.5%。
-**把 cuDNN 内核精度从 bf16 换成 fp16 就消除了发散，验证了 §4.1 的机制。**
-
-### 4.5 步时：attention 内核本身不解释差距
-
-```
-路径                                   s/step      来源
-显式 attention                          1.57–1.64   8-31 报告 / 今日复测 1.60–1.63
-cuDNN bf16（发散）                       1.05        8-31 报告
-cuDNN fp16 v1（反向重算前向）             1.34–1.37   3000 步验收 run
-cuDNN fp16 v2（直接调 fwd/bwd 规则）      1.31–1.35   400 步（loss 逐点与 v1 一致）
-cuDNN bf16，今日同 config 复测            1.32–1.35   400 步
-```
-
-去掉反向重算只省约 2%。单卡微基准（`scripts/cudnn_attention_microbench.py`，层 0 真实激活，
-T=S=1018，含 `nothing_saveable` remat）说明原因：
-
-```
-变体                fwd      fwd+bwd(remat)   ×18 层
-cudnn_bf16         1.35 ms   5.24 ms          0.094 s/step
-cudnn_fp16_v2      1.54 ms   5.53 ms          0.099 s/step
-explicit_bf16      4.46 ms   8.49 ms          0.153 s/step
-```
-
-18 层 attention 前向+反向合计只占每步约 **0.1 s**；fp16 内核与 bf16 同速，动态缩放开销可忽略，
-显式路径也只多 0.06 s。**所以 attention 内核解释不了 1.33 vs 1.05 的差距，也解释不了显式的 1.6 s**
-——差距在整图层面（XLA 的 fusion/remat 决策、显存压力）。今日同 config 的 bf16 复测已是 1.32–1.35，
-与 fp16 持平，说明 8-31 那个 1.05 s 不是同一编译环境下的可比数字。
-
-### 4.6 一个必须记住的负面结论：单步指标不区分好坏
-
-`scripts/attention_precision_update_probe.py`，锚点是 fp16 验收 run 的 step-1000 train state
-（参数 + Adam 矩），以 fp32 显式 attention 为真值：
-
-```
-                   原始梯度 rel / sign_flips     Adam 更新 rel / cos / sign_flips
-explicit_bf16      3.37e-2 / 1.53e-2            3.13e-2 / 0.99951 / 7.47e-3
-cudnn_bf16         3.18e-2 / 1.43e-2            2.93e-2 / 0.99957 / 6.97e-3     ← 会发散的那个
-cudnn_fp16         3.01e-2 / 1.38e-2            2.81e-2 / 0.99961 / 6.73e-3
-```
-
-**会发散的 cudnn_bf16 在这个指标上与两条稳定路径不可区分，甚至略好于 explicit_bf16。**
-D 项误差在范数上只占 ~1e-3，危害来自方向的结构性（同一低维子空间、跨 step 累积）。
-判定 fused attention 后端好坏要看 §4.2 的 `cos(err, H1)` 与 peaked 行误差，最终靠 3000 步 strict-order run。
-**不要再用单步梯度/更新距离做验收指标。**
-
----
-
-## 5. 阶段五：H100 端到端 A/B（09-04，commit `cf07bbe`）
-
-**配置**：`pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml`
-（pi05、cuDNN fp16、`batch_image_views`、全局 batch 256、64 workers、fsdp 8）
-**数据**：`Xense/earbud_case_insertion_teleop_0515_left_right`，776 集 / 1,145,613 帧
-**口径**：每个变体从 pi05_base 起跑 400 步，`strict_batch_order: true`，关 checkpoint 和 W&B；
-主指标是日志 `Step 100` 到 `Step 300` 的时间戳之差除以 200（含日志开销的端到端墙钟）；
-`loss@300` 用来确认数值没变。
-
-
-| 变体                                              | s/step    | 相对基线    | loss@300 | 结论                               |
-| ------------------------------------------------- | --------- | ----------- | -------- | ---------------------------------- |
-| E0 基线                                           | 1.328     | —          | 0.7356   | 等数据 0.03 s，数据侧不是瓶颈      |
-| E1`max_token_len` 200→128                        | 1.246     | **−6.2%**  | 0.7356   | 采纳                               |
-| E2 关闭 remat                                     | OOM       | —          | —       | XLA 试图分配 280 GiB               |
-| E2b gemma remat`dots_with_no_batch_dims_saveable` | 编译失败  | —          | —       | cuDNN 分片器报 q/k/v 分片不一致    |
-| E2c 同上 + q/k/v 显式分片约束                     | OOM       | —          | —       | 编译过了，单块残差 84.6 GB         |
-| E3 XLA 集合通信标志                               | 1.281     | **−3.5%**  | 0.7357   | 采纳                               |
-| E4`fsdp_devices` 8→4                             | 1.335     | +0.5%       | 0.7354   | 无收益                             |
-| **E5 = E1 + E3**                                  | **1.193** | **−10.2%** | 0.7356   | 生产采用；192.8 → 214.6 samples/s |
-
-60000 步按稳态估算：22.1 h → 19.9 h，省约 2.2 h。
-
-### 5.1 E1：token 长度按数据集实测收紧
-
-`scripts/token_len_scan.py` 走真实 transform 链（repack → BiFlexivInputs → Normalize → TokenizePrompt），
-对 4037 个状态（含每一维的最小/最大值样本）做 tokenize：min 110、p50 113、p99 115、**max 116**。
-原来 200 里有 84 个是纯 padding，但 attention/MLP 照算。128 留 12 个余量。
-
-```bash
-JAX_PLATFORMS=cpu uv run scripts/token_len_scan.py \
-  --config-name pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100
-```
-
-> **换数据集或改 prompt 时必须重扫**：tokenizer 超长只会 warning 并截断，不会报错。
-
-### 5.2 E3：XLA 标志
-
-```
---xla_gpu_enable_latency_hiding_scheduler=true
---xla_gpu_all_gather_combine_threshold_bytes=1073741824
---xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824
---xla_gpu_all_reduce_combine_threshold_bytes=1073741824
---xla_gpu_enable_pipelined_all_gather=true
---xla_gpu_enable_pipelined_reduce_scatter=true
---xla_gpu_enable_while_loop_double_buffering=true
-```
-
-FSDP 每步的参数 all-gather / 梯度 reduce-scatter 合并成大块并与计算流水。
-**没有逐项拆分测哪一个起作用。** 通过 `XLA_FLAGS` 环境变量传入，已写在 0904 yaml 头部注释里。
-
-### 5.3 E2：remat 在 80 GB 上没有空间
-
-`gemma.py` / `siglip.py` 每层用 `nothing_saveable` 全量重算。本次给 `Pi0Config` 加了
-`gemma_remat_policy` / `siglip_remat_policy`（默认不变），并在 cuDNN 调用前给 q/k/v/mask 加了统一的
-batch 分片约束——否则非默认 remat 策略下编译失败：
-
-```python
-def _cudnn_attention_call(q, k, v, attn_mask):
-    q = _stop_gradient_for_fully_masked_queries(q, attn_mask)
-    # cuDNN 的 custom partitioner 要求 q/k/v 分片一致。显式钉到 batch 轴，
-    # 否则 XLA 分片传播可能逐 operand 不同（非默认 remat 策略下出现过），编译报 "should have same sharding"。
-    q, k, v, attn_mask = sharding.activation_sharding_constraint((q, k, v, attn_mask))
-    return jax.nn.dot_product_attention(q, k, v, mask=attn_mask, scale=1.0, implementation="cudnn")
-```
-
-但 batch 256（每卡 32）时：完全不 remat 要 280 GiB 单块，只保存无 batch 维 matmul 也要 84.6 GB，都放不下。
-**在 H200 141 GB 或每卡 batch 更小时这两个开关才可能兑现**，理论上限约为一次前向的重算（15–25%）。
-
-### 5.4 其他观察
-
-- worker 启动一次约 330 s（spawn 上下文，64 进程各自 import jax/torch/lerobot 并反序列化 1.1M 行数据集），
-  是每次启动的一次性开销。
-- `ResizeImages` 实际调用的是 `xense_client.image_tools`（PIL），worker 里**不跑 JAX resize**
-  ——阶段一怀疑的头号嫌疑人被排除。
-- fsdp 4 相对 8 没有收益：参数 all-gather 的量只差 12.5%，而每卡持有的优化器状态翻倍。
-- 短跑结束后偶有 DataLoader worker 成为孤儿进程（ppid 1，每个约 3 GB RSS，一次残留 39 个）；
-  批量实验后用 `ps -eo pid,ppid,cmd | grep spawn_main` 检查。
-- 该配置 `num_train_steps: 60000` 但 `lr_schedule.decay_steps` 仍是默认 30000，后半程一直在底线 LR。
-  **这是训练质量问题，本次未动。**
-
----
-
-## 6. 新 attention 后端的验收标准
-
-任何新的全层 fused/flash attention backend、custom VJP 或稳定化方案必须：
-
-1. 从 **step 0** 开始；
-2. `strict_batch_order: true`，固定 seed、数据集、batch 和超参；
-3. 跑满 **3000 个更新**，不得只跑 40、60 或 500 步；
-4. 与显式基线逐点比较 step 200/500/1000/1500/2000/2900；
-5. 全程 loss、梯度、参数和 optimizer state 有限；
-6. `grad_norm` 不得超过基线同点的 3 倍；
-7. 通过后再进入有 checkpoint 的受监控长跑。
-
-**以下均不是有效验收**：单次前向/反向有限；40-step 吞吐短跑；从已收敛 checkpoint 短跑；
-200–500 个连续有限 step；只检查最终是否 NaN；单步梯度/Adam 更新距离（§4.6）；
-跨不同 batch 顺序的 run 直接比较单点 grad norm。
-
-**建议增加发散护栏**：当前只在日志步算 `grad_norm`。在 host 端维护滚动中位数，
-连续多个窗口超过固定倍数时中止——按 §3.3 的曲线可在 step 1200–1300 停下，而不是污染数万个 step。
-
-**checkpoint 恢复规则**：一旦参数或 Adam `mu`/`nu` 出现 NaN，后续 checkpoint 全部视为污染，
-不得通过切换 attention、降 LR 或改 optimizer 继续；必须回退到最后一个确定干净的 checkpoint，
-用独立实验名验证（关 W&B 和所有保存）后才能恢复生产保存。
-
----
-
-## 7. 不要重复的弯路
-
-**测量方法**
-
-1. **先看停顿的周期再动手。** 周期 = `num_workers` 是队头阻塞，周期 = 一个 epoch 的步数是边界重建，
-   周期 = `save_interval` 是 checkpoint——都不是「解码太慢」。
-2. **`buffered` 千金难买。** 停顿时它是 0 说明真的算不出来，是 41 说明算出来了过不去，两种修法完全相反。
-3. **别把别的机器上的预估当结论。** 「跳过触觉解码约 1.8×」在 h200 箱 batch 512 下是对的，
-   搬到 batch 256 + 192 核就一点不剩。
-4. **别拿两个不同任务的日志比绝对时长。** 要比就比归一化指标，或者跑受控 A/B。
-5. **别单看 `dispatch`。** JAX 异步，收益可能只是从 dispatch 挪到 compute_sync。
-6. **`nanmean` 会掩盖 NaN。** 窗口 loss 有限不证明窗口内每个更新有限。
-
-**cuDNN attention**
-
-7. 不要再把全空 mask 的瞬时 NaN 与全层缓慢发散当成同一问题。
-8. 不要再把 cuDNN 9.15.1 的非 128 倍序列长度缺陷或 `T=1024` 当作本任务的修复
-   （实测 T=1018 与 T=1024 的 dQ 误差同为 3.2e-3，NaN=0）。
-9. 不要再尝试传完整物理长度的 `query_seq_lengths`/`key_value_seq_lengths`
-   （编译确认进了 custom call，但数值指纹 `3.311% / 0.776%` 与已知发散路径完全一致，3000 步照样发散）。
-10. 不要再改 dummy-key / stop-gradient 的 mask 语义来解决发散——两者在有效行上数值等价。
-11. 不要只调 `peak_lr`（2.5e-5→1.5e-5 只推迟约 300 步）、Adam `b1`（0.9→0.99 推迟到约 1400 步）、
-    `b2`（0.95→0.99 照样发散）或 `eps`（更新差异几乎不变）。
-12. 不要用独立随机噪声代替真实 cuDNN 差分做实验——1.5% 乘性噪声能稳定跑满 3000 步。
-13. 不要试「V 减去 sink 键的 V」：peaked 行大多不在看 BOS，而是在看 action/text/image 的其他键。
-14. 不要只看 `cuda_versions.cudnn_get_version()` 判断动态库是否干净（它由 engine 返回，看不见 dispatcher 混装）。
-15. 用 `jax.lax.reduce_precision` 模拟低精度舍入，`astype` 往返会被 XLA 删掉。
-
----
-
-## 8. 脚本与配置索引
-
-**诊断脚本**
-
-
-| 脚本                                          | 用途                                                                            |
-| --------------------------------------------- | ------------------------------------------------------------------------------- |
-| `scripts/check_cuda_stack.py`                 | 枚举`/proc/self/maps` 里实际加载的 cuDNN/cuBLAS/NCCL，+ 真实形状数值 smoke test |
-| `scripts/token_len_scan.py`                   | 走真实 transform 链扫数据集的 tokenized prompt 长度分布                         |
-| `scripts/attention_dterm_probe.py`            | 抓真实激活 + 单层五变体误差 + H1/H2 预测（阶段四主力）                          |
-| `scripts/attention_dterm_cpu_analysis.py`     | peaked 行的 float64 逐行解剖                                                    |
-| `scripts/cudnn_fp16_vjp_equivalence.py`       | fp16 VJP v2（直接调规则）与 v1（反向重算）的等价性                              |
-| `scripts/cudnn_attention_microbench.py`       | 单层 attention 前向/反向微基准                                                  |
-| `scripts/attention_precision_update_probe.py` | 真实 Adam 状态上的更新空间对比（**注意 §4.6：该指标不区分好坏**）              |
-| `scripts/grad_ab_probe.py`                    | 完整模型显式/cuDNN 梯度 A/B                                                     |
-| `scripts/grad_ab_direction_probe.py`          | 多 batch 差分方向和随机分量                                                     |
-| `scripts/hybrid_update_ab_probe.py`           | 不同 cuDNN 层段的更新空间误差                                                   |
-| `scripts/optimizer_update_ab_probe.py`        | 从真实 Adam 状态比较参数更新误差                                                |
-
-**诊断配置**：`configs/diag_cudnn_strict_order.yaml`（全层 bf16，会发散）、
-`configs/diag_explicit_strict_order.yaml`（基线）、`configs/diag_cudnn_last3.yaml`（后 3 层折中）、
-`configs/diag_cudnn_fp16.yaml`（fp16 验收）、`configs/diag_cudnn_b1_099.yaml`。
-
-**代码位置**
-
-
-| 文件                                 | 相关改动                                                                                                                                  |
-| ------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------- |
-| `src/openpi/training/data_loader.py` | `SelectiveVideoLeRobotDataset`、`_resolve_decode_video_keys`、`InfiniteSampler`、`in_order` 接线、Torch Tensor collate、`close()`         |
-| `src/openpi/training/config.py`      | `DataConfig.tactile`、`TrainConfig.strict_batch_order`                                                                                    |
-| `src/openpi/models/pi0_config.py`    | `batch_image_views`、`use_cudnn_attention`、`cudnn_attention_{layer_start,num_layers,dtype}`、`gemma_remat_policy`、`siglip_remat_policy` |
-| `src/openpi/models/pi0.py`           | `_encode_image_views_as_batch`                                                                                                            |
-| `src/openpi/models/gemma.py`         | `_stop_gradient_for_fully_masked_queries`、`_cudnn_attention_call`、`_cudnn_attention_in_dtype`、`_cudnn_static_args`、逐层 cuDNN 开关    |
-| `scripts/train.py`                   | `compute_metrics` 静态分支、XProf trace 开关、数据管线 profiling                                                                          |
-
----
-
-## 9. 生产启动
-
-```bash
-env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py    # 先确认环境单源
 
-XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true \
-  --xla_gpu_all_gather_combine_threshold_bytes=1073741824 \
-  --xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824 \
-  --xla_gpu_all_reduce_combine_threshold_bytes=1073741824 \
-  --xla_gpu_enable_pipelined_all_gather=true \
-  --xla_gpu_enable_pipelined_reduce_scatter=true \
-  --xla_gpu_enable_while_loop_double_buffering=true" \
-XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
+env -u LD_LIBRARY_PATH \
+  XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true \
+    --xla_gpu_all_gather_combine_threshold_bytes=1073741824 \
+    --xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824 \
+    --xla_gpu_all_reduce_combine_threshold_bytes=1073741824 \
+    --xla_gpu_enable_pipelined_all_gather=true \
+    --xla_gpu_enable_pipelined_reduce_scatter=true \
+    --xla_gpu_enable_while_loop_double_buffering=true" \
+  XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
   python scripts/train.py <config> --exp-name=<name>
 ```
 
-换数据集时记得：重扫 `max_token_len`（§5.1）、重算 norm stats、检查 `lr_schedule.decay_steps`
-是否匹配 `num_train_steps`（§5.4）。
+这些 XLA 标志作为整体做过验证，未逐项测量贡献。继续已有实验使用 `--resume`；
+仅在确实需要覆盖目标实验目录时使用 `--overwrite`。
diff --git a/scripts/attention_dterm_cpu_analysis.py b/scripts/attention_dterm_cpu_analysis.py
deleted file mode 100644
index 19d55ec..0000000
--- a/scripts/attention_dterm_cpu_analysis.py
+++ /dev/null
@@ -1,176 +0,0 @@
-"""CPU-only follow-up to attention_dterm_probe.py: anatomy of the peaked attention rows.
-
-Loads the captured q/k/v/mask/d_out of every layer (dterm_capture.npz) and, in float64,
-reconstructs the exact single-layer backward for the peaked rows to answer:
-
-  * how big are K/V of the BOS sink key (position 768) relative to other keys?
-  * which rows are peaked (maxP > 0.99): BOS-attending rows or other keys?
-  * on BOS rows, how much cancellation does the true dQ have, and how large is the
-    D-term error from rounding the stored O to bf16 (H1) relative to it -- with and
-    without shifting V by V_sink?
-
-Pure numpy; run anywhere the npz is:  python scripts/attention_dterm_cpu_analysis.py
-"""
-
-import sys
-
-import ml_dtypes
-import numpy as np
-
-CAPTURE_FILE = sys.argv[1] if len(sys.argv) > 1 else "/root/localstorage/logs/dterm_capture.npz"
-BIG_NEG = -2.3819763e38
-SINK = 768
-MAX_ROWS = 3000
-rng = np.random.default_rng(0)
-
-
-def region(pos):
-    if pos < 768:
-        return f"img{pos // 256}"
-    if pos < 968:
-        return "txt"
-    return "act"
-
-
-def bf16_round(x):
-    return x.astype(ml_dtypes.bfloat16).astype(np.float64)
-
-
-def load():
-    data = np.load(CAPTURE_FILE)
-    layers = {}
-    for key in data.files:
-        layer, name = key.split("/")
-        arr = data[key]
-        if arr.dtype.kind == "V" and arr.dtype.itemsize == 2:  # bf16 saved by np.savez as raw bytes
-            arr = arr.view(ml_dtypes.bfloat16)
-        layers.setdefault(int(layer), {})[name] = arr
-    return layers
-
-
-def row_anatomy(q_rows, k, v, mask_rows, d_rows, shift=None):
-    """q_rows/d_rows: R x H ; k, v: S x H (single kv head) ; mask_rows: R x S bool. float64."""
-    s = q_rows @ k.T  # R S
-    s = np.where(mask_rows, s, -np.inf)
-    s -= s.max(axis=1, keepdims=True)
-    p = np.exp(s)
-    p /= p.sum(axis=1, keepdims=True)
-    v_eff = v if shift is None else v - shift[None, :]
-    o = p @ v_eff  # R H   (fp64 "true" output of the kernel-internal V)
-    dp = d_rows @ v_eff.T  # R S
-    d_true = np.sum(p * dp, axis=1)  # R
-    ds = p * (dp - d_true[:, None])
-    dq = ds @ k  # R H
-    # H1: kernel D from bf16-rounded O
-    delta = np.sum(d_rows * (bf16_round(o) - o), axis=1)  # R
-    pk = p @ k  # R H
-    err_h1 = -delta[:, None] * pk
-    return dict(p=p, o=o, dp=dp, d_true=d_true, ds=ds, dq=dq, delta=delta, pk=pk, err_h1=err_h1)
-
-
-def norms(x):
-    return np.sqrt(np.sum(x * x, axis=-1))
-
-
-def main():
-    layers = load()
-    print(f"loaded {len(layers)} layers from {CAPTURE_FILE}", flush=True)
-    for layer in sorted(layers):
-        e = layers[layer]
-        q = e["q"].astype(np.float64)  # B T N H
-        k = e["k"].astype(np.float64)[:, :, 0]  # B S H
-        v = e["v"].astype(np.float64)[:, :, 0]
-        mask = e["mask"][:, 0]  # B T S
-        d = e["d_out"].astype(np.float64)  # B T N H
-        B, T, N, H = q.shape
-        valid_q = mask.any(-1)  # B T
-        valid_k = mask.any(1)  # B S
-
-        kn, vn = norms(k), norms(v)  # B S
-        med_k = np.median(kn[valid_k])
-        med_v = np.median(vn[valid_k])
-        print(
-            f"\n=== layer {layer:2d} ===  |K_768|={kn[:, SINK].mean():6.2f} (median valid |K|={med_k:6.2f})"
-            f"   |V_768|={vn[:, SINK].mean():6.2f} (median valid |V|={med_v:6.2f})",
-            flush=True,
-        )
-
-        # per-row softmax stats for all valid rows (float32 logits are exact enough for argmax/maxP)
-        s = np.einsum("btnh,bsh->btns", q.astype(np.float32), k.astype(np.float32))
-        s = np.where(mask[:, :, None, :], s, BIG_NEG)
-        s -= s.max(-1, keepdims=True)
-        p = np.exp(s)
-        p /= p.sum(-1, keepdims=True)
-        max_p = p.max(-1)  # B T N
-        arg_p = p.argmax(-1)
-        del s, p
-        vrow = np.broadcast_to(valid_q[:, :, None], max_p.shape)
-        peaked = (max_p > 0.99) & vrow
-        bos_rows = (arg_p == SINK) & vrow
-        n_peaked = int(peaked.sum())
-        n_peaked_bos = int((peaked & bos_rows).sum())
-        hist = {}
-        for pos in arg_p[peaked & ~bos_rows]:
-            hist[region(int(pos))] = hist.get(region(int(pos)), 0) + 1
-        print(
-            f"  peaked rows (maxP>0.99): {n_peaked}  of which BOS(768): {n_peaked_bos}; other peaked keys by region: {hist}"
-            f"   BOS-argmax rows: {int(bos_rows.sum())} ({bos_rows.sum() / vrow.sum():.1%} of valid rows),"
-            f" P>0.9: {int((bos_rows & (max_p > 0.9)).sum())}, P>0.99: {n_peaked_bos}",
-            flush=True,
-        )
-
-        def analyse(sel, label, shift_by_sink):
-            idx = np.argwhere(sel)
-            if len(idx) == 0:
-                print(f"  {label}: no rows", flush=True)
-                return
-            if len(idx) > MAX_ROWS:
-                idx = idx[rng.choice(len(idx), MAX_ROWS, replace=False)]
-            out = {"plain": [], "shift": []}
-            for b in np.unique(idx[:, 0]):
-                rows = idx[idx[:, 0] == b]
-                t_idx, n_idx = rows[:, 1], rows[:, 2]
-                q_rows = q[b, t_idx, n_idx]
-                d_rows = d[b, t_idx, n_idx]
-                m_rows = mask[b, t_idx]
-                for key, shift in (("plain", None), ("shift", v[b, SINK] if shift_by_sink else None)):
-                    if key == "shift" and not shift_by_sink:
-                        continue
-                    r = row_anatomy(q_rows, k[b], v[b], m_rows, d_rows, shift=shift)
-                    eps = 1.0 - r["p"].max(1)
-                    dq_n = norms(r["dq"])
-                    naive = eps * norms(d_rows) * med_v * med_k  # first-order scale of a peaked row's dQ
-                    ds_sink = np.abs(np.take_along_axis(r["ds"], r["p"].argmax(1)[:, None], 1)[:, 0])
-                    h1_n = norms(r["err_h1"])
-                    k_arg = k[b][r["p"].argmax(1)]  # R H
-                    cos_k = np.sum(r["err_h1"] * k_arg, 1) / np.maximum(h1_n * norms(k_arg), 1e-300)
-                    out[key].append(
-                        np.stack(
-                            [eps, dq_n, naive, ds_sink * norms(k_arg), np.abs(r["d_true"]), np.abs(r["delta"]), h1_n, cos_k, norms(r["o"])],
-                            1,
-                        )
-                    )
-            for key, chunks in out.items():
-                if not chunks:
-                    continue
-                a = np.concatenate(chunks)
-                eps, dq_n, naive, dssink_k, d_true, delta, h1_n, cos_k, o_n = a.T
-                gm = lambda x: float(np.exp(np.mean(np.log(np.maximum(x, 1e-300)))))
-                print(
-                    f"  {label:34s} [{key:5s}] rows={len(a):5d}  eps={gm(eps):.2e}  |dQ_true|={gm(dq_n):.2e}"
-                    f"  naive eps|dO||V||K|={gm(naive):.2e}  |dS_sink||K|={gm(dssink_k):.2e}  |D|={gm(d_true):.2e}"
-                    f"  |delta_H1|={gm(delta):.2e}  |O|={gm(o_n):.2e}"
-                    f"  |errH1|/|dQ_true|: gm={gm(h1_n / np.maximum(dq_n, 1e-300)):.2e} med={np.median(h1_n / np.maximum(dq_n, 1e-300)):.2e}"
-                    f"  cos(errH1,K_argmax)={np.mean(np.abs(cos_k)):.2f}",
-                    flush=True,
-                )
-
-        analyse(bos_rows & (max_p > 0.99), "BOS rows P>0.99", True)
-        analyse(bos_rows & (max_p > 0.9) & (max_p <= 0.99), "BOS rows 0.9<P<=0.99", True)
-        analyse(bos_rows & (max_p <= 0.9), "BOS rows P<=0.9", True)
-        analyse(peaked & ~bos_rows, "other peaked rows P>0.99", False)
-        analyse(vrow & ~bos_rows & (max_p <= 0.9), "ordinary rows (argmax!=768,P<=0.9)", False)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/attention_dterm_probe.py b/scripts/attention_dterm_probe.py
deleted file mode 100644
index 2d1347f..0000000
--- a/scripts/attention_dterm_probe.py
+++ /dev/null
@@ -1,449 +0,0 @@
-"""Test the flash-attention D-term hypothesis on real Pi0.5 activations.
-
-Mechanism under test (confirmed in docs/training-optimization.md, section 4): the cuDNN
-fused backward forms ``dS = P * (dP - D)`` with ``D_i = rowsum(dO_i * O_i)`` taken
-from the *bf16-rounded* stored output ``O``. On attention-sink rows (P nearly
-one-hot) ``dP`` and ``D`` cancel almost exactly, so the rounding error of ``O`` is
-a large fraction of ``dS`` and breaks the softmax invariant ``sum_j dS_ij = 0``.
-The explicit path forms ``D = sum_j P_ij dP_ij`` from the same rounded ``dP`` and
-keeps the invariant.
-
-If that is the mechanism, the cuDNN gradient error must be *predictable*::
-
-    delta_i   = dO_i . (bf16(O_i) - O_i)               (per query row / head)
-    err dQ_i  = -delta_i * sum_j P_ij K_j
-    err dK_j  = -sum_i delta_i P_ij Q_i
-
-This script captures q/k/v/mask and the true backward cotangent dO of every Gemma
-layer from a real training step (real batch, pi05_base weights), then on each layer
-compares, against an fp32 explicit reference:
-
-    explicit_bf16      production path (known stable)
-    cudnn_bf16         historical cuDNN path (known to diverge)
-    cudnn_fp16         cuDNN kernel in float16 + dynamic scaling (candidate)
-    cudnn_bf16_shift   cuDNN with V shifted by the sink key's V (candidate; makes
-                       O of sink rows small so its rounding error shrinks)
-    cudnn_fp16_shift   both
-
-and reports rel error, the cosine between each variant's error and the predicted
-D-term error, and the residual after subtracting the prediction. It also prints
-where the attention sinks are (position histogram) and the fp16 range headroom.
-
-Run on ONE GPU:
-    CUDA_VISIBLE_DEVICES=7 XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 PYTHONPATH=$PWD/src \
-        python scripts/attention_dterm_probe.py
-"""
-
-import dataclasses
-import functools
-import logging
-import sys
-import time
-
-import einops
-import flax.nnx as nnx
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-import openpi.models.gemma as _gemma
-import openpi.training.config as _config
-import openpi.training.data_loader as _data_loader
-import openpi.training.sharding as sharding
-
-sys.path.insert(0, "scripts")
-import train as train_mod  # noqa: E402
-
-CONFIG = "diag_cudnn_strict_order"
-CAPTURE_FILE = "/root/localstorage/logs/dterm_capture.npz"
-BATCH = 8
-WORKERS = 8
-BIG_NEG = -2.3819763e38
-NUM_LAYERS = 18
-
-# --------------------------------------------------------------------------------------
-# Capture q/k/v/mask/d_out of every layer from a real train step.
-# --------------------------------------------------------------------------------------
-_CAPTURE: dict[int, dict[str, np.ndarray]] = {}
-_LAYER = None  # tracer of the layer index of the Block currently being traced
-
-_orig_block_call = _gemma.Block.__call__
-_orig_cudnn_call = _gemma._cudnn_attention_call
-
-
-@functools.wraps(_orig_block_call)
-def _block_call(self, xs, kv_cache, positions, attn_mask, adarms_cond, layer_index=None, deterministic=True):
-    global _LAYER
-    _LAYER = layer_index
-    return _orig_block_call(self, xs, kv_cache, positions, attn_mask, adarms_cond, layer_index, deterministic)
-
-
-def _save_fwd(layer, q, k, v, mask):
-    entry = _CAPTURE.setdefault(int(layer), {})
-    entry.update(q=np.asarray(q), k=np.asarray(k), v=np.asarray(v), mask=np.asarray(mask))
-
-
-def _save_bwd(layer, d_out):
-    _CAPTURE.setdefault(int(layer), {})["d_out"] = np.asarray(d_out)
-
-
-@jax.custom_vjp
-def _tap(q, k, v, mask, layer):
-    return _orig_cudnn_call(q, k, v, mask)
-
-
-def _tap_fwd(q, k, v, mask, layer):
-    jax.debug.callback(_save_fwd, layer, q, k, v, mask)
-    return _tap(q, k, v, mask, layer), (q, k, v, mask, layer)
-
-
-def _tap_bwd(res, d_out):
-    q, k, v, mask, layer = res
-    jax.debug.callback(_save_bwd, layer, d_out)
-    _, vjp_fn = jax.vjp(lambda a, b, c: _orig_cudnn_call(a, b, c, mask), q, k, v)
-    dq, dk, dv = vjp_fn(d_out)
-    return dq, dk, dv, None, None
-
-
-_tap.defvjp(_tap_fwd, _tap_bwd)
-
-
-def _tapped_cudnn_call(q, k, v, attn_mask):
-    return _tap(q, k, v, attn_mask, _LAYER)
-
-
-def capture_real_activations(cfg):
-    _gemma.Block.__call__ = _block_call
-    _gemma._cudnn_attention_call = _tapped_cudnn_call
-    try:
-        mesh = sharding.make_mesh(cfg.fsdp_devices)
-        data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
-        rng = jax.random.key(cfg.seed)
-        train_rng, init_rng = jax.random.split(rng)
-        state, _ = train_mod.init_train_state(cfg, init_rng, mesh, resume=False)
-        params, model_def = state.params, state.model_def
-        del state  # free the Adam moments
-        jax.block_until_ready(params)
-        print("weights loaded", flush=True)
-
-        loader = _data_loader.create_data_loader(cfg, sharding=data_sharding, shuffle=False)
-        obs, act = next(iter(loader))
-        print("batch loaded", flush=True)
-
-        def loss_fn(model, r, o, a):
-            return jnp.mean(model.compute_loss(r, o, a, train=True))
-
-        @jax.jit
-        def run(params, r, o, a):
-            model = nnx.merge(model_def, params)
-            model.train()
-            return nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, cfg.trainable_filter))(model, r, o, a)
-
-        with sharding.set_mesh(mesh):
-            loss, grads = run(params, jax.random.fold_in(train_rng, 0), obs, act)
-            jax.block_until_ready(grads)
-        print(f"real train step: loss={float(loss):.6f}", flush=True)
-        del grads, params
-    finally:
-        _gemma.Block.__call__ = _orig_block_call
-        _gemma._cudnn_attention_call = _orig_cudnn_call
-    missing = [i for i in range(NUM_LAYERS) if i not in _CAPTURE or "d_out" not in _CAPTURE[i]]
-    if missing:
-        raise RuntimeError(f"capture incomplete, missing layers {missing}; got {sorted(_CAPTURE)}")
-
-
-# --------------------------------------------------------------------------------------
-# Single-layer attention variants (mirror gemma.Attention exactly).
-# --------------------------------------------------------------------------------------
-def explicit_attention(q, k, v, mask, compute_dtype, probs_dtype, out_dtype):
-    num_kv = k.shape[2]
-    q, k, v = (x.astype(compute_dtype) for x in (q, k, v))
-    grouped_q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=num_kv)
-    logits = jnp.einsum("BTKGH,BSKH->BKGTS", grouped_q, k, preferred_element_type=jnp.float32)
-    masked = jnp.where(mask[:, :, None, :, :], logits, BIG_NEG)
-    probs = jax.nn.softmax(masked, axis=-1).astype(probs_dtype)
-    enc = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v).astype(out_dtype)
-    return einops.rearrange(enc, "B T K G H -> B T (K G) H")
-
-
-def _shift_v(v, sink_idx):
-    """Return v - v[:, sink] (rounded once) and the per-(B,K) shift c (stop-gradient)."""
-    c = jax.lax.stop_gradient(jnp.take_along_axis(v, sink_idx[:, None, :, None], axis=1))  # B 1 K H
-    v_shift = (v.astype(jnp.float32) - c.astype(jnp.float32)).astype(v.dtype)
-    return v_shift, c
-
-
-def cudnn_shift_attention(q, k, v, mask, sink_idx, compute_dtype):
-    v_shift, c = _shift_v(v, sink_idx)
-    if jnp.dtype(compute_dtype) == q.dtype:
-        out = _orig_cudnn_call(q, k, v_shift, mask)
-    else:
-        out = _gemma._cudnn_attention_in_dtype(q, k, v_shift, mask, compute_dtype)
-    has_key = jnp.any(mask, axis=-1)[:, 0, :, None, None]  # B T 1 1
-    c_full = einops.repeat(c, "B 1 K H -> B 1 (K G) H", G=q.shape[2] // k.shape[2])
-    return (out.astype(jnp.float32) + jnp.where(has_key, c_full.astype(jnp.float32), 0.0)).astype(out.dtype)
-
-
-VARIANT_NAMES = ("explicit_bf16", "cudnn_bf16", "cudnn_fp16", "cudnn_bf16_shift", "cudnn_fp16_shift")
-
-
-def _variant_output(name, q, k, v, mask, sink_idx):
-    if name == "explicit_bf16":
-        return explicit_attention(q, k, v, mask, jnp.bfloat16, jnp.bfloat16, jnp.bfloat16)
-    if name == "cudnn_bf16":
-        return _orig_cudnn_call(q, k, v, mask)
-    if name == "cudnn_fp16":
-        return _gemma._cudnn_attention_in_dtype(q, k, v, mask, jnp.float16)
-    if name == "cudnn_bf16_shift":
-        return cudnn_shift_attention(q, k, v, mask, sink_idx, jnp.bfloat16)
-    if name == "cudnn_fp16_shift":
-        return cudnn_shift_attention(q, k, v, mask, sink_idx, jnp.float16)
-    raise ValueError(name)
-
-
-def grads_of(fn, q, k, v, mask, d_out):
-    d32 = d_out.astype(jnp.float32)
-
-    def objective(q_, k_, v_):
-        return jnp.sum(fn(q_, k_, v_, mask).astype(jnp.float32) * d32)
-
-    return jax.grad(objective, argnums=(0, 1, 2))(q, k, v)
-
-
-@functools.partial(jax.jit, static_argnums=0)
-def variant_grads(name, q, k, v, mask, d_out, sink_idx):
-    return grads_of(lambda a, b, c, m: _variant_output(name, a, b, c, m, sink_idx), q, k, v, mask, d_out)
-
-
-@jax.jit
-def reference_and_prediction(q, k, v, mask, d_out):
-    """fp32 explicit gradients + the predicted D-term error of a bf16 flash backward."""
-    num_kv = k.shape[2]
-    q32, k32, v32, d32 = (x.astype(jnp.float32) for x in (q, k, v, d_out))
-    ref = grads_of(
-        lambda a, b, c, m: explicit_attention(a, b, c, m, jnp.float32, jnp.float32, jnp.float32), q32, k32, v32, mask, d_out
-    )
-
-    grouped_q = einops.rearrange(q32, "B T (K G) H -> B T K G H", K=num_kv)
-    logits = jnp.einsum("BTKGH,BSKH->BKGTS", grouped_q, k32, preferred_element_type=jnp.float32)
-    masked = jnp.where(mask[:, :, None, :, :], logits, BIG_NEG)
-    probs = jax.nn.softmax(masked, axis=-1)  # B K G T S
-    o32 = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v32)
-    # reduce_precision emulates the bf16 rounding; a plain astype round-trip is deleted by XLA
-    # (xla_allow_excess_precision) and would make the prediction identically zero.
-    o16 = jax.lax.reduce_precision(o32, exponent_bits=8, mantissa_bits=7)
-    d_o = einops.rearrange(d32, "B T (K G) H -> B T K G H", K=num_kv)
-    # H1: D formed from the bf16-rounded stored output O.
-    delta = jnp.sum(d_o * (o16 - o32), axis=-1)  # B T K G
-    pk = jnp.einsum("BKGTS,BSKH->BTKGH", probs, k32)
-    pred_dq = einops.rearrange(-delta[..., None] * pk, "B T K G H -> B T (K G) H")
-    pred_dk = -jnp.einsum("BTKG,BKGTS,BTKGH->BSKH", delta, probs, grouped_q)
-    # H2: dP = dO V^T rounded to bf16 before dS = P (dP - D) with an exact D.
-    dp = jnp.einsum("BTKGH,BSKH->BKGTS", d_o, v32)
-    ddp = jax.lax.reduce_precision(dp, exponent_bits=8, mantissa_bits=7) - dp
-    pred2_dq = einops.rearrange(jnp.einsum("BKGTS,BSKH->BTKGH", probs * ddp, k32), "B T K G H -> B T (K G) H")
-
-    valid_q = jnp.any(mask, axis=-1)[:, 0, :]  # B T
-    max_p = jnp.max(probs, axis=-1)  # B K G T
-    arg_p = jnp.argmax(probs, axis=-1)  # B K G T
-    mass = jnp.einsum("BKGTS,BT->BKS", probs, valid_q.astype(jnp.float32))  # B K S
-    return ref, (pred_dq, pred_dk, pred2_dq), dict(valid_q=valid_q, max_p=max_p, arg_p=arg_p, mass=mass, delta=delta)
-
-
-# --------------------------------------------------------------------------------------
-# Metrics (float64 on host).
-# --------------------------------------------------------------------------------------
-def _f64(x):
-    return np.asarray(jax.device_get(x)).astype(np.float64)
-
-
-def _norm(x):
-    return float(np.sqrt(np.sum(x * x)))
-
-
-def _cos(a, b):
-    return float(np.sum(a * b) / max(_norm(a) * _norm(b), 1e-300))
-
-
-def _region(pos):
-    if pos < 768:
-        return f"img{pos // 256}[{pos % 256}]"
-    if pos < 968:
-        return f"txt[{pos - 768}]"
-    return f"act[{pos - 968}]"
-
-
-def analyse_layer(layer, entry):
-    q, k, v, mask, d_out = (jnp.asarray(entry[n]) for n in ("q", "k", "v", "mask", "d_out"))
-    ref, (pred_dq, pred_dk, pred2_dq), stats = reference_and_prediction(q, k, v, mask, d_out)
-    ref = [_f64(g) for g in ref]
-    pred = [_f64(pred_dq), _f64(pred_dk)]
-    pred2 = _f64(pred2_dq)
-    k_np = _f64(k)  # B S K H
-    valid_q = np.asarray(stats["valid_q"])
-    max_p = np.asarray(stats["max_p"])  # B K G T
-    arg_p = np.asarray(stats["arg_p"])
-    mass = np.asarray(stats["mass"])  # B K S
-
-    # ---- sink statistics ----
-    d_np = np.asarray(entry["d_out"]).astype(np.float32)
-    masked_rows_dout = float(np.abs(d_np[~valid_q]).max()) if (~valid_q).any() else 0.0
-    vq = np.broadcast_to(valid_q[:, None, None, :], max_p.shape)
-    mp = max_p[vq]
-    ap = arg_p[vq]
-    frac = {t: float(np.mean(mp > t)) for t in (0.9, 0.99, 0.999)}
-    # per (B, K) sink key = argmax of attention mass
-    sink_idx = np.argmax(mass, axis=-1)  # B K
-    sink_share = mass.max(axis=-1) / mass.sum(axis=-1)  # fraction of all valid attention mass on the sink key
-    pos_counts = np.bincount(ap, minlength=q.shape[1])
-    top = np.argsort(pos_counts)[::-1][:4]
-    top_desc = ", ".join(f"{_region(int(p))}={pos_counts[p] / ap.size:.2%}" for p in top)
-    print(
-        f"\n=== layer {layer:2d} ===  valid rows maxP>0.9/0.99/0.999: {frac[0.9]:.3f}/{frac[0.99]:.3f}/{frac[0.999]:.3f}"
-        f"   sink key per (b,kv): {sorted(set(sink_idx.ravel().tolist()))} share={sink_share.mean():.2%}"
-        f"   argmax top: {top_desc}"
-        f"   max|q|={float(jnp.abs(q).max()):.1f} |k|={float(jnp.abs(k).max()):.1f} |v|={float(jnp.abs(v).max()):.1f}"
-        f" |dO|={float(jnp.abs(d_out).max()):.2e} |dO| on masked rows={masked_rows_dout:.1e}",
-        flush=True,
-    )
-
-    # ---- gradient variants ----
-    sink_idx_dev = jnp.asarray(sink_idx)
-
-    sink_rows = np.transpose(max_p > 0.99, (0, 3, 1, 2)).reshape(max_p.shape[0], max_p.shape[3], -1)  # B T N
-    arg_rows = np.transpose(arg_p, (0, 3, 1, 2)).reshape(sink_rows.shape)  # B T N: each row's own peak key
-    valid_rows = np.broadcast_to(valid_q[:, :, None], sink_rows.shape)
-    print(
-        f"  {'variant':18s} {'dQ rel':>9s} {'cos(err,pred)':>13s} {'resid':>9s} {'sinkrow rel':>11s} {'other rel':>9s} |"
-        f" {'dK rel':>9s} {'cos':>7s} {'resid':>9s} | {'dV rel':>9s} | nonfinite",
-        flush=True,
-    )
-    results = {}
-    for name in VARIANT_NAMES:
-        t0 = time.time()
-        g = variant_grads(name, q, k, v, mask, d_out, sink_idx_dev)
-        g = [_f64(x) for x in g]
-        nonfinite = sum(int(np.count_nonzero(~np.isfinite(x))) for x in g)
-        err = [g[i] - ref[i] for i in range(3)]
-        row = {}
-        for i, nm in enumerate(("dq", "dk", "dv")):
-            row[f"{nm}_rel"] = _norm(err[i]) / _norm(ref[i])
-            if i < 2:
-                row[f"{nm}_cos"] = _cos(err[i], pred[i])
-                row[f"{nm}_resid"] = _norm(err[i] - pred[i]) / _norm(ref[i])
-        # dQ error split by sink rows / other valid rows
-        e_rows = np.sqrt(np.sum(err[0] ** 2, axis=-1))  # B T N
-        r_rows = np.sqrt(np.sum(ref[0] ** 2, axis=-1))
-        sel_s = sink_rows & valid_rows
-        sel_o = (~sink_rows) & valid_rows
-        row["dq_sink_rel"] = _norm(e_rows[sel_s]) / max(_norm(r_rows[sel_s]), 1e-300)
-        row["dq_other_rel"] = _norm(e_rows[sel_o]) / max(_norm(r_rows[sel_o]), 1e-300)
-            # --- peaked-row structure: is the error along the row's peak key (a softmax row-sum violation)? ---
-        if sel_s.any():
-            e_s = err[0][sel_s]  # R H
-            b_idx = np.nonzero(sel_s)[0]
-            ks = k_np[b_idx, arg_rows[sel_s], 0, :]  # R H: K of each peaked row's own argmax key (single kv head)
-            ks_unit = ks / np.linalg.norm(ks, axis=-1, keepdims=True)
-            along = np.sum(e_s * ks_unit, axis=-1)
-            row["sink_frac_along_ksink"] = float(np.sum(along**2) / max(np.sum(e_s**2), 1e-300))
-            row["sink_cos_h1"] = _cos(e_s, pred[0][sel_s])
-            row["sink_cos_h2"] = _cos(e_s, pred2[sel_s])
-            row["sink_cos_h12"] = _cos(e_s, pred[0][sel_s] + pred2[sel_s])
-            row["sink_resid_h1"] = _norm(e_s - pred[0][sel_s]) / max(_norm(r_rows[sel_s]), 1e-300)
-            row["sink_resid_h12"] = _norm(e_s - pred[0][sel_s] - pred2[sel_s]) / max(_norm(r_rows[sel_s]), 1e-300)
-        else:
-            for key in ("sink_frac_along_ksink", "sink_cos_h1", "sink_cos_h2", "sink_cos_h12", "sink_resid_h1", "sink_resid_h12"):
-                row[key] = float("nan")
-        results[name] = row
-        print(
-            f"  {name:18s} {row['dq_rel']:9.3e} {row['dq_cos']:13.4f} {row['dq_resid']:9.3e} {row['dq_sink_rel']:11.3e}"
-            f" {row['dq_other_rel']:9.3e} | {row['dk_rel']:9.3e} {row['dk_cos']:7.4f} {row['dk_resid']:9.3e} |"
-            f" {row['dv_rel']:9.3e} | {nonfinite} ({time.time() - t0:.1f}s)"
-            f"\n      sink rows: frac(err along K_peak)={row['sink_frac_along_ksink']:.3f}  cos(err,H1)={row['sink_cos_h1']:+.3f}"
-            f"  cos(err,H2)={row['sink_cos_h2']:+.3f}  cos(err,H1+H2)={row['sink_cos_h12']:+.3f}"
-            f"  resid H1={row['sink_resid_h1']:.3e}  resid H1+H2={row['sink_resid_h12']:.3e}",
-            flush=True,
-        )
-    pred_rel = [_norm(pred[i]) / _norm(ref[i]) for i in range(2)]
-    n_sink = int(sel_s.sum())
-    sink_pred = _norm(pred[0][sel_s]) / max(_norm(r_rows[sel_s]), 1e-300) if n_sink else float("nan")
-    sink_pred2 = _norm(pred2[sel_s]) / max(_norm(r_rows[sel_s]), 1e-300) if n_sink else float("nan")
-    print(
-        f"  predicted error norm (rel to fp32 grads): H1 dQ {pred_rel[0]:.3e} dK {pred_rel[1]:.3e};"
-        f" on {n_sink} sink rows: H1 {sink_pred:.3e}  H2 {sink_pred2:.3e}",
-        flush=True,
-    )
-    results["_pred_rel"] = pred_rel
-    results["_sink"] = dict(frac=frac, sink_idx=sink_idx.ravel().tolist(), share=float(sink_share.mean()))
-    return results
-
-
-def main():
-    logging.basicConfig(level=logging.WARNING)
-    cfg = _config.get_config(CONFIG)
-    cfg = dataclasses.replace(
-        cfg,
-        exp_name="attention_dterm_probe",
-        batch_size=BATCH,
-        num_workers=WORKERS,
-        fsdp_devices=1,
-        wandb_enabled=False,
-        model=dataclasses.replace(
-            cfg.model,
-            use_cudnn_attention=True,
-            cudnn_attention_layer_start=0,
-            cudnn_attention_num_layers=None,
-            cudnn_attention_dtype="bfloat16",
-            explicit_attention_fp32=False,
-        ),
-    )
-    assert jax.device_count() == 1, "run with CUDA_VISIBLE_DEVICES=<one gpu>"
-    t0 = time.time()
-    import os
-
-    if os.path.exists(CAPTURE_FILE):
-        data = np.load(CAPTURE_FILE)
-        import ml_dtypes
-
-        for key in data.files:
-            layer, name = key.split("/")
-            arr = data[key]
-            if arr.dtype.kind == "V" and arr.dtype.itemsize == 2:  # bf16 saved by np.savez as raw bytes
-                arr = arr.view(ml_dtypes.bfloat16)
-            _CAPTURE.setdefault(int(layer), {})[name] = arr
-        print(f"loaded capture from {CAPTURE_FILE}", flush=True)
-    else:
-        capture_real_activations(cfg)
-        np.savez(CAPTURE_FILE, **{f"{l}/{n}": a for l, e in _CAPTURE.items() for n, a in e.items()})
-        print(f"captured {len(_CAPTURE)} layers in {time.time() - t0:.0f}s -> {CAPTURE_FILE}", flush=True)
-    e = _CAPTURE[0]
-    print(f"shapes q{e['q'].shape} k{e['k'].shape} v{e['v'].shape} mask{e['mask'].shape} dO{e['d_out'].shape}", flush=True)
-
-    per_layer = {}
-    for layer in range(NUM_LAYERS):
-        per_layer[layer] = analyse_layer(layer, _CAPTURE[layer])
-
-    names = [n for n in per_layer[0] if not n.startswith("_")]
-    print("\n=== SUMMARY over 18 layers (geometric mean of rel errors, mean of cosines) ===", flush=True)
-    print(f"  {'variant':18s} {'dQ rel':>9s} {'cos(err,pred)':>13s} {'resid':>9s} {'sinkrow rel':>11s} | {'dK rel':>9s} {'cos':>7s} {'resid':>9s} | {'dV rel':>9s}")
-    for name in names:
-        gm = lambda key: float(np.exp(np.mean([np.log(max(per_layer[l][name][key], 1e-300)) for l in per_layer])))
-        mean = lambda key: float(np.mean([per_layer[l][name][key] for l in per_layer]))
-        print(
-            f"  {name:18s} {gm('dq_rel'):9.3e} {mean('dq_cos'):13.4f} {gm('dq_resid'):9.3e} {gm('dq_sink_rel'):11.3e} |"
-            f" {gm('dk_rel'):9.3e} {mean('dk_cos'):7.4f} {gm('dk_resid'):9.3e} | {gm('dv_rel'):9.3e}",
-            flush=True,
-        )
-    for name in names:
-        mean = lambda key: float(np.nanmean([per_layer[l][name][key] for l in per_layer]))
-        print(
-            f"  {name:18s} sink rows: frac along K_peak={mean('sink_frac_along_ksink'):.3f}"
-            f" cos H1={mean('sink_cos_h1'):+.3f} cos H2={mean('sink_cos_h2'):+.3f} cos H1+H2={mean('sink_cos_h12'):+.3f}"
-        )
-    print("  predicted D-term rel norm per layer (dQ):", " ".join(f"{per_layer[l]['_pred_rel'][0]:.1e}" for l in per_layer))
-    print("  sink share per layer:", " ".join(f"{per_layer[l]['_sink']['share']:.2f}" for l in per_layer))
-    print("  sink keys per layer:", [sorted(set(per_layer[l]["_sink"]["sink_idx"])) for l in per_layer])
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/attention_precision_update_probe.py b/scripts/attention_precision_update_probe.py
deleted file mode 100644
index b1a04ec..0000000
--- a/scripts/attention_precision_update_probe.py
+++ /dev/null
@@ -1,219 +0,0 @@
-"""Compare attention implementations against an fp32 attention reference at a real Adam state.
-
-The 2026-08-31 report showed that the full-18-layer cuDNN path diverges after
-~step 1000 although its raw gradient is only ~1.5% away from the explicit bf16
-path. This probe restores the real step-1000 train state (params + Adam moments),
-computes the gradient and the actual AdamW update for the same real batch/rng under
-
-  explicit_fp32ref  explicit attention with fp32 q/k/v/probs   (numerical reference)
-  explicit_bf16     production explicit path                    (known stable)
-  cudnn_bf16        historical cuDNN path                       (known to diverge)
-  cudnn_fp16        cuDNN kernel in float16 + dynamic loss scaling (candidate)
-
-and reports, for every variant, raw-gradient and Adam-update distance to the
-fp32 reference and to the explicit bf16 path. A candidate is promising when its
-update-space fingerprint (relative delta / sign-flip fraction) is at or below the
-explicit bf16 path's own distance from the fp32 reference.
-"""
-
-import dataclasses
-import functools
-import logging
-import sys
-import time
-
-import flax.nnx as nnx
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-import openpi.training.checkpoints as _checkpoints
-import openpi.training.config as _config
-import openpi.training.data_loader as _data_loader
-import openpi.training.sharding as sharding
-
-sys.path.insert(0, "scripts")
-import train as train_mod
-
-# Usage: python scripts/attention_precision_update_probe.py [config] [exp_name] [checkpoint_step]
-# The checkpoint only supplies the parameters and Adam moments; every variant is
-# evaluated from that same state, so any clean step-1000 run works as the anchor.
-CONFIG = sys.argv[1] if len(sys.argv) > 1 else "diag_cudnn_strict_order"
-EXP_NAME = sys.argv[2] if len(sys.argv) > 2 else CONFIG
-CHECKPOINT_STEP = int(sys.argv[3]) if len(sys.argv) > 3 else 1000
-BATCH = 32
-WORKERS = 16
-NUM_BATCHES = 2
-
-
-def _loss_grads_update(model_def, trainable_filter, tx, opt_state, params, rng, obs, act):
-    model = nnx.merge(model_def, params)
-    model.train()
-
-    def loss_fn(m, r, o, a):
-        return jnp.mean(m.compute_loss(r, o, a, train=True))
-
-    loss, grads = nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, trainable_filter))(model, rng, obs, act)
-    trainable_params = params.filter(trainable_filter)
-    updates, _ = tx.update(grads, opt_state, trainable_params)
-    return loss, grads, updates
-
-
-def _flat(tree):
-    pure = tree.to_pure_dict()
-    out = []
-
-    def rec(value, prefix):
-        if isinstance(value, dict):
-            for key, child in value.items():
-                rec(child, f"{prefix}/{key}" if prefix else str(key))
-        else:
-            out.append((prefix, value))
-
-    rec(pure, "")
-    return out
-
-
-def _to_host(tree):
-    return {path: np.asarray(jax.device_get(value), np.float32) for path, value in _flat(tree)}
-
-
-def _dot(a, b):
-    return float(np.vdot(a.astype(np.float64), b.astype(np.float64)))
-
-
-def _metrics(reference, candidate):
-    totals = {"ALL": [0.0, 0.0, 0.0, 0.0, 0, 0]}
-    for path, ref in reference.items():
-        cand = candidate[path]
-        delta = cand - ref
-        group = path.split("/", 1)[0]
-        totals.setdefault(group, [0.0, 0.0, 0.0, 0.0, 0, 0])
-        for bucket in (totals["ALL"], totals[group]):
-            bucket[0] += _dot(ref, ref)
-            bucket[1] += _dot(cand, cand)
-            bucket[2] += _dot(delta, delta)
-            bucket[3] += _dot(ref, cand)
-            nonzero = (ref != 0) | (cand != 0)
-            bucket[4] += int(np.count_nonzero(nonzero & (np.signbit(ref) != np.signbit(cand))))
-            bucket[5] += int(np.count_nonzero(nonzero))
-    return totals
-
-
-def _summary(reference, candidate):
-    ref2, cand2, delta2, dot, flips, count = _metrics(reference, candidate)["ALL"]
-    return (
-        np.sqrt(delta2 / max(ref2, 1e-300)),
-        dot / max(np.sqrt(ref2 * cand2), 1e-300),
-        flips / max(count, 1),
-    )
-
-
-def _print_comparison(label, reference, candidate, per_group=False):
-    metrics = _metrics(reference, candidate)
-    ref2, cand2, delta2, dot, flips, count = metrics["ALL"]
-    print(
-        f"  {label:44s} rel_delta={np.sqrt(delta2 / max(ref2, 1e-300)):.4e} "
-        f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.7f} sign_flips={flips / max(count, 1):.4e}",
-        flush=True,
-    )
-    if per_group:
-        for group in sorted(key for key in metrics if key != "ALL"):
-            ref2, cand2, delta2, dot, flips, count = metrics[group]
-            print(
-                f"      {group:22s} rel={np.sqrt(delta2 / max(ref2, 1e-300)):.4e} "
-                f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.7f} sign_flips={flips / max(count, 1):.4e}",
-                flush=True,
-            )
-
-
-def _nonfinite(host):
-    return sum(int(np.count_nonzero(~np.isfinite(v))) for v in host.values())
-
-
-def main():
-    logging.basicConfig(level=logging.WARNING)
-    cfg = _config.get_config(CONFIG)
-    cfg = dataclasses.replace(
-        cfg,
-        exp_name=EXP_NAME,
-        batch_size=BATCH,
-        num_workers=WORKERS,
-        resume=True,
-        overwrite=False,
-        wandb_enabled=False,
-    )
-
-    def variant(**kw):
-        return dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, **kw))
-
-    variant_cfgs = {
-        "explicit_fp32ref": variant(use_cudnn_attention=False, explicit_attention_fp32=True),
-        "explicit_bf16": variant(use_cudnn_attention=False),
-        "cudnn_bf16": variant(use_cudnn_attention=True, cudnn_attention_dtype="bfloat16"),
-        "cudnn_fp16": variant(use_cudnn_attention=True, cudnn_attention_dtype="float16"),
-    }
-
-    mesh = sharding.make_mesh(cfg.fsdp_devices)
-    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
-    rng = jax.random.key(cfg.seed)
-    train_rng, init_rng = jax.random.split(rng)
-
-    cfg_ckpt = variant_cfgs["cudnn_bf16"]
-    manager, resuming = _checkpoints.initialize_checkpoint_dir(
-        cfg_ckpt.checkpoint_dir, keep_period=cfg_ckpt.keep_period, overwrite=False, resume=True
-    )
-    if not resuming or CHECKPOINT_STEP not in manager.all_steps():
-        raise RuntimeError(f"checkpoint {CHECKPOINT_STEP} unavailable; found {manager.all_steps()}")
-    state_shape, _ = train_mod.init_train_state(cfg_ckpt, init_rng, mesh, resume=True)
-    state = _checkpoints.restore_state(manager, state_shape, None, step=CHECKPOINT_STEP)
-    jax.block_until_ready(state)
-    print(f"restored checkpoint step={CHECKPOINT_STEP}; train_state.step={int(state.step)}", flush=True)
-
-    print("building loader ...", flush=True)
-    loader = _data_loader.create_data_loader(variant_cfgs["explicit_bf16"], sharding=data_sharding, shuffle=True)
-    it = iter(loader)
-    batches = [next(it) for _ in range(NUM_BATCHES)]
-
-    graphdefs = {}
-    for name, vcfg in variant_cfgs.items():
-        graphdefs[name] = nnx.graphdef(vcfg.model.create(init_rng))
-    print("graphdefs built", flush=True)
-
-    fns = {
-        name: jax.jit(functools.partial(_loss_grads_update, graphdefs[name], cfg.trainable_filter, state.tx))
-        for name in variant_cfgs
-    }
-
-    for bi, (obs, act) in enumerate(batches):
-        step_rng = jax.random.fold_in(train_rng, state.step + bi)
-        grads_host, updates_host, losses = {}, {}, {}
-        for name, fn in fns.items():
-            t0 = time.time()
-            with sharding.set_mesh(mesh):
-                loss, grads, updates = fn(state.opt_state, state.params, step_rng, obs, act)
-                losses[name] = float(loss)
-                grads_host[name] = _to_host(grads)
-                updates_host[name] = _to_host(updates)
-            del grads, updates
-            print(
-                f"batch {bi} {name:18s} loss={losses[name]:.8f} "
-                f"nonfinite grads={_nonfinite(grads_host[name])} updates={_nonfinite(updates_host[name])} "
-                f"({time.time() - t0:.1f}s)",
-                flush=True,
-            )
-
-        for ref_name in ("explicit_fp32ref", "explicit_bf16"):
-            print(f"\n=== batch {bi}: RAW GRADIENT vs {ref_name} ===", flush=True)
-            for name in variant_cfgs:
-                if name != ref_name:
-                    _print_comparison(name, grads_host[ref_name], grads_host[name])
-            print(f"\n=== batch {bi}: ADAM UPDATE vs {ref_name} ===", flush=True)
-            for name in variant_cfgs:
-                if name != ref_name:
-                    _print_comparison(name, updates_host[ref_name], updates_host[name], per_group=(ref_name == "explicit_fp32ref"))
-        print("", flush=True)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/check_cuda_stack.py b/scripts/check_cuda_stack.py
index c6d5e4e..cf55224 100644
--- a/scripts/check_cuda_stack.py
+++ b/scripts/check_cuda_stack.py
@@ -57,12 +57,12 @@ def report(name, pattern):
 
 def main():
     import jax
+    from jax._src.lib import cuda_versions
     import jax.numpy as jnp
     import numpy as np
     import torch
     import torchcodec
     import torchvision
-    from jax._src.lib import cuda_versions
 
     section("1. reported versions")
     rt, build = cuda_versions.cudnn_get_version(), cuda_versions.cudnn_build_version()
@@ -125,36 +125,36 @@ def main():
     section("3. production attention shape: forward + backward vs fp32 reference")
     # pi0.5 training shape: 3x256 image + 200 prompt + 50 action tokens, GQA 8:1,
     # head_dim 256, bf16, real block mask including fully-masked padding rows.
-    B, NH, NKV, HD = 2, 8, 1, 256
-    T = 768 + 200 + 50
+    batch_size, num_heads, num_kv_heads, head_dim = 2, 8, 1, 256
+    seq_len = 768 + 200 + 50
 
-    im = np.ones((B, T), dtype=bool)
+    im = np.ones((batch_size, seq_len), dtype=bool)
     for b, plen in enumerate((120, 165)):
         im[b, 768 + plen : 768 + 200] = False
-    ar = np.zeros((T,), dtype=bool)
+    ar = np.zeros((seq_len,), dtype=bool)
     ar[768 + 200] = True
     im_j = jnp.asarray(im)
     cs = jnp.broadcast_to(jnp.cumsum(jnp.asarray(ar), -1), im_j.shape)
     mask = jnp.logical_and(cs[:, None, :] <= cs[:, :, None], im_j[:, None, :] * im_j[:, :, None])[:, None]
     valid = np.asarray(jnp.any(mask, -1))[0, 0]
     n_empty = int((~np.asarray(jnp.any(mask, -1))).sum())
-    print(f"  shape B={B} T={T} heads={NH} kv={NKV} head_dim={HD} bf16")
+    print(f"  shape batch_size={batch_size} seq_len={seq_len} heads={num_heads} kv={num_kv_heads} head_dim={head_dim} bf16")
     print(f"  fully-masked query rows in batch: {n_empty}")
 
     rng = np.random.default_rng(0)
-    q = jnp.asarray(rng.normal(0, 10 / np.sqrt(HD), (B, T, NH, HD)), jnp.bfloat16)
-    k = jnp.asarray(rng.normal(0, 1, (B, T, NKV, HD)), jnp.bfloat16)
-    v = jnp.asarray(rng.normal(0, 1, (B, T, NKV, HD)), jnp.bfloat16)
-    cot = np.asarray(rng.normal(0, 1, (B, T, NH, HD)), np.float32)
-    cot[:, : T - 50] = 0.0  # only action tokens carry loss, as in Pi0
+    q = jnp.asarray(rng.normal(0, 10 / np.sqrt(head_dim), (batch_size, seq_len, num_heads, head_dim)), jnp.bfloat16)
+    k = jnp.asarray(rng.normal(0, 1, (batch_size, seq_len, num_kv_heads, head_dim)), jnp.bfloat16)
+    v = jnp.asarray(rng.normal(0, 1, (batch_size, seq_len, num_kv_heads, head_dim)), jnp.bfloat16)
+    cot = np.asarray(rng.normal(0, 1, (batch_size, seq_len, num_heads, head_dim)), np.float32)
+    cot[:, : seq_len - 50] = 0.0  # only action tokens carry loss, as in Pi0
     cot = jnp.asarray(cot)
 
     def reference(q, k, v, m):  # explicit path, fp32 throughout
         q, k, v = (x.astype(jnp.float32) for x in (q, k, v))
-        gq = q.reshape(B, T, NKV, NH // NKV, HD)
+        gq = q.reshape(batch_size, seq_len, num_kv_heads, num_heads // num_kv_heads, head_dim)
         lg = jnp.einsum("BTKGH,BSKH->BKGTS", gq, k, preferred_element_type=jnp.float32)
         lg = jnp.where(m[:, :, None], lg, -2.3819763e38)
-        return jnp.einsum("BKGTS,BSKH->BTKGH", jax.nn.softmax(lg, -1), v).reshape(B, T, NH, HD)
+        return jnp.einsum("BKGTS,BSKH->BTKGH", jax.nn.softmax(lg, -1), v).reshape(batch_size, seq_len, num_heads, head_dim)
 
     def cudnn(q, k, v, m):  # what training runs with use_cudnn_attention=true
         has_key = jnp.any(m, -1)[:, 0, :, None, None]
@@ -167,7 +167,7 @@ def loss(fn):
     try:
         gr = [np.asarray(x.astype(jnp.float32), np.float64) for x in jax.grad(loss(reference), (0, 1, 2))(q, k, v, mask)]
         gc = [np.asarray(x.astype(jnp.float32), np.float64) for x in jax.grad(loss(cudnn), (0, 1, 2))(q, k, v, mask)]
-    except Exception as exc:  # noqa: BLE001
+    except Exception as exc:
         FAIL.append(f"cuDNN attention raised {type(exc).__name__}: {exc}")
         print(f"  !! cuDNN attention FAILED: {type(exc).__name__}: {str(exc)[:300]}")
         gr = gc = None
diff --git a/scripts/cudnn_attention_microbench.py b/scripts/cudnn_attention_microbench.py
deleted file mode 100644
index ddb344f..0000000
--- a/scripts/cudnn_attention_microbench.py
+++ /dev/null
@@ -1,115 +0,0 @@
-"""Time the attention variants used by gemma.Attention at the training shape on one GPU.
-
-Answers "where does the float16 cuDNN path lose time against the bfloat16 cuDNN
-path?" by timing forward and forward+backward of each variant on real captured
-activations (dterm_capture.npz, B=8, T=S=1018, 8 q heads / 1 kv head, head_dim 256),
-optionally tiled along batch to the per-device training batch.
-
-    CUDA_VISIBLE_DEVICES=7 PYTHONPATH=$PWD/src python scripts/cudnn_attention_microbench.py \
-        /root/localstorage/logs/dterm_capture.npz [batch_tile]
-"""
-
-import sys
-import time
-
-import einops
-import jax
-import jax.numpy as jnp
-import ml_dtypes
-import numpy as np
-
-import openpi.models.gemma as _gemma
-
-CAPTURE_FILE = sys.argv[1] if len(sys.argv) > 1 else "/root/localstorage/logs/dterm_capture.npz"
-TILE = int(sys.argv[2]) if len(sys.argv) > 2 else 4  # 8 * 4 = 32 = per-device batch at global batch 256 on 8 GPUs
-LAYER = 0
-ITERS = 20
-
-
-def explicit_bf16(q, k, v, mask):
-    num_kv_heads = k.shape[2]
-    grouped_q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=num_kv_heads)
-    logits = jnp.einsum("BTKGH,BSKH->BKGTS", grouped_q, k, preferred_element_type=jnp.float32)
-    masked_logits = jnp.where(mask[:, :, None, :, :], logits, -2.3819763e38)
-    probs = jax.nn.softmax(masked_logits, axis=-1).astype(v.dtype)
-    encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v).astype(q.dtype)
-    return einops.rearrange(encoded, "B T K G H -> B T (K G) H")
-
-
-def cudnn_bf16(q, k, v, mask):
-    return _gemma._cudnn_attention_call(q, k, v, mask)
-
-
-def cudnn_fp16_plain(q, k, v, mask):
-    # float16 kernel through jax's own custom_vjp, no loss scaling: isolates kernel cost.
-    q16, k16, v16 = (x.astype(jnp.float16) for x in (q, k, v))
-    return _gemma._cudnn_attention_call(q16, k16, v16, mask).astype(q.dtype)
-
-
-def cudnn_fp16_v2(q, k, v, mask):
-    return _gemma._cudnn_attention_in_dtype(q, k, v, mask, jnp.float16)
-
-
-VARIANTS = {
-    "cudnn_bf16": cudnn_bf16,
-    "cudnn_fp16_plain": cudnn_fp16_plain,
-    "cudnn_fp16_v2": cudnn_fp16_v2,
-    "explicit_bf16": explicit_bf16,
-}
-
-
-def _timeit(fn, *args):
-    out = fn(*args)
-    jax.block_until_ready(out)
-    t0 = time.perf_counter()
-    for _ in range(ITERS):
-        out = fn(*args)
-    jax.block_until_ready(out)
-    return (time.perf_counter() - t0) / ITERS * 1e3
-
-
-def main():
-    data = np.load(CAPTURE_FILE)
-    e = {}
-    for name in ("q", "k", "v", "mask", "d_out"):
-        arr = data[f"{LAYER}/{name}"]
-        if arr.dtype.kind == "V" and arr.dtype.itemsize == 2:
-            arr = arr.view(ml_dtypes.bfloat16)
-        if TILE > 1:
-            arr = np.concatenate([arr] * TILE, axis=0)
-        e[name] = jnp.asarray(arr)
-    q, k, v, mask, d_out = e["q"], e["k"], e["v"], e["mask"], e["d_out"]
-    print(f"shape q={q.shape} k={k.shape} mask={mask.shape} dtype={q.dtype} device={jax.devices()[0]}", flush=True)
-    d32 = d_out.astype(jnp.float32)
-
-    for name, fn in VARIANTS.items():
-        fwd = jax.jit(fn)
-
-        def objective(q_, k_, v_, fn=fn):
-            return jnp.sum(fn(q_, k_, v_, mask).astype(jnp.float32) * d32)
-
-        fwd_bwd = jax.jit(jax.grad(objective, argnums=(0, 1, 2)))
-        # Same structure as training: flax remat(nothing_saveable) recomputes the forward in the backward.
-        fwd_bwd_remat = jax.jit(
-            jax.grad(
-                lambda q_, k_, v_, fn=fn: jnp.sum(
-                    jax.checkpoint(lambda a, b, c: fn(a, b, c, mask), policy=jax.checkpoint_policies.nothing_saveable)(
-                        q_, k_, v_
-                    ).astype(jnp.float32)
-                    * d32
-                ),
-                argnums=(0, 1, 2),
-            )
-        )
-        t_fwd = _timeit(fwd, q, k, v, mask)
-        t_fb = _timeit(fwd_bwd, q, k, v)
-        t_fbr = _timeit(fwd_bwd_remat, q, k, v)
-        print(
-            f"{name:18s} fwd {t_fwd:7.2f} ms | fwd+bwd {t_fb:7.2f} ms | fwd+bwd(remat) {t_fbr:7.2f} ms"
-            f" | x18 layers: fwd+bwd(remat) {t_fbr * 18 / 1e3:.3f} s/step",
-            flush=True,
-        )
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/cudnn_fp16_vjp_equivalence.py b/scripts/cudnn_fp16_vjp_equivalence.py
deleted file mode 100644
index e75f20f..0000000
--- a/scripts/cudnn_fp16_vjp_equivalence.py
+++ /dev/null
@@ -1,112 +0,0 @@
-"""Check that the direct-rule float16 cuDNN VJP in gemma.py matches the recompute-based one.
-
-`gemma._cudnn_attention_in_dtype` now calls JAX's cuDNN fwd/bwd rules directly
-(keeping softmax stats + output as residuals). This script re-implements the
-previous version (jax.vjp over jax.nn.dot_product_attention inside the backward)
-and compares forward outputs and dQ/dK/dV on real captured activations. The two
-backwards run the same kernel on the same inputs, so their difference must be at
-the level of the kernel's own run-to-run (atomic accumulation) noise.
-
-    CUDA_VISIBLE_DEVICES=7 PYTHONPATH=$PWD/src python scripts/cudnn_fp16_vjp_equivalence.py \
-        /root/localstorage/logs/dterm_capture.npz
-"""
-
-import sys
-import time
-
-import jax
-import jax.numpy as jnp
-import ml_dtypes
-import numpy as np
-
-import openpi.models.gemma as _gemma
-
-CAPTURE_FILE = sys.argv[1] if len(sys.argv) > 1 else "/root/localstorage/logs/dterm_capture.npz"
-LAYERS = (0, 4, 8, 14, 17)
-
-
-def _reference_in_dtype(q, k, v, attn_mask, compute_dtype):
-    """The previous implementation: recompute the forward inside the backward via jax.vjp."""
-    out_dtype = q.dtype
-
-    @jax.custom_vjp
-    def attention(q, k, v, attn_mask):
-        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
-        return _gemma._cudnn_attention_call(qc, kc, vc, attn_mask).astype(out_dtype)
-
-    def attention_fwd(q, k, v, attn_mask):
-        return attention(q, k, v, attn_mask), (q, k, v, attn_mask)
-
-    def attention_bwd(residuals, d_out):
-        q, k, v, attn_mask = residuals
-        qc, kc, vc = (x.astype(compute_dtype) for x in (q, k, v))
-        d_out32 = d_out.astype(jnp.float32)
-        max_abs = jnp.max(jnp.abs(d_out32))
-        tiny = jnp.finfo(jnp.float32).tiny
-        exponent = jnp.floor(jnp.log2(jnp.maximum(max_abs, tiny)))
-        scale = jnp.where(max_abs > 0, jnp.exp2(-(exponent + 2.0)), 1.0)
-        _, vjp_fn = jax.vjp(lambda a, b, c: _gemma._cudnn_attention_call(a, b, c, attn_mask), qc, kc, vc)
-        dq, dk, dv = vjp_fn((d_out32 * scale).astype(compute_dtype))
-        inv_scale = 1.0 / scale
-        return (*((g.astype(jnp.float32) * inv_scale).astype(out_dtype) for g in (dq, dk, dv)), None)
-
-    attention.defvjp(attention_fwd, attention_bwd)
-    return attention(q, k, v, attn_mask)
-
-
-def _fwd_and_grads(fn, q, k, v, mask, d_out):
-    d32 = d_out.astype(jnp.float32)
-
-    def objective(q_, k_, v_):
-        out = fn(q_, k_, v_, mask, jnp.float16)
-        return jnp.sum(out.astype(jnp.float32) * d32), out
-
-    (_, out), grads = jax.value_and_grad(objective, argnums=(0, 1, 2), has_aux=True)(q, k, v)
-    return out, grads
-
-
-new_fn = jax.jit(lambda *a: _fwd_and_grads(_gemma._cudnn_attention_in_dtype, *a))
-old_fn = jax.jit(lambda *a: _fwd_and_grads(_reference_in_dtype, *a))
-
-
-def _f64(x):
-    return np.asarray(jax.device_get(x)).astype(np.float64)
-
-
-def _rel(a, b):
-    return float(np.sqrt(np.sum((a - b) ** 2) / max(np.sum(b**2), 1e-300)))
-
-
-def main():
-    data = np.load(CAPTURE_FILE)
-    for layer in LAYERS:
-        e = {}
-        for name in ("q", "k", "v", "mask", "d_out"):
-            arr = data[f"{layer}/{name}"]
-            if arr.dtype.kind == "V" and arr.dtype.itemsize == 2:
-                arr = arr.view(ml_dtypes.bfloat16)
-            e[name] = jnp.asarray(arr)
-        args = (e["q"], e["k"], e["v"], e["mask"], e["d_out"])
-        t0 = time.time()
-        out_new, g_new = new_fn(*args)
-        _, g_new2 = new_fn(*args)
-        out_old, g_old = old_fn(*args)
-        _, g_old2 = old_fn(*args)
-        jax.block_until_ready((g_new, g_new2, g_old, g_old2))
-        valid_q = np.asarray(jnp.any(e["mask"], axis=-1)[:, 0, :])
-        dq_new = _f64(g_new[0])
-        masked_dq = float(np.abs(dq_new[~valid_q]).max()) if (~valid_q).any() else 0.0
-        nonfinite = sum(int(np.count_nonzero(~np.isfinite(_f64(x)))) for x in (*g_new, out_new))
-        print(
-            f"layer {layer:2d}: fwd new-vs-old rel={_rel(_f64(out_new), _f64(out_old)):.2e}"
-            f" (bit-identical: {bool(np.array_equal(np.asarray(out_new), np.asarray(out_old)))})"
-            f" | new-vs-old dQ {_rel(dq_new, _f64(g_old[0])):.2e} dK {_rel(_f64(g_new[1]), _f64(g_old[1])):.2e}"
-            f" dV {_rel(_f64(g_new[2]), _f64(g_old[2])):.2e}"
-            f" | new run-to-run dQ {_rel(dq_new, _f64(g_new2[0])):.2e} old run-to-run dQ {_rel(_f64(g_old[0]), _f64(g_old2[0])):.2e}"
-            f" | masked-row |dQ| max={masked_dq:.1e} nonfinite={nonfinite} ({time.time() - t0:.1f}s)",
-            flush=True,
-        )
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/grad_ab_direction_probe.py b/scripts/grad_ab_direction_probe.py
deleted file mode 100644
index e6745cb..0000000
--- a/scripts/grad_ab_direction_probe.py
+++ /dev/null
@@ -1,257 +0,0 @@
-"""Measure whether cuDNN-vs-explicit gradient deltas align across real batches.
-
-For each batch, one explicit-attention gradient is compared with the average of
-three cuDNN-attention gradients. Averaging suppresses the known nondeterministic
-atomic-accumulation component before measuring cross-batch directionality.
-"""
-
-import collections
-import dataclasses
-import functools
-import logging
-import sys
-
-import flax.nnx as nnx
-import jax
-import jax.numpy as jnp
-from jax._src.lib import cuda_versions
-import numpy as np
-import optax
-
-import openpi.training.config as _config
-import openpi.training.data_loader as _data_loader
-import openpi.training.sharding as sharding
-
-sys.path.insert(0, "scripts")
-import train as train_mod
-
-
-CONFIG = "pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100"
-BATCH = 32
-WORKERS = 16
-N_BATCHES = 8
-CUDNN_REPEATS = 3
-
-
-def _grad_fn(model_def, trainable_filter, params, rng, obs, act):
-    model = nnx.merge(model_def, params)
-    model.train()
-
-    def loss_fn(m, r, o, a):
-        return jnp.mean(m.compute_loss(r, o, a, train=True))
-
-    return nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, trainable_filter))(model, rng, obs, act)
-
-
-def _flat(state):
-    pure = state.to_pure_dict()
-    out = []
-
-    def rec(d, prefix):
-        if isinstance(d, dict):
-            for key, value in d.items():
-                rec(value, f"{prefix}/{key}" if prefix else str(key))
-        else:
-            out.append((prefix, d))
-
-    rec(pure, "")
-    return out
-
-
-def _to_host(state_tree):
-    return {key: np.asarray(jax.device_get(value), np.float32) for key, value in _flat(state_tree)}
-
-
-def _dot(a, b):
-    # Accumulate each tensor in float64 without materializing a float64 copy of
-    # the whole 3B-parameter model at once.
-    return float(np.vdot(a.astype(np.float64), b.astype(np.float64)))
-
-
-def _module(path):
-    return path.split("/", 1)[0]
-
-
-def _cos(dot, norm_a_sq, norm_b_sq):
-    return dot / max(np.sqrt(norm_a_sq * norm_b_sq), 1e-300)
-
-
-def main():
-    logging.basicConfig(level=logging.WARNING)
-    print("cuDNN runtime:", cuda_versions.cudnn_get_version(), flush=True)
-    print(f"probe: batches={N_BATCHES}, cudnn_repeats={CUDNN_REPEATS}, local_batch={BATCH}", flush=True)
-
-    cfg = _config.get_config(CONFIG)
-    cfg = dataclasses.replace(cfg, batch_size=BATCH, num_workers=WORKERS, wandb_enabled=False)
-    cfg_off = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=False))
-    cfg_on = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=True))
-
-    mesh = sharding.make_mesh(cfg.fsdp_devices)
-    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
-    rng = jax.random.key(cfg.seed)
-    train_rng, init_rng = jax.random.split(rng)
-
-    print("loading pi05_base weights ...", flush=True)
-    state, _ = train_mod.init_train_state(cfg_off, init_rng, mesh, resume=False)
-    host_params = _to_host(state.params)
-    param_norm_sq = sum(_dot(value, value) for value in host_params.values())
-    print("  trainable param norm:", np.sqrt(param_norm_sq), flush=True)
-
-    gdef_off = state.model_def
-    gdef_on = nnx.graphdef(cfg_on.model.create(init_rng))
-    f_off = jax.jit(functools.partial(_grad_fn, gdef_off, cfg.trainable_filter))
-    f_on = jax.jit(functools.partial(_grad_fn, gdef_on, cfg.trainable_filter))
-
-    print("building data loader ...", flush=True)
-    dl = _data_loader.create_data_loader(cfg_off, sharding=data_sharding, shuffle=True)
-    iterator = iter(dl)
-
-    deltas = []
-    delta_norm_sq = []
-    pairwise_dot = np.zeros((N_BATCHES, N_BATCHES), np.float64)
-    per_batch_summary = []
-    per_module_summary = collections.defaultdict(list)
-
-    for batch_idx in range(N_BATCHES):
-        obs, act = next(iterator)
-        step_rng = jax.random.fold_in(train_rng, batch_idx)
-
-        with sharding.set_mesh(mesh):
-            loss_off, g_off = f_off(state.params, step_rng, obs, act)
-            loss_off = float(loss_off)
-            host_off = _to_host(g_off)
-        del g_off
-
-        host_on_sum = None
-        cudnn_losses = []
-        repeat_diffs_sq = []
-        first_on = None
-        for repeat in range(CUDNN_REPEATS):
-            with sharding.set_mesh(mesh):
-                loss_on, g_on = f_on(state.params, step_rng, obs, act)
-                cudnn_losses.append(float(loss_on))
-                host_on = _to_host(g_on)
-            del g_on
-            if host_on_sum is None:
-                host_on_sum = {key: value.copy() for key, value in host_on.items()}
-                first_on = {key: value.copy() for key, value in host_on.items()}
-            else:
-                repeat_diffs_sq.append(sum(_dot(host_on[key] - first_on[key], host_on[key] - first_on[key]) for key in host_on))
-                for key in host_on_sum:
-                    host_on_sum[key] += host_on[key]
-            del host_on
-
-        host_on_avg = {key: value / CUDNN_REPEATS for key, value in host_on_sum.items()}
-        del host_on_sum
-        delta = {key: host_on_avg[key] - host_off[key] for key in host_off}
-
-        totals = collections.defaultdict(lambda: collections.defaultdict(float))
-        previous_dots = np.zeros(batch_idx, np.float64)
-        for path, grad in host_off.items():
-            cudnn_grad = host_on_avg[path]
-            diff = delta[path]
-            param = host_params[path]
-            group = _module(path)
-            for bucket in (totals["ALL"], totals[group]):
-                bucket["g2"] += _dot(grad, grad)
-                bucket["c2"] += _dot(cudnn_grad, cudnn_grad)
-                bucket["d2"] += _dot(diff, diff)
-                bucket["p2"] += _dot(param, param)
-                bucket["dg"] += _dot(diff, grad)
-                bucket["dp"] += _dot(diff, param)
-            for previous_idx, previous_delta in enumerate(deltas):
-                previous_dots[previous_idx] += _dot(diff, previous_delta[path])
-
-        current_d2 = totals["ALL"]["d2"]
-        pairwise_dot[batch_idx, batch_idx] = current_d2
-        for previous_idx, dot in enumerate(previous_dots):
-            pairwise_dot[batch_idx, previous_idx] = dot
-            pairwise_dot[previous_idx, batch_idx] = dot
-
-        deltas.append(delta)
-        delta_norm_sq.append(current_d2)
-        random_rel = np.mean(np.sqrt(repeat_diffs_sq)) / max(np.sqrt(totals["ALL"]["g2"]), 1e-300)
-        all_stats = totals["ALL"]
-        summary = {
-            "rel_delta": np.sqrt(all_stats["d2"] / all_stats["g2"]),
-            "cos_d_g": _cos(all_stats["dg"], all_stats["d2"], all_stats["g2"]),
-            "proj_d_g": all_stats["dg"] / max(all_stats["g2"], 1e-300),
-            "cos_d_p": _cos(all_stats["dp"], all_stats["d2"], all_stats["p2"]),
-            "proj_d_p": all_stats["dp"] / max(all_stats["p2"], 1e-300),
-            "random_rel": random_rel,
-        }
-        per_batch_summary.append(summary)
-
-        print(f"\n=== batch {batch_idx} ===", flush=True)
-        print(
-            f"loss explicit={loss_off:.6f} cudnn_mean={np.mean(cudnn_losses):.6f} "
-            f"loss_delta={np.mean(cudnn_losses) - loss_off:+.3e}",
-            flush=True,
-        )
-        print(
-            f"|g|={np.sqrt(all_stats['g2']):.6f} |delta_mean|/|g|={summary['rel_delta']:.4e} "
-            f"repeat_random/|g|={random_rel:.4e}",
-            flush=True,
-        )
-        print(
-            f"delta vs grad: cos={summary['cos_d_g']:+.5f} projection/|g|={summary['proj_d_g']:+.4e}; "
-            f"delta vs params: cos={summary['cos_d_p']:+.5f} projection/|p|={summary['proj_d_p']:+.4e}",
-            flush=True,
-        )
-        if batch_idx:
-            cosines = [
-                _cos(pairwise_dot[batch_idx, previous_idx], current_d2, delta_norm_sq[previous_idx])
-                for previous_idx in range(batch_idx)
-            ]
-            print("delta cosine vs prior batches: " + " ".join(f"{value:+.5f}" for value in cosines), flush=True)
-
-        print("module metrics (rel_delta, cos(delta,grad), cos(delta,param)):", flush=True)
-        for group in sorted(key for key in totals if key != "ALL"):
-            stats = totals[group]
-            values = (
-                np.sqrt(stats["d2"] / max(stats["g2"], 1e-300)),
-                _cos(stats["dg"], stats["d2"], stats["g2"]),
-                _cos(stats["dp"], stats["d2"], stats["p2"]),
-            )
-            per_module_summary[group].append(values)
-            print(f"  {group:20s} {values[0]:.4e} {values[1]:+.5f} {values[2]:+.5f}", flush=True)
-        del host_off, host_on_avg, first_on
-
-    cosine_matrix = np.eye(N_BATCHES, dtype=np.float64)
-    off_diagonal = []
-    for row in range(N_BATCHES):
-        for col in range(row):
-            value = _cos(pairwise_dot[row, col], delta_norm_sq[row], delta_norm_sq[col])
-            cosine_matrix[row, col] = cosine_matrix[col, row] = value
-            off_diagonal.append(value)
-
-    print("\n=== cross-batch delta cosine matrix ===", flush=True)
-    print("        " + " ".join(f"b{i:02d}" for i in range(N_BATCHES)), flush=True)
-    for idx, row in enumerate(cosine_matrix):
-        print(f"b{idx:02d} " + " ".join(f"{value:+.5f}" for value in row), flush=True)
-    print(
-        f"off-diagonal cosine: mean={np.mean(off_diagonal):+.6f} "
-        f"std={np.std(off_diagonal, ddof=1):.6f} min={np.min(off_diagonal):+.6f} "
-        f"max={np.max(off_diagonal):+.6f} n={len(off_diagonal)}",
-        flush=True,
-    )
-
-    print("\n=== aggregate projections across batches (mean +/- sample std) ===", flush=True)
-    for key in ("rel_delta", "random_rel", "cos_d_g", "proj_d_g", "cos_d_p", "proj_d_p"):
-        values = np.asarray([row[key] for row in per_batch_summary])
-        print(f"{key:12s} {values.mean():+.6e} +/- {values.std(ddof=1):.6e}", flush=True)
-
-    print("\n=== aggregate module metrics (mean rel_delta, cos(delta,grad), cos(delta,param)) ===", flush=True)
-    for group in sorted(per_module_summary):
-        values = np.asarray(per_module_summary[group])
-        means = values.mean(axis=0)
-        stds = values.std(axis=0, ddof=1)
-        print(
-            f"{group:20s} rel={means[0]:.4e}+/-{stds[0]:.2e} "
-            f"cos_g={means[1]:+.5f}+/-{stds[1]:.2e} cos_p={means[2]:+.5f}+/-{stds[2]:.2e}",
-            flush=True,
-        )
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/grad_ab_probe.py b/scripts/grad_ab_probe.py
deleted file mode 100644
index 1748a04..0000000
--- a/scripts/grad_ab_probe.py
+++ /dev/null
@@ -1,162 +0,0 @@
-"""Full-model gradient A/B: explicit attention vs cuDNN fused attention.
-
-Identical weights, identical batch, identical rng. The only difference is
-Pi0Config.use_cudnn_attention. Reports per-parameter relative gradient error so a
-systematic discrepancy can be localised to a module.
-"""
-
-import collections
-import dataclasses
-import functools
-import logging
-import sys
-
-import flax.nnx as nnx
-import jax
-import jax.numpy as jnp
-from jax._src.lib import cuda_versions
-import numpy as np
-import optax
-
-import openpi.training.config as _config
-import openpi.training.data_loader as _data_loader
-import openpi.training.sharding as sharding
-
-sys.path.insert(0, "scripts")
-import train as train_mod
-
-CONFIG = "pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100"
-BATCH = 32
-WORKERS = 16
-N_BATCHES = 3
-
-
-def _grad_fn(model_def, trainable_filter, params, rng, obs, act):
-    model = nnx.merge(model_def, params)
-    model.train()
-
-    def loss_fn(m, r, o, a):
-        return jnp.mean(m.compute_loss(r, o, a, train=True))
-
-    return nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, trainable_filter))(model, rng, obs, act)
-
-
-def _flat(state):
-    pure = state.to_pure_dict()
-    out = []
-
-    def rec(d, prefix):
-        if isinstance(d, dict):
-            for k, v in d.items():
-                rec(v, f"{prefix}/{k}" if prefix else str(k))
-        else:
-            out.append((prefix, d))
-
-    rec(pure, "")
-    return out
-
-
-def main():
-    logging.basicConfig(level=logging.WARNING)
-    print("cuDNN runtime:", cuda_versions.cudnn_get_version())
-
-    cfg = _config.get_config(CONFIG)
-    cfg = dataclasses.replace(cfg, batch_size=BATCH, num_workers=WORKERS, wandb_enabled=False)
-    cfg_off = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=False))
-    cfg_on = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=True))
-
-    mesh = sharding.make_mesh(cfg.fsdp_devices)
-    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
-
-    rng = jax.random.key(cfg.seed)
-    train_rng, init_rng = jax.random.split(rng)
-
-    print("loading pi05_base weights ...")
-    state, _ = train_mod.init_train_state(cfg_off, init_rng, mesh, resume=False)
-    print("  param global norm:", float(optax.global_norm(state.params)))
-
-    # Same parameter values, two attention implementations.
-    gdef_off = state.model_def
-    gdef_on = nnx.graphdef(cfg_on.model.create(init_rng))
-
-    f_off = jax.jit(functools.partial(_grad_fn, gdef_off, cfg.trainable_filter))
-    f_on = jax.jit(functools.partial(_grad_fn, gdef_on, cfg.trainable_filter))
-
-    print("building data loader (this spawns workers) ...")
-    dl = _data_loader.create_data_loader(cfg_off, sharding=data_sharding, shuffle=True)
-    it = iter(dl)
-
-    agg = collections.defaultdict(lambda: [0.0, 0.0])
-
-    def to_host(state_tree):
-        """Pull one gradient set to numpy so the device copy can be freed."""
-        return {k: np.asarray(jax.device_get(v), np.float32) for k, v in _flat(state_tree)}
-
-    for b in range(N_BATCHES):
-        obs, act = next(it)
-        step_rng = jax.random.fold_in(train_rng, b)
-
-        # One gradient set at a time: 3B params in fp32 is ~12 GB, two live copies
-        # plus the tree arithmetic does not fit alongside the activations.
-        with sharding.set_mesh(mesh):
-            loss_off, g_off = f_off(state.params, step_rng, obs, act)
-            loss_off = float(loss_off)
-            h_off = to_host(g_off)
-        del g_off
-        with sharding.set_mesh(mesh):
-            loss_on, g_on = f_on(state.params, step_rng, obs, act)
-            loss_on = float(loss_on)
-            h_on = to_host(g_on)
-        del g_on
-        # Second identical cuDNN call. cuDNN's backward accumulates with atomics, so
-        # this measures the RANDOM part of the cuDNN/explicit gap. Whatever is left
-        # over is systematic, and only a systematic part can accumulate over steps.
-        with sharding.set_mesh(mesh):
-            _, g_on2 = f_on(state.params, step_rng, obs, act)
-            h_on2 = to_host(g_on2)
-        del g_on2
-
-        sq_off = sq_on = sq_d = 0.0
-        rows = []
-        for path, av in h_off.items():
-            av = av.astype(np.float64)
-            bv = h_on[path].astype(np.float64)
-            na = float(np.linalg.norm(av))
-            nb = float(np.linalg.norm(bv))
-            nd = float(np.linalg.norm(bv - av))
-            sq_off += na**2
-            sq_on += nb**2
-            sq_d += nd**2
-            rows.append((nd / max(na, 1e-30), path, na))
-            key = path.split("/")[0]
-            agg[key][0] += nd**2
-            agg[key][1] += na**2
-
-        sq_rand = 0.0
-        for path, av in h_on.items():
-            sq_rand += float(np.linalg.norm(h_on2[path].astype(np.float64) - av.astype(np.float64)) ** 2)
-        n_off, n_on, n_d = np.sqrt(sq_off), np.sqrt(sq_on), np.sqrt(sq_d)
-        n_rand = np.sqrt(sq_rand)
-        print(f"\n=== batch {b} ===")
-        print(f"  loss    explicit={loss_off:.6f}  cudnn={loss_on:.6f}  diff={loss_on - loss_off:+.3e}")
-        print(f"  |grad|  explicit={n_off:.6f}  cudnn={n_on:.6f}  ratio={n_on / max(n_off, 1e-30):.4f}")
-        rel_total = n_d / max(n_off, 1e-30)
-        rel_rand = n_rand / max(n_off, 1e-30)
-        # Total gap and the random part are orthogonal in expectation.
-        rel_sys = np.sqrt(max(rel_total**2 - rel_rand**2, 0.0))
-        print(f"  cuDNN vs explicit (total)   = {rel_total:.4e}")
-        print(f"  cuDNN vs cuDNN    (random)  = {rel_rand:.4e}")
-        print(f"  implied SYSTEMATIC component= {rel_sys:.4e}")
-        rows.sort(reverse=True)
-        print("  worst 10 parameters by relative gradient difference:")
-        for rel, path, na in rows[:10]:
-            print(f"    {rel:10.3e}  |g|={na:11.4e}  {path}")
-        del h_off, h_on, h_on2
-
-    print("\n=== aggregate by top-level module ===")
-    for k, (dn, an) in sorted(agg.items(), key=lambda kv: -(kv[1][0] / max(kv[1][1], 1e-30))):
-        print(f"  {np.sqrt(dn / max(an, 1e-30)):10.3e}  {k}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/hybrid_update_ab_probe.py b/scripts/hybrid_update_ab_probe.py
deleted file mode 100644
index 622f930..0000000
--- a/scripts/hybrid_update_ab_probe.py
+++ /dev/null
@@ -1,192 +0,0 @@
-"""Compare optimizer-space updates at a saved training state.
-
-Raw gradient L2 can hide errors in coordinates with a small Adam second moment.
-This probe restores the real step-1000 optimizer state, computes gradients for
-the same real batch/rng, and compares the actual AdamW updates produced by
-explicit attention, cuDNN attention, and the prior multiplicative-noise control.
-"""
-
-import dataclasses
-import functools
-import logging
-import sys
-
-import flax.nnx as nnx
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-import openpi.training.checkpoints as _checkpoints
-import openpi.training.config as _config
-import openpi.training.data_loader as _data_loader
-import openpi.training.sharding as sharding
-
-sys.path.insert(0, "scripts")
-import train as train_mod
-
-
-CONFIG = "diag_cudnn_strict_order"
-EXP_NAME = "diag_cudnn_strict_order"
-CHECKPOINT_STEP = 1000
-BATCH = 32
-WORKERS = 16
-
-
-def _loss_and_update(model_def, trainable_filter, tx, noise_scale, opt_state, params, rng, obs, act):
-    model = nnx.merge(model_def, params)
-    model.train()
-
-    def loss_fn(m, r, o, a):
-        return jnp.mean(m.compute_loss(r, o, a, train=True))
-
-    loss, grads = nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, trainable_filter))(
-        model, rng, obs, act
-    )
-    if noise_scale:
-        noise_rng = jax.random.fold_in(rng, 0x6E6F6973)
-        grads, _ = train_mod._add_relative_gradient_noise(grads, noise_rng, noise_scale)
-    trainable_params = params.filter(trainable_filter)
-    updates, _ = tx.update(grads, opt_state, trainable_params)
-    return loss, updates
-
-
-def _flat(tree):
-    pure = tree.to_pure_dict()
-    out = []
-
-    def rec(value, prefix):
-        if isinstance(value, dict):
-            for key, child in value.items():
-                rec(child, f"{prefix}/{key}" if prefix else str(key))
-        else:
-            out.append((prefix, value))
-
-    rec(pure, "")
-    return out
-
-
-def _to_host(tree):
-    return {path: np.asarray(jax.device_get(value), np.float32) for path, value in _flat(tree)}
-
-
-def _dot(a, b):
-    return float(np.vdot(a.astype(np.float64), b.astype(np.float64)))
-
-
-def _metrics(reference, candidate):
-    totals = {"ALL": [0.0, 0.0, 0.0, 0.0, 0, 0]}
-    for path, ref in reference.items():
-        cand = candidate[path]
-        delta = cand - ref
-        group = path.split("/", 1)[0]
-        totals.setdefault(group, [0.0, 0.0, 0.0, 0.0, 0, 0])
-        for bucket in (totals["ALL"], totals[group]):
-            bucket[0] += _dot(ref, ref)
-            bucket[1] += _dot(cand, cand)
-            bucket[2] += _dot(delta, delta)
-            bucket[3] += _dot(ref, cand)
-            nonzero = (ref != 0) | (cand != 0)
-            bucket[4] += int(np.count_nonzero(nonzero & (np.signbit(ref) != np.signbit(cand))))
-            bucket[5] += int(np.count_nonzero(nonzero))
-    return totals
-
-
-def _print_comparison(label, reference, candidate):
-    metrics = _metrics(reference, candidate)
-    ref2, cand2, delta2, dot, flips, count = metrics["ALL"]
-    print(f"\n=== {label} ===", flush=True)
-    print(
-        f"|update_ref|={np.sqrt(ref2):.8e} |update_candidate|={np.sqrt(cand2):.8e} "
-        f"rel_update_delta={np.sqrt(delta2 / ref2):.6e} "
-        f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.8f} "
-        f"sign_flip_fraction={flips / max(count, 1):.6e}",
-        flush=True,
-    )
-    for group in sorted(key for key in metrics if key != "ALL"):
-        ref2, cand2, delta2, dot, flips, count = metrics[group]
-        print(
-            f"  {group:20s} rel={np.sqrt(delta2 / max(ref2, 1e-300)):.6e} "
-            f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.8f} "
-            f"sign_flips={flips / max(count, 1):.6e}",
-            flush=True,
-        )
-
-
-def main():
-    logging.basicConfig(level=logging.WARNING)
-    cfg = _config.get_config(CONFIG)
-    cfg = dataclasses.replace(
-        cfg,
-        exp_name=EXP_NAME,
-        batch_size=BATCH,
-        num_workers=WORKERS,
-        resume=True,
-        overwrite=False,
-        wandb_enabled=False,
-    )
-    cfg_on = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=True))
-    cfg_off = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=False))
-
-    mesh = sharding.make_mesh(cfg.fsdp_devices)
-    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
-    rng = jax.random.key(cfg.seed)
-    train_rng, init_rng = jax.random.split(rng)
-
-    manager, resuming = _checkpoints.initialize_checkpoint_dir(
-        cfg_on.checkpoint_dir,
-        keep_period=cfg_on.keep_period,
-        overwrite=False,
-        resume=True,
-    )
-    if not resuming or CHECKPOINT_STEP not in manager.all_steps():
-        raise RuntimeError(f"checkpoint {CHECKPOINT_STEP} unavailable; found {manager.all_steps()}")
-    state_shape, state_sharding = train_mod.init_train_state(cfg_on, init_rng, mesh, resume=True)
-    state = _checkpoints.restore_state(manager, state_shape, None, step=CHECKPOINT_STEP)
-    jax.block_until_ready(state)
-    print(f"restored checkpoint directory step={CHECKPOINT_STEP}; train_state.step={int(state.step)}", flush=True)
-
-    print("building one-batch loader ...", flush=True)
-    loader = _data_loader.create_data_loader(cfg_off, sharding=data_sharding, shuffle=True)
-    obs, act = next(iter(loader))
-    step_rng = jax.random.fold_in(train_rng, state.step)
-
-    gdef_off = nnx.graphdef(cfg_off.model.create(init_rng))
-    hybrid_ranges = [
-        ("full_18", 0, None),
-        ("first_6", 0, 6),
-        ("middle_6", 6, 6),
-        ("last_6", 12, 6),
-        ("first_3", 0, 3),
-        ("middle_3", 7, 3),
-        ("last_3", 15, 3),
-    ]
-    graphdefs = {}
-    for name, start, count in hybrid_ranges:
-        model_config = dataclasses.replace(
-            cfg_on.model,
-            use_cudnn_attention=True,
-            cudnn_attention_layer_start=start,
-            cudnn_attention_num_layers=count,
-        )
-        graphdefs[name] = nnx.graphdef(model_config.create(init_rng))
-
-    fn = jax.jit(functools.partial(_loss_and_update, gdef_off, cfg.trainable_filter, state.tx, 0.0))
-    with sharding.set_mesh(mesh):
-        loss, updates = fn(state.opt_state, state.params, step_rng, obs, act)
-        reference = _to_host(updates)
-    del updates, fn
-    print(f"computed explicit reference: loss={float(loss):.8f}", flush=True)
-
-    for name, _, _ in hybrid_ranges:
-        fn = jax.jit(functools.partial(_loss_and_update, graphdefs[name], cfg.trainable_filter, state.tx, 0.0))
-        with sharding.set_mesh(mesh):
-            loss, updates = fn(state.opt_state, state.params, step_rng, obs, act)
-            candidate = _to_host(updates)
-        del updates, fn
-        print(f"computed {name}: loss={float(loss):.8f}", flush=True)
-        _print_comparison(name, reference, candidate)
-        del candidate
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/optimizer_update_ab_probe.py b/scripts/optimizer_update_ab_probe.py
deleted file mode 100644
index 487e205..0000000
--- a/scripts/optimizer_update_ab_probe.py
+++ /dev/null
@@ -1,194 +0,0 @@
-"""Compare optimizer-space updates at a saved training state.
-
-Raw gradient L2 can hide errors in coordinates with a small Adam second moment.
-This probe restores the real step-1000 optimizer state, computes gradients for
-the same real batch/rng, and compares the actual AdamW updates produced by
-explicit attention, cuDNN attention, and the prior multiplicative-noise control.
-"""
-
-import dataclasses
-import functools
-import logging
-import sys
-
-import flax.nnx as nnx
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-import openpi.training.checkpoints as _checkpoints
-import openpi.training.config as _config
-import openpi.training.data_loader as _data_loader
-import openpi.training.sharding as sharding
-
-sys.path.insert(0, "scripts")
-import train as train_mod
-
-
-CONFIG = "diag_cudnn_strict_order"
-EXP_NAME = "diag_cudnn_strict_order"
-CHECKPOINT_STEP = 1000
-BATCH = 32
-WORKERS = 16
-
-
-def _loss_and_update(model_def, trainable_filter, tx, noise_scale, opt_state, params, rng, obs, act):
-    model = nnx.merge(model_def, params)
-    model.train()
-
-    def loss_fn(m, r, o, a):
-        return jnp.mean(m.compute_loss(r, o, a, train=True))
-
-    loss, grads = nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, trainable_filter))(
-        model, rng, obs, act
-    )
-    if noise_scale:
-        noise_rng = jax.random.fold_in(rng, 0x6E6F6973)
-        grads, _ = train_mod._add_relative_gradient_noise(grads, noise_rng, noise_scale)
-    trainable_params = params.filter(trainable_filter)
-    updates, _ = tx.update(grads, opt_state, trainable_params)
-    return loss, updates
-
-
-def _flat(tree):
-    pure = tree.to_pure_dict()
-    out = []
-
-    def rec(value, prefix):
-        if isinstance(value, dict):
-            for key, child in value.items():
-                rec(child, f"{prefix}/{key}" if prefix else str(key))
-        else:
-            out.append((prefix, value))
-
-    rec(pure, "")
-    return out
-
-
-def _to_host(tree):
-    return {path: np.asarray(jax.device_get(value), np.float32) for path, value in _flat(tree)}
-
-
-def _dot(a, b):
-    return float(np.vdot(a.astype(np.float64), b.astype(np.float64)))
-
-
-def _metrics(reference, candidate):
-    totals = {"ALL": [0.0, 0.0, 0.0, 0.0, 0, 0]}
-    for path, ref in reference.items():
-        cand = candidate[path]
-        delta = cand - ref
-        group = path.split("/", 1)[0]
-        totals.setdefault(group, [0.0, 0.0, 0.0, 0.0, 0, 0])
-        for bucket in (totals["ALL"], totals[group]):
-            bucket[0] += _dot(ref, ref)
-            bucket[1] += _dot(cand, cand)
-            bucket[2] += _dot(delta, delta)
-            bucket[3] += _dot(ref, cand)
-            nonzero = (ref != 0) | (cand != 0)
-            bucket[4] += int(np.count_nonzero(nonzero & (np.signbit(ref) != np.signbit(cand))))
-            bucket[5] += int(np.count_nonzero(nonzero))
-    return totals
-
-
-def _print_comparison(label, reference, candidate):
-    metrics = _metrics(reference, candidate)
-    ref2, cand2, delta2, dot, flips, count = metrics["ALL"]
-    print(f"\n=== {label} ===", flush=True)
-    print(
-        f"|update_ref|={np.sqrt(ref2):.8e} |update_candidate|={np.sqrt(cand2):.8e} "
-        f"rel_update_delta={np.sqrt(delta2 / ref2):.6e} "
-        f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.8f} "
-        f"sign_flip_fraction={flips / max(count, 1):.6e}",
-        flush=True,
-    )
-    for group in sorted(key for key in metrics if key != "ALL"):
-        ref2, cand2, delta2, dot, flips, count = metrics[group]
-        print(
-            f"  {group:20s} rel={np.sqrt(delta2 / max(ref2, 1e-300)):.6e} "
-            f"cos={dot / max(np.sqrt(ref2 * cand2), 1e-300):+.8f} "
-            f"sign_flips={flips / max(count, 1):.6e}",
-            flush=True,
-        )
-
-
-def main():
-    logging.basicConfig(level=logging.WARNING)
-    cfg = _config.get_config(CONFIG)
-    cfg = dataclasses.replace(
-        cfg,
-        exp_name=EXP_NAME,
-        batch_size=BATCH,
-        num_workers=WORKERS,
-        resume=True,
-        overwrite=False,
-        wandb_enabled=False,
-    )
-    cfg_on = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=True))
-    cfg_off = dataclasses.replace(cfg, model=dataclasses.replace(cfg.model, use_cudnn_attention=False))
-
-    mesh = sharding.make_mesh(cfg.fsdp_devices)
-    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
-    rng = jax.random.key(cfg.seed)
-    train_rng, init_rng = jax.random.split(rng)
-
-    manager, resuming = _checkpoints.initialize_checkpoint_dir(
-        cfg_on.checkpoint_dir,
-        keep_period=cfg_on.keep_period,
-        overwrite=False,
-        resume=True,
-    )
-    if not resuming or CHECKPOINT_STEP not in manager.all_steps():
-        raise RuntimeError(f"checkpoint {CHECKPOINT_STEP} unavailable; found {manager.all_steps()}")
-    state_shape, state_sharding = train_mod.init_train_state(cfg_on, init_rng, mesh, resume=True)
-    state = _checkpoints.restore_state(manager, state_shape, None, step=CHECKPOINT_STEP)
-    jax.block_until_ready(state)
-    print(f"restored checkpoint directory step={CHECKPOINT_STEP}; train_state.step={int(state.step)}", flush=True)
-
-    print("building one-batch loader ...", flush=True)
-    loader = _data_loader.create_data_loader(cfg_off, sharding=data_sharding, shuffle=True)
-    obs, act = next(iter(loader))
-    step_rng = jax.random.fold_in(train_rng, state.step)
-
-    gdef_on = state.model_def
-    gdef_off = nnx.graphdef(cfg_off.model.create(init_rng))
-
-    # The optimizer state structure is unchanged when only Adam epsilon changes.
-    eps_values = (1e-8, 1e-6, 1e-5)
-    tx_by_eps = {1e-8: state.tx}
-    for eps in eps_values[1:]:
-        cfg_eps = dataclasses.replace(cfg_on, optimizer=dataclasses.replace(cfg_on.optimizer, eps=eps))
-        shape_eps, _ = train_mod.init_train_state(cfg_eps, init_rng, mesh, resume=True)
-        tx_by_eps[eps] = shape_eps.tx
-
-    host_updates = {}
-    variants = [
-        ("explicit_eps1e-8", gdef_off, tx_by_eps[1e-8], 0.0),
-        ("cudnn_eps1e-8", gdef_on, tx_by_eps[1e-8], 0.0),
-        ("explicit_noise015_eps1e-8", gdef_off, tx_by_eps[1e-8], 0.015),
-        ("explicit_eps1e-6", gdef_off, tx_by_eps[1e-6], 0.0),
-        ("cudnn_eps1e-6", gdef_on, tx_by_eps[1e-6], 0.0),
-        ("explicit_eps1e-5", gdef_off, tx_by_eps[1e-5], 0.0),
-        ("cudnn_eps1e-5", gdef_on, tx_by_eps[1e-5], 0.0),
-    ]
-    for name, graphdef, tx, noise_scale in variants:
-        fn = jax.jit(functools.partial(_loss_and_update, graphdef, cfg.trainable_filter, tx, noise_scale))
-        with sharding.set_mesh(mesh):
-            loss, updates = fn(state.opt_state, state.params, step_rng, obs, act)
-            loss = float(loss)
-            host_updates[name] = _to_host(updates)
-        del updates, fn
-        print(f"computed {name}: loss={loss:.8f}", flush=True)
-
-    _print_comparison("real cuDNN delta, Adam eps=1e-8", host_updates["explicit_eps1e-8"], host_updates["cudnn_eps1e-8"])
-    _print_comparison(
-        "prior multiplicative 1.5% noise, Adam eps=1e-8",
-        host_updates["explicit_eps1e-8"],
-        host_updates["explicit_noise015_eps1e-8"],
-    )
-    _print_comparison("real cuDNN delta, Adam eps=1e-6", host_updates["explicit_eps1e-6"], host_updates["cudnn_eps1e-6"])
-    _print_comparison("real cuDNN delta, Adam eps=1e-5", host_updates["explicit_eps1e-5"], host_updates["cudnn_eps1e-5"])
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/token_len_scan.py b/scripts/token_len_scan.py
deleted file mode 100644
index 432c79c..0000000
--- a/scripts/token_len_scan.py
+++ /dev/null
@@ -1,115 +0,0 @@
-"""Measure the real tokenized prompt length of a training config's dataset.
-
-For Pi0.5 the language prefix is ``Task: <prompt>, State: <discretized state>;\\nAction: `` and
-everything up to ``max_token_len`` is padded (and masked) -- but attention/MLP still run over the
-padding. On 2026-09-04 trimming ``max_token_len`` 200 -> 128 gave -6.2% step time with identical
-loss (see docs/training-optimization.md). This script runs the *actual* transform chain
-(repack -> data transforms -> model transforms) on real states with dummy images and reports the
-token-length distribution so ``max_token_len`` can be set with a known margin.
-
-Note: the tokenizer only warns and truncates when the prompt exceeds ``max_token_len``.
-
-Example:
-    JAX_PLATFORMS=cpu uv run scripts/token_len_scan.py --config-name pi05_droid --samples 4000
-"""
-
-import dataclasses
-import glob
-import json
-import logging
-import os
-import pathlib
-
-import numpy as np
-import pyarrow.parquet as pq
-import tyro
-
-import openpi.training.config as _config
-import openpi.transforms as _transforms
-
-
-@dataclasses.dataclass
-class Args:
-    config_name: str
-    # Random states to sample; per-dimension min/max states are always added so the worst case is covered.
-    samples: int = 4000
-    # Override the LeRobot dataset root (default: $HF_LEROBOT_HOME/<repo_id>).
-    dataset_root: str | None = None
-    # Image shape handed to the transform chain (channels-first uint8, like LeRobot video frames).
-    image_shape: tuple[int, int, int] = (3, 480, 640)
-    seed: int = 0
-
-
-def _lerobot_home() -> pathlib.Path:
-    if "HF_LEROBOT_HOME" in os.environ:
-        return pathlib.Path(os.environ["HF_LEROBOT_HOME"]).expanduser()
-    return pathlib.Path(os.environ.get("HF_HOME", "~/.cache/huggingface")).expanduser() / "lerobot"
-
-
-def _load_tasks(root: pathlib.Path) -> list[str]:
-    jsonl = root / "meta" / "tasks.jsonl"
-    if jsonl.exists():
-        return sorted({json.loads(line)["task"] for line in jsonl.read_text().splitlines() if line.strip()})
-    parquet = root / "meta" / "tasks.parquet"
-    if parquet.exists():
-        table = pq.read_table(parquet).to_pandas().reset_index()
-        col = "task" if "task" in table.columns else table.columns[0]
-        return sorted(set(table[col].astype(str)))
-    raise FileNotFoundError(f"No tasks metadata under {root / 'meta'}")
-
-
-def main(args: Args) -> None:
-    cfg = _config.get_config(args.config_name)
-    data_config = cfg.data.create(cfg.assets_dirs, cfg.model)
-    if data_config.repo_id is None:
-        raise ValueError("config has no repo_id")
-    chain = _transforms.compose(
-        [
-            *data_config.repack_transforms.inputs,
-            *data_config.data_transforms.inputs,
-            *data_config.model_transforms.inputs,
-        ]
-    )
-    root = pathlib.Path(args.dataset_root) if args.dataset_root else _lerobot_home() / data_config.repo_id
-    tasks = _load_tasks(root)
-    logging.info("tasks: %s", tasks)
-
-    # Discover the raw column/key names from the repack transform so this works for any robot config.
-    repack = next(t for t in data_config.repack_transforms.inputs if isinstance(t, _transforms.RepackTransform))
-    structure = repack.structure
-    image_keys = list(structure["images"].values())
-    state_key, action_key, prompt_key = structure["state"], structure["actions"], structure["prompt"]
-
-    files = sorted(glob.glob(str(root / "data" / "*" / "*.parquet")))
-    states = np.concatenate(
-        [np.stack(pq.read_table(f, columns=[state_key]).column(0).to_pylist()) for f in files]
-    ).astype(np.float32)
-    logging.info("states: %s", states.shape)
-
-    rng = np.random.default_rng(args.seed)
-    idx = rng.choice(len(states), size=min(args.samples, len(states)), replace=False)
-    idx = np.unique(np.concatenate([idx, states.argmin(0), states.argmax(0)]))
-
-    dummy_img = np.zeros(args.image_shape, np.uint8)
-    action_dim = states.shape[1]
-    lengths = []
-    for i in idx:
-        for task in tasks:
-            sample = dict.fromkeys(image_keys, dummy_img)
-            sample[state_key] = states[i]
-            sample[action_key] = np.zeros((cfg.model.action_horizon, action_dim), np.float32)
-            sample[prompt_key] = task
-            out = chain(sample)
-            lengths.append(int(np.asarray(out["tokenized_prompt_mask"]).sum()))
-    lengths = np.asarray(lengths)
-    cap = cfg.model.max_token_len
-    print(
-        f"n={len(lengths)} min={lengths.min()} p50={np.percentile(lengths, 50):.0f} "
-        f"p99={np.percentile(lengths, 99):.0f} max={lengths.max()} "
-        f"(max_token_len={cap}; samples_at_cap={(lengths >= cap).sum()})"
-    )
-
-
-if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO)
-    main(tyro.cli(Args))
diff --git a/scripts/train.py b/scripts/train.py
index b7ddd38..2e05621 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -10,9 +10,9 @@
 from flax.training import common_utils
 import flax.traverse_util as traverse_util
 import jax
+from jax._src.lib import cuda_versions
 import jax.experimental
 import jax.numpy as jnp
-from jax._src.lib import cuda_versions
 import numpy as np
 import optax
 import tqdm_loggable.auto as tqdm
@@ -135,30 +135,6 @@ def init(rng: at.KeyArrayLike, partial_params: at.Params | None = None) -> train
     return train_state, state_sharding
 
 
-def _add_relative_gradient_noise(grads, rng: at.KeyArrayLike, relative_scale: float):
-    """Adds reproducible random noise with an exact tree-wide relative L2 norm.
-
-    The raw noise is weighted by each gradient element before global
-    normalization. This keeps the perturbation energy distributed like the
-    gradient energy instead of concentrating it in the largest parameter leaf.
-    """
-    leaves, tree_def = jax.tree.flatten(grads)
-    keys = jax.random.split(rng, len(leaves))
-    raw_noise = jax.tree.unflatten(
-        tree_def,
-        [grad * jax.random.normal(key, grad.shape, dtype=grad.dtype) for grad, key in zip(leaves, keys)],
-    )
-
-    grad_norm = optax.global_norm(grads)
-    raw_noise_norm = optax.global_norm(raw_noise)
-    tiny = jnp.finfo(grad_norm.dtype).tiny
-    multiplier = jnp.asarray(relative_scale, grad_norm.dtype) * grad_norm / jnp.maximum(raw_noise_norm, tiny)
-    noise = jax.tree.map(lambda value: value * multiplier, raw_noise)
-    noisy_grads = jax.tree.map(lambda grad, delta: grad + delta, grads, noise)
-    actual_relative_norm = multiplier * raw_noise_norm / jnp.maximum(grad_norm, tiny)
-    return noisy_grads, actual_relative_norm
-
-
 @at.typecheck
 def train_step(
     config: _config.TrainConfig,
@@ -184,18 +160,6 @@ def loss_fn(
     diff_state = nnx.DiffState(0, config.trainable_filter)
     loss, grads = nnx.value_and_grad(loss_fn, argnums=diff_state)(model, train_rng, observation, actions)
 
-    clean_grads = grads
-    gradient_noise_rel_norm = None
-    if config.gradient_noise_scale > 0:
-        # Keep the model RNG identical to the clean explicit-attention path.
-        # fold_in gives the diagnostic noise an independent, step-reproducible stream.
-        noise_rng = jax.random.fold_in(train_rng, 0x6E6F6973)
-        grads, gradient_noise_rel_norm = _add_relative_gradient_noise(
-            grads,
-            noise_rng,
-            config.gradient_noise_scale,
-        )
-
     params = state.params.filter(config.trainable_filter)
     updates, new_opt_state = state.tx.update(grads, state.opt_state, params)
     new_params = optax.apply_updates(params, updates)
@@ -225,13 +189,11 @@ def loss_fn(
             ),
         )
         grad_norm = optax.global_norm(grads)
-        clean_grad_norm = optax.global_norm(clean_grads) if gradient_noise_rel_norm is not None else None
         param_norm = optax.global_norm(kernel_params)
     else:
         # Keep a stable output pytree for both JIT variants. The host-side logging
         # reduction ignores these placeholders with nanmean.
         grad_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
-        clean_grad_norm = jnp.asarray(jnp.nan, dtype=loss.dtype) if gradient_noise_rel_norm is not None else None
         param_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
 
     info = {
@@ -239,9 +201,6 @@ def loss_fn(
         "grad_norm": grad_norm,
         "param_norm": param_norm,
     }
-    if gradient_noise_rel_norm is not None:
-        info["clean_grad_norm"] = clean_grad_norm
-        info["gradient_noise_rel_norm"] = gradient_noise_rel_norm
     return new_state, info
 
 
@@ -251,13 +210,6 @@ def main(config: _config.TrainConfig):
     cudnn_runtime_version = cuda_versions.cudnn_get_version() if cuda_versions is not None else None
     logging.info(f"JAX cuDNN runtime version: {cudnn_runtime_version}")
 
-    if not 0.0 <= config.gradient_noise_scale < 1.0:
-        raise ValueError(f"gradient_noise_scale must be in [0, 1), got {config.gradient_noise_scale}")
-    attention_backend = "cuDNN" if getattr(config.model, "use_cudnn_attention", False) else "explicit"
-    logging.info(
-        f"Attention backend: {attention_backend}; gradient_noise_scale={config.gradient_noise_scale:.6f}"
-    )
-
     if config.batch_size % jax.device_count() != 0:
         raise ValueError(
             f"Batch size {config.batch_size} must be divisible by the number of devices {jax.device_count()}."
@@ -332,26 +284,10 @@ def main(config: _config.TrainConfig):
     )
 
     infos = []
-    xprof_active = False
     # --- stall diagnostics ---
     stall_threshold_s = 3.0  # log a warning when any phase exceeds this
     t_prev_loop_end = time.monotonic()
     for step in pbar:
-        if config.xprof_trace_dir is not None and step == config.xprof_start_step:
-            # Exclude work queued before the requested range and make the trace boundaries
-            # correspond to complete training iterations.
-            jax.block_until_ready((train_state, batch))
-            logging.info(
-                f"Starting JAX/XProf trace at step {step} for {config.xprof_num_steps} steps: "
-                f"{config.xprof_trace_dir}"
-            )
-            jax.profiler.start_trace(
-                config.xprof_trace_dir,
-                create_perfetto_link=False,
-                create_perfetto_trace=True,
-            )
-            xprof_active = True
-
         t_loop_start = time.monotonic()
 
         t0 = time.monotonic()
@@ -364,15 +300,6 @@ def main(config: _config.TrainConfig):
             )
         t_dispatch = time.monotonic() - t0
 
-        # JAX dispatch is asynchronous. Fine-grained H2D profiling must first finish the
-        # current train step; otherwise block_until_ready(batch) also measures time spent
-        # queued behind model compute on the same devices.
-        t_compute_sync = 0.0
-        if config.profile_data_pipeline:
-            t0 = time.monotonic()
-            jax.block_until_ready((train_state, info))
-            t_compute_sync = time.monotonic() - t0
-
         infos.append(info)
 
         t_log = 0.0
@@ -393,12 +320,9 @@ def main(config: _config.TrainConfig):
         t0 = time.monotonic()
         batch = next(data_iter)
         t_next_batch = time.monotonic() - t0
-        data_profile = data_loader.profile_stats()
 
         t_ckpt = 0.0
-        if (step % config.save_interval == 0 and step > start_step) or (
-            config.save_final_checkpoint and step == config.num_train_steps - 1
-        ):
+        if (step % config.save_interval == 0 and step > start_step) or step == config.num_train_steps - 1:
             t0 = time.monotonic()
             _checkpoints.save_state(checkpoint_manager, train_state, data_loader, step)
             t_ckpt = time.monotonic() - t0
@@ -423,32 +347,6 @@ def main(config: _config.TrainConfig):
                 f"dispatch={t_dispatch:.2f}s next_batch={t_next_batch:.2f}s "
                 f"log={t_log:.2f}s ckpt={t_ckpt:.2f}s"
             )
-        if (
-            config.profile_data_pipeline
-            and data_profile is not None
-            and step % config.profile_log_interval == 0
-        ):
-            pbar.write(
-                f"[DATA_PROFILE step={step}] "
-                f"dispatch={t_dispatch:.4f}s compute_sync={t_compute_sync:.4f}s "
-                f"main_queue_wait={data_profile['main_queue_wait_s']:.4f}s "
-                f"worker_getitem={data_profile['worker_getitem_s']:.4f}s "
-                f"worker_collate={data_profile['worker_collate_s']:.4f}s "
-                f"jax_array_construct={data_profile['jax_array_construct_s']:.4f}s "
-                f"h2d_wait={data_profile['h2d_wait_s']:.4f}s "
-                f"next_batch_total={t_next_batch:.4f}s"
-            )
-
-        if xprof_active and step + 1 == config.xprof_start_step + config.xprof_num_steps:
-            jax.block_until_ready((train_state, info, batch))
-            jax.profiler.stop_trace()
-            xprof_active = False
-            logging.info(f"Finished JAX/XProf trace at step {step}")
-
-    if xprof_active:
-        jax.block_until_ready((train_state, batch))
-        jax.profiler.stop_trace()
-
     logging.info("Shutting down data loader")
     data_loader.close()
     logging.info("Waiting for checkpoint manager to finish")
diff --git a/scripts/train_test.py b/scripts/train_test.py
index e5cba2d..6165930 100644
--- a/scripts/train_test.py
+++ b/scripts/train_test.py
@@ -6,41 +6,11 @@
 
 os.environ["JAX_PLATFORMS"] = "cpu"
 
-import jax
-import jax.numpy as jnp
-import numpy as np
-
 from openpi.training import config as _config
 
 from . import train
 
 
-def test_add_relative_gradient_noise():
-    grads = {
-        "a": jnp.arange(1, 13, dtype=jnp.float32).reshape(3, 4),
-        "b": jnp.asarray([-3.0, 0.0, 5.0], dtype=jnp.float32),
-    }
-    scale = 0.015
-    key = jax.random.key(7)
-
-    noisy, reported_rel_norm = train._add_relative_gradient_noise(grads, key, scale)
-    noise = jax.tree.map(lambda perturbed, clean: perturbed - clean, noisy, grads)
-    measured_rel_norm = np.sqrt(sum(float(jnp.sum(x**2)) for x in jax.tree.leaves(noise))) / np.sqrt(
-        sum(float(jnp.sum(x**2)) for x in jax.tree.leaves(grads))
-    )
-
-    assert float(reported_rel_norm) == pytest.approx(scale, rel=1e-6)
-    assert measured_rel_norm == pytest.approx(scale, rel=1e-6)
-    same_noise, _ = train._add_relative_gradient_noise(grads, key, scale)
-    different_noise, _ = train._add_relative_gradient_noise(grads, jax.random.key(8), scale)
-    for actual, expected in zip(jax.tree.leaves(same_noise), jax.tree.leaves(noisy)):
-        np.testing.assert_array_equal(actual, expected)
-    assert any(
-        not np.array_equal(actual, expected)
-        for actual, expected in zip(jax.tree.leaves(different_noise), jax.tree.leaves(noisy))
-    )
-
-
 @pytest.mark.parametrize("config_name", ["debug_pi05"])
 def test_train(tmp_path: pathlib.Path, config_name: str):
     config = dataclasses.replace(
diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index f50aa39..490f179 100644
--- a/src/openpi/models/gemma.py
+++ b/src/openpi/models/gemma.py
@@ -211,7 +211,7 @@ def _cudnn_attention_in_dtype(q, k, v, attn_mask, compute_dtype):
     rescaled to a power of two near 1 (dynamic loss scaling) before the fused
     backward and unscaled in float32 afterwards. Forward values are unchanged up
     to the output rounding; q/k/v are exactly representable in float16 unless
-    they exceed its range, which the probe script checks.
+    they exceed its range. Validate activation ranges when changing the model or data.
     """
     out_dtype = q.dtype
 
@@ -292,11 +292,9 @@ class Attention(nn.Module):
     # Compute dtype handed to the cuDNN kernel: "bfloat16" (historical) or "float16"
     # (see _cudnn_attention_in_dtype). Ignored by the explicit path.
     cudnn_attention_dtype: str = "bfloat16"
-    # Diagnostic only: run the explicit path with fp32 q/k/v/probs as a reference.
-    explicit_attention_fp32: bool = False
 
     @nn.compact
-    def __call__(self, xs, positions, attn_mask, kv_cache, use_cudnn_attention=None):
+    def __call__(self, xs, positions, attn_mask, kv_cache):
         # all experts must share the same head dim, num heads, and num kv heads for self-attention to work
         assert all(config.head_dim == self.configs[0].head_dim for config in self.configs)
         assert all(config.num_heads == self.configs[0].num_heads for config in self.configs)
@@ -367,8 +365,6 @@ def cudnn_attention(operands):
 
         def explicit_attention(operands):
             q, k, v, mask = operands
-            if self.explicit_attention_fp32:
-                q, k, v = (x.astype(jnp.float32) for x in (q, k, v))
             grouped_q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=self.configs[0].num_kv_heads)
             logits = jnp.einsum("BTKGH,BSKH->BKGTS", grouped_q, k, preferred_element_type=jnp.float32)
 
@@ -381,15 +377,11 @@ def explicit_attention(operands):
             encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v).astype(dtype)
             return einops.rearrange(encoded, "B T K G H -> B T (K G) H")
 
-        if use_cudnn_attention is None:
-            use_cudnn_attention = self.use_cudnn_attention
         operands = (q, k, v, attn_mask)
-        if kv_cache is not None:
-            encoded = explicit_attention(operands)
-        elif isinstance(use_cudnn_attention, bool):
-            encoded = cudnn_attention(operands) if use_cudnn_attention else explicit_attention(operands)
+        if self.use_cudnn_attention and kv_cache is None:
+            encoded = cudnn_attention(operands)
         else:
-            encoded = jax.lax.cond(use_cudnn_attention, cudnn_attention, explicit_attention, operands)
+            encoded = explicit_attention(operands)
 
         out = []
         start = 0
@@ -447,16 +439,13 @@ class Block(nn.Module):
 
     configs: tuple[Config, ...]
     use_cudnn_attention: bool = False
-    cudnn_attention_layer_start: int = 0
-    cudnn_attention_num_layers: int | None = None
     cudnn_attention_dtype: str = "bfloat16"
-    explicit_attention_fp32: bool = False
 
     dropout: float = 0.0
     dropout_bdims: tuple[int, ...] = ()
 
     @nn.compact
-    def __call__(self, xs, kv_cache, positions, attn_mask, adarms_cond, layer_index=None, deterministic=True):
+    def __call__(self, xs, kv_cache, positions, attn_mask, adarms_cond, deterministic=True):
         xs = sharding.activation_sharding_constraint(xs)
         drop = nn.Dropout(self.dropout, self.dropout_bdims) if self.dropout else lambda x, _: x
 
@@ -464,16 +453,8 @@ def __call__(self, xs, kv_cache, positions, attn_mask, adarms_cond, layer_index=
             configs=self.configs,
             use_cudnn_attention=self.use_cudnn_attention,
             cudnn_attention_dtype=self.cudnn_attention_dtype,
-            explicit_attention_fp32=self.explicit_attention_fp32,
             name="attn",
         )
-        layer_uses_cudnn = self.use_cudnn_attention
-        if layer_uses_cudnn and self.cudnn_attention_num_layers is not None:
-            layer_uses_cudnn = jnp.logical_and(
-                layer_index >= self.cudnn_attention_layer_start,
-                layer_index < self.cudnn_attention_layer_start + self.cudnn_attention_num_layers,
-            )
-
         pre_attn = []
         gates = []
         for i, x in enumerate(xs):
@@ -485,7 +466,7 @@ def __call__(self, xs, kv_cache, positions, attn_mask, adarms_cond, layer_index=
             pre_attn.append(x)
 
         pre_attn = sharding.activation_sharding_constraint(pre_attn)
-        post_attn, kv_cache = attn(pre_attn, positions, attn_mask, kv_cache, use_cudnn_attention=layer_uses_cudnn)
+        post_attn, kv_cache = attn(pre_attn, positions, attn_mask, kv_cache)
         post_attn = jax.tree.map(lambda x: drop(x, deterministic), post_attn)
         post_attn = sharding.activation_sharding_constraint(post_attn)
         xs = [_gated_residual(x, y, gate) for x, y, gate in zip(xs, post_attn, gates, strict=True)]
@@ -529,11 +510,7 @@ class Module(nn.Module):
     dropout_bdims: tuple[int, ...] = ()  # Every float is dropped independently.
     adarms: bool = False
     use_cudnn_attention: bool = False
-    cudnn_attention_layer_start: int = 0
-    cudnn_attention_num_layers: int | None = None
     cudnn_attention_dtype: str = "bfloat16"
-    explicit_attention_fp32: bool = False
-    remat_policy: str = "nothing_saveable"
 
     def setup(self):
         # all experts must have the same depth
@@ -544,15 +521,12 @@ def setup(self):
             embed_dim=self.configs[0].width,  # embedder for first expert only
             name="embedder",
         )
-        if self.remat_policy == "none":
-            block_cls = Block
-        else:
-            block_cls = nn.remat(
-                Block,
-                prevent_cse=False,
-                static_argnums=(6,),  # 0=self, 7=deterministic
-                policy=getattr(jax.checkpoint_policies, self.remat_policy),
-            )
+        block_cls = nn.remat(
+            Block,
+            prevent_cse=False,
+            static_argnums=(5,),  # 0=self, 6=deterministic
+            policy=jax.checkpoint_policies.nothing_saveable,
+        )
         self.layers = nn.scan(
             block_cls,
             variable_axes={"params": 0},
@@ -562,17 +536,13 @@ def setup(self):
                 nn.broadcast,
                 nn.broadcast,
                 nn.broadcast,
-                0,
                 nn.broadcast,
-            ),  # 0=kv_cache, 1=positions, 2=mask, 3=adarms_cond, 4=layer_index, 5=deterministic
+            ),  # 0=kv_cache, 1=positions, 2=mask, 3=adarms_cond, 4=deterministic
             length=self.configs[0].depth,
         )(
             configs=self.configs,
             use_cudnn_attention=self.use_cudnn_attention,
-            cudnn_attention_layer_start=self.cudnn_attention_layer_start,
-            cudnn_attention_num_layers=self.cudnn_attention_num_layers,
             cudnn_attention_dtype=self.cudnn_attention_dtype,
-            explicit_attention_fp32=self.explicit_attention_fp32,
             dropout=self.dropout,
             dropout_bdims=self.dropout_bdims,
         )
@@ -599,8 +569,7 @@ def __call__(
         if adarms_cond is None:
             adarms_cond = [None] * len(self.configs)
 
-        layer_indices = jnp.arange(self.configs[0].depth, dtype=jnp.int32)
-        embedded, kv_cache = self.layers(embedded, kv_cache, positions, mask, adarms_cond, layer_indices, deterministic)
+        embedded, kv_cache = self.layers(embedded, kv_cache, positions, mask, adarms_cond, deterministic)
 
         assert all(e.dtype == jnp.dtype(self.embed_dtype) for e in embedded if e is not None)
 
diff --git a/src/openpi/models/pi0.py b/src/openpi/models/pi0.py
index 284fc79..d9f74c4 100644
--- a/src/openpi/models/pi0.py
+++ b/src/openpi/models/pi0.py
@@ -95,11 +95,7 @@ def __init__(
                 embed_dtype=config.dtype,
                 adarms=config.pi05,
                 use_cudnn_attention=config.use_cudnn_attention,
-                cudnn_attention_layer_start=config.cudnn_attention_layer_start,
-                cudnn_attention_num_layers=config.cudnn_attention_num_layers,
                 cudnn_attention_dtype=config.cudnn_attention_dtype,
-                explicit_attention_fp32=config.explicit_attention_fp32,
-                remat_policy=config.gemma_remat_policy,
             )
         )
         llm.lazy_init(
@@ -113,7 +109,6 @@ def __init__(
                 variant="So400m/14",
                 pool_type="none",
                 scan=True,
-                remat_policy=config.siglip_remat_policy,
                 dtype_mm=config.dtype,
             )
         )
diff --git a/src/openpi/models/pi0_config.py b/src/openpi/models/pi0_config.py
index 95ebad7..a5ba5a9 100644
--- a/src/openpi/models/pi0_config.py
+++ b/src/openpi/models/pi0_config.py
@@ -36,24 +36,9 @@ class Pi0Config(_model.BaseModelConfig):
     # Use cuDNN fused attention for training-time Gemma attention. This requires
     # a cuDNN runtime that supports the model's BF16 GQA shapes and attention mask.
     use_cudnn_attention: bool = False
-    # Diagnostic hybrid mode. None keeps the historical all-layer behavior;
-    # otherwise only this contiguous layer range uses cuDNN.
-    cudnn_attention_layer_start: int = 0
-    cudnn_attention_num_layers: int | None = None
     # Compute dtype for the cuDNN kernel: "bfloat16" or "float16". float16 uses a
     # dynamically loss-scaled custom VJP (see gemma._cudnn_attention_in_dtype).
     cudnn_attention_dtype: str = "bfloat16"
-    # Diagnostic only: explicit attention with fp32 q/k/v/probs as a numerical reference.
-    explicit_attention_fp32: bool = False
-    # Activation rematerialization policy for the Gemma transformer blocks and the SigLIP
-    # encoder blocks. "nothing_saveable" (default) recomputes every block in the backward
-    # pass; "none" disables remat entirely (fastest, most memory); any other name is looked
-    # up in jax.checkpoint_policies (e.g. "dots_with_no_batch_dims_saveable").
-    # NOTE: on 8xH100 80GB with global batch 256 neither "none" (XLA wants a 280 GiB buffer)
-    # nor "dots_with_no_batch_dims_saveable" (84.6 GB) fits; see docs/training-optimization.md.
-    gemma_remat_policy: str = "nothing_saveable"
-    siglip_remat_policy: str = "nothing_saveable"
-
     # training-time RTC config
     enable_training_time_rtc: bool = False
     max_delay: int = 10  # steps 330ms @ 30fps
@@ -63,10 +48,6 @@ def __post_init__(self):
             object.__setattr__(self, "max_token_len", 200 if self.pi05 else 48)
         if self.discrete_state_input is None:
             object.__setattr__(self, "discrete_state_input", self.pi05)
-        if self.cudnn_attention_layer_start < 0:
-            raise ValueError("cudnn_attention_layer_start must be non-negative")
-        if self.cudnn_attention_num_layers is not None and self.cudnn_attention_num_layers < 0:
-            raise ValueError("cudnn_attention_num_layers must be non-negative or None")
         if self.cudnn_attention_dtype not in ("bfloat16", "float16"):
             raise ValueError("cudnn_attention_dtype must be 'bfloat16' or 'float16'")
 
diff --git a/src/openpi/models/siglip.py b/src/openpi/models/siglip.py
index 5b02eca..fd45c29 100644
--- a/src/openpi/models/siglip.py
+++ b/src/openpi/models/siglip.py
@@ -124,15 +124,12 @@ def __call__(self, x, deterministic=True):
         out = {}
 
         if self.scan:
-            if self.remat_policy == "none":
-                block = Encoder1DBlock
-            else:
-                block = nn.remat(
-                    Encoder1DBlock,
-                    prevent_cse=False,
-                    static_argnums=(2,),  # 0=self, 2=deterministic
-                    policy=getattr(jax.checkpoint_policies, self.remat_policy, None),
-                )
+            block = nn.remat(
+                Encoder1DBlock,
+                prevent_cse=False,
+                static_argnums=(2,),  # 0=self, 2=deterministic
+                policy=getattr(jax.checkpoint_policies, self.remat_policy, None),
+            )
             x, scan_out = nn.scan(
                 block,
                 variable_axes={"params": 0},
diff --git a/src/openpi/policies/dobot_nova5_policy.py b/src/openpi/policies/dobot_nova5_policy.py
index 7780f50..6683790 100644
--- a/src/openpi/policies/dobot_nova5_policy.py
+++ b/src/openpi/policies/dobot_nova5_policy.py
@@ -21,7 +21,6 @@
 
 from openpi import transforms
 
-
 TCP = slice(0, 9)
 GRIPPER = slice(9, 10)
 STATE_DIM = 10
diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py
index 70b30b9..9f5b059 100644
--- a/src/openpi/training/config.py
+++ b/src/openpi/training/config.py
@@ -644,10 +644,6 @@ class TrainConfig:
     optimizer: _optimizer.OptimizerConfig = dataclasses.field(default_factory=_optimizer.AdamW)
     ema_decay: float | None = 0.99
 
-    # Diagnostic-only relative L2 perturbation applied to gradients before the
-    # optimizer update. Zero preserves the normal training path exactly.
-    gradient_noise_scale: float = 0.0
-
     # Specifies which weights should be frozen.
     freeze_filter: tyro.conf.Suppress[Filter] = dataclasses.field(default_factory=nnx.Nothing)
 
@@ -673,17 +669,6 @@ class TrainConfig:
     # workers finish them. The sampler still visits every index exactly once per pass, so this
     # only reorders batches that were already randomly shuffled.
     strict_batch_order: bool = False
-    # If true, collect fine-grained data-pipeline timings. This adds per-sample clocks in
-    # workers and synchronizes device transfers, so it is intended only for short profiling
-    # runs rather than production training.
-    profile_data_pipeline: bool = False
-    # Number of steps between fine-grained data-pipeline timing lines.
-    profile_log_interval: int = 1
-    # If set, write a JAX/XProf trace to this directory. The trace begins at
-    # xprof_start_step and covers xprof_num_steps complete training iterations.
-    xprof_trace_dir: str | None = None
-    xprof_start_step: int = 0
-    xprof_num_steps: int = 20
     # Number of train steps (batches) to run.
     num_train_steps: int = 30_000
 
@@ -691,9 +676,6 @@ class TrainConfig:
     log_interval: int = 100
     # How often (in steps) to save checkpoints.
     save_interval: int = 1000
-    # Whether to save the last step even when it is not on save_interval. Useful to disable
-    # for short profiling runs restored from an existing checkpoint.
-    save_final_checkpoint: bool = True
     # If set, any existing checkpoints matching step % keep_period == 0 will not be deleted.
     keep_period: int | None = 5000
 
@@ -734,8 +716,6 @@ def trainable_filter(self) -> nnx.filterlib.Filter:
     def __post_init__(self) -> None:
         if self.resume and self.overwrite:
             raise ValueError("Cannot resume and overwrite at the same time.")
-        if self.xprof_num_steps <= 0:
-            raise ValueError("xprof_num_steps must be positive.")
 
 
 # YAML lookup paths, in priority order. First match wins.
diff --git a/src/openpi/training/data_loader.py b/src/openpi/training/data_loader.py
index ad7a01e..6888ebc 100644
--- a/src/openpi/training/data_loader.py
+++ b/src/openpi/training/data_loader.py
@@ -1,5 +1,4 @@
 from collections.abc import Iterator, Sequence
-import functools
 import logging
 import multiprocessing
 import os
@@ -90,29 +89,10 @@ def data_config(self) -> _config.DataConfig:
     def __iter__(self) -> Iterator[T_co]:
         raise NotImplementedError("Subclasses of DataLoader should implement __iter__.")
 
-    def profile_stats(self) -> dict[str, float] | None:
-        """Return timings for the most recently yielded batch, when profiling is enabled."""
-        return None
-
     def close(self) -> None:
         """Release worker processes and other loader resources."""
 
 
-class _TimedDataset:
-    """Attach worker-side __getitem__ latency to samples for profiling."""
-
-    def __init__(self, dataset: Dataset):
-        self._dataset = dataset
-
-    def __getitem__(self, index: SupportsIndex):
-        started = time.monotonic()
-        sample = self._dataset[index]
-        return sample, time.monotonic() - started
-
-    def __len__(self) -> int:
-        return len(self._dataset)
-
-
 class TransformedDataset(Dataset[T_co]):
     def __init__(self, dataset: Dataset, transforms: Sequence[_transforms.DataTransformFn]):
         self._dataset = dataset
@@ -408,7 +388,6 @@ def create_data_loader(
         skip_norm_stats=skip_norm_stats,
         framework=framework,
         strict_batch_order=config.strict_batch_order,
-        profile_data_pipeline=config.profile_data_pipeline,
     )
 
 
@@ -426,7 +405,6 @@ def create_torch_data_loader(
     seed: int = 0,
     framework: str = "jax",
     strict_batch_order: bool = False,
-    profile_data_pipeline: bool = False,
 ) -> DataLoader[tuple[_model.Observation, _model.Actions]]:
     """Create a data loader for training.
 
@@ -479,7 +457,6 @@ def create_torch_data_loader(
         seed=seed,
         framework=framework,
         strict_batch_order=strict_batch_order,
-        profile_data_pipeline=profile_data_pipeline,
     )
 
     return DataLoaderImpl(data_config, data_loader)
@@ -577,7 +554,6 @@ def __init__(
         seed: int = 0,
         framework: str = "jax",
         strict_batch_order: bool = False,
-        profile_data_pipeline: bool = False,
     ):
         """Create a PyTorch data loader.
 
@@ -599,9 +575,6 @@ def __init__(
         if jax.process_count() > 1:
             raise NotImplementedError("Data loading with multiple processes is not supported.")
 
-        if profile_data_pipeline:
-            dataset = _TimedDataset(dataset)
-
         if len(dataset) < local_batch_size:
             raise ValueError(f"Local batch size ({local_batch_size}) is larger than the dataset size ({len(dataset)}).")
 
@@ -614,8 +587,6 @@ def __init__(
                 jax.sharding.PartitionSpec("B"),
             )
         self._num_batches = num_batches
-        self._profile_data_pipeline = profile_data_pipeline
-        self._last_profile_stats: dict[str, float] | None = None
         self._active_iterator = None
 
         mp_context = None
@@ -640,7 +611,7 @@ def __init__(
             prefetch_factor=16
             if num_workers > 0
             else None,  # stall-fix: 32 over-buffered for cold start, 4 was too low vs slow workers, 16 is the sweet spot
-            collate_fn=functools.partial(_collate_fn, profile=profile_data_pipeline),
+            collate_fn=_collate_fn,
             worker_init_fn=_worker_init_fn,
             drop_last=True,
             # Hand batches over as workers finish them. In strict order, the one worker that is
@@ -674,14 +645,10 @@ def __iter__(self):
                     except StopIteration:
                         break  # We've exhausted the dataset. Create a new iterator and start over.
                     dt = time.monotonic() - t_batch
-                    worker_stats: dict[str, float] = {}
-                    if self._profile_data_pipeline:
-                        batch, worker_stats = batch
                     if num_items < 3 or dt > 3.0:
                         logging.info(f"[TorchDataLoader] batch #{num_items} next() took {dt:.2f}s")
                     num_items += 1
                     # For JAX, convert to sharded arrays; for PyTorch, return torch tensors
-                    t_array = time.monotonic()
                     if self._sharding is not None:
                         batch = jax.tree.map(
                             lambda x: jax.make_array_from_process_local_data(self._sharding, _to_numpy_view(x)),
@@ -689,28 +656,10 @@ def __iter__(self):
                         )
                     else:
                         batch = jax.tree.map(torch.as_tensor, batch)
-                    array_construct_s = time.monotonic() - t_array
-
-                    h2d_wait_s = 0.0
-                    if self._profile_data_pipeline and self._sharding is not None:
-                        t_h2d = time.monotonic()
-                        jax.block_until_ready(batch)
-                        h2d_wait_s = time.monotonic() - t_h2d
-
-                    if self._profile_data_pipeline:
-                        self._last_profile_stats = {
-                            "main_queue_wait_s": dt,
-                            **worker_stats,
-                            "jax_array_construct_s": array_construct_s,
-                            "h2d_wait_s": h2d_wait_s,
-                        }
                     yield batch
         finally:
             self.close()
 
-    def profile_stats(self) -> dict[str, float] | None:
-        return self._last_profile_stats
-
     def close(self) -> None:
         """Shut down the active multiprocessing iterator before interpreter teardown."""
         data_iter = self._active_iterator
@@ -745,27 +694,15 @@ def _to_numpy_view(value):
     return np.asarray(value)
 
 
-def _collate_fn(items, *, profile: bool = False):
+def _collate_fn(items):
     """Collate batch elements into CPU tensors so workers use shared-memory IPC."""
-    getitem_s = 0.0
-    if profile:
-        samples, durations = zip(*items, strict=True)
-        items = samples
-        getitem_s = sum(durations)
-    started = time.monotonic()
     # Some inputs are JAX arrays, so stack through NumPy first. Returning Torch tensors is
     # important here: the DataLoader multiprocessing reducer transfers their storage through
     # shared memory instead of pickling each ~900 MiB NumPy batch into the result queue.
-    batch = jax.tree.map(
+    return jax.tree.map(
         lambda *xs: torch.from_numpy(np.stack([np.asarray(x) for x in xs], axis=0)),
         *items,
     )
-    if not profile:
-        return batch
-    return batch, {
-        "worker_getitem_s": getitem_s,
-        "worker_collate_s": time.monotonic() - started,
-    }
 
 
 def _worker_init_fn(worker_id: int) -> None:
@@ -848,10 +785,6 @@ def __iter__(self):
         for batch in self._data_loader:
             yield _model.Observation.from_dict(batch), batch["actions"]
 
-    def profile_stats(self) -> dict[str, float] | None:
-        profile_stats = getattr(self._data_loader, "profile_stats", None)
-        return profile_stats() if profile_stats is not None else None
-
     def close(self) -> None:
         close = getattr(self._data_loader, "close", None)
         if close is not None:
diff --git a/src/openpi/training/data_loader_test.py b/src/openpi/training/data_loader_test.py
index 7562494..8f6e475 100644
--- a/src/openpi/training/data_loader_test.py
+++ b/src/openpi/training/data_loader_test.py
@@ -100,32 +100,6 @@ def test_torch_data_loader_close_shuts_down_workers():
     assert all(not worker.is_alive() for worker in workers)
 
 
-def test_torch_data_loader_profile_stats():
-    config = pi0_config.Pi0Config(action_dim=24, action_horizon=50, max_token_len=48)
-    batch_size = jax.device_count()
-    dataset = _data_loader.FakeDataset(config, batch_size * 2)
-
-    loader = _data_loader.TorchDataLoader(
-        dataset,
-        local_batch_size=batch_size,
-        num_batches=1,
-        num_workers=2,
-        profile_data_pipeline=True,
-    )
-    _ = next(iter(loader))
-    stats = loader.profile_stats()
-
-    assert stats is not None
-    assert set(stats) == {
-        "main_queue_wait_s",
-        "worker_getitem_s",
-        "worker_collate_s",
-        "jax_array_construct_s",
-        "h2d_wait_s",
-    }
-    assert all(value >= 0 for value in stats.values())
-
-
 def test_with_fake_dataset():
     config = _config.get_config("debug_pi05")
 

From 0b8330d78b01c185a43dfeb840b3ab01cccd3d49 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=9A=20=E8=83=A1?= <1195132365@qq.com>
Date: Thu, 10 Sep 2026 17:34:37 +0800
Subject: [PATCH 22/28] Update Flexiv inference recipes for serial grippers

---
 .../recipes/diagonal-02.yaml                  | 27 +++++++------------
 .../recipes/forward-04.yaml                   | 22 +++++++--------
 2 files changed, 20 insertions(+), 29 deletions(-)

diff --git a/examples/bi_flexiv_rizon4_rt/recipes/diagonal-02.yaml b/examples/bi_flexiv_rizon4_rt/recipes/diagonal-02.yaml
index 5a51215..d4bd534 100644
--- a/examples/bi_flexiv_rizon4_rt/recipes/diagonal-02.yaml
+++ b/examples/bi_flexiv_rizon4_rt/recipes/diagonal-02.yaml
@@ -18,7 +18,7 @@ robot:
   right_home_position_degree: [31.91, -26.72, -27.42, 122.58, -13.93, 89.42, -1.31]
 
   # ── Cameras ─────────────────────────────────────────────────────────────────
-  # Only the head is pinned: the taccap gripper sniffs its own wrist + GSPS
+  # Only the head is pinned: the serial gripper sniffs its own wrist + GSPS
   # tactile cameras off its USB hub at connect (gripper.auto_discover_cameras).
   cameras:
     head:
@@ -31,23 +31,16 @@ robot:
       use_depth: false
 
   # ── Gripper ── One block for both arms; `side` is stamped per arm at build time.
-  # Both grippers -> xense.taccap FollowerGripper (matched pair, single switch).
+  # Both grippers use the serial gripper interface (matched pair, single switch).
   gripper:
-    type: taccap_follower
-    kp: 8.0                 # MIT impedance stiffness (Nm/rad)
-    kd: 1.0                 # MIT impedance damping (Nm·s/rad)
-    # ⚠️ HIGH GRIP FORCE — constant closing feed-forward torque. NEGATIVE = clamp.
-    # -3.0 Nm is ~10x the SDK gentle-grasp abort (0.30 Nm) and near the motor's
-    # ceiling: it is a hard crush. Being a CONSTANT bias, the jaw also settles
-    # ~3/kp≈0.375 rad short of full open and holds the closed stop under load even
-    # empty (runs hot). Back off toward -0.3 if it stalls / can't open / overheats.
-    feedforward_torque: -3.0    # Nm, negative = closing
-    control_hz: 200             # ControlLoop resubmit rate
-    auto_discover_cameras: true
-    require_calibrated: true
+    type: serial
+    gripper_v_max: 100.0    # mm/s
+    gripper_f_max: 30.0     # N
+    gripper_min_pos: 0.0    # mm — fully closed
+    gripper_max_pos: 85.0   # mm — fully open
+    serial_timeout: 1.0
     init_open: true
-    # Everything else worth tuning, at its default — uncomment to change.
-    # side: left        # stamped automatically per arm — do not set here
-    # mcu_device: null  # explicit MCU path; the arm fills this from discovery
+    auto_discover_cameras: true   # sniff wrist + tactile off this gripper's USB hub
+    enable_tactile: false
   left_use_gripper: true
   right_use_gripper: true
diff --git a/examples/bi_flexiv_rizon4_rt/recipes/forward-04.yaml b/examples/bi_flexiv_rizon4_rt/recipes/forward-04.yaml
index 2ed2175..d6102f4 100644
--- a/examples/bi_flexiv_rizon4_rt/recipes/forward-04.yaml
+++ b/examples/bi_flexiv_rizon4_rt/recipes/forward-04.yaml
@@ -18,7 +18,7 @@ robot:
   right_home_position_degree: [83.07, 67.85, 24.64, 107.57, 6.39, 95.96, 15.24]
 
   # ── Cameras ─────────────────────────────────────────────────────────────────
-  # Only the head is pinned: the taccap gripper sniffs its own wrist + GSPS
+  # Only the head is pinned: the serial gripper sniffs its own wrist + GSPS
   # tactile cameras off its USB hub at connect (gripper.auto_discover_cameras).
   cameras:
     head:
@@ -31,18 +31,16 @@ robot:
       use_depth: false
 
   # ── Gripper ── One block for both arms; `side` is stamped per arm at build time.
-  # Both grippers -> xense.taccap FollowerGripper (matched pair, single switch).
+  # Both grippers use the serial gripper interface (matched pair, single switch).
   gripper:
-    type: taccap_follower
-    kp: 8.0                     # MIT impedance stiffness (Nm/rad)
-    kd: 1.0                     # MIT impedance damping (Nm·s/rad)
-    feedforward_torque: 0.0     # Nm; negative = clamp harder. |ff| <= 3.5
-    control_hz: 200             # ControlLoop resubmit rate
-    auto_discover_cameras: true
-    require_calibrated: true
+    type: serial
+    gripper_v_max: 100.0    # mm/s
+    gripper_f_max: 30.0     # N
+    gripper_min_pos: 0.0    # mm — fully closed
+    gripper_max_pos: 85.0   # mm — fully open
+    serial_timeout: 1.0
     init_open: true
-    # Everything else worth tuning, at its default — uncomment to change.
-    # side: left        # stamped automatically per arm — do not set here
-    # mcu_device: null  # explicit MCU path; the arm fills this from discovery
+    auto_discover_cameras: true   # sniff wrist + tactile off this gripper's USB hub
+    enable_tactile: false
   left_use_gripper: true
   right_use_gripper: true

From bbccfaed4b0ca5ab89d54d4c1171b3a63e1e0fd3 Mon Sep 17 00:00:00 2001
From: Vertax42 <1007644658@qq.com>
Date: Fri, 11 Sep 2026 14:11:15 +0800
Subject: [PATCH 23/28] fix(configs): restore three configs that callers still
 reference

4606754 and e13ac0f deleted pi05_droid, pi05_base_xtac_umi_pick_up_cube_0807_h200
and pi05_base_bi_flexiv_pack_6_cosmetic_bottles_lora without updating anything that
names them, and get_config() raises on a missing name rather than falling back:

- xtac_umi_policy_test.py:112,126 call get_config() on the XTac-UMI config, so two
  tests fail outright.
- serve_policy.py:70 maps EnvMode.DROID to pi05_droid; that was the one entry in
  DEFAULT_CHECKPOINT that still resolved.
- README.md, examples/bi_flexiv_rizon4_rt/README.md,
  examples/xtac_umi_bi_flexiv_rizon4_rt/README.md and _FULL_REFERENCE.yaml link to
  all three.

Restore them verbatim from 7c51dc5. The earbuds 0520_a100 config stays deleted: the
0904 h100/h200 configs supersede it and only the changelog mentions it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 ...i_flexiv_pack_6_cosmetic_bottles_lora.yaml | 24 +++++++++++
 ..._base_xtac_umi_pick_up_cube_0807_h200.yaml | 42 +++++++++++++++++++
 configs/_examples/pi05_droid.yaml             | 12 ++++++
 3 files changed, 78 insertions(+)
 create mode 100644 configs/_examples/pi05_base_bi_flexiv_pack_6_cosmetic_bottles_lora.yaml
 create mode 100644 configs/_examples/pi05_base_xtac_umi_pick_up_cube_0807_h200.yaml
 create mode 100644 configs/_examples/pi05_droid.yaml

diff --git a/configs/_examples/pi05_base_bi_flexiv_pack_6_cosmetic_bottles_lora.yaml b/configs/_examples/pi05_base_bi_flexiv_pack_6_cosmetic_bottles_lora.yaml
new file mode 100644
index 0000000..d9dfebd
--- /dev/null
+++ b/configs/_examples/pi05_base_bi_flexiv_pack_6_cosmetic_bottles_lora.yaml
@@ -0,0 +1,24 @@
+model:
+  type: Pi0Config
+  max_token_len: 200
+  paligemma_variant: gemma_2b_lora
+  action_expert_variant: gemma_300m_lora
+  pi05: true
+  discrete_state_input: true
+  enable_training_time_rtc: true
+weight_loader:
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+ema_decay: null
+data:
+  type: LeRobotBiFlexivDataConfig
+  repo_id: Xense/pack_6_cosmetic_bottles_into_carton
+  base_config:
+    prompt_from_task: true
+  default_prompt: Pick up six cosmetic bottles one by one and pack them into the carton
+    box. The box is narrow, so align each bottle carefully and insert it precisely.
+batch_size: 64
+num_train_steps: 20000
+# freeze_filter is intentionally absent. The loader sees the `lora` substring in the
+# model variants above and derives the flax filter via Pi0Config.get_freeze_filter();
+# see test_lora_freeze_filter_auto_derived in src/openpi/training/yaml_loader_test.py.
diff --git a/configs/_examples/pi05_base_xtac_umi_pick_up_cube_0807_h200.yaml b/configs/_examples/pi05_base_xtac_umi_pick_up_cube_0807_h200.yaml
new file mode 100644
index 0000000..3fd461d
--- /dev/null
+++ b/configs/_examples/pi05_base_xtac_umi_pick_up_cube_0807_h200.yaml
@@ -0,0 +1,42 @@
+# XTac-UMI bimanual dataset, recorded on the handheld bi_taccap_gripper rig.
+#
+# State/action are 20D in the rig's own per-side-grouped feature order:
+#   [left_tcp(0-8), left_gripper(9), right_tcp(10-18), right_gripper(19)]
+# No head camera: base_0_rgb is black with image_mask=False; the two wrist
+# cameras stay in their matching wrist slots. The dataset also carries four
+# tactile streams (left/right x tactile_left/tactile_right) that this config
+# does not consume - pi05 takes wrist RGB and the 20D vector only.
+#
+# Deploy target: bi_flexiv (dual Flexiv Rizon4 RT with taccap grippers), which
+# assembles the same per-side-grouped order — see
+# examples/xtac_umi_bi_flexiv_rizon4_rt and src/openpi/policies/xtac_umi_policy.py.
+model:
+  type: Pi0Config
+  max_token_len: 200
+  pi05: true
+  discrete_state_input: true
+  # enable_training_time_rtc: true
+weight_loader:
+  type: CheckpointWeightLoader
+  params_path: gs://openpi-assets/checkpoints/pi05_base/params
+ema_decay: null
+data:
+  type: LeRobotXtacUmiDataConfig
+  repo_id: TacVerse/taccap-g1-pick-up-cube-0807
+  use_delta_cartesian_actions: true
+  # use_head_camera: true   # only for a bench that actually recorded one
+  base_config:
+    prompt_from_task: true
+  # Only a fallback: prompt_from_task above means the dataset's own task string
+  # is what actually reaches the model.
+  default_prompt: Use the left arm to pick the red cube and place it into the blue
+    plastic box first, then use the right arm to grasp the green cube and put it into
+    the same blue plastic box after the first action is fully finished
+batch_size: 256
+num_workers: 64
+# 150 episodes / 43447 frames, i.e. ~290 frames (~9.7 s) per episode - roughly
+# half the frames of the 0710 sort-defective-parts set. 30k steps at batch 256 is
+# ~177 passes over the data, in the same range as the other single-task configs
+# here; raise it only if the loss is still moving.
+num_train_steps: 30000
+fsdp_devices: 8
diff --git a/configs/_examples/pi05_droid.yaml b/configs/_examples/pi05_droid.yaml
new file mode 100644
index 0000000..5ad3728
--- /dev/null
+++ b/configs/_examples/pi05_droid.yaml
@@ -0,0 +1,12 @@
+model:
+  type: Pi0Config
+  action_horizon: 15
+  max_token_len: 200
+  pi05: true
+  discrete_state_input: true
+data:
+  type: DroidInferenceDataConfig
+  assets:
+    asset_id: droid
+  base_config:
+    prompt_from_task: true

From 1db0ec9d750cdf2cf6fa91eef175b1f8874e1059 Mon Sep 17 00:00:00 2001
From: Vertax42 <1007644658@qq.com>
Date: Fri, 11 Sep 2026 14:11:26 +0800
Subject: [PATCH 24/28] fix(training): keep a NaN loss visible in the logged
 metrics

Computing the whole-model norms only on logging steps means the other steps carry NaN
placeholders, and the logging reduction switched wholesale to jnp.nanmean to skip them.
That also applies to `loss`, so a diverged step is now dropped from the average instead
of showing up as loss=nan -- exactly the signal the FP16 cuDNN path needs, given that
the per-step finite-update guard is gone and grad_norm is only sampled every
log_interval steps.

Reduce only grad_norm and param_norm with nanmean; everything else keeps a plain mean.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 scripts/train.py | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/scripts/train.py b/scripts/train.py
index 2e05621..77c9758 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -135,6 +135,12 @@ def init(rng: at.KeyArrayLike, partial_params: at.Params | None = None) -> train
     return train_state, state_sharding
 
 
+# Metrics that `train_step` fills with NaN on the steps it skips them on, and that the
+# logging reduction therefore has to skip too. Everything else is reduced with a plain
+# mean so that a NaN reaches the log instead of disappearing.
+_NAN_PLACEHOLDER_METRICS = frozenset({"grad_norm", "param_norm"})
+
+
 @at.typecheck
 def train_step(
     config: _config.TrainConfig,
@@ -192,7 +198,7 @@ def loss_fn(
         param_norm = optax.global_norm(kernel_params)
     else:
         # Keep a stable output pytree for both JIT variants. The host-side logging
-        # reduction ignores these placeholders with nanmean.
+        # reduction drops these placeholders with nanmean; see _NAN_PLACEHOLDER_METRICS.
         grad_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
         param_norm = jnp.asarray(jnp.nan, dtype=loss.dtype)
 
@@ -306,7 +312,16 @@ def main(config: _config.TrainConfig):
         if step % config.log_interval == 0:
             t0 = time.monotonic()
             stacked_infos = common_utils.stack_forest(infos)
-            reduced_info = jax.device_get(jax.tree.map(jnp.nanmean, stacked_infos))
+            # Only the norms carry NaN placeholders from the steps that skipped the
+            # whole-model reductions, so only they may be reduced with nanmean. `loss`
+            # keeps a plain mean: a NaN loss is a divergence signal and has to reach
+            # the log rather than being silently dropped from the average.
+            reduced_info = jax.device_get(
+                {
+                    key: jnp.nanmean(value) if key in _NAN_PLACEHOLDER_METRICS else jnp.mean(value)
+                    for key, value in stacked_infos.items()
+                }
+            )
             info_str = ", ".join(f"{k}={v:.4f}" for k, v in reduced_info.items())
             pbar.write(f"Step {step}: {info_str}")
             # tqdm_loggable can drop pbar.write() lines when stdout is redirected.

From bc477a6602866c17ec2a10c8a9c7db4c0bb0219d Mon Sep 17 00:00:00 2001
From: Vertax42 <1007644658@qq.com>
Date: Fri, 11 Sep 2026 14:11:26 +0800
Subject: [PATCH 25/28] fix(models): validate the private cuDNN symbols instead
 of importing them eagerly

gemma.py imported jax._src.cudnn.fused_attention_stablehlo at module scope. The module
is not the fragile part -- jax.nn imports it too, so it is loaded before anything here
runs -- but the *names* are: _dot_product_attention_bwd_rule and friends are private and
only validated against the pinned jax 0.5.3. A rename would have surfaced as an
AttributeError from inside a traced VJP, mid-training.

Resolve the module through a cached helper on first use and check the seven symbols the
FP16 path calls, so a JAX upgrade fails once with an actionable message naming the
missing symbols and the use_cudnn_attention=false fallback. Explicit attention, CPU-only
use and inference never reach the helper.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 src/openpi/models/gemma.py | 59 ++++++++++++++++++++++++++++++++++----
 1 file changed, 53 insertions(+), 6 deletions(-)

diff --git a/src/openpi/models/gemma.py b/src/openpi/models/gemma.py
index 490f179..d91ce24 100644
--- a/src/openpi/models/gemma.py
+++ b/src/openpi/models/gemma.py
@@ -27,17 +27,12 @@
 
 from collections.abc import Sequence
 import dataclasses
+import functools
 from typing import Literal, TypeAlias
 
 import einops
 import flax.linen as nn
 import jax
-
-# Private JAX API (pinned jax 0.5.3): the cuDNN fused-attention fwd/bwd rules that
-# jax.nn.dot_product_attention(implementation="cudnn") wires into its custom_vjp.
-# Used by _cudnn_attention_in_dtype to run the float16 kernel with loss scaling
-# without recomputing the forward inside the backward.
-from jax._src.cudnn import fused_attention_stablehlo as _cudnn_fa
 import jax.numpy as jnp
 
 import openpi.models.lora as lora
@@ -166,6 +161,56 @@ def decode(self, x):
         return jnp.dot(x, self.input_embedding_table.T)
 
 
+# The private symbols _cudnn_attention_in_dtype and _cudnn_static_args call directly.
+_CUDNN_FA_REQUIRED = (
+    "_dot_product_attention_fwd_rule",
+    "_dot_product_attention_bwd_rule",
+    "_normalize_layout",
+    "check_cudnn_version",
+    "get_large_negative_number",
+    "should_export_dbias",
+    "MaskType",
+)
+
+
+@functools.cache
+def _cudnn_fused_attention():
+    """Resolve the private JAX cuDNN fused-attention module, on first use only.
+
+    These are the fwd/bwd rules that jax.nn.dot_product_attention(implementation="cudnn")
+    wires into its own custom_vjp; _cudnn_attention_in_dtype calls them directly so the
+    float16 kernel can run under loss scaling without recomputing the forward inside the
+    backward. They are private and only validated against the pinned jax 0.5.3.
+
+    The module itself is not the fragile part -- jax.nn imports it, so it is already loaded
+    by the time anything here runs. What a JAX upgrade actually breaks is the *names*: a
+    renamed `_dot_product_attention_bwd_rule` would otherwise surface as an AttributeError
+    from inside a traced VJP, mid-training. Checking them here turns that into one
+    actionable error, and only for configs that opted into cuDNN attention -- explicit
+    attention, CPU-only use and inference never reach this function.
+    """
+    try:
+        from jax._src.cudnn import fused_attention_stablehlo
+    except ImportError as error:  # pragma: no cover - depends on the installed jax
+        raise ImportError(
+            "cuDNN fused attention needs jax._src.cudnn.fused_attention_stablehlo, a private API "
+            f"validated against jax 0.5.3 and missing from the installed jax ({jax.__version__}). "
+            "Set model.use_cudnn_attention=false to use explicit attention, or revalidate the "
+            "fused path against this jax version. See docs/training-optimization.md."
+        ) from error
+
+    missing = [name for name in _CUDNN_FA_REQUIRED if not hasattr(fused_attention_stablehlo, name)]
+    if missing:
+        raise RuntimeError(
+            f"jax._src.cudnn.fused_attention_stablehlo (jax {jax.__version__}) no longer exposes "
+            f"{', '.join(missing)}. These are private APIs validated against jax 0.5.3. Set "
+            "model.use_cudnn_attention=false to use explicit attention, or revalidate the fused "
+            "path -- including convergence -- against this jax version. "
+            "See docs/training-optimization.md."
+        )
+    return fused_attention_stablehlo
+
+
 def _stop_gradient_for_fully_masked_queries(q, attn_mask):
     """Cut fully masked query rows out of the backward pass, keeping the mask intact.
 
@@ -214,6 +259,7 @@ def _cudnn_attention_in_dtype(q, k, v, attn_mask, compute_dtype):
     they exceed its range. Validate activation ranges when changing the model or data.
     """
     out_dtype = q.dtype
+    _cudnn_fa = _cudnn_fused_attention()
 
     # The forward and backward call JAX's cuDNN fwd/bwd rules directly (the same
     # functions jax.nn.dot_product_attention's own custom_vjp uses) instead of
@@ -268,6 +314,7 @@ def unscale(g):
 
 def _cudnn_static_args(bias_shape, query_shape):
     """Static parameters jax.nn.dot_product_attention(..., scale=1.0, implementation="cudnn") uses."""
+    _cudnn_fa = _cudnn_fused_attention()
     layout = _cudnn_fa._normalize_layout("BTNH")
     has_dbias = _cudnn_fa.should_export_dbias(bias_shape, query_shape, layout.value)
     # (scale, seed, dropout_rate, variadic_args, mask_type, layout, sliding_window_length, cudnn_version)

From 9cbea2ce0e78a4d7651d1909242e2b8830e6b265 Mon Sep 17 00:00:00 2001
From: Vertax42 <1007644658@qq.com>
Date: Fri, 11 Sep 2026 14:11:43 +0800
Subject: [PATCH 26/28] fix(scripts): gate the FP16 cuDNN path, and accept
 patch releases

The README calls check_cuda_stack.py the required gate before production training, but
section 3 only exercised jax.nn.dot_product_attention in bf16 -- the path this PR itself
documents as unsafe. The configuration it tells you to run,
`use_cudnn_attention: true` + `cudnn_attention_dtype: float16`, went through
gemma._cudnn_attention_in_dtype and was never checked.

Add a section 4 that differentiates the FP16 custom VJP at the same production shape
against the same fp32 reference, and assert the zero (not NaN, not garbage) q-gradient
it promises on the fully-masked rows padding produces.

Also compare versions on major.minor: pyproject pins ranges (torch >=2.11,<2.12 and
friends), so pinning 2.11.0+cu128 exactly here failed the gate on any upstream patch
release. The +cu128 suffix is still checked explicitly, since a CPU-only wheel installs
cleanly and silently removes everything this script exists to verify.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 README.md                     |  5 +-
 docs/training-optimization.md |  3 +-
 scripts/check_cuda_stack.py   | 88 ++++++++++++++++++++++++++++++-----
 3 files changed, 81 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index ec60af2..414bb36 100644
--- a/README.md
+++ b/README.md
@@ -243,8 +243,9 @@ env -u LD_LIBRARY_PATH \
 Do **not** prepend `$CONDA_PREFIX/lib` to `LD_LIBRARY_PATH`. The supported
 environment uses the pip CUDA 12.8 stack installed by `lerobot-xense`:
 PyTorch 2.11 pins cuDNN 9.19 and JAX shares that same runtime. Before training,
-run `python scripts/check_cuda_stack.py`; it rejects mixed library sources and
-checks the real BF16 GQA forward/backward shape.
+run `python scripts/check_cuda_stack.py`; it rejects mixed library sources and checks a
+real production-shape forward/backward for both the raw BF16 kernel and the FP16 custom
+VJP, including the fully-masked query rows.
 
 Full-layer cuDNN attention with the default `bfloat16` compute dtype is unsafe:
 it diverged from the explicit-attention baseline after roughly 1,000 steps.
diff --git a/docs/training-optimization.md b/docs/training-optimization.md
index 93cab98..bef77b4 100644
--- a/docs/training-optimization.md
+++ b/docs/training-optimization.md
@@ -83,7 +83,8 @@ q/k/v 与 mask/bias 显式使用一致的 batch 分片，避免 cuDNN partitione
 
 保持 pip CUDA 12.8 库来源一致，避免将 `$CONDA_PREFIX/lib` 加到 `LD_LIBRARY_PATH` 前面。
 当前依赖配置使用 PyTorch 2.11 对应的 cuDNN 9.19；训练启动时记录 JAX cuDNN runtime 版本。
-运行 `check_cuda_stack.py` 检查实际加载的库来源与 BF16 GQA 前后向；该检查不替代 FP16 收敛验证。
+运行 `check_cuda_stack.py` 检查实际加载的库来源,以及生产 shape 下 BF16 内核与 FP16 custom VJP
+的前后向(含全空 mask 行的 dQ)。该检查只覆盖单次前后向数值,不替代 FP16 收敛验证。
 
 ```bash
 env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py
diff --git a/scripts/check_cuda_stack.py b/scripts/check_cuda_stack.py
index cf55224..134068b 100644
--- a/scripts/check_cuda_stack.py
+++ b/scripts/check_cuda_stack.py
@@ -16,14 +16,26 @@
 import sys
 
 FAIL = []
+# Major.minor only: pyproject pins these as ranges (torch >=2.11,<2.12 and friends), so an
+# exact pin here would fail the gate on every upstream patch release even though the
+# environment is still the supported one.
 EXPECTED = {
-    "torch": "2.11.0+cu128",
-    "torchvision": "0.26.0+cu128",
-    "torchcodec": "0.11.1+cu128",
-    "nvidia-cudnn-cu12": "9.19.0.56",
+    "torch": "2.11",
+    "torchvision": "0.26",
+    "torchcodec": "0.11",
+    "nvidia-cudnn-cu12": "9.19",
     "cuda": "12.8",
-    "cudnn": 91900,
 }
+# The three torch wheels must come from the cu128 index. A CPU-only wheel installs cleanly
+# and silently removes everything this script is here to check.
+CU_SUFFIX = "+cu128"
+CU_WHEELS = ("torch", "torchvision", "torchcodec")
+EXPECTED_CUDNN = (9, 19)
+
+
+def major_minor(version):
+    """'2.11.0+cu128' -> '2.11'. Also handles '9.19.0.56' and '12.8'."""
+    return ".".join(re.split(r"[.+]", str(version))[:2])
 
 
 def loaded_libs(pattern):
@@ -64,6 +76,8 @@ def main():
     import torchcodec
     import torchvision
 
+    import openpi.models.gemma as _gemma
+
     section("1. reported versions")
     rt, build = cuda_versions.cudnn_get_version(), cuda_versions.cudnn_build_version()
     cudnn_package_version = importlib.metadata.version("nvidia-cudnn-cu12")
@@ -85,11 +99,17 @@ def main():
         "torchcodec": torchcodec.__version__,
         "nvidia-cudnn-cu12": cudnn_package_version,
         "cuda": torch.version.cuda,
-        "cudnn": rt,
     }
     for name, expected in EXPECTED.items():
-        if actual[name] != expected:
-            FAIL.append(f"{name} must be {expected}, got {actual[name]}")
+        if major_minor(actual[name]) != expected:
+            FAIL.append(f"{name} must be {expected}.x, got {actual[name]}")
+    FAIL.extend(
+        f"{name} must be a {CU_SUFFIX} wheel, got {actual[name]}"
+        for name in CU_WHEELS
+        if not str(actual[name]).endswith(CU_SUFFIX)
+    )
+    if (rt // 10000, rt // 100 % 100) != EXPECTED_CUDNN:
+        FAIL.append(f"cuDNN runtime must be {EXPECTED_CUDNN[0]}.{EXPECTED_CUDNN[1]}.x, got {rt}")
     if torch_rt != rt:
         FAIL.append(f"PyTorch reports cuDNN {torch_rt}, while JAX reports {rt}")
 
@@ -138,7 +158,9 @@ def main():
     mask = jnp.logical_and(cs[:, None, :] <= cs[:, :, None], im_j[:, None, :] * im_j[:, :, None])[:, None]
     valid = np.asarray(jnp.any(mask, -1))[0, 0]
     n_empty = int((~np.asarray(jnp.any(mask, -1))).sum())
-    print(f"  shape batch_size={batch_size} seq_len={seq_len} heads={num_heads} kv={num_kv_heads} head_dim={head_dim} bf16")
+    print(
+        f"  shape batch_size={batch_size} seq_len={seq_len} heads={num_heads} kv={num_kv_heads} head_dim={head_dim} bf16"
+    )
     print(f"  fully-masked query rows in batch: {n_empty}")
 
     rng = np.random.default_rng(0)
@@ -154,7 +176,9 @@ def reference(q, k, v, m):  # explicit path, fp32 throughout
         gq = q.reshape(batch_size, seq_len, num_kv_heads, num_heads // num_kv_heads, head_dim)
         lg = jnp.einsum("BTKGH,BSKH->BKGTS", gq, k, preferred_element_type=jnp.float32)
         lg = jnp.where(m[:, :, None], lg, -2.3819763e38)
-        return jnp.einsum("BKGTS,BSKH->BTKGH", jax.nn.softmax(lg, -1), v).reshape(batch_size, seq_len, num_heads, head_dim)
+        return jnp.einsum("BKGTS,BSKH->BTKGH", jax.nn.softmax(lg, -1), v).reshape(
+            batch_size, seq_len, num_heads, head_dim
+        )
 
     def cudnn(q, k, v, m):  # what training runs with use_cudnn_attention=true
         has_key = jnp.any(m, -1)[:, 0, :, None, None]
@@ -165,7 +189,9 @@ def loss(fn):
         return lambda q, k, v, m: jnp.sum(fn(q, k, v, m).astype(jnp.float32) * cot)
 
     try:
-        gr = [np.asarray(x.astype(jnp.float32), np.float64) for x in jax.grad(loss(reference), (0, 1, 2))(q, k, v, mask)]
+        gr = [
+            np.asarray(x.astype(jnp.float32), np.float64) for x in jax.grad(loss(reference), (0, 1, 2))(q, k, v, mask)
+        ]
         gc = [np.asarray(x.astype(jnp.float32), np.float64) for x in jax.grad(loss(cudnn), (0, 1, 2))(q, k, v, mask)]
     except Exception as exc:
         FAIL.append(f"cuDNN attention raised {type(exc).__name__}: {exc}")
@@ -182,13 +208,51 @@ def loss(fn):
                 FAIL.append(f"{nm}: rel-L2={rel:.3e} NaN={nan}")
             print(f"  {status} {nm}: rel-L2 vs fp32 = {rel:.3e}   NaN count = {nan}   (bf16 noise is ~3e-3)")
 
+    section("4. production FP16 custom VJP: forward + backward vs fp32 reference")
+
+    # Section 3 exercises the raw bf16 kernel, which docs/training-optimization.md marks
+    # unsafe for production. This is the path a production config actually runs:
+    #   model: {use_cudnn_attention: true, cudnn_attention_dtype: float16}
+    def cudnn_fp16(q, k, v, m):
+        return _gemma._cudnn_attention_in_dtype(q, k, v, m, jnp.float16)
+
+    gf = None
+    try:
+        gf = [
+            np.asarray(x.astype(jnp.float32), np.float64) for x in jax.grad(loss(cudnn_fp16), (0, 1, 2))(q, k, v, mask)
+        ]
+    except Exception as exc:
+        FAIL.append(f"FP16 cuDNN custom VJP raised {type(exc).__name__}: {exc}")
+        print(f"  !! FP16 cuDNN custom VJP FAILED: {type(exc).__name__}: {str(exc)[:300]}")
+
+    if gf is not None and gr is not None:
+        for i, nm in enumerate(("dQ", "dK", "dV")):
+            a, b = gf[i][:, valid], gr[i][:, valid]
+            rel = np.linalg.norm(a - b) / max(np.linalg.norm(b), 1e-30)
+            nan = int(np.isnan(gf[i]).sum())
+            status = "OK" if (nan == 0 and rel < 2e-2) else "!!"
+            if nan or rel >= 2e-2:
+                FAIL.append(f"FP16 {nm}: rel-L2={rel:.3e} NaN={nan}")
+            print(f"  {status} {nm}: rel-L2 vs fp32 = {rel:.3e}   NaN count = {nan}   (expect below the bf16 row)")
+
+        # The custom VJP promises a zero -- never NaN, never garbage -- q-gradient on the
+        # fully masked rows that padding produces. Check that at the production shape.
+        empty = ~valid
+        if empty.any():
+            nonzero = int(np.count_nonzero(gf[0][:, empty]))
+            status = "OK" if nonzero == 0 else "!!"
+            if nonzero:
+                FAIL.append(f"FP16 dQ is nonzero on {nonzero} fully-masked query entries")
+            print(f"  {status} dQ on {int(empty.sum())} fully-masked query rows: {nonzero} nonzero entries")
+
     section("verdict")
     if FAIL:
         print("  FAIL — do not start production training:")
         for f in FAIL:
             print(f"    - {f}")
         return 1
-    print("  PASS — single-source cuDNN stack, production-shape backward finite and accurate.")
+    print("  PASS — single-source cuDNN stack, production-shape backward finite and accurate,")
+    print("         for both the raw bf16 kernel and the FP16 custom VJP.")
     print("  NOTE: this gates the kernel only. It does NOT prove training converges;")
     print("        the 2026-08-30 divergence passed every check of this kind and still")
     print("        diverged between step 1000 and 1200. Watch the loss curve.")

From fef66f24b0cdc761625a58ccc323e0356f825a85 Mon Sep 17 00:00:00 2001
From: Vertax42 <1007644658@qq.com>
Date: Fri, 11 Sep 2026 14:11:43 +0800
Subject: [PATCH 27/28] docs(configs): fix the copy-paste in the
 cube-color-sorting config

The header block was copied from the earbuds config and kept its launch command, so it
told you to train pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.

Its max_token_len comment was copied too, and claims 128 was scanned "over this
dataset" -- but the 4037-state scan was the earbuds dataset, and
docs/training-optimization.md is explicit that 128 was validated there and nowhere else.
Say that 128 is carried over and not re-scanned against cube-color-sorting-0821, whose
prompt and state layout differ. @Hubo1231: if you did scan it, replace this with the
numbers; the tokenizer only warns and truncates on overflow, so it would not have shown
up in training.

Also repair the broken sentence in the earbuds comment the copy came from.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 ...se_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml | 11 ++++++-----
 ...v_earbuds_case_insertion_teleop_rtc_0904_h100.yaml |  4 ++--
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml
index e912566..58fb20d 100644
--- a/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml
@@ -12,14 +12,15 @@
 #     --xla_gpu_enable_pipelined_reduce_scatter=true \
 #     --xla_gpu_enable_while_loop_double_buffering=true" \
 #   XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
-#   python scripts/train.py pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100 \
+#   python scripts/train.py pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100 \
 #     --exp-name=<experiment-name>
 model:
   type: Pi0Config
-  # Scanned 2026-09-04 over this dataset (4037 states incl. per-dim extremes) with
-  # The longest tokenized "Task: ..., State: ...;\nAction: " prompt
-  # is 116 tokens, so 128 leaves 12 spare. Down from 200: -6.2% step time, identical loss
-  # (padding tokens were masked anyway). Re-scan when the dataset or prompt changes.
+  # ⚠️ Carried over from the earbuds config, where 128 was measured against that dataset's
+  # prompts (longest tokenized prompt 116 tokens). It has NOT been re-scanned against
+  # cube-color-sorting-0821, whose prompt and state layout differ. Scan the full transform
+  # chain before trusting it here -- the tokenizer only warns and truncates on overflow --
+  # or put 200 back. See docs/training-optimization.md.
   max_token_len: 128
   pi05: true
   discrete_state_input: true
diff --git a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
index c8eb631..f977abb 100644
--- a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
+++ b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml
@@ -16,8 +16,8 @@
 #     --exp-name=<experiment-name>
 model:
   type: Pi0Config
-  # Scanned 2026-09-04 over this dataset (4037 states incl. per-dim extremes) with
-  # The longest tokenized "Task: ..., State: ...;\nAction: " prompt
+  # Scanned 2026-09-04 over this dataset (4037 states incl. per-dim extremes): the longest
+  # tokenized "Task: ..., State: ...;\nAction: " prompt
   # is 116 tokens, so 128 leaves 12 spare. Down from 200: -6.2% step time, identical loss
   # (padding tokens were masked anyway). Re-scan when the dataset or prompt changes.
   max_token_len: 128

From 7ac2a6a54b83ee0ed5867251669decb97d206232 Mon Sep 17 00:00:00 2001
From: Vertax42 <1007644658@qq.com>
Date: Fri, 11 Sep 2026 14:11:43 +0800
Subject: [PATCH 28/28] docs(data-loader): tie the torch-internal shutdown to
 the version pin

close() drives _shutdown_workers, _workers and _iterator, none of which are public.
Record that this is what the torch>=2.11,<2.12 pin in pyproject.toml protects, and that
the getattr guards degrade to the old __del__ reaping rather than raising.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 src/openpi/training/data_loader.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/openpi/training/data_loader.py b/src/openpi/training/data_loader.py
index 6888ebc..0bcb577 100644
--- a/src/openpi/training/data_loader.py
+++ b/src/openpi/training/data_loader.py
@@ -661,7 +661,13 @@ def __iter__(self):
             self.close()
 
     def close(self) -> None:
-        """Shut down the active multiprocessing iterator before interpreter teardown."""
+        """Shut down the active multiprocessing iterator before interpreter teardown.
+
+        `_shutdown_workers`, `_workers` and `_iterator` are torch internals with no public
+        equivalent, so this is tied to the `torch>=2.11,<2.12` pin in pyproject.toml. Every
+        access is getattr-guarded: on a torch that renamed them the loader degrades to the
+        old behaviour (workers reaped by `__del__`) rather than raising.
+        """
         data_iter = self._active_iterator
         if data_iter is None:
             return