diff --git a/README.md b/README.md index bec1e46..414bb36 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,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/ +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/ +``` + +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: `; +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//...`). 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%). diff --git a/configs/_examples/_FULL_REFERENCE.yaml b/configs/_examples/_FULL_REFERENCE.yaml index 6519dc1..5763168 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: "..." @@ -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 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..58fb20d --- /dev/null +++ b/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= +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 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_0520_a100.yaml deleted file mode 100644 index 4d3abc7..0000000 --- a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0520_a100.yaml +++ /dev/null @@ -1,28 +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 --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 - base_config: - prompt_from_task: true - 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 -fsdp_devices: 8 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..f977abb --- /dev/null +++ b/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= +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 diff --git a/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h200.yaml b/configs/_examples/pi05_base_bi_flexiv_earbuds_case_insertion_teleop_rtc_0904_h200.yaml new file mode 100644 index 0000000..b787476 --- /dev/null +++ b/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 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_optical_module_0611_h100.yaml similarity index 61% rename from configs/_examples/pi05_base_bi_flexiv_assemble_box_with_phone_stand_lora_0422_merged_fixed_h100.yaml rename to configs/_examples/pi05_base_bi_flexiv_optical_module_0611_h100.yaml index e39b6e6..cfe9b36 100644 --- 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_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 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_0824_h100.yaml b/configs/_examples/pi05_base_bi_flexiv_unscrew_bottle_cap_rtc_0824_h100.yaml similarity index 92% 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..b3951c5 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,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 diff --git a/configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0902_h100.yaml b/configs/_examples/pi05_base_dobot_nova5_loreal_returns_sorting_0902_h100.yaml new file mode 100644 index 0000000..4c511cc --- /dev/null +++ b/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 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/ - 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..bef77b4 --- /dev/null +++ b/docs/training-optimization.md @@ -0,0 +1,105 @@ +# 训练优化与 cuDNN attention + +本文保留训练配置、实现约束和历史验证结论。详细实验脚本及排障过程可从本分支 +`72ace2a` 及更早的提交中查阅；下述性能数字来自历史实验，并非每台机器上的保证。 + +## 生产配置 + +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: + 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 +strict_batch_order: false +``` + +`max_token_len: 128` 只针对上述耳机盒数据集做过验证：历史采样的 4037 个状态 +（含各维极值样本）最长为 116 tokens。更换数据集、prompt 或状态维度后需重新检查 +完整 transform 链的 token 长度；tokenizer 超长会警告并截断。不要直接套用 128。 +训练前也需重新计算 norm stats，并检查 `lr_schedule.decay_steps` 是否符合训练计划。 + +## 保留的优化 + +- 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，复用前向残差。 + +main 已包含的跳过未使用触觉视频、乱序 batch 交付和无限 sampler 继续保留。 +`strict_batch_order: false` 可避免队头阻塞，但交付顺序不保证跨运行一致；对比训练数值时设为 `true`。 +Gemma 与 SigLIP 保持默认的全量 rematerialization；历史 8×H100、batch 256 实验中， +关闭 remat 或保存更多残差均超出显存容量。 + +## cuDNN 数值约束 + +全层 BF16 cuDNN attention 曾在约 1000 步后偏离显式基线并发散，不能依据短跑有限就用于生产。 +当前显式路径仍是默认回退方案：`use_cudnn_attention: false`。 +启用 cuDNN 时使用 `cudnn_attention_dtype: float16`。 + +排障中发现两个独立问题： + +- 全空 attention mask 行可能产生 NaN query 梯度。实现将这些行的 query 梯度置零，保留原 mask。 +- peaked attention 行的 backward 对存储输出的舍入误差敏感。FP16 提供更多尾数位； + 反向将余切按 2 的幂缩放后送入内核，再以 FP32 还原梯度，缓解 FP16 指数范围限制。 + +当前 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 也继续使用显式实现。 + +## 历史验证与性能 + +2026-09-04，8×H100、全局 batch 256、固定 seed 和严格 batch 顺序，从 pi05_base step 0 +起跑的 FP16 VJP 初版通过 3000 步验收。历史 BF16 分叉点附近，FP16 loss 与显式基线差小于 +0.5%，记录的梯度范数未超过比较基线同点的 1.3 倍。 + +当前直接调用 forward/backward rule 的版本另做过真实激活等价性检查与 400 步对照： +前向、dK/dV 逐位相同，dQ 差异与内核自身跨运行的原子累加差异同阶，loss 与初版逐点一致。 +这与初版的 3000 步训练验证是不同的验证范围。 + +同日 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 的提速比例直接相乘。 + +新的 attention 后端或 VJP 改动，应从 step 0 固定数据、seed、batch 顺序及超参跑满 +3000 次更新，与显式基线对比 loss 和梯度趋势；确认参数与 optimizer state 有限， +梯度范数不超过基线同点的 3 倍后，再进入有 checkpoint 的受监控长跑。 +单次前后向、吞吐短跑或单步梯度距离不能替代收敛验证。 + +## 环境与启动 + +保持 pip CUDA 12.8 库来源一致，避免将 `$CONDA_PREFIX/lib` 加到 `LD_LIBRARY_PATH` 前面。 +当前依赖配置使用 PyTorch 2.11 对应的 cuDNN 9.19；训练启动时记录 JAX cuDNN runtime 版本。 +运行 `check_cuda_stack.py` 检查实际加载的库来源,以及生产 shape 下 BF16 内核与 FP16 custom VJP +的前后向(含全空 mask 行的 dQ)。该检查只覆盖单次前后向数值,不替代 FP16 收敛验证。 + +```bash +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 --exp-name= +``` + +这些 XLA 标志作为整体做过验证，未逐项测量贡献。继续已有实验使用 `--resume`； +仅在确实需要覆盖目标实验目录时使用 `--overwrite`。 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 diff --git a/pyproject.toml b/pyproject.toml index 77091ee..a104f90 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/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", "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..134068b --- /dev/null +++ b/scripts/check_cuda_stack.py @@ -0,0 +1,263 @@ +"""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 = [] +# 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", + "torchvision": "0.26", + "torchcodec": "0.11", + "nvidia-cudnn-cu12": "9.19", + "cuda": "12.8", +} +# 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): + 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 + from jax._src.lib import cuda_versions + import jax.numpy as jnp + import numpy as np + import torch + 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") + 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, + } + for name, expected in EXPECTED.items(): + 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}") + + # 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. + batch_size, num_heads, num_kv_heads, head_dim = 2, 8, 1, 256 + seq_len = 768 + 200 + 50 + + 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((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 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(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(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 + ) + + 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: + 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("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(" 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.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/train.py b/scripts/train.py index 09d48e3..77c9758 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -10,6 +10,7 @@ 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 import numpy as np @@ -134,12 +135,19 @@ 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, 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 +183,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 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) + 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 @@ -195,6 +213,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( @@ -258,6 +278,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) @@ -277,7 +298,12 @@ def main(config: _config.TrainConfig): 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 infos.append(info) @@ -286,9 +312,22 @@ 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)) + # 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. + # 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 @@ -323,7 +362,8 @@ 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" ) - + 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/models/gemma.py b/src/openpi/models/gemma.py index 023b497..d91ce24 100644 --- a/src/openpi/models/gemma.py +++ b/src/openpi/models/gemma.py @@ -27,6 +27,7 @@ from collections.abc import Sequence import dataclasses +import functools from typing import Literal, TypeAlias import einops @@ -160,11 +161,184 @@ 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. + + 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 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)) + + +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") + + +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. 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 + # 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): + return attention_fwd(q, k, v, attn_mask)[0] + + 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) + ) + 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)) + 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) + 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): + 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) + + +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) + 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.""" 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" @nn.compact def __call__(self, xs, positions, attn_mask, kv_cache): @@ -219,22 +393,42 @@ 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) - - 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") + 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. + # Fully masked query rows make cuDNN produce NaN q-gradients; drop + # those rows from the backward pass instead of rebuilding the mask. + 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 + 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(mask[:, :, None, :, :], logits, big_neg) + + probs = jax.nn.softmax(masked_logits, axis=-1).astype(v.dtype) + + encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v).astype(dtype) + return einops.rearrange(encoded, "B T K G H -> B T (K G) H") + + operands = (q, k, v, attn_mask) + if self.use_cudnn_attention and kv_cache is None: + encoded = cudnn_attention(operands) + else: + encoded = explicit_attention(operands) out = [] start = 0 @@ -291,6 +485,8 @@ class Block(nn.Module): """Transformer block.""" configs: tuple[Config, ...] + use_cudnn_attention: bool = False + cudnn_attention_dtype: str = "bfloat16" dropout: float = 0.0 dropout_bdims: tuple[int, ...] = () @@ -300,8 +496,12 @@ 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, + cudnn_attention_dtype=self.cudnn_attention_dtype, + name="attn", + ) pre_attn = [] gates = [] for i, x in enumerate(xs): @@ -356,6 +556,8 @@ 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 + cudnn_attention_dtype: str = "bfloat16" def setup(self): # all experts must have the same depth @@ -386,6 +588,8 @@ def setup(self): length=self.configs[0].depth, )( configs=self.configs, + use_cudnn_attention=self.use_cudnn_attention, + cudnn_attention_dtype=self.cudnn_attention_dtype, dropout=self.dropout, dropout_bdims=self.dropout_bdims, ) diff --git a/src/openpi/models/pi0.py b/src/openpi/models/pi0.py index b57c487..d9f74c4 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. @@ -83,6 +94,8 @@ def __init__( configs=[paligemma_config, action_expert_config], embed_dtype=config.dtype, adarms=config.pi05, + use_cudnn_attention=config.use_cudnn_attention, + cudnn_attention_dtype=config.cudnn_attention_dtype, ) ) llm.lazy_init( @@ -116,6 +129,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 +139,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..a5ba5a9 100644 --- a/src/openpi/models/pi0_config.py +++ b/src/openpi/models/pi0_config.py @@ -30,7 +30,15 @@ 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 + # 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 + # 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" # training-time RTC config enable_training_time_rtc: bool = False max_delay: int = 10 # steps 330ms @ 30fps @@ -40,6 +48,8 @@ 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_dtype not in ("bfloat16", "float16"): + raise ValueError("cudnn_attention_dtype must be 'bfloat16' or 'float16'") @property @override diff --git a/src/openpi/models/pi0_test.py b/src/openpi/models/pi0_test.py index d3d70e8..ef4fbb0 100644 --- a/src/openpi/models/pi0_test.py +++ b/src/openpi/models/pi0_test.py @@ -7,6 +7,58 @@ 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 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( + [ + [ + [ + [True, False, False, False], + [True, True, False, False], + [False, False, False, False], + [False, False, False, False], + ] + ] + ], + dtype=jnp.bool_, + ) + + q = jnp.arange(16, dtype=jnp.float32).reshape(1, 4, 2, 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) + + # 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:])) + + 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/policies/dobot_nova5_policy.py b/src/openpi/policies/dobot_nova5_policy.py new file mode 100644 index 0000000..6683790 --- /dev/null +++ b/src/openpi/policies/dobot_nova5_policy.py @@ -0,0 +1,106 @@ +"""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/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 diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py index d580cc5..9f5b059 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/data_loader.py b/src/openpi/training/data_loader.py index 4bd5ef2..0bcb577 100644 --- a/src/openpi/training/data_loader.py +++ b/src/openpi/training/data_loader.py @@ -89,6 +89,9 @@ def data_config(self) -> _config.DataConfig: def __iter__(self) -> Iterator[T_co]: raise NotImplementedError("Subclasses of DataLoader should implement __iter__.") + def close(self) -> None: + """Release worker processes and other loader resources.""" + class TransformedDataset(Dataset[T_co]): def __init__(self, dataset: Dataset, transforms: Sequence[_transforms.DataTransformFn]): @@ -584,6 +587,7 @@ def __init__( jax.sharding.PartitionSpec("B"), ) self._num_batches = num_batches + self._active_iterator = None mp_context = None if num_workers > 0: @@ -625,38 +629,86 @@ 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) + 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 + 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: + 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) + yield batch + finally: + self.close() + + def close(self) -> None: + """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 + + 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): - """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) + """Collate batch elements into CPU tensors so workers use shared-memory IPC.""" + # 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. + return jax.tree.map( + lambda *xs: torch.from_numpy(np.stack([np.asarray(x) for x in xs], axis=0)), + *items, + ) def _worker_init_fn(worker_id: int) -> None: @@ -738,3 +790,8 @@ def data_config(self) -> _config.DataConfig: def __iter__(self): for batch in self._data_loader: yield _model.Observation.from_dict(batch), batch["actions"] + + 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..8f6e475 100644 --- a/src/openpi/training/data_loader_test.py +++ b/src/openpi/training/data_loader_test.py @@ -54,6 +54,52 @@ 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_with_fake_dataset(): config = _config.get_config("debug_pi05") 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.