Optimize JAX training throughput and add FP16 cuDNN attention

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,6 +188,9 @@ model:
   paligemma_variant: gemma_2b
   action_expert_variant: gemma_300m
   enable_training_time_rtc: true
+  # Explicit attention is the default. For validated FP16 cuDNN training,
+  # see the configuration and validation notes below.
+  use_cudnn_attention: false
   max_delay: 10
 
 data:
@@ -208,10 +215,58 @@ 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>
+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>
+```
+
+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,
+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.
+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: true
+  cudnn_attention_dtype: float16
 ```
 
+At startup, `scripts/train.py` logs `JAX cuDNN runtime version: <version>`;
+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/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
 upstream URLs (`gs://openpi-assets/...`) or paths that every contributor can
@@ -246,7 +301,10 @@ 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
+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%).

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: "..."
@@ -239,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

configs/_examples/pi05_base_bi_flexiv_cube_color_sorting_rtc_0829_h100.yaml

@@ -0,0 +1,53 @@
+# 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_cube_color_sorting_rtc_0829_h100 \
+#     --exp-name=<experiment-name>
+model:
+  type: Pi0Config
+  # ⚠️ 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
+  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:
+  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
+# 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

configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h100.yaml

@@ -0,0 +1,51 @@
+# 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
+  # 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
+  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:
+  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: 80000
+# 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

configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_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/training-optimization.md。
+# 代价：batch 的交付顺序不再逐次可复现（每个 index 每轮仍恰好出现一次）。
+strict_batch_order: false

configs/_examples/pi05_base_bi_flexiv_assemble_box_with_phone_stand_lora_0422_merged_fixed_h100.yaml → configs/_examples/pi05_base_bi_flexiv_optical_module_0611_h100.yaml

@@ -10,12 +10,12 @@ weight_loader:
 ema_decay: null
 data:
   type: LeRobotBiFlexivDataConfig
-  repo_id: Xense/assemble_box_with_phone_stand0410_merged_fixed
+  repo_id: Xense/optical-module-insertion-0731
   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.
+  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: 80000
+num_train_steps: 60000
 fsdp_devices: 8

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

configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_0824_h100.yaml → configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml

@@ -23,10 +23,12 @@ 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
-# 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

configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0902_h100.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