From 72834e7de5f6d8ed95f4e7addd8a8696c4fb28dd Mon Sep 17 00:00:00 2001 From: Starrick Liu <73152103+StarrickLiu@users.noreply.github.com> Date: Mon, 15 Jun 2026 11:40:00 +0800 Subject: [PATCH] Update Wall-X to 1.1.0 (#104) --- .gitattributes | 1 - .github/ISSUE_TEMPLATE/bug_report.md | 35 +- .github/workflows/lint.yml | 32 +- .gitignore | 289 +- .gitmodules | 3 - .pre-commit-config.yaml | 14 +- 3rdparty/cutlass | 1 - CONTRIBUTING.md | 29 +- LICENSE | 202 ++ MANIFEST.in | 6 + README.md | 142 +- README.zh-CN.md | 152 ++ assets/cot_example_frame.png | Bin 136402 -> 0 bytes csrc/README.md | 28 - csrc/dual_asym_grouped_gemm.cu | 366 --- csrc/dual_asym_grouped_gemm.h | 10 - csrc/ops.cu | 21 - csrc/permute.cu | 942 ------- csrc/permute.h | 31 - csrc/rope.cu | 571 ---- csrc/rope.h | 14 - csrc/rope_index.h | 13 - csrc/rot_pos.cu | 332 --- csrc/rot_pos.h | 6 - csrc/window_index.h | 9 - pyproject.toml | 26 + requirements-libero.txt | 16 + requirements.txt | 23 +- scripts/README.md | 142 + scripts/compute_norm_stats.py | 838 +++++- scripts/draw_openloop_plot.py | 1215 ++++++++- scripts/fake_inference.py | 179 +- scripts/infer_libero.py | 449 +-- scripts/infer_robochallenge.py | 1164 -------- scripts/merge_sharded_weights.py | 72 +- scripts/merge_tokenizer.py | 121 +- scripts/normalize.py | 161 -- scripts/run_libero.sh | 273 ++ scripts/run_serving.sh | 285 ++ scripts/vqa_inference.py | 104 - setup.py | 107 +- tests/test_ops_fallback.py | 32 + train_qact.py | 142 - wall_x/_vendor/__init__.py | 4 + wall_x/_vendor/harrix/__init__.py | 5 + wall_x/_vendor/harrix/adapters/__init__.py | 6 + wall_x/_vendor/harrix/adapters/base.py | 51 + wall_x/_vendor/harrix/adapters/qwen_vlact.py | 465 ++++ wall_x/_vendor/harrix/adapters/registry.py | 39 + .../harrix/adapters/variants/__init__.py | 15 + .../harrix/adapters/variants/qwen2_5.py | 17 + .../harrix/drivers}/__init__.py | 0 .../_vendor/harrix/drivers/inproc/__init__.py | 197 ++ .../harrix/drivers/inproc/env_handle.py | 152 ++ .../harrix/drivers/inproc/model_handle.py | 25 + wall_x/_vendor/harrix/drivers/job_state.py | 343 +++ wall_x/_vendor/harrix/envs/__init__.py | 8 + wall_x/_vendor/harrix/envs/base.py | 119 + wall_x/_vendor/harrix/envs/libero.py | 342 +++ wall_x/_vendor/harrix/envs/libero_common.py | 359 +++ wall_x/_vendor/harrix/envs/libero_sim.py | 181 ++ wall_x/_vendor/harrix/envs/registry.py | 53 + wall_x/_vendor/harrix/eval_config.py | 217 ++ .../{ => _vendor/harrix}/serving/__init__.py | 0 .../harrix/serving/_wallx_infer/__init__.py | 1 + .../serving/_wallx_infer}/base_dataclass.py | 192 +- .../serving/_wallx_infer/infer_config.py | 434 +++ .../harrix/serving/_wallx_infer}/logger.py | 94 +- .../serving/_wallx_infer/model_wrapper.py | 760 ++++++ .../harrix/serving/_wallx_infer/robot.py | 1181 ++++++++ .../serving/_wallx_infer/socket_controller.py | 423 +++ .../harrix/serving/_wallx_infer}/utils.py | 81 +- .../_vendor/harrix/serving/launch_serving.py | 280 ++ .../harrix}/serving/policy/__init__.py | 2 + .../harrix/serving/policy/_smoothing.py | 104 + .../harrix/serving/policy/wall_x_policy.py | 328 +++ wall_x/_vendor/harrix/serving/scheduler.py | 160 ++ .../serving/websocket_policy_server.py | 86 +- wall_x/_vendor/harrix/utils/__init__.py | 0 wall_x/_vendor/harrix/utils/ckpt_load.py | 171 ++ wall_x/_vendor/harrix/utils/normalizer.py | 135 + wall_x/_vendor/harrix/utils/seed.py | 23 + wall_x/_vendor/harrix/utils/train_config.py | 371 +++ wall_x/_vendor/x2robot_utils/__init__.py | 5 + wall_x/_vendor/x2robot_utils/geometry.py | 309 +++ wall_x/_vendor/x2robot_utils/grounding.py | 172 ++ .../_vendor/x2robot_utils/text_templates.py | 235 ++ wall_x/config/__init__.py | 62 + wall_x/config/data_config.py | 63 + wall_x/config/hyperparams_config.py | 129 + wall_x/config/infra_config.py | 98 + wall_x/config/loader.py | 285 ++ wall_x/config/model_config.py | 42 + wall_x/config/registry.py | 99 + wall_x/config/task_config.py | 17 + wall_x/config/train_config.py | 81 + wall_x/data/__init__.py | 31 + wall_x/data/_bundle.py | 39 + wall_x/data/_protocol.py | 37 + wall_x/data/_registry.py | 263 ++ wall_x/data/backends/__init__.py | 16 + wall_x/data/backends/lerobot/__init__.py | 26 + wall_x/data/backends/lerobot/build.py | 337 +++ wall_x/data/backends/lerobot/config.py | 134 + wall_x/data/backends/lerobot/loader.py | 895 ++++++ .../data/backends/lerobot/rotation_layout.py | 95 + wall_x/data/{ => backends/lerobot}/utils.py | 1101 ++++---- wall_x/data/config.py | 126 - wall_x/data/load_lerobot_dataset.py | 715 ----- wall_x/fusions/backend.py | 433 --- wall_x/fusions/ops.py | 742 ----- wall_x/infer/env.py | 143 - wall_x/infer/env_libero.py | 744 ----- wall_x/infer/infer_config.py | 587 ---- wall_x/infer/utils_libero.py | 323 --- wall_x/model/action_head.py | 808 ------ wall_x/model/core/__init__.py | 1 + wall_x/model/core/action/__init__.py | 10 + wall_x/model/core/action/head.py | 27 + wall_x/model/core/action/moe.py | 126 + wall_x/model/core/action/normalizer.py | 447 +++ wall_x/model/core/action/processor.py | 317 +++ wall_x/model/core/attention/__init__.py | 9 + .../attention/joint.py} | 507 +++- wall_x/model/core/attention/mask.py | 259 ++ wall_x/model/core/attention/selector.py | 136 + wall_x/model/core/ops/__init__.py | 17 + wall_x/model/core/ops/_cuda_ext.py | 35 + wall_x/model/core/ops/_cuda_wrappers.py | 750 ++++++ wall_x/model/core/ops/base.py | 143 + wall_x/model/core/ops/csrc/binding.cu | 69 + .../core/ops/csrc/common/activation_types.h | 11 + .../core/ops/csrc/common/activations.cuh | 30 + .../model/core/ops/csrc/common/cuda_utils.h | 111 + .../ops/csrc/get_rope_index/get_rope_index.cu | 467 ++-- .../core/ops/csrc/m_rope/m_rope_kernel.cu | 840 ++++++ .../ops/csrc/permute_unpermute/permute.cu | 581 ++++ .../model/core/ops/csrc/rope/rope_kernel.cu | 1098 ++++++++ wall_x/model/core/ops/csrc/rot_pos/rot_pos.cu | 284 ++ .../ops/csrc/window_index}/window_index.cu | 106 +- wall_x/model/core/ops/index.py | 340 +++ wall_x/model/core/ops/moe.py | 97 + wall_x/model/core/ops/norm.py | 42 + wall_x/model/core/ops/rope.py | 234 ++ wall_x/model/core/vla_mixin.py | 746 +++++ wall_x/model/model_utils.py | 319 --- wall_x/model/qact/__init__.py | 1 + wall_x/model/qact/qwen2_5/__init__.py | 2 + wall_x/model/qact/qwen2_5/adapter.py | 56 + .../qwen2_5}/configuration_qwen2_5_vl.py | 80 +- wall_x/model/qact/qwen2_5/inference_mixin.py | 1039 +++++++ .../qwen2_5}/modeling_qwen2_5_vl.py | 628 ++++- .../qwen2_5}/modeling_qwen2_5_vl_act.py | 2399 ++++++++++------- wall_x/model/qact/tokenizer_mixin.py | 962 +++++++ wall_x/model/qwen2_5_based/__init__.py | 8 - wall_x/model/registry.py | 64 + wall_x/model/vla_mixin.py | 987 ------- wall_x/serving/README.md | 263 -- wall_x/serving/client.py | 397 --- wall_x/serving/launch_serving.py | 220 -- wall_x/serving/policy/utils.py | 265 -- wall_x/serving/policy/wall_x_policy.py | 182 -- wall_x/trainer/adapters/__init__.py | 128 + wall_x/trainer/adapters/base_adapter.py | 494 ++++ wall_x/trainer/adapters/vla_model_adapter.py | 990 +++++++ wall_x/trainer/fsdp_trainer/__init__.py | 34 + wall_x/trainer/fsdp_trainer/base_trainer.py | 430 +++ wall_x/trainer/fsdp_trainer/checkpoint_io.py | 1033 +++++++ .../fsdp_trainer/distribution_strategy.py | 400 +++ wall_x/trainer/fsdp_trainer/fsdp_trainer.py | 963 +++++++ .../trainer/fsdp_trainer/metrics/__init__.py | 3 + wall_x/trainer/fsdp_trainer/metrics/norms.py | 123 + wall_x/trainer/fsdp_trainer/metrics_logger.py | 111 + wall_x/trainer/fsdp_trainer/train_fsdp.py | 181 ++ wall_x/trainer/optimizer/__init__.py | 9 + wall_x/trainer/optimizer/dmuon/__init__.py | 16 + wall_x/trainer/optimizer/dmuon/utils.py | 156 ++ wall_x/trainer/optimizer/utils.py | 230 ++ wall_x/trainer/qwen_vl_act_trainer.py | 1144 -------- wall_x/trainer/scheduler/scheduler.py | 71 + wall_x/trainer/trainer_utils.py | 574 ++++ wall_x/trainer/utils/__init__.py | 6 + wall_x/trainer/utils/data.py | 46 + wall_x/trainer/utils/diagnostics.py | 40 + wall_x/utils/constant.py | 420 +-- wall_x/utils/cudagraph_wrapper.py | 278 ++ wall_x/utils/logger.py | 95 + wall_x/utils/metrics.py | 95 + wall_x/utils/timers.py | 217 +- workspace/README.md | 479 +++- .../lerobot/qwen2_5_lerobot_template.yml | 93 + workspace/example/libero.yml | 125 + workspace/example/maniparena_example.yml | 134 + workspace/lerobot_example/config_qact.yml | 151 -- .../lerobot_example/config_qact_from_vlm.yml | 152 -- .../evaluation/lerobot_openloop.png | 3 - .../libero/config_qact_libero_from_vlm.yml | 100 - workspace/lerobot_example/qwen25_config.json | 106 - workspace/lerobot_example/run.sh | 23 - .../models_config/qwen2_5_moe_flash.json | 108 + 200 files changed, 33916 insertions(+), 16771 deletions(-) delete mode 100644 .gitattributes delete mode 100644 .gitmodules delete mode 160000 3rdparty/cutlass create mode 100644 LICENSE create mode 100644 MANIFEST.in create mode 100644 README.zh-CN.md delete mode 100644 assets/cot_example_frame.png delete mode 100644 csrc/README.md delete mode 100644 csrc/dual_asym_grouped_gemm.cu delete mode 100644 csrc/dual_asym_grouped_gemm.h delete mode 100644 csrc/ops.cu delete mode 100644 csrc/permute.cu delete mode 100644 csrc/permute.h delete mode 100644 csrc/rope.cu delete mode 100644 csrc/rope.h delete mode 100644 csrc/rope_index.h delete mode 100644 csrc/rot_pos.cu delete mode 100644 csrc/rot_pos.h delete mode 100644 csrc/window_index.h create mode 100644 requirements-libero.txt create mode 100644 scripts/README.md mode change 100644 => 100755 scripts/fake_inference.py mode change 100644 => 100755 scripts/infer_libero.py delete mode 100644 scripts/infer_robochallenge.py delete mode 100644 scripts/normalize.py create mode 100755 scripts/run_libero.sh create mode 100755 scripts/run_serving.sh delete mode 100644 scripts/vqa_inference.py create mode 100644 tests/test_ops_fallback.py delete mode 100644 train_qact.py create mode 100644 wall_x/_vendor/__init__.py create mode 100644 wall_x/_vendor/harrix/__init__.py create mode 100644 wall_x/_vendor/harrix/adapters/__init__.py create mode 100644 wall_x/_vendor/harrix/adapters/base.py create mode 100644 wall_x/_vendor/harrix/adapters/qwen_vlact.py create mode 100644 wall_x/_vendor/harrix/adapters/registry.py create mode 100644 wall_x/_vendor/harrix/adapters/variants/__init__.py create mode 100644 wall_x/_vendor/harrix/adapters/variants/qwen2_5.py rename wall_x/{fusions => _vendor/harrix/drivers}/__init__.py (100%) create mode 100644 wall_x/_vendor/harrix/drivers/inproc/__init__.py create mode 100644 wall_x/_vendor/harrix/drivers/inproc/env_handle.py create mode 100644 wall_x/_vendor/harrix/drivers/inproc/model_handle.py create mode 100644 wall_x/_vendor/harrix/drivers/job_state.py create mode 100644 wall_x/_vendor/harrix/envs/__init__.py create mode 100644 wall_x/_vendor/harrix/envs/base.py create mode 100644 wall_x/_vendor/harrix/envs/libero.py create mode 100644 wall_x/_vendor/harrix/envs/libero_common.py create mode 100644 wall_x/_vendor/harrix/envs/libero_sim.py create mode 100644 wall_x/_vendor/harrix/envs/registry.py create mode 100644 wall_x/_vendor/harrix/eval_config.py rename wall_x/{ => _vendor/harrix}/serving/__init__.py (100%) create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py rename wall_x/{infer => _vendor/harrix/serving/_wallx_infer}/base_dataclass.py (62%) create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py rename wall_x/{infer => _vendor/harrix/serving/_wallx_infer}/logger.py (74%) create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/robot.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py rename wall_x/{infer => _vendor/harrix/serving/_wallx_infer}/utils.py (75%) create mode 100644 wall_x/_vendor/harrix/serving/launch_serving.py rename wall_x/{ => _vendor/harrix}/serving/policy/__init__.py (56%) create mode 100644 wall_x/_vendor/harrix/serving/policy/_smoothing.py create mode 100644 wall_x/_vendor/harrix/serving/policy/wall_x_policy.py create mode 100644 wall_x/_vendor/harrix/serving/scheduler.py rename wall_x/{ => _vendor/harrix}/serving/websocket_policy_server.py (56%) create mode 100644 wall_x/_vendor/harrix/utils/__init__.py create mode 100644 wall_x/_vendor/harrix/utils/ckpt_load.py create mode 100644 wall_x/_vendor/harrix/utils/normalizer.py create mode 100644 wall_x/_vendor/harrix/utils/seed.py create mode 100644 wall_x/_vendor/harrix/utils/train_config.py create mode 100644 wall_x/_vendor/x2robot_utils/__init__.py create mode 100644 wall_x/_vendor/x2robot_utils/geometry.py create mode 100644 wall_x/_vendor/x2robot_utils/grounding.py create mode 100644 wall_x/_vendor/x2robot_utils/text_templates.py create mode 100644 wall_x/config/__init__.py create mode 100644 wall_x/config/data_config.py create mode 100644 wall_x/config/hyperparams_config.py create mode 100644 wall_x/config/infra_config.py create mode 100644 wall_x/config/loader.py create mode 100644 wall_x/config/model_config.py create mode 100644 wall_x/config/registry.py create mode 100644 wall_x/config/task_config.py create mode 100644 wall_x/config/train_config.py create mode 100644 wall_x/data/_bundle.py create mode 100644 wall_x/data/_protocol.py create mode 100644 wall_x/data/_registry.py create mode 100644 wall_x/data/backends/__init__.py create mode 100644 wall_x/data/backends/lerobot/__init__.py create mode 100644 wall_x/data/backends/lerobot/build.py create mode 100644 wall_x/data/backends/lerobot/config.py create mode 100644 wall_x/data/backends/lerobot/loader.py create mode 100644 wall_x/data/backends/lerobot/rotation_layout.py rename wall_x/data/{ => backends/lerobot}/utils.py (72%) delete mode 100644 wall_x/data/config.py delete mode 100644 wall_x/data/load_lerobot_dataset.py delete mode 100644 wall_x/fusions/backend.py delete mode 100644 wall_x/fusions/ops.py delete mode 100644 wall_x/infer/env.py delete mode 100644 wall_x/infer/env_libero.py delete mode 100644 wall_x/infer/infer_config.py delete mode 100644 wall_x/infer/utils_libero.py delete mode 100644 wall_x/model/action_head.py create mode 100644 wall_x/model/core/__init__.py create mode 100644 wall_x/model/core/action/__init__.py create mode 100644 wall_x/model/core/action/head.py create mode 100644 wall_x/model/core/action/moe.py create mode 100644 wall_x/model/core/action/normalizer.py create mode 100644 wall_x/model/core/action/processor.py create mode 100644 wall_x/model/core/attention/__init__.py rename wall_x/model/{joint_attention.py => core/attention/joint.py} (56%) create mode 100644 wall_x/model/core/attention/mask.py create mode 100644 wall_x/model/core/attention/selector.py create mode 100644 wall_x/model/core/ops/__init__.py create mode 100644 wall_x/model/core/ops/_cuda_ext.py create mode 100644 wall_x/model/core/ops/_cuda_wrappers.py create mode 100644 wall_x/model/core/ops/base.py create mode 100644 wall_x/model/core/ops/csrc/binding.cu create mode 100644 wall_x/model/core/ops/csrc/common/activation_types.h create mode 100644 wall_x/model/core/ops/csrc/common/activations.cuh create mode 100644 wall_x/model/core/ops/csrc/common/cuda_utils.h rename csrc/rope_index.cu => wall_x/model/core/ops/csrc/get_rope_index/get_rope_index.cu (53%) create mode 100644 wall_x/model/core/ops/csrc/m_rope/m_rope_kernel.cu create mode 100644 wall_x/model/core/ops/csrc/permute_unpermute/permute.cu create mode 100644 wall_x/model/core/ops/csrc/rope/rope_kernel.cu create mode 100644 wall_x/model/core/ops/csrc/rot_pos/rot_pos.cu rename {csrc => wall_x/model/core/ops/csrc/window_index}/window_index.cu (77%) create mode 100644 wall_x/model/core/ops/index.py create mode 100644 wall_x/model/core/ops/moe.py create mode 100644 wall_x/model/core/ops/norm.py create mode 100644 wall_x/model/core/ops/rope.py create mode 100644 wall_x/model/core/vla_mixin.py delete mode 100644 wall_x/model/model_utils.py create mode 100644 wall_x/model/qact/__init__.py create mode 100644 wall_x/model/qact/qwen2_5/__init__.py create mode 100644 wall_x/model/qact/qwen2_5/adapter.py rename wall_x/model/{qwen2_5_based => qact/qwen2_5}/configuration_qwen2_5_vl.py (79%) create mode 100644 wall_x/model/qact/qwen2_5/inference_mixin.py rename wall_x/model/{qwen2_5_based => qact/qwen2_5}/modeling_qwen2_5_vl.py (82%) rename wall_x/model/{qwen2_5_based => qact/qwen2_5}/modeling_qwen2_5_vl_act.py (60%) create mode 100644 wall_x/model/qact/tokenizer_mixin.py delete mode 100644 wall_x/model/qwen2_5_based/__init__.py create mode 100644 wall_x/model/registry.py delete mode 100644 wall_x/model/vla_mixin.py delete mode 100644 wall_x/serving/README.md delete mode 100644 wall_x/serving/client.py delete mode 100644 wall_x/serving/launch_serving.py delete mode 100644 wall_x/serving/policy/utils.py delete mode 100644 wall_x/serving/policy/wall_x_policy.py create mode 100644 wall_x/trainer/adapters/__init__.py create mode 100644 wall_x/trainer/adapters/base_adapter.py create mode 100644 wall_x/trainer/adapters/vla_model_adapter.py create mode 100644 wall_x/trainer/fsdp_trainer/__init__.py create mode 100644 wall_x/trainer/fsdp_trainer/base_trainer.py create mode 100644 wall_x/trainer/fsdp_trainer/checkpoint_io.py create mode 100644 wall_x/trainer/fsdp_trainer/distribution_strategy.py create mode 100644 wall_x/trainer/fsdp_trainer/fsdp_trainer.py create mode 100644 wall_x/trainer/fsdp_trainer/metrics/__init__.py create mode 100644 wall_x/trainer/fsdp_trainer/metrics/norms.py create mode 100644 wall_x/trainer/fsdp_trainer/metrics_logger.py create mode 100644 wall_x/trainer/fsdp_trainer/train_fsdp.py create mode 100644 wall_x/trainer/optimizer/__init__.py create mode 100644 wall_x/trainer/optimizer/dmuon/__init__.py create mode 100644 wall_x/trainer/optimizer/dmuon/utils.py create mode 100644 wall_x/trainer/optimizer/utils.py delete mode 100644 wall_x/trainer/qwen_vl_act_trainer.py create mode 100644 wall_x/trainer/scheduler/scheduler.py create mode 100644 wall_x/trainer/trainer_utils.py create mode 100644 wall_x/trainer/utils/__init__.py create mode 100644 wall_x/trainer/utils/data.py create mode 100644 wall_x/trainer/utils/diagnostics.py create mode 100644 wall_x/utils/cudagraph_wrapper.py create mode 100644 wall_x/utils/logger.py create mode 100644 wall_x/utils/metrics.py create mode 100644 workspace/example/lerobot/qwen2_5_lerobot_template.yml create mode 100644 workspace/example/libero.yml create mode 100644 workspace/example/maniparena_example.yml delete mode 100644 workspace/lerobot_example/config_qact.yml delete mode 100644 workspace/lerobot_example/config_qact_from_vlm.yml delete mode 100644 workspace/lerobot_example/evaluation/lerobot_openloop.png delete mode 100644 workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml delete mode 100644 workspace/lerobot_example/qwen25_config.json delete mode 100644 workspace/lerobot_example/run.sh create mode 100644 workspace/models_config/qwen2_5_moe_flash.json diff --git a/.gitattributes b/.gitattributes deleted file mode 100644 index f8f1794..0000000 --- a/.gitattributes +++ /dev/null @@ -1 +0,0 @@ -workspace/lerobot_example/evaluation/lerobot_openloop.png filter=lfs diff=lfs merge=lfs -text diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md index 865f451..75f5e5a 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.md +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -1,30 +1,27 @@ --- -name: Bug report -about: Create a report to help us improve -title: '' -labels: '' -assignees: '' - ---- - ---- -name: 'Bug Report (English)' -about: Report a bug encountered while using or reproducing the Wall-X model -title: '[Bug] ' -labels: 'bug, needs-triage' +name: Bug Report +about: Report a bug encountered while using Wall-X +title: "[Bug] " +labels: bug, needs-triage +assignees: "" --- **Describe the bug** A clear and concise description of what the bug is. **To Reproduce** -Steps to reproduce the behavior: -1. Go to '...' -2. Run command '....' -3. See error +1. Run command `...` +2. Use config `...` +3. See error `...` **Expected behavior** A clear and concise description of what you expected to happen. -**Logs & Screenshots** -If applicable, add the complete error message (traceback) and screenshots to help explain your problem. +**Environment** +- OS: +- Python: +- PyTorch: +- CUDA: + +**Logs** +Paste the relevant error message or traceback. diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index ce67e84..bfe308a 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -2,25 +2,25 @@ name: Pre-commit on: push: - branches: [ main, master, develop ] + branches: [main, master, develop] pull_request: - branches: [ main, master, develop ] + branches: [main, master, develop] jobs: pre-commit: runs-on: ubuntu-latest - + steps: - - name: Checkout code - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: '3.11' - - - name: Install pre-commit - run: pip install pre-commit - - - name: Run pre-commit - run: pre-commit run --all-files + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + + - name: Install pre-commit + run: pip install pre-commit + + - name: Run pre-commit + run: pre-commit run --all-files diff --git a/.gitignore b/.gitignore index 9c8f49b..b243f82 100644 --- a/.gitignore +++ b/.gitignore @@ -1,270 +1,39 @@ -# ============================================================================= -# OPERATING SYSTEM FILES -# ============================================================================= - -# macOS -.DS_Store -.DS_Store? -._* -.Spotlight-V100 -.Trashes -ehthumbs.db -Thumbs.db - -# Windows -Thumbs.db -ehthumbs.db -Desktop.ini -$RECYCLE.BIN/ -*.cab -*.msi -*.msix -*.msm -*.msp - -# Linux -*~ - -# ============================================================================= -# PYTHON -# ============================================================================= - -# Byte-compiled / optimized / DLL files +# Runtime artifacts __pycache__/ *.py[cod] -*$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -pip-wheel-metadata/ -share/python-wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.nox/ +.ruff_cache/ +.pytest_cache/ +.mypy_cache/ .coverage .coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -*.py,cover -.hypothesis/ -.pytest_cache/ +htmlcov/ -# Translations -*.mo -*.pot +# Local environments and build outputs +.venv/ +venv/ +build/ +dist/ +*.egg-info/ -# Django stuff: -local_settings.py -db.sqlite3 -db.sqlite3-journal - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -target/ - -# Jupyter Notebook -.ipynb_checkpoints -*.ipynb - -# IPython -profile_default/ -ipython_config.py - -# pyenv -.python-version - -# pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -#Pipfile.lock - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow -__pypackages__/ - -# Celery stuff -celerybeat-schedule -celerybeat.pid - -# SageMath parsed files -*.sage.py - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy -.mypy_cache/ -.dmypy.json -dmypy.json - -# Pyre type checker -.pyre/ - -# ============================================================================= -# DEVELOPMENT ENVIRONMENT -# ============================================================================= - -# VS Code -.vscode/ -*.code-workspace - -# PyCharm -.idea/ - -# Vim -*.swp -*.swo -*~ - -# Emacs -*~ -\#*\# -/.emacs.desktop -/.emacs.desktop.lock -*.elc -auto-save-list -tramp -.\#* - -# ============================================================================= -# PROJECT SPECIFIC -# ============================================================================= - -# Model checkpoints and weights -*.pth -*.pt -*.ckpt -ckpt/ -checkpoints/ - -# Output directories -outputs/ -results/ -logs/ -bin/ - -# Media files -*.jpg -*.jpeg -*.gif -*.mp4 -*.avi -*.mov -*.wav -*.mp3 - -# Configuration files with sensitive data -fuse.cfg -ray/auth.json - -# AI/ML related -*.ai - -# ============================================================================= -# LOGGING AND MONITORING -# ============================================================================= - -# General logs -*.log -*.err -*.out - -# Weights & Biases +# Training / inference outputs wandb/ -_wandb/ +logs/ +run_logs/ +outputs/ +videos/ +ckpt/ +*.log -# TensorBoard -runs/ -tensorboard/ +# Large local assets +*.pth +*.safetensors +*.mp4 -# MLflow -mlruns/ +# Editor / OS noise +.DS_Store +.vscode/ -# ============================================================================= -# TEMPORARY AND CACHE FILES -# ============================================================================= - -# Temporary files -*.tmp -*.temp -*.bak -*.backup - -# Cache directories -.cache/ -cache/ -.ruff_cache/ - -# ============================================================================= -# FONTS AND ASSETS (if not part of the project) -# ============================================================================= - -# Font files (uncomment if fonts should not be tracked) -# *.ttf -# *.otf -# *.woff -# *.woff2 - -# ============================================================================= -# PERSONAL/LOCAL FILES -# ============================================================================= - -# Personal notes and documentation -code_update_info.md -TODO.md -NOTES.md - -# Local configuration -.env -.env.local -.env.*.local \ No newline at end of file +# Keep exported CUDA kernel sources and Python wrappers tracked: +# wall_x/model/core/ops/csrc/ +# wall_x/model/core/ops/_cuda_ext.py +# wall_x/model/core/ops/_cuda_wrappers.py diff --git a/.gitmodules b/.gitmodules deleted file mode 100644 index 741dda4..0000000 --- a/.gitmodules +++ /dev/null @@ -1,3 +0,0 @@ -[submodule "3rdparty/cutlass"] - path = 3rdparty/cutlass - url = https://github.com/NVIDIA/cutlass.git diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8a9d4b3..10c8227 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,13 +1,19 @@ -# See https://pre-commit.com for more information -# See https://pre-commit.com/hooks.html for more hooks -exclude: ".git" +exclude: | + (?x)^( + wall_x/_vendor/| + wall_x/model/core/ops/csrc/| + wall_x/model/core/ops/_cuda_ext\.py| + wall_x/model/core/ops/_cuda_wrappers\.py| + wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl\.py| + wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act\.py + ) repos: - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.2.2 hooks: - id: ruff - args: [ --fix, --exit-non-zero-on-fix ] + args: [--fix, --exit-non-zero-on-fix] - repo: https://github.com/psf/black rev: 24.2.0 diff --git a/3rdparty/cutlass b/3rdparty/cutlass deleted file mode 160000 index e51efbf..0000000 --- a/3rdparty/cutlass +++ /dev/null @@ -1 +0,0 @@ -Subproject commit e51efbfe18fe4f4cbb66ab814c55bf4aa0185491 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 26e430c..91419ae 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,17 +1,34 @@ -# Contributing to Wall-x +# Contributing to Wall-X -## Submit a Pull Request +Thanks for your interest in contributing to Wall-X. -Before opening a pull request, please make sure your code passes the lint checks. +## Development Setup + +Install the project dependencies and the package in editable mode: ```bash -# Install pre-commit hooks (run once) +pip install -r requirements.txt +pip install -e . +``` + +Install pre-commit hooks: + +```bash +pip install pre-commit pre-commit install ``` -Or +Run checks before opening a pull request: ```bash -# Manually run all checks pre-commit run --all-files +python -m pytest tests -q ``` + +## Pull Requests + +- Keep changes focused and include reproduction steps for bug fixes. +- Update documentation when changing user-facing behavior or commands. +- Do not include private paths, credentials, internal service URLs, or large + checkpoint artifacts. +- Use placeholder paths such as `` in examples. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..d645695 --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..9e25001 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,6 @@ +include README.md +include README.zh-CN.md +include LICENSE +include requirements.txt +include requirements-libero.txt +recursive-include wall_x/model/core/ops/csrc *.cu *.cuh *.h diff --git a/README.md b/README.md index eff4a16..c29cb56 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,6 @@
- Hugging Face @@ -10,7 +9,6 @@ Project Page -
Python 3.10 PyTorch @@ -22,20 +20,33 @@
## Building General-Purpose Robots Based on Embodied Foundation Model -We are building the embodied foundation model to capture and compress the world's most valuable data: the continuous, high-fidelity stream of physical interaction. -By creating a direct feedback loop between the model's decisions and the body's lived experience, we enable the emergence of a truly generalizable intelligence—one that understands not just how the world works, but how to act effectively within it. +We are building embodied foundation models to capture and compress the world's +most valuable data: continuous, high-fidelity physical interaction. + +By creating a direct feedback loop between model decisions and the body's lived +experience, we enable generalizable intelligence that understands not just how +the world works, but how to act effectively within it. ## Repository -This repository provides the training and inference code that supports our WALL series open-source embodied foundation models. It includes end-to-end pipelines for data preparation (LeRobot), model configuration, flow-matching and FAST action branches, and evaluation utilities for real and simulated robots. + +This repository provides the training and inference code for the WALL series +open-source embodied foundation models. It includes LeRobot data preparation, +model configuration, flow-matching and FAST action branches, public serving and +evaluation utilities, and exported CUDA operator sources that compile during +package installation. ## News -- [May 2026] We introduce [**WALL-WM: Carving World Action Modeling at the Event Joints**](https://x2robot.com/api/files/file/WALL-WM.pdf), a World Action Model that couples future-video imagination with action prediction at their semantic event boundaries, delivering state-of-the-art real-robot manipulation and physically grounded video generation from a single event-pretrained backbone *(Code coming soon!)*. -- [May 2026] We introduce [**Wall-OSS-0.5: A Deployment-Ready VLA with Gradient-Bridged Pretraining**](https://x2robot.com/api/files/file/wall_oss_05.pdf), an open-source 4B model that delivers directly deployable, zero-shot real-robot manipulation capabilities while serving as a powerful prior for downstream adaptation *(Code coming soon!)*. -- [Sept 2025] We introduce [**WALL-OSS: Igniting VLMs toward the Embodied Space**](https://x2robot.com/en/research/68bc2cde8497d7f238dde690), an end-to-end embodied foundation model that leverages large-scale multimodal pretraining to achieve (1) embodiment-aware vision–language understanding, (2) strong language–action association, and (3) robust manipulation capability. +- [June 2026] Wall-X 1.1.0 updates the open-source training and inference + stack for Wall-OSS-0.5, including the public serving/evaluation runtime, + DMuon training support, and install-time CUDA operator builds. +- [May 2026] We introduce [**WALL-WM: Carving World Action Modeling at the Event Joints**](https://x2robot.com/api/files/file/WALL-WM.pdf), a World Action Model that couples future-video imagination with action prediction at semantic event boundaries. +- [May 2026] We introduce [**Wall-OSS-0.5: A Deployment-Ready VLA with Gradient-Bridged Pretraining**](https://x2robot.com/api/files/file/wall_oss_05.pdf), an open-source model for directly deployable real-robot manipulation and downstream adaptation. +- [Sept 2025] We introduce [**WALL-OSS: Igniting VLMs toward the Embodied Space**](https://x2robot.com/en/research/68bc2cde8497d7f238dde690), an end-to-end embodied foundation model that leverages large-scale multimodal pretraining to achieve embodiment-aware vision-language understanding, language-action association, and robust manipulation capability. ## Models + - WALL-OSS-0.5: https://huggingface.co/x-square-robot/wall-oss-0.5 - WALL-OSS-FLOW-0.1: https://huggingface.co/x-square-robot/wall-oss-flow-0.1 - WALL-OSS-FLOW: https://huggingface.co/x-square-robot/wall-oss-flow @@ -43,103 +54,140 @@ This repository provides the training and inference code that supports our WALL ## Environment Setup -Create and activate conda environment: +Create and activate a conda environment: + ```bash conda create --name wallx python=3.10 conda activate wallx ``` Install requirements: + ```bash pip install -r requirements.txt MAX_JOBS=4 pip install flash-attn==2.7.4.post1 --no-build-isolation ``` -Install lerobot: +Install DMuon, which is used by the default training configs: + ```bash -git clone https://github.com/huggingface/lerobot.git -git checkout c66cd401767e60baece16e1cf68da2824227e076 -cd lerobot -pip install -e . +pip install "dmuon @ git+https://github.com/X-Square-Robot/dmuon.git" ``` -Install wall_x: +Install LeRobot: + ```bash -git submodule update --init --recursive -MAX_JOBS=4 pip install --no-build-isolation --verbose -e . +git clone https://github.com/huggingface/lerobot.git +cd lerobot +git checkout c66cd401767e60baece16e1cf68da2824227e076 +pip install --no-deps -e . ``` +Use `--no-deps` for LeRobot so it does not override the Wall-X dependency +versions installed from `requirements.txt`. + +Install Wall-X: + +```bash +MAX_JOBS=8 pip install --no-build-isolation -e . +``` + +Public helper scripts live under `scripts/`; the examples below use the +repository-root form, such as `python scripts/fake_inference.py`. + +The exported CUDA operator sources are included in +`wall_x/model/core/ops/csrc/`. `setup.py` builds them with PyTorch +`CUDAExtension` when Wall-X is installed. `ninja` is included in +`requirements.txt` for parallel builds, and `MAX_JOBS` controls compile +parallelism. `--no-build-isolation` is required so the build can use the torch +package already installed in the active environment. + ## Training ### Finetune on LeRobot Datasets -Before training, please refer to `workspace/README.md` for detailed configuration instructions including: - -Training script path configuration +Before training, see `workspace/README.md` for configuration details, +including: +- Training script configuration - GPU setup - Model and data paths - Robot DOF configuration - Training hyperparameters -Download the Flow/FAST pretrained model and run: +Download the pretrained checkpoint, copy +`workspace/example/lerobot/qwen2_5_lerobot_template.yml`, replace the +placeholder paths, and launch training with: + ```bash -bash ./workspace/lerobot_example/run.sh +python -m wall_x.trainer.fsdp_trainer.train_fsdp --config ``` +For Wall-OSS-0.5 fine-tuning, normalization, LIBERO evaluation, and open-loop +WebSocket evaluation instructions, see `workspace/README.md`. + ## Inference ### Basic Action Inference -For model inference, please refer to: +For a minimal end-to-end example, run: ```bash -python ./scripts/fake_inference.py +python scripts/fake_inference.py --checkpoint-path ``` This script demonstrates how to: -- Load the Wall-OSS model using `Qwen2_5_VLMoEForAction.from_pretrained()` -- Prepare input data including proprioceptive information, attention masks, and dataset specifications -- Run inference in validation mode with proper data types (bfloat16) -- Validate model outputs and check for numerical stability -### Open-Loop Evaluation +- Load the Wall-OSS model with `Qwen2_5_VLMoEForAction.from_pretrained()` +- Prepare proprioceptive inputs, attention masks, and dataset specs +- Run inference in `validate` mode at bfloat16 +- Validate output shape and check numerical stability -To generate an open-loop comparison plot, please follow: +### Simulator Evaluation + +Convenience launchers for closed-loop simulator evaluation live under +`scripts/`. LIBERO simulator setup is optional and documented with the helper +scripts; see `scripts/README.md`. ```bash -python ./scripts/draw_openloop_plot.py +bash scripts/run_libero.sh ``` -### VQA Inference and Chain-of-Thought Testing +### WebSocket Serving -To run VQA inference and test the model's Chain-of-Thought (COT) reasoning capabilities, please follow: +Start a Wall-X WebSocket server with: ```bash -python ./scripts/vqa_inference.py +bash scripts/run_serving.sh \ + --checkpoint-path \ + --train-config-path \ + --port 32195 ``` -This script can be used to test the model's COT reasoning abilities for embodied tasks. Below is an example of COT testing: +The wrapper has no built-in checkpoint path. It returns raw model action chunks +by default, which is suitable for open-loop evaluation. Pass +`--serialize-actions` for clients that expect robot-serialized actions. -**Input Image:** +### Open-Loop WebSocket Evaluation -![COT Example Frame](assets/cot_example_frame.png) +To compare predictions from a running Wall-X WebSocket server against LeRobot +ground truth, run: -**Input Text:** -``` -To move the red block in the plate with same color, what should you do next? Think step by step. -``` - -**Model Output (COT Reasoning):** -``` -To move the red block in the plate with the same color, you should first locate the red block. It is currently positioned on the table, not in the plate. Then, you should carefully grasp the red block using your fingers. Next, you should use your hand to lift the red block from the table and place it into the plate that is also red in color. Ensure that the red block is securely placed in the plate without slipping or falling. +```bash +python scripts/draw_openloop_plot.py \ + --uri ws://127.0.0.1:32195 \ + --dataset-root \ + --train-config \ + --episode-indices 0,1,2 ``` ## Join Our Community -- Scan the QR code on WeChat to join the discussion group, where you can engage in in-depth exchanges with community developers and the official team. + +Scan the QR code on WeChat to join the discussion group. + QR Code -## 📚 Cite Us +## Cite Us If you find WALL-OSS models useful, please cite: diff --git a/README.zh-CN.md b/README.zh-CN.md new file mode 100644 index 0000000..2902f05 --- /dev/null +++ b/README.zh-CN.md @@ -0,0 +1,152 @@ +# Wall-X + +Wall-X 开源仓库提供 WALL 系列开源具身基础模型的训练与推理代码，当前公开范围包含 LeRobot 数据链路、Qwen2.5 模型、Wall-OSS-0.5 训练配置、最小 fake inference、LIBERO 仿真评测、open-loop WebSocket 评测，以及安装时编译的导出 CUDA 算子源码。 + +## 最新动态 + +- 2026 年 6 月：Wall-X 1.1.0 更新 Wall-OSS-0.5 训练与推理栈，包含公开 serving/evaluation runtime、DMuon 训练支持，以及安装期 CUDA 算子编译。 +- 2026 年 5 月：发布 [WALL-WM: Carving World Action Modeling at the Event Joints](https://x2robot.com/api/files/file/WALL-WM.pdf)。 +- 2026 年 5 月：发布 [Wall-OSS-0.5: A Deployment-Ready VLA with Gradient-Bridged Pretraining](https://x2robot.com/api/files/file/wall_oss_05.pdf)。 +- 2025 年 9 月：发布 [WALL-OSS: Igniting VLMs toward the Embodied Space](https://x2robot.com/en/research/68bc2cde8497d7f238dde690)。 + +## 模型 + +- WALL-OSS-0.5: https://huggingface.co/x-square-robot/wall-oss-0.5 +- WALL-OSS-FLOW-0.1: https://huggingface.co/x-square-robot/wall-oss-flow-0.1 +- WALL-OSS-FLOW: https://huggingface.co/x-square-robot/wall-oss-flow +- WALL-OSS-FAST: https://huggingface.co/x-square-robot/wall-oss-fast + +## 环境安装 + +创建 Python 环境： + +```bash +conda create --name wallx python=3.10 +conda activate wallx +``` + +安装依赖： + +```bash +pip install -r requirements.txt +MAX_JOBS=4 pip install flash-attn==2.7.4.post1 --no-build-isolation +``` + +安装默认训练配置使用的 DMuon： + +```bash +pip install "dmuon @ git+https://github.com/X-Square-Robot/dmuon.git" +``` + +安装 LeRobot： + +```bash +git clone https://github.com/huggingface/lerobot.git +cd lerobot +git checkout c66cd401767e60baece16e1cf68da2824227e076 +pip install -e . +``` + +安装 Wall-X： + +```bash +MAX_JOBS=8 pip install --no-build-isolation -e . +``` + +公开辅助脚本位于 `scripts/` 下，文档中的命令保持仓库根目录形式，例如 `python scripts/fake_inference.py` 或 `bash scripts/run_libero.sh`。 + +导出的 CUDA 算子源码位于 `wall_x/model/core/ops/csrc/`。安装 Wall-X 时，`setup.py` 会通过 PyTorch `CUDAExtension` 编译这些算子。`requirements.txt` 已包含 `ninja` 用于并行编译，`MAX_JOBS` 控制编译并发。这里需要 `--no-build-isolation`，这样构建过程会复用当前环境里已经安装好的 torch。 + +## 使用 LeRobot 训练 + +仓库内提供了 Qwen2.5 + LeRobot 训练模板： + +```text +workspace/example/lerobot/qwen2_5_lerobot_template.yml +``` + +复制该模板，替换数据集、归一化统计、模型和 checkpoint 输出路径后运行： + +```bash +python -m wall_x.trainer.fsdp_trainer.train_fsdp --config +``` + +完整 Wall-OSS-0.5 微调、归一化统计、LIBERO 评测和 open-loop WebSocket 评测说明见 `workspace/README.md`。 + +## Fake Inference + +`scripts/fake_inference.py` 提供最小推理链路检查，用于确认 checkpoint、训练配置、processor、normalizer 和模型调用可以连通： + +```bash +python scripts/fake_inference.py --checkpoint-path +``` + +如果 checkpoint 旁边没有 `config.yml` 或 `config.yaml`，可以显式指定训练配置： + +```bash +python scripts/fake_inference.py \ + --checkpoint-path \ + --train-config-path +``` + +## LIBERO 评测 + +使用封装脚本运行 LIBERO 评测： + +```bash +bash scripts/run_libero.sh +``` + +常用环境变量： + +```bash +CHECKPOINT_PATH= +TRAIN_CONFIG_PATH= +TASK_SUITE_NAME=libero_spatial +TASK_INDICES=0,1,2 +NUM_TRIALS_PER_TASK=50 +CUDA_ID=0 +SMOKE=1 +MAX_INFER_TIMES=52 +``` + +`MAX_INFER_TIMES` 可以不传；默认会按 suite 使用与 LIBERO 评测对齐的 action chunk 数：spatial 22、object 28、goal 30、libero_10 52、libero_90 40。 + +也可以直接调用 Python 入口： + +```bash +python scripts/infer_libero.py \ + --checkpoint-path \ + --task-suite-name libero_spatial \ + --num-trials-per-task 50 \ + --driver-mode in_process +``` + +## Open-Loop WebSocket 评测 + +```bash +python scripts/draw_openloop_plot.py \ + --uri ws://127.0.0.1:32195 \ + --dataset-root \ + --train-config \ + --episode-indices 0,1,2 +``` + +## 社区 + +欢迎扫码加入社区讨论群。 + +QR Code + +## 引用 + +如果 WALL-OSS 对你的研究或项目有帮助，请引用： + +```bibtex +@article{zhai2025igniting, + title = {Igniting VLMs Toward the Embodied Space}, + author = {Zhai, Andy and Liu, Brae and Fang, Bruno and Cai, Chalse and Ma, Ellie and Yin, Ethan and Wang, Hao and Zhou, Hugo and Wang, James and Shi, Lights and Liang, Lucy and Wang, Make and Wang, Qian and Gan, Roy and Yu, Ryan and Li, Shalfun and Liu, Starrick and Chen, Sylas and Chen, Vincent and Xu, Zach}, + journal = {arXiv preprint arXiv:2509.11766}, + year = {2025} +} +``` diff --git a/assets/cot_example_frame.png b/assets/cot_example_frame.png deleted file mode 100644 index 02b3450456c3c512019bd95a71f7f4983bc03240..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 136402 zcmV)HK)t_-P)32T`a&i3V^@@ z1a=6J)5Dpb={w!`wz~3@N0&$y5$}MD+{S_9FnG+VO*`3b=YbK)9|yA8s0i5BbjaiHMmSodFd3b91f+|E%%_Q@u-~aw+pMCc5;c=R#QKHS~ z%jd%InK{d{G)>>Ud2@Dl7RAHSXfz&=kB*K8gMl&Td*A!sXP*;j*?u#!0%;(eDZ1(N9-=3eJPbL!)dGzQ>k|arzFtag+nZrV4K07)(I(l$C z9*>8^VH8E9(IASV+|8>wFK~kZASWiG6|5$yO{P@$8G%ZeKnem7k%1*5RT>NiCvQ*w z@P|MA{ont6nx^m&5y?qJNC<$60U)y6DS4g(c=GhYlP6COjvk1}``Ln-9VjfKG3M~_ za5NfmG>M{UI82fxiDELw1R^FN5#h+zQ43vo8z^TOMZncMLe?o;xbK_;AS6VjEa;*? z)z((k(51y}TJJzu9U?+@Fr7}HKmX{5Km6ejzW)>F++vaCdH&bG{>!&--%e&zYpu1m z+D%E4v@i-SOk)60OGHps)i1vI>g%t+p1eP`)`lTW7t^OtpMLt;$B!O8IzB!Ikmq@x z=R{gz&XSmmyuR;|B)^2S-OoaU6$s z0vS;>>0U7{vfP_(Dq1HLVNsaD%pj>cS*&2-Sv49Jwz(4*KC_&oRpQV<4Ek69?O)Yz zt>?k$XTY7XD+XMRus)1%?YHB zVWjm&>I}Hv!}^L>Yga8wqbDNb%668olZc4RXWsiZLk*J?UOYZNe*F0HaF{sf*t+nk z*7K%fz+pI($z-uu%oq7!Fc^(S!{LyK7K=raBu7U_&p&z*dhWenEEbE!VltUzSr+F` zRTqodVzD?oJInJNfVK9)gGa`gD2l4?g}!DJ4>({SUm4v+a+1yueL2a0sz`~6&d7g^M_#intI*O8E zmSs~X%K5WK-D{QZ1Qx6Zo!>u&c z66u0C0)T4tEzD1tD0WBZ=jX;m4<9~E;<5KW&%O74mcPsM+^a-U6juKt+SS2;0Yo;j z)=nm~lamwnS`b6;$g2!a96i}_i>q9xQBaDAkGbOsP2sx#m^#pV{F^MR63$Nue{P5M{=?O(Nj zZY=Pn88Co~0b^CI>0`ju>UzCP)$2=b|4q;uw>qY^< zYRwixE3B?A%+?td044%p3E|xa8>wOM_g)K|>7 z=x1$38sdnEI4LMpr+m}sZN*$0m{Y;(OpAz+imHO7^=rZ8g6eFHF@^wuvOpkeU0JgK zfDx7C>b+ND@4Yc5FyJ&z^E@Xa<>qk0qC`vp@;nU^12e~QY^^O8T7#f2n}<2a6& zGf?zRD!pW7LX;cR)yIH~L1Bds08k-SEk(K|=7gO~#cVJU5u$2-w6uaa=ZJk2MWU2t z8OI|4N?hxI&}keIfimyCBCo#g1OS4crY}F3FLS(Ha zBBm8gBO)ft%(0D~bMsm5y(bDni~=yRF~*pPRW&SekgI~EbrCQC5({89nqz!yu);<9Mte~USMSD$0iE+u zV0H_!I`OP;3ECR4^ByJg1@0W^h5-?gu|z~cTCX@$UDB#IrK&!(BP+eC7^L^UVP0YNy*y{_6q0P2ILnh4UE4Q3|g`aVbL9H|esA%69->PqY=pAnRz zgSPBI^|#9Ls;YuA0HR9`iOJSRN>y)*Su=LASmb$L@=+BLYi;IT{VP?UgTPu5gWAxz zlI`kikxO20gqcO1ss>sW1dA1%+>!yir6sI^vbDS*h@ zptDvhyW)%}3`GneX#Ajplo1h&C88urLR{`=Y4G5BZ!XP)>u+6}eRdkg?-uk8feMx;n>@S*sAg2y zWS*(V3f-Z*2TU^o?}1Xqa)2@HO+-YwF?HirX9r63B~=YOx4{?dk2P$1Z3f&J#d7Q% z{`7U_tUca9tW}kC(9lqe2YNP)Q)zSbHL$I$X<8CcLkg*p;JU2c#-F7g99|6mHUa0z zC6=lkGDJkMP5cR*An2^tE*V>%2Nv65*$N*FOudpz%2jBZ4XEI@ZMnSmWX--J1;wLq zy(Cbj4#dn13f}0T*2i(2Wp3HHS^!!KXVzLkDvdVagW<$`FMcWKdC5g2n4O6{fWUx@ zuWQe~Ue#!^>LjXj(`t>D^B-&*mH`k7pnb|fNU5cN4Ju@|bUyu@&?A|gz@lua2#6h%p% z=b!>_gQOlsQR-5sPDz8R9ag^+-$MJ9`=B{{rYU_E^Ao0~)(q9=hguBSQ(*9$T0m_> zfNd{cr45(J)ZI|~#ZuDQlph2&q?b|M>%MED#a%A;0e4O*m2|nfQ?9E6>_%2h&5{I6wIDSg2X&*jRq^vGPBk!D4~%eJ`MsP7$ob22{kkI zD!wiKqgq-u!u?7M$ECJ&IRgs7qmJ#ULivbIa4`k~9So|=cl`<)IOj?mVqjWJjPkx9 z-{Ey;ZjjU}%&;y%)S*3Bjf*h<^@l4cc3;r6Q1ZsU4ZAv@n-bfqhpYfZrIMK{;?RS@ z)VL;ee6f1B)aHGKBwlZ{x_L#Y>`rT-6lgB) zUYqsy8CV*is$pKq)Ep21Dui0nSa^l13Zx`+OEpNLWSCZIdVpN=25mycV7PK@LIGd% zxHLOi_Nq?$3k``FrJmGypfx7|Kz8|904en?Rc$!nR8?!j=tZ!UayEND#2iG;R%$KO zbs(_==UkkGE!+_W9e7E@mvnMAvDU^EheoPq8+=4b+fnWXWh~fev9gu1F_^N&<*Ev< zb=VZ2u55oRU8?z>%US9*L942o73@`23GVGYi=xN{t=(vNbEfJ;3jtq@MWxde^BBGr z0|f#ZjVLPm>_xOE7Vt{ zj1_3&C)f9*tY|e@n$Wjlz-y%Vt3V~lY(s;-i%@AH*dTgq!IUbT(Kg zBIOl}2!ZO_YQqZMX8e_4+a76htVW+HDgh~R9k@wq`dt=!I?W4FiPEP z^Zn8%0dzs|S@!A^mH3Zxb45|41wn5Qvrs^tbM^FrkP4!zLH%0aF(L{K#@b~FYw2%8 z{q^vzJkQk^v0C+yR;^WCs&3Y?;6errEW@MfZK^5*S(XJ93IV+$lBQ`8B};5s_PLZN zV6823sPd)WS;_Z#o{KoD#X111Awn$?vvV4&Yp(UbzFF$lqw1zHrjm^6Tch3Vg><8{ zOJRAL*~XESF?(TVy2jOJz(Ka+HM?x-=?E_Ka_NbcR)eq#iqX#{M8v6t7D5aMGyBzO z2`atC!OY%hp=(2TKw+8C&|q^Od@Uv>;^nWUVo<0C8+-=<=wJPxe^qt8x?A~F!k>-n z>Y4hBeLI=jQ9WgM8)lP&h-CnH)$Xq>&i0KLYaI+{L(_dWB{vltKLIt^n`Q;pW6%C zxG87{#ndC(TK0xg!)u|LG$>ejMybXgXqmEA|5g#4wF;^lBA){;m8(E~&-cMnqeVgA z>fOFD_@nal_Zi6=I3O(xUp;T7u~k%6TUiX1siS*5H49}$daE%meaXvDtVju!F~4AP zt!RIvO@>$LnoW5wSZG#vz;c6dhY1=0q^V_)Ge1br8_nNAU8KVeNj*Gsycu=#GPH@>U6 z>eGOC2z?!{^Qy~zm-w7bRC2c6SFZ zVfpLPfa$^+Yy8^zXWc%xa(_z=xM09-{JdK8y znp3M~X+6{V;)ZZ7%nZfMP=^V9EmZSqf-SWClfg9-4-{VThDE;ilf+AA*M+u_E$%@zfn8yN*|IG>`ZcX|<=wir+EV=Lnab3#|5b6_9s2skK41<}YlShEw2X0DdAA+* zRM}uq#7>9)7do``C@?)ct~Vy5v-?!$N7Q}QRcPDNrnXK4*4q90N2tEO8wGm-muA2_ zls7*xCHN`~`zEL2IvjHQ4~=NG@8AsdK^q^+DY~T&^}x+9+$zr8oCaJ{gFRb-9YG(y zZT|d+0b4WR%lXDHZ0)-ld$sDMya0A({%S-#Ua_@3hYNT0sLzWqkd<_o_DGA5bPjpj$fdffu#peAP!35|!@p`7fKx%223+XDb(ONxLMx4@ z&tX@OC+NYG-qYb)ck0#V3#PvEtJGQx|1$s_#g(!vb#HD>m#*KJA_}*b=tau{oV*%g zqfp@)y2PyBebah!=}Fhl2ofrPZLnJaRIFnS-?#(c6k!FsWUn&&zwfi_Zm$Hy_1y26 znOZbJJHnh=_t(Sz72>#ud$@-d^fLH6Za=zEaM@cKvFunkgIWL9R)g9$$9GrXy@zXs zuogBj>sn!yl?eC_1)>>~jov521Fb~n@NlJtmESFWZ>v$MrKq^#$y;e>O9Ip

N{* z>`ZFB3bZt^KHYR_4BTPhTRYpYomW@;sqcQb>4K>{ZE;O=2luV!-1%HBJ_E2vANx7(BT&1P{vc8J78r2IZTJA@@xvI!cO2@He&1Sdo*UhUJ_}8U-D7)TX01ohZeOY6)?rN$ z;@#R@7u(VS1st4`7J427L35&pBf8u6)2UV0Ae8V7*LBj=SFCgi_oS>Jq-W}OGhO!@ z@Ri_N5qG!oZujs(Lc5yoCb0hmcb~g=9kUnKarJbaB4(`ZD%u40e~GUHb~x&Kt*(up z5C7}8JatCD-ZTU5N$tPGu-hfyNjo>zhx>eS#<_0TNe&NMbQ3*9 zmR2fOE+ITDB1E9oWBAtRjIKPa%M%_~a|G6n@dC{GP^i|8QQ*#mu4eKbD^1TMU zGidk>wyN{eZfduEyN7*-zVx;Ys)#H#;LA~;J0UpkNe>%yb-7L2=PamJv-u2d=i~2d zzb%NgiPa4t!=1pl>wpVlQU(4&&$2vpmcU1;!lkK70BQ0*}U+ILE zf$uJ9c$1B{Ngi|GTbtxLtf%sIl`Z_4V<_Y9ZDqCT%W2 z)SK;qo{_Aflbyt>uhz7z?gb6`&V;_>={;7D>ngjkq5y4xMmYYjG!VV6R}e)`A_Ipl>&GW>09}oeK=!$-K22scJI? z>2ec(eCMWaCvZ6v^bJ6@YPa5#I|tW7i*{k9yIt555`BGec^`1YCS27^+YoM-_1^k- zuLEv!ESK9DZ4Kipp6@#cb+7;WdDuhP{~qoZ+Vt&KIzVU!qTZPQbx)vei>|{RnuMJY zH`{3IitgACo8VqI+%U(*F7&O2N2imHR7&ugcG@un`?(bx2KKhse@!6X{=nuPp#>ePruu~T?aI@s488gsage-)ZPqae|mEBa5LPO%1os1@>E|4ed_O? zORfQSXz^N_pc~kEX=J@YxsL&=OMG`)|GvG!-a{Ys(Xsk6^+j3l7AnTv!+W|tIQ7WqYwA_44;{y&UDJ2(I_xUR+*p&uWwa{*=&;9uwhcz?y0fVj1+-OMcZE_Q) zrE_kGNXdl@G*lnd5xOol6fZWvQoBK|!(;2&pG2zA)iXO+zq|{ewS{Ure5hrA8kSh^ zeXZx+M%wjIX79D5QZL5W`s1rvfSa~ZWtVRFp&~D?PRx zq?=yScU8;Ju8*(2Y^sS~h%uA0BLw;KQj2qZ#8 zfhq=-H>@C$QPtwKUBkNc-RZKw%YRWn;C+?Kb?@bs^M+~V>-A&5m;G(u2(R8t@H4 zpGkRFE%)u5JK-&J(Y*s4%eMG_hF%9D;=XqP3)zx`0uFQe@QpIj&u^M3zm%x+DWeVe_i0x_2VQcLchF&u?liZifNib#8PI_i#D5I=5>W%`O+Ibn}6@Yt*&N z^<)o~V60u=YR>DzNBi{+gWG!_P*fcV|K1hdcm0rW-cEBfU%Y7se9wCC;U4zHC44QL z@q&!il|WTAg7G(Le*c*6h7e}kL#2$jEiEXlO2uh8$_VRn*-_p_)_d!)ZX;YXzTKOk z?-<;531%O7&ko@d0kl^SG;P?T*F3p-<2NSqS1!r*UtHgrj&!DdQiv~Hlz=Vk6Iw2@ zQ?#CZ)3|?SxW{Ca5yzwDlb zy%1I;ObCd`Lx`w;$}cUGfO}W4o?pAfrP>4ad9M1g`qVc?R|L5PhYzReDfOHUdh50TrR2L>h0_msA?#DQ5Va$5)PKa$F+6xy93(_ zURVDXtyaw5!+m;vP)VgX^SQ9IcKEs5wOws42=2Crjn>u&*{xaMzFvJeCg_GbzC+VQ z=l$z$|H7TKVF%iXd-XlJld(5=@Fw&8ftjzjg1)`ItqJ=289-gfYOfTv9@vd#?*_!~ zodnP=DNUab5(gsMI}P|Ab_;iu{b=i)lRCv* zqX)OC-rm{(Zn2=!J=_@VmeUhaV7GD&%5l}zOPue1iElOUg#XOA;*zOXOQ4mh~Q z&TJ0%4ETni-_+dweCLMg46L{W^zFvV6nuTOxZy{mn|yY)UBW)f6kN- z^!BD95f>j4L764IelYDQ9-zvjtq19>?Z?)zWt4L5e9e^w=#BE-wm;Pi4H>Arp_Y78 z5YswUU45pL+OFy4I#FzW2-QY;s4UCga(Jc8m|V?=E}q}gHlaOHruB1$+h+_nJK4{o zR`Qr@K-@hKE2sFj{kt6SI|xjRaNGVRQa|fEh1DNxf(v|ebdL` zvBaQ_wnFQ)rABWX*GA2nhuemV?n)M&p9th<0tk9g|{O{FIw>ov< z*Qw`3>iT)uS|#OU1^4P1(x%T_-PMIE>bvv9D{og~S#Q+0v>ygQ)%oivt9Ta2P|y?0XJ=eUY|#G4;R8+oVS%e@4MJ~50`{Wl5X5@ zaNPY-T<(gY_2%I_Vw_(OH-jFRVZfL6{&pJ5?jlFINm#EkuBmZzIq1x#xkXZ?wSVy; zvhPjvx%-0kw7Zer=BeLPOVD{RXbpS6%Li^g zPxKz{G&CsEwc-!a%Hc6p2he|4YnN{^VO1t$%m4BDv z((|lc$Isn>fo}7p?Z%%Ud_ATs!v2dReX$6Q9xtoR?(b0c0FvK^)t9-RjqU@3vi}SJ_BwhI5I6{uicD*P+pD&~)p) z($#Z(t5+In+_e#1K!`hk?aK`AgVhpU;rW%hy-;=N`{}|QU$juM^Kj7X76f}V32G!m zZ#%IUKC@Bj<<~`8~r!NblzVN9YjMq!eJVBwz8{lsHFBytnW_7C9at|`x;+Q zPwF>2SFLls>QJ{&G=iSiycIBzR!Qi#Y`FS_ueDiRiUROZP_YX-^0uj!`&);VoZzbp zH0yC@TViqI`#HD&7pT!_wDO>V*_+P)8?R?e^F}{AJ?3YvQpQgj_W<7aJ^u3 z;~|R|6R!$3;oNrx_CzJU+>ucGp!3~~UUjWhYjefxLf8`(^UlHMB2<9;!x(N9+H&Z& z&tcwnPK|OWz_y@yKuL*X=050VB4e{>6={l*U53Bb&CgNK6;SH>TRZ1a}%GF45b&HNY#v{tR~SG~h;};jSzy6wDGHHV~68S3n6RWqU8d}j|E z%wSiDPDH9pN{x-lVy|Yrx21IU0~+&onHg5amL3*kOm!aiK2XU!1Pw}stpwciqCo~P zG5a-(*IuvPh1j(dd20=E=6y!)@0vum|8N)es{MqW?7My5yqAEg;LL}V@^BCLum&vF z`p>qJTbt+d_UT3q*fxC-7X|oS@!fXxX1j$>!dJGY+l5or&o^D^PY-Pk5?QBeLLK+( zixQq83QBw@nG4-7Y<^-P~_!1wmEquZ#?$8Al#_9VJ$N7!YU-l{W=3Lz_#e=&<% zO6nW=GE{wg>WEZq$BQDu+J_8UeYozFp9xfY_bSt_FpK z1z`nLf@##LU^VjD)b`gJGFbg2tqI*kr)-Ch6cO?Oni&wyx7vMMMcQWINceG+{?uP= zaDhT{q#|p%0c#rHFB9}$!6jnK?+UCjiPVSXz5f7e2ktzM?>hqRdAk{$Z+v%u${u|O z;nI_T>B!A}gX?zr+pBvtZy0W@x7!l#N}jq9WvePb;l0jhI=#HOej>=s!Tr{-Ox!s< zxiBvJ27sB%(mrdSX{B6$H=(WBSD}}8osL(q?I`#PQS0j;sM7Y9-_qK@AXR06i5*|k z8%RXZ4`=PJ1MZjxd}GA4Kj!EbU=x%3b@aB%scmF2-NOw-JHq+qDO%%;w-vSi-GTM_ z4%;f1SKPe!c8F_&YiGcHD6%a!v~8hf%eyx)?hmTnJj2tnTxF#P-%+M}htfqVSTog~ zd3QS;-G=#Zb-Tj)t$oGx;tKJC5`3XRe76?*CZU>{#$GP+K5Eluboo=k;CGZGa4(AW z!JW~7Z&LRS7rVunHYcMpoT?Kp)O7ozV|LQ*dL4G^&m`U}=-lKtw)6Ycl&uQ=Z}_2A zF1V|PFNUpVvYzv{?Ynyj7sh7p9Qq`fIzzogo11LpO1-{koYxI^h5_FUw2Lck%Xrq{ zYbW;inU^=wNcy4>ud-*tyCMErMAl3~@s(YdK9*C=S{P zF^Q$P^%1W*NTgtJ;t#Hj0Q1sfneHEym zTY4*7e%t+G13#nXlmlHymoy#-(6)`N=QKNEX&W}uX40LyQRDd=h%v_MMD=OC?Vj`* z^x71*ubs`L>E*k1y1z-VZr|MY!6viaWoZvH(EX%RUi^uM$E&Y>(|LdWMrw@C7?bzq zeNs2)*ovccCZmUo(hSiQ92;11X~b?5?!E@AUt31eb(ZVh(Tk#J2Q*d+*Vr!BV+ zYa9KqUTJgU^%H*GCIZvSlMWrxG8%cQLt;AxS5UrGrl$)b*uwy5&-p$5lixAP z8?ID=$*r%fZinU(t`#oz7}dJ_yq;$E*%}+41^_SEjFJbiQ4|GsTh)I1?!dMT;}THU zZ?<}ChmN1!!iE~vuB;RdMyRU5U5yL70k}Sg{f2Gcm52fKb@I0N*Un+Rh+iLWr;t$^ z9QOhmvMFAAGKJt(m#%S_4(?O5T*MS6wj1Fmup}abFzZvasHLb>Ye6SgQf=3DixC z7g_>$#fBRV?iKWo?wWhr>38_xQ~aUwyXW6swdj5!oZH zs|9_u^iFVJ)Pu0izuWkxKRGkV4zSgy# zlI?p0Ygli?nz#=#tBs-VQjaVqmAXBww+k0yb{}zH9}u_|?3-M2 zlh8MfJ)xeL!JiA4Zxgslysfjkx&dZAk`#eFTGxj<9;o0FD0k^eSc&)ru&Qyc;J|9% zHu++=rK+VC3$3(_%WAzlw3uEPOCKJ%L)^Do3+tQ%z|8d^teuSHI%wscqrW#`>)1nX zy+gQjD}E0*34I~HSA-kJx=Saae83wYHOh7gb8JI0yUB^FVy7DX`l*yRLk?WC zw@NGG-c|`tb>-j^W_Hdg6*}C~CfS2;t};ELr3K2;3dqbLsDj9-Dok^V zs@5;;eE2=i#28~hBI2D>)hLP#C=&&cn`QJ<=w5Z6G|{pM5`|s2?EY4JnY3=R*U>p+il>s?XIU>j^)16358h>+QOqpcBTq9|f!OYFU;vR=ftxlic#G;g>rH$du3 zJlR20x++ZHX=s?!3kDp}x9w1=V5!4x{H2Cgv=i8$=(SsfPIZ6tjn{~_Xux$7)@?m@ zTlI&CDv9$ujwE*->H^2gbo|%Ze_w|Fzt^yrjEa#*#-g-P~-7LC^>wU!#cvAx0t+R~v z0_*)@47UIkLFp5I+nVC?X6y0w5_CJru0n_n+4gm4JeQW-?_m>YtnBq!8ZL~ExjqmP zfK`>R=2+PZiU6)F6nO~DlDTMSX>i_njVkIixOOcpGku%_Z;l>vtVG z9YcG^fcLHH`mUhjeD&H=o4{^KTSEhG<3z*RzjMO6s~y#eCshI1cKf{p**2tM&lKyv z{uopE_Ug*YZV5EH{8l$ZF@uR%bXn78)$hJ3Wp4wQ{jc_E8FfJbq%51iM8Ay(N(yFP zh$Hn8696MnXh?qWP|!xLRT_D{`g9S!Ohrsducx+rzPht!uIfv*xl5Uo@{BW zE8n`*hi?|@M?qi%w~?yu9Rt3HKF~5wtn>cn_MeL$yL=Bffp%{{4Dohvp{~As7*2N) zp%RNkgxR2E$we8(YmkxZT4vzEmzYWKeFy?24oxf*frxmk2CGops#jb42seYaH3co! zvSH}JicH_{?reOkfIEV>eI*|nR)lKQ_YC-!fptk7QC$(y%K4jDg35xut!7Yd?;XOm zG5}@{L1FhM?%RT0&2#-2u9XOI0l+8@t0oY6QD#;lRR|OM@*K8=KHt!Wo1&8t6A{z^ za4JkhTKh^*o54F<&sq-626uNAWuvPa;`&(ql?dt!N7Be%+2HvLS_^FW%I)kkw2k(C zHvD@Ad<(E9; zzT%*v;topl(~9AJ<0%bql{IL^d(TQC9&1}@ivUYB+rG66;aUwR8~QBnobjeAt+Bq_ zw&6?beEM>Jukx+!W^I#VZ;fAf;;pQIrEVtOw%NYvuG$GMoVE+Lx+X}K9BJya7xE&& z00<3x&N82}9*`3FI1%-os9Jrxali3r{edP;N4Ei7`%Fx$zQcxDZ%rSz0&Oi?UkSio zuUOBr_vsd<=Y@uxap9mYjDhGhWUYS;t5~tVpBqAZ#tI~<*>SMCz8AJxf2Mxq+{&DY z5Ro!ek&;r8ZK)*PdR3n-q<6ejR#D)}B^!2Yh%~Ncsjf}jYPe1T5V2HAwZZ@#0SbmN zGeTOsWff6^DpW&S3IoWfIe|u(U#UZG3h#0!_z*%K5L}_EA*#H+--){w{r1hy^LJ}@ zf+k_1+I4p?FM`#npsE}766o~>+Qewl=zQb;<}R~_z}-n0wH9-CuL0jQ+*p)ZA#1PR z+{PK))KIHtw^BGdzt|bMa}S$9BM;=l)x0siVZAeNLmu3QzWq`;V%M8AJ1O9r&$THy zYQ$7aIsn z7})7*t&xbkZO-MEt5me`SkQgfU}@inc5AD-aArCm*o4|PM9$kXG4%@vyXnkMOwjGg z+(eAKrugSgpibF${pk>#T>V)#-iLC!MUHG;gU8au-5u$@H824nHOt8u(WcU)cX^Jz$TIs|Lu<>S*n$}|#g082HggI5L)wu3Ri+3m&;Hx5R@PU&(Kl}rgk_(Pd7r|x9GCv{+_Hf)c4f7CiWHXmJ*Eb zAOo($b#UK4+h7$-QjgsrqDFAglKXB`)wW%?uaU)Z#RG{OD)ovgd{I9e@9vLpW(^S0 zwGyDKOu{CMSPxvPE5Uuw@5BXAw0{@YcS9&;2w-=TL@T8VVW*o9v;fg`^?vS z1@)NB3a(W|SZ{#Qp3w0Iu5S3DMj-C>VfY$H$xdA0>!Aw5>i)!4V(DAqLf3&Ch3__8 zg9X6}e5L8Cqjisy_coy+hHr=)n~)JA+D|38&meT?HQ@cE>1TJCpi=4k3FggZsBQTJ5fK+J6H)zB75^q$sTElp za(l4lw}N%CteLY+2*y(9uIsF{%)slEs; zt(V2PdW-BYr>0lbZDoaOAKPxuT!|GO9VP%#520A?;j&S`dOFkkh@mwhMcR`!E*%0{ zwZy2O#vM{zSh{_jT}u{Lp?ii`jAASAcH30?LJInN(g(gZ-0-;XQ5w1ZLJt%RRKu=- zGg6>Jk87z%;Omq3xHg?ZPbFmlF8G6qnOSLh9%7FRxZlm&PReSvP&&vdo(#r9w}t<)nqq)Pr5B2J7GD!pjbncvZg>pjiES zZ#z|*s$*K+STR90LUXDXwlhKZ?Twq%8}?sEX~X6k@EX$GXTEN~sI9sd>t)OjNlV^y z*k|m$bl+}lPP!*zgFmr!nyM%=&M5?Yb-RpO9e-4;uTwuIxCWI zkd#BwI?+x|f=aPXReK_!Lttfv0p^Am{y+P7AJ*zW7~qmGwtl*;jTZCwO)LBB!5P$`+$$t*wvoj6S4O$!W}Y)TrtD#9s;cJ z#qDDYs@FRsMlW2KPT$JaSM#nM0=69kUUfmPD090SVbxAU-RUzv>@XUo6YBXqYpUWk zu6N#EorFwLUyEof`NR{>fLL0bOVYcJZUM*rTzb@X}PHrqZ! zpNU0pd`oT;FtH{J?jE*JMUY_`BfVRar>da7y|;DG`y}6$&wVJcrZ#VLl-msYLgT@8&*6F=+@PUr z94p)UfWK?dIK`-fo8(d}62NtSsPX|*&FkC@U%pw`S_AIPg}7Fzm$29qSQGuRVLvyF zh+cBO8?-~-3{->K?LKc?Df$B9?g|OM-QMny0biI*af#OUY7svGA#Yca@N3GyT<#NV zlBO<%YYKeXiHy?AfO|FLu4mTEs+a3oBMG=|Cu#7ctv=M~d3zuiR*%|wL{(5L8>MLL zJh~C$QAJE2VzlklE*4DOLDO7op=+o_+zp$6B z_B&(Qwxw^4&TJKgg_S{V_n|)MI}f`nD|)%g)PAOC8jLZ4Y&QE-eWw+&(^+j|YZG)Q zlHECRuiZ#(bTkvXh2 zt)0w$n`c!Gp~0R0xXn>_AM=1Wi*>prtk-d?n-}Wz{SWMb;ED(!%uGFc;f6p&r4--R zt2S+=#(#^< zxKa@JjR$;qKtyW@(={zY(dL%;RbfSdrI&EauV0A)Uk5F|ne^W~!93~*C|!tr{hC(s z>xv!wvT)aF(9MU-cH0Rrb&74L&9GhLx}Ybmn9qPq{+cIXTc7b7ll$i6rCxgQ9`*|Q z3i|e0bT$l*!9LNViauY0yzE8;zJ$G~&uQQOsd4AZ8sT09vusFQw&KA(U(|*QbvcRt zu4j7mq3Eq_=7tOWLfC4ts+NfeEvzd-#J4%@@`hoh|9cO83d8=ehr0r$Knbj!y*+dx z)n#1L<}2$xV*6-J*CNC+_>HX~ez+%MlSE)IzxfG;y!sn&?I4 zp)}!0Vl_je%Wa4Txxy^Km!v0e5P#oxQmP6(+4LGXU3ASRnw~bw*8tvxX>H?Pbh; z(4L&$8D{3TO|yQM$mX?ecLnY6%W79tR?2MyZLSa6k<6{rnL6MSeXb3;&@q}lS&hzJ z^1p4rv?cVJ=*1+B|;5IbG zpZ&K@?XZW%kT-j>bIG-p@N>JQWLICmurS$%uy6jPZDG&1!=8lfUEliNUCg-|XluF` zJ-oljya{ZJ4=W6KTUYBQ&(sfKTz!7!U8zD-O-EE`b@D6`qYONuP$v%?1+~^xpQWpO zTMFK%IjvY8LL%X+A*C^-df?208Kov&;jDGUe-f&lR0aS2ZbwU`0q1^`0lZVXkC za6Be^Mai~FI9yWMlcM940Z1B?vp0h{91=KnZfH}^*q|{>(xGOM?+RUf{djNpaI_%2+HN<{ea6m+m)*-ss=!{^*`9(Y0C>ZU86mx zorGZfWe6IN2N2Yxz<8->pHWIxb>WXR=)sIqIu81xETt|YC1_!Ka12!ys_TDCn}Hw{ zXD}!=*FnB4@8XC7u#XY!1|n7fKs<=TONIX32Zc@?MiKN{EH$lqPHKDZOi#lX%=0bb z#SRcsElP_>9g6wk;QMdAGdVISx47d%&P;cq$^}W45kYr5;5fO!A8^Lf?nKNkh*vBnRx)Qt|+J47&Z!H-B1VSV*P&x*N9#T>nQl(2! znir5#YG?_O?i@f$RHVCx8A(NymH|Y%^Si(A`IE(3uIAV4ls{v)e@QT&_Kux|V>CTgzA01^05V{z0w)61OPWjAvk z@-Xtg+Unp7Eiadt;NJ#yC3^#(OdX~kZ)aJvch7LhoXImjtRDCkB7GU-wJS6S*&3Zw zX2d;IeI-x4(&F^_lz4}l=%tZ&q4kY(5%~9a_J6HhV?>5*Q`3sfrV6``aS6n>j!}eUM_5V&(sPX?p=M^8Q`#_{v{UZ+q^;zEk!>XY+(Q=Ok*avB6=PX z$wu!$jUhMOOaMf0skatepuswY>T7s}U z)?%ZGppVcC`Q+>KXAWC@dI*!&T{KV2JJgUhb6EgkQXm%ehvA97-_}XoYf**pcijiYpazpg0HRQi* z{U}fNzK5kBtD`q#>)$VeDC2k_(n5L&rh?Czh^QX4@Eb2t;!!4)ZC>r*{2p0-Kt;M> zsB1%CcXww12h01taF9MVv4eJYpc$R=kC;dQF(gD}T9NMG$C*v8OZ~Yna99c;{%811 z{;7YJ-^#-A#nRc&Kbb{1^k+NwBx=mwtW5q`oB6kH998AmOMA(u>0R?TYti|CBO8*Q zx<}8xk6V~Yb}Qs)<3blz#$;*K4H1vmJsVQrrOV1*=E>5gf^xZIF!;G%5n3si#JCZH zCC#l46`utwMIQ0PEFX~BBQQ+wF+Q)*DBpMWm!*3OS*d(7w-4dcUzM;_r#aPPG##Bi z$W3AFFhm?hM@Of$e-0cEHAHhPH>ZOu9JPA(0+BGV0zU=)Pmw|3t`F8AiQa?~HJlA2 zek)b;T(7Z3_E#Ze&DI;(n4~kOrimzVx8An`l8DET$Fsyzcwb_CIqM#;0 z`>ohAM5UQc8Qz5^y_2Cy5Y_(ompqj6G}N|Hld1>#GP97XkuanJt!a@*4NycG!H+A@ z>QLx|Saeeyi0AAW88<8;rs$mcy+A61<4`Ce!6akjez{?yQ*C<9hhkRWp_%h1*$@v6 z-C4M$lChj2{ltOaT8qGjUdfiDdXmd~2g_V`6CH6Lq|8T0%`N_E`VTE%JknnsXkT{z zfMuSHQ@v|eCMZu~KQ%C|gvV;Nw0r_zE%3&2x^Dyd~WcQ&Dv-44R-s z34b~30(>eJQCl5^sZ_l)m7NjRDY_{A%EA4$U>=f8NSNVL#L<3-fa`>CzZ z^tI}e1u=sff=~A;A>?afHg|!7a~NL%ieaSVbk^zc{A%Q{=K=h`=0iBnDL?=^2!jS^ z8G$T>z;_xX2C4sbv-x7fS%>x=QkK2WKz_0+xgFt!_@K1H;-gz{_{Z6%7W7t%GoWN~ zSA{!CYk|ApJH+t3RrFa$cnx|=bo?;H5GM85Bm4X>8MlP5GZbpYrvdVoCqSl_-cdLE8Hjk#8&7Sv+=PD^YV2b+VptzB`4$m+{6P{Vl~0ySC6&w5 z-b@7Y1OR;DqGZ)QeZ(6AVuThFnTU?jb8(;yZUK2Y63|;Cndy+sfk=KvqG*$nJtKGI zW;vhGC;~wVF79<(A@Hh-`1Cc8X0A&@jX>5;D!=nYx$5kRf8~DM%b>4U#+98dBdnjd z=A#AX^j4~uJ`pF-)X^$l{W7b_=y)?OvCP4~_dBQQsoG-I_sbJRB$^ik(F8{bK|y6u zs9TPqK>FtV_vX3nwOp}0-{JLzON6iczhO29{Q^AmYcqSWH#TTDpC1$hjEtTgzyS@6 zg;%RjnKWhN5m8x560N>&InI8xwzNtx*oZ*ieqbI$xf`V_qWu)x8f^pGqXGpA>N066 zuvxlGNn)to6D8sdOl6+lsXYz5?Bi7?-jFsl8~lLIbo}N`B!JSf?bc`M=83O@YxD7C z-OJ>EY<%V#S^PQsTynepT z69?PBr7?F1U)#Xy0<2}*eYx7-IfZ?0n4zvMe)WE*dp$i!lV_{yhySrrePfo{!*IDb z4a%UR9M+OfQ937D+YGGf@eiZAmK1!p?~YV>r-kT>+%6W2Jvg=bAss5O!kW6;9PNvB zX+_iCZkP?)Y9L$n`>rdcyvK9n~(B3se=I3cZT$+coG(Mv&sjE+)Asx;{ z%?rem$Je?ER}SFsj#Q;EU?%mn`-S|2h>u^3{^$;RxfT>>}|sivHKJ)gU(K z$54BDXuJfOsuu=gM@aI7GQwW~x`IawKUkxQ4SoOyNyMNK>=7rT3*)tzM zn7Q*dbP;X7_d}m(UP#CzDExRNp8oW>QWE}z(PQO%yDk#lh~hk03ZSd@?>4{w&bjp0 zpZ&vw7vIC3c6_e3DNkm}d<~3~{rCe)J16t2ASVSaEZsuvnkLfF+=|W5U6(^6j`q(w z3&x$Z+6le78BU(tW=A*mi4foh+1&4sMj)Y0v&mYeeQMIV+eR*`dFX)jjQYID%y~49>tBh<8SR9~;Vwpx9U62i9 zyWcf`{$9<9B*KS;vnvL9mmle`T&wy3x*(P)ks2-GOivSsqUUsAy9=ePlp%ZUTY*Mn z`(W=gwWvU77DvdD(9i$YLUA^|d|}M4Y|CQWgU3xR`x6{z{wWtZ8;|ZpD=hz%Su&j4 z>wWo8s#Y7NH)%3s{8Z?yFaEi)#SK+NcHtT$Gn9YNO zR&%+-bLcUhXr2Q97`SFR5y6C7&Bi|6AzK*qhU5;zTqq>8BQ?|9Y4h8DvFkRAjx}ea zl@C2jg{*!QORWE1`oVU$ue8T@?_%Z$n^Y+MP`sFw?$o*IUH5+;Z;K{_2{MA3&izGa zUTA1W#D}||8PP+T7M?sZc^!xh{d8sof{F}gnS8We?*{_^2~SFiE7N1hP%IHp#Xz+Y zq%Nw zY=1^foyv~%eB?M$g2adNaq7%^{~}*Kr{C*Z<}{m8AbKg4du$wDn~s&1-5kA&g$t9! zJ9KB#*tNXTHi@h~vX{Fs4qc9if7xBO*7=G9o;33-mJVAre;Zl=?Id&>PIGO0TXd{x8+fBn<~?s~wc_45)0?s@!pSL< zBm`U^s_{Ocv>;_7=p|CSgiw({Z@OmdMPp!%4_~WfdVKk3W^8C(wZm!?->I)k$dUN*J*%x9}=I9@QWS zZQ|t_>m*fu07!!n(px|k^0)zwBI1Z|C{azL7|fMQr3<#Ll{}4 z;a90QUYKD%;oBF)`?!6J~pbM5>T34;83v>nxgqG6mTe20=McpJzJ-zvl!| zW$-`ws2tmk)PkzS)%9VK1TP5CduEE&0Yy zcD&%?sE*(UX`3Mjnj}b{wV#&CA1t`*ghp)~7_+>j)h3Ae8@v}EPV<6GwFdXm+FIFc zA+_M}cMCq)+N}c}-MQlM()qN$7^cblZ1UyN-_NbADU8t+4;DHe(446&t<&C9urgCw z=$}BI4<<@DztaNl*{GS?^@9s4ZN3z!47LDBwn={cdS}r9>l2*o# z9@{1DyI^k~aB3WiiJ4AgsQnb0)+P;T*e%F+_w)q917>5+{z)6$%O#NimX$pvEi>3J zTYv)u-Ly@ZG);^PJ@{pn%SQp;#FV?C4?W}|Sm)-U z6z=d_G8I4}Zz6@eki$W1-5OcG-Ez9nW{QlH1i#iNXvilJFA&KP5>^OGIk52QX&R2x zGkI=`RLVlN^E)KQ<)N+I)8%zZjZhn#Q@}ejbQlRJ(uG~;0A*9B(lQYDL=P9}lDW|6 zz7k;`=0Q%gy00e-ZNFf6@gJ3xC9&>yLff4X4SZgLX1YSK(945A56{<6 z1q{Z|rLj$p0R0vO&}^z0NyRh!m$g&84cU#z)f@&%VNAH;PB%V#HEpp5pSo zQ>+v7>CO8+t%s|wOT%|)?UU+~>_?gsfB$lM5dT6li?)`X`s(bhzH}6G>iL?qOr9{i z!}K6J$4Uf4C5Ml)Z(^FWV**d`B-fq|e8$5cmPX&LR4)7zv;&^!JD0YB@R1p1VFxPl zE+FTZJ9&zC?nO*TpokF&S-zqj>rh(u=vx9Lv8P?8+#o^%h8G&sY;#}|7YV|c%LOI}nrKsG(lUG;wg>=&*BZ=H! zKrEqb5#*mBnL4qf+9D=A1im@V(=O*yO7F--?yBwGL77 zygmH(Ih+HBwK7Ynf*XKUVwz&uc zwCAYEnBTcD@@_w~;_5=~Hlx;p3Mw{*aCZMDCQJeXv+(Dx(mq&kao=*cOeatpAH^bN z{5?zex+K`i0#0MCnH5@6M6DL*3g;5kW#Hw7y4RtdEoHOy7b2&m%YRVEYBWxSD1ZMb z{;u3DXI}k3y3o5NqGPCGZE95$QX)nq|KAgqQuJH^7IqBsFpgvDnFK}LRffYu0`@xo z7Z^@>xM*VwX}TUEQ~v|R5qUk4GQ}juZx6(us|&S(_3x-7z?PstYE;MZDRp#HF%{)}o2E1v)+b`s3Z)2h_Wx-%UDMVIDQE|2YaKPTXDa+ZC z1-h(6gVN!7xBu8QJNt_&=t)15T;XI*U$D04x8Ukbx;q%y16;nZ3DX@eE@hYk8*MCb z2D2j7Jg94Zlm_7mo)Wgl%ChI7=XmK#j&rhcys6(lc-qg9);pNllr^T*f zL_1SXUA|D0do{u!riz@w?eDj6?tX+{o|x*zroB2MrQKnD7uSh7ygIuc@l#0Zyk-be z4nD`0-D~>~D#-w18&*el;E*y?ky#NainST~VNR{o-DXCoKJv)BQ!@T^lE7}r4*l%Mnge z(lfqF&d5&hcX^C|LJFX4A#fh&k-H$2J$U3@6NZKb)t(e9+mt$1;W8J++B~d%(DohM z8o74no%Ut9bki@#sSijrLV_jJ{2-JTUq&~>1mB~G4>C3{9mw6r<_;7ZRH+HRsRG$$ zbK2^X__4c3-uMAL9DcqvQL4=yo}>KgWSgw;3j0%q=E^hQH?sM*3M~*!)14)UK&%^- zI-%?Y_6T0{kO|2=HLenMCU{F5F^NOpg+k*{I`^RNY7$_3UDJUC65DZeaUD!6nzEC4 z=}ZI~@8at~z0Zp9%LqAeH0O6ToUT?a{@&TG%6(Nh7=GFJ%6?$6mB%L*h;(H5SY@gJX9SAHg%?SLI=?JMAAItKVD#4%>%GQxEw}Q_Zt=#UsCkO#l(0SBBe+I z!pUajE|lr@-S72~<)n>WT@x~xnmgA=a1De62A+-W-X3gDRpOrMmkQb6w9R>By5P|Y zLd;^q4$K3utONoWL9#KUImn)!CWfz%5Rp;gRQd8%<3f?7#sXk)cd>A$+323$PouqD z$Dj{rOJ#|{{=PtQHVa6i_G+oi$xMxTCPQ~G_g9AQrHzJ&DF-0-{$)8JCl z+we;Z8!tY6FG~1h&3yhA)6B)35VtuEqezr%&;u3_7)1s)wWCf>CU3wM+%c?-r7Kj1 zDJouG9FGB-!~f>0b;Ru&_(h<~RD2t4*kH~I)Zp<4`j8>&FE1K8UT(}s<*M4`^ypDA z+k+~OKB8Ihbrmcii}x!1Rhpb67;NHF12pid_LeE)m6A9jE2ed*!h#JLFc zLWilF_!$DqYIukgr?3RS9${C;NxI*s);1|a<@wecBmZ;o|P&sYpb@kwaC)$L>2PMCEIc#n=%~nmd%o2J;{jI$~Js_a_A^JZXBn~0o>I__`#|T(QwnEAyzQwy32k> z`{)IeCIyawe4LO+*?W z^kNxLv`bbP=5x05yG?sB*~rLf(U4^mI*zQ9NKG*-8vS}z#T?gWmam6EI7(&%U${wui=}gu=v$Dpz;4Y47CVlC*Dpsg zy-Z9`@K1*O`RyGN#hQ8E^l2d@RljA{_We~JQ}2GB>6C#V%;M8c{{I71O}RMaj24Gx zDgom}ApFy;v*r5abdIR&XWlFrq=9j$(dr>0n(zgMDmkTZx%uYOQa3dzk@I^_VfSl5 zhj?gz`zj+_$r$NM_!i&ln~ZAzjNY0dse6xlg3^;2_>p}sOl3Wkq<-}?jrB-36L`dK zIh`uuo^_%7D+7S2b@+w7G=vt1y%9E;f`@A$QicGuX*eH-k1wNj=%%KL^=7(Tr)&bd z7mZK# zJD;6W?B{6wv1T?p?^_U#=&bs*i0&M=h%@+@Bu-NT1U<;1A*GO-oF2y;&#bNjBaV?n zBbD=$`P19Iq55ZCVM>Jnem^{D!gPO=CG!Jo0_`kQ1N_StYRY$&IE4Qg?m(-(EoMX8kX4rOr{*#cT3WTvRKC>L#v z7*~>^>&@0oP2Ya^EAF|w_+JhBdNfuRZKR{2}OpH4B!dH)gdp*e#1OIEK^ydf_5+dUrLn=5Ux@S{@hYkW#4fAA(Q z{Cz>C$!-07~$A0DbJntX`&IV|(Ywpyw1l0x=a57rD8+5|+tCg>>rGwh6X6Fxz)VsHu zw!TXgo0yfG%$G5yQbKA(UmzhM#&4-)`c5C7yC{a*%#~9Fr?>vy@uPV;IrCj`+K;wu zqVenK-@J!Y(q=GudDMEU?@NVzz&rrgc{K$Hv|r0DR&C|^Ad|+j7DUGqztg|@;n8f$ zQ`^EbatKFEe(-A(a3l`p@c*hol5n!1O|;s9n>yOw_Sir9_t=>uk(Q3mO$Zn5*jSU4<%*C% zAjrXWWe6r{p>7=-pa6*Zc8<*&w=qty?NRVMuz4g>q)u9Gl_C%ro+eFWpx@x*AfT7Q zV1x|ipyc4xq&(?U9aTdI%yQ4NSlE7KY3bclxr;yVZ7MEKfKL)gjxH`PKqD#6CM2Fa zo;U-*JHKQM2fseBxF}9gv8y|d0NU4NNjp!8rvOA|RM2sDw(5cjJQQ%M9~z4M70oa= zfMrM>z2RKHQYTqrxU(=K@^v(&Z~Uu&86!~VddT|v8l5Aj_$zYG88){s|_ zM;!|cCb#nr^?}4;xh{}CVkceYZ2S4pD?zozm0G$A@pt5*bn=I@V$f%nCsw=PTFgS* zgs0P$K%i7I@M{-@-`jE!@HSYu+c*Ygloy7dUye937zm;{IxbEE6_SjMvSBhM2HqR*bYF;ql>$;~l3{n~B-?^hTVXC;2AVHaNy3>{5FWsm6q_CG-&X z@fxRAV*yCN=WSnmWrKOd1k!aga4`YmvYv;h_1tHR;RhjCU1ys$Mk?h-wBf$N!8wIj zHZUArtf#!)1PVnSqgqDsx*q~IG1Nw^=>wAEVCd*S!fv%NRRo~`J-rB%!1OUlukO9@ zeKuPPlW(<@Z(o!BU@(ayfJ;O4sx+(UQKrdp>^Ht3L?MCgok-NZ+8j?pz|2A+dm0Cw%C+|E5Db?+>rik z>Za*xns=)ogrC*hndhsV6V5H=i5%@ix?5c;9`=m0{Ls8?!CP-(TD8K70LsquFsCI=ln=WEx`5_hfr= zwT;3wxUwZPB-1Hd!WqtWS25^#Yo_mey^*2e${c`I+$~x5Vso*~bV?Y2#zSQwY4=1P zYoWL^Cw(y(rg)BcVv>w(U0sB}ZcsMN!`YoNpp&FaBt$~v?4w^nYMOdZq8m@9}!i(Np<;_O+TVUnbi$46tmo6Q$}ek zT~c^^HL%xH8h$ktesysC*Ld~!E1XFlpKkZ9Y=|xVY|z$a-rjvaRI1*Yj`u7=DkIHd z$!A*)5jmjfd;=c`g@oNqcH9b|_wB}3R;CD6_RS?WM@tCUHO;0K3rSYPjAx{AhfkJt zX5^V?;93lgG*OdMJq^o0V0b%399)7XD^Z<7Vu~le_mf0%I0iE)-kecv0cBv~Vn#)j zOe@NTe5MCq9(3GZjBm;SB{Q}reKxnIGbO$3$&&7u1ekyR(d z0!8~N|Cd*hAm5bgU;3IG(pq@C2cQPH&RgB1H46IyznqSs_b$7!cu$5@HQDbWTz6AS zv)XkE;_keaFa}SkDae)40hPqiZBV9>z!cw#w85<(6T8F09OhLx7oCnydnGWCk9T@)rUIQ`{=&!~XJ~(ycCJ!}?iPNEp(ov72W(P~T#1WFn zawy|(JOB6I-(I>=GAS>WsA^aJHQnmeH>TOnE2j4^Ar6e7c@@!Zq(&S{-q6#CmXdy4 z1!V@ONeJJvc9DW|2vEm`q;QZV*=c`RK16`~_Mk^4JSX&|&GYuCGW_!5-%#K3p3IHC zgr(@-)m-4!%YRok`97y6;hw{s_D&A^wQIz5v9WYq@UwjvBP{l+ZRFfbJ~7>~a~LJW zPMsHQBEchEGc~_tT^H>a{YKD9x=2j%n>DTM6Q=3apI?~nD}LV>`TP>i3Yy{>s*V7= zQAS8<#c-=6*d&Hs4T+`P9961xDB}qJ(&H20*Q&MHOTx<%0#6c=+<%$$>d0nMDsb1t z(eL**VX%;6VO<8KL0Z4~0h~^Ik4peBQc^F7>9xSrmfvLHTa)fele_P2a;$dHDHsp8 z&`@f7>6(u8L(`nHQ5KDA?%CtcbMv1u`iK7mQn<6$mX>%-PmMwf zD%{8$X{_FC9BlaShNkh6njw4~vd)3BMw)}B>$F4As35D7&X(hdAv~T0~)u`=HjAvEqu_G=xiV` zUZfX>_yw#)vT}o@VxO-ej4erh$Z`!}hrA~A2qyLUH-%s5{Xj*zhe$c9vB6LK&epx3 z;*I4`MdQ^kEWLxJ!MYXN+9SB|o4JXJiN)LfKc5%F6k)jt6XL1+F2DSbKQGpVcC3DN z11vh04L8=!E$mQwc$<8~(t?g2AaoH3HhPx@n2Z7i-IXu{VX+ZwFbx zZ{xzll?w}oFU_4B<*oEPPW#|;ZlmJ5(!-`c_2Bf==ZXB?3MB2FieqsFT-h2kr^3nS zd^RW@Nj44#XO0?lH#oJT49d>t?UA7q-Cd^MPfX+ch<>atwTg4=t7in!r6Sf@jH~nd z9G4QdfBx{G-nstPGi$u$OwL=sQHj^6m75?2bLKe1yhdGVLt6m7Y&dax+DM``NR5#?{ehGFWrMu+&i8HTUP&(PjndgKv(5d$VZUyMlKyh&|uZ2JLS?4aL;Ws0N+>qv* z%kbNIbI;_$_CPRQVSZko>s8h4VX%^@+Y#jBnAh7UA&zV%4{&5UT=2j>0KxA)yZCp> zqlU%%8cYGxW({&n3V|P+1kO!7Me<|pkS}L?FYt!(jfwI{9`p_VF=>Yej+X}>>Qj-oXs_jZw{WH@3j~8N1dzM(4+>tVrkQWj)PhK7mmI^&Wjznjd zXJFCfxoxzu)@&XN}MAp?-Y~A z7a*HEK@3Wa(~KIH@Rk3h|A<@D0B6tNGKaEiJsNq>X4Gu}!GrbE&&OoiSrPSVIJ6 zzV4Gu-rwH-UwaXLlb)h-u>rs%@`{R_Sp1XBH5wbbP8_v748}gx7cgd3e*S34ps4Qy zBR*xvaCGo)LY_~?I5X}&F6td+IT^uPLbk-9mVKQ}UL^TedTj~) z67++IHO;M)(n|2)vjiKs3V=1Y*$4urY>!UiMc1@yms4+l8m$L?Xk7!QO?!LkM#?a@ zAu*eEz>N#%%LZU;vRVju2>Dq(K8F!pgn{(-{?(m}`V%$H^|JTCdFi{N(7DDYeQN^pCnottai$pSK|&Uk25}B$mE^bxev2zt&Nxq z{L)cHG_ryw#B&TQg8>4$ykh8!#Peo2u4y~`WQV)!Y$YMm=32pEN^m7Bdx*#(5!^V- z!0yZy1x3!%rNAAi|~}RV@YfGI)K)YYUhy1AmMBH z5(d);>FeAprZ-R`j_=)n_6TGGR3qyXMDBITg8+NyWRp(2G&MPyKs1m*73xs;l&=;I zdxexSN+yf!2?Tz%mIH7ALk!xYHXYOQjTXDY!_Nlhfh5)Q-P5Ic&UQz$zg<3O+_|xF z_UDFQ0N_uEuN+8?$K=GF7BCNHRfo%PrP$F9oG<>@ z)KMM^dQRqivv|FB1YCBI*Hz7ceb3zO$SajKPt%#0V5Ge4pu7fVyOa^Cx2|HlZd=92 z2BnOl+`7l}d>JI<19H0KL%`E+5fBndjE;GHH{=evnU}B(mKQ`rVd#SzRYp{N#f=OZ zYyOEf)13HB9RWomQyxPJoiLc4Vw|r;m+?&3#gNPG`r=itP3ZO6|C;&L@a2y1ae#1= zFF1S%cgW?_&07aaGw}&w%~_xf0d{wgGQZ3VO&9BMeHX;>ChG*EP19_LHrHdfqip(> zL77gixI~HBw-Pbzh<^&B9zaCTd{BPOMua!)h|7vAse~ONa9Fr5ezc6+gn@U&LO>#oG3lbyiDcdJ=*3{~R z`jHqB!9fACMQXJsqLyU?Nd>iXp9=;*8lFtAZ60C*DK5S+7AU?JZ~WZ&K0&@crlfDkX95yJwuW(pCE zYC@oO)B3#u?zfcn82O*P2>Imfpq@M@%bYi;_6s;}&H}3L7&${A% zRuc>=KmN-PC6S9gwb%^&Omk1(STkyD@7PiNW+El*qlIYa6aB7>t(vt<%QU{oK>9}j zwEpHPJN4?L)mlPO2Iwf5c{tbUhG?VZM;?Qt4@;(!Wq~R!| zM=iw>$cqa`7UM%wRmT5(`{9CjITW)gxVrehIhUQ+$xTffApQ}N0UK2s5@?DU)(F`` zp+vO@y}WJP-b1Sgtt1Cwtn(-orLIl?D?7g3nufUr6fN3Vm6`Z1?2;t6m5wqzG>90jICT7Xpg}zrbMzp%3p&s|8 zObiFv=*BRXzZgk>#;CJagNxL04iRtJB?s+Ec}nUN-{A^7ZqJPhRf3Pq-GHj^w(x8^ z>oMq=6tqU#I07uiK1sloXXQ*=;1(^)uSudxQd9F^EV*%1)g{mO_c!fHF;fvar=JmF z{Lr^G%6d;;NjBsyHPZfsF@o-QzC^Jzkl9hF7S}y8B&O^r*GJ}KY4tx605CEwoa>*~k@Bffg}KqVZbn7#4dTWXd`wg=OR_ZTwv%ZU*k$~_O^Oi>e-gBY-d_C8*KwqGzoa0Tl zCb6b2%ssIYvXa+>L2mI%^@)w&!9a#2+{Vu!vLS9q^x%7c=>foX#rCe`;9Ch0!rd{{ zKkQ>-+;J11*Yazgzd;A+t&s}190P%c%h#U)+%u1hi6exGH{L8OJJxaGRrtZhKLB;9 zF|P;U(2Z@UoQjo$t-oq*1}02$hH1p5;Dt^aw;nqVjo0_pO#~G5n_5^9hvyIpWuW9a zi~>?3_aizD<(LSE8^9J7tkp|6v;%vlperbFFQ>QSTt1ECH znqQVaf_S2b(A1M=ZvAQTQ|dXD@Ew*MpL-7|l1S(Ok_~m)*hzkoXV%jO<48(QNtueKaz;9T@G$J@9L2N_J|9!Q7R<9+p~JMG>1G|Oq7gNZ7xVR$gA_hW;9vx0mtE6(GOx$Jf!0XV@t_W=m&P!6BzL;@}H<^O!Rns3=#d} zS?41?f(w_2=DRo?_Uf-xI__rse8@;S4As1ut$3?Y_-Yq_)NNHGfIST1U^ns_&WHL4 z?@=a0HCp8|+vkWywl18MGTe&LbbqqenZ|G2m^5#dk+xuGuWbDwM{160y`i=yRFlcZ7)=&=@!f>}qkM7uSK7Oh5_2kTSZ1saEaI zj$hBlhH1lZ_FQl`4^M6LuG+#+jc)f&uIN)%OK+j4=W~wo5LzR7GWvc28%!dvaf4|i zshne@mVf@4I=NJO1%MO+$d9~y(PsajpqFtpvyx~ellseNKMVgRJ3IS(lHc)!Ui4i;yujLnjwy0L9#33Wa?cuh z+;gNiMrWE}h!BpKw2T_vY5+;g;J6uYSID*mZqEiDrdmd6IBfj-((=0__T>*n(qHeV zf~>bZ-%+vsdJ9Zm_Schq4M&5K@6sg=8QfoNu@qbAGimDU11%5(m>x^tg0!g5lqnRc zy5vRcR;@hLpwAqdoSMs2`WYPgVJ(t{!$>I1GJFz-$NM6cLft@tGT7F{jOoWr zV(VC#9NixZ64h_+QvRS}))BUzg_m_Qy7ndHZos}`q%$N&D>{H53rl`8I(~Q_k}CG* z;HYi!ddD}p0HA!NpCORxrjyb%Bjf($=|%&uAL{NAyKkCzO4CpD^z@vc0~F2w#gniL z0DumA{>SrwOt1!cPqpzVGJEW9b5#{{qM7TPS;saIkudcE+^1^_FUmy3M|(n!7w6ghPqdTmSg_L!3D5437U zbAIvs>N?Qgb#Eh6V575g{u2!dA>Q*{1cL#mx5!F}lWsy0h9K1sWoQml+LaR_6s1b? zDM`NFI!cYW(ulS+;OEG zd{+I3hjeC=lYqH%dDz7p5}tOhGYPaS*`bF{H{<4@P;%`FR^c&CbB84Q{`% ze9jorGZ-lLtUvNjgdKz4N*u~O8XfahM;}AY^6rn_T}`9x$R5kr?LWndp=pnXH2G^w zw6xrppZ;mTKL)iTtVZFAAW+f)C+Ozk``K4r^!p2nkDz+BcI!ZY^dpNbrc@{%g`6p; z+P!YP>s$acLE4jbYsFY!es{|Y&XbeiNW18!W8PKLNuuJ)(I$PNYzZx%~8+m_ExDX4}m9As3a zG4Upu=$6`kC&LqwW)Y4cAzzMipjT~VMy*u?i={fN z;rkB{cW>SHYizEClR8h2HXr|8`N4-*7DRgx>LCh5L3ayOg??!4NqE*dRHGGcm-aR` zOG`bT>6`9R(vfSJ_#atD$9>{NCfO<-EtA77`vmLf%^eJ%6L9CbOnq3hhAa4I;qO zeIYsCjw-q6jigQo4q9U=>m=&_BH|9w>B7U*4RwHFHT4<9r$Z3~!Zb;v5Lnnq zNKOS2C)h+NI#1Dl-Ag%?+{>WjE|kv!5BUQ#oy!7t)%(dn$#)f0R8(-3I2*{2z%(69 zS(zvJ`r>$Qj-4)j*$W^~7Xm=}etLKqKqT1)@g%RDhdhIHatOTDdzXWnGp{5k!!1_y zSRlG3=Vizr0MjvPj?j~}W~2WDgh6}0KmF90-@SVE3X`{?$`vz%!Z}fis)o~cLn;fE z?QmI}tnn>f4K-@JQL44(Wxe~Lg?0P$)(ombr=@~^5tJs1;yX>hV9=67dN(hR#L-|V z0*OePE=HpffIQEw&qK#WZ`BbE*2*zCv$!rD#u+~nBD{igTZh( zOcI=*7h$@gB6JxA9dZ`5^mndQs%mtsCiP!$pDzt{$8fbe%&-g|z*>tGKrzM&;k6b7 z(AR7*MmUFX`Ux{92_8Rw@adh zL-xtF?`avTQ#UuYHRF|LHX2T$O#&cMCK6TvgFGjq#0!A2pV;~0(e&M7vA~~x<$wQA z4;Lq~#LlJ6JX1$ZN+iI!&j9%Ivm{BFIk>V^9CJiOip)D7M zQ7&?np#W*vHbIL8K|yeUsuu!OJq6d9YH8fA*$h|#rDD8=(T#{8IaCv!6H%diQ33$2 zW`QVuU8>T*Xu~S?a-c9OLpo2oeGvWjPTI&wu~<=b!&CzdwC+ z{N*tH*C(gnu>8{?9<#&%2*>1D^$-9tF|k4bQl0}4h!Qa!C{Z}+wX7J-6iPD_6-iSy z_Cuu1EdGcM1MsRSBWx?a-eGR7<;AY}v$a-dYUR-5Vn2$M*6Sd|1R*TUgE)kQB?Tcu zh?5}l{2l=KL;=PDm=SPTg6=*tfQcT!GbtjDVWBER{TjrN22+bK|C@jL%Rf8*v;X6N z|F^5#)i{BcF%Y!$_}bbQuSU|ay|~&@1YD9!r}% z)jR^C%b?2;y=c5il4LZ3jS(?2v6Lc-jS*zo(t|@~f?~kB0yCL3q%LVl;ehu!Lt{cZ zk;jMp%b))E{nu|!rsqd-Vhl}5R8@!o$sm@k@!eC?rrQNa@ZA|Pee@dBQzoZ1Db8 zy@EnSL|w(ODddUPQLr$X7HKo88V;kb6lxK1PE4syP+933A?ev#AkDJuJWq`=<222( zER^vpBchp^XWx8tdU_gO5f^%EthJNL>8GFm#a}%A%fI}~SFc_@eE2ZW^TA*srL_qFqNqLO21t#(u2RSy z6=j-;oYOdtRgvep86*b>=EaK_VS}JXu=_Jo`9W4Ee!RHsQGv|#LYYFAaZcJi~ z33VEsUl^|%>Km6;hnQKbh$LoGFtb--=5REE7CD;1*IRBp1&aqzx{XCu8F z$#5__@aGE~Ckea=B%iA)Ywo>YEIhnh7>M^|jjEE^AfYIk2R5I4`+FYnCqMYXFAp9? z5075I`XW0!ohJ;yBPOCnG_pi)k#jCO2B~L|A#%AhQB+be2B^a$r3~e8noxvXM691! zYya%C$xlB$e?0uhSFc{lBk`NK&2mlve;V^0oKw&{pHBgKglFn^3d2uxOU;SsF&A$Ho|KV@`ZV-L8 zpaW*HHdN`#01ze2T>2*K^`b8Wqht^p3oW(@GZ(phAWLLO48e%AY>fh{!B`OL8`DZ% zr(iM&TJVxt75&?o8aFB^E${c(YWrjCmUN8E-Kq*tU?L?}Aft$g;CxP+1I+M%6Rs); zn8+lhjjc@1P#=&%iv|Ipbb+V}sSaW@pF2{O46(&$pMUb|^|$Be=jvUISVVH-$XMl@ zFx2;24sPRk#7lhi182aSiz-VCO!-z&=HfVBEEaFxym|fl^`q%@liYvQ>>kTzOB;oV zC?s&^@T-U@RVoV%ac~^VYU%0r(pHUE#=4Gf(CVK+c+%H&M=i4rR2v(wxF({qGo?wY z^&#(lD2>Ug#u!s(>p~zIN({5B*|Hx=ehG-79#ec>ZYm1lydDTCyJ6h&%gNM zi!96J(;tsUqeqV(y?y&Obg6EH2r;g35dbF3%o2%+1VXT8qlyR>e&u2+k|g=~rTBhOGHIkL50oZix5u@WkO3Ft&%DnfWV zo;1X$VQ;;kOQVjiJB^Bhe^rQ$N%NeUKl$X7pZ@fx|DWIgZa$wcqJ_1V218>Eae>bw zz1cdAqCFSVyDEtwphuPS8dj#E4X=S}ydt7C=gzJwuk8>*59`XP_kOWRjWJ}c_g+y1 zbyh&-(NCxajnp?+^*j~P(cW6?9K6S)M}ue2p1sTS`Fy^ji>rvp_Ic7fg9gpK+%k0{ zSY>0leCJB|z!@-FJUEEfL9|s`o-7oy7$bEeGR`w&%tHu){^H%_4Z`79doq z@XisDH%wFpWCH*U)s{>$sB$5@`iB2FjJmp*BqAwgM08hg@@4+=&UJ$=ftNJ3Bif3Zs)U3P@SK_snrH2t-jd8YD@QkZ78w z66^W-c?=WBaawxTjWH2LAs?yA#|`%P;$T;{%&g8?Yb!NB^s)AD(P@-yDU5wI96f#d z6gWLQJHyFa5jh?X;yA8~?ly)T{EPMVrP2{jM8USaT1>fbZo9)n>JUg|SCZbE<}Zeg zc1NiPuLyaCy=3Ve3H(CP?4-+cI~YV1h3dhdeDcXJe({Sx{_&5q*({2pL6Qyz1FYH% z*7jvn?px13XRdfjRk?em%{6<6$_HdM+297tKW_xpDM`VByQ&8li&E;c9j{kJFc1D9yQ|hz|gW4pqdF-=+S4#5O7-KLYb)%yZ5mC@YlVJ$6^32JA#5tRqfgPB55Ji!bd}iEJ zmQzh%c_Jar57`))3;>J5IS&%33b4 zp`A6+v_M=31bQFnAFx?5WB{wVr|pw#Ydl1n=LUw8iJ3+KNSKI%0h65pASMF*O0E#SOwM1lONJMi(V>1SrBMZ&rIQqqpKL5=}PyYC)86k4BQy&vWnPJY!8)I@= zX8lT82`kLiMnf)PmM${y4v|=5B0(AY7ku$W1iS`hsO=1LAPWBbc7${Nj{0@x_szkS z6{(u>alLQUA8VjXQ86VaVdCI{F)9FDU>*tpm8@RHP=|4q`M;#$bR$)gCEE@ylQC&7 zWOy_lq?wM#^!WJ5sJL`77z`pdq#)slW2YH_>#mTsMhUJza7Ff67d>~yhQD(RcsGUX$0YolRRmv&&XzLzURjsx`aBYTPgQDEtqO_Wo>zoVR8ZAk@F6xUc zqrF6gNxb*uLq4(8JCAi+w|n5><@#2>1d`RAK*3`=o6Vy6JQTAOk$gG{DsmJ>Lky!R ziV4P;c^osd0eSDMI%QDF_ zV+=V*L`A`q9;nO@Q4}2?Cl8MwPX9DD#>9g`@VQ!Rks;4>L`~27<>>DUooyA_Xv}=o z^IO*{g2GQNADC7V;N_o<7lZvE&+~xi-noxI{`hA<``KT7>uyVhYGYT1u2_~fczh*bVfjPP zs|K~>$!Y;Xsx`w1miMKVnHR2YgHn}WF7h3Dne~{s(9XbnKc=bo{`>j2qtWPl|MYj? ze*5iHpacPV zEk>x2fY<<%D5^)bs0KZhg7V8@=ZG2r6j-}3k*KmE8%}0bC3QwEfjkX}NMQkpl+nvT zC<7w31ItK*vDAjt5iJW$2AwHvSBIr**%%?BNJ}kB;W+|89P}X&k&l^~rm;xmlXPw% z{!ibI4iCr4z^9oKOk~aYKvhYtCz{U5Id_y#MFg1|06NI?{5;R2C^}0=X*xU}&l6)~ zGn<^hI`@yBJbChC^8(*Y4jnP`>)5NR$kbZvlt7To<$e)~nlvSmcrb9I5fq-naTlfE z9mwM|0K@5lwRUv!`m@hI`(OY3o3l56KhFOjI-e-Ljw3^cGnW%B5F_#n07Nd-OEOHA z;TZtK5fL;alx2Z}Mxmj;RR~wDNA}kK4X1B~0H2cKh8A>b^0fd|W0di+M3WlZtH)<4 zN)I%2RNa*o&8c1|A$-zRs~>g>f*GXo*dNO-}~a-pZ@i~{?|tb|J&)eFDLnfPd@o1Q_sX4V$HPF(NNd` zxD3{#RenjjQd+f|gSlJG-C{ro*86yo#c_Pf^8aV=&zdDkvhzUjJLjm{;_i2eh#Qf6 ztpKV3qDgjB$Pt--;KPW_Pnh|C^VBjTGs1qFX_4+hqX~j25(I%-P|VEAJ!8LdS6|Fb zRnM6RH8Z!nMMh*sRAyB{A3{9*nwzQVcDC<)Cs;^g5X!bM8Iy6&x6_{U*`Wt*^s_MH zu4ova+}sW7*G*^RO$yZ5=gagdJoVg6e=rfNldt>ZmXwIx1TT{JbLOT=gZ>dac**NF zzgmCPG9lSYL<%coN{JXm$tWe|zD5?xx=OHhi?R|mym7m(UGezQ;laTHTp43@zFMtX znLK~`ny*TpaR!vBZQ}|36RyE3n%OkJBk#DL4{=hB_*ymWl6iAU`MuM<1WJAhP7h>T19E()CBgB^Bqw**6K>G4ogx*UVCk%*<_D zmZkH4HkFEE?U4h{}Pj%*|`@a&w+L(~82 z%DFK$`fnV%uTpNeE97^@zJA*}aeSJ(BM z(KhGjI5~F9^YwC-YxDwu8q}PWMOjDJrKaI7sA@JSqghJHaV`$u*J%(DFpa#-dFNlu zFS=nNLi^m%{#dr0{J=Sx)szi^7(mCBnJpqQGm|GGH4i}ZJuI7uz;+4wmH+FLx|p;`Y#}tuSW4w?0}=@d)`YHY zTkbZ@jIK4ab?TgRiP_9rQYu6i^SD^d_jky9f8MOl&&y@vpz|)w+HT?9Oqr9BDyLKs zwWcvfawLGMK@%vZu0_`-VT&C*$K*8S3gyhE0DN=^R$;R~U(Xloci${Mov=O&;Yc-> zluFmf@`FnNL?>sEWjiTr`8Oyo@Y@OvBbPAsyk^_0o+k|QowrZX99O%JMt|G23*gVG zwsie7QeXHX$Dy~N1MaaATeECjMKA+TPfRB2v9fv1s&1TfBr76f7FAVB04SY-8P2MO z{^=jR`^g78_kVe~x3?!DHYca9+LM$r*0^Ue>&?21mbmPjw!uYq`I{UJvbbs2rYJ^C z20>N3)R8Ou=}<%OU)qlN6gWp0#lj`gd#%j4Q5(DV+l!{*R0gLTIKXe>l)l91GN|wU zbJN6-7*ryYi>1SIrIdN?0JK0%plC47x6V6%c=omr8s!Jyd-vC$e0q9%7)7e=LMqN<-V6l*8|jk_IqUa|yJ_04F2T{m31DDsxL|cgzti~Nw8Lf=L)dIK(UwG%Qp)k& z@rkUOQp%MBatKhk5|LtWcRrsNbxi}UasF+LvFo~3x5{OhFGbl$!9zq@2SR4Z8lvn| zZ$$pDbIzF~qWOHjSS*^RS+CbVB_bM0&44YYG{4>Sq`NRdmu%vxcj`iv;li9xuDRy> zad=<%Ve=v}&b+EeYk--DN=+9=)rQP$21o=TX3m-7nvhaLN<{Si?|=WTx89OpuK=`d zdvrHs?8?;}q|I@zZ)^%Mm)fw*wqdVVH6zml*xBZht@mIGdVTnU(|q8wP+*0HD-k@!gNrYUmHn&hg;k=JA&^vob+MB6CV9W^n=| zg0XN_v-2Ed3US0avt&j+O^yQ_kqQ_vkTKJcv)dcNqx17>KeKz8dn4FJp-WH@jF^Z_ z3Su@Y4OluaBE*%dCa##-)pIF|ma_n^M1%{UB?o8vM`q7`S%}faO&S51r5NLSz3$pJ zHXBv-Y)KM#x#*YzU=3g>62N3>Jp1x~Ev%E2D8QN9=*ld{w)5Uo>Ku|eCUb@v5s@#P zS!+A+KJDz3WyyCB%Ch8Y-T`$yYnrA#UY%jR?+UKCN&NI=^_!-BvN~PQYa$I$A*v!$ zB89R@=E{ySDY3R)V&N3Yd>%xzL?D5rl!~-*&T)I>yY;)b&hO11@!9*fe4z5kIp>Pt z3Xqm&2B$d=4glt0OypsJ`Kxy2RwSE4Z~GdI%`ntx%68in3!U5bTbRG_dF5>?QNFfxxzxAm-KPKD-g)vTn%tr5w&>{IRyL%R;I2oWWoX>Ye#GGelIH7^zE8opc%W(rlu%SkCMm&@aG+iW&A#0&GJyIl7_4vE5#nqA_S#C~YOMF$kx_7oQ=1I$btcm}lzfx7fpc^Xwer0>U1$@&jZ+1a%Z+pQMeih|} zM&Wv)CtGN}UG>CuUFQZGjeZT^8~`OE?@1L02M7E6`>U0Xr_0C?*Nz~XGPv$OL@>NcmFz0aTQeg335IxEYP9RY9#pM=$za43!O!`GmM1-p)quGgVyHN;Ss%=~oOo}HcDr=5oFv;>iLoJ3F>#}K%v`spP3H0npJ&zU9z>Wlh z&rbRgfQHFoI&e$nD|7L&90}HO9gcJU0_!(N=(lV5|D##>5`pWxGW$M-MnYnG-lu{pr9!NP#i@!DE-MI*z!vgI{#9)^$Cb%@((Ad+*(B zpP7r}1=W-RD>CO9J zS$@au{L*EN&NZuUtMS}WmHBtE@kNg1OUrm%0~rvDlrZkT9s+4*7Gu0Uzn5eGIRT+v zwJ9aJ=&l%DM=#s56`oQO>3fp2_h?09bnhFb{cID}Y7U5Yc<-^pEb`-TB}*|M>Ht#V>#5 z!^XQfE3G+K5pjhhB6tR1%rd#frCE$fnqiraS%4i3s&KbXyC<$O;>0@o>}V@+mO zD5aEQ*9C;IS*8RLHya=G@efufmZduz`x0s^~ov2!;rPLGa` z!g>whJSqUhmWXB+74QIoVxBBDm=@e{!?=Nwm2uH+|M-A|g^2 z;gTIG+GdTNxtX(p^(-%yTB_B_axZpo&g*^u=yZL$Tz#HWa?C=M)BzBxh{)xVk%AF} zLsr88@&GYsPI1XrFAKm75nCm$U*;f$A?I=s<6eWvy9Y&N^Xhj`)Wm!Rt@8RGZLwWQByY~?n}oLK>ZTEqs;aW3yTuRZpy!|~ zFf%5B-b+I~9%KV<`5H9Dv@_eT_W$9nx7lpG_x0`lot>S%ciz2o z=T0ncd+*7YX4dHzYLQ{X`OeO{Y_G%^i=rSR=Om?+Q)2QfWD?D6Tqd;`dxJRNBnQRo zogEQbot&)KYX_H7!W5`t(_?)*8m1~DM18z=oTD73bvp+Z z9v?SHM@Q$4h7guZX695Em(<%!=GUd!&qe6+=h1WDo0&DXa)aZ^O} zbJ}-S6SoPE)pNE?zVebaHbFlCxdQa%JKv}=XU_x?&FAy7EVWe;K|qX=hN4ejW6&;5 z;_#HdduD;QWyEfL)8Br((Q^34ZJf1>x7`Y$K6BF4Img6Al;VZ;Ocw(yS9=#@-c4K- z{@~!?{{8!elv28h(d@OhtE-N_D&+TVwQolNe_O*rL=FZ(nTV)JT}tUr>WZTHA<6l} z$ETlum<}I0v(mYGu>-)@+3L*#CQohFh(Tmb23XRBp^~U%>&Qqom+^&w2qerjQP-wE z$^6X}vp_^ zS+jmHpELVM2MZ!PElU91L~)4|WtgXoX#^vbr613Y-nWS~s=d|R&vr^LiL7Odute%e z9oWMdRsaMF0Hcp<`u;h&o|Wa;{HQ{C7%t`MzT|>`fw)Yct@|QOOq0*H_5BQO(BgfQhkwQeK9%g-AVM-t(ItNHl8w_X+rn&~uro;JSarcc| zi(RBAXCz?hx)9MdnkwN*9#{4X*p{vO#fP+QTh*RI2YvX1QnVCJt}K2Uz>kw;{ol5P z-W~-nc}_k@4ZM&!*2jIfe;tn_@4a)9VmJA|*IcK{?so9ih|s03D2fzN75Ls8Z@mBW zpYQC86cbv*e`?pYZyAI&bwt1NVDQy;$xfxAIp@yKo;`W;#GRfpb5&K|dzXWtCK##U zGVc`u7m3ckGC;-5Hn9swpU>3ecj;ddmt&vI3>M3-Y@q8pWrftx7n-`vEig>$^%na;yc6N%Q$jdCx zLUz6Il>_{y4vf0gbQz4ZJf;HYJ#}#nTyRQ#H;mzJduiE!^J2Rcvma+<9FP&_;{8yQ z`#MM6wd;A5=*+l%`*wDfeOs;tx{j;-ITz;(n0ZP|`P*sVngG5!H^Xbc7BM_dK+$;s ztydzt8yYkFv-6ePvKS>-CzvnMvXpqoiPD+_MI$YBU5ja%@6IK#b~;5S0KZ!IGJAmQsYwzRAp7 zvU3iOl0wai_erd6TO=r0s=8oF={&S;*KWq+PmBQ++Fka^CnaX6Qc_1G!X(1XD@Q=; z3Ppi=#emK=vFE4F@!_LSfAQeSYPI^*Ih4i8ysoNhQ!boyLA~ixRij4d zT=q3YQtyRIao)rEL>$PB;lxzc7$g+HfvB44VwZ)_V8`Cpq!AQBAoD##FbKTFNIGEx z&|m>OfIuBUFb@!A`&bL0hreI9|EHV$`RCL$mGZSKsh3)>CQiA0a|YeyrTU%QXCkpB z0JjA|GF62;A&B9Wh=?LX$r_lM1pwkVfZxAe)N^avv!pePs>Stsy|c69r=H@W1LKK_ zt9`f<<;e%@YPAXhIg~?vN|b6}H06TlsmDCg7y#NvOn;}_Rl(VVu{ zWYM7v%rI8lkdpwaT04S8;evp7->hrT>+|D64ApOY#p2@HxG4zo9ch0+X-K0+oKiBk zb^HQ|F+P6$c(q#HB67|FA?Ga=MG?EktjFtKx?%6~MZY^~5(mwe==j|6Cpi*HIX`l% zHfq)urxx(2lVUgog}E%tyLa!t@x~k9``-7kw}-Ms2pPxQs;a8R?L#wddSN%rYJ|=V*$P=}V2RkJ=Wlsw!D5n7NDHXuDjBm|Av#13+)5=0Ok< z>ur50ilSX_%CaP)lu}*S%)EZMzJLGzlPAXjy1EvTtzmQ}>-5_H68G+n9N)!cPfib% zd=l|Mc80H|0zYSy|IS&!c9AXh+6%5qq~RyM9Z5tYoUdf&WEX#bVRt0&z4tydoslDE z&Y!RKU|tf*&sW^@w(7brgkJO#)2RMlJJB?_pufZC_BU-;=Xbh605BtcH8WV>L_}lA z9)2^FCHD9CkI$P|+2faNb}`N5TkOg6+aHhs#vp^4Igy#Eq~N{ZJzF(RGdq0BO*8A- zQj>v9$(%KzVO3+YWM(eoHoK6z zR2s9fOQs6MYQgs;=wB?YrvE zLI`oQVNGCYrvNq2|zYSk^yi|*r+oi5?7V8#p2MJGB?cj}8V2A~|<^oN+<=kh%P04W&}L~3Foa*SZMA#-dD11xTRq`8JPaco=1 zQWYhsAt{j;bD`ZpBWdF25J0@hI6#E4X_~g{h&%v-62TE8X*5*=gj7uoWS}1H!PA$` z68W{tXVMlC@2a7_blHD8Q{R8nK*g?vgFYkdtQ-4 zLLh_?cNXrwciuSq`G>2M6TEqQ`m9T4$?Fh}<2|pv`PJ&q)Bfbqo~kA{s%WpM^WTDh z;2XBVDfKm9el#>k4~(j2`rn+epKqQ+hE3DhP*ew3Mtn*sHNErr-|3api!b&6D?KFH zJ}Zi1HkID}-wtzFGGp@|hbRsUz$ah(5{wD3F<@ycVrGs^JO7{Tin>TDm z>1Qv+c&q78zD6gnDhFWzN#r(Kz|TLd{}9@>FKz~rt@|y1sW&8uPH&ZxZZ;cLEEWqQ z%I>u3gI<$2*e30!zrK#=e)Hz+N~y``kC9$f497NV!C4M$0TIvDIHs+i?@3|#3)e$3 zd~xDf0Dw%AGd-01BcPAvOn+hkNR)_fwG@K=s9QPba_?W_Dj2)-#N?W zxnsbQ^F?A#8oLB#HZw7RITuM{LQ%}fjG>|35Rrh$l+=hoT}GcqHiN3A#D>Zjg=2=E z`TctLUO|P?v(uxavMgs3XR}#~OCW2{unD&IwA7*~A|u9l?|bj;{@@3o>hQq>K00fU zkFi|cx_58cHGlt)zx?-KJUEW_ur8T7MJFOI5_hKS)*y{%?sIRE91)QwWZ*6-MaLk7 z#jIQ`x>-$yFW7;o$c2#97#tW07GncMSke?onUpyN4>el=g`pr~g8)QWs3MnGC5m7s zjUXmoGc$Ca|Bwuz001=B10K=E*VU}WfbYFwUp)W&0ckv~m^7^Q?T;s%SyllUinQn2 zO8M!15e_p0_ul}>MKa9HV8k5}k%5RZJm(^h7eQ6rQdJ`b79#p79H`aMMJGo9jFOqD zKQ=Q`24JZzilR$R=vIZ_VTzm0lil52F|n|Q1@1|v`={kB+pweH`UjAAOkf>LPQNzvFD7r(`YGaCz{98EMrRyV3-XF zhYVojUDtK{3vOevI>kT#v!8tQ!Oy?^@{@As&jECrymP@3Gk4ahYTb`Po~|uiE02#& z{6@p6e*Ieydf#)0CBrsyU!Utw$=6(K^-e(smz0SM3g1m0uc3#RUgQtR1wA+!vf|&* zT}r8K+iZB58IU8S#sdp<{$Vy8FL;l?YAar{hM7e~bYcP1mFut7R*0x9%l-ZR{X2Jd z_x75yBtq)2TCG;+=S|aS7XVPo!rnp%1MwlsvJCB}2Q+2_HpUBwX08D>SVK7p*vuSD zN-3t!25_00rj%H>G1QRZX)%h(?(S~ab(_K99^i!6jXzP>b$(tS92^|ny9efH4<3I0 z`R9E8jEJhLQq{+gAK$-!f4yFpWjVD`#w6w#Bi!Zq;H8+BbB@chEKBl!OadOjpIhY4 zxUZ(V-S4pljA5X%^VzQWB)_sv{h2rYbv&VLsETt{=8P~1aZW=Y-QNOq4VR9&uh&LL z*O&@C?dMG|PIUc=|NGbavuN;!n3)<>d)HAW?xg{?=A29M#(iOCs(tmOs^Yibe*5Xu zr(FnDRe9-4mgSihk#8w7^J>~FJRdFy;Me3Ic(v^Xx8)^IBZd<;f9r^qxG;i&HjS;; zEC5K&h|qz|%vl7W>O_Po=Rr=+=)Rgwt_vhv6Py%;J~btGgULC^u3+azi#&@cX720z zu^15s02hEj4nSZhBGullor8OegFD#Ub5i2o7B1#%V2pol(tD-2Y>SVnF zNJCUrjmZ>DOiV&h%@}U5Awxn{6>0=PBg69dDO*UE*ou%iD`x3cC5?!J{ufSsx$Iqm@-PFPvyNTb`eupS7xG zS&nFt?1aeW1!Ih?GK$sSX_+JL5F!Ezi+8>%i`sn2%;V&nS%!mw*vy!$@M0<{63hvg zu*x(l&XT2Edxziw%6Bl$4|NkjGG7VhY>)NpXz;(yHb`Fs45S0z=y|^Y&_}X+y!VY~ zXI_c`jG!>i){;5GQ239+R{I(Aq(!nG+;xp=6i0?c7RUl=x{{`QLz{AhO`4G?7{aS& zLj)0GRiR!H(=2ZY1_+Ezud_~n$J$HD21Q1aX1$bEkqF5$x15MenuEPIPHv7V=+wzV zL(Jd_Qf+)-X;G_jerUCYc#kW$KCFut(ndY1+ONt37~ zDYE%B;f`izlFJ*rhw{?ozhsvzlAM%W?Cl*K9MrqJC`%&5n3|^97Ahvq>S~Np6Lnq3 zF;h5SfD*(QH6{QMVrGJgh-Chwq6__xJbRsv)AbZI{dC+4)%xHfIfE z9HqRpQeU??jS;4-#ZQ5q&bhKIxh$P?X<|{SUZ@;AorsYodzw!A`r0K2$!mSqZ>^~+ zJ5DK0SUbMLmQ+;`W%JhtK=JsnfO3lT(wCl!AFz}q-GW`?22x4}{(RTbHS zLdNU+L+(0=9}&fvib6bsO0bkPtq=_c7XTE!y01WwYNkwQkt5JKA8rEFRNLCV z8UB-tH)KmUN7OE&w=#oLwxWocj37c$7DQw_`)+sF-FtWO=J&9F7mih}>$;v`GQS?DJm0KWupID9DUn=ADu^+1l9bYb)x`Nv^qW6G#ec? zp%8+W^So<TfJEyo`F(2U(cU&(DeAy+>>@sA1XfW0IX}a$_v}|^1p7jk!y#EnJp%uuQHT-=$spw zl|~&LfRJ43qz>iB%`B50W+)0_E8B8?Bw`-ZPjl(O!2;7u7MS%E{J=z8&!Zfgf5EIXV-GRuay^_iFAnAuFT zhd@L$Ff-Fo2>sH^HnR|dro_%=>~v;zwj1#NbzP^FHqEAO+nm8{Nl#BtiEM;kk4zvr zGZ+^XoqKw{T0l#j`TQ62Fy0mS- zdCT_erP+VZ*_yqG%jGf`pv}|zsybsMITumab?7EI*n`+TJnVyTOepTZ0mA6;kt83W8URKAv>Ec)A zdH;HC+NJ5)jlxuwf@X&*05q*_8d_}}KuG{3msgmHkW*j;jP$T002wO@o`_6Q2umz{ zmVZN9x((RqZD;a_zJR>X;W#ER^${!wkouZL)XkFACLZiMMb`y&MC6&ww9WBIGQfdo*6gWk5bgwc5$=T5Ij(`>rz77`Vuo&?9b=cb)6v+ zTOlpfhPcsI8Pa4ak~*XqW9OX9_;@iDqX-y5M7;Op91&^%wg8Y|M3gBrhfz?ZxL&sB zC+m<%d>yF*37{zjPsRYj0Z=f&td7CPf-MVle&HPByod`P|FtdZco-HNzyGbCHKb~q ziz@c9aOgA?a=G7Vy^hpyr4J6glooFyb`SX}a{LX6DKnj1P}a&R3hP zN_Jf&qGY5PQ$)z+m}xXJU7U1gV?0~?D!SS$>JJ!tdo7w=>z>r zmMwkC)epbT_WG9AOPNwrC2U}+zu=l?MhG`5YBaEX#2AZ0%mBk|HmmFU`1m;cr4T|s zWG-#zDaqyU-9Gz@c8vgj@da_ymD9n!$<=)Ix4p%zy>WT@dfrab^Y3<(DwMxhm=kHd zrj$-&tn2!XAAj#r+5lbetYOdZ)R3sWMBUYOZuKPo{@ zK}abXlbRYMK_P5xN=ift21ADD8j^uv0>H=_B=|z=x}BH)jXRq~eRy`Pi(L|;q$za^ zU(R-YNB->WOj9b0!na#HCgzd@6DIYZ!O?VF`fLEmWd?i=W|p8tBv2;~DTWY2yXm@4 zn2D&0iDJxm?lYcA&8%tGw{GphGtQ2^_c2BjBbvq$iN&cZNW>L!dERv$Ckt&T4Bk5p zX??!(KJM=9Y&IJ+tLL*WhOQK@t2}EdlN4FX=V?kwkU&Jx6uWHUQ;KA)CeAt8+eu~F z`AkSmM2wto4K?d==EgEESDL-?*^|vOy=Bho#UGuXo`$w2s^EItj)uC;}Mr{wLo| zGeS%N(6-hSsIR${)AQoEt(c2#G~L`_QUlZwTdsR;oMXPGs^qf{q-Hry*ED4r5S)_*X2-7P*t-p=hA-PYFEVT&3aFYcGEn4 zy1sK~mNbY6MJh;Cb-iBimNR0cq=cT*E<>*`rY&7RMPIMIa-q9bR&;y%J>Mwo{|W#| z48UCPB?1u{OP5qw$O*T~1_cw8bD(#m!y3<3(tNDV<`pH)_Z5+U{bxV;fBfzL@$Mi0 zG@YHo3rNq;&ligY+JJ45o19yb^5;wA^L#Zz$LlYm_V@~GdM}BF+wfR!EJ&g|{B>1Hf#n^4nl8V+hucpD*1rM{0Q}Dl)QA`2rz?xM=|(ge;BPwsomz z3n8MEOjRvu4ys`gz=F%-A%qZPge3rE%FIbqN(m{B49PiJwyLW2yuPAvS zCNDdTIw=YDBF}wIj*3?-fWJG-=*ICPRgHn25k2EtlDxgRXYb7BA+>F*PuH816JDNe zIzk(WXm>AHWJ$>A&y)-*&h_#ML_|i!BAHgohH%KX4>>NFB>@o)@HEGn|5=H0{Fibm zJR>$1C9rp8S$khHbKAC|$wHS?QfH*9)S;w`L0!|ukVd&^swqYRY3O~<6qA@SB?56@ z9G!JoQ~%$_2O=P%fFO;OAl=gS#Q`cMF%T)KA&qo{l%rF+q#4cVMrp>7W^^}sshLeM7G3|=5HQF&uagD7E#osxu02~O_U{Ee{c`T zBsLnBd^7`H&7zc_%uP{7+ZZyBg zMEJ8&%`P)X)~g1Dj+CR3JtZXKi(PxgsxoP+fXPNM$L7weZQ!V{bH{lFxPdumJ2uvU!6xLIANvzN z1gNaAv{UTo>;D5?Q}TfrYCn8vRA8TTE=bMR5RJk`@kBBD$0CYuT;iS?ZRnd3bj$-S z^U2lDp0hT2j+}bvYXxq^EEVXGPU+V}S>2+?(@XQ&@dhMGIHf5Bf@!%q@N6xpo?l!f z8qNeBJI9h~*5gmnhOWo2^~gvuiqNo7KwcZtQg0I%`F%h=EW#CGH!Sfl@1usQu>c8m zcFUETMTxm?;=;>SA6#HJbdW`5QJo^8;BBjfq1F7D8h=~eOVsi8IW#)Wl4$mPAsQVY z7-9_KKQXEDr$V6;2hE{#XK3M8TGG9BDe1O026|y$TK&)GfjARdoecF+DO7Ozg!(Mfkuz+$r1Say=!}M;bTja<9Gcs z`gt3q`N&NHv}iNZHs_;XL%Eu|na?RR>yc}#_fv~Z-s{kCOvv7^eXeZxYIzJ`ker_C z$Vf;@{W=pQSyomlM?YtM3$npwW@Jk42+}g~k4s^nowVzVtRMr`Eyp;bekVoEhMQ}D zH((v^QBW7tb!F%+T6mWeQc$a+2xjSvWk&b7S;CHA#O8N=Vq#$uR!Y^_lfxf>3kKvm zFp>}b<~mjb{Fa3U)TUwD_!KM6|KyW7QoawWL!Rm+m|p>5Hu%V&YFWSt`jxB&N-WEI zz7qfGBle(5!hPM{eC}|Kgk~PH2|zzk&I@G>D*p~9Lj$r1w2dy=V}#wa#ZEuIi~&mK z&w#QlET!bB#k@IACqIxB_FnHQs_ZO0G7tD}-RXKcR9wFrx)?pm?4dk_g{GgTeD*Jv zKiH)^n%96hp&byQQ!sMt%R!1SV zd=cKzPUR15sjRqdVS^VMY2vEkIifN~wcG=75J-BFl}3>jAqU~2Vz{BBsGs7@kifWMu0d3~L=NrTS3IGENu)L)DCoz%o#DlC)BC^kvvm~*}H4qpqLLrxpr zO*-2leSLeXM0gXz?+}3gEwuU!id}S%smbAs`T^Q)wYJ}(Y)PEv81r<=v{xE|EgeNp z20<7*QvNruy)-GTCog(ETg7lB3{#xKo67;gYLlIO-nsn71r>LQ;S06(qnKbxfe+nv zFG(d_h!w z@OcQJq3cTD{e+cWeY`N`VvxvEXjg-I*@YBm<1x)TS*^jtV6SmMTBb9_tFc&)}Hlf`9$nH>;C!4RNMGx|52|+ zA(h!Lwb|^B_b&_R9u|bC3#5a#TK{(WombcIhy!z1Gdqym_a{CSGYvLvWzOhDCoJzF zs^7$}bo!*Fzp>lgH89wpb@>D-_WKKNPE0oZPT$lIKhJvk^uD*(z=n&1Bsn3}i6GL7 zHSH9<1_K8P_PZM+8(Cd&Q>hOj2(O0!%LDJksX9X`pH+_MEleVixg-gS-Myz z^f$|S?;OnsH~H}i| zh`)c_r@c)c9I^~)tP-gbKHGQYw)dRde#m7Mce*o@ z_bwM95Os|C)62S^(cm)6!8j#P9Wm2sSxwhXcK8?LXMdh`e{y|mQ&sgTo=9 z+7?k6^KBeN8ceOlTc_n{f;OA{JMc=4=YCs{u& z7MX$w9d|4ec&ObKUj@Iz7Yibm`xJ@Pp8#;%xae=;YfbNP=D(yo+7ojd7TfIScd`Mr zL$$WTd?1qIUC+8^ulu!r;GdA}nvUR}tt*p1ZkE_I*>VKpzoWtnER#JfEDGZQf0*2! zlj^bc$`{VfkEo8|&Y*6}6rYd9$!l(H4S!zAu?;#R(CJxOX}SelRsmGz@1wAxdu zrx<6%htvV3+65*4WmPJ*Ey=hQNBAqO%OQK;lm0hZ#I00v{-4XBIjk9!FTc%`tYi*g5UG+z#@kyZ`*H2&(HVXo2Y$d(5G(EGUD5h zGa{#k5d&SxSLaquXMcvyJbPazJg6XC5W7P?eg*9lG|mHQ2NehTIg;+u1gAS?+EG~U z2LCPPFxILK(1D0*9BWH{9d5)NQpS)SP{qU=dHad=m}*IO;DXNJC*)_>;uybm;` zvvq7?EoIZP{-3Y{3#Nm=o@hWCZIgAIlQpXURTd-&1$}K}dc&?&fByWSr&)ty86@AG_`+3=5r6VsOOiUp^!XqL3=;X*V z<8w(y(-+IS!4y#vnko{iSh3gIqv_)EggBza313q`k#TO8j+IiTs8^^fm$s&BqM#sg zW|Aie?e4i&^zbY1%A#4BhZC9h))SfPjAa{Jg=chetR-1uo6z0IBvdY4Hsc^OEjD~D zzLUP^3yK;~f^5Y}b>}PO{(f3JKPH$r3>@h>&@AKOT~9wN(HwM-`{1!WP@-qD@_W0n zsX(uqCM&tuILGHtj5@zATtgIelW*w-p^3d5eZw0jgyCa zo%YLaquaA(fm??%fWV@J(`aR<;ou?imC^a`AhRd$33T=at0&3h$S9E92T?t4>-a*V ztjv~^NgDgh-V#a^LY@~R{Z~|C#TNP4rC_1{)1DYoZ<1yN9g3oC()qW+naJEY4zH)V z-{yy-Sw;#|A>89!JPZBj6lMxCv(#M*pyl+BoCEkZoJfMAxN3P7g=*~y5_I{tIU+a zo0TZsQ1ug+$ICh(@s6>eeei;qHbI1FMe)rblnxTwNQ= zH9w+8Khxx0^zoNisd@j1N*M~1o`6g=GA#E#b*5wYL6MlWq+9mqJOJ7}y1;7Fw(M!>@^U)F|l5VH>Ds z3(J`+1L|w!p^^%-hEK~$SB3xgLIbNt{87;*P53kY|4yPFBKf0czqT66c)3>V^DGwc z*4h63_hS!6xh>!BDP>?Y>~`TZm8ya&tFXX& ztDgO;JL#svA#P3T40T+t<3*FkE+FKTLP7FkCK2mtnJrLSO*fK7!allY($Z>tjC8g6nDXWe z*B6a_laB@j5Xh5nt%hDJP34Uaw$1&z-AMQo*}Ys8Sb`tokXL7DFT{drb(nXK^!W04 zUwL=FwZ84)`Qk$N(MUY6^u;bI$b1KcZ{hNJNZEQ_Jl6R=JKAHMr62BqBl4903-ZZ0 zyROM_x%0xwl>j)v)Iek&ZvBYUqp+(hC&*?s%pnR#?QtbZgzUl%oMR`2EMH@oJ_n8s zdyn0j-Be3d{fr}oQnM~OXDF@Q!<9oV>TuBqgNLH4@~g1bvrA<96A5AaF^YWuERmC# zRH%2CId)CZI)<__>ewvYfG^kZyu)W7RTWCu03zxF*ACcLzU8gqAF#v8c?*&YU|At$ zigRo=ECcx?-3?I}(!hw3;^A_8EPkd@Y|6`EsurGqp`l%#sktgN!mZ?-lJf- z|MEGl$!5CJFA2S}S3s~uyo`^9o3w@SIhgJ@Mq%e)axM?F@Z32oU)!t&!Z!untP2x2a!Vs7*$lvX z2ji4r0)sV^GzdP(n+c~Gq<$Hs%pu_-ptQGsW}e27L}xRr4N04JF|c3+^;!5;tAeY4 zFQw07YyWiNW@2}R6O8HGs8`GLk&T%@1&=3b3MHqwsX}JYppE+E;Dw7?7O45v9vWfN zb_wM3wK!qQ9K3H2kH+e34e-%7B43KYrSQfCJrMk_C!ner6Do6NNbki8w z{k>f4XUw6*#vKm2XhWC!os|&LOL^v9`;b&&>F&3>$uFlPjPFR*qfa`O!2G=)y&bZJ zaNS3SJV|ODlfGr2Sm%uVw4%6*)! z!Sp~^=A1D5v~7 zdaU|zbRPq;aTw|L4}U11ro09=P)}>b;(~E?{sZ?hfgY!%eneUFPk{e0n2>{h>wa-w zq52hnR$T<2Z41lP&L@uTCX<6%F=@OfQZfgT6Uxq1T0Q?m_Q{4R4y6%5rLF$iwRQN{ zjR0o*0vh&9uT z02ug);k7Gq!=J|>{79FG#3Iw#%5BuGZP=u(>QG$JGk-FXG>60;-El!4j`f35J)YI3 zm%;(eP(pRFmA0A}eM?BL6j?D!vA6-i$eIK-cjk&3B2>=RO>C@4>HBKdGrl_9$`+K9 zm|8kk-#v&hpYs;Gv6so5Z}oC36ql2XC>ehHosrR*RnR=gc|H-?#lB|3l_&n;)H=Hj z&Ai`XAx+sUC`+}wyx*-K*?H_q`F+^Q`RLzo{x{RfDQ^Wo@IQKg+J(fi4#0lH_o5tda^|;6zeQi~#q0ph={bTjVf4*T^WbVj|{dIh&ySW*wX8hiGU?C65l%z4BuLEzAl=6AFyTB~A z+e6z9(JvFepFtoWbi>~Bgoqrgs&Truj~8uF_1`lJ0{-atz%aEV3kjWQ;u1+yS@_Rj z3Me)h{1wW}tB81cei#BGdaPR)iEn=mdI;7%SW2wQ54_d2jDj^|;D(+}DWMLtA?dWS zWBg%Q%*W!mClHmfsmj21LoaUZ%VzJ>wd*#aD*qFPf`v$B2CMW(l{WEj>aptg?DiLZ zFZUWk>+AxYRAz;df}81%KfPrH2Wk1A^kGx7nsjyH3V&7}P;buWlJ5owC73g)kJI`` zR+(=qlQF3F^J~tc3r0-dG0DH&Y_swH4OiIW`X2E*JsInALLk3u220ioi3|;UlI9Rb zPHO!kP{(D@do)=*v`T3txpr4>s=MwHF&VS~vaH`&9{ya^4TTxU3@vbCT7Hb347W8b z(n?Pdx5!%@;&k-2R61)At2=((_&9|0j=(pw9e1%#i$c=dzbZ^%V^Kew6mSY1^`fsa z?aJP_nwJz4w`RIEoN&qbh-xQ)!T*AoDj8OWS4d8bmDGdmE=|Xj1!HRoVg1{?A6U)G zuh=l9qg+X4w=`=0G*P(-c4-*;(|)`1WFO=0_h3-An2ParaCRQAiA@Q7C7%2}B$Opb zzpqgy+lkFMn@}Cm7IONpZuTtIyLazT!}rBGXWACerwx9Fto5m%ca)#!KK11LbYi4+ zlZqGHsazTfr_Ov4xE=v3a+(QgYES{oJWO0$JRBj`9TmauRf~YG=C&nOsWrxvvAdkp zf*bosf}OIAS1c6z>jd^esy>$$&$G};CgU0IKZyvQnlr&ypqJgj|!j`A!ZY0m64 z1Kp7`9aa>z!obNpwIRiS{X^fsslux^#A0GfnjZm$jNqepvVF7nG-`2&P!^>+Dq}){ zZ)vZ0P?I`LtT|}{=54Mw)Q(4^YQEOTyBCw#zTa23t#TQ+D|w%+&8osFI7(pLyo8{BBpI1$9wS-{6031O+eIlqfvo1Sd}p~(x>Ul^ zO(|D7{vtM)&>@PQWjGj?%or3Cx54mjwl0KB4xjyny4CCuG`?(YVf-oMPv7vIWbO?F zdN0xoQyLX5^KcjlEPC(sMO)D`q|28{9VmlKZ51Xc z^TYXtRh(R1>%Np})w1pI*QDpVnW6;#30_l6K*9!97%X1Y#Ru;h{n(Xy29~#yE1Lgr@9})7;!K#*2fP8Bo!5{J8crSI{}%K^;ZK% z2I6%EGRB)TLj_!xwO{V8fwb?_F|}4d5x^9|_GqK5y!ApVf6Gjw6~--a>Y+69YttYI zu+Hav-Jm3q>q9O8qAFm8nrPmM3`$qd()gH67B>7G^If>fywE1(Eqkv#3{W+cmt-d8Yx&I@s%_^!q%wfc#s2L;Qzbh$-%v zNT{36US;$bZ|4b|U{A<7_=86aTjD^IvdE=vaaV%_Ag$Le<#ODt^k?ll+?> zG~kHKkngO#k#5f&8ZzI#l5w%Xx!6FFG|1b2dAqg6=3>!?!y7#cMc~{1V4`=dasa(2 zXY4UiS@8O74V^9>bzBu5Y_uYIY1ta~yL^{mPDrIXF(q*Gy=O~LMl;@p)Wb?H9(@+A z%Rg5sW@L_;2`S%y!Ju!rnLl5HKfqw6Jg%;8>ysw^cQrA5>$S#&rkVn*{vmm~k7BYJ zF01=ey++NdSEx&_ysUG)&RhL1Pr0`H{T>#905-|X&H9Dw|4ur94c*zQSx^HuLol9} zEjQD97Q*bJf$Y@K->*Q7UpbMivylZi9vU8QO)flZWQPBR`Eo;xGB3jL+*W^OC?%8- zsV}@}HCX7yt(Sc5k#Ui^Bb`|Y3Qfo6Q~khvMD1T-09sk#D^*1;UWT_{hCggRIC}dY zkhb zdGAuwo(X3FP*Z+;dU4Qc^Dr+$FXLXCEq8m>?@=YlIaqzVmsfETU)B+XLw$ZWZjzfM zkwB_qT4m(ZvoOwc4Jr6ST3Mi9b%vbia0 z`m$5M^#jZ^6a3#VMiW3hN=AHbx3}SMh_WiHI;MqOf0tV$Y(KtGar{>yFu6P+YjBGR zrjN+EW~D!_HT5mKq>+@;3y3<}SZ&v_*kPseiO)G4;94r?Yb$mFL>DSGzm%lttNzKW zi*6FfFplIw5 z_Vq`$+I zlFxKD>M33Y7Qaq)20L$(RTrvd`5tvs4CSEb=gTH<0Q>9+;B;?_N4?Z(6lt+CCjJYiZ?K>SXhJ55s}F0}Uh^^>OuC6+{TmM9A@2JjcxyAOwM|Knl1 zIXael^1E!PQm-4!aYy_Cr=gGL<+>p#S309_Q2mfq-($Ui|HaZo~WzREgEpW60o%Dwy9Y zf*M9I)cCI{UmZK`l;^ zpnqlqU3qjM{*ZhK!EZJw`>$e_W^Q?PGC`8>`8vNM^OIY0@5*mGQHQwII(}4R5{GVSAXpg=RPV%^{Jmlft_I^siU(r^*;D+1c~7X!f3O z?6SGl9M@{xblK$CdXE< z{qx>J(?rV|Y88)2T2%7hyU{XizM4yKGc7s5AAhAQIZ_hDGK`F6WQvspaahYNvb$b1 za5?b_42JGC9L!H?T7O;1XGj3+4zoZ9p5G-1As3%9O4>$^T3nCmHOpKbsX^_><HmRol0w_32<1XyE7DHO_WMc`hjM%}v#faiA=?RT=? zn?Nvm3Kxm$*4lbrO;GZ4fge4;T9{-D$(?zo=A4hR0d?;_-ck4QP(F@$qRgnE)b~pXsSFD$Vmn;xH>PiV1WfL@srA(q zE&8aQ&DxELJfriCpYlUxR6zFHJiNNe51sLGBkv(n%X;6?szc~}bdTI~beNgjk%P^d zYV<8B=1!Da;1lU)-_f_wVk_x=(;Zl!k}(u4+2gFps~Eu#VhrkHA$cp0b}=x}G?qX$B0foTS1ZizvgbaW=+%UIigNswP)aNcRYQe%qJ%xYy#XryCQwU7W46ll>SH!aCOA|XbSiOVC!&dJ z+Y+tHv1`gi-qr?MZY@j?V-g6gDa1k)>KF0{uTS%g*VA<0Ghz}FMy(JL{09fxVCUP@ z?Zw-nMeoxQ#j9$7Z!}{Lfpp$fa}C0e7fLRe59r}Cj20E45>0Z4$fIVHF zoX^t__#JQ5-f<_;l7GMkMN>Fo?!@yxU$0EIN&A>s;JU5WA>a5WQV9@$HP^*n>oCID zAbZ^ex+uGw6JB(H*Ty@EwXnR^ggEdW_i$(q?<*%(am;1roXrsHJwEQ+34cQm~W2n6cC}PJdM#-T}R?PP?w3gX^u|y+yt6 z@f(@OT0b+-8hyXHw@6a<|4w@4J-xz*JHN>sR(`SCx_Z5ZW6cOQ^DybdSJCQeNkon)s zf4Ig(Kh!_$-&-hRKp5IVDL!$TW}jW*;SD?KLiDSN+C+IW^k+Un0jj_`k&F=Mm!B1r zqPLg&kpkh`Bv=goShES;74*{oMq8VUzWuUGI%&QMk1s^H%z`yH8`%g*$bT^Vo$0pM zx@ztnZdkY7BNX_diKdH`e~A-|pN+ccK(BgR{}RapZhL{8lfVBvpe1S5S`y;EPlg(m z=O$?dmCWDIAOxE@PW;YHAR&x1cTF>#q#c@BH`dq@KAg6dbg^QE$-Glx3^_lpTZmYq zHrrrkti_Vl=~zjoka=I&bwC%iAT-12ia~vT#`}*6v>)(?+Xy4)0MJXf0$b$@1E6g3a3UuY6C zjer1;iMoCC5gLl{KuEW_>`v0SZMey=t)YRfISM~1;xPTx1IIrvL7`KoZSX7B*tFWT zl9^;Herj zv$NJ_JxGt3Mz(aP-#Wh(Fda(R>2~f(*3rPp0_f$Zc5C6~+9i~$M3}w15&FBscwUUi z|EQT^ImGNpaJpD8dHc|lYjPSvM)Shw0S40cxjFL! zJ-45i5SJntIn;w}s>>bvNaK!~KQJ_y>ec>bu8UM+a+J6MA%DOeU#A%}X%KWsYSA&) z>~A=?qHQvzynxmka)HbZJxh?9t0& z*90LPB`c`J%lU|)?NAL1-ru_54P;o*FGk|lFn(tf?}cVm`0cM<8A%yqQ)jFaMQr4DQ2 zz}5(IC;ar_`o0cIP$k=!mD8u5+Mc(Kwv5Dia%PF-Q~k(vzOZB;2qS* z%dPGC5IVNysAnZoWIZ)pmW=eA=)<*y;O zW2`~{ERD2k;qXz>o~>!%yBW|_{&c9?c2;E5M&``o%sYYy|kAjHr+KJ-3Z?ixLK9u4}@o9g#b%Q69J z>?!aAXF!Rh>IOyXYXdzfulh$uM#d)k zJzEWJC)xA6_z?Q`8LwBl2{zN?!(dus`KMex>y93C&ip;3AXyG260vhUXyOOT96_c5 zY*B(x{Pc^3+@cI86h;f;q5)`!CT!Zdl*pV&s*vVAbuOR6&4zCqq!AqUhO-_~fSWR+BGy~Sm!OTMZRIav+nReK1!*Q~=u}RQ7uoMFVc9w5i z5GxY7!=)Uy4fA|BJ7Z7Z6C@XO!7CrQK9+-Flbf3*69Y}ZV4^AJv6DhWfr$8<3yk0Q zOZP_oOy?3x(*YLzU?i7W2`n3)`U{r>lm`Y0%Ns*g65~*z_~365(dW)zx{!O-v%1wF z)Jw-OK*9m%8fynf2o0g36#%sKk-}8@w;ul0D&iwEH%K(! z4NZxTQf)f1xH+MFs}T-|dpwF^&Y??V%uh~EVX1>omdSlAo+}nu=|>8^NE`ZSr1LWH z1xYi$2F5s1V9q{s99d&)8j&_(IVQv;z-iTraIZF`mk|TDvFMwhYCyoF_oxuO$r0C~ zfH>4YRbhxWB{0^)hgLlv^eq@B+9(H3g%iiFQ6O#1`^CnR@S{F-1rahFy~gIrqnPtVqOR?N>L0w1DTc@h`Q@ntWRryS0riGgA+WeI<-**$ zUC_%IK1>x_OJ@S<7hV|D;y(ufdXS-R^BO`uwB5u;VDSp1A1_3}DVv=Li%5ZS?MIiO zVn913PYkMz@P@A85U0nlCM(B8v35|1B9n56G7PkA4Xm;LV7J#^b3`*+TbUu<{os<~FZWwd|o z&8?%ADZIYJ{#hWqg8AGg-T$ppD*T%nKSuO->-_BiX+a#3$;b9ykNOciNp z#nj78=2;%2l%$gw8>ixVt=G&IaV6g)(?oj!GYdL6D|!83gMZk~x$pD06&r>?mETrT z-{zf9W`?Z1Py5FJ%NFl}pYOxcQnqOzep05A{hyu$J|AWp?wOL!L#5g@@bJmd;^qT+ zJpPfj2AN^X<68sr)?rJi2x9RO<)k?-M~Rg>S$H>7kbhk#=)z$Ez}jF+5;Q{e%{zPu z8S}?eU&_Zi(k)`GX09e9X;%|16d@vM&qRNjH_q9>;o{Sa?LIek?X{$G$-{1`1>4#* z)C?T(*`fHqNDp$)HIk8O+f&$o{0K_ZDHm!|!Kr*u-&;G9IJhfNOAT%80rxvmBotB2 zg(WE^-UJ|}3Zv-(bTZyPU?zUC{5I4x`CffGp zZ@VDzF(M}pbIDXkZ;!_9r|YDf`WdE)hg#2{mU^#faG^VC@e?Te^}Hh z*1O<^Sq)OOPY7vPf!u$gqhZPnY*?qUD)!-Y3@i6<>2b{k-QO90aIZ>X$f@NW+G`WV zo5{kIE|~e9=9*qii^L@l)_f;?!vx=S#Qh=i6xWWc=Y7TthGN;UITsM^Q$bP$)%w5> z*7ter?4S_l*I`&eViYx2iP2gAJl~Z(7Ocm|%|B0v2_*5p>-fqZ%GRv_8T{;l@tJy| zu2BSEOrb(li*Ck!P?dJ&DjW)hB}i44MDm9eMDmlr$*8rh{^YLa_Wah{_hIjL`R~t% ztF2`K9-5hHpYF7oA01?_1wDY)e=D*uiRyhn`}M~Y68`lDmwx020;s}U(%1$;HB!sL zRLOkUISZovS_}%uA3q9NL`^wAw^bsGO`(5Zw=s4ZZPealhHGe>jcQL++vc z@+~YxecG)_gPG4VIY2Nn6|g=7Q(X3>&XWTQWAcZ-4yE}owG@s4_MbX<8FC7d8N-2Z z(crEHoOg6wFHz_;FxN_kQr{2EW1veS)vA>zd5X9-{qvm&o^O(8%}plZp4= z4H?2N&NhRCUqdV0TEAu+%@I@^ZeL>ksX^*dnb@*MA!bW)&8L{|KreSMCjC#pdEOgA zp}S7VI<MQmohbhF^76v-G6!~QcQ%a? zB}2M2O&L}jO04IMf0cti}BXLnwQO|*H4kxPMr)epxMqQsznKCR$XM8 zd%@mJP%_KEMR3tqzukKijXDSWnc?$I2ivj1$>a_pRHl+X=Rg%>bA|(*)X;QX}Y-l?D{;UUC0!_f$v%q|Kl6DBJ^DhbTZ^2XA)-% z3x=t~LwmrK#~l#F$SkDS*p}#bQR;o0a*54*)B7shHnjRXFUVvM1`~UzjXoCN)K?F~ zhYxryCbE#kUNHTU7kk+ug#UGjul^5(JeJW*`TyjB#8$gi{`I63ZroS90xE2dne7u5 zBS|Of#Qm3D|G!U~?zj#%T2Yf!^8FZM|S^z4w zmf)zN8==lEx>+5M*2rn|ysM-x9zOSA2^pt+N2gZOa#&s*;BL%o!@acNArm2+DGeq6 zguA>MTm!HV?|eUQi?m-28QrfN-JdOT4G-&7Uy^RliQ{a}*+`e{iqE#it+vI_RUVla zj&v%$3QwDvXdiv;?mhNucrgipo4w9ALz?R{1+T8I9_*2I0Q{H*Z|OI&1(vF%CkxH3 zaJW8X&(#3Hcc&_Nm#5H~YQuQtrP>T|JFHtuMGUUv?vNG9FHDp@4qB2Drk~#3Znodf zK3p0}-yZ%QLLURIu*8X&e^Xs2Ui?8ygclG){I~CiybD345>G*|nSGU@la8zD z@qfZ$C$&XdePH?lw@ArS30t^jki}9-B5XAv*^sU%>P?mf39m zA%FnqUo^(#@{&4w?QMVs5E z+5rRz<9tA6XjN2Q=h@CyUU@;ulR!o3mS#!TRC0=U{`q)${|*gja;@GRtUQ6cPV4G= zs?{Ij{3zCo)Vtw(nx<={TeMha^^TCWmrT>uKiVC&T=AL#8utR6YN8Q2)OOLvwNUph z%Vz`F{an{STph~X{ah_PQe0OpuJ@g(+b=#qLIpD)rr`t7S`<`h(K_R)7HF%1fr$}8 zYfd|egBBJUcMP6+!n)ec(gUHAyzAk2Gw4M^AOJ`uF2`!M3Crhp$@u+cznZUOW&Ny+ zWutEHv$Tc&#W$cDcyq7X;)Gxor8I_ zP4^6zi+R~MJ+eJvuh#*v{{`hx#ylJW{7_N8z?X$3ho^^VziT3yo2@&wjgu=>MvCd? zk{xEjcM1C@E}a+=tF-h3uy_7_;l(4+l2uD{?N_$!l`8Ok5rhQ-eE|Qb)%p8OpX)jt zKbW8J_Lt7aPc=V+uV85-PT~rRj^PnmqfZfcD33~B+o}D{&^t0yq$M}-u?VNF(|$rZ zY_?6-#5xpYr7xVs8OHGvV}RUfX|IHo#!o~HE*yY}7)lNS1i*nnIL%fo1;az2A6$YHAnu-nL#6#afA2M}u$+~de4*c^~S#-#DfC8{YD zCLpn?fhr;gl1yF+orBE`aj%q+d#%Id4sikqyEBxrpgz~MUE*@h)ZTKacWX;5PuWCx z)ULNwD2z(WrBOf8T$d84{|H9BWGbwgu7CPn{>;tIT~WJj#|#>u;n_vf1)=M0HUxgRc7JnS3- zMVa?ywr)i;`kpJ_+0}KUbrFp&-WBJxLZ45Bx3EzvdL1K+dsnR5D>x=8+A%U}yi33T zKwM2n4-5EvJfuz*>=R20EC=-CMO_(#;^Y6j!fyST-5-MagpkwpK+j~@Q@aQLZMOPy#F)kA|MmxHzuPn* z^KV*ACTp0z7`8UFa)8Phq-~f@w2$t;ZOgurK)ri!>i6LHu(R%K z%xfXkUXa`+8s#ZCgH5r2O!C9=YM_79vdzGr^m9LQGTp_Xp8lX@Mhawus=&-*|iEe`3G(Pg_B$_9D0SnuP@tddvhNbZvU@$Yh(8$T8k;E7@zDFUy zY!Q`G<*;yL?gkj#eDAvfwYK#7&?wT*ZQrqx>s{&)lxm4t{<#M$lSUn7t{5J}#^VEP z5_WGECVBO|1QRW$fkd_vK|yVC35gOad3782ci`VVQ@bwitLOi#!Xq#I9)O>Rt99N3 zfOMC$m+mES0{Z%rd_+I_difBpbpeFB|^v8*~fb|FHl<@A`Yh(PWDUpcJ|9CR&|uc3YgDnwbbc z_zPSm;tvA(Ye%w>x33=Yeeu9-*76B9E8saBs-mgIQ*vi*fE9|o=UT~}vZj`wwxWB2 z#XG1~6ap7C@C=j`oR~|Q>+*$6UVzVOi?E}#;g^O$8J5jn6r6lQib-608St?h-7k0k z|9$xH4WNj&I{Hhw`egoWq1Uq=&XGo+U|xJOE3clO5PY%lTJx&}Qvk}JAN7n-nz6fX z_^UBA!_mON_?|jyWKekM@JBDRgkMSo=v29)QC#*tx)?|{h?=c?w%_CBXJLT>76tB@ zgH9kKu2~n0SO^4Vaxqv$m+Vki{ooLC6&nSEgA?SJ<$(J9+$PfcUujkaXGY6HM$~@} zhTW=Z2A5Z^Lu>WTaC5`l!L?GmM=`&{uXfY%^hGI5S3+@d9cJrlc}3Ee1Z9@2PVR3o z5A%z6cfg26t6>z0ztXDsJWLFeE-`zu;I$lAwAdCh_>uA`(fPc#&8yw32%NqJqT=H4U1|Wn0XCb{F*Jb=?v14@QT!iAR~;4Q z+eDY{mX_`i=?)d78|ji-y1PMIx>-_b>Fy;2>5%RgkX-3l;(LGJbM`L|2hQ%h&&-{f zJNN#%$fXJ?we?%%B+z0kJD#_86NCQ)>^IBS z<3<_R;KpfomXa&fgR`vgXeub$-RQLqUcWYGv_|mYQ*Y8hG*dub*g?DJIpvy)AsV zTy8GNM}n|57bCxz7YKMH=M3R(1YAx8pU?VBJ#OPcZ{h!`&8^(u-|wH)L-v#vy)eJo z3W5e&m)Qy{csxN3XdfbZ)mK~FW>#Yj0euC}<1y&q+mQFX4^&zKuA>*bwalLqZ+ij_n-8=6tB|$|`%*w}(iwJh`{L8x3AYs%lai9M z6+AeeZ`@6<5X=BnbdG=_6@|Arktq z15FqB`UQ-2LA@s)qptBj{HtbHRBdl=I2|G(x3uuA%M$44C>T~3)p}#vl3Eqzd8H!5 z8#APxIE#$gnLe9!C(aE2gu$@(HQu6tX}vxc(qNq_Gy;tr>a_G#_LM7uF*ppaS8mIF&f zhW2bsjd6opHuXxzMA=gX;u3i0QLrw7f+^C=KyXW`?TVNC0`;5zzUp>&CmoUvC4Gbt z(kYAsobihzdUJYk9cwQXQlG z$R;oS%NtRZ{R1P6ydco^C9~&o=s(|JpgKjT1NQvy12#~sz{F*f8rxgB#J_q8I3#jz zhpFPL`fuNd9=sdqHp@wv`LtugYi(-VNr<%JwOfGv&waI-j~6v5iyWzq>%VpOTxp>C zg}rq~lIMazrh^n?jsatp$v(ViA8&6`%~nVFucx=b4{$RN99Q!PP~W}cZM?4d%7co- zR${d@Kz4}rX6(XwWgjS5pIcu}%7Pz&*VM~t&lA8RJOlce)RF;9iv;xUlRQIY7MY?1 zF?9GEz#x8A6nWQdU|bpT>O2SVOf#uX|K-04A^zz4Z&ECl=J9i2(6e{&t@nR_sP%OK z4&4BLF3>GOZ0ye%nipReUp61ZXxoT&pCkoXp+WdNtI*FL*mzf)l8e=GT=&eg?H`YLH7?;Rd>wVjD0iZ~=>CwbxA?rY>2}zLUYF_AB*JH4}0N5YZVMxMgKMExLZ}xVB10JppgP-n){#$x| zl4Ih7ou0U;CvP{-2TIzi?qaK zj+k%YJ+Xu3!e!8v64pIvT4}xQ%FMy9r0b{0UQe5=FELBNb&uBMSC!cTGHA4WP0d)^ z$=%9`pnX+$_aRUgRn4YS*<)=Av=nGIF^3$F`R54P_g{ zLZOi0I&pxgCwM@w*+B(fxX+DF2$=Rq0?~yEHx5wJ)BEhrtia!FL4! zjs3nKz^mPX@T6i}iZ|?FTnUVs43n^tc(IJ|7*?JbbfuATFY8T_p*<6gR~^Aa{%<|O z4mUCbHT3lEfObl&J6Ji-sYkK{QYKR~g#+~<<8eRak_jH_T>jW9!Z(XgzRtxmTECGW z@CvluK*py8E`_XU6+twPW{Ci7^*d(xOBzIL7H~tl{m+Uyhe;tp{QXH8_>B5IL7cr3 zv;0HRG;}sMZ@UUTg8mWzG5_(^tLKl)By?6@5uFS(t5A2Y-@+8y<$Q_WGHp3*FY0{m zo_y|ji9YANrU*X}E$=V6TUW0H=l2th7bSYQO#{w$JTVpxHY0{9rBRe|yd+-o!RLYT!xkkd@9w5QWRiRn}7mery15qz0E>j#O1E>y1?CSASh4+z^v z4LB|o0Wnq1j75n?_cPOnZGiLbsW#pPUGp-w&c+FmaP820yxFgAZh^tqoes!yCQ$R? z-NBdp6S1Uk|4g0--{HM%M*)%PeDg8K^yx_V>n@zUtpY@HEQd!FB>vsi+uyQ-p4M9tE;P>LPCH^r_}QefWxMO9*-`2o{xK;b){gd zErCJjj~6>LK^MADzhU0hw?W>QMYkJbsNCU?m4w6YfK9rrtW4U=$(P0~iZxp2EOHbi z9u37x!=A;(#c_vhdkbv%@l()SlHb91o6oC&>`BhcQ%}Itf6WE>1Cjly7kqhsetz8x z;GXc6+ER3`8`!Vb0QBniBS1nMP|}iAP2)ku9VAQmWBGnRtbw4)XhB(sn>b>~9PAN= zdq6Rbc?@-~a?Bp|cdjN@!(x8BzglW59t{CL#{YN}k9XIL1wjvoZzW&uSqcJX9D(?C z3g%b}f$VJxoymb5V_O`BDDA7)fiEb37|()1tZNxDd(NC8h*qhbU_Zq~9wLIM{7SkR z+#`E26}O%cWwx|__R#%uy8bi)oYfrXo~Ik2cDV=SYl82#(fYV@CS>u4xv^Dd9e^=Z z&V)oG&!E1N44WAFn?c^9@_>bdvu*A&P#AqUm|1m!S3Gj`;cWt_R^JjD9@uCNUTmcf zikfV+HsZN!t#e$=J`kh#TXf6r()zAT(ed!p`T9*DX*f&qZ2nXyHABQ zfI@%3?GSz<0SX(PC`=(_?l$V``-SlcG;`oj&(9ORQnitx3iTN z+wA>|3%AzQ(+3k_UV?A6wYBXm?S@_JPEQkOzhX(La@|z}%zJ>^wC4qQ{bC{Qd6=*5 ze!NhkU+KBm;iAFo?|oBNDuSAMzClTm?)+KdWY!0k17^4tFzxa=2lzByzL)zimo(4Y z!S|-mq(_7H+>Wc5=fIeI2+og@Mjg%g2@yMv$02-k)5BPf@TB8-rG*&7)`o&UyHS21 zFe_CL-P5R*MQL+iWVlL_S4xNJ9H(bt5=Xw+U;9F=gFw zV`n=0`}&TwI(I)DRkrrP&*m$2UoH&0{H~9jpY2|NeUE~`1%na)`IZ5a4Y5IPZ0Wb- z2wk^HpUBOtuGzC_4kX6&_e|}pBi33;Vt+#xaQceVKB?f-P6gnHUJv|<8(S*qWtSU!^;C!OoiFUYKi%Uv zrzo}U;^-+V2Qw1A0VPthh122m{62mzZI`*mq{-VAYNo1MHo+V`pEHYz`WoBbsh#?+ z>c8mtl$zq48T)_Re${=v9|x=BeuY3&@i8a;3%H1j)&9_8q|Tz->$ml~#09dG`dSU$ z1_s~nbNDX7Zgpl>eY_@P3her63rUK?Z|^78$1>l(0M^|1kN1H5(oH4c_2(G$WD4XS zDc@V*ZhUuMy~$e#nxe}l!0EXY3E|KPXr~JC?xZ|e1A;Uy7M1 zGaU|f$VROe7(}*D8t-vTNn!s}Opn)VGv*}IH%Jzx5sh(Y>7oMI*Uc#!FaC1qRO0?< zxZ6!wQo|&aq6g!?#!heSF!%-(?%J~{ zb?Xe=z^Qyk08l>P9ApAd_ilNr!1TYjcLG>%tP82Q*S{o3$3^fAe&GDQb9MWMPps7f zTn|fejY&o4iy;ZoQ&@jqK0S3@lmN;re}DT;sQ>L$&;9Odt24lSy_m=ma_$D0yiTp1 zQ2UM=`<2QG@MZRMhb;x~e~DcJ=tr%LmD(!b0wk?9eS61E3i^~k4K}Jf1%!HY@a_HmeE=c?u@dOc7i<7;>S2fmdK*qjZx^551ON*Q3Lue(nuky& z|BSb-ZZ>a^`hgZaFK+Ln-d^c>k3u0X?LCy~u4BR#P}P93*aPOvC5S*lR4z=;25*7E zb!MuTe2L;Oq-9^Y-do%VFeQigw0&AIU|Dt|XbSg3V6D)cZQc3IX{}dP{Zpy~4lVb~ zF2%?UK-8eAT`)?8_Vm0w!Gd3wtqY!VtpB6znXGp|opr*p;W~qCcy-!FEZBKc ztk7ul?tlz?V?+rK-KF>+xiE%z7$6(3LLpn>Z*IAq0QR`mKzG8G_HoSkD(b zq73ga8@nZcpy+zJDwBE|>j}8s1uj|ZW`MW5`{jD8Grap{-NmsQ1z4jgwXMxA<+&+5 z?x1Q}FjzHea>V+tU~0~kkvRnbDvbh+g8WOg`9oB0<)muHu-Ye^l;kCkn@@k~z!}aL z_e*PWEG0A$;>6$W`kFk zJThi`77~K2Q&rjLBLdQ?KMb{lguJ6c@@X(!wjIGXGnjsll(4qejK&%hi;1pnk!6tOM_#IGJ6^_u?b)s?P!kDlzA;}CeA-{!e6sCaclwQ3_qkVa zO44@q=H_C4BkSeiv~%<6%==|is>|2WB*2@l;X|fu{q4@AH0MhuC8YP=!eD1t*D~;# zuU2|B_T98Q8rs!Nx10_(<(V}M8lV+)HOb)SQ>zzXH+|S2iL3uc#pvUj-5#B2o{`h^ z`irnkpWZLAhxiX@e){NyF&Hw_wOP^M?pSISe^OY)`*dSyX5=(ShEGv--9|E*WTCy= zaEnIA-o$Y@ofM-t$3m0KQGz{@)OK&47XGiC59?u}w}|Kh5)@a56@b8CH=XNeaM;G$+S=n?|0!U2l~@b-9+rs1%&!+Sb5kWNo1fp83XYA!c zz9U2+$-4xKn4lgO80ac6v0S>NZOOxUsGEtN3`{SFY`@H(U!FQF_5Macb|5fe*H7H>%eX<`h{$M(nj+@`_VzZ>LJx>R#=hyi4*Y#;K%F1{7n z*ya6YN*hQ$hs7AkYrjW3HUOhGHZQbrwn%7jBhuj;6eU1K4z_x?TfWN{y-63 zsd5WM-^{nBPOYo4D!`}T(V1gLCRk$+@QeEP_V&hN4?vt*kb1yyIw-jF1&CzNvv`{? zK%AoZ>$}=QRh8S%9v((8BhRlg``gRDx;Y1c$TvG~o$D9-4UUFOu{Fm|hUZ>21AEh1 zOe$0ykCg;+gTlvw&gQS^h!CRIzC&0l9E|6G>S{F!XW$8uS~Srn`=x&a8_7q!62ZNZ zVf=e{bTsOp2=DBMRl6^&j7>B*w?I4HCr@wIZlCHdh5kFO8}Q&|D8g@ji@0Schgp8o zoCSl<^^MqS0+7W4)Y@wUF6111wQ=h&gxx;@+#Mk0kw0`>{}*x^9kcfV5Jkw97G3ou zliyim;XW}Qmk}-aHow)8W6D8L#xf@%dYB&+K=wnYGk%{@Urynz#^?P}Mr)IgSl_0$ z+(v8>ns?G_eH2Jj@#DzeC}omj$t!&~pr%mB(b=VLb>fsU2L zRb(p1OFm~kT?-v^WNX2woyb9)1q_OTf%1AM=;ongZ0s!ZuTVb~U$FcAIN4>ob^rdR ziNyT*r8xmP6%P{v4o=eqJuwbU67X;4Bz)=}QqnU$2!Ro@|)lCpd?D6pf;o_wSP=Yo52}Jv$`_+pCW&^A(oTN9l4$ zQtu^{$k1q(c$MF%4f6LgMkD9E0RgK@8^$1;)+vNFJb&0z&J!*!r9TE%lLRPM>?VIr zf+yDV)XSR>yZ`y7y65&e?{Xv6cVc`4Z&GB+`(0d>l2ln1F5xOp%fG33fN%7?;cvH6 zbL`-)BI9oj?{utN&c`juY}Mb*-=aM@|V27bUW#9K#`aDYh&oL z1#I!N%QV~=5C-QI^>Gr;V84;>ao?*G4Q(PiY`JxFO!xvHXps;Qxs&E(%lQ08-@>qS zxoP~h_=jAy&kD^YUJcy0mFhh-_Ihd|V>z|F#DsPdwk;45`6d}l>?*T&^s^YT4uygV z3>0Y)SI&*Qbl=5wVfnCla@}GB%~~Q+jS zA%6lu>C|8^a<(Cp)eYZRiZRv&TVTtwj7X#CU)*=q8VUeUP6O?M^Rs;*n8H_KKqSZ` z-b(_aGL?J<^#8|tf3b5Nh`YB!Ft5&aD&x(~%~sw25L~DsxZ54t+)_evP9`MoSMLXe z0k|<%j=@z^xwJWA@|q}8DfY~AsQDq^v-n2+r#5_Hou%%Rp%Q$C4JxyGX2WZ= z8UFUCg9#?hujuuzfuW=rC?$u4(?2C^xibTT7FBB70?br!=lJ7o-vW8>l$hXj0E(X1 zyx^7MS%I=;&_BMR@aJ90VvM`FW1AR>NJg?+b!|AJqK&b}=Jkb~!K@KZCB5k01e11I z%A(e`@PN7=zqWx(lg8>cbg}b;*AJPGSSALe;MX^gTUjD&s|nW(J2n`0^>#}v1Y2%` z#dz_>J#k0KU%&Hxn5v$RAdv23l;@KwlfI~m*<-hD3xH{Tx`FV{S0tp{)6k<`@~d{z zhH?*ngcY*LgzP@5`plGpC_L>3vd_)GHwkR`=i8ZlV9GXC!p??)LO^5rGDHG}cq8|_ z=5dsu4S_lwN3N{DVrx8?d%sXgeOA33SD$No8W$v;&!Ybxjfx(axQSu+JGEY(m^#h(#{7ln{C5`rcz3Rc^aZ*>JE7|(C#-I49|*BlPN5lCT9vxXY!nW8Wd2^1%FrB}NlFYR0>0$@LXRLo#{}nS6-v@QvyR z9JZ<&(Gg|K4|5|~{@N}H@;*Aku1xTZ+*+31F8eQq%`v5z8Oee{b4Pd$M86SjIDKJx zLL0sot#j&)fe2d@HMa}AtsDMhB=787K+#kLF$G<&5+HNbMjA%BRDhzl(&S;U(Ug85o5X2VTk#CMOW0_uFUdl;XryiDYyP0wUMqu+-et6`iuGn-Kd4inc7*95}d)^vak+BT_|1Q`WC`) ztr$?v<0j>N($&?yb=tgz(ge)%PZdbfM4J8MLO~`$#Q!ic!Zd?_(C)apFpLs$+Tv)q z$#1;a-jWV+oy2)BNIX!zrXQ<6vHyNoSY}5zW79!!14x^%-hM1@9{wGFE2)L7@+t3Q z@CTxga%26<1mETkHGQ-D=ST?Z2X?@STWh^M4NY5`?QE)+vo6D-kP#Syy36(mlAXaS zqvm#7y1m<^h)`W~OFd9z#vOvNPawVrsNn83CxBf@S7WZrhiW2>jsRg z1G!>5YRGk4ucvnWyNMIVpk`6Gz1gyl>1pp(Xamvc=#L>XzU&|dK{cbcc)GrJeAl#a zK=igGsY<^TTphA=I8^&7f(Y?7%}!wJ5%(qCaUhw&D{^C7fpFrCt?sI3EdfJ^Ze>0h zRXm+q$*Dg(6XC}37K>W5*S{vQekxGSdr069A*ye8sdMnP`W}75Gg&UURNd?KED7-B z;NxqYEqhg(%mzl@L*(-pX${Z9HM!GYE=|?oX80KW$}w{rL0!bRSF#x5PCI}N`V7_s z{W~=T(b2cAg4%>+WIDk}Tc@#VSt5h*YniJHY8ieqw5<*95Haei=KOXTJO?Nm=bKkI zZvr<*=}tm@a#a2}pARt4_L4P1>lQQ{zqfo`{BC5n$eS(XyB$~}k;08ip7_VQUTLK(It#(EWub|0%@QpTTQYSqK1L~Of7 zl>{Y9`s2!|E@5?|$+A1bJ$X%b&lc}EdD9YSzE&#A-4 ztwt}8yu*#V3$-qC+R+d`qWw~kMLEl*vT-r~*D%H39ArT1#ur_o++?0>H_LwRR_^*? zuV-*Y=ElaLdtf^EfoO#n_A!wWt(KE#4#e%jDr z3$X&z;^JcJN>Ijcbs^b~S*o6*0FK}-^H7X<|N<>O{3@%zr7GFjX1ettUoswTaljVMJuSh(&nqg{5~np$SboS|)|lX!GsIi)-`_FkSN=UGKhm&fjC zag_A}N*bJh?AR@87B+E}Y0bFHf__-cxLqOc;iflqY2$vr3j{Y7EOJclbKY)A5xau= zL7g8V;N$&-HP!_(q;v#X@2wg~glI#ZI`u-ccShhWz>!K7vXl69O^iVHVP7}-V{wf4 zntmhfw;*`7&b{9$apKn&j%>^tL&Y~0i=39di3pxZ2c?C`+>6L}mgOi!gL=TdiNw)- z{>xG&m;GGgn`5fbF|bC-Jm4!m_@zpra3dOd$byrj@M`pxXS+%VgkAGnT5L+ktzw~e zQn@5SeLXfm9nBo8LaG@I>Nr5qkI#yH&!3V5Mfj$>EN22L+A)^TT&j@(wcubA<vASGy)lS&fioz$_=vFaBhJ2p?3cBE$9Y&t9yXe zK~u_V?$yA{^2Z_3Wp5m1jF6MBC9*>%ayLTH5XvxVgd+7iDrnyjq*&oKm?n>_pJn9Q zA=xauTXRbCWu69;Gqr>0LoBjBlstAg7;}S`eLLy zcIhe1eL%g)&nUu~S9r2$6_7mApdrcHM5B=n5oSeo#ux}DChupR`o3jmp-{}LX|_rM zMm0QGYdP(@9n`qkZb}AGT8Q@<9CyawJrE{;^HGrb%OIfWiPo+d6O(l%le)nqzereI zm&HDNP-w2g3w;G4t;iUN;P2PSM*g~5s4Zo8|GP5IVuZ^k4g+~y?TXdFn(M|EOk!T8 zOLS!VBs*9>zyE`8k8HM!$;tTTQ{JhtG9;ZniZ1J4*!SBFxag7NL42HLGDKl3hHaXxk)_m>Xh!Fqe27{&BdA2EP|NK^DHwCEcw43lgij|IYe1o6*>=7C59AV$MtiGH zX%0dX*x=Kf3nS6Eh~|x#_o1a%lb~H-W;=<)2{5I(K<<9hVYZGIsWM&&^L&>1)U+kHs7zxV8~*E}HezivLJ6Lr zp>YT9(P(utck5bEo$I*676j5%!9aojC*4L8G?4UnH^9%19NPyq>;r!=ua418%8}J; z%)uwyLZ&Rp5Atg?$r5cv0~?KG{a-Ynuh>#+Uw6?-M`kUP_ma7y5+2v;Xbp~JjtLg> zmEXh0r4_snSkGT0%rc_c*`?0}B8tntQ(}ItL6lS{hA=+}8;prZd(+_bqg?|BsrYec zq&=yUSL>1fyw+Su;Ulac9^9Q?(O~=|OuE1r4I(6@EGZ%AvG&Sy+aPiT?`5au z70mhhsai9*)<)+TTXKH2oFYRK9KjMg-})-Y01H(7k{3zKI$yCwl+s2o%S-6D0E*V7 zl@3OHik_6Y2s5JLJz-eL!)DJR%3hRZj?L7aP3P}DEW5J4;9@SZE%&n>S>nT>rC;;2 z=Cw^pt1*-lSxERDgU-WuE?-!YM^IgJ79)-j@|6WeLX?|SXs7Ax$4dN9!Ln*76PYe) zyToYZ?9)!=8}OvTt@B{jg48pt{{)Y0l9qIoDA|Y=-)eV~%eLt1> z#To1{u(Kq*`%^^*H2x;_&OnhIPt8HvCK&93gg;DZX8k!XIcG>v@koZTMH@Nx&?hCv zFJ$MG!xPG~lp`m?NvY~}#BAe=*nXs{Mq`z z*?r+EG7iL5EcW6Yoy65GHP1T{Mf?P4kJijUixeilw-S-Omnw_mru!(F8xf=CN3638YlC)U9x$vp5-Kfj4}U{& zw|O=`yZ8%u->%l9C>v+yW*1oi=V01cfX8t14^FbX6MDvH;RBsB9E&)toqcvE0|xbM z(e2Di$Tp&4i{+tgfE2GZA3fDXF~Yt^YJs|pSq$|GQ9Njk$Ubfs3eftX7_pono}Ul4`wW+GyGeV9`)fnt zE3XOg-rZYrA}-S?vA)PY$!z(s-{cR)>}5Vk#Gtyycj%63(W5@5`fM0wCcb?S4NRC$ zPogt!7l%ydo3FBIFk6%@aL_IYssv=_>W;WIGYMCyJWo<0rsYovW=`HPuvaoLQRySc zABU3Z4##;r&W_s<*Yhw#LL)v1*OB*2@RFOiY|Ef)7Ap* zk-j%{IJrjBq^p*H>sEY-pxpRB{mWQY`vTYfYxkb}#6TGb;Ud{Yy62zFXG0tSG5Ltds~C(`SL(J-RNd zrr-SHos7BCWD5>~f*0+6l9YE7zC~IJgA^kbk6+ukLtF=qg zYX4-BU(>Vx3bsKKYiM15dDQwrj=598du$Ch7n1K1;=Pn zjntf{EHckZcm~%I38Kz<_spuDgnOEU+c>jMc`r;WTqckS!JX~>Xiq?MKsS zzYJ&;aE5_9E2XTsEu#EpddE}Nooaf&2=1UG2GrOTB%)~_U6+4Tc)KIQSOQqT?ua3e zeEtw7AXuwp%qIbnR?t=il*p8hLNcnP_!uMrR>1XAs{@C+9P0TKfGBu`` zJNs!JO^p_@a4UtdEM%AQJxIIUNC$DQFi|AdXe`R=lZz`7dwe-r(rx)s8T}sXn}JRIB z3b$zaRSQTPHK_ZAs?#kISMLrVD@T@|O6QOG9CGHlp>Mum@Z6`&U6lgs8MV?s@|MDD3Ok?Y_I@mMD25%>#oP^8 zRsR)POtUcK+}obPM9q%rnc#7zYDOI4F->2UhtG;g2vifvfa16QuX(QTdSS*Z|{fu7eA zo9RC5OHO^_y;dr{-^}P&NaAi@Ksl?5Rp4Tn$w0AGZvTG$%lJ66Oeud5wj?K~ffa4> ztKtr0dUMNZ5oRh5s=f+ZGE~Kp($4oXisjNR-m*{O$Y6w|4{e}v2XfF2P470dk0nyH zP^$Knbwi1^+4jHDi;3hU0WtQBX*+CMYGdj02juX7!h#RFe5oDnhAk8}VGn&#yBJmf z;M;m-rn=9xj1^w8I$-jBt`n;3dvrPW@y!C{T!Lbj$*tEYwzLZdnGBNbvEPM5A<Qz58DmQ3;DU2%V)5t0erlD#cKXHk`j#Zw~=DYCflsr{WoPOM}|HUoaWS#z|{-PG!BW%l^E82EcC^TxNlq5=2Z&!yNj zos3mJDlAcJO1Ys)JtSYxbNhEo>uE0>vPN8fIu+S*6-eK+>uDm;-V-Y&h9;c~5WE57 zW_lpqmza_LWMdU=se4CmOILC+e|Z93y#2=iMD%k%*WM9{e#DxI9=s|Dx`9Kdv=+%h zgZ7;it{0+UX7wL$+1NH8F(~fy2Q=lMV%xneu2)I5kuq;sX|)i|{|Mid2gorXq>8XY)rp)d8ds z-0DS%cp2w^q_%QTK{*apD5vEp+MZ_mERm=N*$%}1b{JTULh@$`^8GKKxDu&-!EhK) zQNfX4Xd8Q-`B#~WiHVV9?QP{EIBU#4O};CsEuoUkKLU<6rXL}5^ z+}bX)j+TIilv#XjSX`CL}*b6)H~qxWgX4x^7y*=y6fUOv5UT02^0A# zJZ>N!`ib_Q^$?eF4;j>3xCF@bSA|B8fuc`5!dxj(8aY)zhSOh5S~6uVCpmovA&Y^( zB4S6ZjeH+r{7BB6-o`ybG?)%b{jk*FZ)?R zEHZhC`aGEqP2c1dP=(8ADsOW$@2E(my)nRcndr_-z0=}|KoofqWXEpibt&AMfhvlc zDTI6<`gxKXerli5RrSn2_|SV&{4S;4^R~RMIy!_SAacOz)t5eTK}M_zJ|3KXS}_?I zj_K@No8c+(eyd_Up>DR`WePZ-VVs;w{fH=zhl{oRK(v`;f4HI!{(4N@wf%mJu<9H) zXLr5Q!ud+5Ov8R(?P{33SOwye*1>Rus zzJJixG7a>!7J0Y-I$VNNGY^+ineEc9`1NP*)WTSvL(wdKC;kb2zkS+kQ}R^dc--dI zhxS*Jjr=PY6+gE4PnQ1r5DWAY31tlI{LP-N)s|4fgeS=a=$EG0UCjB)9fVIK3**Rf zUJ=ssqqxiyo4W$AWMTfNb=1t5=F~r4cXf8D{-|R@=hxxwzF3qEtZ$c7nF+ZQ0}c7U z7e5=OhDg_LNKJi;s5l-U&2vY*R^b*p8cGjoJbWG_2XRr;2&!|R2f zE=HMGQh6z?2wnE!E;oVZ%5C#S_6d(kqDj08@AW-iOEdJ2-d3;=*-tt!O~^M1>YD@- zbh#Ut2nfZqP^X>LgkE@Pt`tkl!4hFMrt_zoa7$Re)WyuKymA4*@oXH7+!r_ z*u+zn`n93Tu7HBQm<+Ectm86SQDK)TF3i72;Oc!Hs&>1gmAP-ZcXbRe{(ao4V$Rcy zBv(>lyNQ=jBC_~YEs`bnPtAz*+(n)DJxaw*u|Q|2TFUoDSmX7vThc2tzl(?mLQW7O zlU5UtETkkR8sF^-5#9Ie!j;|Gq+bd;wAQ$AVFUgXjs)`4-Poq1{lkzQ7XF^_S4BEq zNY>@pN3%ioG7gp5Vz|7Y4k@||dwbPPfyFdpjRvoMsh3GP&I1MU!H3`godM-^k=}6? z#)^LarRutL9k}T`u3ftvsMntn`pxt=mV*(lHUa|I^~l0`g;5N8M5zwc#RG$6bBGs~ zm3)xY`?EuGS$=DAL9^1IafkxB=vho8JAXK~Rz2uc`8i zM>fL5$BVKG7mGq`wlTEoujE=wELPzOMUf-)Dvd|aPr0#69X-g~zg<%!FAs%?kd*yN|3w0)9!NZI zi?B`>!-}~78f(Xt`kk&`W?oXg>tj5dXa%MPO#g9S-7!-{nAnbJd~`*J#A!uDGVmt# z=5vXA%2|DCr)p$QD&PL6=&0(|*_H%#EQ<<6MNufD#Q68>8D4%};p2E_gJnvB_X#r$ zJaR+J7iEzA@W&sRGN%fC2&6JRUH<>VO=@%bX#Apud(ft=~5Vg|b0JL-%$GBjjM>1fk@EcAA5r z%zT)il8cqac&Pe&WvjNN0)Gk~_SB!(1%3s5iT5!O?4h)d_Xu0>WVm0s&8WV9>{Gac zc+Dh|w}w~zNHr1r0Zw@vV4#H~ess~Xv3CfKxv&37tUm({70T>T!TQZ#LbrlA zE6q^C;253u3m(Q|GKBh*fVM%ZGdvn0jz@VwTk(6-Pcvt+L}6V6xDW2lWv|lMMH?d|E zg)UJ2Mcbe}M*N$HhV3b2d2Z>;xX+dRKwImRu^fRRH4kWATJbl>OZ;|gjlbd@e4 zUSMay?yTESV?CvuGO1NUGUrk7q9`gbgQ@2S8Z>Pr#2@FdN~U@7htT0;rt_OqX~D%? z#8}_b4N`?F4!%e@S~7TsDkBDo%X_^|NETyp_p}VMDm#%5lW9c_NzxULV?EOcoqD{) z!F}6X6G3AF*28`{&Q0plH*W~6HimcAE(r}qh1G@q`0rFlZi-T@^C{ZLWXBdV&!-v) z@91)O{@Bf8t7df-(3O`UzeWxEbNYyZWfV?Qo3(leiqU5)PrEeTb58hj(qE)kBjVFR z=rZCN)H^Bunh+-=c($6*7-W|?dG?doQNmK~b?yPSfYv$nD4esXt$xca`P>rywY`RS zb}EX+lr6qowv<|M^|vs&LUM9#J%xfMk4ga&iZP(rSlH3bl$V$aFqD9`zoMT1ZqK@G z5#m&Cv``d4O|umLfLEsaff*T?dqwWE#=SpgHNT&S%N*i4D6QJKFwb!O^T7C5Kx6Zh zL4%Fk_?JL}F1Y*M3W+^#KVnpyW+v-r`)NyTa#__7Eg>R?j$wft&y+>+ILbQtV*iQ#SkIrL?%K48rY)4X`t#;UuGX{7)y1BJMN#G34HC zGQM!x?6OMg`#1NOF|IHOcjiWBah@GBi>9l_6`ZIgv>U_sCy)Mb1kQJppQ>>mTvv1o zL4<5lH0W1b&X; zNP3h~)E>(e!NWCZHRe%5|B?Nl{8=i%7Ywcoo1wZrR@zgm6&`ieG@{7Pco|yc;jl2? zEV^fP29jIJX)^awU}8!NnUv+6+z){Ud?XE43%mv?9E*I9wT&HJ|07fMvOk%=qMV$p zO9wXn{BSceyl|LHVK?cW<}|eBCr0`ARtsb!It}JltKwG1>XL){7XX;adjM^?yph*m zN_%tR8$$%EqHk7Tr@<)S`|(>2Jqg-sBBGRfVN5PWbI7hT6UysU(y>9cGOC$w2bm&Rf7$IQwM-~(Daw~0o&u}CCP3V) zs@_0ze&FRn%1UO8#bh3Ixz*=!j&v;)!^|{_#3OfVbyX$#JXuqHvbig5?5Mt8`{Ba9 z&w^UkdrL&lIp!Ym;wEZhDJ!tb<5#;N=!NZ;5Gu6Fz zddSPfzwm%zst@ZmJURhBm*o1C)tJk$yi<=I*W5LvjFhg-+EE6?sPLXMsy}&SL1okn z;Vz@7H?3q=!mGnvgy#|rBWQmYCHR%^%U|_8kP=jkdQP-jy_OO796VBl-#W$6**Xmte}rqkF+{E;m2lX`!K`%MCO^ihZQtxtQn2yY%fPlWp7G(#*;|2$OH0;7K}k}Pb*11qV5F@|PtJ%e`8zXRb_Hr&a4) z&QDH#Pla;$a}4N50{n#FjKplAs)K^*@|*lurD2|9Cf4-iTq^o1i=8B@<&p>MJQqy7 zT9~yWIuVPNDnWh1wJm@PPv!;L?(Cq^&GKfb_9WLgsF!n*5VMlZ9Cb5=OjCR!lV4?3 zGt*!u*35NS_dMHs?|}h&R9FbNax$#pGBdUCY13$oBvCjZZ9im~n1<5cU(=OCkeh5@d45>T|b>*<2JkzDerI-+|Tqr`ZlnBgk^MC@Jak7u7e z7CnW#G;62hNT{lLvjUU9A59TGZG^J|X1b}F$GlQ2bw4z8cYQ=83y1n>Q^$scdogGQT%iv@odF|*iFsBYZm_Io2c87}9H33?JSy=w@ zU~`?1o7Zs}N>b<4XjL{xM202xY)IBy6(1)gWc8d!6uGka5@gv?CO^gHC@Im59OldD zdI^b?^sFrojO1iZ<02xFi-N#E^7v|_@R%kz;qK6)v8hznr;H$lkxmH-(da=MvzBJ{ zOnsE%b2iMQFp(~73WZ_l7OyA4-|cSz+~*&S8qOq~4Db z@J|C0@;%686b05&yt$fFyHy;0l;nYI6F-6tcn+3FSAEX9RI1#AGR@Rc%u%werak4a zrNo9vd9zXdqy;H$tgVWY3ttIIOZS%+zA9IC8kIdKGijeoV#GvjXDO~t`ELDdB!y9t_4kywStoAyiv?Pv1OYWPTif0A%4*<1dwN`ypbvfL1Jep}{X z$Mi3ml(5iYDU?0wmV;V=l&CH(dYtNUVjD$P2lkp|R{LMoC&{xWvDRA8Lz2h{ zAaF(qMtM9W1k^o~ZIH_gTIn;(;kQQx0ofdybWKkCha$SpKxrUS$Xts#>PLC{$Z( z(9nXUuCa#qKC0nH+tNx$swaX97eI@T3yUKEmZMI*$wCwBSEP!LM-GpIWUR}S^qYt{ z%4*4$oN20`{<d#(4Jm;bNv4PxUJBG&;Bbp<4Ly=vUgS)h0elz$e1& zX*u6QmRh(B8Kw6Jy;$`|k5yIbFQKZ2o%?D|u%4i&Be%O#pQ^8T&aF1c9F`dM7ObSyrzWA^V#mijDJ9Xw(m3GfGoWn+~* z==+_{arqJ|1#bb9l-F}4{9qqgai9lDZ~4=)gue_`j&6wRWr$YzpnJ7r7W6yeC*7Hc>+A^|UzPp+HljGPA>@H(M+|}jq z3n!4Mac?MsMALqjovEvEe3j_7KT|c2ZN=u68JY+7C!@X#@R0yN zQAm2K^ro_z-tQ4jomt*k8858=J*v4_zCzyCv_4;9<>g9k~{9Na#~K?jTsYAsuR{w-_%cvE&@QBnHfJ`r~Zs>THEm8pi*El z2M=l@BV8<}h_il%pDMl%@~;Y$UJ!!?mB3?AsYiwNL#~7J=Fxof5){3@MlJVe{y?t4(8$%YW8_f44uQgnC47f-;Lyyn4Cp3 zU8rvKU~aDkDtNH`q}an2T5GejmnSIq&Q7^YsEJ3DNn4g36J(~xm7SSm{+1W+1chD} zGdZ3190%TP1uohv&gy+sT&Vh$R9foXCuIKoh}D<;ONfOGr#|ok?;W% zSLOZB2x;ue7!%Lf0oyIK7E1+rFVAGw?dq{vRyJm^MkkhG%!9Ezs%UFr!Y2{0lF4V9 zB}sp7xS}&i(u~%GhX+c$X)2%QTvZkCYwO;t;X&5X=kv)ly>C5oDzPJaqL>*qr@6Fe z^f{84O=h|#X;|@CjOMY?KQr3TwOo$EmL@blM6~ukQOBxQrWD=Q94I0ZEwE>1&+Nxl zmAtd*I|3SGz4xrN2$|YbbzA-vGwiq_g*!9=7#`jjB7oy-xKQ@-8c(S?YOo33v_>e*}TK&8mX~V6ssTFdPSCoN9&GdX7)i`zj(> zX?N5pzTECu#2_oXI`=K6qbMj}-shEdrWbQR9pv0kWB#fWqS1YMIfSFi^k9o4zK#SY zBM);45;T@PI2Iwhn68V4iX<^*-!G}`v1s+u=^l)+WC^*Qzls(}xq|7mGRd-`aCy6)-R zx4Kb9!jI~LNV==MG`0M_qS#`ikw88?yUQqW6n5pRy7=V!QxkLUGjotOOm(X;M!N~8 zP7-VpCX9%{0FUJfBr&s?xgo=Viwb-*6}$*w?`8WMA|?Wevs>mK0KD2)dTViUEwDW4 zRT24V!Z?rzkk|D%M>0uA_7WI16P1z)7U*kQR;{62=*)N)7W-L}k=^%qW00@*!C=bTtjgk!Y~@4eOitf{gAG||GD-pPZOcn+S}Y^Z$5XYypwJ)f7+ zE6n9-O<<9?SU#Sa!Z2f)%5O8x!J5sVHbAn!lwbLIjYU%VW_xKtwU}7Og|UkS%wAlh z(U1tip%RKEz#bqnhCGw!WSmd5R{pji){%nMsLK)IVH%*H91sx)pL0q#-HFD%#L5(+ zH`uYbyOc!Nmlu{X&&HN;%etmhvG()p2{T#JO>n4zOA4H~GVlLn0em=38I?QC)qA z6q0?YWoSLmN9h|RvNQ{6^*-bkwg!N;HdD!H7DS$T2FvSDCTjDiq^^4pP%ig`QoOHj zEGU82z)^;PYUS@q0&O98>sD7!^EsEZ!wGRwJHX;(i zHQdsZNUD+N1wfvONP9kq*$A(4SDn=Q%n8wf8dfDHE0K5sGb&T6)F$AZOZQC=LRQ1b zPt1H>?z9tFE1G^Yw@=E`JqHsd?To4Az-}xCQ7vCh{WjJyRe@*)d#W=njbRWO7MFMq zHCx*xfsCT#t~FM6X$UtR`bG;G=SF#baR^V`;=2H|nzreNDHl=H&MRB-FbpPUFzShl z<}-*d109E?C&K!g@_$zMfS*2`gUz2k%^yTtwNg7;rNL2|yksnxQ_!vK&&r;B5J)`P zQKV7(D4k%`m}3a$y0M0=>iXyY)U~boynMNU%JrHayS_f=P|}d}(z5-$_j;G=wo-Fe z5)(#iw5x@{9M|qAD`!+goc?*h!5ygKEG!W-HyJ;2u;F9*Xj1&hRE*#mzM`IgR~Mc3 z?KoTSJU6OYVpZN22T2!PEyJ7QKAC`{HCtJ%K>#4>iOc0L^Ug-{wp-{L7w?bF4{f}e z+d>Nv`fnQfEMYepm5CCt^oNWYg1A9M!E)<%NAmpn&Dv1E3&YKj$O5)VVE)B*`HJNca zCRe|7^V9paV=*lh#`27=+py%g4L643JYe*b()>Qb#xYLi=lCrp&m-pPG{g!*Gd?ql z2iam;I+ZDJNhQ`HJ*P`e#u%?OA6_ChOK)ASxHiTZV|*AT6G3daF`kK!WYslI6ET}N z%m5_rn4GB)#6vW!G%kxF$?lcSL$Wm*jN5Z$hO!vuYM&g(;)H!${4!^2pY*WClpw=S zMdUTH=_Yg-eB5H_Fp~1j00#CTh_*uhWrT9_z|1Dz+ZL-lgQP^%qB#JKj)PKkMtU(T zh*t-tf5_OxsRgZ8d{VyAZy&?fxM0hyHun=T<9>t+oP7%kI?IHxfaW|#Bn}{=*NlCe z5%nIWw^dj#K`w(9?A9o_R;q%Rfr$FPua=+Ftjx(Rvx~o(ph*sSfd{MONI}lfRrL=J>u0Kn|tGBc5F znuh2C0MT|B2EZF*c$6kHhSa1#i5q}m?iV^=9oh*J|3SzwuUE;K!=L11JXud7=CH(c z0x;DN6>l^g)njIMs$)r%Bq9T^g@%K4w=71H50XE#A(s$hE!q5nc(yL7tp#G16vgFu zDN30c&2~X^Fgj^VnJP1s$aWkiBhGx#43zSd2)Psc<@-?cRykClzEQkCRzbF^#RG+9 zHAN2A$}9p$otX&7kHky(7z(s{@bm3AeYp%QDrn7MBljewv=ZG2sNaiS03up1 zW(oQeE2)44JhkBWyx;4`m_ z)bwe?6P1xhFap!GM@1QBgTUj080jgV^DU%3VH5irGm-6aZB(C$msuT&DO6DNUaua)7)q?Q~fUqW^mOGb>CCOL%^ zOLI_d{@7rQNiyilZ9hkPf4Xp9F|}?3{_~irhrmyW1xc!pC$(zpH7gojTSgm)0L|)RnOSv{y${DEE5K; zP}MZ)A)ct2;I*=E(rzT8hx?yF^jezGR7(@PY*t>Iy;EZg<@HZ`_EIIftC`4={lrtV zt_Nq~L6X8;>v*y*$yml;78UXki zmHu*thCvAbn_F1vm76N|6?PtPLro?1iaIV`)M4f}Q`oyS%h-zsI|G?_R&VzrDS^{r+aR+uhy? z!5+PKlfKk@&%*#f8k&80@GHl|W0_C==Nphk|Nue54pbhKul365#+S!+G}ME;GL z*nLqPjoijR9o^Kld<1jNwIC-%l!+DNsM8h2+V;qro)dqWE0?0GqbAib%1shi>I3Q} z=UBZ5!h`@Xs|t0)>t~|6ehaSLe`7llUoDDmNb@4Q`x;3cXX52vL6o6f~bMyN3a&dFYxx=#Sb?YxcZg*W#9L|`QtZV zfBp53U*Elb`*ydxYui?CaEcId{ruTxvl&{`Gz}XI0OXA^a&Ohz(A$I9jnV2srboT% z6Wu$!MAfefUKt{u+x!6b`hu1E#a1K?snb^v9Vx|~{&051)#n-qN$jC=?4Boj;NU0h z2J+*5+{sRZ0pM7XE*S#{gRAbC_uf!Q6>N-wN)rnPMX51_RVxe`i3(QdnVI+dy)|u~ zn25+nx|_*!qwLEdB15v{Oe*Z08-^k5!!RKg>XJGj84ts#Y3!%A)yq7|?P0GcGW>Tb zC1~x8PaM0LS#}~JIzPrKT_8Vm#%aCNDnF@*G_h115SopxJ0OJqpr)f} zv=elex~cK($xzBP<&35I(ZYIG38p;3Tw615a>{zz*)HW3U*tNR16?3}uCM1-3u_@V z6Rnc!8&j3C)UQKLDe_tGl1xjMV^LONPF60*lBy44I)`rPlExU2UM8-d7miZ9FlwnM zLKG@ujEt8j{Nh2uW`4I!W&lqqV4Q}j$Z@)IJLBqR3aOIws_&dIuV_&>`W{r!!TbSW z9TF%~aDPw`rYJ;c()(YPNf~<<;3!+Hna{}9(1`n3*?pI*U)Vb7xvHj1FejZa$jW7H zA`=jiHQS;0*7~+JX4o~pfA`%t%-q`hPoF*eI@#^Q#p0%$o;630x5nUNF z1`WY^_&v9-xBDFrgR$0F`^jeeYJ2^~#r4mgUw(4Ye)ZK?fBfSg_it|B-uGSCUA=f| zz3;ebHXAn1T8q{K@XD5%X5a))G}SOSEmaKxrpWySY;q9iR*Q+nz0RN4I-F6U<1RAA z8ffZ59I1HzNl+Sx?@WFjHXxERzS6jRC*E6e5Q#QtZIw5C=4h}!l4(K$Cr6=EqVO<~ zwYK5rXFq%S>eb6Hzx?v~^XK1w`|Y>ie81c6L;`ly0U+A~fB_LD%`QbKhqE3sGF+XF zBw=JP_4^F_JdiM#oYt|mX%OykrFfKr7^qC?6G~H6*=Z%gQN*59gpDxFD7(zLSWB@) zhUc8ZGQ;zzw~(eo2S$yBTskTVts922+t~w@D2zh)8^=X)aN|p3Oh|AtPgmA}_kJ!6 z$lP+wHFqwZYB(Tqde4GgvuoT%ivVEaY?rGe&ksCJ9|`aQ2=)Vc1`s7(77mnwn1eQp zhMHMT<|1m%*@n*o6vHK~*(D1inmlj>D(p!pB)M9P1^^ox7)=A`+=kc^vw!#bt55#V zzj*n{C!c)r>Y_1t20skLB^opU9Ad;LVk^c$X72as_sndqwHtHgy=Syw7-k+|fiHjd z=`X+h?B?d~kAM97n{U4P$AA3aZ-4ym&Fk0K*Viw<_{D|2HoP%pb`)GumK6X;&EE7d zk29u!<-#DR(SerC&c(?|Eq)BVUHP?qXKN_*H05MA#zeeI^{MyXU?O6k)*Pne(p6Or z(#9A!OgNz=xR->-d+(fML&lgyfb~_Yk_Q^O3X*3zHK71ci^;5@I7*Z4M3ZsJPR7(I zdos)Ary;dJ)D`3^D>;tn)WflokWX2x@FwHbSIQn$_(*_H1ik5viq<5HVN_P|)9TO2 z`S|CIlh;wOreo>cYRaJxK%9ZG%z_3lfweZ7fs-ut^5x6F`I}c)9kzf2B>u3*{1f@x zg%6!`KDzddAp+E}E-%q_7oU9cw>LL`^Tijx_|O0OKmWJi{^924rg`^{h_+<5+byuR z=j)-QC?{+GMofC-KFZV^X=rq$Oj62l$9@|=QBvBX-fx+QTLOxX)_xiYS{}4N*K*AZ zyis;4MT;ZOP?amoGt5@%z)B5MwZagYt31o`d3m1K!@2mshXYX_o~6W5YoMxfIIbJ# zgO3FG0BGB7yvzqg)kG_#7nG#o$P06{;fHT$*LuZjt`s=d?8UTuOw2GdhheL@pwc%Q zt%(ydJlFa~*WV7(fyTID7&bKg&6lq(Ecbh7{jIh3RcktQe!si9x#1TtiAd7k%FW$@ zU=8ls+;g`B49qDc*nHdt(j6Gh^C#ruIlx$>nAreY~OA1#yo^ zhiGnZNN@bqkTLGM>J&_2q~f)oj&=H5F7t5mtuVT8t!*xG z>EW4&J?$G{v)SBFY6$1K?>)0AEmhV|eE6=et^n+I4B+;5_~esMe);7mx3~ZHfBldD z?H~W~A78(IeRXwZE}kun>$D+Pk;Oq+=~}!VN+I~+b*@&~&2~wb)VNi1KL-)(Q6Lm9 zO>D|jBm}RV_VsghP?FBxSZjT8Eiidv%R}Zwo!Nj7!sJf&({p2f2qgCJ*quCO><$;J zy?iW6)I2tY^E_h2_49U+PUsen>d(ZcwW34qU{LdawhjPT(uFd7?<4~x zj~*LFJ8&9!gvfBAQR$Nl~P{6GKy_Scu~vuB3F zZI(e2t}rtvI-V;DjZqPxChwDmljB0j*p+Y`XCu`)J}3EGmDFDu&Yl?#L(ur!FB>>)xJ!u99z{+h#^5d`M0rb9sh4BVJ&uEM>&DJ^RVk~X|9UOq( zdS=D|248SXL^dV?1hBUrfHMHF8C)8UWGFJmFxb@2j`Euwvp5sT7_(xKC6}hxSb- zlrUI3U(A(E^{*f3hyjITu$FrU(6%knN;fe8kOLO%#H+zMkv~E`_xJZHYT&unYmH9g zjPOX!nZ=xd>!B*jXK^Sq2tlTDP9yE-{H3hg7`Y-s=|i4bXT0~I$vG_$>iG zYc>`nDQ=07s~<-xRul@}A-SqndZSp`up;CE3?8NtzLSlmAAbC9x4UawyN~&K#L}G6 z;}I{<+dk!uCN0@)YqVP}Q`UJ9Qyko2bNxIB+3=kuX0{#x+JOOhGs-K|F9n%2ke|JO z(GLJRe-FTVZwG66Ks2CA`8S&_0H6iX+m@M!hKcBEFC7>F7<~t(fsX2)6K`n=Vl;-bn1{YTu3wCmih0(}XqsPJ=bZ@QO%(dWV!h zMfuw!L*NnTq()bCEwB7wnWsX47vM3-@VPdqX0G3l%vP!vc$L$Qh}yR8y3QCF)*dn0=P9rtL^J1{1J(*^saVrAiQ9Si?j6U-e+dJo z`dI1u)>NKp5!Uo-ec_URmxLTOs<5V^g`UTsPCqG;&9VR*6WkTkuZc}eMix;jYIouZ zQ3)?AgW~A4RdyiJuw}|b(?DA3WmZgoPSc=Q`aMObFzt7y33;gv(hOW?kjBi?V_wD; z^Y+4ce{wK411HymLpr2vt&bu@S+%7*`iPX)RQDJ#hd3U4$IM9v#}fFSlAs{YxktAf zVfJty1)BP5T$4qM&9JKTBS$aq437@dCof7imWZl6$T!q8`qnhy!5aVk*~KT%woUk| zt^i~eYn2(p;QPL3?l7XK1Dydd3`5(t;xI4_`1adxx7%&gG-+twdsgZ!H+5zix9ixG z$w@v|Rd}*CtV(?|$U1Wrlaw}(JhaND)jH{UgQ#5)OJVB=01}mVPtO57-`@lH`P~~L zdN$k>(I@-w5l7wkz4v~%+YQ6;yQ{xy+xGU^H4%Ni?Evh#Bci(oh})t@=v7Dj!J z>yD2Kq0e;6>BAxFTNW<$F8WvXJo)h8G5_!qFi%y%@0oFXNJD&9Tkl+LTwIwr88UGv zqpw;>m50J1^TbB8ho++ z6Ax0mWIvl@Sg67=~fk@Av!t{;hwz*=%-g zW36?qW#)n8VmA?yC6n!+@^*N5h$`svW=?gym=|FM%Ad-)P2Liv@*rbC(@%4gJdyP^ zLZ8t<=teV&o>AY75}y%ey&pK2?3>kj0>M6N>Q-{t&0Ie|>^hY!#$)v7LDW}0yE=Xd zP~I>@$rSVWEtOyKR3fYSA})`l{FTd|gu_3US-#S+up-Rk^`}u#U>Gt?Ees5M=iGMV zKDp|yx0^90{G$Ww{6?*U%j%{?lu=h^xefX@kKCXwju{J z0PI!h(DabPu$1>ZORfchO3hAb7~WiZpXrS0sD$!HtFp`a7k=8u2%V{t%ut>5^WxQC z3fa|u64m3JWn5Z%cyDsf(}9H{m$QT%566wflkDr|jIY;keRv7@ApCjkk~2yeM6*f0 zFdqS@#6hKQDoK;2sx2ch2^J)2Va7`_nurj6t-G%4xlo<;nP&aIucVZx~Xp9R1ejX+?MYXZY{C@$|iKRp_3*Mtp$Ev zht60K#EGbB8f&cu^?iT2xxBr-y}i9P#=Psg?RMLn27upq00WqaItZW;g~TGNvAU|Z(i(lh z-t)ddrefKtzJ#MN@M#$P*h4)_fJ+oS_v>7lxsuD}%L3<>eUGipM75``0)V&mdfLKB z<-yEMY}k9lqnR@tNE2_yB$5Nh)CR){7jv1}g0*~^`WE{M@4aYdNw-oFC?^Gwlvggd zT4m@1Q{=@)Wr3pDdjM=ilI7@K=EfMQR-oVp*#a1HbXzOUBirj`ltW+*=Ii^V?Zf@88I`oT2Hxx zhnPoZykVT)r)O!k^j(j>!IQV~+fld~Q6vp0FnjhecQmw2w&sMEXFTy_sj9oJB)%{J zTHgS$-ZS&o-vM~W&KUFa`#S(v{lLsWzxx`%RljHEpWoeh?=SZ|0H5}MVrEATGdt>- z*`WdO+VvbopH^$ zn5d@ro(-!qNmiCit>d;ytdgXP#|kuv4cvh&tT>#(dKs+tXez0A}+riT#m?YOGQd*CK6tPLDWPOY2I5kZ} z1n+Twf4|*s+qNx>Eetjypa(gp@VcP(Rf)DFt8-$_X@@bR$0B6Ns}x1Rdyj~unx<)d zkR~(1j+LTS6ZJ`ynCLXrb=~uB+ceE{|NQRm?#-JwcXxO9&#s!L8M+OC#HLA-=hP?MFiaSTC*YPBdG^q%h3GZ8L^`fN0G8|*fFbq=mJcjzq)@1 z;FG@h-v8~7H{Sc_w|f9>^UnY_eb3C#`+H{Ix&gq(ftiVWM5%&}W9AO`06u9uXWgZJ z_xAoRzW?_A{{G!B|8CP??T0&KjNiPVuEXY{XKr%_X}O~%i{iacf>wqov_LFfNQ)}!gU9g65#T>32nQzqNsG4I)9Tj`M^WDlPIla_ zO6_x!8I%gH>$=TmlbT)WJna36Ou8H`=Wug#^ZNDc7cX8=NU>hgPacF5`jF(OxdNr+ zI%g+dXq}^YdNKjdGy&;~3q{-Qh5?}2H^VSUNR2_|xShzmia;)Ohnb~@Zqr>}US5$K zhC%W-zoVvYTj$#Set%8Y7-PuJ?X)T>TAqbCx-lGM8a_`o$yuWkMvZ8~b#YWOFYN!U z!Zby5t54A^yO#IE;c}mq4&MD7r-pO3_`#VZm=jUd9br%-=zFkUw!rV?b~P1o+X9=$ME%~ z%1+a&1=pUnn&(9qAzh(H0afJwaA4*ayE_1%_dOB)>zlVk^z#0WnP0v8lXLF!?jFF; zZ*Kuy4SN8NZV+e5kY{Fw0dQ+YlhFXMkoZnd#;?#_){hW=@2m=v+4TR-C+CzKDsTp{WfrU8%~n z4x;td>!QWwltzi8Z`r%G^hH>rK=QeXNtUwVOdl#cC*_BDI#8v8#u)Es zUCfRJ?VWzICMm(n4?hy%M}aE!We(}JHLG6bq!zGzS?nxq;0ahe+C&4k(z?5Yr19e7 zLT&&9A(`WHTd9QhhdK(HrZEO^c>Vgd$lp+OF4_N%sxuQ`>XO;{)<{<_LYVm1R9(}L zt#fV%CXJ{FcFwu|e$UMN{T^xt#($@&%MUgxGt+KcV~pQy06c5A!!TSr|N8aoo0}Ww z+zoEpw)OY-o6V-Vs@Gp`q4=hr)T@+|#XLk*xptp0ayqanc9j>$iI%=rsPR3}DCiGu1n8XgTMrof?7!8H=0fnk|09gvtGe)KBA_F{XF7mly5D zW3~neD#N~3^niG-kN5cMZG>ssgow!OCA_Pgfp#$CVqrPtT8nFHrHWn>SYvI)<`K(tH~2n_NXvBJ_Q#U^93 zG2`Zu6$wrlM4pU=)Nhhq7{M=h!mgzU8I-;|6siSzpRPr5klUkK7c*3|N+F=|MwKbi z>t5Rxv7FL%hE;s5w$1u3p zuU`Xr{`~nc3?c`g5l)`W6Iu}oI4KxQH`j+ZAs|*>X72kwIXq8@M6E?;*5;Tw>q6VM zZQFXHzVC1P-WcP$t!$6&cDuzUaT|9xH^!K&X|U+@@ThYx7GEeo25omJtpLftx_Y`( zRl+}MX4+3PI}s6K_U`W-lfgBW%xw9&wR=i1*PjJ;nIjymGpE@Z5?<^^6=NSYdySWT zBSo}?sZB+7gY%N7!Q8;ioKp-HS;~t3Ov~5%Ml;t)=Juc(&+_8-i&SA_d3pEcUoO7HNO}3qU!yG&}eMsUVR^G}ESa)rP7gZ=_;nWuS`0P9d z#AFCEVq#bVSmie)+7p?V*U!KBeA5vOz}y0`$oHG)2EgSoc<+4&BJ{w{`$5tihntXU zo_(F06w%yzG!4G~=Jk)?z5VR7FN|rx4u*7|W&-r>U-f=8kjwhk&sa+RtA6w}M;lz6 z5JAVWIqqUxVPt~Jy&w0L=vNxs8S-woXJe>sZQJUGW^JD{?70aN{8+#Do&n<=Ghg1m zVdks-Z;dfO+kMB(8`oKDcSG;JAGk5bhyuU2Jpcn{W@Ehf-pD;{4*IV_}M?-eft~#Er9FoOCqxO7aQ#RZ-4*l=RbGcF!POV zz4r}u-aE1_Vl$EUGXUf3=#^xb%BbV3aq08{0s$BXy+0>8k;EAb01Sx~41m3?dr%7- zV;Z8SX^3p!_a4N6CliCJN|+msxGtW+Q%#;$esIWk95e((Aoc)=!3@LR zd(VC_fs{&~UZ{#u_vvPHe}BKZ zzC4X-U4Zox@I#na5kp0PPE8}K%X}>7^TWzAEsD%1F3^G3s(mFAvsUsG(RRBfB4g0^ z{hK#$`o6zpIOlfciHN94JvS01%@2LwCxk1Fj;ydslTf741Xe8T%GI9}@hRSiF-89D zOx#R`jSbfY)3qYGAeSoCK8D5^=1O<~<`6VS-b^wnty}gv`R}I)BU805gINZe)mUT4 z(FKxfM80*-QP<2BuEIwGd}c_xvC_ECF1 zA*m~WLrkx(x=UcgV9$-evDT6|0LV4~$ijPg2WAXyv)_A1uokw(#nt5wt#gj37;n%< z6a!!|gFU*Ie*Noz-tGGCV%u&vY^d*h$wa#_?fTBf2?NKBa%GM_rslLkIx`HDy8RNu zd8rf^5so9kb8L+>aE#tI%-oxfhz8ye2t!YVW@E^jhM3t?$Vlw~c-M=7VuEKk8HO>` zHE?bi`eDD@?RLBUeosus7~c|sM9XEpVdmERtEPG3+PnQ=c6%Z+#$kYDbnue9kSs*E zM4+HCSi460xJ~rHGbQIL2{br5WXXFX3IVcH6iJYGgN1F05>GQC2X)trlp{0%&-tBK z;c47V5Raqar%rvHb7p3r&vEJ{X0B>bnVG`B<&~|mP8%H3!~4dPJ5~j~zUW6tFvlqI zA)7&29++^@Rf4*vs)xe^C8v>%m{yPx41vvNLj-Fv3<2-RVQ3m6A`f_veeb;YZqNJu zevh_kn%-EBfdveiBv{1dkO?fu1!j=PFyQv~_Tu8=;^M-4-x_n!o4QoQowg9AsC`gY zUcNCb|5|J1@{GncG9hz8a!YPt4uk^`c~_e`Y3f4+;cz%agb7HEh++&Txs%fx+xPwL z?XCEZNwKc$>@Yg_>SM*sZQE`dzu)h*KMUp?DgZ=LZ3rnWg+LvRx0Q*?Z4GriVH z@DkOtewo0pmIasLyUPZ^v!P?=?=~+0TyL3~f9`I)_s$K*7~|eCGlqePE*emoXFR?zO*in>ZqLm7t~Wd~c^S=% zLgZn&FtI+VhKEq!bgmPXUHJ*-4}0pEj{(C|KgINr^){U^@4cxx$IW$&rp}M4*xV>M zdcd9cTHRBCFjqV-fhhzk0X;c-x*{{`h^X{Tn%iK1aRm5WFn4(L+X(KQ0{0015_4nTU)rlC1)azUTY>Fbu<8@0rewa6`r@e=>IlFLt3@u}FL#=%eE&x7=4+)2qiB?-8C z{aP!5G;e+E-or-%yaLkH=LP@-=T&*bA<(=R)-~j`LHLb1awo9Uub1$G2o*8|uG&V&5)fvFZ8n!z&tJWCggXxagKu2#h8^~MfA8QN4d{$9WDS5l z9NK0WdS-Sa+cAi!5rrhNwJBPV%GRO?CBZQB-q1gP{rlZ+cX4qcG~iv|!4#_GM^`a~V`lXfHELfw5b7-N3ib^w0dZ2`RAZUFpve*xh0n;nCj23^-}yzBek z_#S|9J)9ecJu{2{l)fv~8cQHUlKsj6*cq{g8UU6Z1tFGN*8_NN&EDdIF?fb|09#`k zYZ+F&u>_BLo=6<}e5D(`L{!ZQhH0+tV`IGNN!C08;woenrhS!d0E$+a6l!MfR{9(U z2u+*gmR@g_gRwq-5@4qJrj~)D=_Edn?@84*7SJEoAiAcgg_FHGvmk$21Xy#*=Yo^< zSIg|u$#u%Z=)ZGHsN60j3J%GUv0LW-G z!B7d;6i;Kk_fd(IW4E@2b53ruO816gXxp}GGg}V2yArC}Jgtz|#g2*rQf&wdnKJX7 zfAVRdyQevQYi-wcTWpDFxbKDN0vP(e_rC9Y8Lx@vc0o^JdbzSs^xDa{)>_MBKA$i; z3cqGSqpH5pB@l&>pC$P$d6$(3d_sMv;A3`POm98V)OfGTjLMR1<_~tGZ2h#IJU5&< z74*4yzS+(kS+m}yVGg+-ZvclVt9+iRF>t)Dr3$kIcv(lZ9(2s)iOi>C&8YG^64R=- z=I*)H?gyvp4sO$tlBq6B6^4Q7)O#G+#kg>`T>z&Jd?{XkxmR{s*W_-X(E|f?%Y<}%HCvu4m01R#UZSB=-5Cn=N?m4Vm6s{j3;+hia=HmL<^aGz zKffI|j#*-|N&FdMS2^aXO{Y&Y7R>Am^H$#IIIblqnT+7B7ucE*T$c5a^HwK#wVbH( zzsw90E|KONPYSAlqh4#i5&BTg~1A$t4v?Qeg(2jF(MUDq|X z-E212T4Rhg8xg9EF>Tu#L&E@TL+D8*JxFbVgfNGntTmDl+k2M#D$Fb~4esxoH*elF zP512CvpZ(z+-9?tfX3$?V|CO?!MJ>qr6bcnpCY8bIp^;0@9%IYW?HcXGfy&rRVDs| zknFxq(`>ifzVAhz7Lie8=%j4a+Lt6U3X4Cw=>ee86@SWaRTLEqx*UD#=ay-1l}n6K zr9&g(p#z_%_o$^jo-LA?i9#b29)FegRq@n!f6 zCBP-j&RT_|Q5GW}^N$Yo*F->Dgt_a7l9`Tk2k!>0@BvGIZzW(a&wl}YyeR6T}^7j7Td2BB(c6%>ZW+pSh zyoNVl14MZms!U~y)h4$G>bq8cTNdl70dU(`06%Uz0N-u4%=}sN2B4dJf@eGIj4{@e zXF~%h25&-)(2ICxbZn-0A`GJJE3R8kc&Wxawe8J1-)YsRdhYM<Rzca zrnO{kaFLZeEmCLKSRvISNjFW?8d!@Te)!?@&p&_m>}3*(BoeQib`R>#&jV%91^q#F zvEr7fRP`S4@3El80%qRt_jcc>ewKqct^aZ&Hk-y;TYA<^6-%btu0C!HQt9%`y0Z7A z?Xh6&uhwj0-XsC7l05-<#z9ZB7kb~$IW(n89;hTsjn2(U!hW<)O|4%lPYyAbndG>q zJHRzxgbbEd>!-<70%s$@^LyfHA&x3MguT_&S4+cKwCA0SQcF>|tR>rFCp4tVaB_62 zqdx|6)}ex}ybqqL&UUasdjVqxcFa69U>llq6k_su??rYoUZP(IAR=}S4(R2@s6&WrOB zDU#KYNst-o1&|IN#b`FB$cM4A#RV#&S^R9$au!@^1Rp*RsY~zYda>-oY~4?$PbCvm zzFY*4j}oxJX*)^N+On#g22XiCGY#Im5+_#Mn)qiGOf(%CAGHYJKDnImKsn49LlQ`j zz4v~%-vJnSL_+{hlC=EBz~4hZMJ2V>PYPTo0NTX_{uU z*)*m{jH@+K3^F_CY+ytg=^_dy9VsN*!#g>3I~{ZVCRC;J@2-I`b>Xn`!r6a_R8WPi zEsi$J^d@?e+eF<O?>Lwcxr4o!y7a{UkUv^;5U5)%{l&0-@Zp4Y8B;+ zZISWC%##>B<(g>|pr*m8hLJog%6jY;e7FQ$Lb^(!iApZn2zOk#jmSe(GWY5Gw(FRg z_d`DnS65d6L;?;UGMIT!#(NKMj4=pd7YzUd@4dHn3n03L6EpjvmkDEr_vAf^GJo4$ z^=JvLH*n|(4V%l|&~)Zr<~SSgxzJxaoO5?N=K>p*Ncc;>SRjK+30IYaR(D;T!94&7xHIB&*&o<2`+wEoV z?7p{_tR?G+nH|`HLrBuW5fLb8VhkC!OdzY$LS=X|30`9TrFdo+D58E|zdfoIuKb)8 zWZ#};f2K-G|4q$*W6a(betP(=Im9JMj7Uj?R}3KDvSpSydCgo5gLU|Cl zDE!Lt17>&*5{!}jHgdNTRytrJzOx`}8$|Ash?>T9-4(NCX2~P#y$69j6OnU6V+DbwiCJ)_x#LF~OZVgEP)Z#NBq|M>h@00w&I zy}#=B0DgIMM;6q30MB+zMBWVE`@2m;MB?GR9f0WP{I&xC>nXhnZAb>C!7$^(v0;Dd zfBgKS|Fpk*_wL=RTleXm|90yoW}C4Vrg7{nF~Ek5WU~b@NLK0D{NDo*b4c^J&zWhm z%T=^U#L1TsF~a{DChxtmnmoHSmKdr_o{8Dr?rg$1coCznWp{hFqh_IgsDoo_|?k=XqR_B{B6duMQkb0 zG!<`%3Cu=JHcWEW1b~Yei@^a!k@L(DowJc8F_J_}$}qEu^o58F296SdkJFcrg4Hak zfz0~S&tQy?D(Dj-Q{4#B#dJSuKJk5Wi)XhUW&7+XnIl)+Emn>w6IXpN=in$yOesu2 zp^Sr{C#&b0!lq1e^doRb868XYOaiggX06?r&E4HyY75FuNnU0fV9hKA zx7f+&xor}~oQoY=9J8HUIu3f9MF}#WoTCoZ^4u~rok^vW$++c1DT}ARrH9j0IB8S8 zx2-eOsoM(>SX{Jg7sJ47%*PI55wSkOGMiN9krj^ z3pc#}#*%<@aF=&i2geaHDHew3^-p$i0brLBC-zQ_lFt1Byh(gYV+=9Em->1{6 zizP$;!HaxSzurGmPYj?<=R6d}8+%ghwxQG4%+ za$mD$NL;372PKVC(Dg-WJ2WqZO1@R~ENe&fK!QB2oOG&WPY!!>6;Ac2KPfA}jJ2yf zp=7%s0t@@?yzv88Hu{ty!yfvE@na0iGW%4l$vLg>lUIIXjekJ#SwPY+)I1B@T-dbbK-gHddwO5xHp0?d!X$RjJ z14GOpz*_5eJ&*z=Iz(R;TlkKYPUaH=Nx{Q_qMb4=V%XON201@)A0$*}p2kzliTYGK zQorz~#ppFevAT_kMuz~3i3=ZcNoEaD^j^+(`0^<$x!5pjew1`fL`2g>WHg$(1V4Nj zMbuv$0oHIDP7FNtG)9$3z~Vi$(2y4Is!H$l)0_aEVRNpjpi8QM!Z6o^?sHw)rEtNV zY{60CoMEN1q3n5~jHZa^orIE6(1IFb7Fh2+nV%%}6-Wp@ixh@PCd=|w#fH$Wl`zMB z8e=XmFSQU1wf6#pcqDj1^f37?7Z!x2WRA(qsKS#& z7)!1_A)%i9*k^i0)pSmqldDt}BbV!#)tQh&&{ggSoCC_vlmld?HczQ=SDxXt2hO4! ztm2Vb;i=vk%N*4}N_u|s+ z44?}>Rg+LBL?qH`JX0l$HmqU{tqn$*$XRC0 zeK}pB=LN$z%YgD^Cb39Uk7e@oIn|9W>}M`4YZS4W{^Wz-dS4g<5hXSn4~u(M!h`^8 zBs=%P;~|c)mbL5?3ymIyMLD`kl@(T-;Hf<(NN+Wjl;WG};+i(bl+zn?w#40u@yG>H zHQo@V_-phzVZE*8DGJ?V!AY0o0Gp}VHiDqA$*~vNB1+85tT!_ctZ$Rga>^ODw6x?7cmrH+ie%7&*0qf>kcl}|0 zlH+Jfk~}wonm;(SZ|diEZJ+}AYB%pw3Awn!fJk)W!PiFH6IN6?M;5(mb)rfo|u zEwNUa{zg57&6y@rBEAV-m*Y>s(D$3xg4xO)Vlo{4)I#pX;Lp#L;(`^jV8)ChJ66d1 zmen9IN9`upD+*RsI9bG9nsugHVfj;e_76f;{P18920wrhQ-GU9KRm!-$a+DdGla~? zKrS_m3_~?z4zW~qPxP(QJQQRu6pE;mVk9F=OxBrNszDB8Jjo?=1>@^-N`kN z3B9*6`kf9W(plokWx*rneR3&vL=@a6gMU(lVsYZ7)=o!^|;$GDoh$A-?=dFFM(-`#O%CbK)JpZMylo{aTbY_*Qvp?rX&8$QCaHJbSWRSQLESh z7%y8IW7s?Oerz}OPkG@T2C2EK_OI6PAQkl45AYL#2{D^an)eK0nmPQBSS+uB2Md27 zxJhf}9TAad6M5b!i!afza`A+Z(2gzM zH>I>I)mP=rPB+{5yX?BgS^!4PArYmRn(0mALZWa?`q!e(dNpYIf{_|yc zxb;Gl7X4;ICJ&i{=2+2#eXBM~L@JTV0Z3OtXeP;V*<#66ER0LbL**hK4&1jI&U+lP zhf^%@+)i!$0DUiPQF5lOe}Isg(ayzG|uzG)igCvp--=N+Vf1>{jn zw9{X5eX^NJE}I!^0VKcu^f@b?F!LmdFAy?%8B8)xYV`o3tQ@$^*mD$m!Sp#)nK{ss zy8e9OllKW)+Bt8{3-}?aI?_C+PpV6=Umlw6;}io;b^Iq_>T^lM=h_(y)UJHa=6#ub z#s^NPI-MU)pUZ|aRYqK=a7KO11*RKaAZcTZE>}?7(G`2*eCOGlD1iGIjyZ$VLfIw3 zRK|jFKhOe+Dj^de^?NA!bzzHZe46;&hi^_bih7F&`6q$*fcTbGUNQD~{Oi!nX!%~A z2NKda0N9fO@V4y`PtE^v^$Wga*7!RL$m6oz^brKXA4kJpNGn)4a z>@zJk*np!C2ghllO8ts0C?a$c8DpX=!?1v@Y`yZP^-l(h_ijb5ZvM<*_+fvLLAAjlH1SkXr|wn=$lP zej4R2fwuCNN3Nvh`ro-S<@1_pET!8?yULY}m9=u2jeIO~(% zF)KQ^3eF<3tcNw|xSYEnhj~U!yc50M#?wgZLy?{2tj{@yR7air9EJRsqpmGR|4}2 z{P@N2LweM6;{Ke_vxdMdb45HYwh}11uB1O(CXPPgUM#!R!Q$JimjIqU>u=1>ySDS* zd)fnd+d2R{+X49D`k682hwUX1{mb>wiRkOADmf6?hSe6BCBP36y$_Q81P@=u>PB2B ztlL}3dF2A7idIAM<*19UT|>fjW-?-_*)U%Of?w#HlsCou3GJDYW#=j^*N!Vv!Wm#C zo8_FV>I|0MI5X^S0jQ%2$?5YqYRc*3Wq&$;U2}i9key>wOI7ih|XfY;`iJIZGog5s_LuY=`2C z*l3*-N9A_)zQzaNZy!V)Ux~;bg6~$&ZL=(OO(G=$0F0ZO^x1&n9prXk>9M`siAqWG zqDGt)TN*;0*q5yailvw#9ho+qF5#(DJH$I^itM#DJ#}c zkRi#*Nf_Cd!Z}77 z_0cw#Obf+kVBQ-1y6{>g1}7(i_FUa(Ie1;1AA}EkM#*RI%h5)akv)p&XNGdp*WPgp0!bpeXGnJ^rw; zaKX1mTn|7d0iTyi&lH9S)izHJ65+2-c;ENw^6qA{AtLL%WbQM@=rkS_;gr^A{zR@- zJ`8jAf|GZ3sawu^QW6ql49`R@%@kVcIo*#+=1fF8*W^zUJ+{@&E$`=_&{lkcaO~{O z65zx_j&H($C;PH`kz8-+xem>qoFlirQ+?ER+G(>3$8vfe$~WpmX6GE#8v#Gf%79T0 z1^|=%HHGyhuno%L^a(2&E=AIi7F9!IqZ6mCVopfece!Hz2%KZ)Hx%c;Kp79fL_kojVg(}GpYr)kj!=ISo}u>6glhl8kv znZJ5*0pPv?fV<{Z-}mJ1nfYzk5Ycy6m(2XmTmbOYqH2^0Go$|G+{rmJF+fRt@~4Vu zF{hcRD0Pm4)7B-CK<&(!I=ry?vx!R5tu6x~iRSuTq=)ns&sc>Qq>Bdv^BD1U@{*)R=ffFTQy(iL}ryvPc;o*xH$ew6&re`2gUjfOH62qQ_VnoCYYb|DuAJjqO z=wA(W9)w0~6!UPRHGQCr#fw!RrECqYk(ouMtogqaFw8-`R5H*X9FASMm2hmG4^lx_G1OCoatCK9 zZ=vhVYywRT)He~}kco&6Z$4d{RuENwEnJ7_n9W_KoyjVaZ8!#NjFXOK3O0!cENg}< z<=;pwP4@oaxXRJ7Qe*08teCD7MxjM#cP4f80s@^YoSx6pft4^Qt$}XY)Yguf-?q;X zE!j8rIWrG%ec#`28X|gqVVT+W5)-Nc@M73k+s#ZQi={p;K})bbUN^nA_WC&_mWO- z`tfmXofr!`Yu}x zRZ|(7I>@)_Q^+m0Qp{}FOVSJgmPaqvWVm7uNf+37BJvEx+$pJ~VpbUl_jJ>78f{u35>F_s-4Pif z2Qc)ZivQ9i=-PO06B*{H%ERJK)+~a-c_DRMA=9M*Z5a=NQx9hack{%XA*r~BB7puf zEUQ*Msm=@O!kDfj)j`QAj7JMds$+d;9v8&l8x3Ok$z6BhoZG`QvvmeQ+Y3UFbl0ph~@h=IK+Na|}xL z%unTn2f+#Ta0L7W&7da|;1ORxY{rk$!r>UCueh@r$~4bZYn6#5l~T$7QaR<3xKI}GBF6( zyPf>N%y&%#zzvp(BnR-7BLK$V0Wih^*rQYRkL|h{f{JPBaj>;K;#e^W2Mbh*)A*muCIfe<-yXd`%rFf8`^p7Ykr#?)y>{ z8?s&{3q$%O#9dG?*DzlCu#{?ZD4`|uZb`j54Gi{GE&6PoTz8t7EB_Sl0VgOqrem+J z>OeE`z$~>=Zo#yQ8D5S;te5|c1?C8_+7?{%j~KVtbn9du(KKT<)DkllDrgi3jAvBP z8llYH+nLB8Wr^cqK~~mdO!n8*G3`NmF71=j{*4uC`dX!M%f7604XrOc$K#jDGv=g( z_X9~DE#EJ6spQGtk1lbgw7pQc9MlPFk}BVm4n^Ys)^s5KnefJXDSr*jeLJ^57rvV1 ziD$koAE0(S5n;AKF-51&z}b&gQ`UGuX``Q~ROGxpdpImybQ8$>tQW?Qi5B90T<-|5 zfq;U%12cKY(`!hE^HxNx7Gtd=CW0qOxGIWqiAZGy7E`kc7j^TvuUaD1#7X#Qd9R6q zr{qzk$4g0&TGogX5vHUFkFX05f~L^f4X!g@Jt|w&4Ta}g;RU|1%+gavmz%V^vVXB- zTx9h*rh;>Vj7v0XOyHP3*f-wG#dBtl0RS})0BE3a7(i)Fy%+V<%0{b{`!V+CvTop! zfSFU!$_(b2V^QAa{Nj_s&!;iXE2q6gW%19)Rj{81M@zs}Oie`0r+XF7HI$wcH@9pZ z)NrooMOB%YnN1Xdc?yD4CkKF3n2VVvG#**-8|3LlrpcTrv~`xTXP1>9^#F5S8C!u; zbCCU1Vlb5pY)C7$e(>;)Sb& z#CTzbA7j!^Q9;kq;B)k6Y(gH+8wIf5iq$ z37o8gems#V6MWBGt=f&c+4K!z&xyH_&6CLrfPt4&u4BE51gN4 z*UPKyZHf9$92;xZXeoK2^1v!#Lez~h;8B5kZ20@f=21>pK|dZ1eh)A?LUOdk7ttpj zulz6ARK&$evj4L41`wl`cr1J85r0h;j%|U-OnhYwhyPLZ_ME7!#j029FrswJ(dn9s zGmuj(mCt2rB<2)tBV$BzzGPxNe%~;a+H?O#C>6p(IyvV;C0t7;J(L|=xHi2W;_@Bw zNmA-Q1&Eo%;$O@XMT?yL2LQkw9C{UGgHWH+W8qO%&;pt@A9uo`#3~Pj>^PU4O4+@G zbOV!bTB(istYbzK+k*CrM# z>qd6HBuT)`V`S#+*QzwEi$P0HGyRcJXXS*%8VUvY!d#>=^5opk&2~mM-UCPfaHyQX z=}#ZTS>^WEbGZzYGMsVIoRS+y5hNF5_fM+ z4UYs1)HlPf<#n>TU37we)EEQzWCGU!0GjUTTT&CP7fro7v;29?ETQIt7Icz;6R2&$ zpvDiPD#8Y); zG?x&@f5Wt}EE_UlcE8tLMFy@^$`1F{ufvi?>*cY z0ba_;&Tk@;(lR~8jgi&KF+u~SYDe7V(@_)|NNzlbgv{i<_e|``3;yF2Qw~u(7&ihZ z>OϗTqk%nW!fxdAX3DzgL3nY8A1`$O*m%Tzr96=gX*%8JbjPX#g4Wy-C<-h1yM zIg3R66*rUUjXZX&yuG7EDGhfh-81#u(=Ys5@+BCv7ZdonFElOM>K$e7L^dEGbXoQ~&J?zWou!h|)^ zjwFdW(58X+VA@~sf_@Bwx2DdA+ivIDiW-7RD^>sGGrFWc2GKeh@+~*r00+0(L(Gt~ z=%VC_*VT4xPv$Pwl;m@v<^`gn@qpGymHRLX{WFr;kt8V-chbzz;&EF)&spKoUySo& zq4lU4N-Ha$STm9AOX2E~;MjfgpnSkkoW)Ua7N?bIPwG5WI`(5=iZb(gRBmvQ{!ipo z0p~)@4Ih(;F|!ZZf7K_`WMO$fRUS}MLFuwHsz2c|7t`dWF0-%(k4jKK!b9LMWL*+% zNQV(cmJ1Eb`KHb*!IZ6-nJi2vKz4_%gk@yCEO-cbmDyMP5Qkwn7iv8SC8ezNc1AbD zHG-<)AdRR9ArVVD=PdGHpC(kUAuR)Qg}@L|uvR^|LvxO%$3jVc7g$lK?atExlVTea z*_YuMEcbwYtUonmCJ~!KL|*K{?1+r@R%~`gEC{PdA=DSO16gu{h;l>e*qmObKQYxw zgxdZxsKN(_KUeCTNp;2;A9XxLWXx!$U1kld83XkmKCOGBD&l>iL6uRJYT)h^O_{e z(aqjFpD_uQfnoxzxQ3)@@EW&|HPCgPb1rGQwUY0B0_z)D>&mYzpmAyTYes9+J(TDP zY5;&;S@;RUDFis1kvKlaF-YZrO(Nu|8EZCr()HB>)VWLMadYTJ4c>IhmrvbAX9jx~ z-zpzvS;yWJJE@MCnH_=Q2l5{55dB(11v;CJaRg#dWDLUr#7xW_Hix8R6UiV=BC99w zO}dVpZN9{ZipJX*aG^{y2I5#4au1IVo>W|#)HGPQ$@4N+i4Vxu{m+K#m+190NHt{6 z7IQ5cwV$hOC2|?!K|Nfc)3mGulM}}TKzk6+;(=L)ThMET;U0gj75Lz=yf@auK#al?%*uKDT&)+{lS9RJ`L^BL;CZ-T5dJT-AGBT%V-f&3iRdGCFBkbV+E?LmC; zWHq%_4n4sMGoPW7-7z7e~Fq7r!UYIChc>FO)$;x=8 z4>0lQTpAz1E=2YUe$E~MoFyP$xyx1|h$xR3T`9XqI7W=)$;}_NLzZ~zO8twBi~efP zYjhdrkWsd!$7;g>(}oL5OFGnU={TM7d^TR#%zCAZ?RGl=P17WfP@h4~BQBto>=Q0r z3MTDNKO`bzM?@~nl2^6KEa`i?NDVDaOxZE{2mp_IpcA)n$yn1&Ohv!Snu&UI58Q+`k$%;Ti}J)GbYuqMx8U~C1)j3 zfgs?jHyV`qI}E+|o;dhIOS7u-AO{;`lAYD3NHoC}iDuJD)*zj+T#Vj{1vO z6A(h>aBMdgTC7SoGZgOO4FWR@nmL+jL)Zz^o~qJyX4|ja230dt>chbflqh;`V#5wE zVp<(LKn)e$K!MK@u_dVlBx$zIx!oF4i)&`D`eF6`Lj!R<8Wa{y*;PS`B}BP4BQWC* zD?*wHM6r0(mPmjLI8T8aiVjxZF|+ZGy$3Vc6T6}B$qhDPH=mpODH@rXJ)t3jC&$d} z1q(X{v$A8^DFepdO9d&iE~n@n@~_10N`op6TLu)m=YixfQuoYy&!9yLA_9+c!`Y9` zdw|*bKS|Cm<*khRrsksez+~Qth{4tx@sB0Z>M)=pRju3rq{2(B-!mD$CL)O;hY3eA zX6c!`MU7K8!vjA*3uKdU#(|`o&!y5|*^`AP&eeM5n&@JeQt&KY{@_ME!GMBOx#wm0 z@C5jv#I_#NjU6K5={x_j;8FY3O^N=08}5XoVHjIa>xQgl9Ulvr77tsU zwd9?vY96*Okpl*542P!rC*aBxK0Eb5DZ!0PH|_h zA>Jz~+KfNZ!8>|Eam zVjN{(MEYoS3~XETu~VdjH=QJC&+Ro7Jkp(Y^T*h*-@;qJe$>JGZy#`Xv zX}--xqbdN$9srBGjj3z(i^~~np-T2ntm{-r%c>NP4@!1Y!UrS3bJRCJuhn&=^Soxk zSm+|)y-V9Vd^@Js9WjOXUc8{)Fr>S`6lK$e9>)Dm7|#zSNjfZ{Hj=5hEu{n;4pq_a zBG)qIKUDb1B zJLEJnImZ^1N!JfnUY1$z?dm?xf-$BaT$Sb3tz>ifkQURTl+CSVkE*nMp>TCqlrUfH z@Tg0R??jXewfuSFfhVMb4$q78TvA`pIq401w0M(BdGE=O##Qpn@@ zmAxuifeb|!G-hqz;)rIviwk$+^#r~Y6W7n!8djqU+4Sn#^t_!`?fPGnsp9On5F;x; zX?NtCYg7~C9Tr{B>=4#3)7aK@za`9d2Cd2TzDx!cZn%b56U?o}hbO=Z=0sSnf-dd2 zbCqyS%_8#TT%#NzJp&kqfg`C-FPHl;)2n5k^r1Mqe2&8(ZFpg5^Cy#6hb&*lCb}Mt zxIe5on^BGkKQYdF&Bi+ck*zFt-qsJoQEbN<+g%BMK3L|%p%f&tpO(MPn6?vIe>mfa z`L50~tOqos^{T2P6?FZ?RcL1;z^0g=z=$p6?+n2Df&6%Pf>h0t2UOb*gQ6JmhC}=} zdP;*Mcg<1Noa$!$ zoVf!`GxMA4HKz=?m03-5KH1Du#hofV?J|HEU)NyNG|b@|3VBZr zXzV((mduE0Npxiu;XJ?D$rEQ(R{q#&XyNF^nf-D%r&FY6 zpe$TZ?CHi^4632&pW`kiSjCj5e9U04F;C{Hn|FFf@X%e>JhT>VHMcE)vPpS^rW%Ws zb0%rX0A{wc2~Ii5lcXa#hjT{&3>okIsKXd3C;^as>N<*pF6Ie{UGiZ^$iM2IcZh*F zVmza>affO>0R$hI{+RomqxLZfdA@K(}M7Ch@(q{nPh~q5)5PO6;|B%ZB zvZ)6$sc-t;isO37Ql2q}KOqTtJ*?;K(!zW1sgTU7TYvP-k^n1;_g~48t%pqFTt1@J zU$ZVPiKGdMoIA|qW5aosX^+hmpT*R7b~w2|C7-uYOc>wwHK&Cbf;Im_XkK297RPaGPAsmRZc4; z4&a1|ASlmh&s?rv5Dl~P@*>y$U@s%}Og&@%G=%HUsS-C+r@7e|)*b#y`^SX}c^^zj zan|O7=geVbf9m&M(?c+AdjTK4H^%RUKeJI)O`bS_{u6NX2-{ z{G91iFTjNNPiO3KRZcA*AC^I9!un~SZw<;k4iD-%ur6mP&}1swc79Wu*=$~uIB(yg z_D~3B7fyucA;}c4{8#!|6se`RqI#JDFLbQ8k!c}&|LUL0EK{o6SVsfy?hgB+#C~RC z<{VrU`FIk$TaP|OyNiXV_IdW)*E8VR9P*5+87PqxFX)d#;}-+FENH@mUF*cMY_j3L zMvN4iB#PaqFiqZkU7uS6Kh0g5&g{)F9)Z&+V8a1DWm>)C;V{cO@1mpDCG@4e-{)6p z%NWC(Hg3bR9Ik_2{XrJBwZ!ka!&-vAOK<*^27Gl$ii^-lu`}rB=mtv_OnhzxxIB&)w|bzQv`6}WDQGqfe|0QKrGGdBSUhcvflR3$+!x$}MSj?=q{ ze+q_V$dk+UydRQ&)c!Izcu3A+eV_d8FqF#u%=u`&zQzw#G~=K>8vEAZmEJ~G37hSj zar7N}wJ7sa`nzDLX&4QZI_BD+2yS(#_rl6^!k$s4&7IWaG6VjCuIJWn44i4ee$w>2 zs?C&AM(a7SGG^QNRqk7r`AL#Jq?L_0Zw4PPKgQ(aMHT$a)$&ZC07 z?_`LG-G=mx5Fr*~R(2u$6cLgrF;aYz>b{bKW+bodS~b*5b{j;Juj6PBM%-?M7CCp8tSz^`x4+DkWIR}(V~qP%f10E0z`*BUjTSD1 zB%MV!qwen~pk-GR`q>EkuT?QDzxTBIY2cSYaG^bukgD%J@;`|Obz3bLc;u-Nf1#X< zj5$aOYj5u2THjq#Yo%`G=e2*Ih!Cj{1BVdwf<^3bW?$pQ)$8W(96aN38PqhI8$ zGoOP6G>Ln=`oa>>ch~90wy$m|&Jq;`<4;$47vNQVO5U~U3jm0O5GXtUy2Q7dpsz~4 zJCnYeOmkEzCocP)%pog*ZDMZ2_gDA8q>P1AS>XkxXS&2!e=~Es=7=@WRU~GS$l2{# zy`v>v)qq1+S}bbokWA2_r$5_G*$@uHC$V-l!={sDWpf(ojx7d6?@#rb^SCF8YT1z} z`!{AR4s=7DIy|UMB6el$EuiOXDW3)gKsW%1llgpWW>Y%iIebslU%i2NO!Wxn%FF zrqSrW;eN~0UiuZ+uxigs`u$q_$}S>l&wdVj9NEeG_da%U1M8O!m+B*PDeQZz^e|w1 zbFcy%Cj-Bg1K=0Uyc$y_0Nh12WW=TE>KQVFO+bK%4j#OR(f=|y>FkMMY}^iRev&J} zD5teL(~F2!tbm13^RxQ`#$1j#&52KLt+>w=|L) zX|(4?aDMGFTw1f$;`F>`EPXPs`|2yxefjwmksz!rSu!*=>O{YEv1>()+|&^YiV~5T zmZeG0E?|f~w3B<=>tu+VZ57i8Y=fKA>HS(}ZyEt;mi}J75EJWUuMh2CMB(SgJ0sqs zG%M<~Pll{9iE$&B;D+@!w7W^Uxm(|tv(IHw&yyaX*+ToM8^N&BhB6F+>$^H`seSl! zhD=YI`f~c#HIrvndqqXh+fh|kwB-|2s!@<*7UFM&8iEA|?k(*Vf3r&7XHrMxau z5&@JM>!A^*o-ycA;AadRnOVRbBgt=wP-MxGmIA_o)UCmHZWrMz{~pJMW@e4`gT|Vj zBt-Flt4-fpGiELtp)pP(;?cEemw>KL*zbcK``()@s=op7ozlJ+-XGS!!xtUWPozej zJsrNrAYWM;9EICA$giQ&)Ek4!TpG~z8*gICvmPF{PQ?`DLlflhY`z7XD{!g*3!VaR z94dep9_GydC;d! zrfEZ|K8O0w9hD&xsBiEK5Noxo!vP9Pp!+|8CEE)d697TJwd-GYi zw5#W#pW8!~rP1BL>^PI>i|b?ymBpnSeDyOK7p+HM0NzNx>5zXqIxo+bA&9eqo46n% z-dd0jJDl}xGl>1Ywhc$zFYOe3cP*6Z=I5IX$p*su$Lx+j0UoHWTfNBnCY+ik>}P&6 z5wx%LS&Y%;OKU#FviOMwkwv0L*x-EGkz^}$kFD`wTy(|q3x=|OU`KR_tW`J%P_DN3-1pRArXBGk;7Ipn^m4`?;x_?pZ*gPj0j`-X>M zve`z!7f-oa6TBco&P9L7YZqeE?J`pzAKNSn#)FMRue-8Qcoqpx;nKvD`6l{ZRs3+6 zwPm~qd@2s#f{JN-r*e=%HF(KKsv>Z%%uNcPLeg6#OyY#wq3!iXVX6g}Fs7v!63q+m`9#O)@6#^db?vd5B4$ zY~uXs{@k&*uX!*Lse@0S%ubp~XM(;;$}hd&EU+IDWy6x8-?N|9_?qcrF^&5IJ=d37 z4?ZnAopf{$jjo_dx8EvSAPIqK%&1xHVG5T zu7>i`5CW-VvpTqk(6CrNubfQZT<-@1#(;?kI$qHk#`R793qV2vaGo@^;%cbc&~*js9`-W5 zrm#XjH2HgZ#dk91YI@&MvynYz(q2=##J^t?m}o$;u7P4pPPfUG-ot(G+G2QIoi__H zyeA67Pxg~*WABboyby~TWs&Ly1NSq==zWq)&U|B%&*d9(;~Eo>K#d+8!ZNYiHNbgW z+wO{xP6K9U94pa5DmGu>pTJRHg5exkBsRU-YfJCDr~X+mbBx-&@8}k|#GxDe)-dl7 z=ajyF`MskYLO9(6-hrNad~rE7(SDukF8GUx`~=!o_JFa$)%`T>WmprDe1CDirq@7Y zOC>>Av!tsW1wgQy9h<6i1gL@ui@6e|GrOuR>Jedfx@_bi0*S(m_#4b5QMY?ZL^y(3 z_%{J0vC)7fFSFFBY-DSuIT5L5ET)f&5HTd*i3^Bupo3Tf2@&}-vPELB151jn>G5S; zvI3D>o+R`{-qJB;^@~nfi70u3tNx830_gp|aiDd+3#E zxZXoOxcWOKa8jMu38*P^qowsmRiU>9pFyT{j5ADPUP)a-c~5{&W19?A(z|v@ieRXg zKhJ;@6$MaS%vub~EV7IYO3O*_pYv_Zt=Z`$%{x7S~Qay(E;BT5Kz~tBx&32Xl<^sHeXAJnLT? zkmoz0tjRnV8})kj{MuO5rA?*C5U$Rd6)5}v`A}%kHo9*boqlW}#+N~yanFv}?ZEfm zz-q8I7WX&!HZB9zCWSuxw`i#)no+2Vq&mb{J=HCDS#wf;^ zh=nQ6TDm`RJKCf2r-N*LYRNIi<2Z<;>5X3APst=%2=uuZmN8=9Mbk1XXhmM^3_}vA zzj*P<(!{rbMmGO;ShqslJrh373;gC57BwQ%lb)P%qvGP%mwYN!3_EK77}D&JOb?se z1gI(4I(^-`6tL{H{C()@8nb8cSzHu)2satai z<~Wcp8Dv(*3qWr@V!L+%Xe`f7Ghc9QDO5P)L-`<+sSBQQ+O!bY7UaHqA3_KNv$Ovs z9L+ade#$lsszd*8RMeE&3zWlV8ok|V?^l{PGLnT({;aup1`l+s_fmc%1j85+0PMm$ z6d>!tsUbB3 z9(F?Nfmg|#XSDk4<*U3u-QJ`AFPuWphVA*32F>O4XoUx>W8m)%_dJ1cwuXiY#*%2CcRa3$$`IK;ss`ZdLU}n zhyJ;5vwy!X;YA?U&E_wL$Bg@~H%NN$c8`Ac6z=PEW#++UPy1bLEVfV3=)P%tII%!w z*yYAw%f6WFy%M~s5Hm9!DN>bFce5?pm1f=2SO`w5SmSI=cQY3`fb;(Nmp23H=-2eq zI%5qllMJuM>F@p%76RqPYEL%~nx@@luz3x7HoRjMsGBg(Ybih7z@Y|Wq?O+Y#@9oP zu_?_O*o0z*O|*%EqB|yt~*HMj2jf0LNNbi z=($=zA`>8eMcJaz)Me*-%RwV{@fuU5Pj$pEhl`#O%ppD{M%9vu|88#l4RH9M4WthF zirRKwow-Z`^zso>r)n025GWLjxvKTxWEF?Q&3Z=Yg15^WjJf^YVKRkghwcE2tajPSZs{W(lBbPBWz0*fOS#;k}0MsXp zTmzCV#nnZ*3fuQ3mzbEJ&u4|t_eoF8Ox-~XoORIx8XE>Vmm4!an+ZCEAYRiY>OG=E z2>0^9f6U6+Gzko;+o^r8v9By>v~Auu2abo>N#_c8YW3-HXqd+-!}xkQ<;1G_sk)|D zoY{yx9*4ef+SObpJyec;MVsAXxtcQj+Vx@JWb-N{dPvliW+*kmzle$0g1XcPkkG30iJz92X zDzCpr`=bgwYAJ8nr9)6L!U%C?lc3W1P&aQ$xH7Zt&{K90qkM#{SfTZhXb>Qu(gOt$ ziQaHrMf);rXH#$1klAECC8k}OD~{s(Ge4i3C{GcaEjEvsIni*wz;#G-kvo5cYE^Qx z${MTPtR<&2s?|BNSE`Asn(vT&mzN@J&AiT!H={V^ms2j~a2{{{+DrAocsw3Hi{^Sh zpR>{QTyQ*gljiu8M?G^b~PuY4|BO?EJ6krB1)nO_)*oJvsE1ptU3 z2$P5&xFJY&XL&y@WxiioX~3I8{Z3!(L&JCm%L%|}3VJlW7qq1P-KU69x@bWeGshS` z6T%f)aPSRg01;v0B<0iW2yI=l#wuF+ zmcCAMd9LZ=)zqulcj*KTk9JRBhj(WgSMv$H3zcuXryf@FaU81ND4XFDpkA7~hB$lh z6u4=ZFQF&zeY=JUu0q5APSdMJB(rT45;)14v>gh}jBb2q?d{?58X1x^@8trV7>Tky z0-~CQ^pZK_&La6QoPm+cu=XlHyWGYyCL-%hlr_ninW+Z#HvC{MIv8ylGJvb# z>rFni;enACZcg7es@LgHwR4enFWi6Z5wrJxEezO0&1k)!S7CeRlR3I#fkuwRP5S2E z-;?R&@($5%&}j7OUQqa6$im;Qv4FO$d0-5n|Fq!2G}&p&S`|Ls9-(K+`=Ka(fxhRZGc|3>kwebC7*k*V zTE_B%RIH^9 zM z(!p0?F~4$+-TP%}_AnzyeSM9z(U_2WZh;Zq=v|-8qo<>Ih5HWYfELmZQ2%q$`92VR zY`O?CvkDaZ|YbG zfc^6(3=8za@A$z>PjxsgDsb+zT_K$c0E)nRQtd3om__0V*%m|5*i; z#1j0RE9T-LAd$d{X_UN4M~uh}3{t$HExU*aKOv&VXyOPFicexDJeiqA?_5$$HG6w8 zfh5aJ-57ud`v{7g9TERK+6&)d)oYhCStDYhi)d=w)zncjeDwM60Ny1nWkc; zc?w>hV-%oFG?8SyC^+;MNP*qR(Q0NLJr})I+VW;XoajJp_>fegXk>AHS*eVzcWCZp zY0E#n`&c(;hF<88<3MdqrLr!p-lA8Kv%%_g&i)Vmk~2)4U$Ik->e&)(WOb=^+5#d# z`6Q&7sQ`q|$0F&_{0wPJ(j7{yZSxfllDw?Vj+lI6Ciex+=zM!Rz=*z??5uqF9`5ZH zV%=ze0i*D2X@azGJoEx$NSsBa-Ab>ywk*G|@ql(0LYh-eQ4NtCu-SxXDk>n7++Qc2;IOEH<-@27k3wZ%ec^@B8705V+fan)SK2xa`A+M(~!d zcx@Fl1hKK?Un)_Bp8QZLuj>{K2Wv>U?Neirm};dp%Qvd)cxY|g$PYuD_l>Uq4e*KH z7-+z=9jsGmtrp8<%rS(Z`?ZppWe5RWW=Rkr+h(&49ieuA5qT!97XT1BM#LwWp(lX( zw=f5Z;7G!fHwNxGH97;ka?L#^2-#&Z%)DOWj1M4p=@dyo?4vkI97Q_>+3cA zJ5-Zq#_Br)@i!m$lXe&iqi@c6pq2dIl;dWwhU=x#q#FeMPxT?UqKtEQy13=-Fm&9# zpGdya#lW44gIA$1JXtGUnjRm^^kc<{kH><}>^2KxZbHNu{m~%mN|4%b%RgkBUMeke zG~HLHpw$>cFz)+ZATsq1ZRn=RWuxmU45&bMztX=_r@Up&+*s5<0PJkZYaPOM^vX&% zs&;}3^e;Wx8i2IrtkWABEva=Rdl+MKR@w~h8^PET5iLAdd+0?9$Ioi5V>G_i7RDr5 zFF*PI?6u*cjK$2V`0V#*CCrV$Mp=4}_qFh2ndm##tSJK?QthE2rZ8la7>yTsQKb_D z6@obKn|2z~d(Jlusve9SY?7Em)E(srAk9QDGY2ld{;D$EtOswvj372kYqo$X?^B_s zBbH!2Hz~8*;nyZ|R-0>%*bju6P4A1*ANE+I^Nry;o=rt}jnrCWo|)qs64C36_A}Fm z@pyf1%4OnH#PrG@PS-SSk3CVkt&(H!>QSo3NJXJ`x{k-x%RPOxB29Zqn(2)b!(3z$ zLg;)uZ>S146N4`9yUiw>Zif4g{WIuPhMP6uN{$22L}^wS*`vz1r@w?<|&Bacr{%6ksIGt+v)U3nqYDWyI;r|G@lsbBta ziNn2wvMP-RjY=nc@=2cDSNVRY(dr&0zM*=!dHnqUvrah!BzOk9`GrI5T5bby!5e1ycSC?f$sH^pH-O> zSUKV&KU1piqX&OK3NI+KmeJp)K$JpoH2c1#imd+n_szq91FRc6wZHArqCT8q^Ob_z zRuBPj9ErV!5b{~P!Zz$YN>6HM4Y{P9OqDG<;9@!51z$$f^)AAifhhwXgubyf{apwc zi6*GW$nwGk>sSm>b{-9K?62Erv3g%Pb3AKWd4Cr3+@qlx`@};9?_Q{A z;Ib25=yMDMaiCiUUTz{OhpymyN4)V}lSfv{VTmctKYqp->Y(b(OqC%zBF84t7Xd8s znRH*)y5p-99+Pv$5RMFh#T!hjzBB#8_mF(kS{ZS>%J$Hg96uZLdYkm8kHl~eHP zAVPPh%SJEZfIttJ35DX^tm&ojW2TZ9%tt=6pGPE!1RA?eS1frkYkE2|6bRIo&{*>( zMeZyapLZ-c5nD-q#P=}CokWxc8a z*>Ga2CLv8Uo{n)EDd`hacEVcM=;4=MJQw*_VH*TMW8T(BHCrZA{ev#Il<^%s zi<|w^h|juVX(hbV3FJ(94KhRAb`Y*miiN28oHIN1Rq{I9lOE|htMZ%aN<>N1k|Gn> z8fWGdg5L4k6rVuy(y$=U6w=Xrh=^zetD=f*2m7w5?%8g_@O{mLsi)_di2hU`1=cD3 zd_H50L?=VmYf}yKL&p3Xgb*+hl;`q?K4{b`{i{m%ohrn>5`5<#aHBMP|2ocB1lpUF zSk%2UDti_Y4HBvLSE$$(%dr6}KF$fUp~jhImO{c%_npn}5Q6F)Lx$wvxaToZ`vm7$ z5YU>kWY6k4aq)JpmF9Z|i+|w~G~`J{m@8=*dxNs6m-}&l!Eo}G)1&9Co{0|c*|jvi z2<`==X5lo6XSNV868VR;vNCmA^P2{}inS*zXQf`-SFbwX3khLRJ$O%n7C20SMao;moR)5Ww#A8wof!M zKA)<;P3@o4Q-7vFAR<6JE~yjVz_*v*FSXuCr+?Xy@}WFB^ek;CzR18UV$}-j-h2#v zljXkZOg3$ZujdWMgH4GvPP5>_J~0|A3$-nJ$ECENQ4K~7dt>J3^J&V4)E3Yk#(5(& z3f?S1C4PMxry4>Cy0`6Lu1VHM7%+F5ZYAt^%~j)sdo5Pe3#-e6OBqXn%w7zPF_N;| z$sMR3Y>{{y@@)1D4~~T8jl7pkVN2f_PRdK+T{c0~$Jy^Ls4jkezhTJR@PAGS=1bXs$^D< z^4ft#{U9PwKx^)MptD4&CLO{c2FTS0u3@;&KMI%T1{{WOBH4EDj32xTLyw*wdZPS3 z?;NJNn$f&=ZDaMR2c@0vSoPan#Q(Gg!WeTExKaSnb>RRFK4&A_QZc1la=@ILdoP8>; z8~J&@1wfE)R679CTbRiE4djyTY@QdYW{o{(z^w>;VZltrB0Z;TrHRvh)lwZfs+z1{ znfX8e`Tt^!F~Xuib3RCLiq8(LtvXYEH-n+^H7sH+-EZ=MMkk)8`B$~wh7T#eFPKSA zvz#f2)Azl)Q)xL=m2j#dZK9;GW3{?erq^8Um9Gxs7A0KOXiff2^Z%v(_|PgI!ee+Y3RSO3 zFaHh%bAj51sT@60FTqH{j3`Lqt0X>ik#=e}xoC+<01Jc!9FubvaU|gX{onrs;2Hmh zsPrL-&w>WA;(*bD4^HEZN-XM8UC!^(>FYPOZhhok- zq$P!5r;`OB97si%&CX^L^dfu@J5#%&P}p6r3P*c#bpU>_Lsz6fp_m}LzzAp1>JXE2 ztNvbFwN)r4MZ;4q2LKpg?RFAEWN(4A(woy$4~y&6tX#v#5D>l5ITfRb7CaFFM2eCx zq_qi|_%4LiIVGsZ)4B)|k?0wRPDm9I??&VGq~X24z+fzLkEus%aj?o_eLHHsZNqE8 zk}4@_^!b=DU5}@Rv1E?FIf4*5N_@geu>u$sWai^Ie*bK^CJQw|sf2WjS5p@GDc?8jF0KHHNJ#Tsae!p%n$GcykQQwr8j0;| zNWGy8unI7BYAv|a1l^cKQcu-UojJ#->d{J`(ozgcrgMu(06m}2s2AGn0lWL;lpPJ@ zWygn4!^`UF&~7}mB%3VRr}-zH0KNr`RJ~Kk%u4y1M85#DLfsyObg29KPn#w9LPS)! zY&%)mk9zfALXurD+SxiXdtNJAO$zktve^d6Aa4w*TC4=xj@Xj(#>(X+)f4qe+(zC~ zyC57(yHSnXY4{Mb1o_Va6LfMvlvTNJnfGwQfc-#biE=Bw{pA=lXB9{a*NiNw>GUEN z)~g9K!&a>qkqLlAVSb84jsYy3?jj7S+1z6#qu?8-EgimpjlF4UNgAT$PY56+{iAZ} zdf&iXiP&0!?aS7GL&)l~W3E(hZUO0j)Ym{J0#?TL)pDC+W}WtR)}sNuY;D^GolC-v z{(7MRue{dg)>d~+8OG(#EJ8^Gok1wGo z2NLn0W5>e=F2Bn-y$1VeN3=J{R1j2*jdiibN!EiBB0a_hn|t zQ*=iu>B_>-h|G-xn1v(%X5pZ7F>BJ_0?kTpsBJC8+=HSwdU#WC#VNBbt}IlaJvl$8 z{a``*1%l9vcO*AByf#~Igjwx(2|x=&rMC*se1)71UOme)43_)Y<0KjjVTbV!I(O=oKEx2|b^zq1S#Qr7gMK1k;!) zo2fnbb-Bb?W}_Dl8|4KieK_)-v82es)}PZ3X<7cfy4lYjUNB<&7Ey1}`7ktx*4CYm z*G$mEwJQAu(K-p-vlHvm%;4Ax(bj!2M%3>c!^_Wyrd=Q60mcGVm-XEDqiS~#-^+S` z0u#H9JtkYV_PEYI4B~f0Odxp&J>uYq^NpJ>$gB z@+B7jj9*iiPtbc}_igcYwE>n#{BBTS2OEQ_wL#w}tHST`b@|aENKjlW_;04(;;(n(Ux;7#n$H6%a4iSlz zo#2$*&KDpp?!-cBVpm8li)6c3N>E#kTXgL9o}~%e6;LU-K3BsNjA1k-Z_e8{kViE5E;@KvlxO)|7juvF|dJ9CUt%_Y9=NuA|#0Lk*| z8~ceoY08wk1IahF!%ORP!Np%jBZMHM>_z!?sxXonvwUo@*U6neIT~jSI8e)XZybv+ z0N+tp&64{vQzL{8ap+By3eAja?1!E{-o=k;axsG^+2s6kXQ+CFFWVk3?x7AdEEeNb zP^-NedSgTLW;R&ss%^yW`3MIbV=KGYqgo8i2yE#?ciN8JCcVdMHd}h3=aO!)$zT>) zUMOJ!4S)G4*nrw`tVZ{>xX`B!yug4Z-E%VD?@a7g|9T0Fg;>Cm@`>KWo+W}K@o!;P zshK3EDX(8efCz|aP!2un#aRX3*!>+hp32uD(K+u2XGQmKvaCHQ$5S+B26n@GcvRDZ zq0|Uwrf`UJV`1iS90kLrt%kxQgmi>h=ui(XyYhudkT#!^X6;OZjp^+0NLtUlU{SaP zGs@W4*rdl9-)%N=XY~&2(yu=6;ciH^NkBy0T{}3e$mC8>E>#)wx}CC8>vke4s_l9j zbO;Q2ZQh9#<|xvc8Ec}yU`_)}IQpKk*HqS#1iw!Elr$J<^E>Q~Wxl1offsJ86@CiU zh!-iUZBHpgi!~1>)z83BAl;4(;m0wcrM~Jo`btEStqlGzKpgo9zkGU*$S%6$CiFey z-}s3)Tx&J;bK|0dqoUW2K(|yT3_qlpS0&t-cI$%lXbfuaaVp{{JAnAmC1i%Oa zi)N7~o<8OkYFhI6KCtO^gbUU^lcx;eG*@<>&pIT!m#ZNY*=xYcoU(z7nIk_X-;y*M z3}8-w?;`Z*Cf{AfzJm;$6xOc|Q(G5m_i%JW2>FJ{P6M9Y{W7x^-mdX~Bg`22Eaj=# ztFyTovIT4Jr445rEG9(rm0j1~U;pL4ZVc;4|C(b!vtOHyGe!R^b8j0EotE=sTom}-=zw0Fbf)L-*@Sa%Pm+lRqd%&4g(Y3=JUQPn=rF? zrHB8^k(qO_j<;YjbrzrAO#C@WIOGcz45lf&mkH_};#FOp`uh|5;VR51(KnuH_5z|O zn(OqYV z9Y<7yvFg*lA|k!g@u#-(>tNsP{Ow@3^w@UeFj;QJ(1?p1sG;Si(j6~sNcKJno-z1g zGHEs2OyfYU8Gu|IE2*;d)eWJi5BspD$f!DVij}5Ue%!c5%ieX9-f7teg{bSRk$@~(TY`PwHQVte&XuMoZiSpf%f#H;>&RlB(!IIg@9Rw@R3s{U?aoPlGB zV3ioN+wrCB&c}uIs_%4S-Xe12k&>P;ZLYeZ-R4^u6y;qm(_cNOO;@31?(4(kTFUCH z*PiRs!@fEzs+Xy=1oO4M@G4BT-k0vv?V;DqY`hjIc_|Rfc{u!}fklV!cW_=*WHefr zu;D3nX!o9^E_ON>zPNqWv%Sxu_^|ey+W{HbY`5464PFlELr!1q4L2!GkXqE+;ORRp zgg{(|#@j!ATQ)ULkKuY|Hr%(#IIM0qHwHyX$+Fl3b}z-1sgA=xmq--`+*=`fkdZ~+ zzeyg^oNn`l2(jwDPu&<6PPc*RwLXXpPAQ?En1wkIAhuh@MCs`7i(951XSLuC*3Kjs z1}eo;!0$?}PF{v#wYhm42rWVH^hLf3scf1{04xxEAUdREDU8|P)UM8fIzD z7)>osdQOd`zgNKPY92kmy1igDP-aks@vqcQX7e+xah*k9h))KkpIpC@e%(DhF`7!F z?p*IFNQU(6ma+GUhD+9f&ei?4ei~$A)@zU-v7$kOHaoR%gxSZ2-GPSYhVF6&kArEn zp6L!#_9g$5AjWvLpBM0QM{RQ7USS&2%AM-;|J?_Oh=s`Sw|oTQBD1@duc}@LetFbSfx&$>LEyD1HkPfbv%+{xxiOGuZeL=c$s z#bi1^^RSXCrX7Dw6=^q01k!2F(BRPFLtYE4%b6`~)>FTMOi6|7NvQ^iP)@1m#4Q3{ z*ck-qeK@)NB#~@P;YMj`6iSOsRIJB3Sd?yPD2cPZzp1w>|CtzY!$kkj-{qd*s=X`$ z$qR%v-kAzp2QxF3yNU8aGLjG|hyYUGt?M|2=vfLTLXdjbfS|kA8c8j>rdi}fbe71% zs;#8_5M;dZA@0UGYc6MG^nPEzR6@*lFMnF8rC6@&v{L%;pGBn7)4re}B0_S;^$TPfW|gBhM}~n89^#?TTkRX51NpJ1Xi+$Uz_5wNLellCt4A!EwJ=6D@Mh zxobXu0em~24=(`vh1aYHoZ*@$M-<^84ADbGV2%(GiNZ{*(s6_cKN=?>t& zl62$mpJG**?ye0@&}E;1Xz#KwmGy#d^{NL|;ij2AhY;vEDBpU7`m!Uh0jF#=GdBCP zar$qcsc%BhQ5q6C(O+LI|5BFM>MACYOBTFNqW5IIaidk`K%7YvBR0^KaGKITcve@H z2IhRhlocBeILFRm&$PD}_&v*}-}*NqGW7XSM90-OXwM~;^fIKNRJsN$$J$g}s}*9@ zlYONL{}BTgk%j>qK7xs6sg{hov>pMsCDEXf5EFDA;>)bdn_!ZQ>I4YXdTT$a*07vP zZvjK>mbtYq`RGxRtN}~CmvCNdKGzFrLC1-5QKtLmlEuFMw<}p2FtDU`ywezJ`sBBM z)1R^*ysoG!$H7l#$+Dri7_IbI2mV+szt^a4#I{-KL-sl30sbpHv-5)W&Ym02?Zu0_ zQetL3dIZN=o8k6sxZfRslPN3iDqXF)kGRepMWOyZu6 z;AR!R*c-Q6zn++Uj{ELZZO@P9LqSV$%6*fW8UUdBuH^SN*>%$>2$+TRDivlDp%h~h zBuIYIU&4K_cKSRN+~oZyDL4^@;P`94V!_^6c#TVCdJWj}dII++YQQs!?;*x`pacZA zlvf$sM?1%F=nQwD-?^JX#pRy+HVDjEa?f>UQGZR`Sz*%gdvo6vkr2Y;@rX~||IRiX z5fvr1-jJ7H>YB*tc(6=tq$2WsKEoq0^YJ(=prbC09g*%TG}iv!7|?eMOtu((A|iAt zjh;7UdZ@Zx2AK5R7xifRl+E&KlJ)(!=+fQZmK>KVLDdh#V!xgH>chj&MMI034;nk0 z3>zjQq#dntpNP1K99C>Ee3@VLy^Z`A6s|H7vpm0$_s7EKN^4Z%MarMhGk&CAhO{#0 z+}FNh!XTT(zSus?kX|pcBg+;V&nIsQ2hYok*fime%_$v>WJ=1;Ym(JgaAQxag2!bA}Xk2qr!F_72CfAA-gcs6PwiQw}G<%$CrI1*Z+^MUa=_8@aE? z0RS^Utf%r~OxGbqM5u>kjcY%xf-^bWQ;ae2Q$z%jC3!gO-zH^2D9!J5pCXzG{gHn7 z2C%Mwm3jt%m~?lR1nEYj^zRVSD8mQj5P=XqfTSR!HbIx*8z3J~c!u-rr@>F)Oq<|| z=@9MatS~x1iSKHxKdI6qc?y-2j`G)h{fT*Cq$wJsWTzgOlq`j00_E-+vIHbe;&Q)@ zpEmUr>VnSF|Nig){`J>ie;p4Tc+xY`+{h{PrbTrV)*Gx_N|82R$ZO<{UApNR)wYRF zPCw7U6lFVoyBL?jsgu#p^eueiRC4Om_Zslbowj-C+4=A$U|<^2Ayjk3onIQEp*cKW zR{YRHsVwf>xvz)T3)&cAemTg~&kQkrx7K=x$&eXvsrM)K$ag{Fii@uU`Okc%|5!2i763U>-aY5jfX_nK`0_mRb`;0IzcYTj#Ut+dJL?p{k7 zWArQ7lUGlaXM}i&|L}PRIzRvOpZ^fi^ZESy@4uhVCpqUq>k!Tt|H1W)WSy^f!tzo} zZ!A?cdtW}DO{CYv>%O$BldaKrzn$4=Kb+o4@G@Z;W?3iWAe-*qUz?2q0Mu3@vg#xc`pQ#>Q5HLAIebk*MM5g+PG1-`G)1%hX6mYw*3;ObH0E1ph4>AKMEUcC zNXXmhbD0B)XlQ$-oXa;Jq0P2j4oK-|KW7N+4d5=p;3BD+lq@30gq{yboZi^Y)k|g$ zF7mbwN|R~4000z$NklN!p4`EDT{`NQww;(V-=Cdl3;tJ!dE)Zl(M; zg4aYBX@G-3IVzD}_6rL6#@`?kgaS#(WQa8CO^JK#k0B)heIrNNgw0j-QJ{Ql&G7xV zpyXjDKvcv!X)xUbAOdD4cAL1!lDHDqZ1F0e+;dWFTuO3zD>m;mq5gP`lC}fDP_x(n z9f2shzE%tbWjJb%z(e8Q7^-5>My!^^~K1%1HDtvtX?Bh z-krOpl}<-#KG~TcD^RjQa!2b$Hd^`vU7H$AqpxE=UOrHUCkwrKxo$O3_gm9U4Gcc6 zF^d+>d+PS~Q1{Oz*wN&!7V?QX|4)S%w0{Ge2`|)pfdMO8p`OYja)8t@q2H#rj9<1dM_jYtt=tEV1>G>?pPYT90zjvYtuW0mam^G?-0?3-8< zxt_#K>@5WS`xy?4*;syV`Si%}ab2mOIkf=lJR5vV6C2(@$H6NaFx(`FF~Wzh>$F?x zo~8H8sMGu_wGF|uWPi)=6kHl(+Zqe{kJEsAnQTATC*bP^6eH-1up2YtTRmiUTj=K)W~R!$l<&qoJSsQEb%4%SEx-@Bjfyc%Zt!~%S+OvxhPYw#%NwDP}G zEQ%DvLM(~>Mn-1DN7_Xak#w7i@RNig5h5B3nk3G2b`$RY6!zeO*f~mM=ZW&Q)@#W& zvVf3*6xyUHvr(r+zmECEE%xhz2T@4MD(S4&*8=uo@-AvG}WR7zGo**QK-LlgBMr-c$o*LV$ z(}3!;Lw(zK;5GhHoyV2(CP`)+C8XPpw6tzVwl85{Aowmg?MW}Zz5ssC6ffwVikp;D zc};JbL560oVJ&|m_x&;$n5HvJZr)3NXktoywFdVdyr5frT1nkxfuEbnLpD6?D8EP> zt%BUwndUlMenv%-rae{dlw3Sl?JNppmnswW=PA3Yn)ekU9x5*O zUSE3erzY7c_w6fMz|P{Qyd7v&rw3wdn&MtA8BS*sbs?wBB#~KfDb0#3SR|{k5XoU> zlP5D1YCSdKK+edSSdRe6H=el&S=tGaYtA6Law*5*^4&4+C7q{<<(D`!L?n`ZYLC{` zd?b@NfONNapuYf+q~LM<6o2a%BI}a*^?3MkiOJH}BB(M2Q`Y`zzy)V_30F#UU9vue zaAZS$ewjn)V=`^@eq3`8gRRP=O!$2b1 zr@X!V{^d^)%oS$yL`M5hr#Px;pTZqPMAyZeCTKLF zDECb$CJIU&zGDQ;90@GKAO^8;6cK`-%>!o8EZXm~NMD_Zr)U}vUbjdSa41*&DjtVE zha3(pB3xG5=!o0x{fco(z2SWkXBLk4pzj2|{@th9p>*Qz zvb_YZ;^cnWWbds_ z&?2H+xqX+8)|~)gKp@1NF8(qZL=GiFX0}Z&b@{XLg#pM>0-rJS-+%wD=l6~y*`p!On-(1xsjR3y`-)jJ?>%>wwZ;~jukB8U1?O`Kx5d{1>Ep|LwFDcCWIrf#>2`r zT3!ojS@G%bvnX5zsZ~i019_tWv4|d(RG(U+DhnYW>#{7&+y}|J3Vu-xeXROqJ76({ zk!P2k)5BJ5QtR5C%K@710z8;5HKOtLJF6gz#a@7AC!GZGO%fU3bk zas%f%;?`-9bJ^U1T} z8E_p^XHod%!!*P2^$%xAS({D9WSMBOUvmtmL2UT`t|ej9Pih_XqvK^Y8VNpa&X*lP zLf*XGTo*;xH11m;E}H4S>1BPr8EhcxYx~M`Wr@v)gPs3#n;sWPt)byBeFB9tNl@CD zvs5vktwsei*jgH*VlYr~DwS5-WBv5A)5Z&&XP<`e-(dhVW;WBDv@UOs_LILJUwt^w zTcI&`r|4fY7f-#es97J5+x0Ruou5|CcJ|rAoV|bG4b#>oC?n+R3hGn+!%QQ-H-@M! z$lq&xA^S_~Uj_l;u(G_`noZ3I(E%zNcADr3=RhDJ0AU&{b?S}$t`2uE7$MM5O#q1K zIF9EeIQ}I#<(gBy!5Km^x=HCnp!MpubW)OPRJpk^rdCfyB&lsVfrtbyY_$Y}*7VF8 zSR;kk$U!h5LV5Gp2y+&)qy7X^=n8aOFmTj>wB-G`OY{UF5{QCG0Ei+e1l?ZFmx((+ zEcHYsib-`&QQug~j$MjMO%(EX>k#d2s&9tqXGL|(at+x{I6*6k)MfUH?+#_Qo3Dl$ z({!zgEFnc>UX*?dQ<&Wcj9OiuS)RXuF?kr+{BM%;XAP+|;C*G+EHeN?2&dOQZo<8e ze4m3wFNxDLa)^M!M*|i6bACVeXVwZV;*~vn%ctkcQn6>s;{0G&*t>Ln3k;RlMVD%4 z+8nz}_g(PIo||gFsaMYTLcC16cS;9El@13Nty~HwkCN zAR~+FS@*O}wHZFlt0XN%MeiUYjOP)Ky!eYGP5_XapE`qt z2uyl%F;ED4(-A0u2m}?^m7|s#nJ5JHU`$-Z4*BMPLJ-K+=vTA;*(V;(7GP%n z>#x7`&DBj{F9zwuj#HJ?hPYq9e*OOa+qYsdRrZ@%v4^EU@bBNh)c^paxiVILQm18Z zT0UjnRL}i0bBvt&Gw zK@8i%4CrB=6*zFNGf#n}>nI z&>qvMaY=M`y0Uji4N?M%&HBS~8P3&bmV<(DlpOx7a9I&?G4}hJ=}%+l@v5Bv?U}ii zL3hNdszxV5+1wax541(h=()!yl$7L!RBcO%m^Q*NhTvD9AyRFFI(=1a?2xZGK+{q; zzB70m%OrX@U45NSel~BNzK~m-uy(-v?s{klv?Z=I zRMqOTKQl9@SnkdcEe_CL6J7`9sGYDR>y+Y%Y&8L?tm}j=m-IW3%;3sUlaS=l3p&2# z7?ST9xSz?8p!k^S2tdT+@vt=35XHN3bU)Q^_F>`u+faapuQ$oUab)kaacs*nG?`;y zWM^&AD88{>n)M#AUnRr&*;G2@`ftK!nIE;k&dAPVr;#tPAF~eGGPb{dzenX(KPEic zvz+(Nm=Mp>f4!U16r4;as@H~x5K$lm!l7}j4k98Pfl`oXqC*X6R89jQtnsmOP%Y^e zt+zktI{W_F>Y2MNz&l{`hlqat`jteAyN}rRuwwzj&v1FGkI(=k4cN0!bB8nr93pET zW+o}|WivBw-Lfba2dg7i61Ir2JoPt#u%IpFIARrP=p0gdy=x>u7zb1VCL+Wq0r`Ef z^0SGJ01!%Ky2Mr~0F-282_l#6L&Kd;`zemqH4)sz-6G{X`BLo_CZP%#Vl&xWjMUdj z``>HiGbY-@kfq5`1MZlv>AuEn_qS2jm;inl#!UAs4evWp0if@xQ}?GG;uPC%Ek zWNG@Rb-@>aXNpDEPC|LrdtNgh>cRdVB+0(eJ-fA{?$vDDdMt2M9U zu%59^EsZ?av)=s+60Sm{_pFCNq;kBHM*XeXbBvMeOV_>6v#HK~hl_~SV5LFj00=r# zj>*6owtd&Fg`12RhY*BQNDvW0>Z|k~rh!#`=vBj1dpTdAcLUd_hegIQlYbykJ^+o$ z0(s$V@wdQT=>70!R@V{sX6N@IEl!cKeDq%izK1hmeGrjI$me|n(IbQ)=`N!b|2exf z*NheiBv8J<1w>3~yQ{aA)#~qCn9tsP0_-RV@|^UocxRSYa`ARsnma9gX-&xJokMi{ z&|r8vl-~X%oeA2wuQ6KtYP?w=f+V+RuR24uTgQNl)|LM@4X5s9;sjj0raj2yXhaN5 zX`Tn@ZY~8WTc;u-M8eDgfrwb|L2_De62FS?1sD=z)V*E`M{2n90IkNkxjA5D{55L=cMOTXzv3y+3+2+fPpQ3%2Vvk z0{>hfszO2nl?H5ShKqa=VyQ7WH=l0TfPLYwWt>f_S%JoS8{T%;S;<+~HFypg-Y1Ci zJ5aH}w9i-i&5Od3V{`b>zSd!|3jg*@8fAYsE524ytrHuV&(sjxeNvMUa;#KwCLN=C zwh)opSktx_>CJ(nz9;wqk?t9Lb2NF=GzEkMx>cNR|GjhH^ozNU z5-F>cyBlpUL6b{5{l-|uPLCO(e*u$A@LKrNhCv#OHfYaaT6kic( z&-J>Snnc7dK+|qdr49CgiKrJ(y#Y;9=#1bZbtf$r$P?eo0eb3JwP6w>G9M(EXILPJ zs3)CGxx&cv47rhj> -#include -#include -#include - -#include "cutlass/bfloat16.h" -#include "cutlass/complex.h" -#include "cutlass/gemm/kernel/gemm_grouped.h" -#include "cutlass/gemm/kernel/default_gemm_grouped.h" -#include "cutlass/gemm/device/gemm_grouped.h" - -#define NUM_STREAM 4 - -#define CUDA_CALL(code) \ - do \ - { \ - cudaError_t status = code; \ - std::string err = cudaGetErrorString(status); \ - TORCH_CHECK(status == cudaSuccess, err); \ - } while (0) - -#define CUBLAS_CALL(code) \ - do \ - { \ - cublasStatus_t status = code; \ - TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, "CuBLAS Error"); \ - } while (0) - -#define GROUPED_GEMM_STRINGIFY_HELPER(x) #x -#define GROUPED_GEMM_STRINGIFY(x) \ -GROUPED_GEMM_STRINGIFY_HELPER(x) - -template -torch::Tensor CopyToDevice(const std::vector &x, const torch::Device &device) -{ - size_t bytes = x.size() * sizeof(T); - auto options = torch::TensorOptions().dtype(torch::kInt8).device(device); - torch::Tensor out = torch::empty(bytes, options); - - CUDA_CALL(cudaMemcpyAsync(out.data_ptr(), - x.data(), bytes, - cudaMemcpyHostToDevice, - c10::cuda::getCurrentCUDAStream())); - return out; -} - -using DualExpertGemmKernelNN = typename cutlass::gemm::kernel::DefaultGemmGrouped< - ::cutlass::bfloat16_t, - ::cutlass::layout::RowMajor, - ::cutlass::ComplexTransform::kNone, - 8, - ::cutlass::bfloat16_t, - ::cutlass::layout::RowMajor, - ::cutlass::ComplexTransform::kNone, - 8, - ::cutlass::bfloat16_t, - ::cutlass::layout::RowMajor, - float, - ::cutlass::arch::OpClassTensorOp, - ::cutlass::arch::Sm80, - ::cutlass::gemm::GemmShape<128, 128, 32>, - ::cutlass::gemm::GemmShape<64, 64, 32>, - ::cutlass::gemm::GemmShape<16, 8, 16>, - ::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>, - ::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, - 4>::GemmKernel; - -using DualExpertGemmKernelTN = typename cutlass::gemm::kernel::DefaultGemmGrouped< - ::cutlass::bfloat16_t, - ::cutlass::layout::ColumnMajor, - ::cutlass::ComplexTransform::kNone, - 8, - ::cutlass::bfloat16_t, - ::cutlass::layout::RowMajor, - ::cutlass::ComplexTransform::kNone, - 8, - ::cutlass::bfloat16_t, - ::cutlass::layout::RowMajor, - float, - ::cutlass::arch::OpClassTensorOp, - ::cutlass::arch::Sm80, - ::cutlass::gemm::GemmShape<128, 128, 32>, - ::cutlass::gemm::GemmShape<64, 64, 32>, - ::cutlass::gemm::GemmShape<16, 8, 16>, - ::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>, - ::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, - 4>::GemmKernel; - -using DualExpertGemmKernelNT = typename cutlass::gemm::kernel::DefaultGemmGrouped< - ::cutlass::bfloat16_t, - ::cutlass::layout::RowMajor, - ::cutlass::ComplexTransform::kNone, - 8, - ::cutlass::bfloat16_t, - ::cutlass::layout::ColumnMajor, - ::cutlass::ComplexTransform::kNone, - 8, - ::cutlass::bfloat16_t, - ::cutlass::layout::RowMajor, - float, - ::cutlass::arch::OpClassTensorOp, - ::cutlass::arch::Sm80, - ::cutlass::gemm::GemmShape<128, 128, 32>, - ::cutlass::gemm::GemmShape<64, 64, 32>, - ::cutlass::gemm::GemmShape<16, 8, 16>, - ::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>, - ::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, - 4>::GemmKernel; - -using DualExpertGemmKernelTT = typename cutlass::gemm::kernel::DefaultGemmGrouped< - ::cutlass::bfloat16_t, - ::cutlass::layout::ColumnMajor, - ::cutlass::ComplexTransform::kNone, - 8, - ::cutlass::bfloat16_t, - ::cutlass::layout::ColumnMajor, - ::cutlass::ComplexTransform::kNone, - 8, - ::cutlass::bfloat16_t, - ::cutlass::layout::RowMajor, - float, - ::cutlass::arch::OpClassTensorOp, - ::cutlass::arch::Sm80, - ::cutlass::gemm::GemmShape<128, 128, 32>, - ::cutlass::gemm::GemmShape<64, 64, 32>, - ::cutlass::gemm::GemmShape<16, 8, 16>, - ::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>, - ::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, - 4>::GemmKernel; - -using DualExpertGemmNN = ::cutlass::gemm::device::GemmGrouped; -using DualExpertGemmTN = ::cutlass::gemm::device::GemmGrouped; -using DualExpertGemmNT = ::cutlass::gemm::device::GemmGrouped; -using DualExpertGemmTT = ::cutlass::gemm::device::GemmGrouped; - -template -typename Gemm::Arguments MakeAsymmetricArgumentsSeparated( - torch::Tensor input_expert0, - torch::Tensor input_expert1, - torch::Tensor weight_expert0, - torch::Tensor weight_expert1, - torch::Tensor output_expert0, - torch::Tensor output_expert1) -{ - - TORCH_CHECK(input_expert0.dim() == 2 && input_expert1.dim() == 2, - "Input tensors must be 2D"); - TORCH_CHECK(weight_expert0.dim() == 2 && weight_expert1.dim() == 2, - "Weight tensors must be 2D"); - TORCH_CHECK(output_expert0.dim() == 2 && output_expert1.dim() == 2, - "Output tensors must be 2D"); - - using LayoutA = typename Gemm::LayoutA; - using LayoutB = typename Gemm::LayoutB; - using LayoutC = typename Gemm::LayoutC; - - bool a_is_column_major = std::is_same_v; - bool b_is_column_major = std::is_same_v; - - int64_t m0, k0, n0; - if (a_is_column_major) - { - m0 = input_expert0.size(1); - k0 = input_expert0.size(0); - } - else - { - m0 = input_expert0.size(0); - k0 = input_expert0.size(1); - } - - if (b_is_column_major) - { - n0 = weight_expert0.size(0); - TORCH_CHECK(weight_expert0.size(1) == k0, "Expert 0: k dimensions must match"); - } - else - { - n0 = weight_expert0.size(1); - TORCH_CHECK(weight_expert0.size(0) == k0, "Expert 0: k dimensions must match"); - } - - int64_t m1, k1, n1; - if (a_is_column_major) - { - m1 = input_expert1.size(1); - k1 = input_expert1.size(0); - } - else - { - m1 = input_expert1.size(0); - k1 = input_expert1.size(1); - } - - if (b_is_column_major) - { - n1 = weight_expert1.size(0); - TORCH_CHECK(weight_expert1.size(1) == k1, "Expert 1: k dimensions must match"); - } - else - { - n1 = weight_expert1.size(1); - TORCH_CHECK(weight_expert1.size(0) == k1, "Expert 1: k dimensions must match"); - } - - std::vector problem_sizes_host(2); - problem_sizes_host[0] = cutlass::gemm::GemmCoord(m0, n0, k0); - problem_sizes_host[1] = cutlass::gemm::GemmCoord(m1, n1, k1); - - int64_t num_experts = 2; - int threadblock_count = Gemm::sufficient(problem_sizes_host.data(), num_experts); - if (!threadblock_count) - { - TORCH_CHECK(false, "Dual Expert Grouped GEMM execution not possible with HW"); - } - - std::vector lda_host(num_experts); - std::vector ldb_host(num_experts); - std::vector ldc_host(num_experts); - - using ElementA = typename Gemm::ElementA; - using ElementB = typename Gemm::ElementB; - using ElementC = typename Gemm::ElementC; - - std::vector ptr_a_host(num_experts); - std::vector ptr_b_host(num_experts); - std::vector ptr_c_host(num_experts); - - auto problem_0 = problem_sizes_host[0]; - lda_host[0] = LayoutA::packed({problem_0.m(), problem_0.k()}).stride(0); - ldb_host[0] = LayoutB::packed({problem_0.k(), problem_0.n()}).stride(0); - ldc_host[0] = LayoutC::packed({problem_0.m(), problem_0.n()}).stride(0); - - ptr_a_host[0] = (ElementA *)input_expert0.data_ptr(); - ptr_b_host[0] = (ElementB *)weight_expert0.data_ptr(); - ptr_c_host[0] = (ElementC *)output_expert0.data_ptr(); - - auto problem_1 = problem_sizes_host[1]; - lda_host[1] = LayoutA::packed({problem_1.m(), problem_1.k()}).stride(0); - ldb_host[1] = LayoutB::packed({problem_1.k(), problem_1.n()}).stride(0); - ldc_host[1] = LayoutC::packed({problem_1.m(), problem_1.n()}).stride(0); - - ptr_a_host[1] = (ElementA *)input_expert1.data_ptr(); - ptr_b_host[1] = (ElementB *)weight_expert1.data_ptr(); - ptr_c_host[1] = (ElementC *)output_expert1.data_ptr(); - - torch::Tensor lda = CopyToDevice(lda_host, input_expert0.device()); - torch::Tensor ldb = CopyToDevice(ldb_host, input_expert0.device()); - torch::Tensor ldc = CopyToDevice(ldc_host, input_expert0.device()); - torch::Tensor ptr_a = CopyToDevice(ptr_a_host, input_expert0.device()); - torch::Tensor ptr_b = CopyToDevice(ptr_b_host, input_expert0.device()); - torch::Tensor ptr_c = CopyToDevice(ptr_c_host, input_expert0.device()); - torch::Tensor problem_sizes = CopyToDevice(problem_sizes_host, input_expert0.device()); - - typename Gemm::EpilogueOutputOp::Params epilogue_op(/*alpha=*/1.0f, /*beta=*/0.0f); - typename Gemm::Arguments arguments( - (cutlass::gemm::GemmCoord *)problem_sizes.data_ptr(), - (int)num_experts, - (int)threadblock_count, - epilogue_op, - (ElementA **)ptr_a.data_ptr(), - (ElementB **)ptr_b.data_ptr(), - (ElementC **)ptr_c.data_ptr(), - (ElementC **)ptr_c.data_ptr(), - (int64_t *)lda.data_ptr(), - (int64_t *)ldb.data_ptr(), - (int64_t *)ldc.data_ptr(), - (int64_t *)ldc.data_ptr(), - (cutlass::gemm::GemmCoord *)problem_sizes_host.data()); - - return arguments; -} - -template -void executeDualExpertGemm( - torch::Tensor input_expert0, - torch::Tensor input_expert1, - torch::Tensor weight_expert0, - torch::Tensor weight_expert1, - torch::Tensor output_expert0, - torch::Tensor output_expert1) -{ - - Gemm gemm; - - auto arguments = MakeAsymmetricArgumentsSeparated( - input_expert0, input_expert1, - weight_expert0, weight_expert1, - output_expert0, output_expert1); - - int64_t workspace_size = gemm.get_workspace_size(arguments); - auto options = torch::TensorOptions().dtype(torch::kInt8).device(input_expert0.device()); - torch::Tensor workspace = torch::empty(workspace_size, options); - - if (gemm.initialize(arguments, workspace.data_ptr()) != cutlass::Status::kSuccess) - { - TORCH_CHECK(false, "Failed to initialize CUTLASS Asymmetric Dual Expert GEMM"); - } - - if (gemm.run(c10::cuda::getCurrentCUDAStream()) != cutlass::Status::kSuccess) - { - TORCH_CHECK(false, "Failed to run CUTLASS Asymmetric Dual Expert GEMM"); - } -} - -void AsymmetricDualExpertGemm( - torch::Tensor input_expert0, - torch::Tensor input_expert1, - torch::Tensor weight_expert0, - torch::Tensor weight_expert1, - torch::Tensor output_expert0, - torch::Tensor output_expert1, - bool trans_a, bool trans_b) -{ - - TORCH_CHECK(input_expert0.device() == input_expert1.device() && - input_expert0.device() == weight_expert0.device() && - input_expert0.device() == weight_expert1.device() && - input_expert0.device() == output_expert0.device() && - input_expert0.device() == output_expert1.device(), - "All tensors must be on the same device"); - - TORCH_CHECK(input_expert0.device().is_cuda(), - "All tensors must be on CUDA device for CUTLASS GEMM"); - - TORCH_CHECK(input_expert0.dtype() == torch::kBFloat16, - "All tensors must be BFloat16 for this kernel"); - - torch::Tensor input0_contiguous = input_expert0.contiguous(); - torch::Tensor input1_contiguous = input_expert1.contiguous(); - torch::Tensor weight0_contiguous = weight_expert0.contiguous(); - torch::Tensor weight1_contiguous = weight_expert1.contiguous(); - torch::Tensor output0_contiguous = output_expert0.contiguous(); - torch::Tensor output1_contiguous = output_expert1.contiguous(); - - if (!trans_a && !trans_b) - { - using Gemm = DualExpertGemmNN; - executeDualExpertGemm(input0_contiguous, input1_contiguous, - weight0_contiguous, weight1_contiguous, - output0_contiguous, output1_contiguous); - } - else if (trans_a && !trans_b) - { - using Gemm = DualExpertGemmTN; - executeDualExpertGemm(input0_contiguous, input1_contiguous, - weight0_contiguous, weight1_contiguous, - output0_contiguous, output1_contiguous); - } - else if (!trans_a && trans_b) - { - using Gemm = DualExpertGemmNT; - executeDualExpertGemm(input0_contiguous, input1_contiguous, - weight0_contiguous, weight1_contiguous, - output0_contiguous, output1_contiguous); - } - else - { - using Gemm = DualExpertGemmTT; - executeDualExpertGemm(input0_contiguous, input1_contiguous, - weight0_contiguous, weight1_contiguous, - output0_contiguous, output1_contiguous); - } -} diff --git a/csrc/dual_asym_grouped_gemm.h b/csrc/dual_asym_grouped_gemm.h deleted file mode 100644 index 4d7abde..0000000 --- a/csrc/dual_asym_grouped_gemm.h +++ /dev/null @@ -1,10 +0,0 @@ -#include - -void AsymmetricDualExpertGemm( - torch::Tensor input_expert0, - torch::Tensor input_expert1, - torch::Tensor weight_expert0, - torch::Tensor weight_expert1, - torch::Tensor output_expert0, - torch::Tensor output_expert1, - bool trans_a, bool trans_b); diff --git a/csrc/ops.cu b/csrc/ops.cu deleted file mode 100644 index 80e2f74..0000000 --- a/csrc/ops.cu +++ /dev/null @@ -1,21 +0,0 @@ -#include "dual_asym_grouped_gemm.h" -#include "permute.h" -#include "rope.h" -#include "rope_index.h" -#include "rot_pos.h" -#include "window_index.h" - -#include - - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("asym_dual_gmm", &AsymmetricDualExpertGemm, "Asymmetric Dual Expert Grouped GEMM."); - m.def("permute", &moe_permute_topK_op, "Token permutation kernel"); - m.def("unpermute", &moe_recover_topK_op, "Token un-permutation kernel"); - m.def("unpermute_bwd", &moe_recover_topK_bwd_op, "Token un-permutation backward kernel"); - m.def("rope", &launch_multimodal_rope_forward, "Multimodal RoPE forward kernel"); - m.def("rope_bwd", &launch_multimodal_rope_backward, "Multimodal RoPE backward kernel"); - m.def("rope_index", &get_rope_index, "Get RoPE index kernel"); - m.def("rot_pos_emb", &fused_rot_pos_emb_cuda, "Fused Rotary Position Embedding kernel"); - m.def("get_window_index", &get_window_index_cuda, "Get window index kernel"); -} diff --git a/csrc/permute.cu b/csrc/permute.cu deleted file mode 100644 index 7eb7c26..0000000 --- a/csrc/permute.cu +++ /dev/null @@ -1,942 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include "permute.h" - -#include -#include -#include - -#include "cuda_runtime.h" -#include "device_launch_parameters.h" - -#include "ATen/cuda/CUDAContext.h" - -#include "cutlass/arch/memory.h" -#include "cutlass/arch/cache_operation.h" -#include "cutlass/array.h" -#include "cutlass/numeric_conversion.h" - - -using torch::Tensor; - - -template -inline T *get_ptr(torch::Tensor &t) -{ - return reinterpret_cast(t.data_ptr()); -} - -///////////////////////////////////////////////////////////////////////////////////////////////// -// -// Top K -// -///////////////////////////////////////////////////////////////////////////////////////////////// - -static __global__ void moe_permute_topK_row_map( - const int *sorted_row_id, - int *row_id_map, - const int num_rows, - const int num_topK, - const int num_out_tokens) -{ - // Each block corresponds to one source token - // row_id_map[num_topK][num_rows] - const int bid = blockIdx.x; - const int tid = threadIdx.x; - const int idx = bid * blockDim.x + tid; - - if (idx >= num_rows * num_topK) - return; - - int source_row = sorted_row_id[idx]; - int source_token_id = source_row / num_topK; - int source_topK_id = source_row % num_topK; - - if (idx >= num_out_tokens) - { - row_id_map[source_topK_id * num_rows + source_token_id] = -1; - } - else - { - row_id_map[source_topK_id * num_rows + source_token_id] = idx; - } -} - -template -__global__ void moe_recover_topK_kernel(const T *input, - T *unpermuted_output, - const int *row_id_map, - const float *prob, - const int num_rows, - const int num_topK, - const int num_cols) -{ - extern __shared__ int8_t s_mem[]; - TCompute *s_prob = reinterpret_cast(s_mem); - - using FragmentLoadStore = cutlass::Array; - using FragmentCompute = cutlass::Array; - - cutlass::NumericArrayConverter src_converter; - cutlass::NumericArrayConverter dst_converter; - - // each block corresponds to one source token - const int source_token = blockIdx.x; - const int tid = threadIdx.x; - - if (hasProb) - { - for (int i = tid; i < num_topK; i += blockDim.x * blockDim.y) - { - s_prob[i] = TCompute(prob[source_token * num_topK + i]); - } - __syncthreads(); - } - - for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess) - { - FragmentLoadStore frag_load_store; - FragmentCompute frag_elem; - FragmentCompute frag_sum; - - int source_row = row_id_map[source_token]; - - if (source_row != -1) - { - const T *source_row_ptr = input + source_row * num_cols; - - cutlass::arch::global_load( - frag_load_store, (source_row_ptr + i), true); - frag_sum = src_converter(frag_load_store); - - if (hasProb) - { - frag_sum = frag_sum * s_prob[0]; - } - } - else - { - frag_sum.clear(); - } - - for (int k = 1; k < num_topK; k++) - { - source_row = row_id_map[k * num_rows + source_token]; - - if (source_row == -1) - continue; - - const T *source_row_ptr = input + source_row * num_cols; - - cutlass::arch::global_load( - frag_load_store, (source_row_ptr + i), true); - frag_elem = src_converter(frag_load_store); - - if (hasProb) - { - frag_elem = frag_elem * s_prob[k]; - } - - for (int e = 0; e < kElementsPerAccess; e++) - { - frag_sum.at(e) = frag_sum.at(e) + frag_elem.at(e); - } - } - - T *dest_row_ptr = unpermuted_output + source_token * num_cols; - frag_load_store = dst_converter(frag_sum); - *(float4 *)(dest_row_ptr + i) = *(float4 *)(frag_load_store.data()); - } -} - -template -__global__ void moe_permute_topK_kernel(const T *input_bwd, - const T *input_fwd, - T *act_grad, - const float *prob, - float *prob_grad, - const int *row_id_map, - const int num_rows, - const int num_topK, - const int num_cols) -{ - extern __shared__ int8_t s_mem[]; - TCompute *s_prob = reinterpret_cast(s_mem); - - using FragmentLoadStore = cutlass::Array; - using FragmentCompute = cutlass::Array; - - cutlass::NumericArrayConverter src_converter; - cutlass::NumericArrayConverter dst_converter; - - const int source_token = blockIdx.x; - const int tid = threadIdx.x; - - if (hasProb) - { - for (int i = tid; i < num_topK; i += blockDim.x) - { - s_prob[i] = TCompute(prob[source_token * num_topK + i]); - } - __syncthreads(); - } - - float accum[topKTile] = {0.0f}; - FragmentLoadStore frag_load_store; - - const T *source_row_ptr = input_bwd + source_token * num_cols; - for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess) - { - cutlass::arch::global_load( - frag_load_store, (source_row_ptr + i), true); - FragmentCompute frag_src = src_converter(frag_load_store); - - int index = source_token; - - for (int k = 0; k < topKTile; k++) - { - if (k == num_topK) break; - - int dest_row = row_id_map[index]; - index += num_rows; - - if (dest_row == -1) - continue; - - if (hasProb) - { - frag_load_store = dst_converter(frag_src * s_prob[k]); - } - else - { - frag_load_store = dst_converter(frag_src); - } - - T *dest_row_ptr = act_grad + dest_row * num_cols; - *(float4 *)(dest_row_ptr + i) = *(float4 *)(frag_load_store.data()); - - if (hasProb) - { - const T *input_fwd_ptr = input_fwd + dest_row * num_cols; - cutlass::arch::global_load( - frag_load_store, (input_fwd_ptr + i), true); - FragmentCompute frag_input_fwd = src_converter(frag_load_store); - - for (int e = 0; e < kElementsPerAccess; e++) - { - accum[k] += float(frag_src.at(e) * frag_input_fwd.at(e)); - } - } - } - } - - if (hasProb) - { - for (int k = 0; k < topKTile; k++) - { - if (k == num_topK) break; - - for (int mask = 16; mask > 0; mask /= 2) - { - accum[k] = accum[k] + __shfl_xor_sync(0xffffffff, accum[k], mask, 32); - } - } - - if (tid == 0) - { - for (int k = 0; k < topKTile; k++) - { - if (k == num_topK) break; - prob_grad[source_token * num_topK + k] = accum[k]; - } - } - } -} - - -template -void moe_permute_topK_kernel_launcher( - const T *input, - T *output, - const int *sorted_row_id, - int *row_id_map, - const float *prob, - const int num_rows, - const int num_topK, - const int num_cols, - const int num_out_tokens, - cudaStream_t stream, - float *prob_grad = nullptr, - const T *input_fwd = nullptr) -{ - if (FWD) - { - if (prob_grad == nullptr) - { - // permute_topK fwd - int threads = 64; - int blocks = (num_rows * num_topK + threads - 1) / threads; - moe_permute_topK_row_map<<>>( - sorted_row_id, - row_id_map, - num_rows, - num_topK, - num_out_tokens); - - blocks = num_rows; - threads = std::min(num_cols / kElementsPerAccess, 1024); - moe_permute_topK_kernel<<>>( - input, - nullptr, - output, - nullptr, - nullptr, - row_id_map, - num_rows, - num_topK, - num_cols); - } - else - { - // unpermute_topK bwd - int blocks = num_rows; - int threads = 32; - size_t smem_bytes = num_topK * sizeof(TCompute); - - if (num_topK == 1) - { - moe_permute_topK_kernel<<>>( - input, - input_fwd, - output, - prob, - prob_grad, - row_id_map, - num_rows, - num_topK, - num_cols); - } - else if (num_topK <= 8) - { - moe_permute_topK_kernel<<>>( - input, - input_fwd, - output, - prob, - prob_grad, - row_id_map, - num_rows, - num_topK, - num_cols); - } - else if (num_topK <= 16) - { - moe_permute_topK_kernel<<>>( - input, - input_fwd, - output, - prob, - prob_grad, - row_id_map, - num_rows, - num_topK, - num_cols); - } - else if (num_topK <= 32) - { - moe_permute_topK_kernel<<>>( - input, - input_fwd, - output, - prob, - prob_grad, - row_id_map, - num_rows, - num_topK, - num_cols); - } - else if (num_topK <= 64) - { - moe_permute_topK_kernel<<>>( - input, - input_fwd, - output, - prob, - prob_grad, - row_id_map, - num_rows, - num_topK, - num_cols); - } - else if (num_topK <= 128) - { - moe_permute_topK_kernel<<>>( - input, - input_fwd, - output, - prob, - prob_grad, - row_id_map, - num_rows, - num_topK, - num_cols); - } - else - { - throw std::runtime_error("num_topK cannot exceed 128."); - } - } - } - else - { - int blocks = num_rows; - int threads = std::min(num_cols / kElementsPerAccess, 1024); - size_t smem_bytes = num_topK * sizeof(TCompute); - - - if (num_topK == 1) - { - // permute_topK bwd with topK==1 - moe_recover_topK_kernel<<>>( - input, - output, - row_id_map, - prob, - num_rows, - num_topK, - num_cols); - } - else if (prob == nullptr) - { - // permute_topK bwd - moe_recover_topK_kernel<<>>( - input, - output, - row_id_map, - prob, - num_rows, - num_topK, - num_cols); - } - else - { - // unpermute_topK fwd - moe_recover_topK_kernel<<>>( - input, - output, - row_id_map, - prob, - num_rows, - num_topK, - num_cols); - } - } -} - -///////////////////////////////////////////////////////////////////////////////////////////////// -// -// Permute_topK OP -// -///////////////////////////////////////////////////////////////////////////////////////////////// - -std::tuple> moe_permute_topK_op( - Tensor input, - Tensor indices, - int64_t num_out_tokens, - std::vector workspace, - int64_t max_expanded_token_num) -{ - const int num_tokens = input.size(0); - const int num_cols = input.size(1); - const int num_topK = indices.size(1); - - // initialize the workspace on the first run - if (workspace.empty()) { - auto options = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false); - - Tensor sorted_indices = torch::empty(max_expanded_token_num, options); - Tensor row_id = torch::range(0, max_expanded_token_num - 1, 1, options); - Tensor sorted_row_id = - torch::empty(max_expanded_token_num, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false)); - - size_t temp_storage_bytes = 0; - int *temp_ptr = nullptr; - cub::DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, - temp_ptr, temp_ptr, - temp_ptr, temp_ptr, max_expanded_token_num); - Tensor temp_storage = - torch::empty(temp_storage_bytes, torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false)); - - workspace.push_back(sorted_indices); - workspace.push_back(row_id); - workspace.push_back(sorted_row_id); - workspace.push_back(temp_storage); - } - - int *indices_ptr = get_ptr(indices); - int *sorted_indices_ptr = get_ptr(workspace[0]); - int *row_id_ptr = get_ptr(workspace[1]); - int *sorted_row_id_ptr = get_ptr(workspace[2]); - - void *d_temp_storage = get_ptr(workspace[3]); - size_t temp_storage_bytes = std::numeric_limits::max(); - - cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, - indices_ptr, sorted_indices_ptr, - row_id_ptr, sorted_row_id_ptr, num_tokens * num_topK); - - // activations type - const at::ScalarType _st = input.scalar_type(); - - // Output buffer alloc - num_out_tokens = (num_out_tokens > 0) ? num_out_tokens : num_tokens * num_topK; - Tensor permuted_output = - torch::empty({num_out_tokens, num_cols}, torch::dtype(_st).device(torch::kCUDA).requires_grad(false)); - Tensor row_id_map = - torch::empty({num_tokens * num_topK}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false)); - - int *row_id_map_ptr = get_ptr(row_id_map); - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - switch (_st) - { - case at::ScalarType::Float: - { - using dType = float; - using dTypeCompute = float; - - dType *input_ptr = get_ptr(input); - dType *permuted_output_ptr = get_ptr(permuted_output); - - moe_permute_topK_kernel_launcher( - input_ptr, - permuted_output_ptr, - sorted_row_id_ptr, - row_id_map_ptr, - nullptr, - num_tokens, - num_topK, - num_cols, - num_out_tokens, - stream); - - break; - } - case at::ScalarType::Half: - { - using dType = cutlass::half_t; - using dTypeCompute = cutlass::half_t; - - dType *input_ptr = get_ptr(input); - dType *permuted_output_ptr = get_ptr(permuted_output); - - moe_permute_topK_kernel_launcher( - input_ptr, - permuted_output_ptr, - sorted_row_id_ptr, - row_id_map_ptr, - nullptr, - num_tokens, - num_topK, - num_cols, - num_out_tokens, - stream); - - break; - } -#ifdef ENABLE_BF16 - case at::ScalarType::BFloat16: - { - using dType = cutlass::bfloat16_t; - using dTypeCompute = cutlass::bfloat16_t; - - dType *input_ptr = get_ptr(input); - dType *permuted_output_ptr = get_ptr(permuted_output); - - moe_permute_topK_kernel_launcher( - input_ptr, - permuted_output_ptr, - sorted_row_id_ptr, - row_id_map_ptr, - nullptr, - num_tokens, - num_topK, - num_cols, - num_out_tokens, - stream); - - break; - } -#endif -#ifdef ENABLE_FP8 - case at::ScalarType::Float8_e5m2: - { - using dType = cutlass::float_e5m2_t; - using dTypeCompute = cutlass::half_t; - - dType *input_ptr = get_ptr(input); - dType *permuted_output_ptr = get_ptr(permuted_output); - - moe_permute_topK_kernel_launcher( - input_ptr, - permuted_output_ptr, - sorted_row_id_ptr, - row_id_map_ptr, - nullptr, - num_tokens, - num_topK, - num_cols, - num_out_tokens, - stream); - - break; - } - case at::ScalarType::Float8_e4m3fn: - { - using dType = cutlass::float_e4m3_t; - using dTypeCompute = cutlass::half_t; - - dType *input_ptr = get_ptr(input); - dType *permuted_output_ptr = get_ptr(permuted_output); - - moe_permute_topK_kernel_launcher( - input_ptr, - permuted_output_ptr, - sorted_row_id_ptr, - row_id_map_ptr, - nullptr, - num_tokens, - num_topK, - num_cols, - num_out_tokens, - stream); - - break; - } -#endif - default: - throw std::runtime_error("Wrong activation tensor type."); - } - - return std::make_tuple(permuted_output, row_id_map, workspace); -} - -///////////////////////////////////////////////////////////////////////////////////////////////// -// -// Unpermute_topK OP -// -///////////////////////////////////////////////////////////////////////////////////////////////// - -Tensor moe_recover_topK_op( - Tensor input, - Tensor row_id_map, - Tensor prob, - int64_t num_tokens, - int64_t num_topK) -{ - const int num_cols = input.size(1); - - // activations type - const at::ScalarType _st = input.scalar_type(); - - // Output buffer alloc - Tensor unpermuted_output = - torch::empty({num_tokens, num_cols}, torch::dtype(_st).device(torch::kCUDA).requires_grad(false)); - - int *row_id_map_ptr = get_ptr(row_id_map); - float *prob_ptr = (prob.defined()) ? get_ptr(prob) : nullptr; - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - switch (_st) - { - case at::ScalarType::Float: - { - using dType = float; - using dTypeCompute = float; - - dType *input_ptr = get_ptr(input); - dType *unpermuted_output_ptr = get_ptr(unpermuted_output); - - moe_permute_topK_kernel_launcher( - input_ptr, - unpermuted_output_ptr, - nullptr, - row_id_map_ptr, - prob_ptr, - num_tokens, - num_topK, - num_cols, - 0, - stream); - - break; - } - case at::ScalarType::Half: - { - using dType = cutlass::half_t; - using dTypeCompute = cutlass::half_t; - - dType *input_ptr = get_ptr(input); - dType *unpermuted_output_ptr = get_ptr(unpermuted_output); - - moe_permute_topK_kernel_launcher( - input_ptr, - unpermuted_output_ptr, - nullptr, - row_id_map_ptr, - prob_ptr, - num_tokens, - num_topK, - num_cols, - 0, - stream); - - break; - } -#ifdef ENABLE_BF16 - case at::ScalarType::BFloat16: - { - using dType = cutlass::bfloat16_t; - using dTypeCompute = cutlass::bfloat16_t; - - dType *input_ptr = get_ptr(input); - dType *unpermuted_output_ptr = get_ptr(unpermuted_output); - - moe_permute_topK_kernel_launcher( - input_ptr, - unpermuted_output_ptr, - nullptr, - row_id_map_ptr, - prob_ptr, - num_tokens, - num_topK, - num_cols, - 0, - stream); - - break; - } -#endif -#ifdef ENABLE_FP8 - case at::ScalarType::Float8_e5m2: - { - using dType = cutlass::float_e5m2_t; - using dTypeCompute = cutlass::half_t; - - dType *input_ptr = get_ptr(input); - dType *unpermuted_output_ptr = get_ptr(unpermuted_output); - - moe_permute_topK_kernel_launcher( - input_ptr, - unpermuted_output_ptr, - nullptr, - row_id_map_ptr, - prob_ptr, - num_tokens, - num_topK, - num_cols, - 0, - stream); - - break; - } - case at::ScalarType::Float8_e4m3fn: - { - using dType = cutlass::float_e4m3_t; - using dTypeCompute = cutlass::half_t; - - dType *input_ptr = get_ptr(input); - dType *unpermuted_output_ptr = get_ptr(unpermuted_output); - - moe_permute_topK_kernel_launcher( - input_ptr, - unpermuted_output_ptr, - nullptr, - row_id_map_ptr, - prob_ptr, - num_tokens, - num_topK, - num_cols, - 0, - stream); - - break; - } -#endif - default: - throw std::runtime_error("Wrong activation tensor type."); - } - - return unpermuted_output; -} - -std::tuple moe_recover_topK_bwd_op( - Tensor input_bwd, - Tensor input_fwd, - Tensor row_id_map, - Tensor prob) -{ - const int num_tokens = prob.size(0); - const int num_topK = prob.size(1); - const int num_cols = input_bwd.size(1); - - int *row_id_map_ptr = get_ptr(row_id_map); - float *prob_ptr = get_ptr(prob); - - // activations type - const at::ScalarType _st = input_bwd.scalar_type(); - - // Output buffer alloc - Tensor act_grad = - torch::empty({input_fwd.size(0), num_cols}, torch::dtype(_st).device(torch::kCUDA).requires_grad(false)); - Tensor prob_grad = - torch::empty({num_tokens, num_topK}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false)); - float *prob_grad_ptr = get_ptr(prob_grad); - - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - switch (_st) - { - case at::ScalarType::Float: - { - using dType = float; - using dTypeCompute = float; - - dType *input_bwd_ptr = get_ptr(input_bwd); - dType *input_fwd_ptr = get_ptr(input_fwd); - dType *act_grad_ptr = get_ptr(act_grad); - - moe_permute_topK_kernel_launcher( - input_bwd_ptr, - act_grad_ptr, - nullptr, - row_id_map_ptr, - prob_ptr, - num_tokens, - num_topK, - num_cols, - 0, - stream, - prob_grad_ptr, - input_fwd_ptr); - - break; - } - case at::ScalarType::Half: - { - using dType = cutlass::half_t; - using dTypeCompute = cutlass::half_t; - - dType *input_bwd_ptr = get_ptr(input_bwd); - dType *input_fwd_ptr = get_ptr(input_fwd); - dType *act_grad_ptr = get_ptr(act_grad); - - moe_permute_topK_kernel_launcher( - input_bwd_ptr, - act_grad_ptr, - nullptr, - row_id_map_ptr, - prob_ptr, - num_tokens, - num_topK, - num_cols, - 0, - stream, - prob_grad_ptr, - input_fwd_ptr); - - break; - } -#ifdef ENABLE_BF16 - case at::ScalarType::BFloat16: - { - using dType = cutlass::bfloat16_t; - using dTypeCompute = cutlass::bfloat16_t; - - dType *input_bwd_ptr = get_ptr(input_bwd); - dType *input_fwd_ptr = get_ptr(input_fwd); - dType *act_grad_ptr = get_ptr(act_grad); - - moe_permute_topK_kernel_launcher( - input_bwd_ptr, - act_grad_ptr, - nullptr, - row_id_map_ptr, - prob_ptr, - num_tokens, - num_topK, - num_cols, - 0, - stream, - prob_grad_ptr, - input_fwd_ptr); - - break; - } -#endif -#ifdef ENABLE_FP8 - case at::ScalarType::Float8_e5m2: - { - using dType = cutlass::float_e5m2_t; - using dTypeCompute = cutlass::half_t; - - dType *input_bwd_ptr = get_ptr(input_bwd); - dType *input_fwd_ptr = get_ptr(input_fwd); - dType *act_grad_ptr = get_ptr(act_grad); - - moe_permute_topK_kernel_launcher( - input_bwd_ptr, - act_grad_ptr, - nullptr, - row_id_map_ptr, - prob_ptr, - num_tokens, - num_topK, - num_cols, - 0, - stream, - prob_grad_ptr, - input_fwd_ptr); - - break; - } - case at::ScalarType::Float8_e4m3fn: - { - using dType = cutlass::float_e4m3_t; - using dTypeCompute = cutlass::half_t; - - dType *input_bwd_ptr = get_ptr(input_bwd); - dType *input_fwd_ptr = get_ptr(input_fwd); - dType *act_grad_ptr = get_ptr(act_grad); - - moe_permute_topK_kernel_launcher( - input_bwd_ptr, - act_grad_ptr, - nullptr, - row_id_map_ptr, - prob_ptr, - num_tokens, - num_topK, - num_cols, - 0, - stream, - prob_grad_ptr, - input_fwd_ptr); - - break; - } -#endif - default: - throw std::runtime_error("Wrong activation tensor type."); - } - - return std::make_tuple(act_grad, prob_grad); -} diff --git a/csrc/permute.h b/csrc/permute.h deleted file mode 100644 index 45a8965..0000000 --- a/csrc/permute.h +++ /dev/null @@ -1,31 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#pragma once - -#include - -using torch::Tensor; - -std::tuple> moe_permute_topK_op( - Tensor input, - Tensor indices, - int64_t num_out_tokens, - std::vector workspace, - int64_t max_expanded_token_num); - -torch::Tensor moe_recover_topK_op( - torch::Tensor input, - torch::Tensor row_id_map, - torch::Tensor prob_opt, - int64_t num_tokens, - int64_t num_topK); - -std::tuple moe_recover_topK_bwd_op( - Tensor input_bwd, - Tensor input_fwd, - Tensor row_id_map, - Tensor prob); diff --git a/csrc/rope.cu b/csrc/rope.cu deleted file mode 100644 index 0420ea6..0000000 --- a/csrc/rope.cu +++ /dev/null @@ -1,571 +0,0 @@ -#undef __CUDA_NO_HALF_OPERATORS__ -#undef __CUDA_NO_HALF_CONVERSIONS__ -#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ -#undef __CUDA_NO_HALF2_OPERATORS__ - -#include -#include -#include -#include - -#include -#include -#include -#include - -// Type traits for CUDA types -template -struct CudaTypeTraits -{ - static constexpr bool is_supported = false; -}; - -template <> -struct CudaTypeTraits -{ - static constexpr bool is_supported = true; - using type = float; - static constexpr const char *name = "float"; -}; - -template <> -struct CudaTypeTraits -{ - static constexpr bool is_supported = true; - using type = half; - static constexpr const char *name = "half"; -}; - -template <> -struct CudaTypeTraits<__nv_bfloat16> -{ - static constexpr bool is_supported = true; - using type = __nv_bfloat16; - static constexpr const char *name = "bfloat16"; -}; - -// Optimized forward kernel template -template -__global__ void multimodal_rope_forward_kernel( - const T *__restrict__ q, // [batch, q_heads, seq_len, head_dim] - const T *__restrict__ k, // [batch, kv_heads, seq_len, head_dim] - const T *__restrict__ cos, // [3, batch, seq_len, head_dim] - const T *__restrict__ sin, // [3, batch, seq_len, head_dim] - T *__restrict__ q_out, // [batch, q_heads, seq_len, head_dim] - T *__restrict__ k_out, // [batch, kv_heads, seq_len, head_dim] - const int *__restrict__ mrope_section_doubled, // [32, 48, 48] - int batch_size, - int q_heads, - int kv_heads, - int seq_len, - int head_dim) -{ - static_assert(CudaTypeTraits::is_supported, "Unsupported data type"); - - // Block organization: batch_size * (q_heads + kv_heads) blocks - int batch_idx = blockIdx.x; - int head_idx = blockIdx.y; - int seq_paral_size = gridDim.z; - int seq_paral_idx = blockIdx.z; - - bool is_q_head = head_idx < q_heads; - int actual_head_idx = is_q_head ? head_idx : (head_idx - q_heads); - int total_heads = is_q_head ? q_heads : kv_heads; - - // Each warp processes one token (seq position) - int warp_id = threadIdx.x / 32; - int lane_id = threadIdx.x % 32; - int warps_per_block = blockDim.x / 32; - - // Process tokens in batches across warps - - for (int seq_base = warp_id + seq_paral_idx * warps_per_block; seq_base < seq_len; seq_base += seq_paral_size * warps_per_block) - { - int seq_idx = seq_base; - if (seq_idx >= seq_len) - break; - - // Each lane processes multiple dimensions (head_dim=128, 32 lanes) - constexpr int dims_per_lane = 4; - for (int dim_batch = 0; dim_batch < (head_dim + 31) / 32; ++dim_batch) - { - int dim_start = dim_batch * 32 + lane_id; - if (dim_start >= head_dim) - break; - -#pragma unroll - for (int dim_offset = 0; dim_offset < dims_per_lane; ++dim_offset) - { - int dim_idx = dim_start + dim_offset * 32; - if (dim_idx >= head_dim) - break; - - // Determine which section this dimension belongs to - int section_idx; - int cos_sin_d = dim_idx; - if (dim_idx < mrope_section_doubled[0]) - { - section_idx = 0; - } - else if (dim_idx < mrope_section_doubled[0] + mrope_section_doubled[1]) - { - section_idx = 1; - } - else - { - section_idx = 2; - } - - // Load cos/sin values (coalesced access) - int cos_sin_idx = section_idx * batch_size * seq_len * head_dim + - batch_idx * seq_len * head_dim + - seq_idx * head_dim + cos_sin_d; - - T cos_val = cos[cos_sin_idx]; - T sin_val = sin[cos_sin_idx]; - - // Calculate tensor indices - int tensor_idx = batch_idx * total_heads * seq_len * head_dim + - actual_head_idx * seq_len * head_dim + - seq_idx * head_dim + dim_idx; - - // Get input value - T input_val = is_q_head ? q[tensor_idx] : k[tensor_idx]; - - // Calculate rotate_half value - int half_dim = head_dim / 2; - int rotate_dim = (dim_idx < half_dim) ? (dim_idx + half_dim) : (dim_idx - half_dim); - int rotate_tensor_idx = batch_idx * total_heads * seq_len * head_dim + - actual_head_idx * seq_len * head_dim + - seq_idx * head_dim + rotate_dim; - - T rotate_val = is_q_head ? q[rotate_tensor_idx] : k[rotate_tensor_idx]; - if (dim_idx < half_dim) - { - if constexpr (std::is_same_v) - { - rotate_val = __hneg(rotate_val); - } - else if constexpr (std::is_same_v) - { -#if __CUDA_ARCH__ >= 800 // BFloat16 support requires Ampere or newer - rotate_val = __hneg(rotate_val); // __hneg works for bfloat16 in newer CUDA -#else - rotate_val = __float2bfloat16(-__bfloat162float(rotate_val)); -#endif - } - else - { - rotate_val = -rotate_val; - } - } - - // Apply RoPE: output = input * cos + rotate_half(input) * sin - T output_val = input_val * cos_val + rotate_val * sin_val; - - // Store result - if (is_q_head) - { - q_out[tensor_idx] = output_val; - } - else - { - k_out[tensor_idx] = output_val; - } - } - } - } -} - -// Optimized backward kernel template -template -__global__ void multimodal_rope_backward_kernel( - const T *__restrict__ grad_q_out, // [batch, q_heads, seq_len, head_dim] - const T *__restrict__ grad_k_out, // [batch, kv_heads, seq_len, head_dim] - const T *__restrict__ q, // [batch, q_heads, seq_len, head_dim] - const T *__restrict__ k, // [batch, kv_heads, seq_len, head_dim] - const T *__restrict__ cos, // [3, batch, seq_len, head_dim] - const T *__restrict__ sin, // [3, batch, seq_len, head_dim] - T *__restrict__ grad_q, // [batch, q_heads, seq_len, head_dim] - T *__restrict__ grad_k, // [batch, kv_heads, seq_len, head_dim] - const int *__restrict__ mrope_section_doubled, - int batch_size, - int q_heads, - int kv_heads, - int seq_len, - int head_dim) -{ - int batch_idx = blockIdx.x; - int head_idx = blockIdx.y; - int seq_paral_size = gridDim.z; - int seq_paral_idx = blockIdx.z; - - bool is_q_head = head_idx < q_heads; - int actual_head_idx = is_q_head ? head_idx : (head_idx - q_heads); - int total_heads = is_q_head ? q_heads : kv_heads; - - int warp_id = threadIdx.x / 32; - int lane_id = threadIdx.x % 32; - int warps_per_block = blockDim.x / 32; - - // Process tokens - for (int seq_base = warp_id + seq_paral_idx * warps_per_block; seq_base < seq_len; seq_base += seq_paral_size * warps_per_block) - { - int seq_idx = seq_base; - if (seq_idx >= seq_len) - break; - - // Process dimensions in chunks - much simpler now! - constexpr int dims_per_lane = 4; - for (int dim_batch = 0; dim_batch < (head_dim + 31) / 32; ++dim_batch) - { - int dim_start = dim_batch * 32 + lane_id; - if (dim_start >= head_dim) - break; - -#pragma unroll - for (int dim_offset = 0; dim_offset < dims_per_lane; ++dim_offset) - { - int dim_idx = dim_start + dim_offset * 32; - if (dim_idx >= head_dim) - break; - - // Get section info for cos/sin reconstruction - int section_idx; - int cos_sin_d = dim_idx; - if (dim_idx < mrope_section_doubled[0]) - { - section_idx = 0; - } - else if (dim_idx < mrope_section_doubled[0] + mrope_section_doubled[1]) - { - section_idx = 1; - } - else - { - section_idx = 2; - } - - // Global cos/sin index - int cos_sin_idx = section_idx * batch_size * seq_len * head_dim + - batch_idx * seq_len * head_dim + - seq_idx * head_dim + cos_sin_d; - - // Tensor index for current position - int tensor_idx = batch_idx * total_heads * seq_len * head_dim + - actual_head_idx * seq_len * head_dim + - seq_idx * head_dim + dim_idx; - - // Load values - T cos_val = cos[cos_sin_idx]; - T sin_val = sin[cos_sin_idx]; - T grad_out_val = is_q_head ? grad_q_out[tensor_idx] : grad_k_out[tensor_idx]; - - // Calculate paired dimension for rotate_half - int half_dim = head_dim / 2; - int rotate_dim = (dim_idx < half_dim) ? (dim_idx + half_dim) : (dim_idx - half_dim); - int rotate_tensor_idx = batch_idx * total_heads * seq_len * head_dim + - actual_head_idx * seq_len * head_dim + - seq_idx * head_dim + rotate_dim; - - // Get gradient from the paired dimension - T paired_grad_out = is_q_head ? grad_q_out[rotate_tensor_idx] : grad_k_out[rotate_tensor_idx]; - - // Get paired sin value - int paired_section_idx; - int paired_cos_sin_d = rotate_dim; - if (rotate_dim < mrope_section_doubled[0]) - { - paired_section_idx = 0; - } - else if (rotate_dim < mrope_section_doubled[0] + mrope_section_doubled[1]) - { - paired_section_idx = 1; - } - else - { - paired_section_idx = 2; - } - - int paired_cos_sin_idx = paired_section_idx * batch_size * seq_len * head_dim + - batch_idx * seq_len * head_dim + - seq_idx * head_dim + paired_cos_sin_d; - T paired_sin_val = sin[paired_cos_sin_idx]; - - // === Compute input gradients (the only thing we need!) === - - // Direct term: grad_input = grad_out * cos - T grad_input_direct = grad_out_val * cos_val; - - // Cross term from rotate_half - T grad_input_from_rotate; - if (dim_idx < half_dim) - { - // First half: gets contribution from second half (positive) - grad_input_from_rotate = paired_grad_out * paired_sin_val; - } - else - { - // Second half: gets contribution from first half (negative due to rotate_half) - grad_input_from_rotate = -paired_grad_out * paired_sin_val; - } - - // Total input gradient - T total_grad_input = grad_input_direct + grad_input_from_rotate; - - // Store input gradients (no atomic operations needed!) - if (is_q_head) - { - grad_q[tensor_idx] = total_grad_input; - } - else - { - grad_k[tensor_idx] = total_grad_input; - } - - // No cos/sin gradient computation - they're not learnable parameters! - } - } - } -} - -// Template-based host functions -template -void launch_multimodal_rope_forward_template( - const T *q, const T *k, const T *cos, const T *sin, - T *q_out, T *k_out, - const int *mrope_section_doubled, - int batch_size, int q_heads, int kv_heads, int seq_len, int head_dim, - cudaStream_t stream) -{ - static_assert(CudaTypeTraits::is_supported, "Unsupported data type for multimodal RoPE"); - - // Grid: batch_size x (q_heads + kv_heads) - dim3 grid(batch_size, q_heads + kv_heads, 8); - - // Block: enough threads to handle seq_len with multiple warps - int threads_per_block = min(512, ((seq_len + 3) / 4) * 32); // 4 warps max - threads_per_block = ((threads_per_block + 31) / 32) * 32; // Round to warp size - - multimodal_rope_forward_kernel<<>>( - q, k, cos, sin, q_out, k_out, mrope_section_doubled, - batch_size, q_heads, kv_heads, seq_len, head_dim); -} - -template -void launch_multimodal_rope_backward_template( - const T *grad_q_out, const T *grad_k_out, - const T *q, const T *k, const T *cos, const T *sin, - T *grad_q, T *grad_k, const int *mrope_section_doubled, - int batch_size, int q_heads, int kv_heads, int seq_len, int head_dim, - cudaStream_t stream) -{ - static_assert(CudaTypeTraits::is_supported, "Unsupported data type for multimodal RoPE"); - - dim3 grid(batch_size, q_heads + kv_heads, 8); - int threads_per_block = min(512, ((seq_len + 3) / 4) * 32); - threads_per_block = ((threads_per_block + 31) / 32) * 32; - - multimodal_rope_backward_kernel<<>>( - grad_q_out, grad_k_out, q, k, cos, sin, - grad_q, grad_k, mrope_section_doubled, - batch_size, q_heads, kv_heads, seq_len, head_dim); -} - -// Enum for data type dispatch -enum class DataType -{ - FLOAT32, - FLOAT16, - BFLOAT16 -}; - -// Template dispatcher -template -struct DataTypeDispatcher; - -template <> -struct DataTypeDispatcher -{ - using type = float; -}; - -template <> -struct DataTypeDispatcher -{ - using type = half; -}; - -template <> -struct DataTypeDispatcher -{ - using type = __nv_bfloat16; -}; - -// Type-safe host interface -template -void launch_multimodal_rope_forward_typed( - const void *q, const void *k, const void *cos, const void *sin, - void *q_out, void *k_out, - const int *mrope_section_doubled, - int batch_size, int q_heads, int kv_heads, int seq_len, int head_dim, - cudaStream_t stream) -{ - using T = typename DataTypeDispatcher

::type; - - launch_multimodal_rope_forward_template( - static_cast(q), - static_cast(k), - static_cast(cos), - static_cast(sin), - static_cast(q_out), - static_cast(k_out), - mrope_section_doubled, - batch_size, q_heads, kv_heads, seq_len, head_dim, - stream); -} - -template -void launch_multimodal_rope_backward_typed( - const void *grad_q_out, const void *grad_k_out, - const void *q, const void *k, const void *cos, const void *sin, - void *grad_q, void *grad_k, - const int *mrope_section_doubled, - int batch_size, int q_heads, int kv_heads, int seq_len, int head_dim, - cudaStream_t stream) -{ - using T = typename DataTypeDispatcher
::type; - - launch_multimodal_rope_backward_template( - static_cast(grad_q_out), - static_cast(grad_k_out), - static_cast(q), - static_cast(k), - static_cast(cos), - static_cast(sin), - static_cast(grad_q), - static_cast(grad_k), - mrope_section_doubled, - batch_size, q_heads, kv_heads, seq_len, head_dim, - stream); -} - -void launch_multimodal_rope_forward( - torch::Tensor q, torch::Tensor k, torch::Tensor cos, torch::Tensor sin, - torch::Tensor q_out, torch::Tensor k_out, - std::vector mrope_section_doubled) -{ - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - - int batch_size = q.size(0); - int q_heads = q.size(1); - int seq_len = q.size(2); - int head_dim = q.size(3); - int kv_heads = k.size(1); - - int data_type; - if (q.scalar_type() == torch::kFloat32) - { - data_type = 0; - } - else if (q.scalar_type() == torch::kFloat16) - { - data_type = 1; - } - else if (q.scalar_type() == torch::kBFloat16) - { - data_type = 2; - } - - auto mrope_tensor = torch::from_blob( - mrope_section_doubled.data(), - {3}, - torch::TensorOptions().dtype(torch::kInt32) - ).to(q.device(), /*non_blocking=*/true); - - int *d_mrope_section_doubled = static_cast(mrope_tensor.data_ptr()); - - switch (data_type) - { - case 0: // float32 - launch_multimodal_rope_forward_typed( - q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(), - q_out.data_ptr(), k_out.data_ptr(), d_mrope_section_doubled, - batch_size, q_heads, kv_heads, seq_len, head_dim, stream); - break; - case 1: // float16 - launch_multimodal_rope_forward_typed( - q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(), - q_out.data_ptr(), k_out.data_ptr(), d_mrope_section_doubled, - batch_size, q_heads, kv_heads, seq_len, head_dim, stream); - break; - case 2: // bfloat16 - launch_multimodal_rope_forward_typed( - q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(), - q_out.data_ptr(), k_out.data_ptr(), d_mrope_section_doubled, - batch_size, q_heads, kv_heads, seq_len, head_dim, stream); - break; - } -} - -void launch_multimodal_rope_backward( - torch::Tensor grad_q_out, torch::Tensor grad_k_out, - torch::Tensor q, torch::Tensor k, torch::Tensor cos, torch::Tensor sin, - torch::Tensor grad_q, torch::Tensor grad_k, - std::vector mrope_section_doubled) -{ - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - - int batch_size = q.size(0); - int q_heads = q.size(1); - int seq_len = q.size(2); - int head_dim = q.size(3); - int kv_heads = k.size(1); - - int data_type; - if (q.scalar_type() == torch::kFloat32) - { - data_type = 0; - } - else if (q.scalar_type() == torch::kFloat16) - { - data_type = 1; - } - else if (q.scalar_type() == torch::kBFloat16) - { - data_type = 2; - } - - auto mrope_tensor = torch::from_blob( - mrope_section_doubled.data(), - {3}, - torch::TensorOptions().dtype(torch::kInt32) - ).to(q.device(), /*non_blocking=*/true); - - int *d_mrope_section_doubled = static_cast(mrope_tensor.data_ptr()); - - switch (data_type) - { - case 0: // float32 - launch_multimodal_rope_backward_typed( - grad_q_out.data_ptr(), grad_k_out.data_ptr(), - q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(), - grad_q.data_ptr(), grad_k.data_ptr(), d_mrope_section_doubled, - batch_size, q_heads, kv_heads, seq_len, head_dim, stream); - break; - case 1: // float16 - launch_multimodal_rope_backward_typed( - grad_q_out.data_ptr(), grad_k_out.data_ptr(), - q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(), - grad_q.data_ptr(), grad_k.data_ptr(), d_mrope_section_doubled, - batch_size, q_heads, kv_heads, seq_len, head_dim, stream); - break; - case 2: // bfloat16 - launch_multimodal_rope_backward_typed( - grad_q_out.data_ptr(), grad_k_out.data_ptr(), - q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(), - grad_q.data_ptr(), grad_k.data_ptr(), d_mrope_section_doubled, - batch_size, q_heads, kv_heads, seq_len, head_dim, stream); - break; - } -} diff --git a/csrc/rope.h b/csrc/rope.h deleted file mode 100644 index 930c83e..0000000 --- a/csrc/rope.h +++ /dev/null @@ -1,14 +0,0 @@ -#include - -void launch_multimodal_rope_forward( - torch::Tensor q, torch::Tensor k, torch::Tensor cos, torch::Tensor sin, - torch::Tensor q_out, torch::Tensor k_out, - std::vector mrope_section_doubled -); - -void launch_multimodal_rope_backward( - torch::Tensor grad_q_out, torch::Tensor grad_k_out, - torch::Tensor q, torch::Tensor k, torch::Tensor cos, torch::Tensor sin, - torch::Tensor grad_q, torch::Tensor grad_k, - std::vector mrope_section_doubled -); diff --git a/csrc/rope_index.h b/csrc/rope_index.h deleted file mode 100644 index 643d78a..0000000 --- a/csrc/rope_index.h +++ /dev/null @@ -1,13 +0,0 @@ -#include - -std::tuple get_rope_index( - const torch::optional &input_ids, - const torch::optional &image_grid_thw, - const torch::optional &video_grid_thw, - const torch::optional &second_per_grid_ts, - const torch::optional &attention_mask, - int spatial_merge_size, - int image_token_id, - int video_token_id, - int vision_start_token_id, - float tokens_per_second); diff --git a/csrc/rot_pos.cu b/csrc/rot_pos.cu deleted file mode 100644 index 4db08ac..0000000 --- a/csrc/rot_pos.cu +++ /dev/null @@ -1,332 +0,0 @@ -#include -#include -#include -#include - -// CUDA kernel for fused rotary position embedding computation - int32 version -__global__ void fused_rot_pos_emb_kernel_int32( - const float *__restrict__ inv_freq, // [dim/2] - precomputed inverse frequencies - const int32_t *__restrict__ grid_thw, // [num_grids, 3] - (t, h, w) for each grid - float *__restrict__ output, // [total_tokens, dim] - output rotary embeddings - const int32_t *__restrict__ cumsum_tokens, // [num_grids+1] - cumulative sum of tokens per grid - const int dim_half, // dim/2 (size of inv_freq) - const int spatial_merge_size, // spatial merge size - const int num_grids // number of grids -) -{ - const int32_t tid = blockIdx.x * blockDim.x + threadIdx.x; - const int32_t total_tokens = cumsum_tokens[num_grids]; - - if (tid >= total_tokens * dim_half) - return; - - const int32_t token_idx = tid / dim_half; - const int freq_idx = tid % dim_half; - - // Find which grid this token belongs to - int grid_idx = 0; - int32_t local_token_idx = token_idx; - for (int g = 0; g < num_grids; g++) - { - if (token_idx < cumsum_tokens[g + 1]) - { - grid_idx = g; - local_token_idx = token_idx - cumsum_tokens[g]; - break; - } - } - - // Get grid dimensions - const int32_t h = grid_thw[grid_idx * 3 + 1]; - const int32_t w = grid_thw[grid_idx * 3 + 2]; - - // Calculate spatial dimensions after merging - const int32_t h_merged = h / spatial_merge_size; - const int32_t w_merged = w / spatial_merge_size; - const int32_t spatial_tokens = h_merged * w_merged * spatial_merge_size * spatial_merge_size; - - // Get spatial index - const int32_t spatial_idx = local_token_idx % spatial_tokens; - - // Decompose spatial index to get merged block and position within block - const int32_t tokens_per_block = spatial_merge_size * spatial_merge_size; - const int32_t block_idx = spatial_idx / tokens_per_block; - const int32_t within_block_idx = spatial_idx % tokens_per_block; - - // Get block coordinates in merged grid - const int32_t block_h = block_idx / w_merged; - const int32_t block_w = block_idx % w_merged; - - // Get position within block - const int32_t within_h = within_block_idx / spatial_merge_size; - const int32_t within_w = within_block_idx % spatial_merge_size; - - // Calculate actual h and w positions - const int32_t h_pos = block_h * spatial_merge_size + within_h; - const int32_t w_pos = block_w * spatial_merge_size + within_w; - - // Compute rotary embedding - float freq_val = inv_freq[freq_idx]; - - // Output has shape [total_tokens, dim] where dim = 2 * dim_half - int32_t out_idx = token_idx * dim_half * 2 + freq_idx; - output[out_idx] = h_pos * freq_val; // h_pos frequencies - output[out_idx + dim_half] = w_pos * freq_val; // w_pos frequencies -} - -// CUDA kernel for fused rotary position embedding computation - int64 version -__global__ void fused_rot_pos_emb_kernel_int64( - const float *__restrict__ inv_freq, // [dim/2] - precomputed inverse frequencies - const int64_t *__restrict__ grid_thw, // [num_grids, 3] - (t, h, w) for each grid - float *__restrict__ output, // [total_tokens, dim] - output rotary embeddings - const int64_t *__restrict__ cumsum_tokens, // [num_grids+1] - cumulative sum of tokens per grid - const int dim_half, // dim/2 (size of inv_freq) - const int spatial_merge_size, // spatial merge size - const int num_grids // number of grids -) -{ - const int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; - const int64_t total_tokens = cumsum_tokens[num_grids]; - - if (tid >= total_tokens * dim_half) - return; - - const int64_t token_idx = tid / dim_half; - const int freq_idx = tid % dim_half; - - // Find which grid this token belongs to - int grid_idx = 0; - int64_t local_token_idx = token_idx; - for (int g = 0; g < num_grids; g++) - { - if (token_idx < cumsum_tokens[g + 1]) - { - grid_idx = g; - local_token_idx = token_idx - cumsum_tokens[g]; - break; - } - } - - // Get grid dimensions - const int64_t h = grid_thw[grid_idx * 3 + 1]; - const int64_t w = grid_thw[grid_idx * 3 + 2]; - - // Calculate spatial dimensions after merging - const int64_t h_merged = h / spatial_merge_size; - const int64_t w_merged = w / spatial_merge_size; - const int64_t spatial_tokens = h_merged * w_merged * spatial_merge_size * spatial_merge_size; - - // Get spatial index - const int64_t spatial_idx = local_token_idx % spatial_tokens; - - // Decompose spatial index to get merged block and position within block - const int64_t tokens_per_block = spatial_merge_size * spatial_merge_size; - const int64_t block_idx = spatial_idx / tokens_per_block; - const int64_t within_block_idx = spatial_idx % tokens_per_block; - - // Get block coordinates in merged grid - const int64_t block_h = block_idx / w_merged; - const int64_t block_w = block_idx % w_merged; - - // Get position within block - const int64_t within_h = within_block_idx / spatial_merge_size; - const int64_t within_w = within_block_idx % spatial_merge_size; - - // Calculate actual h and w positions - const int64_t h_pos = block_h * spatial_merge_size + within_h; - const int64_t w_pos = block_w * spatial_merge_size + within_w; - - // Compute rotary embedding - float freq_val = inv_freq[freq_idx]; - - // Output has shape [total_tokens, dim] where dim = 2 * dim_half - int64_t out_idx = token_idx * dim_half * 2 + freq_idx; - output[out_idx] = h_pos * freq_val; // h_pos frequencies - output[out_idx + dim_half] = w_pos * freq_val; // w_pos frequencies -} - -// Parallel computation of token counts per grid - int32 version -__global__ void compute_token_counts_kernel_int32( - const int32_t *__restrict__ grid_thw, - int32_t *__restrict__ token_counts, - const int spatial_merge_size, - const int num_grids) -{ - const int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= num_grids) - return; - - int32_t t = grid_thw[idx * 3 + 0]; - int32_t h = grid_thw[idx * 3 + 1]; - int32_t w = grid_thw[idx * 3 + 2]; - int32_t h_merged = h / spatial_merge_size; - int32_t w_merged = w / spatial_merge_size; - token_counts[idx] = t * h_merged * w_merged * spatial_merge_size * spatial_merge_size; -} - -// Parallel computation of token counts per grid - int64 version -__global__ void compute_token_counts_kernel_int64( - const int64_t *__restrict__ grid_thw, - int64_t *__restrict__ token_counts, - const int spatial_merge_size, - const int num_grids) -{ - const int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= num_grids) - return; - - int64_t t = grid_thw[idx * 3 + 0]; - int64_t h = grid_thw[idx * 3 + 1]; - int64_t w = grid_thw[idx * 3 + 2]; - int64_t h_merged = h / spatial_merge_size; - int64_t w_merged = w / spatial_merge_size; - token_counts[idx] = t * h_merged * w_merged * spatial_merge_size * spatial_merge_size; -} - -// Implementation for int32 -torch::Tensor fused_rot_pos_emb_cuda_int32( - torch::Tensor inv_freq, // [dim/2] - torch::Tensor grid_thw, // [num_grids, 3] - int spatial_merge_size) -{ - TORCH_CHECK(inv_freq.dim() == 1, "inv_freq must be 1-dimensional"); - TORCH_CHECK(inv_freq.is_cuda(), "inv_freq must be a CUDA tensor"); - TORCH_CHECK(inv_freq.scalar_type() == torch::kFloat32, "inv_freq must be float32"); - - TORCH_CHECK(grid_thw.dim() == 2, "grid_thw must be 2-dimensional"); - TORCH_CHECK(grid_thw.size(1) == 3, "grid_thw must have shape [num_grids, 3]"); - TORCH_CHECK(grid_thw.is_cuda(), "grid_thw must be a CUDA tensor"); - TORCH_CHECK(grid_thw.scalar_type() == torch::kInt32, "grid_thw must be int32"); - - TORCH_CHECK(spatial_merge_size > 0, "spatial_merge_size must be positive"); - - const int dim_half = inv_freq.size(0); - const int num_grids = grid_thw.size(0); - - auto token_counts = torch::zeros({num_grids}, torch::TensorOptions().dtype(torch::kInt32).device(grid_thw.device())); - const int threads = 256; - const int blocks = (num_grids + threads - 1) / threads; - - compute_token_counts_kernel_int32<<>>( - grid_thw.data_ptr(), - token_counts.data_ptr(), - spatial_merge_size, - num_grids); - - auto cumsum_tokens = torch::cat({torch::zeros({1}, torch::TensorOptions().dtype(torch::kInt32).device(grid_thw.device())), - token_counts.cumsum(0).to(torch::kInt32)}, - 0); - - cudaDeviceSynchronize(); - - int64_t total_tokens = cumsum_tokens[-1].item(); - TORCH_CHECK(total_tokens > 0, "total_tokens must be positive"); - - auto output = torch::zeros({total_tokens, dim_half * 2}, - torch::TensorOptions().dtype(torch::kFloat32).device(inv_freq.device())); - - const int threads_per_block = 256; - const int64_t num_elements = total_tokens * dim_half; - const int num_blocks = static_cast((num_elements + threads_per_block - 1) / threads_per_block); - - fused_rot_pos_emb_kernel_int32<<>>( - inv_freq.data_ptr(), - grid_thw.data_ptr(), - output.data_ptr(), - cumsum_tokens.data_ptr(), - dim_half, - spatial_merge_size, - num_grids); - - cudaDeviceSynchronize(); - - TORCH_CHECK(output.scalar_type() == torch::kFloat32, "Output must be float32"); - TORCH_CHECK(output.size(0) == total_tokens, "Output token count mismatch"); - TORCH_CHECK(output.size(1) == dim_half * 2, "Output dimension mismatch"); - - return output; -} - -// Implementation for int64 -torch::Tensor fused_rot_pos_emb_cuda_int64( - torch::Tensor inv_freq, // [dim/2] - torch::Tensor grid_thw, // [num_grids, 3] - int spatial_merge_size) -{ - TORCH_CHECK(inv_freq.dim() == 1, "inv_freq must be 1-dimensional"); - TORCH_CHECK(inv_freq.is_cuda(), "inv_freq must be a CUDA tensor"); - TORCH_CHECK(inv_freq.scalar_type() == torch::kFloat32, "inv_freq must be float32"); - - TORCH_CHECK(grid_thw.dim() == 2, "grid_thw must be 2-dimensional"); - TORCH_CHECK(grid_thw.size(1) == 3, "grid_thw must have shape [num_grids, 3]"); - TORCH_CHECK(grid_thw.is_cuda(), "grid_thw must be a CUDA tensor"); - TORCH_CHECK(grid_thw.scalar_type() == torch::kInt64, "grid_thw must be int64"); - - TORCH_CHECK(spatial_merge_size > 0, "spatial_merge_size must be positive"); - - const int dim_half = inv_freq.size(0); - const int num_grids = grid_thw.size(0); - - auto token_counts = torch::zeros({num_grids}, torch::TensorOptions().dtype(torch::kInt64).device(grid_thw.device())); - const int threads = 256; - const int blocks = (num_grids + threads - 1) / threads; - - compute_token_counts_kernel_int64<<>>( - grid_thw.data_ptr(), - token_counts.data_ptr(), - spatial_merge_size, - num_grids); - - auto cumsum_tokens = torch::cat({torch::zeros({1}, torch::TensorOptions().dtype(torch::kInt64).device(grid_thw.device())), - token_counts.cumsum(0).to(torch::kInt64)}, - 0); - - cudaDeviceSynchronize(); - - int64_t total_tokens = cumsum_tokens[-1].item(); - TORCH_CHECK(total_tokens > 0, "total_tokens must be positive"); - - auto output = torch::zeros({total_tokens, dim_half * 2}, - torch::TensorOptions().dtype(torch::kFloat32).device(inv_freq.device())); - - const int threads_per_block = 256; - const int64_t num_elements = total_tokens * dim_half; - const int num_blocks = static_cast((num_elements + threads_per_block - 1) / threads_per_block); - - fused_rot_pos_emb_kernel_int64<<>>( - inv_freq.data_ptr(), - grid_thw.data_ptr(), - output.data_ptr(), - cumsum_tokens.data_ptr(), - dim_half, - spatial_merge_size, - num_grids); - - cudaDeviceSynchronize(); - - TORCH_CHECK(output.scalar_type() == torch::kFloat32, "Output must be float32"); - TORCH_CHECK(output.size(0) == total_tokens, "Output token count mismatch"); - TORCH_CHECK(output.size(1) == dim_half * 2, "Output dimension mismatch"); - - return output; -} - -// Main function that dispatches based on grid_thw scalar type -torch::Tensor fused_rot_pos_emb_cuda( - torch::Tensor inv_freq, - torch::Tensor grid_thw, - int spatial_merge_size) -{ - if (grid_thw.scalar_type() == torch::kInt32) - { - return fused_rot_pos_emb_cuda_int32(inv_freq, grid_thw, spatial_merge_size); - } - else if (grid_thw.scalar_type() == torch::kInt64) - { - return fused_rot_pos_emb_cuda_int64(inv_freq, grid_thw, spatial_merge_size); - } - else - { - TORCH_CHECK(false, "Unsupported grid_thw scalar type: ", grid_thw.scalar_type()); - } -} diff --git a/csrc/rot_pos.h b/csrc/rot_pos.h deleted file mode 100644 index 1ed7edd..0000000 --- a/csrc/rot_pos.h +++ /dev/null @@ -1,6 +0,0 @@ -#include - -torch::Tensor fused_rot_pos_emb_cuda( - torch::Tensor inv_freq, - torch::Tensor grid_thw, - int spatial_merge_size); diff --git a/csrc/window_index.h b/csrc/window_index.h deleted file mode 100644 index ea049bf..0000000 --- a/csrc/window_index.h +++ /dev/null @@ -1,9 +0,0 @@ -#include -#include - -std::tuple get_window_index_cuda( - torch::Tensor grid_thw, - int spatial_merge_size, - int vit_merger_window_size, - int patch_size, - int spatial_merge_unit); diff --git a/pyproject.toml b/pyproject.toml index 96581bb..d14afb6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,2 +1,28 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.black] +line-length = 88 +target-version = ["py310"] + [tool.ruff] +line-length = 88 +target-version = "py310" +exclude = [ + ".git", + ".pytest_cache", + ".ruff_cache", + "build", + "dist", + "*.egg-info", + "wall_x/_vendor", + "wall_x/model/core/ops/csrc", + "wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py", + "wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py", +] + +[tool.ruff.lint] +select = ["E", "F", "I"] +ignore = ["E501"] per-file-ignores = { "__init__.py" = ["F401", "E402"] } diff --git a/requirements-libero.txt b/requirements-libero.txt new file mode 100644 index 0000000..fff0046 --- /dev/null +++ b/requirements-libero.txt @@ -0,0 +1,16 @@ +# Optional simulator runtime dependencies for scripts/run_libero.sh. +# Install this file after requirements.txt. Do not install the full LIBERO +# requirements.txt into the Wall-X environment because it pins older torch, +# transformers, and numpy versions for LIBERO's training stack. +robosuite==1.4.1 +bddl==1.0.1 +easydict==1.9 +cloudpickle==2.1.0 +gym==0.25.2 +future==0.18.2 +h5py>=3.11,<4 +imageio[ffmpeg]==2.37.3 +numpy>=1.26,<2 +opencv-python>=4.11,<4.12 +mujoco>=2.3.7,<3.0 +PyOpenGL>=3.1.0 diff --git a/requirements.txt b/requirements.txt index 65b6121..e44b1a8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,26 @@ torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 -transformers==4.49.0 -accelerate==1.10.1 -peft==0.17.1 +transformers==5.2.0 +accelerate==1.13.0 +peft==0.18.1 +numpy>=1.26,<2 +numba>=0.60,<0.62 +datasets>=2.20,<5 +pyarrow>=14,<23 +opencv-python-headless>=4.11,<4.12 +Pillow>=10,<12 +packaging>=24,<26 +ninja>=1.11,<2 scipy==1.15.3 +pyyaml>=6.0 +safetensors>=0.8,<0.9 torchdiffeq==0.2.5 qwen_vl_utils==0.0.11 +websockets>=12.0,<16 +msgpack-numpy>=0.4.8,<0.5 +numpydantic>=1.8,<2 +tqdm>=4.67,<5 +tyro>=0.9.0,<1.0 +wandb>=0.23,<0.25 +matplotlib>=3.10,<4 diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 0000000..7a35df2 --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,142 @@ +# Scripts + +This directory contains the public Wall-X command-line helpers. Run the examples +below from the repository root, using `python scripts/...` and `bash scripts/...`. +Pass file and directory paths explicitly. + +## Inference smoke test + +Use `fake_inference.py` to verify that a checkpoint can be loaded and can +produce one action chunk from a synthetic LIBERO-style observation. + +```bash +python scripts/fake_inference.py --checkpoint-path /path/to/checkpoint +``` + +If the training config is not stored next to the checkpoint as `config.yml` or +`config.yaml`, pass it explicitly: + +```bash +python scripts/fake_inference.py \ + --checkpoint-path /path/to/checkpoint \ + --train-config-path /path/to/config.yml +``` + +## LIBERO evaluation + +`run_libero.sh` is a small shell wrapper around `infer_libero.py`. It requires +the optional LIBERO simulator stack: + +```bash +pip install -r requirements-libero.txt +mkdir -p third_party +git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git third_party/LIBERO +``` + +The launcher checks for LIBERO, robosuite, MuJoCo, PyOpenGL, BDDL, Gym, and +h5py before loading the model. If LIBERO is cloned elsewhere, pass +`LIBERO_PATH=/path/to/LIBERO`. + +```bash +bash scripts/run_libero.sh /path/to/checkpoint +``` + +Useful environment variables: + +```bash +CHECKPOINT_PATH=/path/to/checkpoint +TRAIN_CONFIG_PATH=/path/to/config.yml +TASK_SUITE_NAME=libero_spatial +TASK_INDICES=0,1,2 +NUM_TRIALS_PER_TASK=50 +CUDA_ID=0 +SMOKE=1 +MAX_INFER_TIMES=52 +``` + +`MAX_INFER_TIMES` is optional. When omitted, the launcher uses suite-specific +defaults aligned with the LIBERO evaluator: spatial 22, object 28, goal 30, +libero_10 52, and libero_90 40 action chunks. + +For full control, call the Python entry directly: + +```bash +python scripts/infer_libero.py \ + --checkpoint-path /path/to/checkpoint \ + --task-suite-name libero_spatial \ + --num-trials-per-task 50 \ + --driver-mode in_process +``` + +You can also pass a complete eval config: + +```bash +python scripts/infer_libero.py --config /path/to/eval_config.yml +``` + +## WebSocket serving + +`run_serving.sh` launches the Wall-X WebSocket server through the public +vendored serving runtime. Pass paths explicitly; the script has no built-in +checkpoint path. + +```bash +bash scripts/run_serving.sh \ + --checkpoint-path /path/to/checkpoint \ + --train-config-path /path/to/config.yml \ + --port 32195 +``` + +By default the script returns raw model action chunks, which is the expected +mode for open-loop plotting. Pass `--serialize-actions` when your client expects +robot-serialized actions. + +Useful options: + +```bash +CUDA_ID=0 +ACTION_HORIZON=32 +IMAGE_PASSING_MODE=base64 +MAX_BATCH_SIZE=1 +``` + +Additional `launch_serving.py` arguments can be forwarded after `--`: + +```bash +bash scripts/run_serving.sh --checkpoint-path /path/to/checkpoint -- \ + --model-config.norm-key libero_all +``` + +## Open-loop WebSocket evaluation + +`draw_openloop_plot.py` compares predicted action chunks from a running +WebSocket server against LeRobot dataset ground truth. `--dataset-root` and +`--train-config` are required and have no built-in default. + +```bash +python scripts/draw_openloop_plot.py \ + --uri ws://127.0.0.1:32195 \ + --dataset-root /path/to/lerobot_dataset \ + --train-config /path/to/train_config.yml \ + --episode-indices 0,1,2 \ + --save-dir ./openloop_plots +``` + +## Dataset and checkpoint utilities + +- `compute_norm_stats.py`: compute action normalization statistics for a + local LeRobot v3 dataset. The script reads state/action parquet columns + directly when available, so image and video columns are not decoded. +- `merge_sharded_weights.py`: merge FSDP sharded checkpoint files into a single + checkpoint directory. +- `merge_tokenizer.py`: merge FAST action tokens into a Qwen2.5-VL processor + tokenizer. + +```bash +python scripts/merge_tokenizer.py \ + --processor-path /path/to/Qwen2.5-VL-3B-Instruct \ + --action-tokenizer-path /path/to/fast_tokenizer \ + --output-dir /path/to/merged_processor +``` + +Most scripts support `--help` for their command-line options. diff --git a/scripts/compute_norm_stats.py b/scripts/compute_norm_stats.py index a820d95..a7acf3f 100644 --- a/scripts/compute_norm_stats.py +++ b/scripts/compute_norm_stats.py @@ -1,182 +1,752 @@ #!/usr/bin/env python3 +"""Compute LeRobot normalization stats (mean, std, q01, q99) for training. +Writes JSON in the format expected by wall-x training configs:: + + {"norm_stats": { + "observation.state": {"mean": [...], "std": [...], "q01": [...], "q99": [...]}, + "action": {"mean": [...], "std": [...], "q01": [...], "q99": [...]} + }} + +When ``--train_config`` is provided, the script reads ``data.lerobot_config.repo_id``, +``norm_stats_path``, ``task.dof_config``, ``task.agent_pos_config``, and +``task.action_horizon`` from the YAML. Per-DOF slices are aggregated separately; +keys ending with ``_relative`` use the same relative-pose logic as the LeRobot loader. + +Usage +----- +Recommended: pass a finetune YAML (paths in the config can be placeholders; override +with CLI flags if needed):: + + python scripts/compute_norm_stats.py \\ + --train_config /path/to/train_config.yml + +Multi-task example:: + + python scripts/compute_norm_stats.py \\ + --train_config /path/to/multitask_config.yml + +Override dataset or output path from the command line:: + + python scripts/compute_norm_stats.py \\ + --train_config /path/to/train_config.yml \\ + --data_root /path/to/repo_id \\ + --output_path /path/to/norm_stats_path + +Without a train config (global stats only, no per-DOF relative slices):: + + python scripts/compute_norm_stats.py \\ + --data_root /path/to/lerobot_dataset \\ + --output_path /path/to/norm_stats.json + +Requirements +------------ +- Local LeRobot v3 dataset at ``--data_root`` (or ``data.lerobot_config.repo_id``) +- ``lerobot>=0.3``, ``datasets``, ``pyarrow``, ``numpy``, ``pyyaml``, ``tqdm`` + +After running, set ``norm_stats_path`` in your training YAML to the generated JSON. +""" + +import argparse import json import logging -from collections import defaultdict +from dataclasses import dataclass from pathlib import Path -from typing import Dict, List -from tqdm import tqdm +from typing import Any import numpy as np +import yaml +from numba import jit, prange +from tqdm import tqdm -from lerobot.datasets.lerobot_dataset import LeRobotDataset +SKIP_DOF_KEYS = frozenset( + {"velocity_decomposed", "height", "head_actions", "action_padding"} +) + +_GEOMETRY_FUNCS_LOADED = False -def write_json(path: Path, data: Dict) -> None: +def _ensure_geometry_funcs() -> None: + global _GEOMETRY_FUNCS_LOADED + global canonicalize_euler_zyx_batch_nb + global euler_to_matrix_zyx_batch_nb + global matrix_to_euler_zyx_batch_nb + global so3_to_matrix_batch_nb + + if _GEOMETRY_FUNCS_LOADED: + return + + try: + from wall_x._vendor.x2robot_utils.geometry import ( + canonicalize_euler_zyx_batch_nb, + euler_to_matrix_zyx_batch_nb, + matrix_to_euler_zyx_batch_nb, + so3_to_matrix_batch_nb, + ) + except ImportError as exc: + raise RuntimeError( + "compute_norm_stats.py requires the vendored x2robot_utils geometry " + "helpers. Install Wall-X first so wall_x._vendor is available." + ) from exc + + _GEOMETRY_FUNCS_LOADED = True + + +@dataclass +class TrainConfigContext: + data_root: Path + output_path: Path + state_key: str + action_key: str + propri_ranges: dict[str, list[int]] + action_ranges: dict[str, list[int]] + action_chunk: int + dof_config: dict[str, int] + agent_pos_config: dict[str, int] + + +def write_norm_stats(path: Path, norm_stats: dict[str, dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) path.write_text( - json.dumps(data, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + json.dumps({"norm_stats": norm_stats}, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) -def compute_action_statistics( - action_data_by_robot: Dict[str, Dict[str, List]] -) -> Dict[str, Dict[str, Dict]]: +def compute_vector_stats(values: np.ndarray) -> dict[str, list[float]]: + if values.ndim == 1: + values = values.reshape(-1, 1) + return { + "mean": np.mean(values, axis=0).tolist(), + "std": np.std(values, axis=0).tolist(), + "q01": np.quantile(values, 0.01, axis=0).tolist(), + "q99": np.quantile(values, 0.99, axis=0).tolist(), + } + + +def _apply_slice_stats( + full_stats: dict[str, list[float]], + index_range: list[int], + slice_stats: dict[str, list[float]], +) -> None: + start, end = index_range + for field in ("mean", "std", "q01", "q99"): + full_stats[field][start:end] = slice_stats[field] + + +def config_to_index_ranges(config: dict[str, int]) -> dict[str, list[int]]: + ranges: dict[str, list[int]] = {} + cur = 0 + for key, dim in config.items(): + if key in SKIP_DOF_KEYS: + continue + ranges[key] = [cur, cur + int(dim)] + cur += int(dim) + return ranges + + +def load_train_config(path: Path) -> dict[str, Any]: + with path.open(encoding="utf-8") as f: + config = yaml.load(f, Loader=yaml.SafeLoader) + if not isinstance(config, dict): + raise ValueError(f"train config must be a YAML mapping, got {type(config)}") + return config + + +def parse_train_config(config: dict[str, Any]) -> TrainConfigContext: + task = config.get("task") + if not isinstance(task, dict): + raise ValueError("train config must contain a 'task' section") + + dof_config = task.get("dof_config") + agent_pos_config = task.get("agent_pos_config") + if not isinstance(dof_config, dict) or not dof_config: + raise ValueError("task.dof_config is required in train config") + if not isinstance(agent_pos_config, dict) or not agent_pos_config: + raise ValueError("task.agent_pos_config is required in train config") + + data_cfg = config.get("data") + if not isinstance(data_cfg, dict): + raise ValueError("train config must contain a 'data' section") + + lerobot_config = data_cfg.get("lerobot_config") + if not isinstance(lerobot_config, dict): + raise ValueError("data.lerobot_config is required in train config") + + repo_id = lerobot_config.get("repo_id") + if not repo_id: + raise ValueError("data.lerobot_config.repo_id is required in train config") + + norm_stats_path = config.get("norm_stats_path") or data_cfg.get("norm_stats_path") + if not norm_stats_path: + raise ValueError( + "norm_stats_path is required in train config " "(top-level or under data)" + ) + + key_mappings = data_cfg.get("key_mappings") + if not isinstance(key_mappings, dict): + raise ValueError("data.key_mappings is required in train config") + + state_key = key_mappings.get("state", "observation.state") + action_key = key_mappings.get("action", "action") + + action_chunk = int( + task.get("action_horizon") + or task.get("action_horizon_flow") + or data_cfg.get("action_horizon") + or 32 + ) + + return TrainConfigContext( + data_root=Path(repo_id), + output_path=Path(norm_stats_path), + state_key=state_key, + action_key=action_key, + propri_ranges=config_to_index_ranges(agent_pos_config), + action_ranges=config_to_index_ranges(dof_config), + action_chunk=action_chunk, + dof_config=dof_config, + agent_pos_config=agent_pos_config, + ) + + +def layout_vector_dim(layout_config: dict[str, int]) -> int: + """Config vector width excluding virtual padding keys.""" + return sum( + int(dim) for key, dim in layout_config.items() if key not in SKIP_DOF_KEYS + ) + + +def _prepare_arrays_for_layout( + states: np.ndarray, + actions: np.ndarray, + agent_pos_config: dict[str, int], + dof_config: dict[str, int], +) -> tuple[np.ndarray, np.ndarray]: """ - Compute statistics (min, q01, q99, max) for each action type and dimension. - - Args: - action_data_by_robot: Dict[robot_id][action_type] -> list of arrays/lists - - Returns: - Dict[robot_id][action_type] -> { - "min": [min for each dim], - "q01": [quantile 1% for each dim], - "q99": [quantile 99% for each dim], - "max": [max for each dim], - "delta": [max - min for each dim] - "delta_q99_q01": [q99 - q01 for each dim] - } + Match LeRobot training loader: convert 3D Euler slices to 6D when config + uses rotation_6D keys but the dataset stores 14-dim Euler vectors. """ - stats = {} + from wall_x.data.backends.lerobot.rotation_layout import ( + euler_layout_dim, + layout_uses_6d_rotation, + maybe_convert_euler_to_6d, + ) - for robot_id, action_data in action_data_by_robot.items(): - stats[robot_id] = {} + convert_state = layout_uses_6d_rotation(agent_pos_config) + convert_action = layout_uses_6d_rotation(dof_config) - for action_type, values_list in action_data.items(): - if not values_list: - continue + if convert_state: + raw_dim = euler_layout_dim(agent_pos_config) + target_dim = layout_vector_dim(agent_pos_config) + if states.shape[-1] == raw_dim: + states = maybe_convert_euler_to_6d(states, agent_pos_config, True) + logging.info( + "Converted state Euler->6D (%d -> %d dims)", raw_dim, states.shape[-1] + ) + elif states.shape[-1] != target_dim: + raise ValueError( + f"State dim {states.shape[-1]} does not match Euler raw dim " + f"{raw_dim} or 6D layout dim {target_dim} from agent_pos_config" + ) - # Convert to numpy array: shape (num_samples, num_dims) - try: - values_array = np.array(values_list) - if values_array.size == 0: - continue + if convert_action: + raw_dim = euler_layout_dim(dof_config) + target_dim = layout_vector_dim(dof_config) + if actions.shape[-1] == raw_dim: + actions = maybe_convert_euler_to_6d(actions, dof_config, True) + logging.info( + "Converted action Euler->6D (%d -> %d dims)", raw_dim, actions.shape[-1] + ) + elif actions.shape[-1] != target_dim: + raise ValueError( + f"Action dim {actions.shape[-1]} does not match Euler raw dim " + f"{raw_dim} or 6D layout dim {target_dim} from dof_config" + ) - # Handle both 1D and 2D cases - if values_array.ndim == 1: - values_array = values_array.reshape(-1, 1) - elif values_array.ndim == 2: - pass - else: - logging.warning( - f"Unexpected shape for {robot_id}/{action_type}: {values_array.shape}" - ) - continue + expected_state_dim = layout_vector_dim(agent_pos_config) + expected_action_dim = layout_vector_dim(dof_config) + if states.shape[-1] != expected_state_dim: + raise ValueError( + f"State dim {states.shape[-1]} != expected layout dim {expected_state_dim}" + ) + if actions.shape[-1] != expected_action_dim: + raise ValueError( + f"Action dim {actions.shape[-1]} != expected layout dim {expected_action_dim}" + ) + return states, actions - # Compute statistics for each dimension - min_vals = np.min(values_array, axis=0).tolist() - max_vals = np.max(values_array, axis=0).tolist() - q01_vals = np.quantile(values_array, 0.01, axis=0).tolist() - q99_vals = np.quantile(values_array, 0.99, axis=0).tolist() - delta_vals = (np.array(max_vals) - np.array(min_vals)).tolist() - delta_q99_q01_vals = (np.array(q99_vals) - np.array(q01_vals)).tolist() - stats[robot_id][action_type] = { - "min": min_vals, - "q01": q01_vals, - "q99": q99_vals, - "max": max_vals, - "delta": delta_vals, - "delta_q99_q01": delta_q99_q01_vals, - } +def resolve_lerobot_dataset_paths(dataset_root: Path) -> tuple[str, Path]: + """Return ``(repo_id, root)`` for a local LeRobot dataset directory.""" + root = dataset_root.expanduser().resolve() + if not root.is_dir(): + raise FileNotFoundError(f"LeRobot dataset root not found: {root}") + return root.name, root - except Exception as e: - logging.warning( - f"Error computing statistics for {robot_id}/{action_type}: {e}" + +def _load_parquet_state_action_only( + data_root: Path, + state_key: str, + action_key: str, +): + """ + Read only state/action columns from LeRobot v3 parquet. + + Does not load image/video columns or decode mp4 files. + """ + import pyarrow.dataset as pa_ds + from datasets import Dataset + + root = data_root.expanduser().resolve() + paths = sorted((root / "data").glob("*/*.parquet")) + if not paths: + raise FileNotFoundError(f"No parquet files under {root / 'data'}") + + logging.info( + "Reading parquet columns %r, %r only (no video/images)", + state_key, + action_key, + ) + arrow_dataset = pa_ds.dataset([str(path) for path in paths], format="parquet") + table = arrow_dataset.to_table(columns=[state_key, action_key]) + return Dataset(table) + + +def _load_state_action_table( + data_root: Path, + state_key: str, + action_key: str, +): + root = data_root.expanduser().resolve() + if root.is_dir() and (root / "meta" / "info.json").is_file(): + table = _load_parquet_state_action_only(root, state_key, action_key) + return table, state_key, action_key + + try: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + except ImportError as exc: + raise RuntimeError( + "compute_norm_stats.py requires LeRobot. Install it first, for example " + "`pip install lerobot==0.4.4` or follow the repository README." + ) from exc + + dataset = LeRobotDataset(str(data_root), root=None, video_backend="pyav") + non_image_columns = [ + col for col in dataset.features if "image" not in col and col not in {"task"} + ] + if state_key not in non_image_columns or action_key not in non_image_columns: + raise ValueError( + f"Expected keys {state_key!r} and {action_key!r} in dataset columns, " + f"got {non_image_columns}" + ) + table = dataset.hf_dataset.select_columns([state_key, action_key]) + return table, state_key, action_key + + +def _table_to_arrays( + table, + state_key: str, + action_key: str, +) -> tuple[np.ndarray, np.ndarray]: + """Load full columns once; avoids O(N*chunk) random row access.""" + logging.info("Loading state/action columns into memory...") + try: + states = np.asarray(table[state_key], dtype=np.float32) + actions = np.asarray(table[action_key], dtype=np.float32) + except (KeyError, TypeError, ValueError) as exc: + logging.warning( + "Column-wise load failed (%s); falling back to per-row stack.", exc + ) + states = np.stack( + [ + np.asarray(table[i][state_key], dtype=np.float32) + for i in range(len(table)) + ] + ) + actions = np.stack( + [ + np.asarray(table[i][action_key], dtype=np.float32) + for i in range(len(table)) + ] + ) + if states.ndim == 1: + states = states.reshape(-1, 1) + if actions.ndim == 1: + actions = actions.reshape(-1, 1) + logging.info( + " frames=%d state_dim=%d action_dim=%d", + len(states), + states.shape[1], + actions.shape[1], + ) + return states, actions + + +def _collect_relative_cartesian( + actions: np.ndarray, + states: np.ndarray, + index_range: list[int], + action_chunk: int, +) -> np.ndarray: + start, end = index_range + max_start = max(0, len(actions) - action_chunk) + chunks = [] + anchor_states = states[:max_start, start:end] + for offset in range(action_chunk): + chunks.append(actions[offset : offset + max_start, start:end] - anchor_states) + return np.concatenate(chunks, axis=0) + + +def _compute_delta_from_state_and_abs_rot( + rotations: np.ndarray, state: np.ndarray +) -> np.ndarray: + """Relative rotation: R_rel = R_abs @ R_state^T (same convention as the loader).""" + _ensure_geometry_funcs() + + if rotations.shape[-1] == 3: + rotations_matrix = euler_to_matrix_zyx_batch_nb(rotations) + out_is_euler = True + elif rotations.shape[-1] == 6: + rotations_matrix = so3_to_matrix_batch_nb(rotations) + out_is_euler = False + else: + raise ValueError( + f"Only 3D euler or 6D rotation supported, got {rotations.shape[-1]}D" + ) + + if state.shape[-1] == 3: + state_matrix = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0] + elif state.shape[-1] == 6: + state_matrix = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] + else: + raise ValueError( + f"Only 3D euler or 6D rotation supported, got {state.shape[-1]}D" + ) + + return _abs_rot_to_delta(rotations_matrix, state_matrix, out_is_euler) + + +@jit(nopython=True, parallel=True) +def _abs_rot_to_delta( + rotations_matrix: np.ndarray, + state_matrix: np.ndarray, + out_is_euler: bool, +) -> np.ndarray: + st = np.empty((3, 3), dtype=np.float64) + st[0, 0] = state_matrix[0, 0] + st[0, 1] = state_matrix[1, 0] + st[0, 2] = state_matrix[2, 0] + st[1, 0] = state_matrix[0, 1] + st[1, 1] = state_matrix[1, 1] + st[1, 2] = state_matrix[2, 1] + st[2, 0] = state_matrix[0, 2] + st[2, 1] = state_matrix[1, 2] + st[2, 2] = state_matrix[2, 2] + + n = rotations_matrix.shape[0] + r_rel = np.empty((n, 3, 3), dtype=np.float64) + for i in prange(n): + a00 = rotations_matrix[i, 0, 0] + a01 = rotations_matrix[i, 0, 1] + a02 = rotations_matrix[i, 0, 2] + a10 = rotations_matrix[i, 1, 0] + a11 = rotations_matrix[i, 1, 1] + a12 = rotations_matrix[i, 1, 2] + a20 = rotations_matrix[i, 2, 0] + a21 = rotations_matrix[i, 2, 1] + a22 = rotations_matrix[i, 2, 2] + + r_rel[i, 0, 0] = a00 * st[0, 0] + a01 * st[1, 0] + a02 * st[2, 0] + r_rel[i, 0, 1] = a00 * st[0, 1] + a01 * st[1, 1] + a02 * st[2, 1] + r_rel[i, 0, 2] = a00 * st[0, 2] + a01 * st[1, 2] + a02 * st[2, 2] + r_rel[i, 1, 0] = a10 * st[0, 0] + a11 * st[1, 0] + a12 * st[2, 0] + r_rel[i, 1, 1] = a10 * st[0, 1] + a11 * st[1, 1] + a12 * st[2, 1] + r_rel[i, 1, 2] = a10 * st[0, 2] + a11 * st[1, 2] + a12 * st[2, 2] + r_rel[i, 2, 0] = a20 * st[0, 0] + a21 * st[1, 0] + a22 * st[2, 0] + r_rel[i, 2, 1] = a20 * st[0, 1] + a21 * st[1, 1] + a22 * st[2, 1] + r_rel[i, 2, 2] = a20 * st[0, 2] + a21 * st[1, 2] + a22 * st[2, 2] + + if out_is_euler: + d_euler = matrix_to_euler_zyx_batch_nb(r_rel) + return canonicalize_euler_zyx_batch_nb(d_euler) + + out6 = np.empty((n, 6), dtype=np.float64) + for i in prange(n): + out6[i, 0] = r_rel[i, 0, 0] + out6[i, 1] = r_rel[i, 0, 1] + out6[i, 2] = r_rel[i, 0, 2] + out6[i, 3] = r_rel[i, 1, 0] + out6[i, 4] = r_rel[i, 1, 1] + out6[i, 5] = r_rel[i, 1, 2] + return out6 + + +def _collect_relative_rotation( + actions: np.ndarray, + states: np.ndarray, + index_range: list[int], + action_chunk: int, +) -> np.ndarray: + """ + Per-anchor action chunk relative to anchor state (matches lerobot loader). + + Unlike cartesian relative, each anchor processes a full [chunk, dim] action + clip against a single proprio rotation at the anchor frame. + """ + start, end = index_range + max_start = max(0, len(actions) - action_chunk) + if max_start == 0: + return np.empty((0, end - start), dtype=np.float32) + + chunks = [] + for anchor_idx in tqdm( + range(max_start), + desc=" relative rotation anchors", + leave=False, + ): + action_clip = actions[anchor_idx : anchor_idx + action_chunk, start:end] + proprio_clip = states[anchor_idx, start:end] + rel = _compute_delta_from_state_and_abs_rot( + action_clip.astype(np.float64), proprio_clip.astype(np.float64) + ).astype(np.float32) + chunks.append(rel) + return np.concatenate(chunks, axis=0) + + +def collect_dof_vectors_from_arrays( + states: np.ndarray, + actions: np.ndarray, + propri_ranges: dict[str, list[int]], + action_ranges: dict[str, list[int]], + action_chunk: int = 32, +) -> dict[str, np.ndarray]: + vectors: dict[str, np.ndarray] = {} + + absolute_action_keys = { + key: index_range + for key, index_range in action_ranges.items() + if not key.endswith("_relative") + } + relative_action_keys = { + key: index_range + for key, index_range in action_ranges.items() + if key.endswith("_relative") + } + + for sub_key, index_range in propri_ranges.items(): + start, end = index_range + vectors[sub_key] = states[:, start:end] + + for sub_key, index_range in absolute_action_keys.items(): + start, end = index_range + vectors[sub_key] = actions[:, start:end] + + if relative_action_keys: + logging.info( + "Computing relative action slices (chunk=%d, anchors=%d)...", + action_chunk, + max(0, len(actions) - action_chunk), + ) + for sub_key, index_range in tqdm( + relative_action_keys.items(), desc="Relative action keys" + ): + if "rotation" in sub_key: + vectors[sub_key] = _collect_relative_rotation( + actions, states, index_range, action_chunk + ) + else: + vectors[sub_key] = _collect_relative_cartesian( + actions, states, index_range, action_chunk ) - continue - return stats + return vectors -def load_lerobot_dataset( - repo_id: str, - trajectory_keys: Dict, - base_dir: Path, -) -> None: +def compute_norm_stats_with_dof_config( + data_root: Path, + output_path: Path, + propri_ranges: dict[str, list[int]], + action_ranges: dict[str, list[int]], + state_key: str = "observation.state", + action_key: str = "action", + action_chunk: int = 32, + dof_config: dict[str, int] | None = None, + agent_pos_config: dict[str, int] | None = None, +) -> dict[str, dict]: + table, state_key, action_key = _load_state_action_table( + data_root, state_key, action_key + ) + states, actions = _table_to_arrays(table, state_key, action_key) + if dof_config is not None and agent_pos_config is not None: + states, actions = _prepare_arrays_for_layout( + states, actions, agent_pos_config, dof_config + ) + norm_stats = { + state_key: compute_vector_stats(states), + action_key: compute_vector_stats(actions), + } - # Load local or remote dataset - dataset = LeRobotDataset(base_dir) + vectors = collect_dof_vectors_from_arrays( + states=states, + actions=actions, + propri_ranges=propri_ranges, + action_ranges=action_ranges, + action_chunk=action_chunk, + ) - # Iterate through all data - frames: Dict[str, Dict[str, List]] = defaultdict(lambda: defaultdict(list)) + for sub_key, index_range in propri_ranges.items(): + if sub_key not in vectors: + logging.warning("No samples collected for propri key %s, skipping", sub_key) + continue + slice_stats = compute_vector_stats(vectors[sub_key]) + _apply_slice_stats(norm_stats[state_key], index_range, slice_stats) + logging.info(" %s (agent_pos): dim=%d", sub_key, len(slice_stats["mean"])) - all_features = dataset.features - non_image_columns = [col for col in all_features if "image" not in col] + for sub_key, index_range in action_ranges.items(): + if sub_key not in vectors: + logging.warning("No samples collected for action key %s, skipping", sub_key) + continue + slice_stats = compute_vector_stats(vectors[sub_key]) + _apply_slice_stats(norm_stats[action_key], index_range, slice_stats) + mode = "relative" if sub_key.endswith("_relative") else "absolute" + logging.info(" %s (dof, %s): dim=%d", sub_key, mode, len(slice_stats["mean"])) - print(f"Reading the following fields:{non_image_columns}") - fast_dataset = dataset.hf_dataset.select_columns(non_image_columns) - - for i in tqdm(range(len(fast_dataset))): - sample = fast_dataset[i] - action = sample["action"] # torch.Tensor - propri = sample["observation.state"] - - for key, action_keys in trajectory_keys.items(): - for action_key, action_range in action_keys.items(): - if key == "action": - frames[repo_id][action_key].append( - action[action_range[0] : action_range[1]].numpy().tolist() - ) - else: - frames[repo_id][action_key].append( - propri[action_range[0] : action_range[1]].numpy().tolist() - ) - - return frames + write_norm_stats(output_path, norm_stats) + return norm_stats -def compute_action_normalizer( - repo_id: str, trajectory_keys: Dict, base_dir: Path, output_dir: Path -) -> None: - """ - Compute action normalizer statistics for all robot_ids. - """ - logging.info("Starting action normalizer computation...") +def load_vectors( + data_root: Path, + state_key: str, + action_key: str, +) -> tuple[np.ndarray, np.ndarray]: + table, state_key, action_key = _load_state_action_table( + data_root, state_key, action_key + ) + return _table_to_arrays(table, state_key, action_key) - frames = load_lerobot_dataset(repo_id, trajectory_keys, base_dir) - # Compute statistics - stats = compute_action_statistics(frames) +def compute_norm_stats( + data_root: Path, + output_path: Path, + state_key: str = "observation.state", + action_key: str = "action", + train_ctx: TrainConfigContext | None = None, +) -> dict[str, dict]: + if train_ctx is None: + states, actions = load_vectors(data_root, state_key, action_key) + norm_stats = { + state_key: compute_vector_stats(states), + action_key: compute_vector_stats(actions), + } + write_norm_stats(output_path, norm_stats) + return norm_stats - # Save statistics for each robot_id - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) + return compute_norm_stats_with_dof_config( + data_root=data_root, + output_path=output_path, + propri_ranges=train_ctx.propri_ranges, + action_ranges=train_ctx.action_ranges, + state_key=state_key, + action_key=action_key, + action_chunk=train_ctx.action_chunk, + dof_config=train_ctx.dof_config, + agent_pos_config=train_ctx.agent_pos_config, + ) - # for robot_id, robot_stats in stats.items(): - # output_file = output_dir / f"{robot_id}_action_stats.json" - # write_json(output_file, robot_stats) - # logging.info(f"Saved action statistics for {robot_id} to {output_file}") - # Also save a combined file - combined_output = output_dir / "all_robots_action_stats.json" - write_json(combined_output, stats) - logging.info(f"Saved combined action statistics to {combined_output}") +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Compute norm stats for a local LeRobot v3 dataset.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog="""\ +examples: + %(prog)s --train_config /path/to/train_config.yml + %(prog)s --train_config /path/to/multitask_config.yml \\ + --data_root /path/to/repo_id --output_path /path/to/norm_stats_path + %(prog)s --data_root /path/to/lerobot_dataset --output_path /path/to/norm_stats.json +""", + ) + parser.add_argument( + "--train_config", + type=str, + default=None, + help=( + "Training YAML config (e.g. cvpr_example.yml). When set, reads " + "data.lerobot_config.repo_id, data.norm_stats_path, task.dof_config, " + "task.agent_pos_config and task.action_horizon. Action keys ending " + "with '_relative' use the same relative-pose logic as lerobot loader." + ), + ) + parser.add_argument( + "--data_root", + type=str, + default=None, + help="Local LeRobot dataset directory (overrides train config)", + ) + parser.add_argument( + "--output_path", + type=str, + default=None, + help="Output json path (overrides train config)", + ) + parser.add_argument( + "--state_key", + type=str, + default=None, + help="Dataset column for proprioception (overrides train config)", + ) + parser.add_argument( + "--action_key", + type=str, + default=None, + help="Dataset column for action (overrides train config)", + ) + return parser.parse_args() def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") + args = parse_args() - repo_id = "xxx" # your dataset name - data_root_path = "/path/to/lerobot/dataset" - output_stats_dir = "/path/to/save/action_stats" - trajectory_keys = { # your dataset keys - "propri": { - "follow_right_ee_cartesian_pos": [0, 3], - "follow_right_ee_rotation": [3, 6], - "follow_right_gripper": [6, 7], - }, - "action": { - "master_right_ee_cartesian_pos": [0, 3], - "master_right_ee_rotation": [3, 6], - "master_right_gripper": [6, 7], - }, - } + train_ctx: TrainConfigContext | None = None + if args.train_config: + config_path = Path(args.train_config) + if not config_path.exists(): + raise FileNotFoundError(f"train config not found: {config_path}") + train_ctx = parse_train_config(load_train_config(config_path)) + logging.info("train_config: %s", config_path) + logging.info(" dof_config keys: %s", list(train_ctx.action_ranges)) + logging.info(" agent_pos_config keys: %s", list(train_ctx.propri_ranges)) + logging.info(" action_chunk: %d", train_ctx.action_chunk) - compute_action_normalizer( - repo_id, trajectory_keys, data_root_path, output_stats_dir + data_root = Path(args.data_root or (train_ctx.data_root if train_ctx else "")) + + output_path = Path(args.output_path or (train_ctx.output_path if train_ctx else "")) + state_key = args.state_key or ( + train_ctx.state_key if train_ctx else "observation.state" ) - logging.info("Action normalizer computation completed.") + action_key = args.action_key or (train_ctx.action_key if train_ctx else "action") + + if not data_root.exists(): + raise FileNotFoundError(f"Dataset not found: {data_root}") + + logging.info("dataset: %s", data_root) + logging.info("output: %s", output_path) + + norm_stats = compute_norm_stats( + data_root=data_root, + output_path=output_path, + state_key=state_key, + action_key=action_key, + train_ctx=train_ctx, + ) + + for key, stats in norm_stats.items(): + logging.info(" %s: dim=%d", key, len(stats["mean"])) + + logging.info("Saved norm stats to %s", output_path) if __name__ == "__main__": diff --git a/scripts/draw_openloop_plot.py b/scripts/draw_openloop_plot.py index 07b3d4f..cbff1ef 100644 --- a/scripts/draw_openloop_plot.py +++ b/scripts/draw_openloop_plot.py @@ -1,125 +1,1106 @@ -import os -import yaml -import torch +#!/usr/bin/env python3 +"""Open-loop evaluation over WebSocket using LeRobot-format datasets. + +Loads episodes from a local LeRobot v3 dataset, sends observations to a running +Wall-X websocket server, collects predicted action chunks, and plots them against +ground-truth trajectories. + +Reference: ``infer_openloop_websocket.py`` (websocket client + open-loop loop). + +For single-arm LIBERO checkpoints, start the server with raw model output:: + + python -m wall_x._vendor.harrix.serving.launch_serving \\ + --env X2ROBOT --port 32194 \\ + --no-serialize-actions \\ + model-config:server-model-config \\ + --model-config.checkpoint-path /path/to/ckpt \\ + --model-config.train-config-path /path/to/libero.yml \\ + --model-config.action-horizon 10 \\ + --model-config.robot-type desktop +""" + +from __future__ import annotations + import argparse -from tqdm import tqdm -import matplotlib.pyplot as plt -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction -from wall_x.data.load_lerobot_dataset import load_test_dataset, get_data_configs -from wall_x.model.model_utils import register_normalizers -import copy +import asyncio +import base64 +import json +import logging +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import yaml + +logger = logging.getLogger(__name__) -def load_config(config_path): - """Load configuration from YAML file.""" - with open(config_path, "r") as f: - config = yaml.load(f, Loader=yaml.FullLoader) +def require_cv2(): + try: + import cv2 + except ImportError as exc: + raise RuntimeError( + "Open-loop image encoding requires opencv-python-headless. " + "Install the repository requirements first." + ) from exc + return cv2 - config["data"]["model_type"] = config.get("model_type") - return config +def require_msgpack(): + try: + import msgpack + import msgpack_numpy as m + except ImportError as exc: + raise RuntimeError( + "Open-loop websocket evaluation requires msgpack-numpy. " + "Install the repository requirements first." + ) from exc + m.patch() + return msgpack + + +def require_websockets(): + try: + import websockets + except ImportError as exc: + raise RuntimeError( + "Open-loop websocket evaluation requires websockets. " + "Install the repository requirements first." + ) from exc + return websockets + + +_DEFAULT_CAM_WS_KEYS = { + "observation.images.faceImg": "face_view", + "observation.images.rightImg": "right_wrist_view", + "observation.images.leftImg": "left_wrist_view", + "observation.images.move1Img": "move1_view", +} + + +@dataclass +class EpisodeArrays: + """Episode state/action plus lazily decoded camera frames.""" + + episode_index: int + instruction: str + states: np.ndarray # (T, D_state) + actions: np.ndarray # (T, D_action) + camera_keys: list[str] + frame_offset: int = 0 + num_steps: int = 0 + _dataset: Any = field(default=None, repr=False) + _image_cache: dict[int, dict[str, np.ndarray]] = field( + default_factory=dict, repr=False + ) + + def get_frame_images(self, frame_idx: int) -> dict[str, np.ndarray]: + """Decode camera frames for one observation index (cached).""" + if frame_idx in self._image_cache: + return self._image_cache[frame_idx] + if self._dataset is None: + raise RuntimeError("Episode image loader is not initialized.") + if not self._image_cache: + logger.info( + "Decoding camera frames on demand (first obs_idx=%d)", frame_idx + ) + item = self._dataset[frame_idx] + images = { + cam_key: tensor_to_rgb_uint8(item[cam_key]) + for cam_key in self.camera_keys + if cam_key in item + } + self._image_cache[frame_idx] = images + return images + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def _add_source_root_if_needed() -> None: + repo_root = _repo_root() + if (repo_root / "wall_x").is_dir() and str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + + +def load_train_config(path: str | Path) -> dict[str, Any]: + with open(path, "r", encoding="utf-8") as f: + cfg = yaml.load(f, Loader=yaml.FullLoader) or {} + if "data" not in cfg: + cfg["data"] = {} + cfg["data"]["model_type"] = cfg.get("model_type") + return cfg + + +def parse_int_list(value: str | None) -> list[int] | None: + if value is None or not value.strip(): + return None + return [int(x.strip()) for x in value.split(",") if x.strip()] + + +def _task_block(train_config: dict[str, Any]) -> dict[str, Any]: + return train_config.get("task") or {} + + +def _dof_config(train_config: dict[str, Any]) -> dict[str, int]: + return ( + train_config.get("dof_config") + or _task_block(train_config).get("dof_config") + or {} + ) + + +def _agent_pos_config(train_config: dict[str, Any]) -> dict[str, int]: + return ( + train_config.get("agent_pos_config") + or _task_block(train_config).get("agent_pos_config") + or {} + ) + + +def real_vector_dim(config_block: dict[str, int]) -> int: + """Sum dof/agent dims excluding ``action_padding``.""" + return sum(d for k, d in config_block.items() if k != "action_padding" and d > 0) + + +def slice_by_config(vec: np.ndarray, config_block: dict[str, int]) -> np.ndarray: + """Keep only non-padding dims from a 1-D vector (state / single action row).""" + vec = np.asarray(vec, dtype=np.float32).reshape(-1) + if not config_block: + return vec + parts: list[np.ndarray] = [] + start = 0 + for key, dim in config_block.items(): + if key == "action_padding": + start += dim + continue + end = start + dim + if end <= vec.shape[0]: + parts.append(vec[start:end]) + start = end + if parts: + return np.concatenate(parts, axis=0).astype(np.float32) + return vec + + +def strip_padding_columns(arr: np.ndarray, config_block: dict[str, int]) -> np.ndarray: + """Drop ``action_padding`` columns from a [T, D] action chunk.""" + arr = np.asarray(arr, dtype=np.float32) + if arr.ndim == 1: + return slice_by_config(arr, config_block) + if not config_block: + return arr + parts: list[np.ndarray] = [] + start = 0 + for key, dim in config_block.items(): + if key == "action_padding": + start += dim + continue + parts.append(arr[:, start : start + dim]) + start += dim + if parts: + return np.concatenate(parts, axis=1).astype(np.float32) + return arr + + +def tensor_to_rgb_uint8(img: Any) -> np.ndarray: + arr = np.asarray(img) + if arr.ndim == 3 and arr.shape[0] in (1, 3) and arr.shape[0] != arr.shape[-1]: + arr = np.transpose(arr, (1, 2, 0)) + if np.issubdtype(arr.dtype, np.floating): + arr = np.clip(arr, 0.0, 1.0) * 255.0 + return arr.astype(np.uint8) + + +def encode_image_rgb(image: np.ndarray) -> str: + """JPEG base64; input must be RGB uint8.""" + cv2 = require_cv2() + bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + ok, buffer = cv2.imencode(".jpg", bgr) + if not ok: + raise RuntimeError("cv2.imencode failed") + return base64.b64encode(buffer).decode("utf-8") + + +def build_camera_ws_mapping(train_config: dict[str, Any]) -> dict[str, str]: + key_mappings = (train_config.get("data") or {}).get("key_mappings") or {} + cam_map = key_mappings.get("camera") or {} + if cam_map: + return dict(cam_map) + return dict(_DEFAULT_CAM_WS_KEYS) + + +def resolve_feature_keys( + train_config: dict[str, Any], + state_key: str | None, + action_key: str | None, +) -> tuple[str, str]: + key_mappings = (train_config.get("data") or {}).get("key_mappings") or {} + sk = state_key or key_mappings.get("state") or "observation.state" + ak = action_key or key_mappings.get("action") or "action" + return sk, ak + + +def _arm_follow_pos_from_agent_cfg( + state_vec: np.ndarray, + agent_cfg: dict[str, int], + arm_prefix: str, +) -> list[float] | None: + """Pack one arm into websocket ``follow{1,2}_pos`` layout: pos3 + rpy3 + grip1.""" + from wall_x._vendor.x2robot_utils import geometry as geom + + pos = rot = grip = None + start = 0 + for key, dim in agent_cfg.items(): + if key == "action_padding": + start += int(dim) + continue + end = start + int(dim) + chunk = ( + state_vec[start:end] + if end <= state_vec.shape[0] + else np.zeros(int(dim), dtype=np.float32) + ) + start = end + if not key.startswith(arm_prefix): + continue + if "cartesian_pos" in key: + pos = chunk.reshape(-1) + elif "rotation_6d" in key.lower(): + rot = geom.so3_to_euler_zyx_batch_nb(chunk.reshape(1, -1)).reshape(-1) + elif "rotation" in key: + rot = chunk.reshape(-1) + elif "gripper" in key: + grip = chunk.reshape(-1) + + if pos is None or rot is None or grip is None: + return None + packed = np.concatenate([pos[:3], rot[:3], grip[:1]], axis=0).astype(np.float32) + if packed.shape[0] != 7: + return None + return packed.tolist() + + +def build_state_payload( + state_vec: np.ndarray, + train_config: dict[str, Any], +) -> dict[str, list[float]]: + """Build websocket ``state`` dict from a flat proprio vector.""" + state_vec = np.asarray(state_vec, dtype=np.float32).reshape(-1) + + # LeRobot dual-arm export in this workflow is already raw follow-pos 14D: + # [L_pos3, L_rot3, L_grip1, R_pos3, R_rot3, R_grip1]. + # Prefer direct passthrough and avoid parsing with agent_pos_config + # (which may be 6D/padded and would mis-slice a 14D state). + if state_vec.shape[0] >= 14: + return { + "follow1_pos": state_vec[:7].tolist(), + "follow2_pos": state_vec[7:14].tolist(), + } + if state_vec.shape[0] >= 7: + return {"follow2_pos": state_vec.tolist()} + + agent_cfg = _agent_pos_config(train_config) + + if agent_cfg: + has_left = any(k.startswith("follow_left_") for k in agent_cfg) + has_right = any(k.startswith("follow_right_") for k in agent_cfg) + payload: dict[str, list[float]] = {} + if has_left: + left = _arm_follow_pos_from_agent_cfg(state_vec, agent_cfg, "follow_left_") + if left is not None: + payload["follow1_pos"] = left + if has_right: + right = _arm_follow_pos_from_agent_cfg( + state_vec, agent_cfg, "follow_right_" + ) + if right is not None: + payload["follow2_pos"] = right + if payload: + return payload + + # Legacy single-arm LIBERO path: only follow_right_* in agent_pos_config. + real_dim = real_vector_dim(agent_cfg) + if real_dim > 0 and state_vec.shape[0] >= real_dim: + sliced = slice_by_config(state_vec, agent_cfg) + if sliced.shape[0] == 0: + sliced = state_vec[:real_dim] + right = _arm_follow_pos_from_agent_cfg( + state_vec, agent_cfg, "follow_right_" + ) + if right is not None: + return {"follow2_pos": right} + return {"follow2_pos": sliced.tolist()} + raise ValueError(f"Unsupported state dimension: {state_vec.shape[0]}") + + +def decode_predict_action( + predict_action: Any, + train_config: dict[str, Any], +) -> np.ndarray: + """Convert raw ``predict_action`` [H, D_model] to flat action rows per step.""" + if hasattr(predict_action, "detach"): + predict_action = predict_action.detach().cpu().numpy() + pa = np.asarray(predict_action, dtype=np.float32) + if pa.ndim == 3: + pa = pa[0] + + dof_cfg = _dof_config(train_config) + if dof_cfg: + parts: list[np.ndarray] = [] + start = 0 + for key, dim in dof_cfg.items(): + if key == "action_padding": + start += dim + continue + parts.append(pa[:, start : start + dim]) + start += dim + if parts: + return np.concatenate(parts, axis=1).astype(np.float32) + + try: + from wall_x._vendor.harrix.envs.libero_common import decode_chunk + + chunk = decode_chunk(pa, train_config) + if chunk.ndim == 2 and chunk.shape[1] > 0: + return chunk.astype(np.float32) + except Exception: + pass + + return pa.astype(np.float32) + + +def extract_follow_pos_14d_from_response(result: dict[str, Any]) -> np.ndarray: + """Read serialized ``follow{1,2}_pos`` action rows (skip state row 0) as ``[H, 14]``.""" + if "follow2_pos" not in result or "follow1_pos" not in result: + raise KeyError("Response missing follow1_pos/follow2_pos for 14D eval.") + left = np.asarray(result["follow1_pos"], dtype=np.float32) + right = np.asarray(result["follow2_pos"], dtype=np.float32) + if left.ndim == 1: + left = left.reshape(1, -1) + if right.ndim == 1: + right = right.reshape(1, -1) + if left.shape[0] < 2 or right.shape[0] < 2: + raise ValueError( + f"follow_pos response must include state+action rows, got " + f"left={left.shape}, right={right.shape}" + ) + horizon = min(left.shape[0] - 1, right.shape[0] - 1) + return np.concatenate([left[1 : 1 + horizon], right[1 : 1 + horizon]], axis=1) + + +def extract_action_chunk_from_response( + result: dict[str, Any], + train_config: dict[str, Any], +) -> np.ndarray: + if "predict_action" in result: + return decode_predict_action(result["predict_action"], train_config) + + if "follow2_pos" in result and "follow1_pos" in result: + return extract_follow_pos_14d_from_response(result) + + if "follow2_pos" in result: + right = np.asarray(result["follow2_pos"], dtype=np.float32) + if right.ndim == 1: + right = right.reshape(1, -1) + return right[1:] if right.shape[0] > 1 else right + + if "action" in result: + action = np.asarray(result["action"], dtype=np.float32) + if action.ndim == 1: + action = action[np.newaxis, :] + return action + + raise KeyError( + "Response has no action fields. For single-arm LIBERO, restart the server with " + "`--no-serialize-actions` so responses include `predict_action`." + ) + + +def _hf_row_to_numpy(row: Any, key: str) -> np.ndarray: + value = row[key] + if hasattr(value, "numpy"): + value = value.numpy() + return np.asarray(value, dtype=np.float32).reshape(-1) + + +def _episode_instruction(ds: Any, row_index: int = 0) -> str: + row = ds.hf_dataset[row_index] + task_idx = row["task_index"] + if hasattr(task_idx, "item"): + task_idx = task_idx.item() + return str(ds.meta.tasks.iloc[task_idx].name) + + +def resolve_lerobot_dataset_paths(dataset_root: str | Path) -> tuple[str, Path]: + """Return ``(repo_id, root)`` for a local LeRobot dataset directory.""" + root = Path(dataset_root).expanduser().resolve() + if not root.is_dir(): + raise FileNotFoundError(f"LeRobot dataset root not found: {root}") + return root.name, root + + +def open_local_lerobot_dataset( + dataset_root: str | Path, + *, + episodes: list[int] | None = None, + video_backend: str = "pyav", +) -> Any: + """Open a local LeRobot v3 dataset without HuggingFace Hub fallback.""" + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + repo_id, root = resolve_lerobot_dataset_paths(dataset_root) + try: + return LeRobotDataset( + repo_id, + root=root, + episodes=episodes, + video_backend=video_backend, + ) + except TypeError: + pass + + import torch + from lerobot.datasets.lerobot_dataset import ( + CODEBASE_VERSION, + LeRobotDatasetMetadata, + ) + from lerobot.datasets.utils import ( + get_hf_features_from_features, + hf_transform_to_torch, + load_nested_dataset, + ) + from lerobot.datasets.video_utils import get_safe_default_codec + + meta = LeRobotDatasetMetadata(repo_id, root=root) + + features = get_hf_features_from_features(meta.features) + hf_dataset = load_nested_dataset( + root / "data", features=features, episodes=episodes + ) + hf_dataset.set_transform(hf_transform_to_torch) + + if episodes is not None: + available = { + int(ep.item()) if hasattr(ep, "item") else int(ep) + for ep in hf_dataset.unique("episode_index") + } + missing = set(episodes) - available + if missing: + raise ValueError( + f"Episodes {sorted(missing)} not found under {root}. " + f"Available in loaded parquet: {sorted(available)[:10]}" + f"{'...' if len(available) > 10 else ''}" + ) + + if meta.video_keys: + check_eps = ( + episodes if episodes is not None else list(range(meta.total_episodes)) + ) + for ep_idx in check_eps: + for vid_key in meta.video_keys: + video_path = root / meta.get_video_file_path(ep_idx, vid_key) + if not video_path.exists(): + raise FileNotFoundError( + "LeRobot camera videos are missing for open-loop eval.\n" + f" dataset root: {root}\n" + f" first missing file: {video_path}\n" + "Parquet state/action may exist, but this script also needs " + "mp4 files under videos/." + ) + + ds = LeRobotDataset.__new__(LeRobotDataset) + ds.repo_id = repo_id + ds.root = root + ds.image_transforms = None + ds.delta_timestamps = None + ds.episodes = episodes + ds.tolerance_s = 1e-4 + ds.revision = CODEBASE_VERSION + ds.video_backend = video_backend or get_safe_default_codec() + ds.delta_indices = None + ds.meta = meta + ds.hf_dataset = hf_dataset + ds._lazy_loading = False + ds._absolute_to_relative_idx = None + if episodes is not None: + ds._absolute_to_relative_idx = { + abs_idx.item() if isinstance(abs_idx, torch.Tensor) else abs_idx: rel_idx + for rel_idx, abs_idx in enumerate(hf_dataset["index"]) + } + ds.image_writer = None + ds.episode_buffer = None + ds.writer = None + ds.latest_episode = None + ds._current_file_start_frame = None + ds._streaming_encoder = None + ds.batch_encoding_size = 1 + ds.episodes_since_last_encoding = 0 + return ds + + +def load_lerobot_metadata(dataset_root: str | Path) -> Any: + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + + repo_id, root = resolve_lerobot_dataset_paths(dataset_root) + return LeRobotDatasetMetadata(repo_id, root=root) + + +def plan_eval_frame_range( + num_steps: int, + start_idx: int, + step_stride: int, + action_horizon: int, + max_inferences: int | None, +) -> tuple[int, int, int]: + """Return ``(tabular_start, tabular_end, expected_video_decodes)``.""" + tabular_start = max(0, int(start_idx)) + if max_inferences is None: + return tabular_start, num_steps, -1 + + tabular_end = tabular_start + idx = tabular_start + infer_count = 0 + video_decodes = 0 + while idx <= num_steps - action_horizon - 1: + if infer_count >= max_inferences: + break + record_n = min(step_stride, action_horizon, num_steps - idx) + if record_n <= 0: + break + tabular_end = max(tabular_end, idx + record_n) + video_decodes += 1 + idx += record_n + infer_count += 1 + return tabular_start, min(num_steps, tabular_end), video_decodes + + +def load_episode_arrays( + dataset_root: str | Path, + episode_index: int, + state_key: str, + action_key: str, + camera_keys: list[str], + *, + frame_start: int = 0, + frame_end: int | None = None, + preload_all_images: bool = False, +) -> EpisodeArrays: + ds = open_local_lerobot_dataset( + dataset_root, + episodes=[episode_index], + video_backend="pyav", + ) + num_steps = len(ds.hf_dataset) + if num_steps == 0: + raise ValueError(f"Episode {episode_index} is empty under {dataset_root}") + + start = max(0, int(frame_start)) + end = num_steps if frame_end is None else min(num_steps, int(frame_end)) + states = [_hf_row_to_numpy(ds.hf_dataset[i], state_key) for i in range(start, end)] + actions = [ + _hf_row_to_numpy(ds.hf_dataset[i], action_key) for i in range(start, end) + ] + instruction = _episode_instruction(ds, start) + + episode = EpisodeArrays( + episode_index=episode_index, + instruction=instruction, + states=np.stack(states, axis=0), + actions=np.stack(actions, axis=0), + camera_keys=list(camera_keys), + frame_offset=start, + num_steps=num_steps, + _dataset=ds, + ) + if preload_all_images: + for frame_idx in range(start, end): + episode.get_frame_images(frame_idx) + return episode + + +def build_obs_payload( + episode: EpisodeArrays, + frame_idx: int, + train_config: dict[str, Any], + cam_ws_mapping: dict[str, str], + extra_view_keys: list[str], +) -> dict[str, Any]: + local_idx = frame_idx - episode.frame_offset + if local_idx < 0 or local_idx >= episode.states.shape[0]: + raise IndexError( + f"Frame {frame_idx} is outside loaded tabular range " + f"[{episode.frame_offset}, {episode.frame_offset + episode.states.shape[0]})." + ) + state = episode.states[local_idx] + frame_images = episode.get_frame_images(frame_idx) + views: dict[str, str] = {} + + ref_shape: tuple[int, int, int] | None = None + for cam_key, ws_key in cam_ws_mapping.items(): + if cam_key not in frame_images: + continue + rgb = frame_images[cam_key] + ref_shape = rgb.shape + views[ws_key] = encode_image_rgb(rgb) + + for ws_key in extra_view_keys: + if ws_key in views: + continue + if ref_shape is None: + ref_shape = (256, 256, 3) + views[ws_key] = encode_image_rgb(np.zeros(ref_shape, dtype=np.uint8)) + + return { + "state": build_state_payload(state, train_config), + "views": views, + "instruction": episode.instruction, + } + + +def _dim_label(i: int, dim: int) -> str: + arm7 = [ + "pos_x", + "pos_y", + "pos_z", + "rot_x", + "rot_y", + "rot_z", + "gripper", + ] + if dim == 14: + prefix = "L_" if i < 7 else "R_" + return prefix + arm7[i % 7] + if dim == 7: + return arm7[i] if i < len(arm7) else f"dim_{i}" + if dim == 12: + prefix = "L_" if i < 6 else "R_" + return f"{prefix}rot6_{i % 6}" + if dim == 10: + names = [ + "pos_x", + "pos_y", + "pos_z", + "rot6_0", + "rot6_1", + "rot6_2", + "rot6_3", + "rot6_4", + "rot6_5", + "gripper", + ] + return names[i] if i < len(names) else f"dim_{i}" + return f"dim_{i}" + + +def follow_pos_14d_to_rotation_6d(rows: list | np.ndarray) -> np.ndarray: + """Convert ``follow1_pos + follow2_pos`` rows from euler rotation to rot6D only.""" + from wall_x._vendor.x2robot_utils import geometry as geom + + arr = np.asarray(rows, dtype=np.float32) + if arr.ndim == 1: + arr = arr[np.newaxis, :] + if arr.shape[1] < 14: + raise ValueError(f"Expected follow_pos_14d rows, got shape {arr.shape}") + + left_rot6 = geom.euler_to_matrix_zyx_6d_nb(arr[:, 3:6]).astype(np.float32) + right_rot6 = geom.euler_to_matrix_zyx_6d_nb(arr[:, 10:13]).astype(np.float32) + return np.concatenate([left_rot6, right_rot6], axis=1).astype(np.float32) + + +def plot_openloop( + action_pred_list: list | np.ndarray, + action_gt_list: list | np.ndarray, + save_path: Path, + *, + title: str = "", + action_l1: float | None = None, +) -> None: + """Plot aligned GT vs predicted action trajectories.""" + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + gt_rows = np.asarray(action_gt_list, dtype=np.float32) + pred_rows = np.asarray(action_pred_list, dtype=np.float32) + assert ( + gt_rows.shape == pred_rows.shape + ), f"shape mismatch: pred={pred_rows.shape} gt={gt_rows.shape}" + n, dim = gt_rows.shape + + fig, axes = plt.subplots(dim, 1, figsize=(12, 3.5 * dim), sharex=True) + if dim == 1: + axes = [axes] + + for i, ax in enumerate(axes): + ax.plot(gt_rows[:, i], label="Ground Truth", color="blue", linewidth=1.5) + ax.plot(pred_rows[:, i], label="Model Output", color="orange", linewidth=1.5) + ax.set_title(_dim_label(i, dim)) + ax.set_ylabel("Action Value") + ax.grid(True, alpha=0.25) + ax.legend(loc="best") + + axes[-1].set_xlabel("Time Step") + axes[0].set_xticks(np.arange(0, n, step=max(1, min(10, max(n // 15, 1))))) + l1_note = f"Mean L1: {action_l1:.6f}" if action_l1 is not None else "" + if title and l1_note: + fig.suptitle(f"{title}\n{l1_note}") + elif title: + fig.suptitle(title) + elif l1_note: + fig.suptitle(l1_note) + fig.tight_layout(rect=[0, 0, 1, 0.97]) + + save_path.parent.mkdir(parents=True, exist_ok=True) + out = save_path.with_suffix(".jpg") if save_path.suffix != ".jpg" else save_path + fig.savefig(out, dpi=200) + plt.close(fig) + logger.info("Saved plot -> %s", out) + + +def compute_second_diff(arr: np.ndarray) -> np.ndarray: + d1 = arr[1:] - arr[:-1] + return d1[1:] - d1[:-1] + + +async def run_openloop_eval( + uri: str, + dataset_root: str, + train_config_path: str, + save_dir: str, + episode_indices: list[int] | None, + start_ratio: float, + stride: int | None, + max_inferences: int | None, + state_key: str | None, + action_key: str | None, + extra_view_keys: list[str], + preload_all_images: bool = False, + plot_rotation_6d: bool = False, +) -> None: + train_config = load_train_config(train_config_path) + state_key, action_key = resolve_feature_keys(train_config, state_key, action_key) + cam_ws_mapping = build_camera_ws_mapping(train_config) + camera_keys = list(cam_ws_mapping.keys()) + + action_horizon = ( + train_config.get("action_horizon_flow") + or _task_block(train_config).get("action_horizon_flow") + or _task_block(train_config).get("action_horizon") + or train_config.get("action_horizon") + or 10 + ) + step_stride = stride if stride is not None else int(action_horizon) + + meta = load_lerobot_metadata(dataset_root) + if episode_indices is None: + episode_indices = [0] + episode_indices = [i for i in episode_indices if 0 <= i < meta.total_episodes] + if not episode_indices: + raise ValueError("No valid episode indices to evaluate.") + + save_root = Path(save_dir) + save_root.mkdir(parents=True, exist_ok=True) + msgpack = require_msgpack() + websockets = require_websockets() + + async with websockets.connect( + uri, + ping_interval=None, + ping_timeout=None, + max_size=None, + ) as websocket: + metadata = msgpack.unpackb(await websocket.recv()) + logger.info("Connected to %s, server metadata: %s", uri, metadata) + + for ep_idx in episode_indices: + ep_meta = meta.episodes[ep_idx] + num_steps = int( + ep_meta.get("length") + if isinstance(ep_meta, dict) + else getattr(ep_meta, "length", 0) + ) + if num_steps <= 0: + raise ValueError(f"Episode {ep_idx} has invalid length metadata.") + + start_idx = int(start_ratio * num_steps) + tabular_start, tabular_end, expected_decodes = plan_eval_frame_range( + num_steps, + start_idx, + step_stride, + int(action_horizon), + max_inferences, + ) + logger.info( + "Loading episode %d tabular frames [%d, %d) / %d " + "(images on demand%s)...", + ep_idx, + tabular_start, + tabular_end, + num_steps, + f", ~{expected_decodes} obs frames" if expected_decodes >= 0 else "", + ) + episode = load_episode_arrays( + dataset_root, + ep_idx, + state_key=state_key, + action_key=action_key, + camera_keys=camera_keys, + frame_start=tabular_start, + frame_end=tabular_end, + preload_all_images=preload_all_images, + ) + logger.info( + "Episode %d tabular ready: %d frames loaded, video decode per inference", + ep_idx, + episode.states.shape[0], + ) + + gt_full = np.asarray(episode.actions, dtype=np.float32) + if gt_full.ndim == 1: + gt_full = gt_full[np.newaxis, :] + pred_full = np.full_like(gt_full, np.nan) + aligned_gt: list[list[float]] = [] + aligned_pred: list[list[float]] = [] + obs_infer_points: list[tuple[int, int]] = [] + + idx = start_idx + infer_count = 0 + while idx <= num_steps - action_horizon - 1: + if max_inferences is not None and infer_count >= max_inferences: + break + + payload = build_obs_payload( + episode, + idx, + train_config, + cam_ws_mapping, + extra_view_keys=extra_view_keys, + ) + await websocket.send(msgpack.packb(payload, use_bin_type=True)) + raw = await websocket.recv() + if isinstance(raw, str): + hint = "" + if "get_serialized_actions" in raw: + hint = ( + "\n\nHint: restart the server with `--no-serialize-actions` " + "so responses include `predict_action`." + ) + raise RuntimeError(f"Server error at frame {idx}:\n{raw}{hint}") + + result = msgpack.unpackb(raw, raw=False) + pred_chunk = extract_action_chunk_from_response(result, train_config) + if pred_chunk.ndim == 1: + pred_chunk = pred_chunk[np.newaxis, :] + if pred_chunk.shape[1] != gt_full.shape[1]: + raise ValueError( + f"Prediction action dim {pred_chunk.shape[1]} does not match " + f"dataset action dim {gt_full.shape[1]}." + ) + record_n = min( + step_stride, action_horizon, num_steps - idx, pred_chunk.shape[0] + ) + if record_n <= 0: + break + + local_idx = idx - episode.frame_offset + gt_chunk = gt_full[local_idx : local_idx + record_n] + pred_chunk = pred_chunk[:record_n] + record_n = min(record_n, pred_chunk.shape[0], gt_chunk.shape[0]) + if record_n <= 0: + break + + row0 = len(aligned_gt) + obs_infer_points.append((row0, idx)) + + pred_full[local_idx : local_idx + record_n] = pred_chunk[:record_n] + aligned_gt.extend(gt_chunk[:record_n].tolist()) + aligned_pred.extend(pred_chunk[:record_n].tolist()) + + logger.info( + "episode=%d obs_idx=%d row0=%d record_n=%d pred_shape=%s gt_dim=%d", + ep_idx, + idx, + row0, + record_n, + pred_chunk.shape, + gt_chunk.shape[1], + ) + + idx += record_n + infer_count += 1 + + if not obs_infer_points: + logger.warning("Episode %d: no inference steps executed.", ep_idx) + continue + + ep_tag = f"ep{ep_idx}" + valid_mask = ~np.isnan(pred_full).any(axis=1) + action_l1 = float( + np.mean(np.abs(pred_full[valid_mask] - gt_full[valid_mask])) + ) + summary = { + "episode_index": ep_idx, + "instruction": episode.instruction, + "num_steps": num_steps, + "obs_infer_points": obs_infer_points, + "action_horizon": action_horizon, + "stride": step_stride, + "mean_l1": action_l1, + "state_key": state_key, + "action_key": action_key, + } + summary_path = save_root / f"{ep_tag}_summary.json" + summary_path.write_text( + json.dumps(summary, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + logger.info("Episode %d mean L1 = %.6f", ep_idx, action_l1) + + ml = min(len(aligned_gt), len(aligned_pred)) + if ml < 3: + logger.warning( + "Episode %d: too few aligned rows (%d) to plot.", ep_idx, ml + ) + continue + + gt_ml = np.asarray(aligned_gt[:ml], dtype=np.float32) + pred_ml = np.asarray(aligned_pred[:ml], dtype=np.float32) + title = f"Episode {ep_idx}: {episode.instruction[:80]}" + plot_openloop( + pred_ml, + gt_ml, + save_root / ep_tag, + title=title, + action_l1=action_l1, + ) + if plot_rotation_6d: + gt_rot6d = follow_pos_14d_to_rotation_6d(gt_ml) + pred_rot6d = follow_pos_14d_to_rotation_6d(pred_ml) + rot6d_l1 = float(np.mean(np.abs(pred_rot6d - gt_rot6d))) + plot_openloop( + pred_rot6d, + gt_rot6d, + save_root / f"{ep_tag}_rot6d", + title=f"{title} (rotation 6D)", + action_l1=rot6d_l1, + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Open-loop websocket inference on LeRobot datasets.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--uri", default="ws://127.0.0.1:32194") + parser.add_argument( + "--dataset-root", + required=True, + help="Local LeRobot dataset root.", + ) + parser.add_argument( + "--train-config", + required=True, + help="Training YAML for key mappings / dof layout.", + ) + parser.add_argument( + "--save-dir", + default="./openloop_lerobot_plots", + ) + parser.add_argument( + "--episode-indices", + default="0", + help="Comma-separated episode indices, e.g. 0,1,2", + ) + parser.add_argument( + "--start-ratio", + type=float, + default=0.0, + help="Start open-loop from this fraction of the episode length.", + ) + parser.add_argument( + "--stride", + type=int, + default=None, + help="Frames to advance between inferences (default: action_horizon).", + ) + parser.add_argument( + "--max-inferences", + type=int, + default=None, + help=( + "Cap inference requests per episode; also limits tabular/video loading " + "to the evaluated frame range." + ), + ) + parser.add_argument( + "--preload-all-images", + action="store_true", + help="Decode all camera frames up front (slow; old behavior).", + ) + parser.add_argument( + "--plot-rotation-6d", + action="store_true", + help="Also save a rotation-6D-only plot as ep*_rot6d.jpg.", + ) + parser.add_argument( + "--state-key", + default=None, + help="LeRobot state feature key (default: from train config key_mappings).", + ) + parser.add_argument( + "--action-key", + default=None, + help="LeRobot action feature key (default: from train config key_mappings).", + ) + parser.add_argument( + "--extra-view-keys", + default="left_wrist_view", + help="Websocket view keys to fill with black images when missing from data.", + ) + parser.add_argument("--log-level", default="INFO") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + logging.basicConfig( + level=getattr(logging, args.log_level.upper(), logging.INFO), + format="%(asctime)s - %(levelname)s - %(message)s", + ) + + _add_source_root_if_needed() + + extra_view_keys = [k.strip() for k in args.extra_view_keys.split(",") if k.strip()] + episode_indices = parse_int_list(args.episode_indices) + + try: + asyncio.run( + run_openloop_eval( + uri=args.uri, + dataset_root=args.dataset_root, + train_config_path=args.train_config, + save_dir=args.save_dir, + episode_indices=episode_indices, + start_ratio=args.start_ratio, + stride=args.stride, + max_inferences=args.max_inferences, + state_key=args.state_key, + action_key=args.action_key, + extra_view_keys=extra_view_keys, + preload_all_images=args.preload_all_images, + plot_rotation_6d=args.plot_rotation_6d, + ) + ) + except KeyboardInterrupt: + logger.info("Stopped by user.") + return 130 + return 0 if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--pred_horizon", type=int, default=32) - parser.add_argument("--origin_action_dim", type=int, default=14) - args = parser.parse_args() - - origin_action_dim = args.origin_action_dim - pred_horizon = args.pred_horizon - - # get train config - model_path = "/path/to/your/checkpoint" - action_tokenizer_path = "/path/to/Models/fast" - save_dir = "/path/to/save/dir" - path = f"{model_path}/config.yml" - config = load_config(path) - - normalizer_action, normalizer_propri = register_normalizers(config, model_path) - - # load model with customized robot config - model = Qwen2_5_VLMoEForAction.from_pretrained( - model_path, train_config=config, action_tokenizer_path=action_tokenizer_path - ) - - model.set_normalizer( - copy.deepcopy(normalizer_action), copy.deepcopy(normalizer_propri) - ) - model.eval() - model = model.to("cuda") - model.to_bfloat16_for_selected_params() - - # get test dataloader - dataload_config = get_data_configs(config["data"]) - lerobot_config = dataload_config.get("lerobot_config", {}) - dataset = load_test_dataset( - config, lerobot_config, normalizer_action, normalizer_propri, seed=42 - ) - dataloader = dataset.get_dataloader() - # dataloader = dataset.get_train_dataloader() - - total_frames = len(dataloader) - - predict_mode = "fast" if config.get("use_fast_tokenizer", False) else "diffusion" - action_dim = 14 if predict_mode == "diffusion" else origin_action_dim - gt_traj = torch.zeros((total_frames, origin_action_dim)) - pred_traj = torch.zeros((total_frames, origin_action_dim)) - - # use tqdm to show the progress - for idx, batch in tqdm( - enumerate(dataloader), total=total_frames, desc="predicting" - ): - if idx % pred_horizon == 0 and idx + pred_horizon < total_frames: - batch = batch.to("cuda") - with torch.no_grad(): - outputs = model( - **batch, - action_dim=action_dim, - action_horizon=pred_horizon, - mode="predict", - predict_mode=predict_mode, - ) - pred_traj[idx : idx + pred_horizon] = ( - outputs["predict_action"][:, :, :origin_action_dim] - .detach() - .cpu() - .squeeze(0) - ) - - # Denormalize ground truth actions - gt_action_chunk = batch["action_chunk"][:, :, :origin_action_dim] - dof_mask = batch["dof_mask"].to(gt_action_chunk.dtype) - denormalized_gt = ( - model.action_preprocessor.normalizer_action.unnormalize_data( - gt_action_chunk, - [lerobot_config.get("repo_id", "physical-intelligence/libero")], - dof_mask, - ).squeeze(0) - ) - gt_traj[idx : idx + pred_horizon] = denormalized_gt.detach().cpu() - - gt_traj_np = gt_traj.numpy() - pred_traj_np = pred_traj.numpy() - - timesteps = gt_traj.shape[0] - - fig, axs = plt.subplots( - origin_action_dim, 1, figsize=(15, 5 * origin_action_dim), sharex=True - ) - fig.suptitle("Action Comparison for lerobot", fontsize=16) - - for i in range(origin_action_dim): - axs[i].plot(range(timesteps), gt_traj_np[:, i], label="Ground Truth") - axs[i].plot(range(timesteps), pred_traj_np[:, i], label="Prediction") - axs[i].set_ylabel(f"Action Dim {i+1}") - axs[i].legend() - axs[i].grid(True) - - axs[-1].set_xlabel("Timestep") - plt.tight_layout(rect=[0, 0.03, 1, 0.95]) - os.makedirs(save_dir, exist_ok=True) - save_path = os.path.join(save_dir, "lerobot_comparison.png") - plt.savefig(save_path) - print(f"Saved plot to {save_path}") - plt.close() + raise SystemExit(main()) diff --git a/scripts/fake_inference.py b/scripts/fake_inference.py old mode 100644 new mode 100755 index 428aec7..d14d731 --- a/scripts/fake_inference.py +++ b/scripts/fake_inference.py @@ -1,85 +1,124 @@ -import torch -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction +#!/usr/bin/env python3 +"""Run one Wall-X VLA inference pass through the harrix adapter. -model_path = "/path/to/model" -model = Qwen2_5_VLMoEForAction.from_pretrained(model_path) -model.eval() +This is a lightweight smoke test for the inference path. It builds a synthetic +LIBERO-style observation, loads the checkpoint through harrix, and prints the +predicted action chunk shape. +""" -# Gen Fake data -batch_size = 1 -seq_length = 50 +from __future__ import annotations -torch.manual_seed(0) -fake_input_ids = torch.randint( - 0, len(model.processor.tokenizer), (batch_size, seq_length), dtype=torch.long -) -fake_attention_mask = torch.ones((batch_size, seq_length), dtype=torch.long) -fake_moe_token_types = torch.zeros((batch_size, seq_length), dtype=torch.long) -fake_position_ids = ( - torch.arange(seq_length, dtype=torch.long).unsqueeze(0).expand(batch_size, -1) -) -fake_proprioception = torch.randn((batch_size, 1, 20), dtype=torch.float32) -fake_agent_pos_mask = torch.ones((batch_size, 1, 20), dtype=torch.float32) -fake_dof_mask = torch.ones((batch_size, 32, 20), dtype=torch.float32) -fake_dataset_names = ["x2_normal"] +import argparse +import sys +from pathlib import Path + +import numpy as np -device = "cuda" +def _ensure_local_harrix_on_path() -> None: + repo_root = Path(__file__).resolve().parents[1] + harrix_python = repo_root / "third_party" / "harrix" / "python" + if harrix_python.is_dir(): + sys.path.insert(0, str(harrix_python)) -model = model.to(device) -model = model.bfloat16() -fake_input_ids = fake_input_ids.to(device) -fake_attention_mask = fake_attention_mask.to(device) -fake_moe_token_types = fake_moe_token_types.to(device) -fake_position_ids = fake_position_ids.to(device) -fake_proprioception = fake_proprioception.to(device).bfloat16() -fake_agent_pos_mask = fake_agent_pos_mask.to(device).bfloat16() -fake_dof_mask = fake_dof_mask.to(device).bfloat16() +def _right_gripper_dim_from_config(train_config: dict) -> int: + layout = train_config.get("agent_pos_config") or train_config.get("task", {}).get( + "agent_pos_config", {} + ) + if not isinstance(layout, dict): + return 1 -try: - with torch.no_grad(): - outputs = model( - input_ids=fake_input_ids, - attention_mask=fake_attention_mask, - moe_token_types=fake_moe_token_types, - position_ids=fake_position_ids, - proprioception=fake_proprioception, - agent_pos_mask=fake_agent_pos_mask, - dof_mask=fake_dof_mask, - dataset_names=fake_dataset_names, - mode="validate", - ) + gripper_dim = 1 + for key, dim in layout.items(): + bare = key.replace("follow_", "").replace("master_", "") + if bare == "right_gripper": + gripper_dim = int(dim) + break - print("✅ Fake inference test successful!") - print(f"Output logits shape: {outputs.logits.shape}") - print(f"Output logits dtype: {outputs.logits.dtype}") - print(f"Output logits device: {outputs.logits.device}") + norm_dim = int(train_config.get("_libero_proprio_norm_dim") or 0) + real_dim = sum(int(v) for k, v in layout.items() if k != "action_padding") + if norm_dim == real_dim + 1: + gripper_dim += 1 + return max(1, gripper_dim) - # Check if output is reasonable - if outputs.logits.shape == (batch_size, seq_length, model.config.vocab_size): - print("✅ Output shape correct") - else: - print("❌ Output shape incorrect") - if not torch.isnan(outputs.logits).any(): - print("✅ Output contains no NaN values") - else: - print("❌ Output contains NaN values") +def _build_fake_observation(seed: int, image_size: int, gripper_dim: int) -> dict: + rng = np.random.default_rng(seed) + return { + "eef_pos": rng.normal(size=(3,)).astype(np.float32), + "eef_axisangle": rng.normal(size=(3,)).astype(np.float32), + "gripper": rng.normal(size=(gripper_dim,)).astype(np.float32), + "face_view": rng.integers(0, 256, (image_size, image_size, 3), dtype=np.uint8), + "wrist_view": rng.integers(0, 256, (image_size, image_size, 3), dtype=np.uint8), + } - if not torch.isinf(outputs.logits).any(): - print("✅ Output contains no infinity values") - else: - print("❌ Output contains infinity values") - print("Output logits statistics:") - print(f" Min value: {outputs.logits.min().item():.4f}") - print(f" Max value: {outputs.logits.max().item():.4f}") - print(f" Mean: {outputs.logits.mean().item():.4f}") - print(f" Standard deviation: {outputs.logits.std().item():.4f}") +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--checkpoint-path", required=True, help="Checkpoint directory." + ) + parser.add_argument( + "--train-config-path", + default=None, + help="Optional training config path. Defaults to config.yml/config.yaml next to the checkpoint.", + ) + parser.add_argument("--norm-key", default="libero_all") + parser.add_argument("--architecture", default="qwen2_5") + parser.add_argument("--action-mode", default="flow") + parser.add_argument( + "--cam-names", + nargs="+", + default=["face_view", "right_wrist_view"], + help="Camera names expected by the checkpoint.", + ) + parser.add_argument("--action-horizon", type=int, default=None) + parser.add_argument("--instruction", default="pick up the object") + parser.add_argument("--image-size", type=int, default=128) + parser.add_argument("--seed", type=int, default=0) + return parser.parse_args() -except Exception as e: - print(f"❌ Fake inference test failed: {e}") - import traceback - traceback.print_exc() +def main() -> int: + args = parse_args() + _ensure_local_harrix_on_path() + + import wall_x._vendor.harrix.adapters # noqa: F401 register model adapters + from wall_x._vendor.harrix.adapters.registry import build_adapter + from wall_x._vendor.harrix.eval_config import ( + EvalConfig, + LiberoEnvParams, + autofill_from_checkpoint, + ) + + cfg = EvalConfig() + cfg.model.checkpoint_path = args.checkpoint_path + cfg.model.train_config_path = args.train_config_path + cfg.model.norm_key = args.norm_key + cfg.model.cam_names = list(args.cam_names) + cfg.model.action_horizon = args.action_horizon + cfg.model.architecture = args.architecture + cfg.model.action_mode = args.action_mode + cfg.env.libero = LiberoEnvParams(num_trials_per_task=1, task_indices=[0]) + cfg = autofill_from_checkpoint(cfg) + + adapter = build_adapter(cfg) + gripper_dim = _right_gripper_dim_from_config(getattr(adapter, "_train_config", {})) + payload = { + "observation": _build_fake_observation(args.seed, args.image_size, gripper_dim), + "instruction": args.instruction, + "noise": None, + } + actions = adapter.predict_batch([payload]) + action = np.asarray(actions[0]) + + print("Fake inference succeeded.") + print(f"action shape: {action.shape}") + print(f"action dtype: {action.dtype}") + print(f"action min/max: {float(action.min()):.6f} / {float(action.max()):.6f}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/infer_libero.py b/scripts/infer_libero.py old mode 100644 new mode 100755 index 4b10725..8d22438 --- a/scripts/infer_libero.py +++ b/scripts/infer_libero.py @@ -1,222 +1,235 @@ -import argparse -import time -import os +#!/usr/bin/env python3 +"""Run LIBERO evaluation through harrix. -# from wall_x.utils.baseline_utils import check_baseline_dump, update_baseline -from wall_x.infer.utils_libero import set_seed_everywhere, TaskSuite, TASK_MAX_STEPS -from wall_x.infer.infer_config import InferConfig -from wall_x.infer.env_libero import LiberoRobotEnv +The script accepts either a full harrix EvalConfig YAML or a checkpoint path +plus common command-line overrides. It intentionally bypasses the legacy +Wall-X inference stack. +""" + +from __future__ import annotations + +import argparse +import sys +import tempfile +from pathlib import Path + +import yaml + +LIBERO_DEFAULT_MAX_INFER_TIMES = { + "libero_spatial": 22, + "libero_object": 28, + "libero_goal": 30, + "libero_10": 52, + "libero_90": 40, +} + + +def _ensure_local_harrix_on_path() -> None: + repo_root = Path(__file__).resolve().parents[1] + harrix_python = repo_root / "third_party" / "harrix" / "python" + if harrix_python.is_dir(): + sys.path.insert(0, str(harrix_python)) + + +def _parse_task_indices(value: str | None) -> list[int] | None: + if value is None or value.strip() == "": + return None + return [int(x) for x in value.split(",") if x.strip()] + + +def _resolve_max_infer_times( + task_suite_name: str | None, max_infer_times: int | None +) -> int: + if max_infer_times is not None: + return max_infer_times + suite = task_suite_name or "libero_spatial" + return LIBERO_DEFAULT_MAX_INFER_TIMES.get(suite, 22) + + +def _load_or_build_raw_config(args: argparse.Namespace) -> dict: + if args.config is not None: + with open(args.config, "r") as f: + raw = yaml.safe_load(f) or {} + model = raw.setdefault("model", {}) + env = raw.setdefault("env", {}) + libero = env.setdefault("libero", {}) + runtime = raw.setdefault("runtime", {}) + debug = raw.setdefault("debug", {}) + if args.checkpoint_path is not None: + model["checkpoint_path"] = args.checkpoint_path + if args.train_config_path is not None: + model["train_config_path"] = args.train_config_path + task_indices = _parse_task_indices(args.task_indices) + if task_indices is not None: + libero["task_indices"] = task_indices + if args.max_infer_times is not None or libero.get("max_infer_times") is None: + libero["max_infer_times"] = _resolve_max_infer_times( + libero.get("task_suite_name", args.task_suite_name), + args.max_infer_times, + ) + if args.smoke: + libero["task_indices"] = [0] + libero["num_trials_per_task"] = 5 + runtime["num_workers"] = 1 + runtime["max_batch_size"] = 1 + if args.deterministic_model: + debug["deterministic_model"] = True + return raw + else: + if args.checkpoint_path is None: + raise ValueError("--checkpoint-path is required when --config is not set") + max_infer_times = _resolve_max_infer_times( + args.task_suite_name, args.max_infer_times + ) + raw = { + "model": { + "checkpoint_path": args.checkpoint_path, + "norm_key": args.norm_key, + "cam_names": args.cam_names, + "architecture": args.architecture, + "action_mode": args.action_mode, + }, + "env": { + "type": "libero", + "seed": args.seed, + "libero": { + "task_suite_name": args.task_suite_name, + "initial_states_path": args.initial_states_path, + "num_trials_per_task": args.num_trials_per_task, + "max_infer_times": max_infer_times, + "skip_intermediate_render": args.skip_intermediate_render, + }, + }, + "runtime": { + "num_workers": args.num_workers, + "max_batch_size": args.max_batch_size, + "ws_port": args.ws_port, + "log_dir": args.log_dir, + "driver_mode": args.driver_mode, + }, + "debug": {"deterministic_model": args.deterministic_model}, + } + + model = raw.setdefault("model", {}) + env = raw.setdefault("env", {}) + libero = env.setdefault("libero", {}) + runtime = raw.setdefault("runtime", {}) + debug = raw.setdefault("debug", {}) + + if args.checkpoint_path is not None: + model["checkpoint_path"] = args.checkpoint_path + if args.train_config_path is not None: + model["train_config_path"] = args.train_config_path + if args.norm_key is not None: + model["norm_key"] = args.norm_key + if args.cam_names is not None: + model["cam_names"] = args.cam_names + if args.action_horizon is not None: + model["action_horizon"] = args.action_horizon + if args.architecture is not None: + model["architecture"] = args.architecture + if args.action_mode is not None: + model["action_mode"] = args.action_mode + + env["type"] = "libero" + env["seed"] = args.seed + libero["task_suite_name"] = args.task_suite_name + libero["initial_states_path"] = args.initial_states_path + libero["num_trials_per_task"] = args.num_trials_per_task + libero["max_infer_times"] = _resolve_max_infer_times( + args.task_suite_name, args.max_infer_times + ) + libero["skip_intermediate_render"] = args.skip_intermediate_render + task_indices = _parse_task_indices(args.task_indices) + if task_indices is not None: + libero["task_indices"] = task_indices + if args.smoke: + libero["task_indices"] = [0] + libero["num_trials_per_task"] = 5 + runtime["num_workers"] = 1 + runtime["max_batch_size"] = 1 + + runtime["num_workers"] = ( + args.num_workers if not args.smoke else runtime["num_workers"] + ) + runtime["max_batch_size"] = ( + args.max_batch_size if not args.smoke else runtime["max_batch_size"] + ) + runtime["ws_port"] = args.ws_port + runtime["log_dir"] = args.log_dir + runtime["driver_mode"] = args.driver_mode + debug["deterministic_model"] = args.deterministic_model + return raw + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", default=None, help="Optional harrix EvalConfig YAML." + ) + parser.add_argument("--checkpoint-path", default=None) + parser.add_argument("--train-config-path", default=None) + parser.add_argument("--norm-key", default="libero_all") + parser.add_argument("--architecture", default="qwen2_5") + parser.add_argument("--action-mode", default="flow") + parser.add_argument( + "--cam-names", nargs="+", default=["face_view", "right_wrist_view"] + ) + parser.add_argument("--action-horizon", type=int, default=None) + parser.add_argument("--task-suite-name", default="libero_spatial") + parser.add_argument("--initial-states-path", default="DEFAULT") + parser.add_argument("--num-trials-per-task", type=int, default=50) + parser.add_argument( + "--task-indices", default=None, help="Comma-separated task ids." + ) + parser.add_argument( + "--max-infer-times", + type=int, + default=None, + help=( + "Number of model action chunks per episode. Defaults are suite-specific " + "and match the internal LIBERO evaluator." + ), + ) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument("--max-batch-size", type=int, default=1) + parser.add_argument("--ws-port", type=int, default=8765) + parser.add_argument("--log-dir", default="/tmp/harrix_libero_eval") + parser.add_argument("--driver-mode", choices=["in_process"], default="in_process") + parser.add_argument("--smoke", action="store_true") + parser.add_argument("--deterministic-model", action="store_true") + parser.add_argument( + "--skip-intermediate-render", + action=argparse.BooleanOptionalAction, + default=True, + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + _ensure_local_harrix_on_path() + + from wall_x._vendor.harrix.eval_config import ( + autofill_from_checkpoint, + load_eval_config, + ) + + raw = _load_or_build_raw_config(args) + + with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as f: + yaml.safe_dump(raw, f, sort_keys=False) + tmp_config = f.name + + cfg = autofill_from_checkpoint(load_eval_config(tmp_config)) + if cfg.runtime.driver_mode != "in_process": + raise ValueError("Only driver_mode='in_process' is supported") + from wall_x._vendor.harrix.drivers.inproc import run + + run(cfg) + return 0 if __name__ == "__main__": - args = argparse.ArgumentParser(description="Wall-X Libero evaluation script") - args.add_argument("--seed", type=int, default=42, help="Random seed") - args.add_argument("--id", type=int, default=None, help="Unique index id") - args.add_argument("--name", type=int, default=None, help="Launch command name") - args.add_argument( - "--baseline_path", type=str, default=None, help="Path to baseline record table" - ) - args.add_argument( - "--update_baseline", - type=bool, - default=False, - help="Whether to update baseline table", - ) - args.add_argument( - "--mode", type=str, default="flow", choices=["flow", "ar"], help="Running mode" - ) - args.add_argument( - "--checkpoint_path", type=str, required=True, help="Model checkpoint path" - ) - args.add_argument( - "--train_config_path", - type=str, - required=False, - default=None, - help="Path to training config .yml file", - ) - args.add_argument( - "--norm_key", - type=str, - default="physical-intelligence/libero", - help="Key for normalization statistics", - ) - args.add_argument( - "--cam_names", - nargs="+", - default=["face_view", "right_wrist_view"], - help="List of camera names (e.g., --cam_names face_view right_wrist_view)", - ) - args.add_argument( - "--task_suite_name", - type=str, - default=TaskSuite.LIBERO_SPATIAL, - choices=[e.value for e in TaskSuite], - help="Libero task suite to load", - ) - args.add_argument( - "--initial_states_path", - type=str, - default="DEFAULT", - help="Path to initial states .json file, or 'DEFAULT' to use default states.", - ) - args.add_argument( - "--num_trials_per_task", - type=int, - default=50, - help="Number of evaluation episodes to run per task", - ) - args.add_argument( - "--rollout_dir", - type=str, - default="./rollouts", - help="Directory to save rollout videos", - ) - args = args.parse_args() - - print(f"Using random seed: {args.seed}") - set_seed_everywhere(args.seed) - - print("Initializing InferConfig...") - if args.train_config_path is None: - args.train_config_path = os.path.join(args.checkpoint_path, "config.yml") - - config = InferConfig( - checkpoint_path=args.checkpoint_path, - train_config_path=args.train_config_path, - norm_key=args.norm_key, - cam_names=args.cam_names, - ) - if args.mode == "flow": - config.action_horizon = config.train_config.get("data", {}).get( - "action_horizon_flow", 10 - ) - elif args.mode == "ar": - config.action_horizon = config.train_config.get("data", {}).get( - "action_horizon_ar", 10 - ) - else: - raise ValueError(f"Invalid mode: {args.mode}") - config.model_device = "cuda" - - print("Initializing LiberoRobotEnv (Evaluator)...") - - config.action_dim = 7 - config.pred_horizon = 10 - - evaluator = LiberoRobotEnv( - config=config, - task_suite_name=args.task_suite_name, - initial_states_path=args.initial_states_path, - rollout_dir=args.rollout_dir, - seed=args.seed, - ) - - print(f"\n{'='*20} Starting Evaluation {'='*20}") - print(f"Task suite: {args.task_suite_name}") - print(f"Number of tasks: {evaluator.num_tasks}") - print(f"Trials per task: {args.num_trials_per_task}") - print(f"Initial states: {args.initial_states_path}") - print(f"Videos will be saved to: {evaluator.rollout_dir}") - print(f"{'='*50}\n") - - total_successes = 0 - total_episodes_run = 0 - start_time = time.time() - - for task_id in range(evaluator.num_tasks): - task_successes = 0 - task_episodes_attempted = 0 - - libero_env_instance = None - task_desc = "" - initial_states = None - - max_infer_times = TASK_MAX_STEPS[args.task_suite_name] - print( - f"{args.task_suite_name} TASK_MAX_STEPS: {TASK_MAX_STEPS[args.task_suite_name]}" - ) - for ep_idx in range(args.num_trials_per_task): - print(f" > Running trial {ep_idx + 1} / {args.num_trials_per_task}...") - - try: - print(f"\nCreating environment for Task {task_id}...") - libero_env_instance, task_desc, initial_states = ( - evaluator.create_env_for_task(task_id) - ) - print( - f"--- Starting task {task_id + 1} / {evaluator.num_tasks}: {task_desc} ---" - ) - except Exception as e: - print( - f"\n[CRITICAL ERROR] Failed to create environment for task {task_id}: {e}. Skipping entire task." - ) - continue - - task_episodes_attempted += 1 - total_episodes_run += 1 - - success = False - try: - if args.mode == "flow": - success = evaluator.run_infer_flow_action( - env=libero_env_instance, - task_id=task_id, - task_desc=task_desc, - default_initial_states=initial_states, - episode_idx=ep_idx, - max_infer_times=max_infer_times, - ) - elif args.mode == "ar": - success = evaluator.run_infer_ar_action( - env=libero_env_instance, - task_id=task_id, - task_desc=task_desc, - default_initial_states=initial_states, - episode_idx=ep_idx, - max_infer_times=max_infer_times, - ) - except Exception as e: - print(f" [EXCEPTION] Episode run error: {e}") - - if success: - task_successes += 1 - total_successes += 1 - print(" > Trial result: SUCCESS") - else: - print(" > Trial result: FAILURE") - - if task_episodes_attempted > 0: - print( - f" > Task {task_id} current success rate: {task_successes / task_episodes_attempted * 100:.1f}% ({task_successes}/{task_episodes_attempted})" - ) - if total_episodes_run > 0: - print( - f" > Overall current success rate: {total_successes / total_episodes_run * 100:.1f}% ({total_successes}/{total_episodes_run})" - ) - - task_success_rate = ( - task_successes / task_episodes_attempted - if task_episodes_attempted > 0 - else 0 - ) - print(f"\n--- Task {task_id} ({task_desc}) Summary ---") - print( - f"Success rate: {task_success_rate * 100:.1f}% ({task_successes}/{task_episodes_attempted})" - ) - print(f"{'-'*40}\n") - - end_time = time.time() - total_time = end_time - start_time - final_success_rate = ( - total_successes / total_episodes_run if total_episodes_run > 0 else 0 - ) - - print(f"\n{'='*20} Final Evaluation Summary {'='*20}") - print(f"Total runtime: {total_time:.2f} seconds ({total_time / 60:.1f} minutes)") - print(f"Total trials run: {total_episodes_run}") - print(f"Total successes: {total_successes}") - print(f"Overall success rate: {final_success_rate * 100:.2f}%") - print(f"{'='*56}") - - print("Evaluation completed.") + raise SystemExit(main()) diff --git a/scripts/infer_robochallenge.py b/scripts/infer_robochallenge.py deleted file mode 100644 index f25b3d3..0000000 --- a/scripts/infer_robochallenge.py +++ /dev/null @@ -1,1164 +0,0 @@ -import os -from scipy.fft import dct -from scipy.fft import idct -import yaml -import torch -import numpy as np -import dataclasses -import copy -import json -from PIL import Image -from safetensors.torch import load_file -from qwen_vl_utils.vision_process import smart_resize -from transformers import BatchFeature, AutoProcessor - -from wall_x.model.action_head import Normalizer -from wall_x.utils.constant import action_statistic_dof as default_action_statistic_dof -from numba import jit, prange - -try: - from spatial_tokenizer.spatial_tokenizer_kdisk import SpatialActionTokenizer -except ImportError: - SpatialActionTokenizer = None - -device = "cuda" - -dof_config = { - "follow_left_ee_cartesian_pos": 3, - "follow_left_ee_rotation": 3, - "follow_left_gripper": 1, - "follow_right_ee_cartesian_pos": 3, - "follow_right_ee_rotation": 3, - "follow_right_gripper": 1, - "head_actions": 2, - "height": 1, - "car_pose": 3, - "velocity_decomposed": 3, -} -_CAM_NAME_MAPPING = { - "face_view": "front view", - "left_wrist_view": "left wrist view", - "right_wrist_view": "right wrist view", - "move1_view": "move view", - "move2_view": "move view", - "wall_view": "wall view", - "top_view": "top view", - "side_view": "side view", - "global_view": "global view", -} -camera_to_view_mapping = { - "camera_front": "face_view", - "camera_left": "left_wrist_view", - "camera_right": "right_wrist_view", - "camera_side": "side_view", - "camera_global": "global_view", -} - -action_key_mapping = { - "follow_left_ee_cartesian_pos": "follow1_pos[:3]", - "follow_left_ee_rotation": "follow1_pos[3:6]", - "follow_left_gripper": "follow1_pos[6:7]", - "follow_right_ee_cartesian_pos": "follow2_pos[:3]", - "follow_right_ee_rotation": "follow2_pos[3:6]", - "follow_right_gripper": "follow2_pos[6:7]", - "head_actions": "head_pos", - "height": "lift", - "velocity_decomposed": "velocity_decomposed", - "follow_left_arm_joint_cur": "follow1_joints_cur[-1:]", - "follow_right_arm_joint_cur": "follow2_joints_cur[-1:]", - "follow_left_arm_joint_pos": "follow1_pos", - "follow_right_arm_joint_pos": "follow2_pos", -} - -dim_dof_config = { - "right_xyz": {"rpy": (0, 3), "so3": (0, 3)}, - "right_rot": {"rpy": (3, 6), "so3": (3, 9)}, - "right_gripper": {"rpy": (6, 7), "so3": (9, 10)}, -} -SINGLE_ARM_DIM = 7 - - -@jit(nopython=True, parallel=True) -def euler_to_matrix_zyx_batch_nb(eulers): - N = eulers.shape[0] - R = np.empty((N, 3, 3), dtype=np.float64) - for i in prange(N): - roll = eulers[i, 0] - pitch = eulers[i, 1] - yaw = eulers[i, 2] - - cy, sy = np.cos(yaw), np.sin(yaw) - cp, sp = np.cos(pitch), np.sin(pitch) - cr, sr = np.cos(roll), np.sin(roll) - - R[i, 0, 0] = cy * cp - R[i, 0, 1] = cy * sp * sr - sy * cr - R[i, 0, 2] = cy * sp * cr + sy * sr - - R[i, 1, 0] = sy * cp - R[i, 1, 1] = sy * sp * sr + cy * cr - R[i, 1, 2] = sy * sp * cr - cy * sr - - R[i, 2, 0] = -sp - R[i, 2, 1] = cp * sr - R[i, 2, 2] = cp * cr - return R - - -@jit(nopython=True, parallel=True) -def compose_state_and_delta_to_abs_rpy(delta, state): - """ - Input: - delta: (N,3) -> Δrpy(ZYX) or (N,6) -> Δ6D (first two rows flattened) - state: (3,) -> rpy(ZYX) or (6,) -> 6D (first two rows flattened) - Output: - abs_rpy: (N,3) Absolute pose rpy(ZYX, radians), normalized to (-π, π] - """ - if delta.shape[-1] == 3: - R_delta = euler_to_matrix_zyx_batch_nb(delta) # (N,3,3) - elif delta.shape[-1] == 6: - R_delta = so3_to_matrix_batch_nb(delta) # (N,3,3) - else: - raise ValueError(f"delta last dim must be 3 or 6, got {delta.shape[-1]}") - - if state.shape[-1] == 3: - R_state = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0] # (3,3) - elif state.shape[-1] == 6: - R_state = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] # (3,3) - else: - raise ValueError(f"state last dim must be 3 or 6, got {state.shape[-1]}") - - N = R_delta.shape[0] - R_abs = np.empty((N, 3, 3), dtype=np.float64) - - S00 = R_state[0, 0] - S01 = R_state[0, 1] - S02 = R_state[0, 2] - S10 = R_state[1, 0] - S11 = R_state[1, 1] - S12 = R_state[1, 2] - S20 = R_state[2, 0] - S21 = R_state[2, 1] - S22 = R_state[2, 2] - - for i in prange(N): - A00 = R_delta[i, 0, 0] - A01 = R_delta[i, 0, 1] - A02 = R_delta[i, 0, 2] - A10 = R_delta[i, 1, 0] - A11 = R_delta[i, 1, 1] - A12 = R_delta[i, 1, 2] - A20 = R_delta[i, 2, 0] - A21 = R_delta[i, 2, 1] - A22 = R_delta[i, 2, 2] - - R_abs[i, 0, 0] = A00 * S00 + A01 * S10 + A02 * S20 - R_abs[i, 0, 1] = A00 * S01 + A01 * S11 + A02 * S21 - R_abs[i, 0, 2] = A00 * S02 + A01 * S12 + A02 * S22 - - R_abs[i, 1, 0] = A10 * S00 + A11 * S10 + A12 * S20 - R_abs[i, 1, 1] = A10 * S01 + A11 * S11 + A12 * S21 - R_abs[i, 1, 2] = A10 * S02 + A11 * S12 + A12 * S22 - - R_abs[i, 2, 0] = A20 * S00 + A21 * S10 + A22 * S20 - R_abs[i, 2, 1] = A20 * S01 + A21 * S11 + A22 * S21 - R_abs[i, 2, 2] = A20 * S02 + A21 * S12 + A22 * S22 - - -@jit(nopython=True, parallel=True) -def so3_to_matrix_batch_nb(batch_so3): - N = batch_so3.shape[0] - R_all = np.empty((N, 3, 3), dtype=np.float64) - eps = 1e-12 - for i in prange(N): - r1x, r1y, r1z = batch_so3[i, 0], batch_so3[i, 1], batch_so3[i, 2] - r2x, r2y, r2z = batch_so3[i, 3], batch_so3[i, 4], batch_so3[i, 5] - - # normalize r1 - n1 = np.sqrt(r1x * r1x + r1y * r1y + r1z * r1z) + eps - r1x /= n1 - r1y /= n1 - r1z /= n1 - - # orthogonalize r2 to r1, then normalize - dot12 = r1x * r2x + r1y * r2y + r1z * r2z - r2x -= dot12 * r1x - r2y -= dot12 * r1y - r2z -= dot12 * r1z - n2 = np.sqrt(r2x * r2x + r2y * r2y + r2z * r2z) + eps - r2x /= n2 - r2y /= n2 - r2z /= n2 - - # r3 = r1 x r2 - r3x = r1y * r2z - r1z * r2y - r3y = r1z * r2x - r1x * r2z - r3z = r1x * r2y - r1y * r2x - - R_all[i, 0, 0] = r1x - R_all[i, 0, 1] = r1y - R_all[i, 0, 2] = r1z - R_all[i, 1, 0] = r2x - R_all[i, 1, 1] = r2y - R_all[i, 1, 2] = r2z - R_all[i, 2, 0] = r3x - R_all[i, 2, 1] = r3y - R_all[i, 2, 2] = r3z - return R_all - - -@jit(nopython=True, parallel=True) -def matrix_to_euler_zyx_batch_nb(Rs): - """ - R = Rz(yaw) * Ry(pitch) * Rx(roll) - extract: - pitch = asin(-R[2,0]) - roll = atan2(R[2,1], R[2,2]) - yaw = atan2(R[1,0], R[0,0]) - """ - N = Rs.shape[0] - eulers = np.empty((N, 3), dtype=np.float64) - for i in prange(N): - r00 = Rs[i, 0, 0] - # r01 = Rs[i, 0, 1] - # r02 = Rs[i, 0, 2] - r10 = Rs[i, 1, 0] - # r11 = Rs[i, 1, 1] - # r12 = Rs[i, 1, 2] - r20 = Rs[i, 2, 0] - r21 = Rs[i, 2, 1] - r22 = Rs[i, 2, 2] - - x = -r20 - if x > 1.0: - x = 1.0 - elif x < -1.0: - x = -1.0 - - pitch = np.arcsin(x) - roll = np.arctan2(r21, r22) - yaw = np.arctan2(r10, r00) - - eulers[i, 0] = roll - eulers[i, 1] = pitch - eulers[i, 2] = yaw - return eulers - - -@jit(nopython=True, parallel=True) -def canonicalize_euler_zyx_batch_nb(rpy_batch): - """ - Batch ZYX Euler Angle Normalization (Parallel Version) - Input: rpy_batch (N, 3) [roll, pitch, yaw] (radians) - Output: out (N, 3) Constrained to the same branch with each component in (-π, π] - Rules: - 1) First, wrap each component to (-π, π] - 2) If p > π/2: p = π - p; r += π; y += π - If p <= -π/2: p = -π - p; r += π; y += π - 3) Finally, wrap each component to (-π, π] again. - """ - N = rpy_batch.shape[0] - out = np.empty_like(rpy_batch) - two_pi = 2.0 * np.pi - - for i in prange(N): - r = rpy_batch[i, 0] - p = rpy_batch[i, 1] - y = rpy_batch[i, 2] - - r = (r + np.pi) % two_pi - np.pi - p = (p + np.pi) % two_pi - np.pi - y = (y + np.pi) % two_pi - np.pi - - if p > np.pi / 2.0: - p = np.pi - p - r = r + np.pi - y = y + np.pi - elif p <= -np.pi / 2.0: - p = -np.pi - p - r = r + np.pi - y = y + np.pi - - r = (r + np.pi) % two_pi - np.pi - p = (p + np.pi) % two_pi - np.pi - y = (y + np.pi) % two_pi - np.pi - - out[i, 0] = r - out[i, 1] = p - out[i, 2] = y - - return out - - -def so3_to_euler_zyx_batch_nb(batch_so3): - matrix = so3_to_matrix_batch_nb(batch_so3) - eulers = matrix_to_euler_zyx_batch_nb(matrix) - return canonicalize_euler_zyx_batch_nb(eulers) - - -def update_model_config(train_config, model_config): - model_config.use_state_string_representation = train_config["data"].get( - "use_state_string_representation", False - ) - model_config.flow_loss_weight = train_config.get("flow_loss_weight", 1.0) - - model_config.dof_config = train_config["dof_config"] - model_config.agent_pos_config = train_config["agent_pos_config"] - - model_config.action_horizon_flow = train_config["data"].get( - "action_horizon_flow", 32 - ) - - if train_config.get("_attn_implementation", None) is not None: - model_config._attn_implementation = train_config["_attn_implementation"] - - return model_config - - -def move_to_cuda(obj, device=device): - if isinstance(obj, torch.Tensor): - return obj.to(device) - elif isinstance(obj, (dict, BatchFeature)): - return {k: move_to_cuda(v, device) for k, v in obj.items()} - elif isinstance(obj, list): - return [move_to_cuda(v, device) for v in obj] - elif isinstance(obj, tuple): - return tuple(move_to_cuda(v, device) for v in obj) - else: - return obj - - -def extract_components(data, config, is_rpy): - mode = "rpy" if is_rpy else "so3" - result = {} - for key, slice_config in config.items(): - slice_range = slice_config[mode] - result[key] = data[:, slice_range[0] : slice_range[1]] - return result - - -@dataclasses.dataclass -class WallxInferArgs: - config_path: str | None = None - checkpoint_path: str | None = None - action_mode: str = "diffusion" # "ar" or "diffusion" - - max_time_step: int = 1000 - action_start_ratio: float = 0 - action_end_ratio: float = 0.6 - - model_action_dim: int = 14 - action_horizon: int = 32 - - action_dim: int = 14 - - interpolate_action: bool = False - interpolate_multiplier: int = 1 - turtle_as_desktop: bool = False - generate_subtask: bool = False - subtask_interval: int = 0 - with_cur: bool = False - state_str: bool = True ### NOTE - wostate: bool = False - delta_action: bool = False - state_rpy: bool = True - action_rpy: bool = True - - dataset_name: str = "robochallenge_aloha" - use_hard_prompt: bool = True - dct_scale: float = -1 - - -class WallxModelWrapper: - def __init__(self, args: WallxInferArgs): - self.args = args - self.get_model_and_processor() - self.action_predict_mode = "ar" if args.action_mode == "ar" else "diffusion" - print("action_predict_mode", self.action_predict_mode, flush=True) - - def get_model_and_processor(self): - if self.args.config_path is None: - self.args.config_path = os.path.join( - self.args.checkpoint_path, "config.yml" - ) - with open(self.args.config_path, "r") as f: - config = yaml.load(f, Loader=yaml.FullLoader) - self.config = config - self.dof_config = config["dof_config"] - self.agent_pos_config = config["agent_pos_config"] - self.obs_action_keys = [ - "follow_left_ee_cartesian_pos", - "follow_left_ee_rotation", - "follow_left_gripper", - "follow_right_ee_cartesian_pos", - "follow_right_ee_rotation", - "follow_right_gripper", - ] - print("obs_action_keys", self.obs_action_keys, flush=True) - self.action_tokenizer_type = config.get("action_tokenizer_type", None) - config_path = config["qwen_vl_act_config_path"] - self.processor = AutoProcessor.from_pretrained( - config["processor_path"], use_fast=True - ) - self.processor.tokenizer.padding_side = "left" - new_tokens = ["<|propri|>", "<|action|>"] - - # load fast tokenizer - self.action_tokenizer_type = config.get("action_tokenizer_type", None) - print("self.action_tokenizer_type", self.action_tokenizer_type, flush=True) - self.action_tokenizer = None - self.action_mapper = None - if self.action_tokenizer_type: - # fast - if self.action_tokenizer_type == "fast": - print("Using fast tokenizer") - self.action_tokenizer = AutoProcessor.from_pretrained( - config["action_tokenizer_path"], trust_remote_code=True - ) - elif self.action_tokenizer_type == "spatialvla": - print("Using spatialvla tokenizer") - assert ( - SpatialActionTokenizer is not None - ), "SpatialActionTokenizer is not installed" - self.action_tokenizer = SpatialActionTokenizer() - else: - raise ValueError( - f"Unsupported action tokenizer type: {self.action_tokenizer_type}" - ) - new_tokens += [ - f"<|action_token_{i}|>" for i in range(self.action_tokenizer.vocab_size) - ] - - # num_added_tokens = self.processor.tokenizer.add_tokens(new_tokens) - - # define action mapper - if self.action_tokenizer_type: - self.action_mapper = {} - for i in range(self.action_tokenizer.vocab_size): - token = f"<|action_token_{i}|>" - token_id = self.processor.tokenizer.convert_tokens_to_ids(token) - self.action_mapper[token_id] = i - - # action & propri normalizer - self._register_normalizers() - - model_type = config["model_type"] - - if model_type == "qwen2_5": - print("Using qwen2_5 model as base model") - from wall_x.model.qwen2_5_based import ( - Qwen2_5_VLMoEForAction, - Qwen2_5_VLConfig, - ) - - ModelClass = Qwen2_5_VLMoEForAction - ConfigClass = Qwen2_5_VLConfig - - model_config = ConfigClass.from_pretrained(config_path) - model_config = update_model_config(config, model_config) - - print("model_config", model_config, flush=True) - - # if self.args.action_mode == "ar": - # model_config._attn_implementation = "flash_attention_2" - # else - model_config._attn_implementation = "sdpa" - model_config.vision_config._attn_implementation = "flash_attention_2" - - if model_config.model_type == "qwen3_vl": - self.MAX_PIXELS = 16384 * 32 * 32 - self.MIN_PIXELS = 4 * 32 * 32 - self.IMAGE_FACTOR = 32 - elif model_config.model_type == "qwen2_5_vl": - self.MAX_PIXELS = 16384 * 28 * 28 - self.MIN_PIXELS = 4 * 28 * 28 - self.IMAGE_FACTOR = 28 - - model = ModelClass( - model_config, - self.action_tokenizer_type, - self.processor, - self.action_tokenizer, - self.action_mapper, - ) - model.resize_token_embeddings(len(self.processor.tokenizer)) - state_dict = load_file( - self.args.checkpoint_path + "/model.safetensors", device="cpu" - ) - if os.path.exists(os.path.join(self.args.checkpoint_path, "global_step.pth")): - global_step = torch.load( - os.path.join(self.args.checkpoint_path, "global_step.pth") - )["global_step"] - print("Loaded global step:", global_step) - msg = model.load_state_dict(state_dict, strict=False) - print(msg) - model.eval() - model.set_normalizer( - copy.deepcopy(self.normalizer_action), copy.deepcopy(self.normalizer_propri) - ) - model.to(device) - model.to_bfloat16_for_selected_params() - - self.model = model - print("self.args.dataset_name", self.args.dataset_name, flush=True) - print( - "normalizer_action min", - self.normalizer_action.min.__getattr__(self.args.dataset_name), - flush=True, - ) - print( - "normalizer_action delta", - self.normalizer_action.delta.__getattr__(self.args.dataset_name), - flush=True, - ) - - def _register_normalizers(self): - if self.config.get("customized_action_statistic_dof", None): - action_statistic_dof = json.load( - open(self.config["customized_action_statistic_dof"], "r") - ) - else: - action_statistic_dof = default_action_statistic_dof - - if os.path.exists(self.args.checkpoint_path + "/normalizer_action.pth"): - print( - "Loading normalizer_action from checkpoint", - self.args.checkpoint_path + "/normalizer_action.pth", - flush=True, - ) - self.normalizer_action = Normalizer.from_ckpt( - self.args.checkpoint_path + "/normalizer_action.pth" - ) - else: - self.normalizer_action = Normalizer( - action_statistic_dof, - self.config["dof_config"], - min_key=self.config.get("min_key", "min"), - delta_key=self.config.get("delta_key", "delta"), - ) - - # print("action_statistic_dof",action_statistic_dof) - - if os.path.exists(self.args.checkpoint_path + "/normalizer_propri.pth"): - print( - "Loading normalizer_propri from checkpoint", - self.args.checkpoint_path + "/normalizer_propri.pth", - flush=True, - ) - self.normalizer_propri = Normalizer.from_ckpt( - self.args.checkpoint_path + "/normalizer_propri.pth" - ) - else: - self.normalizer_propri = Normalizer( - action_statistic_dof, - self.config["agent_pos_config"], - min_key=self.config.get("min_key", "min"), - delta_key=self.config.get("delta_key", "delta"), - ) - - print("self.args.dataset_name", self.args.dataset_name, flush=True) - print( - "normalizer_propri min", - self.normalizer_propri.min.__getattr__(self.args.dataset_name), - flush=True, - ) - print( - "normalizer_propri delta", - self.normalizer_propri.delta.__getattr__(self.args.dataset_name), - flush=True, - ) - print( - "normalizer_action min", - self.normalizer_action.min.__getattr__(self.args.dataset_name), - flush=True, - ) - print( - "normalizer_action delta", - self.normalizer_action.delta.__getattr__(self.args.dataset_name), - flush=True, - ) - - def get_text_ar(self, instruction, camera_names, norm_state=None, state_mask=None): - role_start_symbol = "<|im_start|>" - role_end_symbol = "<|im_end|>" - vision_start_symbol = "<|vision_start|>" - vision_end_symbol = "<|vision_end|>" - image_pad_symbol = "<|image_pad|>" - propri_symbol = "<|propri|>" - - prologue = f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n" - user_request = f"{role_start_symbol}user\nObservation:" - print("camera_names", camera_names, flush=True) - for cam_name in camera_names: - user_request += f" {_CAM_NAME_MAPPING[cam_name]}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" - user_request += "\nInstruction:" - if self.args.state_str: - assert norm_state is not None - if isinstance(norm_state, torch.Tensor): - if state_mask is not None: - if isinstance(state_mask, torch.Tensor): - mask_1d = state_mask[0, 0].to( - dtype=torch.bool, device=norm_state.device - ) - else: - mask_1d = torch.as_tensor(state_mask, device=norm_state.device)[ - 0, 0 - ].to(dtype=torch.bool) - - norm_state = norm_state[..., mask_1d] - - norm_state = norm_state.detach().cpu().numpy() - print("norm_state", norm_state, flush=True) - discretized_state = ( - np.digitize(norm_state, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1 - ) - propri = " ".join(map(str, discretized_state[0, 0])) - elif self.args.wostate: - propri = "" - else: - propri = propri_symbol - text_prompt = ( - f"\nPredict the next action in robot action.\nProprioception: {propri}\n" - ) - user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" - assistant_message = f"{role_start_symbol}assistant\n" - text = prologue + user_message + assistant_message - - return text - - def get_text_flow( - self, - instruction, - camera_names, - action_chunk_size, - norm_state=None, - state_mask=None, - ): - role_start_symbol = "<|im_start|>" - role_end_symbol = "<|im_end|>" - vision_start_symbol = "<|vision_start|>" - vision_end_symbol = "<|vision_end|>" - image_pad_symbol = "<|image_pad|>" - propri_symbol = "<|propri|>" - action_symbol = "<|action|>" - action_space = "Rel EEF" if self.args.delta_action else "Abs EEF" - _camera = ", ".join([_CAM_NAME_MAPPING[cam_name] for cam_name in camera_names]) - prologue = f"<|im_start|>system\nYou are an embodied vision-language-action (VLA) model controlling the robot with language instructions.\n Embodiment: {self.args.dataset_name.split('_')[-1]}\n Camera Setup: {_camera},\n Frequency: 32HZ\n Action Space: {action_space}\n<|im_end|>\n" - - user_request = f"{role_start_symbol}user\nObservation:" - print("camera_names", camera_names, flush=True) - for cam_name in camera_names: - user_request += f" {_CAM_NAME_MAPPING[cam_name]}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" - user_request += "\nInstruction:" - if self.args.state_str: - assert norm_state is not None - if isinstance(norm_state, torch.Tensor): - if state_mask is not None: - if isinstance(state_mask, torch.Tensor): - mask_1d = state_mask[0, 0].to( - dtype=torch.bool, device=norm_state.device - ) - else: - mask_1d = torch.as_tensor(state_mask, device=norm_state.device)[ - 0, 0 - ].to(dtype=torch.bool) - norm_state = norm_state[..., mask_1d] - norm_state = norm_state.detach().cpu().numpy() - discretized_state = ( - np.digitize(norm_state, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1 - ) - propri = " ".join(map(str, discretized_state[0, 0])) - elif self.args.wostate: - propri = "" - else: - propri = propri_symbol - text_prompt = ( - f"\nPredict the next action in robot action.\nProprioception: {propri}\n" - ) - user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" - assistant_message = f"{role_start_symbol}assistant\n" - action = f"{action_symbol * action_chunk_size}" - text = prologue + user_message + assistant_message + action - - return text - - def resize_images(self, observation): - image_inputs = [] - view_candidates = [ - "face_view", - "left_wrist_view", - "right_wrist_view", - "side_view", - "global_view", - ] - for key in observation.keys(): - if key not in view_candidates: - print("!!! key not in view_candidates", key, flush=True) - continue - # 1. Get the original image - current_obs = observation[key] - img_pil = Image.fromarray(current_obs) - orig_width, orig_height = img_pil.size - - # 2. Apply resolution limits (if the configuration is not -1) - target_size = 256 - if target_size != -1: - # Logic for maintaining aspect ratio constraints - if orig_width > orig_height: - new_width = target_size - new_height = int(target_size * orig_height / orig_width) - else: - new_height = target_size - new_width = int(target_size * orig_width / orig_height) - img_pil = img_pil.resize((new_width, new_height)) - - # 3. Apply intelligent scaling - current_width, current_height = img_pil.size - resized_height, resized_width = smart_resize( - current_height, - current_width, - factor=self.IMAGE_FACTOR, - min_pixels=self.MIN_PIXELS, - max_pixels=self.MAX_PIXELS, - ) - resized_img = img_pil.resize((resized_width, resized_height)) - print("resized_img", resized_img.size, flush=True) - - image_inputs.append(resized_img) - - return image_inputs - - def _construct_input( - self, - observation, - instruction, - camera_names, - valid_action_dim=7, - mode="ar", - single_image=False, - ): - additional_inputs = {} - - agent_pos = torch.from_numpy(observation["agent_pos"]) - agent_pos_mask = torch.from_numpy(observation["agent_pos_mask"]) - dof_mask = torch.from_numpy(observation["dof_mask"]) - additional_inputs["dof_mask"] = dof_mask - print("before normalizing agent_pos", agent_pos, flush=True) - - if self.normalizer_propri is not None: - - agent_pos = self.normalizer_propri.normalize_data( - agent_pos, [self.args.dataset_name] - ) - additional_inputs["proprioception"] = agent_pos - additional_inputs["agent_pos_mask"] = agent_pos_mask - - print( - f"normalizing agent_pos: {agent_pos}, {self.args.dataset_name}", - flush=True, - ) - print("agent_pos_mask", agent_pos_mask, flush=True) - print("dof_mask", dof_mask[0, 0], flush=True) - if mode == "ar": - text = self.get_text_ar( - instruction, camera_names, agent_pos, agent_pos_mask - ) - elif mode == "diffusion": - text = self.get_text_flow( - instruction, - camera_names, - self.args.action_horizon, - agent_pos, - agent_pos_mask, - ) - elif mode == "subtask": - text = self.get_text_subtask(instruction, single_view=single_image) - else: - raise ValueError(f"Invalid mode: {mode}") - text = [text] - - image_inputs = self.resize_images(observation) - if single_image: - image_inputs = [ - image_inputs[0] - ] # single view subtask/vqa use head view only - image_inputs = self.processor.image_processor( - images=image_inputs, videos=None, return_tensors="pt" - ) - image_grid_thw = image_inputs["image_grid_thw"] - # Processing image placeholder tokens in the text - if image_grid_thw is not None: - merge_length = self.processor.image_processor.merge_size**2 - index = 0 - for i in range(len(text)): - while "<|image_pad|>" in text[i]: - # Replace image placeholders with actual quantities. - text[i] = text[i].replace( - "<|image_pad|>", - "<|placeholder|>" - * (image_grid_thw[index].prod() // merge_length), - 1, - ) - index += 1 - text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>") - - text_inputs = self.processor.tokenizer( - text, return_tensors="pt", padding=True, truncation=True, max_length=1024 - ) - inputs = BatchFeature(data={**text_inputs, **image_inputs}) - - action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") - additional_inputs["moe_token_types"] = inputs.input_ids == action_token_id - additional_inputs["dataset_names"] = [self.args.dataset_name] - - inputs.update(additional_inputs) - inputs = move_to_cuda(inputs, device) - print("inputs", inputs.keys(), flush=True) - for k in inputs.keys(): - if isinstance(inputs[k], torch.Tensor): - print(k, inputs[k].shape, flush=True) - return inputs - - def preprocess(self, state, views, valid_action_dim): - # print("state",state, flush=True) - # state: dict, keys: follow1_pos, follow2_pos - model_action_dim = sum(self.dof_config.values()) - - if valid_action_dim not in (SINGLE_ARM_DIM, 2 * SINGLE_ARM_DIM): - raise ValueError( - f"Invalid valid_action_dim: {valid_action_dim}, expect 7 or 14" - ) - - # 1) First, prepare a container of size (1, 1, D), where D = model_action_dim. - agent_data = np.zeros((1, 1, model_action_dim), dtype=np.float32) - agent_pos_mask = np.zeros((1, 1, model_action_dim), dtype=np.float32) - dof_mask = np.zeros( - (1, self.args.action_horizon, model_action_dim), dtype=np.float32 - ) - - # 2) Determine the interval to be filled [start:end) - if valid_action_dim == SINGLE_ARM_DIM: - start = 0 if model_action_dim == SINGLE_ARM_DIM else SINGLE_ARM_DIM - end = start + SINGLE_ARM_DIM - - if end > model_action_dim: - raise ValueError( - f"model_action_dim={model_action_dim} too small for valid_action_dim=7 " - f"(need end={end})" - ) - - follow2 = np.asarray(state["follow2_pos"], dtype=np.float32).reshape( - 1, 1, SINGLE_ARM_DIM - ) - agent_data[:, :, start:end] = follow2 - agent_pos_mask[:, :, start:end] = 1 - dof_mask[:, :, start:end] = 1 - - else: # valid_action_dim == 14 - end = 2 * SINGLE_ARM_DIM - if end > model_action_dim: - raise ValueError( - f"model_action_dim={model_action_dim} too small for valid_action_dim=14" - ) - - follow1 = np.asarray(state["follow1_pos"], dtype=np.float32).reshape( - 1, 1, SINGLE_ARM_DIM - ) - follow2 = np.asarray(state["follow2_pos"], dtype=np.float32).reshape( - 1, 1, SINGLE_ARM_DIM - ) - agent_data[:, :, :end] = np.concatenate([follow1, follow2], axis=-1) - agent_pos_mask[:, :, :end] = 1 - dof_mask[:, :, :end] = 1 - - observation = { - camera_to_view_mapping[key]: views[key][0] for key in views.keys() - } - observation["agent_pos"] = agent_data - observation["agent_pos_mask"] = agent_pos_mask - observation["dof_mask"] = dof_mask - return observation - - def model_output_process(self, action_pred, state): - - if not self.args.delta_action and self.args.action_rpy: - return action_pred - - pred_components = extract_components( - action_pred, dim_dof_config, self.args.action_rpy - ) - - pred_right_xyz = pred_components["right_xyz"] - pred_right_rot = pred_components["right_rot"] - pred_right_gripper = pred_components["right_gripper"] - if "left_xyz" in pred_components: - pred_left_xyz = pred_components["left_xyz"] - pred_left_rot = pred_components["left_rot"] - pred_left_gripper = pred_components["left_gripper"] - else: - pred_left_xyz = np.zeros((self.args.action_horizon, 3)) - pred_left_rot = np.zeros((self.args.action_horizon, 3)) - pred_left_gripper = np.zeros((self.args.action_horizon, 1)) - - post_action_pred = np.zeros((self.args.action_horizon, self.args.action_dim)) - if self.args.delta_action: - assert ( - self.args.action_dim == 14 - ), "Delta robot support and testing are not yet available." - - state_components = extract_components( - state, dim_dof_config, self.args.state_rpy - ) - left_xyz = state_components["left_xyz"] - left_rot = state_components["left_rot"] - right_xyz = state_components["right_xyz"] - right_rot = state_components["right_rot"] - - post_action_pred[:, :3] = pred_left_xyz + left_xyz - post_action_pred[:, 3:6] = compose_state_and_delta_to_abs_rpy( - pred_left_rot, left_rot[0] - ) - post_action_pred[:, 6:7] = pred_left_gripper - post_action_pred[:, 7:10] = pred_right_xyz + right_xyz - post_action_pred[:, 10:13] = compose_state_and_delta_to_abs_rpy( - pred_right_rot, right_rot[0] - ) - post_action_pred[:, 13:14] = pred_right_gripper - - elif not self.args.action_rpy: - post_action_pred[:, :3] = pred_left_xyz - post_action_pred[:, 3:6] = so3_to_euler_zyx_batch_nb(pred_left_rot) - post_action_pred[:, 6:7] = pred_left_gripper - post_action_pred[:, 7:10] = pred_right_xyz - post_action_pred[:, 10:13] = so3_to_euler_zyx_batch_nb(pred_right_rot) - post_action_pred[:, 13:14] = pred_right_gripper - else: - post_action_pred = action_pred - return post_action_pred - - def postprocess(self, action_pred, interpolate_multiplier=None): - - if interpolate_multiplier is None: - interpolate_multiplier = self.args.interpolate_multiplier - - if isinstance(action_pred, torch.Tensor): - action_pred = action_pred.to(torch.float32).cpu().squeeze(0).numpy() - left_action_pred = action_pred[:, :7] # (32, 7) - right_action_pred = action_pred[:, 7:14] # (32, 7) - - start_frame = int(self.args.action_start_ratio * len(left_action_pred)) - end_frame = int(self.args.action_end_ratio * len(left_action_pred)) - left_action_pred = left_action_pred[start_frame:end_frame] - right_action_pred = right_action_pred[start_frame:end_frame] - - print("left_action_pred", left_action_pred[-1], flush=True) - print("right_action_pred", right_action_pred[-1], flush=True) - - left_action_pred = left_action_pred.tolist() - right_action_pred = right_action_pred.tolist() - - serialized_actions = { - "follow1_pos": left_action_pred, - "follow2_pos": right_action_pred, - ## for joint-control - # "follow1_joints":left_action_pred, - # "follow2_joints":right_action_pred, - } - - return serialized_actions - - def predict_action_rtc( - self, - state, - views, - instruction=None, - valid_action_dim=7, - update_subtask=False, - action_predict_mode=None, - ): - if action_predict_mode is not None: - self.action_predict_mode = action_predict_mode - - observation = self.preprocess(state, views, valid_action_dim) - print("use instruction", instruction, flush=True) - camera_names = [camera_to_view_mapping[key] for key in views.keys()] - # camera_names = ["right_wrist_view", "global_view", "side_view"] - print("mode:", self.action_predict_mode, flush=True) - inputs = self._construct_input( - observation, - instruction, - camera_names=camera_names, - valid_action_dim=valid_action_dim, - mode=self.action_predict_mode, - ) - model_action_dim = sum(self.dof_config.values()) - padding = torch.zeros((1, model_action_dim)) - norm_padding = self.normalizer_action.normalize_data( - padding, [self.args.dataset_name] - ) - inputs["padding_action"] = norm_padding - inputs = move_to_cuda(inputs, device="cuda:0") - agent_data = observation["agent_pos"][..., : self.args.model_action_dim][0] - print("before generate_flow_action", flush=True) - print(inputs.keys(), flush=True) - print(inputs["dataset_names"], flush=True) - print(self.processor.tokenizer.decode(inputs["input_ids"][0]), flush=True) - print(inputs["agent_pos_mask"][0], flush=True) - print(inputs["dof_mask"][0, 0], flush=True) - if self.action_predict_mode == "ar": - action_pred = self.generate_ar_action(inputs) - else: - action_pred = self.generate_flow_action(inputs) - print("after generate_flow_action", flush=True) - if isinstance(action_pred, torch.Tensor): - action_pred = action_pred.float().cpu().squeeze(0).numpy() - - if action_pred is None: - return None - - if self.args.dct_scale > 0: - scale = self.args.dct_scale - dct_coeff = dct(action_pred, axis=0, norm="ortho") - dct_coeff = np.around(dct_coeff * scale) - action_pred = idct(dct_coeff / scale, axis=0, norm="ortho") - - if action_pred.shape[-1] == 7: - print("Before concat action_pred", action_pred.shape, flush=True) - right_action_pred = action_pred - left_action_pred = np.zeros_like(right_action_pred) - action_pred = np.concatenate([left_action_pred, right_action_pred], axis=1) - print("After concat action_pred", action_pred.shape, flush=True) - # unnorm action_pred - # print("action_pred", action_pred[:, 3], flush=True) - action_pred = ( - self.normalizer_action.unnormalize_data( - torch.tensor(action_pred).unsqueeze(0), [self.args.dataset_name] - ) - .squeeze(0) - .numpy() - ) - print("After unnormalize_data action_pred", action_pred.shape, flush=True) - - action_pred = self.model_output_process(action_pred, agent_data) - action_pred = self.postprocess(action_pred) - - return action_pred - - def generate_flow_action( - self, - inputs, - last_action_chunk=None, - max_guidance_weight=20.0, - num_inference_timesteps=10, - sigma_action=0.2, - ): - model_action_dim = sum(self.dof_config.values()) - if last_action_chunk is None: - output = self.model.generate_flow_action( - action_horizon=self.args.action_horizon, - action_dim=model_action_dim, - num_inference_timesteps=num_inference_timesteps, - unnorm=False, - **inputs, - ) - else: - output = self.model.generate_flow_action_rtc( - action_horizon=self.args.action_horizon, - action_dim=model_action_dim, - num_inference_timesteps=num_inference_timesteps, - inference_delay=self.args.rtc_inference_delay, - execution_horizon=self.args.rtc_execution_horizon - 1, - max_guidance_weight=max_guidance_weight, - last_action_chunk=last_action_chunk, - sigma_action=sigma_action, - unnorm=False, - **inputs, - ) - action_pred = output["predict_action"] # (b, action_horizon, action_dim) - return action_pred - - def generate_text(self, inputs): - return self.model.generate_text(**inputs) - - def _preprocess_ar_batch(self, batch): - input_ids = batch["input_ids"] - attention_mask = batch["attention_mask"] - moe_token_types = batch["moe_token_types"] - labels = batch.get("labels", None) - prefix_length = batch.get("prefix_length", None) - - generation_prompt_ids = torch.tensor( - [151644, 77091], device=input_ids.device, dtype=input_ids.dtype - ) # <|im_start|>assistant - matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & ( - input_ids[0, 1:] == generation_prompt_ids[1] - ) - if matches.any(): - split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item() - # construct output ids - gt_output_ids = input_ids[:, split_pos + 3 : prefix_length] - # remove output part from input - input_ids = input_ids[:, : split_pos + 3] - moe_token_types = moe_token_types[:, : split_pos + 3] - if attention_mask is not None: - attention_mask = attention_mask[:, : split_pos + 3] - if labels is not None: - labels = labels[:, split_pos + 3 : prefix_length] - - batch.update( - { - "input_ids": input_ids, - "attention_mask": attention_mask, - "moe_token_types": moe_token_types, - "labels": labels, - "gt_output_ids": gt_output_ids, - "prefix_length": split_pos + 3, - } - ) - - return batch - - def generate_ar_action(self, inputs): - # batch = self._preprocess_ar_batch(batch=inputs) - action_pred = None - count = 0 - while action_pred is None: - if count > 5: - # raise ValueError("re-generate ar action failed") - return None - count += 1 - output = self.model.generate_ar_action( - # action_dim=args.action_dim, - action_dim=14, - action_horizon=self.args.action_horizon, - unnorm=False, - **inputs, - ) - action_pred = output["predict_action"] - - action_pred = action_pred[0] - return action_pred - - -class WallxInfer: - def __init__(self, args: WallxInferArgs): - self.args = args - self.model_wrapper = WallxModelWrapper(args) - - def run_infer_robochallenge( - self, state, views, instruction, valid_action_dim=7, action_predict_mode=None - ): - action_pred = self.model_wrapper.predict_action_rtc( - state=state, - views=views, - instruction=instruction, - valid_action_dim=valid_action_dim, - action_predict_mode=action_predict_mode, - ) - return action_pred - - -if __name__ == "__main__": - args = WallxInferArgs() - Infer = WallxInfer(args) - Infer.run_infer() diff --git a/scripts/merge_sharded_weights.py b/scripts/merge_sharded_weights.py index 314e228..1f45b57 100644 --- a/scripts/merge_sharded_weights.py +++ b/scripts/merge_sharded_weights.py @@ -6,9 +6,10 @@ Works around the StorageMeta compatibility issue between PyTorch versions. import os import sys -import torch from pathlib import Path from typing import Dict + +import torch from safetensors.torch import save_file @@ -70,65 +71,16 @@ def load_sharded_checkpoint(checkpoint_dir: str) -> Dict[str, torch.Tensor]: except AttributeError as e: if "StorageMeta" in str(e): - print(f"[ERROR] StorageMeta compatibility issue: {e}") - print("[INFO] Attempting alternative loading method...") - return load_checkpoint_alternative(checkpoint_dir) - else: - raise - - -def load_checkpoint_alternative(checkpoint_dir: str) -> Dict[str, torch.Tensor]: - """ - Alternative method to load checkpoint by directly reading shard files. - - Args: - checkpoint_dir: Path to directory containing .distcp files - - Returns: - Dictionary of merged model state - """ - checkpoint_path = Path(checkpoint_dir) - - # Find all shard files - shard_files = sorted(checkpoint_path.glob("*.distcp")) - - if not shard_files: - raise FileNotFoundError(f"No .distcp files found in {checkpoint_dir}") - - print(f"[INFO] Found {len(shard_files)} shard files") - - # Load all shards - merged_state = {} - - for shard_file in shard_files: - print(f"[INFO] Loading shard: {shard_file.name}") - try: - shard_data = torch.load(shard_file, map_location="cpu") - - # Merge the shard into the state dict - if isinstance(shard_data, dict): - for key, value in shard_data.items(): - if isinstance(value, torch.Tensor): - if key in merged_state: - # Handle duplicates - concatenate or overwrite based on shape - print(f"[WARNING] Duplicate key found: {key}") - merged_state[key] = value - elif isinstance(value, dict): - # Nested dict structure - for subkey, subvalue in value.items(): - full_key = f"{key}.{subkey}" if key else subkey - if isinstance(subvalue, torch.Tensor): - merged_state[full_key] = subvalue - - except Exception as e: - print(f"[WARNING] Failed to load shard {shard_file.name}: {e}") - continue - - if not merged_state: - raise RuntimeError("Failed to load any checkpoint data from shards") - - print(f"[INFO] Loaded {len(merged_state)} tensors from shards") - return merged_state + raise RuntimeError( + "Unable to load this DCP checkpoint because its metadata uses " + "StorageMeta from a different PyTorch version. The previous " + "manual .distcp fallback was removed because FSDP shards cannot " + "be reconstructed by directly loading shard files and overwriting " + "duplicate keys. Please run this script with a PyTorch version " + "compatible with the checkpoint writer, or re-save the checkpoint " + "with the current PyTorch version." + ) from e + raise def save_merged_checkpoint( diff --git a/scripts/merge_tokenizer.py b/scripts/merge_tokenizer.py index 52ac1ae..8478fff 100644 --- a/scripts/merge_tokenizer.py +++ b/scripts/merge_tokenizer.py @@ -1,26 +1,107 @@ -from transformers import AutoProcessor -import os +#!/usr/bin/env python3 +"""Merge Wall-X action tokens into a Qwen2.5-VL processor tokenizer.""" -processor_path = "/path/to/Qwen2.5-VL-3B-Instruct" -action_tokenizer_path = "/path/to/fast" -use_fast_tokenizer = True +from __future__ import annotations -processor = AutoProcessor.from_pretrained(processor_path, use_fast=True) -processor.tokenizer.padding_side = "left" +import argparse +from pathlib import Path -action_tokenizer = AutoProcessor.from_pretrained( - action_tokenizer_path, trust_remote_code=True -) -new_tokens = ["<|propri|>", "<|action|>"] -new_tokens += [f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)] -num_added_tokens = processor.tokenizer.add_tokens(new_tokens) +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Create a Wall-X processor directory by adding FAST action tokens " + "to a Qwen2.5-VL processor tokenizer." + ) + ) + parser.add_argument( + "--processor-path", + required=True, + help="Base Qwen2.5-VL processor directory or Hugging Face repo id.", + ) + parser.add_argument( + "--action-tokenizer-path", + required=True, + help="FAST/action tokenizer processor directory or Hugging Face repo id.", + ) + parser.add_argument( + "--output-dir", + required=True, + help="Directory where the merged processor will be written.", + ) + parser.add_argument( + "--use-fast", + action=argparse.BooleanOptionalAction, + default=True, + help="Use the fast tokenizer implementation when loading the base processor.", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Allow custom code when loading the action tokenizer processor.", + ) + return parser.parse_args() -begin_idx_token = "<|action_token_0|>" -token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token) -processor.tokenizer.init_kwargs["action_token_start_index"] = token_id -processor.tokenizer.init_kwargs["action_token_vocab_size"] = action_tokenizer.vocab_size -new_tokenizer_dir = "/path/to/new_tokenizer" -os.makedirs(new_tokenizer_dir, exist_ok=True) -processor.save_pretrained(new_tokenizer_dir) +def _resolve_action_vocab_size(action_processor) -> int: + vocab_size = getattr(action_processor, "vocab_size", None) + if vocab_size is None and hasattr(action_processor, "tokenizer"): + vocab_size = getattr(action_processor.tokenizer, "vocab_size", None) + if vocab_size is None: + raise AttributeError( + "Could not determine action tokenizer vocab size from the loaded processor." + ) + return int(vocab_size) + + +def merge_tokenizer( + *, + processor_path: str, + action_tokenizer_path: str, + output_dir: str, + use_fast: bool, + trust_remote_code: bool, +) -> None: + from transformers import AutoProcessor + + processor = AutoProcessor.from_pretrained(processor_path, use_fast=use_fast) + processor.tokenizer.padding_side = "left" + + action_processor = AutoProcessor.from_pretrained( + action_tokenizer_path, + trust_remote_code=trust_remote_code, + ) + action_vocab_size = _resolve_action_vocab_size(action_processor) + + new_tokens = ["<|propri|>", "<|action|>"] + new_tokens += [f"<|action_token_{i}|>" for i in range(action_vocab_size)] + num_added_tokens = processor.tokenizer.add_tokens(new_tokens) + + begin_idx_token = "<|action_token_0|>" + token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token) + processor.tokenizer.init_kwargs["action_token_start_index"] = token_id + processor.tokenizer.init_kwargs["action_token_vocab_size"] = action_vocab_size + + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + processor.save_pretrained(output_path) + + print(f"Saved merged processor to {output_path}") + print(f"Added {num_added_tokens} tokenizer tokens") + print(f"action_token_start_index={token_id}") + print(f"action_token_vocab_size={action_vocab_size}") + + +def main() -> None: + args = parse_args() + merge_tokenizer( + processor_path=args.processor_path, + action_tokenizer_path=args.action_tokenizer_path, + output_dir=args.output_dir, + use_fast=args.use_fast, + trust_remote_code=args.trust_remote_code, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/normalize.py b/scripts/normalize.py deleted file mode 100644 index 75a3f98..0000000 --- a/scripts/normalize.py +++ /dev/null @@ -1,161 +0,0 @@ -# This file is copied from openpi -import json -import pathlib - -import numpy as np -import numpydantic -import pydantic - - -@pydantic.dataclasses.dataclass -class NormStats: - mean: numpydantic.NDArray - std: numpydantic.NDArray - q01: numpydantic.NDArray | None = None # 1st quantile - q99: numpydantic.NDArray | None = None # 99th quantile - - -class RunningStats: - """Compute running statistics of a batch of vectors.""" - - def __init__(self): - self._count = 0 - self._mean = None - self._mean_of_squares = None - self._min = None - self._max = None - self._histograms = None - self._bin_edges = None - self._num_quantile_bins = 5000 # for computing quantiles on the fly - - def update(self, batch: np.ndarray) -> None: - """ - Update the running statistics with a batch of vectors. - - Args: - vectors (np.ndarray): An array where all dimensions except the last are batch dimensions. - """ - batch = batch.reshape(-1, batch.shape[-1]) - num_elements, vector_length = batch.shape - if self._count == 0: - self._mean = np.mean(batch, axis=0) - self._mean_of_squares = np.mean(batch**2, axis=0) - self._min = np.min(batch, axis=0) - self._max = np.max(batch, axis=0) - self._histograms = [ - np.zeros(self._num_quantile_bins) for _ in range(vector_length) - ] - self._bin_edges = [ - np.linspace( - self._min[i] - 1e-10, - self._max[i] + 1e-10, - self._num_quantile_bins + 1, - ) - for i in range(vector_length) - ] - else: - if vector_length != self._mean.size: - raise ValueError( - "The length of new vectors does not match the initialized vector length." - ) - new_max = np.max(batch, axis=0) - new_min = np.min(batch, axis=0) - max_changed = np.any(new_max > self._max) - min_changed = np.any(new_min < self._min) - self._max = np.maximum(self._max, new_max) - self._min = np.minimum(self._min, new_min) - - if max_changed or min_changed: - self._adjust_histograms() - - self._count += num_elements - - batch_mean = np.mean(batch, axis=0) - batch_mean_of_squares = np.mean(batch**2, axis=0) - - # Update running mean and mean of squares. - self._mean += (batch_mean - self._mean) * (num_elements / self._count) - self._mean_of_squares += (batch_mean_of_squares - self._mean_of_squares) * ( - num_elements / self._count - ) - - self._update_histograms(batch) - - def get_statistics(self) -> NormStats: - """ - Compute and return the statistics of the vectors processed so far. - - Returns: - dict: A dictionary containing the computed statistics. - """ - if self._count < 2: - raise ValueError("Cannot compute statistics for less than 2 vectors.") - - variance = self._mean_of_squares - self._mean**2 - stddev = np.sqrt(np.maximum(0, variance)) - q01, q99 = self._compute_quantiles([0.01, 0.99]) - return NormStats(mean=self._mean, std=stddev, q01=q01, q99=q99) - - def _adjust_histograms(self): - """Adjust histograms when min or max changes.""" - for i in range(len(self._histograms)): - old_edges = self._bin_edges[i] - new_edges = np.linspace( - self._min[i], self._max[i], self._num_quantile_bins + 1 - ) - - # Redistribute the existing histogram counts to the new bins - new_hist, _ = np.histogram( - old_edges[:-1], bins=new_edges, weights=self._histograms[i] - ) - - self._histograms[i] = new_hist - self._bin_edges[i] = new_edges - - def _update_histograms(self, batch: np.ndarray) -> None: - """Update histograms with new vectors.""" - for i in range(batch.shape[1]): - hist, _ = np.histogram(batch[:, i], bins=self._bin_edges[i]) - self._histograms[i] += hist - - def _compute_quantiles(self, quantiles): - """Compute quantiles based on histograms.""" - results = [] - for q in quantiles: - target_count = q * self._count - q_values = [] - for hist, edges in zip(self._histograms, self._bin_edges, strict=True): - cumsum = np.cumsum(hist) - idx = np.searchsorted(cumsum, target_count) - q_values.append(edges[idx]) - results.append(np.array(q_values)) - return results - - -class _NormStatsDict(pydantic.BaseModel): - norm_stats: dict[str, NormStats] - - -def serialize_json(norm_stats: dict[str, NormStats]) -> str: - """Serialize the running statistics to a JSON string.""" - return _NormStatsDict(norm_stats=norm_stats).model_dump_json(indent=2) - - -def deserialize_json(data: str) -> dict[str, NormStats]: - """Deserialize the running statistics from a JSON string.""" - return _NormStatsDict(**json.loads(data)).norm_stats - - -def save(directory: pathlib.Path | str, norm_stats: dict[str, NormStats]) -> None: - """Save the normalization stats to a directory.""" - path = pathlib.Path(directory) / "norm_stats.json" - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(serialize_json(norm_stats)) - - -def load(directory: pathlib.Path | str) -> dict[str, NormStats]: - """Load the normalization stats from a directory.""" - path = pathlib.Path(directory) / "norm_stats.json" - if not path.exists(): - raise FileNotFoundError(f"Norm stats file not found at: {path}") - return deserialize_json(path.read_text()) diff --git a/scripts/run_libero.sh b/scripts/run_libero.sh new file mode 100755 index 0000000..c78b46f --- /dev/null +++ b/scripts/run_libero.sh @@ -0,0 +1,273 @@ +#!/usr/bin/env bash +set -euo pipefail + +# LIBERO evaluation launcher. Configure with environment variables: +# +# CHECKPOINT_PATH=/path/to/checkpoint bash scripts/run_libero.sh +# bash scripts/run_libero.sh /path/to/checkpoint +# CHECKPOINT_PATH=/path/to/checkpoint SMOKE=1 bash scripts/run_libero.sh +# CONFIG=/path/to/eval.yaml bash scripts/run_libero.sh +# +# Optional knobs: +# LIBERO_PATH=/path/to/LIBERO # optional when not cloned to third_party/LIBERO +# CUDA_ID=0 +# DRIVER_MODE=in_process +# NUM_WORKERS=1 +# MAX_BATCH_SIZE=1 +# ALL_SUITES=1 # run all standard LIBERO suites (40 tasks) +# TASK_SUITES="libero_spatial ..." # custom suite list (space- or comma-separated) +# TASK_SUITE_NAME=libero_spatial # single suite when ALL_SUITES=0 and TASK_SUITES unset +# NUM_TRIALS_PER_TASK=50 +# TASK_INDICES=0,1,2 # omit to run every task in the suite +# MAX_INFER_TIMES=52 +# NORM_KEY=libero_all +# ROLLOUT_BASE=/path/to/rollout # per-suite logs under ${ROLLOUT_BASE}/${suite}/ +# LOG_DIR=/tmp/harrix_libero_eval # overrides ROLLOUT_BASE when set (single suite) +# SKIP_LIBERO_DEP_CHECK=1 # bypass dependency preflight + +# export ALL_SUITES=1 +# export TASK_SUITE_NAME=libero_10 +# export CUDA_ID=0 +# export NUM_WORKERS=10 +# export NUM_TRIALS_PER_TASK=20 +# export CHECKPOINT_PATH=/path/to/checkpoint +# export ROLLOUT_BASE=/path/to/rollout + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SOURCE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" + +if [[ -d "${SOURCE_ROOT}/third_party/LIBERO" ]]; then + export PYTHONPATH="${SOURCE_ROOT}/third_party/LIBERO:${PYTHONPATH:-}" +fi +if [[ -n "${LIBERO_PATH:-}" ]]; then + export PYTHONPATH="${LIBERO_PATH}:${PYTHONPATH:-}" +fi + +if [[ -f "${SCRIPT_DIR}/infer_libero.py" ]]; then + INFER_LIBERO_CMD=("${SCRIPT_DIR}/infer_libero.py") +elif [[ -f "${SOURCE_ROOT}/scripts/infer_libero.py" ]]; then + INFER_LIBERO_CMD=(python "${SOURCE_ROOT}/scripts/infer_libero.py") +elif command -v infer_libero.py >/dev/null 2>&1; then + INFER_LIBERO_CMD=(infer_libero.py) +else + echo "infer_libero.py is not available. Install Wall-X first." >&2 + exit 2 +fi + +DEFAULT_ALL_SUITES=( + libero_spatial + libero_object + libero_goal + libero_10 +) + +if [[ $# -gt 1 ]]; then + echo "Usage: bash scripts/run_libero.sh [CHECKPOINT_PATH]" >&2 + exit 2 +fi +if [[ $# -eq 1 ]]; then + CHECKPOINT_PATH="$1" +fi + +export CUDA_VISIBLE_DEVICES="${CUDA_ID:-0}" + +# MuJoCo offscreen rendering via NVIDIA EGL (required on headless GPU nodes). +export MUJOCO_GL=egl +export PYOPENGL_PLATFORM=egl +# LIBERO init-state files are trusted simulator assets. PyTorch 2.6 changed +# torch.load() defaults in a way that breaks LIBERO's upstream loader. +export TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD="${TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD:-1}" +EGL_VENDOR_DIR="${HOME}/.config/glvnd/egl_vendor.d" +EGL_VENDOR_FILE="${EGL_VENDOR_DIR}/10_nvidia.json" +if [[ ! -f "${EGL_VENDOR_FILE}" ]]; then + mkdir -p "${EGL_VENDOR_DIR}" + cat > "${EGL_VENDOR_FILE}" <<'EOF' +{ + "file_format_version" : "1.0.0", + "ICD" : { + "library_path" : "libEGL_nvidia.so.0" + } +} +EOF +fi +export __EGL_VENDOR_LIBRARY_FILENAMES="${EGL_VENDOR_FILE}" +export LD_LIBRARY_PATH="/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib:/usr/local/nvidia/lib64:${LD_LIBRARY_PATH:-}" + +# After CUDA_VISIBLE_DEVICES remapping, MuJoCo only sees devices from index 0. +export MUJOCO_EGL_DEVICE_ID="${MUJOCO_EGL_DEVICE_ID:-0}" +if [[ -d "${SOURCE_ROOT}/third_party/harrix/python" ]]; then + export PYTHONPATH="${SOURCE_ROOT}/third_party/harrix/python:${PYTHONPATH:-}" +fi +if [[ -d "${SOURCE_ROOT}/wall_x" ]]; then + export PYTHONPATH="${SOURCE_ROOT}:${PYTHONPATH:-}" +fi + +check_libero_dependencies() { + local missing + if missing="$( + python - <<'PY' +import importlib.util + +checks = [ + ("libero.libero", "LIBERO"), + ("robosuite", "robosuite"), + ("mujoco", "mujoco"), + ("OpenGL", "PyOpenGL"), + ("bddl", "bddl"), + ("gym", "gym"), + ("h5py", "h5py"), +] + +missing = [label for module, label in checks if importlib.util.find_spec(module) is None] +if missing: + print(", ".join(missing)) + raise SystemExit(1) +PY + )"; then + return 0 + fi + + cat >&2 <&2 + exit 2 + fi + args+=(--checkpoint-path "${CHECKPOINT_PATH}") + fi + + if [[ -n "${TRAIN_CONFIG_PATH:-}" ]]; then + args+=(--train-config-path "${TRAIN_CONFIG_PATH}") + fi + if [[ "${SMOKE:-0}" == "1" ]]; then + args+=(--smoke) + fi + if [[ "${DET:-0}" == "1" ]]; then + args+=(--deterministic-model) + fi + if [[ -n "${TASK_INDICES:-}" ]]; then + args+=(--task-indices "${TASK_INDICES}") + fi + if [[ -n "${MAX_INFER_TIMES:-}" ]]; then + args+=(--max-infer-times "${MAX_INFER_TIMES}") + fi + + args+=(--driver-mode "${DRIVER_MODE:-in_process}") + args+=(--num-workers "${NUM_WORKERS:-1}") + args+=(--max-batch-size "${MAX_BATCH_SIZE:-1}") + args+=(--task-suite-name "${suite}") + args+=(--num-trials-per-task "${NUM_TRIALS_PER_TASK:-50}") + args+=(--norm-key "${NORM_KEY:-libero_all}") + args+=(--log-dir "${LOG_DIR_RESOLVED}") + + echo "=== LIBERO eval: suite=${suite} log_dir=${LOG_DIR_RESOLVED} ===" + "${INFER_LIBERO_CMD[@]}" "${args[@]}" +} + +suites=($(resolve_task_suites)) +multi_suite=0 +if [[ "${#suites[@]}" -gt 1 ]]; then + multi_suite=1 +fi + +for suite in "${suites[@]}"; do + run_suite "${suite}" "${multi_suite}" +done diff --git a/scripts/run_serving.sh b/scripts/run_serving.sh new file mode 100755 index 0000000..14bfeeb --- /dev/null +++ b/scripts/run_serving.sh @@ -0,0 +1,285 @@ +#!/usr/bin/env bash +set -euo pipefail + +usage() { + cat <<'EOF' +Usage: + bash scripts/run_serving.sh --checkpoint-path /path/to/checkpoint [options] + +Required: + --checkpoint-path PATH Checkpoint directory or checkpoint file. + +Common options: + --train-config-path PATH Training config used by the checkpoint. + --port PORT WebSocket port. Default: 32195. + --host HOST Bind host. Default: 0.0.0.0. + --env X2ROBOT|LIBERO Serving environment. Default: X2ROBOT. + --cuda-id ID Sets CUDA_VISIBLE_DEVICES. Default: 0. + --image-passing-mode MODE base64 or numpy. Default: base64. + --action-horizon N Model action horizon. Default: 32. + --robot-type TYPE desktop, turtle, or ex001. Default: desktop. + --raw-actions Alias for --no-serialize-actions. + --serialize-actions Return robot-serialized actions. + --no-serialize-actions Return raw model action chunks. Default. + --max-batch-size N Enable dynamic batching. + --enable-cuda-graph Enable CUDA graph in the serving runtime. + --enable-experimental-engine Enable the experimental inference engine. + --debug Enable debug logging. + --dry-run Print the command without running it. + +Additional arguments after "--" are forwarded to launch_serving.py, for example: + bash scripts/run_serving.sh --checkpoint-path /ckpt -- \ + --model-config.norm-key libero_all + +Environment variables can also be used, e.g. CHECKPOINT_PATH, +TRAIN_CONFIG_PATH, PORT, CUDA_ID, ACTION_HORIZON, WALLX_ENV. +EOF +} + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SOURCE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" + +if [[ -d "${SOURCE_ROOT}/wall_x" ]]; then + export PYTHONPATH="${SOURCE_ROOT}:${PYTHONPATH:-}" +fi + +PYTHON_BIN="${PYTHON_BIN:-python}" +CHECKPOINT_PATH="${CHECKPOINT_PATH:-}" +TRAIN_CONFIG_PATH="${TRAIN_CONFIG_PATH:-}" +PORT="${PORT:-32195}" +HOST="${HOST:-0.0.0.0}" +WALLX_ENV="${WALLX_ENV:-X2ROBOT}" +CUDA_ID="${CUDA_ID:-0}" +IMAGE_PASSING_MODE="${IMAGE_PASSING_MODE:-base64}" +ACTION_HORIZON="${ACTION_HORIZON:-32}" +ROBOT_TYPE="${ROBOT_TYPE:-desktop}" +ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${ROBOT_ACTION_INTERPOLATE_MULTIPLIER:-1}" +ROBOT_ACTION_END_RATIO="${ROBOT_ACTION_END_RATIO:-1.0}" +MODEL_DEVICE="${MODEL_DEVICE:-cuda}" +MAX_BATCH_SIZE="${MAX_BATCH_SIZE:-}" +DEFAULT_PROMPT="${DEFAULT_PROMPT:-}" +SERIALIZE_ACTIONS="${SERIALIZE_ACTIONS:-0}" +ENABLE_CUDA_GRAPH="${ENABLE_CUDA_GRAPH:-0}" +ENABLE_EXPERIMENTAL_ENGINE="${ENABLE_EXPERIMENTAL_ENGINE:-0}" +DEBUG="${DEBUG:-0}" +DRY_RUN=0 +EXTRA_ARGS=() + +while [[ $# -gt 0 ]]; do + case "$1" in + -h|--help) + usage + exit 0 + ;; + --checkpoint-path) + CHECKPOINT_PATH="${2:?missing value for --checkpoint-path}" + shift 2 + ;; + --checkpoint-path=*) + CHECKPOINT_PATH="${1#*=}" + shift + ;; + --train-config-path) + TRAIN_CONFIG_PATH="${2:?missing value for --train-config-path}" + shift 2 + ;; + --train-config-path=*) + TRAIN_CONFIG_PATH="${1#*=}" + shift + ;; + --port) + PORT="${2:?missing value for --port}" + shift 2 + ;; + --port=*) + PORT="${1#*=}" + shift + ;; + --host) + HOST="${2:?missing value for --host}" + shift 2 + ;; + --host=*) + HOST="${1#*=}" + shift + ;; + --env) + WALLX_ENV="${2:?missing value for --env}" + shift 2 + ;; + --env=*) + WALLX_ENV="${1#*=}" + shift + ;; + --cuda-id) + CUDA_ID="${2:?missing value for --cuda-id}" + shift 2 + ;; + --cuda-id=*) + CUDA_ID="${1#*=}" + shift + ;; + --image-passing-mode) + IMAGE_PASSING_MODE="${2:?missing value for --image-passing-mode}" + shift 2 + ;; + --image-passing-mode=*) + IMAGE_PASSING_MODE="${1#*=}" + shift + ;; + --action-horizon) + ACTION_HORIZON="${2:?missing value for --action-horizon}" + shift 2 + ;; + --action-horizon=*) + ACTION_HORIZON="${1#*=}" + shift + ;; + --robot-type) + ROBOT_TYPE="${2:?missing value for --robot-type}" + shift 2 + ;; + --robot-type=*) + ROBOT_TYPE="${1#*=}" + shift + ;; + --robot-action-interpolate-multiplier) + ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${2:?missing value for --robot-action-interpolate-multiplier}" + shift 2 + ;; + --robot-action-interpolate-multiplier=*) + ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${1#*=}" + shift + ;; + --robot-action-end-ratio) + ROBOT_ACTION_END_RATIO="${2:?missing value for --robot-action-end-ratio}" + shift 2 + ;; + --robot-action-end-ratio=*) + ROBOT_ACTION_END_RATIO="${1#*=}" + shift + ;; + --model-device) + MODEL_DEVICE="${2:?missing value for --model-device}" + shift 2 + ;; + --model-device=*) + MODEL_DEVICE="${1#*=}" + shift + ;; + --max-batch-size) + MAX_BATCH_SIZE="${2:?missing value for --max-batch-size}" + shift 2 + ;; + --max-batch-size=*) + MAX_BATCH_SIZE="${1#*=}" + shift + ;; + --default-prompt) + DEFAULT_PROMPT="${2:?missing value for --default-prompt}" + shift 2 + ;; + --default-prompt=*) + DEFAULT_PROMPT="${1#*=}" + shift + ;; + --serialize-actions) + SERIALIZE_ACTIONS=1 + shift + ;; + --no-serialize-actions|--raw-actions) + SERIALIZE_ACTIONS=0 + shift + ;; + --enable-cuda-graph) + ENABLE_CUDA_GRAPH=1 + shift + ;; + --enable-experimental-engine) + ENABLE_EXPERIMENTAL_ENGINE=1 + shift + ;; + --debug) + DEBUG=1 + shift + ;; + --dry-run) + DRY_RUN=1 + shift + ;; + --) + shift + EXTRA_ARGS+=("$@") + break + ;; + *) + EXTRA_ARGS+=("$1") + shift + ;; + esac +done + +if [[ -z "${CHECKPOINT_PATH}" ]]; then + echo "error: --checkpoint-path is required." >&2 + echo >&2 + usage >&2 + exit 2 +fi + +export CUDA_VISIBLE_DEVICES="${CUDA_ID}" +export ENABLE_FAST_PREPROCESS="${ENABLE_FAST_PREPROCESS:-true}" + +CMD=( + "${PYTHON_BIN}" -m wall_x._vendor.harrix.serving.launch_serving + --env "${WALLX_ENV}" + --host "${HOST}" + --port "${PORT}" + --image-passing-mode "${IMAGE_PASSING_MODE}" +) + +if [[ "${SERIALIZE_ACTIONS}" == "1" ]]; then + CMD+=(--serialize-actions) +else + CMD+=(--no-serialize-actions) +fi +if [[ -n "${MAX_BATCH_SIZE}" ]]; then + CMD+=(--max-batch-size "${MAX_BATCH_SIZE}") +fi +if [[ -n "${DEFAULT_PROMPT}" ]]; then + CMD+=(--default-prompt "${DEFAULT_PROMPT}") +fi +if [[ "${ENABLE_CUDA_GRAPH}" == "1" ]]; then + CMD+=(--enable-cuda-graph) +fi +if [[ "${ENABLE_EXPERIMENTAL_ENGINE}" == "1" ]]; then + CMD+=(--enable-experimental-engine) +fi +if [[ "${DEBUG}" == "1" ]]; then + CMD+=(--debug) +fi + +CMD+=( + model-config:server-model-config + --model-config.checkpoint-path "${CHECKPOINT_PATH}" + --model-config.action-horizon "${ACTION_HORIZON}" + --model-config.robot-type "${ROBOT_TYPE}" + --model-config.robot-action-interpolate-multiplier "${ROBOT_ACTION_INTERPOLATE_MULTIPLIER}" + --model-config.robot-action-end-ratio "${ROBOT_ACTION_END_RATIO}" + --model-config.model-device "${MODEL_DEVICE}" +) + +if [[ -n "${TRAIN_CONFIG_PATH}" ]]; then + CMD+=(--model-config.train-config-path "${TRAIN_CONFIG_PATH}") +fi + +CMD+=("${EXTRA_ARGS[@]}") + +printf 'Launching Wall-X serving:\n' +printf ' %q' "${CMD[@]}" +printf '\n' + +if [[ "${DRY_RUN}" == "1" ]]; then + exit 0 +fi + +exec "${CMD[@]}" diff --git a/scripts/vqa_inference.py b/scripts/vqa_inference.py deleted file mode 100644 index adbef89..0000000 --- a/scripts/vqa_inference.py +++ /dev/null @@ -1,104 +0,0 @@ -import torch -from PIL import Image -from transformers import AutoProcessor -import yaml -import os - -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction - - -class VQAWrapper(object): - def __init__(self, model_path: str, train_config: dict = None): - - self.device = self._setup_device() - if train_config is None: - try: - with open(os.path.join(model_path, "config.yml"), "r") as f: - train_config = yaml.load(f, Loader=yaml.FullLoader) - except Exception as e: - print(f"load train_config.yml fail: {e}") - self.processor = self._load_processor(train_config["processor_path"]) - self.model = self._load_model(model_path, train_config) - - def _setup_device(self) -> str: - if torch.cuda.is_available(): - return "cuda" - else: - return "cpu" - - def _load_processor(self, model_path: str) -> AutoProcessor: - return AutoProcessor.from_pretrained(model_path, trust_remote_code=True) - - def _load_model( - self, model_path: str, train_config: dict - ) -> Qwen2_5_VLMoEForAction: - model = Qwen2_5_VLMoEForAction.from_pretrained( - model_path, train_config=train_config - ) - if self.device == "cuda": - model = model.to(self.device, dtype=torch.bfloat16) - else: - model.to(self.device) - model.eval() - return model - - def generate(self, image: Image.Image, text: str, **kwargs) -> str: - messages = [ - { - "role": "user", - "content": [{"type": "image"}, {"type": "text", "text": text}], - } - ] - text_prompt = self.processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - inputs = self.processor(text=[text_prompt], images=[image], return_tensors="pt") - inputs = {k: v.to(self.device) for k, v in inputs.items()} - - generation_params = { - "max_new_tokens": 1024, # default value, can be overridden by kwargs - "do_sample": False, - "eos_token_id": self.processor.tokenizer.eos_token_id, - "pad_token_id": self.processor.tokenizer.pad_token_id, - **kwargs, - } - - with torch.no_grad(): - generated_ids = self.model.generate(**inputs, **generation_params) - - generated_ids = [ - output_ids[len(input_ids) :] - for input_ids, output_ids in zip(inputs["input_ids"], generated_ids) - ] - response = self.processor.batch_decode( - generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False - )[0] - return response - - -if __name__ == "__main__": - MODEL_PATH_FOR_MODULE_TEST = "/path/to/model_path" - train_config_path = "/path/to/model_path/config.yml" - with open(train_config_path, "r") as f: - train_config = yaml.load(f, Loader=yaml.FullLoader) - wrapper = VQAWrapper( - model_path=MODEL_PATH_FOR_MODULE_TEST, train_config=train_config - ) - - try: - test_question = "To move the red block in the plate with same color, what should you do next? Think step by step." - - # Local Image - img = Image.open( - "/x2robot_v2/yangping/github/wall-x/assets/cot_example_frame.png" - ).convert("RGB") - # Internet Image - # import requests - # test_image_url = "https://www.ilankelman.org/stopsigns/australia.jpg" - # img = Image.open(requests.get(test_image_url, stream=True).raw).convert("RGB") - - answer = wrapper.generate(img, test_question) - - print("model answer:", answer) - except Exception as e: - print(f"model answer fail: {e}") diff --git a/setup.py b/setup.py index 52fe219..ca6b4f8 100644 --- a/setup.py +++ b/setup.py @@ -1,63 +1,78 @@ -import os -import torch from pathlib import Path -from setuptools import setup, find_packages + +from setuptools import find_packages, setup from torch.utils.cpp_extension import BuildExtension, CUDAExtension -cwd = Path(os.path.dirname(os.path.abspath(__file__))) +ROOT = Path(__file__).resolve().parent +OPS_DIR = Path("wall_x") / "model" / "core" / "ops" +CSRC_DIR = OPS_DIR / "csrc" -nvcc_flags = [ - "-std=c++17", # NOTE: CUTLASS requires c++17 - "-DENABLE_BF16", # Enable BF16 for cuda_version >= 11 +PUBLIC_SCRIPTS = [ + "scripts/compute_norm_stats.py", + "scripts/draw_openloop_plot.py", + "scripts/fake_inference.py", + "scripts/infer_libero.py", + "scripts/merge_sharded_weights.py", + "scripts/merge_tokenizer.py", + "scripts/run_libero.sh", + "scripts/run_serving.sh", ] -env_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST", None) -if env_arch_list: - # Let PyTorch builder to choose device to target for. - device_capability = "" -else: - device_capability = torch.cuda.get_device_capability() - device_capability = f"{device_capability[0]}{device_capability[1]}" +def read_readme() -> str: + readme = ROOT / "README.md" + return readme.read_text(encoding="utf-8") if readme.exists() else "" -if device_capability: - nvcc_flags.extend( - [ - f"--generate-code=arch=compute_{device_capability},code=sm_{device_capability}", - f"-DGROUPED_GEMM_DEVICE_CAPABILITY={device_capability}", - ] + +def build_ext_modules(): + binding = CSRC_DIR / "binding.cu" + if not (ROOT / binding).exists(): + raise RuntimeError( + "CUDA operator sources are missing. The OSS export must include " + "wall_x/model/core/ops/csrc/binding.cu." + ) + + cuda_sources = [binding] + sorted( + path for path in (ROOT / CSRC_DIR).rglob("*.cu") if path.name != "binding.cu" ) + cuda_sources = [ + path if path == binding else path.relative_to(ROOT) for path in cuda_sources + ] + + return [ + CUDAExtension( + name="wall_x.model.core.ops._cuda_ext_bin", + sources=[str(path) for path in cuda_sources], + include_dirs=[str(CSRC_DIR), str(CSRC_DIR / "common")], + extra_compile_args={ + "cxx": ["-O3", "-std=c++17"], + "nvcc": ["-O3", "--use_fast_math", "-std=c++17"], + }, + ) + ] -ext_modules = [ - CUDAExtension( - "wallx_csrc", - [ - "csrc/ops.cu", - "csrc/dual_asym_grouped_gemm.cu", - "csrc/permute.cu", - "csrc/rope.cu", - "csrc/rope_index.cu", - "csrc/rot_pos.cu", - "csrc/window_index.cu", - ], - include_dirs=[f"{cwd}/3rdparty/cutlass/include/", f"{cwd}/csrc"], - extra_compile_args={ - "cxx": ["-fopenmp", "-fPIC", "-Wno-strict-aliasing"], - "nvcc": nvcc_flags, - }, - ) -] setup( name="wall_x", - version="1.0.1", - author="X2Robot Team", + version="1.1.0", + description="Training and inference code for WALL open-source embodied models.", + long_description=read_readme(), + long_description_content_type="text/markdown", + author="X-Square Robot", + url="https://github.com/X-Square-Robot/wall-x", + python_requires=">=3.10", + packages=find_packages(exclude=("tests", "tests.*")), + scripts=PUBLIC_SCRIPTS, + ext_modules=build_ext_modules(), + cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)}, classifiers=[ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3", - "License :: OSI Approved :: BSD License", - "Operating System :: Unix", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Topic :: Scientific/Engineering :: Artificial Intelligence", ], - packages=find_packages(), - ext_modules=ext_modules, - cmdclass={"build_ext": BuildExtension}, ) diff --git a/tests/test_ops_fallback.py b/tests/test_ops_fallback.py new file mode 100644 index 0000000..b6a36e5 --- /dev/null +++ b/tests/test_ops_fallback.py @@ -0,0 +1,32 @@ +"""Public fallback tests for operator proxies.""" + +import pytest +import torch + +from wall_x.model.core.ops import rmsnorm + + +def test_rmsnorm_pytorch_fallback_matches_reference(): + x = torch.randn(2, 8, dtype=torch.float32) + weight = torch.randn(8, dtype=torch.float32) + eps = 1e-6 + + out = rmsnorm.call_with_backend("pytorch", x, weight, eps) + + ref = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) + ref = ref * weight + torch.testing.assert_close(out, ref) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +def test_rmsnorm_cuda_inline_matches_pytorch_when_available(): + if "cuda_inline" not in rmsnorm.available_backends(): + pytest.skip("CUDA inline backend is not available") + + x = torch.randn(2, 16, device="cuda", dtype=torch.float16) + weight = torch.randn(16, device="cuda", dtype=torch.float16) + eps = 1e-6 + + out = rmsnorm.call_with_backend("cuda_inline", x, weight, eps) + ref = rmsnorm.call_with_backend("pytorch", x, weight, eps) + torch.testing.assert_close(out, ref, rtol=1e-3, atol=1e-3) diff --git a/train_qact.py b/train_qact.py deleted file mode 100644 index 1375f79..0000000 --- a/train_qact.py +++ /dev/null @@ -1,142 +0,0 @@ -import os -import json -import time -import yaml -import wandb -import accelerate -from argparse import ArgumentParser -from accelerate import ( - Accelerator, - DistributedDataParallelKwargs, - DataLoaderConfiguration, -) - -from wall_x.trainer.qwen_vl_act_trainer import QwenVlAct_Trainer - - -def setup_environment(): - """Set up environment variables for training.""" - os.environ["TOKENIZERS_PARALLELISM"] = "false" - - -def load_config(config_path): - """Load configuration from YAML file.""" - with open(config_path, "r") as f: - config = yaml.load(f, Loader=yaml.FullLoader) - - # Set model_type in data config if not already set - config["data"]["model_type"] = config.get("model_type") - - return config - - -def setup_accelerator(config): - """Initialize and configure the accelerator for distributed training.""" - print( - f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Preparing accelerator" - ) - - ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) - accelerator_dataloader_config = DataLoaderConfiguration(dispatch_batches=False) - - if config.get("FSDP2", False): - # Use Fully Sharded Data Parallel (FSDP) version 2 - fsdp_plugin = accelerate.utils.dataclasses.FullyShardedDataParallelPlugin( - fsdp_version=2, reshard_after_forward=True - ) - print("[INFO] Using FSDP version 2 for distributed training") - else: - fsdp_plugin = None - - if config.get("torch_compile", False): - # Use Torch Dynamo for compilation - dynamo_plugin = accelerate.utils.TorchDynamoPlugin( - backend="inductor", - mode="default", - fullgraph=False, - dynamic=False, - ) - print("[INFO] Using Torch Dynamo for compilation") - else: - dynamo_plugin = None - - accelerator = Accelerator( - kwargs_handlers=[ddp_kwargs], - mixed_precision="bf16", - fsdp_plugin=fsdp_plugin, - dynamo_plugin=dynamo_plugin, - dataloader_config=accelerator_dataloader_config, - gradient_accumulation_steps=config.get("gradient_accumulation_steps", 1), - ) - - print( - f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Accelerator initialization complete" - ) - - return accelerator - - -def setup_logging(config, accelerator): - """Set up logging with wandb for the main process.""" - if not accelerator.is_main_process: - return None - - # Create save directory if it doesn't exist - save_path = config["save_path"] - if not os.path.exists(save_path): - print(f"Save path {save_path} does not exist, creating directory.") - os.makedirs(save_path, exist_ok=True) - - print("Configuration:") - print("=" * 50) - print(json.dumps(config, indent=2, ensure_ascii=False)) - print("=" * 50) - - # Initialize wandb logger - logger = wandb.init( - project=config["log_project"], - name=config["log_name"], - save_code=False, - force=False, - ) - - return logger - - -def main(args): - """Main training function.""" - setup_environment() - - # Load configuration - config = load_config(args.config) - - # Set up accelerator - accelerator = setup_accelerator(config) - - # Set up logging - logger = setup_logging(config, accelerator) - - # Initialize trainer - trainer = QwenVlAct_Trainer( - config=config, - logger=logger, - accelerator=accelerator, - seed=args.seed, - data_config_path=args.config, - ) - - # Start training - trainer.fit() - - -if __name__ == "__main__": - parser = ArgumentParser(description="Training script for Wall-X model") - parser.add_argument( - "--config", type=str, required=True, help="Path to configuration YAML file" - ) - parser.add_argument( - "--seed", type=int, default=42, help="Random seed for reproducibility" - ) - - args = parser.parse_args() - main(args) diff --git a/wall_x/_vendor/__init__.py b/wall_x/_vendor/__init__.py new file mode 100644 index 0000000..e2e697e --- /dev/null +++ b/wall_x/_vendor/__init__.py @@ -0,0 +1,4 @@ +"""Vendored sister-package sources, bundled at export time so the exported +package is self-contained. Internal development uses ``pip install -e`` +against the live repos instead. +""" diff --git a/wall_x/_vendor/harrix/__init__.py b/wall_x/_vendor/harrix/__init__.py new file mode 100644 index 0000000..1db9b1b --- /dev/null +++ b/wall_x/_vendor/harrix/__init__.py @@ -0,0 +1,5 @@ +"""Public harrix inference and evaluation runtime.""" + +__version__ = "0.1.0" + +__all__ = ["__version__"] diff --git a/wall_x/_vendor/harrix/adapters/__init__.py b/wall_x/_vendor/harrix/adapters/__init__.py new file mode 100644 index 0000000..5791ac8 --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/__init__.py @@ -0,0 +1,6 @@ +"""Adapter package entry point. + +Importing ``variants`` triggers variant registration side effects. +""" + +from wall_x._vendor.harrix.adapters import variants # noqa: F401 diff --git a/wall_x/_vendor/harrix/adapters/base.py b/wall_x/_vendor/harrix/adapters/base.py new file mode 100644 index 0000000..68b7f91 --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/base.py @@ -0,0 +1,51 @@ +"""Common inference adapter abstraction. + +The adapter owns all architecture-specific model setup and exposes a single +``predict_batch`` entry point. Environment drivers pass payloads through without +interpreting the env-adapter schema. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +import numpy as np + +from wall_x._vendor.harrix.eval_config import EvalConfig + + +class BaseInferAdapter(ABC): + @abstractmethod + def __init__(self, cfg: EvalConfig) -> None: + """Load processors, model config, checkpoints, and normalizers. + + Subclasses should reject unsupported ``cfg.model.action_mode`` values + during construction. + """ + + @property + @abstractmethod + def chunk_horizon(self) -> int: + """Number of action steps returned by each ``predict_batch`` call.""" + + @property + @abstractmethod + def action_mode(self) -> str: + """Configured inference algorithm, fixed at construction time.""" + + @abstractmethod + def predict_batch(self, payloads: list[dict]) -> list[np.ndarray]: + """Run one batched inference call and return one chunk per payload. + + ``payloads[i]`` schema is defined by the env-adapter pair: + { + "observation": dict[str, np.ndarray], + # Env-defined ndarray bundle, for example LIBERO: + # {"eef_pos":(3,), "eef_axisangle":(3,), "gripper":(1,), + # "face_view":(H,W,3), "wrist_view":(H,W,3)} + "instruction": str, + "noise": np.ndarray | None, + # Flow may pass (chunk_horizon, action_dim); other modes + # may leave this as None. + } + """ diff --git a/wall_x/_vendor/harrix/adapters/qwen_vlact.py b/wall_x/_vendor/harrix/adapters/qwen_vlact.py new file mode 100644 index 0000000..a7f2118 --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/qwen_vlact.py @@ -0,0 +1,465 @@ +"""Shared adapter base for Qwen-VL action models. + +Variant-specific subclasses live in ``harrix.adapters.variants`` and provide +the Wall-X training adapter class via ``_training_adapter``. The shared base +handles checkpoint loading, normalizers, LIBERO observation encoding, prompt +construction, and flow-action inference. +""" + +from __future__ import annotations + +import copy +import logging + +import numpy as np +import torch +from PIL import Image + +from qwen_vl_utils.vision_process import smart_resize + +from wall_x._vendor.x2robot_utils.text_templates import ( + get_prologue_with_embodied_information, + preprocesser_call, +) + +from wall_x._vendor.harrix.adapters.base import BaseInferAdapter +from wall_x._vendor.harrix.envs.libero_common import decode_chunk, encode_proprio +from wall_x._vendor.harrix.eval_config import EvalConfig +from wall_x._vendor.harrix.utils.ckpt_load import ( + load_state_dict, + resolve_checkpoint_dir, +) +from wall_x._vendor.harrix.utils.normalizer import build_normalizers +from wall_x._vendor.harrix.utils.train_config import ( + build_data_config, + build_model_config, + load_train_config_with_ckpt_overlay, + normalize_train_config_for_inference, + register_data_backend, + resolve_state_bins, + resolve_use_state_string_representation, +) +from wall_x.trainer.trainer_utils import load_wallx_processors + + +_SUPPORTED_MODES = {"flow", "ar", "dllm", "vqa", "subtask"} +_IMPLEMENTED_MODES = {"flow"} + +# Special tokens used by the Qwen-VL action model. +_ROLE_START = "<|im_start|>" +_ROLE_END = "<|im_end|>" +_VISION_START = "<|vision_start|>" +_VISION_END = "<|vision_end|>" +_IMAGE_PAD = "<|image_pad|>" +_PROPRI = "<|propri|>" +_ACTION = "<|action|>" +logger = logging.getLogger(__name__) + + +_PUBLIC_CAMERA_LABELS = { + "face_view": "front view", + "right_wrist_view": "right wrist view", + "left_wrist_view": "left wrist view", +} + + +def _camera_label(cam_name: str) -> str: + return _PUBLIC_CAMERA_LABELS.get(cam_name, cam_name.replace("_", " ")) + + +def _normalizer_width(normalizer, norm_key: str) -> int | None: + if normalizer is None or norm_key not in normalizer.delta: + return None + return int(normalizer.delta[norm_key].shape[0]) + + +def _normalize_real_prefix(normalizer, tensor, dataset_names): + if normalizer is None: + return tensor + widths = [_normalizer_width(normalizer, name) for name in dataset_names] + if any(width is None for width in widths) or len(set(widths)) != 1: + return normalizer.normalize_data(tensor, dataset_names) + width = widths[0] + if tensor.shape[-1] == width: + return normalizer.normalize_data(tensor, dataset_names) + if tensor.shape[-1] < width: + raise ValueError( + f"normalizer width {width} exceeds tensor dim {tensor.shape[-1]}" + ) + out = tensor.clone() + out[..., :width] = normalizer.normalize_data( + tensor[..., :width], dataset_names + ) + out[..., width:] = 0 + return out + +class QwenVLActInferAdapter(BaseInferAdapter): + """Shared constructor and batched flow inference implementation.""" + + # ---- subclass hook ---- + + @classmethod + def _training_adapter(cls): + """Return the training-side ModelAdapter subclass.""" + raise NotImplementedError(f"{cls.__name__} must override _training_adapter()") + + # ---- ctor ---- + + def __init__(self, cfg: EvalConfig) -> None: + mode = cfg.model.action_mode + if mode not in _SUPPORTED_MODES: + raise ValueError( + f"{type(self).__name__}: unknown action_mode={mode!r}, " + f"supported={sorted(_SUPPORTED_MODES)}" + ) + if mode not in _IMPLEMENTED_MODES: + raise NotImplementedError( + f"{type(self).__name__}: action_mode={mode!r} is not implemented; " + "currently supported=" + f"{sorted(_IMPLEMENTED_MODES)}" + ) + self._action_mode = mode + + ta = self._training_adapter() + device = "cuda" + self._device = device + self._checkpoint_path = resolve_checkpoint_dir(cfg.model.checkpoint_path) + + # 1) train_config with checkpoint overlays, then data backend registration. + train_config = load_train_config_with_ckpt_overlay( + cfg.model.train_config_path, self._checkpoint_path + ) + train_config = normalize_train_config_for_inference( + train_config, cfg.model.train_config_path + ) + register_data_backend(train_config) + self._train_config = train_config + + # 2) normalizers first - action_tokenizer setup needs them during processor load. + normalizer_action, normalizer_propri, resolved_norm_key = build_normalizers( + self._checkpoint_path, train_config, cfg.model.norm_key + ) + self._normalizer_action = normalizer_action + self._normalizer_propri = normalizer_propri + self._norm_key = resolved_norm_key + if normalizer_propri is not None and resolved_norm_key in normalizer_propri.delta: + train_config["_libero_proprio_norm_dim"] = int( + normalizer_propri.delta[resolved_norm_key].shape[0] + ) + + # 3) data_config. Image resizing needs resolution/image_factor/min/max. + self._data_config = build_data_config(cfg.model.train_config_path, train_config) + + # 4) HF model config + ConfigClass = ta.config_class() + self._model_config = build_model_config( + ConfigClass, + self._checkpoint_path, + train_config, + cfg.model.train_config_path, + ) + + # 5) processor + tokenizer_mixin (may extend vocab via action_tokenizer) + procs = load_wallx_processors( + train_config, normalizer=normalizer_action, device=device + ) + self._processor = procs["processor"] + self._tokenizer_mixin = procs.get("tokenizer_mixin") + logger.info( + "processor vocab size: %s (action_tokenizer_type=%r)", + len(self._processor.tokenizer), + train_config.get("action_tokenizer_type"), + ) + + # 6) model + ModelClass = ta.inference_model_class() + self._model_class = ModelClass + model = ModelClass(self._model_config, self._processor, self._tokenizer_mixin) + model.resize_token_embeddings(len(self._processor.tokenizer)) + model.to_bfloat16_for_selected_params() + + # 7) checkpoint weights and model finalization. + state_dict = load_state_dict(self._checkpoint_path, ModelClass) + embed_key = "model.embed_tokens.weight" + if embed_key in state_dict: + ckpt_vocab = state_dict[embed_key].shape[0] + cur_vocab = model.model.embed_tokens.weight.shape[0] + if cur_vocab != ckpt_vocab: + logger.info( + "resize_token_embeddings from %d to %d to match checkpoint", + cur_vocab, + ckpt_vocab, + ) + model.resize_token_embeddings(ckpt_vocab) + msg = model.load_state_dict(state_dict, strict=False) + logger.info( + "%s load_state_dict: missing=%s unexpected=%s", + type(self).__name__, + len(msg.missing_keys), + len(msg.unexpected_keys), + ) + model.set_normalizer( + copy.deepcopy(normalizer_action), + copy.deepcopy(normalizer_propri), + ) + model.eval() + model.to(device) + model.to_bfloat16_for_selected_params() + self._model = model + + # 8) cached runtime fields + self._cam_names = list(cfg.model.cam_names) + self._action_horizon = int(cfg.model.action_horizon) + self._action_dim = sum(train_config["dof_config"].values()) + self._num_inference_timesteps = 10 + self._robot_id = "10000" + + # ---- BaseInferAdapter API ---- + + @property + def chunk_horizon(self) -> int: + return self._action_horizon + + @property + def action_mode(self) -> str: + return self._action_mode + + def predict_batch(self, payloads: list[dict]) -> list[np.ndarray]: + if self._action_mode == "flow": + return self._flow_batch(payloads) + raise NotImplementedError( + f"{type(self).__name__}.predict_batch: action_mode={self._action_mode!r} " + "is not dispatched" + ) + + # ---- batched flow inference ---- + + def _flow_batch(self, payloads: list[dict]) -> list[np.ndarray]: + # 1) Encode each observation into proprioception, masks, and views. + observations: list[dict] = [] + instructions: list[str] = [] + noises: list[np.ndarray | None] = [] + any_noise = False + for p in payloads: + observations.append( + encode_proprio( + p["observation"], self._train_config, self._action_horizon + ) + ) + instructions.append(p["instruction"]) + n = p.get("noise") + if n is not None: + any_noise = True + noises.append(n) + + # 2) Stack noise when provided; all-None delegates sampling to the model. + if any_noise: + if any(n is None for n in noises): + raise ValueError( + "payload noise must be provided for every payload or for none" + ) + batch_noise = torch.stack( + [torch.from_numpy(n).to(dtype=torch.float32) for n in noises], dim=0 + ) + else: + batch_noise = None + + # 3) Prompts and model inputs. + prefix_list, postfix_list = [], [] + for ins in instructions: + prefix, postfix = self._get_flow_prompt(ins) + prefix_list.append(prefix) + postfix_list.append(postfix) + batch_inputs = self._construct_model_input( + observations, prefix_list, postfix_list + ) + + # 4) Normalized zero action as the flow starting point. + padding = torch.zeros( + ( + len(batch_inputs["dataset_names"]), + 1, + self._action_dim, + ), + dtype=torch.float32, + ) + padding_action = _normalize_real_prefix( + self._normalizer_action, padding, batch_inputs["dataset_names"] + ).to(batch_inputs["input_ids"].device) + + # 5) Flow forward. + model_output = self._model.generate_flow_action( + action_horizon=self._action_horizon, + action_dim=self._action_dim, + num_inference_timesteps=self._num_inference_timesteps, + padding_action=padding_action, + noise=batch_noise, + **batch_inputs, + ) + + # 6) Decode one action chunk per payload. + predict_action = model_output["predict_action"] # (B, H, D_action) + if isinstance(predict_action, torch.Tensor): + predict_action = predict_action.detach().cpu().numpy() + return [ + decode_chunk(predict_action[i : i + 1], self._train_config) + for i in range(len(payloads)) + ] + + # ---- prompt template ---- + + def _get_flow_prompt(self, instruction: str) -> tuple[str, str]: + """Build the flow-action prompt.""" + if self._train_config["data"].get("use_embodied_system_prompt_ratio", 0) > 0: + robot_id = ( + self._robot_id if self._norm_key in ("x2_normal", "ex_normal") else 0 + ) + cam_name_mapping = {cn: cn for cn in self._cam_names} + prologue = get_prologue_with_embodied_information( + dataset_name=self._norm_key, + cam_mapping=cam_name_mapping, + robot_id=robot_id, + uid="", + config=self._data_config, + ) + else: + prologue = f"{_ROLE_START}system\nYou are a helpful assistant.{_ROLE_END}\n" + user_request = f"{_ROLE_START}user\nObservation:" + for cn in self._cam_names: + user_request += ( + f" {_camera_label(cn)}: " f"{_VISION_START}{_IMAGE_PAD}{_VISION_END}" + ) + user_request += "\nInstruction:" + text_prompt = ( + f"\nPredict the next action in robot action.\nProprioception: {_PROPRI}\n" + ) + user_message = f"{user_request} {instruction}{text_prompt}{_ROLE_END}\n" + assistant_message = f"{_ROLE_START}assistant\n" + flow_action = _ACTION * self._action_horizon + + prefix_text = prologue + user_message + assistant_message + postfix_text = flow_action + return prefix_text, postfix_text + + # ---- batched model input construction ---- + + def _construct_model_input( + self, + observations: list[dict], + prefix_list: list[str], + postfix_list: list[str], + ) -> dict: + """Build model inputs for the flow inference path.""" + batch_size = len(observations) + dataset_names = [self._norm_key] * batch_size + + # Proprioception and masks are prepared as ndarrays in encode_proprio. + agent_pos = torch.cat( + [torch.from_numpy(o["proprioception"]) for o in observations], dim=0 + ) + agent_pos_mask = torch.cat( + [torch.from_numpy(o["agent_pos_mask"]) for o in observations], dim=0 + ) + dof_mask = torch.cat( + [torch.from_numpy(o["dof_mask"]) for o in observations], dim=0 + ) + agent_pos = _normalize_real_prefix( + self._normalizer_propri, agent_pos, dataset_names + ) + + # Resize images per sample. + image_inputs: list[torch.Tensor] = [] + all_image_sizes: list[tuple[int, int]] = [] + for o in observations: + for cn in self._cam_names: + if cn not in o: + continue + tensor = self._resize_image(o[cn], cn) + image_inputs.append(tensor) + # Tensor (H, W, C) to PIL-compatible (W, H). + all_image_sizes.append((tensor.shape[1], tensor.shape[0])) + + inputs = preprocesser_call( + processor=self._processor, + prefix_text=prefix_list, + postfix_text=postfix_list, + images=image_inputs, + videos=None, + padding=True, + truncation=True, + return_tensors="pt", + max_length=1000, + pad_to_128_multiple=False, + pad_prefix_to_same_length=False, + norm_state=( + agent_pos + if resolve_use_state_string_representation(self._train_config) + else None + ), + agent_pos_mask=agent_pos_mask, + state_augmentation_prob=0.0, + state_drop_prob=0.0, + state_augmentation_ratio=0.0, + state_bins=resolve_state_bins(self._train_config), + inference_mode=True, + ) + + action_token_id = self._processor.tokenizer.convert_tokens_to_ids(_ACTION) + moe_token_types = inputs["input_ids"] == action_token_id + + extra = { + "proprioception": agent_pos.detach(), + "agent_pos_mask": agent_pos_mask, + "dof_mask": dof_mask, + "image_size": all_image_sizes, + "moe_token_types": moe_token_types, + "dataset_names": dataset_names, + } + inputs.update(extra) + return _move_to_device(inputs, self._device) + + # ---- image resizing ---- + + def _resize_image(self, img: np.ndarray, cam_name: str) -> torch.Tensor: + """Resize one image with the train-time image config.""" + if isinstance(img, np.ndarray): + pil = Image.fromarray(img) + elif isinstance(img, Image.Image): + pil = img + else: + raise ValueError(f"unsupported image type: {type(img)}") + orig_w, orig_h = pil.size + + target = self._data_config.resolution.get(cam_name, -1) + if target != -1: + if orig_w > orig_h: + new_w, new_h = target, int(target * orig_h / orig_w) + else: + new_h, new_w = target, int(target * orig_w / orig_h) + pil = pil.resize((new_w, new_h)) + + cur_w, cur_h = pil.size + resized_h, resized_w = smart_resize( + cur_h, + cur_w, + factor=self._data_config.image_factor, + min_pixels=self._data_config.min_pixels, + max_pixels=self._data_config.max_pixels, + ) + resized = pil.resize((resized_w, resized_h)) + return torch.from_numpy(np.array(resized)).to(self._device) + + +def _move_to_device(obj, device): + """Recursively move tensors inside common containers to ``device``.""" + from transformers import BatchFeature + + if isinstance(obj, torch.Tensor): + return obj.to(device) + if isinstance(obj, (dict, BatchFeature)): + return {k: _move_to_device(v, device) for k, v in obj.items()} + if isinstance(obj, list): + return [_move_to_device(v, device) for v in obj] + if isinstance(obj, tuple): + return tuple(_move_to_device(v, device) for v in obj) + return obj diff --git a/wall_x/_vendor/harrix/adapters/registry.py b/wall_x/_vendor/harrix/adapters/registry.py new file mode 100644 index 0000000..ddc0ec4 --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/registry.py @@ -0,0 +1,39 @@ +"""Inference adapter registry and factory.""" + +from __future__ import annotations + +from typing import Type + +from wall_x._vendor.harrix.adapters.base import BaseInferAdapter +from wall_x._vendor.harrix.eval_config import EvalConfig + + +ADAPTER_REGISTRY: dict[str, Type[BaseInferAdapter]] = {} + + +def register_adapter(name: str): + def deco(cls: Type[BaseInferAdapter]): + if name in ADAPTER_REGISTRY: + raise ValueError( + f"adapter {name!r} is already registered " + f"(cls={ADAPTER_REGISTRY[name].__name__})" + ) + ADAPTER_REGISTRY[name] = cls + return cls + + return deco + + +def build_adapter(cfg: EvalConfig) -> BaseInferAdapter: + arch = cfg.model.architecture + cls = ADAPTER_REGISTRY.get(arch) + if cls is None: + raise ValueError( + f"unknown architecture={arch!r}, registered={sorted(ADAPTER_REGISTRY)}" + ) + return cls(cfg) + + +def registered_architectures() -> list[str]: + """Return registered architecture names for diagnostics.""" + return sorted(ADAPTER_REGISTRY) diff --git a/wall_x/_vendor/harrix/adapters/variants/__init__.py b/wall_x/_vendor/harrix/adapters/variants/__init__.py new file mode 100644 index 0000000..7219446 --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/variants/__init__.py @@ -0,0 +1,15 @@ +"""Register inference adapter variants through import side effects.""" + +from __future__ import annotations + +import importlib +import pkgutil + + +for module_info in pkgutil.iter_modules(__path__): + if module_info.name.startswith("_"): + continue + try: + importlib.import_module(f"{__name__}.{module_info.name}") + except ImportError: + continue diff --git a/wall_x/_vendor/harrix/adapters/variants/qwen2_5.py b/wall_x/_vendor/harrix/adapters/variants/qwen2_5.py new file mode 100644 index 0000000..036409b --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/variants/qwen2_5.py @@ -0,0 +1,17 @@ +"""Qwen2.5 inference adapter registration. + +The shared inference logic lives in ``harrix.adapters.qwen_vlact``. This module +binds the Qwen2.5 variant key to Wall-X's training-side adapter class. +""" + +from wall_x._vendor.harrix.adapters.qwen_vlact import QwenVLActInferAdapter +from wall_x._vendor.harrix.adapters.registry import register_adapter + + +@register_adapter("qwen2_5") +class Qwen2_5InferAdapter(QwenVLActInferAdapter): + @classmethod + def _training_adapter(cls): + from wall_x.model.qact.qwen2_5.adapter import Qwen2_5Adapter + + return Qwen2_5Adapter diff --git a/wall_x/fusions/__init__.py b/wall_x/_vendor/harrix/drivers/__init__.py similarity index 100% rename from wall_x/fusions/__init__.py rename to wall_x/_vendor/harrix/drivers/__init__.py diff --git a/wall_x/_vendor/harrix/drivers/inproc/__init__.py b/wall_x/_vendor/harrix/drivers/inproc/__init__.py new file mode 100644 index 0000000..5813d0c --- /dev/null +++ b/wall_x/_vendor/harrix/drivers/inproc/__init__.py @@ -0,0 +1,197 @@ +"""In-process driver: DirectModelHandle + N x SubprocEnvHandle. + +The model lives in the driver process while simulator envs live in subprocesses. +Evaluation proceeds in chunk-level lockstep: + 1. Seed the driver process and construct the model handle. + 2. Start one subprocess env handle per worker. + 3. Claim a task-local frame from JobState. + 4. Reset envs, batch active observations, run model.predict_batch, then + fan out action chunks to env subprocesses. + 5. Complete all episodes in the frame and move to the next frame. +""" + +from __future__ import annotations + +import json +import logging +import os +import time + +import numpy as np + +from wall_x._vendor.harrix.eval_config import EvalConfig +from wall_x._vendor.harrix.drivers.inproc.env_handle import SubprocEnvHandle +from wall_x._vendor.harrix.drivers.inproc.model_handle import DirectModelHandle +from wall_x._vendor.harrix.drivers.job_state import JobState + +logger = logging.getLogger(__name__) + + +def run(cfg: EvalConfig) -> None: + import wall_x._vendor.harrix.envs # noqa: F401 trigger env register + import wall_x._vendor.harrix.adapters # noqa: F401 trigger adapter register + from wall_x._vendor.harrix.envs.registry import enumerate_episodes_for + from wall_x._vendor.harrix.utils.seed import set_seed_everywhere + + # 1) Seed the driver process. + set_seed_everywhere(cfg.env.seed) + + # 2) Build DirectModelHandle; the model is loaded in the driver process. + logger.info("Constructing DirectModelHandle in the driver process") + t_model = time.time() + model_handle = DirectModelHandle(cfg) + logger.info("Model loaded in %.1fs", time.time() - t_model) + + # 3) Start env subprocesses. + logger.info("Starting %s SubprocEnvHandle(s)", cfg.runtime.num_workers) + env_handles = [ + SubprocEnvHandle(cfg, worker_id=i) for i in range(cfg.runtime.num_workers) + ] + + # 4) JobState runs in frame-sync mode for lockstep evaluation. + os.makedirs(cfg.runtime.log_dir, exist_ok=True) + log_path = os.path.join(cfg.runtime.log_dir, "state.jsonl") + report_path = os.path.join(cfg.runtime.log_dir, "report.json") + episodes = enumerate_episodes_for(cfg) + logger.info("env.type=%r; scheduled %s episodes", cfg.env.type, len(episodes)) + state = JobState( + episodes, + log_path, + batch_sync_mode=True, + batch_size=cfg.runtime.num_workers, + ) + + # 5) Main loop, one frame at a time. + t0 = time.time() + frame_idx = 0 + while not state.is_drained(): + frame_eps = state.claim_frame() + if not frame_eps: + # Previous frame still has in-flight episodes. + time.sleep(0.05) + continue + + results = _run_frame(frame_eps, model_handle, env_handles, cfg, frame_idx) + for ep, res in zip(frame_eps, results): + if "_error" in res: + state.fail(ep, str(res["_error"])) + else: + state.complete(ep, res) + frame_idx += 1 + + elapsed = time.time() - t0 + logger.info("All episodes finished in %.1fs (%.1f min)", elapsed, elapsed / 60) + + state.dump_final(report_path) + with open(report_path) as f: + report = json.load(f) + overall = report["overall"] + logger.info( + "attempted=%s, successes=%s, success_rate=%.2f%%, failed=%s", + overall["attempted"], + overall["successes"], + overall["success_rate"] * 100, + overall["failed"], + ) + logger.info("Report: %s", report_path) + logger.info("State log: %s", log_path) + + for h in env_handles: + h.shutdown() + model_handle.shutdown() + + +def _run_frame( + frame_eps: list, + model_handle: DirectModelHandle, + env_handles: list, + cfg: EvalConfig, + frame_idx: int, +) -> list[dict]: + """Run one task-local frame with chunk-level lockstep. + + Active workers are batched together at each chunk boundary. Workers that + already finished no longer participate in later forwards. + """ + from wall_x._vendor.harrix.envs.libero_common import encode_raw_obs + + n = len(frame_eps) + t_frame_start = time.time() + + # ---- a) reset: fan out, then gather ---- + for i in range(n): + env_handles[i].submit_reset(tuple(frame_eps[i])) + initials = [env_handles[i].wait_reset() for i in range(n)] + + obs_list = [r["obs"] for r in initials] + instr_list = [r["instruction"] for r in initials] + status = [ + { + "done": False, + "success": False, + "steps": 0, + "task_desc": initials[i].get("task_desc", ""), + } + for i in range(n) + ] + + max_rounds = cfg.env.libero.max_infer_times + + # ---- b) chunk lockstep ---- + for round_idx in range(max_rounds): + active = [i for i in range(n) if not status[i]["done"]] + if not active: + break + + payloads = [ + { + "observation": encode_raw_obs(obs_list[i]), + "instruction": instr_list[i], + "noise": None, + } + for i in active + ] + chunks = model_handle.predict_batch(payloads) + + # fan-out submit + for k, i in enumerate(active): + env_handles[i].submit_execute_chunk(chunks[k]) + # gather + for k, i in enumerate(active): + try: + r = env_handles[i].wait_execute_chunk() + except Exception as e: + status[i]["_error"] = str(e) + status[i]["done"] = True + continue + obs_list[i] = r["obs"] + status[i]["steps"] += r["steps"] + if r["done"]: + status[i]["done"] = True + status[i]["success"] = True + + for i in range(n): + env_handles[i].submit_finalize_episode(status[i]["success"]) + for i in range(n): + env_handles[i].wait_finalize_episode() + + elapsed_frame = time.time() - t_frame_start + logger.info( + "frame=%s n=%s succ=%s/%s elapsed=%.1fs", + frame_idx, + n, + sum(s["success"] for s in status), + n, + elapsed_frame, + ) + + return [ + { + "success": s["success"], + "steps": s["steps"], + "elapsed_sec": round(elapsed_frame / max(1, n), 3), + "task_desc": s["task_desc"], + **({"_error": s["_error"]} if "_error" in s else {}), + } + for s in status + ] diff --git a/wall_x/_vendor/harrix/drivers/inproc/env_handle.py b/wall_x/_vendor/harrix/drivers/inproc/env_handle.py new file mode 100644 index 0000000..64bbec5 --- /dev/null +++ b/wall_x/_vendor/harrix/drivers/inproc/env_handle.py @@ -0,0 +1,152 @@ +"""Subprocess env handle used by the in-process driver. + +The model remains in the driver process. Each env subprocess receives reset and +execute-chunk commands through a pipe. +""" + +from __future__ import annotations + +import multiprocessing as mp +import os +import random + +import numpy as np + + +def _subproc_main(cfg, worker_id, child_conn): + """Env subprocess entry point.""" + # Spawned children inherit env vars, but set these explicitly for launchers + # that did not configure them. Robosuite validates the EGL id against the + # CUDA_VISIBLE_DEVICES environment string, so keep the same visible id here. + cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") or "0" + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices + # EGL device index is always 0 after CUDA_VISIBLE_DEVICES remapping. + os.environ["MUJOCO_EGL_DEVICE_ID"] = os.environ.get("MUJOCO_EGL_DEVICE_ID") or "0" + + # Env workers run robosuite only. Do not import torch here; otherwise many + # subprocesses may initialize CUDA contexts and compete with the driver model. + seed = cfg.env.seed + worker_id + np.random.seed(seed) + random.seed(seed) + os.environ["PYTHONHASHSEED"] = str(seed) + + import wall_x._vendor.harrix.envs # noqa: F401 trigger register + from wall_x._vendor.harrix.envs.registry import build_env + + env = build_env(cfg, worker_id) + try: + while True: + try: + cmd, args = child_conn.recv() + except EOFError: + break + try: + if cmd == "reset_episode": + result = env.reset_episode(tuple(args)) + elif cmd == "execute_chunk": + result = env.execute_chunk(args) + elif cmd == "finalize_episode": + env.finalize_episode(bool(args)) + result = None + elif cmd == "shutdown": + env.shutdown() + child_conn.send(("ok", None)) + break + else: + raise ValueError(f"unknown cmd {cmd!r}") + child_conn.send(("ok", result)) + except Exception as e: + import traceback + + child_conn.send(("err", f"{e}\n{traceback.format_exc()}")) + finally: + try: + child_conn.close() + except Exception: + pass + + +class SubprocEnvHandle: + """Synchronous pipe wrapper around one env subprocess. + + Usage: + h.submit_reset(ep_id) + ... + h.wait_reset() # -> {"obs", "instruction", "task_desc"} + h.submit_execute_chunk(actions) + ... + h.wait_execute_chunk() # -> {"obs", "done", "steps"} + """ + + def __init__(self, cfg, worker_id: int): + # Use spawn instead of fork because the driver may already hold a CUDA + # context for the model. + ctx = mp.get_context("spawn") + self._parent_conn, child_conn = ctx.Pipe() + self._proc = ctx.Process( + target=_subproc_main, + args=(cfg, worker_id, child_conn), + daemon=False, + name=f"infer-subproc-w{worker_id}", + ) + self._proc.start() + child_conn.close() + self._has_pending = False + + def _send(self, cmd: str, args): + self._parent_conn.send((cmd, args)) + self._has_pending = True + + def _recv(self): + if not self._has_pending: + raise RuntimeError("no pending request to wait for") + status, payload = self._parent_conn.recv() + self._has_pending = False + if status == "err": + raise RuntimeError(f"subproc env error (w={self._proc.name}): {payload}") + return payload + + def submit_reset(self, ep_id): + self._send("reset_episode", ep_id) + + def wait_reset(self): + return self._recv() + + def submit_execute_chunk(self, actions): + self._send("execute_chunk", np.asarray(actions, dtype=np.float32)) + + def wait_execute_chunk(self): + return self._recv() + + def submit_finalize_episode(self, success: bool): + self._send("finalize_episode", success) + + def wait_finalize_episode(self): + return self._recv() + + def finalize_episode(self, success: bool): + """Blocking convenience wrapper.""" + self.submit_finalize_episode(success) + return self.wait_finalize_episode() + + def reset_episode(self, ep_id): + """Blocking convenience wrapper.""" + self.submit_reset(ep_id) + return self.wait_reset() + + def execute_chunk(self, actions): + """Blocking convenience wrapper.""" + self.submit_execute_chunk(actions) + return self.wait_execute_chunk() + + def shutdown(self) -> None: + try: + self._send("shutdown", None) + self._recv() + except Exception: + pass + if self._proc.is_alive(): + self._proc.join(timeout=5) + if self._proc.is_alive(): + self._proc.terminate() + self._proc.join(timeout=2) diff --git a/wall_x/_vendor/harrix/drivers/inproc/model_handle.py b/wall_x/_vendor/harrix/drivers/inproc/model_handle.py new file mode 100644 index 0000000..ada020b --- /dev/null +++ b/wall_x/_vendor/harrix/drivers/inproc/model_handle.py @@ -0,0 +1,25 @@ +"""In-process model handle. + +The model is constructed in the driver process and calls the adapter directly. +""" + +from __future__ import annotations + + +class DirectModelHandle: + def __init__(self, cfg): + # Trigger adapter registration. + import wall_x._vendor.harrix.adapters # noqa: F401 + from wall_x._vendor.harrix.adapters.registry import build_adapter + + self._adapter = build_adapter(cfg) + + @property + def chunk_horizon(self) -> int: + return self._adapter.chunk_horizon + + def predict_batch(self, payloads): + return self._adapter.predict_batch(payloads) + + def shutdown(self) -> None: + pass diff --git a/wall_x/_vendor/harrix/drivers/job_state.py b/wall_x/_vendor/harrix/drivers/job_state.py new file mode 100644 index 0000000..0a9b11c --- /dev/null +++ b/wall_x/_vendor/harrix/drivers/job_state.py @@ -0,0 +1,343 @@ +"""Episode work queue with JSONL persistence. + +Drivers enumerate episode ids into ``pending``; env handles claim work and +return results through ``complete`` or ``fail``. Every state transition is +appended to ``state.jsonl``, so completed episodes can be skipped on restart. + +Two scheduling modes are supported: +- FIFO: workers claim the next pending episode. +- batch_sync: episodes are grouped into task-local frames. A new frame is not + released until the previous frame is complete, which gives deterministic + lockstep batches at the cost of possible idle workers. +""" + +from __future__ import annotations + +import json +import os +import threading +import time +from typing import Optional + + +class JobState: + def __init__( + self, + all_episodes: list[tuple], + log_path: str, + batch_sync_mode: bool = False, + batch_size: int = 1, + ): + """ + all_episodes: list of (suite_name, task_idx, ep_idx) tuples + log_path: state.jsonl path; existing completed episodes are skipped + batch_sync_mode: enable task-local frame barriers + batch_size: number of episodes per frame + """ + self._lock = threading.Lock() + self._log_path = log_path + self._log_fh = None + self._batch_sync_mode = bool(batch_sync_mode) + self._batch_size = int(batch_size) + + completed_set = ( + self._load_completed(log_path) if os.path.exists(log_path) else set() + ) + remaining: list[tuple] = [ + tuple(ep) for ep in all_episodes if tuple(ep) not in completed_set + ] + + self._in_progress: dict[tuple, int] = {} + self._completed: dict[tuple, dict] = {} + self._failed: dict[tuple, str] = {} + + if self._batch_sync_mode: + self._frames: list[dict] = self._build_frames(remaining, self._batch_size) + self._cur_frame_idx: int = 0 + self._cur_frame_inflight: int = 0 + self._pending = None + else: + self._frames = [] + self._cur_frame_idx = 0 + self._cur_frame_inflight = 0 + self._pending: list[tuple] = remaining + + os.makedirs(os.path.dirname(log_path) or ".", exist_ok=True) + self._log_fh = open(log_path, "a") + + log_entry = { + "event": "session_start", + "pending": ( + len(self._pending) + if not self._batch_sync_mode + else sum(len(f["eps"]) for f in self._frames) + ), + "skipped_completed": len(completed_set), + "batch_sync_mode": self._batch_sync_mode, + } + if self._batch_sync_mode: + log_entry["num_frames"] = len(self._frames) + log_entry["batch_size"] = self._batch_size + self._append_log(log_entry) + + @staticmethod + def _build_frames(remaining: list[tuple], batch_size: int) -> list[dict]: + """Group episodes by ``(suite, task_idx)`` into fixed-size frames. + + Partial tail frames are kept. + """ + frames: list[dict] = [] + cur_key: Optional[tuple] = None + cur_bucket: list[tuple] = [] + for ep in remaining: + key = (ep[0], ep[1]) + if key != cur_key and cur_bucket: + for i in range(0, len(cur_bucket), batch_size): + frames.append( + {"eps": cur_bucket[i : i + batch_size], "next_slot": 0} + ) + cur_bucket = [] + cur_key = key + cur_bucket.append(ep) + if cur_bucket: + for i in range(0, len(cur_bucket), batch_size): + frames.append({"eps": cur_bucket[i : i + batch_size], "next_slot": 0}) + return frames + + @staticmethod + def _load_completed(log_path: str) -> set[tuple]: + completed = set() + with open(log_path, "r") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except json.JSONDecodeError: + continue + if rec.get("event") == "completed" and "ep_id" in rec: + completed.add(tuple(rec["ep_id"])) + return completed + + def _append_log(self, extra: dict): + rec = {"ts": time.strftime("%Y-%m-%dT%H:%M:%S"), **extra} + if "ep_id" in rec and isinstance(rec["ep_id"], tuple): + rec["ep_id"] = list(rec["ep_id"]) + self._log_fh.write(json.dumps(rec, ensure_ascii=False) + "\n") + self._log_fh.flush() + + # ---- claim API ---- + + def claim(self, worker_id: int) -> Optional[list]: + """Claim one episode id. + + FIFO returns ``None`` when the pending queue is empty. In batch-sync + mode, ``None`` can also mean the current frame has been fully issued but + still has in-flight episodes. + """ + with self._lock: + if self._batch_sync_mode: + return self._claim_sync(worker_id) + return self._claim_fifo(worker_id) + + def claim_frame(self) -> list[list]: + """Claim one complete frame for lockstep in-process evaluation. + + In FIFO mode this returns up to ``batch_size`` pending episodes. + """ + with self._lock: + if self._batch_sync_mode: + return self._claim_frame_sync() + return self._claim_frame_fifo() + + def _claim_fifo(self, worker_id: int) -> Optional[list]: + if not self._pending: + return None + ep = self._pending.pop(0) + self._in_progress[ep] = worker_id + self._append_log({"event": "claimed", "ep_id": ep, "worker": worker_id}) + return list(ep) + + def _claim_sync(self, worker_id: int) -> Optional[list]: + while self._cur_frame_idx < len(self._frames): + frame = self._frames[self._cur_frame_idx] + if frame["next_slot"] < len(frame["eps"]): + ep = frame["eps"][frame["next_slot"]] + frame["next_slot"] += 1 + self._cur_frame_inflight += 1 + self._in_progress[ep] = worker_id + self._append_log( + { + "event": "claimed", + "ep_id": ep, + "worker": worker_id, + "frame": self._cur_frame_idx, + } + ) + return list(ep) + if self._cur_frame_inflight > 0: + return None + self._append_log( + { + "event": "frame_done", + "frame": self._cur_frame_idx, + "size": len(frame["eps"]), + } + ) + self._cur_frame_idx += 1 + return None + + def _claim_frame_sync(self) -> list[list]: + # Advance completed frames. + while self._cur_frame_idx < len(self._frames): + frame = self._frames[self._cur_frame_idx] + if frame["next_slot"] < len(frame["eps"]): + break + if self._cur_frame_inflight > 0: + # The previous frame still has in-flight episodes. + return [] + self._append_log( + { + "event": "frame_done", + "frame": self._cur_frame_idx, + "size": len(frame["eps"]), + } + ) + self._cur_frame_idx += 1 + if self._cur_frame_idx >= len(self._frames): + return [] + frame = self._frames[self._cur_frame_idx] + out: list[list] = [] + while frame["next_slot"] < len(frame["eps"]): + ep = frame["eps"][frame["next_slot"]] + frame["next_slot"] += 1 + self._cur_frame_inflight += 1 + self._in_progress[ep] = -1 + self._append_log( + {"event": "claimed", "ep_id": ep, "frame": self._cur_frame_idx} + ) + out.append(list(ep)) + return out + + def _claim_frame_fifo(self) -> list[list]: + if not self._pending: + return [] + n = min(self._batch_size, len(self._pending)) + out: list[list] = [] + for _ in range(n): + ep = self._pending.pop(0) + self._in_progress[ep] = -1 + self._append_log({"event": "claimed", "ep_id": ep}) + out.append(list(ep)) + return out + + # ---- complete / fail ---- + + def complete(self, ep_id: list, result: dict) -> None: + ep = tuple(ep_id) + with self._lock: + if ep in self._in_progress: + self._in_progress.pop(ep, None) + if self._batch_sync_mode: + self._cur_frame_inflight = max(0, self._cur_frame_inflight - 1) + self._completed[ep] = result + self._append_log({"event": "completed", "ep_id": ep, **result}) + + def fail(self, ep_id: list, error: str) -> None: + ep = tuple(ep_id) + with self._lock: + if ep in self._in_progress: + self._in_progress.pop(ep, None) + if self._batch_sync_mode: + self._cur_frame_inflight = max(0, self._cur_frame_inflight - 1) + self._failed[ep] = error + self._append_log({"event": "failed", "ep_id": ep, "error": error}) + + # ---- status queries ---- + + def get_frame_inflight(self) -> int: + with self._lock: + return self._cur_frame_inflight if self._batch_sync_mode else 0 + + def is_drained(self) -> bool: + with self._lock: + if self._batch_sync_mode: + return ( + self._cur_frame_idx >= len(self._frames) + and self._cur_frame_inflight == 0 + ) + return len(self._pending) == 0 and len(self._in_progress) == 0 + + def progress(self) -> dict: + with self._lock: + base = { + "in_progress": len(self._in_progress), + "completed": len(self._completed), + "failed": len(self._failed), + } + if self._batch_sync_mode: + pending = sum( + len(f["eps"]) - f["next_slot"] + for f in self._frames[self._cur_frame_idx :] + ) + base["pending"] = pending + base["frame"] = f"{self._cur_frame_idx}/{len(self._frames)}" + base["frame_inflight"] = self._cur_frame_inflight + else: + base["pending"] = len(self._pending) + return base + + def dump_final(self, report_path: str) -> None: + with self._lock: + per_task: dict[tuple, dict] = {} + for ep, result in self._completed.items(): + key = (ep[0], ep[1]) + d = per_task.setdefault( + key, {"attempted": 0, "successes": 0, "steps": []} + ) + d["attempted"] += 1 + if result.get("success"): + d["successes"] += 1 + if "steps" in result: + d["steps"].append(result["steps"]) + for ep in self._failed: + key = (ep[0], ep[1]) + d = per_task.setdefault( + key, {"attempted": 0, "successes": 0, "steps": []} + ) + d["attempted"] += 1 + + total_attempted = sum(d["attempted"] for d in per_task.values()) + total_successes = sum(d["successes"] for d in per_task.values()) + overall_rate = total_successes / max(1, total_attempted) + + report = { + "overall": { + "attempted": total_attempted, + "successes": total_successes, + "success_rate": overall_rate, + "failed": len(self._failed), + }, + "per_task": { + f"{suite}_t{task_idx}": { + **d, + "success_rate": d["successes"] / max(1, d["attempted"]), + "avg_steps": ( + (sum(d["steps"]) / max(1, len(d["steps"]))) + if d["steps"] + else None + ), + } + for (suite, task_idx), d in sorted(per_task.items()) + }, + } + with open(report_path, "w") as f: + json.dump(report, f, indent=2, ensure_ascii=False) + self._append_log( + { + "event": "session_end", + "report_path": report_path, + "overall_success_rate": overall_rate, + } + ) diff --git a/wall_x/_vendor/harrix/envs/__init__.py b/wall_x/_vendor/harrix/envs/__init__.py new file mode 100644 index 0000000..605263f --- /dev/null +++ b/wall_x/_vendor/harrix/envs/__init__.py @@ -0,0 +1,8 @@ +"""Environment registrations for harrix.""" + +_LIBERO_IMPORT_ERROR = None + +try: + from wall_x._vendor.harrix.envs import libero # noqa: F401 +except ModuleNotFoundError as exc: + _LIBERO_IMPORT_ERROR = exc diff --git a/wall_x/_vendor/harrix/envs/base.py b/wall_x/_vendor/harrix/envs/base.py new file mode 100644 index 0000000..5c3aef3 --- /dev/null +++ b/wall_x/_vendor/harrix/envs/base.py @@ -0,0 +1,119 @@ +"""Base environment abstraction. + +Two execution granularities are supported: + +- Episode-level ``run_episode``: caller supplies a predict callback and the env + owns the full episode loop. +- Chunk-level ``reset_episode`` + ``execute_chunk``: caller runs the model + between chunks and feeds action chunks back to the env. + +Subclasses must implement the chunk-level primitives. The default episode loop +is built on top of those primitives. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Callable + +import numpy as np + +from wall_x._vendor.harrix.eval_config import EvalConfig + + +class BaseEnv(ABC): + + @abstractmethod + def __init__(self, cfg: EvalConfig, worker_id: int) -> None: + """Perform env-specific setup.""" + + @classmethod + @abstractmethod + def enumerate_episodes(cls, cfg: EvalConfig) -> list[tuple]: + """Return episode ids to seed JobState before env instances are built. + + JobState treats the returned tuples as opaque ids. + """ + + @property + @abstractmethod + def robot_spec(self) -> dict: + """Return robot metadata for driver/checkpoint validation. + + Fields: + - dof_layout: dict[str, int] + - cam_names: list[str] + - norm_key: str + """ + + # ---- chunk-level primitives ---- + + @abstractmethod + def reset_episode(self, ep_id: tuple) -> dict: + """Start a new episode and return the first fresh observation. + + Returns {"obs": dict, "instruction": str, "task_desc": str (optional)}. + """ + + @abstractmethod + def execute_chunk(self, actions: np.ndarray) -> dict: + """Execute one action chunk with shape ``(H, action_dim)``. + + Returns {"obs": dict, "done": bool, "steps": int}. + """ + + def shutdown(self) -> None: + """Release resources. Subclasses may override.""" + + # ---- episode-level default implementation ---- + + def run_episode( + self, + ep_id: tuple, + predict: Callable[[dict, str, int], np.ndarray], + ) -> dict: + """Episode loop built from reset, predict, and execute_chunk. + + ``predict(observation, instruction, step)`` returns one action chunk. + The env remains unaware of the model transport. + """ + import time + + from wall_x._vendor.harrix.envs.libero_common import encode_raw_obs + + t_ep_start = time.time() + initial = self.reset_episode(ep_id) + obs = initial["obs"] + instruction = initial["instruction"] + task_desc = initial.get("task_desc", "") + + max_rounds = self._max_infer_rounds() + success = False + steps_total = 0 + for round_idx in range(max_rounds): + encoded = encode_raw_obs(obs) + chunk = predict(encoded, instruction, round_idx) + result = self.execute_chunk(chunk) + obs = result["obs"] + steps_total += result["steps"] + if result["done"]: + success = True + break + + self.finalize_episode(success) + return { + "success": bool(success), + "steps": steps_total, + "elapsed_sec": round(time.time() - t_ep_start, 3), + "task_desc": task_desc, + } + + def finalize_episode(self, success: bool) -> None: + """Hook for env-specific cleanup after an episode (e.g. save rollouts).""" + + def _max_infer_rounds(self) -> int: + """Return the maximum number of model chunks for one episode.""" + raise NotImplementedError( + f"{type(self).__name__} must override _max_infer_rounds when using " + "the default run_episode implementation" + ) diff --git a/wall_x/_vendor/harrix/envs/libero.py b/wall_x/_vendor/harrix/envs/libero.py new file mode 100644 index 0000000..885d51e --- /dev/null +++ b/wall_x/_vendor/harrix/envs/libero.py @@ -0,0 +1,342 @@ +"""LIBERO environment implementation. + +The chunk-level API mirrors standard LIBERO rollout semantics: +- reset_episode enables rendering, sets initial state, and performs warmup steps. +- execute_chunk may skip intermediate image rendering but re-enables rendering + before returning an observation. +""" + +from __future__ import annotations + +import gc +import logging +import os +from typing import Optional + +import numpy as np + +from wall_x._vendor.harrix.envs.base import BaseEnv +from wall_x._vendor.harrix.envs.libero_common import ( + get_rollout_frame, + model_action_to_libero_env, + save_rollout_video, +) +from wall_x._vendor.harrix.envs.libero_sim import ( + create_libero_engine, + find_image_observables, + get_instruction, + get_libero_dummy_action, + get_task_suite, + load_initial_states, + pick_initial_state, + resolve_task_info, + set_render_enabled, +) +from wall_x._vendor.harrix.envs.registry import register_env +from wall_x._vendor.harrix.eval_config import EvalConfig + +logger = logging.getLogger(__name__) + + +@register_env("libero") +class LiberoEnv(BaseEnv): + + def __init__(self, cfg: EvalConfig, worker_id: int) -> None: + self._cfg = cfg + self._libero_cfg = cfg.env.libero + self._worker_id = worker_id + self._seed = cfg.env.seed + self._task_suite_name = self._libero_cfg.task_suite_name + + # Task-suite metadata. + self._task_suite = get_task_suite(self._task_suite_name) + self._num_tasks = self._task_suite.n_tasks + self._custom_initial_states = load_initial_states( + self._libero_cfg.initial_states_path + ) + + # Robosuite engine, lazily rebuilt on task changes. + self._libero_env = None + self._current_task_id: Optional[int] = None + self._rebuild_env_per_episode = self._libero_cfg.rebuild_env_per_episode + + # Render-skip state. + self._skip_intermediate_render = self._libero_cfg.skip_intermediate_render + self._force_render_task_ids = set( + self._libero_cfg.force_render_task_indices or [] + ) + self._effective_skip_render = self._skip_intermediate_render + self._image_obs: list = [] + self._chunk_granular_render_toggle = ( + self._libero_cfg.chunk_granular_render_toggle + ) + self._render_enabled_state = False + + # Optional bit-alignment dump for debugging. + self._bit_dump_dir = os.environ.get("WALLX_BIT_DUMP_DIR", "").strip() or None + if self._bit_dump_dir: + os.makedirs(self._bit_dump_dir, exist_ok=True) + + rollout_dir = (self._libero_cfg.rollout_dir or "").strip() + if not rollout_dir: + rollout_dir = os.environ.get("WALLX_ROLLOUT_DIR", "").strip() + if rollout_dir and os.environ.get("WALLX_DISABLE_ROLLOUT", "0") == "1": + rollout_dir = "" + self._rollout_dir = rollout_dir or None + self._rollout_fps = int(self._libero_cfg.rollout_fps) + if self._rollout_dir: + worker_subdir = f"worker{worker_id}" if cfg.runtime.num_workers > 1 else "" + self._rollout_dir = os.path.join( + self._rollout_dir, + self._task_suite_name, + worker_subdir, + ) + os.makedirs(self._rollout_dir, exist_ok=True) + logger.info("Rollout MP4 saving enabled: %s", self._rollout_dir) + + # Per-episode state consumed by execute_chunk. + self._current_ep: Optional[tuple] = None + self._current_task_desc: str = "" + self._current_instruction: str = "" + self._chunk_counter_in_ep: int = 0 + self._last_obs_for_dump: Optional[dict] = None + self._replay_images: list[np.ndarray] = [] + self._rollout_saved = False + + # ---- BaseEnv API ---- + + @classmethod + def enumerate_episodes(cls, cfg: EvalConfig) -> list[tuple]: + libero_cfg = cfg.env.libero + suite = libero_cfg.task_suite_name + + if libero_cfg.task_indices is not None: + task_indices = [int(x) for x in libero_cfg.task_indices] + else: + ts = get_task_suite(suite) + task_indices = list(range(ts.n_tasks)) + + eps = [] + for tid in task_indices: + for epi in range(libero_cfg.num_trials_per_task): + eps.append((suite, tid, epi)) + return eps + + @property + def robot_spec(self) -> dict: + """Return robot spec for driver/adapter validation.""" + from wall_x._vendor.harrix.utils.train_config import ( + load_train_config_with_ckpt_overlay, + ) + + train_cfg = load_train_config_with_ckpt_overlay( + self._cfg.model.train_config_path, + self._cfg.model.checkpoint_path, + ) + return { + "dof_layout": train_cfg.get("dof_config", {}), + "cam_names": list(self._cfg.model.cam_names), + "norm_key": self._cfg.model.norm_key, + } + + def _max_infer_rounds(self) -> int: + return self._libero_cfg.max_infer_times + + def reset_episode(self, ep_id: tuple) -> dict: + suite, task_id, ep_idx = ep_id + if suite != self._task_suite_name: + raise ValueError( + f"env bound to suite={self._task_suite_name}, got ep with suite={suite}" + ) + + need_rebuild = self._rebuild_env_per_episode or self._current_task_id != task_id + if need_rebuild: + self._rebuild_env(task_id) + + task_desc, default_states = resolve_task_info(self._task_suite, task_id) + init_state = pick_initial_state( + self._libero_cfg.initial_states_path, + self._custom_initial_states, + task_desc, + default_states, + ep_idx, + ) + + if not need_rebuild: + self._libero_env.reset() + obs = self._libero_env.set_init_state(init_state) + if obs is None: + raise RuntimeError("set_init_state returned None") + + set_render_enabled(self._image_obs, True) + self._render_enabled_state = True + + dummy_action = get_libero_dummy_action() + for _ in range(10): + obs, _, _, _ = self._libero_env.step(dummy_action) + + self._current_ep = (task_id, ep_idx) + self._current_task_desc = task_desc + self._current_instruction = get_instruction(task_desc) + self._chunk_counter_in_ep = 0 + self._last_obs_for_dump = obs + self._begin_rollout_capture(obs) + + return { + "obs": obs, + "instruction": self._current_instruction, + "task_desc": task_desc, + } + + def execute_chunk(self, actions: np.ndarray) -> dict: + actions = np.asarray(actions, dtype=np.float32) + H = actions.shape[0] + + if self._bit_dump_dir and self._current_ep is not None: + self._dump_chunk_npz( + (self._task_suite_name, *self._current_ep), + self._chunk_counter_in_ep, + self._last_obs_for_dump or {}, + actions, + ) + self._chunk_counter_in_ep += 1 + + skip_render = self._chunk_skip_render() + # Disable rendering at the chunk start when render-skip is enabled. + if skip_render and self._image_obs: + if self._render_enabled_state: + set_render_enabled(self._image_obs, False) + self._render_enabled_state = False + + last_obs = None + done = False + steps = 0 + for step_idx in range(H): + # Re-enable rendering before the final step to return a fresh image. + if ( + skip_render + and self._image_obs + and step_idx == H - 1 + and not self._render_enabled_state + ): + set_render_enabled(self._image_obs, True) + self._render_enabled_state = True + + action = model_action_to_libero_env(actions[step_idx].reshape(-1)) + obs, _, done_flag, _ = self._libero_env.step(action) + last_obs = obs + steps += 1 + self._append_rollout_frame(obs) + if bool(done_flag): + done = True + # Do not add an extra simulator step on early success; just make + # sure future rendering is enabled. + if skip_render and self._image_obs and not self._render_enabled_state: + set_render_enabled(self._image_obs, True) + self._render_enabled_state = True + break + + self._last_obs_for_dump = last_obs + return {"obs": last_obs, "done": done, "steps": steps} + + def finalize_episode(self, success: bool) -> None: + self._save_episode_rollout(success) + + def shutdown(self) -> None: + if self._libero_env is not None: + try: + self._libero_env.close() + except Exception: + pass + self._libero_env = None + + # ---- internals ---- + + def _chunk_skip_render(self) -> bool: + """Skip intermediate renders unless rollout MP4 saving needs every frame.""" + return self._effective_skip_render and self._rollout_dir is None + + def _begin_rollout_capture(self, obs: dict | None) -> None: + self._replay_images = [] + self._rollout_saved = False + if self._rollout_dir and obs is not None: + if self._image_obs and not self._render_enabled_state: + set_render_enabled(self._image_obs, True) + self._render_enabled_state = True + self._replay_images.append(get_rollout_frame(obs)) + + def _append_rollout_frame(self, obs: dict | None) -> None: + if self._rollout_dir and obs is not None: + self._replay_images.append(get_rollout_frame(obs)) + + def _save_episode_rollout(self, success: bool) -> None: + if ( + not self._rollout_dir + or self._rollout_saved + or not self._replay_images + or self._current_ep is None + ): + return + task_id, ep_idx = self._current_ep + try: + mp4_path = save_rollout_video( + self._rollout_dir, + self._replay_images, + task_id=task_id, + episode_idx=ep_idx, + success=success, + task_description=self._current_task_desc, + fps=self._rollout_fps, + ) + self._rollout_saved = True + if mp4_path: + logger.info("Saved rollout MP4: %s", mp4_path) + except Exception as exc: + logger.warning( + "Failed to save rollout MP4 for task%d ep%d: %s", + task_id, + ep_idx, + exc, + ) + + def _rebuild_env(self, task_id: int) -> None: + if self._libero_env is not None: + try: + self._libero_env.close() + except Exception: + pass + self._libero_env = None + gc.collect() + + self._libero_env = create_libero_engine( + task_id=task_id, + task_suite_name=self._task_suite_name, + resolution=256, + seed=self._seed, + ) + self._current_task_id = task_id + self._image_obs = find_image_observables(self._libero_env) + self._effective_skip_render = ( + self._skip_intermediate_render + and task_id not in self._force_render_task_ids + ) + + def _dump_chunk_npz( + self, ep_id, chunk_idx: int, raw_obs: dict, chunk_actions: np.ndarray + ) -> None: + _, task_id, ep_idx = ep_id + path = os.path.join( + self._bit_dump_dir, f"t{task_id}_ep{ep_idx}_c{chunk_idx}.npz" + ) + fields = {"action_chunk": np.asarray(chunk_actions, dtype=np.float32)} + for k in ( + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "agentview_image", + "robot0_eye_in_hand_image", + ): + v = raw_obs.get(k) + if v is not None: + fields[k] = np.asarray(v) + np.savez_compressed(path, **fields) diff --git a/wall_x/_vendor/harrix/envs/libero_common.py b/wall_x/_vendor/harrix/envs/libero_common.py new file mode 100644 index 0000000..2bfa2f3 --- /dev/null +++ b/wall_x/_vendor/harrix/envs/libero_common.py @@ -0,0 +1,359 @@ +"""Env-adapter IO helpers for single-arm LIBERO tasks. + +Env code extracts a compact ndarray payload from raw LIBERO observations. +Adapters then build proprioception, masks, and 7-dof right-arm action chunks +from that payload. This module intentionally depends only on NumPy so adapter +processes can import it without importing robosuite or LIBERO. +""" + +from __future__ import annotations + +import math +import os + +import numpy as np + + +# Fallback single-arm LIBERO dof_config when train_config does not provide one. +_LIBERO_FALLBACK_DOF_CONFIG = { + "follow_right_ee_cartesian_pos": 3, + "follow_right_ee_rotation": 3, + "follow_right_gripper": 1, +} + +_LIBERO_FALLBACK_AGENT_POS_CONFIG = dict(_LIBERO_FALLBACK_DOF_CONFIG) +_VIRTUAL_TAIL_KEYS = frozenset(("action_padding",)) + + +def _resolve_dof_config(train_config: dict) -> dict: + return ( + train_config.get("dof_config") + or train_config.get("task", {}).get("dof_config") + or _LIBERO_FALLBACK_DOF_CONFIG + ) + + +def _resolve_agent_pos_config(train_config: dict) -> dict: + return ( + train_config.get("agent_pos_config") + or train_config.get("task", {}).get("agent_pos_config") + or _LIBERO_FALLBACK_AGENT_POS_CONFIG + ) + + +def _move_virtual_keys_to_tail(layout: dict) -> dict: + head = {k: v for k, v in layout.items() if k not in _VIRTUAL_TAIL_KEYS} + tail = {k: v for k, v in layout.items() if k in _VIRTUAL_TAIL_KEYS} + return {**head, **tail} + + +def _effective_agent_pos_config(train_config: dict, state_values: dict) -> dict: + config = _move_virtual_keys_to_tail(dict(_resolve_agent_pos_config(train_config))) + gripper_key = next( + ( + key + for key in config + if key.replace("follow_", "").replace("master_", "") == "right_gripper" + ), + None, + ) + if gripper_key is None: + return config + + old_dim = int(config[gripper_key]) + target_dim = None + override = os.environ.get("WALLX_LIBERO_STATE_GRIPPER_DIM") + if override: + target_dim = int(override) + elif os.environ.get("WALLX_LIBERO_AUTO_STATE_GRIPPER_DIM", "1") != "0": + norm_dim = int(train_config.get("_libero_proprio_norm_dim") or 0) + real_dim = sum(v for k, v in config.items() if k not in _VIRTUAL_TAIL_KEYS) + if norm_dim == real_dim + 1: + target_dim = old_dim + 1 + + available_dim = state_values["right_gripper"].shape[1] + if target_dim is None or target_dim == old_dim or target_dim > available_dim: + return config + + config[gripper_key] = target_dim + delta = target_dim - old_dim + if "action_padding" in config: + config["action_padding"] = max(0, int(config["action_padding"]) - delta) + return config + + +def _build_right_arm_state_values(obs_ndarrays: dict) -> dict[str, np.ndarray]: + """Bare-key state tensors for single-arm LIBERO proprio construction.""" + rot3 = np.asarray(obs_ndarrays["eef_axisangle"], dtype=np.float32).reshape(1, 3) + values: dict[str, np.ndarray] = { + "right_ee_cartesian_pos": np.asarray( + obs_ndarrays["eef_pos"], dtype=np.float32 + ).reshape(1, 3), + "right_ee_rotation": rot3, + "right_gripper": np.asarray(obs_ndarrays["gripper"], dtype=np.float32).reshape( + 1, -1 + ), + } + from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb + + rot6d = euler_to_matrix_zyx_6d_nb(rot3.astype(np.float64)).reshape(1, 6) + values["right_ee_rotation_6D"] = rot6d.astype(np.float32) + return values + + +# Auxiliary action keys that should be masked out for single-arm LIBERO. +_DOF_MASK_ZERO_KEYS = frozenset( + ( + "follow_left_ee_cartesian_pos", + "follow_left_ee_rotation", + "follow_left_ee_rotation_6D", + "follow_left_gripper", + "head_actions", + "height", + "velocity_decomposed", + "action_padding", + ) +) + + +# ============================================================ +# LIBERO raw observation decoding helpers. +# ============================================================ + + +def _get_libero_image(obs: dict) -> np.ndarray: + """Return the third-person camera image, rotated to match preprocessing.""" + return obs["agentview_image"][::-1, ::-1] + + +def get_rollout_frame(obs: dict) -> np.ndarray: + """Return one RGB frame for rollout MP4 saving.""" + return np.asarray(_get_libero_image(obs), dtype=np.uint8) + + +def _get_libero_wrist_image(obs: dict) -> np.ndarray: + """Return the wrist camera image, rotated to match preprocessing.""" + return obs["robot0_eye_in_hand_image"][::-1, ::-1] + + +def _quat2axisangle(quat) -> np.ndarray: + """Convert an xyzw quaternion to a 3D axis-angle vector.""" + if quat[3] > 1.0: + quat[3] = 1.0 + elif quat[3] < -1.0: + quat[3] = -1.0 + + den = np.sqrt(1.0 - quat[3] * quat[3]) + if math.isclose(den, 0.0): + return np.zeros(3) + + return (quat[:3] * 2.0 * math.acos(quat[3])) / den + + +# ============================================================ +# Env-side to adapter-side observation encoding. +# ============================================================ + + +def encode_raw_obs(raw_obs: dict) -> dict: + """Extract the minimal ndarray payload from a raw LIBERO observation. + + The payload contains three 1-D state arrays and two rotated image arrays. + """ + if "agentview_image" not in raw_obs: + raise KeyError( + "agentview_image missing in raw_obs; render-skip may have returned " + "a stale observation" + ) + return { + "eef_pos": np.asarray(raw_obs["robot0_eef_pos"], dtype=np.float32), + "eef_axisangle": np.asarray( + _quat2axisangle(raw_obs["robot0_eef_quat"]), dtype=np.float32 + ), + "gripper": np.asarray(raw_obs["robot0_gripper_qpos"], dtype=np.float32), + "face_view": _get_libero_image(raw_obs), + "wrist_view": _get_libero_wrist_image(raw_obs), + } + + +# ============================================================ +# Adapter-side proprioception and mask construction. +# ============================================================ + + +def encode_proprio( + obs_ndarrays: dict, + train_config: dict, + action_horizon: int, +) -> dict: + """Convert a single-arm LIBERO ndarray payload into model input fields. + + Returned fields include proprioception, agent_pos_mask, dof_mask, + face_view, and right_wrist_view. + """ + state_values = _build_right_arm_state_values(obs_ndarrays) + agent_pos_config = _effective_agent_pos_config(train_config, state_values) + dof_config = _move_virtual_keys_to_tail(dict(_resolve_dof_config(train_config))) + + propri_parts: list[np.ndarray] = [] + mask_parts: list[np.ndarray] = [] + for key, dim in agent_pos_config.items(): + bare = key.replace("follow_", "").replace("master_", "") + if bare in state_values: + v = state_values[bare] + if bare == "right_gripper" and v.shape[1] > dim: + v = v[:, :dim] + if v.shape[1] != dim: + raise ValueError( + f"agent_pos_config[{key!r}]={dim} does not match " + f"observation dimension {v.shape[1]}" + ) + propri_parts.append(v) + mask_parts.append(np.ones((1, dim), dtype=np.float32)) + else: + propri_parts.append(np.zeros((1, dim), dtype=np.float32)) + mask_parts.append(np.zeros((1, dim), dtype=np.float32)) + + # (1, 1, D) + proprioception = np.concatenate(propri_parts, axis=1)[None] + agent_pos_mask = np.concatenate(mask_parts, axis=1)[None] + + # dof_mask: (1, T, D_action) + total_dof = sum(dof_config.values()) + dof_mask = np.ones((1, action_horizon, total_dof)) + start = 0 + for key, dim in dof_config.items(): + if key in _DOF_MASK_ZERO_KEYS: + dof_mask[:, :, start : start + dim] = 0 + start += dim + + return { + "proprioception": proprioception.astype(np.float32), + "agent_pos_mask": agent_pos_mask.astype(np.float32), + "dof_mask": dof_mask, + "face_view": obs_ndarrays["face_view"], + "right_wrist_view": obs_ndarrays["wrist_view"], + } + + +# ============================================================ +# Adapter-side action decoding. +# ============================================================ + + +def decode_chunk(predict_action: np.ndarray, train_config: dict) -> np.ndarray: + """Extract a 7-dof right-arm chunk from model action output.""" + if predict_action.ndim == 3: + predict_action = predict_action[0] + + dof_config = _resolve_dof_config(train_config) + slices: dict[str, slice] = {} + start = 0 + for key, dim in dof_config.items(): + bare = key.replace("follow_", "").replace("master_", "") + slices[bare] = slice(start, start + dim) + start += dim + + pos = predict_action[:, slices["right_ee_cartesian_pos"]] + grip = predict_action[:, slices["right_gripper"]] + + if "right_ee_rotation_6D" in slices: + from wall_x._vendor.x2robot_utils.geometry import so3_to_euler_zyx_batch_nb + + rot6d = np.asarray( + predict_action[:, slices["right_ee_rotation_6D"]], dtype=np.float64 + ) + rot = so3_to_euler_zyx_batch_nb(rot6d).astype(np.float32) + elif "right_ee_rotation" in slices: + rot = predict_action[:, slices["right_ee_rotation"]] + else: + raise KeyError( + "dof_config has no right-arm rotation slice " + f"(keys={list(slices.keys())})" + ) + + return np.concatenate([pos, rot, grip], axis=1) + + +def gripper_model_to_libero_osc(grip_2d: np.ndarray) -> float: + """Map model gripper output to robosuite OSC_POSE gripper command in [-1, 1].""" + g = np.asarray(grip_2d, dtype=np.float64).reshape(-1) + if g.size == 0: + raise ValueError("empty gripper action") + cmd = float(g[0]) + if os.environ.get("WALLX_LIBERO_GRIPPER_BINARIZE", "0") == "1": + if abs(cmd) < 1e-6: + cmd = -1.0 + else: + cmd = float(np.sign(cmd)) + if os.environ.get("WALLX_LIBERO_INVERT_GRIPPER", "0") == "1": + cmd *= -1.0 + return cmd + + +def _sanitize_task_description(task_description: str, max_len: int = 50) -> str: + return ( + task_description.lower() + .replace(" ", "_") + .replace("\n", "_") + .replace(".", "_")[:max_len] + ) + + +def save_rollout_video( + rollout_dir: str, + rollout_images: list[np.ndarray], + *, + task_id: int, + episode_idx: int, + success: bool, + task_description: str, + fps: int = 30, +) -> str | None: + """Save an MP4 replay of one LIBERO episode.""" + if not rollout_images: + return None + + import imageio + + os.makedirs(rollout_dir, exist_ok=True) + task_slug = _sanitize_task_description(task_description) + mp4_path = os.path.join( + rollout_dir, + f"task{task_id}_ep{episode_idx}--success={int(success)}--{task_slug}.mp4", + ) + writer = imageio.get_writer(mp4_path, fps=fps, macro_block_size=1) + try: + for img in rollout_images: + writer.append_data(np.asarray(img, dtype=np.uint8)) + finally: + writer.close() + return mp4_path + + +def model_action_to_libero_env(action: np.ndarray) -> np.ndarray: + """Convert model output to the robosuite OSC_POSE 7D action. + + Internal LIBERO evaluation passes the model's 7D chunk directly to + ``env.step``. Keep that as the public default; conversion modes remain + available only for ablations through environment variables. + """ + from scipy.spatial.transform import Rotation as R + + a = np.asarray(action, dtype=np.float64).reshape(-1) + if a.size not in (7, 8): + raise ValueError(f"expected 7D or 8D model action, got shape {a.shape}") + + pos_delta = a[:3] + rot_mode = os.environ.get("WALLX_LIBERO_ROT_MODE", "direct").strip() + if rot_mode == "euler_zyx_to_rotvec": + rot_aa = R.from_euler("zyx", a[3:6]).as_rotvec() + elif rot_mode == "direct": + rot_aa = a[3:6] + else: + raise ValueError( + "WALLX_LIBERO_ROT_MODE must be 'euler_zyx_to_rotvec' or 'direct', " + f"got {rot_mode!r}" + ) + grip = gripper_model_to_libero_osc(a[6:8] if a.size == 8 else a[6:7]) + return np.concatenate([pos_delta, rot_aa, np.array([grip], dtype=np.float64)]) diff --git a/wall_x/_vendor/harrix/envs/libero_sim.py b/wall_x/_vendor/harrix/envs/libero_sim.py new file mode 100644 index 0000000..ae51046 --- /dev/null +++ b/wall_x/_vendor/harrix/envs/libero_sim.py @@ -0,0 +1,181 @@ +"""LIBERO benchmark and robosuite engine helpers. + +This module may import robosuite and LIBERO. Adapter-side code should use +``libero_common.py`` instead, which only depends on NumPy. +""" + +from __future__ import annotations + +import json +import logging +import os +from typing import Any, Optional + +import numpy as np +from robosuite.wrappers import VisualizationWrapper + +logger = logging.getLogger(__name__) + + +# ============================================================ +# One-time side effect: auto-create ~/.libero/config.yaml so importing LIBERO +# does not trigger an interactive prompt. +# ============================================================ + + +def _ensure_libero_config() -> None: + import yaml as _yaml + + libero_config_path = os.environ.get( + "LIBERO_CONFIG_PATH", os.path.expanduser("~/.libero") + ) + config_file = os.path.join(libero_config_path, "config.yaml") + + if not os.path.exists(config_file): + os.makedirs(libero_config_path, exist_ok=True) + import libero.libero as _libero_pkg + + benchmark_root = os.path.dirname(os.path.abspath(_libero_pkg.__file__)) + default_paths = { + "benchmark_root": benchmark_root, + "bddl_files": os.path.join(benchmark_root, "./bddl_files"), + "init_states": os.path.join(benchmark_root, "./init_files"), + "datasets": os.path.join(benchmark_root, "../datasets"), + "assets": os.path.join(benchmark_root, "./assets"), + } + with open(config_file, "w") as f: + _yaml.dump(default_paths, f) + logger.info("Auto-created LIBERO config: %s", config_file) + + +_ensure_libero_config() + + +# ============================================================ +# task-suite entry point +# ============================================================ + + +def get_task_suite(task_suite_name: str): + """Load a LIBERO task suite.""" + from libero.libero import benchmark + + return benchmark.get_benchmark_dict()[task_suite_name]() + + +# ============================================================ +# actions +# ============================================================ + + +def get_libero_dummy_action() -> list[float]: + """Return the 7-dof dummy action used for episode warmup.""" + return [0, 0, 0, 0, 0, 0, -1] + + +# ============================================================ +# robosuite engine factory +# ============================================================ + + +def create_libero_engine( + task_id: int, + task_suite_name: str, + resolution: int = 256, + seed: int = 7, +) -> Any: + """Construct one LIBERO robosuite engine.""" + from libero.libero import get_libero_path + from libero.libero.envs import OffScreenRenderEnv + + task_suite = get_task_suite(task_suite_name) + task = task_suite.get_task(task_id) + task_bddl_file = os.path.join( + get_libero_path("bddl_files"), task.problem_folder, task.bddl_file + ) + env = OffScreenRenderEnv( + bddl_file_name=task_bddl_file, + camera_heights=resolution, + camera_widths=resolution, + ) + # The seed still affects object poses even when an initial state is fixed. + env.seed(seed) + env.env = VisualizationWrapper(env.env) + env.env.set_visualization_setting(setting="grippers", visible=False) + return env + + +# ============================================================ +# task metadata / initial states +# ============================================================ + + +def load_initial_states(initial_states_path: str) -> Optional[dict]: + """Load custom initial states, or return None for suite defaults.""" + if initial_states_path == "DEFAULT": + return None + with open(initial_states_path, "r") as f: + return json.load(f) + + +def resolve_task_info(task_suite, task_id: int) -> tuple[str, Any]: + """Return ``(task_desc, default_initial_states)`` for one task id.""" + num_tasks = task_suite.n_tasks + if task_id < 0 or task_id >= num_tasks: + raise ValueError(f"invalid task_id={task_id}, num_tasks={num_tasks}") + task = task_suite.get_task(task_id) + return task.language, task_suite.get_task_init_states(task_id) + + +def pick_initial_state( + initial_states_path: str, + custom_initial_states: Optional[dict], + task_desc: str, + default_states: Any, + episode_idx: int, +) -> np.ndarray: + """Pick one initial state from suite defaults or a custom states file.""" + if initial_states_path == "DEFAULT": + if default_states is None: + raise ValueError("default states missing for DEFAULT mode") + return default_states[episode_idx] + + if custom_initial_states is None: + raise ValueError(f"custom initial states not loaded for {initial_states_path}") + key = task_desc.replace(" ", "_") + ep_key = f"demo_{episode_idx}" + record = custom_initial_states[key][ep_key] + if not record["success"]: + raise ValueError(f"expert demo failed for {ep_key}") + return np.array(record["initial_state"]) + + +def get_instruction(task_desc: str) -> str: + """Return the instruction text for a LIBERO task description.""" + return task_desc + + +# ============================================================ +# render-skip: directly assign obs._enabled to avoid set_enabled() side effects. +# ============================================================ + + +def find_image_observables(env) -> list: + """Find image observables along the env.env wrapper chain.""" + cur = env + seen = set() + while cur is not None and id(cur) not in seen: + seen.add(id(cur)) + if hasattr(cur, "_observables") and isinstance(cur._observables, dict): + return [ + obs + for obs in cur._observables.values() + if getattr(obs, "modality", None) == "image" + ] + cur = getattr(cur, "env", None) + return [] + + +def set_render_enabled(image_obs_list, enabled: bool) -> None: + for obs in image_obs_list: + obs._enabled = enabled diff --git a/wall_x/_vendor/harrix/envs/registry.py b/wall_x/_vendor/harrix/envs/registry.py new file mode 100644 index 0000000..78d455c --- /dev/null +++ b/wall_x/_vendor/harrix/envs/registry.py @@ -0,0 +1,53 @@ +"""Environment registry and factory.""" + +from __future__ import annotations + +from typing import Type + +from wall_x._vendor.harrix.eval_config import EvalConfig +from wall_x._vendor.harrix.envs.base import BaseEnv + + +_REGISTRY: dict[str, Type[BaseEnv]] = {} + + +def register_env(name: str): + def deco(cls: Type[BaseEnv]): + if name in _REGISTRY: + raise ValueError( + f"env {name!r} already registered (cls={_REGISTRY[name].__name__})" + ) + _REGISTRY[name] = cls + return cls + + return deco + + +def _get_class(cfg: EvalConfig) -> Type[BaseEnv]: + t = cfg.env.type + cls = _REGISTRY.get(t) + if cls is None: + if t == "libero": + import wall_x._vendor.harrix.envs as _envs + + exc = getattr(_envs, "_LIBERO_IMPORT_ERROR", None) + if exc is not None: + raise RuntimeError( + "LIBERO evaluation dependencies are not installed. " + "Install LIBERO/robosuite and their simulator dependencies " + "before using env.type='libero'." + ) from exc + raise ValueError(f"unknown env type={t!r}, registered: {sorted(_REGISTRY)}") + return cls + + +def build_env(cfg: EvalConfig, worker_id: int) -> BaseEnv: + return _get_class(cfg)(cfg, worker_id) + + +def enumerate_episodes_for(cfg: EvalConfig) -> list[tuple]: + return _get_class(cfg).enumerate_episodes(cfg) + + +def registered_envs() -> list[str]: + return sorted(_REGISTRY) diff --git a/wall_x/_vendor/harrix/eval_config.py b/wall_x/_vendor/harrix/eval_config.py new file mode 100644 index 0000000..1ab59d6 --- /dev/null +++ b/wall_x/_vendor/harrix/eval_config.py @@ -0,0 +1,217 @@ +"""Typed YAML-driven configuration for inference and evaluation. + +The driver, model handle, and environment workers share one ``EvalConfig`` so +configuration is parsed once and then passed through explicitly. + +YAML schema: + model: + checkpoint_path: # checkpoint directory + train_config_path: null # null = /config.yml + norm_key: libero_all + cam_names: [face_view, right_wrist_view] + action_horizon: null # null = read data.action_horizon_flow + architecture: qwen2_5 # adapter registry key + action_mode: flow # flow / ar / dllm / vqa / subtask + env: + type: libero # env registry key + seed: 42 + libero: + task_suite_name: libero_spatial + initial_states_path: DEFAULT + num_trials_per_task: 50 + task_indices: null # null = all tasks + max_infer_times: 22 + skip_intermediate_render: true + force_render_task_indices: null # e.g. [5] always render that task + chunk_granular_render_toggle: false + rebuild_env_per_episode: false + rollout_dir: null # save third-person MP4 replays per episode + rollout_fps: 30 + runtime: + num_workers: 1 + max_batch_size: 1 + ws_port: 8765 + log_dir: /path/to/wallx_log + batch_sync_mode: false + debug: + deterministic_model: false + +Driver flow: + cfg = load_eval_config(yaml_path) + cfg = autofill_from_checkpoint(cfg) + +Unknown YAML fields are rejected to avoid silent typos. +""" + +import dataclasses +import os +from dataclasses import dataclass, field +from typing import Optional + +import yaml + + +@dataclass +class ModelSection: + checkpoint_path: str = "" + train_config_path: Optional[str] = None + norm_key: str = "libero_all" + cam_names: list = field(default_factory=lambda: ["face_view", "right_wrist_view"]) + action_horizon: Optional[int] = None + # Adapter implementation key in harrix.adapters.registry.ADAPTER_REGISTRY. + architecture: str = "qwen2_5" + # Inference algorithm. Each adapter validates its supported subset. + action_mode: str = "flow" + + +@dataclass +class LiberoEnvParams: + """LIBERO-specific env settings used when ``env.type == "libero"``.""" + + task_suite_name: str = "libero_spatial" + initial_states_path: str = "DEFAULT" + num_trials_per_task: int = 50 + task_indices: Optional[list] = None + max_infer_times: int = 22 + # Render-skip is enabled by default. Listed task ids always render every + # simulator step for contact-sensitive tasks. + skip_intermediate_render: bool = True + force_render_task_indices: Optional[list] = None + # If enabled, render observables are toggled only when the chunk boundary + # actually changes the desired state. + chunk_granular_render_toggle: bool = False + # Rebuild the simulator for every episode instead of only on task changes. + # This is mainly a debugging option because it changes the simulator RNG path. + rebuild_env_per_episode: bool = False + # When set, save a third-person MP4 replay for each episode under this directory. + # Falls back to the ``WALLX_ROLLOUT_DIR`` environment variable when null. + rollout_dir: Optional[str] = None + rollout_fps: int = 30 + + +@dataclass +class EnvSection: + # Env implementation key in harrix.envs.registry. + type: str = "libero" + seed: int = 42 + libero: LiberoEnvParams = field(default_factory=LiberoEnvParams) + + +@dataclass +class RuntimeSection: + num_workers: int = 1 + max_batch_size: int = 1 + ws_port: int = 8765 + log_dir: str = "/path/to/wallx_log" + # Run fixed task-local frames instead of dynamic work stealing. This improves + # reproducibility at the cost of possible idle workers inside a frame. + batch_sync_mode: bool = False + # Public Wall-X evaluation uses the in-process driver. Other drivers may be + # enabled by downstream/internal integrations. + driver_mode: str = "in_process" + + +@dataclass +class DebugSection: + # Enable deterministic torch backend options for variance debugging. + deterministic_model: bool = False + + +@dataclass +class EvalConfig: + model: ModelSection = field(default_factory=ModelSection) + env: EnvSection = field(default_factory=EnvSection) + runtime: RuntimeSection = field(default_factory=RuntimeSection) + debug: DebugSection = field(default_factory=DebugSection) + + +def _build_dataclass(cls, raw): + """Build a dataclass from a dict and reject unknown fields.""" + if raw is None: + return cls() + if not isinstance(raw, dict): + raise ValueError(f"Expected dict for {cls.__name__}, got {type(raw).__name__}") + field_names = {f.name for f in dataclasses.fields(cls)} + field_types = {f.name: f.type for f in dataclasses.fields(cls)} + unknown = set(raw.keys()) - field_names + if unknown: + raise ValueError( + f"Unknown field(s) {sorted(unknown)} in {cls.__name__}; " + f"expected one of {sorted(field_names)}" + ) + built = {} + for k, v in raw.items(): + ft = field_types[k] + if ( + isinstance(ft, type) + and dataclasses.is_dataclass(ft) + and isinstance(v, dict) + ): + built[k] = _build_dataclass(ft, v) + else: + built[k] = v + return cls(**built) + + +def load_eval_config(yaml_path: str) -> EvalConfig: + """Load ``EvalConfig`` from YAML and validate checkpoint_path.""" + with open(yaml_path, "r") as f: + raw = yaml.safe_load(f) + if raw is None: + raise ValueError(f"Empty YAML: {yaml_path}") + + cfg = _build_dataclass(EvalConfig, raw) + + if not cfg.model.checkpoint_path: + raise ValueError(f"model.checkpoint_path is required in {yaml_path}") + if not os.path.isdir(cfg.model.checkpoint_path): + raise FileNotFoundError( + f"model.checkpoint_path does not exist: {cfg.model.checkpoint_path}" + ) + + return cfg + + +def _load_train_config_yaml(path: str) -> dict: + """Load a checkpoint-side train YAML. + + Training checkpoints may contain PyYAML-specific tags such as + ``!!python/tuple``; ``safe_load`` cannot parse those. + """ + with open(path, "r") as f: + return yaml.load(f, Loader=yaml.FullLoader) or {} + + +def _read_action_horizon_flow(train_yml: dict) -> int: + task = train_yml.get("task") or {} + data = train_yml.get("data") or {} + return int( + train_yml.get("action_horizon_flow") + or task.get("action_horizon_flow") + or data.get("action_horizon_flow") + or 32 + ) + + +def autofill_from_checkpoint(cfg: EvalConfig) -> EvalConfig: + """Fill omitted model fields from the checkpoint-side train config. + + This resolves ``model.train_config_path`` and ``model.action_horizon``. + """ + if cfg.model.train_config_path is None: + for fname in ("config.yml", "config.yaml"): + cand = os.path.join(cfg.model.checkpoint_path, fname) + if os.path.exists(cand): + cfg.model.train_config_path = cand + break + else: + raise FileNotFoundError( + f"No config.yml/config.yaml in {cfg.model.checkpoint_path}; " + "set model.train_config_path explicitly in YAML" + ) + + if cfg.model.action_horizon is None: + train_yml = _load_train_config_yaml(cfg.model.train_config_path) + cfg.model.action_horizon = _read_action_horizon_flow(train_yml) + + return cfg diff --git a/wall_x/serving/__init__.py b/wall_x/_vendor/harrix/serving/__init__.py similarity index 100% rename from wall_x/serving/__init__.py rename to wall_x/_vendor/harrix/serving/__init__.py diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py b/wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py new file mode 100644 index 0000000..cb5d7d6 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py @@ -0,0 +1 @@ +"""Wall-X inference helpers vendored for harrix serving.""" diff --git a/wall_x/infer/base_dataclass.py b/wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py similarity index 62% rename from wall_x/infer/base_dataclass.py rename to wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py index b550c3b..dccdacc 100644 --- a/wall_x/infer/base_dataclass.py +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py @@ -1,10 +1,10 @@ -from wall_x.infer.infer_config import InferConfig +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig from typing import Optional, List from dataclasses import dataclass, field import numpy as np import torch -import wall_x.infer.data_utils as data_utils -from wall_x.infer.logger import InferLogger +import wall_x._vendor.x2robot_utils.geometry as data_utils +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger dof_dims = { "left_ee_cartesian_pos": 3, @@ -13,29 +13,57 @@ dof_dims = { "left_ee_rotation_relative": 3, "left_ee_rotation_6D": 6, "left_ee_rotation_6D_relative": 6, - "left_arm_joint_pos": 7, + "left_arm_joint_pos": 7, # 6 joint + 1 gripper "left_gripper": 1, "left_gripper_cur": 1, - "left_arm_joint_cur": 1, + "left_arm_joint_cur": 7, # 1 -> 7 "right_ee_cartesian_pos": 3, "right_ee_cartesian_pos_relative": 3, "right_ee_rotation": 3, "right_ee_rotation_relative": 3, "right_ee_rotation_6D": 6, "right_ee_rotation_6D_relative": 6, - "right_arm_joint_pos": 7, + "right_arm_joint_pos": 7, # 6 joint + 1 gripper "right_gripper": 1, "right_gripper_cur": 1, - "right_arm_joint_cur": 1, + "right_arm_joint_cur": 7, # 1 -> 7 "head_actions": 2, "height": 1, "car_pose": 3, "velocity_decomposed": 3, + "velocity_decomposed_odom": 3, + "head_rotation": 2, # match ex001 + "left_joint": 6, + "left_joint_gripper": 1, + "right_joint": 6, + "right_joint_gripper": 1, + "left_rotation_quat": 4, + "right_rotation_quat": 4, + "left_quaternion": 4, + "right_quaternion": 4, + "left_wrench_ext_local_force": 3, + "left_wrench_ext_local_torque": 3, + "right_wrench_ext_local_force": 3, + "right_wrench_ext_local_torque": 3, + "left_wrench_ext_local_force_from_joint": 3, + "left_wrench_ext_local_torque_from_joint": 3, + "right_wrench_ext_local_force_from_joint": 3, + "right_wrench_ext_local_torque_from_joint": 3, + "left_wrench_ext_world_force": 3, + "left_wrench_ext_world_torque": 3, + "left_wrench_ext_world_force_from_joint": 3, + "left_wrench_ext_world_torque_from_joint": 3, + "right_wrench_ext_world_force": 3, + "right_wrench_ext_world_torque": 3, + "right_wrench_ext_world_force_from_joint": 3, + "right_wrench_ext_world_torque_from_joint": 3, + "left_arm_joint_dev": 7, + "right_arm_joint_dev": 7, } class ComputedDict(dict): - """Smart dictionary that supports registering computation rules and auto-computes None values on get""" + """Dict that registers compute rules and auto-computes None values on get""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -43,45 +71,45 @@ class ComputedDict(dict): def register_compute_rule(self, key, compute_func): """ - Register a computation rule. + Register a compute rule Args: - key: The key that needs computation - compute_func: Computation function that takes self as argument and returns the computed result + key: Key to compute + compute_func: Callable taking self, returns computed value """ self._compute_rules[key] = compute_func def get(self, key, default=None): - """Override get method to support auto-computation""" + """Override get to support auto-compute""" value = super().get(key, default) - # If value is None and there's a compute rule, try to compute + # If value is None and a rule exists, try to compute if value is None and key in self._compute_rules: try: computed_value = self._compute_rules[key](self) if computed_value is not None: - # Cache the computed result + # Cache computed value self[key] = computed_value return computed_value except Exception: - pass # If computation fails, return None or default + pass # On failure return None or default return value if value is not None else default def __getitem__(self, key): - """Override [] operator to support auto-computation""" + """Override [] to support auto-compute""" value = super().get(key, None) - # If value is None and there's a compute rule, try to compute + # If value is None and a rule exists, try to compute if value is None and key in self._compute_rules: try: computed_value = self._compute_rules[key](self) if computed_value is not None: - # Cache the computed result + # Cache computed value self[key] = computed_value return computed_value except Exception: - pass # If computation fails, raise original KeyError or return None + pass # On failure raise KeyError or return None if key in self: return super().__getitem__(key) @@ -94,7 +122,7 @@ class RobotStateActionData: data: ComputedDict = field( default_factory=lambda: ComputedDict( { - # State (formerly pose) - using state_ prefix + # State (formerly pose) - state_ prefix "state_left_ee_cartesian_pos": None, # (1, 3) "state_left_ee_rotation": None, # (1, 3) "state_left_ee_rotation_6D": None, @@ -110,10 +138,42 @@ class RobotStateActionData: "state_right_gripper_cur": None, "state_right_arm_joint_cur": None, # (1, 1) "state_head_actions": None, + "state_head_rotation": None, # match ex001 "state_height": None, "state_car_pose": None, "state_velocity_decomposed": None, - # Action - using action_ prefix + "state_velocity_decomposed_odom": None, + # support joint control + "state_left_joint": None, + "state_left_joint_gripper": None, + "state_right_joint": None, + "state_right_joint_gripper": None, + # support quaternion control + "state_left_quaternion": None, + "state_right_quaternion": None, + "state_left_rotation_quat": None, + "state_right_rotation_quat": None, + # support wrench/force-torque observation + "state_left_wrench_ext_local_force": None, + "state_left_wrench_ext_local_torque": None, + "state_right_wrench_ext_local_force": None, + "state_right_wrench_ext_local_torque": None, + "state_left_wrench_ext_local_force_from_joint": None, + "state_left_wrench_ext_local_torque_from_joint": None, + "state_right_wrench_ext_local_force_from_joint": None, + "state_right_wrench_ext_local_torque_from_joint": None, + "state_left_wrench_ext_world_force": None, + "state_left_wrench_ext_world_torque": None, + "state_left_wrench_ext_world_force_from_joint": None, + "state_left_wrench_ext_world_torque_from_joint": None, + "state_right_wrench_ext_world_force": None, + "state_right_wrench_ext_world_torque": None, + "state_right_wrench_ext_world_force_from_joint": None, + "state_right_wrench_ext_world_torque_from_joint": None, + # support joint deviation + "state_left_arm_joint_dev": None, + "state_right_arm_joint_dev": None, + # Action - action_ prefix "action_left_ee_cartesian_pos": None, "action_left_ee_cartesian_pos_relative": None, "action_left_ee_rotation": None, @@ -131,9 +191,21 @@ class RobotStateActionData: "action_right_gripper": None, "action_right_arm_joint_pos": None, "action_head_actions": None, + "action_head_rotation": None, # match ex001 "action_height": None, "action_car_pose": None, "action_velocity_decomposed": None, + "action_velocity_decomposed_odom": None, + # support joint control + "action_left_joint": None, + "action_left_joint_gripper": None, + "action_right_joint": None, + "action_right_joint_gripper": None, + # support quaternion control + "action_left_quaternion": None, + "action_right_quaternion": None, + "action_left_rotation_quat": None, + "action_right_rotation_quat": None, } ) ) @@ -141,8 +213,8 @@ class RobotStateActionData: logger = InferLogger.get_robot_logger("RobotStateActionData") def __post_init__(self): - """Register computation rules""" - # State computation rules - euler angles -> 6D rotation + """Register compute rules""" + # State rules: euler angles -> 6D rotation self.data.register_compute_rule( "state_left_ee_rotation_6D", lambda d: ( @@ -160,7 +232,7 @@ class RobotStateActionData: ), ) - # Action computation rules - absolute position computed from relative + state + # Action rules: absolute position from relative + state self.data.register_compute_rule( "action_left_ee_cartesian_pos", lambda d: ( @@ -182,7 +254,7 @@ class RobotStateActionData: ), ) - # Action computation rules - get absolute rpy + # Action rules: compute absolute rpy self.data.register_compute_rule( # delta rpy -> abs rpy "action_left_ee_rotation", lambda d: ( @@ -203,10 +275,10 @@ class RobotStateActionData: else None ), ) - self.data.register_compute_rule( # delta 6D -> abs 6D -> abs rpy + self.data.register_compute_rule( # delta 6D -> abs 6D "action_left_ee_rotation_6D", lambda d: ( - data_utils.compose_state_and_delta_to_abs_rpy( + data_utils.compose_state_and_delta_to_abs_6d( d["action_left_ee_rotation_6D_relative"], d["state_left_ee_rotation_6D"][0], ) @@ -236,10 +308,10 @@ class RobotStateActionData: else None ), ) - self.data.register_compute_rule( # delta 6D -> abs 6D -> abs rpy + self.data.register_compute_rule( # delta 6D -> abs 6D "action_right_ee_rotation_6D", lambda d: ( - data_utils.compose_state_and_delta_to_abs_rpy( + data_utils.compose_state_and_delta_to_abs_6d( d["action_right_ee_rotation_6D_relative"], d["state_right_ee_rotation_6D"][0], ) @@ -251,11 +323,17 @@ class RobotStateActionData: def get_agent_pos(self, obs_action_keys=None): if obs_action_keys is None: - obs_action_keys = self.config.train_config["data"]["obs_action_keys"] + obs_action_keys = self.config.train_config["agent_pos_config"].keys() agent_pose_data = [] for key in obs_action_keys: - # Remove follow_ or master_ prefix + # action_padding is a virtual key used to pad agent_pos width. + if key == "action_padding": + dim = self.config.train_config["agent_pos_config"][key] + agent_pose_data.append(np.zeros((1, dim))) + continue + + # Strip follow_ or master_ prefix if key.startswith("follow_"): key = key.replace("follow_", "") elif key.startswith("master_"): @@ -265,10 +343,10 @@ class RobotStateActionData: state_key = f"state_{key}" if state_key in self.data: - # Use get method, which will auto-handle None value computation + # get() auto-computes when value is None value = self.data.get(state_key) if value is None: - # If still None after computation, use zero vector + # If still None after compute, use zeros agent_pose_data.append(np.zeros((1, dof_dims[key]))) else: agent_pose_data.append(value) @@ -281,11 +359,17 @@ class RobotStateActionData: def get_agent_pos_mask(self, obs_action_keys=None): if obs_action_keys is None: - obs_action_keys = self.config.train_config["data"]["obs_action_keys"] + obs_action_keys = self.config.train_config["agent_pos_config"].keys() agent_pos_mask_data = [] for key in obs_action_keys: - # Remove follow_ or master_ prefix + # action_padding carries no information and should stay masked out. + if key == "action_padding": + dim = self.config.train_config["agent_pos_config"][key] + agent_pos_mask_data.append(np.zeros((1, dim))) + continue + + # Strip follow_ or master_ prefix if key.startswith("follow_"): key = key.replace("follow_", "") elif key.startswith("master_"): @@ -295,19 +379,19 @@ class RobotStateActionData: state_key = f"state_{key}" if state_key in self.data: - # Use get method, which will auto-handle None value computation + # get() auto-computes when value is None value = self.data.get(state_key) if value is None: agent_pos_mask_data.append(np.zeros((1, dof_dims[key]))) else: - agent_pos_mask_data.append(np.ones((1, dof_dims[key]))) + agent_pos_mask_data.append(np.ones((1, value.shape[1]))) else: raise ValueError(f"Key {state_key} not found in data") return np.concatenate(agent_pos_mask_data, axis=1)[None] # (1, 1, D) - def save_state_data_with_key(self, value, key): - # Remove follow_ or master_ prefix + def save_state_data_with_key(self, value, key, gt_dim=None): + # Strip follow_ or master_ prefix key = key.replace("follow_", "") key = key.replace("master_", "") @@ -316,15 +400,16 @@ class RobotStateActionData: value = value.detach().cpu().numpy() if f"state_{key}" not in self.data: # TODO: joint angle control - self.logger.warning(f"{key} is not a valid state key, not recorded") + self.logger.warning(f"{key} is not a valid state key; not recorded") return + gt_dim = dof_dims[key] if gt_dim is None else gt_dim - # Shape validation for value, expected shape is (1, D) - if value.shape == (1, dof_dims[key]): + # Shape check: expected (1, D) + if value.shape == (1, gt_dim): self.data[f"state_{key}"] = value - elif value.shape == (1, 1, dof_dims[key]): + elif value.shape == (1, 1, gt_dim): self.data[f"state_{key}"] = value[0] - elif value.shape == (dof_dims[key],): + elif value.shape == (gt_dim,): self.data[f"state_{key}"] = value[None] else: raise ValueError(f"Value shape {value.shape} is not legal") @@ -345,7 +430,16 @@ class RobotStateActionData: self, predict_action, predict_action_keys: Optional[List[str]] = None ): if predict_action_keys is None: - predict_action_keys = self.config.data_config["predict_action_keys"] + predict_action_keys = getattr( + self.config.data_config, "predict_action_keys", None + ) + if predict_action_keys is None: + try: + predict_action_keys = self.config.data_config["predict_action_keys"] + except (KeyError, AttributeError): + predict_action_keys = list( + self.config.train_config["dof_config"].keys() + ) if isinstance(predict_action, torch.Tensor): predict_action = predict_action.detach().cpu().numpy() @@ -355,6 +449,12 @@ class RobotStateActionData: dof_start = 0 for action_key in predict_action_keys: + # action_padding is a virtual key that only advances dof_start; no data is written. + if action_key == "action_padding": + dof_dim = self.config.train_config["dof_config"]["action_padding"] + dof_start += dof_dim + continue + action_key = action_key.replace("follow_", "") action_key = action_key.replace("master_", "") dof_dim = dof_dims[action_key] @@ -362,7 +462,7 @@ class RobotStateActionData: self.data[action_key] = predict_action[:, dof_start : dof_start + dof_dim] dof_start += dof_dim - # For compatibility, provide convenient property access + # Convenience properties for compatibility @property def agent_pos(self): return self.get_agent_pos() diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py b/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py new file mode 100644 index 0000000..5f7f70b --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py @@ -0,0 +1,434 @@ +import logging +import os +from dataclasses import dataclass, field +from typing import Any + +import yaml +from qwen_vl_utils.vision_process import IMAGE_FACTOR, MAX_PIXELS, MIN_PIXELS + +logger = logging.getLogger(__name__) + + +def _env_bool(name: str, default: bool) -> bool: + v = os.environ.get(name) + if v is None: + return default + return v.strip().lower() in ("1", "true", "yes", "on") + + +@dataclass +class InferenceDataConfig: + """Minimal data config needed by online inference. + + Serving does not build datasets, so it should not require closed-source + data backends just to read image resize settings from a checkpoint yaml. + """ + + resolution: dict[str, int] = field(default_factory=dict) + model_type: str = "qwen2_5" + max_pixels: int = MAX_PIXELS + min_pixels: int = MIN_PIXELS + image_factor: int = IMAGE_FACTOR + use_relative_action: bool = False + predict_action_keys: list[str] = field(default_factory=list) + + def __getitem__(self, key: str) -> Any: + if not hasattr(self, key): + raise KeyError(key) + return getattr(self, key) + + +class InferConfig: + def __init__( + self, + checkpoint_path: str | None = None, + train_config_path: str | None = None, + robot_host: str = "0.0.0.0", + robot_port: int = 41776, + robot_type: str = "desktop", # ["desktop", "turtle", "ex001"] + robot_action_start_ratio: float = 0, # Start ratio for trimming executed actions + robot_action_end_ratio: float = 0.8, # End ratio for trimming executed actions + robot_action_interpolate_multiplier: int = 10, # Action interpolation multiplier + robot_use_joint_angle_control: bool = False, # Joint control (model must predict joints) + turtle_as_desktop: bool = False, # Use turtle as desktop: fixed base/head motion, head camera, base height + action_horizon: int = 32, # Set to the model's action horizon + action_dim: int | None = None, + ar_action_dim: int | None = None, + model_device: str = "cuda:0", + num_inference_timesteps: int = 10, + num_inference_steps: int | None = None, + cfg_scale: float | None = None, + seed: int | None = None, + norm_key: str = "x2_normal", # ["x2_normal", "ex_normal"] + cam_names: list[str] | None = None, + camera_front_key: str = "camera_front", + camera_left_key: str = "camera_left", + camera_right_key: str = "camera_right", + default_instruction: str | None = None, + prompt_template: str | None = None, + qwen25_prompt_template: str | None = None, + prompt_priority_order: str | None = None, + save_video_dir: str = "./videos", + robot_id: str = "10000", + model_type: str = "wallx", # ["wallx", "vga"] + smooth_action: bool = False, + smooth_gripper: bool = True, + ): + # Private attributes for paths + assert checkpoint_path is not None + from wall_x._vendor.harrix.utils.ckpt_load import resolve_checkpoint_dir + + if not os.path.isdir(checkpoint_path) and not os.path.isfile(checkpoint_path): + raise FileNotFoundError( + f"Checkpoint path not found: {checkpoint_path!r}. " + "Serving requires a directory containing model.safetensors and " + "normalizer_action.pth / normalizer_propri.pth (or norm_stats.json)." + ) + self._checkpoint_path = resolve_checkpoint_dir(checkpoint_path) + action_pth = os.path.join(self._checkpoint_path, "normalizer_action.pth") + propri_pth = os.path.join(self._checkpoint_path, "normalizer_propri.pth") + if os.path.exists(action_pth): + self.normalizer_action_path = action_pth + if os.path.exists(propri_pth): + self.normalizer_propri_path = propri_pth + + # Other config attributes + self.robot_host = robot_host + self.robot_port = robot_port + self.robot_type = robot_type # ["desktop", "turtle", "ex001"] + self.robot_action_start_ratio = robot_action_start_ratio + self.robot_action_end_ratio = robot_action_end_ratio + self.robot_action_interpolate_multiplier = robot_action_interpolate_multiplier + self.robot_use_joint_angle_control = ( + robot_use_joint_angle_control # Joint-angle control + ) + self.turtle_as_desktop = turtle_as_desktop + self.robot_id = robot_id + self._action_horizon = ( + action_horizon # Default from train config flow action horizon + ) + self._action_dim = action_dim # Default from train config dof config + self.ar_action_dim = ar_action_dim # Default from train config ar dof config + self.model_device = model_device + self.num_inference_timesteps = ( + num_inference_timesteps # flow matching related config + ) + self.num_inference_steps = num_inference_steps + self.cfg_scale = cfg_scale + self.seed = seed + self.save_video_dir = save_video_dir + self.model_type = model_type # ["wallx", "vga"] + # Initialize config objects + self.train_config: dict = {} + self.model_config = None + self.data_config = None + self.norm_key = norm_key + self.cam_names = cam_names or [ + "face_view", + "left_wrist_view", + "right_wrist_view", + ] + self.camera_front_key = camera_front_key + self.camera_left_key = camera_left_key + self.camera_right_key = camera_right_key + self.default_instruction = default_instruction + self.prompt_template = prompt_template + self.qwen25_prompt_template = qwen25_prompt_template + self.prompt_priority_order = prompt_priority_order + self.smooth_action = _env_bool("WALLX_SMOOTH_ACTION", smooth_action) + self.smooth_gripper = _env_bool("WALLX_SMOOTH_GRIPPER", smooth_gripper) + # Load all configs + self._load_all_configs(train_config_path) + self._apply_cam_names_from_train_config(cam_names) + + def _apply_cam_names_from_train_config( + self, cam_names: list[str] | None + ) -> None: + """Use train YAML camera mapping when CLI did not override cam_names.""" + if cam_names is not None: + return + from wall_x._vendor.harrix.utils.train_config import ( + resolve_cam_names_from_train_config, + ) + + resolved = resolve_cam_names_from_train_config(self.train_config) + if resolved: + self.cam_names = resolved + logger.info( + "[InferConfig] cam_names from train key_mappings: %s", + self.cam_names, + ) + + @property + def checkpoint_path(self) -> str | None: + return self._checkpoint_path + + @checkpoint_path.setter + def checkpoint_path(self, value: str | None): + """Reload all configs when checkpoint_path is updated""" + if self._checkpoint_path != value: + self._checkpoint_path = value + self._load_all_configs() + + @property + def action_horizon(self) -> int: + return self._action_horizon + + @action_horizon.setter + def action_horizon(self, value: int): + self._action_horizon = value + + @property + def action_dim(self) -> int | None: + return self._action_dim + + @action_dim.setter + def action_dim(self, value: int | None): + self._action_dim = value + + @property + def ar_action_dim(self) -> int | None: + return self._ar_action_dim + + @ar_action_dim.setter + def ar_action_dim(self, value: int | None): + self._ar_action_dim = value + + def _load_all_configs(self, train_config_path=None): + """Unified entry to load all configs""" + self._load_train_config(train_config_path) + if self.model_type != "vga": + self._load_model_config_for_wallx() + self._load_data_config() + + # Update action_horizon and action_dim if needed + if self._action_horizon is None: + self._action_horizon = self.train_config.get("data", {}).get( + "action_horizon_flow", 32 + ) + assert self._action_horizon is not None and self._action_horizon > 0 + + if self._action_dim is None: + dof_config = ( + self.train_config.get("dof_config") + or self.train_config.get("task", {}).get("dof_config") + or self.train_config.get("data", {}).get("dof_config", {}) + ) + self._action_dim = sum(dof_config.values()) + + if self._ar_action_dim is None: + ar_dof_config = self.train_config.get("ar_dof_config") or self.train_config.get( + "task", {} + ).get("ar_dof_config", {}) + self._ar_action_dim = sum(ar_dof_config.values()) + + def _load_train_config(self, train_config_path): + if train_config_path is None: + for fname in ("config.yml", "config.yaml"): + candidate = os.path.join(self._checkpoint_path, fname) + if os.path.exists(candidate): + train_config_path = candidate + break + else: + raise FileNotFoundError( + f"No config.yml/config.yaml found in {self._checkpoint_path}" + ) + with open(train_config_path, "r") as f: + self.train_config = yaml.load(f, Loader=yaml.FullLoader) + + ckpt_dir = self._checkpoint_path + preprocessor_file = os.path.join(ckpt_dir, "preprocessor_config.json") + if os.path.exists(preprocessor_file): + logger.info( + "[LoadConfig] Found %s, override processor_path.", + preprocessor_file, + ) + self.train_config["processor_path"] = ckpt_dir + + tokenizer_file = os.path.join(ckpt_dir, "tokenizer.json") + tokenizer_config_file = os.path.join(ckpt_dir, "tokenizer_config.json") + if self.train_config.get( + "action_tokenizer_path", None + ) is not None and not os.path.exists( + self.train_config.get("action_tokenizer_path", None) + ): + if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file): + logger.info( + "[LoadConfig] Found tokenizer files in %s, override action_tokenizer_path.", + ckpt_dir, + ) + self.train_config["action_tokenizer_path"] = ckpt_dir + else: + logger.warning("[LoadConfig] Cannot load action tokenizer! ") + + from wall_x._vendor.harrix.utils.train_config import ( + normalize_train_config_for_inference, + strip_action_tokenizer_fields, + ) + + self.train_config = normalize_train_config_for_inference( + self.train_config, train_config_path + ) + # Flow serving does not load AR action tokenizers; crop embed_tokens instead. + strip_action_tokenizer_fields(self.train_config) + self._train_config_path = train_config_path + + def _load_model_config_for_wallx(self): + model_type = self.train_config["model_type"] + if not os.path.isdir(self._checkpoint_path): + return + # For Qwen models + ckpt_config_path = os.path.join(self._checkpoint_path, "config.json") + resolved_cfg_path = None + + if os.path.exists(ckpt_config_path): + # Prefer checkpoint config + resolved_cfg_path = ckpt_config_path + logger.info( + "[LoadModelConfig] Using checkpoint config.json: %s", + ckpt_config_path, + ) + else: + # Fallback to original config path + fallback_cfg = self.train_config.get("qwen_vl_act_config_path", None) + if fallback_cfg is not None: + resolved_cfg_path = fallback_cfg + logger.info( + "[LoadModelConfig] Using fallback act config: %s", + fallback_cfg, + ) + + if resolved_cfg_path is None or (not os.path.exists(resolved_cfg_path)): + raise ValueError( + f"[LoadModelConfig] Cannot load model config! " + f"Checked:\n" + f" - Checkpoint config.json: {ckpt_config_path}\n" + f" - Fallback path: {self.train_config.get('qwen_vl_act_config_path', None)}" + ) + + # Save back to config for consistency + self.train_config["qwen_vl_act_config_path"] = resolved_cfg_path + + from wall_x.trainer.adapters import resolve_adapter + + adapter_cls = resolve_adapter(model_type) + ConfigClass = adapter_cls.config_class() + + logger.info( + "[LoadModelConfig] Loading model config from: %s", resolved_cfg_path + ) + if resolved_cfg_path.endswith(".json"): + self.model_config = ConfigClass.from_json_file(resolved_cfg_path) + else: + self.model_config = ConfigClass.from_pretrained(resolved_cfg_path) + + self.model_config.update_model_config(self.train_config) + + self.model_config._attn_implementation = "sdpa" + vision_attn = os.environ.get("WALLX_VISION_ATTN_IMPLEMENTATION", "flash_attention_2") + self.model_config.vision_config._attn_implementation = vision_attn + logger.info( + "[LoadModelConfig] vision _attn_implementation=%s (override via WALLX_VISION_ATTN_IMPLEMENTATION)", + vision_attn, + ) + + logger.info("[LoadModelConfig] Model config loaded and updated successfully.") + + def _load_data_config(self): + # Prefer typed TrainConfig path (handles new 8-section schema where + # dof_config lives under task:). Fall back to legacy raw-dict path + # for old flat yamls. See trainer/adapters/base_adapter.py. + typed_dcfg = self._try_typed_data_config() + if typed_dcfg is not None: + self.data_config = typed_dcfg + elif self.model_type != "vga": + self.data_config = self._build_inference_data_config() + else: + # TEMPORARY: direct call to the private _set_data_backend. + # VGA inference still uses the legacy x2robot data config object. + # Wall-X online inference uses the lightweight InferenceDataConfig + # above so serving does not require excluded data backends. + from wall_x.data._registry import _set_data_backend + + _set_data_backend(self.train_config.get("dataset_type", "x2robot_v1")) + from wall_x.model.vga.openloop_visualization import get_data_configs + + dataload_config = get_data_configs(self.train_config.get("data", {})) + dataload_config["predict_action_keys"] = list( + dataload_config.get("dof_config", {}).keys() + ) + self.data_config = dataload_config + + self._ensure_predict_action_keys() + + def _ensure_predict_action_keys(self) -> None: + dof_keys = list((self.train_config.get("dof_config") or {}).keys()) + if not dof_keys: + return + if isinstance(self.data_config, InferenceDataConfig): + if not self.data_config.predict_action_keys: + self.data_config.predict_action_keys = dof_keys + elif isinstance(self.data_config, dict): + self.data_config.setdefault("predict_action_keys", dof_keys) + + def _build_inference_data_config(self) -> InferenceDataConfig: + data = self.train_config.get("data", {}) or {} + + def get(key: str, default: Any) -> Any: + return data.get(key, self.train_config.get(key, default)) + + resolution = get("resolution", {}) or {} + dof_config = self.train_config.get("dof_config") or {} + return InferenceDataConfig( + resolution=dict(resolution), + model_type=get( + "model_type", self.train_config.get("model_type", "qwen2_5") + ), + max_pixels=get("max_pixels", MAX_PIXELS), + min_pixels=get("min_pixels", MIN_PIXELS), + image_factor=get("image_factor", IMAGE_FACTOR), + use_relative_action=get("use_relative_action", False), + predict_action_keys=list(dof_config.keys()), + ) + + def _try_typed_data_config(self): + """Attempt to load TrainConfig and build X2RDataConfig via typed path. + + Returns the X2RDataConfig on success, or None if the yaml is not in + TrainConfig schema (legacy flat yaml). + """ + yml_path = getattr(self, "_train_config_path", None) + if yml_path is None: + for fname in ("config.yml", "config.yaml"): + candidate = os.path.join(self._checkpoint_path, fname) + if os.path.exists(candidate): + yml_path = candidate + break + if yml_path is None: + return None + try: + from wall_x.config.loader import load_config + from wall_x.trainer.adapters.base_adapter import load_trainer_data_config + + typed_cfg = load_config(yml_path) + except Exception as e: + logger.warning( + "[InferConfig] Typed config load failed for %s: %s", + yml_path, + e, + ) + return None + try: + return load_trainer_data_config(typed_cfg) + except Exception as e: + logger.warning( + "[InferConfig] Typed config load failed for %s: %s", + yml_path, + e, + ) + return None + + +if __name__ == "__main__": + config = InferConfig() + logger.info("%s", config.train_config) diff --git a/wall_x/infer/logger.py b/wall_x/_vendor/harrix/serving/_wallx_infer/logger.py similarity index 74% rename from wall_x/infer/logger.py rename to wall_x/_vendor/harrix/serving/_wallx_infer/logger.py index ec8bf3f..ebc6079 100644 --- a/wall_x/infer/logger.py +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/logger.py @@ -1,28 +1,4 @@ -""" -Hierarchical Inference Logging System - -Level structure: -- ENV: Environment layer (RealRobotEnv) -- ROBOT: Robot layer (Robot subclasses) -- CONTROLLER: Controller layer (RobotController, RobotCommunication) -- MODEL: Model layer (WallxModelWrapper) -- UTILS: Utility layer (various utility classes) - -Usage examples: - # Method 1: Auto-detect level - from wall_x.infer.logger import get_logger - logger = get_logger(__name__) - logger.info("This is an info message") - - # Method 2: Manually specify level - logger = get_logger(__name__, "ROBOT") - logger.debug("Robot state updated") - - # Method 3: Use shortcut methods - from wall_x.infer.logger import InferLogger - logger = InferLogger.get_robot_logger("DesktopRobot") - logger.warning("Action out of bounds") -""" +"""Layered logging helpers for inference components.""" import logging import sys @@ -36,25 +12,27 @@ try: HAS_COLORLOG = True except ImportError: HAS_COLORLOG = False - print("[WARNING] colorlog not installed. Install with: pip install colorlog") + logging.getLogger(__name__).warning( + "colorlog not installed. Install with: pip install colorlog" + ) class InferLogger: """ - Hierarchical inference logging system + Layered inference logging system """ _loggers = {} _initialized = False - # Level definitions + # Layer identifiers LEVEL_ENV = "ENV" LEVEL_ROBOT = "ROBOT" LEVEL_CONTROLLER = "CONTROLLER" LEVEL_MODEL = "MODEL" LEVEL_UTILS = "UTILS" - # Level color mapping (for terminal output) + # Layer colors (terminal output) LEVEL_COLORS = { LEVEL_ENV: "cyan", LEVEL_ROBOT: "green", @@ -78,9 +56,9 @@ class InferLogger: Args: log_level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL) log_dir: Log file directory - console_output: Whether to output to console - file_output: Whether to output to file - colorful: Whether to use colorful output (requires colorlog) + console_output: Whether to log to console + file_output: Whether to log to file + colorful: Colored console output (requires colorlog) """ if cls._initialized: return @@ -104,11 +82,11 @@ class InferLogger: @classmethod def get_logger(cls, name: str, level: str = None) -> logging.Logger: """ - Get logger for specified level + Get a logger for the given layer Args: - name: Logger name (usually module name or class name) - level: Level identifier (ENV, ROBOT, CONTROLLER, MODEL, UTILS) + name: Logger name (usually module or class name) + level: Layer tag (ENV, ROBOT, CONTROLLER, MODEL, UTILS) Returns: Configured logger instance @@ -116,7 +94,7 @@ class InferLogger: if not cls._initialized: cls.setup() - # Auto-detect level + # Auto-detect layer if level is None: level = cls._detect_level(name) @@ -139,7 +117,7 @@ class InferLogger: console_handler.setLevel(cls.log_level) if cls.colorful: - # Colorful formatting + # Colored formatter color = cls.LEVEL_COLORS.get(level, "white") console_format = ( f"%(log_color)s[%(asctime)s]%(reset)s " @@ -161,7 +139,7 @@ class InferLogger: }, ) else: - # Plain formatting + # Plain formatter console_format = ( f"[%(asctime)s] [{level:^10}] [%(name)s] " f"%(levelname)-8s %(message)s" @@ -191,7 +169,7 @@ class InferLogger: @classmethod def _detect_level(cls, name: str) -> str: - """Auto-detect level based on name""" + """Auto-detect layer from name""" name_lower = name.lower() if "env" in name_lower: @@ -211,32 +189,32 @@ class InferLogger: @classmethod def get_env_logger(cls, name: str = "Environment") -> logging.Logger: - """Get environment layer logger""" + """Get ENV layer logger""" return cls.get_logger(name, cls.LEVEL_ENV) @classmethod def get_robot_logger(cls, name: str = "Robot") -> logging.Logger: - """Get robot layer logger""" + """Get ROBOT layer logger""" return cls.get_logger(name, cls.LEVEL_ROBOT) @classmethod def get_controller_logger(cls, name: str = "Controller") -> logging.Logger: - """Get controller layer logger""" + """Get CONTROLLER layer logger""" return cls.get_logger(name, cls.LEVEL_CONTROLLER) @classmethod def get_model_logger(cls, name: str = "Model") -> logging.Logger: - """Get model layer logger""" + """Get MODEL layer logger""" return cls.get_logger(name, cls.LEVEL_MODEL) @classmethod def get_utils_logger(cls, name: str = "Utils") -> logging.Logger: - """Get utility layer logger""" + """Get UTILS layer logger""" return cls.get_logger(name, cls.LEVEL_UTILS) @classmethod def set_level(cls, level: str): - """Dynamically modify log level for all loggers""" + """Change log level for all loggers""" new_level = getattr(logging, level.upper()) cls.log_level = new_level for logger in cls._loggers.values(): @@ -258,19 +236,19 @@ class InferLogger: # Convenience functions def get_logger(name: str, level: str = None) -> logging.Logger: """ - Convenience function to get logger + Convenience wrapper to get a logger Args: - name: Logger name (usually use __name__) - level: Level identifier (optional, will auto-detect) + name: Logger name (usually __name__) + level: Layer tag (optional; auto-detected if omitted) Returns: Configured logger instance - Usage examples: - from wall_x.infer.logger import get_logger - logger = get_logger(__name__) # Auto-detect level - logger = get_logger(__name__, "ROBOT") # Manually specify level + Example: + from wall_x._vendor.harrix.serving._wallx_infer.logger import get_logger + logger = get_logger(__name__) # auto-detect layer + logger = get_logger(__name__, "ROBOT") # explicit layer """ return InferLogger.get_logger(name, level) @@ -283,17 +261,17 @@ def setup_logger( colorful: bool = True, ): """ - Convenience function to setup logging system + Convenience wrapper to configure logging Args: log_level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL) log_dir: Log file directory - console_output: Whether to output to console - file_output: Whether to output to file - colorful: Whether to use colorful output + console_output: Whether to log to console + file_output: Whether to log to file + colorful: Colored console output - Usage examples: - from wall_x.infer.logger import setup_logger + Example: + from wall_x._vendor.harrix.serving._wallx_infer.logger import setup_logger setup_logger(log_level="DEBUG", log_dir="./logs") """ InferLogger.setup(log_level, log_dir, console_output, file_output, colorful) diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py new file mode 100644 index 0000000..187ea8a --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py @@ -0,0 +1,760 @@ +import os +import torch +import copy +from safetensors.torch import load_file +import numpy as np +from PIL import Image +from qwen_vl_utils.vision_process import smart_resize +from transformers import BatchFeature + +from wall_x.trainer.trainer_utils import load_wallx_processors + +from wall_x._vendor.x2robot_utils.text_templates import ( + preprocesser_call, + get_prologue_with_embodied_information, +) +from wall_x._vendor.x2robot_utils.grounding import ( + reverse_grounding_points, + extract_grounding_points, +) + +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig +from wall_x._vendor.harrix.utils.ckpt_load import reshape_compatible_state_dict +from wall_x._vendor.harrix.utils.train_config import ( + resolve_camera_label, + resolve_max_length, + resolve_state_bins, + resolve_use_state_string_representation, +) +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger +from wall_x.utils.timers import timer, ScopeTimer + + +ENABLE_FAST_PREPROCESS = os.getenv("ENABLE_FAST_PREPROCESS", "False").lower() == "true" + + +def move_to_cuda(obj, device="cuda"): + if isinstance(obj, torch.Tensor): + return obj.to(device) + elif isinstance(obj, (dict, BatchFeature)): + return {k: move_to_cuda(v, device) for k, v in obj.items()} + elif isinstance(obj, list): + return [move_to_cuda(v, device) for v in obj] + elif isinstance(obj, tuple): + return tuple(move_to_cuda(v, device) for v in obj) + else: + return obj + + +class WallxModelWrapper: + def __init__(self, config: InferConfig): + self.config = config + self.logger = InferLogger.get_model_logger("WallxModelWrapper") + self.norm_key = self.config.norm_key + self._register_normalizers() + self.logger.info(f"normalizers {self.norm_key} registered") + self._load_processor() + self._load_model() + self.load_ckpt() + self.logger.info(f"model {self.config.checkpoint_path} loaded") + + self.norm_key = self.config.norm_key + self._register_normalizers() + self.logger.info(f"normalizers {self.norm_key} registered") + + # Initialize robot_type_id (for v3.1 delta tokenizer) + self.robot_type_id = None + if self.tokenizer_mixin is not None: + robot_type = getattr(self.config, "robot_type", None) + if robot_type: + self.tokenizer_mixin.init_inference(robot_type=robot_type) + if hasattr(self.tokenizer_mixin, "robot_type_id"): + self.robot_type_id = self.tokenizer_mixin.robot_type_id + if self.robot_type_id is not None: + self.logger.info( + f"Initialized robot_type_id: {self.robot_type_id} from robot_type: {robot_type}" + ) + + self.role_start_symbol = "<|im_start|>" + self.role_end_symbol = "<|im_end|>" + self.vision_start_symbol = "<|vision_start|>" + self.vision_end_symbol = "<|vision_end|>" + self.image_pad_symbol = "<|image_pad|>" + self.propri_symbol = "<|propri|>" + self.action_symbol = "<|action|>" + + self.cam_names = self.config.cam_names + + def _camera_name_mapping(self): + return self.config.train_config.get("data", {}).get("camera_name_mapping") + + def _load_processor(self): + # Load tokenizer on model_device for inference + device = getattr(self.config, "model_device", "cuda") + processors_dict = load_wallx_processors(self.config.train_config, device=device) + self.processor = processors_dict["processor"] + self.action_mapper = processors_dict["action_mapper"] + self.tokenizer_mixin = processors_dict.get("tokenizer_mixin") + + def _load_model(self): + from wall_x.trainer.adapters import resolve_adapter + + model_type = self.config.train_config["model_type"] + adapter_cls = resolve_adapter(model_type) + ModelClass = adapter_cls.inference_model_class() + self.ModelClass = ModelClass + + self.logger.info(f"initializing model: {model_type} ({ModelClass.__name__})") + + self.model = ModelClass( + self.config.model_config, + self.processor, + self.tokenizer_mixin, + ) + + # log attention implementation - variant-specific layout dispatched + # via the adapter. + adapter_cls.log_attention_implementation(self.logger, self.model) + self.model_type = model_type + + self.logger.info("resizing model token embeddings") + self.model.resize_token_embeddings(len(self.processor.tokenizer)) + self.logger.info("token embedding resize done") + self.logger.info("casting selected params to bfloat16") + self.model.to_bfloat16_for_selected_params() + self.logger.info("bfloat16 cast done") + + def load_ckpt(self, checkpoint_path: str = None): + if checkpoint_path is None: + checkpoint_path = self.config.checkpoint_path + + if os.path.exists(os.path.join(checkpoint_path, "global_step.pth")): + global_step = torch.load(os.path.join(checkpoint_path, "global_step.pth"))[ + "global_step" + ] + self.logger.info(f"checkpoint global_step: {global_step}") + + fsdp_ckpt = os.path.join(checkpoint_path, "pytorch_model_fsdp.bin") + safetensor_ckpt = os.path.join(checkpoint_path, "model.safetensors") + if os.path.exists(fsdp_ckpt): + self.logger.info(f"Loading FSDP checkpoint: {fsdp_ckpt}") + state_dict = torch.load(fsdp_ckpt, map_location="cpu") + # Unwrap outer dict (e.g. {'state_dict': {...}}) + if isinstance(state_dict, dict) and "state_dict" in state_dict: + self.logger.info( + "Using nested state_dict inside pytorch_model_fsdp.bin" + ) + state_dict = state_dict["state_dict"] + elif os.path.exists(safetensor_ckpt): + self.logger.info(f"loading safetensors checkpoint: {safetensor_ckpt}") + state_dict = load_file(safetensor_ckpt, device="cpu") + else: + raise FileNotFoundError( + f"ERROR: No checkpoint found under {checkpoint_path}. " + "Expecting either pytorch_model_fsdp.bin or model.safetensors." + ) + + # Qwen models need fused weight conversion + if not self.ModelClass.is_fused(state_dict): + self.logger.info( + "Converting non-fused weights to fused format...", + ) + state_dict = self.ModelClass.convert_to_fused(state_dict) + else: + self.logger.info( + "The weights is fused, skipping conversion.", + ) + + state_dict = reshape_compatible_state_dict( + state_dict, self.model.state_dict(), log_fn=self.logger.info + ) + msg = self.model.load_state_dict(state_dict, strict=False) + self.model.set_normalizer( + copy.deepcopy(self.normalizer_action), + copy.deepcopy(self.normalizer_propri), + ) + self.logger.info(f"load_state_dict result: {msg}") + self.model.eval() + self.model.to(self.config.model_device) + self.model.to_bfloat16_for_selected_params() + + if hasattr(self.model, "load_optimized_weights"): + self.model.load_optimized_weights(state_dict) + + def _register_normalizers(self): + from wall_x._vendor.harrix.utils.normalizer import build_normalizers + + self.normalizer_action, self.normalizer_propri, resolved = build_normalizers( + self.config.checkpoint_path, + self.config.train_config, + self.config.norm_key, + ) + self.norm_key = resolved + self.config.norm_key = resolved + self._log_norm_debug(self.norm_key) + + def _log_norm_debug(self, norm_key): + if not hasattr(self, "normalizer_action") or not hasattr( + self, "normalizer_propri" + ): + self.logger.warning("[NormDebug] normalizer is not initialized yet") + return + + if ( + norm_key in self.normalizer_action.min + and norm_key in self.normalizer_propri.min + ): + self.logger.debug( + "[NormDebug] norm_key=%s action(min/delta) shape=%s/%s", + norm_key, + tuple(self.normalizer_action.min[norm_key].shape), + tuple(self.normalizer_action.delta[norm_key].shape), + ) + self.logger.debug( + "[NormDebug] action min(all)=%s delta(all)=%s", + self.normalizer_action.min[norm_key].detach().cpu().tolist(), + self.normalizer_action.delta[norm_key].detach().cpu().tolist(), + ) + self.logger.debug( + "[NormDebug] propri(min/delta) shape=%s/%s", + tuple(self.normalizer_propri.min[norm_key].shape), + tuple(self.normalizer_propri.delta[norm_key].shape), + ) + self.logger.debug( + "[NormDebug] propri min(all)=%s delta(all)=%s", + self.normalizer_propri.min[norm_key].detach().cpu().tolist(), + self.normalizer_propri.delta[norm_key].detach().cpu().tolist(), + ) + else: + self.logger.warning( + "[NormDebug] norm_key=%s not found in normalizer", norm_key + ) + + @timer + def construct_model_input( + self, observation, prefix_text, postfix_text, pad_prefix=False + ): + batch_size = len(observation) + dataset_names = [self.norm_key] * batch_size + self.logger.debug("[NormDebug] dataset_names=%s", dataset_names) + + additional_inputs = {} + + # -------- proprioception / masks (batch) -------- + agent_pos_list = [] + agent_pos_mask_list = [] + dof_mask_list = [] + for obs in observation: + if "robot_state_action_data" in obs: + robot_state_action_data = obs["robot_state_action_data"] + + agent_pos = torch.from_numpy(robot_state_action_data.agent_pos) + # Normalize to [1, T, D] for batch cat + if agent_pos.dim() == 2: + agent_pos = agent_pos.unsqueeze(0) + agent_pos_list.append(agent_pos) + + agent_pos_mask = torch.from_numpy( + robot_state_action_data.agent_pos_mask + ) + if agent_pos_mask.dim() == 2: + agent_pos_mask = agent_pos_mask.unsqueeze(0) + agent_pos_mask_list.append(agent_pos_mask) + + dof_mask = torch.from_numpy(robot_state_action_data.dof_mask) + if dof_mask.dim() == 1: + dof_mask = dof_mask.unsqueeze(0) + dof_mask_list.append(dof_mask) + + # cat: [B, T, D] / [B, ...] + if len(agent_pos_list) > 0: + agent_pos = torch.cat(agent_pos_list, dim=0) + agent_pos_mask = torch.cat(agent_pos_mask_list, dim=0) + dof_mask = torch.cat(dof_mask_list, dim=0) + + if self.normalizer_propri is not None: + agent_pos = self.normalizer_propri.normalize_data( + agent_pos, dataset_names + ) + + additional_inputs["proprioception"] = agent_pos.detach() + additional_inputs["agent_pos_mask"] = agent_pos_mask + additional_inputs["dof_mask"] = dof_mask + + # -------- images (flattened, in placeholder scan order) -------- + with ScopeTimer("resize_images"): + # TODO[KC]: optimize this using batch processing + image_inputs = [] + all_image_sizes = [] + if ENABLE_FAST_PREPROCESS: + for obs in observation: + current_image_inputs = self._resize_images_fast(obs) + image_inputs.extend(current_image_inputs) + # tensor shape is (H, W, C), convert to (W, H) to match PIL.size format + all_image_sizes.extend( + [ + (image_i.shape[1], image_i.shape[0]) + for image_i in current_image_inputs + ] + ) + else: + for obs in observation: + current_image_inputs = self._resize_images(obs) + image_inputs.extend(current_image_inputs) + # tensor shape is (H, W, C), convert to (W, H) to match PIL.size format + all_image_sizes.extend( + [ + (image_i.shape[1], image_i.shape[0]) + for image_i in current_image_inputs + ] + ) + + additional_inputs["image_size"] = all_image_sizes + + with ScopeTimer("preprocesser_call"): + inputs = preprocesser_call( + processor=self.model.processor, + prefix_text=prefix_text, + postfix_text=postfix_text, + images=image_inputs, + videos=None, + padding=True, + truncation=True, + return_tensors="pt", + max_length=resolve_max_length(self.config.train_config), + pad_to_128_multiple=False, + pad_prefix_to_same_length=pad_prefix, + norm_state=( + additional_inputs["proprioception"] + if "proprioception" in additional_inputs + and resolve_use_state_string_representation( + self.config.train_config + ) + else None + ), + agent_pos_mask=( + additional_inputs["agent_pos_mask"] + if "agent_pos_mask" in additional_inputs + else None + ), + state_augmentation_prob=0.0, + state_drop_prob=0.0, + state_augmentation_ratio=0.0, + state_bins=resolve_state_bins(self.config.train_config), + inference_mode=True, + ) + + with ScopeTimer("convert_action_token_id and post "): + action_token_id = self.model.processor.tokenizer.convert_tokens_to_ids( + "<|action|>" + ) + flow_action_mask = inputs["input_ids"] == action_token_id + additional_inputs["moe_token_types"] = flow_action_mask + additional_inputs["dataset_names"] = dataset_names + + inputs.update(additional_inputs) + inputs = move_to_cuda(inputs, self.config.model_device) + + return inputs + + def get_text_for_dllm_action(self, instruction): + if ( + self.config.train_config["data"].get("use_embodied_system_prompt_ratio", 0) + > 0 + ): + if self.norm_key != "x2_normal" and self.norm_key != "ex_normal": + self.config.robot_id = 0 + cam_name_mapping = {cam_name: cam_name for cam_name in self.cam_names} + prologue = get_prologue_with_embodied_information( + dataset_name=self.norm_key, + cam_mapping=cam_name_mapping, + robot_id=self.config.robot_id, + uid="", + config=self.config.data_config, + ) + else: + prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n" + user_request = f"{self.role_start_symbol}user\nObservation:" + camera_name_mapping = self._camera_name_mapping() + for cam_name in self.cam_names: + user_request += ( + f" {resolve_camera_label(cam_name, camera_name_mapping)}:" + f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}" + ) + user_request += "\nInstruction:" + text_prompt = f"\nPredict the next action in robot action.\nProprioception: {self.propri_symbol}\n" + user_message = ( + f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n" + ) + placeholder_seq = self.tokenizer_mixin.get_placeholder_for_dllm() + ar_token = "".join(placeholder_seq) + "<|im_end|>\n" + assistant_message = f"{self.role_start_symbol}assistant\n{ar_token}" + flow_action = f"{self.action_symbol * self.config.action_horizon}" + + prefix_text = prologue + user_message + assistant_message + postfix_text = flow_action + + return prefix_text, postfix_text + + @timer + def get_text_for_action(self, instruction): + if ( + self.config.train_config["data"].get("use_embodied_system_prompt_ratio", 0) + > 0 + ): + if self.norm_key not in ["x2_normal", "ex_normal"]: + self.config.robot_id = 0 + cam_name_mapping = {cam_name: cam_name for cam_name in self.cam_names} + prologue = get_prologue_with_embodied_information( + dataset_name=self.norm_key, + cam_mapping=cam_name_mapping, + robot_id=self.config.robot_id, + uid="", + config=self.config.data_config, + ) + else: + prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n" + user_request = f"{self.role_start_symbol}user\nObservation:" + camera_name_mapping = self._camera_name_mapping() + for cam_name in self.cam_names: + user_request += ( + f" {resolve_camera_label(cam_name, camera_name_mapping)}:" + f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}" + ) + user_request += "\nInstruction:" + text_prompt = f"\nPredict the next action in robot action.\nProprioception: {self.propri_symbol}\n" + user_message = ( + f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n" + ) + assistant_message = f"{self.role_start_symbol}assistant\n" + flow_action = f"{self.action_symbol * self.config.action_horizon}" + + prefix_text = prologue + user_message + assistant_message + postfix_text = flow_action + + return prefix_text, postfix_text + + @timer + def get_text_for_subtask_generation(self, instruction): + prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n" + user_request = f"{self.role_start_symbol}user\nObservation:" + camera_name_mapping = self._camera_name_mapping() + for cam_name in self.cam_names: + user_request += ( + f" {resolve_camera_label(cam_name, camera_name_mapping)}:" + f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}" + ) + user_request += "\nInstruction:" + text_prompt = "\nPredict the next action in language.\n" + user_message = ( + f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n" + ) + assistant_message = f"{self.role_start_symbol}assistant\n" + + prefix_text = prologue + user_message + assistant_message + postfix_text = "" + + return prefix_text, postfix_text + + @timer + def _resize_images(self, observation): + image_inputs = [] + for key in self.cam_names: + if key not in observation: + continue + current_obs = observation[key] + if isinstance(current_obs, np.ndarray): + img_pil = Image.fromarray(current_obs) + elif isinstance(current_obs, Image.Image): + img_pil = current_obs + else: + raise ValueError(f"Unsupported image type: {type(current_obs)}") + orig_width, orig_height = img_pil.size + + target_size = self.config.data_config.resolution.get(key, -1) + if target_size != -1: + # Aspect-ratio-preserving resize + if orig_width > orig_height: # landscape + new_width = target_size + new_height = int(target_size * orig_height / orig_width) + else: # portrait + new_height = target_size + new_width = int(target_size * orig_width / orig_height) + img_pil = img_pil.resize((new_width, new_height)) + + # Apply smart resize (Qwen logic) + current_width, current_height = img_pil.size + resized_height, resized_width = smart_resize( + current_height, + current_width, + factor=self.config.data_config.image_factor, + min_pixels=self.config.data_config.min_pixels, + max_pixels=self.config.data_config.max_pixels, + ) + resized_img = img_pil.resize((resized_width, resized_height)) + resized_img = torch.from_numpy(np.array(resized_img)).to( + self.config.model_device + ) + image_inputs.append(resized_img) + + return image_inputs + + def _resize_images_fast(self, observation): + import cv2 + + image_inputs = [] + for key in self.cam_names: + if key not in observation: + continue + current_obs = observation[key] + orig_height, orig_width, _ = current_obs.shape + + target_size = self.config.data_config.resolution.get(key, -1) + current_width, current_height = orig_width, orig_height + if target_size != -1: + # Aspect-ratio-preserving resize + if orig_width > orig_height: # landscape + new_width = target_size + new_height = int(target_size * orig_height / orig_width) + else: # portrait + new_height = target_size + new_width = int(target_size * orig_width / orig_height) + current_width = new_width + current_height = new_height + + # Apply smart resize (Qwen logic) + resized_height, resized_width = smart_resize( + current_height, + current_width, + factor=self.config.data_config.image_factor, # FIXME + min_pixels=self.config.data_config.min_pixels, # FIXME + max_pixels=self.config.data_config.max_pixels, # FIXME + ) + + resized_img = cv2.resize( + current_obs, + (resized_width, resized_height), + interpolation=cv2.INTER_CUBIC, + ) + resized_img = torch.from_numpy(resized_img).to(self.config.model_device) + image_inputs.append(resized_img) + + return image_inputs + + def infer_flow_action(self, observation, instruction): + self.logger.info("generating flow action") + self.logger.info(f"flow action instruction: {instruction}") + + prefix_text, postfix_text = self.get_text_for_action(instruction) + model_input = self.construct_model_input( + [observation], [prefix_text], [postfix_text] + ) + + padding = ( + torch.zeros_like( + self.normalizer_action.delta[model_input["dataset_names"][0]] + ) + .unsqueeze(0) + .to("cpu") + ) + padding_action = self.normalizer_action.normalize_data( + padding, model_input["dataset_names"] + ).to(model_input["input_ids"].device) + + self.logger.info( + "generate_flow_action start (horizon=%s, flow_steps=%s, device=%s)", + self.config.action_horizon, + self.config.num_inference_timesteps, + self.config.model_device, + ) + if torch.cuda.is_available(): + try: + free_b, total_b = torch.cuda.mem_get_info( + torch.device(self.config.model_device) + ) + self.logger.info( + "CUDA mem before flow: free=%.2f GiB / total=%.2f GiB", + free_b / (1024**3), + total_b / (1024**3), + ) + except Exception as e: + self.logger.warning("CUDA mem_get_info failed: %s", e) + + with ScopeTimer("generate_flow_action"): + model_output = self.model.generate_flow_action( + action_horizon=self.config.action_horizon, + action_dim=self.config.action_dim, + num_inference_timesteps=self.config.num_inference_timesteps, + padding_action=padding_action, + **model_input, + ) + + self.logger.info("flow action generation done") + + model_output["robot_state_action_data"] = observation["robot_state_action_data"] + model_output["robot_state_action_data"].save_action_data( + model_output["predict_action"] + ) + self.logger.info("saved flow action to robot_state_action_data") + return model_output + + def infer_flow_action_batch(self, observations, instructions): + """ + Batch flow action inference: + - observations: List[Dict], same format as single inference + - instructions: List[str], aligned with observations + Returns List[model_output] of length batch size + """ + assert len(observations) == len( + instructions + ), "observations and instructions must have the same length" + batch_size = len(observations) + + prefix_list = [] + postfix_list = [] + for ins in instructions: + prefix_text, postfix_text = self.get_text_for_action(ins) + prefix_list.append(prefix_text) + postfix_list.append(postfix_text) + + # Build batch input in one preprocesser_call; avoid per-sample cat + batch_inputs = self.construct_model_input( + observations, prefix_list, postfix_list + ) + + padding_list = [] + for ds_name in batch_inputs["dataset_names"]: + padding = ( + torch.zeros_like(self.normalizer_action.delta[ds_name]) + .unsqueeze(0) + .to("cpu") + ) + padding_list.append(padding) + padding = torch.cat(padding_list, dim=0) + padding_action = self.normalizer_action.normalize_data( + padding, batch_inputs["dataset_names"] + ).to(batch_inputs["input_ids"].device) + + with ScopeTimer("generate_flow_action_batch"): + model_output = self.model.generate_flow_action( + action_horizon=self.config.action_horizon, + action_dim=self.config.action_dim, + num_inference_timesteps=self.config.num_inference_timesteps, + padding_action=padding_action, + **batch_inputs, + ) + + predict_action = model_output["predict_action"] # [B, H, D] + + outputs = [] + for i in range(batch_size): + single_action = predict_action[i : i + 1] + single_output = { + "predict_action": single_action, + "robot_state_action_data": observations[i]["robot_state_action_data"], + } + single_output["robot_state_action_data"].save_action_data( + single_output["predict_action"] + ) + outputs.append(single_output) + + return outputs + + def infer_ar_action(self, observation, instruction): + self.logger.info("generating ar action") + self.logger.info(f"ar action instruction: {instruction}") + + prefix_text, _ = self.get_text_for_action(instruction) + model_input = self.construct_model_input([observation], [prefix_text], [""]) + + model_output = self.model.generate_ar_action( + action_horizon=self.config.action_horizon, + action_dim=self.config.ar_action_dim, + num_inference_timesteps=self.config.num_inference_timesteps, + robot_type_id=self.robot_type_id, + **model_input, + ) + self.logger.info("ar action generation done") + + model_output["robot_state_action_data"] = observation["robot_state_action_data"] + model_output["robot_state_action_data"].save_action_data( + model_output["predict_action"] + ) + self.logger.info("saved ar action to robot_state_action_data") + return model_output + + def infer_subtask(self, observation, instruction): + self.logger.info("generating subtask") + self.logger.info(f"subtask instruction: {instruction}") + + prefix_text, postfix_text = self.get_text_for_subtask_generation(instruction) + model_input = self.construct_model_input([observation], [prefix_text], [""]) + + model_output = self.model.generate_text(**model_input) + subtask = model_output["predict_output_text"][0].split("<|im_end|>")[0].strip() + self.logger.info(f"subtask generation done, subtask: {subtask}") + return subtask + + def infer_vqa(self, observation, instruction): + self.logger.info("generating vqa answer") + + if isinstance(observation["multi_modal"], list): + orig_size = observation["multi_modal"][0].size + else: + orig_size = observation["multi_modal"].size + + prefix_text = instruction + model_input = self.construct_model_input([observation], [prefix_text], [""]) + + model_output = self.model.generate_text(**model_input) + answer = model_output["predict_output_text"][0].split("<|im_end|>")[0].strip() + self.logger.info(f"vqa answer done, answer: {answer}") + answer = reverse_grounding_points( + answer, + orig_size[1], + orig_size[0], + model_input["image_size"][0][1], + model_input["image_size"][0][0], + self.config.data_config.model_type, + ) + points = extract_grounding_points(answer) + + return { + "answer": answer, + "points": points, + } + + def infer_dllm_action( + self, observation, instruction, use_ar_action=False, dataset_name="x2_normal" + ): + assert self.model_type in ["qwen2_5"], "DLLM only supports qwen2_5" + prefix_text, postfix_text = self.get_text_for_dllm_action(instruction) + model_input = self.construct_model_input( + [observation], [prefix_text], [postfix_text], pad_prefix=True + ) + model_input = self.model.update_infer_dllm_position_mask(model_input) + total_ar_step = self.tokenizer_mixin.inference_ar_steps_for_dllm + cnt = 0 + while cnt < 3: # fast mode: at most 3 retries + model_output = self.model.generate_dllm_action( + action_horizon=self.config.action_horizon, + action_dim=self.config.action_dim, + ar_action_dim=self.config.ar_action_dim, + num_inference_timesteps=self.config.num_inference_timesteps, + use_ar_action=use_ar_action, + total_ar_step=total_ar_step, + robot_type_id=self.robot_type_id, # for v3.1 delta decode + **model_input, + ) + if model_output["predict_action"] is not None: + break + cnt += 1 + self.logger.warning(f"dllm action generation failed, retry {cnt}") + self.logger.info("dllm action generation done") + + model_output["robot_state_action_data"] = observation["robot_state_action_data"] + model_output["robot_state_action_data"].save_action_data( + model_output["predict_action"] + ) + self.logger.info("saved dllm action to robot_state_action_data") + return model_output diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py b/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py new file mode 100644 index 0000000..641bbc4 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py @@ -0,0 +1,1181 @@ +import numpy as np +from abc import ABC +import re + +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig +from wall_x._vendor.harrix.serving._wallx_infer.utils import VehiclePoseHandler, UnifiedTrajectoryProcessor +from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import ( + RobotStateActionData, + dof_dims, +) +from wall_x._vendor.harrix.serving._wallx_infer.socket_controller import RobotController + +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger + +robot_action_key_mapping = { + "follow_left_ee_cartesian_pos": "follow1_pos[:3]", + "follow_left_ee_rotation": "follow1_pos[3:6]", + "follow_left_gripper": "follow1_pos[6:7]", + "follow_right_ee_cartesian_pos": "follow2_pos[:3]", + "follow_right_ee_rotation": "follow2_pos[3:6]", + "follow_right_gripper": "follow2_pos[6:7]", + "head_actions": "head_pos", + "head_rotation": "head_pos", # match ex001 + "height": "lift", + "velocity_decomposed": "velocity_decomposed", + "velocity_decomposed_odom": "velocity_decomposed_odom", # match ex001 + "follow_left_arm_joint_cur": "follow1_joints_cur", # 1 -> 7 + "follow_right_arm_joint_cur": "follow2_joints_cur", # 1 -> 7 + "follow_left_gripper_cur": "follow1_joints_cur[-1:]", + "follow_right_gripper_cur": "follow2_joints_cur[-1:]", + "follow_left_arm_joint_pos": "follow1_joints", + "follow_right_arm_joint_pos": "follow2_joints", + "follow_left_wrench_ext_local_force": "follow1_end_effort_force", + "follow_left_wrench_ext_local_torque": "follow1_end_effort_torque", + "follow_right_wrench_ext_local_force": "follow2_end_effort_force", + "follow_right_wrench_ext_local_torque": "follow2_end_effort_torque", + "follow_left_wrench_ext_local_force_from_joint": "follow1_end_effort_force_from_joint", + "follow_left_wrench_ext_local_torque_from_joint": "follow1_end_effort_torque_from_joint", + "follow_right_wrench_ext_local_force_from_joint": "follow2_end_effort_force_from_joint", + "follow_right_wrench_ext_local_torque_from_joint": "follow2_end_effort_torque_from_joint", + "follow_left_wrench_ext_world_force": "follow1_wrench_ext_world_force", + "follow_left_wrench_ext_world_torque": "follow1_wrench_ext_world_torque", + "follow_left_wrench_ext_world_force_from_joint": "follow1_wrench_ext_world_force_from_joint", + "follow_left_wrench_ext_world_torque_from_joint": "follow1_wrench_ext_world_torque_from_joint", + "follow_right_wrench_ext_world_force": "follow2_wrench_ext_world_force", + "follow_right_wrench_ext_world_torque": "follow2_wrench_ext_world_torque", + "follow_right_wrench_ext_world_force_from_joint": "follow2_wrench_ext_world_force_from_joint", + "follow_right_wrench_ext_world_torque_from_joint": "follow2_wrench_ext_world_torque_from_joint", + "follow_left_arm_joint_dev": "follow1_joints_dev", + "follow_right_arm_joint_dev": "follow2_joints_dev", +} + + +def _parse_follow_pos(follow_pos) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Parse websocket ``follow{1,2}_pos``: pos3 + euler3 + one or more gripper dims.""" + arr = np.asarray(follow_pos, dtype=np.float64).reshape(-1) + if arr.shape[0] < 7: + raise ValueError(f"follow_pos expects 7 dims, got shape {arr.shape}") + return arr[:3], arr[3:6], arr[6:] + + +def ingest_websocket_follow_pos( + state: dict, + config: InferConfig, + robot_state_action_data: RobotStateActionData, +) -> None: + """Convert client ``follow1_pos``/``follow2_pos`` (14D) into train-config proprio.""" + from wall_x._vendor.harrix.utils.train_config import resolve_agent_pos_config + from wall_x._vendor.x2robot_utils import geometry as geom + + agent_cfg = resolve_agent_pos_config(config.train_config) + arm_ws_keys = ( + ("left", "follow1_pos", "follow_left_"), + ("right", "follow2_pos", "follow_right_"), + ) + for _side, ws_key, cfg_prefix in arm_ws_keys: + if ws_key not in state or state[ws_key] is None: + continue + pos, euler, grip = _parse_follow_pos(state[ws_key]) + rot6d = geom.euler_to_matrix_zyx_6d_nb(euler.reshape(1, 3)) + for key in agent_cfg: + if key == "action_padding" or not key.startswith(cfg_prefix): + continue + if "cartesian_pos" in key: + robot_state_action_data.save_state_data_with_key(pos, key) + elif "rotation_6d" in key.lower(): + robot_state_action_data.save_state_data_with_key( + rot6d.reshape(1, -1), key + ) + elif "rotation" in key: + robot_state_action_data.save_state_data_with_key(euler, key) + elif "gripper" in key: + robot_state_action_data.save_state_data_with_key( + grip[: int(agent_cfg[key])], key, gt_dim=int(agent_cfg[key]) + ) + + for key, dim in agent_cfg.items(): + if key == "action_padding" or key.startswith(("follow_left_", "follow_right_")): + continue + norm_key = key.replace("follow_", "").replace("master_", "") + if robot_state_action_data.data.get(f"state_{norm_key}") is None: + robot_state_action_data.save_state_data_with_key( + np.zeros(int(dim), dtype=np.float64), key + ) + + +def export_websocket_follow_pos( + robot_state_action_data: RobotStateActionData, +) -> dict[str, list]: + """Convert model ``robot_state_action_data`` back to ``follow1_pos``/``follow2_pos``.""" + left_arm_action = _stack_state_action_series(robot_state_action_data, side="left") + right_arm_action = _stack_state_action_series(robot_state_action_data, side="right") + return { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + + +def _zero_state_series(key: str) -> np.ndarray: + """Fallback proprio when websocket state omits an arm (e.g. only follow2_pos).""" + dim = dof_dims.get(key, 1) + return np.zeros((1, dim), dtype=np.float64) + + +def _identity_state_rotation_6d() -> np.ndarray: + from wall_x._vendor.x2robot_utils import geometry as geom + + return geom.euler_to_matrix_zyx_6d_nb(np.zeros((1, 3), dtype=np.float64)) + + +def _as_2d_series(value, name: str) -> np.ndarray: + """Normalize state/action arrays to (T, D) for follow*_pos serialization.""" + arr = np.asarray(value, dtype=np.float64) + if arr.ndim == 0: + arr = arr.reshape(1, 1) + elif arr.ndim == 1: + arr = arr.reshape(1, -1) if arr.size <= 7 else arr.reshape(-1, 1) + elif arr.ndim > 2: + arr = arr.reshape(arr.shape[0], -1) + if arr.ndim != 2: + raise ValueError(f"{name}: expected 2D series after normalize, got {arr.shape}") + return arr + + +def _stack_state_action_series( + robot_state_action_data: RobotStateActionData, + *, + side: str, + pos_key: str = "ee_cartesian_pos", + rot_key: str = "ee_rotation", + grip_key: str = "gripper", +) -> np.ndarray: + """Build (1+horizon, 7) follow arm pose: pos3 + rot3 + grip1.""" + from wall_x._vendor.x2robot_utils import geometry as geom + + data = robot_state_action_data.data + prefix = "left" if side == "left" else "right" + + state_pos = data.get(f"state_{prefix}_{pos_key}") + action_pos = data.get(f"action_{prefix}_{pos_key}") + rel_pos = data.get(f"action_{prefix}_{pos_key}_relative") + if action_pos is None and rel_pos is not None: + base_pos = ( + _as_2d_series(state_pos, f"state_{prefix}_{pos_key}") + if state_pos is not None + else _zero_state_series(f"{prefix}_{pos_key}") + ) + action_pos = base_pos + np.asarray(rel_pos, dtype=np.float64) + if state_pos is None: + state_pos = _zero_state_series(f"{prefix}_{pos_key}") + + state_rot = data.get(f"state_{prefix}_{rot_key}") + action_rot = data.get(f"action_{prefix}_{rot_key}") + + rot6d = data.get(f"action_{prefix}_{rot_key}_6D") + if action_rot is None and rot6d is not None: + action_rot = geom.so3_to_euler_zyx_batch_nb(np.asarray(rot6d, dtype=np.float64)) + if action_rot is None: + rel6d = data.get(f"action_{prefix}_{rot_key}_6D_relative") + state6d = data.get(f"state_{prefix}_{rot_key}_6D") + if state6d is None and state_rot is not None: + state6d = geom.euler_to_matrix_zyx_6d_nb( + _as_2d_series(state_rot, f"state_{prefix}_{rot_key}") + ) + if state6d is None: + state6d = _identity_state_rotation_6d() + if rel6d is not None: + rot6d = geom.compose_state_and_delta_to_abs_6d( + np.asarray(rel6d, dtype=np.float64), + np.asarray(state6d, dtype=np.float64).reshape(-1), + ) + action_rot = geom.so3_to_euler_zyx_batch_nb(rot6d) + + if state_rot is None: + state6d = data.get(f"state_{prefix}_{rot_key}_6D") + if state6d is not None: + state_rot = geom.so3_to_euler_zyx_batch_nb( + np.asarray(state6d, dtype=np.float64) + ) + else: + state_rot = _zero_state_series(f"{prefix}_{rot_key}") + + state_grip = data.get(f"state_{prefix}_{grip_key}") + if state_grip is None: + state_grip = _zero_state_series(f"{prefix}_{grip_key}") + action_grip = data.get(f"action_{prefix}_{grip_key}") + + if action_pos is None or action_rot is None or action_grip is None: + missing = [ + name + for name, value in ( + (f"action_{prefix}_{pos_key}", action_pos), + (f"action_{prefix}_{rot_key}", action_rot), + (f"action_{prefix}_{grip_key}", action_grip), + ) + if value is None + ] + raise ValueError( + f"Cannot serialize {prefix} arm action; missing action fields: {missing}" + ) + + pos = np.concatenate( + [ + _as_2d_series(state_pos, f"state_{prefix}_{pos_key}"), + _as_2d_series(action_pos, f"action_{prefix}_{pos_key}"), + ], + axis=0, + ) + rot = np.concatenate( + [ + _as_2d_series(state_rot, f"state_{prefix}_{rot_key}"), + _as_2d_series(action_rot, f"action_{prefix}_{rot_key}"), + ], + axis=0, + ) + grip = np.concatenate( + [ + _as_2d_series(state_grip, f"state_{prefix}_{grip_key}"), + _as_2d_series(action_grip, f"action_{prefix}_{grip_key}"), + ], + axis=0, + ) + return np.concatenate([pos, rot, grip], axis=1) + + +def _views_to_camera_observation( + config: InferConfig, views: dict +) -> dict[str, np.ndarray]: + """Map websocket camera keys to model observation keys. + + Only cameras listed in ``config.cam_names`` are required (e.g. LIBERO uses + face_view + right_wrist_view without left_wrist_view). + """ + camera_mappings = { + "face_view": [config.camera_front_key, "face_view", "front_view"], + "left_wrist_view": [ + config.camera_left_key, + "left_wrist_view", + "left_view", + ], + "right_wrist_view": [ + config.camera_right_key, + "right_wrist_view", + "right_view", + ], + } + observation: dict[str, np.ndarray] = {} + for obs_key in config.cam_names: + possible_keys = camera_mappings.get(obs_key, [obs_key]) + for view_key in possible_keys: + if view_key in views and views[view_key] is not None: + view_data = views[view_key] + if ( + isinstance(view_data, np.ndarray) + and view_data.ndim == 4 + and view_data.shape[0] == 1 + ): + observation[obs_key] = view_data[0] + else: + observation[obs_key] = view_data + break + else: + raise KeyError(f"Missing view for {obs_key!r}, tried keys: {possible_keys}") + return observation + + +class Robot(ABC): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + + self.robot_controller = RobotController( + robot_id=robot_id, host=config.robot_host, port=config.robot_port + ) + + self.robot_controller.connect() + + self.is_received = False # blocking flag + self.current_robot_state_action = None + + self.logger = InferLogger.get_robot_logger("Robot") + + def _get_views_and_state(self): + state = self.robot_controller.recv_action() + views = self.robot_controller.recv_image( + [ + self.config.camera_left_key, + self.config.camera_front_key, + self.config.camera_right_key, + ] + ) + self.is_received = True + return state, views + + def _get_dof_mask(self): + dof_config = self.config.train_config["dof_config"] + total_dof = sum(dof_config.values()) + dof_mask = np.ones((1, self.config.action_horizon, total_dof)) + return dof_mask + + @staticmethod + def _try_compose_arm_joint_pos_with_gripper(state: dict, key: str): + """Build 7D arm joint position from runtime state when possible. + + Why this is needed: + - Joint-control training expects `*_arm_joint_pos` to be 7D + (6 arm joints + 1 gripper). + - Some online robot states expose arm joints as 6D (`follow*_joints`) + and gripper separately in `follow*_pos[6]`. + - If we pass raw 6D joints directly, `save_state_data_with_key` will fail + shape validation against dof_dims=7. + + Behavior: + - Only handles `follow_left_arm_joint_pos` / `follow_right_arm_joint_pos`. + - Returns a composed 7D vector when both required sources exist and shapes + match expected runtime format. + - Returns None for all other keys or incompatible state shapes, so caller + falls back to original mapping logic. + """ + + if key == "follow_left_arm_joint_pos": + joints_key = "follow1_joints" + ee_key = "follow1_pos" + elif key == "follow_right_arm_joint_pos": + joints_key = "follow2_joints" + ee_key = "follow2_pos" + else: + return None + + if joints_key not in state or ee_key not in state: + return None + + joints = np.asarray(state[joints_key]).reshape(-1) + ee = np.asarray(state[ee_key]).reshape(-1) + if joints.shape[0] == 6 and ee.shape[0] >= 7: + return np.concatenate([joints, ee[6:7]], axis=0) + + return None + + def get_observation(self): + state, views = self._get_views_and_state() + + robot_state_action_data = RobotStateActionData(config=self.config) + for key in robot_action_key_mapping.keys(): + state_key_str = robot_action_key_mapping[key] + value = None + # Prefer a robust 7D composition path for arm_joint_pos: + # online state often provides `follow*_joints` as 6D plus gripper + # in `follow*_pos[6]`. Compose first to match model/dof schema. + composed_value = self._try_compose_arm_joint_pos_with_gripper(state, key) + if composed_value is not None: + value = composed_value + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + continue + + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + return { + "robot_state_action_data": robot_state_action_data, + **_views_to_camera_observation(self.config, views), + } + + def go_home(self): + if self.current_robot_state_action is None: + self.get_observation() + + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "left_gripper" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "right_gripper" + ) + + action_dict = {"robot_state_action_data": self.current_robot_state_action} + self.apply_action( + action_dict, + robot_action_interpolate_multiplier=150, + robot_action_start_ratio=0, + robot_action_end_ratio=1, + ) + self.logger.info("robot returned to initial pose") + + def _get_left_arm_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + if not self.config.robot_use_joint_angle_control: + return _stack_state_action_series(robot_state_action_data, side="left") + left_arm_joint_pos = np.concatenate( + [ + _as_2d_series( + robot_state_action_data.data["state_left_arm_joint_pos"], + "state_left_arm_joint_pos", + ), + _as_2d_series( + robot_state_action_data.data["action_left_arm_joint_pos"], + "action_left_arm_joint_pos", + ), + ], + axis=0, + ) + return left_arm_joint_pos + + def _get_right_arm_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + if not self.config.robot_use_joint_angle_control: + return _stack_state_action_series(robot_state_action_data, side="right") + right_arm_joint_pos = np.concatenate( + [ + _as_2d_series( + robot_state_action_data.data["state_right_arm_joint_pos"], + "state_right_arm_joint_pos", + ), + _as_2d_series( + robot_state_action_data.data["action_right_arm_joint_pos"], + "action_right_arm_joint_pos", + ), + ], + axis=0, + ) + return right_arm_joint_pos + + def get_serialized_actions( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> dict: + """Build trimmed+interpolated serialized actions from model_output (with robot_state_action_data); does not send. + + Used by apply_action and serving policy; returns a dict ready for robot or client. + """ + assert "robot_state_action_data" in input + + left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) + right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) + + # Trim action + action_length = len(left_arm_action) + if robot_action_start_ratio is None: + robot_action_start_ratio = self.config.robot_action_start_ratio + if robot_action_end_ratio is None: + robot_action_end_ratio = self.config.robot_action_end_ratio + start_frame = int(robot_action_start_ratio * action_length) + end_frame = int(robot_action_end_ratio * action_length) + + left_arm_action = left_arm_action[start_frame:end_frame] + right_arm_action = right_arm_action[start_frame:end_frame] + + # Interpolate + if robot_action_interpolate_multiplier is None: + robot_action_interpolate_multiplier = ( + self.config.robot_action_interpolate_multiplier + ) + target_length = robot_action_interpolate_multiplier * len(left_arm_action) + left_arm_action, right_arm_action = ( + UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [left_arm_action, right_arm_action], target_length + ) + ) + + if not self.config.robot_use_joint_angle_control: + return { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + return { + "follow1_joints": left_arm_action.tolist(), + "follow2_joints": right_arm_action.tolist(), + } + + def apply_action( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> None: + serialized_actions = self.get_serialized_actions( + input, + robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, + robot_action_start_ratio=robot_action_start_ratio, + robot_action_end_ratio=robot_action_end_ratio, + ) + self._send_actions(serialized_actions) + + def _send_actions(self, serialized_actions: dict) -> None: + """Send actions to robot controller""" + self.robot_controller.robot_comm.send_dict(serialized_actions) + self.is_received = False + self.current_robot_state_action = None + + +class DesktopRobot(Robot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.is_received = False # blocking flag + self.current_robot_state_action = None + + self.logger = InferLogger.get_robot_logger("Robot") + + def _get_dof_mask(self): + from wall_x._vendor.harrix.utils.train_config import resolve_dof_config + + dof_config = resolve_dof_config(self.config.train_config) + total_dof = sum(dof_config.values()) + dof_mask = np.ones((1, self.config.action_horizon, total_dof)) + + # For desktop, mask out head_actions, height, and velocity_decomposed + mask_keys = ["head_actions", "height", "velocity_decomposed"] + start_idx = 0 + for key, dof_size in dof_config.items(): + if key in mask_keys: + # Set mask for these dims to 0 + dof_mask[:, :, start_idx : start_idx + dof_size] = 0 + start_idx += dof_size + + return dof_mask + + +class DesktopRobotVGA(Robot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + super().__init__(config, robot_id) + self.is_received = False # blocking flag + self.current_robot_state_action = None + + self.logger = InferLogger.get_robot_logger("Robot") + + def _get_dof_mask(self): + dof_config = self.config.train_config["data"]["dof_config"] + total_dof = sum(dof_config.values()) + dof_mask = np.ones((1, self.config.action_horizon, total_dof)) + + # For desktop, mask out head_actions, height, and velocity_decomposed + mask_keys = ["head_actions", "height", "velocity_decomposed"] + start_idx = 0 + for key, dof_size in dof_config.items(): + if key in mask_keys: + # Set mask for these dims to 0 + dof_mask[:, :, start_idx : start_idx + dof_size] = 0 + start_idx += dof_size + + return dof_mask + + +class DesktopRobotPreprocessor(DesktopRobot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + + self.is_received = False # blocking flag + self.current_robot_state_action = None + self.logger = InferLogger.get_robot_logger("Robot") + + def _send_actions(self, serialized_actions: dict) -> None: + # bypass + return + + def get_observation(self, state, views): + robot_state_action_data = RobotStateActionData(config=self.config) + if "follow1_pos" in state or "follow2_pos" in state: + ingest_websocket_follow_pos(state, self.config, robot_state_action_data) + for key in robot_action_key_mapping.keys(): + if key.startswith(("follow_left_", "follow_right_")): + continue + state_key_str = robot_action_key_mapping[key] + value = None + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + return { + "robot_state_action_data": robot_state_action_data, + **_views_to_camera_observation(self.config, views), + } + + +class TurtleRobot(Robot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.vehicle_pose_handler = VehiclePoseHandler() + self.last_speed = [0, 0, 0] + + def _get_state_from_controller(self): + state = super()._get_state_from_controller() + + self.vehicle_pose_handler.update_pose(state["car_pose"]) + state["velocity_decomposed"] = self.last_speed + + if self.config.turtle_as_desktop: + state["head_pos"] = [0, -1] + state["lift"] = [0.4] + + return state + + def _calculate_car_pose(self, base_velocity_pred): + # Vectorized integrate velocity to pose + dt = 1 / 20 + current_pose = ( + self.vehicle_pose_handler.current_pose.copy() + if self.vehicle_pose_handler.current_pose is not None + else np.array([0.0, 0.0, 0.0]) + ) + + # Batch integrate positions + poses_frames = [] + for i in range(len(base_velocity_pred)): + current_pose = self.vehicle_pose_handler.velocity_to_pose( + base_velocity_pred[i, 0], + base_velocity_pred[i, 1], + base_velocity_pred[i, 2], + dt, + current_pose, + ) + poses_frames.append(current_pose.copy()) + + return poses_frames + + def _get_car_velocity( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + car_velocity = robot_state_action_data.data["action_velocity_decomposed"] + self.last_speed = car_velocity[-1, :].copy().tolist() + return car_velocity + + def _get_head_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + head_actions = np.concatenate( + [ + robot_state_action_data.data["state_head_actions"], + robot_state_action_data.data["action_head_actions"], + ], + axis=0, + ) + return head_actions + + def _get_height_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + height = np.concatenate( + [ + robot_state_action_data.data["state_height"], + robot_state_action_data.data["action_height"], + ], + axis=0, + ) + return height + + def go_home(self): + if self.current_robot_state_action is None: + self.get_observation() + + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "left_gripper" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "right_gripper" + ) + + self.current_robot_state_action.save_action_data_with_key( + np.array([[0, -1]]), "head_actions" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([[0.4]]), "height" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([self.last_speed]), "velocity_decomposed" + ) + + action_dict = {"robot_state_action_data": self.current_robot_state_action} + self.apply_action( + action_dict, + robot_action_interpolate_multiplier=150, + robot_action_start_ratio=0, + robot_action_end_ratio=1, + ) + self.logger.info("robot returned to initial pose") + + def get_serialized_actions( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> dict: + """Build serialized Turtle actions from model_output (dual arms + head/lift/car_pose); does not send.""" + assert "robot_state_action_data" in input + + left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) + right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) + action_length = len(left_arm_action) + if robot_action_start_ratio is None: + robot_action_start_ratio = self.config.robot_action_start_ratio + if robot_action_end_ratio is None: + robot_action_end_ratio = self.config.robot_action_end_ratio + start_frame = int(robot_action_start_ratio * action_length) + end_frame = int(robot_action_end_ratio * action_length) + + left_arm_action = left_arm_action[start_frame:end_frame] + right_arm_action = right_arm_action[start_frame:end_frame] + if robot_action_interpolate_multiplier is None: + robot_action_interpolate_multiplier = ( + self.config.robot_action_interpolate_multiplier + ) + target_length = robot_action_interpolate_multiplier * action_length + ( + left_arm_action, + right_arm_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [left_arm_action, right_arm_action], + target_length, + ) + serialized_actions = { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + + if not self.config.turtle_as_desktop: + head_action = self._get_head_action(input["robot_state_action_data"]) + height_action = self._get_height_action(input["robot_state_action_data"]) + car_velocity_action = self._get_car_velocity( + input["robot_state_action_data"] + ) + head_action = head_action[start_frame:end_frame] + height_action = height_action[start_frame:end_frame] + car_velocity_action = car_velocity_action[start_frame:end_frame] + car_pose_action = np.array(self._calculate_car_pose(car_velocity_action)) + ( + head_action, + height_action, + car_pose_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [head_action, height_action, car_pose_action], + target_length, + ) + serialized_actions["head_pos"] = head_action.tolist() + serialized_actions["lift"] = height_action.tolist() + serialized_actions["car_pose"] = car_pose_action.tolist() + else: + serialized_actions["head_pos"] = [ + [0, -1] for _ in range(len(left_arm_action)) + ] + serialized_actions["lift"] = [0.4 for _ in range(len(left_arm_action))] + serialized_actions["car_pose"] = [ + [0.0, 0.0, 0.0] for _ in range(len(left_arm_action)) + ] + return serialized_actions + + def apply_action( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> None: + serialized_actions = self.get_serialized_actions( + input, + robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, + robot_action_start_ratio=robot_action_start_ratio, + robot_action_end_ratio=robot_action_end_ratio, + ) + self._send_actions(serialized_actions) + + +class TurtleRobotPreprocessor(TurtleRobot): + """Turtle preprocessor: no real robot; builds observations and get_serialized_actions for serving.""" + + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + self.is_received = False + self.current_robot_state_action = None + self.logger = InferLogger.get_robot_logger("Robot") + self.vehicle_pose_handler = VehiclePoseHandler() + self.last_speed = [0, 0, 0] + + def _send_actions(self, serialized_actions: dict) -> None: + return + + def get_observation(self, state, views): + robot_state_action_data = RobotStateActionData(config=self.config) + for key in robot_action_key_mapping.keys(): + state_key_str = robot_action_key_mapping[key] + value = None + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + return { + "robot_state_action_data": robot_state_action_data, + **_views_to_camera_observation(self.config, views), + } + + +class EX001Robot(Robot): + """EX001 robot: no pose integration; predict from velocity_decomposed_odom state and send to client.""" + + def __init__(self, config: InferConfig, robot_id=10) -> None: + super().__init__(config, robot_id) + self.last_speed = [0, 0, 0] + + def _get_velocity_decomposed( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + """Get decomposed velocity actions: concat state and action velocity_decomposed_odom""" + velocity_decomposed = np.concatenate( + [ + robot_state_action_data.data["state_velocity_decomposed_odom"], + robot_state_action_data.data["action_velocity_decomposed_odom"], + ], + axis=0, + ) + self.last_speed = velocity_decomposed[-1, :].copy().tolist() + return velocity_decomposed + + def _get_head_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + """Get head actions: concat state and action head sequence; supports head_rotation / head_actions""" + if robot_state_action_data.data.get("state_head_rotation") is not None: + state_key, action_key = "state_head_rotation", "action_head_rotation" + else: + state_key, action_key = "state_head_actions", "action_head_actions" + head_actions = np.concatenate( + [ + robot_state_action_data.data[state_key], + robot_state_action_data.data[action_key], + ], + axis=0, + ) + return head_actions + + def _get_height_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + """Get lift height actions: concat state and action height sequence""" + height = np.concatenate( + [ + robot_state_action_data.data["state_height"], + robot_state_action_data.data["action_height"], + ], + axis=0, + ) + return height + + def go_home(self): + """Return to initial pose""" + if self.current_robot_state_action is None: + self.get_observation() + + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "left_gripper" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "right_gripper" + ) + + self.current_robot_state_action.save_action_data_with_key( + np.array([[0, -1]]), "head_actions" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([[0.4]]), "height" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([self.last_speed]), "velocity_decomposed_odom" + ) + + action_dict = {"robot_state_action_data": self.current_robot_state_action} + self.apply_action( + action_dict, + robot_action_interpolate_multiplier=150, + robot_action_start_ratio=0, + robot_action_end_ratio=1, + ) + self.logger.info("robot returned to initial pose") + + def get_serialized_actions( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> dict: + """Build serialized EX001 actions from model_output (dual arms + head/lift/velocity_decomposed_odom); does not send. + + Uses predicted velocity_decomposed_odom directly; no pose integration. + """ + assert "robot_state_action_data" in input + + left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) + right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) + action_length = len(left_arm_action) + + # action_length is 1 state frame + action_horizon action frames + # Actual action length is action_length - 1 + actual_action_length = action_length - 1 + + if robot_action_start_ratio is None: + robot_action_start_ratio = self.config.robot_action_start_ratio + if robot_action_end_ratio is None: + robot_action_end_ratio = self.config.robot_action_end_ratio + + # Trim indices in action section (skip frame 0 state) + start_frame = 1 + int(robot_action_start_ratio * actual_action_length) + end_frame = 1 + int(robot_action_end_ratio * actual_action_length) + + left_arm_action = left_arm_action[start_frame:end_frame] + right_arm_action = right_arm_action[start_frame:end_frame] + + if robot_action_interpolate_multiplier is None: + robot_action_interpolate_multiplier = ( + self.config.robot_action_interpolate_multiplier + ) + target_length = robot_action_interpolate_multiplier * len(left_arm_action) + + ( + left_arm_action, + right_arm_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [left_arm_action, right_arm_action], + target_length, + ) + + serialized_actions = { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + + # Get head, lift, and velocity actions + head_action = self._get_head_action(input["robot_state_action_data"]) + height_action = self._get_height_action(input["robot_state_action_data"]) + velocity_decomposed_action = self._get_velocity_decomposed( + input["robot_state_action_data"] + ) + + # Trim actions (same start_frame and end_frame) + head_action = head_action[start_frame:end_frame] + height_action = height_action[start_frame:end_frame] + velocity_decomposed_action = velocity_decomposed_action[start_frame:end_frame] + + # Interpolate + ( + head_action, + height_action, + velocity_decomposed_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [head_action, height_action, velocity_decomposed_action], + target_length, + ) + + serialized_actions["head_pos"] = head_action.tolist() + serialized_actions["lift"] = height_action.tolist() + # Send velocity_decomposed_odom directly; no pose integration + serialized_actions["velocity_decomposed_odom"] = ( + velocity_decomposed_action.tolist() + ) + return serialized_actions + + def apply_action( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> None: + serialized_actions = self.get_serialized_actions( + input, + robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, + robot_action_start_ratio=robot_action_start_ratio, + robot_action_end_ratio=robot_action_end_ratio, + ) + self._send_actions(serialized_actions) + + +class EX001RobotPreprocessor(EX001Robot): + """EX001 preprocessor: no real robot; builds observations and get_serialized_actions for serving.""" + + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + self.is_received = False + self.current_robot_state_action = None + self.logger = InferLogger.get_robot_logger("Robot") + self.last_speed = [0, 0, 0] + + def _send_actions(self, serialized_actions: dict) -> None: + """No action send; preprocessing only""" + return + + def get_observation(self, state, views): + """Build observation from external state and views""" + robot_state_action_data = RobotStateActionData(config=self.config) + for key in robot_action_key_mapping.keys(): + state_key_str = robot_action_key_mapping[key] + value = None + + # Keep the same 7D adaptation behavior as Robot.get_observation(), + # so offline/serving preprocessor path is consistent with runtime path. + composed_value = self._try_compose_arm_joint_pos_with_gripper(state, key) + if composed_value is not None: + value = composed_value + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + continue + + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + + # Build observation; handle key name mismatches + observation = { + "robot_state_action_data": robot_state_action_data, + } + + # Try multiple possible key names + camera_mappings = { + "face_view": [ + self.config.camera_front_key, + "face_view", + "front_view", + ], + "left_wrist_view": [ + self.config.camera_left_key, + "left_wrist_view", + "left_view", + ], + "right_wrist_view": [ + self.config.camera_right_key, + "right_wrist_view", + "right_view", + ], + } + + for obs_key in self.config.cam_names: + possible_view_keys = camera_mappings.get(obs_key, [obs_key]) + found = False + for view_key in possible_view_keys: + if view_key in views and views[view_key] is not None: + view_data = views[view_key] + # If view_data is (1, H, W, 3), take [0] + if ( + isinstance(view_data, np.ndarray) + and len(view_data.shape) == 4 + and view_data.shape[0] == 1 + ): + observation[obs_key] = view_data[0] + else: + observation[obs_key] = view_data + found = True + self.logger.info( + f"EX001RobotPreprocessor: using '{view_key}' as '{obs_key}'" + ) + break + if not found: + self.logger.error( + f"EX001RobotPreprocessor: no view for '{obs_key}', tried keys: {possible_view_keys}" + ) + # Black image placeholder + observation[obs_key] = np.zeros((480, 640, 3), dtype=np.uint8) + + return observation diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py b/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py new file mode 100644 index 0000000..aa8f72a --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py @@ -0,0 +1,423 @@ +import socket +import struct +import json +import cv2 +import numpy as np +import time +import select +import errno +import os +import pickle + +import threading + +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger + +_DEFAULT_ROBOT_CONFIG = { + 0: {"host": "127.0.0.1", "action_port": 57749, "keyboard_port": 58849}, + 1: {"host": "127.0.0.1", "action_port": 57750, "keyboard_port": 58850}, + 2: {"host": "127.0.0.1", "action_port": 57751, "keyboard_port": 58851}, + 3: {"host": "127.0.0.1", "action_port": 57761, "keyboard_port": 58861}, +} + + +class RobotRegistry: + def __init__(self): + self.registry = _DEFAULT_ROBOT_CONFIG + self.logger = InferLogger.get_controller_logger("RobotRegistry") + self.logger.debug(f"init robot registry, registered {len(self.registry)} robots") + + def register_robot(self, robot_id, host, action_port, keyboard_port): + if robot_id in self.registry: + self.logger.warning(f"robot ID {robot_id} already registered") + return False + else: + self.registry[robot_id] = { + "host": host, + "action_port": action_port, + "keyboard_port": keyboard_port, + } + self.logger.info( + f"register robot ID={robot_id}: host={host}, action_port={action_port}, keyboard_port={keyboard_port}" + ) + return True + + def get_robot_info(self, robot_id): + if robot_id in self.registry: + info = self.registry[robot_id] + self.logger.debug(f"get robot ID={robot_id} info: {info}") + return info + else: + self.logger.error(f"robot ID={robot_id} not found") + return None + + def exist(self, robot_id): + return robot_id in self.registry + + +class RobotCommunication: + def __init__(self, robot_id, host=None, action_port=None, keyboard_port=None): + self.logger = InferLogger.get_controller_logger("RobotCommunication") + self.logger.debug(f"init robot comms: robot_id={robot_id}") + + self.robot_register = RobotRegistry() + if self.robot_register.exist(robot_id): + self.robot_info = self.robot_register.get_robot_info(robot_id) + self.logger.debug(f"using registered robot config: {robot_id}") + else: + self.robot_register.register_robot( + robot_id, host, action_port, keyboard_port + ) + self.robot_info = self.robot_register.get_robot_info(robot_id) + self.logger.debug(f"register new robot config: {robot_id}") + + self.action_sock = None + self.action_conn = None + + self.keyboard_sock = None + self.keyboard_conn = None + + self.client_socks = [] + + def connect(self): + self.logger.info("establishing socket connection...") + + self.action_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.action_sock.setblocking(True) + self.action_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + + host = self.robot_info["host"] + port = self.robot_info["action_port"] + self.action_sock.bind((host, port)) + self.action_sock.listen(1) + self.logger.info(f"listening on action port: {host}:{self.action_sock.getsockname()[1]}") + + action_thread = threading.Thread(target=self.handle_action_client) + + self.logger.debug("starting connection handler thread...") + action_thread.start() + + action_thread.join() + self.logger.info("socket connection established") + + def handle_action_client(self): + self.logger.debug("waiting for client connection...") + self.action_conn, addr = self.action_sock.accept() + self.logger.info(f"accepted connection from {addr}") + + def recv_image(self, index): + self.logger.debug(f"receiving image index={index}...") + image_size = struct.unpack(" dict: + if not self.views_path: + raise RuntimeError( + "DummyRobotController requires views_path or " + "WALL_X_DUMMY_VIEWS_PATH." + ) + self.logger.debug(f"load views from pickle: {cam_names}") + with open(self.views_path, "rb") as f: + all_views = pickle.load(f) + self.logger.debug(f"views loaded: {len(all_views)} views") + return all_views + + def recv_action(self): + if not self.state_path: + raise RuntimeError( + "DummyRobotController requires state_path or " + "WALL_X_DUMMY_STATE_PATH." + ) + self.logger.debug("load actions from pickle") + with open(self.state_path, "rb") as f: + all_actions = pickle.load(f) + self.logger.debug(f"actions loaded: {len(all_actions)} fields") + return all_actions + + +class RobotController: + def __init__( + self, + robot_id: int, + host: str = None, + port: str = None, + max_time_step: int = 10000, + ): + self.logger = InferLogger.get_controller_logger("RobotController") + self.logger.info( + f"init robot controller: robot_id={robot_id}, max_time_step={max_time_step}" + ) + + self.robot_comm = RobotCommunication(robot_id, host, port) + self.max_time_step = max_time_step + self.global_step = 0 + + def connect(self): + self.logger.info("connecting robot controller...") + self.robot_comm.connect() + self.logger.info("robot controller connected") + + def close(self): + self.logger.info("closing robot controller...") + self.robot_comm.close() + self.logger.info("robot controller closed") + + def prediction(self, views, actions) -> dict: + raise NotImplementedError + + def recv_image( + self, cam_names: list = ["camera_left", "camera_front", "camera_right"] + ) -> dict: + self.logger.debug(f"receive images: {cam_names}") + views = {} + for i, name in enumerate(cam_names): + image = np.array(self.robot_comm.recv_image(i)) + views[name] = image[None, :] + self.logger.debug(f"received {name}: shape={image.shape}") + self.logger.debug(f"all images received: {len(views)} views") + return views + + def recv_action(self): + self.logger.debug("receiving action data...") + action_data = self.robot_comm.recv_action_data() + return action_data + + def recv_keyboard_input(self): + self.robot_comm.accept_connections() + json_data = self.robot_comm.recv_keyboard_input() + if json_data is not None: + self.logger.debug( + f"keyboard: motionStatus={json_data.get('motionStatus')}, armMode={json_data.get('armMode')}" + ) + return json_data["motionStatus"], json_data["armMode"] + + return None, None + + def reset(self): + self.logger.info(f"reset controller: global_step {self.global_step} -> 0") + self.global_step = 0 + + def record_start(self): + self.logger.info("start recording...") + record_signal = {"cmd": "RECORD_START"} + self.robot_comm.send_dict(record_signal) + + def record_continue(self): + self.logger.debug("continue recording...") + record_signal = {"cmd": "RECORD_CONTINUE"} + self.robot_comm.send_dict(record_signal) + + def record_start_process(self): + if self.global_step == 0: + self.logger.info("waiting for START signal...") + action = None + while action != "START": + action, _ = self.recv_keyboard_input() + self.logger.info("received START signal") + self.record_start() + + self.record_continue() + + def set_zero(self): + self.logger.info("setting zero pose...") + record_signal = {"cmd": "INIT_ZERO", "gripper": [0.0, 0.0]} + self.robot_comm.send_dict(record_signal) + self.reset() + self.logger.info("zero pose set") + + def record_stop(self): + self.logger.info("stop recording...") + record_signal = {"cmd": "RECORD_STOP"} + self.robot_comm.send_dict(record_signal) + time.sleep(0.01) + + def recover_from_failure(self): + self.logger.warning("recovering from failure...") + record_signal = {"cmd": "TO_MASTER_SLAVE"} + self.robot_comm.send_dict(record_signal) + time.sleep(0.01) + + self.logger.info("waiting for START signal...") + action = None + while action != "START": + action, _ = self.recv_keyboard_input() + + self.logger.info("recovery done, resume recording") + self.record_start() + + def reset_to_zero(self): + self.logger.info("reset to zero pose...") + self.record_stop() + self.set_zero() + + def to_slave(self): + self.logger.info("switching to slave mode...") + record_signal = {"cmd": "TO_SLAVE"} + self.robot_comm.send_dict(record_signal) + time.sleep(0.01) + + def run( + self, + record_mode=False, + cam_names: list = ["camera_left", "camera_front", "camera_right"], + ): + + while self.global_step <= self.max_time_step: + if record_mode: + self.record_start_process() + + self.global_step += 1 + + action = self.recv_action() + view = self.recv_image(cam_names) + + pred = self.prediction(view, action) + self.robot_comm.send_dict(pred) + + if record_mode: + action, arm_mode = self.recv_keyboard_input() + + if action is not None and arm_mode is not None: + self.logger.info("action: %s, arm_mode: %s", action, arm_mode) + if action == "STOP" and arm_mode == "ARM_TEST_MODE_MS": + self.record_stop() + self.recover_from_failure() + action, arm_mode = None, None + while action != "STOP" or arm_mode != "ARM_TEST_MODE_S": + action, arm_mode = self.recv_keyboard_input() + self.record_stop() + self.to_slave() + self.reset() + elif arm_mode == "ARM_TEST_MODE_S" and action == "INIT": + self.reset_to_zero() diff --git a/wall_x/infer/utils.py b/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py similarity index 75% rename from wall_x/infer/utils.py rename to wall_x/_vendor/harrix/serving/_wallx_infer/utils.py index e4fba55..4163768 100644 --- a/wall_x/infer/utils.py +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py @@ -1,21 +1,21 @@ import numpy as np from scipy.signal import savgol_filter -from scipy.spatial.transform import Rotation as R # TODO: Convert to numba functions +from scipy.spatial.transform import Rotation as R # TODO: convert to numba from collections import deque import threading -from wall_x.infer.logger import InferLogger +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger class KeyboardThread(threading.Thread): """ - Simple keyboard listening thread that provides stop and reset functionality + Simple keyboard listener thread with stop and reset """ def __init__(self): self.should_reset = False self.should_stop = False - self.new_instruction_index = None # Used to store new instruction index + self.new_instruction_index = None # Stores new instruction index self.logger = InferLogger.get_utils_logger("KeyboardThread") super(KeyboardThread, self).__init__(name="keyboard-thread", daemon=True) @@ -23,45 +23,41 @@ class KeyboardThread(threading.Thread): self.start() def run(self): - """Listen to keyboard input""" + """Listen for keyboard input""" while True: try: user_input = input().strip().lower() if user_input in ["s", "stop"]: self.should_stop = not self.should_stop - self.logger.info("[Keyboard] Stop signal sent") + self.logger.info("[keyboard] stop signal sent") elif user_input in ["r", "reset"]: - self.logger.info("[Keyboard] Executing reset...") + self.logger.info("[keyboard] resetting...") self.should_reset = True - self.logger.info("[Keyboard] Reset signal sent") + self.logger.info("[keyboard] reset signal sent") elif user_input.isdigit(): - # Handle digit input, switch to corresponding instruction index + # Numeric input: switch instruction index index = int(user_input) self.new_instruction_index = index - self.logger.info( - f"[Keyboard] Switched to instruction index: {index}" - ) + self.logger.info(f"[keyboard] switched to instruction index: {index}") else: - self.logger.info( - f"[Keyboard] Received input: {user_input}. No action taken." - ) + self.logger.info(f"[keyboard] input received: {user_input}. no action taken.") except EOFError: break except Exception as e: - self.logger.error(f"[Keyboard] Error: {e}") + self.logger.error(f"[keyboard] error: {e}") def show_help(self): self.logger.info( - "[Keyboard] Keyboard control: Enter 's' to stop, 'r' to reset, 'number' to switch instruction index" + "[keyboard] controls: 's' stop, 'r' reset, digit switches instruction index" ) -# Robot arm trajectory parameters +# Arm trajectory parameters ARM_MAX_VELOCITY = 0.02 ARM_EXECUTION_HZ = 20 ARM_MIN_EXECUTION_TIME = 5.0 @@ -74,9 +70,9 @@ class UnifiedTrajectoryProcessor: @staticmethod def interpolate_trajectory_batch(trajectories, target_length, smooth=True): """ - Batch interpolate multiple trajectories to unified length + Interpolate multiple trajectories to a common length Args: - trajectories: list of np.array, each array with shape (N, D) + trajectories: list of np.array, each shape (N, D) target_length: int, target length smooth: bool, whether to smooth Returns: @@ -99,8 +95,8 @@ class UnifiedTrajectoryProcessor: original_indices = np.linspace(0, len(traj) - 1, len(traj)) target_indices = np.linspace(0, len(traj) - 1, target_length) - # Handle different types of data - if traj.shape[1] == 7: # Robot arm data [x,y,z,rx,ry,rz,gripper] + # Handle different data types + if traj.shape[1] == 7: # Arm data [x,y,z,rx,ry,rz,gripper] interpolated = UnifiedTrajectoryProcessor._interpolate_arm_trajectory( traj, original_indices, target_indices, target_length ) @@ -111,7 +107,7 @@ class UnifiedTrajectoryProcessor: target_indices, original_indices, traj[:, i] ) - # Smooth processing + # Smoothing if smooth and len(interpolated) >= 5: interpolated = UnifiedTrajectoryProcessor._smooth_trajectory( interpolated @@ -125,10 +121,10 @@ class UnifiedTrajectoryProcessor: def _interpolate_arm_trajectory( traj, original_indices, target_indices, target_length ): - """Optimized robot arm trajectory interpolation""" + """Optimized arm trajectory interpolation""" interpolated = np.zeros((target_length, 7)) - # Vectorized interpolation for position and gripper + # Vectorized interp for position and gripper for i in [0, 1, 2, 6]: # x, y, z, gripper interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i]) @@ -140,11 +136,11 @@ class UnifiedTrajectoryProcessor: target_indices, original_indices, quaternions[:, i] ) - # Batch normalization + # Batch normalize norms = np.linalg.norm(interpolated_quats, axis=1, keepdims=True) interpolated_quats = interpolated_quats / norms - # Batch convert back to Euler angles + # Batch convert back to euler interpolated[:, 3:6] = R.from_quat(interpolated_quats).as_euler("xyz") return interpolated @@ -161,12 +157,12 @@ class UnifiedTrajectoryProcessor: @staticmethod def _smooth_trajectory(trajectory): - """Vectorized smooth processing""" + """Vectorized smoothing""" if len(trajectory) < 5: return trajectory try: - # Batch smooth all dimensions + # Smooth all dimensions in batch smoothed = np.zeros_like(trajectory) for dim in range(trajectory.shape[1]): smoothed[:, dim] = savgol_filter( @@ -188,9 +184,9 @@ class UnifiedTrajectoryProcessor: @staticmethod def calculate_optimal_trajectory_length(left_traj, right_traj): - """Calculate optimal trajectory length""" + """Compute optimal trajectory length""" - # Vectorized distance calculation + # Vectorized distance computation def calc_distance(traj): if len(traj) < 2: return 0.0 @@ -213,12 +209,13 @@ class UnifiedTrajectoryProcessor: class VehiclePoseHandler: - """Vehicle pose and velocity calculation""" + """Vehicle pose and velocity computation""" def __init__(self): self.current_pose = None self.previous_pose = None self.pose_history = deque(maxlen=10) + self.logger = InferLogger.get_utils_logger("VehiclePoseHandler") def update_pose(self, new_pose): """Update vehicle pose""" @@ -226,11 +223,11 @@ class VehiclePoseHandler: self.previous_pose = self.current_pose self.current_pose = np.array(new_pose) self.pose_history.append(self.current_pose.copy()) - print("current_pose", self.current_pose, flush=True) + self.logger.info("current_pose %s", self.current_pose) return self.current_pose def velocity_to_pose(self, vx_body, vy_body, vyaw, dt, start_pose=None): - """Convert body frame velocity to global frame position""" + """Convert body-frame velocity to global pose""" if start_pose is None: if self.current_pose is not None: start_pose = self.current_pose.copy() @@ -239,21 +236,21 @@ class VehiclePoseHandler: x, y, theta = start_pose - # Convert body frame velocity to global frame displacement + # Body velocity to global displacement cos_theta = np.cos(theta) sin_theta = np.sin(theta) - # Coordinate transformation: body frame -> global frame + # Transform: body frame -> global frame dx_global = (vx_body * cos_theta - vy_body * sin_theta) * dt dy_global = (vx_body * sin_theta + vy_body * cos_theta) * dt dtheta = vyaw * dt - # Calculate new position + # Compute new pose x_new = x + dx_global y_new = y + dy_global theta_new = theta + dtheta - # Constrain angle to [-pi, pi] range + # Wrap angle to [-pi, pi] theta_new = (theta_new + np.pi) % (2 * np.pi) - np.pi return np.array([x_new, y_new, theta_new]) @@ -261,21 +258,21 @@ class VehiclePoseHandler: def compute_body_velocities_from_poses( self, current_pose, previous_pose, dt=1 / 20 ): - """Compute body frame velocity from pose changes""" + """Compute body-frame velocity from pose delta""" if current_pose is None or previous_pose is None: return np.array([0.0, 0.0, 0.0]) - # Calculate displacement in global frame + # Global-frame displacement dx_global = current_pose[0] - previous_pose[0] dy_global = current_pose[1] - previous_pose[1] dtheta = current_pose[2] - previous_pose[2] - # Use previous frame's angle for coordinate transformation + # Use previous frame angle for transform theta = previous_pose[2] cos_theta = np.cos(theta) sin_theta = np.sin(theta) - # Convert global frame displacement to body frame velocity + # Global displacement to body velocity vx_body = (dx_global * cos_theta + dy_global * sin_theta) / dt vy_body = (-dx_global * sin_theta + dy_global * cos_theta) / dt vyaw = dtheta / dt diff --git a/wall_x/_vendor/harrix/serving/launch_serving.py b/wall_x/_vendor/harrix/serving/launch_serving.py new file mode 100644 index 0000000..9b9c0f4 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/launch_serving.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python3 +""" +Server script for Wall-X model. + +This script serves a Wall-X model using a websocket server, allowing +clients to connect and get action predictions from observations. + +""" + +import dataclasses +import enum +import inspect +import logging +import os +import socket +import sys +import yaml +import traceback +import tyro + +from wall_x._vendor.harrix.serving.websocket_policy_server import WebsocketPolicyServer + + +def _server_model_config_to_infer_kwargs(model_config) -> dict: + from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig + + infer_params = inspect.signature(InferConfig.__init__).parameters + return { + k: v + for k, v in vars(model_config).items() + if v is not None and k in infer_params + } + + +def get_wallx_policy(model_config, image_passing_mode, serialize_actions=True): + from wall_x._vendor.harrix.serving.policy.wall_x_policy import WallXPolicy + from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig + + config = InferConfig(**_server_model_config_to_infer_kwargs(model_config)) + return WallXPolicy(config=config, image_passing_mode=image_passing_mode, serialize_actions=serialize_actions) + + + + +logger = logging.getLogger(__name__) + + +class EnvMode(enum.Enum): + """Supported environments/datasets.""" + + X2ROBOT = "x2robot" + LIBERO = "libero" + + +@dataclasses.dataclass +class ServerModelConfig: + """Configuration for loading a Wall-X model.""" + + checkpoint_path: str | None = None + train_config_path: str | None = None + # robot_host: str = '0.0.0.0' + # robot_port: int = 33723 + robot_type: str = "desktop" # ["desktop", "turtle"] + robot_action_start_ratio: float = ( + 0.0 # proportion of action execution to start from + ) + robot_action_end_ratio: float = 1.0 # proportion of action execution to end at + robot_action_interpolate_multiplier: int = 10 # action interpolation multiplier + robot_use_joint_angle_control: bool = ( + False # use joint angle control (model must predict joints) + ) + turtle_as_desktop: bool = ( + False # use turtle platform as desktop with fixed base/head/camera/height + ) + action_horizon: int = 32 # specify the correct horizon for the model + action_dim: int | None = None + model_device: str = "cuda" + num_inference_timesteps: int = 10 + num_inference_steps: int | None = None + cfg_scale: float | None = None + seed: int | None = None + save_video_dir: str = "./videos" + + # Please specify explicitly if the checkpoint was not trained on the x2robot dataset. + norm_key: str | None = None + # Model cameras; None = infer from train config ``data.key_mappings.camera``. + cam_names: list[str] | None = None + # Robot camera keys in incoming websocket observations. + camera_front_key: str = "camera_front" + camera_left_key: str = "camera_left" + camera_right_key: str = "camera_right" + # Serving prompt controls. If the client request has no instruction, + # default_instruction is used. prompt_template follows train-config semantics. + default_instruction: str | None = None + prompt_template: str | None = None + qwen25_prompt_template: str | None = None + prompt_priority_order: str | None = None + +@dataclasses.dataclass +class Args: + """Arguments for the serve_wall_x script.""" + + # Environment mode (used for default configurations) + env: EnvMode = EnvMode.X2ROBOT + + # Model configuration. If not provided, uses default config for the environment + model_config: ServerModelConfig | None = None + + # Default text prompt to use if not provided in observation + default_prompt: str | None = None + + # Port to serve the policy on + port: int = 43007 + + # Host to bind the server to + host: str = "0.0.0.0" + + # Enable debug logging + debug: bool = False + + # Image passing mode + image_passing_mode: str = "base64" # ["numpy", "base64"] + + # Model type + model_type: str = "wallx" # OSS supports qwen2.5 Wall-X only + + # Serialize actions via robot_preprocessor (True for robot control, False for raw output) + serialize_actions: bool = True + + # -- Dynamic batching ----------------------------------------- + # Set max_batch_size to enable dynamic batching. None = single mode. + max_batch_size: int | None = None + max_wait_time_ms: float = 0 + max_queue_size: int = 100 + timeout_ms: float = 30000 + + # -- Engine flags --------------------------------------------- + enable_experimental_engine: bool = False + enable_cuda_graph: bool = False + + +# Default model configurations for each environment +DEFAULT_CONFIGS: dict[EnvMode, ServerModelConfig] = { + EnvMode.X2ROBOT: ServerModelConfig( + checkpoint_path=None, + train_config_path=None, + robot_action_start_ratio=0.0, + robot_action_end_ratio=1.0, + robot_action_interpolate_multiplier=10, + robot_use_joint_angle_control=False, + turtle_as_desktop=False, + action_horizon=32, + action_dim=None, + model_device="cuda", + num_inference_timesteps=10, + ), + EnvMode.LIBERO: ServerModelConfig( + checkpoint_path=None, + train_config_path=None, + robot_type="desktop", + robot_action_start_ratio=0.0, + robot_action_end_ratio=1.0, + robot_action_interpolate_multiplier=1, + action_horizon=10, + action_dim=None, + model_device="cuda", + num_inference_timesteps=10, + cam_names=None, + ), +} + + +def get_model_config(args: Args) -> ServerModelConfig: + """Get model configuration from args or defaults.""" + if args.model_config is not None: + return args.model_config + + if config := DEFAULT_CONFIGS.get(args.env): + logger.info(f"Using default configuration for {args.env.value}") + return config + + raise ValueError( + f"No default configuration for {args.env.value}. " + f"Please provide --model-config with model_path and action_tokenizer_path." + ) + + +def create_policy(args: Args): + """Create a policy from the given arguments.""" + model_config = get_model_config(args) + if args.model_type != "wallx": + raise ValueError( + f"Unsupported model type: {args.model_type!r}. " + "The public package only supports model_type='wallx'." + ) + policy = get_wallx_policy( + model_config, args.image_passing_mode, args.serialize_actions + ) + return policy + + +def main(args: Args) -> None: + """Main function to start the Wall-X model server.""" + log_level = logging.DEBUG if args.debug else logging.INFO + logging.basicConfig( + level=log_level, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + + # Set engine environment variables + if args.enable_experimental_engine: + os.environ["ENABLE_EXPERIMENTAL_INFERENCE_ENGINE"] = "true" + logger.info("ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=true") + + if args.enable_cuda_graph: + os.environ["ENABLE_CUDA_GRAPH"] = "true" + logger.info("ENABLE_CUDA_GRAPH=true") + + logger.info("Starting model server") + logger.info(f"Model type: {args.model_type}") + logger.info(f"Environment: {args.env.value}") + logger.info(f"Port: {args.port}") + logger.info(f"Host: {args.host}") + logger.info(f"Serialize actions: {args.serialize_actions}") + + # Create policy + try: + policy = create_policy(args) + except Exception as e: + logger.error(f"Failed to create policy: {e}") + logger.error(traceback.format_exc()) + sys.exit(1) + + # Get policy metadata + policy_metadata = policy.metadata + policy_metadata["env"] = args.env.value + + # Get network info + hostname = socket.gethostname() + try: + local_ip = socket.gethostbyname(hostname) + except Exception: + local_ip = "unknown" + + logger.info(f"Server hostname: {hostname}") + logger.info(f"Server IP: {local_ip}") + logger.info(f"Server will be available at: ws://{args.host}:{args.port}") + logger.info(f"Health check endpoint: http://{args.host}:{args.port}/healthz") + + batching_str = ( + f"batch_size={args.max_batch_size}, wait={args.max_wait_time_ms}ms" + if args.max_batch_size + else "disabled" + ) + logger.info(f"Batching: {batching_str}") + + # Create and start server + server = WebsocketPolicyServer( + policy=policy, + host=args.host, + port=args.port, + metadata=policy_metadata, + max_batch_size=args.max_batch_size, + max_wait_time_ms=args.max_wait_time_ms, + max_queue_size=args.max_queue_size, + timeout_ms=args.timeout_ms, + ) + + logger.info("Starting server...") + try: + server.serve_forever() + except KeyboardInterrupt: + logger.info("Server stopped by user") + except Exception as e: + logger.error(f"Server error: {e}") + sys.exit(1) + + +if __name__ == "__main__": + main(tyro.cli(Args)) diff --git a/wall_x/serving/policy/__init__.py b/wall_x/_vendor/harrix/serving/policy/__init__.py similarity index 56% rename from wall_x/serving/policy/__init__.py rename to wall_x/_vendor/harrix/serving/policy/__init__.py index 261bf30..e59a287 100644 --- a/wall_x/serving/policy/__init__.py +++ b/wall_x/_vendor/harrix/serving/policy/__init__.py @@ -1,3 +1,5 @@ +"""Wall-X policy for harrix websocket serving.""" + from .wall_x_policy import WallXPolicy __all__ = ["WallXPolicy"] diff --git a/wall_x/_vendor/harrix/serving/policy/_smoothing.py b/wall_x/_vendor/harrix/serving/policy/_smoothing.py new file mode 100644 index 0000000..9b6c127 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/policy/_smoothing.py @@ -0,0 +1,104 @@ +"""Test-time Laplacian smoothing for serving model outputs. + +Mirrors the open-loop implementation in +`wallx_bus2604/wall-x/run_scripts/infer_openloop.py` +(_laplacian_smooth @ 203-209, smooth_action / smooth_gripper fields @ 186-187, +application @ 458-463). Invoked by WallXPolicy before downstream 6D->euler +conversion so that smoothed rotations propagate through serialization. +""" + +from __future__ import annotations + +from typing import List, Optional, Sequence + +import numpy as np +import torch + +from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import dof_dims + + +def _laplacian_smooth(a: np.ndarray, lam: float = 1.0, iters: int = 30) -> np.ndarray: + """Iterative Laplacian smoothing along axis 0; endpoints pinned.""" + a = a.copy() + orig = a.copy() + for _ in range(iters): + a[1:-1] = (orig[1:-1] + lam * (a[:-2] + a[2:])) / (1 + 2 * lam) + return a + + +def _gripper_column_indices( + predict_action_keys: Sequence[str], + action_padding_dof: Optional[int] = None, +) -> List[int]: + """Column indices in the flat (T, D) layout that correspond to gripper dofs.""" + cols: List[int] = [] + dof_start = 0 + for key in predict_action_keys: + if key == "action_padding": + dof_start += action_padding_dof or 0 + continue + short = key.replace("follow_", "").replace("master_", "") + width = dof_dims[short] + if "gripper" in short: + cols.extend(range(dof_start, dof_start + width)) + dof_start += width + return cols + + +_LAZY_ACTION_KEYS = ( + "action_left_ee_cartesian_pos", + "action_right_ee_cartesian_pos", + "action_left_ee_rotation", + "action_right_ee_rotation", + "action_left_ee_rotation_6D", + "action_right_ee_rotation_6D", +) + + +def apply_smoothing( + model_output: dict, + smooth_action: bool, + smooth_gripper: bool, + predict_action_keys: Sequence[str], + action_padding_dof: Optional[int] = None, + action_dim: Optional[int] = None, +) -> None: + """Smooth `model_output['predict_action']` in place and refresh per-arm keys.""" + if not smooth_action: + return + pa = model_output.get("predict_action") + if pa is None: + return + + was_tensor = isinstance(pa, torch.Tensor) + arr = pa.detach().cpu().numpy() if was_tensor else np.asarray(pa) + + orig_ndim = arr.ndim + if orig_ndim == 3: + if arr.shape[0] != 1: + return + arr = arr[0] + if arr.ndim != 2 or arr.shape[0] < 3: + return + if action_dim is not None and arr.shape[-1] != action_dim: + return + + orig = arr.copy() + smoothed = _laplacian_smooth(arr) + + if not smooth_gripper: + for c in _gripper_column_indices(predict_action_keys, action_padding_dof): + if c < smoothed.shape[-1]: + smoothed[:, c] = orig[:, c] + + out = smoothed[None] if orig_ndim == 3 else smoothed + if was_tensor: + out = torch.from_numpy(out).to(device=pa.device, dtype=pa.dtype) + model_output["predict_action"] = out + + rsd = model_output.get("robot_state_action_data") + if rsd is not None: + for k in _LAZY_ACTION_KEYS: + if k in rsd.data: + rsd.data[k] = None + rsd.save_action_data(out) diff --git a/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py b/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py new file mode 100644 index 0000000..155d602 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py @@ -0,0 +1,328 @@ +import base64 +import logging +from typing import Dict, Any, List + +import cv2 +import numpy as np +import torch +from wall_x._vendor.harrix.serving.websocket_policy_server import BasePolicy +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig +from wall_x._vendor.harrix.serving._wallx_infer.model_wrapper import WallxModelWrapper +from wall_x._vendor.harrix.serving._wallx_infer.robot import ( + DesktopRobotPreprocessor, + TurtleRobotPreprocessor, + EX001RobotPreprocessor, +) +from wall_x._vendor.harrix.serving.policy._smoothing import apply_smoothing +from wall_x.utils.timers import ScopeTimer + +logger = logging.getLogger(__name__) + +# Inference modes aligned with model_wrapper for vqa / batch extensions +INFER_MODE_FLOW = "flow" +INFER_MODE_AR = "ar" +INFER_MODE_FLOW_WITH_SUBTASK = "flow_with_subtask" +INFER_MODE_DLLM_FLOW = "dllm" +INFER_MODE_DLLM_DD = "discrete_diffusion" +ACTION_INFER_MODES = ( + INFER_MODE_FLOW, + INFER_MODE_AR, + INFER_MODE_FLOW_WITH_SUBTASK, + INFER_MODE_DLLM_FLOW, + INFER_MODE_DLLM_DD, +) + + +class WallXPolicy(BasePolicy): + """Policy wrapper for Wall-X model that implements the BasePolicy interface.""" + + def __init__( + self, + config: InferConfig, + image_passing_mode: str = "base64", + default_infer_mode: str = INFER_MODE_FLOW, + serialize_actions: bool = True, + ): + """Initialize the Wall-X policy. + + Args: + config: Inference configuration dataclass. + image_passing_mode: How images are passed from client ('base64' or 'numpy'). + default_infer_mode: Default inference mode ('flow', 'ar', 'flow_with_subtask', etc.). + serialize_actions: If True, actions are serialized via robot_preprocessor; + if False, raw model output is returned directly. + """ + self.config = config + self.model_wrapper = WallxModelWrapper(config) + self.robot_preprocessor = self._register_robot_preprocessor() + self.image_passing_mode = image_passing_mode + self.default_infer_mode = default_infer_mode + self.serialize_actions = serialize_actions + logger.info( + "Image passing mode: %s, robot_type: %s, default_infer_mode: %s, " + "serialize_actions: %s, smooth_action: %s, smooth_gripper: %s", + self.image_passing_mode, + config.robot_type, + self.default_infer_mode, + self.serialize_actions, + config.smooth_action, + config.smooth_gripper, + ) + + def _register_robot_preprocessor(self): + """Select preprocessor by config.robot_type; mirrors env._register_robot.""" + if self.config.robot_type == "desktop": + return DesktopRobotPreprocessor(self.config) + if self.config.robot_type == "turtle": + return TurtleRobotPreprocessor(self.config) + if self.config.robot_type == "ex001": + return EX001RobotPreprocessor(self.config) + raise ValueError(f"Invalid robot_type: {self.config.robot_type!r}") + + @staticmethod + def _decode_base64_image(image_b64: str) -> np.ndarray: + """Decode client base64 JPEG/PNG payloads to RGB images.""" + img_bytes = base64.b64decode(image_b64) + img_array = np.frombuffer(img_bytes, dtype=np.uint8) + decoded_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) + if decoded_img is None: + raise ValueError("Failed to decode base64 image payload") + return cv2.cvtColor(decoded_img, cv2.COLOR_BGR2RGB) + + def reset(self) -> None: + """Reset the policy state.""" + self.action_buffer = [] + self.buffer_index = 0 + logger.debug("Policy reset") + + def _get_dof_config(self) -> dict: + train_config = self.config.train_config or {} + return ( + train_config.get("dof_config") + or train_config.get("task", {}).get("dof_config", {}) + or {} + ) + + def _is_single_arm_right_only(self) -> bool: + """True for single-arm LIBERO-style configs (right arm only, no left arm).""" + train_config = self.config.train_config or {} + agent_pos = ( + train_config.get("agent_pos_config") + or train_config.get("task", {}).get("agent_pos_config") + or {} + ) + skip = {"action_padding"} + has_left = any(k.startswith("follow_left_") for k in agent_pos if k not in skip) + has_right = any( + k.startswith("follow_right_") for k in agent_pos if k not in skip + ) + return has_right and not has_left + + def _pack_action_chunk_response( + self, model_output: Dict[str, Any] + ) -> Dict[str, Any]: + """Return a msgpack-safe action chunk for websocket clients.""" + predict_action = model_output["predict_action"] + if isinstance(predict_action, torch.Tensor): + predict_action = predict_action.detach().cpu().numpy() + response: Dict[str, Any] = { + "predict_action": np.asarray(predict_action, dtype=np.float32), + } + if "subtask" in model_output: + response["subtask"] = model_output["subtask"] + return response + + def _get_predict_action_keys(self) -> List[str]: + """Resolve flat action key order for smoothing / serialization.""" + data_cfg = self.config.data_config + keys = None + if isinstance(data_cfg, dict): + keys = data_cfg.get("predict_action_keys") + else: + keys = getattr(data_cfg, "predict_action_keys", None) + if keys: + return list(keys) + return list(self._get_dof_config().keys()) + + def _apply_smoothing(self, model_output: Dict[str, Any]) -> None: + """Apply Laplacian smoothing on predict_action before 6D->euler conversion.""" + cfg = self.config + if not getattr(cfg, "smooth_action", False): + return + dof_config = self._get_dof_config() + apply_smoothing( + model_output, + smooth_action=True, + smooth_gripper=cfg.smooth_gripper, + predict_action_keys=self._get_predict_action_keys(), + action_padding_dof=dof_config.get("action_padding"), + action_dim=cfg.action_dim, + ) + + def _run_action_infer( + self, observation: Dict, instruction: str, mode: str + ) -> Dict[str, Any]: + """Run model_wrapper action inference by mode; returns model_output with robot_state_action_data. + + Supports flow | ar | flow_with_subtask. VQA can be added here later (different return format). + """ + if mode == INFER_MODE_FLOW: + return self.model_wrapper.infer_flow_action(observation, instruction) + if mode == INFER_MODE_AR: + return self.model_wrapper.infer_ar_action(observation, instruction) + if mode == INFER_MODE_FLOW_WITH_SUBTASK: + with ScopeTimer("infer_subtask"): + subtask = self.model_wrapper.infer_subtask(observation, instruction) + with ScopeTimer("infer_flow_action"): + model_output = self.model_wrapper.infer_flow_action( + observation, subtask + ) + model_output["subtask"] = subtask + return model_output + if mode == INFER_MODE_DLLM_FLOW: + return self.model_wrapper.infer_dllm_action( + observation, instruction, use_ar_action=False + ) + if mode == INFER_MODE_DLLM_DD: + return self.model_wrapper.infer_dllm_action( + observation, instruction, use_ar_action=True + ) + raise ValueError( + f"Unsupported infer_mode={mode!r}, expected one of {ACTION_INFER_MODES}" + ) + + def infer(self, obs: Dict) -> Dict: + """Infer action from observation. + + Args: + obs: Dictionary containing: + - 'state': Robot state + - 'views': Camera views (keyed by camera name) + - 'instruction': Task instruction + - Optional: 'infer_mode' - one of 'flow' | 'ar' | 'flow_with_subtask' + - Optional: 'robot_action_start_ratio' / 'robot_action_end_ratio' / + 'robot_action_interpolate_multiplier' - override config for action trim/interpolate + + Returns: + When serialize_actions=True (default): Serialized action dict + (e.g. follow1_pos/follow2_pos or follow1_joints/follow2_joints). + When serialize_actions=False: Raw model_output dict. + If infer_mode is flow_with_subtask, also includes 'subtask'. + """ + state = obs["state"] + views = obs["views"] + instruction = obs.get("instruction") or self.config.default_instruction or "" + if self.image_passing_mode == "base64": + for k, v in views.items(): + views[k] = np.expand_dims(self._decode_base64_image(v), axis=0) + + with ScopeTimer("get_observation"): + observation = self.robot_preprocessor.get_observation(state, views) + + mode = obs.get("infer_mode", self.default_infer_mode) + with ScopeTimer(f"infer_{mode}"): + model_output = self._run_action_infer(observation, instruction, mode) + + self._apply_smoothing(model_output) + if ( + "robot_state_action_data" in model_output + and "predict_action" in model_output + ): + model_output["robot_state_action_data"].save_action_data( + model_output["predict_action"] + ) + + # Single-arm LIBERO has no left-arm EE to serialize; dual-arm follow1/follow2 + # layout would fail in get_serialized_actions. Return raw action chunks instead. + if not self.serialize_actions or self._is_single_arm_right_only(): + return self._pack_action_chunk_response(model_output) + + return self.robot_preprocessor.get_serialized_actions( + model_output, robot_action_interpolate_multiplier=1 + ) # interpolation on robot websocket client + + # -- Batch inference ---------------------------------------------- + + def _preprocess_obs(self, obs: Dict): + """Preprocess a single observation dict into (observation, instruction). + + Handles both base64 and raw image modes. + """ + state = obs["state"] + views = obs["views"] + instruction = obs.get("instruction") or self.config.default_instruction or "" + + if self.image_passing_mode == "base64": + for k, v in views.items(): + views[k] = np.expand_dims(self._decode_base64_image(v), axis=0) + else: + # raw numpy: ensure (1, H, W, C) + for k, img in views.items(): + if isinstance(img, np.ndarray) and img.ndim == 3: + views[k] = np.expand_dims(img, axis=0) + + observation = self.robot_preprocessor.get_observation(state, views) + return observation, instruction + + def infer_batch( + self, obs_list: List[Dict[str, Any]], skip_serialize: bool = False + ) -> List[Dict[str, Any]]: + """Perform batch inference. + + Args: + obs_list: List of observations, each containing: + - "views": dict of camera images + - "state": robot state dict or array + - "instruction": text instruction + + Returns: + List of action dicts. + """ + batch_size = len(obs_list) + logger.info(f"WallXPolicy.infer_batch: processing {batch_size} observations") + + try: + observations = [] + instructions = [] + for obs in obs_list: + observation, instruction = self._preprocess_obs(obs) + observations.append(observation) + instructions.append(instruction) + + mode = obs_list[0].get("infer_mode", self.default_infer_mode) + + with torch.no_grad(): + if mode == INFER_MODE_FLOW: + model_outputs = self.model_wrapper.infer_flow_action_batch( + observations, instructions + ) + else: + # AR / flow_with_subtask: fall back to per-sample inference + model_outputs = [] + for obs_dict, instruction in zip(observations, instructions): + output = self._run_action_infer(obs_dict, instruction, mode) + model_outputs.append(output) + + if skip_serialize: + return [{}] * len(model_outputs) + + results = [] + for model_output in model_outputs: + self._apply_smoothing(model_output) + if self.serialize_actions and not self._is_single_arm_right_only(): + action = self.robot_preprocessor.get_serialized_actions( + model_output, robot_action_interpolate_multiplier=1 + ) + results.append(action) + else: + results.append(self._pack_action_chunk_response(model_output)) + + return results + + except Exception as e: + logger.error(f"Batch inference failed: {e}", exc_info=True) + return [{"action": {}, "error": str(e)} for _ in obs_list] + + @property + def metadata(self) -> Dict[str, Any]: + return {"batch_enabled": True, "model": "wall-x"} diff --git a/wall_x/_vendor/harrix/serving/scheduler.py b/wall_x/_vendor/harrix/serving/scheduler.py new file mode 100644 index 0000000..8d36fdb --- /dev/null +++ b/wall_x/_vendor/harrix/serving/scheduler.py @@ -0,0 +1,160 @@ +import asyncio +import logging +import time +import uuid +from collections import deque +from dataclasses import dataclass +from typing import Any, Dict, List + +logger = logging.getLogger(__name__) + + +@dataclass +class Request: + request_id: str + obs: Dict[str, Any] + future: asyncio.Future + timestamp: float + + +class RequestScheduler: + def __init__( + self, + policy, + max_batch_size: int = 8, + max_wait_time_ms: float = 100, + max_queue_size: int = 128, + timeout_ms: float = 5000, + ): + self.policy = policy + self.max_batch_size = max_batch_size + self.max_wait_time = max_wait_time_ms / 1000.0 + self.max_queue_size = max_queue_size + self.timeout = timeout_ms / 1000.0 + + self.queue = deque() + self.queue_lock = asyncio.Lock() + self.queue_not_empty = asyncio.Condition(self.queue_lock) + + self.running = False + self.batch_task = None + + async def start(self): + self.running = True + self.batch_task = asyncio.create_task(self._batch_loop()) + logger.info( + f"RequestScheduler started: max_batch={self.max_batch_size}, max_wait={self.max_wait_time*1000}ms" + ) + + async def stop(self): + self.running = False + # Process all remaining requests before shutting down + async with self.queue_lock: + remaining = len(self.queue) + if remaining: + logger.info( + f"Graceful shutdown: processing {remaining} remaining request(s)" + ) + while True: + async with self.queue_lock: + if len(self.queue) == 0: + break + batch = [] + while self.queue and len(batch) < self.max_batch_size: + batch.append(self.queue.popleft()) + if batch: + await self._process_batch(batch) + # Now stop the batch loop + async with self.queue_lock: + self.queue_not_empty.notify_all() + if self.batch_task: + await self.batch_task + logger.info("RequestScheduler stopped") + + async def add_request(self, obs: Dict[str, Any]) -> Dict[str, Any]: + request_id = str(uuid.uuid4()) + future = asyncio.Future() + request = Request( + request_id=request_id, obs=obs, future=future, timestamp=time.monotonic() + ) + + async with self.queue_lock: + if len(self.queue) >= self.max_queue_size: + raise RuntimeError( + f"Queue full: {len(self.queue)}/{self.max_queue_size}" + ) + self.queue.append(request) + self.queue_not_empty.notify() + + try: + result = await asyncio.wait_for(future, timeout=self.timeout) + return result + except asyncio.TimeoutError: + logger.error(f"Request {request_id} timeout after {self.timeout}s") + raise + + async def _batch_loop(self): + while self.running: + try: + batch = await self._collect_batch() + if batch: + await self._process_batch(batch) + except Exception as e: + logger.error(f"Unexpected error in batch loop: {e}", exc_info=True) + + async def _collect_batch(self) -> List[Request]: + async with self.queue_not_empty: + while self.running and len(self.queue) == 0: + await self.queue_not_empty.wait() + + if not self.running: + return [] + + batch = [] + deadline = time.monotonic() + self.max_wait_time + + while len(batch) < self.max_batch_size: + if len(self.queue) > 0: + batch.append(self.queue.popleft()) + + if len(batch) >= self.max_batch_size: + break + + if len(self.queue) == 0: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + await asyncio.wait_for( + self.queue_not_empty.wait(), timeout=remaining + ) + except asyncio.TimeoutError: + break + + return batch + + async def _process_batch(self, batch: List[Request]): + if not batch: + return + + start_time = time.monotonic() + logger.info(f"Processing batch of {len(batch)} requests") + + try: + obs_list = [req.obs for req in batch] + results = await asyncio.to_thread(self.policy.infer_batch, obs_list) + + for req, result in zip(batch, results): + if not req.future.done(): + req.future.set_result(result) + + infer_time = time.monotonic() - start_time + logger.info( + f"Batch processed in {infer_time*1000:.1f}ms, throughput: {len(batch)/infer_time:.1f} req/s" + ) + + except Exception as e: + logger.error(f"Batch processing failed: {e}", exc_info=True) + for req in batch: + if not req.future.done(): + req.future.set_exception(e) diff --git a/wall_x/serving/websocket_policy_server.py b/wall_x/_vendor/harrix/serving/websocket_policy_server.py similarity index 56% rename from wall_x/serving/websocket_policy_server.py rename to wall_x/_vendor/harrix/serving/websocket_policy_server.py index 719a110..54ae810 100644 --- a/wall_x/serving/websocket_policy_server.py +++ b/wall_x/_vendor/harrix/serving/websocket_policy_server.py @@ -3,7 +3,7 @@ import http import logging import time import traceback -from typing import Any, Dict, Optional +from typing import Dict, Any, List try: import msgpack @@ -29,6 +29,10 @@ class BasePolicy: """Infer actions from observations.""" raise NotImplementedError + def infer_batch(self, obs_list: List[Dict]) -> List[Dict]: + """Batch inference. Default implementation calls infer() per sample.""" + return [self.infer(obs) for obs in obs_list] + def reset(self) -> None: """Reset the policy to its initial state.""" pass @@ -45,8 +49,11 @@ class WebsocketPolicyServer: Implements a websocket server that: 1. Sends policy metadata on connection 2. Receives observations - 3. Returns predicted actions + 3. Returns predicted actions (single or batched) 4. Tracks timing information + + When batching parameters are provided, dynamic batching is enabled: + requests from multiple clients are queued and processed in batches. """ def __init__( @@ -54,30 +61,66 @@ class WebsocketPolicyServer: policy: BasePolicy, host: str = "0.0.0.0", port: int = 8000, - metadata: Optional[Dict] = None, + metadata: Dict | None = None, + # Dynamic batching parameters (None = disabled) + max_batch_size: int | None = None, + max_wait_time_ms: float | None = None, + max_queue_size: int = 100, + timeout_ms: float = 30000, ) -> None: self._policy = policy self._host = host self._port = port self._metadata = metadata or {} + + # Dynamic batching + self._scheduler = None + if max_batch_size is not None: + from .scheduler import RequestScheduler + + self._scheduler = RequestScheduler( + policy=policy, + max_batch_size=max_batch_size, + max_wait_time_ms=( + max_wait_time_ms if max_wait_time_ms is not None else 0 + ), + max_queue_size=max_queue_size, + timeout_ms=timeout_ms, + ) + logging.getLogger("websockets.server").setLevel(logging.INFO) + @property + def batching_enabled(self) -> bool: + return self._scheduler is not None + def serve_forever(self) -> None: asyncio.run(self.run()) async def run(self): - async with _server.serve( - self._handler, - self._host, - self._port, - compression=None, - max_size=None, - ping_interval=None, # Disable automatic ping for long-running inference - ping_timeout=None, # Disable ping timeout - process_request=_health_check, - ) as server: - logger.info(f"Server started on {self._host}:{self._port}") - await server.serve_forever() + # Start the scheduler if batching is enabled + if self._scheduler is not None: + await self._scheduler.start() + + try: + async with _server.serve( + self._handler, + self._host, + self._port, + compression=None, + max_size=None, + ping_interval=None, + ping_timeout=None, + process_request=_health_check, + ) as server: + mode_str = "batched" if self.batching_enabled else "single" + logger.info( + f"Server started on {self._host}:{self._port} (mode={mode_str})" + ) + await server.serve_forever() + finally: + if self._scheduler is not None: + await self._scheduler.stop() async def _handler(self, websocket: _server.ServerConnection): logger.info(f"Connection from {websocket.remote_address} opened") @@ -99,7 +142,14 @@ class WebsocketPolicyServer: obs = msgpack.unpackb(await websocket.recv()) infer_time = time.monotonic() - action = self._policy.infer(obs) + + if self._scheduler is not None: + # Dynamic batching path + action = await self._scheduler.add_request(obs) + else: + # Single inference path + action = await asyncio.to_thread(self._policy.infer, obs) + infer_time = time.monotonic() - infer_time action["server_timing"] = { @@ -126,7 +176,9 @@ class WebsocketPolicyServer: def _health_check( connection: _server.ServerConnection, request: _server.Request -) -> Optional[_server.Response]: +) -> _server.Response | None: if request.path == "/healthz": return connection.respond(http.HTTPStatus.OK, "OK\n") + if request.path == "/v2/health/ready": + return connection.respond(http.HTTPStatus.OK, "OK\n") return None diff --git a/wall_x/_vendor/harrix/utils/__init__.py b/wall_x/_vendor/harrix/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/wall_x/_vendor/harrix/utils/ckpt_load.py b/wall_x/_vendor/harrix/utils/ckpt_load.py new file mode 100644 index 0000000..4e00350 --- /dev/null +++ b/wall_x/_vendor/harrix/utils/ckpt_load.py @@ -0,0 +1,171 @@ +"""Load checkpoint weights and apply fused-format conversion when needed. + +Model-instance operations such as ``load_state_dict`` and ``set_normalizer`` +are intentionally left to the adapter. +""" + +from __future__ import annotations + +import os +from typing import Callable, Optional + +import torch +from safetensors.torch import load_file + + +def _noop_log(_msg: str, **_kw) -> None: + pass + + +def _align_checkpoint_tensor( + param: torch.Tensor, + target: torch.Tensor, + name: str, + log_fn: Callable, +) -> torch.Tensor | None: + """Crop or pad a checkpoint tensor to match the current model parameter.""" + if param.shape == target.shape: + return param + if param.ndim != target.ndim: + log_fn( + f"Skipping '{name}': ndim mismatch " + f"checkpoint={param.ndim} model={target.ndim}" + ) + return None + + overlap = tuple( + slice(0, min(src, dst)) for src, dst in zip(param.shape, target.shape) + ) + + if all(src >= dst for src, dst in zip(param.shape, target.shape)): + aligned = param[overlap].contiguous() + log_fn( + f"Cropped '{name}': checkpoint {tuple(param.shape)} " + f"-> model {tuple(aligned.shape)}" + ) + return aligned + + if all(src <= dst for src, dst in zip(param.shape, target.shape)): + aligned = target.detach().clone() + aligned[overlap] = param[overlap] + log_fn( + f"Padded '{name}': checkpoint {tuple(param.shape)} " + f"-> model {tuple(aligned.shape)} (tail keeps model init)" + ) + return aligned + + aligned = target.detach().clone() + aligned[overlap] = param[overlap] + log_fn( + f"Partially aligned '{name}': checkpoint {tuple(param.shape)} " + f"-> model {tuple(aligned.shape)} (non-overlap keeps model init)" + ) + return aligned + + +def reshape_compatible_state_dict( + state_dict: dict, model_sd: dict, log_fn: Optional[Callable] = None +) -> dict: + """Align checkpoint tensors to the target model shapes via crop / pad.""" + log_fn = log_fn or _noop_log + out = {} + for name, param in state_dict.items(): + if name not in model_sd: + log_fn(f"Not used parameter: {name}") + continue + target = model_sd[name] + if param.shape == target.shape: + out[name] = param + continue + aligned = _align_checkpoint_tensor(param, target, name, log_fn) + if aligned is not None: + out[name] = aligned + return out + + +def load_state_dict(checkpoint_path: str, model_class) -> dict: + """Load a state dict from a checkpoint directory. + + Supported formats: + - pytorch_model_fsdp.bin, optionally wrapped as {"state_dict": ...} + - model.safetensors + + If the model class reports that the state dict is not fused, it is converted + through ``model_class.convert_to_fused``. + """ + fsdp_ckpt = os.path.join(checkpoint_path, "pytorch_model_fsdp.bin") + safetensor_ckpt = os.path.join(checkpoint_path, "model.safetensors") + + if os.path.exists(fsdp_ckpt): + state_dict = torch.load(fsdp_ckpt, map_location="cpu") + if isinstance(state_dict, dict) and "state_dict" in state_dict: + state_dict = state_dict["state_dict"] + elif os.path.exists(safetensor_ckpt): + state_dict = load_file(safetensor_ckpt, device="cpu") + else: + raise FileNotFoundError( + "checkpoint contains neither pytorch_model_fsdp.bin nor model.safetensors: " + f"{checkpoint_path}" + ) + + if not model_class.is_fused(state_dict): + state_dict = model_class.convert_to_fused(state_dict) + + return state_dict + + +def read_global_step(checkpoint_path: str) -> int | None: + """Read ``global_step.pth`` when present.""" + p = os.path.join(checkpoint_path, "global_step.pth") + if not os.path.exists(p): + return None + payload = torch.load(p) + return int(payload["global_step"]) + + +def _dir_has_weights(path: str) -> bool: + return os.path.exists( + os.path.join(path, "pytorch_model_fsdp.bin") + ) or os.path.exists(os.path.join(path, "model.safetensors")) + + +def resolve_checkpoint_dir(checkpoint_path: str) -> str: + """Return a directory that directly contains model weights. + + Training saves under a root such as ``libero6/`` with step subdirs + ``libero6/0/``, ``libero6/3/``, etc. Inference callers may pass either the + root or a concrete step directory. + """ + if os.path.isfile(checkpoint_path): + checkpoint_path = os.path.dirname(checkpoint_path) + + if _dir_has_weights(checkpoint_path): + return checkpoint_path + + if not os.path.isdir(checkpoint_path): + raise FileNotFoundError(f"checkpoint path does not exist: {checkpoint_path}") + + candidates: list[tuple[int, float, str]] = [] + for entry in os.listdir(checkpoint_path): + sub = os.path.join(checkpoint_path, entry) + if not os.path.isdir(sub) or not _dir_has_weights(sub): + continue + step = read_global_step(sub) + sort_step = step if step is not None else -1 + candidates.append((sort_step, os.path.getmtime(sub), sub)) + + if not candidates: + return checkpoint_path + + candidates.sort() + resolved = candidates[-1][2] + if resolved != checkpoint_path: + import logging + + logging.getLogger(__name__).info( + "Resolved checkpoint root %s -> %s (global_step=%s)", + checkpoint_path, + resolved, + read_global_step(resolved), + ) + return resolved diff --git a/wall_x/_vendor/harrix/utils/normalizer.py b/wall_x/_vendor/harrix/utils/normalizer.py new file mode 100644 index 0000000..2494b83 --- /dev/null +++ b/wall_x/_vendor/harrix/utils/normalizer.py @@ -0,0 +1,135 @@ +"""Build action/proprio normalizers and resolve the effective norm key. + +Public inference artifacts must carry their own normalization data. This module +uses checkpoint-local ``norm_stats.json`` first, then checkpoint-side normalizer +state dicts, and finally an explicit ``customized_action_statistic_dof`` path. +It does not fall back to internal default action statistics. +""" + +from __future__ import annotations + +import json +import logging +import os + +import torch + +from wall_x.data.backends.lerobot.utils import NormStats +from wall_x.model.core.action.normalizer import Normalizer, pad_normalizer_to_dim +from wall_x._vendor.harrix.utils.train_config import ( + resolve_agent_pos_config, + resolve_dof_config, +) + +logger = logging.getLogger(__name__) + + +def _load_norm_stats(norm_stats_path: str, action_key: str) -> NormStats: + with open(norm_stats_path, "r") as f: + norm_stats = json.load(f) + q01 = torch.tensor(norm_stats["norm_stats"][action_key]["q01"]) + q99 = torch.tensor(norm_stats["norm_stats"][action_key]["q99"]) + return NormStats(min=q01, max=q99, delta=q99 - q01) + + +def _load_custom_action_stats(train_config: dict) -> dict | None: + custom = train_config.get("customized_action_statistic_dof", None) + if not custom: + return None + with open(custom, "r") as f: + return json.load(f) + + +def _normalizer_from_stats(action_stats: dict, train_config: dict, key: str) -> Normalizer: + return Normalizer( + action_stats, + train_config[key], + min_key=train_config.get("min_key", "min"), + delta_key=train_config.get("delta_key", "delta"), + ) + + +def _missing_normalizer_error(checkpoint_path: str, train_config: dict) -> FileNotFoundError: + custom = train_config.get("customized_action_statistic_dof", None) + return FileNotFoundError( + "Public inference requires normalization data. Expected one of: " + f"{os.path.join(checkpoint_path, 'norm_stats.json')}; checkpoint-side " + "normalizer_action.pth and normalizer_propri.pth; or an explicit " + f"customized_action_statistic_dof path. Current customized_action_statistic_dof={custom!r}." + ) + + +def build_normalizers( + checkpoint_path: str, + train_config: dict, + norm_key: str, +) -> tuple[Normalizer, Normalizer, str]: + """Return action/proprio normalizers and the resolved norm key.""" + norm_stats_path = os.path.join(checkpoint_path, "norm_stats.json") + if os.path.exists(norm_stats_path): + propri_stats = _load_norm_stats(norm_stats_path, "observation.state") + action_stats = _load_norm_stats(norm_stats_path, "action") + normalizer_propri = Normalizer.from_lerobot_norm_stats(propri_stats, norm_key) + normalizer_action = Normalizer.from_lerobot_norm_stats(action_stats, norm_key) + else: + action_pth = os.path.join(checkpoint_path, "normalizer_action.pth") + propri_pth = os.path.join(checkpoint_path, "normalizer_propri.pth") + custom_stats = _load_custom_action_stats(train_config) + if custom_stats is None and (not os.path.exists(action_pth) or not os.path.exists(propri_pth)): + raise _missing_normalizer_error(checkpoint_path, train_config) + + if os.path.exists(action_pth): + normalizer_action = Normalizer.from_ckpt(action_pth) + else: + normalizer_action = _normalizer_from_stats(custom_stats, train_config, "dof_config") + + if os.path.exists(propri_pth): + normalizer_propri = Normalizer.from_ckpt(propri_pth) + else: + normalizer_propri = _normalizer_from_stats(custom_stats, train_config, "agent_pos_config") + + action_dim = sum(resolve_dof_config(train_config).values()) + propri_dim = sum(resolve_agent_pos_config(train_config).values()) + pad_normalizer_to_dim(normalizer_action, action_dim, "action") + pad_normalizer_to_dim(normalizer_propri, propri_dim, "propri") + + resolved = _resolve_norm_key(norm_key, normalizer_action, normalizer_propri) + return normalizer_action, normalizer_propri, resolved + + +def _resolve_norm_key( + norm_key: str, + normalizer_action: Normalizer, + normalizer_propri: Normalizer, +) -> str: + """Resolve a requested norm key against normalizer keys.""" + available = sorted( + set(normalizer_action.min.keys()) & set(normalizer_propri.min.keys()) + ) + if norm_key in available: + return norm_key + if not available: + return norm_key + + prefix_matches = [k for k in available if k.startswith(f"{norm_key}_")] + if len(prefix_matches) == 1: + logger.warning( + "norm_key=%r not found; using prefix fallback %r", + norm_key, + prefix_matches[0], + ) + return prefix_matches[0] + if len(available) == 1: + logger.warning( + "norm_key=%r not found; using the only available key %r", + norm_key, + available[0], + ) + return available[0] + + logger.warning( + "norm_key=%r not found; available=%s; returning the requested key unchanged", + norm_key, + available, + ) + return norm_key diff --git a/wall_x/_vendor/harrix/utils/seed.py b/wall_x/_vendor/harrix/utils/seed.py new file mode 100644 index 0000000..9c4fd8d --- /dev/null +++ b/wall_x/_vendor/harrix/utils/seed.py @@ -0,0 +1,23 @@ +"""Seed helpers for driver and environment worker processes.""" + +from __future__ import annotations + +import os +import random + +import numpy as np +import torch + + +def set_seed_everywhere(seed: int) -> None: + """Set random seeds and deterministic backend options.""" + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + random.seed(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ["PYTHONHASHSEED"] = str(seed) + # Deterministic cuBLAS workspace plus warning-only deterministic op checks. + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + torch.use_deterministic_algorithms(True, warn_only=True) diff --git a/wall_x/_vendor/harrix/utils/train_config.py b/wall_x/_vendor/harrix/utils/train_config.py new file mode 100644 index 0000000..435c28a --- /dev/null +++ b/wall_x/_vendor/harrix/utils/train_config.py @@ -0,0 +1,371 @@ +"""Load checkpoint-side train config and construct runtime config objects. + +Inference adapters use this module to load the train config, apply checkpoint +overrides for moved processor/tokenizer files, build model config, and build the +data config required by preprocessing. +""" + +from __future__ import annotations + +import os +import logging + +import yaml + +logger = logging.getLogger(__name__) + + +_TASK_INFERENCE_KEYS = ( + "dof_config", + "agent_pos_config", + "ar_dof_config", + "action_horizon", + "action_horizon_flow", + "noise_scheduler", + "use_state_string_representation", +) + +_VIRTUAL_TAIL_KEYS = frozenset(("action_padding",)) + + +def _move_virtual_keys_to_tail(layout: dict | None) -> dict | None: + """Keep real LeRobot dims before virtual padding dims in inference layouts.""" + if not isinstance(layout, dict): + return layout + head = {k: v for k, v in layout.items() if k not in _VIRTUAL_TAIL_KEYS} + tail = {k: v for k, v in layout.items() if k in _VIRTUAL_TAIL_KEYS} + if not tail: + return layout + return {**head, **tail} + + +def _canonicalize_task_layouts(target: dict) -> None: + if os.environ.get("WALLX_CANONICALIZE_VIRTUAL_DOF_ORDER", "1") == "0": + return + task = target.get("task") + data = target.get("data") + for key in ("dof_config", "agent_pos_config", "ar_dof_config"): + value = _move_virtual_keys_to_tail(target.get(key)) + if value is not None: + target[key] = value + if isinstance(task, dict): + task_value = _move_virtual_keys_to_tail(task.get(key)) + if task_value is not None: + task[key] = task_value + if isinstance(data, dict): + data_value = _move_virtual_keys_to_tail(data.get(key)) + if data_value is not None: + data[key] = data_value + + +def resolve_use_state_string_representation(train_config: dict) -> bool: + """Read ``use_state_string_representation`` with task YAML as authority.""" + task = train_config.get("task") or {} + if task.get("use_state_string_representation") is not None: + return bool(task["use_state_string_representation"]) + if train_config.get("use_state_string_representation") is not None: + return bool(train_config["use_state_string_representation"]) + data = train_config.get("data") or {} + if data.get("use_state_string_representation") is not None: + return bool(data["use_state_string_representation"]) + return False + + +def resolve_state_bins(train_config: dict, default: int = 256) -> int: + """Read discretization bin count from flat or nested train config.""" + if train_config.get("state_bins") is not None: + return int(train_config["state_bins"]) + data = train_config.get("data") or {} + if data.get("state_bins") is not None: + return int(data["state_bins"]) + return default + + +def resolve_agent_pos_config(train_config: dict) -> dict: + """Return ``agent_pos_config`` from flat or ``task``-nested train YAML.""" + task = train_config.get("task") or {} + agent_pos_config = train_config.get("agent_pos_config") or task.get( + "agent_pos_config" + ) + if not agent_pos_config: + raise KeyError( + "agent_pos_config missing from train config " + "(expected top-level or task.agent_pos_config)" + ) + return dict(agent_pos_config) + + +def resolve_dof_config(train_config: dict) -> dict: + """Return ``dof_config`` from flat or ``task``-nested train YAML.""" + task = train_config.get("task") or {} + dof_config = train_config.get("dof_config") or task.get("dof_config") + if not dof_config: + raise KeyError( + "dof_config missing from train config (expected top-level or task.dof_config)" + ) + return dict(dof_config) + + +def resolve_cam_names_from_train_config(train_config: dict) -> list[str] | None: + """Infer model camera keys from ``data.key_mappings.camera`` (e.g. LIBERO).""" + data = train_config.get("data") or {} + key_mappings = data.get("key_mappings") or {} + camera = key_mappings.get("camera") or {} + if not camera: + return None + names: list[str] = [] + for value in camera.values(): + name = str(value) + if name not in names: + names.append(name) + return names or None + + +def resolve_camera_label(cam_name: str, camera_name_mapping: dict | None = None) -> str: + """Match training ``get_wallx_normal_text`` camera display names.""" + mapping = camera_name_mapping or {} + return mapping.get(cam_name, cam_name) + + +def resolve_max_length(train_config: dict, default: int = 768) -> int: + """Read ``max_length`` from flat or nested train YAML (collator default 768).""" + data = train_config.get("data") or {} + if train_config.get("max_length") is not None: + return int(train_config["max_length"]) + if data.get("max_length") is not None: + return int(data["max_length"]) + raw = train_config.get("_raw_data") or {} + if raw.get("max_length") is not None: + return int(raw["max_length"]) + return default + + +# Fields ``load_wallx_processors`` and model construction read from a flat dict. +_ACTION_TOKENIZER_KEYS = ( + "action_tokenizer_type", + "action_tokenizer_path", + "action_tokenizer_checkpoint_path", + "action_tokenizer_config_dir", + "action_tokenizer", +) + +_INFERENCE_MODEL_KEYS = ( + "processor_path", + "pretrained_path", + "config_path", + *_ACTION_TOKENIZER_KEYS, + "ar_loss_weight", + "attn_deterministic", + "flow_loss_weight", + "use_ema", +) + + +def _merge_model_fields(target: dict, source: dict, *, overwrite: bool = False) -> None: + """Copy model-related keys from ``source`` into flat ``target``.""" + if not source: + return + for key in _INFERENCE_MODEL_KEYS: + value = source.get(key) + if value is None: + continue + if overwrite or target.get(key) in (None, ""): + target[key] = value + + +def _mirror_task_fields(target: dict, task: dict) -> None: + """Copy typed task fields into the flat/data legacy inference mirrors.""" + if not isinstance(task, dict): + return + data = target.setdefault("data", {}) + if not isinstance(data, dict): + data = {} + target["data"] = data + for key in _TASK_INFERENCE_KEYS: + value = task.get(key) + if value is None: + continue + target[key] = value + data[key] = value + + +def _load_ckpt_config_overlay(checkpoint_path: str) -> dict: + """Read ``config.yml`` saved beside a training checkpoint.""" + ckpt_yml = os.path.join(checkpoint_path, "config.yml") + if not os.path.exists(ckpt_yml): + return {} + with open(ckpt_yml, "r") as f: + raw = yaml.load(f, Loader=yaml.FullLoader) or {} + + overlay: dict = {} + _merge_model_fields(overlay, raw.get("model") or {}, overwrite=True) + raw_yaml = raw.get("_raw_yaml") or {} + _merge_model_fields(overlay, raw_yaml.get("model") or {}, overwrite=False) + _merge_model_fields(overlay, raw, overwrite=False) + return overlay + + +def strip_action_tokenizer_fields(train_config: dict) -> None: + """Drop action-tokenizer fields so flow inference can skip AR tokenizer setup.""" + for key in _ACTION_TOKENIZER_KEYS: + train_config.pop(key, None) + model = train_config.get("model") + if isinstance(model, dict): + for key in _ACTION_TOKENIZER_KEYS: + model.pop(key, None) + + +def load_train_config_with_ckpt_overlay( + train_config_path: str, + checkpoint_path: str, +) -> dict: + """Load train YAML and apply checkpoint-local processor/tokenizer overlays. + + This lets a checkpoint remain portable when the original training-machine + processor or action-tokenizer paths are no longer available. + """ + with open(train_config_path, "r") as f: + train_config = yaml.load(f, Loader=yaml.FullLoader) + + preprocessor_file = os.path.join(checkpoint_path, "preprocessor_config.json") + if os.path.exists(preprocessor_file): + train_config["processor_path"] = checkpoint_path + + orig_action_tok = train_config.get("action_tokenizer_path", None) + if orig_action_tok is not None and not os.path.exists(orig_action_tok): + tokenizer_file = os.path.join(checkpoint_path, "tokenizer.json") + tokenizer_config_file = os.path.join(checkpoint_path, "tokenizer_config.json") + if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file): + train_config["action_tokenizer_path"] = checkpoint_path + + ckpt_overlay = _load_ckpt_config_overlay(checkpoint_path) + _merge_model_fields(train_config, ckpt_overlay, overwrite=False) + if isinstance(train_config.get("model"), dict): + _merge_model_fields(train_config["model"], ckpt_overlay, overwrite=False) + + return train_config + + +def normalize_train_config_for_inference( + train_config: dict, + train_config_path: str, +) -> dict: + """Flatten typed TrainConfig YAML into the legacy dict inference expects. + + Serving and ``update_model_config`` still read top-level ``dof_config`` and + nested ``data.*`` fields. New training yamls store those under ``task:``. + """ + try: + from wall_x.config.loader import load_config + + typed_cfg = load_config(train_config_path) + except Exception as e: + logger.debug( + "Train config is not typed TrainConfig schema (%s): %s", + train_config_path, + e, + ) + _mirror_task_fields(train_config, train_config.get("task") or {}) + _canonicalize_task_layouts(train_config) + return train_config + + import dataclasses + + normalized = typed_cfg.build_data_loader_dict() + model_dict = dataclasses.asdict(typed_cfg.model) + _merge_model_fields(normalized, model_dict, overwrite=True) + + normalized["data"] = dict(normalized.get("data") or {}) + if isinstance(train_config.get("model"), dict): + _merge_model_fields(normalized, train_config["model"], overwrite=False) + for key, value in train_config.items(): + if key in _INFERENCE_MODEL_KEYS or key == "qwen_vl_act_config_path": + if value is not None: + normalized[key] = value + elif key == "data" and isinstance(value, dict): + normalized["data"].update(value) + + _mirror_task_fields(normalized, dataclasses.asdict(typed_cfg.task)) + _canonicalize_task_layouts(normalized) + return normalized + + +def register_data_backend(train_config: dict) -> None: + """Register the data backend before typed data config construction.""" + from wall_x.data._registry import _set_data_backend + + data_section = train_config.get("data", {}) + dataset_type = train_config.get("dataset_type") or data_section.get( + "dataset_type", "lerobot" + ) + _set_data_backend(dataset_type) + + +def build_model_config( + config_class, + checkpoint_path: str, + train_config: dict, + train_config_path: str | None = None, +): + """Build the HF model config and inject train_config-derived fields. + + ``config_class`` is supplied by the variant adapter. If the checkpoint does + not include ``config.json``, this falls back to + ``train_config["qwen_vl_act_config_path"]``. + """ + ckpt_config_path = os.path.join(checkpoint_path, "config.json") + if os.path.exists(ckpt_config_path): + resolved = ckpt_config_path + else: + resolved = train_config.get("qwen_vl_act_config_path") + if resolved is None or not os.path.exists(resolved): + raise ValueError( + f"cannot load model config: checkpoint file {ckpt_config_path} " + f"does not exist, and fallback qwen_vl_act_config_path={resolved!r} " + "does not exist either" + ) + train_config["qwen_vl_act_config_path"] = resolved + + if resolved.endswith(".json"): + model_config = config_class.from_json_file(resolved) + else: + model_config = config_class.from_pretrained(resolved) + + legacy_train_config = ( + normalize_train_config_for_inference(train_config, train_config_path) + if train_config_path + else train_config + ) + model_config.update_model_config(legacy_train_config) + model_config._attn_implementation = "sdpa" + model_config.vision_config._attn_implementation = "flash_attention_2" + + return model_config + + +def build_data_config(train_config_path: str, train_config: dict): + """Build the data config used by image resizing and preprocessing. + + New checkpoints use the typed schema. Older checkpoints may still use a flat + legacy schema, so this falls back to the active data backend. + """ + from wall_x.config.loader import load_config + from wall_x.data import data_backend + + try: + typed_cfg = load_config(train_config_path) + backend = data_backend() + if backend.supports("load_trainer_data_config"): + return backend.load_trainer_data_config(typed_cfg) + except Exception as e: + logger.warning( + "typed TrainConfig loading failed; falling back to raw data config: %s", + e, + ) + + backend = data_backend() + if backend.supports("load_trainer_data_config_from_yaml_dict"): + return backend.load_trainer_data_config_from_yaml_dict(train_config) + raise RuntimeError( + f"active data backend {backend!r} cannot build a trainer data config" + ) diff --git a/wall_x/_vendor/x2robot_utils/__init__.py b/wall_x/_vendor/x2robot_utils/__init__.py new file mode 100644 index 0000000..535937f --- /dev/null +++ b/wall_x/_vendor/x2robot_utils/__init__.py @@ -0,0 +1,5 @@ +"""Vendored subset of x2robot_utils for Wall-X.""" + +__version__ = "0.2.0-vendored" + +__all__ = ["__version__"] diff --git a/wall_x/_vendor/x2robot_utils/geometry.py b/wall_x/_vendor/x2robot_utils/geometry.py new file mode 100644 index 0000000..3816ce1 --- /dev/null +++ b/wall_x/_vendor/x2robot_utils/geometry.py @@ -0,0 +1,309 @@ +"""Rotation / pose geometry utilities (pure numpy + numba). + +Shared between wall-x and internal_dataset_backend: +- ``euler_to_matrix_zyx_6d_nb``: ZYX Euler -> flattened top-2-rows of R (N, 6) +- ``so3_to_euler_zyx_batch_nb``: 6D rotation -> ZYX Euler (canonicalized) +- ``compose_state_and_delta_to_abs_{rpy,6d}``: state + delta -> absolute pose +""" + +import numpy as np +from numba import jit, prange + + +@jit(nopython=True) +def euler_to_matrix_zyx_6d_nb(eulers): + """Euler angles (N, 3) -> flattened top two rows of rotation matrix (N, 6).""" + N = eulers.shape[0] + R6 = np.empty((N, 6), dtype=np.float64) + for i in prange(N): + roll = eulers[i, 0] + pitch = eulers[i, 1] + yaw = eulers[i, 2] + + cy, sy = np.cos(yaw), np.sin(yaw) + cp, sp = np.cos(pitch), np.sin(pitch) + cr, sr = np.cos(roll), np.sin(roll) + + r00 = cy * cp + r01 = cy * sp * sr - sy * cr + r02 = cy * sp * cr + sy * sr + + r10 = sy * cp + r11 = sy * sp * sr + cy * cr + r12 = sy * sp * cr - cy * sr + + R6[i, 0] = r00 + R6[i, 1] = r01 + R6[i, 2] = r02 + R6[i, 3] = r10 + R6[i, 4] = r11 + R6[i, 5] = r12 + return R6 + + +@jit(nopython=True) +def euler_to_matrix_zyx_batch_nb(eulers): + N = eulers.shape[0] + R = np.empty((N, 3, 3), dtype=np.float64) + for i in prange(N): + roll = eulers[i, 0] + pitch = eulers[i, 1] + yaw = eulers[i, 2] + + cy, sy = np.cos(yaw), np.sin(yaw) + cp, sp = np.cos(pitch), np.sin(pitch) + cr, sr = np.cos(roll), np.sin(roll) + + R[i, 0, 0] = cy * cp + R[i, 0, 1] = cy * sp * sr - sy * cr + R[i, 0, 2] = cy * sp * cr + sy * sr + + R[i, 1, 0] = sy * cp + R[i, 1, 1] = sy * sp * sr + cy * cr + R[i, 1, 2] = sy * sp * cr - cy * sr + + R[i, 2, 0] = -sp + R[i, 2, 1] = cp * sr + R[i, 2, 2] = cp * cr + return R + + +@jit(nopython=True) +def matrix_to_euler_zyx_batch_nb(Rs): + """R = Rz(yaw) * Ry(pitch) * Rx(roll) -> (roll, pitch, yaw).""" + N = Rs.shape[0] + eulers = np.empty((N, 3), dtype=np.float64) + for i in prange(N): + r00 = Rs[i, 0, 0] + r10 = Rs[i, 1, 0] + r20 = Rs[i, 2, 0] + r21 = Rs[i, 2, 1] + r22 = Rs[i, 2, 2] + + x = -r20 + if x > 1.0: + x = 1.0 + elif x < -1.0: + x = -1.0 + + pitch = np.arcsin(x) + roll = np.arctan2(r21, r22) + yaw = np.arctan2(r10, r00) + + eulers[i, 0] = roll + eulers[i, 1] = pitch + eulers[i, 2] = yaw + return eulers + + +@jit(nopython=True) +def so3_to_matrix_batch_nb(batch_so3): + N = batch_so3.shape[0] + R_all = np.empty((N, 3, 3), dtype=np.float64) + eps = 1e-12 + for i in prange(N): + r1x, r1y, r1z = batch_so3[i, 0], batch_so3[i, 1], batch_so3[i, 2] + r2x, r2y, r2z = batch_so3[i, 3], batch_so3[i, 4], batch_so3[i, 5] + + n1 = np.sqrt(r1x * r1x + r1y * r1y + r1z * r1z) + eps + r1x /= n1 + r1y /= n1 + r1z /= n1 + + dot12 = r1x * r2x + r1y * r2y + r1z * r2z + r2x -= dot12 * r1x + r2y -= dot12 * r1y + r2z -= dot12 * r1z + n2 = np.sqrt(r2x * r2x + r2y * r2y + r2z * r2z) + eps + r2x /= n2 + r2y /= n2 + r2z /= n2 + + r3x = r1y * r2z - r1z * r2y + r3y = r1z * r2x - r1x * r2z + r3z = r1x * r2y - r1y * r2x + + R_all[i, 0, 0] = r1x + R_all[i, 0, 1] = r1y + R_all[i, 0, 2] = r1z + R_all[i, 1, 0] = r2x + R_all[i, 1, 1] = r2y + R_all[i, 1, 2] = r2z + R_all[i, 2, 0] = r3x + R_all[i, 2, 1] = r3y + R_all[i, 2, 2] = r3z + return R_all + + +@jit(nopython=True) +def canonicalize_euler_zyx_batch_nb(rpy_batch): + """Canonicalize ZYX Euler angles so each component falls in (-pi, pi].""" + N = rpy_batch.shape[0] + out = np.empty_like(rpy_batch) + two_pi = 2.0 * np.pi + + for i in prange(N): + r = rpy_batch[i, 0] + p = rpy_batch[i, 1] + y = rpy_batch[i, 2] + + r = (r + np.pi) % two_pi - np.pi + p = (p + np.pi) % two_pi - np.pi + y = (y + np.pi) % two_pi - np.pi + + if p > np.pi / 2.0: + p = np.pi - p + r = r + np.pi + y = y + np.pi + elif p <= -np.pi / 2.0: + p = -np.pi - p + r = r + np.pi + y = y + np.pi + + r = (r + np.pi) % two_pi - np.pi + p = (p + np.pi) % two_pi - np.pi + y = (y + np.pi) % two_pi - np.pi + + out[i, 0] = r + out[i, 1] = p + out[i, 2] = y + + return out + + +def so3_to_euler_zyx_batch_nb(batch_so3): + matrix = so3_to_matrix_batch_nb(batch_so3) + eulers = matrix_to_euler_zyx_batch_nb(matrix) + return canonicalize_euler_zyx_batch_nb(eulers) + + +@jit(nopython=True) +def compose_state_and_delta_to_abs_rpy(delta, state): + """Compose a delta (ZYX rpy or 6D) with an absolute state -> absolute rpy(ZYX). + + delta: (N, 3) deltarpy or (N, 6) delta6D. state: (3,) rpy or (6,) 6D. + Output: (N, 3) rpy canonicalized into (-pi, pi]. + """ + if delta.shape[-1] == 3: + R_delta = euler_to_matrix_zyx_batch_nb(delta) + elif delta.shape[-1] == 6: + R_delta = so3_to_matrix_batch_nb(delta) + else: + raise ValueError(f"delta last dim must be 3 or 6, got {delta.shape[-1]}") + + if state.shape[-1] == 3: + R_state = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0] + elif state.shape[-1] == 6: + R_state = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] + else: + raise ValueError(f"state last dim must be 3 or 6, got {state.shape[-1]}") + + N = R_delta.shape[0] + R_abs = np.empty((N, 3, 3), dtype=np.float64) + + S00 = R_state[0, 0] + S01 = R_state[0, 1] + S02 = R_state[0, 2] + S10 = R_state[1, 0] + S11 = R_state[1, 1] + S12 = R_state[1, 2] + S20 = R_state[2, 0] + S21 = R_state[2, 1] + S22 = R_state[2, 2] + + for i in prange(N): + A00 = R_delta[i, 0, 0] + A01 = R_delta[i, 0, 1] + A02 = R_delta[i, 0, 2] + A10 = R_delta[i, 1, 0] + A11 = R_delta[i, 1, 1] + A12 = R_delta[i, 1, 2] + A20 = R_delta[i, 2, 0] + A21 = R_delta[i, 2, 1] + A22 = R_delta[i, 2, 2] + + R_abs[i, 0, 0] = A00 * S00 + A01 * S10 + A02 * S20 + R_abs[i, 0, 1] = A00 * S01 + A01 * S11 + A02 * S21 + R_abs[i, 0, 2] = A00 * S02 + A01 * S12 + A02 * S22 + + R_abs[i, 1, 0] = A10 * S00 + A11 * S10 + A12 * S20 + R_abs[i, 1, 1] = A10 * S01 + A11 * S11 + A12 * S21 + R_abs[i, 1, 2] = A10 * S02 + A11 * S12 + A12 * S22 + + R_abs[i, 2, 0] = A20 * S00 + A21 * S10 + A22 * S20 + R_abs[i, 2, 1] = A20 * S01 + A21 * S11 + A22 * S21 + R_abs[i, 2, 2] = A20 * S02 + A21 * S12 + A22 * S22 + + abs_rpy = matrix_to_euler_zyx_batch_nb(R_abs) + abs_rpy = canonicalize_euler_zyx_batch_nb(abs_rpy) + + return abs_rpy + + +@jit(nopython=True) +def compose_state_and_delta_to_abs_6d(delta, state): + """Compose a 6D delta with a 6D state -> absolute 6D rotation. + + delta: (N, 6). state: (6,). Output: (N, 6). + """ + R_delta = so3_to_matrix_batch_nb(delta) + R_state = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] + + N = R_delta.shape[0] + R_abs = np.empty((N, 3, 3), dtype=np.float64) + + S00 = R_state[0, 0] + S01 = R_state[0, 1] + S02 = R_state[0, 2] + S10 = R_state[1, 0] + S11 = R_state[1, 1] + S12 = R_state[1, 2] + S20 = R_state[2, 0] + S21 = R_state[2, 1] + S22 = R_state[2, 2] + + for i in prange(N): + A00 = R_delta[i, 0, 0] + A01 = R_delta[i, 0, 1] + A02 = R_delta[i, 0, 2] + A10 = R_delta[i, 1, 0] + A11 = R_delta[i, 1, 1] + A12 = R_delta[i, 1, 2] + A20 = R_delta[i, 2, 0] + A21 = R_delta[i, 2, 1] + A22 = R_delta[i, 2, 2] + + R_abs[i, 0, 0] = A00 * S00 + A01 * S10 + A02 * S20 + R_abs[i, 0, 1] = A00 * S01 + A01 * S11 + A02 * S21 + R_abs[i, 0, 2] = A00 * S02 + A01 * S12 + A02 * S22 + + R_abs[i, 1, 0] = A10 * S00 + A11 * S10 + A12 * S20 + R_abs[i, 1, 1] = A10 * S01 + A11 * S11 + A12 * S21 + R_abs[i, 1, 2] = A10 * S02 + A11 * S12 + A12 * S22 + + R_abs[i, 2, 0] = A20 * S00 + A21 * S10 + A22 * S20 + R_abs[i, 2, 1] = A20 * S01 + A21 * S11 + A22 * S21 + R_abs[i, 2, 2] = A20 * S02 + A21 * S12 + A22 * S22 + + abs_6d = np.empty((N, 6), dtype=np.float64) + for i in prange(N): + abs_6d[i, 0] = R_abs[i, 0, 0] + abs_6d[i, 1] = R_abs[i, 0, 1] + abs_6d[i, 2] = R_abs[i, 0, 2] + abs_6d[i, 3] = R_abs[i, 1, 0] + abs_6d[i, 4] = R_abs[i, 1, 1] + abs_6d[i, 5] = R_abs[i, 1, 2] + + return abs_6d + + +__all__ = [ + "euler_to_matrix_zyx_6d_nb", + "euler_to_matrix_zyx_batch_nb", + "matrix_to_euler_zyx_batch_nb", + "so3_to_matrix_batch_nb", + "canonicalize_euler_zyx_batch_nb", + "so3_to_euler_zyx_batch_nb", + "compose_state_and_delta_to_abs_rpy", + "compose_state_and_delta_to_abs_6d", +] diff --git a/wall_x/_vendor/x2robot_utils/grounding.py b/wall_x/_vendor/x2robot_utils/grounding.py new file mode 100644 index 0000000..7e7eb5d --- /dev/null +++ b/wall_x/_vendor/x2robot_utils/grounding.py @@ -0,0 +1,172 @@ +"""Grounding-point helpers (point / bbox coordinate remap, pure regex).""" + +import re +from typing import List, Optional + + +def process_grounding_points( + text: str, orig_height, orig_width, resized_height, resized_width, model_type +) -> str: + """Remap // coordinates inside ``text`` from the original + image size to the resized space used by the given model type. + """ + point_pattern = re.compile(r"<(point|box|bbox)>(.*?)") + + def process_match(match): + tag_name = match.group(1) + coords_str = match.group(2) + try: + coords = list(map(int, re.findall(r"\d+", coords_str))) + + scale_w = resized_width / orig_width + scale_h = resized_height / orig_height + + if len(coords) == 2: + x, y = coords + if model_type == "qwen2_5": + new_x = max(0, min(round(x * scale_w), resized_width - 1)) + new_y = max(0, min(round(y * scale_h), resized_height - 1)) + elif model_type in ["qwen2"]: + new_x = max(0, min(999.999, (x / orig_width) * 1000)) + new_y = max(0, min(999.999, (y / orig_height) * 1000)) + else: + raise ValueError("Unsupported model type") + coords = [new_x, new_y] + + if len(coords) == 4: + x1, y1, x2, y2 = coords + if model_type == "qwen2_5": + new_x1 = max(0, min(round(x1 * scale_w), resized_width - 1)) + new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1)) + new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1)) + new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1)) + elif model_type in ["qwen2"]: + new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000)) + new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000)) + new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000)) + new_y2 = max(0, min(999.999, (y2 / orig_height) * 1000)) + else: + raise ValueError("Unsupported model type") + coords = [new_x1, new_y1, new_x2, new_y2] + + return f'<{tag_name}>[{", ".join(map(str, coords))}]' + + except (ValueError, TypeError): + return match.group(0) + + return point_pattern.sub(process_match, text) + + +def extract_grounding_points(text: str) -> List[List[float]]: + """Extract all // coordinates from ``text`` as list-of-list.""" + point_pattern = re.compile(r"<(point|box|bbox)>\s*\[([^\]]+)\]\s*") + + points: List[List[float]] = [] + for match in point_pattern.finditer(text): + coords_str = match.group(2) + raw_values = re.findall(r"-?\d+\.?\d*", coords_str) + converted: List[float] = [] + for value in raw_values: + number = float(value) + converted.append(int(number) if number.is_integer() else number) + if converted: + points.append(converted) + + return points + + +def reverse_grounding_points( + text: str, orig_height, orig_width, resized_height, resized_width, model_type +) -> str: + """Inverse of ``process_grounding_points`` - map resized coords back to original.""" + point_pattern = re.compile(r"<(point|box|bbox)>(.*?)") + + def reverse_match(match): + tag_name = match.group(1) + coords_str = match.group(2) + try: + coords = list(map(float, re.findall(r"-?\d+\.?\d*", coords_str))) + + scale_w = resized_width / orig_width + scale_h = resized_height / orig_height + + if len(coords) == 2: + x, y = coords + if model_type == "qwen2_5": + orig_x = max(0, min(orig_width - 1, round(x / scale_w))) + orig_y = max(0, min(orig_height - 1, round(y / scale_h))) + elif model_type in ["qwen2"]: + orig_x = max( + 0, min(orig_width - 1, round((x / 1000) * orig_width)) + ) + orig_y = max( + 0, min(orig_height - 1, round((y / 1000) * orig_height)) + ) + else: + raise ValueError("Unsupported model type") + coords = [orig_x, orig_y] + + if len(coords) == 4: + x1, y1, x2, y2 = coords + if model_type == "qwen2_5": + orig_x1 = max(0, min(orig_width - 1, round(x1 / scale_w))) + orig_y1 = max(0, min(orig_height - 1, round(y1 / scale_h))) + orig_x2 = max(0, min(orig_width - 1, round(x2 / scale_w))) + orig_y2 = max(0, min(orig_height - 1, round(y2 / scale_h))) + elif model_type in ["qwen2"]: + orig_x1 = max( + 0, min(orig_width - 1, round((x1 / 1000) * orig_width)) + ) + orig_y1 = max( + 0, min(orig_height - 1, round((y1 / 1000) * orig_height)) + ) + orig_x2 = max( + 0, min(orig_width - 1, round((x2 / 1000) * orig_width)) + ) + orig_y2 = max( + 0, min(orig_height - 1, round((y2 / 1000) * orig_height)) + ) + else: + raise ValueError("Unsupported model type") + coords = [orig_x1, orig_y1, orig_x2, orig_y2] + + return f'<{tag_name}>[{", ".join(map(str, map(int, coords)))}]' + + except (ValueError, TypeError): + return match.group(0) + + return point_pattern.sub(reverse_match, text) + + +def calculate_point_l1_distance(gt_text: str, pred_text: str) -> Optional[float]: + """Average L1 distance between 2D points extracted from ``gt_text`` / ``pred_text``. + + Returns None if either side has no points or counts differ. + """ + point_pattern = re.compile(r"\[(\d+),\s*(\d+)\]") + + gt_matches = point_pattern.findall(gt_text) + pred_matches = point_pattern.findall(pred_text) + + if not gt_matches or not pred_matches or len(gt_matches) != len(pred_matches): + return None + + total_l1_distance = 0.0 + for (gt_x, gt_y), (pred_x, pred_y) in zip(gt_matches, pred_matches): + try: + gt_x, gt_y = int(gt_x), int(gt_y) + pred_x, pred_y = int(pred_x), int(pred_y) + l1_dist = abs(gt_x - pred_x) + abs(gt_y - pred_y) + total_l1_distance += l1_dist + except ValueError: + continue + + return total_l1_distance / len(gt_matches) if gt_matches else None + + +__all__ = [ + "process_grounding_points", + "extract_grounding_points", + "reverse_grounding_points", + "calculate_point_l1_distance", +] diff --git a/wall_x/_vendor/x2robot_utils/text_templates.py b/wall_x/_vendor/x2robot_utils/text_templates.py new file mode 100644 index 0000000..ae1ca20 --- /dev/null +++ b/wall_x/_vendor/x2robot_utils/text_templates.py @@ -0,0 +1,235 @@ +"""Multimodal text preprocessing helpers. + +This file is generated by ``scripts/export_opensource.py``. It keeps only the +processor wrapper needed by harrix inference and uses a generic system prompt. +Internal robot-id, dataset, camera, and frequency maps are not bundled here. +""" + +from __future__ import annotations + +import numpy as np +import torch +from transformers import BatchFeature +from transformers.tokenization_utils_base import BatchEncoding + + +def pad_text_input_to_target_length( + text_inputs, target_length, pad_token_id=151643, padding_side="right" +): + """Pad or truncate tokenized text to ``target_length``.""" + batch_size, current_length = text_inputs.input_ids.shape + if current_length < target_length: + padding_size = target_length - current_length + padding = torch.full( + (batch_size, padding_size), + pad_token_id, + dtype=text_inputs.input_ids.dtype, + device=text_inputs.input_ids.device, + ) + attention_padding = torch.zeros( + (batch_size, padding_size), + dtype=text_inputs.attention_mask.dtype, + device=text_inputs.attention_mask.device, + ) + if padding_side == "right": + text_inputs["input_ids"] = torch.cat([text_inputs.input_ids, padding], dim=1) + text_inputs["attention_mask"] = torch.cat( + [text_inputs.attention_mask, attention_padding], dim=1 + ) + else: + text_inputs["input_ids"] = torch.cat([padding, text_inputs.input_ids], dim=1) + text_inputs["attention_mask"] = torch.cat( + [attention_padding, text_inputs.attention_mask], dim=1 + ) + elif current_length > target_length: + if padding_side == "right": + text_inputs["input_ids"] = text_inputs.input_ids[:, :target_length] + text_inputs["attention_mask"] = text_inputs.attention_mask[:, :target_length] + else: + text_inputs["input_ids"] = text_inputs.input_ids[:, -target_length:] + text_inputs["attention_mask"] = text_inputs.attention_mask[:, -target_length:] + return text_inputs + + +def _replace_media_placeholders(text, grid_thw, token, merge_length): + if grid_thw is None: + return text + index = 0 + for i in range(len(text)): + while token in text[i]: + if index >= len(grid_thw): + raise ValueError( + f"More {token} placeholders than media tensors in sample {i}" + ) + token_count = int(grid_thw[index].prod() // merge_length) + text[i] = text[i].replace(token, "<|placeholder|>" * token_count, 1) + index += 1 + text[i] = text[i].replace("<|placeholder|>", token) + return text + + +_PUBLIC_CAMERA_LABELS = { + "face_view": "front view", + "right_wrist_view": "right wrist view", + "left_wrist_view": "left wrist view", +} + + +def _camera_label(cam_name): + return _PUBLIC_CAMERA_LABELS.get(str(cam_name), str(cam_name).replace("_", " ")) + + +def preprocesser_call( + processor, + norm_state=None, + agent_pos_mask=None, + images=None, + prefix_text=None, + postfix_text=None, + videos=None, + padding=False, + padding_side="left", + truncation=None, + max_length=None, + return_tensors="pt", + pad_prefix_to_same_length=False, + pad_to_128_multiple=True, + state_augmentation_prob=0.0, + state_augmentation_ratio=0.0, + state_bins=256, + inference_mode=False, + **_, +): + """Build a ``BatchFeature`` for Wall-X VLA inference. + + This is the inference subset of the internal preprocessing helper: text, + image/video placeholder expansion, optional discretized proprioception + strings, padding, and labels=None for inference. + """ + if prefix_text is None: + raise ValueError("prefix_text is required") + if postfix_text is None: + postfix_text = [""] * len(prefix_text) + if not isinstance(prefix_text, list): + prefix_text = [prefix_text] + if not isinstance(postfix_text, list): + postfix_text = [postfix_text] + batch_size = len(prefix_text) + + if images is not None and len(images) > 0: + image_inputs = processor.image_processor(images=images, return_tensors=return_tensors) + image_grid_thw = image_inputs["image_grid_thw"] + else: + image_inputs = {} + image_grid_thw = None + + if videos is not None: + if hasattr(processor, "video_processor"): + videos_inputs = processor.video_processor(videos=videos, return_tensors=return_tensors) + else: + videos_inputs = processor.image_processor( + images=None, videos=videos, return_tensors=return_tensors + ) + video_grid_thw = videos_inputs["video_grid_thw"] + else: + videos_inputs = {} + video_grid_thw = None + + merge_length = processor.image_processor.merge_size**2 + prefix_text = _replace_media_placeholders( + list(prefix_text), image_grid_thw, "<|image_pad|>", merge_length + ) + prefix_text = _replace_media_placeholders( + prefix_text, video_grid_thw, "<|video_pad|>", merge_length + ) + + if norm_state is not None: + norm_state = norm_state.cpu().numpy() if isinstance(norm_state, torch.Tensor) else norm_state + agent_pos_mask = ( + agent_pos_mask[:, 0, :].cpu().numpy().astype(bool) + if isinstance(agent_pos_mask, torch.Tensor) + else agent_pos_mask[:, 0, :].astype(bool) + ) + discretized = np.digitize(norm_state, bins=np.linspace(-1, 1, state_bins + 1)[:-1]) - 1 + discretized = discretized[:, 0, :] + for i in range(batch_size): + if "<|propri|>" not in prefix_text[i]: + continue + state_str = " ".join(map(str, discretized[i, agent_pos_mask[i]])) + prefix_text[i] = prefix_text[i].replace("<|propri|>", state_str) + + if not pad_prefix_to_same_length: + text = [pre + post for pre, post in zip(prefix_text, postfix_text)] + text_inputs = processor.tokenizer( + text, + return_tensors=return_tensors, + padding=padding, + padding_side=padding_side, + truncation=truncation, + max_length=max_length, + ) + text_inputs["prefix_length"] = None + else: + prefix_inputs = processor.tokenizer( + prefix_text, + return_tensors=return_tensors, + padding=padding, + padding_side="left", + truncation=truncation, + max_length=max_length, + ) + postfix_inputs = processor.tokenizer( + postfix_text, + return_tensors=return_tensors, + padding=padding, + padding_side="right", + truncation=truncation, + max_length=max_length, + ) + text_inputs = BatchEncoding( + data={ + "input_ids": torch.cat([prefix_inputs.input_ids, postfix_inputs.input_ids], dim=1), + "attention_mask": torch.cat( + [prefix_inputs.attention_mask, postfix_inputs.attention_mask], dim=1 + ), + "prefix_length": prefix_inputs.input_ids.shape[1], + } + ) + + pad_token_id = processor.tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = processor.tokenizer.eos_token_id + if pad_to_128_multiple: + target_length = 128 * ((max(len(t) for t in text_inputs.input_ids) + 127) // 128) + text_inputs = pad_text_input_to_target_length( + text_inputs, target_length, pad_token_id=pad_token_id, padding_side=padding_side + ) + + text_inputs["labels"] = None if inference_mode else None + return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}) + + +def get_prologue_with_embodied_information(dataset_name, cam_mapping, robot_id, uid, config): + """Return a generic VLA system prompt without private robot maps.""" + role_start = "<|im_start|>" + role_end = "<|im_end|>" + prologue = ( + f"{role_start}system\n" + "You are an embodied vision-language-action model controlling a robot " + "with language instructions." + ) + if cam_mapping: + cameras = ", ".join(_camera_label(name) for name in cam_mapping.values()) + prologue += f"\nCamera Setup: {cameras}" + if not getattr(config, "use_relative_action", False): + prologue += "\nAction Space: Abs EEF" + else: + prologue += "\nAction Space: Rel EEF" + return f"{prologue}\n{role_end}\n" + + +__all__ = [ + "preprocesser_call", + "get_prologue_with_embodied_information", + "pad_text_input_to_target_length", +] diff --git a/wall_x/config/__init__.py b/wall_x/config/__init__.py new file mode 100644 index 0000000..329780a --- /dev/null +++ b/wall_x/config/__init__.py @@ -0,0 +1,62 @@ +"""Typed training configuration system.""" + +from .data_config import DataConfig, LeRobotDataConfig +from .hyperparams_config import ( + AdamWConfig, + ConstantSchedulerConfig, + CosineSchedulerConfig, + DMuonConfig, + OptimizerConfig, + SchedulerConfig, + StepSchedulerConfig, + TrainHyperParams, +) +from .infra_config import ( + CheckpointConfig, + DebugConfig, + DistributedConfig, + LoggingConfig, +) +from .loader import load_config, save_config +from .model_config import ModelConfig, QActModelConfig +from .registry import ( + register_data_config, + register_model_config, + register_optimizer_config, + register_scheduler_config, +) +from .task_config import TaskConfig +from .train_config import TrainConfig + +__all__ = [ + # Core + "TrainConfig", + "TaskConfig", + # Model configs + "ModelConfig", + "QActModelConfig", + # Data configs + "DataConfig", + "LeRobotDataConfig", + # Hyperparams + "TrainHyperParams", + "OptimizerConfig", + "AdamWConfig", + "DMuonConfig", + "SchedulerConfig", + "CosineSchedulerConfig", + "ConstantSchedulerConfig", + "StepSchedulerConfig", + # Infra configs + "DistributedConfig", + "LoggingConfig", + "CheckpointConfig", + "DebugConfig", + # Functions + "load_config", + "save_config", + "register_data_config", + "register_model_config", + "register_optimizer_config", + "register_scheduler_config", +] diff --git a/wall_x/config/data_config.py b/wall_x/config/data_config.py new file mode 100644 index 0000000..9ed7f94 --- /dev/null +++ b/wall_x/config/data_config.py @@ -0,0 +1,63 @@ +"""Public data config dataclasses. + +Only data backends shipped in the public package should define config classes +here. Internal backends register their dataclasses from their own packages via +``wall_x.config.registry.register_data_config``. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + +from .registry import register_data_config + + +@dataclass +class DataConfig: + """Base fields shared by data backends. + + ``normalizer_config`` may contain: + - ``min_key``: stats key for the minimum value. + - ``delta_key``: stats key for the value range. + - ``customized_action_statistic_dof``: explicit action-stats JSON path. + """ + + dataset_type: str = "lerobot" + resolution: Dict[str, int] = field( + default_factory=lambda: { + "face_view": 256, + "left_wrist_view": 256, + "right_wrist_view": 256, + } + ) + train_test_split: float = 0.95 + normalizer_config: Optional[Dict[str, Any]] = None + + +@register_data_config("lerobot") +@dataclass +class LeRobotDataConfig(DataConfig): + """LeRobot data config. + + ``lerobot_config`` is expected to contain fields such as ``repo_id`` and + ``root`` for a HuggingFace LeRobot dataset. ``norm_stats_path`` points to + explicit action normalizer stats; the core package does not bundle private + defaults. + """ + + dataset_type: str = "lerobot" + lerobot_config: Optional[Dict[str, Any]] = None + key_mappings: Optional[Dict[str, Any]] = None + norm_stats_path: Optional[str] = None + priority_order: Optional[Dict[str, float]] = None + camera_name_mapping: Optional[Dict[str, str]] = None + num_workers: int = 4 + action_tokenizer_path: Optional[str] = None + use_fast_tokenizer: bool = False + padding_side: str = "left" + noise_scheduler: Optional[Dict[str, Any]] = None + + +__all__ = [ + "DataConfig", + "LeRobotDataConfig", +] diff --git a/wall_x/config/hyperparams_config.py b/wall_x/config/hyperparams_config.py new file mode 100644 index 0000000..10a57b6 --- /dev/null +++ b/wall_x/config/hyperparams_config.py @@ -0,0 +1,129 @@ +"""Training hyperparameter config dataclasses.""" + +from dataclasses import dataclass, field +from typing import List, Optional, Tuple + +from .registry import register_optimizer_config, register_scheduler_config + + +@dataclass +class LRGroupConfig: + """Named optimizer LR group matched by parameter-name substrings.""" + + name: str + lr: float + include: List[str] = field(default_factory=list) + fail_on_empty: bool = True + + +@dataclass +class OptimizerConfig: + """Base optimizer config. ``optimizer_type`` selects a registered subclass.""" + + optimizer_type: str = "adamw" + learning_rate: float = 1e-4 + max_grad_norm: float = 1.0 + enable_grad_clip: bool = True + # Named parameter groups with independent learning rates. Unmatched + # trainable parameters remain in the base group using ``learning_rate``. + lr_groups: Optional[List[LRGroupConfig]] = None + # Optional action-expert LR split. + train_action_expert_only: bool = False + action_expert_learning_rate: Optional[float] = None + action_lr_keywords: Optional[List[str]] = None + + +@register_optimizer_config("adamw") +@dataclass +class AdamWConfig(OptimizerConfig): + """AdamW optimizer config.""" + + optimizer_type: str = "adamw" + betas: Tuple[float, float] = (0.9, 0.98) + weight_decay: float = 1e-8 + eps: float = 1e-8 + fused: bool = True + foreach: Optional[bool] = None + + +@register_optimizer_config("dmuon") +@dataclass +class DMuonConfig(OptimizerConfig): + """DMuon optimizer config.""" + + optimizer_type: str = "dmuon" + muon_lr: float = 0.02 + momentum: float = 0.95 + ns_steps: int = 5 + muon_weight_decay: float = 0.0 + adamw_lr: float = 1e-3 + adamw_betas: Tuple[float, float] = (0.9, 0.999) + adamw_weight_decay: float = 0.01 + adamw_eps: float = 1e-8 + ns_backend: str = "gram" + ns_coefficients: str = "default" + nesterov: bool = True + + +@dataclass +class SchedulerConfig: + """Base scheduler config. ``scheduler_type`` selects a registered subclass.""" + + scheduler_type: str = "cosine" + # Optional training-step cap. When > 0, the trainer sets + # loss_guard_should_stop=True once global_step >= num_training_steps, + # regardless of scheduler type. Cosine reads it for its own decay + # horizon; constant / step schedulers use it only for the stop signal. + num_training_steps: int = 0 + + +@register_scheduler_config("cosine") +@dataclass +class CosineSchedulerConfig(SchedulerConfig): + """Cosine annealing with warmup.""" + + scheduler_type: str = "cosine" + num_warmup_steps: int = 0 + num_training_steps: int = 0 + min_lr: Optional[float] = None # None means 0.1 * learning_rate at runtime. + + +@register_scheduler_config("constant") +@dataclass +class ConstantSchedulerConfig(SchedulerConfig): + """Constant learning rate with no decay.""" + + scheduler_type: str = "constant" + + +@register_scheduler_config("step") +@dataclass +class StepSchedulerConfig(SchedulerConfig): + """Step decay scheduler.""" + + scheduler_type: str = "step" + step_size: int = 10000 + gamma: float = 0.1 + + +@dataclass +class TrainHyperParams: + num_epoch: int = 1 + batch_size_per_gpu: int = 1 + gradient_accumulation_steps: int = 1 + seed: int = 42 + optimizer: OptimizerConfig = field(default_factory=AdamWConfig) + scheduler: SchedulerConfig = field(default_factory=CosineSchedulerConfig) + + +__all__ = [ + "AdamWConfig", + "ConstantSchedulerConfig", + "CosineSchedulerConfig", + "DMuonConfig", + "LRGroupConfig", + "OptimizerConfig", + "SchedulerConfig", + "StepSchedulerConfig", + "TrainHyperParams", +] diff --git a/wall_x/config/infra_config.py b/wall_x/config/infra_config.py new file mode 100644 index 0000000..64c9cec --- /dev/null +++ b/wall_x/config/infra_config.py @@ -0,0 +1,98 @@ +"""Infrastructure config: distributed runtime, logging, checkpoints, and debug.""" + +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class DistributedConfig: + # FSDP + use_fsdp: bool = False + fsdp_sharding_strategy: str = "full_shard" + fsdp_backward_prefetch: str = "backward_pre" + fsdp_cpu_offload: bool = False + fsdp_use_orig_params: bool = True + fsdp_limit_all_gathers: bool = True + fsdp_forward_prefetch: bool = False + fsdp_sync_module_states: bool = True + fsdp_save_policy: str = "full" + fsdp_hsdp_replicate_size: Optional[int] = None + # Mixed precision + use_mixed_precision: bool = True + bf16: bool = True + fsdp_reduce_dtype: str = "bf16" + use_amp: bool = False + use_gradient_checkpointing: bool = False + use_gradient_checkpointing_offload: bool = False + use_selective_recompute: bool = False + # DDP fallback + find_unused_parameters: bool = False + broadcast_buffers: bool = True + bucket_cap_mb: int = 25 + + +@dataclass +class LoggingConfig: + log_name: str = "exp" + log_project: str = "wallx" + log_entity: Optional[str] = None + use_wandb: bool = True + wandb_offline: bool = False + log_interval: int = 1 + save_interval: int = 1000 + val_interval: int = 4000 + epoch_save_interval: int = 1 + gc_interval_steps: int = 1000 + ignore_until_interval: int = 0 + # Rolling-window for smoothing per-step training metrics displayed on the + # console (and reused by tqdm). 1 = raw per-step (historical behavior). + # 10 = DZ-style 10-step rolling average - diffusion losses are dominated + # by timestep-sampling noise per step; the smoothing only changes display, + # not training. Independent of log_interval (which buffers for wandb). + loss_log_smooth_window: int = 1 + + +@dataclass +class CheckpointConfig: + """Checkpoint save and resume options. + + ``resume_from`` is shorthand for setting every component-specific resume + path to the same checkpoint. Component-specific fields take precedence. + """ + + save_path: str = "./ckpt" + validate_first: bool = False + # Shorthand path for all components. + resume_from: Optional[str] = None + # Component-specific resume paths. + resume_model: Optional[str] = None + resume_optimizer: Optional[str] = None + resume_scheduler: Optional[str] = None + resume_ema: Optional[str] = None + resume_rng: Optional[str] = None + resume_data: Optional[str] = None + resume_step: Optional[str] = None + + def get_resume_path(self, component: str) -> Optional[str]: + """Return the resume path for one checkpoint component.""" + specific = getattr(self, f"resume_{component}", None) + if specific is not None: + return specific + return self.resume_from + + +@dataclass +class DebugConfig: + profile: bool = False + profile_save_path: str = "./profile" + profile_wait_iters: int = 1 + profile_warmup_iters: int = 1 + profile_active_iters: int = 3 + show_time_details: bool = False + visualize_sample: bool = False + save_debug_batch_path: Optional[str] = None + nvtx: bool = False + # Formula MFU uses the local FLOPs estimate and measured step time. + enable_mfu: bool = False + # Optional FLOPs profiling runs an extra forward at step 0. + enable_mfu_profile: bool = False diff --git a/wall_x/config/loader.py b/wall_x/config/loader.py new file mode 100644 index 0000000..c22bcd5 --- /dev/null +++ b/wall_x/config/loader.py @@ -0,0 +1,285 @@ +"""Wall-X config loader.""" + +import dataclasses +import os +import shutil +from typing import Any, Type, TypeVar + +import yaml + +from .data_config import LeRobotDataConfig +from .hyperparams_config import ( + AdamWConfig, + LRGroupConfig, + OptimizerConfig, + SchedulerConfig, + TrainHyperParams, +) +from .infra_config import ( + CheckpointConfig, + DebugConfig, + DistributedConfig, + LoggingConfig, +) +from .model_config import ModelConfig, QActModelConfig +from .registry import ( + get_data_config, + get_model_config, + get_optimizer_config, + get_scheduler_config, + registered_data_configs, + registered_model_configs, + registered_optimizer_configs, + registered_scheduler_configs, +) +from .task_config import TaskConfig +from .train_config import TrainConfig + +T = TypeVar("T") + +_CONFIG_PLUGINS_LOADED = False + + +class _TrainConfigSafeLoader(yaml.SafeLoader): + pass + + +def _construct_python_tuple(loader: yaml.SafeLoader, node: yaml.Node) -> tuple: + return tuple(loader.construct_sequence(node)) + + +_TrainConfigSafeLoader.add_constructor( + "tag:yaml.org,2002:python/tuple", _construct_python_tuple +) + + +def _ensure_config_plugins_loaded() -> None: + """Load optional internal config plugins when they are present.""" + global _CONFIG_PLUGINS_LOADED + if _CONFIG_PLUGINS_LOADED: + return + _CONFIG_PLUGINS_LOADED = True + try: + from .internal_plugins import register_internal_config_plugins + except ImportError: + return + register_internal_config_plugins() + + +def load_config(config_path: str, cli_args: Any = None) -> TrainConfig: + """Load a ``TrainConfig`` from a YAML file.""" + _ensure_config_plugins_loaded() + + with open(config_path, "r") as f: + raw = yaml.load(f, Loader=_TrainConfigSafeLoader) + + if raw is None: + raise ValueError(f"Config file is empty: {config_path}") + + model_type = raw.get("model_type") + if model_type is None: + raise ValueError(f"Missing required field 'model_type' in {config_path}") + + config = TrainConfig( + model_type=model_type, + task=_build_dataclass(TaskConfig, raw.get("task", {})), + model=_build_model_config(model_type, raw.get("model", {})), + data=_build_data_config(raw.get("data", {})), + hyperparams=_build_hyperparams(raw.get("hyperparams", {})), + distributed=_build_dataclass(DistributedConfig, raw.get("distributed", {})), + logging=_build_dataclass(LoggingConfig, raw.get("logging", {})), + checkpoint=_build_dataclass(CheckpointConfig, raw.get("checkpoint", {})), + debug=_build_dataclass(DebugConfig, raw.get("debug", {})), + _raw_data=raw.get("data", {}), + _raw_yaml=raw, + dataset_path=raw.get("dataset_path"), + ) + + if cli_args is not None: + _apply_cli_overrides(config, cli_args) + + _validate(config) + + # Register the active data backend now that cfg is fully resolved + # (post-CLI-override, post-validate). This is the single source of + # truth for "which dataset backend is this run using". + from wall_x.data._registry import _set_data_backend + + _set_data_backend(config.data.dataset_type) + + return config + + +def save_config(config: TrainConfig, save_dir: str) -> str: + """Save ``TrainConfig`` to ``config.yml`` under ``save_dir``.""" + os.makedirs(save_dir, exist_ok=True) + config_path = os.path.join(save_dir, "config.yml") + + data = _sanitize_for_yaml(dataclasses.asdict(config)) + with open(config_path, "w") as f: + yaml.dump( + data, f, default_flow_style=False, allow_unicode=True, sort_keys=False + ) + + dataset_config_path = getattr(config.data, "dataset_config_path", None) + if dataset_config_path and os.path.exists(dataset_config_path): + dst = os.path.join(save_dir, "dataset_config.yml") + shutil.copy(dataset_config_path, dst) + + return config_path + + +def _sanitize_for_yaml(obj: Any) -> Any: + """Convert dataclass output to YAML-safe containers.""" + if isinstance(obj, dict): + return {k: _sanitize_for_yaml(v) for k, v in obj.items()} + elif isinstance(obj, (list, tuple)): + return [_sanitize_for_yaml(v) for v in obj] + return obj + + +def _build_dataclass(cls: Type[T], raw: dict) -> T: + """Build a dataclass from a dict, ignoring unknown keys.""" + if not raw: + return cls() + + field_names = {f.name for f in dataclasses.fields(cls)} + field_types = {f.name: f.type for f in dataclasses.fields(cls)} + filtered = {} + + for k, v in raw.items(): + if k not in field_names: + continue + ft = field_types[k] + if ( + isinstance(ft, type) + and dataclasses.is_dataclass(ft) + and isinstance(v, dict) + ): + filtered[k] = _build_dataclass(ft, v) + else: + filtered[k] = v + + return cls(**filtered) + + +def _build_model_config(model_type: str, raw: dict) -> ModelConfig: + """Build the registered model config for ``model_type``.""" + cls = get_model_config(model_type) + if cls is None: + raise ValueError( + f"Unknown model_type: '{model_type}'. " + f"Supported: {registered_model_configs()}" + ) + return _build_dataclass(cls, raw) + + +def _build_data_config(raw: dict): + """Build the registered data config for ``dataset_type``.""" + if not raw: + return LeRobotDataConfig() + + dataset_type = raw.get("dataset_type", "lerobot") + cls = get_data_config(dataset_type) + if cls is None: + raise ValueError( + f"Unknown dataset_type: '{dataset_type}'. " + f"Supported: {registered_data_configs()}" + ) + + return _build_dataclass(cls, raw) + + +def _build_optimizer_config(raw: dict) -> OptimizerConfig: + """Build the registered optimizer config for ``optimizer_type``.""" + if not raw: + return AdamWConfig() + + raw = dict(raw) + optimizer_type = raw.get("optimizer_type", "adamw") + cls = get_optimizer_config(optimizer_type) + if cls is None: + raise ValueError( + f"Unknown optimizer_type: '{optimizer_type}'. " + f"Supported: {registered_optimizer_configs()}" + ) + + if "betas" in raw and isinstance(raw["betas"], list): + raw["betas"] = tuple(raw["betas"]) + if "adamw_betas" in raw and isinstance(raw["adamw_betas"], list): + raw["adamw_betas"] = tuple(raw["adamw_betas"]) + if raw.get("lr_groups") is not None: + raw["lr_groups"] = [ + _build_dataclass(LRGroupConfig, group) for group in raw["lr_groups"] + ] + + return _build_dataclass(cls, raw) + + +def _build_scheduler_config(raw: dict) -> SchedulerConfig: + """Build the registered scheduler config for ``scheduler_type``.""" + if not raw: + cls = get_scheduler_config("cosine") + if cls is None: + raise ValueError("Scheduler config 'cosine' is not registered") + return cls() + + scheduler_type = raw.get("scheduler_type", "cosine") + cls = get_scheduler_config(scheduler_type) + if cls is None: + raise ValueError( + f"Unknown scheduler_type: '{scheduler_type}'. " + f"Supported: {registered_scheduler_configs()}" + ) + return _build_dataclass(cls, raw) + + +def _build_hyperparams(raw: dict) -> TrainHyperParams: + """Build ``TrainHyperParams`` with polymorphic optimizer/scheduler config.""" + if not raw: + return TrainHyperParams() + + raw = dict(raw) # shallow copy to avoid mutating caller's dict + optimizer_raw = raw.pop("optimizer", {}) + scheduler_raw = raw.pop("scheduler", {}) + + optimizer = _build_optimizer_config(optimizer_raw) + scheduler = _build_scheduler_config(scheduler_raw) + + hp = _build_dataclass(TrainHyperParams, raw) + hp.optimizer = optimizer + hp.scheduler = scheduler + return hp + + +def _apply_cli_overrides(config: TrainConfig, args: Any) -> None: + """Apply CLI overrides to the loaded config.""" + if getattr(args, "fsdp_sharding_strategy", None) is not None: + config.distributed.fsdp_sharding_strategy = args.fsdp_sharding_strategy + + if getattr(args, "debug", False): + config.logging.log_name = "debug" + config.logging.log_project = "debug" + config.checkpoint.save_path = "./ckpt/debug" + + if getattr(args, "visualize", False): + config.debug.visualize_sample = True + + if getattr(args, "wandb_offline", None) is not None: + config.logging.wandb_offline = args.wandb_offline in ("true", "True", "1") + + +def _validate(config: TrainConfig) -> None: + """Validate required fields and model-specific config.""" + if not config.model_type: + raise ValueError("model_type is required") + + if not config.task.dof_config: + raise ValueError("task.dof_config is required (cannot be empty)") + + if isinstance(config.model, QActModelConfig): + model = config.model + if not model.config_path: + raise ValueError("model.config_path is required for QAct models") + if not model.processor_path: + raise ValueError("model.processor_path is required for QAct models") diff --git a/wall_x/config/model_config.py b/wall_x/config/model_config.py new file mode 100644 index 0000000..4edaf44 --- /dev/null +++ b/wall_x/config/model_config.py @@ -0,0 +1,42 @@ +"""Base model config dataclasses. + +Optional model variants register their config dataclasses from their own +packages via ``wall_x.config.registry.register_model_config``. +""" + +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +from .registry import register_model_config + + +@dataclass +class ModelConfig: + """Base model architecture fields.""" + + use_ema: bool = False + attn_implementation: Optional[str] = None + attn_deterministic: Optional[bool] = None + ar_loss_weight: float = 1.0 + + +@register_model_config("qwen2_5") +@dataclass +class QActModelConfig(ModelConfig): + """QAct model config for the Qwen2.5 VLA path.""" + + config_path: str = "" + processor_path: str = "" + pretrained_path: Optional[str] = None + backbone: str = "qwen2_5" + action_tokenizer_type: Optional[str] = None + action_tokenizer_path: Optional[str] = None + action_tokenizer_checkpoint_path: Optional[str] = None + action_tokenizer_config_dir: Optional[str] = None + new_special_tokens: Optional[List[str]] = None + flow_loss_weight: float = 1.0 + enable_customized_robot_config: bool = False + customized_robot_config: Optional[Dict[str, Any]] = None + + +__all__ = ["ModelConfig", "QActModelConfig"] diff --git a/wall_x/config/registry.py b/wall_x/config/registry.py new file mode 100644 index 0000000..d34debf --- /dev/null +++ b/wall_x/config/registry.py @@ -0,0 +1,99 @@ +"""Runtime registries for typed config variants. + +Core Wall-X keeps only base config types in ``wall_x.config``. Optional +datasets, model families, and optimizer variants register their dataclasses +from their own packages, so trimmed distributions do not leave dead config +entries behind. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, TypeVar + +T = TypeVar("T") + +_DATA_CONFIG_REGISTRY: dict[str, type[Any]] = {} +_MODEL_CONFIG_REGISTRY: dict[str, type[Any]] = {} +_OPTIMIZER_CONFIG_REGISTRY: dict[str, type[Any]] = {} +_SCHEDULER_CONFIG_REGISTRY: dict[str, type[Any]] = {} + + +def _register( + registry: dict[str, type[Any]], names: tuple[str, ...], cls: type[T] +) -> type[T]: + for name in names: + if not name: + raise ValueError("Config registry name must be non-empty") + existing = registry.get(name) + if existing is not None and existing is not cls: + raise ValueError( + f"Config name {name!r} is already registered by " + f"{existing.__module__}.{existing.__name__}" + ) + registry[name] = cls + return cls + + +def register_data_config(*names: str) -> Callable[[type[T]], type[T]]: + return lambda cls: _register(_DATA_CONFIG_REGISTRY, names, cls) + + +def register_model_config(*names: str) -> Callable[[type[T]], type[T]]: + return lambda cls: _register(_MODEL_CONFIG_REGISTRY, names, cls) + + +def register_optimizer_config(*names: str) -> Callable[[type[T]], type[T]]: + return lambda cls: _register(_OPTIMIZER_CONFIG_REGISTRY, names, cls) + + +def register_scheduler_config(*names: str) -> Callable[[type[T]], type[T]]: + return lambda cls: _register(_SCHEDULER_CONFIG_REGISTRY, names, cls) + + +def get_data_config(name: str) -> type[Any] | None: + return _DATA_CONFIG_REGISTRY.get(name) + + +def get_model_config(name: str) -> type[Any] | None: + return _MODEL_CONFIG_REGISTRY.get(name) + + +def get_optimizer_config(name: str) -> type[Any] | None: + return _OPTIMIZER_CONFIG_REGISTRY.get(name) + + +def get_scheduler_config(name: str) -> type[Any] | None: + return _SCHEDULER_CONFIG_REGISTRY.get(name) + + +def registered_data_configs() -> list[str]: + return sorted(_DATA_CONFIG_REGISTRY) + + +def registered_model_configs() -> list[str]: + return sorted(_MODEL_CONFIG_REGISTRY) + + +def registered_optimizer_configs() -> list[str]: + return sorted(_OPTIMIZER_CONFIG_REGISTRY) + + +def registered_scheduler_configs() -> list[str]: + return sorted(_SCHEDULER_CONFIG_REGISTRY) + + +__all__ = [ + "get_data_config", + "get_model_config", + "get_optimizer_config", + "get_scheduler_config", + "register_data_config", + "register_model_config", + "register_optimizer_config", + "register_scheduler_config", + "registered_data_configs", + "registered_model_configs", + "registered_optimizer_configs", + "registered_scheduler_configs", +] diff --git a/wall_x/config/task_config.py b/wall_x/config/task_config.py new file mode 100644 index 0000000..1cd02b8 --- /dev/null +++ b/wall_x/config/task_config.py @@ -0,0 +1,17 @@ +"""Task config shared by model and data code.""" + +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + + +@dataclass +class TaskConfig: + """Robot task definition shared by model construction and data slicing.""" + + dof_config: Dict[str, int] = field(default_factory=dict) + agent_pos_config: Dict[str, int] = field(default_factory=dict) + ar_dof_config: Optional[Dict[str, int]] = None + action_horizon: int = 32 + action_horizon_flow: int = 32 + noise_scheduler: Optional[Dict[str, Any]] = None + use_state_string_representation: bool = False diff --git a/wall_x/config/train_config.py b/wall_x/config/train_config.py new file mode 100644 index 0000000..979622e --- /dev/null +++ b/wall_x/config/train_config.py @@ -0,0 +1,81 @@ +"""Top-level training config.""" + +import dataclasses +from dataclasses import dataclass, field +from typing import Any, Dict + +from .data_config import DataConfig +from .hyperparams_config import TrainHyperParams +from .infra_config import ( + CheckpointConfig, + DebugConfig, + DistributedConfig, + LoggingConfig, +) +from .model_config import ModelConfig +from .task_config import TaskConfig + + +@dataclass +class TrainConfig: + """Top-level Wall-X training config.""" + + model_type: str = "qwen2_5" + task: TaskConfig = field(default_factory=TaskConfig) + model: ModelConfig = field(default_factory=ModelConfig) + data: DataConfig = field(default_factory=DataConfig) + hyperparams: TrainHyperParams = field(default_factory=TrainHyperParams) + distributed: DistributedConfig = field(default_factory=DistributedConfig) + logging: LoggingConfig = field(default_factory=LoggingConfig) + checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig) + debug: DebugConfig = field(default_factory=DebugConfig) + # Raw YAML sections preserved verbatim for backend APIs that read fields + # not captured by the typed DataConfig dataclass. + _raw_data: Dict[str, Any] = field(default_factory=dict) + # Full raw YAML dict for backend-specific compatibility paths. + _raw_yaml: Dict[str, Any] = field(default_factory=dict) + # dataset_path lives at top level in YAML, consumed by data loaders directly + dataset_path: Any = None + + @property + def action_dim(self) -> int: + return sum(self.task.dof_config.values()) + + @property + def propri_dim(self) -> int: + return sum(self.task.agent_pos_config.values()) + + def build_data_loader_dict(self) -> Dict[str, Any]: + """Build the raw dict consumed by backend-specific compatibility paths. + + Merges ``_raw_data`` (verbatim YAML ``data:`` section) with task + fields (dof_config, action_horizon, etc.) and hyperparams + (batch_size). Backend compatibility APIs may read fields that the + typed DataConfig dataclass does not carry. + + This keeps legacy flat configs working while the main config surface + stays typed. + """ + # Start with the raw YAML data section, then add typed task defaults. + data_dict = dict(self._raw_data) + task_dict = dataclasses.asdict(self.task) + for key in ( + "dof_config", + "agent_pos_config", + "action_horizon", + "action_horizon_flow", + "ar_dof_config", + "use_state_string_representation", + ): + if key in task_dict and task_dict[key] is not None: + data_dict.setdefault(key, task_dict[key]) + data_dict.setdefault("batch_size_per_gpu", self.hyperparams.batch_size_per_gpu) + data_dict.setdefault("batch_size", self.hyperparams.batch_size_per_gpu) + result: Dict[str, Any] = { + "model_type": self.model_type, + "data": data_dict, + **task_dict, + } + if self.dataset_path is not None: + result["dataset_path"] = self.dataset_path + return result diff --git a/wall_x/data/__init__.py b/wall_x/data/__init__.py index e69de29..b4661af 100644 --- a/wall_x/data/__init__.py +++ b/wall_x/data/__init__.py @@ -0,0 +1,31 @@ +"""Public data facade for backend selection and dataset construction.""" + +from wall_x.data import backends # noqa: F401 (side-effect registration) +from wall_x.data._bundle import DataBundle +from wall_x.data._protocol import BuildContext, DatasetBackend +from wall_x.data._registry import ( + MissingOperationError, + available_backends, + backend_for, + build_data, + data_backend, + has_data_backend, + register, + register_module, + resolve_dataset_type, +) + +__all__ = [ + "DataBundle", + "BuildContext", + "DatasetBackend", + "MissingOperationError", + "data_backend", + "available_backends", + "backend_for", + "build_data", + "has_data_backend", + "register", + "register_module", + "resolve_dataset_type", +] diff --git a/wall_x/data/_bundle.py b/wall_x/data/_bundle.py new file mode 100644 index 0000000..9a250aa --- /dev/null +++ b/wall_x/data/_bundle.py @@ -0,0 +1,39 @@ +"""Dataset bundle returned by data backends.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Iterable, Optional + + +def _noop_set_epoch(_epoch: int) -> None: + """Default ``set_epoch`` for backends with epoch-agnostic shuffling.""" + return None + + +@dataclass +class DataBundle: + """Container returned by every backend ``build()``. + + Attributes: + dataset: backend-private; trainer code should treat it as opaque. + train_loader: anything iterable that yields training batches. + val_loader: val iterable or None if the backend does not split val. + train_iters: one-epoch step count for the train loader. Backends + with dynamic resizing should set this to a stable snapshot + and expose the live value separately on ``dataset``. + val_iters: one-epoch step count for the val loader; 0 if no val. + set_epoch: per-epoch seed hook. Called before each epoch by + the trainer. Backends that don't need per-epoch reshuffle + should use ``_noop_set_epoch``. + """ + + dataset: Any + train_loader: Iterable + val_loader: Optional[Iterable] = None + train_iters: int = 0 + val_iters: int = 0 + set_epoch: Callable[[int], None] = field(default=_noop_set_epoch) + + +__all__ = ["DataBundle"] diff --git a/wall_x/data/_protocol.py b/wall_x/data/_protocol.py new file mode 100644 index 0000000..50694eb --- /dev/null +++ b/wall_x/data/_protocol.py @@ -0,0 +1,37 @@ +"""Shared data backend protocol definitions.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional, Protocol, runtime_checkable + +from wall_x.data._bundle import DataBundle + + +@dataclass +class BuildContext: + """Shared runtime state passed to every backend ``build()``. + + Fields are Optional so backends not using a particular piece can + simply leave it ``None``. The trainer populates whatever it has. + """ + + rank: int = 0 + world_size: int = 1 + tokenizer: Optional[Any] = None + processor: Optional[Any] = None + tokenizer_mixin: Optional[Any] = None + normalizer_action: Optional[Any] = None + normalizer_propri: Optional[Any] = None + model_config: Optional[Any] = None + resume_state: Optional[dict] = None + + +@runtime_checkable +class DatasetBackend(Protocol): + """Callable every backend registers under its ``dataset_type`` name.""" + + def __call__(self, cfg: Any, ctx: BuildContext) -> DataBundle: ... + + +__all__ = ["BuildContext", "DatasetBackend"] diff --git a/wall_x/data/_registry.py b/wall_x/data/_registry.py new file mode 100644 index 0000000..23bfe01 --- /dev/null +++ b/wall_x/data/_registry.py @@ -0,0 +1,263 @@ +"""Data backend registry + multi-verb dispatch. + +Each backend registers an entire **module** (not just a single function). +Consumers obtain a backend via :func:`backend_for` and then call any +operation the backend supports (``backend.build`` or optional helper +operations published by that backend). Operations a backend does *not* +implement raise +:class:`MissingOperationError` with a list of backends that do. + +Three failure shapes, each pointing at the right next step: + +- **Unknown name** (typo / never registered): :class:`KeyError` + ``Unknown backend 'foo_bar'. Known: [...]`` +- **Known but failed to import** (optional dependency missing): + :class:`RuntimeError` + ``Backend 'example' failed to import: . Install its + dependency or switch to one of [...]`` +- **Backend exists but operation missing**: :class:`MissingOperationError` + ``Backend 'example' does not implement 'make_processor'. + Supported by: ['other_backend']`` +""" + +from __future__ import annotations + +import logging +from types import ModuleType +from typing import Any, Dict + +from wall_x.data._bundle import DataBundle +from wall_x.data._protocol import BuildContext + +logger = logging.getLogger(__name__) + +_BACKENDS: Dict[str, ModuleType] = {} +_import_errors: Dict[str, BaseException] = {} + +# The single, process-wide active backend. Set exactly once at config load +# time by :func:`_set_data_backend`; all consumer code reads it via +# :func:`data_backend`. Keeping this as module-level state (rather than +# threading cfg through every callsite) is the whole point of the design. +_DATA_BACKEND: str | None = None + + +# --- Registration -------------------------------------------------------- + + +def register_module(dataset_type: str, module: ModuleType) -> None: + """Register a backend module under ``dataset_type``. + + The module must expose a ``build(cfg, ctx) -> DataBundle`` callable. + Other optional operations are + discovered lazily via ``getattr`` when ``backend_for(name).`` is + accessed; backends only declare what they support. + """ + if not hasattr(module, "build"): + raise TypeError( + f"Backend module for {dataset_type!r} must expose a " + f"``build(cfg, ctx) -> DataBundle`` callable; got {module!r}." + ) + if dataset_type in _BACKENDS: + logger.warning("overwriting existing backend registration for %r", dataset_type) + _BACKENDS[dataset_type] = module + + +def register(dataset_type: str, build_callable) -> None: + """Legacy single-callable registration. Wraps ``build_callable`` in a + minimal module so the new dispatch path still works. + + New backends should prefer :func:`register_module` so they can publish + multiple operations. + """ + shim = ModuleType(f"_legacy_backend_shim_{dataset_type}") + shim.build = build_callable # type: ignore[attr-defined] + register_module(dataset_type, shim) + + +def record_import_error(dataset_type: str, error: BaseException) -> None: + """Stash the exception that prevented a backend from registering.""" + _import_errors[dataset_type] = error + + +def available_backends() -> list[str]: + """Return the names of currently-registered backends.""" + return sorted(_BACKENDS) + + +# --- Lookup -------------------------------------------------------------- + + +class MissingOperationError(NotImplementedError): + """Raised when a backend module does not implement a requested op.""" + + def __init__(self, backend_name: str, op_name: str) -> None: + impls = [n for n, mod in _BACKENDS.items() if hasattr(mod, op_name)] + if impls: + hint = f"Supported by: {impls}." + else: + hint = ( + f"No registered backend implements {op_name!r} - check the " + f"spelling or add it to the backend module." + ) + super().__init__( + f"Backend {backend_name!r} does not implement {op_name!r}. {hint}" + ) + self.backend_name = backend_name + self.op_name = op_name + + +class _BackendProxy: + """Thin wrapper that forwards ``proxy.(...)`` to the backend module. + + Wrapping (instead of returning the module directly) lets us emit + ``MissingOperationError`` with a useful "supported by" hint instead + of plain ``AttributeError``. + """ + + __slots__ = ("_name", "_module") + + def __init__(self, name: str, module: ModuleType) -> None: + self._name = name + self._module = module + + def __repr__(self) -> str: + return f"" + + def __getattr__(self, op: str): + attr = getattr(self._module, op, None) + if attr is None: + raise MissingOperationError(self._name, op) + return attr + + def supports(self, op: str) -> bool: + """Cheap predicate: does this backend implement ``op``?""" + return hasattr(self._module, op) + + +def backend_for(name: str) -> _BackendProxy: + """Look up a backend by ``dataset_type`` name. + + Raises :class:`RuntimeError` (with chained ImportError) if the + backend is known but failed to import; :class:`KeyError` if the + name was never registered. + """ + if name in _BACKENDS: + return _BackendProxy(name, _BACKENDS[name]) + if name in _import_errors: + err = _import_errors[name] + raise RuntimeError( + f"Backend {name!r} failed to import: {err}. " + f"Install its dependency or switch to one of " + f"{available_backends()}." + ) from err + raise KeyError(f"Unknown backend {name!r}. Known: {available_backends()}.") + + +# --- Convenience helpers ------------------------------------------------- + + +def resolve_dataset_type(cfg_or_yaml: Any, default: str = "lerobot") -> str: + """Pull ``dataset_type`` out of either a typed TrainConfig or a raw yaml dict. + + Lookup order: + 1. ``cfg.data.dataset_type`` (typed TrainConfig) + 2. ``yaml_dict["dataset_type"]`` (legacy flat yaml) + 3. ``default`` + """ + data = getattr(cfg_or_yaml, "data", None) + if data is not None: + v = getattr(data, "dataset_type", None) + if v: + return v + if isinstance(cfg_or_yaml, dict): + v = cfg_or_yaml.get("dataset_type") + if v: + return v + return default + + +def build_data(cfg: Any, ctx: BuildContext) -> DataBundle: + """Dispatch ``backend.build(cfg, ctx)`` to the backend named by cfg. + + ``cfg`` is expected to be a typed :class:`TrainConfig` - that's what + ``wall_x.config.load_config`` returns and what every trainer entry + point passes in. Raw yaml dicts are not supported here on purpose: + the typed schema is the contract that gives backends their + ``cfg.data.dataset_type`` access. If you have a raw dict, use + :func:`resolve_dataset_type` + :func:`backend_for` directly. + """ + if not hasattr(cfg, "data") or not hasattr(cfg.data, "dataset_type"): + raise TypeError( + f"build_data() expects a typed TrainConfig, got {type(cfg).__name__}. " + f"Load via wall_x.config.load_config() instead of passing a raw dict." + ) + return backend_for(cfg.data.dataset_type).build(cfg, ctx) + + +# --- Active backend (process-global) ------------------------------------ + + +def _set_data_backend(name: str) -> None: + """Internal - only config loaders should call this. + + Strict semantics: first call wins; same-value re-set is a no-op; + different-value re-set raises. The error is loud on purpose - it means + two config loaders disagreed about which backend to use, which is + almost always a bug (e.g. business code calling this directly, or two + yamls being loaded in the same process). + """ + global _DATA_BACKEND + if _DATA_BACKEND is None: + _DATA_BACKEND = name + return + if _DATA_BACKEND == name: + return # idempotent on same value + raise RuntimeError( + f"Active backend already set to {_DATA_BACKEND!r}, refusing to " + f"overwrite with {name!r}. This usually means a config loader was " + f"called twice with different dataset_type, or business code called " + f"_set_data_backend directly. Use _reset_data_backend() in tests " + f"if you need to switch." + ) + + +def _reset_data_backend() -> None: + """Internal - clear the active backend. Tests only.""" + global _DATA_BACKEND + _DATA_BACKEND = None + + +def has_data_backend() -> bool: + """Whether a backend has been registered for this process.""" + return _DATA_BACKEND is not None + + +def data_backend() -> _BackendProxy: + """Return the active backend proxy. Raises if no config has been loaded. + + Direct construction of ``TrainConfig(...)`` does **not** register a + backend - that's intentional. Business code must go through + :func:`wall_x.config.load_config`, which registers the backend before + returning. + """ + if _DATA_BACKEND is None: + raise RuntimeError( + "No active data backend. Load a TrainConfig via " + "wall_x.config.load_config() before accessing backend verbs. " + "Direct TrainConfig() construction does not register a backend." + ) + return backend_for(_DATA_BACKEND) + + +__all__ = [ + "register", + "register_module", + "record_import_error", + "available_backends", + "backend_for", + "build_data", + "resolve_dataset_type", + "MissingOperationError", + "data_backend", + "has_data_backend", +] diff --git a/wall_x/data/backends/__init__.py b/wall_x/data/backends/__init__.py new file mode 100644 index 0000000..5b9a128 --- /dev/null +++ b/wall_x/data/backends/__init__.py @@ -0,0 +1,16 @@ +"""Data backend registration. + +Importing this package registers the default shipped backends. Additional +backends are loaded by plugin modules so their names and dependencies do not +have to appear in the default import path. +""" + +import importlib + +from wall_x.data._registry import record_import_error + +for _name in ("lerobot",): + try: + importlib.import_module(f"wall_x.data.backends.{_name}") + except ImportError as _e: + record_import_error(_name, _e) diff --git a/wall_x/data/backends/lerobot/__init__.py b/wall_x/data/backends/lerobot/__init__.py new file mode 100644 index 0000000..d85849b --- /dev/null +++ b/wall_x/data/backends/lerobot/__init__.py @@ -0,0 +1,26 @@ +"""LeRobot backend registration.""" + +from __future__ import annotations + +import sys + +try: + from wall_x.data._registry import register_module + from wall_x.data.backends.lerobot.build import ( + build, + load_trainer_data_config, + load_trainer_data_config_from_yaml_dict, + ) + + register_module("lerobot", sys.modules[__name__]) +except ImportError as _e: + from wall_x.data._registry import record_import_error + + record_import_error("lerobot", _e) + + +__all__ = [ + "build", + "load_trainer_data_config", + "load_trainer_data_config_from_yaml_dict", +] diff --git a/wall_x/data/backends/lerobot/build.py b/wall_x/data/backends/lerobot/build.py new file mode 100644 index 0000000..36edd4f --- /dev/null +++ b/wall_x/data/backends/lerobot/build.py @@ -0,0 +1,337 @@ +"""LeRobot data loading bridge for typed training configs.""" + +from __future__ import annotations + +import logging +import multiprocessing as mp +from typing import Any, Dict, Tuple + +import torch +import torch.distributed as dist + +from wall_x.data.backends.lerobot.config import LerobotConfig +from wall_x.data.backends.lerobot.utils import load_norm_stats +from wall_x.model.core.action.normalizer import ( + create_normalizers_from_lerobot_norm_stats, +) + +logger = logging.getLogger(__name__) + + +def load_lerobot_normalizers(cfg: Any): + """Create model normalizers from the LeRobot norm stats configured for a run.""" + data = cfg.data + + norm_stats_path = getattr(data, "norm_stats_path", None) + if not norm_stats_path: + return None + + key_mappings = getattr(data, "key_mappings", None) + if not key_mappings: + raise ValueError( + "LeRobot normalizer from norm_stats_path requires data.key_mappings" + ) + + lerobot_config = getattr(data, "lerobot_config", None) + if not isinstance(lerobot_config, dict) or not lerobot_config.get("repo_id"): + raise ValueError( + "LeRobot normalizer from norm_stats_path requires " + "data.lerobot_config.repo_id" + ) + + dataset_name = str(lerobot_config["repo_id"]) + norm_stats = load_norm_stats( + norm_stats_path, + key_mappings, + dof_config=dict(cfg.task.dof_config or {}), + agent_pos_config=dict(cfg.task.agent_pos_config or {}), + ) + normalizer_action, normalizer_propri = create_normalizers_from_lerobot_norm_stats( + norm_stats, + dataset_name, + cfg.action_dim, + cfg.propri_dim, + ) + return normalizer_action, normalizer_propri, norm_stats_path, dataset_name + + +class _LerobotDatasetWrapper: + """Trainer-facing wrapper aligning PreprocessedDataset with v1 API. + + PreprocessedDataset internally switches ``self._dataset`` between + its train/val splits via ``_train()`` / ``_eval()``. Its + ``get_train_dataloader`` / ``get_val_dataloader`` return + ``(dataloader, sampler)`` tuples and no-argument calls are supported + (they read rank/world_size/seed from the inner object itself). + + This wrapper: + - Caches the rebuilt train dataloader / sampler so + ``set_epoch(epoch)`` can reset shuffling per-epoch. + - Owns the val dataloader so the trainer's ``val_loop`` can do + ``self.dataset.get_val_dataloader()`` and iterate directly (matching + what the v1/v2 wrappers return). + """ + + def __init__( + self, + inner, + train_dataloader: torch.utils.data.DataLoader, + train_sampler, + train_num: int, + val_dataloader: torch.utils.data.DataLoader = None, + val_num: int = 0, + ): + self._inner = inner + self._train_dataloader = train_dataloader + self._train_sampler = train_sampler + self._train_num = train_num + self._val_dataloader = val_dataloader + self.global_train_iters = mp.Value("i", train_num) + self.global_val_iters = mp.Value("i", val_num) + + def __len__(self) -> int: + return self._train_num + + def _activate_train_split(self) -> None: + if hasattr(self._inner, "_train"): + self._inner._train() + + def get_train_dataloader(self): + self._activate_train_split() + return self._train_dataloader + + def get_val_dataloader(self): + # PreprocessedDataset shares one ``_dataset`` pointer between its + # train and val splits (flipped by ``_train()`` / ``_eval()``). + # The val DataLoader's DistributedSampler caches total_size sized + # to the val split but ``__iter__`` reads ``len(self.dataset)`` + # live - if a preceding train_loop left the pointer at train, that + # live len is ~20x total_size and DistributedSampler asserts. + # Rebuild each time so ``_eval()`` runs and a fresh sampler is + # snapped to the current (val) split length. Mirrors the train-side + # rebuild-on-every-epoch pattern. + if self._val_dataloader is None: + return None + self._val_dataloader, _ = self._inner.get_val_dataloader() + return self._val_dataloader + + def set_epoch(self, epoch: int) -> None: + """Seed the per-epoch shuffle in the train DistributedSampler.""" + self._activate_train_split() + if self._train_sampler is not None and hasattr( + self._train_sampler, "set_epoch" + ): + self._train_sampler.set_epoch(epoch) + + +def load_trainer_data_config(cfg: Any) -> LerobotConfig: + """Build the inference/trainer data config from a typed TrainConfig.""" + raw_yaml = dict(getattr(cfg, "_raw_yaml", {}) or {}) + raw_data = dict(getattr(cfg, "_raw_data", {}) or {}) + data = getattr(cfg, "data", None) + + data_section = dict(raw_yaml.get("data", {}) or {}) + data_section.update(raw_data) + + if data is not None: + for key in ( + "resolution", + "train_test_split", + "priority_order", + "camera_name_mapping", + ): + value = getattr(data, key, None) + if value is not None: + data_section.setdefault(key, value) + + raw_yaml["data"] = data_section + raw_yaml.setdefault("model_type", getattr(cfg, "model_type", "qwen2_5")) + return load_trainer_data_config_from_yaml_dict(raw_yaml) + + +def load_trainer_data_config_from_yaml_dict(yaml_dict: Dict[str, Any]) -> LerobotConfig: + """Build the LeRobot runtime config from a raw training YAML dict.""" + return LerobotConfig.from_yaml_dict(yaml_dict) + + +def _build_flat_config(cfg: Any) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """Map typed TrainConfig -> (flat_config, lerobot_config) for legacy entry. + + ``load_lerobot_data`` expects a 2509-style flat dict plus a separate + ``lerobot_config`` carrying ``repo_id`` / ``root``. This function is + the one place that translation lives; keep it surgical so future + field additions on ``LeRobotDataConfig`` do not require touching the + legacy loader. + """ + model = cfg.model + data = cfg.data + hp = cfg.hyperparams + raw = getattr(cfg, "_raw_yaml", {}) or {} + raw_data = dict(getattr(cfg, "_raw_data", {}) or {}) + + lerobot_cfg = dict(data.lerobot_config or {}) + if "repo_id" not in lerobot_cfg: + raise ValueError( + "lerobot requires data.lerobot_config.repo_id to be set " + "(HuggingFace LeRobot dataset id)." + ) + + data_section: Dict[str, Any] = { + "key_mappings": data.key_mappings, + "action_horizon": cfg.task.action_horizon, + "train_test_split": data.train_test_split, + "seed": hp.seed, + "resolution": data.resolution, + } + if raw_data.get("max_length") is not None: + data_section["max_length"] = raw_data["max_length"] + if data.priority_order is not None: + data_section["priority_order"] = data.priority_order + if data.camera_name_mapping is not None: + data_section["camera_name_mapping"] = data.camera_name_mapping + data_section.setdefault( + "use_state_string_representation", + cfg.task.use_state_string_representation, + ) + data_section.setdefault( + "state_bins", + raw_data.get("state_bins", raw.get("state_bins", 256)), + ) + + # Dof/agent_pos totals for the collator's zero-pad step. When resuming + # from a checkpoint trained on a larger action space, task.dof_config + # should include an ``action_padding`` key that absorbs the diff; the + # collator right-pads action/agent_pos tensors to these totals with + # dof_mask/agent_pos_mask zeroed on padded dims so loss doesn't flow + # through them. + dof_total = int(sum((cfg.task.dof_config or {}).values())) + agent_pos_total = int(sum((cfg.task.agent_pos_config or {}).values())) + + flat: Dict[str, Any] = { + "model_type": cfg.model_type, + "processor_path": getattr(model, "processor_path", "") or "", + "norm_stats_path": data.norm_stats_path or raw.get("norm_stats_path"), + "batch_size_per_gpu": hp.batch_size_per_gpu, + "eval_batch_size_per_gpu": raw.get( + "eval_batch_size_per_gpu", hp.batch_size_per_gpu + ), + "num_workers": data.num_workers, + "padding_side": data.padding_side, + "use_fast_tokenizer": data.use_fast_tokenizer, + "action_tokenizer_path": data.action_tokenizer_path, + "noise_scheduler": data.noise_scheduler or {}, + "dof_total_dim": dof_total, + "agent_pos_total_dim": agent_pos_total, + "dof_config": dict(cfg.task.dof_config or {}), + "agent_pos_config": dict(cfg.task.agent_pos_config or {}), + "use_state_string_representation": cfg.task.use_state_string_representation, + "state_bins": int( + raw_data.get("state_bins") + or data_section.get("state_bins") + or raw.get("state_bins") + or 256 + ), + "data": data_section, + } + return flat, lerobot_cfg + + +def load_lerobot_v2( + cfg: Any, +) -> Tuple[_LerobotDatasetWrapper, torch.utils.data.DataLoader, int]: + """Build lerobot (wrapper, dataloader, train_num) from TrainConfig. + + The third return value ``train_num`` is a snapshot of + ``len(train_dataloader)`` at construction time. It matches + ``wrapper.global_train_iters.value`` initially but does not track + subsequent rebuilds inside ``set_epoch`` - callers doing dynamic + resampling should read from the mp.Value, not from this snapshot. + """ + from wall_x.data.backends.lerobot.loader import load_lerobot_data + + flat_cfg, lerobot_cfg = _build_flat_config(cfg) + + if dist.is_initialized(): + rank = dist.get_rank() + world_size = dist.get_world_size() + else: + rank = 0 + world_size = 1 + + seed = cfg.hyperparams.seed + inner, _ = load_lerobot_data( + flat_cfg, + lerobot_cfg, + rank=rank, + world_size=world_size, + seed=seed, + ) + + # PreprocessedDataset.get_*_dataloader returns (dataloader, sampler). + # Build val first, train second, so the inner ``_dataset`` pointer is + # left at the train split when we finish - workers fork from that + # state on first iteration. + val_dataloader, _ = inner.get_val_dataloader() + val_num = len(val_dataloader) if val_dataloader is not None else 0 + + train_dataloader, train_sampler = inner.get_train_dataloader() + train_num = len(train_dataloader) + + if rank == 0: + logger.info( + "\n%s\nLeRobot Data Loading Configuration:\n" + " RANK: %d\n WORLD SIZE: %d\n" + " BATCH SIZE PER DEVICE: %d\n GLOBAL BATCH SIZE: %d\n" + " TRAIN BATCHES: %d\n VAL BATCHES: %d\n" + " NUM WORKERS: %d\n REPO ID: %s\n%s", + "=" * 50, + rank, + world_size, + flat_cfg["batch_size_per_gpu"], + flat_cfg["batch_size_per_gpu"] * world_size, + train_num, + val_num, + flat_cfg["num_workers"], + lerobot_cfg.get("repo_id"), + "=" * 50, + ) + + wrapper = _LerobotDatasetWrapper( + inner, + train_dataloader, + train_sampler, + train_num, + val_dataloader=val_dataloader, + val_num=val_num, + ) + return wrapper, train_dataloader, train_num + + +def build(cfg, ctx): + """Backend Protocol entry - returns a ``DataBundle``. + + Wraps ``load_lerobot_v2`` (which returns the trainer-facing triple) + into the unified ``DataBundle`` shape every backend exposes. + """ + from wall_x.data._bundle import DataBundle + + wrapper, train_dataloader, train_num = load_lerobot_v2(cfg) + + # PreprocessedDataset shares one ``self._dataset`` pointer between + # train and val splits (flipped by ``_train()`` / ``_eval()``). + # ``wrapper.get_val_dataloader()`` flips the pointer to val. Flip back once + # here so the initial train loop starts from the right split even if callers + # inspect the raw ``train_dataloader`` before invoking ``set_epoch``. + val_loader = wrapper.get_val_dataloader() + inner = wrapper._inner + if hasattr(inner, "_train"): + inner._train() + + return DataBundle( + dataset=wrapper, + train_loader=train_dataloader, + val_loader=val_loader, + train_iters=train_num, + val_iters=wrapper.global_val_iters.value, + set_epoch=wrapper.set_epoch, + ) diff --git a/wall_x/data/backends/lerobot/config.py b/wall_x/data/backends/lerobot/config.py new file mode 100644 index 0000000..0b79d07 --- /dev/null +++ b/wall_x/data/backends/lerobot/config.py @@ -0,0 +1,134 @@ +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + +from qwen_vl_utils.vision_process import IMAGE_FACTOR, MAX_PIXELS, MIN_PIXELS + + +@dataclass +class LerobotConfig: + """Configuration for the LeRobot preprocessing pipeline. + + Dataset-specific camera display names are optional config inputs. Other + dataset behavior is derived from the current LeRobot sample. + """ + + # Image resolution settings for different views + resolution: Dict[str, int] = field( + default_factory=lambda: { + "face_view": -1, + "left_wrist_view": 128, + "right_wrist_view": 128, + } + ) + + # Dataset splitting + train_test_split: float = 0.9 + seed: int = 42 + + # Instruction handling + priority_order: Optional[Dict[str, float]] = None + camera_name_mapping: Optional[Dict[str, str]] = None + + # Vision model parameters + model_type: str = "qwen2_5" + max_pixels: int = MAX_PIXELS + min_pixels: int = MIN_PIXELS + image_factor: int = IMAGE_FACTOR + + generate_subtask_ratio: float = 0.0 + + def __post_init__(self): + """Post-initialization validation and setup.""" + # Validate train/test split + if not 0 < self.train_test_split < 1: + raise ValueError( + f"train_test_split must be between 0 and 1, got {self.train_test_split}" + ) + + def as_dict(self) -> Dict: + """Convert configuration to dictionary format. + + Returns: + Dict: Configuration as dictionary + """ + return self.__dict__ + + def update(self, **kwargs) -> "LerobotConfig": + """Update configuration parameters. + + Args: + **kwargs: Key-value pairs to update + + Returns: + LerobotConfig: Updated configuration instance + """ + for key, value in kwargs.items(): + if hasattr(self, key): + setattr(self, key, value) + else: + raise ValueError(f"Unknown configuration parameter: {key}") + return self + + def __getitem__(self, key: str): + return getattr(self, key) + + @classmethod + def from_yaml_dict(cls, yaml_dict: Dict[str, Any]) -> "LerobotConfig": + """ + Build a LerobotConfig instance from a YAML dictionary. + + Supports two styles: + + 1) Top-level fields: + train_test_split: 0.8 + model_type: qwen2_5 + + 2) Nested under `data:` (higher priority): + data: + train_test_split: 0.8 + model_type: qwen2_5 + + Keys inside `data:` override top-level keys. + """ + + data_config = yaml_dict.get("data", {}) + + def get(key: str, default: Any = None): + """ + Helper function: + Read from `data` first, then fallback to the top-level YAML. + """ + return data_config.get(key, yaml_dict.get(key, default)) + + # Construct only fields that actually exist in LerobotConfig + params: Dict[str, Any] = { + # Action prediction settings + # Image resolution per camera view + "resolution": get( + "resolution", + { + "face_view": -1, + "left_wrist_view": 128, + "right_wrist_view": 128, + }, + ), + # Dataset train/test split configuration + "train_test_split": get("train_test_split", 0.9), + "seed": get("seed", 42), + # Instruction priority ordering (optional) + "priority_order": get("priority_order", None), + "camera_name_mapping": get("camera_name_mapping", None), + # Vision model parameters + "model_type": get("model_type", "qwen2_5"), + "max_pixels": get("max_pixels", MAX_PIXELS), + "min_pixels": get("min_pixels", MIN_PIXELS), + "image_factor": get("image_factor", IMAGE_FACTOR), + # Subtask generation ratio + "generate_subtask_ratio": get("generate_subtask_ratio", 0.0), + } + + # Keep only valid dataclass fields (ignore unknown YAML keys) + valid_fields = {f.name for f in cls.__dataclass_fields__.values()} + filtered_params = {k: v for k, v in params.items() if k in valid_fields} + + return cls(**filtered_params) diff --git a/wall_x/data/backends/lerobot/loader.py b/wall_x/data/backends/lerobot/loader.py new file mode 100644 index 0000000..c000ec8 --- /dev/null +++ b/wall_x/data/backends/lerobot/loader.py @@ -0,0 +1,895 @@ +""" +LeRobot Dataset Loader - Distributed Version +""" + +import logging +import os +from typing import Protocol, SupportsIndex, TypeVar + +import numpy as np +import torch +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from qwen_vl_utils.vision_process import smart_resize +from torch.utils.data import DistributedSampler, random_split +from transformers import AutoProcessor + +from wall_x._vendor.x2robot_utils.geometry import ( + canonicalize_euler_zyx_batch_nb, + euler_to_matrix_zyx_batch_nb, + matrix_to_euler_zyx_batch_nb, + so3_to_matrix_batch_nb, +) +from wall_x.data.backends.lerobot.config import LerobotConfig +from wall_x.data.backends.lerobot.rotation_layout import ( + LAYOUT_SKIP_KEYS, + maybe_convert_euler_to_6d, +) +from wall_x.data.backends.lerobot.rotation_layout import ( + euler_layout_dim as _euler_layout_dim, +) +from wall_x.data.backends.lerobot.rotation_layout import ( + layout_uses_6d_rotation as _layout_uses_6d_rotation, +) +from wall_x.data.backends.lerobot.utils import ( + get_wallx_normal_text, + load_norm_stats, + preprocesser_call, + process_grounding_points, + replace_action_token, +) + +T_co = TypeVar("T_co", covariant=True) +logger = logging.getLogger(__name__) + +RELATIVE_KEYWORD = "relative" +ROTATION_KEYWORD = "rotation" +RELATIVE_SKIP_KEYS = LAYOUT_SKIP_KEYS + + +def _compute_delta_from_state_and_abs_rot( + rotations: np.ndarray, state: np.ndarray +) -> np.ndarray: + """Relative rotation: R_rel = R_abs @ R_state^T.""" + if rotations.shape[-1] == 3: + rotations_matrix = euler_to_matrix_zyx_batch_nb(rotations) + out_is_euler = True + elif rotations.shape[-1] == 6: + rotations_matrix = so3_to_matrix_batch_nb(rotations) + out_is_euler = False + else: + raise ValueError( + f"Only 3D euler or 6D rotation supported, got {rotations.shape[-1]}D" + ) + + if state.shape[-1] == 3: + state_matrix = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0] + elif state.shape[-1] == 6: + state_matrix = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] + else: + raise ValueError( + f"Only 3D euler or 6D rotation supported, got {state.shape[-1]}D" + ) + + r_rel = np.matmul(rotations_matrix, state_matrix.T) + if out_is_euler: + d_euler = matrix_to_euler_zyx_batch_nb(r_rel) + return canonicalize_euler_zyx_batch_nb(d_euler) + return r_rel[:, :2, :].reshape(r_rel.shape[0], 6) + + +# Abstract class for dataset +class Dataset(Protocol[T_co]): + """Interface for a dataset with random access.""" + + def __getitem__(self, index: SupportsIndex) -> T_co: + raise NotImplementedError("Subclasses of Dataset should implement __getitem__.") + + def __len__(self) -> int: + raise NotImplementedError("Subclasses of Dataset should implement __len__.") + + +class PreprocessedDataset(Dataset[T_co]): + def __init__( + self, + dataset, + config, + norm_stats, + dataload_config, + lerobot_config, + seed=42, + rank=0, + world_size=1, + test_only=False, + ): + self.hf_dataset = dataset + + if test_only: + self._dataset = dataset + else: + self._dataset = None + self.train_dataset, self.val_dataset = random_split( + dataset, + [0.95, 0.05], + torch.Generator().manual_seed(seed) if seed is not None else None, + ) + self._train() + + self.seed = seed + self.rank = rank + self.world_size = world_size + + # init configs + self.config = config + self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) + self.dataload_config = dataload_config + self.norm_stats = norm_stats + self.lerobot_config = lerobot_config + + self.data_config = LerobotConfig().update( + train_test_split=self.dataload_config["train_test_split"], + seed=self.dataload_config["seed"], + resolution=self.dataload_config.get("resolution", None), + priority_order=self.dataload_config.get("priority_order", None), + camera_name_mapping=self.dataload_config.get("camera_name_mapping", None), + ) + + self.key_mappings = self.dataload_config["key_mappings"] + + self._cam_key_mapping = self.key_mappings["camera"] + self._state_key_mapping = self.key_mappings + self._action_key_mapping = self.key_mappings + + task_cfg = self.config.get("task") or {} + self._dof_config = self.config.get("dof_config") or task_cfg.get( + "dof_config", {} + ) + self._agent_pos_config = self.config.get("agent_pos_config") or task_cfg.get( + "agent_pos_config", {} + ) + self._use_relative_action = any( + RELATIVE_KEYWORD in key for key in self._dof_config + ) + self._convert_action_euler_to_6d = _layout_uses_6d_rotation(self._dof_config) + self._convert_state_euler_to_6d = _layout_uses_6d_rotation( + self._agent_pos_config + ) + if self._convert_action_euler_to_6d or self._convert_state_euler_to_6d: + logger.info( + "LeRobot loader: Euler->6D rotation enabled " + "(action=%s, state=%s; raw action dim=%s -> %s)", + self._convert_action_euler_to_6d, + self._convert_state_euler_to_6d, + ( + _euler_layout_dim(self._dof_config) + if self._convert_action_euler_to_6d + else "-" + ), + ( + sum( + d + for k, d in self._dof_config.items() + if k not in RELATIVE_SKIP_KEYS + ) + if self._convert_action_euler_to_6d + else "-" + ), + ) + + def _maybe_convert_euler_to_6d(self, vec, layout_config: dict, enabled: bool): + converted = maybe_convert_euler_to_6d(vec, layout_config, enabled) + if ( + enabled + and layout_config + and isinstance(vec, torch.Tensor) + and converted is not vec + ): + return torch.as_tensor(converted, dtype=vec.dtype, device=vec.device) + return converted + + def _to_relative_action(self, action, agent_pos): + """Convert absolute action horizon to deltas w.r.t. current agent_pos.""" + action = np.asarray(action, dtype=np.float64) + agent_pos = np.asarray(agent_pos, dtype=np.float64) + if action.ndim == 1: + action = action[np.newaxis, :] + if agent_pos.ndim > 1: + agent_pos = agent_pos.reshape(-1) + + parts = [] + cur = 0 + for key, dim in self._dof_config.items(): + if key in RELATIVE_SKIP_KEYS: + continue + action_clip = action[:, cur : cur + dim] + agent_pos_clip = agent_pos[cur : cur + dim] + if RELATIVE_KEYWORD not in key: + parts.append(action_clip) + elif ROTATION_KEYWORD in key: + parts.append( + _compute_delta_from_state_and_abs_rot( + action_clip.astype(np.float64), + agent_pos_clip.astype(np.float64), + ) + ) + else: + parts.append(action_clip - agent_pos_clip[np.newaxis, :]) + cur += dim + + if not parts: + return action + return np.concatenate(parts, axis=1).astype(np.float32) + + def _vision_preprocess(self, frames): + processed_frames = [] + for key in self.hf_dataset.meta.camera_keys: + from PIL import Image + + current_obs = frames[key].clone().permute(1, 2, 0) + + img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy()) + orig_width, orig_height = img_pil.size + # 2. Apply resolution constraints (if config is not -1) + target_size = self.data_config.resolution.get( + self._cam_key_mapping[key], -1 + ) + if target_size != -1: + # Maintain aspect ratio logic + if orig_width > orig_height: # Landscape image + new_width = target_size + new_height = int(target_size * orig_height / orig_width) + else: # Portrait image + new_height = target_size + new_width = int(target_size * orig_width / orig_height) + img_pil = img_pil.resize((new_width, new_height)) + + # 3. Apply smart scaling (qwen logic) + current_width, current_height = img_pil.size + resized_height, resized_width = smart_resize( + current_height, + current_width, + factor=self.data_config.image_factor, + min_pixels=self.data_config.min_pixels, + max_pixels=self.data_config.max_pixels, + ) + resized_img = img_pil.resize((resized_width, resized_height)) + processed_frames.append(resized_img) + + return processed_frames, orig_height, orig_width, resized_height, resized_width + + def __getitem__(self, index): + data = self._dataset[index] + image_inputs, h, w, resize_h, resize_w = self._vision_preprocess(data) + agent_pos = data[self._state_key_mapping["state"]] + action = data[self._action_key_mapping["action"]] + agent_pos = self._maybe_convert_euler_to_6d( + agent_pos, self._agent_pos_config, self._convert_state_euler_to_6d + ) + action = self._maybe_convert_euler_to_6d( + action, self._dof_config, self._convert_action_euler_to_6d + ) + if self._use_relative_action: + device = action.device if isinstance(action, torch.Tensor) else None + action = torch.as_tensor( + self._to_relative_action(action, agent_pos), + dtype=torch.float32, + device=device, + ) + frame_index = data["frame_index"] + instruction_info = {"instruction": data["task"]} + generate_subtask_ratio = self.data_config.generate_subtask_ratio + complete_text, generate_subtask = get_wallx_normal_text( + instruction_info, + self.dataload_config.get("action_horizon", 33) - 1, + frame_index, + self.data_config.priority_order, + self._cam_key_mapping, + generate_subtask_ratio=generate_subtask_ratio, + camera_name_mapping=self.data_config.camera_name_mapping, + ) + text = process_grounding_points( + complete_text, h, w, resize_h, resize_w, self.data_config.model_type + ) + result = { + "image_inputs": image_inputs, + "text": text, + "action": action, + "agent_pos": agent_pos, + "frame_index": frame_index, + } + + return result + + def __len__(self) -> int: + return len(self._dataset) + + def _eval(self): + self._dataset = self.val_dataset + + def _train(self): + self._dataset = self.train_dataset + + def get_train_dataloader(self): + """ + Get distributed training dataloader + + Args: + rank: Current process rank + world_size: Total number of processes + seed: Random seed for reproducibility + """ + self._train() + + batch_size = self.config.get("batch_size_per_gpu", 8) + num_workers = self.config.get("num_workers", 4) + + # Create distributed sampler + sampler = DistributedSampler( + self, + num_replicas=self.world_size, + rank=self.rank, + shuffle=True, + seed=self.seed, + drop_last=True, # Ensure all processes have same number of batches + ) + + dataloader = torch.utils.data.DataLoader( + self, + batch_size=batch_size, + sampler=sampler, # Use distributed sampler instead of shuffle=True + num_workers=num_workers, + collate_fn=DataCollator( + self.config, self.dataload_config, self.norm_stats, self.lerobot_config + ), + pin_memory=True, # Enable for GPU training + persistent_workers=num_workers > 0, # Only if num_workers > 0 + prefetch_factor=2, # Reduce memory usage + drop_last=True, # Avoid incomplete batches + ) + + return dataloader, sampler + + def get_val_dataloader(self): + """ + Get distributed evaluation dataloader (no shuffling for consistent evaluation) + """ + self._eval() + + batch_size = self.config.get( + "eval_batch_size_per_gpu", self.config.get("batch_size_per_gpu", 8) + ) + num_workers = self.config.get("num_workers", 4) + + # Create distributed sampler for evaluation (no shuffle) + sampler = DistributedSampler( + self, + num_replicas=self.world_size, + rank=self.rank, + shuffle=False, # No shuffling for evaluation + drop_last=False, # Keep all samples for evaluation + ) + + dataloader = torch.utils.data.DataLoader( + self, + batch_size=batch_size, + sampler=sampler, + num_workers=num_workers, + collate_fn=DataCollator( + self.config, self.dataload_config, self.norm_stats, self.lerobot_config + ), + pin_memory=True, + persistent_workers=num_workers > 0, + prefetch_factor=2, + drop_last=False, + ) + + return dataloader, sampler + + +class DataCollator: + # Class-level cache for processors to avoid reloading + _processor_cache = {} + _action_tokenizer_cache = {} + _norm_stat_alignment_warnings = set() + + def __init__(self, config, dataload_config, stats, lerobot_config): + self.config = config + self.dataload_config = dataload_config + self.stats = stats + self.action_min_stat = stats["action"].min + self.action_delta = stats["action"].delta + self.state_min_stat = stats["state"].min + self.state_delta = stats["state"].delta + self.lerobot_config = lerobot_config + self.np_rng = np.random.default_rng() + + noise_scheduler_config = config.get("noise_scheduler", {}) + self.beta_alpha = noise_scheduler_config.get( + "beta_alpha", 1.5 + ) # alpha parameter of the Beta distribution + self.beta_beta = noise_scheduler_config.get( + "beta_beta", 1.0 + ) # beta parameter of the Beta distribution + self.s = noise_scheduler_config.get("s", 0.999) # scaling factor + self.time_shift = noise_scheduler_config.get( + "time_shift", 1.0 + ) # time shift factor + + self.beta_alpha = float(self.beta_alpha) + self.beta_beta = float(self.beta_beta) + self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) + self.use_state_string_representation = bool( + self.config.get("use_state_string_representation", False) + ) + self.state_bins = int(self.config.get("state_bins", 256)) + self.load_processor() + + def load_processor(self): + processor_path = self.config["processor_path"] + action_tokenizer_path = self.config.get("action_tokenizer_path", None) + + if ( + self.use_fast_tokenizer + and action_tokenizer_path not in self._action_tokenizer_cache + ): + self._action_tokenizer_cache[action_tokenizer_path] = ( + AutoProcessor.from_pretrained( + action_tokenizer_path, trust_remote_code=True + ) + ) + + # Use cached processors if available + if processor_path not in self._processor_cache: + processor = AutoProcessor.from_pretrained(processor_path, use_fast=True) + if self.config.get("padding_side", "left") == "left": + processor.tokenizer.padding_side = "left" + + new_tokens = ["<|propri|>", "<|action|>"] + processor.tokenizer.add_tokens(new_tokens) + if self.use_fast_tokenizer and self.config.get("model_type") == "qwen2_5": + action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path] + new_tokens = [ + f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size) + ] + processor.tokenizer.add_tokens(new_tokens) + begin_idx_token = "<|action_token_0|>" + token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token) + processor.tokenizer.init_kwargs["action_token_start_index"] = token_id + processor.tokenizer.init_kwargs["action_token_vocab_size"] = ( + action_tokenizer.vocab_size + ) + + self._processor_cache[processor_path] = processor + + self.processor = self._processor_cache[processor_path] + + if not self.use_fast_tokenizer: + self.train_action_tokenizer = None + else: + self.train_action_tokenizer = self._action_tokenizer_cache[ + action_tokenizer_path + ] + + @classmethod + def _normalize(cls, action, min_stat, delta): + """ + Normalize action data using min-max normalization. + """ + delta = torch.where(delta == 0, torch.ones_like(delta), delta) + x = (action - min_stat) / delta + x = x * 2 - 1 + x = torch.clamp(x, -1, 1) + return x + + @staticmethod + def _align_norm_stat(stat, value, *, pad_value: float, name: str): + """Align a 1-D norm stat with the current LeRobot tensor width.""" + stat = stat.to(device=value.device, dtype=value.dtype) + target_dim = value.shape[-1] + stat_dim = stat.shape[-1] + if stat_dim == target_dim: + return stat + if stat_dim > target_dim: + warning_key = ("truncate", name, stat_dim, target_dim) + if warning_key not in DataCollator._norm_stat_alignment_warnings: + logger.warning( + "Truncating LeRobot %s norm stat from %s to %s dims", + name, + stat_dim, + target_dim, + ) + DataCollator._norm_stat_alignment_warnings.add(warning_key) + return stat[..., :target_dim] + pad_shape = (*stat.shape[:-1], target_dim - stat_dim) + pad = torch.full(pad_shape, pad_value, device=value.device, dtype=value.dtype) + warning_key = ("pad", name, stat_dim, target_dim) + if warning_key not in DataCollator._norm_stat_alignment_warnings: + logger.warning( + "Padding LeRobot %s norm stat from %s to %s dims", + name, + stat_dim, + target_dim, + ) + DataCollator._norm_stat_alignment_warnings.add(warning_key) + return torch.cat([stat, pad], dim=-1) + + def __call__(self, batch): + additional_inputs = {} + + # Tail-pad widths when dof_config / agent_pos_config (sum) is larger + # than the lerobot action/state - typical when resuming a ckpt that + # was pretrained on a bigger action space. Extra columns are filled + # with zeros and their mask set to 0 so loss is not propagated. + dof_total = int(self.config.get("dof_total_dim", 0) or 0) + agent_pos_total = int(self.config.get("agent_pos_total_dim", 0) or 0) + + # Explicit init so the ``if action is not None`` guard and later + # ``replace_action_token`` call stay well-defined even if a batch + # unexpectedly omits the action / agent_pos keys. Without this the + # loop-local variables would leak NameError on the first miss. + action = None + dof_mask = None + agent_pos = None + agent_pos_mask = None + + for key in batch[0].keys(): + if key == "agent_pos": + agent_pos = torch.stack([item["agent_pos"] for item in batch]) + if agent_pos.dim() == 2: + agent_pos = agent_pos.unsqueeze(1) + agent_pos_mask = (~torch.isnan(agent_pos)).float() + agent_pos.nan_to_num_(nan=0.0) + state_min_stat = self._align_norm_stat( + self.state_min_stat, + agent_pos, + pad_value=0.0, + name="state.min", + ) + state_delta = self._align_norm_stat( + self.state_delta, + agent_pos, + pad_value=1.0, + name="state.delta", + ) + agent_pos = self._normalize(agent_pos, state_min_stat, state_delta) + if agent_pos_total and agent_pos.shape[-1] < agent_pos_total: + pad_w = agent_pos_total - agent_pos.shape[-1] + agent_pos = torch.nn.functional.pad(agent_pos, (0, pad_w)) + agent_pos_mask = torch.nn.functional.pad(agent_pos_mask, (0, pad_w)) + additional_inputs["proprioception"] = agent_pos + additional_inputs["agent_pos_mask"] = agent_pos_mask + elif key == "action": + action = torch.stack([item["action"] for item in batch]) + if action.dim() == 2: + action = action.unsqueeze(1) + dof_mask = (~torch.isnan(action)).float() + action.nan_to_num_(nan=0.0) + action_min_stat = self._align_norm_stat( + self.action_min_stat, + action, + pad_value=0.0, + name="action.min", + ) + action_delta = self._align_norm_stat( + self.action_delta, + action, + pad_value=1.0, + name="action.delta", + ) + action = self._normalize(action, action_min_stat, action_delta) + if dof_total and action.shape[-1] < dof_total: + pad_w = dof_total - action.shape[-1] + action = torch.nn.functional.pad(action, (0, pad_w)) + dof_mask = torch.nn.functional.pad(dof_mask, (0, pad_w)) + additional_inputs["action_chunk"] = action + additional_inputs["dof_mask"] = dof_mask + elif key == "image_inputs": + additional_inputs["image_inputs"] = [ + item["image_inputs"] for item in batch + ] + elif key == "text": + additional_inputs["text"] = [item["text"] for item in batch] + elif key == "frame_index": + additional_inputs["frame_index"] = torch.stack( + [item["frame_index"] for item in batch] + ) + else: + raise NotImplementedError( + f"{key} input not implemented in preprocesser" + ) + + # sample noise time + if action is not None: + sample_time = self.sample_time( + action.shape[0], + device=action.device, + dtype=torch.float32, + ) + additional_inputs["sample_time"] = sample_time + + additional_inputs["text"] = replace_action_token( + additional_inputs["text"], + additional_inputs["action_chunk"], + self.train_action_tokenizer if self.use_fast_tokenizer else None, + additional_inputs["dof_mask"], + ) + + inputs = preprocesser_call( + processor=self.processor, + text=additional_inputs.pop("text"), + images=additional_inputs.pop("image_inputs"), + videos=None, + padding=True, + truncation=True, + return_tensors="pt", + max_length=self.dataload_config.get("max_length", 768), + norm_state=( + additional_inputs["proprioception"] + if self.use_state_string_representation + and "proprioception" in additional_inputs + else None + ), + agent_pos_mask=additional_inputs.get("agent_pos_mask"), + state_bins=self.state_bins, + ) + + action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") + + # Gating token types + additional_inputs["moe_token_types"] = inputs.input_ids == action_token_id + + inputs.update(additional_inputs) + + inputs["dataset_names"] = [self.lerobot_config["repo_id"]] * inputs[ + "action_chunk" + ].shape[0] + + return inputs + + def sample_time(self, batch_size, device, dtype): + """ + Sample timesteps + + Use a Beta distribution to sample values in [0, 1], then scale them. + + Args: + batch_size (int): batch size + device: Device type + dtype: dtype + + Returns: + torch.Tensor: sampled timesteps with shape [batch_size] + """ + + sample_np = self.np_rng.beta( + self.beta_alpha, self.beta_beta, size=(batch_size,) + ).astype(np.float32) + sample = torch.from_numpy(sample_np).to( + device=device, dtype=dtype, non_blocking=True + ) + + # sample = self.beta_dist.sample([batch_size]).to(dtype=dtype) + time = 1 - sample + + # Apply diffusion time shift + if self.time_shift != 1.0: + time = (self.time_shift * time) / (1 + (self.time_shift - 1) * time) + + time = time * self.s # noise should denoise from 0 to 1 here + return time + + +def load_lerobot_data( + config, + lerobot_config, + rank=0, + world_size=1, + seed=42, +): + """ + Load LeRobot dataset with distributed support + + Args: + config: Model configuration + rank: Current process rank (default: 0) + world_size: Total number of processes (default: 1) + seed: Random seed for reproducibility (default: 42) + + Returns: + dataset: Training dataset + train_num: Number of training samples per process + sampler: Distributed sampler (None if world_size=1) + """ + + # Set seed for reproducibility + torch.manual_seed(seed) + + dataload_config = get_data_configs(config["data"]) + key_mappings = dataload_config["key_mappings"] + + repo_id = lerobot_config.get("repo_id", None) + assert repo_id is not None, "repo id is required" + root = lerobot_config.get("root", None) + meta_info = LeRobotDatasetMetadata(repo_id, root=root) + dataset_fps = meta_info.fps + episodes_num = meta_info.total_episodes + + norm_stats_path = config.get("norm_stats_path", None) + assert ( + norm_stats_path is not None + ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats" + task_cfg = config.get("task") or {} + dof_config = config.get("dof_config") or task_cfg.get("dof_config", {}) + agent_pos_config = config.get("agent_pos_config") or task_cfg.get( + "agent_pos_config", {} + ) + norm_stats = load_norm_stats( + norm_stats_path, + key_mappings, + dof_config=dof_config, + agent_pos_config=agent_pos_config, + ) + + delta_timestamps = { + # action chunk + key_mappings["action"]: [ + t / dataset_fps + for t in range(dataload_config.get("action_horizon", 33) - 1) + ], + } + batch_size = config.get("batch_size_per_gpu", 8) + + # Optional episode subset. YAML ``lerobot_config.episodes`` has always + # been present in examples but previously ignored; honour it so smoke + # tests / small-dataset runs don't pay the O(N) LeRobotDataset indexing + # cost on a multi-thousand-episode repo (~10s / episode on some formats). + episodes_override = lerobot_config.get("episodes") + if episodes_override is not None: + episodes = list(episodes_override) + episodes_num_effective = len(episodes) + else: + episodes = np.arange(episodes_num).tolist() + episodes_num_effective = episodes_num + + train_test_split = dataload_config.get("train_test_split", 0.95) + split_idx = int(episodes_num_effective * train_test_split) + # Guard: tiny episode subsets + high train_test_split can floor split_idx + # to 0 (e.g. 1 ep * 0.95 = 0), which would silently hand LeRobotDataset an + # empty episode list and end training after 0 iterations. Fail loud. + if split_idx < 1: + raise ValueError( + f"train_test_split={train_test_split} applied to " + f"{episodes_num_effective} episode(s) yields 0 train episodes. " + f"Use more episodes or a lower train_test_split." + ) + train_episodes = episodes[:split_idx] + test_episodes = episodes[split_idx:] + + global_rank = torch.distributed.get_rank() + local_rank = int(os.environ["LOCAL_RANK"]) + local_world_size = int(os.environ["LOCAL_WORLD_SIZE"]) + # TODO: Some LeRobot formats need to load all metadata before splitting + # episodes; loading from all ranks at once can exhaust memory. + train_dataset = None + + # Sequential loading inside each node + for r in range(local_world_size): + if local_rank == r: + logger.info( + "[Global rank %s] Loading dataset on local_rank=%s", + global_rank, + local_rank, + ) + + train_dataset = LeRobotDataset( + repo_id=repo_id, + root=root, + episodes=train_episodes, + delta_timestamps=delta_timestamps, + video_backend="pyav", + ) + + logger.info( + "[Global rank %s] Finished loading on local_rank=%s", + global_rank, + local_rank, + ) + + # Barrier only within the node + torch.distributed.barrier(device_ids=[local_rank]) + + if rank == 0: + logger.info("Selected train episodes: %s", train_dataset.episodes) + logger.info("Number of train episodes selected: %s", train_dataset.num_episodes) + logger.info("Number of train frames selected: %s", train_dataset.num_frames) + logger.info("Selected test episodes: %s", test_episodes) + + dataset = PreprocessedDataset( + train_dataset, + config, + norm_stats, + dataload_config, + lerobot_config, + seed=seed, + rank=rank, + world_size=world_size, + ) + + # Calculate samples per process + if world_size > 1: + # With DistributedSampler, each process gets approximately len(dataset) // world_size samples + samples_per_process = len(dataset) // world_size + train_num = samples_per_process // batch_size + else: + train_num = len(dataset) // batch_size + + if rank == 0: + lines = [ + "LeRobot Data Loading Configuration:", + f" rank: {rank}", + f" world_size: {world_size}", + f" batch_size_per_gpu: {batch_size}", + f" repo_id: {repo_id}", + f" total_dataset_size: {len(dataset)}", + ] + if world_size > 1: + lines.extend( + [ + f" samples_per_process: {samples_per_process}", + f" batches_per_process: {train_num}", + f" total_batches_all_processes: {train_num * world_size}", + ] + ) + else: + lines.append(f" total_batches: {train_num}") + lines.append(f" seed: {seed}") + logger.info("\n%s", "\n".join(lines)) + + return dataset, train_num + + +def get_distributed_dataloader( + dataset, config, rank=0, world_size=1, seed=42, is_train=True +): + """ + Helper function to get distributed dataloader + + Args: + dataset: PreprocessedDataset instance + config: Configuration dict + rank: Current process rank + world_size: Total number of processes + seed: Random seed + is_train: Whether this is for training (affects shuffling) + + Returns: + dataloader: Distributed DataLoader + sampler: DistributedSampler + """ + if is_train: + return dataset.get_train_dataloader(rank=rank, world_size=world_size, seed=seed) + else: + return dataset.get_val_dataloader(rank=rank, world_size=world_size) + + +def get_data_configs(config): + default_data_config = { + "train_test_split": 0.95, + "seed": 42, + "batch_size": 8, + "action_horizon": 21, + "action_history_length": 0, + "image_horizon": 1, + "image_history_length": 0, + "left_padding": False, + "right_padding": False, + "return_first_obs": False, + "return_last_obs": False, + "randomize_obs_after": None, + "datasets": [], + "labeled_pathes": [], + "camera_name_mapping": None, + } + data_config = default_data_config | config + data_config["action_horizon"] += 1 + + return data_config diff --git a/wall_x/data/backends/lerobot/rotation_layout.py b/wall_x/data/backends/lerobot/rotation_layout.py new file mode 100644 index 0000000..ba8a780 --- /dev/null +++ b/wall_x/data/backends/lerobot/rotation_layout.py @@ -0,0 +1,95 @@ +"""Layout helpers when config expects 6D rotation but LeRobot stores 3D Euler.""" + +from __future__ import annotations + +import numpy as np + +from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb + +LAYOUT_SKIP_KEYS = frozenset( + {"velocity_decomposed", "height", "head_actions", "action_padding"} +) +ROTATION_KEYWORD = "rotation" +ROTATION_6D_KEYWORD = "6D" + + +def layout_uses_6d_rotation(layout_config: dict) -> bool: + for key, dim in layout_config.items(): + if key in LAYOUT_SKIP_KEYS: + continue + if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6: + return True + return False + + +def euler_layout_dim(layout_config: dict) -> int: + """Vector width in LeRobot when rotation slices are still 3D Euler.""" + total = 0 + for key, dim in layout_config.items(): + if key in LAYOUT_SKIP_KEYS: + continue + if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6: + total += 3 + else: + total += int(dim) + return total + + +def convert_euler_to_6d(vec: np.ndarray, layout_config: dict) -> np.ndarray: + """Rewrite [pos, euler(3), tail...] to [pos, rot6d(6), tail...] per layout.""" + vec = np.asarray(vec, dtype=np.float64) + single = vec.ndim == 1 + if single: + vec = vec[np.newaxis, :] + + out_rows = [] + for row in vec: + parts: list[np.ndarray] = [] + raw_cur = 0 + for key, dim in layout_config.items(): + if key in LAYOUT_SKIP_KEYS: + continue + dim = int(dim) + if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6: + euler = row[raw_cur : raw_cur + 3] + rot6d = euler_to_matrix_zyx_6d_nb(euler.reshape(1, 3)).reshape(6) + parts.append(rot6d) + raw_cur += 3 + else: + parts.append(row[raw_cur : raw_cur + dim]) + raw_cur += dim + out_rows.append(np.concatenate(parts, axis=0)) + + out = np.stack(out_rows, axis=0) + return out[0] if single else out + + +def maybe_convert_norm_stats_vector( + values, + layout_config: dict, + enabled: bool | None = None, +): + """Convert a 1D norm-stat vector (q01/q99/mean/std) from Euler layout to 6D.""" + if enabled is None: + enabled = layout_uses_6d_rotation(layout_config) + if not enabled or not layout_config: + return values + arr = np.asarray(values, dtype=np.float64) + if arr.ndim != 1: + return values + raw_dim = euler_layout_dim(layout_config) + if arr.shape[0] != raw_dim: + return values + return convert_euler_to_6d(arr, layout_config).astype(np.float32) + + +def maybe_convert_euler_to_6d( + vec: np.ndarray, layout_config: dict, enabled: bool +) -> np.ndarray: + if not enabled or not layout_config: + return vec + raw_dim = euler_layout_dim(layout_config) + arr = np.asarray(vec) + if arr.shape[-1] != raw_dim: + return vec + return convert_euler_to_6d(arr, layout_config).astype(np.float32) diff --git a/wall_x/data/utils.py b/wall_x/data/backends/lerobot/utils.py similarity index 72% rename from wall_x/data/utils.py rename to wall_x/data/backends/lerobot/utils.py index 9cb0efc..e002840 100644 --- a/wall_x/data/utils.py +++ b/wall_x/data/backends/lerobot/utils.py @@ -1,119 +1,457 @@ -""" -Data processing utilities for Wall-X multimodal robotic learning. - -This module provides utilities for preprocessing text, images, and actions -for multimodal transformer models in robotic learning tasks. -""" - import json +import logging import random import re from collections import OrderedDict from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple, Union -import numpy as np +import numpy as np import torch from transformers import BatchFeature -KEY_MAPPINGS = { - "lerobot/aloha_mobile_cabinet": { - "camera": { - "observation.images.cam_high": "face_view", - "observation.images.cam_left_wrist": "left_wrist_view", - "observation.images.cam_right_wrist": "right_wrist_view", - }, - "state": "observation.state", - "action": "action", - }, - "physical-intelligence/libero": { - "camera": { - "image": "face_view", - "wrist_image": "left_wrist_view", - }, - "state": "state", - "action": "actions", - }, - "x2_normal": { - "camera": { - "observation.images.faceImg": "face_view", - "observation.images.leftImg": "left_wrist_view", - "observation.images.rightImg": "right_wrist_view", - }, - "state": "observation.state", - "action": "action", - }, - "libero": { - "camera": { - "observation.images.faceImg": "face_view", - "observation.images.rightImg": "right_wrist_view", - }, - "state": "observation.state", - "action": "action", - }, - "robochallenge_aloha": { - "camera": { - "observation.images.cam_high_rgb": "face_view", - "observation.images.cam_wrist_left_rgb": "left_wrist_view", - "observation.images.cam_wrist_right_rgb": "right_wrist_view", - }, - "state": "observation.state", - "action": "action", - }, -} - -CAMERA_NAME_MAPPING = { - "face_view": "front view", - "left_wrist_view": "left wrist view", - "right_wrist_view": "right wrist view", - "move1_view": "move view", - "move2_view": "move view", - "wall_view": "wall view", - "top_view": "top view", -} +logger = logging.getLogger(__name__) -MULTIMODAL_DATASET_NAMES = [ - "x2_multimodal_from_action", - "x2_multimodal", - "x2_subtask_generation", - "multimodal_CapsFusion", - "multimodal_Robo2VLM", - "multimodal_RoboPoint", - "multimodal_EQA", - "multimodal_Cambrian", - "multimodal_pixmo", - "multimodal_VQAv2", - "multimodal_COCO", -] +@dataclass +class NormStats: + min: torch.Tensor + max: torch.Tensor + delta: torch.Tensor -FREQUENCY_MAPPING = { - "x2_normal": 32, - "fractal": 5, - "bridge_data_v2": 5, - "droid": 15, - "agibotworld_alpha": 32, - "DobbE": 30, - "RH20T": 10, - "UMI-biarm": 10, - "austin_buds": 20, - "austin_sailor": 20, - "austin_sirius": 20, - "bc_z": 10, - "berkeley_autolab_ur5": 5, - "berkeley_cable_routing": 10, - "berkeley_fanuc_manipulation": 10, - "dlr_edan_shared_control": 5, - "fmb": 10, - "furniture_bench": 10, - "jaco_play": 10, - "nyu_rot": 10, - "stanford_hydra": 10, - "stanford_kuka_multimodal": 20, - "taco_play": 30, - "utaustin_mutex": 20, - "viola": 20, -} + +def load_norm_stats( + norm_stats_path, + key_mappings, + dof_config: dict | None = None, + agent_pos_config: dict | None = None, +): + from wall_x.data.backends.lerobot.rotation_layout import ( + layout_uses_6d_rotation, + maybe_convert_norm_stats_vector, + ) + + with open(norm_stats_path, "r") as f: + norm_stats = json.load(f) + action_key = key_mappings["action"] + state_key = key_mappings["state"] + + action_q01 = maybe_convert_norm_stats_vector( + norm_stats["norm_stats"][action_key]["q01"], dof_config or {} + ) + action_q99 = maybe_convert_norm_stats_vector( + norm_stats["norm_stats"][action_key]["q99"], dof_config or {} + ) + if layout_uses_6d_rotation(dof_config or {}) and len(action_q01) != len( + norm_stats["norm_stats"][action_key]["q01"] + ): + logger.info( + "Converted action norm stats Euler->6D (%d -> %d dims)", + len(norm_stats["norm_stats"][action_key]["q01"]), + len(action_q01), + ) + + q01 = torch.tensor(action_q01) + q99 = torch.tensor(action_q99) + delta = q99 - q01 + action_norm_stats = NormStats( + min=q01, + max=q99, + delta=delta, + ) + + state_q01 = maybe_convert_norm_stats_vector( + norm_stats["norm_stats"][state_key]["q01"], agent_pos_config or {} + ) + state_q99 = maybe_convert_norm_stats_vector( + norm_stats["norm_stats"][state_key]["q99"], agent_pos_config or {} + ) + if layout_uses_6d_rotation(agent_pos_config or {}) and len(state_q01) != len( + norm_stats["norm_stats"][state_key]["q01"] + ): + logger.info( + "Converted state norm stats Euler->6D (%d -> %d dims)", + len(norm_stats["norm_stats"][state_key]["q01"]), + len(state_q01), + ) + + q01 = torch.tensor(state_q01) + q99 = torch.tensor(state_q99) + delta = q99 - q01 + state_norm_stats = NormStats( + min=q01, + max=q99, + delta=delta, + ) + + return {"action": action_norm_stats, "state": state_norm_stats} + + +def get_frame_instruction( + instruction_info: Dict[str, Any], + frame_idx: Optional[int] = None, + truncate_keys: Optional[List[str]] = None, +) -> Tuple[Dict[str, Any], Optional[int]]: + """Extract frame-specific instruction from instruction dictionary. + + Args: + instruction_info: Dictionary containing instruction components + frame_idx: Current frame index + truncate_keys: Keys that trigger truncation when found + + Returns: + Tuple of (frame_instruction_dict, split_end_frame) + """ + if truncate_keys is None: + truncate_keys = [ + "subtask_generation", + "distribute", + "subtask_generation_zh", + "distribute_zh", + ] + + instruction_for_frame = {} + split_end = None + + for key, value in instruction_info.items(): + if isinstance(value, dict): + # Handle frame-range specific instructions + for frame_range, frame_instruction in value.items(): + start_frame, end_frame = map(int, frame_range.split(" ")) + if start_frame <= frame_idx < end_frame or (start_frame == frame_idx): + instruction_for_frame[key] = frame_instruction + if ( + truncate_keys is not None + and split_end is None + and key in truncate_keys + ): + split_end = end_frame + 1 + break + else: + instruction_for_frame[key] = value + + return instruction_for_frame, split_end + + +def get_action_tokens( + normalized_actions: Union[torch.Tensor, List], action_tokenizer +) -> List[List[str]]: + """Convert normalized actions to action token strings. + + Args: + normalized_actions: Normalized action arrays/tensors + action_tokenizer: Tokenizer for converting actions to tokens + + Returns: + List of action token string lists for each sample + """ + if isinstance(normalized_actions, torch.Tensor): + normalized_actions = normalized_actions.cpu().numpy() + + all_action_tokens = [] + for i in range(len(normalized_actions)): + if isinstance(normalized_actions[i], torch.Tensor): + normalized_actions[i] = normalized_actions[i].cpu().numpy() + + token_id = action_tokenizer(normalized_actions[i]) + action_tokens = [f"<|action_token_{j}|>" for j in token_id[0]] + all_action_tokens.append(action_tokens) + + return all_action_tokens + + +def pad_action_token_strs( + actions_token_lists: List[List[str]], pad_token: str = "<|endoftext|>" +) -> List[str]: + """Pad action token lists to same length and join as strings. + + Args: + actions_token_lists: List of action token lists for each sample + pad_token: Token used for padding + + Returns: + List of padded action token strings + """ + max_len = max(len(tokens) for tokens in actions_token_lists) + padded_action_strs = [] + + for tokens in actions_token_lists: + padded_tokens = ( + tokens + ["<|im_end|>\n"] + [pad_token] * (max_len - len(tokens)) + ) + padded_action_strs.append("".join(padded_tokens)) + + return padded_action_strs + + +def get_task_instruction( + frame_instruction_info: Dict[str, Any], priority_order: Optional[OrderedDict] = None +) -> str: + """Construct task instruction from available instruction fields using priority sampling. + + Args: + frame_instruction_info: Dictionary containing instruction fields + priority_order: OrderedDict specifying sampling probability for each field + + Returns: + Combined instruction string with priority components + """ + # Default priority settings + default_priority_order = OrderedDict( + { + "subtask_generation": 0.25, + "subtask_generation_zh": 0.25, + "distribute": 0.25, + "distribute_zh": 0.25, + } + ) + + if priority_order is not None: + priority_order = OrderedDict(priority_order) + else: + priority_order = default_priority_order + + got_instruction = False + task_instruction = "" + + # Sample instruction components based on priority probabilities + for key, prob in priority_order.items(): + if key in frame_instruction_info and frame_instruction_info[key] != "": + if got_instruction: + if random.random() >= prob: + continue + + task_instruction += f"\n{frame_instruction_info[key]}" + got_instruction = True + break + + # Fall back to base instruction if no priority components found + if not got_instruction: + task_instruction = frame_instruction_info.get("instruction", "") + + return task_instruction + + +def process_grounding_points( + text: str, + orig_height: int, + orig_width: int, + resized_height: int, + resized_width: int, + model_type: str, +) -> str: + """Process grounding point coordinates in text based on image resizing. + + Adjusts coordinate values in tags to match resized image dimensions + for different model types (qwen2, qwen2_5). + + Args: + text: Input text containing tags with coordinates + orig_height: Original image height + orig_width: Original image width + resized_height: Resized image height + resized_width: Resized image width + model_type: Model type for coordinate processing ('qwen2' or 'qwen2_5') + + Returns: + Text with adjusted coordinate values + """ + # Regex pattern to match tags and their contents + point_pattern = re.compile(r"(.*?)") + + def process_match(match): + """Process a single point match and adjust coordinates.""" + coords_str = match.group(1) + try: + # Extract coordinates from string + coords = list(map(int, re.findall(r"\d+", coords_str))) + + # Calculate resize scale factors + scale_w = resized_width / orig_width + scale_h = resized_height / orig_height + + if len(coords) == 2: + x, y = coords + if model_type == "qwen2_5": + # Qwen2.5 uses pixel coordinates + new_x = max(0, min(round(x * scale_w), resized_width - 1)) + new_y = max(0, min(round(y * scale_h), resized_height - 1)) + elif model_type == "qwen2": + # Qwen2 normalizes to [0, 1000) range + new_x = max(0, min(999.999, (x / orig_width) * 1000)) + new_y = max(0, min(999.999, (y / orig_height) * 1000)) + else: + raise ValueError(f"Unsupported model type: {model_type}") + coords = [new_x, new_y] + + elif len(coords) == 4: + x1, y1, x2, y2 = coords + if model_type == "qwen2_5": + new_x1 = max(0, min(round(x1 * scale_w), resized_width - 1)) + new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1)) + new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1)) + new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1)) + elif model_type == "qwen2": + new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000)) + new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000)) + new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000)) + new_y2 = max(0, min(999.999, (y2 / orig_height) * 1000)) + else: + raise ValueError(f"Unsupported model type: {model_type}") + coords = [new_x1, new_y1, new_x2, new_y2] + + # Return processed point tag + return f'[{", ".join(map(str, coords))}]' + + except (ValueError, TypeError): + # Return original content if processing fails + return match.group(0) + + # Replace all matching point tags + processed_text = point_pattern.sub(process_match, text) + return processed_text + + +def get_wallx_normal_text( + instruction_info: Dict[str, Any], + action_chunk_size: int, + frame_idx: int, + priority_order: Optional[OrderedDict] = None, + cam_mapping: Optional[Dict[str, str]] = None, + generate_subtask_ratio: float = 0.0, + camera_name_mapping: Optional[Dict[str, str]] = None, +) -> Tuple[str, bool]: + """Construct complete multimodal prompt text for Wall-X model. + + Formats input using special tokens including: + - System message + - User observations (with image placeholders) + - Task instructions + - Proprioception prompts + - Assistant responses (with action tokens) + + Args: + instruction_info: Dictionary containing instruction components + action_chunk_size: Number of action tokens to generate + frame_idx: Current frame index + priority_order: Priority order for instruction sampling + cam_mapping: Camera name mapping dictionary + generate_subtask_ratio: Probability of generating subtask instead of actions + camera_name_mapping: Optional display-name mapping for prompt text + + Returns: + Tuple of (formatted_prompt_text, is_subtask_generation) + """ + # Special tokens for formatting + role_start_symbol = "<|im_start|>" + role_end_symbol = "<|im_end|>" + vision_start_symbol = "<|vision_start|>" + vision_end_symbol = "<|vision_end|>" + image_pad_symbol = "<|image_pad|>" + propri_symbol = "<|propri|>" + action_symbol = "<|action|>" + action_fast_symbol = "<|action_fast|>" + + # System prologue + prologue = ( + f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n" + ) + + # User request with observation + user_request = f"{role_start_symbol}user\nObservation:" + if cam_mapping: + camera_name_mapping = camera_name_mapping or {} + for _, cam_name in cam_mapping.items(): + view_name = camera_name_mapping.get(cam_name, cam_name) + user_request += f" {view_name}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" + user_request += "\nInstruction:" + + # Get frame-specific instruction + frame_instruction_info, _ = get_frame_instruction( + instruction_info, frame_idx=frame_idx + ) + + generate_subtask = False + priority_keys = ["subtask_generation", "distribute"] + + # Decide whether to generate subtask or actions + if ( + bool(set(frame_instruction_info.keys()) & set(priority_keys)) + and random.random() < generate_subtask_ratio + ): + # Generate subtask (equivalent to VQA task) + instruction = frame_instruction_info.get("instruction", "") + text_prompt = "\nPredict the next action in language.\n" + user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" + + # Find output instruction from priority keys + for key in priority_keys: + if key in frame_instruction_info: + output_instruction = frame_instruction_info[key] + break + + assistant_output = ( + f"{role_start_symbol}assistant\n{output_instruction}\n{role_end_symbol}" + ) + generate_subtask = True + else: + # Generate actions + instruction = get_task_instruction( + frame_instruction_info, priority_order=priority_order + ) + text_prompt = f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n" + user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" + assistant_output = f"{role_start_symbol}assistant\n{action_fast_symbol}{role_end_symbol}\n{action_symbol * action_chunk_size}" + + complete_text = prologue + user_message + assistant_output + return complete_text, generate_subtask + + +def replace_action_token( + text: List[str], + norm_action: Optional[torch.Tensor], + action_tokenizer, + dof_masks: Optional[torch.Tensor] = None, +) -> List[str]: + """Replace action placeholders in text with actual action tokens. + + Args: + text: List of text strings with action placeholders + norm_action: Normalized action tensors + action_tokenizer: Tokenizer for converting actions to tokens + dof_masks: Masks for degrees of freedom + + Returns: + List of text strings with action tokens replaced + """ + if action_tokenizer is not None and norm_action is not None: + if dof_masks is not None: + norm_action = [ + action[:, dof_masks[i, 0].bool()] + for i, action in enumerate(norm_action) + ] + + # Convert to action tokens and pad + actions_fast_tokens = get_action_tokens(norm_action, action_tokenizer) + actions_fast_token_strs = pad_action_token_strs(actions_fast_tokens) + + # Replace action placeholders with actual tokens + actions_fast_token_idx = 0 + for i in range(len(text)): + if "<|action_fast|>" in text[i]: + text[i] = text[i].replace( + "<|action_fast|><|im_end|>\n", + actions_fast_token_strs[actions_fast_token_idx], + ) + actions_fast_token_idx += 1 + + # Remove remaining action placeholders + text = [t.replace("<|action|>", "") for t in text] + else: + # Remove action placeholders when no tokenizer available + text = [t.replace("<|action_fast|><|im_end|>\n", "") for t in text] + + return text def preprocesser_call( @@ -125,6 +463,10 @@ def preprocesser_call( truncation: Optional[bool] = None, max_length: Optional[int] = None, return_tensors: str = "pt", + norm_state=None, + agent_pos_mask=None, + state_bins: int = 256, + state_drop_prob: float = 0.0, ) -> BatchFeature: """Unified preprocessing function for Wall-X model handling text, image and video inputs. @@ -154,21 +496,37 @@ def preprocesser_call( - video_grid_thw: Video grid dimensions for LLM - labels: Training labels with masking """ - # Process image inputs + # Process image inputs. transformers>=5.2 split image/video processing + # onto distinct callables and no longer accepts ``videos=`` on + # image_processor (or vice versa); older versions accepted ``videos=None``. + # Only pass the kwargs that correspond to present inputs. if images is not None and len(images) > 0: image_inputs = processor.image_processor( - images=images, videos=None, return_tensors=return_tensors + images=images, return_tensors=return_tensors ) image_grid_thw = image_inputs["image_grid_thw"] else: image_inputs = {} image_grid_thw = None - # Process video inputs if videos is not None: - videos_inputs = processor.image_processor( - images=None, videos=videos, return_tensors=return_tensors - ) + # transformers>=5.2 split video processing onto a dedicated + # ``video_processor``; older versions accepted ``videos=`` on the + # image_processor. Prefer the new API when present and fail loud + # otherwise, since passing ``videos=`` to an old-style + # image_processor silently works but feeds through a different + # preprocessing pipeline than the one the checkpoint was trained + # with. + if hasattr(processor, "video_processor"): + videos_inputs = processor.video_processor( + videos=videos, return_tensors=return_tensors + ) + else: + raise RuntimeError( + "processor has no video_processor attribute - this code path " + "requires transformers>=5.2. Upgrade transformers or pass " + "videos=None." + ) video_grid_thw = videos_inputs["video_grid_thw"] else: videos_inputs = {} @@ -178,6 +536,34 @@ def preprocesser_call( if not isinstance(text, list): text = [text] + # Discretize normalized proprioception into the <|propri|> prompt slot. + if norm_state is not None: + norm_state = ( + norm_state.cpu().numpy() + if isinstance(norm_state, torch.Tensor) + else norm_state + ) + discretized = ( + np.digitize(norm_state, bins=np.linspace(-1, 1, state_bins + 1)[:-1]) - 1 + ) + discretized = discretized[:, 0, :] + if agent_pos_mask is not None: + mask = ( + agent_pos_mask[:, 0, :].cpu().numpy().astype(bool) + if isinstance(agent_pos_mask, torch.Tensor) + else agent_pos_mask[:, 0, :].astype(bool) + ) + else: + mask = np.ones(discretized.shape, dtype=bool) + for i in range(len(text)): + if "<|propri|>" not in text[i]: + continue + if state_drop_prob > 0 and random.random() < state_drop_prob: + text[i] = text[i].replace("<|propri|>", "") + else: + state_str = " ".join(map(str, discretized[i, mask[i]])) + text[i] = text[i].replace("<|propri|>", state_str) + # Process image placeholder tokens in text if image_grid_thw is not None: merge_length = processor.image_processor.merge_size**2 @@ -186,10 +572,11 @@ def preprocesser_call( while "<|image_pad|>" in text[i]: # Add bounds checking to avoid index overflow if index >= len(image_grid_thw): - print( - f"Warning: Number of image placeholders ({index + 1}) " - f"exceeds actual images ({len(image_grid_thw)}), " - f"skipping remaining placeholder processing" + logger.warning( + "Number of image placeholders (%s) exceeds actual images " + "(%s); skipping remaining placeholder processing", + index + 1, + len(image_grid_thw), ) break # Replace image placeholder with actual token count @@ -284,491 +671,3 @@ def preprocesser_call( text_inputs["labels"] = None return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}) - - -def process_grounding_points( - text: str, - orig_height: int, - orig_width: int, - resized_height: int, - resized_width: int, - model_type: str, -) -> str: - """Process grounding point coordinates in text based on image resizing. - - Adjusts coordinate values in tags to match resized image dimensions - for different model types (qwen2, qwen2_5). - - Args: - text: Input text containing tags with coordinates - orig_height: Original image height - orig_width: Original image width - resized_height: Resized image height - resized_width: Resized image width - model_type: Model type for coordinate processing ('qwen2' or 'qwen2_5') - - Returns: - Text with adjusted coordinate values - """ - # Regex pattern to match tags and their contents - point_pattern = re.compile(r"(.*?)") - - def process_match(match): - """Process a single point match and adjust coordinates.""" - coords_str = match.group(1) - try: - # Extract coordinates from string - coords = list(map(int, re.findall(r"\d+", coords_str))) - - # Calculate resize scale factors - scale_w = resized_width / orig_width - scale_h = resized_height / orig_height - - if len(coords) == 2: - x, y = coords - if model_type == "qwen2_5": - # Qwen2.5 uses pixel coordinates - new_x = max(0, min(round(x * scale_w), resized_width - 1)) - new_y = max(0, min(round(y * scale_h), resized_height - 1)) - elif model_type == "qwen2": - # Qwen2 normalizes to [0, 1000) range - new_x = max(0, min(999.999, (x / orig_width) * 1000)) - new_y = max(0, min(999.999, (y / orig_height) * 1000)) - else: - raise ValueError(f"Unsupported model type: {model_type}") - coords = [new_x, new_y] - - elif len(coords) == 4: - x1, y1, x2, y2 = coords - if model_type == "qwen2_5": - new_x1 = max(0, min(round(x1 * scale_w), resized_width - 1)) - new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1)) - new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1)) - new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1)) - elif model_type == "qwen2": - new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000)) - new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000)) - new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000)) - new_y2 = max(0, min(999.999, (y2 / orig_height) * 1000)) - else: - raise ValueError(f"Unsupported model type: {model_type}") - coords = [new_x1, new_y1, new_x2, new_y2] - - # Return processed point tag - return f'[{", ".join(map(str, coords))}]' - - except (ValueError, TypeError): - # Return original content if processing fails - return match.group(0) - - # Replace all matching point tags - processed_text = point_pattern.sub(process_match, text) - return processed_text - - -def get_frame_instruction( - instruction_info: Dict[str, Any], - frame_idx: Optional[int] = None, - truncate_keys: Optional[List[str]] = None, -) -> Tuple[Dict[str, Any], Optional[int]]: - """Extract frame-specific instruction from instruction dictionary. - - Args: - instruction_info: Dictionary containing instruction components - frame_idx: Current frame index - truncate_keys: Keys that trigger truncation when found - - Returns: - Tuple of (frame_instruction_dict, split_end_frame) - """ - if truncate_keys is None: - truncate_keys = [ - "subtask_generation", - "distribute", - "subtask_generation_zh", - "distribute_zh", - ] - - instruction_for_frame = {} - split_end = None - - for key, value in instruction_info.items(): - if isinstance(value, dict): - # Handle frame-range specific instructions - for frame_range, frame_instruction in value.items(): - start_frame, end_frame = map(int, frame_range.split(" ")) - if start_frame <= frame_idx < end_frame or (start_frame == frame_idx): - instruction_for_frame[key] = frame_instruction - if ( - truncate_keys is not None - and split_end is None - and key in truncate_keys - ): - split_end = end_frame + 1 - break - else: - instruction_for_frame[key] = value - - return instruction_for_frame, split_end - - -def get_task_instruction( - frame_instruction_info: Dict[str, Any], priority_order: Optional[OrderedDict] = None -) -> str: - """Construct task instruction from available instruction fields using priority sampling. - - Args: - frame_instruction_info: Dictionary containing instruction fields - priority_order: OrderedDict specifying sampling probability for each field - - Returns: - Combined instruction string with priority components - """ - # Default priority settings - default_priority_order = OrderedDict( - { - "subtask_generation": 0.25, - "subtask_generation_zh": 0.25, - "distribute": 0.25, - "distribute_zh": 0.25, - } - ) - - if priority_order is not None: - priority_order = OrderedDict(priority_order) - else: - priority_order = default_priority_order - - got_instruction = False - task_instruction = "" - - # Sample instruction components based on priority probabilities - for key, prob in priority_order.items(): - if key in frame_instruction_info and frame_instruction_info[key] != "": - if got_instruction: - if random.random() >= prob: - continue - - task_instruction += f"\n{frame_instruction_info[key]}" - got_instruction = True - break - - # Fall back to base instruction if no priority components found - if not got_instruction: - task_instruction = frame_instruction_info.get("instruction", "") - - return task_instruction - - -def get_wallx_normal_text( - instruction_info: Dict[str, Any], - action_chunk_size: int, - frame_idx: int, - priority_order: Optional[OrderedDict] = None, - cam_mapping: Optional[Dict[str, str]] = None, - generate_subtask_ratio: float = 0.0, -) -> Tuple[str, bool]: - """Construct complete multimodal prompt text for Wall-X model. - - Formats input using special tokens including: - - System message - - User observations (with image placeholders) - - Task instructions - - Proprioception prompts - - Assistant responses (with action tokens) - - Args: - instruction_info: Dictionary containing instruction components - action_chunk_size: Number of action tokens to generate - frame_idx: Current frame index - priority_order: Priority order for instruction sampling - cam_mapping: Camera name mapping dictionary - generate_subtask_ratio: Probability of generating subtask instead of actions - - Returns: - Tuple of (formatted_prompt_text, is_subtask_generation) - """ - # Special tokens for formatting - role_start_symbol = "<|im_start|>" - role_end_symbol = "<|im_end|>" - vision_start_symbol = "<|vision_start|>" - vision_end_symbol = "<|vision_end|>" - image_pad_symbol = "<|image_pad|>" - propri_symbol = "<|propri|>" - action_symbol = "<|action|>" - action_fast_symbol = "<|action_fast|>" - - # System prologue - prologue = ( - f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n" - ) - - # User request with observation - user_request = f"{role_start_symbol}user\nObservation:" - if cam_mapping: - for _, cam_name in cam_mapping.items(): - view_name = CAMERA_NAME_MAPPING.get(cam_name, cam_name) - user_request += f" {view_name}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" - user_request += "\nInstruction:" - - # Get frame-specific instruction - frame_instruction_info, _ = get_frame_instruction( - instruction_info, frame_idx=frame_idx - ) - - generate_subtask = False - priority_keys = ["subtask_generation", "distribute"] - - # Decide whether to generate subtask or actions - if ( - bool(set(frame_instruction_info.keys()) & set(priority_keys)) - and random.random() < generate_subtask_ratio - ): - # Generate subtask (equivalent to VQA task) - instruction = frame_instruction_info.get("instruction", "") - text_prompt = "\nPredict the next action in language.\n" - user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" - - # Find output instruction from priority keys - for key in priority_keys: - if key in frame_instruction_info: - output_instruction = frame_instruction_info[key] - break - - assistant_output = ( - f"{role_start_symbol}assistant\n{output_instruction}\n{role_end_symbol}" - ) - generate_subtask = True - else: - # Generate actions - instruction = get_task_instruction( - frame_instruction_info, priority_order=priority_order - ) - text_prompt = f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n" - user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" - assistant_output = f"{role_start_symbol}assistant\n{action_fast_symbol}{role_end_symbol}\n{action_symbol * action_chunk_size}" - - complete_text = prologue + user_message + assistant_output - return complete_text, generate_subtask - - -def get_action_tokens( - normalized_actions: Union[torch.Tensor, List], action_tokenizer -) -> List[List[str]]: - """Convert normalized actions to action token strings. - - Args: - normalized_actions: Normalized action arrays/tensors - action_tokenizer: Tokenizer for converting actions to tokens - - Returns: - List of action token string lists for each sample - """ - if isinstance(normalized_actions, torch.Tensor): - normalized_actions = normalized_actions.cpu().numpy() - - all_action_tokens = [] - for i in range(len(normalized_actions)): - if isinstance(normalized_actions[i], torch.Tensor): - normalized_actions[i] = normalized_actions[i].cpu().numpy() - - token_id = action_tokenizer(normalized_actions[i]) - action_tokens = [f"<|action_token_{j}|>" for j in token_id[0]] - all_action_tokens.append(action_tokens) - - return all_action_tokens - - -def pad_action_token_strs( - actions_token_lists: List[List[str]], pad_token: str = "<|endoftext|>" -) -> List[str]: - """Pad action token lists to same length and join as strings. - - Args: - actions_token_lists: List of action token lists for each sample - pad_token: Token used for padding - - Returns: - List of padded action token strings - """ - max_len = max(len(tokens) for tokens in actions_token_lists) - padded_action_strs = [] - - for tokens in actions_token_lists: - padded_tokens = ( - tokens + ["<|im_end|>\n"] + [pad_token] * (max_len - len(tokens)) - ) - padded_action_strs.append("".join(padded_tokens)) - - return padded_action_strs - - -def replace_action_token( - text: List[str], - norm_action: Optional[torch.Tensor], - action_tokenizer, - dataset_names: List[str], - dof_masks: Optional[torch.Tensor] = None, -) -> List[str]: - """Replace action placeholders in text with actual action tokens. - - Args: - text: List of text strings with action placeholders - norm_action: Normalized action tensors - action_tokenizer: Tokenizer for converting actions to tokens - dataset_names: Names of datasets for each sample - dof_masks: Masks for degrees of freedom - - Returns: - List of text strings with action tokens replaced - """ - # Filter out multimodal dataset names - dataset_names = [ - name for name in dataset_names if name not in MULTIMODAL_DATASET_NAMES - ] - - # Get required action chunk sizes - required_chunk_sizes = [FREQUENCY_MAPPING.get(name, 32) for name in dataset_names] - - if action_tokenizer is not None and norm_action is not None: - # Extract actions based on chunk sizes and DOF masks - norm_action = [ - action[: required_chunk_sizes[i], dof_masks[i, 0].bool()] - for i, action in enumerate(norm_action) - ] - - # Convert to action tokens and pad - actions_fast_tokens = get_action_tokens(norm_action, action_tokenizer) - actions_fast_token_strs = pad_action_token_strs(actions_fast_tokens) - - # Replace action placeholders with actual tokens - actions_fast_token_idx = 0 - for i in range(len(text)): - if "<|action_fast|>" in text[i]: - text[i] = text[i].replace( - "<|action_fast|><|im_end|>\n", - actions_fast_token_strs[actions_fast_token_idx], - ) - actions_fast_token_idx += 1 - - # Remove remaining action placeholders - text = [t.replace("<|action|>", "") for t in text] - else: - # Remove action placeholders when no tokenizer available - text = [t.replace("<|action_fast|><|im_end|>\n", "") for t in text] - - return text - - -@dataclass -class NormStats: - min: torch.Tensor - max: torch.Tensor - delta: torch.Tensor - - -def load_norm_stats(norm_stats_path, dataset_name): - with open(norm_stats_path, "r") as f: - norm_stats = json.load(f) - action_key = KEY_MAPPINGS[dataset_name]["action"] - state_key = KEY_MAPPINGS[dataset_name]["state"] - q01 = torch.tensor(norm_stats["norm_stats"][action_key]["q01"]) - q99 = torch.tensor(norm_stats["norm_stats"][action_key]["q99"]) - delta = q99 - q01 - action_norm_stats = NormStats( - min=q01, - max=q99, - delta=delta, - ) - q01 = torch.tensor(norm_stats["norm_stats"][state_key]["q01"]) - q99 = torch.tensor(norm_stats["norm_stats"][state_key]["q99"]) - delta = q99 - q01 - state_norm_stats = NormStats( - min=q01, - max=q99, - delta=delta, - ) - - return {"action": action_norm_stats, "state": state_norm_stats} - - -def update_action_statistics( - action_statistic_dof: Dict[str, Any], - norm_stats_path: str, - repo_id: str, - dof_config: Dict[str, int] = None, - agent_pos_config: Dict[str, int] = None, - robot_name: str = None, - customized_dof_config: Dict[str, int] = None, - customized_agent_pos_config: Dict[str, int] = None, -) -> None: - """ - Update the action statistics dictionary with new robot configuration. - - Args: - action_statistic_dof (Dict[str, Any]): The dictionary to be updated with statistics - norm_stats_path (str): Path to the normalization statistics file - repo_id (str): Repository ID for the LeRobot configuration - dof_config (Dict[str, int]): Configuration mapping DOF names to their dimensions - agent_pos_config (Dict[str, int]): Configuration mapping agent position names to their dimensions - robot_name (str, optional): Name of the robot. If None, uses repo_id as the key - customized_dof_config (Dict[str, int], optional): Customized DOF configuration for specific robot - customized_agent_pos_config (Dict[str, int], optional): Customized agent position configuration for specific robot - """ - # Load normalization statistics - norm_stats = load_norm_stats(norm_stats_path, repo_id) - - # Extract min and delta values for action and state - action_min = norm_stats["action"].min.numpy().tolist() - action_delta = norm_stats["action"].delta.numpy().tolist() - state_min = norm_stats["state"].min.numpy().tolist() - state_delta = norm_stats["state"].delta.numpy().tolist() - - # Use customized configurations if provided, otherwise use default ones - current_dof_config = ( - customized_dof_config if customized_dof_config is not None else dof_config - ) - current_agent_pos_config = ( - customized_agent_pos_config - if customized_agent_pos_config is not None - else agent_pos_config - ) - - # Prepare keys and values for DOF and agent position configurations - dof_key = [] - agent_pos_key = [] - dof_value = [] - agent_pos_value = [] - stats_dict = {} - - # Extract DOF configuration - for k, v in current_dof_config.items(): - dof_key.append(k) - dof_value.append(v) - - # Extract agent position configuration - for k, v in current_agent_pos_config.items(): - agent_pos_key.append(k) - agent_pos_value.append(v) - - # Calculate DOF indices and extract corresponding min/delta values - dof_idx = np.array([0] + dof_value).cumsum() - for i in range(len(dof_idx) - 1): - stats_dict[dof_key[i]] = { - "min": action_min[dof_idx[i] : dof_idx[i + 1]], - "delta": action_delta[dof_idx[i] : dof_idx[i + 1]], - } - - # Calculate agent position indices and extract corresponding min/delta values - agent_pos_idx = np.array([0] + agent_pos_value).cumsum() - for i in range(len(agent_pos_idx) - 1): - stats_dict[agent_pos_key[i]] = { - "min": state_min[agent_pos_idx[i] : agent_pos_idx[i + 1]], - "delta": state_delta[agent_pos_idx[i] : agent_pos_idx[i + 1]], - } - - # Use provided robot name or repo_id as the key - robot_key = robot_name if robot_name is not None else repo_id - - # Update the action_statistic_dof dictionary - action_statistic_dof.update({robot_key: stats_dict}) diff --git a/wall_x/data/config.py b/wall_x/data/config.py deleted file mode 100644 index 098e58e..0000000 --- a/wall_x/data/config.py +++ /dev/null @@ -1,126 +0,0 @@ -from typing import List, Dict, Optional -from dataclasses import dataclass, field -from qwen_vl_utils.vision_process import MIN_PIXELS, MAX_PIXELS, IMAGE_FACTOR - - -# Tactile sensor file mapping for data processing -TACTILE_FILE_MAPPING = { - "tactile_data_left": "left_tactile", - "tactile_data_right": "right_tactile", -} - -# Supported action datasets -ACTION_DATASET_NAMES = [ - "x2_normal", - "agibotworld_alpha", - "droid", - "fractal", - "bridge_data_v2", - "DobbE", - "RH20T", - "UMI-biarm", - "austin_buds", - "austin_sailor", - "austin_sirius", - "bc_z", - "berkeley_autolab_ur5", - "berkeley_cable_routing", - "berkeley_fanuc_manipulation", - "dlr_edan_shared_control", - "fmb", - "furniture_bench", - "jaco_play", - "nyu_rot", - "stanford_hydra", - "stanford_kuka_multimodal", - "taco_play", - "utaustin_mutex", - "viola", - "physical-intelligence/libero", - "lerobot/aloha_mobile_cabinet", -] - -# Supported multimodal datasets -MULTIMODAL_DATASET_NAMES = [ - "x2_multimodal_from_action", - "x2_multimodal", - "x2_subtask_generation", - "multimodal_CapsFusion", - "multimodal_Robo2VLM", - "multimodal_RoboPoint", - "multimodal_EQA", - "multimodal_Cambrian", - "multimodal_pixmo", - "multimodal_VQAv2", - "multimodal_COCO", -] - - -@dataclass -class X2RDataProcessingConfig: - """Configuration class for X2R data processing pipeline. - - This class contains all the necessary parameters for processing robotic data - including camera mappings, tactile sensor configurations, action predictions, - and various processing options. - """ - - # Action prediction configuration - predict_action_keys: List[str] = field(default_factory=list) - obs_action_keys: List[str] = field(default_factory=list) - - # Image resolution settings for different views - resolution: Dict[str, int] = field( - default_factory=lambda: { - "face_view": -1, - "left_wrist_view": 128, - "right_wrist_view": 128, - } - ) - - # Dataset splitting - train_test_split: float = 0.9 - split_seed: int = 42 - - # Instruction handling - priority_order: Optional[Dict[str, float]] = None - - # Vision model parameters - model_type: str = "qwen2_5" - max_pixels: int = MAX_PIXELS - min_pixels: int = MIN_PIXELS - image_factor: int = IMAGE_FACTOR - - generate_subtask_ratio: float = 0.0 - - def __post_init__(self): - """Post-initialization validation and setup.""" - # Validate train/test split - if not 0 < self.train_test_split < 1: - raise ValueError( - f"train_test_split must be between 0 and 1, got {self.train_test_split}" - ) - - def as_dict(self) -> Dict: - """Convert configuration to dictionary format. - - Returns: - Dict: Configuration as dictionary - """ - return self.__dict__ - - def update(self, **kwargs) -> "X2RDataProcessingConfig": - """Update configuration parameters. - - Args: - **kwargs: Key-value pairs to update - - Returns: - X2RDataProcessingConfig: Updated configuration instance - """ - for key, value in kwargs.items(): - if hasattr(self, key): - setattr(self, key, value) - else: - raise ValueError(f"Unknown configuration parameter: {key}") - return self diff --git a/wall_x/data/load_lerobot_dataset.py b/wall_x/data/load_lerobot_dataset.py deleted file mode 100644 index 6bef686..0000000 --- a/wall_x/data/load_lerobot_dataset.py +++ /dev/null @@ -1,715 +0,0 @@ -""" -LeRobot Dataset Loader - Distributed Version -""" - -import numpy as np -import torch -from torch.utils.data import DistributedSampler, random_split -from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata -from typing import Protocol, SupportsIndex, TypeVar -from qwen_vl_utils.vision_process import smart_resize -from wall_x.data.config import X2RDataProcessingConfig -from wall_x.data.utils import ( - process_grounding_points, - get_wallx_normal_text, - replace_action_token, - preprocesser_call, -) - -from transformers import AutoProcessor -from .utils import KEY_MAPPINGS - -T_co = TypeVar("T_co", covariant=True) - - -# Abstract class for dataset -class Dataset(Protocol[T_co]): - """Interface for a dataset with random access.""" - - def __getitem__(self, index: SupportsIndex) -> T_co: - raise NotImplementedError("Subclasses of Dataset should implement __getitem__.") - - def __len__(self) -> int: - raise NotImplementedError("Subclasses of Dataset should implement __len__.") - - -class PreprocessedDataset(Dataset[T_co]): - def __init__( - self, - dataset, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - seed=42, - rank=0, - world_size=1, - test_only=False, - ): - self.hf_dataset = dataset - - if test_only: - self._dataset = dataset - else: - self._dataset = None - self.train_dataset, self.val_dataset = random_split( - dataset, - [0.95, 0.05], - torch.Generator().manual_seed(seed) if seed is not None else None, - ) - self._train() - - self.seed = seed - self.rank = rank - self.world_size = world_size - - # init configs - self.config = config - self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) - self.dataload_config = dataload_config - self.normalizer_action = (normalizer_action,) - self.normalizer_propri = normalizer_propri - # self.norm_stats = norm_stats - self.lerobot_config = lerobot_config - - self.data_config = X2RDataProcessingConfig().update( - train_test_split=self.dataload_config["train_test_split"], - split_seed=self.dataload_config["split_seed"], - predict_action_keys=self.dataload_config["predict_action_keys"], - obs_action_keys=self.dataload_config["obs_action_keys"], - resolution=self.dataload_config.get("resolution", None), - priority_order=self.dataload_config.get("priority_order", None), - ) - - self._cam_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]["camera"] - self._state_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id] - self._action_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id] - - def _vision_preprocess(self, frames): - processed_frames = [] - for key in self.hf_dataset.meta.camera_keys: - from PIL import Image - - current_obs = frames[key].clone().permute(1, 2, 0) - - img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy()) - orig_width, orig_height = img_pil.size - # 2. Apply resolution constraints (if config is not -1) - target_size = self.data_config.resolution.get( - self._cam_key_mapping[key], -1 - ) - if target_size != -1: - # Maintain aspect ratio logic - if orig_width > orig_height: # Landscape image - new_width = target_size - new_height = int(target_size * orig_height / orig_width) - else: # Portrait image - new_height = target_size - new_width = int(target_size * orig_width / orig_height) - img_pil = img_pil.resize((new_width, new_height)) - - # 3. Apply smart scaling (qwen logic) - current_width, current_height = img_pil.size - resized_height, resized_width = smart_resize( - current_height, - current_width, - factor=self.data_config.image_factor, - min_pixels=self.data_config.min_pixels, - max_pixels=self.data_config.max_pixels, - ) - resized_img = img_pil.resize((resized_width, resized_height)) - processed_frames.append(resized_img) - - return processed_frames, orig_height, orig_width, resized_height, resized_width - - def __getitem__(self, index): - data = self._dataset[index] - image_inputs, h, w, resize_h, resize_w = self._vision_preprocess(data) - agent_pos = data[self._state_key_mapping["state"]] - action = data[self._action_key_mapping["action"]] - frame_index = data["frame_index"] - instruction_info = {"instruction": data["task"]} - generate_subtask_ratio = self.data_config.generate_subtask_ratio - - complete_text, generate_subtask = get_wallx_normal_text( - instruction_info, - self.dataload_config.get("action_horizon", 33) - 1, - frame_index, - self.data_config.priority_order, - self._cam_key_mapping, - generate_subtask_ratio=generate_subtask_ratio, - ) - text = process_grounding_points( - complete_text, h, w, resize_h, resize_w, self.data_config.model_type - ) - result = { - "image_inputs": image_inputs, - "text": text, - "action": action, - "agent_pos": agent_pos, - "frame_index": frame_index, - } - - return result - - def __len__(self) -> int: - return len(self._dataset) - - def _eval(self): - self._dataset = self.val_dataset - - def _train(self): - self._dataset = self.train_dataset - - def get_train_dataloader(self): - """ - Get distributed training dataloader - - Args: - rank: Current process rank - world_size: Total number of processes - seed: Random seed for reproducibility - """ - self._train() - - batch_size = self.config.get("batch_size_per_gpu", 8) - num_workers = self.config.get("num_workers", 4) - - # Create distributed sampler - sampler = DistributedSampler( - self, - num_replicas=self.world_size, - rank=self.rank, - shuffle=True, - seed=self.seed, - drop_last=True, # Ensure all processes have same number of batches - ) - - dataloader = torch.utils.data.DataLoader( - self, - batch_size=batch_size, - sampler=sampler, # Use distributed sampler instead of shuffle=True - num_workers=num_workers, - collate_fn=DataCollator( - self.config, - self.dataload_config, - self.normalizer_action, - self.normalizer_propri, - self.lerobot_config, - ), - pin_memory=True, # Enable for GPU training - persistent_workers=num_workers > 0, # Only if num_workers > 0 - prefetch_factor=2, # Reduce memory usage - drop_last=True, # Avoid incomplete batches - ) - - return dataloader, sampler - - def get_val_dataloader(self): - """ - Get distributed evaluation dataloader (no shuffling for consistent evaluation) - """ - self._eval() - - batch_size = self.config.get( - "eval_batch_size_per_gpu", self.config.get("batch_size_per_gpu", 8) - ) - num_workers = self.config.get("num_workers", 4) - - # Create distributed sampler for evaluation (no shuffle) - sampler = DistributedSampler( - self, - num_replicas=self.world_size, - rank=self.rank, - shuffle=False, # No shuffling for evaluation - drop_last=False, # Keep all samples for evaluation - ) - - dataloader = torch.utils.data.DataLoader( - self, - batch_size=batch_size, - sampler=sampler, - num_workers=num_workers, - collate_fn=DataCollator( - self.config, self.dataload_config, self.norm_stats, self.lerobot_config - ), - pin_memory=True, - persistent_workers=num_workers > 0, - prefetch_factor=2, - drop_last=False, - ) - - return dataloader, sampler - - -class DataCollator: - # Class-level cache for processors to avoid reloading - _processor_cache = {} - _action_tokenizer_cache = {} - - def __init__( - self, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - ): - self.config = config - self.dataload_config = dataload_config - - self.normalizer_action = normalizer_action[0] - self.normalizer_propri = normalizer_propri - self.lerobot_config = lerobot_config - - self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) - self.dataset_name = self.config["data"]["lerobot_config"].get("repo_id", "") - self.dataset_name = [self.dataset_name] * self.config["batch_size_per_gpu"] - self.load_processor() - - def load_processor(self): - processor_path = self.config["pretrained_wallx_path"] - action_tokenizer_path = self.config.get("action_tokenizer_path", None) - - if ( - self.use_fast_tokenizer - and action_tokenizer_path not in self._action_tokenizer_cache - ): - self._action_tokenizer_cache[action_tokenizer_path] = ( - AutoProcessor.from_pretrained( - action_tokenizer_path, trust_remote_code=True - ) - ) - - # Use cached processors if available - if processor_path not in self._processor_cache: - processor = AutoProcessor.from_pretrained(processor_path, use_fast=True) - if self.config.get("padding_side", "left") == "left": - processor.tokenizer.padding_side = "left" - - new_tokens = ["<|propri|>", "<|action|>"] - processor.tokenizer.add_tokens(new_tokens) - if self.use_fast_tokenizer and self.config.get("model_type") == "qwen2_5": - action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path] - new_tokens = [ - f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size) - ] - processor.tokenizer.add_tokens(new_tokens) - begin_idx_token = "<|action_token_0|>" - token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token) - processor.tokenizer.init_kwargs["action_token_start_index"] = token_id - processor.tokenizer.init_kwargs["action_token_vocab_size"] = ( - action_tokenizer.vocab_size - ) - - self._processor_cache[processor_path] = processor - - self.processor = self._processor_cache[processor_path] - - if not self.use_fast_tokenizer: - self.train_action_tokenizer = None - else: - self.train_action_tokenizer = self._action_tokenizer_cache[ - action_tokenizer_path - ] - - @classmethod - def _normalize(cls, action, min_stat, delta): - """ - Normalize action data using min-max normalization. - """ - delta = torch.where(delta == 0, torch.ones_like(delta), delta) - x = (action - min_stat) / delta - x = x * 2 - 1 - x = torch.clamp(x, -1, 1) - return x - - def __call__(self, batch): - additional_inputs = {} - - for key in batch[0].keys(): - if key == "agent_pos": - agent_pos = torch.stack([item["agent_pos"] for item in batch]) - if agent_pos.dim() == 2: - agent_pos = agent_pos.unsqueeze(1) - agent_pos_mask = (~torch.isnan(agent_pos)).float() - # print("agent_pos_mask",agent_pos_mask.shape) - agent_pos.nan_to_num_(nan=0.0) - - # if agent_pos.shape[-1] != 20: - # agent_pos = torch.cat( - # [ - # agent_pos, - # torch.zeros( - # agent_pos.shape[0], - # agent_pos.shape[1], - # 20 - agent_pos.shape[-1], - # ), - # ], - # dim=-1, - # ) - # agent_pos_mask = torch.cat( - # [ - # agent_pos_mask, - # torch.zeros( - # agent_pos_mask.shape[0], - # agent_pos_mask.shape[1], - # 20 - agent_pos_mask.shape[-1], - # ), - # ], - # dim=-1, - # ) - agent_pos = self.normalizer_propri.normalize_data( - agent_pos, self.dataset_name - ) - additional_inputs["proprioception"] = agent_pos - additional_inputs["agent_pos_mask"] = agent_pos_mask - elif key == "action": - action = torch.stack([item["action"] for item in batch]) - if action.dim() == 2: - action = action.unsqueeze(1) - dof_mask = (~torch.isnan(action)).float() - action.nan_to_num_(nan=0.0) - - # if action.shape[-1] != 20: - # action = torch.cat( - # [ - # action, - # torch.zeros( - # action.shape[0], action.shape[1], 20 - action.shape[-1] - # ), - # ], - # dim=-1, - # ) - # dof_mask = torch.cat( - # [ - # dof_mask, - # torch.zeros( - # dof_mask.shape[0], - # dof_mask.shape[1], - # 20 - dof_mask.shape[-1], - # ), - # ], - # dim=-1, - # ) - action = self.normalizer_action.normalize_data( - action, self.dataset_name - ) - additional_inputs["action_chunk"] = action - additional_inputs["dof_mask"] = dof_mask - elif key == "image_inputs": - additional_inputs["image_inputs"] = [ - item["image_inputs"] for item in batch - ] - elif key == "text": - additional_inputs["text"] = [item["text"] for item in batch] - elif key == "frame_index": - additional_inputs["frame_index"] = torch.stack( - [item["frame_index"] for item in batch] - ) - else: - raise NotImplementedError( - f"{key} input not implemented in preprocesser" - ) - - additional_inputs["text"] = replace_action_token( - additional_inputs["text"], - additional_inputs["action_chunk"], - self.train_action_tokenizer if self.use_fast_tokenizer else None, - [self.lerobot_config["repo_id"]] * additional_inputs["text"].__len__(), - additional_inputs["dof_mask"], - ) - - inputs = preprocesser_call( - processor=self.processor, - text=additional_inputs.pop("text"), - images=additional_inputs.pop("image_inputs"), - videos=None, - padding=True, - truncation=True, - return_tensors="pt", - max_length=self.dataload_config.get("max_length", 768), - ) - - action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") - - # Gating token types - additional_inputs["moe_token_types"] = inputs.input_ids == action_token_id - - inputs.update(additional_inputs) - - inputs["dataset_names"] = [self.lerobot_config["repo_id"]] * inputs[ - "action_chunk" - ].shape[0] - - return inputs - - -def load_lerobot_data( - config, - lerobot_config, - normalizer_action, - normalizer_propri, - rank=0, - world_size=1, - seed=42, -): - """ - Load LeRobot dataset with distributed support - - Args: - config: Model configuration - rank: Current process rank (default: 0) - world_size: Total number of processes (default: 1) - seed: Random seed for reproducibility (default: 42) - - Returns: - dataset: Training dataset - train_num: Number of training samples per process - sampler: Distributed sampler (None if world_size=1) - """ - - # Set seed for reproducibility - torch.manual_seed(seed) - - dataload_config = get_data_configs(config["data"]) - - repo_id = lerobot_config.get("repo_id", None) - assert repo_id is not None, "repo id is required" - root = lerobot_config.get("root", None) - meta_info = LeRobotDatasetMetadata(repo_id, root=root) - dataset_fps = meta_info.fps - episodes_num = meta_info.total_episodes - - # norm_stats_path = config.get("norm_stats_path", None) - # assert ( - # norm_stats_path is not None - # ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats" - # norm_stats = load_norm_stats(norm_stats_path, repo_id) - - delta_timestamps = { - # action chunk - KEY_MAPPINGS[repo_id]["action"]: [ - t / dataset_fps - for t in range(dataload_config.get("action_horizon", 33) - 1) - ], - } - batch_size = config.get("batch_size_per_gpu", 8) - episodes = np.arange(episodes_num).tolist() - - train_test_split = dataload_config.get("train_test_split", 0.95) - train_episodes = episodes[: int(episodes_num * train_test_split)] - test_episodes = episodes[int(episodes_num * train_test_split) :] - - train_dataset = LeRobotDataset( - repo_id, - root=root, - episodes=train_episodes, - delta_timestamps=delta_timestamps, - video_backend="pyav", - ) - - if rank == 0: - print(f"Selected train episodes: {train_dataset.episodes}") - print(f"Number of train episodes selected: {train_dataset.num_episodes}") - print(f"Number of train frames selected: {train_dataset.num_frames}") - print(f"Selected test episodes: {test_episodes}") - - dataset = PreprocessedDataset( - train_dataset, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - seed=seed, - rank=rank, - world_size=world_size, - ) - - # Calculate samples per process - if world_size > 1: - # With DistributedSampler, each process gets approximately len(dataset) // world_size samples - samples_per_process = len(dataset) // world_size - train_num = samples_per_process // batch_size - else: - train_num = len(dataset) // batch_size - - if rank == 0: - print("\n" + "=" * 50) - print("LeRobot Data Loading Configuration:") - print(f"✦ RANK: {rank}") - print(f"✦ WORLD SIZE: {world_size}") - print(f"✦ BATCH SIZE PER GPU: {batch_size}") - print(f"✦ REPO ID: {repo_id}") - print(f"✦ TOTAL DATASET SIZE: {len(dataset)}") - if world_size > 1: - print(f"✦ SAMPLES PER PROCESS: {samples_per_process}") - print(f"✦ BATCHES PER PROCESS: {train_num}") - print(f"✦ TOTAL BATCHES (ALL PROCESSES): {train_num * world_size}") - else: - print(f"✦ TOTAL BATCHES: {train_num}") - print(f"✦ SEED: {seed}") - print("=" * 50 + "\n") - - return dataset, train_num - - -def get_distributed_dataloader( - dataset, config, rank=0, world_size=1, seed=42, is_train=True -): - """ - Helper function to get distributed dataloader - - Args: - dataset: PreprocessedDataset instance - config: Configuration dict - rank: Current process rank - world_size: Total number of processes - seed: Random seed - is_train: Whether this is for training (affects shuffling) - - Returns: - dataloader: Distributed DataLoader - sampler: DistributedSampler - """ - if is_train: - return dataset.get_train_dataloader(rank=rank, world_size=world_size, seed=seed) - else: - return dataset.get_val_dataloader(rank=rank, world_size=world_size) - - -def get_data_configs(config): - default_data_config = { - "train_test_split": 0.95, - "split_seed": 42, - "batch_size": 8, - "action_horizon": 21, - "action_history_length": 0, - "image_horizon": 1, - "image_history_length": 0, - "left_padding": False, - "right_padding": False, - "return_first_obs": False, - "return_last_obs": False, - "randomize_obs_after": None, - "datasets": [], - "labeled_pathes": [], - } - data_config = default_data_config | config - data_config["action_horizon"] += 1 - - return data_config - - -class TestDataset(PreprocessedDataset): - def __init__( - self, - dataset, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - seed=42, - ): - super().__init__( - dataset, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - seed=seed, - rank=0, - world_size=1, - test_only=True, - ) - - def get_dataloader(self): - """ - Get distributed evaluation dataloader (no shuffling for consistent evaluation) - """ - - dataloader = torch.utils.data.DataLoader( - self, - batch_size=1, - collate_fn=DataCollator( - self.config, - self.dataload_config, - self.normalizer_action, - self.normalizer_propri, - self.lerobot_config, - ), - ) - - return dataloader - - -def load_test_dataset( - config, - lerobot_config, - normalizer_action, - normalizer_propri, - seed=42, - episode=0, -): - """ - Load test dataset - - Args: - config: Model configuration - seed: Random seed for reproducibility (default: 42) - - Returns: - dataset: Test dataset - """ - - # Set seed for reproducibility - torch.manual_seed(seed) - - repo_id = lerobot_config.get("repo_id", None) - assert repo_id is not None, "repo id is required" - root = lerobot_config.get("root", None) - meta_info = LeRobotDatasetMetadata(repo_id, root=root) - dataset_fps = meta_info.fps - dataload_config = get_data_configs(config["data"]) - - norm_stats_path = config.get("norm_stats_path", None) - assert ( - norm_stats_path is not None - ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats" - # norm_stats = load_norm_stats(norm_stats_path, repo_id) - - delta_timestamps = { - # action chunk - KEY_MAPPINGS[repo_id]["action"]: [ - t / dataset_fps - for t in range(dataload_config.get("action_horizon", 33) - 1) - ], - } - - dataset = LeRobotDataset( - repo_id, - episodes=[episode], - delta_timestamps=delta_timestamps, - video_backend="pyav", - root=root, - ) - - print(f"Selected episodes: {dataset.episodes}") - print(f"Number of episodes selected: {dataset.num_episodes}") - print(f"Number of frames selected: {dataset.num_frames}") - - dataset = TestDataset( - dataset, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - seed=seed, - ) - - return dataset diff --git a/wall_x/fusions/backend.py b/wall_x/fusions/backend.py deleted file mode 100644 index 3f9d2e2..0000000 --- a/wall_x/fusions/backend.py +++ /dev/null @@ -1,433 +0,0 @@ -""" -High-performance C++ backend interface for optimized matrix operations. - -This module provides Python bindings for custom CUDA kernels optimized for -transformer and MoE (Mixture of Experts) operations, including: -- Asymmetric dual expert operations -- Token permutation/unpermutation for MoE routing -- RoPE (Rotary Position Embedding) operations -""" - -import torch -from typing import Tuple, Optional -import wallx_csrc as backend - - -def _allocate_asymmetric_dual_outputs( - input_expert0: torch.Tensor, - input_expert1: torch.Tensor, - weight_expert0: torch.Tensor, - weight_expert1: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Allocate output tensors for asymmetric dual expert GEMM operations. - - This function handles the case where two experts may have different output - dimensions, which is common in heterogeneous MoE architectures. - - Args: - input_expert0 (torch.Tensor): Expert 0 input tensor of shape [m0, k] - input_expert1 (torch.Tensor): Expert 1 input tensor of shape [m1, k] - weight_expert0 (torch.Tensor): Expert 0 weight tensor of shape [k, n0] - weight_expert1 (torch.Tensor): Expert 1 weight tensor of shape [k, n1] - - Returns: - Tuple[torch.Tensor, torch.Tensor]: Pre-allocated output tensors - - output_expert0: Shape [m0, n0] - - output_expert1: Shape [m1, n1] - - Raises: - AssertionError: If tensor dimensions are incompatible - """ - # Validate input tensor dimensions - assert input_expert0.ndim == 2, "Expected 2D tensor for input_expert0" - assert input_expert1.ndim == 2, "Expected 2D tensor for input_expert1" - assert weight_expert0.ndim == 2, "Expected 2D tensor for weight_expert0" - assert weight_expert1.ndim == 2, "Expected 2D tensor for weight_expert1" - - # Verify dimension compatibility for matrix multiplication - assert input_expert0.size(1) == weight_expert0.size( - 0 - ), f"Input expert0 K dimension {input_expert0.size(1)} != weight expert0 K dimension {weight_expert0.size(0)}" - assert input_expert1.size(1) == weight_expert1.size( - 0 - ), f"Input expert1 K dimension {input_expert1.size(1)} != weight expert1 K dimension {weight_expert1.size(0)}" - - # Calculate output shapes: [m, k] × [k, n] = [m, n] - m0, n0 = input_expert0.size(0), weight_expert0.size(1) - m1, n1 = input_expert1.size(0), weight_expert1.size(1) - - # Allocate output tensors with matching device and dtype - output_expert0 = torch.empty( - m0, n0, device=input_expert0.device, dtype=input_expert0.dtype - ) - output_expert1 = torch.empty( - m1, n1, device=input_expert1.device, dtype=input_expert1.dtype - ) - - return output_expert0, output_expert1 - - -def asym_dual_gmm_separated( - input_expert0: torch.Tensor, - input_expert1: torch.Tensor, - weight_expert0: torch.Tensor, - weight_expert1: torch.Tensor, - output_expert0: Optional[torch.Tensor] = None, - output_expert1: Optional[torch.Tensor] = None, - trans_a: bool = False, - trans_b: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Asymmetric dual expert grouped GEMM with separated inputs and outputs. - - This is the recommended interface for maximum flexibility and performance when - dealing with two experts that may have different intermediate dimensions. - The operation is equivalent to: - output_expert0 = input_expert0 @ weight_expert0 - output_expert1 = input_expert1 @ weight_expert1 - But optimized as a single fused kernel call. - - Args: - input_expert0 (torch.Tensor): Expert 0 input tensor of shape [m0, k] - input_expert1 (torch.Tensor): Expert 1 input tensor of shape [m1, k] - weight_expert0 (torch.Tensor): Expert 0 weight tensor of shape [k, n0] - weight_expert1 (torch.Tensor): Expert 1 weight tensor of shape [k, n1] - Note: n0 can be different from n1 - output_expert0 (torch.Tensor, optional): Pre-allocated output for expert 0 [m0, n0] - output_expert1 (torch.Tensor, optional): Pre-allocated output for expert 1 [m1, n1] - trans_a (bool, optional): Whether to transpose input tensors. Defaults to False. - trans_b (bool, optional): Whether to transpose weight tensors. Defaults to False. - - Returns: - Tuple[torch.Tensor, torch.Tensor]: Output tensors (output_expert0, output_expert1) - - Example: - >>> # Two experts with different output dimensions - >>> input0 = torch.randn(512, 1024, device='cuda') # 512 tokens for expert 0 - >>> input1 = torch.randn(256, 1024, device='cuda') # 256 tokens for expert 1 - >>> weight0 = torch.randn(1024, 2048, device='cuda') # Expert 0: 1024->2048 - >>> weight1 = torch.randn(1024, 4096, device='cuda') # Expert 1: 1024->4096 - >>> out0, out1 = asym_dual_gmm_separated(input0, input1, weight0, weight1) - """ - # Allocate outputs if not provided - if output_expert0 is None or output_expert1 is None: - alloc_out0, alloc_out1 = _allocate_asymmetric_dual_outputs( - input_expert0, input_expert1, weight_expert0, weight_expert1 - ) - if output_expert0 is None: - output_expert0 = alloc_out0 - if output_expert1 is None: - output_expert1 = alloc_out1 - - # Call optimized C++ backend kernel - backend.asym_dual_gmm( - input_expert0, - input_expert1, - weight_expert0, - weight_expert1, - output_expert0, - output_expert1, - trans_a, - trans_b, - ) - - return output_expert0, output_expert1 - - -def permute( - input: torch.Tensor, - indices: torch.Tensor, - num_out_tokens: int, - workspace: torch.Tensor, - max_expanded_token_num: int, -) -> torch.Tensor: - """ - Permute input tokens according to expert assignment indices for MoE routing. - - This function reorders tokens based on their assigned experts to enable - efficient grouped processing. Used in the forward pass of MoE layers. - - Args: - input (torch.Tensor): Input tokens to permute - indices (torch.Tensor): Expert assignment indices for each token - num_out_tokens (int): Number of output tokens after expansion - workspace (torch.Tensor): Temporary workspace tensor for intermediate computations - max_expanded_token_num (int): Maximum number of tokens after top-k expansion - - Returns: - torch.Tensor: Permuted tokens grouped by expert assignment - - Note: - This is typically used with top-k expert selection where each token - can be routed to multiple experts. - """ - return backend.permute( - input, indices, num_out_tokens, workspace, max_expanded_token_num - ) - - -def unpermute( - input: torch.Tensor, - row_id_map: torch.Tensor, - prob: torch.Tensor, - max_tokens: int, - num_topK: int, -) -> torch.Tensor: - """ - Unpermute expert outputs back to original token order with probability weighting. - - This function reverses the permutation applied in the forward pass and combines - outputs from multiple experts using their routing probabilities. - - Args: - input (torch.Tensor): Permuted expert outputs to unpermute - row_id_map (torch.Tensor): Mapping from permuted positions to original positions - prob (torch.Tensor): Expert routing probabilities for weighted combination - max_tokens (int): Maximum number of tokens in the sequence - num_topK (int): Number of top experts selected per token - - Returns: - torch.Tensor: Unpermuted tokens in original order with expert outputs combined - - Note: - The output combines multiple expert predictions for each token using - the routing probabilities as weights. - """ - return backend.unpermute(input, row_id_map, prob, max_tokens, num_topK) - - -def unpermute_bwd( - input_bwd: torch.Tensor, - input_fwd: torch.Tensor, - row_id_map: torch.Tensor, - prob: Optional[torch.Tensor], -) -> torch.Tensor: - """ - Backward pass for unpermute operation with gradient flow. - - This function handles the backward pass through the unpermute operation, - ensuring proper gradient flow for training MoE models. - - Args: - input_bwd (torch.Tensor): Backward gradients from the next layer - input_fwd (torch.Tensor): Forward pass inputs (for gradient computation) - row_id_map (torch.Tensor): Row mapping used in forward unpermute - prob (torch.Tensor, optional): Expert probabilities. If None, uniform weights are used. - - Returns: - torch.Tensor: Gradients with respect to the input of unpermute forward pass - - Note: - If prob is None, uniform probabilities are assumed for gradient computation. - """ - # Handle case where probabilities are not provided - if prob is None: - prob = torch.ones( - [input_bwd.size(0), 1], dtype=torch.float32, device=input_bwd.device - ) - - return backend.unpermute_bwd(input_bwd, input_fwd, row_id_map, prob) - - -def rope( - q: torch.Tensor, - k: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - q_out: torch.Tensor, - k_out: torch.Tensor, - mrope_section_doubled: bool, -) -> None: - """ - Apply RoPE (Rotary Position Embedding) to query and key tensors. - - Applies rotary position embeddings to query and key tensors using precomputed - cosine and sine values. Supports both standard RoPE and multi-dimensional RoPE (mRoPE). - - Args: - q (torch.Tensor): Query tensor to apply RoPE to - k (torch.Tensor): Key tensor to apply RoPE to - cos (torch.Tensor): Precomputed cosine values for rotation - sin (torch.Tensor): Precomputed sine values for rotation - q_out (torch.Tensor): Output tensor for rotated queries (in-place operation supported) - k_out (torch.Tensor): Output tensor for rotated keys (in-place operation supported) - mrope_section_doubled (bool): Whether using multi-dimensional RoPE with doubled sections - - Note: - This function performs in-place operations if q_out and k_out point to the same - memory as q and k respectively. The rotation is applied using the standard - RoPE formulation with complex number rotation. - """ - return backend.rope(q, k, cos, sin, q_out, k_out, mrope_section_doubled) - - -def rope_bwd( - grad_q_out: torch.Tensor, - grad_k_out: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - grad_q: torch.Tensor, - grad_k: torch.Tensor, - mrope_section_doubled: bool, -) -> None: - """ - Backward pass for RoPE operation with gradient computation. - - Computes gradients with respect to the input query and key tensors - for the RoPE operation used in transformer attention mechanisms. - - Args: - grad_q_out (torch.Tensor): Gradient with respect to output queries - grad_k_out (torch.Tensor): Gradient with respect to output keys - q (torch.Tensor): Original query tensor from forward pass - k (torch.Tensor): Original key tensor from forward pass - cos (torch.Tensor): Cosine values used in forward pass - sin (torch.Tensor): Sine values used in forward pass - grad_q (torch.Tensor): Output tensor for query gradients - grad_k (torch.Tensor): Output tensor for key gradients - mrope_section_doubled (bool): Whether using multi-dimensional RoPE configuration - - Note: - This function computes the analytical gradient of the RoPE operation, - which involves the inverse rotation compared to the forward pass. - """ - return backend.rope_bwd( - grad_q_out, grad_k_out, q, k, cos, sin, grad_q, grad_k, mrope_section_doubled - ) - - -def get_rope_index( - input_ids: torch.Tensor, - image_grid_thw: Optional[torch.Tensor], - video_grid_thw: Optional[torch.Tensor], - second_per_grid_ts: Optional[torch.Tensor], - attention_mask: Optional[torch.Tensor], - spatial_merge_size: int, - image_token_id: int, - video_token_id: int, - vision_start_token_id: int, - tokens_per_second: float, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Generate position indices for multimodal RoPE (Rotary Position Embedding). - - This function computes 3D position indices for text, image, and video tokens - to enable proper spatial-temporal position encoding in multimodal transformers. - - Args: - input_ids (torch.Tensor): Input token IDs of shape [batch_size, seq_len] - image_grid_thw (torch.Tensor, optional): Image grid specifications of shape [num_images, 3] (T, H, W) - video_grid_thw (torch.Tensor, optional): Video grid specifications of shape [num_videos, 3] (T, H, W) - second_per_grid_ts (torch.Tensor, optional): Temporal scaling per video grid of shape [num_videos] - attention_mask (torch.Tensor, optional): Attention mask of shape [batch_size, seq_len] - spatial_merge_size (int): Spatial dimension merge factor for patch grouping - image_token_id (int): Token ID representing image patches - video_token_id (int): Token ID representing video frames - vision_start_token_id (int): Token ID marking vision sequence start - tokens_per_second (float): Temporal scaling factor for video sequences - - Returns: - Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - - position_ids: 3D position indices of shape [3, batch_size, seq_len] - - mrope_deltas: Position deltas for multimodal RoPE of shape [batch_size, 1] - - Note: - When both image_grid_thw and video_grid_thw are None, returns standard - text-only position indices based on attention_mask or sequence order. - """ - return backend.rope_index( - input_ids, - image_grid_thw, - video_grid_thw, - second_per_grid_ts, - attention_mask, - spatial_merge_size, - image_token_id, - video_token_id, - vision_start_token_id, - tokens_per_second, - ) - - -def rot_pos_emb( - inv_freq: torch.Tensor, - grid_thw: torch.Tensor, - spatial_merge_size: int, -) -> torch.Tensor: - """ - Compute fused rotary position embeddings for multimodal grids. - - This function efficiently computes rotary position embeddings for spatial-temporal - grids using a fused CUDA kernel, supporting both int32 and int64 grid specifications. - - Args: - inv_freq (torch.Tensor): Inverse frequencies for RoPE of shape [dim/2] - Must be float32 dtype on CUDA device - grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W) - Supports int32 or int64 dtype on CUDA device - spatial_merge_size (int): Merge factor for spatial dimensions (must be positive) - - Returns: - torch.Tensor: Computed rotary embeddings of shape [total_tokens, dim] - where total_tokens is determined by grid layouts and spatial_merge_size - - Example: - >>> inv_freq = torch.randn(64, device='cuda', dtype=torch.float32) # 128-dim model - >>> grids = torch.tensor([[8, 14, 14], [16, 7, 7]], device='cuda', dtype=torch.int32) - >>> embeddings = rot_pos_emb(inv_freq, grids, spatial_merge_size=2) - >>> print(embeddings.shape) # [computed_tokens, 128] - - Note: - The function automatically dispatches to int32 or int64 implementations - based on the dtype of grid_thw. Output is always float32. - """ - return backend.rot_pos_emb(inv_freq, grid_thw, spatial_merge_size) - - -def get_window_index( - grid_thw: torch.Tensor, - spatial_merge_size: int, - vit_merger_window_size: int, - patch_size: int, - spatial_merge_unit: int, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Generate window attention indices for Vision Transformer architectures. - - Computes window-based attention indices for hierarchical processing of vision - tokens, enabling efficient sliding window attention patterns in ViT models. - - Args: - grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W) - Must be int32 dtype on CUDA device - spatial_merge_size (int): Spatial dimension merge factor - vit_merger_window_size (int): Size of attention windows for ViT processing - patch_size (int): Size of vision patches in pixels - spatial_merge_unit (int): Unit size for spatial merging operations - - Returns: - Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - - window_indices: Flattened window indices of shape [total_elements] - - cu_window_seqlens: Cumulative window sequence lengths of shape [num_windows + 1] - - Example: - >>> grids = torch.tensor([[1, 14, 14]], device='cuda', dtype=torch.int32) - >>> indices, seqlens = get_window_index( - ... grids, spatial_merge_size=2, vit_merger_window_size=7, - ... patch_size=16, spatial_merge_unit=4 - ... ) - - Note: - Returns empty tensors if input grid is empty or no valid windows can be formed. - The cu_window_seqlens tensor enables efficient batched attention computation. - """ - return backend.get_window_index( - grid_thw, - spatial_merge_size, - vit_merger_window_size, - patch_size, - spatial_merge_unit, - ) diff --git a/wall_x/fusions/ops.py b/wall_x/fusions/ops.py deleted file mode 100644 index 3f993b8..0000000 --- a/wall_x/fusions/ops.py +++ /dev/null @@ -1,742 +0,0 @@ -import torch -import warnings -from wall_x.fusions import backend - - -class AsymmetricDualExpertGemm(torch.autograd.Function): - @staticmethod - def forward( - ctx, input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False - ): - """ - Forward pass for asymmetric dual expert GEMM. - - Args: - input_expert0: Expert 0 input [m0, k] - input_expert1: Expert 1 input [m1, k] - weight_expert0: Expert 0 weight [k, n0] or [n0, k] if trans_b=True - weight_expert1: Expert 1 weight [k, n1] or [n1, k] if trans_b=True - trans_b: Whether to transpose the weight matrices - - Returns: - Tuple of (output_expert0, output_expert1) - """ - # Validate inputs - assert input_expert0.dim() == 2, "input_expert0 must be 2D" - assert input_expert1.dim() == 2, "input_expert1 must be 2D" - assert weight_expert0.dim() == 2, "weight_expert0 must be 2D" - assert weight_expert1.dim() == 2, "weight_expert1 must be 2D" - - # Dimension validation depends on trans_b - if trans_b: - assert input_expert0.size(1) == weight_expert0.size( - 1 - ), "Expert 0 dimension mismatch (trans_b=True)" - assert input_expert1.size(1) == weight_expert1.size( - 1 - ), "Expert 1 dimension mismatch (trans_b=True)" - else: - assert input_expert0.size(1) == weight_expert0.size( - 0 - ), "Expert 0 dimension mismatch (trans_b=False)" - assert input_expert1.size(1) == weight_expert1.size( - 0 - ), "Expert 1 dimension mismatch (trans_b=False)" - - # Save tensors and trans_b for backward pass - ctx.save_for_backward( - input_expert0, input_expert1, weight_expert0, weight_expert1 - ) - ctx.trans_b = trans_b - - # Allocate output tensors - m0 = input_expert0.size(0) - m1 = input_expert1.size(0) - n0 = weight_expert0.size(0) if trans_b else weight_expert0.size(1) - n1 = weight_expert1.size(0) if trans_b else weight_expert1.size(1) - - output_expert0 = torch.empty( - m0, n0, device=input_expert0.device, dtype=input_expert0.dtype - ) - output_expert1 = torch.empty( - m1, n1, device=input_expert1.device, dtype=input_expert1.dtype - ) - - # Call the backend C++ function - backend.asym_dual_gmm_separated( - input_expert0, - input_expert1, - weight_expert0, - weight_expert1, - output_expert0, - output_expert1, - trans_b=trans_b, - ) - - return output_expert0, output_expert1 - - @staticmethod - def backward(ctx, grad_output_expert0, grad_output_expert1): - """ - Optimized backward pass using specialized kernels. - Always computes all gradients to minimize kernel calls. - """ - grad_output_expert0 = grad_output_expert0.contiguous() - grad_output_expert1 = grad_output_expert1.contiguous() - - input_expert0, input_expert1, weight_expert0, weight_expert1 = ctx.saved_tensors - trans_b = ctx.trans_b - - # Always allocate all gradient tensors (no conditional computation) - grad_input_expert0 = torch.empty_like(input_expert0) - grad_input_expert1 = torch.empty_like(input_expert1) - grad_weight_expert0 = torch.empty_like(weight_expert0) - grad_weight_expert1 = torch.empty_like(weight_expert1) - - # Compute input gradients: grad_input = grad_output @ weight^T (if trans_b=False) - # = grad_output @ weight (if trans_b=True) - backend.asym_dual_gmm_separated( - grad_output_expert0, - grad_output_expert1, - weight_expert0, - weight_expert1, - grad_input_expert0, - grad_input_expert1, - trans_a=False, - trans_b=not trans_b, - ) - - # Compute weight gradients - if trans_b: - # When trans_b=True in forward: output = input @ weight^T - # So grad_weight^T = input^T @ grad_output - # Which means grad_weight = grad_output^T @ input - backend.asym_dual_gmm_separated( - grad_output_expert0, - grad_output_expert1, - input_expert0, - input_expert1, - grad_weight_expert0, - grad_weight_expert1, - trans_a=True, - trans_b=False, - ) - else: - # When trans_b=False in forward: output = input @ weight - # So grad_weight = input^T @ grad_output - backend.asym_dual_gmm_separated( - input_expert0, - input_expert1, - grad_output_expert0, - grad_output_expert1, - grad_weight_expert0, - grad_weight_expert1, - trans_a=True, - trans_b=False, - ) - - return ( - grad_input_expert0, - grad_input_expert1, - grad_weight_expert0, - grad_weight_expert1, - None, - ) - - -def asym_dual_gmm( - input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False -): - """ - Convenience function for asymmetric dual expert GEMM. - - Args: - input_expert0: Expert 0 input [m0, k] - input_expert1: Expert 1 input [m1, k] - weight_expert0: Expert 0 weight [k, n0] or [n0, k] if trans_b=True - weight_expert1: Expert 1 weight [k, n1] or [n1, k] if trans_b=True - trans_b: Whether to transpose the weight matrices - - Returns: - Tuple of (output_expert0, output_expert1) - """ - return AsymmetricDualExpertGemm.apply( - input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b - ) - - -################################################################################################ -## -## PermuteMoE topK -## -################################################################################################ - - -class PermuteMoE_topK(torch.autograd.Function): - - workspace_fw = None - dtype = None - max_expanded_token_num = 0 - - @staticmethod - def forward( - ctx, - input_act: torch.Tensor, - indices: torch.Tensor, - num_out_tokens: int, - max_token_num: int, - ): - """ - indices: for topK=1, indices in a 1-d tensor of shape [num_tokens], - otherwise, it's a 2-d tensor of shape [num_tokens, topK] - """ - # Empty input check - if not input_act.numel(): - return input_act, None - - # For top1 case, view the indices as 2D tensor to unify the shape for topk>=2 cases. - if indices.dim() == 1: - indices = indices.view(-1, 1) - - # Device check - if input_act.is_cpu: - raise RuntimeError( - "[Error] The input `input_act` of permute_topK op is on the device: CPU!" - ) - if indices.is_cpu: - warnings.warn( - "The input `indices` of permute_topK op is on the device: CPU!" - ) - - # Shape check - if input_act.size(0) != indices.size(0): - raise RuntimeError( - f"[Error] permute_topK op input `indices` shape mismatch! " - f"Expect {input_act.size(0)}, but got {indices.size(0)}." - ) - - # Data type check - if indices.dtype != torch.int32: - warnings.warn( - f"The data type of the input `indices` of permute_topK op is {indices.dtype}! " - "The recommended type is torch.int32." - ) - indices = indices.to(torch.int32) - - # Contiguous check - if not input_act.is_contiguous(): - warnings.warn("The input `input_act` of permute_topK op is discontiguous!") - input_act = input_act.contiguous() - if not indices.is_contiguous(): - warnings.warn("The input `indices` of permute_topK op is discontiguous!") - indices = indices.contiguous() - - num_topK = indices.size(1) - - input_max_expanded_token_num = max(max_token_num, input_act.size(0)) * num_topK - if PermuteMoE_topK.max_expanded_token_num < input_max_expanded_token_num: - PermuteMoE_topK.max_expanded_token_num = input_max_expanded_token_num - PermuteMoE_topK.workspace_fw = [] - - if PermuteMoE_topK.dtype != input_act.dtype: - PermuteMoE_topK.dtype = input_act.dtype - PermuteMoE_topK.workspace_fw = [] - - permuted_act, row_id_map, PermuteMoE_topK.workspace_fw = backend.permute( - input_act, - indices, - num_out_tokens, - PermuteMoE_topK.workspace_fw, - PermuteMoE_topK.max_expanded_token_num, - ) - - ctx.row_id_map = row_id_map - ctx.num_tokens = indices.size(0) - ctx.num_topK = num_topK - return permuted_act, row_id_map - - @staticmethod - def backward(ctx, permuted_act_grad, _): - # Empty input check - if not permuted_act_grad.numel(): - return permuted_act_grad, None, None, None - - if not permuted_act_grad.is_contiguous(): - permuted_act_grad = permuted_act_grad.contiguous() - - row_id_map = ctx.row_id_map - num_tokens = ctx.num_tokens - num_topK = ctx.num_topK - - unpermuted_act_grad = backend.unpermute( - permuted_act_grad, row_id_map, torch.tensor([]), num_tokens, num_topK - ) - return unpermuted_act_grad, None, None, None - - -################################################################################################ -## -## UnpermuteMoE topK -## -################################################################################################ - - -class UnpermuteMoE_topK(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - input_act: torch.Tensor, - row_id_map: torch.Tensor, - probs: torch.Tensor = None, - ): - # Empty input check - if not input_act.numel(): - ctx.probs = probs - return input_act - - # Device check - if input_act.is_cpu: - raise RuntimeError( - "[Error] The input `input_act` of unpermute_topK op is on the device: CPU!" - ) - if row_id_map.is_cpu: - warnings.warn( - "The input `row_id_map` of unpermute_topK op is on the device: CPU!" - ) - row_id_map = row_id_map.cuda() - if probs is not None and probs.is_cpu: - warnings.warn( - "The input `probs` of unpermute_topK op is on the device: CPU!" - ) - probs = probs.cuda() - - # Shape check - if probs is not None and row_id_map.size(0) != probs.size(0) * probs.size(1): - raise RuntimeError( - f"[Error] unpermute_topK op input `probs` shape mismatch! " - f"Expect {row_id_map.size(0)}, but got {probs.size(0) * probs.size(1)}." - ) - - # Data type check - if row_id_map.dtype != torch.int32: - warnings.warn( - f"The data type of the input `row_id_map` of unpermute_topK op is {row_id_map.dtype}! " - "The recommended type is torch.int32." - ) - row_id_map = row_id_map.to(torch.int32) - if probs is not None and probs.dtype != torch.float32: - warnings.warn( - f"The data type of the input `probs` of unpermute_topK op is {probs.dtype}! " - "The recommended type is torch.float32." - ) - probs = probs.to(torch.float32) - - # Contiguous check - if not input_act.is_contiguous(): - warnings.warn( - "The input `input_act` of unpermute_topK op is discontiguous!" - ) - input_act = input_act.contiguous() - if not row_id_map.is_contiguous(): - warnings.warn( - "The input `row_id_map` of unpermute_topK op is discontiguous!" - ) - row_id_map = row_id_map.contiguous() - if probs is not None and not probs.is_contiguous(): - warnings.warn("The input `probs` of unpermute_topK op is discontiguous!") - probs = probs.contiguous() - - num_tokens = probs.size(0) if probs is not None else input_act.size(0) - num_topK = probs.size(1) if probs is not None else 1 - - unpermuted_output = backend.unpermute( - input_act, - row_id_map, - probs if probs is not None else torch.tensor([]), - num_tokens, - num_topK, - ) - - ctx.save_for_backward(input_act, row_id_map, probs) - return unpermuted_output - - @staticmethod - def backward(ctx, unpermuted_act_grad): - # Empty input check - if not unpermuted_act_grad.numel(): - return unpermuted_act_grad, None, ctx.probs - - if not unpermuted_act_grad.is_contiguous(): - unpermuted_act_grad = unpermuted_act_grad.contiguous() - - input_act, row_id_map, probs = ctx.saved_tensors - - act_grad = None - if ctx.needs_input_grad[0]: - act_grad, prob_grad = backend.unpermute_bwd( - unpermuted_act_grad, input_act, row_id_map, probs - ) - - if not ctx.needs_input_grad[2]: - prob_grad = None - return act_grad, None, prob_grad - - -def permute(input_act, indices, num_out_tokens=None, max_token_num=0): - num_out_tokens = 0 if num_out_tokens is None else num_out_tokens - return PermuteMoE_topK.apply(input_act, indices, num_out_tokens, max_token_num) - - -def unpermute(input_act, row_id_map, probs=None): - return UnpermuteMoE_topK.apply(input_act, row_id_map, probs) - - -################################################################################################ -## -## mutlimodal RoPE -## -################################################################################################ - - -class MultimodalRoPE(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - mrope_section: list, - ): - - # Device check - if q.is_cpu: - raise RuntimeError( - "[Error] The input `q` of multimodal_rope op is on the device: CPU!" - ) - if k.is_cpu: - raise RuntimeError( - "[Error] The input `k` of multimodal_rope op is on the device: CPU!" - ) - if cos.is_cpu: - raise RuntimeError( - "[Error] The input `cos` of multimodal_rope op is on the device: CPU!" - ) - if sin.is_cpu: - raise RuntimeError( - "[Error] The input `sin` of multimodal_rope op is on the device: CPU!" - ) - if len(mrope_section) != 3: - raise RuntimeError( - "[Error] The input `mrope_section` of multimodal_rope op must be a list of 3 integers!" - ) - - # Contiguous check - if not q.is_contiguous(): - warnings.warn("The input `q` of multimodal_rope op is discontiguous!") - q = q.contiguous() - if not k.is_contiguous(): - warnings.warn("The input `k` of multimodal_rope op is discontiguous!") - k = k.contiguous() - if not cos.is_contiguous(): - warnings.warn("The input `cos` of multimodal_rope op is discontiguous!") - cos = cos.contiguous() - if not sin.is_contiguous(): - warnings.warn("The input `sin` of multimodal_rope op is discontiguous!") - sin = sin.contiguous() - - # Prepare mrope_section_doubled - mrope_section_doubled = [x * 2 for x in mrope_section] - - # Create output tensors - q_out = torch.empty_like(q) - k_out = torch.empty_like(k) - backend.rope(q, k, cos, sin, q_out, k_out, mrope_section_doubled) - - ctx.save_for_backward(q, k, cos, sin) - ctx.mrope_section_doubled = mrope_section_doubled - - return q_out, k_out - - @staticmethod - def backward(ctx, grad_q_out, grad_k_out): - - if not grad_q_out.is_contiguous(): - grad_q_out = grad_q_out.contiguous() - if not grad_k_out.is_contiguous(): - grad_k_out = grad_k_out.contiguous() - - q, k, cos, sin = ctx.saved_tensors - - grad_q = None - grad_k = None - if ctx.needs_input_grad[0]: - grad_q = torch.empty_like(q) - if ctx.needs_input_grad[1]: - grad_k = torch.empty_like(k) - - if grad_q is not None or grad_k is not None: - backend.rope_bwd( - grad_q_out, - grad_k_out, - q, - k, - cos, - sin, - grad_q if grad_q is not None else torch.empty_like(q), - grad_k if grad_k is not None else torch.empty_like(k), - ctx.mrope_section_doubled, - ) - - return grad_q, grad_k, None, None, None - - -def multimodal_rope(q, k, cos, sin, mrope_section): - return MultimodalRoPE.apply(q, k, cos, sin, mrope_section) - - -################################################################################################ -## -## RoPE Index 3D -## -################################################################################################ - - -def get_rope_index( - input_ids: torch.Tensor, - spatial_merge_size: int, - image_token_id: int, - video_token_id: int, - vision_start_token_id: int, - tokens_per_second: float, - image_grid_thw: torch.Tensor = None, - video_grid_thw: torch.Tensor = None, - second_per_grid_ts: torch.Tensor = None, - attention_mask: torch.Tensor = None, -): - """ - Generate 3D RoPE position indices for multimodal transformer inputs. - - Computes position indices for text, image, and video tokens to enable proper - spatial-temporal position encoding in multimodal transformers with RoPE. - - Args: - input_ids (torch.Tensor): Input token sequence of shape [batch_size, seq_len] - Must be LongTensor on CUDA device - spatial_merge_size (int): Spatial merge size for patch grouping (must be positive) - image_token_id (int): Token ID representing image patches - video_token_id (int): Token ID representing video frames - vision_start_token_id (int): Token ID marking start of vision sequences - tokens_per_second (float): Temporal scaling factor for video sequences (must be positive) - image_grid_thw (torch.Tensor, optional): Image grid dimensions of shape [num_images, 3] (T, H, W) - video_grid_thw (torch.Tensor, optional): Video grid dimensions of shape [num_videos, 3] (T, H, W) - second_per_grid_ts (torch.Tensor, optional): Video time intervals of shape [num_videos] - attention_mask (torch.Tensor, optional): Attention mask of shape [batch_size, seq_len] - - Returns: - Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - - position_ids: 3D position indices of shape [3, batch_size, seq_len] - - mrope_position_deltas: mRoPE position deltas of shape [batch_size, 1] - - Raises: - TypeError: If input_ids is not a torch.Tensor - ValueError: If input dimensions are incorrect or tensors not on CUDA - """ - # Input validation - if not isinstance(input_ids, torch.Tensor): - raise TypeError("input_ids must be a torch.Tensor") - - if input_ids.dim() != 2: - raise ValueError("input_ids must be 2D tensor (batch_size, seq_len)") - - if not input_ids.is_cuda: - raise ValueError("input_ids must be on CUDA device") - - # Parameter validation - if not isinstance(spatial_merge_size, int) or spatial_merge_size <= 0: - raise ValueError( - f"spatial_merge_size must be positive integer, got {spatial_merge_size}" - ) - - if not isinstance(tokens_per_second, (int, float)) or tokens_per_second <= 0: - raise ValueError( - f"tokens_per_second must be positive number, got {tokens_per_second}" - ) - - return backend.get_rope_index( - input_ids, - image_grid_thw, - video_grid_thw, - second_per_grid_ts, - attention_mask, - spatial_merge_size, - image_token_id, - video_token_id, - vision_start_token_id, - float(tokens_per_second), - ) - - -################################################################################################ -## -## Fused Rotary Position Embedding -## -################################################################################################ - - -def rot_pos_emb( - inv_freq: torch.Tensor, - grid_thw: torch.Tensor, - spatial_merge_size: int, -) -> torch.Tensor: - """ - Compute fused rotary position embeddings using optimized CUDA kernel. - - This function fuses all rotary position embedding computations into a single - CUDA kernel for improved performance with spatial-temporal grids. - - Args: - inv_freq (torch.Tensor): Inverse frequencies tensor of shape [dim/2] - Contains precomputed 1.0 / (theta ** (torch.arange(0, dim, 2) / dim)) - Must be float32 on CUDA device - grid_thw (torch.Tensor): Grid dimensions tensor of shape [num_grids, 3] - Each row contains (T, H, W) for temporal, height, width dimensions - Supports int32 or int64 on CUDA device - spatial_merge_size (int): Spatial merge size for token grouping (must be positive) - - Returns: - torch.Tensor: Rotary position embeddings of shape [total_tokens, dim] - where dim = 2 * len(inv_freq) - First half contains h_pos frequencies, second half contains w_pos frequencies - - Raises: - TypeError: If inputs are not torch.Tensor or spatial_merge_size not int - ValueError: If tensor dimensions incorrect, not on CUDA, or devices mismatch - RuntimeError: If CUDA kernel execution fails - """ - # Type checking - if not isinstance(inv_freq, torch.Tensor): - raise TypeError(f"inv_freq must be a torch.Tensor, got {type(inv_freq)}") - - if not isinstance(grid_thw, torch.Tensor): - raise TypeError(f"grid_thw must be a torch.Tensor, got {type(grid_thw)}") - - # Dimension checking - if inv_freq.dim() != 1: - raise ValueError( - f"inv_freq must be 1-dimensional, got {inv_freq.dim()}D tensor" - ) - - if grid_thw.dim() != 2: - raise ValueError( - f"grid_thw must be 2-dimensional, got {grid_thw.dim()}D tensor" - ) - - if grid_thw.size(1) != 3: - raise ValueError( - f"grid_thw must have shape [num_grids, 3], got shape {list(grid_thw.shape)}" - ) - - # Device checking - if not inv_freq.is_cuda: - raise ValueError("inv_freq must be on CUDA device") - - if not grid_thw.is_cuda: - raise ValueError("grid_thw must be on CUDA device") - - # Ensure both tensors are on the same device - if inv_freq.device != grid_thw.device: - raise ValueError( - f"inv_freq and grid_thw must be on the same device, " - f"got {inv_freq.device} and {grid_thw.device}" - ) - - # Parameter validation - if not isinstance(spatial_merge_size, int): - raise TypeError( - f"spatial_merge_size must be an integer, got {type(spatial_merge_size)}" - ) - - if spatial_merge_size <= 0: - raise ValueError( - f"spatial_merge_size must be positive, got {spatial_merge_size}" - ) - - # Ensure inv_freq is float32 (the kernel expects float) - if inv_freq.dtype != torch.float32: - inv_freq = inv_freq.to(torch.float32) - - # Call the CUDA backend - try: - return backend.rot_pos_emb(inv_freq, grid_thw, spatial_merge_size) - except RuntimeError as e: - raise RuntimeError(f"CUDA kernel execution failed: {str(e)}") - - -################################################################################################ -## -## Fused Window Index Generation -## -################################################################################################ - - -def get_window_index( - grid_thw: torch.Tensor, - window_size: int, - spatial_merge_size: int, - patch_size: int, - spatial_merge_unit: int = 1, -): - """ - Generate window attention indices for Vision Transformer architectures. - - Computes window-based attention indices for hierarchical processing of vision - tokens, enabling efficient sliding window attention patterns in ViT models. - - Args: - grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W) - Must be or will be converted to int32 on CUDA device - window_size (int): Window size for attention computation - spatial_merge_size (int): Spatial merge size for patch grouping - patch_size (int): Size of vision patches in pixels - spatial_merge_unit (int, optional): Spatial merging unit size. Defaults to 1. - - Returns: - Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - - window_index: Window indices tensor of shape [total_elements] - - cu_window_seqlens: Cumulative window sequence lengths of shape [num_windows + 1] - - Raises: - AssertionError: If grid_thw dimensions are incorrect - - Note: - Returns empty tensors if input grid is empty or no valid windows can be formed. - The function automatically converts input to CUDA int32 if needed. - """ - # Input validation - assert ( - grid_thw.dim() == 2 and grid_thw.size(1) == 3 - ), f"grid_thw must have shape (num_grids, 3), got {grid_thw.shape}" - - # Ensure input is on CUDA and int32 type - if not grid_thw.is_cuda: - grid_thw = grid_thw.cuda() - - if grid_thw.dtype != torch.int32: - grid_thw = grid_thw.to(torch.int32) - - # Calculate vit_merger_window_size - vit_merger_window_size = window_size // spatial_merge_size // patch_size - - # Call CUDA backend - window_index, cu_window_seqlens = backend.get_window_index( - grid_thw, - spatial_merge_size, - vit_merger_window_size, - patch_size, - spatial_merge_unit, - ) - - return window_index, cu_window_seqlens diff --git a/wall_x/infer/env.py b/wall_x/infer/env.py deleted file mode 100644 index 35106f7..0000000 --- a/wall_x/infer/env.py +++ /dev/null @@ -1,143 +0,0 @@ -""" -Base Environment Class for Robot Control and Inference -""" - -from typing import Dict, Any, List -from abc import ABC, abstractmethod -import time -from wall_x.infer.infer_config import InferConfig -from wall_x.infer.utils import KeyboardThread -from wall_x.infer.logger import InferLogger - - -class BaseEnv(ABC): - def __init__(self, config: InferConfig): - self.config = config - self.logger = InferLogger.get_env_logger("Env") - - @abstractmethod - def get_observation(self) -> Dict[str, Any]: - raise NotImplementedError - - @abstractmethod - def apply_action(self, input: dict) -> None: - raise NotImplementedError - - @abstractmethod - def get_instruction(self) -> str: - raise NotImplementedError - - def reset(self) -> Dict[str, Any]: - raise NotImplementedError - - def stop(self) -> None: - raise NotImplementedError - - -class RealRobotEnv(BaseEnv): - def __init__( - self, config: InferConfig, instructions: List[str], enable_keyboard: bool = True - ): - """ - Args: - config: Inference configuration - instruction: Task instruction - """ - super().__init__(config) - self.instruction = "test" - self.model = self._register_model() - self.robot = self._register_robot() - - # Keyboard control - self.keyboard_thread = None - if enable_keyboard: - self.keyboard_thread = KeyboardThread() - - # Instruction list - self.instructions = instructions - self.instruction_index = 0 - - # def _register_model(self) -> WallxModelWrapper: - # return WallxModelWrapper(self.config) - - def _register_robot(self): - from wall_x.infer.robot import DesktopRobot, TurtleRobot - - if self.config.robot_type == "desktop": - return DesktopRobot(self.config) - elif self.config.robot_type == "turtle": - return TurtleRobot(self.config) - else: - raise ValueError(f"Invalid robot type: {self.config.robot_type}") - - def get_observation(self): - return self.robot.get_observation() - - def apply_action(self, input: dict): - self.robot.apply_action(input) - - def get_instruction(self) -> str: - """Return task instruction""" - return self.instructions[self.instruction_index] - - def reset(self): - self.robot.go_home() - - def listen_to_keyboard(self): - if self.keyboard_thread is not None: - if self.keyboard_thread.should_stop: - time.sleep(1) - return True - if self.keyboard_thread.should_reset: - self.reset() - self.keyboard_thread.should_reset = False - time.sleep(1) - return True - if self.keyboard_thread.new_instruction_index is not None: - new_index = self.keyboard_thread.new_instruction_index - # Check if index is valid - if 0 <= new_index < len(self.instructions): - self.instruction_index = new_index - self.logger.info( - f"[Keyboard] Instruction index switched to {new_index}: {self.instructions[new_index]}" - ) - else: - self.logger.info( - f"[Keyboard] Invalid instruction index {new_index}, valid range: 0-{len(self.instructions)-1}" - ) - # Reset flag - self.keyboard_thread.new_instruction_index = None - time.sleep(1) - return True - return False - - def run_infer_flow_action(self): - while True: - if self.listen_to_keyboard(): - continue - observation = self.get_observation() - instruction = self.get_instruction() - model_output = self.model.infer_flow_action(observation, instruction) - self.apply_action(model_output) - - def run_infer_flow_action_with_subtask(self, subtask_interval: int = 2): - step = 0 - subtask = "" - while True: - if self.listen_to_keyboard(): - continue - observation = self.get_observation() - instruction = self.get_instruction() - if step == 0 or step % subtask_interval == 0: - subtask = self.model.infer_subtask(observation, instruction) - model_output = self.model.infer_flow_action(observation, subtask) - self.apply_action(model_output) - - def run_infer_ar_action(self): - while True: - if self.listen_to_keyboard(): - continue - observation = self.get_observation() - instruction = self.get_instruction() - model_output = self.model.infer_ar_action(observation, instruction) - self.apply_action(model_output) diff --git a/wall_x/infer/env_libero.py b/wall_x/infer/env_libero.py deleted file mode 100644 index 07fac47..0000000 --- a/wall_x/infer/env_libero.py +++ /dev/null @@ -1,744 +0,0 @@ -import os -import json -import numpy as np -from typing import Dict, Any, Tuple, List -from libero.libero import benchmark -from wall_x.infer.env import BaseEnv, InferConfig -from wall_x.serving.policy.wall_x_policy import WallXPolicy - -from wall_x.infer.base_dataclass import RobotStateActionData -from wall_x.infer.utils_libero import ( - get_libero_env, - get_libero_dummy_action, - get_libero_image, - get_libero_wrist_image, - quat2axisangle, - TaskSuite, - save_rollout_video, -) -from robosuite.wrappers import VisualizationWrapper - - -def _create_libero_env_standalone( - task_id: int, - task_suite_name: str, - model_family: str = "wallx", - resolution: int = 256, - seed: int = 7, -) -> Any: - """ - Standalone function to create a Libero environment, independent of LiberoRobotEnv instance. - Used for creating environments in subprocess during multi-batch inference, - avoiding serialization of large objects containing the model. - - Args: - task_id: Task ID - task_suite_name: Task suite name - model_family: Model family - resolution: Resolution - seed: Random seed - - Returns: - Environment instance - """ - from libero.libero import benchmark - - # Get task suite and task - benchmark_dict = benchmark.get_benchmark_dict() - task_suite = benchmark_dict[task_suite_name]() - task = task_suite.get_task(task_id) - - # Create environment - env, _ = get_libero_env( - task, - model_family=model_family, - resolution=resolution, - seed=seed, - ) - - # Wrap environment - env.env = VisualizationWrapper(env.env) - env.env.set_visualization_setting(setting="grippers", visible=False) - - return env - - -class LiberoRobotEnv(BaseEnv): - def __init__( - self, - config: InferConfig, - task_suite_name: str = TaskSuite.LIBERO_SPATIAL, - initial_states_path: str = "DEFAULT", - rollout_dir: str = "./rollouts", - model_family: str = "wallx", - resolution: int = 256, - seed: int = 7, - ): - - super().__init__(config) - self.logger.info( - f"Initializing LiberoRobotEnv (Stateless), task suite: {task_suite_name}" - ) - - self.model = self._register_model() - self.model_family = model_family - self.resolution = resolution - self.seed = seed - - self.logger.info("Importing Libero and related utils...") - self.RobotStateActionData = RobotStateActionData - - self.rollout_dir = os.path.join(rollout_dir, task_suite_name) - os.makedirs(self.rollout_dir, exist_ok=True) - - if save_rollout_video is not None: - self.save_rollout_video = save_rollout_video - else: - self.save_rollout_video = None - self.logger.warning("save_rollout_video not found, video saving disabled.") - - self.task_suite_name = task_suite_name - benchmark_dict = benchmark.get_benchmark_dict() - self.task_suite = benchmark_dict[self.task_suite_name]() - self.num_tasks = self.task_suite.n_tasks - - self.initial_states_path = initial_states_path - self.all_initial_states = None - if self.initial_states_path != "DEFAULT": - try: - with open(self.initial_states_path, "r") as f: - self.all_initial_states = json.load(f) - self.logger.info( - f"Loaded custom initial states from {self.initial_states_path}" - ) - except Exception as e: - self.logger.error(f"Failed to load initial states file: {e}") - raise - - def _register_model(self) -> WallXPolicy: - - return WallXPolicy( - model_path=self.config.model_path, - train_config=self.config.train_config, - action_tokenizer_path=self.config.action_tokenizer_path, - action_dim=self.config.action_dim, - agent_pos_dim=self.config.action_dim, - pred_horizon=self.config.pred_horizon, - camera_key=self.config.cam_names, - predict_mode=self.config.predict_mode, - ) - - def get_instruction(self, task_desc: str) -> str: - return task_desc - - def get_observation(self, raw_obs: Dict[str, Any]) -> Dict[str, Any]: - if raw_obs is None: - raise ValueError("Raw observation is None") - - data_obj = self.RobotStateActionData(config=self.config) - - pos = raw_obs["robot0_eef_pos"] - rot = quat2axisangle(raw_obs["robot0_eef_quat"]) - grip = raw_obs["robot0_gripper_qpos"][0:1] - - data_obj.save_state_data_with_key(pos[None], "follow_right_ee_cartesian_pos") - data_obj.save_state_data_with_key(rot[None], "follow_right_ee_rotation") - data_obj.save_state_data_with_key(grip[None], "follow_right_gripper") - data_obj.dof_mask = self._get_dof_mask() - - face_view = get_libero_image(raw_obs) - right_wrist_view = get_libero_wrist_image(raw_obs) - - return { - "robot_state_action_data": data_obj, - "face_view": face_view, - "right_wrist_view": right_wrist_view, - } - - def apply_action( - self, input_data: Dict[str, Any], env: Any = None, replay_images: list = None - ) -> bool: - if env is None: - raise ValueError( - "In Stateless mode, apply_action must be called with explicit 'env' parameter" - ) - - action_data = input_data["robot_state_action_data"] - right_arm_traj = self._get_right_arm_action(action_data) - while ( - right_arm_traj is not None - and right_arm_traj.ndim > 2 - and right_arm_traj.shape[0] == 1 - ): - right_arm_traj = right_arm_traj.squeeze(0) - - done = False - t = 0 - - try: - for i in range(len(right_arm_traj)): - if done: - break - - action_7d = right_arm_traj[i] - obs, reward, done, info = env.step(action_7d) - t += 1 - - if obs is not None and replay_images is not None: - replay_images.append(get_libero_image(obs)) - - input_data["_last_obs"] = obs - - except Exception as e: - self.logger.error(f"Env step error: {e}") - return False # Error treated as failure - - return done, t - - def apply_action_batch( - self, - vec_env: Any, - trajectories: List[np.ndarray], - active_indices: List[int], - status_list: List[Dict[str, Any]], - model_outputs: List[Dict[str, Any]], - ) -> None: - """ - Execute action trajectories in parallel batch. - - Uses SubprocVectorEnv to execute actions in parallel for all active environments. - Integrates vec_env.step(batch_actions, id=still_active) in this function. - """ - if not trajectories: - return - - max_traj_len = max(len(traj) for traj in trajectories) - if max_traj_len == 0: - return - - for step_idx in range(max_traj_len): - # Check if there are still active environments - still_active = [ - idx - for idx in active_indices - if (not status_list[idx]["done"]) - and status_list[idx]["count"] > 0 - and step_idx < len(trajectories[active_indices.index(idx)]) - ] - if not still_active: - break - - # Build batch actions (only includes actions for still_active environments) - batch_actions = [] - for idx in still_active: - traj_idx = active_indices.index(idx) - action_7d = trajectories[traj_idx][step_idx] - # Ensure action_7d is numpy array or list - if isinstance(action_7d, np.ndarray): - batch_actions.append(action_7d) - else: - batch_actions.append(np.array(action_7d)) - - # Convert to numpy array with shape (batch_size, action_dim) - batch_actions = np.array(batch_actions) - - # Execute step in parallel (only for still_active environments) - obs_list, reward_list, done_list, info_list = vec_env.step( - batch_actions, id=still_active - ) - - # Process returned results - if obs_list.dtype == object: - obs_list = list(obs_list) - else: - obs_list = [obs_list[i] for i in range(len(obs_list))] - done_list = [bool(done_list[i]) for i in range(len(done_list))] - - # Update each environment's status - for i, idx in enumerate(still_active): - st = status_list[idx] - obs = obs_list[i] - done = done_list[i] - - if obs is not None: - st["current_obs"] = obs - if st["replay_images"] is not None: - st["replay_images"].append(get_libero_image(obs)) - - st["count"] -= 1 - st["success"] = done - st["done"] = done or st["count"] <= 0 - - # Update model_output's _last_obs - model_outputs[active_indices.index(idx)]["_last_obs"] = obs - - def get_task_info(self, task_id: int) -> Tuple[str, Any]: - """ - Get task information (task description and initial states) without creating environment. - - Returns: - Tuple[str, Any]: (task_desc, default_initial_states) - """ - if task_id < 0 or task_id >= self.num_tasks: - raise ValueError(f"Invalid task ID: {task_id}") - - task = self.task_suite.get_task(task_id) - task_desc = task.language - default_initial_states = self.task_suite.get_task_init_states(task_id) - - return task_desc, default_initial_states - - def create_env_for_task(self, task_id: int) -> Tuple[Any, str, Any]: - if task_id < 0 or task_id >= self.num_tasks: - raise ValueError(f"Invalid task ID: {task_id}") - - task = self.task_suite.get_task(task_id) - default_initial_states = self.task_suite.get_task_init_states(task_id) - - env, task_desc = get_libero_env( - task, - model_family=self.model_family, - resolution=self.resolution, - seed=self.seed, - ) - - env.env = VisualizationWrapper(env.env) - env.env.set_visualization_setting(setting="grippers", visible=False) - - return env, task_desc, default_initial_states - - def _get_initial_state_for_episode( - self, task_desc: str, default_states: Any, episode_idx: int - ) -> np.ndarray: - if self.initial_states_path == "DEFAULT": - if default_states is None: - raise ValueError("Default states missing") - return default_states[episode_idx] - else: - if self.all_initial_states is None: - raise ValueError("Custom states not loaded") - initial_states_task_key = task_desc.replace(" ", "_") - episode_key = f"demo_{episode_idx}" - if not self.all_initial_states[initial_states_task_key][episode_key][ - "success" - ]: - raise ValueError(f"Expert demo failed for {episode_key}") - return np.array( - self.all_initial_states[initial_states_task_key][episode_key][ - "initial_state" - ] - ) - - def reset_env( - self, env: Any, task_desc: str, default_states: Any, episode_idx: int - ) -> Any: - try: - if episode_idx >= 0: - state = self._get_initial_state_for_episode( - task_desc, default_states, episode_idx - ) - obs = env.set_init_state(state) - return obs - else: - return env.reset() - except Exception as e: - self.logger.error(f"Reset failed: {e}, falling back to default reset") - return env.reset() - - def _get_dof_mask(self): - dof_config = self.config.train_config["dof_config"] - total_dof = sum(dof_config.values()) - dof_mask = np.ones((1, self.config.action_horizon, total_dof)) - mask_keys = [ - "follow_left_ee_cartesian_pos", - "follow_left_ee_rotation", - "follow_left_gripper", - "head_actions", - "height", - "velocity_decomposed", - ] - start_idx = 0 - for key, dof_size in dof_config.items(): - if key in mask_keys: - dof_mask[:, :, start_idx : start_idx + dof_size] = 0 - start_idx += dof_size - return dof_mask - - def _get_right_arm_action( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - right_ee_cartesian_pos = robot_state_action_data.data[ - "action_right_ee_cartesian_pos" - ] - right_ee_rotation = robot_state_action_data.data["action_right_ee_rotation"] - right_gripper = robot_state_action_data.data["action_right_gripper"] - return np.concatenate( - [right_ee_cartesian_pos, right_ee_rotation, right_gripper], axis=1 - ) - - def _get_left_arm_action( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - left_ee_cartesian_pos = robot_state_action_data.data[ - "action_left_ee_cartesian_pos" - ] - left_ee_rotation = robot_state_action_data.data["action_left_ee_rotation"] - left_gripper = robot_state_action_data.data["action_left_gripper"] - return np.concatenate( - [left_ee_cartesian_pos, left_ee_rotation, left_gripper], axis=1 - ) - - def _save_rollout( - self, - replay_images: List[np.ndarray], - success: bool, - task_id: int, - task_desc: str, - episode_idx: int, - ): - if not self.save_rollout_video or not replay_images: - return - try: - task_name_safe = task_desc.replace(" ", "_").replace(".", "") - filename = f"{episode_idx}{'_SUCCESS' if success else '_FAILURE'}--_{task_name_safe}.mp4" - self.save_rollout_video( - self.rollout_dir, - replay_images, - filename, - success=success, - task_description=task_desc, - log_file=None, - model_family=self.model_family, - ) - self.logger.info(f"Saved video: {filename}") - except Exception as e: - self.logger.error(f"Save video failed: {e}") - - def run_infer_flow_action( - self, - env: Any, - task_id: int, - task_desc: str, - default_initial_states: Any, - episode_idx: int, - max_infer_times: int = 5, - num_steps_wait: int = 10, - ) -> bool: - replay_images = [] - num_steps = 0 - done = False - count = max_infer_times - - current_obs = self.reset_env( - env, task_desc, default_initial_states, episode_idx - ) - if current_obs is None: - return False - - while num_steps < num_steps_wait: - obs, reward, done, info = env.step( - get_libero_dummy_action(self.model_family) - ) - num_steps += 1 - if obs is not None: - current_obs = obs - - while not done and count > 0: - try: - model_input = self.get_observation(current_obs) - instruction = self.get_instruction(task_desc) - model_input["prompt"] = instruction - model_input["dataset_names"] = "libero_all" - - state = np.concatenate( - [ - model_input["robot_state_action_data"].data[ - "state_right_ee_cartesian_pos" - ], - model_input["robot_state_action_data"].data[ - "state_right_ee_rotation" - ], - model_input["robot_state_action_data"].data[ - "state_right_gripper" - ], - ], - axis=-1, - ) - - model_input["state"] = state - model_output = self.model.infer(model_input) - - model_output["robot_state_action_data"] = model_input[ - "robot_state_action_data" - ] - model_output["robot_state_action_data"].save_action_data( - model_output["predict_action"] - ) - - model_output["_last_obs"] = None - - done, delta_t = self.apply_action( - model_output, env=env, replay_images=replay_images - ) - - if model_output.get("_last_obs") is not None: - current_obs = model_output["_last_obs"] - - count -= delta_t - - except Exception as e: - self.logger.error(f"Episode Error: {e}") - import traceback - - traceback.print_exc() - break - - success = done - if count <= 0 and not done: - self.logger.warning( - f"Timeout: reached {max_infer_times} steps without success." - ) - success = False - - self._save_rollout(replay_images, success, task_id, task_desc, episode_idx) - return success - - def run_infer_flow_action_batch( - self, - vec_env: Any, - task_ids: List[int] = None, - task_descs: List[str] = None, - default_initial_states_list: List[Any] = None, - episode_indices: List[int] = None, - max_infer_times: int = 5, - num_steps_wait: int = 10, - ) -> List[bool]: - """ - Support batch inference: model inference in parallel (batch), environment execution in parallel (SubprocVectorEnv). - - Uses SubprocVectorEnv to run environments in subprocess during multi-batch inference, all environments execute actions in parallel. - - Returns a list of success flags for each sample. - """ - if vec_env is None: - raise ValueError("vec_env must be specified") - batch_size = len(vec_env) - if task_ids is not None: - assert len(task_ids) == batch_size, "task_ids length must match envs" - if episode_indices is not None: - assert ( - len(episode_indices) == batch_size - ), "episode_indices length must match envs" - - status_list = [] - for i in range(batch_size): - status_list.append( - { - "vec_env": vec_env, - "env_id": i, # Index in vec_env - "task_desc": task_descs[i], - "replay_images": [], - "num_steps": 0, - "done": False, # Whether episode has ended - "success": False, # Whether successfully completed - "count": max_infer_times, - "current_obs": None, - "default_states": None, - } - ) - - # Initialize/reset: Use SubprocVectorEnv to batch set initial states - init_states_to_set = [] - for i in range(batch_size): - task_desc = task_descs[i] - default_states = default_initial_states_list[i] - status_list[i]["default_states"] = default_states - ep_i = episode_indices[i] - init_state = self._get_initial_state_for_episode( - task_desc, default_states, ep_i - ) - init_states_to_set.append(init_state) - - # Batch set initial states - try: - obs_list = vec_env.set_init_state(init_states_to_set) - if obs_list.dtype == object: - obs_list = list(obs_list) - else: - obs_list = [obs_list[i] for i in range(len(obs_list))] - - for i, obs in enumerate(obs_list): - if obs is None: - raise ValueError( - f"Reset environment returned None, task_id: {task_ids[i]}, episode_idx: {episode_indices[i]}" - ) - status_list[i]["current_obs"] = obs - status_list[i]["done"] = False - status_list[i]["success"] = False - status_list[i]["count"] = max_infer_times - except Exception as e: - self.logger.error(f"Failed to batch set initial states: {e}") - raise - - # Warmup steps (batch execution) - dummy_action = get_libero_dummy_action(self.model_family) - dummy_actions = np.array([dummy_action] * batch_size) - for _ in range(num_steps_wait): - obs_list, _, done_list, _ = vec_env.step(dummy_actions) - # Update current_obs - if obs_list.dtype == object: - obs_list = list(obs_list) - else: - obs_list = [obs_list[i] for i in range(len(obs_list))] - for i, obs in enumerate(obs_list): - if obs is not None: - status_list[i]["current_obs"] = obs - - # Main loop: model parallel inference, environment parallel execution (SubprocVectorEnv) - while any((not st["done"]) and st["count"] > 0 for st in status_list): - active_indices = [ - idx - for idx, st in enumerate(status_list) - if (not st["done"]) and st["count"] > 0 - ] - if not active_indices: - break - print(f"Batch infer loop, active indices: {active_indices}") - - observations = [] - instructions = [] - for idx in active_indices: - st = status_list[idx] - observations.append(self.get_observation(st["current_obs"])) - instructions.append(self.get_instruction(st["task_desc"])) - - # Model batch inference - model_outputs = self.model.infer_flow_action_batch( - observations, instructions - ) - # Extract action trajectories for all active environments - trajectories = [] - for out in model_outputs: - action_data = out["robot_state_action_data"] - right_arm_traj = self._get_right_arm_action(action_data) - while ( - right_arm_traj is not None - and right_arm_traj.ndim > 2 - and right_arm_traj.shape[0] == 1 - ): - right_arm_traj = right_arm_traj.squeeze(0) - if right_arm_traj is None or len(right_arm_traj) == 0: - # If trajectory is empty, create an empty trajectory - right_arm_traj = np.array([]).reshape(0, 7) - trajectories.append(right_arm_traj) - - # Use apply_action_batch to execute action trajectories in parallel - try: - self.apply_action_batch( - vec_env=vec_env, - trajectories=trajectories, - active_indices=active_indices, - status_list=status_list, - model_outputs=model_outputs, - ) - except Exception as e: - self.logger.error(f"Batch parallel action error: {e}") - # Mark all active environments as failed - for idx in active_indices: - status_list[idx]["done"] = True - status_list[idx]["success"] = False - - # Save replay and results - success_list = [] - for i, st in enumerate(status_list): - success = st.get("success", False) - if st["count"] <= 0 and not st["success"]: - self.logger.warning( - f"Batch timeout: reached {max_infer_times} steps without success (idx {i})." - ) - st["replay_images"] = st.get("replay_images", []) - tid_i = task_ids[i] - epi_i = episode_indices[i] - self._save_rollout( - st["replay_images"], - success, - tid_i, - st["task_desc"], - epi_i, - ) - success_list.append(success) - - return success_list - - def run_infer_ar_action( - self, - env: Any, - task_id: int, - task_desc: str, - default_initial_states: Any, - episode_idx: int, - max_infer_times: int = 10, - num_steps_wait: int = 10, - ) -> bool: - replay_images = [] - num_steps = 0 - done = False - count = max_infer_times - - current_obs = self.reset_env( - env, task_desc, default_initial_states, episode_idx - ) - if current_obs is None: - self.logger.error("Environment reset returned None.") - return False - - while num_steps < num_steps_wait: - obs, reward, done, info = env.step( - get_libero_dummy_action(self.model_family) - ) - num_steps += 1 - if obs is not None: - current_obs = obs - - while not done and count > 0: - try: - model_input = self.get_observation(current_obs) - instruction = self.get_instruction(task_desc) - - model_output = self.model.infer_ar_action(model_input, instruction) - - model_output["_last_obs"] = None - - done, delta_t = self.apply_action( - model_output, env=env, replay_images=replay_images - ) - - if model_output.get("_last_obs") is not None: - current_obs = model_output["_last_obs"] - else: - if not done: - self.logger.warning( - "Did not receive new observation after apply_action, but episode is not done." - ) - - count -= delta_t - - except Exception as e: - self.logger.error( - f"AR Episode Run Error (Task {task_id}, Ep {episode_idx}): {e}" - ) - import traceback - - traceback.print_exc() - break - - success = done - if count <= 0 and not done: - self.logger.warning( - f"Timeout: AR policy reached {max_infer_times} steps without success." - ) - success = False - - self._save_rollout(replay_images, success, task_id, task_desc, episode_idx) - - return success diff --git a/wall_x/infer/infer_config.py b/wall_x/infer/infer_config.py deleted file mode 100644 index bf7dc82..0000000 --- a/wall_x/infer/infer_config.py +++ /dev/null @@ -1,587 +0,0 @@ -import yaml -import os -from wall_x.model.model_utils import update_model_config - -# from x2robot_dataset.configs.config import X2RDataConfig - -import json -from typing import List, Dict, Optional, Any -from dataclasses import dataclass, field - - -@dataclass -class X2RDataConfig: - """ - Unified X2Robot data configuration class (reorganized by README's 5 modules): - 1) Data I/O and caching - 2) Visual input and sampling (image/camera) - 3) Action and time series - 4) Instruction and multimodal - 5) Data cleaning and alignment (validation/augmentation/framework constraints) - """ - - # ---------------------------------------------------------------------- - # 1) Data I/O and caching - # ---------------------------------------------------------------------- - cache_dir: str = "~/.cache/dataset_cache" - dataset_config_path: Optional[str] = None - use_cache: bool = True - check_mode: bool = True - preload_size: int = 128 - buffer_size: int = 20000 - batch_size: int = 32 - train_test_split: float = 0.9 - seed: int = 42 - episode_chunk_size: int = ( - 500 # Commonly used on VG side (number of frames for episode chunking) - ) - - # ---------------------------------------------------------------------- - # 2) Visual input and sampling (image/camera) - # ---------------------------------------------------------------------- - # Camera mapping - cam_mapping: Dict[str, str] = field( - default_factory=lambda: { - "faceImg": "face_view", - "leftImg": "left_wrist_view", - "rightImg": "right_wrist_view", - } - ) - # Image and augmentation - resolution: Dict[str, int] = field( - default_factory=lambda: { - "face_view": -1, - "left_wrist_view": 128, - "right_wrist_view": 128, - } - ) - cam_augmentation_list: List[str] = field(default_factory=list) - - # Image time series (history/future) - image_horizon: int = 1 - image_history_length: int = 0 - image_history_interval: int = 1 - future_image_length: int = 0 - future_image_interval: int = 1 - future_image_indices: Optional[List[int]] = ( - None # If provided, length must equal image_horizon - ) - - # Smart scaling - max_pixels: int = field( - default_factory=lambda: 1280 * 28 * 28 - ) # Will be replaced with MAX_PIXELS in __post_init__ - min_pixels: int = field( - default_factory=lambda: 4 * 28 * 28 - ) # Will be replaced with MIN_PIXELS in __post_init__ - image_factor: int = 28 # Will be replaced with IMAGE_FACTOR in __post_init__ - - # ---------------------------------------------------------------------- - # 3) Action and time series - # ---------------------------------------------------------------------- - predict_action_keys: List[str] = field(default_factory=list) - obs_action_keys: List[str] = field(default_factory=list) - - # Action window - action_horizon: int = 21 - action_history_length: int = 0 - action_horizon_flow: int = 32 - action_horizon_ar: int = 0 - - # Padding strategy - left_padding: bool = True - right_padding: bool = True - - # Dimension configuration - dof_config: Dict[str, int] = field(default_factory=dict) # Input degrees of freedom - agent_pos_config: Dict[str, int] = field( - default_factory=dict - ) # Output degrees of freedom - - # State augmentation - state_augmentation_ratio: float = 1.0 # Ratio of augmented states - state_augmentation_prob: float = ( - 0.1 # Random dimension masking probability for state string - ) - state_drop_prob: float = 0.0 # Probability of dropping entire state - - # ---------------------------------------------------------------------- - # 4) Instruction and multimodal - # ---------------------------------------------------------------------- - default_instruction: str = "" - instruction_path: Optional[str] = None - instruction_key: Optional[List[Dict]] = None - - multimodal_chunk_size: int = 500 - generate_subtask_ratio: float = 0.0 - cot_ratio: float = 0.0 - multimodal_data_ratio: float = ( - 0.25 # Multimodal data ratio per batch in VLA dataset - ) - instruction_key_prob: Optional[Dict[str, float]] = None - trunc_action_with_instruction: bool = True - use_embodied_system_prompt_ratio: float = 0.0 - - # ---------------------------------------------------------------------- - # 5) Data cleaning and alignment (validation/augmentation/framework constraints) - # ---------------------------------------------------------------------- - filter_angle_outliers: bool = False - trim_stationary: bool = False - use_state_string_representation: bool = False - pad_prefix_to_same_length: bool = False - put_ar_predict_in_postfix: bool = ( - False # Whether to put ar prediction in postfix, set to True in prediction mode, False in training - ) - pad_to_128_multiple: bool = ( - False # Triton Attention requirement (deprecated, always set to False) - ) - max_seqlen: int = 768 - model_type: Optional[str] = None # qwen2_5, qwen2 - model_config_path: Optional[str] = ( - None # Model config path (used to derive PaddingSide) - ) - low_dim_obs_horizon: int = 1 # To be deprecated - - # ---------------------------------------------------------------------- - # Validation and post-processing - # ---------------------------------------------------------------------- - def __post_init__(self): - # TODO: Determine VGA model type validation here - # assert self.model_type in ["qwen2_5", "qwen3"], f"Unsupported model type: {self.model_type}" - - if self.model_type == "qwen2_5": - self.max_pixels = 16384 * 28 * 28 - self.min_pixels = 4 * 28 * 28 - self.image_factor = 28 - elif self.model_type == "qwen3": - self.max_pixels = 16384 * 32 * 32 - self.min_pixels = 4 * 32 * 32 - self.image_factor = 32 - - # Future image indices validation - if ( - self.future_image_indices - and len(self.future_image_indices) != self.image_horizon - ): - raise ValueError( - f"future_image_indices length must equal image_horizon: " - f"{len(self.future_image_indices)} != {self.image_horizon}" - ) - - # Auto-derive action window - if self.action_horizon == 0: - self.action_horizon = max(self.action_horizon_flow, self.action_horizon_ar) - - # Auto-derive action keys - if not self.obs_action_keys: - self.obs_action_keys = list(self.agent_pos_config.keys()) - if not self.predict_action_keys: - self.predict_action_keys = list(self.dof_config.keys()) - - # Derive PaddingSide - # @Ryan: Only FlashAttention can use RightPadding, other AttnImpl use LeftPadding - if self.model_config_path is not None: - with open(self.model_config_path, "r", encoding="utf-8") as f: - cfg = json.load(f) - - attn_impl = cfg["_attn_implementation"] - - if attn_impl == "flash_attention_2": - self.padding_side = "right" - else: - self.padding_side = "left" - - # Convenience properties - @property - def use_6D_rotation(self) -> bool: - """Whether to use 6D rotation (auto-determined from predict_action_keys)""" - if hasattr(self, "_use_6D_rotation"): - return self._use_6D_rotation - self._use_6D_rotation = any("6D" in key for key in self.predict_action_keys) - return self._use_6D_rotation - - @property - def use_relative_action(self) -> bool: - """Whether to use relative action (auto-determined from predict_action_keys)""" - if hasattr(self, "_use_relative_action"): - return self._use_relative_action - self._use_relative_action = any( - "relative" in key for key in self.predict_action_keys - ) - return self._use_relative_action - - # ---------------------------------------------------------------------- - # YAML initialization - # ---------------------------------------------------------------------- - @classmethod - def from_yaml_dict(cls, yaml_dict: Dict[str, Any]) -> "X2RDataConfig": - """ - Create config object from YAML config dict. Prioritizes data sub-config, then top-level fields. - """ - data_config = yaml_dict.get("data", {}) - params: Dict[str, Any] = {} - - # 1) Data I/O and caching - params.update( - { - "cache_dir": data_config.get( - "cache_dir", yaml_dict.get("cache_dir", "~/.cache/dataset_cache") - ), - "dataset_config_path": data_config.get( - "dataset_config_path", yaml_dict.get("dataset_config_path", None) - ), - "use_cache": data_config.get( - "use_cache", yaml_dict.get("use_cache", True) - ), - "check_mode": data_config.get( - "check_mode", yaml_dict.get("check_mode", True) - ), - "preload_size": data_config.get( - "preload_size", yaml_dict.get("preload_size", 128) - ), - "buffer_size": data_config.get( - "buffer_size", yaml_dict.get("buffer_size", 20000) - ), - "batch_size": data_config.get( - "batch_size", - yaml_dict.get( - "batch_size_per_gpu", yaml_dict.get("batch_size", 32) - ), - ), - "train_test_split": data_config.get("train_test_split", 0.9), - "seed": yaml_dict.get("seed", 42), - "episode_chunk_size": data_config.get("episode_chunk_size", 500), - } - ) - - # 2) Visual input and sampling (image/camera) - params.update( - { - "cam_mapping": data_config.get( - "cam_mapping", - { - "faceImg": "face_view", - "leftImg": "left_wrist_view", - "rightImg": "right_wrist_view", - }, - ), - "resolution": data_config.get( - "resolution", - {"face_view": -1, "left_wrist_view": 128, "right_wrist_view": 128}, - ), - "cam_augmentation_list": data_config.get("cam_augmentation_list", []), - "image_horizon": data_config.get("image_horizon", 1), - "image_history_length": data_config.get("image_history_length", 0), - "image_history_interval": data_config.get("image_history_interval", 1), - "future_image_length": data_config.get("future_image_length", 0), - "future_image_interval": data_config.get("future_image_interval", 1), - "future_image_indices": data_config.get("future_image_indices", None), - "max_pixels": data_config.get("max_pixels", 1280 * 28 * 28), - "min_pixels": data_config.get("min_pixels", 4 * 28 * 28), - "image_factor": data_config.get("image_factor", 28), - } - ) - - # 3) Action and time series - params.update( - { - "predict_action_keys": data_config.get("predict_action_keys", []), - "obs_action_keys": data_config.get("obs_action_keys", []), - "action_horizon": data_config.get("action_horizon", 0), - "action_history_length": data_config.get("action_history_length", 0), - "action_horizon_flow": data_config.get( - "action_horizon_flow", yaml_dict.get("action_horizon_flow", 32) - ), - "action_horizon_ar": data_config.get("action_horizon_ar", 0), - "left_padding": data_config.get("left_padding", True), - "right_padding": data_config.get("right_padding", True), - "dof_config": yaml_dict.get( - "dof_config", data_config.get("dof_config", {}) - ), - "agent_pos_config": yaml_dict.get( - "agent_pos_config", data_config.get("agent_pos_config", {}) - ), - "state_augmentation_prob": data_config.get( - "state_augmentation_prob", 0.05 - ), - "state_drop_prob": data_config.get("state_drop_prob", 0.0), - } - ) - - # 4) Instruction and multimodal - params.update( - { - "default_instruction": data_config.get("default_instruction", ""), - "instruction_path": data_config.get("instruction_path", None), - "instruction_key": data_config.get("instruction_key", None), - "multimodal_chunk_size": data_config.get("multimodal_chunk_size", 500), - "generate_subtask_ratio": data_config.get( - "generate_subtask_ratio", 0.0 - ), - "cot_ratio": data_config.get("cot_ratio", 0.0), - "multimodal_data_ratio": data_config.get("multimodal_data_ratio", 0.25), - "instruction_key_prob": data_config.get("instruction_key_prob", None), - "trunc_action_with_instruction": data_config.get( - "trunc_action_with_instruction", True - ), - "use_embodied_system_prompt_ratio": data_config.get( - "use_embodied_system_prompt_ratio", - yaml_dict.get("use_embodied_system_prompt_ratio", 0.0), - ), - } - ) - - # 5) Data cleaning and alignment (validation/augmentation/framework constraints) - params.update( - { - "filter_angle_outliers": data_config.get( - "filter_angle_outliers", False - ), - "trim_stationary": data_config.get("trim_stationary", False), - "use_state_string_representation": data_config.get( - "use_state_string_representation", - yaml_dict.get("use_state_string_representation", False), - ), - "pad_prefix_to_same_length": data_config.get( - "pad_prefix_to_same_length", False - ), - "put_ar_predict_in_postfix": data_config.get( - "put_ar_predict_in_postfix", False - ), - # "pad_to_128_multiple": data_config.get("pad_to_128_multiple", True), - "padding_side": data_config.get("padding_side", "left"), - "max_seqlen": yaml_dict.get("max_seqlen", 768), - "model_type": yaml_dict.get("model_type", "qwen2_5"), - "model_config_path": yaml_dict.get("qwen_vl_act_config_path", None), - "low_dim_obs_horizon": data_config.get("low_dim_obs_horizon", 1), - } - ) - - # Only keep valid fields defined in dataclass - valid_fields = {f.name for f in cls.__dataclass_fields__.values()} - filtered = {k: v for k, v in params.items() if k in valid_fields} - return cls(**filtered) - - # ---------------------------------------------------------------------- - # Dict-style access (for compatibility with existing calls) - # ---------------------------------------------------------------------- - def __getitem__(self, key: str): - try: - return getattr(self, key) - except AttributeError: - raise KeyError(f"'{key}' not found in {self.__class__.__name__}") - - def __setitem__(self, key: str, value): - setattr(self, key, value) - - def __contains__(self, key: str) -> bool: - return hasattr(self, key) - - def keys(self): - return self.__dict__.keys() - - def values(self): - return self.__dict__.values() - - def items(self): - return self.__dict__.items() - - -class InferConfig: - def __init__( - self, - checkpoint_path: str | None = None, - train_config_path: str | None = None, - robot_host: str = "0.0.0.0", - robot_port: int = 33723, - robot_id: str = "10053", - robot_type: str = "desktop", # ["desktop", "turtle"] - robot_action_start_ratio: float = 0, # Action execution start ratio - robot_action_end_ratio: float = 0.8, # Action execution end ratio - robot_action_interpolate_multiplier: int = 70, # Action interpolation - robot_use_joint_angle_control: bool = False, # Use joint control (model must be joint prediction model) - turtle_as_desktop: bool = False, # Use turtle body for desktop operation, fixed chassis head movement, head camera, and chassis height - action_horizon: int = 10, # Please correctly fill in the model's horizon - action_dim: int | None = None, - model_device: str = "cuda:0", - num_inference_timesteps: int = 10, - norm_key: str = "x2_normal", - cam_names: list[str] = ["face_view", "right_wrist_view"], - ): - # Private attribute for storing path - assert checkpoint_path is not None - self._checkpoint_path = checkpoint_path - if os.path.exists(os.path.join(checkpoint_path, "normalizer_action.pth")): - self.normalizer_action_path = os.path.join( - checkpoint_path, "normalizer_action.pth" - ) - if os.path.exists(os.path.join(checkpoint_path, "normalizer_propri.pth")): - self.normalizer_propri_path = os.path.join( - checkpoint_path, "normalizer_propri.pth" - ) - - self.model_path = checkpoint_path - self.action_tokenizer_path = "/x2robot_v2/Models/fast/" - - # Other configuration attributes - self.robot_host = robot_host - self.robot_port = robot_port - self.robot_type = robot_type # ["desktop", "turtle"] - self.robot_id = robot_id - self.robot_action_start_ratio = robot_action_start_ratio - self.robot_action_end_ratio = robot_action_end_ratio - self.robot_action_interpolate_multiplier = robot_action_interpolate_multiplier - self.robot_use_joint_angle_control = ( - robot_use_joint_angle_control # Use joint angle control - ) - self.turtle_as_desktop = turtle_as_desktop - - self._action_horizon = ( - action_horizon # Default controlled by train config's flow action horizon - ) - self._action_dim = action_dim # Default determined by train config's dof config - - self.action_dim = action_dim - self.pred_horizon = action_horizon - self.predict_mode = "diffusion" - self.camera_key = cam_names - - self.model_device = model_device - self.num_inference_timesteps = ( - num_inference_timesteps # flow matching related config - ) - - # Initialize config objects - self.train_config: dict = {} - self.model_config = None - self.data_config = None - self.norm_key = norm_key - self.cam_names = cam_names - # Load all configs - self._load_all_configs(train_config_path) - - @property - def checkpoint_path(self) -> str | None: - return self._checkpoint_path - - @checkpoint_path.setter - def checkpoint_path(self, value: str | None): - """When checkpoint_path is updated, reload all configs""" - if self._checkpoint_path != value: - self._checkpoint_path = value - self._load_all_configs() - - @property - def action_horizon(self) -> int: - return self._action_horizon - - @action_horizon.setter - def action_horizon(self, value: int): - self._action_horizon = value - - @property - def action_dim(self) -> int | None: - return self._action_dim - - @action_dim.setter - def action_dim(self, value: int | None): - self._action_dim = value - - def _load_all_configs(self, train_config_path=None): - """Unified entry point for loading all configs""" - self._load_train_config(train_config_path) - self._load_model_config() - self._load_data_config() - - # Update action_horizon and action_dim (if needed) - if self._action_horizon is None: - self._action_horizon = self.train_config.get("data", {}).get( - "action_horizon_flow", 32 - ) - assert self._action_horizon is not None and self._action_horizon > 0 - - if self._action_dim is None: - self._action_dim = sum(self.train_config.get("dof_config", {}).values()) - - def _load_train_config(self, train_config_path): - if train_config_path is None: - train_config_path = os.path.join(self._checkpoint_path, "config.yml") - with open(train_config_path, "r") as f: - self.train_config = yaml.load(f, Loader=yaml.FullLoader) - - ckpt_dir = self._checkpoint_path - preprocessor_file = os.path.join(ckpt_dir, "preprocessor_config.json") - if os.path.exists(preprocessor_file): - print(f"[LoadConfig] Found {preprocessor_file}, override processor_path.") - self.train_config["processor_path"] = ckpt_dir - - tokenizer_file = os.path.join(ckpt_dir, "tokenizer.json") - tokenizer_config_file = os.path.join(ckpt_dir, "tokenizer_config.json") - if "action_tokenizer_path" in self.train_config and not os.path.exists( - self.train_config["action_tokenizer_path"] - ): - if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file): - print( - f"[LoadConfig] Found tokenizer files in {ckpt_dir}, override action_tokenizer_path." - ) - self.train_config["action_tokenizer_path"] = ckpt_dir - else: - print("[LoadConfig] Cannot load action tokenizer! ") - - def _load_model_config(self): - ckpt_config_path = os.path.join(self._checkpoint_path, "config.json") - resolved_cfg_path = None - - if os.path.exists(ckpt_config_path): - # Prefer checkpoint config - resolved_cfg_path = ckpt_config_path - print(f"[LoadModelConfig] Using checkpoint config.json: {ckpt_config_path}") - else: - # Fallback to original config path - fallback_cfg = self.train_config.get("qwen_vl_act_config_path", None) - if fallback_cfg is not None: - resolved_cfg_path = fallback_cfg - print(f"[LoadModelConfig] Using fallback act config: {fallback_cfg}") - - if resolved_cfg_path is None or (not os.path.exists(resolved_cfg_path)): - raise ValueError( - f"[LoadModelConfig] Cannot load model config! " - f"Checked:\n" - f" - Checkpoint config.json: {ckpt_config_path}\n" - f" - Fallback path: {self.train_config.get('qwen_vl_act_config_path', None)}" - ) - - # Save back to config for consistency - self.train_config["qwen_vl_act_config_path"] = resolved_cfg_path - - model_type = self.train_config["model_type"] - if model_type == "qwen2_5": - from wall_x.model.qwen2_5_based import Qwen2_5_VLConfig - - ConfigClass = Qwen2_5_VLConfig - - # elif model_type == "qwen3": - # from wall_x.model.qwen3_based import Qwen3VLConfig - - # ConfigClass = Qwen3VLConfig - - else: - raise ValueError(f"[LoadModelConfig] Unsupported model type: {model_type}") - - print(f"[LoadModelConfig] Loading model config from: {resolved_cfg_path}") - self.model_config = ConfigClass.from_pretrained(resolved_cfg_path) - - self.model_config = update_model_config(self.train_config, self.model_config) - - self.model_config._attn_implementation = "sdpa" - self.model_config.vision_config._attn_implementation = "flash_attention_2" - - print("[LoadModelConfig] Model config loaded and updated successfully.") - - def _load_data_config(self): - self.data_config = X2RDataConfig.from_yaml_dict(self.train_config) - - -if __name__ == "__main__": - config = InferConfig() - print(config.train_config) diff --git a/wall_x/infer/utils_libero.py b/wall_x/infer/utils_libero.py deleted file mode 100644 index b6fd5cb..0000000 --- a/wall_x/infer/utils_libero.py +++ /dev/null @@ -1,323 +0,0 @@ -"""Utils for evaluating policies in LIBERO simulation environments.""" - -import math -import os -from enum import Enum -import imageio -import numpy as np -import matplotlib.pyplot as plt -import torch -from transformers import BatchFeature -import random -import time - -from libero.libero import get_libero_path -from libero.libero.envs import OffScreenRenderEnv - - -# Define task suite constants -class TaskSuite(str, Enum): - LIBERO_SPATIAL = "libero_spatial" - LIBERO_OBJECT = "libero_object" - LIBERO_GOAL = "libero_goal" - LIBERO_10 = "libero_10" - LIBERO_90 = "libero_90" - - -# Define max steps for each task suite -TASK_MAX_STEPS = { - TaskSuite.LIBERO_SPATIAL: 220, # longest training demo has 193 steps - TaskSuite.LIBERO_OBJECT: 280, # longest training demo has 254 steps - TaskSuite.LIBERO_GOAL: 300, # longest training demo has 270 steps - TaskSuite.LIBERO_10: 520, # longest training demo has 505 steps - TaskSuite.LIBERO_90: 400, # longest training demo has 373 steps -} - - -# Initialize important constants -ACTION_DIM = 7 -DATE = time.strftime("%Y_%m_%d") -DATE_TIME = time.strftime("%Y_%m_%d-%H_%M_%S") - -# Configure NumPy print settings -np.set_printoptions(formatter={"float": lambda x: "{0:0.3f}".format(x)}) - - -def set_seed_everywhere(seed: int) -> None: - """ - Set random seed for all random number generators for reproducibility. - - Args: - seed: The random seed to use - """ - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - np.random.seed(seed) - random.seed(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False - os.environ["PYTHONHASHSEED"] = str(seed) - - -def normalize_gripper_action(action: np.ndarray, binarize: bool = True) -> np.ndarray: - """ - Normalize gripper action from [0,1] to [-1,+1] range. - - This is necessary for some environments because the dataset wrapper - standardizes gripper actions to [0,1]. Note that unlike the other action - dimensions, the gripper action is not normalized to [-1,+1] by default. - - Normalization formula: y = 2 * (x - orig_low) / (orig_high - orig_low) - 1 - - Args: - action: Action array with gripper action in the last dimension - binarize: Whether to binarize gripper action to -1 or +1 - - Returns: - np.ndarray: Action array with normalized gripper action - """ - # Create a copy to avoid modifying the original - normalized_action = action.copy() - - # Normalize the last action dimension to [-1,+1] - orig_low, orig_high = 0.0, 1.0 - normalized_action[..., -1] = ( - 2 * (normalized_action[..., -1] - orig_low) / (orig_high - orig_low) - 1 - ) - - if binarize: - # Binarize to -1 or +1 - normalized_action[..., -1] = np.sign(normalized_action[..., -1]) - - return normalized_action - - -def invert_gripper_action(action: np.ndarray) -> np.ndarray: - """ - Flip the sign of the gripper action (last dimension of action vector). - - This is necessary for environments where -1 = open, +1 = close, since - the RLDS dataloader aligns gripper actions such that 0 = close, 1 = open. - - Args: - action: Action array with gripper action in the last dimension - - Returns: - np.ndarray: Action array with inverted gripper action - """ - # Create a copy to avoid modifying the original - inverted_action = action.copy() - - # Invert the gripper action - inverted_action[..., -1] *= -1.0 - - return inverted_action - - -def move_to_cuda(obj, device="cuda"): - if isinstance(obj, torch.Tensor): - return obj.to(device) - elif isinstance(obj, (dict, BatchFeature)): - return {k: move_to_cuda(v, device) for k, v in obj.items()} - elif isinstance(obj, list): - return [move_to_cuda(v, device) for v in obj] - elif isinstance(obj, tuple): - return tuple(move_to_cuda(v, device) for v in obj) - else: - return obj - - -def get_libero_env(task, model_family, resolution=256, seed=7): - """Initializes and returns the LIBERO environment, along with the task description.""" - task_description = task.language - task_bddl_file = os.path.join( - get_libero_path("bddl_files"), task.problem_folder, task.bddl_file - ) - env_args = { - "bddl_file_name": task_bddl_file, - "camera_heights": resolution, - "camera_widths": resolution, - } - env = OffScreenRenderEnv(**env_args) - env.seed( - seed - ) # IMPORTANT: seed seems to affect object positions even when using fixed initial state - return env, task_description - - -def get_libero_dummy_action(model_family: str): - """Get dummy/no-op action, used to roll out the simulation while the robot does nothing.""" - return [0, 0, 0, 0, 0, 0, -1] - - -def get_libero_image(obs): - """Extracts third-person image from observations and preprocesses it.""" - img = obs["agentview_image"] - img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing - return img - - -def get_libero_wrist_image(obs): - """Extracts wrist camera image from observations and preprocesses it.""" - img = obs["robot0_eye_in_hand_image"] - img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing - return img - - -def save_rollout_video( - rollout_dir, - rollout_images, - idx, - success, - task_description, - log_file=None, - model_family="openvla_oft", -): - """Saves an MP4 replay of an episode.""" - processed_task_description = ( - task_description.lower() - .replace(" ", "_") - .replace("\n", "_") - .replace(".", "_")[:50] - ) - mp4_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}.mp4" - video_writer = imageio.get_writer(mp4_path, fps=30) - for img in rollout_images: - video_writer.append_data(img) - video_writer.close() - print(f"Saved rollout MP4 at path {mp4_path}") - if log_file is not None: - log_file.write(f"Saved rollout MP4 at path {mp4_path}\n") - return mp4_path - - -def save_rollout_data( - rollout_dir, - rollout_data, - idx, - success, - task_description, - log_file=None, - model_family="openvla_oft", -): - """ - Saves an NPY file of the rollout data. - """ - - # Process task description to make it suitable for filename - processed_task_description = ( - task_description.lower() - .replace(" ", "_") - .replace("\n", "_") - .replace(".", "_")[:50] - ) - - # Build .npy file path - npy_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--action.npy" - - # Save rollout_data as .npy file - np.save(npy_path, rollout_data) - print(f"Saved rollout data at path {npy_path}") - - fig, axes = plt.subplots(nrows=1, ncols=7, figsize=(20, 3)) - - titles = ["x", "y", "z", "roll", "pitch", "yaw", "grasp"] - - for i in range(rollout_data.shape[1]): - ax = axes[i] # Select the i-th subplot - ax.plot(rollout_data[:, i], label=f"Feature {i+1}") # Plot line chart - ax.set_title(titles[i]) # Set subplot title - ax.set_xlabel("Time in one episode") # Set x-axis label - - axes[-1].legend(["predicted action"], loc="upper right") - - plt.tight_layout() - png_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--action.png" - plt.savefig(png_path, dpi=300) # Save image as PNG file - - # If log file is provided, record the save path - if log_file is not None: - log_file.write(f"Saved rollout data at path {npy_path}\n") - - return npy_path - - -def save_rollout_observation( - rollout_dir, - rollout_data, - idx, - success, - task_description, - log_file=None, - model_family="openvla_oft", -): - """ - Saves an NPY file of the rollout data. - """ - - # Process task description to make it suitable for filename - processed_task_description = ( - task_description.lower() - .replace(" ", "_") - .replace("\n", "_") - .replace(".", "_")[:50] - ) - - # Build .npy file path - npy_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--observation.npy" - - # Save rollout_data as .npy file - np.save(npy_path, rollout_data) - print(f"Saved rollout data at path {npy_path}") - - if rollout_data.shape[1] == 7: - fig, axes = plt.subplots(nrows=1, ncols=7, figsize=(20, 3)) - titles = ["x", "y", "z", "roll", "pitch", "yaw", "grasp"] - else: - fig, axes = plt.subplots(nrows=1, ncols=8, figsize=(20, 3)) - titles = ["x", "y", "z", "roll", "pitch", "yaw", "-", "grasp"] - - for i in range(rollout_data.shape[1]): - ax = axes[i] # Select the i-th subplot - ax.plot(rollout_data[:, i], label=f"Feature {i+1}") # Plot line chart - ax.set_title(titles[i]) # Set subplot title - ax.set_xlabel("Time in one episode") # Set x-axis label - - axes[-1].legend(["predicted action"], loc="upper right") - - plt.tight_layout() - png_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--observation.png" - plt.savefig(png_path, dpi=300) # Save image as PNG file - - # If log file is provided, record the save path - if log_file is not None: - log_file.write(f"Saved rollout data at path {npy_path}\n") - - return npy_path - - -def quat2axisangle(quat): - """ - Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55 - - Converts quaternion to axis-angle format. - Returns a unit vector direction scaled by its angle in radians. - - Args: - quat (np.array): (x,y,z,w) vec4 float angles - - Returns: - np.array: (ax,ay,az) axis-angle exponential coordinates - """ - # clip quaternion - if quat[3] > 1.0: - quat[3] = 1.0 - elif quat[3] < -1.0: - quat[3] = -1.0 - - den = np.sqrt(1.0 - quat[3] * quat[3]) - if math.isclose(den, 0.0): - # This is (close to) a zero degree rotation, immediately return - return np.zeros(3) - - return (quat[:3] * 2.0 * math.acos(quat[3])) / den diff --git a/wall_x/model/action_head.py b/wall_x/model/action_head.py deleted file mode 100644 index 2f3ab73..0000000 --- a/wall_x/model/action_head.py +++ /dev/null @@ -1,808 +0,0 @@ -import torch - -import torch.nn as nn - -from typing import Union -import math - -from diffusers.schedulers.scheduling_ddpm import DDPMScheduler -from torch.distributions import Beta - - -def print_rank_last(message): - """If distributed is initialized, print only on last rank.""" - if torch.distributed.is_initialized(): - if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1): - print(message, flush=True) - else: - print(message, flush=True) - - -class Normalizer(nn.Module): - @classmethod - def from_ckpt(cls, ckpt_path): - instance = cls.__new__(cls) - nn.Module.__init__(instance) - - instance.min = nn.ParameterDict() - instance.delta = nn.ParameterDict() - instance.min_key = "min" - instance.delta_key = "delta" - - ckpt = torch.load(ckpt_path, map_location="cpu") - - for key, value in ckpt.items(): - # Parse key: "min.robot_name" -> prefix="min", name="robot_name" - try: - prefix, name = key.split(".", 1) - if hasattr(instance, prefix): - getattr(instance, prefix)[name] = nn.Parameter( - value, requires_grad=False - ) - print("prefix", prefix) - print("name", name) - except ValueError: - continue - - return instance - - def __init__( - self, action_statistic_dof, dof_config, min_key="min", delta_key="delta" - ): - super(Normalizer, self).__init__() - - self.min_key = min_key - self.delta_key = delta_key - - action_statistic = {} - for robot_name in action_statistic_dof.keys(): - action_statistic[robot_name] = {} - all_dof_min = [] - all_dof_delta = [] - for k in dof_config: - if k in action_statistic_dof[robot_name]: - if ( - min_key in action_statistic_dof[robot_name][k] - and delta_key in action_statistic_dof[robot_name][k] - ): - all_dof_min.extend(action_statistic_dof[robot_name][k][min_key]) - all_dof_delta.extend( - action_statistic_dof[robot_name][k][delta_key] - ) - else: - if robot_name == "x2_normal" or "libero" in robot_name: - print_rank_last( - f"Normalizer (Warning): min_key {min_key} or delta_key {delta_key} " - ) - print_rank_last( - f"not in action_statistic_dof[{robot_name}][{k}], use default min 0.0 and delta 1.0" - ) - all_dof_min.extend([0.0] * dof_config[k]) - all_dof_delta.extend([1.0] * dof_config[k]) - else: - if robot_name == "x2_normal" or "libero" in robot_name: - print_rank_last( - f"Normalizer (Warning): Action {k} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" - ) - all_dof_min.extend([0.0] * dof_config[k]) - all_dof_delta.extend([1.0] * dof_config[k]) - all_dof_min = torch.tensor(all_dof_min) - all_dof_delta = torch.tensor(all_dof_delta) - action_statistic[robot_name][min_key] = all_dof_min - action_statistic[robot_name][delta_key] = all_dof_delta - - self.min = nn.ParameterDict( - { - k: nn.Parameter(action_statistic[k][min_key], requires_grad=False) - for k in action_statistic.keys() - } - ) - self.delta = nn.ParameterDict( - { - k: nn.Parameter(action_statistic[k][delta_key], requires_grad=False) - for k in action_statistic.keys() - } - ) - - for k, v in action_statistic.items(): - print_rank_last( - f"Normalizer: {k} min {action_statistic[k][min_key]} delta {action_statistic[k][delta_key]}" - ) - - def normalize_data(self, xs, dataset_names): - new_xs = [] - dataset_names = [name for name in dataset_names if name != "x2_multimodal"] - for x, dataset_name in zip(xs, dataset_names): - x = (x - self.min[dataset_name]) / (self.delta[dataset_name]) - x = x * 2 - 1 - x = torch.clamp(x, -1, 1) - new_xs.append(x) - new_xs = torch.stack(new_xs) - return new_xs - - def unnormalize_data(self, xs, dataset_names, dof_mask=None): - new_xs = [] - dataset_names = [name for name in dataset_names if name != "x2_multimodal"] - dof_mask = dof_mask if dof_mask is not None else [None] * len(xs) - for x, dataset_name, mask in zip(xs, dataset_names, dof_mask): - x = (x + 1) / 2 - if mask is not None: - mask = mask[0].bool() - action_space_delta = self.delta[dataset_name][mask] - action_space_min = self.min[dataset_name][mask] - else: - action_space_delta = self.delta[dataset_name] - action_space_min = self.min[dataset_name] - x = x * action_space_delta + action_space_min - new_xs.append(x) - new_xs = torch.stack(new_xs) - return new_xs - - -class SinusoidalPosEmb(nn.Module): - def __init__(self, dim: int, min_period: float = 4e-3, max_period: float = 4.0): - super().__init__() - self.dim = dim - if dim % 2 != 0: - raise ValueError(f"embedding_dim ({dim}) must be divisible by 2") - self.min_period = min_period - self.max_period = max_period - - def forward(self, x): - device = x.device - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp( - torch.arange(half_dim, device=device, dtype=torch.float32) * -emb - ) - emb = x[:, None] * emb[None, :] - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb - - -class Downsample1d(nn.Module): - def __init__(self, dim): - super().__init__() - self.conv = nn.Conv1d(dim, dim, 3, 2, 1) - - def forward(self, x): - return self.conv(x) - - -class Upsample1d(nn.Module): - def __init__(self, dim): - super().__init__() - self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1) - - def forward(self, x): - return self.conv(x) - - -class Conv1dBlock(nn.Module): - """ - Conv1d --> GroupNorm --> Mish - """ - - def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8): - super().__init__() - - self.block = nn.Sequential( - nn.Conv1d( - inp_channels, out_channels, kernel_size, padding=kernel_size // 2 - ), - nn.GroupNorm(n_groups, out_channels), - nn.Mish(), - ) - - def forward(self, x): - return self.block(x) - - -class ConditionalResidualBlock1D(nn.Module): - def __init__(self, in_channels, out_channels, cond_dim, kernel_size=3, n_groups=8): - super().__init__() - - self.blocks = nn.ModuleList( - [ - Conv1dBlock(in_channels, out_channels, kernel_size, n_groups=n_groups), - Conv1dBlock(out_channels, out_channels, kernel_size, n_groups=n_groups), - ] - ) - - # FiLM modulation https://arxiv.org/abs/1709.07871 - # predicts per-channel scale and bias - cond_channels = out_channels * 2 - self.out_channels = out_channels - self.cond_encoder = nn.Sequential( - nn.Mish(), - nn.Linear(cond_dim, cond_channels), - nn.Dropout(0.1), - nn.Unflatten(-1, (-1, 1)), - ) - - # make sure dimensions compatible - self.residual_conv = ( - nn.Conv1d(in_channels, out_channels, 1) - if in_channels != out_channels - else nn.Identity() - ) - - def forward(self, x, cond): - """ - x : [ batch_size x in_channels x horizon ] - cond : [ batch_size x cond_dim] - - returns: - out : [ batch_size x out_channels x horizon ] - """ - out = self.blocks[0](x) - embed = self.cond_encoder(cond) - - embed = embed.reshape(embed.shape[0], 2, self.out_channels, 1) - scale = embed[:, 0, ...] - bias = embed[:, 1, ...] - out = scale * out + bias - - out = self.blocks[1](out) - out = out + self.residual_conv(x) - return out - - -class ConditionalUnet1D(nn.Module): - def __init__( - self, - input_dim, - global_cond_dim, - diffusion_step_embed_dim=256, - down_dims=[256, 512, 1024], - # down_dims=[512, 1024, 2048], - kernel_size=5, - n_groups=8, - ): - """ - input_dim: Dim of actions. - global_cond_dim: Dim of global conditioning applied with FiLM - in addition to diffusion step embedding. This is usually obs_horizon * obs_dim - diffusion_step_embed_dim: Size of positional encoding for diffusion iteration k - down_dims: Channel size for each UNet level. - The length of this array determines numebr of levels. - kernel_size: Conv kernel size - n_groups: Number of groups for GroupNorm - """ - - super().__init__() - all_dims = [input_dim] + list(down_dims) - start_dim = down_dims[0] - - dsed = diffusion_step_embed_dim - diffusion_step_encoder = nn.Sequential( - SinusoidalPosEmb(dsed), - nn.Linear(dsed, dsed * 4), - nn.Mish(), - nn.Linear(dsed * 4, dsed), - ) - cond_dim = dsed + global_cond_dim - - in_out = list(zip(all_dims[:-1], all_dims[1:])) - mid_dim = all_dims[-1] - self.mid_modules = nn.ModuleList( - [ - ConditionalResidualBlock1D( - mid_dim, - mid_dim, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - ConditionalResidualBlock1D( - mid_dim, - mid_dim, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - ] - ) - - down_modules = nn.ModuleList([]) - for ind, (dim_in, dim_out) in enumerate(in_out): - is_last = ind >= (len(in_out) - 1) - down_modules.append( - nn.ModuleList( - [ - ConditionalResidualBlock1D( - dim_in, - dim_out, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - ConditionalResidualBlock1D( - dim_out, - dim_out, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - Downsample1d(dim_out) if not is_last else nn.Identity(), - ] - ) - ) - - up_modules = nn.ModuleList([]) - for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): - is_last = ind >= (len(in_out) - 1) - up_modules.append( - nn.ModuleList( - [ - ConditionalResidualBlock1D( - dim_out * 2, - dim_in, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - ConditionalResidualBlock1D( - dim_in, - dim_in, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - Upsample1d(dim_in) if not is_last else nn.Identity(), - ] - ) - ) - - final_conv = nn.Sequential( - Conv1dBlock(start_dim, start_dim, kernel_size=kernel_size), - nn.Conv1d(start_dim, input_dim, 1), - ) - - self.diffusion_step_encoder = diffusion_step_encoder - self.up_modules = up_modules - self.down_modules = down_modules - self.final_conv = final_conv - - # print("number of parameters: {:e}".format( - # sum(p.numel() for p in self.parameters())) - # ) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], - global_cond=None, - ): - """ - x: (B,T,input_dim) - timestep: (B,) or int, diffusion step - global_cond: (B,global_cond_dim) - output: (B,T,input_dim) - """ - # (B,T,C) - sample = sample.moveaxis(-1, -2) - # (B,C,T) - - # 1. time - timesteps = timestep - if not torch.is_tensor(timesteps): - # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can - timesteps = torch.tensor( - [timesteps], dtype=torch.long, device=sample.device - ) - elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - - global_feature = self.diffusion_step_encoder(timesteps) - - if global_cond is not None: - global_feature = torch.cat([global_feature, global_cond], axis=-1) - x = sample - h = [] - for idx, (resnet, resnet2, downsample) in enumerate(self.down_modules): - x = resnet(x, global_feature) - x = resnet2(x, global_feature) - h.append(x) - x = downsample(x) # bs, 2048, 5 - - for mid_module in self.mid_modules: - x = mid_module(x, global_feature) # bs, 2048, 5 - - for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules): - x = torch.cat((x, h.pop()), dim=1) - x = resnet(x, global_feature) - x = resnet2(x, global_feature) - x = upsample(x) # bs, 512, 20 - - x = self.final_conv(x) # bs, 14, 20 - - # (B,C,T) - x = x.moveaxis(-1, -2) - # (B,T,C) - return x - - -class DP_Action_head(nn.Module): - def __init__( - self, - action_dim=14, - transformer_dim=896, - global_cond_dim=1806, - load_pretrained=True, - ): - super().__init__() - self.action_dim = action_dim - self.transformer_dim = transformer_dim - self.load_pretrained = load_pretrained - self.noise_scheduler = DDPMScheduler( - num_train_timesteps=132, - beta_schedule="squaredcos_cap_v2", - clip_sample=True, - prediction_type="epsilon", - ) - if self.load_pretrained: - self.global_cond_dim = global_cond_dim - self.condition_proj = nn.Sequential( - nn.Linear(self.transformer_dim, 2 * self.transformer_dim), - nn.ReLU(), - nn.Linear(2 * self.transformer_dim, 2 * self.global_cond_dim), - nn.ReLU(), - nn.Linear(2 * self.global_cond_dim, self.global_cond_dim), - ) - - self.noise_pred_net = ConditionalUnet1D( - input_dim=self.action_dim, - # down_dims=[256,512,1024], - down_dims=[512, 1024, 2048], - global_cond_dim=self.global_cond_dim, - ) - - # load pretrained model - action_pretrained_path = "/x2robot/liangyuxin/workspace/DiffusionPolicy/big_mix_0718_mn/30_noise_pred_net.pth" - print("load noise_pred_net from:", action_pretrained_path, flush=True) - self.noise_pred_net.load_state_dict(torch.load(action_pretrained_path)) - else: - self.noise_pred_net = ConditionalUnet1D( - input_dim=self.action_dim, - down_dims=[256, 512, 1024], - global_cond_dim=self.transformer_dim, - ) - - def forward(self, naction, condition, sample_times): - bs = naction.shape[0] - noise_shape = ( - naction.shape[0] * sample_times, - naction.shape[1], - naction.shape[2], - ) - noise = torch.randn(noise_shape, device=naction.device) - naction = ( - naction.unsqueeze(1) - .repeat(1, sample_times, 1, 1) - .reshape(bs * sample_times, naction.shape[1], naction.shape[2]) - ) - - timesteps = torch.randint( - 0, - self.noise_scheduler.config.num_train_timesteps, - (bs * sample_times,), - device=naction.device, - ).long() - condition = condition.to(self.condition_proj[0].weight.data.dtype) - if self.load_pretrained: - condition = self.condition_proj(condition) - - noisy_actions = self.noise_scheduler.add_noise(naction, noise, timesteps) - noise_pred = self.noise_pred_net( - noisy_actions, timesteps, global_cond=condition - ) - return noise, noise_pred - - @torch.no_grad() - def predict(self, condition, naction=None): - bs = condition.shape[0] - condition = condition.to(self.condition_proj[0].weight.data.dtype) - if self.load_pretrained: - condition = self.condition_proj(condition) - - if naction is not None: - noise_shape = (naction.shape[0], naction.shape[1], naction.shape[2]) - else: - noise_shape = (bs, 16, self.action_dim) # tobe parameterized - noise = torch.randn(noise_shape, device=condition.device) - naction_pred = noise - # init scheduler - self.noise_scheduler.set_timesteps( - self.noise_scheduler.config.num_train_timesteps - ) - - for k in self.noise_scheduler.timesteps: - # predict noise - noise_pred = self.noise_pred_net( - sample=naction_pred, timestep=k, global_cond=condition - ) - - # inverse diffusion step (remove noise) - naction_pred = self.noise_scheduler.step( - model_output=noise_pred, timestep=k, sample=naction_pred - ).prev_sample - - return naction, naction_pred - - -class ActionProcessor(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.dof_config = config.dof_config - self.agent_pos_config = config.agent_pos_config - self.action_dim = sum([v for k, v in self.dof_config.items()]) - self.propri_dim = sum([v for k, v in self.agent_pos_config.items()]) - - print_rank_last( - f"self.dof_config: {self.dof_config}; action_dim: {self.action_dim}; self.agent_pos_config: {self.agent_pos_config}; propri_dim: {self.propri_dim}" - ) - - self.action_hidden_size = config.action_hidden_size - self.state_hidden_size = config.state_hidden_size - self.hidden_size = config.hidden_size - - if not self.config.use_state_string_representation: - if self.config.proj_with_mask: - self.propri_proj = nn.Linear( - self.propri_dim * 2, self.state_hidden_size, bias=False - ) - else: - self.propri_proj = nn.Linear( - self.propri_dim, self.state_hidden_size, bias=False - ) - - # noise scheduler configing - if getattr(self.config, "use_flow_action_expert", True): - noise_scheduler_config = config.noise_scheduler - self.beta_alpha = noise_scheduler_config.get("beta_alpha", 1.5) - self.beta_beta = noise_scheduler_config.get("beta_beta", 1.0) - self.s = noise_scheduler_config.get("s", 0.999) - alpha_tensor = torch.tensor(self.beta_alpha, dtype=torch.float32).to("cuda") - beta_tensor = torch.tensor(self.beta_beta, dtype=torch.float32).to("cuda") - self.beta_dist = Beta(alpha_tensor, beta_tensor) - self.time_embed = SinusoidalPosEmb(self.action_hidden_size) - - # project to hidden space - if self.config.proj_with_mask: - self.w1 = nn.Linear( - self.action_dim * 2, self.action_hidden_size, bias=False - ) - else: - self.w1 = nn.Linear( - self.action_dim, self.action_hidden_size, bias=False - ) - if not self.config.use_adarms: - self.w2 = nn.Linear( - self.action_hidden_size * 2, self.action_hidden_size, bias=False - ) - self.w3 = nn.Linear( - self.action_hidden_size, self.action_hidden_size, bias=False - ) - self.act_fn = nn.SiLU() - else: - self.time_mlp_in = nn.Linear( - self.action_hidden_size, self.action_hidden_size - ) - self.time_mlp_out = nn.Linear( - self.action_hidden_size, self.action_hidden_size - ) - self.act_fn = nn.SiLU() - - # project back to action space - self.action_proj_back = nn.Linear( - self.action_hidden_size, self.action_dim, bias=False - ) - self.mse_loss = nn.MSELoss(reduction="none") - - def set_normalizer(self, normalizer_action, normalizer_propri): - self.normalizer_action = normalizer_action - self.normalizer_propri = normalizer_propri - - # dataset_name = self.config["data"]["lerobot_config"]["repo_id"] - # print("normalizer_propri min", self.normalizer_propri.min.__getattr__(dataset_name), flush=True) - # print("normalizer_propri delta", self.normalizer_propri.delta.__getattr__(dataset_name), flush=True) - # print("normalizer_action min", self.normalizer_action.min.__getattr__(dataset_name), flush=True) - # print("normalizer_action delta", self.normalizer_action.delta.__getattr__(dataset_name), flush=True) - - def sample_time(self, batch_size, device, dtype): - """ - Sampling Time Step - Generates random numbers in the range [0, 1] using a Beta distribution, and then scales them. - - Parameters: - batch_size (int): Batch size - device: Device type - dtype: Data type - - Returns: - torch.Tensor: Sampled time steps, with shape [batch_size] - """ - sample = self.beta_dist.sample([batch_size]).to(device=device, dtype=dtype) - time = (1 - sample) * self.s - return time - - def proprioception_proj( - self, proprioception, dataset_names=None, dof_mask=None, use_history=False - ): - """ - proprioception: [batch_size, 1, action_dim] - dataset_names: [batch_size] - dof_mask: [batch_size, action_dim] - """ - proprioception = proprioception.to(device=self.propri_proj.weight.device).to( - dtype=self.propri_proj.weight.dtype - ) - if dof_mask is not None: - if self.config.proj_with_mask: - proprioception = torch.cat( - [proprioception, dof_mask], dim=-1 - ) # .unsqueeze(1) - proprioception = proprioception.to(device=self.propri_proj.weight.device).to( - dtype=self.propri_proj.weight.dtype - ) - proprio_embed = self.propri_proj( - proprioception - ) # [batch_size, 1, state_hidden_size] - - if self.state_hidden_size < self.hidden_size: - # padding to hidden size - padding_size = self.hidden_size - self.state_hidden_size - padding = torch.zeros( - (proprio_embed.shape[0], 1, padding_size), - device=proprio_embed.device, - dtype=proprio_embed.dtype, - ) - proprio_embed = torch.cat([proprio_embed, padding], dim=-1) - - return proprio_embed # [batch_size, 1, hidden_size] - - def forward(self, action_chunk, dataset_names, dof_mask=None): - """ - Parameters: - action_chunk (torch.Tensor): Action sequence, shape [batch_size, action_chunk_len, action_dim] - dataset_names: [batch_size] - dof_mask: [batch_size, action_dim] - - Returns: - torch.Tensor: Processed action representation, shape [batch_size, seq_len, hidden_size] - """ - with torch.autocast("cuda", dtype=torch.float32): - action_chunk = action_chunk.to(dtype=torch.float32) - batch_size = action_chunk.shape[0] - device = action_chunk.device - dtype = action_chunk.dtype - - # 1. add noise to action_chunk - noise = torch.randn_like(action_chunk) - time = self.sample_time(batch_size, device, dtype) - time_expanded = time.unsqueeze(-1).unsqueeze(-1) - noisy_action = (1 - time_expanded) * noise + time_expanded * action_chunk - flow = action_chunk - noise - - # 2. sinusoidal positional encoding for timesteps - time_embed = self.time_embed(time).to(torch.float32) - - self.noise = noise - self.noisy_action = noisy_action # for new x-pred - - # 3.action_chunk_nosiy + t_pos_emb -> MLP_act_chunk -> action_chunk_nosiy_emb_with_t (dim=trans * chunk) - if dof_mask is not None: - noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) - - noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) - action_embed = self.w1(noisy_action) - - self.time_expanded = time_expanded # for new x-pred - - if not self.config.use_adarms: - time_embed = ( - time_embed.unsqueeze(1) - .repeat(1, action_embed.shape[1], 1) - .to(dtype=self.w2.weight.dtype) - ) - concat_embed = torch.cat([action_embed, time_embed], dim=-1) - concat_embed = self.w2(concat_embed) - action_time_embed = self.w3(self.act_fn(concat_embed)) - adarms_cond = None - else: - time_embed = self.time_mlp_in(time_embed) - time_embed = self.act_fn(time_embed) - time_embed = self.time_mlp_out(time_embed) - time_embed = self.act_fn(time_embed) - action_time_embed = action_embed - adarms_cond = time_embed - - if self.action_hidden_size < self.hidden_size: - # padding to hidden size - padding_size = self.hidden_size - self.action_hidden_size - padding = torch.zeros( - ( - action_time_embed.shape[0], - action_time_embed.shape[1], - padding_size, - ), - device=action_time_embed.device, - dtype=action_time_embed.dtype, - ) - action_time_embed = torch.cat([action_time_embed, padding], dim=-1) - - return action_time_embed, flow, adarms_cond - - def step(self, timestep, noisy_action, dof_mask=None): - # noisy_action: bs, pred_horizon, action_dim - # timestep: bs - with torch.autocast("cuda", dtype=torch.float32): - if dof_mask is not None and self.config.proj_with_mask: - if dof_mask.shape[1] == 1: - dof_mask = dof_mask.unsqueeze(1).repeat(1, noisy_action.shape[1], 1) - noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) - - noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) - time_embed = self.time_embed(timestep).to(torch.float32) # bs,hidden_size - action_embed = self.w1(noisy_action) - - if not self.config.use_adarms: - time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1) - time_embed = time_embed.to(device=noisy_action.device).to( - dtype=noisy_action.dtype - ) - concat_embed = torch.cat([action_embed, time_embed], dim=-1) - concat_embed = self.w2(concat_embed) - embed = self.w3(self.act_fn(concat_embed)) # is this right? - adarms_cond = None - else: - time_embed = time_embed.to(dtype=self.time_mlp_in.weight.dtype) - time_embed = self.time_mlp_in(time_embed) - time_embed = self.act_fn(time_embed) - time_embed = self.time_mlp_out(time_embed) - time_embed = self.act_fn(time_embed) - embed = action_embed - adarms_cond = time_embed - - if self.action_hidden_size < self.hidden_size: - # padding to hidden size - padding_size = self.hidden_size - self.action_hidden_size - padding = torch.zeros( - (embed.shape[0], embed.shape[1], padding_size), - device=embed.device, - dtype=embed.dtype, - ) - embed = torch.cat([embed, padding], dim=-1) - - return embed, adarms_cond - - def flow_loss( - self, - action_hidden_states, - flow, - action_chunk, - dof_mask=None, - flow_loss_mask=None, - ): - with torch.autocast("cuda", dtype=torch.float32): - action_pred = self.action_proj_back( - action_hidden_states[:, : self.action_hidden_size] - ) - v_pred = action_pred - loss = self.mse_loss(v_pred, flow) - if dof_mask is not None: - dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]) - loss = loss * dof_mask - - if flow_loss_mask is not None: - flow_loss_mask = ( - flow_loss_mask.unsqueeze(-1) - .reshape(-1, 1) - .expand(-1, loss.shape[-1]) - ) - loss = loss * flow_loss_mask - return loss diff --git a/wall_x/model/core/__init__.py b/wall_x/model/core/__init__.py new file mode 100644 index 0000000..de478ad --- /dev/null +++ b/wall_x/model/core/__init__.py @@ -0,0 +1 @@ +"""Core building blocks shared across model implementations.""" diff --git a/wall_x/model/core/action/__init__.py b/wall_x/model/core/action/__init__.py new file mode 100644 index 0000000..65582b1 --- /dev/null +++ b/wall_x/model/core/action/__init__.py @@ -0,0 +1,10 @@ +"""Action model components: normalizer, processor, head, MoE.""" + +from wall_x.model.core.action.head import SinusoidalPosEmb +from wall_x.model.core.action.normalizer import ( + Normalizer, + create_normalizers, + normalize_data_with_virtual_tail, + unnormalize_data_with_virtual_tail, +) +from wall_x.model.core.action.processor import ActionProcessor diff --git a/wall_x/model/core/action/head.py b/wall_x/model/core/action/head.py new file mode 100644 index 0000000..30a1229 --- /dev/null +++ b/wall_x/model/core/action/head.py @@ -0,0 +1,27 @@ +"""Action head helpers used by the VLA action processor.""" + +import math + +import torch +import torch.nn as nn + + +class SinusoidalPosEmb(nn.Module): + """Sinusoidal timestep embedding for action flow timesteps.""" + + def __init__(self, dim: int, min_period: float = 4e-3, max_period: float = 4.0): + super().__init__() + if dim % 2 != 0: + raise ValueError(f"embedding_dim ({dim}) must be divisible by 2") + self.dim = dim + self.min_period = min_period + self.max_period = max_period + + def forward(self, x): + half_dim = self.dim // 2 + exponent = math.log(10000) / (half_dim - 1) + frequencies = torch.exp( + torch.arange(half_dim, device=x.device, dtype=torch.float32) * -exponent + ) + emb = x[:, None] * frequencies[None, :] + return torch.cat((emb.sin(), emb.cos()), dim=-1) diff --git a/wall_x/model/core/action/moe.py b/wall_x/model/core/action/moe.py new file mode 100644 index 0000000..a57e77d --- /dev/null +++ b/wall_x/model/core/action/moe.py @@ -0,0 +1,126 @@ +import torch +import torch.nn as nn +import torch.utils.checkpoint as cp +from transformers.activations import ACT2FN + +from wall_x.model.core.ops import permute, unpermute + + +class TokenTypeRouter(nn.Module): + def __init__(self, num_experts: int): + super().__init__() + self.num_experts = num_experts + + def forward(self, token_types: torch.Tensor) -> torch.Tensor: + """ + Route tokens to experts based on token_type. + + Args: + token_types (torch.Tensor): Tensor of shape (batch_size, seq_length) containing each token type. + + Returns: + experts_indices (torch.Tensor): Tensor of shape (batch_size, seq_length) containing each assigned expert index. + """ + # Simple rule: assign by token_type modulo the expert count + experts_indices = token_types % self.num_experts + return experts_indices + + +class BlockSparseMLP(nn.Module): + def __init__(self, config, use_selective_recompute: bool = False): + super().__init__() + self.hidden_size = config["hidden_size"] + self.intermediate_size = config["intermediate_size"] + self.hidden_act = config["hidden_act"] + + self.use_selective_recompute = use_selective_recompute + + self.gate_up_proj = nn.Linear( + self.hidden_size, 2 * self.intermediate_size, bias=False + ) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + + self.act_fn = ACT2FN[self.hidden_act] + + # FIXME: The full MLP is recomputed for now; recomputing only activations can be optimized later. + def _full_mlp(self, hidden_state): + gate_up_out = self.gate_up_proj(hidden_state) + gate_out, up_out = gate_up_out.split( + [self.intermediate_size, self.intermediate_size], dim=-1 + ) + + act_out = self.act_fn(gate_out) * up_out + return self.down_proj(act_out) + + def forward(self, hidden_state): + if self.use_selective_recompute: + # Checkpoint-recompute the whole expert MLP + return cp.checkpoint( + self._full_mlp, + hidden_state, + use_reentrant=False, + ) + else: + return self._full_mlp(hidden_state) + + +class SparseMoeBlock(nn.Module): + def __init__(self, config, num_experts: int, use_selective_recompute: bool = False): + super().__init__() + self.num_experts = num_experts + self.use_selective_recompute = use_selective_recompute + + # Pass use_selective_recompute to each expert + self.experts = nn.ModuleList( + [ + BlockSparseMLP( + config.experts[i], use_selective_recompute=use_selective_recompute + ) + for i in range(num_experts) + ] + ) + + if not hasattr(config, "dim_inputs") or not config.dim_inputs: + raise ValueError("config.dim_inputs must be set") + + self.dim_inputs = config.dim_inputs + self.permuted = config.mot_opt + + def forward( + self, + hidden_states: torch.Tensor, + experts_indices: torch.Tensor, + start_indices: torch.Tensor, + end_indices: torch.Tensor, + ) -> torch.Tensor: + + if self.permuted: + permuted_inputs = hidden_states + else: + batch_size, seq_length, hidden_dim = hidden_states.shape + + flat_hidden = hidden_states.reshape(-1, hidden_dim) + experts_indices = experts_indices.reshape(-1) + probs = torch.ones_like(experts_indices, dtype=torch.float32).reshape(-1, 1) + permuted_inputs, row_id_map = permute(flat_hidden, experts_indices) + + # buffer + final_output = torch.zeros_like(permuted_inputs) + + # Expert forward, including selective recompute + for expert_idx, expert in enumerate(self.experts): + start, end = start_indices[expert_idx], end_indices[expert_idx] + if start == end: + continue + + dim_input = self.dim_inputs[expert_idx] + expert_input = permuted_inputs[start:end, :dim_input] + + partial_output = expert(expert_input) + final_output[start:end, :dim_input] = partial_output[:, :dim_input] + + if self.permuted: + return final_output + else: + final_output = unpermute(final_output, row_id_map, probs) + return final_output.reshape(batch_size, seq_length, hidden_dim) diff --git a/wall_x/model/core/action/normalizer.py b/wall_x/model/core/action/normalizer.py new file mode 100644 index 0000000..5c1a7d8 --- /dev/null +++ b/wall_x/model/core/action/normalizer.py @@ -0,0 +1,447 @@ +import json +import logging + +import torch +import torch.nn as nn + +from wall_x.utils.constant import _ACTION_KEY_FULL_MAPPING as _MODEL_KEY_TO_RAW_KEY + +logger = logging.getLogger(__name__) + + +def resolve_normalizer_dataset_names( + dataset_names, + normalizer, + batch_size, + *, + source="dataset_names", + allow_skip_names=None, +): + """Validate dataset names before indexing a normalizer ParameterDict. + + The normalizer key is data-dependent. Callers should pass the + dataset names carried by the batch/model inputs, matching the qact/VLA + path. This helper intentionally only validates presence, batch-size + alignment, and key membership; it does not remap unknown names or fall + back to a default key, since silent fallback can apply the wrong action + scale without failing loudly. Pass ``allow_skip_names`` explicitly only + for names that should bypass normalizer lookup, such as multimodal-only + rows that do not carry action values. + """ + if dataset_names is None: + raise KeyError(f"Missing {source}; cannot choose normalizer key.") + + if isinstance(dataset_names, str): + names = [dataset_names] + elif isinstance(dataset_names, (list, tuple)): + names = [str(name) for name in dataset_names] + else: + raise TypeError( + f"{source} must be a string or sequence of strings, got " + f"{type(dataset_names).__name__}" + ) + + if len(names) != batch_size: + raise ValueError( + f"{source} count ({len(names)}) does not match batch size " + f"({batch_size}). names={names}" + ) + + available = list(getattr(normalizer, "delta", {}).keys()) + if not available: + raise KeyError("Normalizer has no registered dataset keys.") + + allowed_skip = set(allow_skip_names or ()) + missing = sorted( + {name for name in names if name not in available and name not in allowed_skip} + ) + if missing: + raise KeyError( + f"{source} contains keys not present in normalizer: {missing}. " + f"available={available}" + ) + + return names + + +def _normalizer_dataset_name_list(dataset_names): + if dataset_names is None: + return [] + if isinstance(dataset_names, str): + return [dataset_names] + return list(dataset_names) + + +def _normalizer_width(normalizer, dataset_names): + names = _normalizer_dataset_name_list(dataset_names) + if not names: + return None + if any(name not in normalizer.delta for name in names): + return None + widths = [int(normalizer.delta[name].shape[0]) for name in names] + if not widths or len(set(widths)) != 1: + return None + return widths[0] + + +def _has_uniform_virtual_tail_mask(dof_mask, width): + if dof_mask is None: + return False + mask = dof_mask.reshape(-1, dof_mask.shape[-1]).bool() + if mask.shape[-1] <= width: + return False + prefix_active = mask[:, :width].all(dim=1) + tail_inactive = (~mask[:, width:]).all(dim=1) + return bool((prefix_active & tail_inactive).all().item()) + + +def normalize_data_with_virtual_tail(normalizer, values, dataset_names, dof_mask): + """Normalize real prefix dims when model tensors include a virtual tail.""" + width = _normalizer_width(normalizer, dataset_names) + if width is None or values.shape[-1] == width: + return normalizer.normalize_data(values, dataset_names) + if values.shape[-1] < width or not _has_uniform_virtual_tail_mask(dof_mask, width): + return normalizer.normalize_data(values, dataset_names) + + out = values.clone() + out[..., :width] = normalizer.normalize_data(values[..., :width], dataset_names) + out[..., width:] = 0 + return out + + +def unnormalize_data_with_virtual_tail(normalizer, values, dataset_names, dof_mask): + """Unnormalize real prefix dims when model tensors include a virtual tail.""" + width = _normalizer_width(normalizer, dataset_names) + if width is None or values.shape[-1] == width: + return normalizer.unnormalize_data(values, dataset_names, None) + if values.shape[-1] < width or not _has_uniform_virtual_tail_mask(dof_mask, width): + return normalizer.unnormalize_data(values, dataset_names, dof_mask) + + out = values.clone() + out[..., :width] = normalizer.unnormalize_data( + values[..., :width], dataset_names, None + ) + out[..., width:] = 0 + return out + + +def print_rank_last(message): + """If distributed is initialized, log only on last rank.""" + if torch.distributed.is_initialized(): + if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1): + logger.info(message) + else: + logger.info(message) + + +class Normalizer(nn.Module): + @classmethod + def from_ckpt(cls, ckpt_path): + instance = cls.__new__(cls) + nn.Module.__init__(instance) + + instance.min = nn.ParameterDict() + instance.delta = nn.ParameterDict() + instance.min_key = "min" + instance.delta_key = "delta" + + ckpt = torch.load(ckpt_path, map_location="cpu") + + for key, value in ckpt.items(): + # parse key: "min.robot_name" -> prefix="min", name="robot_name" + try: + prefix, name = key.split(".", 1) + if hasattr(instance, prefix): + getattr(instance, prefix)[name] = nn.Parameter( + value, requires_grad=False + ) + except ValueError: + continue + + return instance + + @classmethod + def from_lerobot_norm_stats( + cls, + action_stats, + dataset_name, + ): + + # Create instance without calling __init__ + instance = cls.__new__(cls) + nn.Module.__init__(instance) + + # Set containers + instance.min = nn.ParameterDict() + instance.delta = nn.ParameterDict() + + # Fill dataset entry + instance.min[dataset_name] = nn.Parameter(action_stats.min, requires_grad=False) + instance.delta[dataset_name] = nn.Parameter( + action_stats.delta, requires_grad=False + ) + + # Record keys + instance.min_key = "min" + instance.delta_key = "delta" + + return instance + + def __init__( + self, + action_statistic_dof, + dof_config, + min_key="min", + delta_key="delta", + name="normalizer", + ): + super(Normalizer, self).__init__() + + self.min_key = min_key + self.delta_key = delta_key + + action_statistic = {} + normalizer_missing_information = [] + for robot_name in action_statistic_dof.keys(): + action_statistic[robot_name] = {} + all_dof_min = [] + all_dof_delta = [] + for k in dof_config: + if ( + k not in action_statistic_dof[robot_name] + and k.replace("master_", "follow_", 1) + in action_statistic_dof[robot_name] + ): + k = k.replace("master_", "follow_", 1) + if k in action_statistic_dof[robot_name]: + if ( + min_key in action_statistic_dof[robot_name][k] + and delta_key in action_statistic_dof[robot_name][k] + ): + all_dof_min.extend(action_statistic_dof[robot_name][k][min_key]) + all_dof_delta.extend( + action_statistic_dof[robot_name][k][delta_key] + ) + else: + normalizer_missing_information.append( + f"Normalizer (Warning): min_key {min_key} or delta_key {delta_key} not in action_statistic_dof[{robot_name}][{k}], use default min 0.0 and delta 1.0" + ) + all_dof_min.extend([0.0] * dof_config[k]) + all_dof_delta.extend([1.0] * dof_config[k]) + else: + # k is a model key; action_statistic_dof may store raw keys - try fallback lookup + raw_key = _MODEL_KEY_TO_RAW_KEY.get(k) + if ( + raw_key is not None + and raw_key in action_statistic_dof[robot_name] + ): + stat = action_statistic_dof[robot_name][raw_key] + if min_key in stat and delta_key in stat: + all_dof_min.extend(stat[min_key]) + all_dof_delta.extend(stat[delta_key]) + else: + normalizer_missing_information.append( + f"Normalizer (Warning): Action {k} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" + ) + all_dof_min.extend([0.0] * dof_config[k]) + all_dof_delta.extend([1.0] * dof_config[k]) + else: + normalizer_missing_information.append( + f"Normalizer (Warning): Action {k} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" + ) + all_dof_min.extend([0.0] * dof_config[k]) + all_dof_delta.extend([1.0] * dof_config[k]) + all_dof_min = torch.tensor(all_dof_min) + all_dof_delta = torch.tensor(all_dof_delta) + action_statistic[robot_name][min_key] = all_dof_min + action_statistic[robot_name][delta_key] = all_dof_delta + + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == ( + torch.distributed.get_world_size() - 1 + ): + # export normalizer_missing_information to file + with open(f"normalizer_missing_information_{name}.txt", "w") as f: + for info in normalizer_missing_information: + f.write(info + "\n") + logger.info( + "Normalizer missing information saved to normalizer_missing_information_%s.txt", + name, + ) + + self.min = nn.ParameterDict( + { + k: nn.Parameter(action_statistic[k][min_key], requires_grad=False) + for k in action_statistic.keys() + } + ) + self.delta = nn.ParameterDict( + { + k: nn.Parameter(action_statistic[k][delta_key], requires_grad=False) + for k in action_statistic.keys() + } + ) + + def normalize_data(self, xs, dataset_names): + new_xs = [] + dataset_names = [name for name in dataset_names if name != "x2_multimodal"] + for x, dataset_name in zip(xs, dataset_names): + # if dataset_name == "ex_normal": + # dataset_name = "x2_normal" + x = (x - self.min[dataset_name]) / (self.delta[dataset_name]) + x = x * 2 - 1 + x = torch.clamp(x, -1, 1) + new_xs.append(x) + new_xs = torch.stack(new_xs) + return new_xs + + def unnormalize_data(self, xs, dataset_names, dof_mask=None): + new_xs = [] + dataset_names = [name for name in dataset_names if name != "x2_multimodal"] + dof_mask = dof_mask if dof_mask is not None else [None] * len(xs) + for x, dataset_name, mask in zip(xs, dataset_names, dof_mask): + x = (x + 1) / 2 + if mask is not None: + mask = mask[0].bool() + action_space_delta = self.delta[dataset_name][mask] + action_space_min = self.min[dataset_name][mask] + else: + action_space_delta = self.delta[dataset_name] + action_space_min = self.min[dataset_name] + x = x * action_space_delta + action_space_min + new_xs.append(x) + new_xs = torch.stack(new_xs) + return new_xs + + +def _pad_lerobot_stats(stats, target_dim, label): + """Tail-pad flat LeRobot stats to the model's configured action/state dim.""" + target_dim = int(target_dim) + current_dim = int(stats.min.numel()) + if current_dim == target_dim: + return stats + if current_dim > target_dim: + raise ValueError( + f"LeRobot {label} norm stats dim ({current_dim}) exceeds configured " + f"{label} dim ({target_dim})" + ) + + pad_dim = target_dim - current_dim + min_stat = torch.cat( + [stats.min.flatten(), torch.zeros(pad_dim, dtype=stats.min.dtype)] + ) + delta = torch.cat( + [stats.delta.flatten(), torch.ones(pad_dim, dtype=stats.delta.dtype)] + ) + max_stat = torch.cat( + [stats.max.flatten(), torch.ones(pad_dim, dtype=stats.max.dtype)] + ) + logger.info( + "Padded LeRobot %s normalizer stats from %d to %d dims", + label, + current_dim, + target_dim, + ) + return type(stats)(min=min_stat, max=max_stat, delta=delta) + + +def pad_normalizer_to_dim(normalizer, target_dim, label): + """Tail-pad every dataset entry to ``target_dim`` (virtual ``action_padding`` tail).""" + target_dim = int(target_dim) + for name in list(normalizer.min.keys()): + current_dim = int(normalizer.min[name].numel()) + if current_dim == target_dim: + continue + if current_dim > target_dim: + raise ValueError( + f"Normalizer {label}[{name!r}] dim ({current_dim}) exceeds configured " + f"{label} dim ({target_dim})" + ) + + pad_dim = target_dim - current_dim + dtype = normalizer.min[name].dtype + normalizer.min[name] = nn.Parameter( + torch.cat( + [ + normalizer.min[name].flatten(), + torch.zeros(pad_dim, dtype=dtype), + ] + ), + requires_grad=False, + ) + normalizer.delta[name] = nn.Parameter( + torch.cat( + [ + normalizer.delta[name].flatten(), + torch.ones(pad_dim, dtype=dtype), + ] + ), + requires_grad=False, + ) + logger.info( + "Padded %s normalizer[%r] from %d to %d dims", + label, + name, + current_dim, + target_dim, + ) + + +def create_normalizers_from_lerobot_norm_stats( + norm_stats, dataset_name, action_dim, propri_dim +): + """Create model normalizers from LeRobot flat action/state norm stats.""" + action_stats = _pad_lerobot_stats(norm_stats["action"], action_dim, "action") + propri_stats = _pad_lerobot_stats(norm_stats["state"], propri_dim, "state") + return ( + Normalizer.from_lerobot_norm_stats(action_stats, dataset_name), + Normalizer.from_lerobot_norm_stats(propri_stats, dataset_name), + ) + + +def create_normalizers(config, action_statistic_dof=None): + """Create action and proprioception normalizers from explicit stats. + + Normalization statistics must come from the config, checkpoint, or dataset. + Public Wall-X builds intentionally do not bundle private default stats. + """ + + if action_statistic_dof is None: + custom_path = config.get("customized_action_statistic_dof") + if custom_path: + with open(custom_path, "r") as f: + action_statistic_dof = json.load(f) + else: + raise ValueError( + "create_normalizers requires action statistics. Provide " + "`customized_action_statistic_dof` in the config or pass " + "`action_statistic_dof` from the checkpoint/dataset." + ) + + min_key = config.get("min_key", "min") + delta_key = config.get("delta_key", "delta") + + # Required keys for constructing action/proprio normalizers. + missing_required = [] + if config.get("dof_config") is None: + missing_required.append("dof_config") + if config.get("agent_pos_config") is None: + missing_required.append("agent_pos_config") + if missing_required: + raise KeyError( + "create_normalizers requires non-None config keys: " + + ", ".join(missing_required) + ) + + normalizer_action = Normalizer( + action_statistic_dof, + config["dof_config"], + min_key=min_key, + delta_key=delta_key, + ) + normalizer_propri = Normalizer( + action_statistic_dof, + config["agent_pos_config"], + min_key=min_key, + delta_key=delta_key, + ) + return normalizer_action, normalizer_propri, action_statistic_dof diff --git a/wall_x/model/core/action/processor.py b/wall_x/model/core/action/processor.py new file mode 100644 index 0000000..358352f --- /dev/null +++ b/wall_x/model/core/action/processor.py @@ -0,0 +1,317 @@ +import torch +import torch.nn as nn + +from wall_x.model.core.action.head import SinusoidalPosEmb +from wall_x.model.core.action.normalizer import print_rank_last + + +class ActionProcessor(nn.Module): + """ + Action processor + + Main responsibilities: + 1. Add noise to action sequences + 2. Generate time encodings + 3. Project actions into the model hidden space + + Uses a Beta distribution to control noise scheduling and provide flexible noise injection. + """ + + def __init__(self, config): + """ + Args: + config: Configuration object containing: + - action_dim: action-space dimension + - hidden_size: model hidden size + - noise_scheduler: noise scheduler configuration + """ + super().__init__() + self.config = config + self.dof_config = config.dof_config + self.agent_pos_config = config.agent_pos_config + self.action_dim = sum([v for k, v in self.dof_config.items()]) + self.propri_dim = sum([v for k, v in self.agent_pos_config.items()]) + + print_rank_last( + f"self.dof_config: {self.dof_config}; action_dim: {self.action_dim}; self.agent_pos_config: {self.agent_pos_config}; propri_dim: {self.propri_dim}" + ) + + self.action_hidden_size = config.action_hidden_size + self.state_hidden_size = config.state_hidden_size + self.hidden_size = getattr(config, "hidden_size", config.dim_inputs[0]) + + if not self.config.use_state_string_representation: + if self.config.proj_with_mask: + self.propri_proj = nn.Linear( + self.propri_dim * 2, self.state_hidden_size, bias=False + ) + else: + self.propri_proj = nn.Linear( + self.propri_dim, self.state_hidden_size, bias=False + ) + + # noise scheduler configing + if getattr(self.config, "use_flow_action_expert", True): + noise_scheduler_config = config.noise_scheduler + self.s = noise_scheduler_config.get("s", 0.999) + self.time_shift = noise_scheduler_config.get( + "time_shift", 1.0 + ) # time shift factor + self.time_embed = SinusoidalPosEmb(self.action_hidden_size) + + # project to hidden space + if self.config.proj_with_mask: + self.w1 = nn.Linear( + self.action_dim * 2, self.action_hidden_size, bias=False + ) + else: + self.w1 = nn.Linear( + self.action_dim, self.action_hidden_size, bias=False + ) + if not self.config.use_adarms: + self.w2 = nn.Linear( + self.action_hidden_size * 2, self.action_hidden_size, bias=False + ) + self.w3 = nn.Linear( + self.action_hidden_size, self.action_hidden_size, bias=False + ) + self.act_fn = nn.SiLU() + else: + self.time_mlp_in = nn.Linear( + self.action_hidden_size, self.action_hidden_size + ) + self.time_mlp_out = nn.Linear( + self.action_hidden_size, self.action_hidden_size + ) + self.act_fn = nn.SiLU() + + # project back to action space + self.action_proj_back = nn.Linear( + self.action_hidden_size, self.action_dim, bias=False + ) + self.mse_loss = nn.MSELoss(reduction="none") + + def set_normalizer(self, normalizer_action, normalizer_propri): + self.normalizer_action = normalizer_action + self.normalizer_propri = normalizer_propri + + def get_inference_times(self, num_steps, device, dtype): + """ + Get inference timesteps + + Apply time shift and scaling + + Args: + num_steps (int): number of inference steps + device: Device type + dtype: dtype + + Returns: + torch.Tensor: inference timestep sequence + """ + times = torch.linspace(0.0, 1.0, num_steps + 1, device=device, dtype=dtype) + if self.time_shift != 1.0: + times = (self.time_shift * times) / (1 + (self.time_shift - 1) * times) + times = times * self.s + return times + + def proprioception_proj( + self, proprioception, dataset_names=None, dof_mask=None, use_history=False + ): + """ + Args: + proprioception: [batch_size, 1, action_dim] + dataset_names: [batch_size] + dof_mask: [batch_size, action_dim] + """ + with torch.autocast("cuda", dtype=torch.float32): + proprioception = proprioception.to( + device=self.propri_proj.weight.device + ).to(dtype=self.propri_proj.weight.dtype) + if dof_mask is not None: + if self.config.proj_with_mask: + proprioception = torch.cat( + [proprioception, dof_mask], dim=-1 + ) # .unsqueeze(1) + proprioception = proprioception.to( + device=self.propri_proj.weight.device + ).to(dtype=self.propri_proj.weight.dtype) + proprio_embed = self.propri_proj( + proprioception + ) # [batch_size, 1, state_hidden_size] + if self.state_hidden_size < self.hidden_size: + # padding to hidden size + padding_size = self.hidden_size - self.state_hidden_size + padding = torch.zeros( + (proprio_embed.shape[0], 1, padding_size), + device=proprio_embed.device, + dtype=proprio_embed.dtype, + ) + proprio_embed = torch.cat([proprio_embed, padding], dim=-1) + return proprio_embed # [batch_size, 1, hidden_size] + + def forward(self, action_chunk, dataset_names, sample_time, dof_mask=None): + """ + Args: + action_chunk (torch.Tensor): action sequence with shape [batch_size, action_chunk_len, action_dim] + dataset_names: [batch_size] + dof_mask: [batch_size, action_dim] + + Returns: + torch.Tensor: processed action representation with shape [batch_size, seq_len, hidden_size] + """ + with torch.autocast("cuda", dtype=torch.float32): + action_chunk = action_chunk.to(dtype=torch.float32) + + # 1. add noise to action_chunk + noise = torch.randn_like(action_chunk) + time_expanded = sample_time.unsqueeze(-1).unsqueeze(-1) + noisy_action = ( + 1 - time_expanded + ) * noise + time_expanded * action_chunk # denoise from 0 to 1; integration does not need a negative sign + flow = action_chunk - noise # used to compute loss + + # 2. sinusoidal positional encoding for timesteps + time_embed = self.time_embed(sample_time).to(torch.float32) + + self.noise = noise + self.noisy_action = noisy_action # for new x-pred + # 3.action_chunk_nosiy + t_pos_emb -> MLP_act_chunk -> action_chunk_nosiy_emb_with_t (dim=trans * chunk) + if dof_mask is not None: + noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) + + noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) + action_embed = self.w1(noisy_action) + + self.time_expanded = time_expanded # for new x-pred + + if not self.config.use_adarms: + time_embed = ( + time_embed.unsqueeze(1) + .repeat(1, action_embed.shape[1], 1) + .to(dtype=self.w2.weight.dtype) + ) + concat_embed = torch.cat([action_embed, time_embed], dim=-1) + concat_embed = self.w2(concat_embed) + action_time_embed = self.w3(self.act_fn(concat_embed)) + adarms_cond = None + else: + time_embed = self.time_mlp_in(time_embed) + time_embed = self.act_fn(time_embed) + time_embed = self.time_mlp_out(time_embed) + time_embed = self.act_fn(time_embed) + action_time_embed = action_embed + adarms_cond = time_embed + + if self.action_hidden_size < self.hidden_size: + # padding to hidden size + padding_size = self.hidden_size - self.action_hidden_size + padding = torch.zeros( + ( + action_time_embed.shape[0], + action_time_embed.shape[1], + padding_size, + ), + device=action_time_embed.device, + dtype=action_time_embed.dtype, + ) + action_time_embed = torch.cat([action_time_embed, padding], dim=-1) + + return action_time_embed, flow, adarms_cond + + def step(self, timestep, noisy_action, dof_mask=None): + # noisy_action: bs, pred_horizon, action_dim + # timestep: bs + with torch.autocast("cuda", dtype=torch.float32): + if dof_mask is not None and self.config.proj_with_mask: + noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) + + noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) + time_embed = self.time_embed(timestep).to(torch.float32) # bs,hidden_size + action_embed = self.w1(noisy_action) + + if not self.config.use_adarms: + time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1) + time_embed = time_embed.to(device=noisy_action.device).to( + dtype=noisy_action.dtype + ) + concat_embed = torch.cat([action_embed, time_embed], dim=-1) + concat_embed = self.w2(concat_embed) + embed = self.w3(self.act_fn(concat_embed)) # is this right? + adarms_cond = None + else: + time_embed = time_embed.to(dtype=self.time_mlp_in.weight.dtype) + time_embed = self.time_mlp_in(time_embed) + time_embed = self.act_fn(time_embed) + time_embed = self.time_mlp_out(time_embed) + time_embed = self.act_fn(time_embed) + embed = action_embed + adarms_cond = time_embed + + if self.action_hidden_size < self.hidden_size: + # padding to hidden size + padding_size = self.hidden_size - self.action_hidden_size + padding = torch.zeros( + (embed.shape[0], embed.shape[1], padding_size), + device=embed.device, + dtype=embed.dtype, + ) + embed = torch.cat([embed, padding], dim=-1) + + return embed, adarms_cond + + def flow_loss( + self, + action_hidden_states, + flow, + action_chunk, + dof_mask=None, + flow_loss_mask=None, + ): + with torch.autocast("cuda", dtype=torch.float32): + action_pred = self.action_proj_back( + action_hidden_states[:, : self.action_hidden_size] + ) + + if getattr(self.config, "use_x_pred", False): + noisy_action_flat = self.noisy_action.reshape( + -1, self.noisy_action.shape[-1] + ) + time_expanded_flat = self.time_expanded.expand( + -1, self.noisy_action.shape[1], -1 + ).reshape(-1, 1) + v_pred = (action_pred - noisy_action_flat) / torch.clamp( + 1 - time_expanded_flat, min=0.05 + ) + x_pred = action_pred + else: + v_pred = action_pred + time_expanded_flat = self.time_expanded.expand( + -1, self.noisy_action.shape[1], -1 + ).reshape(-1, 1) + x_pred = (1 - time_expanded_flat) * v_pred + self.noisy_action.reshape( + -1, self.noisy_action.shape[-1] + ) + + if getattr(self.config, "use_x_loss", False): + loss = self.mse_loss( + x_pred, + action_chunk.reshape(-1, action_chunk.shape[-1]).to( + dtype=x_pred.dtype + ), + ) + else: + loss = self.mse_loss(v_pred, flow) + + if dof_mask is not None: + dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]) + loss = loss * dof_mask + + if flow_loss_mask is not None: + flow_loss_mask = ( + flow_loss_mask.unsqueeze(-1) + .reshape(-1, 1) + .expand(-1, loss.shape[-1]) + ) + loss = loss * flow_loss_mask + return loss diff --git a/wall_x/model/core/attention/__init__.py b/wall_x/model/core/attention/__init__.py new file mode 100644 index 0000000..cd791d1 --- /dev/null +++ b/wall_x/model/core/attention/__init__.py @@ -0,0 +1,9 @@ +"""Attention mechanisms: joint attention, VLA attention, mask builders, backend selector.""" + +from wall_x.model.core.attention.mask import ( + find_first_last_ones, + update_joint_attention_flash_mask, + update_joint_attention_mask_2d, + update_position_ids, +) +from wall_x.model.core.attention.selector import AttentionsSelectorMixin diff --git a/wall_x/model/joint_attention.py b/wall_x/model/core/attention/joint.py similarity index 56% rename from wall_x/model/joint_attention.py rename to wall_x/model/core/attention/joint.py index 5ba9546..aae7433 100644 --- a/wall_x/model/joint_attention.py +++ b/wall_x/model/core/attention/joint.py @@ -1,45 +1,36 @@ -import torch -import torch.nn as nn from typing import Optional, Tuple -from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig -from transformers.cache_utils import Cache -from transformers.utils import logging -from wall_x.fusions import ops -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import ( - apply_multimodal_rotary_pos_emb, -) +import torch +import torch.nn as nn from flash_attn import flash_attn_func +from transformers.cache_utils import Cache from transformers.modeling_flash_attention_utils import ( is_flash_attn_greater_or_equal_2_10, ) +from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2_5_VLRotaryEmbedding, repeat_kv, ) +from transformers.utils import logging + +from wall_x.model.core.attention.mask import find_first_last_ones +from wall_x.model.core.ops import m_rope, permute, unpermute + +try: + from flash_mask.flash_mask_interface import flash_mask_attn_func +except ImportError: + flash_mask_attn_func = None + +try: + from flash_mask.flash_mask_interface import flashmask_attn_func_stop_gradient +except ImportError: + flashmask_attn_func_stop_gradient = None + logger = logging.get_logger(__name__) -# def rotate_half(x): -# x1 = x[..., : x.shape[-1] // 2] -# x2 = x[..., x.shape[-1] // 2 :] -# return torch.cat((-x2, x1), dim=-1) - - -# def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=2): -# mrope_section = mrope_section * 2 -# cos_split = torch.cat( -# [m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1 -# ).unsqueeze(unsqueeze_dim) -# sin_split = torch.cat( -# [m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1 -# ).unsqueeze(unsqueeze_dim) -# q_embed = (q * cos_split) + (rotate_half(q) * sin_split) -# k_embed = (k * cos_split) + (rotate_half(k) * sin_split) -# return q_embed, k_embed - - class JointQwen2VLAttention(nn.Module): def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): super().__init__() @@ -52,7 +43,7 @@ class JointQwen2VLAttention(nn.Module): "when creating this class." ) if not hasattr(config, "dim_inputs") or not config.dim_inputs: - raise ValueError("Configuration must contain a valid dim_inputs") + raise ValueError("config.dim_inputs must be set") self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads @@ -69,30 +60,18 @@ class JointQwen2VLAttention(nn.Module): self.dim_inputs = config.dim_inputs # Tuple[int, ...] - if config.model_type == "qwen2_5_vl": - bias_qkv = True - else: - bias_qkv = False + if config.model_type != "qwen2_5_vl": + raise NotImplementedError(f"Unsupported model type: {config.model_type}") + bias_qkv = True - self.q_proj_experts = nn.ModuleList( - [ - nn.Linear(dim_input, self.num_heads * self.head_dim, bias=bias_qkv) - for dim_input in self.dim_inputs - ] + qkv_out_features = ( + self.num_heads * self.head_dim + + 2 * self.num_key_value_heads * self.head_dim ) - self.k_proj_experts = nn.ModuleList( + + self.qkv_proj_experts = nn.ModuleList( [ - nn.Linear( - dim_input, self.num_key_value_heads * self.head_dim, bias=bias_qkv - ) - for dim_input in self.dim_inputs - ] - ) - self.v_proj_experts = nn.ModuleList( - [ - nn.Linear( - dim_input, self.num_key_value_heads * self.head_dim, bias=bias_qkv - ) + nn.Linear(dim_input, qkv_out_features, bias=bias_qkv) for dim_input in self.dim_inputs ] ) @@ -103,11 +82,7 @@ class JointQwen2VLAttention(nn.Module): ] ) - # Rotary embedding init - if config.model_type == "qwen2_5_vl": - self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) - else: - raise NotImplementedError(f"Unsupported model type: {config.model_type}") + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) def repeat_kv(self, hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ @@ -129,6 +104,10 @@ class JointQwen2VLAttention(nn.Module): return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim) + @property + def _projection_dtype(self): + return self.qkv_proj_experts[0].weight.dtype + def forward( self, hidden_states: torch.Tensor, @@ -147,10 +126,10 @@ class JointQwen2VLAttention(nn.Module): orig_shape: Optional[Tuple[int]] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if token_types is None: - raise ValueError("token_types can not be None") + raise ValueError("token_types must not be empty") if token_types.max() >= len(self.dim_inputs): raise ValueError( - f"token_types contains invalid expert indices: {token_types.max()}" + f"token_types contains an invalid expert index: {token_types.max()}" ) if self.config.mot_opt: @@ -195,24 +174,42 @@ class JointQwen2VLAttention(nn.Module): key_states, value_states, self.layer_idx, cache_kwargs ) else: - past_key_states, past_value_states = past_key_value[self.layer_idx] + # Compatible across transformers versions: + # v5.x: DynamicCache uses .layers[idx].keys/.values + # v4.x: DynamicCache uses .key_cache[idx]/.value_cache[idx] + # old: Cache object is subscriptable, returns (key, value) tuple + if hasattr(past_key_value, "layers"): + past_key_states = past_key_value.layers[self.layer_idx].keys + past_value_states = past_key_value.layers[self.layer_idx].values + elif hasattr(past_key_value, "key_cache"): + past_key_states = past_key_value.key_cache[self.layer_idx] + past_value_states = past_key_value.value_cache[self.layer_idx] + else: + past_key_states, past_value_states = past_key_value[self.layer_idx] key_states = torch.cat([past_key_states, key_states], dim=-2) value_states = torch.cat([past_value_states, value_states], dim=-2) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) + target_dtype = self._projection_dtype + if ( + query_states.dtype != target_dtype + or key_states.dtype != target_dtype + or value_states.dtype != target_dtype + ): + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + causal_mask = attention_mask if attention_mask is not None: - # Ensure that the attention_mask correctly matches across the head dimension. if len(attention_mask.shape) == 2: # [batch_size, seq_len] - # Expanded to a causal mask format of [batch_size, 1, seq_len, seq_len] bsz, seq_len = attention_mask.shape causal_mask = attention_mask.view(bsz, 1, 1, seq_len).expand( bsz, 1, seq_len, seq_len ) elif len(attention_mask.shape) == 3: # [batch_size, seq_len, seq_len] - # add head dimension: [batch_size, 1, seq_len, seq_len] causal_mask = attention_mask.unsqueeze(1) elif ( len(attention_mask.shape) == 4 @@ -220,10 +217,9 @@ class JointQwen2VLAttention(nn.Module): causal_mask = attention_mask else: raise ValueError( - f"Unsupported attention_mask dim: {attention_mask.shape}" + f"Unsupported attention_mask shape: {attention_mask.shape}" ) - # convert the attention mask to bool type causal_mask = causal_mask.to(torch.bool) # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, @@ -261,6 +257,8 @@ class JointQwen2VLAttention(nn.Module): attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, q_len, -1) + if attn_output.dtype != target_dtype: + attn_output = attn_output.to(target_dtype) if self.config.mot_opt: output = self._generate_output_mot_opt( @@ -269,7 +267,32 @@ class JointQwen2VLAttention(nn.Module): else: output = self._generate_output(attn_output, masks) - return output, None, past_key_value + attention_map = None + if output_attentions: + with torch.no_grad(): + action_token_num = int((token_types > 0).sum()) + action_query_states = query_states[:, :, -action_token_num:] + scale = 1.0 / torch.sqrt( + torch.tensor( + self.head_dim, device=hidden_states.device, dtype=torch.float32 + ) + ) + attention_score = ( + torch.matmul(action_query_states, key_states.transpose(-2, -1)) + * scale + ) + mask = causal_mask[:, :, -action_token_num:].expand( + -1, attention_score.shape[1], -1, -1 + ) # Mask only queries used for actions + if mask.dtype != attention_score.dtype: + mask = mask.to(dtype=attention_score.dtype) + attention_score = attention_score.masked_fill( + ~(mask.bool()), float("-inf") + ) + attention_score = torch.softmax(attention_score, dim=-1) + attention_map = attention_score[0].mean(0) + + return output, attention_map, past_key_value def _generate_qkv(self, hidden_states, masks): bsz, q_len, _ = hidden_states.size() @@ -300,29 +323,24 @@ class JointQwen2VLAttention(nn.Module): ) # for expert_idx in range(len(self.dim_inputs)): - for expert_idx, (q_proj, k_proj, v_proj, mask) in enumerate( - zip(self.q_proj_experts, self.k_proj_experts, self.v_proj_experts, masks) + for expert_idx, (qkv_proj, mask) in enumerate( + zip(self.qkv_proj_experts, masks) ): if not mask.any(): continue dim_input = self.dim_inputs[expert_idx] selected_hidden = hidden_states[mask].clone() - - q_out = q_proj(selected_hidden[:, :dim_input]).view( - -1, self.num_heads, self.head_dim + if selected_hidden.dtype != qkv_proj.weight.dtype: + selected_hidden = selected_hidden.to(qkv_proj.weight.dtype) + qkv_out = qkv_proj(selected_hidden[:, :dim_input]).view( + -1, self.num_heads + 2 * self.num_key_value_heads, self.head_dim ) - k_out = k_proj(selected_hidden[:, :dim_input]).view( - -1, self.num_key_value_heads, self.head_dim - ) - v_out = v_proj(selected_hidden[:, :dim_input]).view( - -1, self.num_key_value_heads, self.head_dim - ) - - if self.config.model_type == "qwen3_vl_text": - q_out = self.q_norms[expert_idx](q_out)[0] - k_out = self.k_norms[expert_idx](k_out)[0] - + q_out = qkv_out[:, : self.num_heads, :] + k_out = qkv_out[ + :, self.num_heads : self.num_heads + self.num_key_value_heads, : + ] + v_out = qkv_out[:, self.num_heads + self.num_key_value_heads :, :] query_states[mask] = q_out key_states[mask] = k_out value_states[mask] = v_out @@ -377,26 +395,22 @@ class JointQwen2VLAttention(nn.Module): ) # === Each expert processes its own token slice === - for expert_idx, (q_proj, k_proj, v_proj) in enumerate( - zip(self.q_proj_experts, self.k_proj_experts, self.v_proj_experts) - ): + for expert_idx, qkv_proj in enumerate(self.qkv_proj_experts): start, end = start_indices[expert_idx], end_indices[expert_idx] if start == end: continue dim_input = self.dim_inputs[expert_idx] expert_input = hidden_states[start:end, :dim_input] + if expert_input.dtype != qkv_proj.weight.dtype: + expert_input = expert_input.to(qkv_proj.weight.dtype) # Compute Q/K/V - q_out = q_proj(expert_input) - k_out = k_proj(expert_input) - v_out = v_proj(expert_input) - - if getattr(self.config, "model_type", None) == "qwen3_vl_text": - q_out = self.q_norms[expert_idx](q_out) - q_out = q_out[0] if isinstance(q_out, (tuple, list)) else q_out - k_out = self.k_norms[expert_idx](k_out) - k_out = k_out[0] if isinstance(k_out, (tuple, list)) else k_out + qkv_out = qkv_proj(expert_input) + kv_dim = self.num_key_value_heads * self.head_dim + q_out, k_out, v_out = torch.split( + qkv_out, [self.num_heads * self.head_dim, kv_dim, kv_dim], dim=-1 + ) q_buffer[start:end] = q_out k_buffer[start:end] = k_out @@ -404,9 +418,9 @@ class JointQwen2VLAttention(nn.Module): # === Restore tokens to the original order === # unpermute (using the same unpermute operation) - q_unpermuted = ops.unpermute(q_buffer, row_id_map, probs) - k_unpermuted = ops.unpermute(k_buffer, row_id_map, probs) - v_unpermuted = ops.unpermute(v_buffer, row_id_map, probs) + q_unpermuted = unpermute(q_buffer, row_id_map, probs) + k_unpermuted = unpermute(k_buffer, row_id_map, probs) + v_unpermuted = unpermute(v_buffer, row_id_map, probs) # === Reshape to final form === query_states = q_unpermuted.view( @@ -424,19 +438,14 @@ class JointQwen2VLAttention(nn.Module): def _apply_rotary_pos_embed( self, query_states, key_states, cos, sin, unsqueeze_dim=1 ): - if self.config.model_type == "qwen2_5_vl": - query_states, key_states = apply_multimodal_rotary_pos_emb( - query_states.contiguous(), - key_states.contiguous(), - cos.contiguous(), - sin.contiguous(), - self.rope_scaling["mrope_section"], - unsqueeze_dim, - ) - else: - raise NotImplementedError( - f"Unsupported model type: {self.config.model_type}" - ) + del unsqueeze_dim + query_states, key_states = m_rope( + query_states.contiguous(), + key_states.contiguous(), + cos[..., : (cos.size(3) // 2)].contiguous().float(), + sin[..., : (sin.size(3) // 2)].contiguous().float(), + self.rope_scaling["mrope_section"], + ) return query_states, key_states def _generate_output(self, attn_output, masks): @@ -451,16 +460,18 @@ class JointQwen2VLAttention(nn.Module): continue dim_input = self.dim_inputs[expert_idx] - # Obtain all necessary indexes in a single operation. - mask_indices = mask.nonzero(as_tuple=False) + mask_indices = mask.nonzero( + as_tuple=False + ) # more efficient index retrieval if mask_indices.numel() == 0: continue batch_indices = mask_indices[:, 0] seq_indices = mask_indices[:, 1] - # Use advanced indexing directly to avoid intermediate tensors. selected_attn_output = attn_output[batch_indices, seq_indices] + if selected_attn_output.dtype != o_proj.weight.dtype: + selected_attn_output = selected_attn_output.to(o_proj.weight.dtype) projected_output = o_proj(selected_attn_output) output[batch_indices, seq_indices, :dim_input] = projected_output @@ -497,7 +508,7 @@ class JointQwen2VLAttention(nn.Module): # === 1. Flatten and reorder by expert assignment === flat_attn_output = attn_output.view(-1, hidden_dim) # [B*S, H] flat_expert_indices = experts_indices.reshape(-1) # [B*S] - permuted_inputs, _ = ops.permute(flat_attn_output, flat_expert_indices) + permuted_inputs, _ = permute(flat_attn_output, flat_expert_indices) total_tokens = permuted_inputs.shape[0] # === 2. Initialize output buffer (still in permuted token space) === @@ -513,6 +524,8 @@ class JointQwen2VLAttention(nn.Module): dim_input = self.dim_inputs[expert_idx] expert_input = permuted_inputs[start:end] # [N_e, dim_input] + if expert_input.dtype != o_proj.weight.dtype: + expert_input = expert_input.to(o_proj.weight.dtype) expert_output = o_proj(expert_input) # [N_e, hidden_dim] # Write results into the buffer (overwrite only valid dimension region) @@ -523,13 +536,14 @@ class JointQwen2VLAttention(nn.Module): class JointQwen2VLFlashAttention(JointQwen2VLAttention): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): + super().__init__(config, layer_idx) # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + self.deterministic = config.attn_deterministic def forward( self, @@ -552,7 +566,10 @@ class JointQwen2VLFlashAttention(JointQwen2VLAttention): ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if token_types is None: - raise ValueError("token_types cannot be empty") + raise ValueError("token_types must not be empty") + # This check will lead to cudastreamsync. + # if token_types.max() >= len(self.dim_inputs): + # raise ValueError(f"token_types contains an invalid expert index: {token_types.max()}") if self.config.mot_opt: bsz, q_len, _ = orig_shape @@ -611,7 +628,7 @@ class JointQwen2VLFlashAttention(JointQwen2VLAttention): elif hasattr(self.config, "_pre_quantization_dtype"): target_dtype = self.config._pre_quantization_dtype else: - target_dtype = self.q_proj.weight.dtype + target_dtype = self._projection_dtype logger.warning_once( f"The input hidden states seems to be silently casted in float32, this might be related to" @@ -630,6 +647,7 @@ class JointQwen2VLFlashAttention(JointQwen2VLAttention): dropout_rate, softmax_scale=None, causal=self.is_causal, + deterministic=self.deterministic, ) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() @@ -644,8 +662,265 @@ class JointQwen2VLFlashAttention(JointQwen2VLAttention): return output, None, past_key_value +class JointQwen2VLFlashMaskAttention(JointQwen2VLAttention): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + token_types: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + probs: Optional[torch.Tensor] = None, + row_id_map: Optional[torch.Tensor] = None, + orig_shape: Optional[Tuple[int]] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if token_types is None: + raise ValueError("token_types must not be empty") + if token_types.max() >= len(self.dim_inputs): + raise ValueError( + f"token_types contains an invalid expert index: {token_types.max()}" + ) + + if self.config.mot_opt: + bsz, q_len, _ = orig_shape + query_states, key_states, value_states = self._generate_qkv_mot_opt( + hidden_states, + token_types, + start_indices, + end_indices, + probs, + row_id_map, + bsz, + q_len, + ) + else: + bsz, q_len, _ = hidden_states.size() + masks = [ + (token_types == expert_idx) + for expert_idx in range(len(self.dim_inputs)) + ] + query_states, key_states, value_states = self._generate_qkv( + hidden_states, masks + ) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = self._apply_rotary_pos_embed( + query_states, key_states, cos, sin, unsqueeze_dim=2 + ) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = self.repeat_kv(key_states, self.num_key_value_groups) + value_states = self.repeat_kv(value_states, self.num_key_value_groups) + # dropout_rate = 0.0 if not self.training else self.attention_dropout + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self._projection_dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Expand the attention_mask head dimension from 1 to num_heads + if attention_mask is not None and attention_mask.shape[1] == 1: + attention_mask = attention_mask.expand( + -1, self.num_heads, -1, -1 + ).contiguous() + + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + attn_output = flash_mask_attn_func( + query_states, + key_states, + value_states, + startend_row_indices=attention_mask, + causal=False, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + + if self.config.mot_opt: + output = self._generate_output_mot_opt( + attn_output, token_types, start_indices, end_indices + ) + else: + output = self._generate_output(attn_output, masks) + + return output, None, past_key_value + + +class JointQwen2VLFlashMaskAttention_KI(JointQwen2VLAttention): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + token_types: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + probs: Optional[torch.Tensor] = None, + row_id_map: Optional[torch.Tensor] = None, + orig_shape: Optional[Tuple[int]] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if token_types is None: + raise ValueError("token_types must not be empty") + if token_types.max() >= len(self.dim_inputs): + raise ValueError( + f"token_types contains an invalid expert index: {token_types.max()}" + ) + + if self.config.mot_opt: + bsz, q_len, _ = orig_shape + query_states, key_states, value_states = self._generate_qkv_mot_opt( + hidden_states, + token_types, + start_indices, + end_indices, + probs, + row_id_map, + bsz, + q_len, + ) + else: + bsz, q_len, _ = hidden_states.size() + masks = [ + (token_types == expert_idx) + for expert_idx in range(len(self.dim_inputs)) + ] + query_states, key_states, value_states = self._generate_qkv( + hidden_states, masks + ) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = self._apply_rotary_pos_embed( + query_states, key_states, cos, sin, unsqueeze_dim=2 + ) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + # repeat k/v heads if n_kv_heads < n_heads + key_states = self.repeat_kv(key_states, self.num_key_value_groups) + value_states = self.repeat_kv(value_states, self.num_key_value_groups) + # dropout_rate = 0.0 if not self.training else self.attention_dropout + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self._projection_dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Expand the attention_mask head dimension from 1 to num_heads + if attention_mask is not None and attention_mask.shape[1] == 1: + attention_mask = attention_mask.expand( + -1, self.num_heads, -1, -1 + ).contiguous() + + # has_moe1_token = token_types.any(dim=1) # Check whether each row has nonzero values + # moe0_seq_len = (token_types != 0).int().argmax(dim=1) + # moe0_seq_len = torch.where(has_moe1_token, moe0_seq_len, q_len) # Set to q_len when there are no tokens + flow_mask = token_types == 1 + start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) + + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + attn_output = flashmask_attn_func_stop_gradient( + query_states, + key_states, + value_states, + start_flow_pos, + startend_row_indices=attention_mask, + causal=False, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + + if self.config.mot_opt: + output = self._generate_output_mot_opt( + attn_output, token_types, start_indices, end_indices + ) + else: + output = self._generate_output(attn_output, masks) + + return output, None, past_key_value + + JOINT_QWEN_ATTENTION_CLASSES = { "eager": JointQwen2VLAttention, "flash_attention_2": JointQwen2VLFlashAttention, + # "flash_attention_2_ki": JointQwen2VLFlashAttention_KI, + # "flash_attention_2_triton": JointQwen2VLFlashAttention_Triton, "sdpa": JointQwen2VLAttention, + "flash_mask": JointQwen2VLFlashMaskAttention, + "flash_mask_ki": JointQwen2VLFlashMaskAttention_KI, } diff --git a/wall_x/model/core/attention/mask.py b/wall_x/model/core/attention/mask.py new file mode 100644 index 0000000..ec44608 --- /dev/null +++ b/wall_x/model/core/attention/mask.py @@ -0,0 +1,259 @@ +import torch + + +def find_first_last_ones(tensor): + """ + Input: tensor of shape (bs, seq_len) containing 0s and 1s + Output: (first_indices, last_indices), each of shape (bs,) + first_indices[i] is the first index of 1 in batch i, or -1 if none exists + last_indices[i] is the last index of 1 in batch i, or -1 if none exists + """ + bs, seq_len = tensor.shape + masks = tensor == 1 + has_ones = masks.any(dim=1) + + first = torch.full((bs,), -1, dtype=torch.long, device=tensor.device) + last = first.clone() + + # Compute the first index of 1 + first[has_ones] = torch.argmax(masks[has_ones].float(), dim=1) + + # Compute the last index of 1 + flipped_masks = masks.flip(dims=[1]) + last_argmax = torch.argmax(flipped_masks[has_ones].float(), dim=1) + last[has_ones] = seq_len - 1 - last_argmax + + return first, last + + +def update_position_ids(position_ids, moe_token_types, positional_masks): + """Extracted from ActionModelMixMin._update_position_ids (was @staticmethod).""" + if positional_masks is None or "ar_predict_token_positions" not in positional_masks: + return position_ids + + new_position_ids = position_ids.clone() + ar_predict_token_positions = positional_masks["ar_predict_token_positions"] + flow_mask = moe_token_types == 1 + + start_ar_pos, end_ar_pos = find_first_last_ones(ar_predict_token_positions) + start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) + + for bs_i in range(position_ids.shape[1]): + if start_ar_pos[bs_i] != -1 and end_ar_pos[bs_i] != -1: + start_ar_ids = new_position_ids[:, bs_i, start_ar_pos[bs_i]] + start_flow_ids = new_position_ids[:, bs_i, start_flow_pos[bs_i]] + diff = start_flow_ids - start_ar_ids + new_position_ids[:, bs_i, start_flow_pos[bs_i] :] = position_ids[ + :, bs_i, start_flow_pos[bs_i] : + ] - diff.unsqueeze(-1) + + return new_position_ids + + +def update_joint_attention_mask_2d( + attention_mask, + moe_token_types, + positional_masks, + causal_action_attention_mask=False, +): + """Extracted from ActionModelMixMin._update_joint_attention_mask_2d. + The only self attribute used was self.config.causal_action_attention_mask, now passed as parameter. + """ + if attention_mask.dim() == 3: # bs, seq_len, seq_len + return attention_mask + + bs, seq_len = moe_token_types.shape[0], moe_token_types.shape[1] + # Create a lower-triangular causal mask + causal_mask = torch.tril( + torch.ones( + (seq_len, seq_len), dtype=torch.bfloat16, device=moe_token_types.device + ) + ) + # Expand to the batch dimension + attention_mask = causal_mask.unsqueeze(0).expand(bs, -1, -1) + + if positional_masks is not None and "padding_positions" in positional_masks: + padding_positions = positional_masks["padding_positions"] + # Set padding rows to zero + attention_mask = torch.where( + padding_positions[:, None, :], + torch.zeros_like(attention_mask), + attention_mask, + ) + # Set padding columns to zero + attention_mask = torch.where( + padding_positions[:, :, None], + torch.zeros_like(attention_mask), + attention_mask, + ) + + # Set the moe1 region to 1 and mask it from the fast region + moe1_mask = (moe_token_types[:, :, None]) & (moe_token_types[:, None, :]) + + if ( + not causal_action_attention_mask + ): # If causal action attention mask is disabled, set the whole moe1 region to 1 + attention_mask = torch.where( + moe1_mask, torch.ones_like(attention_mask), attention_mask + ) + + if ( + positional_masks is not None + and "ar_predict_token_positions" in positional_masks + ): + ar_predict_token_positions = positional_masks["ar_predict_token_positions"] + moe1_mask = (moe_token_types[:, :, None]) & ( + ar_predict_token_positions[:, None, :] + ) + attention_mask = torch.where( + moe1_mask, torch.zeros_like(attention_mask), attention_mask + ) + + if ( + positional_masks is not None + and "valid_flow_action_positions" in positional_masks + ): + # true in moe_token_types but false in valid_flow_action_positions + nonvalid_flow_action_positions = ( + moe_token_types & ~positional_masks["valid_flow_action_positions"] + ) + attention_mask = torch.where( + nonvalid_flow_action_positions[:, None, :], + torch.zeros_like(attention_mask), + attention_mask, + ) + attention_mask = torch.where( + nonvalid_flow_action_positions[:, :, None], + torch.zeros_like(attention_mask), + attention_mask, + ) + + # AR and flow are bidirectional + if positional_masks is not None and "ar_action_mask" in positional_masks: + ar_action_mask = positional_masks["ar_action_mask"] != 0 + flow_positions = moe_token_types == 1 + if positional_masks.get("ar_visible", True): + flow_ar_position = ar_action_mask | flow_positions + flow_ar_mask = flow_ar_position[:, :, None] & flow_ar_position[:, None, :] + attention_mask = torch.where( + flow_ar_mask, torch.ones_like(attention_mask), attention_mask + ) + else: + flow_flow_mask = flow_positions[:, :, None] & flow_positions[:, None, :] + ar_ar_mask = ar_action_mask[:, :, None] & ar_action_mask[:, None, :] + flow_ar_mask = flow_flow_mask | ar_ar_mask + + affected = ar_action_mask | flow_positions # (B, N) + affected_pair = affected[:, :, None] & affected[:, None, :] + attention_mask = attention_mask.masked_fill(affected_pair, 0) + + attention_mask = torch.where( + flow_ar_mask, torch.ones_like(attention_mask), attention_mask + ) + + return attention_mask + + +def update_joint_attention_flash_mask( + attention_mask, + moe_token_types, + positional_masks, + causal_action_attention_mask=False, + debug=False, +): + """Extracted from ActionModelMixMin._update_joint_attention_flash_mask. + The only self attribute used was self.config.causal_action_attention_mask, now passed as parameter. + """ + device = moe_token_types.device + B, S = moe_token_types.shape + i32 = torch.int32 + + # ---- Return-vector initialization ---- + LTS = torch.ones((B, S), device=device, dtype=i32) * S + UTE = torch.arange(S, device=device, dtype=i32).unsqueeze(0).expand(B, S).clone() + + # Handle padding positions + if positional_masks is not None and "padding_positions" in positional_masks: + padding_positions = positional_masks["padding_positions"] + LTS[padding_positions] = 0 + UTE[padding_positions] = S + + # Handle AR predict tokens + if ( + positional_masks is not None + and "ar_predict_token_positions" in positional_masks + ): + start_ar_pos, end_ar_pos = find_first_last_ones( + positional_masks["ar_predict_token_positions"] + ) + for bs_i in range(B): + if end_ar_pos[bs_i] != -1: + LTS[bs_i, positional_masks["ar_predict_token_positions"][bs_i]] = ( + end_ar_pos[bs_i].to(i32) + 1 + ) + + flow_mask = moe_token_types == 1 + start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) + if positional_masks is None or "ar_action_mask" not in positional_masks: + # Handle bidirectional flow action masks + if not causal_action_attention_mask: + for bs_i in range(B): + if start_flow_pos[bs_i] != -1: + UTE[bs_i, flow_mask[bs_i]] = start_flow_pos[bs_i].to(i32) + else: + # AR and flow are bidirectional + ar_action_mask = positional_masks["ar_action_mask"] != 0 + flow_mask = moe_token_types == 1 + if positional_masks.get("ar_visible", True): + flow_ar_position = ar_action_mask | flow_mask + for bs_i in range(B): + idx = flow_ar_position[bs_i].nonzero(as_tuple=True)[0] + if idx.numel() == 0: + continue + block_start = idx.min() + block_end = idx.max() + 1 + # Set the visible range of every token in the block to [block_start, block_end) + UTE[bs_i, idx] = block_start.to(i32) + LTS[bs_i, idx] = block_end.to(i32) + else: + for bs_i in range(B): + # 1) Flow sub-block: flow attends only to flow + flow_idx = flow_mask[bs_i].nonzero(as_tuple=True)[0] + if flow_idx.numel() > 0: + flow_start = flow_idx.min() + flow_end = flow_idx.max() + 1 + # Only flow-token rows are set to [flow_start, flow_end) + UTE[bs_i, flow_idx] = flow_start.to(i32) + LTS[bs_i, flow_idx] = flow_end.to(i32) + + # 2) AR sub-block: AR attends only to AR + ar_idx = ar_action_mask[bs_i].nonzero(as_tuple=True)[0] + if ar_idx.numel() > 0: + ar_start = ar_idx.min() + ar_end = ar_idx.max() + 1 + # Only AR-token rows are set to [ar_start, ar_end) + UTE[bs_i, ar_idx] = ar_start.to(i32) + LTS[bs_i, ar_idx] = ar_end.to(i32) + + # Handle validation flow + if ( + positional_masks is not None + and "valid_flow_action_positions" in positional_masks + ): + flow_mask = moe_token_types == 1 + nonvalid_flow_action_positions = ( + flow_mask & ~positional_masks["valid_flow_action_positions"] + ) + if nonvalid_flow_action_positions.any(): + LTS[nonvalid_flow_action_positions] = 0 + UTE[nonvalid_flow_action_positions] = S + + LTS = LTS.unsqueeze(-1) + UTE = UTE.unsqueeze(-1) + + startend_row_indices = torch.cat([LTS, UTE], dim=-1) + + # add num_heads dimension + startend_row_indices = startend_row_indices.unsqueeze(1) + + return startend_row_indices diff --git a/wall_x/model/core/attention/selector.py b/wall_x/model/core/attention/selector.py new file mode 100644 index 0000000..a288935 --- /dev/null +++ b/wall_x/model/core/attention/selector.py @@ -0,0 +1,136 @@ +from typing import Dict, Optional, Union + +import torch +from packaging import version +from transformers.modeling_utils import AttentionInterface +from transformers.utils import is_torch_xla_available, logging + +ALL_ATTENTION_FUNCTIONS: AttentionInterface = AttentionInterface() +logger = logging.get_logger(__name__) + +CUSTOM_ATTENTION_FUNCTIONS = [ + "flash_attention_2_ki", + "flash_attention_2_triton", + "flash_mask", + "flash_mask_ki", +] +ATTENTION_TYPES_WITH_2D_MASK = [ + "flash_attention_2_ki", + "flash_attention_2_triton", + "sdpa", +] +ATTENTION_TYPES_WITH_FLASH_MASK = ["flash_mask", "flash_mask_ki"] + + +class AttentionsSelectorMixin: + + @classmethod + def _autoset_attn_implementation( + cls, + config, + use_flash_attention_2: bool = False, + torch_dtype: Optional[torch.dtype] = None, + device_map: Optional[Union[str, Dict[str, int]]] = None, + check_device_map: bool = True, + ): + """ + Automatically checks and dispatches to a default attention implementation. In order of priority: + 1. An implementation specified in `config._attn_implementation` (due for example to the argument attn_implementation="sdpa" in from_pretrained). + 2. DEPRECATED: if use_flash_attention_2 is set to `True` and `flash_attn` is available, flash attention. (`LlamaFlashAttention` for example) + 3. SDPA implementation, if available and supported by the model type. (`LlamaSdpaAttention` for example) + 4. The default model's implementation otherwise (`LlamaAttention` for example) . + """ + # Here we use config._attn_implementation_internal to check whether the attention implementation was explicitly set by the user. + # The property `PretrainedConfig._attn_implementation` is never `None`, for backward compatibility (always fall back on "eager"). + # The `hasattr` here is used as some Transformers tests for some reason do not call PretrainedConfig __init__ (e.g. test_no_super_init_config_and_model) + requested_attn_implementation = None + if ( + hasattr(config, "_attn_implementation_internal") + and config._attn_implementation_internal is not None + ): + if ( + config._attn_implementation != "flash_attention_2" + and use_flash_attention_2 + ): + raise ValueError( + f'Both attn_implementation="{config._attn_implementation}" and `use_flash_attention_2=True` were used when loading the model, which are not compatible.' + ' We recommend to just use `attn_implementation="flash_attention_2"` when loading the model.' + ) + + if ( + not isinstance(config._attn_implementation, dict) + and config._attn_implementation + not in ["eager"] + + ALL_ATTENTION_FUNCTIONS.valid_keys() + + CUSTOM_ATTENTION_FUNCTIONS + ): + message = f'Specified `attn_implementation="{config._attn_implementation}"` is not supported. The only possible arguments are `attn_implementation="eager"` (manual attention implementation)' + if cls._supports_flash_attn_2: + message += ', `"attn_implementation=flash_attention_2"` (implementation using flash attention 2)' + if cls._supports_sdpa: + message += ', `"attn_implementation=sdpa"` (implementation using torch.nn.functional.scaled_dot_product_attention)' + if cls._supports_flex_attn: + message += ', `"attn_implementation=flex_attention"` (implementation using torch\'s flex_attention)' + raise ValueError(message + ".") + + # If a config is passed with a preset attn_implementation, we skip the automatic dispatch and use the user-provided config, with hard checks that the requested attention implementation is available. + requested_attn_implementation = config._attn_implementation_internal + + if use_flash_attention_2: + logger.warning_once( + 'The model was loaded with use_flash_attention_2=True, which is deprecated and may be removed in a future release. Please use `attn_implementation="flash_attention_2"` instead.' + ) + config._attn_implementation = "flash_attention_2" + + if config._attn_implementation == "flash_attention_2": + cls._check_and_enable_flash_attn_2( + config, + torch_dtype=torch_dtype, + device_map=device_map, + hard_check_only=False, + check_device_map=check_device_map, + ) + elif requested_attn_implementation == "flex_attention": + config = cls._check_and_enable_flex_attn(config, hard_check_only=True) + elif ( + requested_attn_implementation in [None, "sdpa"] + and not is_torch_xla_available() + ): + # use_flash_attention_2 takes priority over SDPA, hence SDPA treated in this elif. + config = cls._check_and_enable_sdpa( + config, + hard_check_only=( + False if requested_attn_implementation is None else True + ), + ) + + if ( + torch.version.hip is not None + and config._attn_implementation == "sdpa" + and torch.cuda.device_count() > 1 + and version.parse(torch.__version__) < version.parse("2.4.1") + ): + logger.warning_once( + "Using the `SDPA` attention implementation on multi-gpu setup with ROCM may lead to performance issues due to the FA backend. Disabling it to use alternative backends." + ) + torch.backends.cuda.enable_flash_sdp(False) + elif requested_attn_implementation in ALL_ATTENTION_FUNCTIONS.valid_keys(): + config._attn_implementation = requested_attn_implementation + elif isinstance(requested_attn_implementation, dict): + config._attn_implementation = None + elif config._attn_implementation in CUSTOM_ATTENTION_FUNCTIONS: + pass + else: + config._attn_implementation = "eager" + + config._attn_implementation_autoset = True + return config + + def _check_and_adjust_attn_implementation( + self, attn_implementation: Optional[str], is_init_check: bool = False + ) -> str: + assert ( + attn_implementation + in ["eager", "flash_attention_2", "sdpa"] + CUSTOM_ATTENTION_FUNCTIONS + ) + return attn_implementation diff --git a/wall_x/model/core/ops/__init__.py b/wall_x/model/core/ops/__init__.py new file mode 100644 index 0000000..1e4a6ad --- /dev/null +++ b/wall_x/model/core/ops/__init__.py @@ -0,0 +1,17 @@ +"""Operator proxy layer with runtime fallback backends.""" + +from wall_x.model.core.ops.index import get_rope_index, get_window_index +from wall_x.model.core.ops.moe import permute, unpermute +from wall_x.model.core.ops.norm import rmsnorm +from wall_x.model.core.ops.rope import m_rope, rope, rot_pos_emb + +__all__ = [ + "rmsnorm", + "rope", + "m_rope", + "rot_pos_emb", + "permute", + "unpermute", + "get_rope_index", + "get_window_index", +] diff --git a/wall_x/model/core/ops/_cuda_ext.py b/wall_x/model/core/ops/_cuda_ext.py new file mode 100644 index 0000000..8f4f7fc --- /dev/null +++ b/wall_x/model/core/ops/_cuda_ext.py @@ -0,0 +1,35 @@ +"""Install-time compiled CUDA kernels for wall_x ops. + +Auto-generated by scripts/extract_cuda_kernels.py. +Do not edit manually. +""" + +_module = None + + +def load(): + """Load the CUDA extension module built by setup.py.""" + global _module + if _module is not None: + return _module + + try: + from wall_x.model.core.ops import _cuda_ext_bin + except ImportError as exc: + raise ImportError( + "Wall-X CUDA operators were not built. Install requirements first, " + "then reinstall Wall-X with: " + "MAX_JOBS=8 pip install --no-build-isolation -e ." + ) from exc + + _module = _cuda_ext_bin + return _module + + +def is_available() -> bool: + """Check whether the install-time CUDA extension can be imported.""" + try: + load() + return True + except Exception: + return False diff --git a/wall_x/model/core/ops/_cuda_wrappers.py b/wall_x/model/core/ops/_cuda_wrappers.py new file mode 100644 index 0000000..0fef8a2 --- /dev/null +++ b/wall_x/model/core/ops/_cuda_wrappers.py @@ -0,0 +1,750 @@ +"""Auto-generated CUDA kernel wrappers with autograd support. + +Generated by scripts/extract_cuda_kernels.py from internal CUDA wrappers. +Do not edit manually. +""" + +import torch +from torch.autograd import Function + + +def _m(): + """Get the compiled CUDA extension module.""" + from wall_x.model.core.ops._cuda_ext import load + return load() + + + +class _RopeFunction(Function): + @staticmethod + def forward( + ctx, + q, + k, + cos, + sin, + interleave, + fwd_kernel=None, + bwd_kernel=None, + ): + q_embed = q + k_embed = k + ctx.save_for_backward(q, k, cos, sin) + ctx.interleave = interleave + ctx.bwd_kernel = bwd_kernel + fwd_kernel(q, k, q_embed, k_embed, cos, sin, interleave) + return q_embed, k_embed + + @staticmethod + def backward(ctx, grad_q_embed, grad_k_embed): + q, k, cos, sin = ctx.saved_tensors + interleave = ctx.interleave + bwd_kernel = ctx.bwd_kernel + grad_q = grad_q_embed + grad_k = grad_k_embed + + bwd_kernel( + grad_q_embed, + grad_k_embed, + grad_q, # output + grad_k, # output + cos, + sin, + interleave, + ) + + return grad_q, grad_k, None, None, None, None, None + + +class _RopePackFunction(Function): + @staticmethod + def forward( + ctx, + qkv, + q_num_heads, + kv_num_heads, + cos, + sin, + interleave, + fwd_kernel=None, + bwd_kernel=None, + ): + ctx.save_for_backward(cos, sin) + ctx.q_num_heads = q_num_heads + ctx.kv_num_heads = kv_num_heads + ctx.interleave = interleave + ctx.bwd_kernel = bwd_kernel + + fwd_kernel(qkv, cos, sin, q_num_heads, kv_num_heads, interleave) + return qkv + + @staticmethod + def backward(ctx, dqkv): + if dqkv is None: + return None, None, None, None, None, None, None, None + + cos, sin = ctx.saved_tensors + if dqkv.stride(-1) != 1: + dqkv = dqkv.contiguous() + + ctx.bwd_kernel( + dqkv, + cos, + sin, + ctx.q_num_heads, + ctx.kv_num_heads, + ctx.interleave, + ) + return dqkv, None, None, None, None, None, None, None + + +class Rope: + def __init__(self): + pass + self.fwd_kernel = _m().rope + self.inplace_kernel = _m().rope_inplace + self.pack_kernel = _m().rope_inplace_pack + self.pack_bwd_kernel = _m().rope_inplace_pack_bwd + self.bwd_kernel = _m().rope_bwd + + def pack( + self, + qkv, + q_num_heads, + kv_num_heads, + cos, + sin, + interleave=False, + inference=False, + ): + """packed qkvq/krope,qkv. + + Split/view/offset calculations are done on the C++ side. Callers consume qkv through the paired + flash_attn pack interface without splitting q/k/v in Python. + autograd,backwarddqkvdq/dk. + + Args: + qkv: [seq_len, q_dim + 2*kv_dim] + q_num_heads: number of query heads + kv_num_heads: number of key/value heads + cos, sin: rotary position encoding [1, seq_len, head_dim/2] + interleave: whether to use interleaved RoPE + inference: kernel,autograd + """ + if inference: + self.pack_kernel(qkv, cos, sin, q_num_heads, kv_num_heads, interleave) + return qkv + + return _RopePackFunction.apply( + qkv, + q_num_heads, + kv_num_heads, + cos, + sin, + interleave, + self.pack_kernel, + self.pack_bwd_kernel, + ) + + def pack_backward( + self, dqkv, q_num_heads, kv_num_heads, cos, sin, interleave=False + ): + """packed dqkvdq/dkrope backward,dqkv.""" + self.pack_bwd_kernel(dqkv, cos, sin, q_num_heads, kv_num_heads, interleave) + return dqkv + + def __call__(self, q, k, cos, sin, interleave=False, inference=False): + if inference: + self.inplace_kernel(q, k, cos, sin, interleave) + return q, k + + return _RopeFunction.apply( + q, + k, + cos, + sin, + interleave, + self.fwd_kernel, + self.bwd_kernel, + ) + + + + +class _MRopeFunction(Function): + @staticmethod + def forward( + ctx, + q, + k, + cos, + sin, + mrope_section, + fwd_kernel=None, + bwd_kernel=None, + ): + first = mrope_section[0] + second = mrope_section[1] + ctx.save_for_backward(q, k, cos, sin) + ctx.mrope_section = mrope_section + ctx.bwd_kernel = bwd_kernel + q_embed = q + k_embed = k + fwd_kernel(q, k, q_embed, k_embed, cos, sin, first, second) + return q_embed, k_embed + + @staticmethod + def backward(ctx, grad_q_embed, grad_k_embed): + if grad_q_embed.stride(-1) != 1: + grad_q_embed = grad_q_embed.contiguous() + if grad_k_embed.stride(-1) != 1: + grad_k_embed = grad_k_embed.contiguous() + q, k, cos, sin = ctx.saved_tensors + mrope_section = ctx.mrope_section + bwd_kernel = ctx.bwd_kernel + grad_q = grad_q_embed + grad_k = grad_k_embed + first = mrope_section[0] + second = mrope_section[1] + + bwd_kernel( + grad_q_embed, + grad_k_embed, + grad_q, # output + grad_k, # output + cos, + sin, + first, + second, + ) + + return grad_q, grad_k, None, None, None, None, None + + +class _MRopePackFunction(Function): + @staticmethod + def forward( + ctx, + qkv, + q_num_heads, + kv_num_heads, + cos, + sin, + mrope_section, + fwd_kernel=None, + bwd_kernel=None, + ): + first = mrope_section[0] + second = mrope_section[1] + ctx.save_for_backward(cos, sin) + ctx.q_num_heads = q_num_heads + ctx.kv_num_heads = kv_num_heads + ctx.mrope_section = mrope_section + ctx.bwd_kernel = bwd_kernel + + fwd_kernel(qkv, cos, sin, q_num_heads, kv_num_heads, first, second) + return qkv + + @staticmethod + def backward(ctx, dqkv): + if dqkv is None: + return None, None, None, None, None, None, None, None + + cos, sin = ctx.saved_tensors + if dqkv.stride(-1) != 1: + dqkv = dqkv.contiguous() + first = ctx.mrope_section[0] + second = ctx.mrope_section[1] + + ctx.bwd_kernel( + dqkv, + cos, + sin, + ctx.q_num_heads, + ctx.kv_num_heads, + first, + second, + ) + return dqkv, None, None, None, None, None, None, None + + +class MRope: + def __init__(self): + pass + self.fwd_kernel = _m().m_rope + self.inplace_kernel = _m().m_rope_inplace + self.pack_kernel = _m().m_rope_inplace_pack + self.pack_bwd_kernel = _m().m_rope_inplace_pack_bwd + self.bwd_kernel = _m().m_rope_bwd + + def pack( + self, + qkv, + q_num_heads, + kv_num_heads, + cos, + sin, + mrope_section, + inference=False, + ): + """packed qkvq/kmrope,qkv.GQA. + + Split/view/offset calculations are done on the C++ side. qkv must be 3D: [bz, seq_len, q_dim+2*kv_dim]. + + Args: + qkv: [bz, seq_len, q_dim + 2*kv_dim] + q_num_heads: number of query heads + kv_num_heads: number of key/value heads + cos, sin: rotary position encoding [3, bz, seq_len, head_dim/2] + mrope_section: (first, second) M-RoPE section parameters + inference: kernel,autograd + """ + if inference: + self.pack_kernel( + qkv, + cos, + sin, + q_num_heads, + kv_num_heads, + mrope_section[0], + mrope_section[1], + ) + return qkv + + return _MRopePackFunction.apply( + qkv, + q_num_heads, + kv_num_heads, + cos, + sin, + mrope_section, + self.pack_kernel, + self.pack_bwd_kernel, + ) + + def pack_backward(self, dqkv, q_num_heads, kv_num_heads, cos, sin, mrope_section): + """packed dqkvdq/dkmrope backward,dqkv.""" + self.pack_bwd_kernel( + dqkv, + cos, + sin, + q_num_heads, + kv_num_heads, + mrope_section[0], + mrope_section[1], + ) + return dqkv + + def __call__(self, q, k, cos, sin, mrope_section, inference=False): + if inference: + first = mrope_section[0] + second = mrope_section[1] + self.inplace_kernel(q, k, cos, sin, first, second) + return q, k + + return _MRopeFunction.apply( + q, + k, + cos, + sin, + mrope_section, + self.fwd_kernel, + self.bwd_kernel, + ) + + + + +class RotPos: + def __init__(self): + pass + + def __call__(self, inv_freq, grid_thw, spatial_merge_size): + assert ( + inv_freq.dtype == torch.float32 + ), f"Expected float32, got {inv_freq.dtype}" + get_token_counts_kernel = _m().get_token_counts + rot_pos_kernel = _m().rot_pos + + num_grids = grid_thw.size(0) + token_counts = torch.zeros( + (num_grids), dtype=grid_thw.dtype, device=grid_thw.device + ) + + get_token_counts_kernel(grid_thw, token_counts, spatial_merge_size) + cumsum_tokens = torch.cat( + [ + torch.zeros(1, dtype=token_counts.dtype, device=token_counts.device), + token_counts.cumsum(dim=0), + ], + dim=0, + ).to(grid_thw.dtype) + + output = torch.empty( + (cumsum_tokens[-1].item(), inv_freq.size(0) * 2), + dtype=torch.float, + device=inv_freq.device, + ) + rot_pos_kernel(inv_freq, grid_thw, output, cumsum_tokens, spatial_merge_size) + + return output + + + + +class GetRopeIndex: + def __init__(self): + pass + + def __call__( + self, + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + spatial_merge_size, + image_token_id, + video_token_id, + vision_start_token_id, + tokens_per_second, + ): + get_workspace = _m().get_rope_index_getworkspace + get_rope_index_kernel = _m().get_rope_index + work_space_size = get_workspace(input_ids, image_grid_thw, video_grid_thw) + + workspace = torch.empty( + work_space_size, dtype=torch.uint8, device=input_ids.device + ) + batch_size = input_ids.size(0) + seq_len = input_ids.size(1) + position_ids = torch.empty( + (3, batch_size, seq_len), dtype=torch.int64, device=input_ids.device + ) + mrope_deltas = torch.empty( + (batch_size, 1), dtype=torch.int64, device=input_ids.device + ) + get_rope_index_kernel( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + position_ids, + mrope_deltas, + workspace, + spatial_merge_size, + image_token_id, + video_token_id, + vision_start_token_id, + tokens_per_second, + ) + + return position_ids, mrope_deltas + + + +get_window_index_kernel = _m().get_window_index +get_totals_kernel = _m().get_totals + + +def get_window_index_cuda( + grid_thw, window_size, spatial_merge_size, patch_size, spatial_merge_unit=1 +): + + if grid_thw.size(0) == 0: + return ( + torch.empty(0, dtype=grid_thw.dtype, device=grid_thw.device), + torch.zeros(1, dtype=grid_thw.dtype, device=grid_thw.device), + ) + + vit_merger_window_size = window_size // spatial_merge_size // patch_size + + grid_info_tensor = torch.empty( + (grid_thw.size(0), 6), dtype=grid_thw.dtype, device=grid_thw.device + ) + global_totals_tensor = torch.zeros( + (2), dtype=grid_thw.dtype, device=grid_thw.device + ) + get_totals_kernel( + grid_thw, + grid_info_tensor, + global_totals_tensor, + spatial_merge_size, + vit_merger_window_size, + ) + total_elements = global_totals_tensor[0].item() + total_windows = global_totals_tensor[1].item() + if total_elements == 0 or total_windows == 0: + return ( + torch.empty(0, dtype=grid_thw.dtype, device=grid_thw.device), + torch.zeros(1, dtype=grid_thw.dtype, device=grid_thw.device), + ) + + window_indices = torch.empty( + total_elements, dtype=grid_thw.dtype, device=grid_thw.device + ) + cu_window_seqlens = torch.empty( + (total_windows + 1), dtype=grid_thw.dtype, device=grid_thw.device + ) + window_counts_tensor = torch.empty( + (total_windows), dtype=grid_thw.dtype, device=grid_thw.device + ) + + max_grid_t = grid_thw[:, 0].max().item() + get_window_index_kernel( + grid_thw, + grid_info_tensor, + window_indices, + cu_window_seqlens, + window_counts_tensor, + max_grid_t, + spatial_merge_size, + vit_merger_window_size, + patch_size, + spatial_merge_unit, + ) + + return window_indices, cu_window_seqlens + + + +################################################################################################ +## +## PermuteMoE topK +## +################################################################################################ + + + +class PermuteMoE_topK(torch.autograd.Function): + + dtype = None + max_expanded_token_num = 0 + + @staticmethod + def forward( + ctx, + input_act: torch.Tensor, + indices: torch.Tensor, + num_out_tokens: int, + max_token_num: int, + ): + """ + indices: for topK=1, indices in a 1-d tensor of shape [num_tokens], + otherwise, it's a 2-d tensor of shape [num_tokens, topK] + """ + if not input_act.numel(): + return input_act, None + + # For top1 case, view the indices as 2D tensor to unify the shape for topk>=2 cases. + if indices.dim() == 1: + indices = indices.view(-1, 1) + + # Device check + if input_act.is_cpu: + raise RuntimeError( + "[Error] The input `input_act` of permute_topK op is on the device: CPU!" + ) + + # Data type check + if indices.dtype != torch.int32: + indices = indices.to(torch.int32) + + # Contiguous check + if not input_act.is_contiguous(): + input_act = input_act.contiguous() + if not indices.is_contiguous(): + indices = indices.contiguous() + + num_topK = indices.size(1) + + input_max_expanded_token_num = max(max_token_num, input_act.size(0)) * num_topK + if PermuteMoE_topK.max_expanded_token_num < input_max_expanded_token_num: + PermuteMoE_topK.max_expanded_token_num = input_max_expanded_token_num + + permute_kernel = _m().permute + sorted_indices = torch.empty( + PermuteMoE_topK.max_expanded_token_num, + dtype=torch.int32, + device=input_act.device, + ) + row_id = torch.arange( + PermuteMoE_topK.max_expanded_token_num, + dtype=torch.int32, + device=input_act.device, + ) + sorted_row_id = torch.empty( + PermuteMoE_topK.max_expanded_token_num, + dtype=torch.int32, + device=input_act.device, + ) + get_storage_bytes_kernel = _m().cub_sort_pair_get_storage_bytes + temp_storage_bytes = get_storage_bytes_kernel( + PermuteMoE_topK.max_expanded_token_num + ) + temp_storage = torch.empty( + temp_storage_bytes, dtype=torch.int8, device=input_act.device + ) + num_out = ( + num_out_tokens if (num_out_tokens > 0) else (indices.size(0) * num_topK) + ) + permuted_output = torch.empty( + (num_out, input_act.size(1)), dtype=input_act.dtype, device=input_act.device + ) + row_id_map = torch.empty( + (indices.size(0) * num_topK), dtype=torch.int32, device=input_act.device + ) + permute_kernel( + input_act, + indices, + sorted_indices, + row_id, + sorted_row_id, + temp_storage, + permuted_output, + row_id_map, + num_out_tokens, + PermuteMoE_topK.max_expanded_token_num, + ) + + ctx.row_id_map = row_id_map + ctx.num_tokens = indices.size(0) + ctx.num_topK = num_topK + + return permuted_output, row_id_map + + @staticmethod + def backward(ctx, permuted_act_grad, _): + if not permuted_act_grad.numel(): + return permuted_act_grad, None, None, None + + unpermute_kernel = _m().unpermute + if not permuted_act_grad.is_contiguous(): + permuted_act_grad = permuted_act_grad.contiguous() + + row_id_map = ctx.row_id_map + num_tokens = ctx.num_tokens + num_topK = ctx.num_topK + num_cols = permuted_act_grad.size(1) + unpermuted_output = torch.empty( + (num_tokens, num_cols), + dtype=permuted_act_grad.dtype, + device=permuted_act_grad.device, + ) + unpermute_kernel( + permuted_act_grad, row_id_map, None, unpermuted_output, num_tokens, num_topK + ) + + return unpermuted_output, None, None, None + + +################################################################################################ +## +## UnpermuteMoE topK +## +################################################################################################ + + +class UnpermuteMoE_topK(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + input_act: torch.Tensor, + row_id_map: torch.Tensor, + probs: torch.Tensor = None, + ): + + if not input_act.numel(): + ctx.probs = probs + return input_act + + # Device check + if input_act.is_cpu: + raise RuntimeError( + "[Error] The input `input_act` of unpermute_topK op is on the device: CPU!" + ) + if row_id_map.is_cpu: + row_id_map = row_id_map.cuda() + if probs is not None and probs.is_cpu: + probs = probs.cuda() + + # Shape check + if probs is not None and row_id_map.size(0) != probs.size(0) * probs.size(1): + raise RuntimeError( + f"[Error] unpermute_topK op input `probs` shape mismatch! " + f"Expect {row_id_map.size(0)}, but got {probs.size(0) * probs.size(1)}." + ) + + # Data type check + if row_id_map.dtype != torch.int32: + row_id_map = row_id_map.to(torch.int32) + if probs is not None and probs.dtype != torch.float32: + probs = probs.to(torch.float32) + + # Contiguous check + if not input_act.is_contiguous(): + input_act = input_act.contiguous() + if not row_id_map.is_contiguous(): + row_id_map = row_id_map.contiguous() + if probs is not None and not probs.is_contiguous(): + probs = probs.contiguous() + + num_tokens = probs.size(0) if probs is not None else input_act.size(0) + num_topK = probs.size(1) if probs is not None else 1 + unpermute_kernel = _m().unpermute + num_cols = input_act.size(1) + unpermuted_output = torch.empty( + (num_tokens, num_cols), dtype=input_act.dtype, device=input_act.device + ) + unpermute_kernel( + input_act, row_id_map, probs, unpermuted_output, num_tokens, num_topK + ) + + ctx.save_for_backward(input_act, row_id_map, probs) + + return unpermuted_output + + @staticmethod + def backward(ctx, unpermuted_act_grad): + + if not unpermuted_act_grad.numel(): + return unpermuted_act_grad, None, ctx.probs + + if not unpermuted_act_grad.is_contiguous(): + unpermuted_act_grad = unpermuted_act_grad.contiguous() + + input_act, row_id_map, probs = ctx.saved_tensors + + act_grad = None + unpermute_bwd_kernel = _m().unpermute_bwd + if ctx.needs_input_grad[0]: + num_cols = unpermuted_act_grad.size(1) + act_grad = torch.empty( + (input_act.size(0), num_cols), + dtype=unpermuted_act_grad.dtype, + device=unpermuted_act_grad.device, + ) + prob_grad = torch.empty( + (probs.size(0), probs.size(1)), + dtype=torch.float32, + device=unpermuted_act_grad.device, + ) + unpermute_bwd_kernel( + unpermuted_act_grad, input_act, row_id_map, probs, act_grad, prob_grad + ) + + if not ctx.needs_input_grad[2]: + prob_grad = None + + return act_grad, None, prob_grad + + +def permute_kernel(input_act, indices, num_out_tokens=None, max_token_num=0): + num_out_tokens = 0 if num_out_tokens is None else num_out_tokens + return PermuteMoE_topK.apply(input_act, indices, num_out_tokens, max_token_num) + + +def unpermute_kernel(input_act, row_id_map, probs=None): + return UnpermuteMoE_topK.apply(input_act, row_id_map, probs) diff --git a/wall_x/model/core/ops/base.py b/wall_x/model/core/ops/base.py new file mode 100644 index 0000000..1a70e17 --- /dev/null +++ b/wall_x/model/core/ops/base.py @@ -0,0 +1,143 @@ +"""Base class for operator proxies with lazy backend resolution.""" + +import importlib +import logging +import threading + +logger = logging.getLogger(__name__) + + +class OpsProxy: + """Callable proxy that lazily resolves to the best available backend. + + Three-level fallback: + Level 1: external_accel (fastest, requires external_accel package) + Level 2: CUDA inline kernel (JIT compiled, requires CUDA) + Level 3: Pure PyTorch (always available) + + Subclasses define _external_accel_name, _pytorch_fallback, and optionally _cuda_kernel. + + Set _external_accel_name to None to skip external_accel resolution (e.g., when the + external_accel API is incompatible). + """ + + def __init__(self): + self._resolved_fn = None + self._backend = None + self._resolve_lock = threading.Lock() + + def _resolve(self): + """Resolve the best available backend. Called once on first use (thread-safe).""" + if self._resolved_fn is not None: + return + with self._resolve_lock: + # Double-check after acquiring lock + if self._resolved_fn is not None: + return + + # Level 1: try external_accel (skip if _external_accel_name is None) + fn = self._probe_external_accel() + if fn is not None: + self._backend = "external_accel" + self._resolved_fn = fn + return + + # Level 2: try CUDA inline kernel (subclass override) + cuda_fn = self._get_cuda_kernel() + if cuda_fn is not None: + self._backend = "cuda_inline" + self._resolved_fn = cuda_fn + return + + # Level 3: PyTorch fallback + self._backend = "pytorch" + self._resolved_fn = self._pytorch_fallback + logger.info(f"{self.__class__.__name__}: using PyTorch fallback") + + def _get_cuda_kernel(self): + """Override in subclass to provide Level 2 CUDA kernel. Returns None if unavailable.""" + return None + + @property + def _external_accel_name(self): + """Return external_accel attribute name, or None to skip external_accel resolution. + + Base returns None (skip external_accel). Subclasses override to return the + attribute name when external_accel integration is desired. + """ + return None + + def _pytorch_fallback(self, *args, **kwargs): + raise NotImplementedError(f"{self.__class__.__name__} has no PyTorch fallback") + + def __call__(self, *args, **kwargs): + if self._resolved_fn is None: + self._resolve() + return self._resolved_fn(*args, **kwargs) + + @property + def backend(self) -> str: + if self._resolved_fn is None: + self._resolve() + return self._backend + + # ------------------------------------------------------------------ + # Multi-backend API (for testing / benchmarking) + # ------------------------------------------------------------------ + + def _probe_external_accel(self): + """Return the external_accel function for this op, or None if unavailable.""" + if self._external_accel_name is None: + return None + try: + mod = importlib.import_module("external_accel") + return getattr(mod, self._external_accel_name) + except (ImportError, AttributeError): + return None + + def available_backends(self): + """Return list of available backend names for this operator.""" + backends = ["pytorch"] + if self._probe_external_accel() is not None: + backends.append("external_accel") + if self._get_cuda_kernel() is not None: + backends.append("cuda_inline") + return backends + + def _get_backend_fn(self, backend): + """Return the callable for a specific backend, or None if unavailable.""" + if backend == "pytorch": + return self._pytorch_fallback + elif backend == "external_accel": + return self._probe_external_accel() + elif backend == "cuda_inline": + return self._get_cuda_kernel() + else: + raise ValueError(f"Unknown backend: {backend}") + + def call_with_backend(self, backend, *args, **kwargs): + """Call this operator using a specific backend. + + This does NOT change the default backend used by ``__call__``. + The ``backend`` property still reflects the auto-resolved default. + + Args: + backend: One of "external_accel", "cuda_inline", "pytorch". + + Returns: + Operator output from the specified backend. + + Raises: + RuntimeError: If the requested backend is not available. + """ + fn = self._get_backend_fn(backend) + if fn is None: + raise RuntimeError( + f"{self.__class__.__name__}: backend '{backend}' not available. " + f"Available: {self.available_backends()}" + ) + return fn(*args, **kwargs) + + def __repr__(self): + backend = self._backend or "unresolved" + return f"<{self.__class__.__name__} backend={backend}>" diff --git a/wall_x/model/core/ops/csrc/binding.cu b/wall_x/model/core/ops/csrc/binding.cu new file mode 100644 index 0000000..7b86061 --- /dev/null +++ b/wall_x/model/core/ops/csrc/binding.cu @@ -0,0 +1,69 @@ +// Auto-generated by scripts/extract_cuda_kernels.py +// This file registers CUDA kernel bindings via pybind11. +// +// Kernel .cu files are compiled separately so ninja can parallelize builds. + +#include +#include +#include + +namespace wallx_cuda_get_rope_index { +void GetRopeIndex(const c10::optional& input_ids, const c10::optional& image_grid_thw, const c10::optional& video_grid_thw, const c10::optional& second_per_grid_ts, const c10::optional& attention_mask, const at::Tensor& position_ids, const at::Tensor& mrope_deltas, const at::Tensor& workspace, int spatial_merge_size, int image_token_id, int video_token_id, int vision_start_token_id, float tokens_per_second); +int64_t GetRopeIndexGetWorkSpace(const c10::optional& input_ids, const c10::optional& image_grid_thw, const c10::optional& video_grid_thw); +} // namespace wallx_cuda_get_rope_index + +namespace wallx_cuda_m_rope { +void MRope(const at::Tensor& q, const at::Tensor& k, const at::Tensor& q_embed, const at::Tensor& k_embed, const at::Tensor& cos, const at::Tensor& sin, const int first, const int second ); +void MRopeInplace(const at::Tensor& q, const at::Tensor& k, const at::Tensor& cos, const at::Tensor& sin, const int first, const int second ); +void MRopeInplacePack(const at::Tensor& qkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, const int first, const int second); +void MRopeInplacePackBackward(const at::Tensor& dqkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, const int first, const int second); +void MRopeBackward( const at::Tensor& grad_q_embed, const at::Tensor& grad_k_embed, const at::Tensor& grad_q, const at::Tensor& grad_k, const at::Tensor& cos, const at::Tensor& sin, const int first, const int second ); +} // namespace wallx_cuda_m_rope + +namespace wallx_cuda_permute { +void MoePermuteTopKOp( const at::Tensor& input, const at::Tensor& indices, const at::Tensor& sorted_indices, const at::Tensor& row_id, const at::Tensor& sorted_row_id, const at::Tensor& temp_storage, const at::Tensor& permuted_output, const at::Tensor& row_id_map, int64_t num_out_tokens, int64_t max_expanded_token_num ); +void MoeRecoverTopKOp( const at::Tensor& input, const at::Tensor& row_id_map, const c10::optional& prob, const at::Tensor& unpermuted_output, int64_t num_tokens, int64_t num_topK ); +void MoeRecoverTopKBwdOp( const at::Tensor& input_bwd, const at::Tensor& input_fwd, const at::Tensor& row_id_map, const at::Tensor& prob, const at::Tensor& act_grad, const at::Tensor& prob_grad ); +size_t CubSortPairGetStorageBytes(int64_t num_items); +} // namespace wallx_cuda_permute + +namespace wallx_cuda_rope { +void Rope(const at::Tensor& q, const at::Tensor& k, const at::Tensor& q_embed, const at::Tensor& k_embed, const at::Tensor& cos, const at::Tensor& sin, bool interleave ); +void RopeInplace(const at::Tensor& q, const at::Tensor& k, const at::Tensor& cos, const at::Tensor& sin, bool interleave ); +void RopeInplacePack(const at::Tensor& qkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, bool interleave); +void RopeInplacePackBackward(const at::Tensor& dqkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, bool interleave); +void RopeBackward( const at::Tensor& grad_q_embed, const at::Tensor& grad_k_embed, const at::Tensor& grad_q, const at::Tensor& grad_k, const at::Tensor& cos, const at::Tensor& sin, bool interleave ); +} // namespace wallx_cuda_rope + +namespace wallx_cuda_rot_pos { +void GetTokenCounts( const at::Tensor& grid_thw, const at::Tensor& token_counts, int spatial_merge_size ); +void RotPosEmb( const at::Tensor& inv_freq, const at::Tensor& grid_thw, const at::Tensor& output, const at::Tensor& cumsum_tokens, int spatial_merge_size ); +} // namespace wallx_cuda_rot_pos + +namespace wallx_cuda_window_index { +void GetWindowIndex( const at::Tensor& grid_thw, const at::Tensor& grid_info_tensor, const at::Tensor& window_indices, const at::Tensor& cu_window_seqlens, const at::Tensor& window_counts_tensor, int max_grid_t, int spatial_merge_size, int vit_merger_window_size, int patch_size, int spatial_merge_unit); +void GetTotals( const at::Tensor& grid_thw, const at::Tensor& grid_info_tensor, const at::Tensor& global_totals_tensor, int spatial_merge_size, int vit_merger_window_size ); +} // namespace wallx_cuda_window_index + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("rope", &wallx_cuda_rope::Rope); + m.def("rope_inplace", &wallx_cuda_rope::RopeInplace); + m.def("rope_inplace_pack", &wallx_cuda_rope::RopeInplacePack); + m.def("rope_inplace_pack_bwd", &wallx_cuda_rope::RopeInplacePackBackward); + m.def("rope_bwd", &wallx_cuda_rope::RopeBackward); + m.def("m_rope", &wallx_cuda_m_rope::MRope); + m.def("m_rope_inplace", &wallx_cuda_m_rope::MRopeInplace); + m.def("m_rope_inplace_pack", &wallx_cuda_m_rope::MRopeInplacePack); + m.def("m_rope_inplace_pack_bwd", &wallx_cuda_m_rope::MRopeInplacePackBackward); + m.def("m_rope_bwd", &wallx_cuda_m_rope::MRopeBackward); + m.def("get_token_counts", &wallx_cuda_rot_pos::GetTokenCounts); + m.def("rot_pos", &wallx_cuda_rot_pos::RotPosEmb); + m.def("get_rope_index", &wallx_cuda_get_rope_index::GetRopeIndex); + m.def("get_rope_index_getworkspace", &wallx_cuda_get_rope_index::GetRopeIndexGetWorkSpace); + m.def("get_window_index", &wallx_cuda_window_index::GetWindowIndex); + m.def("get_totals", &wallx_cuda_window_index::GetTotals); + m.def("permute", &wallx_cuda_permute::MoePermuteTopKOp); + m.def("unpermute", &wallx_cuda_permute::MoeRecoverTopKOp); + m.def("unpermute_bwd", &wallx_cuda_permute::MoeRecoverTopKBwdOp); + m.def("cub_sort_pair_get_storage_bytes", &wallx_cuda_permute::CubSortPairGetStorageBytes); +} diff --git a/wall_x/model/core/ops/csrc/common/activation_types.h b/wall_x/model/core/ops/csrc/common/activation_types.h new file mode 100644 index 0000000..57bc784 --- /dev/null +++ b/wall_x/model/core/ops/csrc/common/activation_types.h @@ -0,0 +1,11 @@ +#pragma once + +namespace wallx_cuda { + +enum class ActivationType { + RELU, + SILU, + GELU +}; + +} // namespace wallx_cuda diff --git a/wall_x/model/core/ops/csrc/common/activations.cuh b/wall_x/model/core/ops/csrc/common/activations.cuh new file mode 100644 index 0000000..4560eba --- /dev/null +++ b/wall_x/model/core/ops/csrc/common/activations.cuh @@ -0,0 +1,30 @@ +#pragma once + +#include +#include +#include "activation_types.h" + +namespace wallx_cuda { + +template +__forceinline__ __device__ T applyActivation(const T &x) { + if constexpr (activation_type == ActivationType::RELU) { + return x > (T)0.0f ? x : (T)0.0f; + } + else if constexpr (activation_type == ActivationType::SILU) { + return (T)((float)x / (1.0f + __expf((float)-x))); + } + else if constexpr (activation_type == ActivationType::GELU) { + // GELU implementation from vllm (gelu_new_kernel) + const float x_f = (float)x; + const float x3 = x_f * x_f * x_f; + const float t = tanhf(0.79788456f * (x_f + 0.044715f * x3)); + return (T)(0.5f * x_f * (1.0f + t)); + } + else { + // No activation matches + assert(false); + } +} + +} // namespace wallx_cuda diff --git a/wall_x/model/core/ops/csrc/common/cuda_utils.h b/wall_x/model/core/ops/csrc/common/cuda_utils.h new file mode 100644 index 0000000..3239d38 --- /dev/null +++ b/wall_x/model/core/ops/csrc/common/cuda_utils.h @@ -0,0 +1,111 @@ +// Re-enable CUDA half operators (PyTorch disables them) +#undef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ + +#pragma once + +#include +#include +#include +#include + +#include +#include +#include + +#define CUDA_CHECK(cmd) do { \ + cudaError_t result = cmd; \ + if (result != cudaSuccess) { \ + printf("[ERROR] CUDA error %s:%d '%s': (%d) %s\n", __FILE__, __LINE__, #cmd, (int)result, cudaGetErrorString(result)); \ + exit(-1); \ + } \ +} while(0) + +inline void syncAndCheck(const char* const file, int const line, bool force_check = false) { +#ifdef DEBUG + force_check = true; +#endif + if (force_check) { + cudaDeviceSynchronize(); + cudaError_t result = cudaGetLastError(); + if (result) { + throw std::runtime_error(std::string("[ST] CUDA runtime error: ") + cudaGetErrorString(result) + " " + + file + ":" + std::to_string(line) + " \n"); + } + } +} + +#define sync_check_cuda_error() syncAndCheck(__FILE__, __LINE__, false) +#define sync_check_cuda_error_force() syncAndCheck(__FILE__, __LINE__, true) + +// #ifdef DEBUG +#define ASSERT_CHECK(__cond) \ + do { \ + const bool __cond_var = (__cond); \ + if (!__cond_var) { \ + ::std::string __err_msg = \ + ::std::string("`") + #__cond + \ + "` check failed at " + \ + __FILE__ + ":" + \ + ::std::to_string(__LINE__); \ + throw std::runtime_error(__err_msg); \ + } \ + } while (0) +// #else +// #define ASSERT_CHECK(__cond) do { } while (0) +// #endif + +// Some stuff for indexing into an 1-D array +#define INDEX_2D(dim1, dim2, index1, index2) \ + (((int64_t)index1) * (dim2) + (index2)) +#define INDEX_3D(dim1, dim2, dim3, index1, index2, index3) \ + (((int64_t)index1) * (dim2) * (dim3) + ((int64_t)index2) * (dim3) + (index3)) +#define INDEX_4D(dim1, dim2, dim3, dim4, index1, index2, index3, index4) \ + (((int64_t)index1) * (dim2) * (dim3) * (dim4) + ((int64_t)index2) * (dim3) * (dim4) + ((int64_t)index3) * (dim4) + (index4)) +#define INDEX_5D(dim1, dim2, dim3, dim4, dim5, index1, index2, index3, index4, index5) \ + (((int64_t)index1) * (dim2) * (dim3) * (dim4) * (dim5) + ((int64_t)index2) * (dim3) * (dim4) * (dim5) + ((int64_t)index3) * (dim4) * (dim5) + (index4) * (dim5) + (index5)) + +template +struct VecTraits; + +template<> +struct VecTraits { + using Type = float4; + static constexpr int vec_size = 4; + __device__ static inline Type load(const float* p) { return *reinterpret_cast(p); } + __device__ static inline void store(float* p, const Type& v) { *reinterpret_cast(p) = v; } +}; + +template<> +struct VecTraits<__half> { + using Type = __half2; + static constexpr int vec_size = 2; + __device__ static inline Type load(const __half* p) { return *reinterpret_cast(p); } + __device__ static inline void store(__half* p, const Type& v) { *reinterpret_cast<__half2*>(p) = v; } +}; + +template<> +struct VecTraits<__nv_bfloat16> { + using Type = __nv_bfloat162; + static constexpr int vec_size = 2; + __device__ static inline Type load(const __nv_bfloat16* p) { return *reinterpret_cast(p); } + __device__ static inline void store(__nv_bfloat16* p, const Type& v) { *reinterpret_cast<__nv_bfloat162*>(p) = v; } +}; + +template +cudaDataType_t getCudaDataType() { + if (std::is_same::value) { + return CUDA_R_16F; + } + else if (std::is_same::value) { + return CUDA_R_16BF; + } + else if (std::is_same::value) { + return CUDA_R_32F; + } + else { + throw std::runtime_error("Cuda data type: Unsupported type"); + } +} diff --git a/csrc/rope_index.cu b/wall_x/model/core/ops/csrc/get_rope_index/get_rope_index.cu similarity index 53% rename from csrc/rope_index.cu rename to wall_x/model/core/ops/csrc/get_rope_index/get_rope_index.cu index fd0402d..bc53190 100644 --- a/csrc/rope_index.cu +++ b/wall_x/model/core/ops/csrc/get_rope_index/get_rope_index.cu @@ -5,38 +5,45 @@ #include #include -#include -#include // for int64_t +#include +#include +#include #include #include #include -#include -#include -#include -#include +#include "../common/cuda_utils.h" +#include +#include + +namespace wallx_cuda_get_rope_index { + +// Constants and macros #define MAX_SEQ_LEN 8192 #define MAX_VISION_TOKENS 64 #define WARP_SIZE 32 #define MAX_THREADS_PER_BLOCK 1024 +// Vision content descriptor struct VisionDescriptor { - int64_t start_pos; - int64_t token_pos; - int64_t patch_count; - int64_t grid_t, grid_h, grid_w; - float time_interval; - int64_t is_video; - int64_t position_offset; + int64_t start_pos; // Starting position in the sequence + int64_t token_pos; // Position of the vision token + int64_t patch_count; // Number of patches + int64_t grid_t, grid_h, grid_w; // Grid dimensions (T, H, W) + float time_interval; // Time interval (for video) + int64_t is_video; // 0 = image, 1 = video + int64_t position_offset; // Positional encoding offset }; +// Device function: fast integer division using reciprocal multiplication __device__ __forceinline__ int64_t fast_div(int64_t a, int64_t b) { return __float2ll_rd(__ll2float_rn(a) * __frcp_rn(__ll2float_rn(b))); } +// Device function: map 3D patch index to (t, h, w) coordinates __device__ __forceinline__ void get_3d_coords(int64_t patch_idx, int64_t H, int64_t W, int64_t &t, int64_t &h, int64_t &w) { @@ -47,11 +54,12 @@ __device__ __forceinline__ void get_3d_coords(int64_t patch_idx, int64_t H, int6 w = remaining - h * W; } +// Stage 1: Kernel to count image and video tokens per batch __global__ void compute_vision_counts( const int64_t *input_ids, // (batch_size, seq_len) const int64_t *attention_mask, // (batch_size, seq_len) - int64_t *image_counts, // (batch_size,) - int64_t *video_counts, // (batch_size,) + int64_t *image_counts, // (batch_size,) - output: number of images per batch + int64_t *video_counts, // (batch_size,) - output: number of videos per batch const int64_t batch_size, const int64_t seq_len, const int64_t image_token_id, @@ -67,16 +75,20 @@ __global__ void compute_vision_counts( __shared__ int64_t shared_image_counts[MAX_THREADS_PER_BLOCK]; __shared__ int64_t shared_video_counts[MAX_THREADS_PER_BLOCK]; + // Initialize per-thread counters int64_t thread_image_count = 0; int64_t thread_video_count = 0; + // Parallel scan over sequence tokens (each thread processes multiple tokens) for (int64_t i = thread_idx; i < seq_len - 1; i += blockDim.x) { + // Skip if masked out if ((attention_mask != nullptr) && attention_mask[batch_idx * seq_len + i] == 0) continue; int64_t token_id = input_ids[batch_idx * seq_len + i]; + // Check if current token is a vision start token followed by image/video token if (token_id == vision_start_token_id && i + 1 < seq_len) { int64_t next_token = input_ids[batch_idx * seq_len + i + 1]; @@ -92,10 +104,12 @@ __global__ void compute_vision_counts( } } + // Store per-thread counts in shared memory shared_image_counts[thread_idx] = thread_image_count; shared_video_counts[thread_idx] = thread_video_count; __syncthreads(); + // Parallel reduction to sum counts across threads for (int64_t stride = blockDim.x / 2; stride > 0; stride /= 2) { if (thread_idx < stride) @@ -106,6 +120,7 @@ __global__ void compute_vision_counts( __syncthreads(); } + // Write final result for this batch to global memory if (thread_idx == 0) { image_counts[batch_idx] = shared_image_counts[0]; @@ -114,6 +129,7 @@ __global__ void compute_vision_counts( } +// Stage 2: Preprocessing kernel - parse and analyze all vision content __global__ void preprocess_vision_tokens( const int64_t *input_ids, // (batch_size, seq_len) const int64_t *attention_mask, // (batch_size, seq_len) @@ -123,9 +139,9 @@ __global__ void preprocess_vision_tokens( const int64_t *image_counts, // (batch_size,) const int64_t *video_counts, // (batch_size,) VisionDescriptor *vision_desc, // (batch_size, MAX_VISION_TOKENS) - int64_t *vision_counts, // (batch_size,) - int64_t *text_lengths, // (batch_size, MAX_VISION_TOKENS+1) - int64_t *position_offsets, // (batch_size, MAX_VISION_TOKENS+1) + int64_t *vision_counts, // (batch_size,) - total vision tokens per batch + int64_t *text_lengths, // (batch_size, MAX_VISION_TOKENS+1) - lengths of text segments + int64_t *position_offsets, // (batch_size, MAX_VISION_TOKENS+1) - cumulative position offsets const int64_t batch_size, const int64_t seq_len, const int64_t spatial_merge_size, @@ -134,10 +150,12 @@ __global__ void preprocess_vision_tokens( const int64_t vision_start_token_id, const float tokens_per_second) { + // Each thread processes one batch sequence int64_t batch_idx = blockIdx.x * blockDim.x + threadIdx.x; if (batch_idx >= batch_size) return; + // Compute cumulative image/video indices before this batch int64_t image_idx = 0, video_idx = 0; for (int64_t i = 0; i < batch_idx; i++) { @@ -145,34 +163,41 @@ __global__ void preprocess_vision_tokens( video_idx += video_counts[i]; } + // Local state int64_t vision_count = 0; int64_t current_pos = 0; int64_t position_offset = 0; + // Sequential scan over the sequence (single-threaded per batch) for (int64_t i = 0; i < seq_len - 1; i++) { + // Skip masked tokens if ((attention_mask != nullptr) && attention_mask[batch_idx * seq_len + i] == 0) continue; int64_t token_id = input_ids[batch_idx * seq_len + i]; + // Check for vision start token if (token_id == vision_start_token_id) { int64_t next_token = input_ids[batch_idx * seq_len + i + 1]; if (next_token == image_token_id || next_token == video_token_id) { + // Record length of preceding text segment int64_t vision_pos = i + 1; text_lengths[batch_idx * (MAX_VISION_TOKENS + 1) + vision_count] = vision_pos - current_pos; position_offsets[batch_idx * (MAX_VISION_TOKENS + 1) + vision_count] = position_offset; position_offset += (vision_pos - current_pos); + // Fetch grid dimensions int64_t T, H, W; float time_interval = 0.0f; int64_t is_video = (next_token == video_token_id) ? 1 : 0; if (is_video == 0) { + // Image case T = image_grid_thw[image_idx * 3 + 0]; H = image_grid_thw[image_idx * 3 + 1]; W = image_grid_thw[image_idx * 3 + 2]; @@ -180,6 +205,7 @@ __global__ void preprocess_vision_tokens( } else { + // Video case T = video_grid_thw[video_idx * 3 + 0]; H = video_grid_thw[video_idx * 3 + 1]; W = video_grid_thw[video_idx * 3 + 2]; @@ -187,10 +213,12 @@ __global__ void preprocess_vision_tokens( video_idx++; } + // Compute number of patches after spatial merging int64_t H_merged = H / spatial_merge_size; int64_t W_merged = W / spatial_merge_size; int64_t patch_count = T * H_merged * W_merged; + // Populate vision descriptor if (vision_count < MAX_VISION_TOKENS) { VisionDescriptor &desc = vision_desc[batch_idx * MAX_VISION_TOKENS + vision_count]; @@ -204,6 +232,7 @@ __global__ void preprocess_vision_tokens( desc.is_video = is_video; desc.position_offset = position_offset; + // Update position offset based on content type if (is_video) { position_offset += max(static_cast((T - 1) * time_interval * tokens_per_second) + 1, @@ -218,16 +247,20 @@ __global__ void preprocess_vision_tokens( vision_count++; } + // Skip over the vision patch tokens i = vision_pos + patch_count - 1; } } } + // Store total vision token count for this batch vision_counts[batch_idx] = vision_count; + // Handle final text segment after last vision token int64_t effective_len = seq_len; if (attention_mask != nullptr) { + // Find last valid (unmasked) token for (int64_t i = seq_len - 1; i >= 0; i--) { if (attention_mask[batch_idx * seq_len + i] != 0) @@ -243,6 +276,7 @@ __global__ void preprocess_vision_tokens( } +// Stage 3: Main kernel - compute 3D position IDs for all tokens in parallel __global__ void compute_3d_positions( const int64_t *input_ids, // (batch_size, seq_len) const int64_t *attention_mask, // (batch_size, seq_len) @@ -250,12 +284,13 @@ __global__ void compute_3d_positions( const int64_t *vision_counts, // (batch_size,) const int64_t *text_lengths, // (batch_size, MAX_VISION_TOKENS+1) const int64_t *position_offsets, // (batch_size, MAX_VISION_TOKENS+1) - int64_t *position_ids, // (3, batch_size, seq_len) - int64_t *mrope_deltas, // (batch_size,) + int64_t *position_ids, // (3, batch_size, seq_len) - output 3D position IDs + int64_t *mrope_deltas, // (batch_size,) - RoPE length adjustment const int64_t batch_size, const int64_t seq_len, const float tokens_per_second) { + // Grid: (batch_size), Block: (threads_per_block) int64_t batch_idx = blockIdx.x; int64_t thread_idx = threadIdx.x; @@ -264,10 +299,11 @@ __global__ void compute_3d_positions( __shared__ VisionDescriptor shared_visions[MAX_VISION_TOKENS]; __shared__ int64_t shared_position_offsets[MAX_VISION_TOKENS + 1]; - __shared__ int64_t shared_max_positions[MAX_THREADS_PER_BLOCK]; + __shared__ int64_t shared_max_positions[MAX_THREADS_PER_BLOCK]; // For block-wide reduction int64_t shared_vision_count = vision_counts[batch_idx]; + // Cooperative loading of vision descriptors into shared memory for (int64_t i = thread_idx; i < MAX_VISION_TOKENS; i += blockDim.x) { if (i < shared_vision_count) @@ -281,40 +317,54 @@ __global__ void compute_3d_positions( shared_position_offsets[i] = position_offsets[batch_idx * (MAX_VISION_TOKENS + 1) + i]; } + // Initialize per-thread max position int64_t thread_max_position = -1; __syncthreads(); + // Parallel processing: each thread handles multiple tokens for (int64_t token_idx = thread_idx; token_idx < seq_len; token_idx += blockDim.x) { + // Check validity via attention mask + int mask_offset = 0; bool is_valid_token = true; if (attention_mask != nullptr) { is_valid_token = attention_mask[batch_idx * seq_len + token_idx] != 0; + for (int i = 0; i < token_idx; i++) { + if (attention_mask[batch_idx * seq_len + i] == 0) { + mask_offset += 1; + } + } } if (!is_valid_token) { + // Set masked tokens to default position (1) position_ids[0 * batch_size * seq_len + batch_idx * seq_len + token_idx] = 1; position_ids[1 * batch_size * seq_len + batch_idx * seq_len + token_idx] = 1; position_ids[2 * batch_size * seq_len + batch_idx * seq_len + token_idx] = 1; continue; } + // Determine which segment this token belongs to int64_t segment_idx = -9999999; int64_t local_pos = token_idx; + // Linear search (vision count is small, so this is efficient) for (int64_t v = 0; v < shared_vision_count; v++) { if (token_idx < shared_visions[v].token_pos) { + // Token belongs to text segment before vision v segment_idx = v; local_pos = token_idx - (v > 0 ? shared_visions[v - 1].token_pos + shared_visions[v - 1].patch_count : 0); break; } else if (token_idx < shared_visions[v].token_pos + shared_visions[v].patch_count) { - segment_idx = -(v + 1); + // Token belongs to vision patch v + segment_idx = -(v + 1); // Negative index indicates vision local_pos = token_idx - shared_visions[v].token_pos; break; } @@ -322,6 +372,7 @@ __global__ void compute_3d_positions( if (segment_idx == -9999999) { + // Token belongs to final text segment segment_idx = shared_vision_count; int64_t last_vision_end = 0; if (shared_vision_count > 0) @@ -336,33 +387,41 @@ __global__ void compute_3d_positions( if (segment_idx >= 0) { + // Text token: use 1D position encoding int64_t offset = shared_position_offsets[segment_idx]; pos_t = pos_h = pos_w = offset + local_pos; } else { + // Vision patch: compute 3D position encoding int64_t vision_idx = -(segment_idx + 1); const VisionDescriptor &desc = shared_visions[vision_idx]; + // Map linear patch index to (t, h, w) int64_t t, h, w; get_3d_coords(local_pos, desc.grid_h, desc.grid_w, t, h, w); + // Compute 3D positions with temporal scaling for video pos_t = static_cast(t * desc.time_interval * tokens_per_second) + desc.position_offset; pos_h = h + desc.position_offset; pos_w = w + desc.position_offset; } - position_ids[0 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_t; - position_ids[1 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_h; - position_ids[2 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_w; + // Write 3D position IDs + position_ids[0 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_t - mask_offset; + position_ids[1 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_h - mask_offset; + position_ids[2 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_w - mask_offset; - int64_t max_pos = max(pos_t, max(pos_h, pos_w)); + // Track per-thread maximum position + int64_t max_pos = max(pos_t - mask_offset, max(pos_h - mask_offset, pos_w - mask_offset)); thread_max_position = max(thread_max_position, max_pos); } + // Store per-thread max into shared memory shared_max_positions[thread_idx] = thread_max_position; __syncthreads(); + // Block-wide reduction to find global max position for (int64_t stride = blockDim.x / 2; stride > 0; stride /= 2) { if (thread_idx < stride) @@ -373,6 +432,7 @@ __global__ void compute_3d_positions( __syncthreads(); } + // Compute mRoPE delta: (max_pos + 1) - seq_len if (thread_idx == 0) { int64_t global_max_position = shared_max_positions[0]; @@ -380,68 +440,231 @@ __global__ void compute_3d_positions( } } +// Fallback kernel when no vision tokens exist: assign sequential positions based on attention mask +__global__ void compute_3d_positions_mask_text( + const int64_t* __restrict__ attention_mask, // (batch_size, seq_len) + int64_t* __restrict__ position_ids, // (3, batch_size, seq_len) - flattened row-major + int64_t* __restrict__ mrope_deltas, // (batch_size,) + const int64_t batch_size, + const int64_t seq_len +) { + // One block per batch + int batch_idx = blockIdx.x; + if (batch_idx >= batch_size) return; -void launch_optimized_3d_rope_kernel( - const int64_t *input_ids, - const int64_t *attention_mask, - const int64_t *image_grid_thw, - const int64_t *video_grid_thw, - const float *second_per_grid_ts, - int64_t *position_ids, - int64_t *mrope_deltas, - int64_t batch_size, - int64_t seq_len, - int64_t spatial_merge_size, - int64_t image_token_id, - int64_t video_token_id, - int64_t vision_start_token_id, - float tokens_per_second) -{ - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + extern __shared__ int64_t shared_mask[]; + int64_t* mask = shared_mask; - torch::Device device(torch::kCUDA, at::cuda::current_device()); + // Load attention mask for current batch into shared memory + for (int i = threadIdx.x; i < seq_len; i += blockDim.x) { + mask[i] = attention_mask[batch_idx * seq_len + i]; + } + __syncthreads(); - auto options = torch::TensorOptions().dtype(torch::kInt64).device(device); + int tid = threadIdx.x; + int64_t base_offset = batch_idx * seq_len; + int64_t total_valid = 0; - auto vision_desc_tensor = torch::empty( - {batch_size * MAX_VISION_TOKENS * static_cast(sizeof(VisionDescriptor))}, - torch::TensorOptions().dtype(torch::kUInt8).device(device) - ); - VisionDescriptor *d_vision_desc = reinterpret_cast(vision_desc_tensor.data_ptr()); + // Compute prefix sum (cumulative count of valid tokens) + // Simple O(seq_len^2) per-thread loop - acceptable for moderate seq_len + for (int idx = tid; idx < seq_len; idx += blockDim.x) { + int64_t sum = 0; + for (int i = 0; i <= idx; ++i) { + sum += mask[i]; + } + int64_t pos_val = (mask[idx] ? (sum - 1) : 1); // Use 0-based if valid, else 1 - auto vision_counts_tensor = torch::empty({batch_size}, options); - auto text_lengths_tensor = torch::empty({batch_size * (MAX_VISION_TOKENS + 1)}, options); - auto position_offsets_tensor = torch::empty({batch_size * (MAX_VISION_TOKENS + 1)}, options); - auto image_counts_tensor = torch::empty({batch_size}, options); - auto video_counts_tensor = torch::empty({batch_size}, options); + // Write same position to all three dimensions + int64_t flat_idx0 = 0 * batch_size * seq_len + base_offset + idx; + int64_t flat_idx1 = 1 * batch_size * seq_len + base_offset + idx; + int64_t flat_idx2 = 2 * batch_size * seq_len + base_offset + idx; - int64_t *d_vision_counts = vision_counts_tensor.data_ptr(); - int64_t *d_text_lengths = text_lengths_tensor.data_ptr(); - int64_t *d_position_offsets = position_offsets_tensor.data_ptr(); - int64_t *d_image_counts = image_counts_tensor.data_ptr(); - int64_t *d_video_counts = video_counts_tensor.data_ptr(); + position_ids[flat_idx0] = pos_val; + position_ids[flat_idx1] = pos_val; + position_ids[flat_idx2] = pos_val; + // Only thread 0 tracks total valid tokens + if (tid == 0) { + total_valid = sum; // After last idx, sum = total number of 1s + } + } + + // Compute mrope_deltas: (total_valid) - seq_len + if (tid == 0) { + // Re-compute total_valid robustly + total_valid = 0; + for (int i = 0; i < seq_len; ++i) { + total_valid += mask[i]; + } + mrope_deltas[batch_idx] = total_valid - seq_len; + } +} + +// Kernel to fill position_ids with arange(seq_len) in 3D layout when no vision or mask +__global__ void arange_3d(int64_t seq_len, + int64_t* out) { + int i = blockIdx.x; // dimension index: 0, 1, 2 + int j = blockIdx.y; // batch index + int64_t k = blockIdx.z * blockDim.x + threadIdx.x; // token index + + if (k >= seq_len) { + return; + } + + // Flatten index: (dim, batch, token) -> linear + int64_t idx = i * gridDim.y * seq_len + j * seq_len + k; + out[idx] = k; +} + +// Kernel to zero-initialize a 1D tensor +__global__ void fill_zeros_1d(int64_t numel, int64_t* data) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < numel) { + data[idx] = 0; + } +} + +// Compute required workspace size (in bytes) +int64_t GetRopeIndexGetWorkSpace(const c10::optional& input_ids, + const c10::optional& image_grid_thw, + const c10::optional& video_grid_thw) { + + if (!image_grid_thw.has_value() && !video_grid_thw.has_value()) { + return 1; // minimal workspace + } else { + const at::Tensor& input_ids_tensor = input_ids.value(); + const int batch_size = input_ids_tensor.size(0); + return (batch_size * MAX_VISION_TOKENS * static_cast(sizeof(VisionDescriptor)) + + 3 * batch_size * sizeof(int64_t) + + 2 * batch_size * (MAX_VISION_TOKENS + 1) * sizeof(int64_t)); + } +} + +// Main entry function: compute RoPE position IDs and mRoPE deltas +void GetRopeIndex(const c10::optional& input_ids, + const c10::optional& image_grid_thw, + const c10::optional& video_grid_thw, + const c10::optional& second_per_grid_ts, + const c10::optional& attention_mask, + const at::Tensor& position_ids, + const at::Tensor& mrope_deltas, + const at::Tensor& workspace, + int spatial_merge_size, + int image_token_id, + int video_token_id, + int vision_start_token_id, + float tokens_per_second) { + + ASSERT_CHECK(input_ids.has_value()); + + const at::Tensor& input_ids_tensor = input_ids.value(); + ASSERT_CHECK(input_ids_tensor.dim() == 2); + ASSERT_CHECK(input_ids_tensor.is_contiguous()); + + const auto batch_size = input_ids_tensor.size(0); + const auto seq_len = input_ids_tensor.size(1); + const auto device = input_ids_tensor.device(); + + const int64_t *input_ids_ptr = static_cast(input_ids_tensor.data_ptr()); + + const int64_t *attention_mask_ptr = nullptr; + + if (attention_mask.has_value()) { + const at::Tensor& attention_mask_tensor = attention_mask.value(); + ASSERT_CHECK(attention_mask_tensor.is_contiguous()); + attention_mask_ptr = static_cast(attention_mask_tensor.data_ptr()); + } + + const int64_t *image_grid_thw_ptr = nullptr; + + if (image_grid_thw.has_value()) { + const at::Tensor& image_grid_thw_tensor = image_grid_thw.value(); + ASSERT_CHECK(image_grid_thw_tensor.dim() == 2 && image_grid_thw_tensor.size(1) == 3); + ASSERT_CHECK(image_grid_thw_tensor.is_contiguous()); + image_grid_thw_ptr = static_cast(image_grid_thw_tensor.data_ptr()); + } + + const int64_t *video_grid_thw_ptr = nullptr; + + if (video_grid_thw.has_value()) { + const at::Tensor& video_grid_thw_tensor = video_grid_thw.value(); + ASSERT_CHECK(video_grid_thw_tensor.is_contiguous()); + ASSERT_CHECK(video_grid_thw_tensor.dim() == 2 && video_grid_thw_tensor.size(1) == 3); + video_grid_thw_ptr = static_cast(video_grid_thw_tensor.data_ptr()); + } + + const float *second_per_grid_ts_ptr = nullptr; + + if (second_per_grid_ts.has_value()) { + const at::Tensor& second_per_grid_ts_tensor = second_per_grid_ts.value(); + ASSERT_CHECK(second_per_grid_ts_tensor.is_contiguous()); + ASSERT_CHECK(second_per_grid_ts_tensor.dim() == 1); + second_per_grid_ts_ptr = static_cast(second_per_grid_ts_tensor.data_ptr()); + } + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + // Case: no vision tokens -> use simple positional encoding + if (!image_grid_thw.has_value() && !video_grid_thw.has_value()) { + if (attention_mask.has_value()) { + // Use attention mask to compute valid positions + int block_size = min(1024, (int)seq_len); + dim3 grid(batch_size); + dim3 block(block_size); + size_t shared_mem_size = seq_len * sizeof(int64_t); + + compute_3d_positions_mask_text<<>>( + attention_mask_ptr, static_cast(position_ids.data_ptr()), static_cast(mrope_deltas.data_ptr()), batch_size, seq_len + ); + return; + } else { + // No mask, no vision -> use simple arange + dim3 thread_num(256); + dim3 grid(3, batch_size, (seq_len + 255) / 256); + arange_3d<<>>(seq_len, static_cast(position_ids.data_ptr())); + + int block_size = 256; + int grid_size = (batch_size + block_size - 1) / block_size; + fill_zeros_1d<<>>(batch_size, static_cast(mrope_deltas.data_ptr())); + return; + } + } + + // Allocate workspace buffers + void* work_ptr = workspace.data_ptr(); + VisionDescriptor *d_vision_desc = reinterpret_cast(work_ptr); + + int64_t *d_vision_counts = reinterpret_cast(reinterpret_cast(work_ptr) + batch_size * MAX_VISION_TOKENS); + int64_t *d_text_lengths = d_vision_counts + batch_size; + int64_t *d_position_offsets = d_text_lengths + batch_size * (MAX_VISION_TOKENS + 1); + int64_t *d_image_counts = d_position_offsets + batch_size * (MAX_VISION_TOKENS + 1); + int64_t *d_video_counts = d_image_counts + batch_size; + + // Stage 1: count vision tokens dim3 index_grid(static_cast(batch_size)); dim3 index_block(256); compute_vision_counts<<>>( - input_ids, attention_mask, + input_ids_ptr, attention_mask_ptr, d_image_counts, d_video_counts, batch_size, seq_len, image_token_id, video_token_id, vision_start_token_id); - int64_t threads_per_block = std::min(batch_size, static_cast(256)); + // Stage 2: preprocess vision tokens + int64_t threads_per_block = std::min(static_cast(batch_size), static_cast(256)); int64_t num_blocks = (batch_size + threads_per_block - 1) / threads_per_block; dim3 preprocess_grid(static_cast(num_blocks)); dim3 preprocess_block(static_cast(threads_per_block)); preprocess_vision_tokens<<>>( - input_ids, attention_mask, image_grid_thw, video_grid_thw, - second_per_grid_ts, d_image_counts, d_video_counts, + input_ids_ptr, attention_mask_ptr, image_grid_thw_ptr, video_grid_thw_ptr, + second_per_grid_ts_ptr, d_image_counts, d_video_counts, d_vision_desc, d_vision_counts, d_text_lengths, d_position_offsets, batch_size, seq_len, spatial_merge_size, image_token_id, video_token_id, vision_start_token_id, tokens_per_second); + // Stage 3: compute 3D positions threads_per_block = std::min(static_cast(seq_len), static_cast(MAX_THREADS_PER_BLOCK)); if (threads_per_block < 32) threads_per_block = 32; @@ -454,116 +677,16 @@ void launch_optimized_3d_rope_kernel( dim3 compute_block(static_cast(threads_per_block)); compute_3d_positions<<>>( - input_ids, attention_mask, d_vision_desc, d_vision_counts, - d_text_lengths, d_position_offsets, position_ids, mrope_deltas, + input_ids_ptr, attention_mask_ptr, d_vision_desc, d_vision_counts, + d_text_lengths, d_position_offsets, + static_cast(position_ids.data_ptr()), static_cast(mrope_deltas.data_ptr()), batch_size, seq_len, tokens_per_second); - AT_CUDA_CHECK(cudaGetLastError()); + sync_check_cuda_error(); } -std::tuple get_rope_index( - const torch::optional &input_ids, - const torch::optional &image_grid_thw, - const torch::optional &video_grid_thw, - const torch::optional &second_per_grid_ts, - const torch::optional &attention_mask, - int spatial_merge_size, - int image_token_id, - int video_token_id, - int vision_start_token_id, - float tokens_per_second) -{ - TORCH_CHECK(input_ids.has_value(), "input_ids cannot be None"); - TORCH_CHECK(input_ids->dim() == 2, "input_ids must be 2D tensor (batch_size, seq_len)"); +} // wallx_cuda_get_rope_index - const auto batch_size = input_ids->size(0); - const auto seq_len = input_ids->size(1); - const auto device = input_ids->device(); - - TORCH_CHECK(device.is_cuda(), "All tensors must be on CUDA device"); - - torch::Tensor input_ids_tensor = input_ids->contiguous(); - const int64_t *input_ids_ptr = input_ids_tensor.data_ptr(); - - torch::Tensor attention_mask_tensor; - const int64_t *attention_mask_ptr = nullptr; - - if (attention_mask.has_value()) { - attention_mask_tensor = attention_mask->contiguous(); - attention_mask_ptr = attention_mask_tensor.data_ptr(); - } - - torch::Tensor image_grid_thw_tensor; - const int64_t *image_grid_thw_ptr = nullptr; - - if (image_grid_thw.has_value()) { - TORCH_CHECK(image_grid_thw->dim() == 2 && image_grid_thw->size(1) == 3, - "image_grid_thw must be shape (num_images, 3)"); - image_grid_thw_tensor = image_grid_thw->contiguous(); - image_grid_thw_ptr = image_grid_thw_tensor.data_ptr(); - } - - torch::Tensor video_grid_thw_tensor; - const int64_t *video_grid_thw_ptr = nullptr; - - if (video_grid_thw.has_value()) { - TORCH_CHECK(video_grid_thw->dim() == 2 && video_grid_thw->size(1) == 3, - "video_grid_thw must be shape (num_videos, 3)"); - video_grid_thw_tensor = video_grid_thw->contiguous(); - video_grid_thw_ptr = video_grid_thw_tensor.data_ptr(); - } - - torch::Tensor second_per_grid_ts_tensor; - const float *second_per_grid_ts_ptr = nullptr; - - if (second_per_grid_ts.has_value()) { - TORCH_CHECK(second_per_grid_ts->dim() == 1, - "second_per_grid_ts must be 1D tensor"); - second_per_grid_ts_tensor = second_per_grid_ts->contiguous(); - second_per_grid_ts_ptr = second_per_grid_ts_tensor.data_ptr(); - } - - if (!image_grid_thw.has_value() && !video_grid_thw.has_value()) { - torch::Tensor position_ids; - torch::Tensor mrope_deltas; - - if (attention_mask.has_value()) { - auto cumsum_result = attention_mask_tensor.to(torch::kInt64).cumsum(-1) - 1; - position_ids = cumsum_result.masked_fill_(attention_mask_tensor.eq(0), 1); - position_ids = position_ids.unsqueeze(0).expand({3, -1, -1}); - - auto max_position_ids = std::get<0>(std::get<0>(position_ids.max(0)).max(-1, true)); - mrope_deltas = max_position_ids + 1 - attention_mask_tensor.size(-1); - mrope_deltas = mrope_deltas.view({batch_size, 1}); - } else { - auto pos_range = torch::arange(seq_len, torch::TensorOptions().dtype(torch::kInt64).device(device)); - position_ids = pos_range.view({1, 1, -1}).expand({3, batch_size, -1}); - mrope_deltas = torch::zeros({batch_size, 1}, torch::TensorOptions().dtype(torch::kInt64).device(device)); - } - - return std::make_tuple(position_ids, mrope_deltas); - } - - auto position_ids = torch::empty({3, batch_size, seq_len}, - torch::TensorOptions().dtype(torch::kInt64).device(device)); - auto mrope_deltas = torch::empty({batch_size, 1}, - torch::TensorOptions().dtype(torch::kInt64).device(device)); - - launch_optimized_3d_rope_kernel( - input_ids_ptr, - attention_mask_ptr, - image_grid_thw_ptr, - video_grid_thw_ptr, - second_per_grid_ts_ptr, - position_ids.data_ptr(), - mrope_deltas.data_ptr(), - static_cast(batch_size), - static_cast(seq_len), - static_cast(spatial_merge_size), - static_cast(image_token_id), - static_cast(video_token_id), - static_cast(vision_start_token_id), - tokens_per_second); - - return std::make_tuple(position_ids, mrope_deltas); -} +// TVM FFI exports +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) diff --git a/wall_x/model/core/ops/csrc/m_rope/m_rope_kernel.cu b/wall_x/model/core/ops/csrc/m_rope/m_rope_kernel.cu new file mode 100644 index 0000000..12422a9 --- /dev/null +++ b/wall_x/model/core/ops/csrc/m_rope/m_rope_kernel.cu @@ -0,0 +1,840 @@ +// Re-enable CUDA half operators (PyTorch disables them) +#undef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ + +#include +#include + +#include + #include +#include +#include +#include +#include "../common/cuda_utils.h" + +namespace wallx_cuda_m_rope { + +// Helper: convert vector of T to float4 +template +__device__ __forceinline__ float4 to_float4(const T* ptr) { + if constexpr (std::is_same_v) { + half2 h2_0 = reinterpret_cast(ptr)[0]; + half2 h2_1 = reinterpret_cast(ptr)[1]; + float2 f2_0 = __half22float2(h2_0); + float2 f2_1 = __half22float2(h2_1); + return make_float4(f2_0.x, f2_0.y, f2_1.x, f2_1.y); + } else if constexpr (std::is_same_v) { + __nv_bfloat162 b2_0 = reinterpret_cast(ptr)[0]; + __nv_bfloat162 b2_1 = reinterpret_cast(ptr)[1]; +#if __CUDA_ARCH__ >= 800 + float2 f2_0 = __bfloat1622float2(b2_0); + float2 f2_1 = __bfloat1622float2(b2_1); + return make_float4(f2_0.x, f2_0.y, f2_1.x, f2_1.y); +#else + assert(false && "Unsupported on arch < 800"); +#endif + } else if constexpr (std::is_same_v) { + return reinterpret_cast(ptr)[0]; + } else { + static_assert(sizeof(T) == 0, "Unsupported type"); + } +} + +// Helper: convert float4 back to vector of T +template +__device__ __forceinline__ void from_float4(T* ptr, const float4& f4) { + if constexpr (std::is_same_v) { + half2 h2_0 = __float22half2_rn(make_float2(f4.x, f4.y)); + half2 h2_1 = __float22half2_rn(make_float2(f4.z, f4.w)); + reinterpret_cast(ptr)[0] = h2_0; + reinterpret_cast(ptr)[1] = h2_1; + } else if constexpr (std::is_same_v) { +#if __CUDA_ARCH__ >= 800 + __nv_bfloat162 b2_0 = __float22bfloat162_rn(make_float2(f4.x, f4.y)); + __nv_bfloat162 b2_1 = __float22bfloat162_rn(make_float2(f4.z, f4.w)); + reinterpret_cast<__nv_bfloat162*>(ptr)[0] = b2_0; + reinterpret_cast<__nv_bfloat162*>(ptr)[1] = b2_1; +#else + assert(false && "Unsupported on arch < 800"); +#endif + } else if constexpr (std::is_same_v) { + reinterpret_cast(ptr)[0] = f4; + } else { + static_assert(sizeof(T) == 0, "Unsupported type"); + } +} + +template +__global__ void MRopeKernel(const float* cos, + const float* sin, + const T* q, + const int q_h, + const T* k, + const int k_h, + T* q_embed, + T* k_embed, + const int first, + const int second, + const int d, + const int qb_stride, + const int qs_stride, + const int qh_stride, + const int kb_stride, + const int ks_stride, + const int kh_stride, + const int qeb_stride, + const int qes_stride, + const int qeh_stride, + const int keb_stride, + const int kes_stride, + const int keh_stride) { + extern __shared__ char cos_sin[]; + const int half_dim = d / 2; + float* cos_smem = reinterpret_cast(cos_sin); + float* sin_smem = cos_smem + half_dim; + int b = blockIdx.x; + int s = blockIdx.y; + + int64_t offset = gridDim.x * gridDim.y * half_dim; + int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t cos_sin_s_stride = half_dim; +#define SIN_GMEM(a, b, c, d) sin[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] +#define COS_GMEM(a, b, c, d) cos[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] + + for (int i = threadIdx.x; i < half_dim; i += blockDim.x) { + if (i < first) { + cos_smem[i] = COS_GMEM(0, b, s, i); + sin_smem[i] = SIN_GMEM(0, b, s, i); + } else if (i < (second + first)) { + cos_smem[i] = COS_GMEM(1, b, s, i); + sin_smem[i] = SIN_GMEM(1, b, s, i); + } else { + cos_smem[i] = COS_GMEM(2, b, s, i); + sin_smem[i] = SIN_GMEM(2, b, s, i); + } + } + __syncthreads(); + +#define Q_GMEM(a, b, c, d) q[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + d] +#define K_GMEM(a, b, c, d) k[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + d] +#define Q_EMBED_GMEM(a, b, c, d) q_embed[(a) * qeb_stride + (b) * qes_stride + (c) * qeh_stride + d] +#define K_EMBED_GMEM(a, b, c, d) k_embed[(a) * keb_stride + (b) * kes_stride + (c) * keh_stride + d] + + for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + float4 q_x0 = to_float4(&Q_GMEM(b, s, h_idx, base_d)); + float4 q_x1 = to_float4(&Q_GMEM(b, s, h_idx, base_d + half_dim)); + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + + float4 q0_rot = make_float4( + q_x0.x * cos_vec.x - q_x1.x * sin_vec.x, + q_x0.y * cos_vec.y - q_x1.y * sin_vec.y, + q_x0.z * cos_vec.z - q_x1.z * sin_vec.z, + q_x0.w * cos_vec.w - q_x1.w * sin_vec.w + ); + + float4 q1_rot = make_float4( + q_x1.x * cos_vec.x + q_x0.x * sin_vec.x, + q_x1.y * cos_vec.y + q_x0.y * sin_vec.y, + q_x1.z * cos_vec.z + q_x0.z * sin_vec.z, + q_x1.w * cos_vec.w + q_x0.w * sin_vec.w + ); + + from_float4(&Q_EMBED_GMEM(b, s, h_idx, base_d), q0_rot); + from_float4(&Q_EMBED_GMEM(b, s, h_idx, base_d + half_dim), q1_rot); + } + + for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + float4 k_x0 = to_float4(&K_GMEM(b, s, h_idx, base_d)); + float4 k_x1 = to_float4(&K_GMEM(b, s, h_idx, base_d + half_dim)); + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + + float4 k0_rot = make_float4( + k_x0.x * cos_vec.x - k_x1.x * sin_vec.x, + k_x0.y * cos_vec.y - k_x1.y * sin_vec.y, + k_x0.z * cos_vec.z - k_x1.z * sin_vec.z, + k_x0.w * cos_vec.w - k_x1.w * sin_vec.w + ); + + float4 k1_rot = make_float4( + k_x1.x * cos_vec.x + k_x0.x * sin_vec.x, + k_x1.y * cos_vec.y + k_x0.y * sin_vec.y, + k_x1.z * cos_vec.z + k_x0.z * sin_vec.z, + k_x1.w * cos_vec.w + k_x0.w * sin_vec.w + ); + + from_float4(&K_EMBED_GMEM(b, s, h_idx, base_d), k0_rot); + from_float4(&K_EMBED_GMEM(b, s, h_idx, base_d + half_dim), k1_rot); + } +} + +template +__global__ void MRopeInplaceKernel(float* cos, float* sin, + T* q, const int q_h, T* k, const int k_h, + const int first, const int second, + const int d, const int qb_stride, const int qs_stride, const int qh_stride, + const int kb_stride, const int ks_stride, const int kh_stride) { + extern __shared__ char cos_sin[]; + const int half_dim = d / 2; + float* cos_smem = reinterpret_cast(cos_sin); + float* sin_smem = cos_smem + half_dim; + int b = blockIdx.x; + int s = blockIdx.y; + + int64_t offset = gridDim.x * gridDim.y * half_dim; + int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t cos_sin_s_stride = half_dim; +#define SIN_GMEM(a, b, c, d) sin[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] +#define COS_GMEM(a, b, c, d) cos[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] + + for (int i = threadIdx.x; i < half_dim; i += blockDim.x) { + if (i < first) { + cos_smem[i] = COS_GMEM(0, b, s, i); + sin_smem[i] = SIN_GMEM(0, b, s, i); + } else if (i < (second + first)) { + cos_smem[i] = COS_GMEM(1, b, s, i); + sin_smem[i] = SIN_GMEM(1, b, s, i); + } else { + cos_smem[i] = COS_GMEM(2, b, s, i); + sin_smem[i] = SIN_GMEM(2, b, s, i); + } + } + __syncthreads(); + +#define Q_GMEM(a, b, c, d) q[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + d] +#define K_GMEM(a, b, c, d) k[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + d] + + for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + float4 q_x0 = to_float4(&Q_GMEM(b, s, h_idx, base_d)); + float4 q_x1 = to_float4(&Q_GMEM(b, s, h_idx, base_d + half_dim)); + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + + float4 q0_rot = make_float4( + q_x0.x * cos_vec.x - q_x1.x * sin_vec.x, + q_x0.y * cos_vec.y - q_x1.y * sin_vec.y, + q_x0.z * cos_vec.z - q_x1.z * sin_vec.z, + q_x0.w * cos_vec.w - q_x1.w * sin_vec.w + ); + + float4 q1_rot = make_float4( + q_x1.x * cos_vec.x + q_x0.x * sin_vec.x, + q_x1.y * cos_vec.y + q_x0.y * sin_vec.y, + q_x1.z * cos_vec.z + q_x0.z * sin_vec.z, + q_x1.w * cos_vec.w + q_x0.w * sin_vec.w + ); + + from_float4(&Q_GMEM(b, s, h_idx, base_d), q0_rot); + from_float4(&Q_GMEM(b, s, h_idx, base_d + half_dim), q1_rot); + } + + for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + float4 k_x0 = to_float4(&K_GMEM(b, s, h_idx, base_d)); + float4 k_x1 = to_float4(&K_GMEM(b, s, h_idx, base_d + half_dim)); + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + + float4 k0_rot = make_float4( + k_x0.x * cos_vec.x - k_x1.x * sin_vec.x, + k_x0.y * cos_vec.y - k_x1.y * sin_vec.y, + k_x0.z * cos_vec.z - k_x1.z * sin_vec.z, + k_x0.w * cos_vec.w - k_x1.w * sin_vec.w + ); + + float4 k1_rot = make_float4( + k_x1.x * cos_vec.x + k_x0.x * sin_vec.x, + k_x1.y * cos_vec.y + k_x0.y * sin_vec.y, + k_x1.z * cos_vec.z + k_x0.z * sin_vec.z, + k_x1.w * cos_vec.w + k_x0.w * sin_vec.w + ); + + from_float4(&K_GMEM(b, s, h_idx, base_d), k0_rot); + from_float4(&K_GMEM(b, s, h_idx, base_d + half_dim), k1_rot); + } +} + +template +__global__ void MRopeKernelBackward( + const float* cos, + const float* sin, + const T* grad_q_embed, + const T* grad_k_embed, + T* grad_q, + T* grad_k, + const int first, const int second, + const int q_h, + const int k_h, + const int d, + const int qb_stride, const int qs_stride, const int qh_stride, + const int kb_stride, const int ks_stride, const int kh_stride) { + + extern __shared__ char cos_sin[]; + const int half_dim = d / 2; + float* cos_smem = reinterpret_cast(cos_sin); + float* sin_smem = cos_smem + half_dim; + + int b = blockIdx.x; + int s = blockIdx.y; + + int64_t offset = gridDim.x * gridDim.y * half_dim; + int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t cos_sin_s_stride = half_dim; + +#define SIN_GMEM(a, b, c, d) sin[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] +#define COS_GMEM(a, b, c, d) cos[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] + + for (int i = threadIdx.x; i < half_dim; i += blockDim.x) { + if (i < first) { + cos_smem[i] = COS_GMEM(0, b, s, i); + sin_smem[i] = SIN_GMEM(0, b, s, i); + } else if (i < (second + first)) { + cos_smem[i] = COS_GMEM(1, b, s, i); + sin_smem[i] = SIN_GMEM(1, b, s, i); + } else { + cos_smem[i] = COS_GMEM(2, b, s, i); + sin_smem[i] = SIN_GMEM(2, b, s, i); + } + } + __syncthreads(); + +#define GQ_EMBED(a, b, c, d) grad_q_embed[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + (d)] +#define GK_EMBED(a, b, c, d) grad_k_embed[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + (d)] +#define GQ(a, b, c, d) grad_q[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + (d)] +#define GK(a, b, c, d) grad_k[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + (d)] + + // Process grad_q + for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + // float path + float4 gq0_rot = to_float4(&GQ_EMBED(b, s, h_idx, base_d)); + float4 gq1_rot = to_float4(&GQ_EMBED(b, s, h_idx, base_d + half_dim)); + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + + float4 gq0 = make_float4( + gq0_rot.x * cos_vec.x + gq1_rot.x * sin_vec.x, + gq0_rot.y * cos_vec.y + gq1_rot.y * sin_vec.y, + gq0_rot.z * cos_vec.z + gq1_rot.z * sin_vec.z, + gq0_rot.w * cos_vec.w + gq1_rot.w * sin_vec.w + ); + float4 gq1 = make_float4( + gq1_rot.x * cos_vec.x - gq0_rot.x * sin_vec.x, + gq1_rot.y * cos_vec.y - gq0_rot.y * sin_vec.y, + gq1_rot.z * cos_vec.z - gq0_rot.z * sin_vec.z, + gq1_rot.w * cos_vec.w - gq0_rot.w * sin_vec.w + ); + + from_float4(&GQ(b, s, h_idx, base_d), gq0); + from_float4(&GQ(b, s, h_idx, base_d + half_dim), gq1); + } + + // Process grad_k + for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + // float path + float4 gk0_rot = to_float4(&GK_EMBED(b, s, h_idx, base_d)); + float4 gk1_rot = to_float4(&GK_EMBED(b, s, h_idx, base_d + half_dim)); + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + + float4 gk0 = make_float4( + gk0_rot.x * cos_vec.x + gk1_rot.x * sin_vec.x, + gk0_rot.y * cos_vec.y + gk1_rot.y * sin_vec.y, + gk0_rot.z * cos_vec.z + gk1_rot.z * sin_vec.z, + gk0_rot.w * cos_vec.w + gk1_rot.w * sin_vec.w + ); + float4 gk1 = make_float4( + gk1_rot.x * cos_vec.x - gk0_rot.x * sin_vec.x, + gk1_rot.y * cos_vec.y - gk0_rot.y * sin_vec.y, + gk1_rot.z * cos_vec.z - gk0_rot.z * sin_vec.z, + gk1_rot.w * cos_vec.w - gk0_rot.w * sin_vec.w + ); + + from_float4(&GK(b, s, h_idx, base_d), gk0); + from_float4(&GK(b, s, h_idx, base_d + half_dim), gk1); + } +} + +void MRope(const at::Tensor& q, // [b, s, h, d] + const at::Tensor& k, // [b, s, h_k, d] + const at::Tensor& q_embed, // [b, s, h, d] + const at::Tensor& k_embed, // [b, s, h_k, d] + const at::Tensor& cos, // [3, b, s, d / 2] + const at::Tensor& sin, // [3, b, s, d / 2] + const int first, + const int second + ) { + int Nthreads = 256; + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + ASSERT_CHECK(q.scalar_type() == k.scalar_type()); + ASSERT_CHECK(q.scalar_type() == q_embed.scalar_type()); + ASSERT_CHECK(k.scalar_type() == k_embed.scalar_type()); + ASSERT_CHECK(q_embed.size(0) == q.size(0)); + ASSERT_CHECK(q_embed.size(1) == q.size(1)); + ASSERT_CHECK(q_embed.size(2) == q.size(2)); + ASSERT_CHECK(q_embed.size(3) == q.size(3)); + ASSERT_CHECK(k_embed.size(0) == k.size(0)); + ASSERT_CHECK(k_embed.size(1) == k.size(1)); + ASSERT_CHECK(k_embed.size(2) == k.size(2)); + ASSERT_CHECK(k_embed.size(3) == k.size(3)); + ASSERT_CHECK(q.size(3) % 8 == 0); + ASSERT_CHECK(cos.size(0) == 3); + ASSERT_CHECK(cos.size(3) == q.size(3) / 2); + ASSERT_CHECK(cos.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.size(0) == 3); + ASSERT_CHECK(q.stride(3) == 1); + ASSERT_CHECK(k.stride(3) == 1); + ASSERT_CHECK(q_embed.stride(3) == 1); + ASSERT_CHECK(k_embed.stride(3) == 1); + ASSERT_CHECK(cos.stride(3) == 1); + ASSERT_CHECK(sin.stride(3) == 1); + ASSERT_CHECK(cos.is_contiguous()); + ASSERT_CHECK(sin.is_contiguous()); + + int64_t batch = q.size(0); + int64_t q_head_num = q.size(2); + int64_t k_head_num = k.size(2); + int64_t seq_len = q.size(1); + int64_t dim = q.size(3); + int64_t qb_stride = q.stride(0); + int64_t kb_stride = k.stride(0); + int64_t qs_stride = q.stride(1); + int64_t ks_stride = k.stride(1); + int64_t qh_stride = q.stride(2); + int64_t kh_stride = k.stride(2); + int64_t qeb_stride = q_embed.stride(0); + int64_t keb_stride = k_embed.stride(0); + int64_t qes_stride = q_embed.stride(1); + int64_t kes_stride = k_embed.stride(1); + int64_t qeh_stride = q_embed.stride(2); + int64_t keh_stride = k_embed.stride(2); + + dim3 grid(batch, seq_len); + + if (q.scalar_type() == at::kHalf) { + const half* q_data = static_cast(q.data_ptr()); + const half* k_data = static_cast(k.data_ptr()); + half* q_embed_data = static_cast(q_embed.data_ptr()); + half* k_embed_data = static_cast(k_embed.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + MRopeKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + q_embed_data, k_embed_data, + first, second, dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); + } else if (q.scalar_type() == at::kBFloat16) { + const __nv_bfloat16* q_data = static_cast(q.data_ptr()); + const __nv_bfloat16* k_data = static_cast(k.data_ptr()); + __nv_bfloat16* q_embed_data = static_cast<__nv_bfloat16*>(q_embed.data_ptr()); + __nv_bfloat16* k_embed_data = static_cast<__nv_bfloat16*>(k_embed.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + MRopeKernel<__nv_bfloat16><<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + q_embed_data, k_embed_data, + first, second, dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); + } else if (q.scalar_type() == at::kFloat) { + const float* q_data = static_cast(q.data_ptr()); + const float* k_data = static_cast(k.data_ptr()); + float* q_embed_data = static_cast(q_embed.data_ptr()); + float* k_embed_data = static_cast(k_embed.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + MRopeKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + q_embed_data, k_embed_data, + first, second, dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride + ); + } else { + throw std::runtime_error("Unsupported data type for m_rope"); + } + + sync_check_cuda_error(); +} + +void MRopeInplace(const at::Tensor& q, // [b, s, h, d] + const at::Tensor& k, // [b, s, h_k, d] + const at::Tensor& cos, // [3, b, s, d / 2] + const at::Tensor& sin, // [3, b, s, d / 2] + const int first, + const int second + ) { + int Nthreads = 256; + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + ASSERT_CHECK(q.scalar_type() == k.scalar_type()); + ASSERT_CHECK(q.size(3) % 8 == 0); + ASSERT_CHECK(cos.size(0) == 3); + ASSERT_CHECK(cos.size(3) == q.size(3) / 2); + ASSERT_CHECK(cos.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.size(0) == 3); + ASSERT_CHECK(q.stride(3) == 1); + ASSERT_CHECK(k.stride(3) == 1); + ASSERT_CHECK(cos.stride(3) == 1); + ASSERT_CHECK(sin.stride(3) == 1); + ASSERT_CHECK(cos.is_contiguous()); + ASSERT_CHECK(sin.is_contiguous()); + + int64_t batch = q.size(0); + int64_t q_head_num = q.size(2); + int64_t k_head_num = k.size(2); + int64_t seq_len = q.size(1); + int64_t dim = q.size(3); + int64_t qb_stride = q.stride(0); + int64_t kb_stride = k.stride(0); + int64_t qs_stride = q.stride(1); + int64_t ks_stride = k.stride(1); + int64_t qh_stride = q.stride(2); + int64_t kh_stride = k.stride(2); + + dim3 grid(batch, seq_len); + + if (q.scalar_type() == at::kHalf) { + half* q_data = static_cast(q.data_ptr()); + half* k_data = static_cast(k.data_ptr()); + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + MRopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + first, second, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); + } else if (q.scalar_type() == at::kBFloat16) { + __nv_bfloat16* q_data = static_cast<__nv_bfloat16*>(q.data_ptr()); + __nv_bfloat16* k_data = static_cast<__nv_bfloat16*>(k.data_ptr()); + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + MRopeInplaceKernel<__nv_bfloat16><<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + first, second, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); + } else if (q.scalar_type() == at::kFloat) { + float* q_data = static_cast(q.data_ptr()); + float* k_data = static_cast(k.data_ptr()); + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + MRopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + first, second, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } else { + throw std::runtime_error("Unsupported data type for m_rope"); + } + + sync_check_cuda_error(); +} + +void MRopeBackward( + const at::Tensor& grad_q_embed, + const at::Tensor& grad_k_embed, + const at::Tensor& grad_q, // output: [b, s, q_h, d] + const at::Tensor& grad_k, // output: [b, s, k_h, d] + const at::Tensor& cos, // [3, b, s, d/2] + const at::Tensor& sin, // [3, b, s, d/2] + const int first, + const int second +) { + int Nthreads = 256; + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + ASSERT_CHECK(grad_q_embed.scalar_type() == grad_k_embed.scalar_type()); + ASSERT_CHECK(grad_q.scalar_type() == grad_k.scalar_type()); + ASSERT_CHECK(grad_q_embed.size(3) == grad_q.size(3)); + ASSERT_CHECK(grad_k_embed.size(3) == grad_k.size(3)); + ASSERT_CHECK(grad_q_embed.size(3) % 8 == 0); + ASSERT_CHECK(grad_q_embed.stride(3) == 1); + ASSERT_CHECK(grad_k_embed.stride(3) == 1); + ASSERT_CHECK(grad_q.stride(0) == grad_q_embed.stride(0)); + ASSERT_CHECK(grad_q.stride(1) == grad_q_embed.stride(1)); + ASSERT_CHECK(grad_q.stride(2) == grad_q_embed.stride(2)); + ASSERT_CHECK(grad_q.stride(3) == grad_q_embed.stride(3)); + ASSERT_CHECK(grad_k.stride(0) == grad_k_embed.stride(0)); + ASSERT_CHECK(grad_k.stride(1) == grad_k_embed.stride(1)); + ASSERT_CHECK(grad_k.stride(2) == grad_k_embed.stride(2)); + ASSERT_CHECK(grad_k.stride(3) == grad_k_embed.stride(3)); + + int64_t batch = grad_q_embed.size(0); + int64_t seq_len = grad_q_embed.size(1); + int64_t q_head_num = grad_q.size(2); + int64_t k_head_num = grad_k.size(2); + int64_t dim = grad_q_embed.size(3); + int64_t qb_stride = grad_q_embed.stride(0); + int64_t kb_stride = grad_k_embed.stride(0); + int64_t qs_stride = grad_q_embed.stride(1); + int64_t ks_stride = grad_k_embed.stride(1); + int64_t qh_stride = grad_q_embed.stride(2); + int64_t kh_stride = grad_k_embed.stride(2); + + dim3 grid(batch, seq_len); + + if (grad_q_embed.scalar_type() == at::kHalf) { + const half* gq_embed = static_cast(grad_q_embed.data_ptr()); + const half* gk_embed = static_cast(grad_k_embed.data_ptr()); + half* gq = static_cast(grad_q.data_ptr()); + half* gk = static_cast(grad_k.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + MRopeKernelBackward<<>>( + cos_data, sin_data, + gq_embed, gk_embed, + gq, gk, + first, second, + q_head_num, k_head_num, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } else if (grad_q_embed.scalar_type() == at::kBFloat16) { + const __nv_bfloat16* gq_embed = static_cast(grad_q_embed.data_ptr()); + const __nv_bfloat16* gk_embed = static_cast(grad_k_embed.data_ptr()); + __nv_bfloat16* gq = static_cast<__nv_bfloat16*>(grad_q.data_ptr()); + __nv_bfloat16* gk = static_cast<__nv_bfloat16*>(grad_k.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + MRopeKernelBackward<__nv_bfloat16><<>>( + cos_data, sin_data, + gq_embed, gk_embed, + gq, gk, + first, second, + q_head_num, k_head_num, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } else if (grad_q_embed.scalar_type() == at::kFloat) { + const float* gq_embed = static_cast(grad_q_embed.data_ptr()); + const float* gk_embed = static_cast(grad_k_embed.data_ptr()); + float* gq = static_cast(grad_q.data_ptr()); + float* gk = static_cast(grad_k.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + MRopeKernelBackward<<>>( + cos_data, sin_data, + gq_embed, gk_embed, + gq, gk, + first, second, + q_head_num, k_head_num, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } else { + throw std::runtime_error("Unsupported data type for m_rope backward"); + } + + sync_check_cuda_error(); +} + +// Pack interface: accepts qkv [bz, seq_len, q_dim + 2*kv_dim], applies mrope +// inplace on q and k without Python-side split/view overhead. +// Supports GQA (q_num_heads != kv_num_heads). CUDA kernels are not modified. +void MRopeInplacePack(const at::Tensor& qkv, // [bz, seq_len, q_dim + 2*kv_dim] + const at::Tensor& cos, // [3, bz, seq_len, head_dim/2] + const at::Tensor& sin, // [3, bz, seq_len, head_dim/2] + int64_t q_num_heads, + int64_t kv_num_heads, + const int first, + const int second) { + int Nthreads = 256; + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + ASSERT_CHECK(qkv.dim() == 3); + ASSERT_CHECK(qkv.stride(2) == 1); + ASSERT_CHECK(cos.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.scalar_type() == at::kFloat); + ASSERT_CHECK(cos.size(0) == 3); + ASSERT_CHECK(sin.size(0) == 3); + ASSERT_CHECK(cos.is_contiguous()); + ASSERT_CHECK(sin.is_contiguous()); + + int64_t bz = qkv.size(0); + int64_t seq_len = qkv.size(1); + int64_t head_dim = cos.size(3) * 2; // cos: [3, bz, seq, half_dim] + int64_t q_dim = q_num_heads * head_dim; + int64_t kv_dim = kv_num_heads * head_dim; + ASSERT_CHECK(qkv.size(2) == q_dim + 2 * kv_dim); + ASSERT_CHECK(head_dim % 8 == 0); + + // q at offset 0, k at offset q_dim; logical [bz, seq, *_num_heads, head_dim] + // with strides [qkv.stride(0), qkv.stride(1), head_dim, 1] + int64_t qkv_b_stride = qkv.stride(0); + int64_t qkv_s_stride = qkv.stride(1); + + dim3 grid(bz, seq_len); + + if (qkv.scalar_type() == at::kHalf) { + half* q_data = static_cast(qkv.data_ptr()); + half* k_data = q_data + q_dim; + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + MRopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_num_heads, k_data, kv_num_heads, + first, second, head_dim, + qkv_b_stride, qkv_s_stride, head_dim, + qkv_b_stride, qkv_s_stride, head_dim); + } else if (qkv.scalar_type() == at::kBFloat16) { + __nv_bfloat16* q_data = static_cast<__nv_bfloat16*>(qkv.data_ptr()); + __nv_bfloat16* k_data = q_data + q_dim; + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + MRopeInplaceKernel<__nv_bfloat16><<>>( + cos_data, sin_data, + q_data, q_num_heads, k_data, kv_num_heads, + first, second, head_dim, + qkv_b_stride, qkv_s_stride, head_dim, + qkv_b_stride, qkv_s_stride, head_dim); + } else if (qkv.scalar_type() == at::kFloat) { + float* q_data = static_cast(qkv.data_ptr()); + float* k_data = q_data + q_dim; + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + MRopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_num_heads, k_data, kv_num_heads, + first, second, head_dim, + qkv_b_stride, qkv_s_stride, head_dim, + qkv_b_stride, qkv_s_stride, head_dim); + } else { + throw std::runtime_error("Unsupported data type for m_rope_inplace_pack"); + } + + sync_check_cuda_error(); +} + +// Pack backward interface: accepts dqkv [bz, seq_len, q_dim + 2*kv_dim] and +// applies the inverse mrope transform inplace on the dq and dk slices. The dv +// slice is left unchanged. +void MRopeInplacePackBackward(const at::Tensor& dqkv, // [bz, seq_len, q_dim + 2*kv_dim] + const at::Tensor& cos, // [3, bz, seq_len, head_dim/2] + const at::Tensor& sin, // [3, bz, seq_len, head_dim/2] + int64_t q_num_heads, + int64_t kv_num_heads, + const int first, + const int second) { + int Nthreads = 256; + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + ASSERT_CHECK(dqkv.dim() == 3); + ASSERT_CHECK(dqkv.stride(2) == 1); + ASSERT_CHECK(cos.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.scalar_type() == at::kFloat); + ASSERT_CHECK(cos.size(0) == 3); + ASSERT_CHECK(sin.size(0) == 3); + ASSERT_CHECK(cos.is_contiguous()); + ASSERT_CHECK(sin.is_contiguous()); + + int64_t bz = dqkv.size(0); + int64_t seq_len = dqkv.size(1); + int64_t head_dim = cos.size(3) * 2; // cos: [3, bz, seq, half_dim] + int64_t q_dim = q_num_heads * head_dim; + int64_t kv_dim = kv_num_heads * head_dim; + ASSERT_CHECK(dqkv.size(2) == q_dim + 2 * kv_dim); + ASSERT_CHECK(head_dim % 8 == 0); + + int64_t dqkv_b_stride = dqkv.stride(0); + int64_t dqkv_s_stride = dqkv.stride(1); + + dim3 grid(bz, seq_len); + + if (dqkv.scalar_type() == at::kHalf) { + half* dq_data = static_cast(dqkv.data_ptr()); + half* dk_data = dq_data + q_dim; + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + MRopeKernelBackward<<>>( + cos_data, sin_data, + dq_data, dk_data, + dq_data, dk_data, + first, second, + q_num_heads, kv_num_heads, head_dim, + dqkv_b_stride, dqkv_s_stride, head_dim, + dqkv_b_stride, dqkv_s_stride, head_dim); + } else if (dqkv.scalar_type() == at::kBFloat16) { + __nv_bfloat16* dq_data = static_cast<__nv_bfloat16*>(dqkv.data_ptr()); + __nv_bfloat16* dk_data = dq_data + q_dim; + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + MRopeKernelBackward<__nv_bfloat16><<>>( + cos_data, sin_data, + dq_data, dk_data, + dq_data, dk_data, + first, second, + q_num_heads, kv_num_heads, head_dim, + dqkv_b_stride, dqkv_s_stride, head_dim, + dqkv_b_stride, dqkv_s_stride, head_dim); + } else if (dqkv.scalar_type() == at::kFloat) { + float* dq_data = static_cast(dqkv.data_ptr()); + float* dk_data = dq_data + q_dim; + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + MRopeKernelBackward<<>>( + cos_data, sin_data, + dq_data, dk_data, + dq_data, dk_data, + first, second, + q_num_heads, kv_num_heads, head_dim, + dqkv_b_stride, dqkv_s_stride, head_dim, + dqkv_b_stride, dqkv_s_stride, head_dim); + } else { + throw std::runtime_error("Unsupported data type for m_rope_inplace_pack_bwd"); + } + + sync_check_cuda_error(); +} + +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) + +} // wallx_cuda_m_rope diff --git a/wall_x/model/core/ops/csrc/permute_unpermute/permute.cu b/wall_x/model/core/ops/csrc/permute_unpermute/permute.cu new file mode 100644 index 0000000..498ebe9 --- /dev/null +++ b/wall_x/model/core/ops/csrc/permute_unpermute/permute.cu @@ -0,0 +1,581 @@ +// Re-enable CUDA half operators (PyTorch disables them) +#undef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ + +#include +#include + +#include +#include +#include +#include "../common/cuda_utils.h" +#include "cuda_runtime.h" +#include "device_launch_parameters.h" + +// Global FP8 guard: use CUDA_VERSION (defined in both host and device passes), +// not __CUDA_ARCH__ (device-pass only - leaves host dispatchers unable to see +// FP8 typedefs/branches). FP8 storage types arrived in CUDA 11.8. +#if defined(CUDA_VERSION) && (CUDA_VERSION >= 11800) +#include +#define XCOMPUTE_HAS_FP8 1 +#else +#define XCOMPUTE_HAS_FP8 0 +#endif + +namespace wallx_cuda_permute { + +// --------------------------------------------------------------------------- +// CUTLASS-free helpers: replace cutlass::Array, NumericArrayConverter, +// arch::global_load, and type aliases with standard CUDA equivalents. +// --------------------------------------------------------------------------- + +// Vectorized array: replaces cutlass::Array +// __align__(16) ensures safe reinterpret_cast in vec_load/store. +template +struct __align__(16) VecArray { + T data[N]; + __device__ __forceinline__ T& at(int i) { return data[i]; } + __device__ __forceinline__ const T& at(int i) const { return data[i]; } + __device__ __forceinline__ T* raw() { return data; } + __device__ __forceinline__ const T* raw() const { return data; } + __device__ __forceinline__ void clear() { + #pragma unroll + for (int i = 0; i < N; i++) data[i] = T(0); + } +}; + +// Scalar multiply: VecArray * scalar +template +__device__ __forceinline__ VecArray operator*(const VecArray& a, T s) { + VecArray r; + #pragma unroll + for (int i = 0; i < N; i++) r.data[i] = a.data[i] * s; + return r; +} + +// Element-wise convert: replaces cutlass::NumericArrayConverter +template +struct ArrayConverter { + __device__ __forceinline__ VecArray operator()(const VecArray& src) const { + VecArray dst; + #pragma unroll + for (int i = 0; i < N; i++) dst.data[i] = To(src.data[i]); + return dst; + } +}; + +// Vectorized global load: replaces cutlass::arch::global_load +template +__device__ __forceinline__ void vec_load(Fragment& frag, const void* ptr) { + static_assert(sizeof(Fragment) == sizeof(float4) || sizeof(Fragment) == 16, + "vec_load expects 16-byte fragment"); + *reinterpret_cast(&frag) = *reinterpret_cast(ptr); +} + +// Type aliases: replace cutlass types with standard CUDA types +using half_t = __half; +using bfloat16_t = __nv_bfloat16; +#if XCOMPUTE_HAS_FP8 +using float_e5m2_t = __nv_fp8_e5m2; +using float_e4m3_t = __nv_fp8_e4m3; // fn variant; CUDA's __nv_fp8_e4m3 IS the "fn" form +#endif + +// --------------------------------------------------------------------------- +// Kernels (unchanged logic, only CUTLASS types/calls replaced) +// --------------------------------------------------------------------------- + +static __global__ void moe_permute_topK_row_map( + const int *sorted_row_id, + int *row_id_map, + const int num_rows, + const int num_topK, + const int num_out_tokens) +{ + const int bid = blockIdx.x; + const int tid = threadIdx.x; + const int idx = bid * blockDim.x + tid; + + if (idx >= num_rows * num_topK) + return; + + int source_row = sorted_row_id[idx]; + int source_token_id = source_row / num_topK; + int source_topK_id = source_row % num_topK; + + if (idx >= num_out_tokens) + { + row_id_map[source_topK_id * num_rows + source_token_id] = -1; + } + else + { + row_id_map[source_topK_id * num_rows + source_token_id] = idx; + } +} + +template +__global__ void moe_recover_topK_kernel(const T *input, + T *unpermuted_output, + const int *row_id_map, + const float *prob, + const int num_rows, + const int num_topK, + const int num_cols) +{ + extern __shared__ int8_t s_mem[]; + TCompute *s_prob = reinterpret_cast(s_mem); + + using FragLS = VecArray; + using FragC = VecArray; + + ArrayConverter src_converter; + ArrayConverter dst_converter; + + const int source_token = blockIdx.x; + const int tid = threadIdx.x; + + if (hasProb) + { + for (int i = tid; i < num_topK; i += blockDim.x * blockDim.y) + { + s_prob[i] = TCompute(prob[source_token * num_topK + i]); + } + __syncthreads(); + } + + for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess) + { + FragLS frag_load_store; + FragC frag_elem; + FragC frag_sum; + + int source_row = row_id_map[source_token]; + + if (source_row != -1) + { + const T *source_row_ptr = input + source_row * num_cols; + + vec_load(frag_load_store, source_row_ptr + i); + frag_sum = src_converter(frag_load_store); + + if (hasProb) + { + frag_sum = frag_sum * s_prob[0]; + } + } + else + { + frag_sum.clear(); + } + + for (int k = 1; k < num_topK; k++) + { + source_row = row_id_map[k * num_rows + source_token]; + + if (source_row == -1) + continue; + + const T *source_row_ptr = input + source_row * num_cols; + + vec_load(frag_load_store, source_row_ptr + i); + frag_elem = src_converter(frag_load_store); + + if (hasProb) + { + frag_elem = frag_elem * s_prob[k]; + } + + for (int e = 0; e < kElementsPerAccess; e++) + { + frag_sum.at(e) = frag_sum.at(e) + frag_elem.at(e); + } + } + + T *dest_row_ptr = unpermuted_output + source_token * num_cols; + frag_load_store = dst_converter(frag_sum); + *(float4 *)(dest_row_ptr + i) = *(float4 *)(frag_load_store.raw()); + } +} + +template +__global__ void moe_permute_topK_kernel(const T *input_bwd, + const T *input_fwd, + T *act_grad, + const float *prob, + float *prob_grad, + const int *row_id_map, + const int num_rows, + const int num_topK, + const int num_cols) +{ + extern __shared__ int8_t s_mem[]; + TCompute *s_prob = reinterpret_cast(s_mem); + + using FragLS = VecArray; + using FragC = VecArray; + + ArrayConverter src_converter; + ArrayConverter dst_converter; + + const int source_token = blockIdx.x; + const int tid = threadIdx.x; + + if (hasProb) + { + for (int i = tid; i < num_topK; i += blockDim.x) + { + s_prob[i] = TCompute(prob[source_token * num_topK + i]); + } + __syncthreads(); + } + + float accum[topKTile] = {0.0f}; + FragLS frag_load_store; + + const T *source_row_ptr = input_bwd + source_token * num_cols; + for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess) + { + vec_load(frag_load_store, source_row_ptr + i); + FragC frag_src = src_converter(frag_load_store); + + int index = source_token; + + for (int k = 0; k < topKTile; k++) + { + if (k == num_topK) + break; + + int dest_row = row_id_map[index]; + index += num_rows; + + if (dest_row == -1) + continue; + + if (hasProb) + { + frag_load_store = dst_converter(frag_src * s_prob[k]); + } + else + { + frag_load_store = dst_converter(frag_src); + } + + T *dest_row_ptr = act_grad + dest_row * num_cols; + *(float4 *)(dest_row_ptr + i) = *(float4 *)(frag_load_store.raw()); + + if (hasProb) + { + const T *input_fwd_ptr = input_fwd + dest_row * num_cols; + vec_load(frag_load_store, input_fwd_ptr + i); + FragC frag_input_fwd = src_converter(frag_load_store); + + for (int e = 0; e < kElementsPerAccess; e++) + { + accum[k] += float(frag_src.at(e) * frag_input_fwd.at(e)); + } + } + } + } + + if (hasProb) + { + for (int k = 0; k < topKTile; k++) + { + if (k == num_topK) + break; + + for (int mask = 16; mask > 0; mask /= 2) + { + accum[k] = accum[k] + __shfl_xor_sync(0xffffffff, accum[k], mask, 32); + } + } + + if (tid == 0) + { + for (int k = 0; k < topKTile; k++) + { + if (k == num_topK) + break; + prob_grad[source_token * num_topK + k] = accum[k]; + } + } + } +} + +// --------------------------------------------------------------------------- +// Launcher (unchanged) +// --------------------------------------------------------------------------- + +template +void moe_permute_topK_kernel_launcher( + const T *input, + T *output, + const int *sorted_row_id, + int *row_id_map, + const float *prob, + const int num_rows, + const int num_topK, + const int num_cols, + const int num_out_tokens, + cudaStream_t stream, + float *prob_grad = nullptr, + const T *input_fwd = nullptr) +{ + if (FWD) + { + if (prob_grad == nullptr) + { + int threads = 64; + int blocks = (num_rows * num_topK + threads - 1) / threads; + moe_permute_topK_row_map<<>>( + sorted_row_id, row_id_map, num_rows, num_topK, num_out_tokens); + + blocks = num_rows; + threads = std::min(num_cols / kElementsPerAccess, 1024); + moe_permute_topK_kernel<<>>( + input, nullptr, output, nullptr, nullptr, row_id_map, + num_rows, num_topK, num_cols); + } + else + { + int blocks = num_rows; + int threads = 32; + size_t smem_bytes = num_topK * sizeof(TCompute); + + if (num_topK == 1) + moe_permute_topK_kernel<<>>( + input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols); + else if (num_topK <= 8) + moe_permute_topK_kernel<<>>( + input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols); + else if (num_topK <= 16) + moe_permute_topK_kernel<<>>( + input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols); + else if (num_topK <= 32) + moe_permute_topK_kernel<<>>( + input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols); + else if (num_topK <= 64) + moe_permute_topK_kernel<<>>( + input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols); + else if (num_topK <= 128) + moe_permute_topK_kernel<<>>( + input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols); + else + throw std::runtime_error("num_topK cannot exceed 128."); + } + } + else + { + int blocks = num_rows; + int threads = std::min(num_cols / kElementsPerAccess, 1024); + size_t smem_bytes = num_topK * sizeof(TCompute); + + if (num_topK == 1) + moe_recover_topK_kernel<<>>( + input, output, row_id_map, prob, num_rows, num_topK, num_cols); + else if (prob == nullptr) + moe_recover_topK_kernel<<>>( + input, output, row_id_map, prob, num_rows, num_topK, num_cols); + else + moe_recover_topK_kernel<<>>( + input, output, row_id_map, prob, num_rows, num_topK, num_cols); + } +} + +// --------------------------------------------------------------------------- +// Host-side ops +// --------------------------------------------------------------------------- + +void MoePermuteTopKOp( + const at::Tensor& input, + const at::Tensor& indices, + const at::Tensor& sorted_indices, + const at::Tensor& row_id, + const at::Tensor& sorted_row_id, + const at::Tensor& temp_storage, + const at::Tensor& permuted_output, + const at::Tensor& row_id_map, + int64_t num_out_tokens, + int64_t max_expanded_token_num +) { + ASSERT_CHECK(input.size(0) == indices.size(0)); + const int num_tokens = input.size(0); + const int num_cols = input.size(1); + const int num_topK = indices.size(1); + + int *indices_ptr = static_cast(indices.data_ptr()); + int *sorted_indices_ptr = static_cast(sorted_indices.data_ptr()); + int *row_id_ptr = static_cast(row_id.data_ptr()); + int *sorted_row_id_ptr = static_cast(sorted_row_id.data_ptr()); + + void *d_temp_storage = static_cast(temp_storage.data_ptr()); + size_t temp_storage_bytes = std::numeric_limits::max(); + + cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, + indices_ptr, sorted_indices_ptr, + row_id_ptr, sorted_row_id_ptr, num_tokens * num_topK); + + num_out_tokens = (num_out_tokens > 0) ? num_out_tokens : num_tokens * num_topK; + int *row_id_map_ptr = static_cast(row_id_map.data_ptr()); + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + if (input.scalar_type() == at::kFloat) { + float *input_ptr = static_cast(input.data_ptr()); + float *out_ptr = static_cast(permuted_output.data_ptr()); + moe_permute_topK_kernel_launcher( + input_ptr, out_ptr, sorted_row_id_ptr, row_id_map_ptr, + nullptr, num_tokens, num_topK, num_cols, num_out_tokens, stream); + } else if (input.scalar_type() == at::kHalf) { + half_t *input_ptr = static_cast(input.data_ptr()); + half_t *out_ptr = static_cast(permuted_output.data_ptr()); + moe_permute_topK_kernel_launcher( + input_ptr, out_ptr, sorted_row_id_ptr, row_id_map_ptr, + nullptr, num_tokens, num_topK, num_cols, num_out_tokens, stream); + } else if (input.scalar_type() == at::kBFloat16) { + bfloat16_t *input_ptr = static_cast(input.data_ptr()); + bfloat16_t *out_ptr = static_cast(permuted_output.data_ptr()); + moe_permute_topK_kernel_launcher( + input_ptr, out_ptr, sorted_row_id_ptr, row_id_map_ptr, + nullptr, num_tokens, num_topK, num_cols, num_out_tokens, stream); +#if XCOMPUTE_HAS_FP8 + } else if (input.scalar_type() == at::kFloat8_e5m2) { + float_e5m2_t *input_ptr = static_cast(input.data_ptr()); + float_e5m2_t *out_ptr = static_cast(permuted_output.data_ptr()); + moe_permute_topK_kernel_launcher( + input_ptr, out_ptr, sorted_row_id_ptr, row_id_map_ptr, + nullptr, num_tokens, num_topK, num_cols, num_out_tokens, stream); + } else if (input.scalar_type() == at::kFloat8_e4m3fn) { + float_e4m3_t *input_ptr = static_cast(input.data_ptr()); + float_e4m3_t *out_ptr = static_cast(permuted_output.data_ptr()); + moe_permute_topK_kernel_launcher( + input_ptr, out_ptr, sorted_row_id_ptr, row_id_map_ptr, + nullptr, num_tokens, num_topK, num_cols, num_out_tokens, stream); +#endif + } else { + throw std::runtime_error("Unsupported data type for MoePermuteTopKOp"); + } +} + +void MoeRecoverTopKOp( + const at::Tensor& input, + const at::Tensor& row_id_map, + const c10::optional& prob, + const at::Tensor& unpermuted_output, + int64_t num_tokens, + int64_t num_topK +) { + const int num_cols = input.size(1); + int *row_id_map_ptr = static_cast(row_id_map.data_ptr()); + float *prob_ptr = (prob.has_value()) ? static_cast(prob.value().data_ptr()) : nullptr; + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + if (input.scalar_type() == at::kFloat) { + float *in_ptr = static_cast(input.data_ptr()); + float *out_ptr = static_cast(unpermuted_output.data_ptr()); + moe_permute_topK_kernel_launcher( + in_ptr, out_ptr, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream); + } else if (input.scalar_type() == at::kHalf) { + half_t *in_ptr = static_cast(input.data_ptr()); + half_t *out_ptr = static_cast(unpermuted_output.data_ptr()); + moe_permute_topK_kernel_launcher( + in_ptr, out_ptr, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream); + } else if (input.scalar_type() == at::kBFloat16) { + bfloat16_t *in_ptr = static_cast(input.data_ptr()); + bfloat16_t *out_ptr = static_cast(unpermuted_output.data_ptr()); + moe_permute_topK_kernel_launcher( + in_ptr, out_ptr, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream); +#if XCOMPUTE_HAS_FP8 + } else if (input.scalar_type() == at::kFloat8_e5m2) { + float_e5m2_t *in_ptr = static_cast(input.data_ptr()); + float_e5m2_t *out_ptr = static_cast(unpermuted_output.data_ptr()); + moe_permute_topK_kernel_launcher( + in_ptr, out_ptr, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream); + } else if (input.scalar_type() == at::kFloat8_e4m3fn) { + float_e4m3_t *in_ptr = static_cast(input.data_ptr()); + float_e4m3_t *out_ptr = static_cast(unpermuted_output.data_ptr()); + moe_permute_topK_kernel_launcher( + in_ptr, out_ptr, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream); +#endif + } else { + throw std::runtime_error("Unsupported data type for MoeRecoverTopKOp"); + } + sync_check_cuda_error(); +} + +void MoeRecoverTopKBwdOp( + const at::Tensor& input_bwd, + const at::Tensor& input_fwd, + const at::Tensor& row_id_map, + const at::Tensor& prob, + const at::Tensor& act_grad, + const at::Tensor& prob_grad +) { + const int num_tokens = prob.size(0); + const int num_topK = prob.size(1); + const int num_cols = input_bwd.size(1); + int *row_id_map_ptr = static_cast(row_id_map.data_ptr()); + float *prob_ptr = static_cast(prob.data_ptr()); + float *prob_grad_ptr = static_cast(prob_grad.data_ptr()); + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + if (input_bwd.scalar_type() == at::kFloat) { + float *bwd = static_cast(input_bwd.data_ptr()); + float *fwd = static_cast(input_fwd.data_ptr()); + float *grad = static_cast(act_grad.data_ptr()); + moe_permute_topK_kernel_launcher( + bwd, grad, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream, prob_grad_ptr, fwd); + } else if (input_bwd.scalar_type() == at::kHalf) { + half_t *bwd = static_cast(input_bwd.data_ptr()); + half_t *fwd = static_cast(input_fwd.data_ptr()); + half_t *grad = static_cast(act_grad.data_ptr()); + moe_permute_topK_kernel_launcher( + bwd, grad, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream, prob_grad_ptr, fwd); + } else if (input_bwd.scalar_type() == at::kBFloat16) { + bfloat16_t *bwd = static_cast(input_bwd.data_ptr()); + bfloat16_t *fwd = static_cast(input_fwd.data_ptr()); + bfloat16_t *grad = static_cast(act_grad.data_ptr()); + moe_permute_topK_kernel_launcher( + bwd, grad, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream, prob_grad_ptr, fwd); +#if XCOMPUTE_HAS_FP8 + } else if (input_bwd.scalar_type() == at::kFloat8_e5m2) { + float_e5m2_t *bwd = static_cast(input_bwd.data_ptr()); + float_e5m2_t *fwd = static_cast(input_fwd.data_ptr()); + float_e5m2_t *grad = static_cast(act_grad.data_ptr()); + moe_permute_topK_kernel_launcher( + bwd, grad, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream, prob_grad_ptr, fwd); + } else if (input_bwd.scalar_type() == at::kFloat8_e4m3fn) { + float_e4m3_t *bwd = static_cast(input_bwd.data_ptr()); + float_e4m3_t *fwd = static_cast(input_fwd.data_ptr()); + float_e4m3_t *grad = static_cast(act_grad.data_ptr()); + moe_permute_topK_kernel_launcher( + bwd, grad, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream, prob_grad_ptr, fwd); +#endif + } else { + throw std::runtime_error("Unsupported data type for MoeRecoverTopKBwdOp"); + } + sync_check_cuda_error(); +} + +size_t CubSortPairGetStorageBytes(int64_t num_items){ + size_t temp_storage_bytes = 0; + int *temp_ptr = nullptr; + cub::DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, + temp_ptr, temp_ptr, + temp_ptr, temp_ptr, num_items); + return temp_storage_bytes; +} + +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) + +} // wallx_cuda_permute diff --git a/wall_x/model/core/ops/csrc/rope/rope_kernel.cu b/wall_x/model/core/ops/csrc/rope/rope_kernel.cu new file mode 100644 index 0000000..f9ef0a6 --- /dev/null +++ b/wall_x/model/core/ops/csrc/rope/rope_kernel.cu @@ -0,0 +1,1098 @@ +// Re-enable CUDA half operators (PyTorch disables them) +#undef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ + +#include +#include + +#include + #include +#include +#include +#include +#include "../common/cuda_utils.h" + +namespace wallx_cuda_rope { + +// Helper: convert vector of T to float4 +template +__device__ __forceinline__ float4 to_float4(const T* ptr) { + if constexpr (std::is_same_v) { + half2 h2_0 = reinterpret_cast(ptr)[0]; + half2 h2_1 = reinterpret_cast(ptr)[1]; + float2 f2_0 = __half22float2(h2_0); + float2 f2_1 = __half22float2(h2_1); + return make_float4(f2_0.x, f2_0.y, f2_1.x, f2_1.y); + } else if constexpr (std::is_same_v) { + __nv_bfloat162 b2_0 = reinterpret_cast(ptr)[0]; + __nv_bfloat162 b2_1 = reinterpret_cast(ptr)[1]; +#if __CUDA_ARCH__ >= 800 + float2 f2_0 = __bfloat1622float2(b2_0); + float2 f2_1 = __bfloat1622float2(b2_1); + return make_float4(f2_0.x, f2_0.y, f2_1.x, f2_1.y); +#else + assert(false && "Unsupported on arch < 800"); +#endif + } else if constexpr (std::is_same_v) { + return reinterpret_cast(ptr)[0]; + } else { + static_assert(sizeof(T) == 0, "Unsupported type"); + } +} + +// Helper: convert float4 back to vector of T +template +__device__ __forceinline__ void from_float4(T* ptr, const float4& f4) { + if constexpr (std::is_same_v) { + half2 h2_0 = __float22half2_rn(make_float2(f4.x, f4.y)); + half2 h2_1 = __float22half2_rn(make_float2(f4.z, f4.w)); + reinterpret_cast(ptr)[0] = h2_0; + reinterpret_cast(ptr)[1] = h2_1; + } else if constexpr (std::is_same_v) { +#if __CUDA_ARCH__ >= 800 + __nv_bfloat162 b2_0 = __float22bfloat162_rn(make_float2(f4.x, f4.y)); + __nv_bfloat162 b2_1 = __float22bfloat162_rn(make_float2(f4.z, f4.w)); + reinterpret_cast<__nv_bfloat162*>(ptr)[0] = b2_0; + reinterpret_cast<__nv_bfloat162*>(ptr)[1] = b2_1; +#else + assert(false && "Unsupported on arch < 800"); +#endif + } else if constexpr (std::is_same_v) { + reinterpret_cast(ptr)[0] = f4; + } else { + static_assert(sizeof(T) == 0, "Unsupported type"); + } +} + +template +__global__ void RopeKernel(const float* cos, + const float* sin, + const T* q, + const int q_h, + const T* k, + const int k_h, + T* q_embed, + T* k_embed, + const int d, + const int qb_stride, + const int qs_stride, + const int qh_stride, + const int kb_stride, + const int ks_stride, + const int kh_stride, + const int qeb_stride, + const int qes_stride, + const int qeh_stride, + const int keb_stride, + const int kes_stride, + const int keh_stride) { + extern __shared__ char cos_sin[]; + const int half_dim = d / 2; + float* cos_smem = reinterpret_cast(cos_sin); + float* sin_smem = cos_smem + half_dim; + int b = blockIdx.x; + int s = blockIdx.y; + + int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t cos_sin_s_stride = half_dim; +#define SIN_GMEM(b, c, d) sin[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] +#define COS_GMEM(b, c, d) cos[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] + + for (int i = threadIdx.x; i < half_dim; i += blockDim.x) { + cos_smem[i] = COS_GMEM(b, s, i); + sin_smem[i] = SIN_GMEM(b, s, i); + } + __syncthreads(); + +#define Q_GMEM(a, b, c, d) q[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + d] +#define K_GMEM(a, b, c, d) k[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + d] +#define Q_EMBED_GMEM(a, b, c, d) q_embed[(a) * qeb_stride + (b) * qes_stride + (c) * qeh_stride + d] +#define K_EMBED_GMEM(a, b, c, d) k_embed[(a) * keb_stride + (b) * kes_stride + (c) * keh_stride + d] + + for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + + if constexpr (interleave) { + float4 q_x0 = to_float4(&Q_GMEM(b, s, h_idx, 2 * base_d)); + float4 q_x1 = to_float4(&Q_GMEM(b, s, h_idx, 2 * base_d + 4)); + float4 q0_rot = make_float4( + q_x0.x * cos_vec.x - q_x0.y * sin_vec.x, + q_x0.y * cos_vec.x + q_x0.x * sin_vec.x, + q_x0.z * cos_vec.y - q_x0.w * sin_vec.y, + q_x0.w * cos_vec.y + q_x0.z * sin_vec.y + ); + float4 q1_rot = make_float4( + q_x1.x * cos_vec.z - q_x1.y * sin_vec.z, + q_x1.y * cos_vec.z + q_x1.x * sin_vec.z, + q_x1.z * cos_vec.w - q_x1.w * sin_vec.w, + q_x1.w * cos_vec.w + q_x1.z * sin_vec.w + ); + + from_float4(&Q_EMBED_GMEM(b, s, h_idx, 2 * base_d), q0_rot); + from_float4(&Q_EMBED_GMEM(b, s, h_idx, 2 * base_d + 4), q1_rot); + } else { + float4 q_x0 = to_float4(&Q_GMEM(b, s, h_idx, base_d)); + float4 q_x1 = to_float4(&Q_GMEM(b, s, h_idx, base_d + half_dim)); + float4 q0_rot = make_float4( + q_x0.x * cos_vec.x - q_x1.x * sin_vec.x, + q_x0.y * cos_vec.y - q_x1.y * sin_vec.y, + q_x0.z * cos_vec.z - q_x1.z * sin_vec.z, + q_x0.w * cos_vec.w - q_x1.w * sin_vec.w + ); + + float4 q1_rot = make_float4( + q_x1.x * cos_vec.x + q_x0.x * sin_vec.x, + q_x1.y * cos_vec.y + q_x0.y * sin_vec.y, + q_x1.z * cos_vec.z + q_x0.z * sin_vec.z, + q_x1.w * cos_vec.w + q_x0.w * sin_vec.w + ); + + from_float4(&Q_EMBED_GMEM(b, s, h_idx, base_d), q0_rot); + from_float4(&Q_EMBED_GMEM(b, s, h_idx, base_d + half_dim), q1_rot); + } + } + + for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + if constexpr (interleave) { + float4 k_x0 = to_float4(&K_GMEM(b, s, h_idx, 2 * base_d)); + float4 k_x1 = to_float4(&K_GMEM(b, s, h_idx, 2 * base_d + 4)); + + float4 k0_rot = make_float4( + k_x0.x * cos_vec.x - k_x0.y * sin_vec.x, + k_x0.y * cos_vec.x + k_x0.x * sin_vec.x, + k_x0.z * cos_vec.y - k_x0.w * sin_vec.y, + k_x0.w * cos_vec.y + k_x0.z * sin_vec.y + ); + + float4 k1_rot = make_float4( + k_x1.x * cos_vec.z - k_x1.y * sin_vec.z, + k_x1.y * cos_vec.z + k_x1.x * sin_vec.z, + k_x1.z * cos_vec.w - k_x1.w * sin_vec.w, + k_x1.w * cos_vec.w + k_x1.z * sin_vec.w + ); + + from_float4(&K_EMBED_GMEM(b, s, h_idx, 2 * base_d), k0_rot); + from_float4(&K_EMBED_GMEM(b, s, h_idx, 2 * base_d + 4), k1_rot); + } else { + float4 k_x0 = to_float4(&K_GMEM(b, s, h_idx, base_d)); + float4 k_x1 = to_float4(&K_GMEM(b, s, h_idx, base_d + half_dim)); + + float4 k0_rot = make_float4( + k_x0.x * cos_vec.x - k_x1.x * sin_vec.x, + k_x0.y * cos_vec.y - k_x1.y * sin_vec.y, + k_x0.z * cos_vec.z - k_x1.z * sin_vec.z, + k_x0.w * cos_vec.w - k_x1.w * sin_vec.w + ); + + float4 k1_rot = make_float4( + k_x1.x * cos_vec.x + k_x0.x * sin_vec.x, + k_x1.y * cos_vec.y + k_x0.y * sin_vec.y, + k_x1.z * cos_vec.z + k_x0.z * sin_vec.z, + k_x1.w * cos_vec.w + k_x0.w * sin_vec.w + ); + + from_float4(&K_EMBED_GMEM(b, s, h_idx, base_d), k0_rot); + from_float4(&K_EMBED_GMEM(b, s, h_idx, base_d + half_dim), k1_rot); + } + } +} + +template +__global__ void RopeInplaceKernel(float* cos, float* sin, + T* q, const int q_h, T* k, const int k_h, + const int d, + const int qb_stride, const int qs_stride, const int qh_stride, + const int kb_stride, const int ks_stride, const int kh_stride) { + extern __shared__ char cos_sin[]; + const int half_dim = d / 2; + float* cos_smem = reinterpret_cast(cos_sin); + float* sin_smem = cos_smem + half_dim; + int b = blockIdx.x; + int s = blockIdx.y; + + int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t cos_sin_s_stride = half_dim; +#define SIN_GMEM(b, c, d) sin[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] +#define COS_GMEM(b, c, d) cos[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] + + for (int i = threadIdx.x; i < half_dim; i += blockDim.x) { + cos_smem[i] = COS_GMEM(b, s, i); + sin_smem[i] = SIN_GMEM(b, s, i); + } + __syncthreads(); + +#define Q_GMEM(a, b, c, d) q[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + d] +#define K_GMEM(a, b, c, d) k[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + d] + + for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + + if constexpr (interleave) { + float4 q_x0 = to_float4(&Q_GMEM(b, s, h_idx, 2 * base_d)); + float4 q_x1 = to_float4(&Q_GMEM(b, s, h_idx, 2 * base_d + 4)); + float4 q0_rot = make_float4( + q_x0.x * cos_vec.x - q_x0.y * sin_vec.x, + q_x0.y * cos_vec.x + q_x0.x * sin_vec.x, + q_x0.z * cos_vec.y - q_x0.w * sin_vec.y, + q_x0.w * cos_vec.y + q_x0.z * sin_vec.y + ); + float4 q1_rot = make_float4( + q_x1.x * cos_vec.z - q_x1.y * sin_vec.z, + q_x1.y * cos_vec.z + q_x1.x * sin_vec.z, + q_x1.z * cos_vec.w - q_x1.w * sin_vec.w, + q_x1.w * cos_vec.w + q_x1.z * sin_vec.w + ); + + from_float4(&Q_GMEM(b, s, h_idx, 2 * base_d), q0_rot); + from_float4(&Q_GMEM(b, s, h_idx, 2 * base_d + 4), q1_rot); + } else { + float4 q_x0 = to_float4(&Q_GMEM(b, s, h_idx, base_d)); + float4 q_x1 = to_float4(&Q_GMEM(b, s, h_idx, base_d + half_dim)); + float4 q0_rot = make_float4( + q_x0.x * cos_vec.x - q_x1.x * sin_vec.x, + q_x0.y * cos_vec.y - q_x1.y * sin_vec.y, + q_x0.z * cos_vec.z - q_x1.z * sin_vec.z, + q_x0.w * cos_vec.w - q_x1.w * sin_vec.w + ); + + float4 q1_rot = make_float4( + q_x1.x * cos_vec.x + q_x0.x * sin_vec.x, + q_x1.y * cos_vec.y + q_x0.y * sin_vec.y, + q_x1.z * cos_vec.z + q_x0.z * sin_vec.z, + q_x1.w * cos_vec.w + q_x0.w * sin_vec.w + ); + + from_float4(&Q_GMEM(b, s, h_idx, base_d), q0_rot); + from_float4(&Q_GMEM(b, s, h_idx, base_d + half_dim), q1_rot); + } + } + + for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + if constexpr (interleave) { + float4 k_x0 = to_float4(&K_GMEM(b, s, h_idx, 2 * base_d)); + float4 k_x1 = to_float4(&K_GMEM(b, s, h_idx, 2 * base_d + 4)); + + float4 k0_rot = make_float4( + k_x0.x * cos_vec.x - k_x0.y * sin_vec.x, + k_x0.y * cos_vec.x + k_x0.x * sin_vec.x, + k_x0.z * cos_vec.y - k_x0.w * sin_vec.y, + k_x0.w * cos_vec.y + k_x0.z * sin_vec.y + ); + + float4 k1_rot = make_float4( + k_x1.x * cos_vec.z - k_x1.y * sin_vec.z, + k_x1.y * cos_vec.z + k_x1.x * sin_vec.z, + k_x1.z * cos_vec.w - k_x1.w * sin_vec.w, + k_x1.w * cos_vec.w + k_x1.z * sin_vec.w + ); + + from_float4(&K_GMEM(b, s, h_idx, 2 * base_d), k0_rot); + from_float4(&K_GMEM(b, s, h_idx, 2 * base_d + 4), k1_rot); + } else { + float4 k_x0 = to_float4(&K_GMEM(b, s, h_idx, base_d)); + float4 k_x1 = to_float4(&K_GMEM(b, s, h_idx, base_d + half_dim)); + + float4 k0_rot = make_float4( + k_x0.x * cos_vec.x - k_x1.x * sin_vec.x, + k_x0.y * cos_vec.y - k_x1.y * sin_vec.y, + k_x0.z * cos_vec.z - k_x1.z * sin_vec.z, + k_x0.w * cos_vec.w - k_x1.w * sin_vec.w + ); + + float4 k1_rot = make_float4( + k_x1.x * cos_vec.x + k_x0.x * sin_vec.x, + k_x1.y * cos_vec.y + k_x0.y * sin_vec.y, + k_x1.z * cos_vec.z + k_x0.z * sin_vec.z, + k_x1.w * cos_vec.w + k_x0.w * sin_vec.w + ); + + from_float4(&K_GMEM(b, s, h_idx, base_d), k0_rot); + from_float4(&K_GMEM(b, s, h_idx, base_d + half_dim), k1_rot); + } + } +} + +template +__global__ void RopeKernelBackward( + const float* cos, + const float* sin, + const T* grad_q_embed, + const T* grad_k_embed, + T* grad_q, + T* grad_k, + const int q_h, + const int k_h, + const int d, + const int qb_stride, const int qs_stride, const int qh_stride, + const int kb_stride, const int ks_stride, const int kh_stride +) { + extern __shared__ char cos_sin[]; + const int half_dim = d / 2; + float* cos_smem = reinterpret_cast(cos_sin); + float* sin_smem = cos_smem + half_dim; + + int b = blockIdx.x; + int s = blockIdx.y; + + int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t cos_sin_s_stride = half_dim; +#define SIN_GMEM(b, c, d) sin[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] +#define COS_GMEM(b, c, d) cos[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] + + for (int i = threadIdx.x; i < half_dim; i += blockDim.x) { + cos_smem[i] = COS_GMEM(b, s, i); + sin_smem[i] = SIN_GMEM(b, s, i); + } + __syncthreads(); + +#define GQ_EMBED(a, b, c, d) grad_q_embed[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + (d)] +#define GK_EMBED(a, b, c, d) grad_k_embed[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + (d)] +#define GQ(a, b, c, d) grad_q[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + (d)] +#define GK(a, b, c, d) grad_k[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + (d)] + + // Process grad_q + for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + if constexpr (interleave) { + float4 gq0_embed = to_float4(&GQ_EMBED(b, s, h_idx, 2 * base_d)); + float4 gq1_embed = to_float4(&GQ_EMBED(b, s, h_idx, 2 * base_d + 4)); + + float4 gq0 = make_float4( + gq0_embed.x * cos_vec.x + gq0_embed.y * sin_vec.x, + gq0_embed.y * cos_vec.x - gq0_embed.x * sin_vec.x, + gq0_embed.z * cos_vec.y + gq0_embed.w * sin_vec.y, + gq0_embed.w * cos_vec.y - gq0_embed.z * sin_vec.y + ); + + float4 gq1 = make_float4( + gq1_embed.x * cos_vec.z + gq1_embed.y * sin_vec.z, + gq1_embed.y * cos_vec.z - gq1_embed.x * sin_vec.z, + gq1_embed.z * cos_vec.w + gq1_embed.w * sin_vec.w, + gq1_embed.w * cos_vec.w - gq1_embed.z * sin_vec.w + ); + + from_float4(&GQ(b, s, h_idx, 2 * base_d), gq0); + from_float4(&GQ(b, s, h_idx, 2 * base_d + 4), gq1); + } else { + float4 gq0_rot = to_float4(&GQ_EMBED(b, s, h_idx, base_d)); + float4 gq1_rot = to_float4(&GQ_EMBED(b, s, h_idx, base_d + half_dim)); + + float4 gq0 = make_float4( + gq0_rot.x * cos_vec.x + gq1_rot.x * sin_vec.x, + gq0_rot.y * cos_vec.y + gq1_rot.y * sin_vec.y, + gq0_rot.z * cos_vec.z + gq1_rot.z * sin_vec.z, + gq0_rot.w * cos_vec.w + gq1_rot.w * sin_vec.w + ); + float4 gq1 = make_float4( + gq1_rot.x * cos_vec.x - gq0_rot.x * sin_vec.x, + gq1_rot.y * cos_vec.y - gq0_rot.y * sin_vec.y, + gq1_rot.z * cos_vec.z - gq0_rot.z * sin_vec.z, + gq1_rot.w * cos_vec.w - gq0_rot.w * sin_vec.w + ); + + from_float4(&GQ(b, s, h_idx, base_d), gq0); + from_float4(&GQ(b, s, h_idx, base_d + half_dim), gq1); + } + } + + // Process grad_k + for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) { + int h_idx = j / half_dim; + int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + float4 cos_vec = to_float4(&cos_smem[base_d]); + float4 sin_vec = to_float4(&sin_smem[base_d]); + if constexpr (interleave) { + float4 gk0_embed = to_float4(&GK_EMBED(b, s, h_idx, 2 * base_d)); + float4 gk1_embed = to_float4(&GK_EMBED(b, s, h_idx, 2 * base_d + 4)); + + float4 gk0 = make_float4( + gk0_embed.x * cos_vec.x + gk0_embed.y * sin_vec.x, + gk0_embed.y * cos_vec.x - gk0_embed.x * sin_vec.x, + gk0_embed.z * cos_vec.y + gk0_embed.w * sin_vec.y, + gk0_embed.w * cos_vec.y - gk0_embed.z * sin_vec.y + ); + + float4 gk1 = make_float4( + gk1_embed.x * cos_vec.z + gk1_embed.y * sin_vec.z, + gk1_embed.y * cos_vec.z - gk1_embed.x * sin_vec.z, + gk1_embed.z * cos_vec.w + gk1_embed.w * sin_vec.w, + gk1_embed.w * cos_vec.w - gk1_embed.z * sin_vec.w + ); + + from_float4(&GK(b, s, h_idx, 2 * base_d), gk0); + from_float4(&GK(b, s, h_idx, 2 * base_d + 4), gk1); + } else { + float4 gk0_rot = to_float4(&GK_EMBED(b, s, h_idx, base_d)); + float4 gk1_rot = to_float4(&GK_EMBED(b, s, h_idx, base_d + half_dim)); + + float4 gk0 = make_float4( + gk0_rot.x * cos_vec.x + gk1_rot.x * sin_vec.x, + gk0_rot.y * cos_vec.y + gk1_rot.y * sin_vec.y, + gk0_rot.z * cos_vec.z + gk1_rot.z * sin_vec.z, + gk0_rot.w * cos_vec.w + gk1_rot.w * sin_vec.w + ); + float4 gk1 = make_float4( + gk1_rot.x * cos_vec.x - gk0_rot.x * sin_vec.x, + gk1_rot.y * cos_vec.y - gk0_rot.y * sin_vec.y, + gk1_rot.z * cos_vec.z - gk0_rot.z * sin_vec.z, + gk1_rot.w * cos_vec.w - gk0_rot.w * sin_vec.w + ); + + from_float4(&GK(b, s, h_idx, base_d), gk0); + from_float4(&GK(b, s, h_idx, base_d + half_dim), gk1); + } + } +} + +void Rope(const at::Tensor& q, // [b, s, h, d] or [s, h, d] + const at::Tensor& k, // [b, s, h_k, d] or [s, h_k, d] + const at::Tensor& q_embed, // [b, s, h, d] or [s, h, d] + const at::Tensor& k_embed, // [b, s, h_k, d] or [s, h_k, d] + const at::Tensor& cos, // [b, s, d / 2] or [s, d / 2] + const at::Tensor& sin, // [b, s, d / 2] or [s, d / 2] + bool interleave + ) { + int Nthreads = 256; + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + bool q_3d = (q.dim() == 3); + + ASSERT_CHECK(q.scalar_type() == k.scalar_type()); + ASSERT_CHECK(q.scalar_type() == q_embed.scalar_type()); + ASSERT_CHECK(k.scalar_type() == k_embed.scalar_type()); + ASSERT_CHECK(cos.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.scalar_type() == at::kFloat); + ASSERT_CHECK(cos.is_contiguous()); + ASSERT_CHECK(sin.is_contiguous()); + + int64_t batch, q_head_num, k_head_num, seq_len, dim; + int64_t qb_stride, kb_stride, qs_stride, ks_stride, qh_stride, kh_stride; + int64_t qeb_stride, keb_stride, qes_stride, kes_stride, qeh_stride, keh_stride; + + if (q_3d) { + ASSERT_CHECK(q_embed.size(0) == q.size(0)); + ASSERT_CHECK(q_embed.size(1) == q.size(1)); + ASSERT_CHECK(q_embed.size(2) == q.size(2)); + ASSERT_CHECK(k_embed.size(0) == k.size(0)); + ASSERT_CHECK(k_embed.size(1) == k.size(1)); + ASSERT_CHECK(k_embed.size(2) == k.size(2)); + ASSERT_CHECK(q.size(2) % 8 == 0); + ASSERT_CHECK(cos.size(-1) == q.size(-1) / 2); + ASSERT_CHECK(q.stride(2) == 1); + ASSERT_CHECK(k.stride(2) == 1); + ASSERT_CHECK(q_embed.stride(2) == 1); + ASSERT_CHECK(k_embed.stride(2) == 1); + batch = 1; + seq_len = q.size(0); + q_head_num = q.size(1); + k_head_num = k.size(1); + dim = q.size(2); + qb_stride = 0; qs_stride = q.stride(0); qh_stride = q.stride(1); + kb_stride = 0; ks_stride = k.stride(0); kh_stride = k.stride(1); + qeb_stride = 0; qes_stride = q_embed.stride(0); qeh_stride = q_embed.stride(1); + keb_stride = 0; kes_stride = k_embed.stride(0); keh_stride = k_embed.stride(1); + } else { + ASSERT_CHECK(q_embed.size(0) == q.size(0)); + ASSERT_CHECK(q_embed.size(1) == q.size(1)); + ASSERT_CHECK(q_embed.size(2) == q.size(2)); + ASSERT_CHECK(q_embed.size(3) == q.size(3)); + ASSERT_CHECK(k_embed.size(0) == k.size(0)); + ASSERT_CHECK(k_embed.size(1) == k.size(1)); + ASSERT_CHECK(k_embed.size(2) == k.size(2)); + ASSERT_CHECK(k_embed.size(3) == k.size(3)); + ASSERT_CHECK(q.size(3) % 8 == 0); + ASSERT_CHECK(cos.size(2) == q.size(3) / 2); + ASSERT_CHECK(q.stride(3) == 1); + ASSERT_CHECK(k.stride(3) == 1); + ASSERT_CHECK(q_embed.stride(3) == 1); + ASSERT_CHECK(k_embed.stride(3) == 1); + batch = q.size(0); + seq_len = q.size(1); + q_head_num = q.size(2); + k_head_num = k.size(2); + dim = q.size(3); + qb_stride = q.stride(0); qs_stride = q.stride(1); qh_stride = q.stride(2); + kb_stride = k.stride(0); ks_stride = k.stride(1); kh_stride = k.stride(2); + qeb_stride = q_embed.stride(0); qes_stride = q_embed.stride(1); qeh_stride = q_embed.stride(2); + keb_stride = k_embed.stride(0); kes_stride = k_embed.stride(1); keh_stride = k_embed.stride(2); + } + + dim3 grid(batch, seq_len); + + if (q.scalar_type() == at::kHalf) { + const half* q_data = static_cast(q.data_ptr()); + const half* k_data = static_cast(k.data_ptr()); + half* q_embed_data = static_cast(q_embed.data_ptr()); + half* k_embed_data = static_cast(k_embed.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + q_embed_data, k_embed_data, + dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); + } else { + RopeKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + q_embed_data, k_embed_data, + dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); + } + } else if (q.scalar_type() == at::kBFloat16) { + const __nv_bfloat16* q_data = static_cast(q.data_ptr()); + const __nv_bfloat16* k_data = static_cast(k.data_ptr()); + __nv_bfloat16* q_embed_data = static_cast<__nv_bfloat16*>(q_embed.data_ptr()); + __nv_bfloat16* k_embed_data = static_cast<__nv_bfloat16*>(k_embed.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeKernel<__nv_bfloat16, true><<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + q_embed_data, k_embed_data, + dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); + } else { + RopeKernel<__nv_bfloat16, false><<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + q_embed_data, k_embed_data, + dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); + } + } else if (q.scalar_type() == at::kFloat) { + const float* q_data = static_cast(q.data_ptr()); + const float* k_data = static_cast(k.data_ptr()); + float* q_embed_data = static_cast(q_embed.data_ptr()); + float* k_embed_data = static_cast(k_embed.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + q_embed_data, k_embed_data, + dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride + ); + } else { + RopeKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + q_embed_data, k_embed_data, + dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride + ); + } + } else { + throw std::runtime_error("Unsupported data type for rope"); + } + + sync_check_cuda_error(); +} + +void RopeInplace(const at::Tensor& q, // [b, s, h, d] or [s, h, d] + const at::Tensor& k, // [b, s, h_k, d] or [s, h_k, d] + const at::Tensor& cos, // [b, s, d / 2] or [s, d / 2] + const at::Tensor& sin, // [b, s, d / 2] or [s, d / 2] + bool interleave + ) { + int Nthreads = 256; + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + bool q_3d = (q.dim() == 3); + + ASSERT_CHECK(q.scalar_type() == k.scalar_type()); + ASSERT_CHECK(cos.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.scalar_type() == at::kFloat); + ASSERT_CHECK(cos.is_contiguous()); + ASSERT_CHECK(sin.is_contiguous()); + + int64_t batch, q_head_num, k_head_num, seq_len, dim; + int64_t qb_stride, kb_stride, qs_stride, ks_stride, qh_stride, kh_stride; + + if (q_3d) { + ASSERT_CHECK(q.size(2) % 8 == 0); + ASSERT_CHECK(cos.size(2) == q.size(2) / 2); + ASSERT_CHECK(q.stride(2) == 1); + ASSERT_CHECK(k.stride(2) == 1); + batch = 1; + seq_len = q.size(0); + q_head_num = q.size(1); + k_head_num = k.size(1); + dim = q.size(2); + qb_stride = 0; qs_stride = q.stride(0); qh_stride = q.stride(1); + kb_stride = 0; ks_stride = k.stride(0); kh_stride = k.stride(1); + } else { + ASSERT_CHECK(q.size(3) % 8 == 0); + ASSERT_CHECK(cos.size(2) == q.size(3) / 2); + ASSERT_CHECK(q.stride(3) == 1); + ASSERT_CHECK(k.stride(3) == 1); + batch = q.size(0); + seq_len = q.size(1); + q_head_num = q.size(2); + k_head_num = k.size(2); + dim = q.size(3); + qb_stride = q.stride(0); qs_stride = q.stride(1); qh_stride = q.stride(2); + kb_stride = k.stride(0); ks_stride = k.stride(1); kh_stride = k.stride(2); + } + + dim3 grid(batch, seq_len); + + if (q.scalar_type() == at::kHalf) { + half* q_data = static_cast(q.data_ptr()); + half* k_data = static_cast(k.data_ptr()); + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); + } else { + RopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); + } + } else if (q.scalar_type() == at::kBFloat16) { + __nv_bfloat16* q_data = static_cast<__nv_bfloat16*>(q.data_ptr()); + __nv_bfloat16* k_data = static_cast<__nv_bfloat16*>(k.data_ptr()); + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeInplaceKernel<__nv_bfloat16, true><<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); + } else { + RopeInplaceKernel<__nv_bfloat16, false><<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); + } + } else if (q.scalar_type() == at::kFloat) { + float* q_data = static_cast(q.data_ptr()); + float* k_data = static_cast(k.data_ptr()); + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } else { + RopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_head_num, + k_data, k_head_num, + dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } + } else { + throw std::runtime_error("Unsupported data type for rope"); + } + + sync_check_cuda_error(); +} + +void RopeBackward( + const at::Tensor& grad_q_embed, + const at::Tensor& grad_k_embed, + const at::Tensor& grad_q, // output: [b, s, q_h, d] or [s, q_h, d] + const at::Tensor& grad_k, // output: [b, s, k_h, d] or [s, k_h, d] + const at::Tensor& cos, // [b, s, d/2] or [s, d/2] + const at::Tensor& sin, // [b, s, d/2] or [s, d/2] + bool interleave +) { + int Nthreads = 256; + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + bool q_3d = (grad_q_embed.dim() == 3); + + ASSERT_CHECK(grad_q_embed.scalar_type() == grad_k_embed.scalar_type()); + + int64_t batch, seq_len, q_head_num, k_head_num, dim; + int64_t qb_stride, kb_stride, qs_stride, ks_stride, qh_stride, kh_stride; + + if (q_3d) { + ASSERT_CHECK(grad_q_embed.size(2) == grad_q.size(2)); + ASSERT_CHECK(grad_k_embed.size(2) == grad_k.size(2)); + ASSERT_CHECK(grad_q_embed.size(2) % 8 == 0); + ASSERT_CHECK(grad_q_embed.stride(2) == 1); + ASSERT_CHECK(grad_k_embed.stride(2) == 1); + ASSERT_CHECK(grad_q.stride(0) == grad_q_embed.stride(0)); + ASSERT_CHECK(grad_q.stride(1) == grad_q_embed.stride(1)); + ASSERT_CHECK(grad_q.stride(2) == grad_q_embed.stride(2)); + ASSERT_CHECK(grad_k.stride(0) == grad_k_embed.stride(0)); + ASSERT_CHECK(grad_k.stride(1) == grad_k_embed.stride(1)); + ASSERT_CHECK(grad_k.stride(2) == grad_k_embed.stride(2)); + batch = 1; + seq_len = grad_q_embed.size(0); + q_head_num = grad_q.size(1); + k_head_num = grad_k.size(1); + dim = grad_q_embed.size(2); + qb_stride = 0; qs_stride = grad_q_embed.stride(0); qh_stride = grad_q_embed.stride(1); + kb_stride = 0; ks_stride = grad_k_embed.stride(0); kh_stride = grad_k_embed.stride(1); + } else { + ASSERT_CHECK(grad_q_embed.size(3) == grad_q.size(3)); + ASSERT_CHECK(grad_k_embed.size(3) == grad_k.size(3)); + ASSERT_CHECK(grad_q_embed.size(3) % 8 == 0); + ASSERT_CHECK(grad_q_embed.stride(3) == 1); + ASSERT_CHECK(grad_k_embed.stride(3) == 1); + ASSERT_CHECK(grad_q.stride(0) == grad_q_embed.stride(0)); + ASSERT_CHECK(grad_q.stride(1) == grad_q_embed.stride(1)); + ASSERT_CHECK(grad_q.stride(2) == grad_q_embed.stride(2)); + ASSERT_CHECK(grad_q.stride(3) == grad_q_embed.stride(3)); + ASSERT_CHECK(grad_k.stride(0) == grad_k_embed.stride(0)); + ASSERT_CHECK(grad_k.stride(1) == grad_k_embed.stride(1)); + ASSERT_CHECK(grad_k.stride(2) == grad_k_embed.stride(2)); + ASSERT_CHECK(grad_k.stride(3) == grad_k_embed.stride(3)); + batch = grad_q_embed.size(0); + seq_len = grad_q_embed.size(1); + q_head_num = grad_q.size(2); + k_head_num = grad_k.size(2); + dim = grad_q_embed.size(3); + qb_stride = grad_q_embed.stride(0); qs_stride = grad_q_embed.stride(1); qh_stride = grad_q_embed.stride(2); + kb_stride = grad_k_embed.stride(0); ks_stride = grad_k_embed.stride(1); kh_stride = grad_k_embed.stride(2); + } + + dim3 grid(batch, seq_len); + + if (grad_q_embed.scalar_type() == at::kHalf) { + const half* gq_embed = static_cast(grad_q_embed.data_ptr()); + const half* gk_embed = static_cast(grad_k_embed.data_ptr()); + half* gq = static_cast(grad_q.data_ptr()); + half* gk = static_cast(grad_k.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeKernelBackward<<>>( + cos_data, sin_data, + gq_embed, gk_embed, + gq, gk, + q_head_num, k_head_num, dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } else { + RopeKernelBackward<<>>( + cos_data, sin_data, + gq_embed, gk_embed, + gq, gk, + q_head_num, k_head_num, dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } + } else if (grad_q_embed.scalar_type() == at::kBFloat16) { + const __nv_bfloat16* gq_embed = static_cast(grad_q_embed.data_ptr()); + const __nv_bfloat16* gk_embed = static_cast(grad_k_embed.data_ptr()); + __nv_bfloat16* gq = static_cast<__nv_bfloat16*>(grad_q.data_ptr()); + __nv_bfloat16* gk = static_cast<__nv_bfloat16*>(grad_k.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeKernelBackward<__nv_bfloat16, true><<>>( + cos_data, sin_data, + gq_embed, gk_embed, + gq, gk, + q_head_num, k_head_num, dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } else { + RopeKernelBackward<__nv_bfloat16, false><<>>( + cos_data, sin_data, + gq_embed, gk_embed, + gq, gk, + q_head_num, k_head_num, dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } + } else if (grad_q_embed.scalar_type() == at::kFloat) { + const float* gq_embed = static_cast(grad_q_embed.data_ptr()); + const float* gk_embed = static_cast(grad_k_embed.data_ptr()); + float* gq = static_cast(grad_q.data_ptr()); + float* gk = static_cast(grad_k.data_ptr()); + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeKernelBackward<<>>( + cos_data, sin_data, + gq_embed, gk_embed, + gq, gk, + q_head_num, k_head_num, dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } else { + RopeKernelBackward<<>>( + cos_data, sin_data, + gq_embed, gk_embed, + gq, gk, + q_head_num, k_head_num, dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + ); + } + } else { + throw std::runtime_error("Unsupported data type for rope backward"); + } + + sync_check_cuda_error(); +} + +// Pack interface: accepts qkv [seq_len, q_dim + 2*kv_dim], applies rope inplace +// on the q and k slices. Supports GQA (q_num_heads != kv_num_heads). +// CUDA kernels are not modified; only host-side pointer arithmetic is added. +void RopeInplacePack(const at::Tensor& qkv, // [seq_len, q_dim + 2*kv_dim] + const at::Tensor& cos, // [1, seq_len, head_dim/2] + const at::Tensor& sin, // [1, seq_len, head_dim/2] + int64_t q_num_heads, + int64_t kv_num_heads, + bool interleave) { + int Nthreads = 256; + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + ASSERT_CHECK(qkv.dim() == 2); + ASSERT_CHECK(qkv.stride(1) == 1); + ASSERT_CHECK(cos.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.scalar_type() == at::kFloat); + ASSERT_CHECK(cos.is_contiguous()); + ASSERT_CHECK(sin.is_contiguous()); + + int64_t seq_len = qkv.size(0); + int64_t head_dim = cos.size(-1) * 2; // cos: [1, seq, half_dim] + int64_t q_dim = q_num_heads * head_dim; + int64_t kv_dim = kv_num_heads * head_dim; + ASSERT_CHECK(qkv.size(1) == q_dim + 2 * kv_dim); + ASSERT_CHECK(head_dim % 8 == 0); + + // q at offset 0, k at offset q_dim; logical [seq, *_num_heads, head_dim] + // with strides [qkv_row_stride, head_dim, 1] + int64_t qkv_row_stride = qkv.stride(0); + int64_t qs_stride = qkv_row_stride; + int64_t qh_stride = head_dim; + int64_t ks_stride = qkv_row_stride; + int64_t kh_stride = head_dim; + + dim3 grid(1, seq_len); + + if (qkv.scalar_type() == at::kHalf) { + half* q_data = static_cast(qkv.data_ptr()); + half* k_data = q_data + q_dim; + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } else { + RopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } + } else if (qkv.scalar_type() == at::kBFloat16) { + __nv_bfloat16* q_data = static_cast<__nv_bfloat16*>(qkv.data_ptr()); + __nv_bfloat16* k_data = q_data + q_dim; + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeInplaceKernel<__nv_bfloat16, true><<>>( + cos_data, sin_data, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } else { + RopeInplaceKernel<__nv_bfloat16, false><<>>( + cos_data, sin_data, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } + } else if (qkv.scalar_type() == at::kFloat) { + float* q_data = static_cast(qkv.data_ptr()); + float* k_data = q_data + q_dim; + float* cos_data = static_cast(cos.data_ptr()); + float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } else { + RopeInplaceKernel<<>>( + cos_data, sin_data, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } + } else { + throw std::runtime_error("Unsupported data type for rope_inplace_pack"); + } + + sync_check_cuda_error(); +} + +// Pack backward interface: accepts dqkv [seq_len, q_dim + 2*kv_dim] and applies +// the inverse rope transform inplace on the dq and dk slices. The dv slice is +// left unchanged. +void RopeInplacePackBackward(const at::Tensor& dqkv, // [seq_len, q_dim + 2*kv_dim] + const at::Tensor& cos, // [1, seq_len, head_dim/2] + const at::Tensor& sin, // [1, seq_len, head_dim/2] + int64_t q_num_heads, + int64_t kv_num_heads, + bool interleave) { + int Nthreads = 256; + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + ASSERT_CHECK(dqkv.dim() == 2); + ASSERT_CHECK(dqkv.stride(1) == 1); + ASSERT_CHECK(cos.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.scalar_type() == at::kFloat); + ASSERT_CHECK(cos.is_contiguous()); + ASSERT_CHECK(sin.is_contiguous()); + + int64_t seq_len = dqkv.size(0); + int64_t head_dim = cos.size(-1) * 2; // cos: [1, seq, half_dim] + int64_t q_dim = q_num_heads * head_dim; + int64_t kv_dim = kv_num_heads * head_dim; + ASSERT_CHECK(dqkv.size(1) == q_dim + 2 * kv_dim); + ASSERT_CHECK(head_dim % 8 == 0); + + int64_t dqkv_row_stride = dqkv.stride(0); + int64_t qs_stride = dqkv_row_stride; + int64_t qh_stride = head_dim; + int64_t ks_stride = dqkv_row_stride; + int64_t kh_stride = head_dim; + + dim3 grid(1, seq_len); + + if (dqkv.scalar_type() == at::kHalf) { + half* dq_data = static_cast(dqkv.data_ptr()); + half* dk_data = dq_data + q_dim; + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeKernelBackward<<>>( + cos_data, sin_data, + dq_data, dk_data, + dq_data, dk_data, + q_num_heads, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } else { + RopeKernelBackward<<>>( + cos_data, sin_data, + dq_data, dk_data, + dq_data, dk_data, + q_num_heads, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } + } else if (dqkv.scalar_type() == at::kBFloat16) { + __nv_bfloat16* dq_data = static_cast<__nv_bfloat16*>(dqkv.data_ptr()); + __nv_bfloat16* dk_data = dq_data + q_dim; + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeKernelBackward<__nv_bfloat16, true><<>>( + cos_data, sin_data, + dq_data, dk_data, + dq_data, dk_data, + q_num_heads, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } else { + RopeKernelBackward<__nv_bfloat16, false><<>>( + cos_data, sin_data, + dq_data, dk_data, + dq_data, dk_data, + q_num_heads, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } + } else if (dqkv.scalar_type() == at::kFloat) { + float* dq_data = static_cast(dqkv.data_ptr()); + float* dk_data = dq_data + q_dim; + const float* cos_data = static_cast(cos.data_ptr()); + const float* sin_data = static_cast(sin.data_ptr()); + if (interleave) { + RopeKernelBackward<<>>( + cos_data, sin_data, + dq_data, dk_data, + dq_data, dk_data, + q_num_heads, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } else { + RopeKernelBackward<<>>( + cos_data, sin_data, + dq_data, dk_data, + dq_data, dk_data, + q_num_heads, kv_num_heads, head_dim, + 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); + } + } else { + throw std::runtime_error("Unsupported data type for rope_inplace_pack_bwd"); + } + + sync_check_cuda_error(); +} + +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) + +} // wallx_cuda_rope diff --git a/wall_x/model/core/ops/csrc/rot_pos/rot_pos.cu b/wall_x/model/core/ops/csrc/rot_pos/rot_pos.cu new file mode 100644 index 0000000..2df3599 --- /dev/null +++ b/wall_x/model/core/ops/csrc/rot_pos/rot_pos.cu @@ -0,0 +1,284 @@ +#include +#include + +#include +#include +#include +#include "../common/cuda_utils.h" + +namespace wallx_cuda_rot_pos { + +__global__ void fused_rot_pos_emb_kernel_int32( + const float *__restrict__ inv_freq, // [dim/2] - precomputed inverse frequencies + const int32_t *__restrict__ grid_thw, // [num_grids, 3] - (t, h, w) for each grid + float *__restrict__ output, // [total_tokens, dim] - output rotary embeddings + const int32_t *__restrict__ cumsum_tokens, // [num_grids+1] - cumulative sum of tokens per grid + const int dim_half, // dim/2 (size of inv_freq) + const int spatial_merge_size, // spatial merge size + const int num_grids // number of grids +) +{ + const int32_t tid = blockIdx.x * blockDim.x + threadIdx.x; + const int32_t total_tokens = cumsum_tokens[num_grids]; + + extern __shared__ float freq_smem[]; + for (int i = threadIdx.x; i < dim_half; i += blockDim.x) { + freq_smem[i] = inv_freq[i]; + } + if (tid >= total_tokens * dim_half) + return; + + const int32_t token_idx = tid / dim_half; + const int freq_idx = tid % dim_half; + + // Find which grid this token belongs to + int grid_idx = 0; + int32_t local_token_idx = token_idx; + for (int g = 0; g < num_grids; g++) + { + if (token_idx < cumsum_tokens[g + 1]) + { + grid_idx = g; + local_token_idx = token_idx - cumsum_tokens[g]; + break; + } + } + + // Get grid dimensions + const int32_t h = grid_thw[grid_idx * 3 + 1]; + const int32_t w = grid_thw[grid_idx * 3 + 2]; + + // Calculate spatial dimensions after merging + const int32_t h_merged = h / spatial_merge_size; + const int32_t w_merged = w / spatial_merge_size; + const int32_t spatial_tokens = h_merged * w_merged * spatial_merge_size * spatial_merge_size; + + // Get spatial index + const int32_t spatial_idx = local_token_idx % spatial_tokens; + + // Decompose spatial index to get merged block and position within block + const int32_t tokens_per_block = spatial_merge_size * spatial_merge_size; + const int32_t block_idx = spatial_idx / tokens_per_block; + const int32_t within_block_idx = spatial_idx % tokens_per_block; + + // Get block coordinates in merged grid + const int32_t block_h = block_idx / w_merged; + const int32_t block_w = block_idx % w_merged; + + // Get position within block + const int32_t within_h = within_block_idx / spatial_merge_size; + const int32_t within_w = within_block_idx % spatial_merge_size; + + // Calculate actual h and w positions + const int32_t h_pos = block_h * spatial_merge_size + within_h; + const int32_t w_pos = block_w * spatial_merge_size + within_w; + + // Compute rotary embedding + float freq_val = freq_smem[freq_idx]; + + // Output has shape [total_tokens, dim] where dim = 2 * dim_half + int32_t out_idx = token_idx * dim_half * 2 + freq_idx; + output[out_idx] = h_pos * freq_val; // h_pos frequencies + output[out_idx + dim_half] = w_pos * freq_val; // w_pos frequencies +} + +__global__ void fused_rot_pos_emb_kernel_int64( + const float *__restrict__ inv_freq, // [dim/2] - precomputed inverse frequencies + const int64_t *__restrict__ grid_thw, // [num_grids, 3] - (t, h, w) for each grid + float *__restrict__ output, // [total_tokens, dim] - output rotary embeddings + const int64_t *__restrict__ cumsum_tokens, // [num_grids+1] - cumulative sum of tokens per grid + const int dim_half, // dim/2 (size of inv_freq) + const int spatial_merge_size, // spatial merge size + const int num_grids // number of grids +) +{ + const int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; + const int64_t total_tokens = cumsum_tokens[num_grids]; + + extern __shared__ float freq_smem[]; + for (int i = threadIdx.x; i < dim_half; i += blockDim.x) { + freq_smem[i] = inv_freq[i]; + } + + if (tid >= total_tokens * dim_half) + return; + + const int64_t token_idx = tid / dim_half; + const int freq_idx = tid % dim_half; + + // Find which grid this token belongs to + int grid_idx = 0; + int64_t local_token_idx = token_idx; + for (int g = 0; g < num_grids; g++) + { + if (token_idx < cumsum_tokens[g + 1]) + { + grid_idx = g; + local_token_idx = token_idx - cumsum_tokens[g]; + break; + } + } + + // Get grid dimensions + const int64_t h = grid_thw[grid_idx * 3 + 1]; + const int64_t w = grid_thw[grid_idx * 3 + 2]; + + // Calculate spatial dimensions after merging + const int64_t h_merged = h / spatial_merge_size; + const int64_t w_merged = w / spatial_merge_size; + const int64_t spatial_tokens = h_merged * w_merged * spatial_merge_size * spatial_merge_size; + + // Get spatial index + const int64_t spatial_idx = local_token_idx % spatial_tokens; + + // Decompose spatial index to get merged block and position within block + const int64_t tokens_per_block = spatial_merge_size * spatial_merge_size; + const int64_t block_idx = spatial_idx / tokens_per_block; + const int64_t within_block_idx = spatial_idx % tokens_per_block; + + // Get block coordinates in merged grid + const int64_t block_h = block_idx / w_merged; + const int64_t block_w = block_idx % w_merged; + + // Get position within block + const int64_t within_h = within_block_idx / spatial_merge_size; + const int64_t within_w = within_block_idx % spatial_merge_size; + + // Calculate actual h and w positions + const int64_t h_pos = block_h * spatial_merge_size + within_h; + const int64_t w_pos = block_w * spatial_merge_size + within_w; + + // Compute rotary embedding + float freq_val = freq_smem[freq_idx]; + + // Output has shape [total_tokens, dim] where dim = 2 * dim_half + int64_t out_idx = token_idx * dim_half * 2 + freq_idx; + output[out_idx] = h_pos * freq_val; // h_pos frequencies + output[out_idx + dim_half] = w_pos * freq_val; // w_pos frequencies +} + +__global__ void compute_token_counts_kernel_int32( + const int32_t *__restrict__ grid_thw, + int32_t *__restrict__ token_counts, + const int spatial_merge_size, + const int num_grids) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_grids) + return; + + int32_t t = grid_thw[idx * 3 + 0]; + int32_t h = grid_thw[idx * 3 + 1]; + int32_t w = grid_thw[idx * 3 + 2]; + int32_t h_merged = h / spatial_merge_size; + int32_t w_merged = w / spatial_merge_size; + token_counts[idx] = t * h_merged * w_merged * spatial_merge_size * spatial_merge_size; +} + +__global__ void compute_token_counts_kernel_int64( + const int64_t *__restrict__ grid_thw, + int64_t *__restrict__ token_counts, + const int spatial_merge_size, + const int num_grids) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_grids) + return; + + int64_t t = grid_thw[idx * 3 + 0]; + int64_t h = grid_thw[idx * 3 + 1]; + int64_t w = grid_thw[idx * 3 + 2]; + int64_t h_merged = h / spatial_merge_size; + int64_t w_merged = w / spatial_merge_size; + token_counts[idx] = t * h_merged * w_merged * spatial_merge_size * spatial_merge_size; +} + +void GetTokenCounts( + const at::Tensor& grid_thw, // [num_grids, 3] + const at::Tensor& token_counts, // [num_grids] + int spatial_merge_size +) { + ASSERT_CHECK(grid_thw.dim() == 2); + ASSERT_CHECK(grid_thw.size(1) == 3); + ASSERT_CHECK(grid_thw.size(0) == token_counts.size(0)); + + ASSERT_CHECK(spatial_merge_size > 0); + + const int num_grids = grid_thw.size(0); + + const int threads = 256; + const int blocks = (num_grids + threads - 1) / threads; + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + if (grid_thw.scalar_type() == at::kInt) { + compute_token_counts_kernel_int32<<>>( + static_cast(grid_thw.data_ptr()), + static_cast(token_counts.data_ptr()), + spatial_merge_size, + num_grids); + } else if (grid_thw.scalar_type() == at::kLong) { + compute_token_counts_kernel_int64<<>>( + static_cast(grid_thw.data_ptr()), + static_cast(token_counts.data_ptr()), + spatial_merge_size, + num_grids); + } else { + throw std::runtime_error("Unsupported data type for GetTokenCounts"); + } + + sync_check_cuda_error(); +} + +void RotPosEmb( + const at::Tensor& inv_freq, // [dim/2] + const at::Tensor& grid_thw, // [num_grids, 3] + const at::Tensor& output, + const at::Tensor& cumsum_tokens, + int spatial_merge_size +) { + ASSERT_CHECK(output.size(0) > 0); + ASSERT_CHECK(inv_freq.dim() == 1); + ASSERT_CHECK(inv_freq.scalar_type() == at::kFloat); + + const int dim_half = inv_freq.size(0); + const int num_grids = grid_thw.size(0); + const int total_tokens = output.size(0); + + const int64_t num_elements = total_tokens * dim_half; + + const int threads_per_block = 256; + const int num_blocks = static_cast((num_elements + threads_per_block - 1) / threads_per_block); + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + if (grid_thw.scalar_type() == at::kInt) { + fused_rot_pos_emb_kernel_int32<<>>( + static_cast(inv_freq.data_ptr()), + static_cast(grid_thw.data_ptr()), + static_cast(output.data_ptr()), + static_cast(cumsum_tokens.data_ptr()), + dim_half, + spatial_merge_size, + num_grids); + } else if (grid_thw.scalar_type() == at::kLong) { + fused_rot_pos_emb_kernel_int64<<>>( + static_cast(inv_freq.data_ptr()), + static_cast(grid_thw.data_ptr()), + static_cast(output.data_ptr()), + static_cast(cumsum_tokens.data_ptr()), + dim_half, + spatial_merge_size, + num_grids); + } else { + throw std::runtime_error("Unsupported data type for RotPosEmb"); + } + + sync_check_cuda_error(); +} + +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) + +} // wallx_cuda_rot_pos diff --git a/csrc/window_index.cu b/wall_x/model/core/ops/csrc/window_index/window_index.cu similarity index 77% rename from csrc/window_index.cu rename to wall_x/model/core/ops/csrc/window_index/window_index.cu index 4a1486b..48f0c75 100644 --- a/csrc/window_index.cu +++ b/wall_x/model/core/ops/csrc/window_index/window_index.cu @@ -1,8 +1,10 @@ -#include #include -#include -#include -#include + +#include +#include +#include "../common/cuda_utils.h" + +namespace wallx_cuda_window_index { __global__ void compute_metadata( const int *grid_thw, // [num_grids, 3] @@ -103,9 +105,10 @@ __global__ void compute_window_counts( __global__ void compute_cu_window_seqlens( const int *window_counts, - int *cu_window_seqlens, // [total_windows + 1] + int *cu_window_seqlens, // [total_windows + 1] int total_windows, - int spatial_merge_unit) + int spatial_merge_unit +) { int tid = blockIdx.x * blockDim.x + threadIdx.x; @@ -205,74 +208,64 @@ __global__ void generate_window_indices( } } -std::tuple get_window_index_cuda( - torch::Tensor grid_thw, +void GetTotals( + const at::Tensor& grid_thw, + const at::Tensor& grid_info_tensor, + const at::Tensor& global_totals_tensor, int spatial_merge_size, - int vit_merger_window_size, - int patch_size, - int spatial_merge_unit) -{ - TORCH_CHECK(grid_thw.is_cuda(), "grid_thw must be a CUDA tensor"); - TORCH_CHECK(grid_thw.dim() == 2 && grid_thw.size(1) == 3); - TORCH_CHECK(grid_thw.dtype() == torch::kInt32); + int vit_merger_window_size +) { + ASSERT_CHECK(grid_thw.dim() == 2 && grid_thw.size(1) == 3); + ASSERT_CHECK(grid_thw.scalar_type() == at::kInt); int num_grids = grid_thw.size(0); - if (num_grids == 0) - { - return std::make_tuple( - torch::empty({0}, grid_thw.options()), - torch::zeros({1}, grid_thw.options())); - } + const int *d_grid_thw = static_cast(grid_thw.data_ptr()); - const int *d_grid_thw = grid_thw.data_ptr(); - auto options = grid_thw.options(); - - auto grid_thw_cpu = grid_thw.cpu(); - int max_grid_t = 0; - for (int i = 0; i < num_grids; i++) - { - max_grid_t = std::max(max_grid_t, grid_thw_cpu[i][0].item()); - } - - auto grid_info_tensor = torch::empty({num_grids, 6}, options); - auto global_totals_tensor = torch::zeros({2}, options); - - int *d_grid_info = grid_info_tensor.data_ptr(); - int *d_global_totals = global_totals_tensor.data_ptr(); + int *d_grid_info = static_cast(grid_info_tensor.data_ptr()); + int *d_global_totals = static_cast(global_totals_tensor.data_ptr()); int threads1 = 256; int blocks1 = (num_grids + threads1 - 1) / threads1; + compute_metadata<<>>( d_grid_thw, d_grid_info, d_global_totals, num_grids, spatial_merge_size, vit_merger_window_size); - auto totals_cpu = global_totals_tensor.cpu(); - int total_elements = totals_cpu[0].item(); - int total_windows = totals_cpu[1].item(); + sync_check_cuda_error(); +} - if (total_elements == 0 || total_windows == 0) - { - return std::make_tuple( - torch::empty({0}, options), - torch::zeros({1}, options)); - } +void GetWindowIndex( + const at::Tensor& grid_thw, + const at::Tensor& grid_info_tensor, + const at::Tensor& window_indices, + const at::Tensor& cu_window_seqlens, + const at::Tensor& window_counts_tensor, + int max_grid_t, + int spatial_merge_size, + int vit_merger_window_size, + int patch_size, + int spatial_merge_unit) { - torch::Tensor window_indices = torch::empty({total_elements}, options); - torch::Tensor cu_window_seqlens = torch::empty({total_windows + 1}, options); - - int *d_window_indices = window_indices.data_ptr(); - int *d_cu_window_seqlens = cu_window_seqlens.data_ptr(); - - auto window_counts_tensor = torch::empty({total_windows}, options); - int *d_window_counts = window_counts_tensor.data_ptr(); + int num_grids = grid_thw.size(0); dim3 blocks2(max_grid_t, num_grids); dim3 threads2(256); + const int *d_grid_thw = static_cast(grid_thw.data_ptr()); + + int *d_grid_info = static_cast(grid_info_tensor.data_ptr()); + + int *d_window_indices = static_cast(window_indices.data_ptr()); + int *d_cu_window_seqlens = static_cast(cu_window_seqlens.data_ptr()); + + int *d_window_counts = static_cast(window_counts_tensor.data_ptr()); + compute_window_counts<<>>( d_grid_thw, d_grid_info, d_window_counts, vit_merger_window_size, spatial_merge_unit, num_grids); + int total_windows = window_counts_tensor.size(0); + int threads4 = 256; int blocks4 = (total_windows + threads4 - 1) / threads4; compute_cu_window_seqlens<<>>( @@ -282,5 +275,10 @@ std::tuple get_window_index_cuda( d_grid_thw, d_grid_info, d_cu_window_seqlens, d_window_indices, vit_merger_window_size, spatial_merge_unit, num_grids); - return std::make_tuple(window_indices, cu_window_seqlens); + sync_check_cuda_error(); } + +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) + +} // wallx_cuda_window_index diff --git a/wall_x/model/core/ops/index.py b/wall_x/model/core/ops/index.py new file mode 100644 index 0000000..1323d49 --- /dev/null +++ b/wall_x/model/core/ops/index.py @@ -0,0 +1,340 @@ +"""Index computation operators for attention.""" + +import logging +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F + +from wall_x.model.core.ops.base import OpsProxy + +logger = logging.getLogger(__name__) + +# Sentinel value for padding in window index computation. +# Valid indices are always >= 0 (they are positional offsets within flattened +# grid tensors), so any negative value is safe. -100 is chosen by convention +# (same as HuggingFace's ignore_index for cross-entropy loss). +_PAD_VALUE = -100 + + +def _compute_vision_position_ids( + input_tokens, + token_set, + st, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + image_index, + video_index, + remain_images, + remain_videos, + image_token_id, + video_token_id, + spatial_merge_size, + tokens_per_second, + device, + llm_pos_ids_list, +): + """Compute position IDs for a single vision token (image or video). + + Returns updated (st, image_index, video_index, remain_images, remain_videos). + """ + if image_token_id in token_set and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in token_set and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + + if ed_image < ed_video: + if image_index >= len(image_grid_thw): + raise IndexError( + f"image_index {image_index} out of range (have {len(image_grid_thw)} image grids)" + ) + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + # Images are single-frame: no temporal progression, so second_per_grid_t = 0. + # This makes all image tokens share t_index = 0 (only spatial positions vary). + second_per_grid_t = 0 + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + if video_index >= len(video_grid_thw): + raise IndexError( + f"video_index {video_index} out of range (have {len(video_grid_thw)} video grids)" + ) + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + if second_per_grid_ts is not None: + second_per_grid_t = second_per_grid_ts[video_index] + else: + second_per_grid_t = 1.0 + video_index += 1 + remain_videos -= 1 + ed = ed_video + + llm_grid_t, llm_grid_h, llm_grid_w = ( + int(t), + int(h) // spatial_merge_size, + int(w) // spatial_merge_size, + ) + text_len = ed - st + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append( + torch.arange(text_len, device=device).view(1, -1).expand(3, -1) + st_idx + ) + range_tensor = torch.arange(llm_grid_t, device=device).view(-1, 1) + expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) + time_tensor = expanded_range * second_per_grid_t * tokens_per_second + t_index = time_tensor.long().flatten() + h_index = ( + torch.arange(llm_grid_h, device=device) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w, device=device) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + return st, image_index, video_index, remain_images, remain_videos + + +def _compute_text_only_positions(input_ids, attention_mask): + """Compute position IDs for text-only inputs (no vision tokens).""" + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = ( + position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) + ) + max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[ + 0 + ] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + return position_ids, mrope_position_deltas + + +class GetRopeIndexOp(OpsProxy): + """Compute 3D RoPE position indices for multimodal inputs. + + Signature: get_rope_index(input_ids, image_grid_thw, video_grid_thw, + second_per_grid_ts, attention_mask, spatial_merge_size, + image_token_id, video_token_id, vision_start_token_id, + tokens_per_second) -> (position_ids, mrope_position_deltas) + """ + + @property + def _external_accel_name(self): + return "get_rope_index" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import GetRopeIndex + + return GetRopeIndex() + except ImportError: + return None + except Exception as e: + logger.warning("GetRopeIndexOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback( + self, + input_ids: torch.LongTensor, + image_grid_thw: Optional[torch.LongTensor], + video_grid_thw: Optional[torch.LongTensor], + second_per_grid_ts: Optional[torch.Tensor], + attention_mask: Optional[torch.Tensor], + spatial_merge_size: int, + image_token_id: int, + video_token_id: int, + vision_start_token_id: int, + tokens_per_second: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if input_ids is None or (image_grid_thw is None and video_grid_thw is None): + return _compute_text_only_positions(input_ids, attention_mask) + + mrope_position_deltas = [] + device = input_ids.device + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids_i in enumerate(total_input_ids): + input_ids_i = input_ids_i[attention_mask[i] == 1] + vision_start_indices = torch.argwhere( + input_ids_i == vision_start_token_id + ).squeeze(1) + # Boundary check: ensure vision_start + 1 doesn't exceed sequence length + valid_mask = (vision_start_indices + 1) < len(input_ids_i) + vision_start_indices = vision_start_indices[valid_mask] + vision_tokens = input_ids_i[vision_start_indices + 1] + image_nums = int((vision_tokens == image_token_id).sum()) + video_nums = int((vision_tokens == video_token_id).sum()) + input_tokens = input_ids_i.tolist() + token_set = set(input_tokens) + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + st, image_index, video_index, remain_images, remain_videos = ( + _compute_vision_position_ids( + input_tokens, + token_set, + st, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + image_index, + video_index, + remain_images, + remain_videos, + image_token_id, + video_token_id, + spatial_merge_size, + tokens_per_second, + device, + llm_pos_ids_list, + ) + ) + if st < len(input_tokens): + st_idx = ( + llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + ) + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len, device=device).view(1, -1).expand(3, -1) + + st_idx + ) + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=total_input_ids.device + ).unsqueeze(1) + return position_ids, mrope_position_deltas + + +class GetWindowIndexOp(OpsProxy): + """Compute window attention indices for ViT. + + Signature: get_window_index(grid_thw, window_size, spatial_merge_size, + patch_size, spatial_merge_unit=1) -> (window_index, cu_window_seqlens) + """ + + @property + def _external_accel_name(self): + return "get_window_index" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import get_window_index_cuda + + return get_window_index_cuda + except ImportError: + return None + except Exception as e: + logger.warning("GetWindowIndexOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback( + self, + grid_thw, + window_size, + spatial_merge_size, + patch_size, + spatial_merge_unit=1, + ): + device = grid_thw.device + vit_merger_window_size = window_size // spatial_merge_size // patch_size + window_index_list = [] + cu_window_seqlens = [0] + window_index_id = 0 + for grid_t, grid_h, grid_w in grid_thw: + llm_grid_h = grid_h // spatial_merge_size + llm_grid_w = grid_w // spatial_merge_size + index = torch.arange( + grid_t * llm_grid_h * llm_grid_w, device=device + ).reshape(grid_t, llm_grid_h, llm_grid_w) + pad_h = ( + vit_merger_window_size - llm_grid_h % vit_merger_window_size + ) % vit_merger_window_size + pad_w = ( + vit_merger_window_size - llm_grid_w % vit_merger_window_size + ) % vit_merger_window_size + num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size + num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size + index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", _PAD_VALUE) + index_padded = ( + index_padded.reshape( + grid_t, + num_windows_h, + vit_merger_window_size, + num_windows_w, + vit_merger_window_size, + ) + .permute(0, 1, 3, 2, 4) + .reshape( + grid_t, + num_windows_h * num_windows_w, + vit_merger_window_size, + vit_merger_window_size, + ) + ) + seqlens = (index_padded != _PAD_VALUE).sum([2, 3]).reshape(-1) + index_padded = index_padded.reshape(-1) + index_new = index_padded[index_padded != _PAD_VALUE] + window_index_list.append(index_new + window_index_id) + cu_seqlens_tmp = ( + seqlens.cumsum(0) * spatial_merge_unit + cu_window_seqlens[-1] + ) + cu_window_seqlens.extend(cu_seqlens_tmp.tolist()) + window_index_id += int(grid_t) * int(llm_grid_h) * int(llm_grid_w) + window_index = torch.cat(window_index_list, dim=0) + cu_window_seqlens = torch.tensor( + cu_window_seqlens, dtype=grid_thw.dtype, device=grid_thw.device + ) + return window_index, cu_window_seqlens + + +get_rope_index = GetRopeIndexOp() +get_window_index = GetWindowIndexOp() diff --git a/wall_x/model/core/ops/moe.py b/wall_x/model/core/ops/moe.py new file mode 100644 index 0000000..c50d946 --- /dev/null +++ b/wall_x/model/core/ops/moe.py @@ -0,0 +1,97 @@ +"""MoE routing operators (permute/unpermute). + +These operators reorder tokens by expert assignment for efficient +Mixture-of-Experts processing. +""" + +import logging + +import torch + +from wall_x.model.core.ops.base import OpsProxy + +logger = logging.getLogger(__name__) + + +class PermuteOp(OpsProxy): + """Reorder tokens by expert assignment for MoE processing. + + Signature: permute(tokens, indices, num_out_tokens=None, max_token_num=0) -> (permuted_tokens, sorted_indices) + """ + + @property + def _external_accel_name(self): + return "permute" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import permute_kernel + + return permute_kernel + except ImportError: + return None + except Exception as e: + logger.warning("PermuteOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, tokens, indices, num_out_tokens=None, max_token_num=0): + """PyTorch fallback for permute. + + Args: + num_out_tokens: If set, truncate output to this many tokens (matches + external_accel behavior of discarding padding expert assignments). + max_token_num: Unused, kept for external_accel API compatibility. + """ + del max_token_num # unused, external_accel API compat + if indices.dim() == 1: + indices = indices.view(-1, 1) + expand_factor = indices.size(1) + flatten_indices = indices.view(-1) + # Keep int64 throughout to avoid precision loss on large token counts + sorted_indices = torch.argsort(flatten_indices, stable=True) + permuted_tokens = tokens.index_select(0, sorted_indices // expand_factor) + if num_out_tokens is not None: + permuted_tokens = permuted_tokens[:num_out_tokens] + sorted_indices = sorted_indices[:num_out_tokens] + return permuted_tokens, sorted_indices + + +class UnpermuteOp(OpsProxy): + """Restore tokens to original order after MoE processing. + + Signature: unpermute(permuted_tokens, sorted_indices, probs=None) -> restored_tokens + """ + + @property + def _external_accel_name(self): + return "unpermute" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import unpermute_kernel + + return unpermute_kernel + except ImportError: + return None + except Exception as e: + logger.warning("UnpermuteOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, permuted_tokens, sorted_indices, probs=None): + if probs is not None: + merge_factor = probs.size(1) + else: + merge_factor = 1 + unpermuted_tokens = torch.zeros_like(permuted_tokens) + unpermuted_tokens.index_copy_(0, sorted_indices.long(), permuted_tokens) + unpermuted_tokens = unpermuted_tokens.reshape( + -1, merge_factor, permuted_tokens.size(-1) + ) + if probs is not None: + unpermuted_tokens = unpermuted_tokens * probs.unsqueeze(-1) + unpermuted_tokens = unpermuted_tokens.sum(dim=1) + return unpermuted_tokens + + +permute = PermuteOp() +unpermute = UnpermuteOp() diff --git a/wall_x/model/core/ops/norm.py b/wall_x/model/core/ops/norm.py new file mode 100644 index 0000000..f410540 --- /dev/null +++ b/wall_x/model/core/ops/norm.py @@ -0,0 +1,42 @@ +"""RMS normalization operator.""" + +import logging + +import torch + +from wall_x.model.core.ops.base import OpsProxy + +logger = logging.getLogger(__name__) + + +class RMSNormOp(OpsProxy): + """RMS normalization: rmsnorm(x, weight, eps). + + Accepts ``rmsnorm(x, weight, eps)``. + """ + + @property + def _external_accel_name(self): + return "rmsnorm" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import Rmsnorm + + return Rmsnorm() + except ImportError: + return None + except Exception as e: + logger.warning("RMSNormOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, hidden_states, weight, eps=1e-6): + """Pure PyTorch RMSNorm. eps default matches external_accel (1e-6, not PyTorch's 1e-5).""" + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + eps) + return (weight * hidden_states).to(input_dtype) + + +rmsnorm = RMSNormOp() diff --git a/wall_x/model/core/ops/rope.py b/wall_x/model/core/ops/rope.py new file mode 100644 index 0000000..09f6f42 --- /dev/null +++ b/wall_x/model/core/ops/rope.py @@ -0,0 +1,234 @@ +"""Rotary position embedding operators.""" + +import logging +from typing import List + +import torch + +from wall_x.model.core.ops.base import OpsProxy + +logger = logging.getLogger(__name__) + + +def _rotate_half(x): + """Rotate half: split last dim in two halves, negate-swap, concat.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _rotate_interleave(x): + """Interleaved rotation: pairs of (even, odd) elements.""" + x_even = x[..., ::2] + x_odd = x[..., 1::2] + return torch.stack((-x_odd, x_even), dim=-1).flatten(-2) + + +class RoPEOp(OpsProxy): + """Standard rotary position embedding. + + Signature: rope(q, k, cos, sin, interleave=False) -> (q_embed, k_embed) + """ + + @property + def _external_accel_name(self): + return "rope" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import Rope + + return Rope() + except ImportError: + return None + except Exception as e: + logger.warning("RoPEOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, q, k, cos, sin, interleave=False, **kwargs): + cos = cos.float() + sin = sin.float() + rotary_dim = cos.size(-1) * 2 + head_dim = q.size(-1) + if rotary_dim > head_dim: + raise ValueError( + f"rotary_dim ({rotary_dim}) > head_dim ({head_dim}): " + f"cos last dim ({cos.size(-1)}) is too large for q" + ) + partial = rotary_dim < head_dim + if partial: + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + else: + q_rot, k_rot = q, k + cos = cos.unsqueeze(-2) # (..., 1, half_dim) + sin = sin.unsqueeze(-2) + if interleave: + cos = cos.repeat_interleave(2, dim=-1) + sin = sin.repeat_interleave(2, dim=-1) + q_embed = q_rot.float() * cos + _rotate_interleave(q_rot.float()) * sin + k_embed = k_rot.float() * cos + _rotate_interleave(k_rot.float()) * sin + else: + cos = torch.cat((cos, cos), dim=-1) + sin = torch.cat((sin, sin), dim=-1) + q_embed = q_rot.float() * cos + _rotate_half(q_rot.float()) * sin + k_embed = k_rot.float() * cos + _rotate_half(k_rot.float()) * sin + if partial: + q_embed = torch.cat([q_embed.to(q.dtype), q_pass], dim=-1) + k_embed = torch.cat([k_embed.to(k.dtype), k_pass], dim=-1) + else: + q_embed = q_embed.to(q.dtype) + k_embed = k_embed.to(k.dtype) + return q_embed, k_embed + + def pack(self, *args, **kwargs): + """Delegate to resolved backend's pack method if available.""" + if self._resolved_fn is None: + self._resolve() + if hasattr(self._resolved_fn, "pack"): + return self._resolved_fn.pack(*args, **kwargs) + return args + + +class MRoPEOp(OpsProxy): + """Multi-head rotary position embedding (used by Qwen2.5-VL models). + + Signature: m_rope(query_states, key_states, cos, sin, mrope_section, interleaved=False) -> (q_embed, k_embed) + + cos/sin shape: (3, B, S, D//2). mrope_section like [16, 24, 24] specifies + how many head dims each of T/H/W occupies. After split by mrope_section * 2, + uses i % 3 to select from corresponding temporal/height/width rows. + """ + + @property + def _external_accel_name(self): + return "m_rope" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import MRope + + return MRope() + except ImportError: + return None + except Exception as e: + logger.warning("MRoPEOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback( + self, + query_states, + key_states, + cos, + sin, + mrope_section: List[int], + interleaved=False, + **kwargs, + ): + cos = cos.float() + sin = sin.float() + # Double cos/sin along last dim (matching external_accel kernel's internal behavior) + cos = torch.cat((cos, cos), dim=-1) + sin = torch.cat((sin, sin), dim=-1) + # List concatenation: [16,24,24] -> [16,24,24,16,24,24] (NOT element-wise multiply) + mrope_section_doubled = mrope_section + mrope_section + cos_split = torch.cat( + [m[i % 3] for i, m in enumerate(cos.split(mrope_section_doubled, dim=-1))], + dim=-1, + ).unsqueeze( + 2 + ) # (B, S, 1, D) + sin_split = torch.cat( + [m[i % 3] for i, m in enumerate(sin.split(mrope_section_doubled, dim=-1))], + dim=-1, + ).unsqueeze(2) + q_embed = (query_states.float() * cos_split) + ( + _rotate_half(query_states.float()) * sin_split + ) + k_embed = (key_states.float() * cos_split) + ( + _rotate_half(key_states.float()) * sin_split + ) + return q_embed.to(query_states.dtype), k_embed.to(key_states.dtype) + + def pack(self, *args, **kwargs): + """Delegate to resolved backend's pack method if available.""" + if self._resolved_fn is None: + self._resolve() + if hasattr(self._resolved_fn, "pack"): + return self._resolved_fn.pack(*args, **kwargs) + return args + + +class RotPosEmbOp(OpsProxy): + """Rotary position embedding computation for ViT (used by Qwen2.5-VL vision encoder). + + Signature: rot_pos_emb(inv_freq, grid_thw, spatial_merge_size) -> rotary_pos_emb + """ + + @property + def _external_accel_name(self): + return "rot_pos_emb" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import RotPos + + return RotPos() + except ImportError: + return None + except Exception as e: + logger.warning("RotPosEmbOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, inv_freq, grid_thw, spatial_merge_size): + if inv_freq.dtype != torch.float32: + inv_freq = inv_freq.to(torch.float32) + pos_ids = [] + for t, h, w in grid_thw: + t, h, w = int(t), int(h), int(w) + if h % spatial_merge_size != 0 or w % spatial_merge_size != 0: + raise ValueError( + f"grid h={h}, w={w} must be divisible by spatial_merge_size={spatial_merge_size}" + ) + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + hpos_ids = ( + hpos_ids.reshape( + h // spatial_merge_size, + spatial_merge_size, + w // spatial_merge_size, + spatial_merge_size, + ) + .permute(0, 2, 1, 3) + .flatten() + ) + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + wpos_ids = ( + wpos_ids.reshape( + h // spatial_merge_size, + spatial_merge_size, + w // spatial_merge_size, + spatial_merge_size, + ) + .permute(0, 2, 1, 3) + .flatten() + ) + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0).to(inv_freq.device) + max_grid_size = grid_thw[:, 1:].max() + seq = torch.arange(max_grid_size, device=inv_freq.device, dtype=inv_freq.dtype) + rotary_pos_emb_full = torch.outer(seq, inv_freq) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb.to(torch.float) + + def pack(self, *args, **kwargs): + """Delegate to resolved backend's pack method if available.""" + if self._resolved_fn is None: + self._resolve() + if hasattr(self._resolved_fn, "pack"): + return self._resolved_fn.pack(*args, **kwargs) + return args + + +rope = RoPEOp() +m_rope = MRoPEOp() +rot_pos_emb = RotPosEmbOp() diff --git a/wall_x/model/core/vla_mixin.py b/wall_x/model/core/vla_mixin.py new file mode 100644 index 0000000..b438589 --- /dev/null +++ b/wall_x/model/core/vla_mixin.py @@ -0,0 +1,746 @@ +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.utils.checkpoint as cp +from peft import LoraConfig, get_peft_model +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import MixedPrecision as MP +from transformers import GenerationMixin +from transformers.utils import logging + +from wall_x.model.core.action.moe import SparseMoeBlock, TokenTypeRouter +from wall_x.model.core.action.processor import ActionProcessor +from wall_x.model.core.attention.mask import ( + update_joint_attention_flash_mask, + update_joint_attention_mask_2d, + update_position_ids, +) +from wall_x.utils.constant import is_action_dataset_name + +logger = logging.get_logger(__name__) + + +class ActionModelMixMin: + # config: Qwen2_5_VLConfig + action_preprocessor: ActionProcessor + router: "TokenTypeRouter" + moe: "SparseMoeBlock" + + def __init__(self, config, action_preprocessor, router, moe): + self.config = config + self.action_preprocessor = action_preprocessor + self.router = router + self.moe = moe + self._mot_opt_warned = False + + def set_normalizer(self, normalizer_action, normalizer_propri): + if hasattr(self, "action_preprocessor"): + self.action_preprocessor.set_normalizer( + normalizer_action, normalizer_propri + ) + else: + # WARNING: normalizer cannot be set when action_preprocessor is missing + logger.warning( + "ActionModelMixMin.set_normalizer is called but action_preprocessor is not set" + ) + + def _apply_mlp_moe(self, hidden_states, token_types, start_indices, end_indices): + if self.config.mlp_moe: + hidden_states = self.moe( + hidden_states, token_types, start_indices, end_indices + ) + else: + hidden_states = self.mlp(hidden_states) + return hidden_states + + def _apply_norm_moe( + self, + hidden_states, + token_types, + adarms_conds, + norms, # list of norm layers (expert-wise) + norm, # shared norm if not norm_moe + start_indices=None, + end_indices=None, + use_selective_recompute=False, + ): + """ + MoE-aware LayerNorm with optional selective activation recomputation. + + Only activation math is recomputed. No GEMM is recomputed. + Safe for FSDP (use_reentrant=False). + """ + + gate = None + gate_mask = None + + # ------------------------- + # Case 1: norm_moe=True (expert-wise norm) + # ------------------------- + if self.config.norm_moe: + + # --------------------------------------------------------- + # Case 1A: mot_opt=True (segments assigned by start/end) + # --------------------------------------------------------- + if self.config.mot_opt: + new_hidden_states = torch.zeros_like(hidden_states) + + for expert_idx, expert_norm in enumerate(norms): + start = start_indices[expert_idx] + end = end_indices[expert_idx] + if start == end: + continue + + dim_input = self.config.dim_inputs[expert_idx] + selected = hidden_states[start:end] # [K, D] + + # ====== reshape if adarms on flow expert ====== + if self.config.use_adarms and expert_idx == 1: + selected = selected.view( + -1, + self.config.action_horizon_flow, + selected.shape[-1], + ) + input_slice = selected[:, :, :dim_input] + cond = adarms_conds[expert_idx] + else: + input_slice = selected[:, :dim_input] + cond = adarms_conds[expert_idx] + + if use_selective_recompute: + + def norm_chunk(t_x, t_cond, expert_norm=expert_norm): + if t_cond is None or ( + isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0 + ): + out, _ = expert_norm(t_x) + else: + out, _ = expert_norm(t_x, t_cond) + return out + + cond_for_cp = ( + cond + if cond is not None + else torch.empty(0, device=input_slice.device) + ) + processed = cp.checkpoint( + norm_chunk, + input_slice, + cond_for_cp, + use_reentrant=False, + ) + else: + processed, gate = expert_norm(input_slice, cond) + + # reshape back if needed + if self.config.use_adarms and expert_idx == 1: + processed = processed.view(-1, dim_input) + + new_hidden_states[start:end, :dim_input] = processed.to( + hidden_states.dtype + ) + + hidden_states = new_hidden_states + + # --------------------------------------------------------- + # Case 1B: mot_opt=False (token-level mask) + # --------------------------------------------------------- + else: + + new_hidden_states = torch.zeros_like(hidden_states) + B, S, D = hidden_states.shape + + for expert_idx, expert_norm in enumerate(norms): + mask = token_types == expert_idx + if mask.sum() == 0: + continue + + dim_input = self.config.dim_inputs[expert_idx] + selected = hidden_states[mask] # [K, D] + + if self.config.use_adarms and expert_idx == 1: + gate_mask = mask + selected = selected.view( + -1, + self.config.action_horizon_flow, + selected.shape[-1], + ) + input_slice = selected[:, :, :dim_input] + cond = adarms_conds[expert_idx] + else: + input_slice = selected[:, :dim_input] + cond = adarms_conds[expert_idx] + + if use_selective_recompute: + + def norm_chunk(t_x, t_cond, expert_norm=expert_norm): + if t_cond is None or ( + isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0 + ): + out, _ = expert_norm(t_x) + else: + out, _ = expert_norm(t_x, t_cond) + return out + + cond_for_cp = ( + cond + if cond is not None + else torch.empty(0, device=input_slice.device) + ) + + processed = cp.checkpoint( + norm_chunk, + input_slice, + cond_for_cp, + use_reentrant=False, + ) + else: + processed, gate = expert_norm(input_slice, cond) + + if self.config.use_adarms and expert_idx == 1: + processed = processed.view(-1, dim_input) + + # scatter back + b_id, s_id = torch.where(mask) + new_hidden_states[b_id, s_id, :dim_input] = processed.to( + hidden_states.dtype + ) + + hidden_states = new_hidden_states + + # ------------------------- + # Case 2: norm_moe=False (single LN) + # ------------------------- + else: + + def norm_chunk_shared(t_x, dummy, norm_module=norm): + out, _ = norm_module(t_x) + return out + + if use_selective_recompute: + dummy = torch.empty(0, device=hidden_states.device) + hidden_states = cp.checkpoint( + norm_chunk_shared, + hidden_states, + dummy, + use_reentrant=False, + ) + else: + hidden_states, gate = norm(hidden_states) + + return hidden_states, gate, gate_mask + + def _gated_residual(self, x, y, gate, start_indices=None, end_indices=None): + """ + Applies gated residual connection with optional gate parameter. + + Args: + x: Input tensor (residual) + y: Output tensor to be added + gate: Optional gate tensor to modulate the addition + + Returns: + x + y if gate is None, otherwise x + y * gate + """ + if x is None and y is None: + return None + if x is None or y is None: + return x if x is not None else y + if gate is None: + return x + y + + new_y = y.clone() + selected_y = y[start_indices[1] : end_indices[1]] + selected_y = selected_y.view( + -1, self.config.action_horizon_flow, selected_y.shape[-1] + )[:, :, : self.config.dim_inputs[1]] + selected_y = selected_y.to(torch.float32) * gate + new_y[start_indices[1] : end_indices[1], : self.config.dim_inputs[1]] = ( + selected_y.view(-1, self.config.dim_inputs[1]).to(new_y.dtype) + ) + + return x + new_y + + def scatter_proprioception_embeddings( + self, input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask + ): + use_state_string_representation = getattr( + self.config, "use_state_string_representation", False + ) + if proprioception is not None and not use_state_string_representation: + proprioception = proprioception.to(inputs_embeds.device).to( + inputs_embeds.dtype + ) + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to( + inputs_embeds.dtype + ) + proprioception = self.action_preprocessor.proprioception_proj( + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + mask = input_ids == self.action_token_id_set["propri_token_id"] + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + proprioception_mask = mask_expanded.to(inputs_embeds.device) + + proprioception = proprioception.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter( + proprioception_mask, proprioception + ) + + return inputs_embeds + + def scatter_flow_action_embeddings( + self, + input_ids, + inputs_embeds, + action_chunk, + dataset_names, + sample_time, + dof_mask, + ): + if not self.config.use_flow_action_expert: + return inputs_embeds, None, None + adarms_cond, flow = None, None + if action_chunk is not None: + action_chunk = action_chunk.to(inputs_embeds.device) + dof_mask = dof_mask.to(inputs_embeds.device) + noisy_action_emb, flow, adarms_cond = self.action_preprocessor( + action_chunk, dataset_names, sample_time, dof_mask + ) + mask = input_ids == self.action_token_id_set["action_token_id"] + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + action_mask = mask_expanded.to(inputs_embeds.device) + + noisy_action_emb = noisy_action_emb.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb) + + return inputs_embeds, flow, adarms_cond + + @staticmethod + def _update_position_ids(position_ids, moe_token_types, positional_masks): + return update_position_ids(position_ids, moe_token_types, positional_masks) + + def _update_joint_attention_mask_2d( + self, attention_mask, moe_token_types, positional_masks + ): + return update_joint_attention_mask_2d( + attention_mask, + moe_token_types, + positional_masks, + causal_action_attention_mask=self.config.causal_action_attention_mask, + ) + + def _update_joint_attention_flash_mask( + self, attention_mask, moe_token_types, positional_masks, debug=False + ): + return update_joint_attention_flash_mask( + attention_mask, + moe_token_types, + positional_masks, + causal_action_attention_mask=self.config.causal_action_attention_mask, + debug=debug, + ) + + +class ActionGenerationMixin(GenerationMixin): + action_preprocessor: ActionProcessor + + def to_bfloat16_for_selected_params(self, fsdp_plugin=None, accelerator=None): + """ + Keep selected model parameters in float32 and cast the rest to bfloat16. + - If fsdp_plugin exists, use FSDP v1 mixed_precision wrapping. + - Otherwise, modify parameter dtypes directly. + """ + + def _assign_child(root_module, dotted_name: str, new_child): + parts = dotted_name.split(".") + parent = root_module + for p in parts[:-1]: + parent = getattr(parent, p) + setattr(parent, parts[-1], new_child) + + if fsdp_plugin: + fsdp_version = getattr(fsdp_plugin, "fsdp_version", None) + if fsdp_version != 1: + raise RuntimeError("Only FSDP v1 (fsdp_version=1) is supported.") + + device = getattr( + accelerator, "device", torch.device("cuda", torch.cuda.current_device()) + ) + if isinstance(device, torch.device) and device.type == "cuda": + if device.index is not None: + torch.cuda.set_device(device.index) + device_id = device.index + + # Move the model to the target device first + self = self.to(device) + + # Define the mixed precision policy + bf16_policy = MP( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.bfloat16, + cast_forward_inputs=False, + cast_root_forward_inputs=False, + ) + + fp32_policy = MP( + param_dtype=torch.float32, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + cast_forward_inputs=False, + cast_root_forward_inputs=False, + ) + + from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl import ( + Qwen2_5_VLVisionBlock, + ) + from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl_act import ( + Qwen2_5_VLDecoderLayer_with_MoE, + ) + + target_classes = ( + Qwen2_5_VLDecoderLayer_with_MoE, + Qwen2_5_VLVisionBlock, + ) + + # Step 1: Find top-level ActionProcessor modules and wrap them separately with FSDP (FP32) + for name, module in list(self.named_modules()): + if isinstance(module, nn.Module) and any( + k in name.lower() + for k in [ + "input_layernorm", + "post_attention_layernorm", + "model.norm", + "action_preprocessor", + ] + ): + if any( + True for _ in module.children() + ): # Wrap only leaves to avoid parent-child duplication + continue + if getattr(module, "_fsdp_wrapped", False): + continue + + logger.info("[FSDP v1] wrapping module in FP32: %s", name) + wrapped = FSDP( + module, + mixed_precision=fp32_policy, + sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, + backward_prefetch="BACKWARD_PRE", + device_id=device_id, + use_orig_params=True, + ) + _assign_child(self, name, wrapped) + setattr(wrapped, "_fsdp_wrapped", True) + + for name, module in list(self.named_modules()): + if isinstance(module, target_classes): + + if getattr(module, "_fsdp_wrapped", False): + continue + + logger.info("[FSDP v1] wrapping module in BF16: %s", name) + wrapped = FSDP( + module, + mixed_precision=bf16_policy, + sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, + backward_prefetch="BACKWARD_PRE", + device_id=device_id, + use_orig_params=True, + ) + _assign_child(self, name, wrapped) + setattr(wrapped, "_fsdp_wrapped", True) + + # Step 2: Wrap the outer module with FSDP using the BF16 policy + logger.info("[FSDP v1] wrapping root model with bf16 mixed precision...") + self = FSDP( + self, + mixed_precision=bf16_policy, + sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, + backward_prefetch="BACKWARD_PRE", + device_id=device_id, + use_orig_params=True, + ) + + return self + + # ----------------- Non-FSDP path ----------------- + else: + logger.info("Running manual dtype conversion (no FSDP).") + params_to_keep_float32 = [] + for name, _ in self.named_parameters(): + if any( + k in name + for k in [ + "input_layernorm", + "post_attention_layernorm", + "model.norm", + "action_preprocessor", + "action_processor", + ] + ): + params_to_keep_float32.append(name) + + for name, param in self.named_parameters(): + if name not in params_to_keep_float32: + param.data = param.data.to(torch.bfloat16) + if name in params_to_keep_float32: + param.data = param.data.to(torch.float32) + + return self + + def define_action_token_id(self): + # Get the action token list through tokenizer_mixin; compatible with pure flow mode + if self.processor is not None: + action_token_list = [] + if self.tokenizer_mixin is not None: + action_token_list = self.tokenizer_mixin.get_action_token_list( + self.processor + ) + + action_token_id = self.processor.tokenizer.convert_tokens_to_ids( + "<|action|>" + ) + propri_token_id = self.processor.tokenizer.convert_tokens_to_ids( + "<|propri|>" + ) + self.action_token_id_set = { + "action_token_list": action_token_list, + "propri_token_id": propri_token_id, + "action_token_id": action_token_id, + } + + def add_lora( + self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1 + ): + """Add LoRA adapters""" + config = LoraConfig( + r=r, + lora_alpha=lora_alpha, + target_modules=target_modules, + lora_dropout=lora_dropout, + bias="none", + task_type="CAUSAL_LM", + ) + self.model = get_peft_model(self.model, config) + # Log trainable parameter information + self.model.print_trainable_parameters() + + def compute_loss( + self, + hidden_states, + logits, + input_ids=None, + dataset_names=None, + labels=None, + action_chunk=None, + dof_mask=None, + flow=None, + flow_loss_mask=None, + _lm_loss_mask=None, + action_hidden_states=None, + **kwargs, + ): + if input_ids is not None: + batch_size, seq_length = input_ids.shape + + loss = 0 + cross_entropy_loss, flow_loss = None, None + + if dataset_names is not None: + unique_datasets_name = list(set(dataset_names)) + _device = hidden_states.device + _flow_channel_names = [ + f"{name}_flow" + for name in unique_datasets_name + if is_action_dataset_name(name) + ] + channel_loss_dict = { + name: torch.tensor(0.0, device=_device) + for name in unique_datasets_name + _flow_channel_names + } + channel_loss_count_dict = { + name: torch.tensor(0, device=_device) + for name in unique_datasets_name + _flow_channel_names + } + else: + unique_datasets_name, channel_loss_dict, channel_loss_count_dict = ( + None, + None, + None, + ) + + if labels is not None: + if _lm_loss_mask is not None: + # ===== Optimized path: logits already gathered for loss tokens only ===== + # logits: [N_loss, V] or None, _lm_loss_mask: [B, S-1] + if logits is not None: + shift_logits = logits.to(torch.float32) # [N_loss, V] + shift_labels = labels[..., 1:].contiguous() + shift_labels_flat = shift_labels[_lm_loss_mask] # [N_loss] + shift_labels_flat = shift_labels_flat.to(shift_logits.device) + + _cross_entropy_loss = self.loss_fct(shift_logits, shift_labels_flat) + cross_entropy_loss = _cross_entropy_loss.mean() + + # compute channel loss + if unique_datasets_name is not None: + batch_idx = ( + torch.arange(batch_size, device=_lm_loss_mask.device) + .unsqueeze(1) + .expand_as(_lm_loss_mask) + ) + loss_batch_idx = batch_idx[_lm_loss_mask] # [N_loss] + for dataset_name_i in unique_datasets_name: + ds_mask = torch.tensor( + [name == dataset_name_i for name in dataset_names], + device=_lm_loss_mask.device, + dtype=torch.bool, + ) + tok_ds_mask = ds_mask[loss_batch_idx] + channel_loss_dict[dataset_name_i] = ( + _cross_entropy_loss[tok_ds_mask].sum() + if tok_ds_mask.any() + else torch.tensor(0.0, device=shift_logits.device) + ) + channel_loss_count_dict[dataset_name_i] += tok_ds_mask.sum() + else: + cross_entropy_loss = torch.tensor(0.0, device=hidden_states.device) + else: + # ===== Original path (inference / no labels optimization) ===== + shift_logits = logits[..., :-1, :].contiguous().to(torch.float32) + shift_labels = labels[..., 1:].contiguous() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + non_ignored_mask = shift_labels != -100 + _cross_entropy_loss = self.loss_fct(shift_logits, shift_labels) + cross_entropy_loss = ( + _cross_entropy_loss[non_ignored_mask].mean() + if non_ignored_mask.any() + else torch.tensor(0.0, device=shift_logits.device) + ) + + # compute channel loss + _cross_entropy_loss = _cross_entropy_loss.view( + batch_size, seq_length - 1 + ) + non_ignored_mask = non_ignored_mask.view(batch_size, seq_length - 1) + if unique_datasets_name is not None: + for dataset_name_i in unique_datasets_name: + dataset_mask = torch.tensor( + [name == dataset_name_i for name in dataset_names], + device=logits.device, + ) + combined_mask = dataset_mask.unsqueeze(1) & non_ignored_mask + channel_loss_dict[dataset_name_i] = ( + _cross_entropy_loss[combined_mask].sum() + if combined_mask.any() + else torch.tensor(0.0, device=shift_logits.device) + ) + channel_loss_count_dict[dataset_name_i] += combined_mask.sum() + + if not torch.isnan(cross_entropy_loss): + loss += cross_entropy_loss * self.config.ar_loss_weight + else: + with torch.no_grad(): + cross_entropy_loss.detach() + + # compute action token accuracy(computed uniformly through tokenizer_mixin) + # if self.tokenizer_mixin is not None and self.action_mapper is not None: + # accuracy_dict = self.tokenizer_mixin.compute_accuracy( + # logits, labels, self.action_mapper, self.action_token_id_set + # ) + # channel_loss_dict.update(accuracy_dict) + + if action_chunk is not None: + action_mask = input_ids == self.action_token_id_set["action_token_id"] + if action_mask.any(): + if action_hidden_states is None: + action_hidden_states = hidden_states[action_mask].to(torch.float32) + else: + action_hidden_states = action_hidden_states.reshape( + -1, action_hidden_states.shape[-1] + ).to(torch.float32) + flow = flow.reshape(-1, flow.shape[-1]) + _flow_loss = self.action_preprocessor.flow_loss( + action_hidden_states, flow, action_chunk, dof_mask, flow_loss_mask + ) + if isinstance(_flow_loss, torch.Tensor): + # Compute the valid-element mask as the intersection of dof_mask and flow_loss_mask + valid_mask = ( + dof_mask.reshape(-1, dof_mask.shape[-1]) + if dof_mask is not None + else None + ) + if flow_loss_mask is not None: + flow_mask_expanded = ( + flow_loss_mask.unsqueeze(-1) + .reshape(-1, 1) + .expand(-1, _flow_loss.shape[-1]) + ) + valid_mask = ( + valid_mask * flow_mask_expanded + if valid_mask is not None + else flow_mask_expanded + ) + if valid_mask is not None and not valid_mask.all(): + flow_loss = _flow_loss.sum() / valid_mask.sum() + else: + flow_loss = _flow_loss.mean() + loss += flow_loss + _flow_loss = _flow_loss.view( + dof_mask.shape[0], dof_mask.shape[1], dof_mask.shape[2] + ) + + # compute flow channel loss + if unique_datasets_name is not None: + B, T, D = _flow_loss.shape + action_ds_names = [ + name + for name in dataset_names + if is_action_dataset_name(name) + ] + if valid_mask is not None: + valid_mask_3d = valid_mask.view(B, T, D) + else: + valid_mask_3d = torch.ones_like( + _flow_loss, dtype=torch.bool, device=_flow_loss.device + ) + for dataset_name_i in unique_datasets_name: + ds_mask = torch.tensor( + [name == dataset_name_i for name in action_ds_names], + device=_flow_loss.device, + dtype=torch.bool, + ) + if not ds_mask.any(): + continue + flow_key = f"{dataset_name_i}_flow" + ds_mask_3d = ds_mask.view(-1, 1, 1).expand(B, T, D) + flow_loss_sum = (_flow_loss * ds_mask_3d).sum() + flow_count = (valid_mask_3d * ds_mask_3d).sum() + channel_loss_dict[flow_key] = ( + channel_loss_dict[flow_key] + flow_loss_sum + ) + channel_loss_count_dict[flow_key] = ( + channel_loss_count_dict[flow_key] + flow_count + ) + + return ( + loss, + cross_entropy_loss, + flow_loss, + channel_loss_dict, + channel_loss_count_dict, + ) diff --git a/wall_x/model/model_utils.py b/wall_x/model/model_utils.py deleted file mode 100644 index 5aaeecc..0000000 --- a/wall_x/model/model_utils.py +++ /dev/null @@ -1,319 +0,0 @@ -import torch -import os -import numpy as np -from transformers import AutoProcessor -from wall_x.model.action_head import Normalizer - - -def update_model_config(train_config, model_config): - model_config.use_state_string_representation = train_config["data"].get( - "use_state_string_representation", False - ) - model_config.flow_loss_weight = train_config.get("flow_loss_weight", 1.0) - - model_config.dof_config = train_config["dof_config"] - model_config.agent_pos_config = train_config["agent_pos_config"] - - model_config.action_horizon_flow = train_config["data"].get( - "action_horizon_flow", 32 - ) - - if train_config.get("_attn_implementation", None) is not None: - model_config._attn_implementation = train_config["_attn_implementation"] - - return model_config - - -def load_wallx_processors(config): - processor = AutoProcessor.from_pretrained(config["processor_path"], use_fast=True) - # pad side = left - processor.tokenizer.padding_side = "left" - - new_tokens = ["<|propri|>", "<|action|>"] - # special_tokens = [] - action_tokenizer_type = config.get("action_tokenizer_type", None) - if action_tokenizer_type == "fast": - train_action_tokenizer = AutoProcessor.from_pretrained( - config["action_tokenizer_path"], trust_remote_code=True - ) - val_action_tokenizer = AutoProcessor.from_pretrained( - config["action_tokenizer_path"], trust_remote_code=True - ) - new_tokens += [ - f"<|action_token_{i}|>" for i in range(train_action_tokenizer.vocab_size) - ] - elif action_tokenizer_type == "spatialvla": - raise NotImplementedError("SpatialActionTokenizer is not implemented") - else: - train_action_tokenizer = None - val_action_tokenizer = None - - num_added_tokens = processor.tokenizer.add_tokens(new_tokens) - - if action_tokenizer_type and train_action_tokenizer.vocab_size > 0: - action_mapper = {} - for i in range(train_action_tokenizer.vocab_size): - token = f"<|action_token_{i}|>" - token_id = processor.tokenizer.convert_tokens_to_ids(token) - action_mapper[token_id] = i - else: - action_mapper = None - - return { - "processor": processor, - "train_action_tokenizer": train_action_tokenizer, - "val_action_tokenizer": val_action_tokenizer, - "action_mapper": action_mapper, - "num_added_tokens": num_added_tokens, - } - - -def register_normalizers(config, model_path): - # if config.get("customized_action_statistic_dof", None): - # action_statistic_dof = json.load(open(config["customized_action_statistic_dof"], "r")) - # else: - # action_statistic_dof = default_action_statistic_dof - - action_statistic_dof = None - - if os.path.exists(model_path + "/normalizer_action.pth"): - print( - "Loading normalizer_action from checkpoint", - model_path + "/normalizer_action.pth", - flush=True, - ) - normalizer_action = Normalizer.from_ckpt(model_path + "/normalizer_action.pth") - else: - normalizer_action = Normalizer( - action_statistic_dof, - config["dof_config"], - min_key=config.get("min_key", "min"), - delta_key=config.get("delta_key", "delta"), - ) - - # print("action_statistic_dof",action_statistic_dof) - - if os.path.exists(model_path + "/normalizer_propri.pth"): - print( - "Loading normalizer_propri from checkpoint", - model_path + "/normalizer_propri.pth", - flush=True, - ) - normalizer_propri = Normalizer.from_ckpt(model_path + "/normalizer_propri.pth") - else: - normalizer_propri = Normalizer( - action_statistic_dof, - config["agent_pos_config"], - min_key=config.get("min_key", "min"), - delta_key=config.get("delta_key", "delta"), - ) - - return normalizer_action, normalizer_propri - - -def find_first_last_ones(tensor): - """ - Input: a tensor of shape (bs, seq_len) containing 0s and 1s - Output: (first_indices, last_indices), each of shape (bs,) - where first_indices[i] is the index of the first 1 in the i-th batch, or -1 if none exists. - last_indices[i] is the index of the last 1 in the i-th batch, or -1 if none exists. - """ - bs, seq_len = tensor.shape - masks = tensor == 1 - has_ones = masks.any(dim=1) - - first = torch.full((bs,), -1, dtype=torch.long, device=tensor.device) - last = first.clone() - - first[has_ones] = torch.argmax(masks[has_ones].float(), dim=1) - - flipped_masks = masks.flip(dims=[1]) - last_argmax = torch.argmax(flipped_masks[has_ones].float(), dim=1) - last[has_ones] = seq_len - 1 - last_argmax - - return first, last - - -def flashmask_to_densemask(startend_row_indices, dtype, causal=True): - if startend_row_indices is None: - return None - bz, num_head, seq_len, bound_num = startend_row_indices.shape - m = np.ones((bz, num_head, seq_len, seq_len), dtype=dtype) - has_end = (causal and bound_num == 2) or ((not causal) and bound_num == 4) - for bi in range(bz): - for hi in range(num_head): - for j in range(seq_len): - downstart = startend_row_indices[bi, hi, j, 0] - if has_end: - downend = startend_row_indices[bi, hi, j, 1] - m[bi, hi, downstart:downend, j] = 0 - else: - m[bi, hi, downstart:, j] = 0 - if causal: - m[bi, hi, :j, j] = 0 - else: - if has_end: - upstart = startend_row_indices[bi, hi, j, 2] - upend = startend_row_indices[bi, hi, j, 3] - m[bi, hi, upstart:upend, j] = 0 - else: - upend = startend_row_indices[bi, hi, j, 1] - m[bi, hi, :upend, j] = 0 - return m - - -def num_floating_point_operations( - args, - batch_size: int, - num_lang_tokens: int, - num_action_tokens: int, - vision_seq_length: int = 756, -): - """ - Accurately estimate the training FLOPs of Transformer + MoE + MoT + Vision. - - Supported: - - expert0 = language tokens - - expert1 = action tokens - - MoE MLP (2 experts) - - MoT Attention (2 experts) - - GQA - - Vision Transformer (full+window attention) - """ - assert args.num_experts == 2, "The current model only supports 2 experts." - - dim_lang, dim_act = args.dim_inputs - - # Number of tokens per layer (flattened across batch) - N_lang = batch_size * num_lang_tokens - N_action = batch_size * num_action_tokens - N_total = N_lang + N_action # Used for non-MoT attention - - # ================================================================ - # Text MLP FLOPs - # ================================================================ - hidden_size = args.hidden_size - ffn_hidden_size = args.intermediate_size - num_layers = args.num_hidden_layers - - use_moe_mlp = getattr(args, "mlp_moe", False) - - # ---------- Forward-only MLP FLOPs ---------- - def forward_mlp_flops(N, d_in, d_ff): - """ - SwiGLU forward: - gate = x @ W1 (2*N*d_in*d_ff) - up = x @ W2 (2*N*d_in*d_ff) - act = silu + mul (~2*N*d_ff) - down = h @ W3 (2*N*d_ff*d_in) - - Forward ≈ 4*N*d_in*d_ff + 2*N*d_ff*d_in = 6*N*d_in*d_ff + 2*N*d_ff - """ - return 6 * N * d_in * d_ff + 2 * N * d_ff - - if not use_moe_mlp: - # Dense MLP - F_fwd = forward_mlp_flops(N_total, hidden_size, ffn_hidden_size) - total_mlp_flops_text = 3 * num_layers * F_fwd # <-- training FLOPs - else: - # MoE: expert0(language) + expert1(action) - hid_lang = args.experts[0]["intermediate_size"] - hid_act = args.experts[1]["intermediate_size"] - - F_lang_fwd = forward_mlp_flops(N_lang, dim_lang, hid_lang) - F_act_fwd = forward_mlp_flops(N_action, dim_act, hid_act) - - total_mlp_flops_text = 3 * num_layers * (F_lang_fwd + F_act_fwd) - - # ================================================================ - # Text Attention FLOPs - # ================================================================ - num_heads = args.num_attention_heads - num_kv = args.num_key_value_heads - H = hidden_size - B = batch_size - S = num_lang_tokens + num_action_tokens - N = B * S - - use_mot = getattr(args, "attention_moe", False) - - # ---------- attention matmul ---------- - F_matmul_fwd = 4 * B * (S**2) * H - - if not use_mot: - # -------- GQA + QKV / O -------- - F_q_fwd = 2 * N * H * H - F_kv_fwd = 4 * N * H * H * (num_kv / num_heads) - F_o_fwd = 2 * N * H * H - - F_attn_fwd = F_q_fwd + F_kv_fwd + F_o_fwd + F_matmul_fwd - - else: - # -------- MoT: expert0 + expert1 QKV -------- - F_lang_qkv = N_lang * dim_lang * H * (2 + 4 * num_kv / num_heads) - F_act_qkv = N_action * dim_act * H * (2 + 4 * num_kv / num_heads) - F_attn_fwd = F_lang_qkv + F_act_qkv + F_matmul_fwd - - # Training FLOPs - total_attn_flops_text = 3 * num_layers * F_attn_fwd - - # ================================================================ - # Logits projection FLOPs - # ================================================================ - vocab_size = getattr(args, "padded_vocab_size", args.vocab_size) - - F_logits_fwd = 2 * N * H * vocab_size - total_logits_flops = 3 * F_logits_fwd - - total_text_flops = total_mlp_flops_text + total_attn_flops_text + total_logits_flops - - # ================================================================ - # Vision Transformer FLOPs - # ================================================================ - total_vision_flops = 0 - - if hasattr(args, "vision_config") and vision_seq_length is not None: - vcfg = args.vision_config - - Bv = batch_size - Sv = vision_seq_length - Nv = Bv * Sv - - Hv = vcfg.hidden_size - Iv = vcfg.intermediate_size - num_heads_v = vcfg.num_heads - window_size = vcfg.window_size - out_hidden = vcfg.out_hidden_size - - depth_v = vcfg.depth - fullatt = set(vcfg.fullatt_block_indexes) - num_full = len(fullatt) - num_local = depth_v - num_full - - # ---------- forward FLOPs ---------- - def forward_vit_mlp(N, H, Inner): - return 6 * N * H * Inner + 2 * N * Inner - - F_mlp_v = forward_vit_mlp(Nv, Hv, Iv) - F_qkv_v = 6 * Nv * Hv * Hv - F_o_v = 2 * Nv * Hv * Hv - - F_full = 4 * Bv * (Sv**2) * Hv - num_windows = Sv / window_size - F_win = 4 * Bv * num_windows * (window_size**2) * (Hv / num_heads_v) - - F_block_full_fwd = F_mlp_v + F_qkv_v + F_o_v + F_full - F_block_local_fwd = F_mlp_v + F_qkv_v + F_o_v + F_win - - # ---------- train FLOPs ---------- - total_vision_flops = 3 * ( - num_full * F_block_full_fwd + num_local * F_block_local_fwd - ) - - # merger - total_vision_flops += 3 * (2 * Nv * Hv * out_hidden) - - # ================================================================ - # TOTAL TRAIN FLOPs - # ================================================================ - return total_text_flops + total_vision_flops diff --git a/wall_x/model/qact/__init__.py b/wall_x/model/qact/__init__.py new file mode 100644 index 0000000..860ac44 --- /dev/null +++ b/wall_x/model/qact/__init__.py @@ -0,0 +1 @@ +"""QAct (Qwen-VLA) model family.""" diff --git a/wall_x/model/qact/qwen2_5/__init__.py b/wall_x/model/qact/qwen2_5/__init__.py new file mode 100644 index 0000000..3040cdd --- /dev/null +++ b/wall_x/model/qact/qwen2_5/__init__.py @@ -0,0 +1,2 @@ +from .configuration_qwen2_5_vl import Qwen2_5_VLConfig +from .modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction, Qwen2_5_VLMoEModel diff --git a/wall_x/model/qact/qwen2_5/adapter.py b/wall_x/model/qact/qwen2_5/adapter.py new file mode 100644 index 0000000..17352f4 --- /dev/null +++ b/wall_x/model/qact/qwen2_5/adapter.py @@ -0,0 +1,56 @@ +"""Qwen2.5 VLA adapter - variant-specific overrides on top of VLAdapter.""" + +from wall_x.model.registry import register_model +from wall_x.trainer.adapters.vla_model_adapter import VLAdapter + + +@register_model("qwen2_5") +class Qwen2_5Adapter(VLAdapter): + MODEL_TYPE = "qwen2_5" + + @classmethod + def model_class(cls): + from wall_x.model.qact.qwen2_5 import Qwen2_5_VLMoEForAction + + return Qwen2_5_VLMoEForAction + + @classmethod + def config_class(cls): + from wall_x.model.qact.qwen2_5 import Qwen2_5_VLConfig + + return Qwen2_5_VLConfig + + @classmethod + def inference_model_class(cls): + return cls.model_class() + + def get_transformer_layer_cls(self): + layer_classes = set() + try: + from transformers.models.qwen2_vl.modeling_qwen2_vl import ( + Qwen2VLDecoderLayer, + ) + + layer_classes.add(Qwen2VLDecoderLayer) + except ImportError: + pass + try: + from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl import ( + Qwen2_5_VLDecoderLayer, + ) + + layer_classes.add(Qwen2_5_VLDecoderLayer) + except ImportError: + pass + return layer_classes if layer_classes else None + + @staticmethod + def log_attention_implementation(logger, model): + logger.info( + f"*** model attention implementation: " + f"{model.model._attn_implementation} ***" + ) + logger.info( + f"*** model.visual attention implementation: " + f"{model.visual.config._attn_implementation} ***" + ) diff --git a/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py b/wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py similarity index 79% rename from wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py rename to wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py index a2d51a1..8956834 100644 --- a/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py +++ b/wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py @@ -1,6 +1,35 @@ +# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! +# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_qwen2_5_vl.py file directly. One of our CI enforces this. +# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! +# coding=utf-8 +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import logging + from transformers.configuration_utils import PretrainedConfig from transformers.modeling_rope_utils import rope_config_validation +logger = logging.getLogger(__name__) + class Qwen2_5_VLVisionConfig(PretrainedConfig): model_type = "qwen2_5_vl" @@ -203,18 +232,24 @@ class Qwen2_5_VLConfig(PretrainedConfig): mlp_moe=False, norm_moe=False, mot_opt=False, - flow_loss_weight=10, + ar_loss_weight=1.0, use_state_string_representation=False, use_adarms=False, proj_with_mask=True, - use_flow_action_expert=True, adarms_cond_dim=None, action_horizon_flow=32, causal_action_attention_mask=False, + use_flow_action_expert=True, use_x_pred=False, attn_deterministic=False, + use_x_loss=False, **kwargs, ): + # Compatibility with newer transformers versions (5.x): + # - Older versions: super() sets self.pad_token_id; override it with the saved value so kwargs are preserved + # - Newer versions: super() does not set self.pad_token_id; assign it afterward + _pad_token_id = kwargs.pop("pad_token_id", None) + self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size @@ -249,7 +284,7 @@ class Qwen2_5_VLConfig(PretrainedConfig): self.mlp_moe = mlp_moe self.norm_moe = norm_moe self.mot_opt = mot_opt - self.flow_loss_weight = flow_loss_weight + self.ar_loss_weight = ar_loss_weight self.use_state_string_representation = use_state_string_representation self.use_adarms = use_adarms @@ -258,9 +293,10 @@ class Qwen2_5_VLConfig(PretrainedConfig): self.use_flow_action_expert = use_flow_action_expert self.action_horizon_flow = action_horizon_flow self.causal_action_attention_mask = causal_action_attention_mask + self.use_flow_action_expert = use_flow_action_expert self.use_x_pred = use_x_pred self.attn_deterministic = attn_deterministic - + self.use_x_loss = use_x_loss # Validate the correctness of rotary position embeddings parameters # BC: if there is a 'type' field, move it to 'rope_type'. # and change type from 'mrope' to 'default' because `mrope` does defeault RoPE calculations @@ -274,6 +310,9 @@ class Qwen2_5_VLConfig(PretrainedConfig): super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + # Assign pad_token_id as described above, overriding older super() values or filling newer missing attributes + self.pad_token_id = _pad_token_id + # move vision config initialization after super init to avoid recursively set in latest transformers version # TODO: make it better if isinstance(vision_config, dict): @@ -281,5 +320,38 @@ class Qwen2_5_VLConfig(PretrainedConfig): elif vision_config is None: self.vision_config = self.sub_configs["vision_config"]() + def update_model_config(self, train_config): + """Update model configuration from training config. + + This method updates the model configuration with training-specific + settings such as action horizon, DOF config, attention implementation, etc. + + Args: + train_config: dict containing training configuration parameters. + """ + self.use_state_string_representation = train_config["data"].get( + "use_state_string_representation", False + ) + self.ar_loss_weight = train_config.get("ar_loss_weight", 1.0) + + self.dof_config = train_config["dof_config"] + self.agent_pos_config = train_config["agent_pos_config"] + + self.action_horizon_flow = train_config["data"].get("action_horizon_flow", 32) + + if train_config.get("_attn_implementation", None) is not None: + self._attn_implementation = train_config["_attn_implementation"] + + if train_config.get("attn_deterministic", None) is not None: + self.attn_deterministic = train_config["attn_deterministic"] + self.vision_config.attn_deterministic = train_config["attn_deterministic"] + logger.debug("Attention is using deterministic kernel for this run") + else: + self.attn_deterministic = True + self.vision_config.attn_deterministic = True + + if train_config.get("noise_scheduler", None) is not None: + self.noise_scheduler = train_config["noise_scheduler"] + __all__ = ["Qwen2_5_VLConfig"] diff --git a/wall_x/model/qact/qwen2_5/inference_mixin.py b/wall_x/model/qact/qwen2_5/inference_mixin.py new file mode 100644 index 0000000..2325a88 --- /dev/null +++ b/wall_x/model/qact/qwen2_5/inference_mixin.py @@ -0,0 +1,1039 @@ +import logging +from typing import List, Optional, Tuple + +import numpy as np +import torch +from torchdiffeq import odeint + +logger = logging.getLogger(__name__) + + +def topk_right_tie_break_1d(x, k): + L = x.size(0) + x_rev = torch.flip(x, [0]) + idx_in_rev = torch.argsort(x_rev, dim=0, descending=True, stable=True) + orig_idx = (L - 1) - idx_in_rev + topk_idx = orig_idx[:k] + topk_vals = x[topk_idx] + return topk_vals, topk_idx + + +def add_gumbel_noise(logits, temperature): + """ + The Gumbel max is a method for sampling categorical distributions. + According to arXiv:2409.02908, for MDM, low-precision Gumbel Max improves perplexity score but reduces generation quality. + Thus, we use float64. + """ + if temperature == 0: + return logits + logits = logits.to(torch.float64) + noise = torch.rand_like(logits, dtype=torch.float64) + gumbel_noise = (-torch.log(noise)) ** temperature + return logits.exp() / gumbel_noise + + +class VLAInferenceMixin: + + # TODO: Integrate with the optimized implementation. + def prepare_inputs_embeds( + self, + input_ids: torch.LongTensor, + inputs_embeds: Optional[torch.FloatTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + dataset_names: Optional[str] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> Tuple[torch.FloatTensor, Optional[torch.Tensor]]: + """ + Prepare model input embeddings, including text, image, video, and proprioception embeddings. + + Args: + input_ids: Input token IDs. + inputs_embeds: Precomputed input embeddings, if provided. + pixel_values: Image pixel values. + pixel_values_videos: Video pixel values. + image_grid_thw: Image grid time, height, and width. + video_grid_thw: Video grid time, height, and width. + proprioception: Proprioception data. + dataset_names: Dataset names. + agent_pos_mask: Agent position mask. + attention_mask: Attention mask. + + Returns: + inputs_embeds: Complete input embeddings. + attention_mask: Processed attention mask. + """ + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + + # Process image embeddings. + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + + mask = input_ids == self.config.image_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + # Process video embeddings. + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + # Process proprioception embeddings. + if proprioception is not None and not getattr( + self.config, "use_state_string_representation", False + ): + proprioception = proprioception.to(inputs_embeds.device) + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device) + proprio_embed = self.action_preprocessor.proprioception_proj( + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + proprioception_mask = ( + input_ids == self.action_token_id_set["propri_token_id"] + ) + inputs_embeds[proprioception_mask] = proprio_embed.reshape( + -1, inputs_embeds.shape[-1] + ).to(inputs_embeds.dtype) + + # Process the attention mask. + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + return inputs_embeds, attention_mask + + def prepare_position_ids( + self, + input_ids: torch.LongTensor, + inputs_embeds: torch.FloatTensor, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + ) -> torch.LongTensor: + """ + Prepare position IDs, including RoPE delta calculation and caching. + + Args: + input_ids: Input token IDs. + inputs_embeds: Input embeddings. + image_grid_thw: Image grid time, height, and width. + video_grid_thw: Video grid time, height, and width. + second_per_grid_ts: Time step for each grid. + attention_mask: Attention mask. + position_ids: Precomputed position IDs, if provided. + cache_position: Cache position. + past_key_values: Previous key/value cache. + + Returns: + position_ids: Calculated position IDs. + """ + # RoPE deltas cannot be calculated once the attention mask is 4D. + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): + # Calculate RoPE indices once per generation in the prefill stage. + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + ) + self.rope_deltas = rope_deltas + # Reuse the previous RoPE deltas to obtain the correct position IDs. + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + return position_ids + + @torch.no_grad() + def generate_dllm_action( + self, + input_ids, + action_horizon, + action_dim, + ar_action_dim, + total_ar_step, + use_ar_action: bool = False, + num_inference_timesteps: int = 10, + prefix_length: Optional[int] = None, # Prefix length. + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + positional_masks: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + unnorm_proprioception: Optional[torch.FloatTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + robot_type_id: Optional[int] = None, # Used for v3.1 delta decoding. + **kwargs, + ): + batch_size = ( + input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] + ) + inputs_embeds, attention_mask = self.prepare_inputs_embeds( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + proprioception=proprioception, + dataset_names=dataset_names, + agent_pos_mask=agent_pos_mask, + attention_mask=attention_mask, + ) + position_ids = self.prepare_position_ids( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + attention_mask=attention_mask, + position_ids=position_ids, + cache_position=cache_position, + past_key_values=past_key_values, + ) + + noisy_action = torch.randn( + size=(batch_size, action_horizon, action_dim), + dtype=torch.float32, + device=inputs_embeds.device, + ) + + times = torch.linspace( + 0.0, + 1.0, + num_inference_timesteps + 1, + device=inputs_embeds.device, + dtype=torch.float32, + ) + + dt = times[1] - times[0] + time_0 = times[0].unsqueeze(0).repeat(noisy_action.shape[0]) + action_embed, adarms_cond = self.action_preprocessor.step( + timestep=time_0, noisy_action=noisy_action, dof_mask=dof_mask + ) + action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]).to( + inputs_embeds.dtype + ) + flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] + inputs_embeds[flow_action_mask] = action_embed + + # Compute the token span for each expert group after permutation. + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + + prefetch_output = self.model( + input_ids=None, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=None, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, + positional_masks=positional_masks, + use_cache=True, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + adarms_conds=[None, adarms_cond], + start_indices=start_indices, + end_indices=end_indices, + ) + + hidden_states = prefetch_output.last_hidden_state + prefix_kv_cache = prefetch_output.past_key_values + + action_hidden_states = hidden_states[flow_action_mask].to(torch.float32) + v_0 = self.action_preprocessor.action_proj_back( + action_hidden_states[:, : self.action_preprocessor.action_hidden_size] + ) + noisy_action = noisy_action + dt * v_0.reshape( + batch_size, action_horizon, action_dim + ) + + ar_prefix_length = positional_masks["ar_action_mask"].nonzero(as_tuple=True)[ + -1 + ][ + 0 + ] # Attention, not tested yet with batch + start_indices[1] -= ar_prefix_length.item() + end_indices -= ar_prefix_length.item() + if hasattr(prefix_kv_cache, "key_cache"): + for layer_i in range(len(prefix_kv_cache.key_cache)): + prefix_kv_cache.key_cache[layer_i] = prefix_kv_cache.key_cache[layer_i][ + :, :, :ar_prefix_length, : + ] + prefix_kv_cache.value_cache[layer_i] = prefix_kv_cache.value_cache[ + layer_i + ][:, :, :ar_prefix_length, :] + else: + for layer_i in range(len(prefix_kv_cache.layers)): + prefix_kv_cache.layers[layer_i].keys = prefix_kv_cache.layers[ + layer_i + ].keys[:, :, :ar_prefix_length, :] + prefix_kv_cache.layers[layer_i].values = prefix_kv_cache.layers[ + layer_i + ].values[:, :, :ar_prefix_length, :] + + ar_postfix_position_ids = position_ids[:, :, ar_prefix_length:] + ar_postfix_inputs_embeds = inputs_embeds[:, ar_prefix_length:, :] + ar_postfix_attention_mask = attention_mask[:, ar_prefix_length:] + ar_postfix_moe_token_types = moe_token_types[:, ar_prefix_length:] + ar_postfix_input_ids = input_ids[:, ar_prefix_length:] + + noisy_steps_per_sample = ( + torch.ceil((1.0 - times[0] - dt) * (total_ar_step + 1)).long() - 1 + ) + noisy_steps_per_sample = noisy_steps_per_sample.clamp(min=0, max=total_ar_step) + ar_step = total_ar_step - noisy_steps_per_sample + ready_step = [False] * total_ar_step + # Always initialize ar_tokens to record generated discrete codes. + # Use a tensor to support multi-token steps in v3.1 delta mode. + ar_tokens = torch.full( + (ar_postfix_inputs_embeds.shape[1],), + -1, + dtype=torch.long, + device=inputs_embeds.device, + ) + ar_postfix_inputs_embeds, ready_step, ar_tokens = self.update_ar_action( + ar_postfix_inputs_embeds, + hidden_states[:, ar_prefix_length:], + positional_masks["ar_action_mask"][:, ar_prefix_length:], + ar_step=ar_step, + ready_step=ready_step, + ar_tokens=ar_tokens, + ) + + pad_token_id = self.processor.tokenizer.pad_token_id + padding_mask = input_ids == pad_token_id + + ar_postfix_length = input_ids.shape[-1] - ar_prefix_length + _ar_postfix_attention_mask = torch.ones( + (batch_size, ar_postfix_length, ar_prefix_length + ar_postfix_length), + dtype=torch.bool, + device=ar_postfix_attention_mask.device, + ) + + ar_postfix_padding_mask = padding_mask[ + :, ar_prefix_length: + ] # [batch_size, postfix_length] + full_padding_mask = padding_mask # [batch_size, prefix_length + postfix_length] + + for batch_idx in range(padding_mask.shape[0]): + # Disable rows for padded query positions. + _ar_postfix_attention_mask[ + batch_idx, ar_postfix_padding_mask[batch_idx], : + ] = False + # Disable columns for padded key positions. + _ar_postfix_attention_mask[batch_idx, :, full_padding_mask[batch_idx]] = ( + False + ) + + # The 3D attention mask already filters padding, but kv-cache generation + # differs from training, so rebuild the mask before passing it to the model. + if not positional_masks.get("ar_visible", True): + ar_action_mask = positional_masks["ar_action_mask"] != 0 + flow_positions = moe_token_types == 1 + flow_flow_mask = flow_positions[:, :, None] & flow_positions[:, None, :] + ar_ar_mask = ar_action_mask[:, :, None] & ar_action_mask[:, None, :] + flow_ar_mask = flow_flow_mask | ar_ar_mask + kvcache_flow_ar_mask = flow_ar_mask[:, ar_prefix_length:] + + affected = ar_action_mask | flow_positions # (B, N) + affected_pair = affected[:, :, None] & affected[:, None, :] + affected_pair = affected_pair[:, ar_prefix_length:] + + _ar_postfix_attention_mask = torch.where( + affected_pair, kvcache_flow_ar_mask, _ar_postfix_attention_mask + ) + + def step_with_kvcache( + timestep, + noisy_action, + ar_postfix_inputs_embeds, + _ar_postfix_attention_mask, + ar_prefix_length, + ready_step, + dt, + ar_tokens, + ): + action_mask = ( + ar_postfix_input_ids == self.action_token_id_set["action_token_id"] + ) + assert action_mask.any(), "No action token found in input_ids" + timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) + action_embed, adarms_cond = self.action_preprocessor.step( + timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask + ) + action_embed = action_embed.reshape(-1, ar_postfix_inputs_embeds.shape[-1]) + + # Clone inputs_embeds locally to keep this step isolated. + temp_inputs_embeds = ar_postfix_inputs_embeds.clone() + temp_inputs_embeds[action_mask] = action_embed.to(temp_inputs_embeds.dtype) + + transformer_outputs = self.model( + input_ids=None, + attention_mask=_ar_postfix_attention_mask, + position_ids=ar_postfix_position_ids, + past_key_values=prefix_kv_cache, + inputs_embeds=temp_inputs_embeds, + moe_token_types=ar_postfix_moe_token_types, + use_cache=False, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + adarms_conds=[None, adarms_cond], + start_indices=start_indices, + end_indices=end_indices, + ) + + hidden_states = transformer_outputs.last_hidden_state + + noisy_steps_per_sample = ( + torch.ceil((1.0 - timestep - dt) * (total_ar_step + 1)).long() - 1 + ) + noisy_steps_per_sample = noisy_steps_per_sample.clamp( + min=0, max=total_ar_step + ) + ar_step = total_ar_step - noisy_steps_per_sample + ar_postfix_inputs_embeds, ready_step, ar_tokens = self.update_ar_action( + ar_postfix_inputs_embeds, + hidden_states, + positional_masks["ar_action_mask"][:, ar_prefix_length:], + ar_step=ar_step, + ready_step=ready_step, + ar_tokens=ar_tokens, + ) + + action_hidden_states = hidden_states[action_mask].to(torch.float32) + v_t = self.action_preprocessor.action_proj_back( + action_hidden_states[:, : self.action_preprocessor.action_hidden_size] + ) + + return v_t.reshape(batch_size, action_horizon, action_dim) + + action_trajectory = odeint( + lambda timestep, noisy_action: step_with_kvcache( + timestep, + noisy_action, + ar_postfix_inputs_embeds, + _ar_postfix_attention_mask, + ar_prefix_length, + ready_step, + dt, + ar_tokens, + ), + noisy_action, + times[1:], + method="euler", + ) + output = {} + if use_ar_action: + # Decode through tokenizer_mixin.decode_action() for both fast and v3.1 delta modes. + # ar_tokens is already a tensor. + ar_tokens_tensor = ar_tokens.unsqueeze(0) # [1, seq_len] + + predict_action, decode_success = self.tokenizer_mixin.decode_action( + output_ids=ar_tokens_tensor, + action_mapper=self.action_mapper, + action_horizon=action_horizon, + action_dim=ar_action_dim, + device=inputs_embeds.device, + proprioception=proprioception, + dof_mask=dof_mask, + robot_type_id=robot_type_id, + ) + + if not decode_success: + logger.warning("Error in DLLM decoding action, predict_action is None") + output["predict_action"] = None + else: + # unnormalize + if isinstance(predict_action, np.ndarray): + predict_action = torch.tensor( + predict_action, device=inputs_embeds.device + ) + elif predict_action.device != inputs_embeds.device: + predict_action = predict_action.to(inputs_embeds.device) + + # Add the batch dimension when needed. + if predict_action.dim() == 2: + predict_action = predict_action.unsqueeze(0) + + # Decide whether dof_mask is needed from the mixin setting. + uses_dof_mask = ( + self.tokenizer_mixin.uses_dof_mask_for_unnorm + if self.tokenizer_mixin is not None + else True + ) + + if uses_dof_mask: + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names, dof_mask + ) + ) + else: + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names, None + ) + ) + output["predict_action"] = predict_action + + else: + predict_action = action_trajectory[-1] + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names + ) + ) + output["predict_action"] = predict_action + + if action_chunk is not None: + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names + ) + ) + return output + + def update_ar_action2( + self, + ar_postfix_inputs_embeds, + ar_hidden_states, + ar_action_mask, + ar_step, + ready_step, + remask=False, + ar_tokens=None, + ): + ready_num = sum(ready_step) + if ready_num >= ar_step and not remask: + return ar_postfix_inputs_embeds, ready_step, ar_tokens + valid_steps = torch.unique( + ar_action_mask[(ar_action_mask != 0) & (ar_action_mask != -1)] + ) + placeholder_seq_embed = self.model.embed_tokens( + torch.Tensor(self.processor.placeholder_seq) + .to(self.model.device) + .to(torch.int) + ) + step_confidences = [] + step_token_ids = [] + for step in valid_steps: + if ready_step[step - 1] and not remask: + step_confidences.append(1.0) + step_token_ids.append(None) + continue + step_mask = ar_action_mask == step + step_hidden_states = ar_hidden_states[step_mask] + logits = self.lm_head(step_hidden_states) + pred_ids = logits.argmax(dim=-1) + step_confidence = torch.softmax(logits, dim=-1) + total_step_confidence = step_confidence.max(dim=-1).values.mean() + step_confidences.append(total_step_confidence) + step_token_ids.append(pred_ids) + step_confidences = torch.tensor(step_confidences).to(self.model.device) + _, top_k_indices = topk_right_tie_break_1d(step_confidences, ar_step) + for i, step in enumerate(valid_steps): + if ready_step[step - 1] and not remask: + continue + step_mask = ar_action_mask == step + step_indices = step_mask.nonzero(as_tuple=True)[1] + if i in top_k_indices: + # update + ar_postfix_inputs_embeds[0, step_indices, :] = self.model.embed_tokens( + step_token_ids[i] + ) + if ar_tokens is not None: + ar_tokens[step_indices] = step_token_ids[i] + ready_step[step - 1] = True + elif remask and ready_step[step - 1]: + # back to placeholder + ar_postfix_inputs_embeds[0, step_indices, :] = placeholder_seq_embed[ + : len(step_indices) + ] + ready_step[step - 1] = False + + return ar_postfix_inputs_embeds, ready_step, ar_tokens + + def update_ar_action( + self, + ar_postfix_inputs_embeds, + ar_hidden_states, + ar_action_mask, + ar_step, + ready_step, + remask=False, + ar_tokens=None, + ): + ready_num = sum(ready_step) + if ready_num >= ar_step and not remask: + return ar_postfix_inputs_embeds, ready_step, ar_tokens + update_step = ar_step - ready_num + valid_steps = torch.unique( + ar_action_mask[(ar_action_mask != 0) & (ar_action_mask != -1)] + ) + placeholder_seq_embed = self.model.embed_tokens( + torch.Tensor(self.processor.placeholder_seq) + .to(self.model.device) + .to(torch.int) + ) + step_confidences = [] + step_token_ids = [] + + logits = self.lm_head(ar_hidden_states) + confidence = torch.softmax(logits, dim=-1) + pred_ids = logits.argmax(dim=-1) + pad_mask = pred_ids == 151668 # pad token id + step_is_ready = [] + step_is_pad = [] + for step in valid_steps: + step_mask = ar_action_mask == step + step_pred_ids = pred_ids[step_mask] + step_token_ids.append(step_pred_ids) + + is_pad = pad_mask[step_mask].all() + step_is_pad.append(is_pad) + + is_ready = ready_step[step - 1] + step_is_ready.append(is_ready) + + step_confidence = confidence[step_mask] + total_step_confidence = step_confidence.max(dim=-1).values.mean() + step_confidences.append(total_step_confidence) + + step_confidences = torch.tensor(step_confidences).to(self.model.device) + step_is_ready = torch.tensor(step_is_ready, device=self.model.device) + step_is_pad = torch.tensor(step_is_pad, device=self.model.device) + + non_ready_mask = ~step_is_ready + nr_steps = valid_steps[non_ready_mask] # 1-based step ID + nr_conf = step_confidences[non_ready_mask] + nr_pad = step_is_pad[non_ready_mask] + + if nr_steps.numel() == 0: + return ar_postfix_inputs_embeds, ready_step, ar_tokens + + chosen_steps = [] + used_mask = torch.zeros_like(nr_steps, dtype=torch.bool) + + # Step 1: choose one non-PAD step by confidence when available. + non_pad_mask = ~nr_pad + if non_pad_mask.any(): + non_pad_conf = nr_conf[non_pad_mask] + non_pad_steps = nr_steps[non_pad_mask] + + best_idx = torch.argmax(non_pad_conf) # choose highest confidence + chosen_step = non_pad_steps[best_idx] + chosen_steps.append(chosen_step.item()) + + # Mark this step as used. + used_mask[nr_steps == chosen_step] = True + + # If enough steps were selected, reuse chosen_steps below. + if len(chosen_steps) >= update_step: + chosen_steps = torch.tensor(chosen_steps, device=self.model.device) + # chosen_steps will be used below. + else: + # ============================================================ + # Step 2: choose PAD steps by descending index, ignoring confidence. + # ============================================================ + remaining = update_step - len(chosen_steps) + + pad_mask_only = nr_pad & (~used_mask) + pad_steps_only = nr_steps[pad_mask_only] + + if pad_steps_only.numel() > 0: + # Sort from right to left. + pad_sorted = torch.argsort(pad_steps_only, descending=True) + pick = pad_steps_only[pad_sorted[:remaining]] + + chosen_steps.extend(pick.tolist()) + used_mask[(nr_steps.unsqueeze(1) == pick).any(dim=-1)] = True + + # Continue to Step 3 if more steps are still needed. + if len(chosen_steps) < update_step: + remaining = update_step - len(chosen_steps) + + # ============================================================ + # Step 3: choose remaining non-PAD steps by confidence. + # ============================================================ + non_pad_left_mask = (~nr_pad) & (~used_mask) + + if non_pad_left_mask.any(): + left_conf = nr_conf[non_pad_left_mask] + left_steps = nr_steps[non_pad_left_mask] + + sorted_idx = torch.argsort(left_conf, descending=True) + pick = left_steps[sorted_idx[:remaining]] + + chosen_steps.extend(pick.tolist()) + used_mask[(nr_steps.unsqueeze(1) == pick).any(dim=-1)] = True + + chosen_steps = torch.tensor(chosen_steps, device=self.model.device) + for i, step in enumerate(valid_steps): + if ready_step[step - 1] and not remask: + continue + step_mask = ar_action_mask == step + step_indices = step_mask.nonzero(as_tuple=True)[1] + if step in chosen_steps: + # update + ar_postfix_inputs_embeds[0, step_indices, :] = self.model.embed_tokens( + step_token_ids[i] + ) + if ar_tokens is not None: + ar_tokens[step_indices] = step_token_ids[i] + ready_step[step - 1] = True + elif remask and ready_step[step - 1]: + # back to placeholder + ar_postfix_inputs_embeds[0, step_indices, :] = placeholder_seq_embed[ + : len(step_indices) + ] + ready_step[step - 1] = False + + return ar_postfix_inputs_embeds, ready_step, ar_tokens + + def update_infer_dllm_position_mask(self, model_input): + positional_masks = self.tokenizer_mixin.update_placeholder_mask( + self.processor, + model_input["prefix_length"], + model_input["input_ids"], + ) + model_input["positional_masks"] = positional_masks + model_input["positional_masks"]["ar_visible"] = True + return model_input + + @torch.no_grad() + def generate_flow_action_no_cache( + self, + input_ids, + action_horizon, + action_dim, + num_inference_timesteps: int = 10, + prefix_length: Optional[int] = None, # Prefix length. + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + unnorm_proprioception: Optional[torch.FloatTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + **kwargs, + ): + # assert self.config._attn_implementation == "sdpa", "generate_flow_action only support sdpa attn implementation" + batch_size = ( + input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] + ) + assert ( + batch_size == 1 + ), "generate_flow_action_no_cache only support batch size 1" + + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + + mask = input_ids == self.config.image_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if proprioception is not None and not getattr( + self.config, "use_state_string_representation", False + ): + proprioception = proprioception.to(inputs_embeds.device) + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device) + proprio_embed = self.action_preprocessor.proprioception_proj( + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + proprioception_mask = ( + input_ids == self.action_token_id_set["propri_token_id"] + ) + inputs_embeds[proprioception_mask] = proprio_embed.reshape( + -1, inputs_embeds.shape[-1] + ).to(inputs_embeds.dtype) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): + # calculate RoPE index once per generation in the pre-fill stage only + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + if start_indices is None or end_indices is None: + # Compute the token span for each expert group after permutation. + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + + if action_chunk is not None: + action_chunk = action_chunk.to(inputs_embeds.device).to(torch.float32) + + output = {} + + noise = torch.randn( + size=(batch_size, action_horizon, action_dim), + dtype=torch.float32, + device=inputs_embeds.device, + ) + noisy_action = noise.clone() + dof_mask = dof_mask.to(inputs_embeds.device).to(torch.float32) + + times = self.action_preprocessor.get_inference_times( + num_inference_timesteps, inputs_embeds.device, torch.float32 + ) + flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] + + if not dof_mask.all(): + padding_action = ( + torch.zeros((1, dof_mask.shape[-1])) + .to(dof_mask.device) + .to(torch.float32) + ) + padding_action = self.action_preprocessor.normalizer_action.normalize_data( + padding_action, dataset_names + ) + v_padding = padding_action - noisy_action + + def step(timestep, noisy_action): + timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) + action_embed, adarms_cond = self.action_preprocessor.step( + timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask + ) + action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]).to( + inputs_embeds.dtype + ) + + inputs_embeds[flow_action_mask] = action_embed + model_output = self.model( + input_ids=None, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=None, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, + start_indices=start_indices, + end_indices=end_indices, + positional_masks=positional_masks, + use_cache=True, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + adarms_conds=[None, adarms_cond], + ) + + hidden_states = model_output.last_hidden_state + action_hidden_states = hidden_states[flow_action_mask].to(torch.float32) + action_pred = self.action_preprocessor.action_proj_back( + action_hidden_states[:, : self.action_preprocessor.action_hidden_size] + ) + if getattr(self.config, "use_x_pred", False): + v_t = (action_pred - noisy_action) / torch.clamp(1 - timestep, min=0.05) + else: + v_t = action_pred + + if not dof_mask.all(): + v_t = (v_padding) * (1 - dof_mask) + v_t * dof_mask + + return v_t.reshape(batch_size, action_horizon, action_dim) + + action_trajectory = odeint(step, noisy_action, times, method="euler") + + predict_action = action_trajectory[-1] + predict_action = self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names + ) + output["predict_action"] = predict_action + # normalize action chunk to get gt_action + if action_chunk is not None: + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names + ) + ) + + return output diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py similarity index 82% rename from wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py rename to wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py index 951bed9..9eff1fe 100644 --- a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py +++ b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py @@ -1,18 +1,51 @@ +# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! +# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_qwen2_5_vl.py file directly. One of our CI enforces this. +# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! +# coding=utf-8 +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import math +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + import torch import torch.nn as nn import torch.nn.functional as F -from dataclasses import dataclass +import torch.utils.checkpoint as cp from torch.nn import CrossEntropyLoss -from typing import Any, Dict, List, Optional, Tuple, Union from transformers.activations import ACT2FN from transformers.cache_utils import ( Cache, DynamicCache, - SlidingWindowCache, StaticCache, ) + +try: + from transformers.cache_utils import SlidingWindowCache +except ImportError: + # Compatibility with newer transformers versions + SlidingWindowCache = StaticCache from transformers.generation import GenerationMixin from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput @@ -27,17 +60,20 @@ from transformers.utils import ( logging, replace_return_docstrings, ) + + from .configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig -from wall_x.fusions import ops +from wall_x.model.core.attention.selector import AttentionsSelectorMixin +from wall_x.model.core.ops import rot_pos_emb, get_window_index, m_rope if is_flash_attn_2_available(): + from flash_attn import flash_attn_func from flash_attn import flash_attn_varlen_func from flash_attn.layers.rotary import apply_rotary_emb - from flash_attn import flash_attn_func + else: flash_attn_varlen_func = None apply_rotary_emb = None - flash_attn_func = None if is_flash_attn_2_available(): @@ -46,26 +82,142 @@ else: flash_attn_varlen_func = None +try: + from flash_mask.flash_mask_interface import flash_mask_attn_func +except ImportError: + flash_mask_attn_func = None + + logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = "Qwen2_5_VLConfig" class Qwen2_5_VLMLP(nn.Module): - def __init__(self, config, bias: bool = False): + # Align intermediate_size to this value for cuBLAS GEMM tile efficiency. + # E.g. 3420 -> 3456 (= 128 * 27), improving MFU from 34.8% to ~60%+. + # Zero-padded weights don't affect forward output or gradient flow. + _GEMM_ALIGN = 128 + + def __init__(self, config, bias=False, use_selective_recompute=False): super().__init__() self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias) + + # Pad intermediate_size to next multiple of _GEMM_ALIGN + align = self._GEMM_ALIGN + self.padded_intermediate_size = ( + (self.intermediate_size + align - 1) // align * align + ) + + self.gate_up_proj = nn.Linear( + self.hidden_size, 2 * self.padded_intermediate_size, bias=bias + ) + self.down_proj = nn.Linear( + self.padded_intermediate_size, self.hidden_size, bias=bias + ) + self.act_fn = ACT2FN[config.hidden_act] + self.use_selective_recompute = use_selective_recompute + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + """Pad checkpoint weights from original intermediate_size to padded size.""" + pad = self.padded_intermediate_size - self.intermediate_size + if pad > 0: + orig = self.intermediate_size + # gate_up_proj.weight: [2*orig, hidden] -> [2*padded, hidden] + key_w = prefix + "gate_up_proj.weight" + if key_w in state_dict and state_dict[key_w].shape[0] == 2 * orig: + w = state_dict[key_w] + gate_w, up_w = w.split([orig, orig], dim=0) + gate_w = F.pad(gate_w, (0, 0, 0, pad)) + up_w = F.pad(up_w, (0, 0, 0, pad)) + state_dict[key_w] = torch.cat([gate_w, up_w], dim=0) + # gate_up_proj.bias: [2*orig] -> [2*padded] + key_b = prefix + "gate_up_proj.bias" + if key_b in state_dict and state_dict[key_b].shape[0] == 2 * orig: + b = state_dict[key_b] + gate_b, up_b = b.split([orig, orig]) + state_dict[key_b] = torch.cat( + [F.pad(gate_b, (0, pad)), F.pad(up_b, (0, pad))] + ) + # down_proj.weight: [hidden, orig] -> [hidden, padded] + key_dw = prefix + "down_proj.weight" + if key_dw in state_dict and state_dict[key_dw].shape[1] == orig: + state_dict[key_dw] = F.pad(state_dict[key_dw], (0, pad)) + # down_proj.bias: [hidden] - no change needed + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def _save_to_state_dict(self, destination, prefix, keep_vars): + """Strip padding so saved checkpoints use the original intermediate_size.""" + super()._save_to_state_dict(destination, prefix, keep_vars) + pad = self.padded_intermediate_size - self.intermediate_size + if pad == 0: + return + orig = self.intermediate_size + padded = self.padded_intermediate_size + + # gate_up_proj.weight: [2*padded, hidden] -> [2*orig, hidden] + key_w = prefix + "gate_up_proj.weight" + if key_w in destination: + w = destination[key_w] + gate_w, up_w = w.split([padded, padded], dim=0) + destination[key_w] = torch.cat([gate_w[:orig], up_w[:orig]], dim=0) + + # gate_up_proj.bias: [2*padded] -> [2*orig] + key_b = prefix + "gate_up_proj.bias" + if key_b in destination: + b = destination[key_b] + gate_b, up_b = b.split([padded, padded]) + destination[key_b] = torch.cat([gate_b[:orig], up_b[:orig]]) + + # down_proj.weight: [hidden, padded] -> [hidden, orig] + key_dw = prefix + "down_proj.weight" + if key_dw in destination: + destination[key_dw] = destination[key_dw][:, :orig] + + def _activation_chunk(self, gate_out, up_out): + # gate_out / up_out are precomputed GEMM results + act = self.act_fn(gate_out) + return act * up_out def forward(self, hidden_state): - return self.down_proj( - self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state) + # ----------------------------- + # 1. Precompute GEMM without checkpointing + # ----------------------------- + gate_up_out = self.gate_up_proj(hidden_state) + gate_out, up_out = gate_up_out.split( + [self.padded_intermediate_size, self.padded_intermediate_size], dim=-1 ) + if self.use_selective_recompute: + # checkpoint only activation, not GEMM + out = cp.checkpoint( + self._activation_chunk, gate_out, up_out, use_reentrant=False + ) + else: + out = self._activation_chunk(gate_out, up_out) + + return self.down_proj(out) + class Qwen2_5_VisionPatchEmbed(nn.Module): def __init__( @@ -99,9 +251,18 @@ class Qwen2_5_VisionPatchEmbed(nn.Module): self.patch_size, self.patch_size, ) - hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view( - -1, self.embed_dim + # hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view( + # -1, self.embed_dim + # ) + weight = self.proj.weight.view(self.embed_dim, -1) + out = hidden_states.view( + -1, + self.in_channels + * self.temporal_patch_size + * self.patch_size + * self.patch_size, ) + hidden_states = F.linear(out.to(target_dtype), weight) return hidden_states @@ -196,7 +357,8 @@ class Qwen2_5_VLPatchMerger(nn.Module): ) def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.mlp(self.ln_q(x)[0].view(-1, self.hidden_size)) + normed_x, _ = self.ln_q(x) + x = self.mlp(normed_x.view(-1, self.hidden_size)) return x @@ -211,11 +373,12 @@ def apply_rotary_pos_emb_flashatt( class Qwen2_5_VLVisionFlashAttention2(nn.Module): - def __init__(self, dim: int, num_heads: int = 16) -> None: + def __init__(self, config: Qwen2_5_VLConfig, dim: int, num_heads: int = 16) -> None: super().__init__() self.num_heads = num_heads self.qkv = nn.Linear(dim, dim * 3, bias=True) self.proj = nn.Linear(dim, dim) + self.deterministic = config.attn_deterministic def forward( self, @@ -251,7 +414,14 @@ class Qwen2_5_VLVisionFlashAttention2(nn.Module): if max_seqlen is None: max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() attn_output = flash_attn_varlen_func( - q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen + q, + k, + v, + cu_seqlens, + cu_seqlens, + max_seqlen, + max_seqlen, + deterministic=self.deterministic, ).reshape(seq_length, -1) attn_output = self.proj(attn_output) return attn_output @@ -279,7 +449,7 @@ def apply_rotary_pos_emb_vision( class Qwen2_5_VLVisionAttention(nn.Module): - def __init__(self, dim: int, num_heads: int = 16) -> None: + def __init__(self, config: Qwen2_5_VLConfig, dim: int, num_heads: int = 16) -> None: super().__init__() self.num_heads = num_heads self.head_dim = dim // num_heads @@ -290,9 +460,9 @@ class Qwen2_5_VLVisionAttention(nn.Module): self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, - max_seqlen: Optional[int] = None, rotary_pos_emb: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, ) -> torch.Tensor: seq_length = hidden_states.shape[0] q, k, v = ( @@ -344,7 +514,7 @@ class Qwen2_5_VLVisionAttention(nn.Module): class Qwen2_5_VLVisionSdpaAttention(nn.Module): - def __init__(self, dim: int, num_heads: int = 16) -> None: + def __init__(self, config: Qwen2_5_VLConfig, dim: int, num_heads: int = 16) -> None: super().__init__() self.num_heads = num_heads self.qkv = nn.Linear(dim, dim * 3, bias=True) @@ -354,7 +524,6 @@ class Qwen2_5_VLVisionSdpaAttention(nn.Module): self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, - max_seqlen: Optional[int] = None, rotary_pos_emb: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: @@ -408,14 +577,33 @@ QWEN2_5_VL_VISION_ATTENTION_CLASSES = { class Qwen2_5_VLVisionBlock(nn.Module): - def __init__(self, config, attn_implementation: str = "sdpa") -> None: + def __init__( + self, + config, + attn_implementation: str = "sdpa", + use_selective_recompute: bool = False, + ): super().__init__() + self.use_selective_recompute = use_selective_recompute + self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) + self.attn = QWEN2_5_VL_VISION_ATTENTION_CLASSES[attn_implementation]( - config.hidden_size, num_heads=config.num_heads + config, config.hidden_size, num_heads=config.num_heads ) - self.mlp = Qwen2_5_VLMLP(config, bias=True) + self.mlp = Qwen2_5_VLMLP( + config, bias=True, use_selective_recompute=use_selective_recompute + ) + + # ----------------------------- + # selective checkpoint for norm + # ----------------------------- + def _norm1_chunk(self, hidden_states): + return self.norm1(hidden_states)[0] # return only normed_hidden_states + + def _norm2_chunk(self, hidden_states): + return self.norm2(hidden_states)[0] def forward( self, @@ -425,14 +613,46 @@ class Qwen2_5_VLVisionBlock(nn.Module): rotary_pos_emb: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: + + # ============================ + # 1) Norm1 with checkpoint + # ============================ + if self.use_selective_recompute: + normed_hidden_states = cp.checkpoint( + self._norm1_chunk, + hidden_states, + use_reentrant=False, + ) + # residual RMS is not used by subsequent computation and is usually unused in large-model training + _ = None + else: + normed_hidden_states, _ = self.norm1(hidden_states) + + # Attention does not recompute GEMM hidden_states = hidden_states + self.attn( - self.norm1(hidden_states)[0], + normed_hidden_states, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, rotary_pos_emb=rotary_pos_emb, position_embeddings=position_embeddings, ) - hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)[0]) + + # ============================ + # 2) Norm2 with checkpoint + # ============================ + if self.use_selective_recompute: + normed_hidden_states = cp.checkpoint( + self._norm2_chunk, + hidden_states, + use_reentrant=False, + ) + _ = None + else: + normed_hidden_states, _ = self.norm2(hidden_states) + + # MLP only recomputes the hooked activation and does not recompute GEMM + hidden_states = hidden_states + self.mlp(normed_hidden_states) + return hidden_states @@ -457,7 +677,7 @@ Qwen2_5_VL_START_DOCSTRING = r""" "The bare Qwen2_5_VL Model outputting raw hidden-states without any specific head on top.", Qwen2_5_VL_START_DOCSTRING, ) -class Qwen2_5_VLPreTrainedModel(PreTrainedModel): +class Qwen2_5_VLPreTrainedModel(AttentionsSelectorMixin, PreTrainedModel): config_class = Qwen2_5_VLConfig base_model_prefix = "model" supports_gradient_checkpointing = True @@ -484,8 +704,11 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): config_class = Qwen2_5_VLVisionConfig _no_split_modules = ["Qwen2_5_VLVisionBlock"] - def __init__(self, config, *inputs, **kwargs) -> None: + def __init__( + self, config, use_selective_recompute=False, *inputs, **kwargs + ) -> None: super().__init__(config, *inputs, **kwargs) + self.use_selective_recompute = use_selective_recompute self.spatial_merge_size = config.spatial_merge_size self.patch_size = config.patch_size self.fullatt_block_indexes = config.fullatt_block_indexes @@ -504,7 +727,9 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): self.blocks = nn.ModuleList( [ - Qwen2_5_VLVisionBlock(config, config._attn_implementation) + Qwen2_5_VLVisionBlock( + config, config._attn_implementation, use_selective_recompute + ) for _ in range(config.depth) ] ) @@ -605,12 +830,13 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): `torch.Tensor`: hidden_states. """ hidden_states = self.patch_embed(hidden_states) - rotary_pos_emb = ops.rot_pos_emb( - self.rotary_pos_emb.inv_freq, grid_thw, self.spatial_merge_size + rotary_pos_emb = rot_pos_emb( + self.rotary_pos_emb.inv_freq.to(torch.float32), + grid_thw, + self.spatial_merge_size, ) - - window_index, cu_window_seqlens = ops.get_window_index( - grid_thw=grid_thw, + window_index, cu_window_seqlens = get_window_index( + grid_thw=grid_thw.to(torch.int32), window_size=self.window_size, spatial_merge_size=self.spatial_merge_size, patch_size=self.patch_size, @@ -658,9 +884,9 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): hidden_states = self._gradient_checkpointing_func( blk.__call__, hidden_states, - cu_seqlens_now, - None, - position_embeddings, + cu_seqlens=cu_seqlens_now, + max_seqlen=max_seqlen_now, + position_embeddings=position_embeddings, ) else: hidden_states = blk( @@ -683,7 +909,7 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module): # BC: "rope_type" was originally "type" if hasattr(config, "rope_scaling") and config.rope_scaling is not None: self.rope_type = config.rope_scaling.get( - "rope_type", config.rope_scaling.get("type") + "rope_type", config.rope_scaling.get("type", "default") ) else: self.rope_type = "default" @@ -691,12 +917,43 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module): self.original_max_seq_len = config.max_position_embeddings self.config = config - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + # Compatibility with newer transformers versions (5.x): + # Newer versions remove the "default" key from ROPE_INIT_FUNCTIONS and compute default RoPE inline. + # Older versions (4.x): ROPE_INIT_FUNCTIONS["default"] exists and can be used normally. + # Newer versions (5.x): when rope_type=="default", use compute_default_rope_parameters instead. + if self.rope_type == "default": + self.rope_init_fn = self.compute_default_rope_parameters + else: + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq + @staticmethod + def compute_default_rope_parameters( + config: Qwen2_5_VLConfig, + device=None, + seq_len=None, + ): + """Compute default RoPE parameters, matching older ROPE_INIT_FUNCTIONS["default"]. + Transformers 5.x moved this logic out of ROPE_INIT_FUNCTIONS and into the model class. + """ + base = config.rope_theta + dim = config.hidden_size // config.num_attention_heads + attention_factor = 1.0 + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, dim, 2, dtype=torch.int64).to( + device=device, dtype=torch.float + ) + / dim + ) + ) + return inv_freq, attention_factor + def _dynamic_frequency_update(self, position_ids, device): """ dynamic RoPE layers should recompute `inv_freq` in the following situations: @@ -763,13 +1020,18 @@ class Qwen2MLP(nn.Module): self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.gate_up_proj = nn.Linear( + self.hidden_size, 2 * self.intermediate_size, bias=False + ) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + gate_up_out = self.gate_up_proj(x) + gate_out, up_out = gate_up_out.split( + [self.intermediate_size, self.intermediate_size], dim=-1 + ) + down_proj = self.down_proj(self.act_fn(gate_out) * up_out) return down_proj @@ -863,15 +1125,13 @@ class Qwen2_5_VLAttention(nn.Module): f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" f" and `num_heads`: {self.num_heads})." ) - self.q_proj = nn.Linear( - self.hidden_size, self.num_heads * self.head_dim, bias=True - ) - self.k_proj = nn.Linear( - self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True - ) - self.v_proj = nn.Linear( - self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True + + qkv_dim = ( + self.num_heads * self.head_dim + + 2 * self.num_key_value_heads * self.head_dim ) + self.qkv_proj = nn.Linear(self.hidden_size, qkv_dim, bias=True) + self.o_proj = nn.Linear( self.num_heads * self.head_dim, self.hidden_size, bias=False ) @@ -890,16 +1150,22 @@ class Qwen2_5_VLAttention(nn.Module): position_embeddings: Optional[ Tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC + **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + qkv = self.qkv_proj(hidden_states) # [bsz, q_len, q_dim + 2 * kv_dim] + kv_dim = self.num_key_value_heads * self.head_dim + q, k, v = torch.split( + qkv, [self.num_heads * self.head_dim, kv_dim, kv_dim], dim=-1 + ) + query_states = q.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = k.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = v.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_multimodal_rotary_pos_emb( @@ -970,13 +1236,14 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): config.max_window_layers layers. """ - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): + super().__init__(config, layer_idx) # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + self.deterministic = config.attn_deterministic def forward( self, @@ -993,19 +1260,25 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): ): bsz, q_len, _ = hidden_states.size() - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) + qkv = self.qkv_proj(hidden_states) # [bsz, q_len, q_dim + 2 * kv_dim] + kv_dim = self.num_key_value_heads * self.head_dim + q, k, v = torch.split( + qkv, [self.num_heads * self.head_dim, kv_dim, kv_dim], dim=-1 + ) - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + query_states = q.view(bsz, q_len, self.num_heads, self.head_dim) + key_states = k.view(bsz, q_len, self.num_key_value_heads, self.head_dim) + value_states = v.view(bsz, q_len, self.num_key_value_heads, self.head_dim) # Because the input can be padded, the absolute sequence length depends on the max position id. cos, sin = position_embeddings - query_states, key_states = ops.multimodal_rope( - query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] + query_states, key_states = m_rope( + query_states.contiguous(), + key_states.contiguous(), + cos[..., : (cos.size(-1) // 2)].contiguous().float(), + sin[..., : (sin.size(-1) // 2)].contiguous().float(), + self.rope_scaling["mrope_section"], ) if past_key_value is not None: cache_kwargs = { @@ -1014,8 +1287,14 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): "cache_position": cache_position, } # Specific to RoPE models key_states, value_states = past_key_value.update( - key_states, value_states, self.layer_idx, cache_kwargs + key_states.transpose(1, 2), + value_states.transpose(1, 2), + self.layer_idx, + cache_kwargs, ) + key_states, value_states = key_states.transpose( + 1, 2 + ), value_states.transpose(1, 2) # repeat k/v heads if n_kv_heads < n_heads # key_states = repeat_kv(key_states, self.num_key_value_groups) @@ -1033,7 +1312,7 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): elif hasattr(self.config, "_pre_quantization_dtype"): target_dtype = self.config._pre_quantization_dtype else: - target_dtype = self.q_proj.weight.dtype + target_dtype = self.qkv_proj.weight.dtype logger.warning_once( f"The input hidden states seems to be silently casted in float32, this might be related to" @@ -1045,11 +1324,6 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): key_states = key_states.to(target_dtype) value_states = value_states.to(target_dtype) - # Reashape to the expected shape for Flash Attention - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - value_states = value_states.transpose(1, 2) - attn_output = flash_attn_func( query_states, key_states, @@ -1057,6 +1331,7 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): dropout_rate, softmax_scale=None, causal=self.is_causal, + deterministic=self.deterministic, ) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() @@ -1064,6 +1339,39 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): if not output_attentions: attn_weights = None + if output_attentions: + with torch.no_grad(): + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + key_states = repeat_kv(key_states, self.num_key_value_groups) + scale = 1.0 / torch.sqrt( + torch.tensor( + self.head_dim, device=hidden_states.device, dtype=torch.float32 + ) + ) + attention_score = ( + torch.matmul(query_states, key_states.transpose(-2, -1)) * scale + ) + causal_mask_local = torch.tril( + torch.ones( + attention_score.size(-2), + attention_score.size(-1), + device=attention_score.device, + dtype=torch.bool, + ) + ) + mask = ( + causal_mask_local.unsqueeze(0) + .unsqueeze(0) + .expand(attention_score.size(0), attention_score.size(1), -1, -1) + ) + if mask.dtype != attention_score.dtype: + mask = mask.to(dtype=attention_score.dtype) + attention_score = attention_score.masked_fill( + ~(mask.bool()), float("-inf") + ) + attention_score = torch.softmax(attention_score, dim=-1) + attn_weights = attention_score[0].mean(0) return attn_output, attn_weights, past_key_value @@ -1108,13 +1416,18 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): bsz, q_len, _ = hidden_states.size() - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + qkv = self.qkv_proj(hidden_states) # [bsz, q_len, q_dim + 2 * kv_dim] + kv_dim = self.num_key_value_heads * self.head_dim + q, k, v = torch.split( + qkv, [self.num_heads * self.head_dim, kv_dim, kv_dim], dim=-1 + ) + query_states = q.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = k.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = v.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_multimodal_rotary_pos_emb( @@ -1132,7 +1445,18 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): key_states, value_states, self.layer_idx, cache_kwargs ) else: - past_key_states, past_value_states = past_key_value[self.layer_idx] + # Compatible across transformers versions: + # v5.x: DynamicCache uses .layers[idx].keys/.values + # v4.x: DynamicCache uses .key_cache[idx]/.value_cache[idx] + # old: Cache object is subscriptable, returns (key, value) tuple + if hasattr(past_key_value, "layers"): + past_key_states = past_key_value.layers[self.layer_idx].keys + past_value_states = past_key_value.layers[self.layer_idx].values + elif hasattr(past_key_value, "key_cache"): + past_key_states = past_key_value.key_cache[self.layer_idx] + past_value_states = past_key_value.value_cache[self.layer_idx] + else: + past_key_states, past_value_states = past_key_value[self.layer_idx] key_states = torch.cat([past_key_states, key_states], dim=-2) value_states = torch.cat([past_value_states, value_states], dim=-2) @@ -1141,15 +1465,15 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): causal_mask = attention_mask if attention_mask is not None: # no matter the length, we just slice it - # Ensure the attention_mask correctly matches the head dimension + # Ensure attention_mask matches the head dimension if len(attention_mask.shape) == 2: # [batch_size, seq_len] - # Expand to [batch_size, 1, seq_len, seq_len] causal mask format + # Expand to causal mask format [batch_size, 1, seq_len, seq_len] bsz, seq_len = attention_mask.shape causal_mask = attention_mask.view(bsz, 1, 1, seq_len).expand( bsz, 1, seq_len, seq_len ) elif len(attention_mask.shape) == 3: # [batch_size, seq_len, seq_len] - # Add head dimension: [batch_size, 1, seq_len, seq_len] + # Add the head dimension: [batch_size, 1, seq_len, seq_len] causal_mask = attention_mask.unsqueeze(1) elif ( len(attention_mask.shape) == 4 @@ -1157,10 +1481,10 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): causal_mask = attention_mask else: raise ValueError( - f"Unsupported attention_mask dim: {attention_mask.shape}" + f"Unsupported attention_mask shape: {attention_mask.shape}" ) - # Convert the attention mask to boolean type + # Convert attention mask to bool causal_mask = causal_mask.to(torch.bool) # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, @@ -1204,10 +1528,132 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): return attn_output, None, past_key_value +class Qwen2_5_VLFlashMaskAttention(Qwen2_5_VLAttention): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC + ): + bsz, q_len, _ = hidden_states.size() + + qkv = self.qkv_proj(hidden_states) # [bsz, q_len, q_dim + 2 * kv_dim] + kv_dim = self.num_key_value_heads * self.head_dim + q, k, v = torch.split( + qkv, [self.num_heads * self.head_dim, kv_dim, kv_dim], dim=-1 + ) + query_states = q.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = k.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = v.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = apply_multimodal_rotary_pos_emb( + query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] + ) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + # dropout_rate = 0.0 if not self.training else self.attention_dropout + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.qkv_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Reashape to the expected shape for Flash Attention + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + # if ( + # self.config.use_sliding_window + # and getattr(self.config, "sliding_window", None) is not None + # and self.layer_idx >= self.config.max_window_layers + # ): + # sliding_window = self.config.sliding_window + # else: + # sliding_window = None + + # Expand the attention_mask head dimension from 1 to num_heads + if attention_mask is not None and attention_mask.shape[1] == 1: + # attention_mask: (bsz, 1, q_len, 2) or (bsz, 1, q_len, 4) + attention_mask = attention_mask.expand( + -1, self.num_heads, -1, -1 + ).contiguous() + + if attention_mask is not None: + attention_mask = attention_mask.contiguous() + + query_states_contiguous = query_states.contiguous() + key_states_contiguous = key_states.contiguous() + value_states_contiguous = value_states.contiguous() + + attn_output = flash_mask_attn_func( + query_states_contiguous, + key_states_contiguous, + value_states_contiguous, + startend_row_indices=attention_mask, + causal=False, + ) + + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + QWEN2_5_VL_ATTENTION_CLASSES = { "eager": Qwen2_5_VLAttention, "flash_attention_2": Qwen2_5_VLFlashAttention2, "sdpa": Qwen2_5_VLSdpaAttention, + "flash_mask": Qwen2_5_VLFlashMaskAttention, } @@ -1765,14 +2211,16 @@ QWEN2_5_VL_INPUTS_DOCSTRING = r""" class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin): - _tied_weights_keys = ["lm_head.weight"] + # Compatibility with newer transformers versions (5.x):_tied_weights_keys changed from list[str] to dict[str, str] (target -> source mapping). + # Older versions (4.x) iterate over dict keys, so the behavior is equivalent and compatible. + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} config_class = Qwen2_5_VLConfig _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] - def __init__(self, config): + def __init__(self, config, use_selective_recompute): super().__init__(config) self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config( - config.vision_config + config.vision_config, use_selective_recompute ) self.model = Qwen2_5_VLModel(config) self.vocab_size = config.vocab_size diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py similarity index 60% rename from wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py rename to wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py index 996f879..4a0f991 100644 --- a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py +++ b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py @@ -1,61 +1,80 @@ -import os +import re import torch -import yaml -import numpy as np -import glob import torch.nn as nn -from torchdiffeq import odeint from dataclasses import dataclass from torch.nn import CrossEntropyLoss -from safetensors.torch import load_file -from peft import LoraConfig, get_peft_model -from typing import Optional, List, Tuple, Any, Dict, Union -import time -from transformers import AutoConfig, AutoProcessor -from transformers.utils import logging, is_torchdynamo_compiling -from transformers.cache_utils import ( - Cache, - DynamicCache, - SlidingWindowCache, - StaticCache, -) -from transformers.modeling_attn_mask_utils import AttentionMaskConverter - -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import ( +from typing import Optional, List, Tuple +from torchdiffeq import odeint +import numpy as np +from typing import Any, Dict, Union +from wall_x.model.qact.qwen2_5.configuration_qwen2_5_vl import Qwen2_5_VLConfig +from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl import ( + QWEN2_5_VL_ATTENTION_CLASSES, Qwen2_5_VLMLP, Qwen2_5_VLRotaryEmbedding, Qwen2_5_VLPreTrainedModel, Qwen2RMSNorm, - Qwen2_5_VLForConditionalGeneration, -) -from transformers.modeling_outputs import ( - ModelOutput, - BaseModelOutputWithPast, -) - -from wall_x.fusions import ops -from wall_x.model.action_head import ActionProcessor -from wall_x.model.qwen2_5_based.configuration_qwen2_5_vl import Qwen2_5_VLConfig -from wall_x.model.model_utils import load_wallx_processors, update_model_config -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import ( + Qwen2_5_VLPatchMerger, + Qwen2_5_VLVisionBlock, Qwen2_5_VisionTransformerPretrainedModel, - Qwen2_5_VLAttention, - Qwen2_5_VLFlashAttention2, - Qwen2_5_VLSdpaAttention, ) -from wall_x.model.vla_mixin import ActionGenerationMixin, ActionModelMixMin -from wall_x.model.vla_mixin import TokenTypeRouter, SparseMoeBlock -from wall_x.model.vla_mixin import ( - ATTENTION_TYPES_WITH_2D_MASK, -) -from wall_x.model.joint_attention import JOINT_QWEN_ATTENTION_CLASSES -from wall_x.data.utils import update_action_statistics -from wall_x.utils.constant import action_statistic_dof -from pprint import pprint +from transformers import AutoConfig + +from transformers.cache_utils import ( + Cache, + DynamicCache, + StaticCache, +) + +try: + from transformers.cache_utils import SlidingWindowCache +except ImportError: + # Compatibility with newer transformers versions + SlidingWindowCache = StaticCache + +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + ModelOutput, +) +from transformers.utils import ( + logging, + is_torchdynamo_compiling, +) +from transformers.modeling_attn_mask_utils import ( + AttentionMaskConverter, +) +from wall_x.model.core.vla_mixin import ActionGenerationMixin, ActionModelMixMin +from wall_x.model.core.action.normalizer import ( + normalize_data_with_virtual_tail, + unnormalize_data_with_virtual_tail, +) +from wall_x.model.core.action.moe import TokenTypeRouter, SparseMoeBlock +from wall_x.model.core.attention.selector import ( + ATTENTION_TYPES_WITH_2D_MASK, + ATTENTION_TYPES_WITH_FLASH_MASK, +) +from wall_x.model.core.attention.joint import JOINT_QWEN_ATTENTION_CLASSES +from wall_x.model.qact.qwen2_5.inference_mixin import VLAInferenceMixin + +from wall_x.model.core.ops import unpermute, permute, get_rope_index logger = logging.get_logger(__name__) +_QWEN25_DMUON_BLOCKED_NAME_PARTS = ( + "embed_tokens", + "pos_embed", + "lm_head", + "norm", +) + + +def _is_qwen25_dmuon_target_param(name: str, param: nn.Parameter) -> bool: + """Select trainable Qwen2.5 VLA matrices for DMuon.""" + if not param.requires_grad or param.ndim != 2 or not name.endswith(".weight"): + return False + return not any(part in name for part in _QWEN25_DMUON_BLOCKED_NAME_PARTS) + @dataclass class Qwen2_5_VLACausalLMOutputWithPast(ModelOutput): @@ -72,13 +91,6 @@ class Qwen2_5_VLACausalLMOutputWithPast(ModelOutput): channel_loss_count_dict: Optional[dict[torch.FloatTensor]] = None -QWEN2_5_VL_ATTENTION_CLASSES = { - "eager": Qwen2_5_VLAttention, - "flash_attention_2": Qwen2_5_VLFlashAttention2, - "sdpa": Qwen2_5_VLSdpaAttention, -} - - class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module, ActionModelMixMin): def __init__( self, @@ -203,9 +215,9 @@ class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module, ActionModelMixMin): Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code into the model """ - residual = hidden_states - hidden_states, gate, _ = self._apply_norm_moe( + residual = hidden_states + hidden_states, gate, gate_mask = self._apply_norm_moe( hidden_states, token_types, adarms_conds, @@ -286,7 +298,7 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): def from_pretrained( cls, pretrained_model_name_or_path, num_experts=None, *args, **kwargs ): - # If `num_experts` is provided, ensure it is added to the config. + # If num_experts is provided, make sure it is added to config config = kwargs.get("config", None) if config is None: config = AutoConfig.from_pretrained(pretrained_model_name_or_path) @@ -360,10 +372,12 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - moe_token_types: Optional[torch.LongTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, # new parameter start_indices: Optional[torch.Tensor] = None, end_indices: Optional[torch.Tensor] = None, - positional_masks: Optional[dict] = None, + positional_masks: Optional[ + dict + ] = None, # stores token position masks needed by each category use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, @@ -393,7 +407,7 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): "You must specify exactly one of input_ids or inputs_embeds" ) if moe_token_types is None: - raise ValueError("moe_token_types must be provided for MoE routing.") + raise ValueError("moe_token_types must be provided for MoE routing") if start_indices is None or end_indices is None: raise ValueError( "start_indices and end_indices must be provided for MoE routing" @@ -456,18 +470,16 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) - # If `mot_opt` is enabled, the tokens from different experts will be permuted first, resulting in a dimension of [Tokens, HiddenSize]. + # When mot_opt is enabled, permute tokens from different experts first; shape becomes [Tokens, HiddenSize] orig_shape = hidden_states.shape + moe_token_types = moe_token_types.contiguous() if self.config.mot_opt: hidden_states = hidden_states.view(-1, hidden_states.size(-1)) - hidden_states, row_id_map = ops.permute( - hidden_states, moe_token_types.view(-1) - ) + hidden_states, row_id_map = permute(hidden_states, moe_token_types.view(-1)) else: row_id_map = None - probs = torch.ones_like(moe_token_types.view(-1), dtype=torch.float32).view( - -1, 1 - ) + # Use reshape instead of view to support non-contiguous tensors + probs = torch.ones_like(moe_token_types, dtype=torch.float32).reshape(-1, 1) # decoder layers all_hidden_states = () if output_hidden_states else None @@ -486,6 +498,16 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): positional_masks=positional_masks, ) + if ( + self.config._attn_implementation in ATTENTION_TYPES_WITH_FLASH_MASK + and self.config.attention_moe is True + ): + causal_mask = self._update_joint_attention_flash_mask( + attention_mask=causal_mask, + moe_token_types=moe_token_types, + positional_masks=positional_masks, + ) + for decoder_layer in self.layers: if output_hidden_states: assert ( @@ -538,9 +560,6 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): next_decoder_cache = layer_outputs[2 if output_attentions else 1] if output_attentions: - assert ( - self.config.mot_opt is False - ), "When using mot_opt, output_hidden_states is not supported yet." all_self_attns += (layer_outputs[1],) hidden_states, _, _ = self._apply_norm_moe( @@ -564,7 +583,7 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): next_cache = next_decoder_cache if use_cache else None if self.config.mot_opt: - hidden_states = ops.unpermute(hidden_states, row_id_map, probs) + hidden_states = unpermute(hidden_states, row_id_map, probs) hidden_states = hidden_states.view(orig_shape) if not return_dict: @@ -658,22 +677,22 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): config=self.config, past_key_values=past_key_values, ) - # Modify the mask to support bidirectional attention. + # Modify the mask to support bidirectional attention if moe_token_types is not None: - # Find the positions of all tokens of type 1. + # Find all token positions with type 1 type1_tokens = ( (moe_token_types == 1).unsqueeze(1).unsqueeze(2) ) # [B, 1, 1, S] - # Create a square mask for the type1 region. + # Create a square mask for the type-1 region type1_mask = torch.zeros_like(causal_mask) # [B, num_heads, S, S] type1_region = type1_tokens & type1_tokens.transpose(-1, -2) # [B, 1, S, S] type1_mask = type1_mask.masked_fill(type1_region, 1.0).to(torch.bool) - # Set the original causal_mask to zero in the type1 region, and then add the type1_mask. + # Zero the original causal_mask in the type-1 region, then add type1_mask causal_mask = torch.where( - type1_mask, - torch.zeros_like(causal_mask), - causal_mask, + type1_mask, # Expand dimensions to match causal_mask + torch.zeros_like(causal_mask), # zero the type-1 region + causal_mask, # keep other regions unchanged ) if ( self.config._attn_implementation == "sdpa" @@ -771,210 +790,59 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): class Qwen2_5_VLMoEForAction( - Qwen2_5_VLForConditionalGeneration, ActionGenerationMixin, ActionModelMixMin + Qwen2_5_VLPreTrainedModel, + ActionGenerationMixin, + ActionModelMixMin, + VLAInferenceMixin, ): - """ - Qwen2.5 Vision-Language Mixture of Experts model for action processing. - - This model extends the base Qwen2.5 VL model with action token processing capabilities - and optional LoRA fine-tuning support. - """ - - _tied_weights_keys = ["lm_head.weight"] + # Compatibility with newer transformers versions (5.x):_tied_weights_keys changed from list[str] to dict[str, str] (target -> source mapping). + # Older versions (4.x) iterate over dict keys, so the behavior is equivalent and compatible. + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} config_class = Qwen2_5_VLConfig _no_split_modules = ["Qwen2_5_VLDecoderLayer_with_MoE", "Qwen2_5_VLVisionBlock"] - @classmethod - def _set_customized_config(cls, config): - """ - Processing norm_stats.json and reconstruct the DoF mapping - """ - dataload_config = config["data"] - if not dataload_config.get("use_lerobot", False): - raise NotImplementedError( - "Not implemented for non-lerobot dataset currently" - ) - - enable_customized_robot_config = config.get( - "enable_customized_robot_config", False - ) - assert ( - enable_customized_robot_config - ), "enable_customized_robot_config must be true when use lerobot dataset" - - customized_dof_config = config["customized_robot_config"][ - "customized_dof_config" - ] - customized_agent_pos_config = config["customized_robot_config"][ - "customized_agent_pos_config" - ] - norm_stats_path = config["norm_stats_path"] - - # Use the compute_action_statistics function from utils - - name = config["customized_robot_config"]["name"] - - update_action_statistics( - action_statistic_dof=action_statistic_dof, # Assuming this is a global variable - norm_stats_path=norm_stats_path, - repo_id=config["data"]["lerobot_config"]["repo_id"], - robot_name=name, - customized_dof_config=customized_dof_config, - customized_agent_pos_config=customized_agent_pos_config, - ) - - print("Customized robot config added") - pprint(action_statistic_dof) - - @classmethod - def from_pretrained( - cls, - pretrained_model_path, - train_config=None, - config_path=None, - processor_path=None, - action_tokenizer_path=None, - is_train=False, - **kwargs, - ): - """ - Load model from pretrained model path. - - Args: - pretrained_model_path (str): Model directory path containing model.safetensors file - config_path (str, optional): Configuration file path, if None will look for qwen25_config.json in pretrained_model_path - processor_path (str, optional): Processor path, if None will load from default config - action_tokenizer_path (str, optional): Action tokenizer path, if None will load from default config - **kwargs: Additional arguments - - Returns: - Qwen2_5_VLMoEForAction: Loaded model instance - """ - # Load model components from pretrained path - - if train_config is None: - try: - with open(os.path.join(pretrained_model_path, "config.yml"), "r") as f: - train_config = yaml.load(f, Loader=yaml.FullLoader) - except Exception as e: - print(f"load train_config.yml fail: {e}") - train_config = None - - model_config_path = os.path.join(pretrained_model_path, "config.json") - model_config = cls.config_class.from_pretrained(model_config_path) - - if train_config is not None: - model_config = update_model_config(train_config, model_config) - processors_dict = load_wallx_processors(train_config) - processor = processors_dict["processor"] - else: - processor = AutoProcessor.from_pretrained( - pretrained_model_path, use_fast=True - ) - - if not is_train: - model_config._attn_implementation = "sdpa" - - if action_tokenizer_path is not None: - processor.action_processor = AutoProcessor.from_pretrained( - action_tokenizer_path, trust_remote_code=True - ) - - # Set the customized robot configuration to ensure consistency between cross-embodiment - # representations and the Wall-X action dimensionality. - # if not train_config: - # cls._set_customized_config(train_config) - # customized_dof_config = train_config["customized_robot_config"][ - # "customized_dof_config" - # ] - # customized_agent_pos_config = train_config["customized_robot_config"][ - # "customized_agent_pos_config" - # ] - # setattr(model_config, "customized_dof_config", customized_dof_config) - # setattr(model_config, "customized_agent_pos_config", customized_agent_pos_config) - - # Initialize model with configuration and processor - model = cls(model_config, processor=processor, **kwargs) - - # Resize token embeddings to match processor tokenizer vocabulary size - model.resize_token_embeddings(len(processor.tokenizer)) - - # Load model state dict from safetensors file - safetensor_files = glob.glob( - os.path.join(pretrained_model_path, "*.safetensors") - ) - state_dict = {} - embed_tokens_size = len(processor.tokenizer) - for file in safetensor_files: - sd = load_file(file, device="cpu") - # filter normalizer statistic params - del_keys = [] - for key in sd.keys(): - if "action_preprocessor.normalizer" in key: - print(f"filter load model weight {key}") - del_keys.append(key) - if "embed_tokens.weight" in key: - embed_tokens_size = sd[key].shape[0] - # if train_config is not None: - for key in del_keys: - del sd[key] - state_dict.update(sd) - if embed_tokens_size != len(processor.tokenizer): - model.resize_token_embeddings(embed_tokens_size) - model.load_state_dict(state_dict, strict=False) - - return model - def __init__( self, - config, - use_fast_tokenizer=False, + config: Qwen2_5_VLConfig, processor=None, - action_tokenizer=None, - action_mapper=None, - flow_loss_weight=1.0, + tokenizer_mixin=None, use_selective_recompute=False, ): - """ - Initialize the Qwen2.5 VLMoE model for action processing. - - Args: - config: Model configuration - use_fast_tokenizer (bool): Whether to use fast tokenizer - processor: Text and image processor - action_tokenizer: Action-specific tokenizer - action_mapper: Action mapping utility - flow_loss_weight (float): Weight for flow loss computation - """ super().__init__(config) - - # Initialize vision transformer and language model components - self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config( - config.vision_config - ) + self.visual = self._build_visual(config, use_selective_recompute) self.model = Qwen2_5_VLMoEModel( config, use_selective_recompute=use_selective_recompute ) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - # Initialize loss function without reduction for channel-wise loss computation - self.loss_fct = CrossEntropyLoss(reduction="none") - self.flow_loss_weight = flow_loss_weight - self.use_fast_tokenizer = use_fast_tokenizer + self.loss_fct = CrossEntropyLoss( + reduction="none" + ) # do not do reduction to compute channel loss + + # Manage through tokenizer_mixin; in pure flow mode the mixin is None self.processor = processor - - # Define action token IDs + self.tokenizer_mixin = tokenizer_mixin + if self.tokenizer_mixin is not None: + self.action_tokenizer_type = self.tokenizer_mixin.tokenizer_type + self.action_tokenizer = self.tokenizer_mixin.tokenizer + self.action_mapper = self.tokenizer_mixin.build_action_mapper(processor) + else: + self.action_tokenizer_type = None + self.action_tokenizer = None + self.action_mapper = None self.define_action_token_id() + self.times_cache = {} # cache times linspace for each num_inference_timesteps + self._infer_stable_cache: dict = {} # cache stable-state inference structures - # Cache for rope deltas - self.rope_deltas = None + self.rope_deltas = None # cache rope_deltas here - # Initialize action preprocessor - self.action_preprocessor = ActionProcessor(config) + from wall_x.model.core.action.processor import ActionProcessor - # Apply LoRA if specified in configuration + self.action_preprocessor = ActionProcessor(config) # action processing + + # Apply LoRA if the configuration contains LoRA settings if hasattr(config, "use_lora") and config.use_lora: self.add_lora( r=config.lora_r, @@ -982,85 +850,37 @@ class Qwen2_5_VLMoEForAction( target_modules=config.lora_target_modules, lora_dropout=config.lora_dropout, ) - # Initialize weights and apply final processing self.post_init() + self._post_init_engine(config) - def define_action_token_id(self): - """ - Define action token IDs based on tokenizer configuration. - - Creates mappings for fast action tokens, proprioception tokens, and general action tokens. - """ - # Create list of fast action token IDs - fast_action_token_list = [] - if self.use_fast_tokenizer: - for i in range( - self.processor.tokenizer.init_kwargs["action_token_vocab_size"] - ): - action_token_id = self.processor.tokenizer.convert_tokens_to_ids( - f"<|action_token_{i}|>" - ) - fast_action_token_list.append(action_token_id) - - # Get special action token IDs - action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") - propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>") - - # Store action token ID mappings - self.action_token_id_set = { - "fast_action_token_list": fast_action_token_list, - "propri_token_id": propri_token_id, - "action_token_id": action_token_id, - } - - def add_lora( - self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1 - ): - """ - Add LoRA (Low-Rank Adaptation) adapters to the model. - - Args: - r (int): Rank of adaptation - lora_alpha (int): LoRA scaling parameter - target_modules (list): List of module names to apply LoRA to - lora_dropout (float): Dropout probability for LoRA layers - """ - config = LoraConfig( - r=r, - lora_alpha=lora_alpha, - target_modules=target_modules, - lora_dropout=lora_dropout, - bias="none", - task_type="CAUSAL_LM", + def _build_visual(self, config, use_selective_recompute): + """Factory method for visual encoder. Subclasses may override to swap implementation.""" + return Qwen2_5_VisionTransformerPretrainedModel( + config=config.vision_config, + use_selective_recompute=use_selective_recompute, ) - self.model = get_peft_model(self.model, config) - # Print information about trainable parameters - self.model.print_trainable_parameters() + def _post_init_engine(self, config): + """Hook called at the end of __init__. Subclasses may override to add engine-specific state.""" + pass def get_input_embeddings(self): - """Get input embeddings layer.""" return self.model.embed_tokens def set_input_embeddings(self, value): - """Set input embeddings layer.""" self.model.embed_tokens = value def get_output_embeddings(self): - """Get output embeddings layer.""" return self.lm_head def set_output_embeddings(self, new_embeddings): - """Set output embeddings layer.""" self.lm_head = new_embeddings def set_decoder(self, decoder): - """Set the decoder model.""" self.model = decoder def get_decoder(self): - """Get the decoder model.""" return self.model def get_rope_index( @@ -1284,6 +1104,173 @@ class Qwen2_5_VLMoEForAction( return position_ids, mrope_position_deltas + def get_attention_maps( + self, + input_ids, + action_horizon, + action_dim, + num_inference_timesteps: int = 10, + prefix_length: Optional[int] = None, # prefix length + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + unnorm_proprioception: Optional[torch.FloatTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + **kwargs, + ): + + if input_ids is not None: + batch_size, seq_length = input_ids.shape + + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if start_indices is None or end_indices is None: + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): + # calculate RoPE index once per generation in the pre-fill stage only + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = get_rope_index( + input_ids=input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + attention_mask=attention_mask, + spatial_merge_size=self.config.vision_config.spatial_merge_size, + image_token_id=self.config.image_token_id, + video_token_id=self.config.video_token_id, + vision_start_token_id=self.config.vision_start_token_id, + tokens_per_second=self.config.vision_config.tokens_per_second, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + # batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(cache_position.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=cache_position.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + img_mask = input_ids == self.config.image_token_id + mask_unsqueezed = img_mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + inputs_embeds = self.scatter_proprioception_embeddings( + input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask + ) + inputs_embeds, flow, adarms_cond = self.scatter_flow_action_embeddings( + input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask + ) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, # pass token_types + start_indices=start_indices, + end_indices=end_indices, + positional_masks=positional_masks, # pass masks for each category + use_cache=False, + output_attentions=True, + output_hidden_states=False, + return_dict=return_dict, + adarms_conds=[None, adarms_cond], + # cache_position=cache_position, + ) + attention_maps = torch.stack(outputs.attentions, dim=0)[ + None, :, -action_horizon: + ] # [flow steps, layer depths, all tokens, all tokens] + attention_maps = attention_maps.mean(2) + image_mask_indices = img_mask[0].nonzero(as_tuple=True)[0] + attention_maps = attention_maps[:, :, image_mask_indices] + return attention_maps + def train_step_forward( self, input_ids: torch.LongTensor = None, @@ -1304,6 +1291,7 @@ class Qwen2_5_VLMoEForAction( cache_position: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, # for vla + sample_time: Optional[torch.FloatTensor] = None, moe_token_types: Optional[torch.LongTensor] = None, start_indices: Optional[torch.Tensor] = None, end_indices: Optional[torch.Tensor] = None, @@ -1334,7 +1322,7 @@ class Qwen2_5_VLMoEForAction( ) if start_indices is None or end_indices is None: - # Calculate the start and end positions of each expert group's tokens after permutation + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here group_size = torch.zeros( self.config.num_experts, dtype=torch.long, device="cpu" ) @@ -1355,12 +1343,17 @@ class Qwen2_5_VLMoEForAction( or self.rope_deltas is None or (past_key_values is None or past_key_values.get_seq_length() == 0) ): - position_ids, rope_deltas = self.get_rope_index( + position_ids, rope_deltas = get_rope_index( input_ids=input_ids, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, second_per_grid_ts=second_per_grid_ts, attention_mask=attention_mask, + spatial_merge_size=self.config.vision_config.spatial_merge_size, + image_token_id=self.config.image_token_id, + video_token_id=self.config.video_token_id, + vision_start_token_id=self.config.vision_start_token_id, + tokens_per_second=self.config.vision_config.tokens_per_second, ) self.rope_deltas = rope_deltas # then use the prev pre-calculated rope-deltas to get the correct position ids @@ -1416,9 +1409,13 @@ class Qwen2_5_VLMoEForAction( inputs_embeds = self.scatter_proprioception_embeddings( input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask ) - inputs_embeds, flow, adarms_cond = self.scatter_flow_action_embeddings( - input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask + input_ids, + inputs_embeds, + action_chunk, + dataset_names, + sample_time, + dof_mask, ) if attention_mask is not None: @@ -1430,10 +1427,10 @@ class Qwen2_5_VLMoEForAction( attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, - moe_token_types=moe_token_types, + moe_token_types=moe_token_types, # pass token_types start_indices=start_indices, end_indices=end_indices, - positional_masks=positional_masks, + positional_masks=positional_masks, # pass masks for each category use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, @@ -1443,7 +1440,23 @@ class Qwen2_5_VLMoEForAction( ) hidden_states = outputs[0] - logits = self.lm_head(hidden_states) + + # --- optimization: loss-token-only LM Head --- + # During training, only compute lm_head for tokens that contribute to CE loss, + # reducing GEMM size from [B*S, V] to [N_loss, V]. + # FSDP fp32 layernorms emit fp32 hidden_states; lm_head runs in bf16. + _lm_dtype = self.lm_head.weight.dtype + _lm_loss_mask = None + if labels is not None and self.training: + shift_labels = labels[..., 1:].contiguous() + _lm_loss_mask = shift_labels != -100 # [B, S-1] + if _lm_loss_mask.any(): + lm_hidden = hidden_states[:, :-1, :][_lm_loss_mask] # [N_loss, H] + logits = self.lm_head(lm_hidden.to(_lm_dtype)) # [N_loss, V] + else: + logits = None + else: + logits = self.lm_head(hidden_states.to(_lm_dtype)) ( loss, @@ -1461,6 +1474,7 @@ class Qwen2_5_VLMoEForAction( dof_mask=dof_mask, flow=flow, flow_loss_mask=flow_loss_mask, + _lm_loss_mask=_lm_loss_mask, ) if not return_dict: @@ -1483,380 +1497,79 @@ class Qwen2_5_VLMoEForAction( ) def predict_action(self, predict_mode: str, **kwargs): - """ - Predict actions using specified prediction mode. + assert predict_mode in ["fast", "spatial_token", "flow", "diffusion", "dllm"] - Args: - predict_mode (str): Prediction mode, either "fast" or "diffusion" - **kwargs: Additional arguments passed to the predict method - - Returns: - tuple: (predicted_action, ground_truth_action) where ground_truth_action may be None - """ - assert predict_mode in ["fast", "diffusion"] - - output = self.predict(predict_mode=predict_mode, **kwargs) + if predict_mode in ["fast", "spatial_token"]: + output = self.generate_ar_action(**kwargs) + elif predict_mode in ["diffusion", "flow"]: + output = self.generate_flow_action(**kwargs) + elif predict_mode == "dllm": + output = self.generate_dllm_action(**kwargs) return output["predict_action"], output.get("gt_action", None) @torch.no_grad() - def predict( + def generate_text( self, - predict_mode: str, - pred_horizon: Optional[int] = None, - action_dim: Optional[int] = None, - input_ids: torch.LongTensor = None, + input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - moe_token_types: Optional[torch.LongTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, # new parameter + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, pixel_values: Optional[torch.Tensor] = None, + pixel_values_side_image: Optional[torch.Tensor] = None, pixel_values_videos: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, + side_image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, - action_chunk: Optional[torch.FloatTensor] = None, - proprioception: Optional[torch.FloatTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, # action chunk + proprioception: Optional[torch.FloatTensor] = None, # joint positions rope_deltas: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, - num_inference_timesteps: Optional[int] = 10, dataset_names: Optional[str] = None, dof_mask: Optional[torch.FloatTensor] = None, agent_pos_mask: Optional[torch.FloatTensor] = None, + gt_output_ids: Optional[torch.LongTensor] = None, + prefix_length: Optional[int] = None, # prefix length + positional_masks: Optional[dict] = None, # position masks for each modality re_generate: bool = False, **kwargs, ): - """ - Multi-modal prediction method supporting text generation, fast action prediction, and diffusion-based action prediction. + # assert self.config._attn_implementation == "flash_attention_2", "generate_text only support flash_attention_2 attn implementation" - This method handles three prediction modes: - 1. "text": Pure text generation using autoregressive decoding - 2. "fast": Fast action prediction using discrete action tokens - 3. "diffusion": Continuous action prediction using diffusion/flow matching - - Args: - predict_mode (str): Prediction mode ("text", "fast", or "diffusion") - pred_horizon (int, optional): Prediction horizon for action sequences - action_dim (int, optional): Dimensionality of action space - input_ids (torch.LongTensor, optional): Input token IDs - attention_mask (torch.Tensor, optional): Attention mask for input tokens - position_ids (torch.LongTensor, optional): Position IDs for tokens - past_key_values (List[torch.FloatTensor], optional): Cached key-value pairs - inputs_embeds (torch.FloatTensor, optional): Pre-computed input embeddings - moe_token_types (torch.LongTensor, optional): Token type assignments for MoE routing - labels (torch.LongTensor, optional): Target labels for evaluation - use_cache (bool, optional): Whether to use key-value caching - output_attentions (bool, optional): Whether to return attention weights - output_hidden_states (bool, optional): Whether to return hidden states - return_dict (bool, optional): Whether to return structured output - pixel_values (torch.Tensor, optional): Image pixel values - pixel_values_videos (torch.FloatTensor, optional): Video pixel values - image_grid_thw (torch.LongTensor, optional): Image grid dimensions - video_grid_thw (torch.LongTensor, optional): Video grid dimensions - action_chunk (torch.FloatTensor, optional): Ground truth action sequences - proprioception (torch.FloatTensor, optional): Proprioceptive sensor data - rope_deltas (torch.LongTensor, optional): RoPE position deltas - cache_position (torch.LongTensor, optional): Cache position indices - second_per_grid_ts (torch.Tensor, optional): Time interval per temporal grid - num_inference_timesteps (int, optional): Number of diffusion inference steps - dataset_names (str, optional): Dataset names for normalization - dof_mask (torch.FloatTensor, optional): Degrees of freedom mask - agent_pos_mask (torch.FloatTensor, optional): Agent position mask - re_generate (bool, optional): Whether to use sampling for regeneration - **kwargs: Additional keyword arguments - - Returns: - dict: Dictionary containing prediction results with keys like: - - 'predict_action': Predicted action sequences - - 'gt_action': Ground truth actions (if available) - - 'input_text': Input text (for text/fast modes) - - 'predict_output_text': Generated text (for text/fast modes) - - 'gt_output_text': Ground truth text (for text/fast modes) - """ - batch_size = ( - input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] - ) - - # Text and fast modes require batch size 1 for autoregressive generation - if predict_mode in ["text", "fast"]: - assert ( - batch_size == 1 - ), "predict only support batch size 1 for ar generation" - - # Set output configuration from model config if not specified - output_attentions = ( - output_attentions - if output_attentions is not None - else self.config.output_attentions - ) - output_hidden_states = ( - output_hidden_states - if output_hidden_states is not None - else self.config.output_hidden_states - ) - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - # Process input embeddings with multi-modal data - if inputs_embeds is None: - inputs_embeds = self.model.embed_tokens(input_ids) - - # Process image embeddings - if pixel_values is not None: - pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) - n_image_tokens = (input_ids == self.config.image_token_id).sum().item() - n_image_features = image_embeds.shape[0] - - # Validate image token and feature count match - if n_image_tokens != n_image_features: - raise ValueError( - f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" - ) - - mask = input_ids == self.config.image_token_id - mask_unsqueezed = mask.unsqueeze(-1) - mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) - image_mask = mask_expanded.to(inputs_embeds.device) - - image_embeds = image_embeds.to( - inputs_embeds.device, inputs_embeds.dtype - ) - inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) - - # Process video embeddings - if pixel_values_videos is not None: - pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - n_video_tokens = (input_ids == self.config.video_token_id).sum().item() - n_video_features = video_embeds.shape[0] - - # Validate video token and feature count match - if n_video_tokens != n_video_features: - raise ValueError( - f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" - ) - - mask = input_ids == self.config.video_token_id - mask_unsqueezed = mask.unsqueeze(-1) - mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) - video_mask = mask_expanded.to(inputs_embeds.device) - - video_embeds = video_embeds.to( - inputs_embeds.device, inputs_embeds.dtype - ) - inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) - - # Process proprioceptive data - if proprioception is not None: - proprioception = proprioception.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - proprio_embed = self.action_preprocessor.proprioception_proj( - proprioception, - dataset_names, - agent_pos_mask, - use_history=proprioception.shape[1] > 1, - ) - proprioception_mask = ( - input_ids == self.action_token_id_set["propri_token_id"] - ) - inputs_embeds[proprioception_mask] = proprio_embed.reshape( - -1, inputs_embeds.shape[-1] - ).to(inputs_embeds.dtype) - - if attention_mask is not None: - attention_mask = attention_mask.to(inputs_embeds.device) - - # Calculate RoPE position IDs if not provided - # Note: Cannot calculate rope deltas with 4D attention mask. TODO: Fix this limitation - if position_ids is None and ( - attention_mask is None or attention_mask.ndim == 2 - ): - # Calculate RoPE index once per generation in the pre-fill stage only - if ( - (cache_position is not None and cache_position[0] == 0) - or self.rope_deltas is None - or (past_key_values is None or past_key_values.get_seq_length() == 0) - ): - position_ids, rope_deltas = ops.get_rope_index( - input_ids=input_ids, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - attention_mask=attention_mask, - spatial_merge_size=self.config.vision_config.spatial_merge_size, - image_token_id=self.config.image_token_id, - video_token_id=self.config.video_token_id, - vision_start_token_id=self.config.vision_start_token_id, - tokens_per_second=self.config.vision_config.tokens_per_second, - ) - self.rope_deltas = rope_deltas - # Use previously calculated rope deltas to get correct position IDs - else: - batch_size, seq_length, _ = inputs_embeds.shape - delta = ( - (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) - if cache_position is not None - else 0 - ) - position_ids = torch.arange(seq_length, device=inputs_embeds.device) - position_ids = position_ids.view(1, -1).expand(batch_size, -1) - if cache_position is not None: # otherwise `deltas` is an int `0` - delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) - position_ids = position_ids.add(delta) - position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) - - # Prepare action chunk data if provided - if action_chunk is not None: - action_chunk = action_chunk.to(inputs_embeds.device).to(inputs_embeds.dtype) - - output = {} - - # Split input sequence for text and fast modes (not needed for diffusion) - if predict_mode == "text" or predict_mode == "fast": - # Look for generation prompt tokens: <|im_start|>assistant - generation_prompt_ids = torch.tensor( - [151644, 77091], device=input_ids.device, dtype=input_ids.dtype + # generate text only need prefix part + if prefix_length is not None: + input_ids, _gt_output_ids = ( + input_ids[:, :prefix_length], + input_ids[:, prefix_length:], ) - matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & ( - input_ids[0, 1:] == generation_prompt_ids[1] + gt_output_ids = ( + gt_output_ids if gt_output_ids is not None else _gt_output_ids ) - - if matches.any(): - split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item() - # Extract ground truth output tokens (including newline) - gt_output_ids = input_ids[:, split_pos + 3 :] - # Remove output part from input, keeping prompt - input_ids = input_ids[:, : split_pos + 3] - inputs_embeds = inputs_embeds[:, : split_pos + 3, :] - if attention_mask is not None: - attention_mask = attention_mask[:, : split_pos + 3] - if labels is not None: - labels = labels[:, split_pos + 3 :] - else: - raise Warning( - "input_ids does not contain the generation prompt tokens <|im_start|>assistant" - ) - - # Decode input text for output - input_text = self.processor.batch_decode( - input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True + attention_mask = ( + attention_mask[:, :prefix_length] + if attention_mask is not None + else None ) - output["input_text"] = input_text - - # Handle text and fast prediction modes using autoregressive generation - if predict_mode == "text" or predict_mode == "fast": - # Initialize MoE token types for generation - moe_token_types = torch.zeros_like(input_ids) - batch = { - "input_ids": input_ids, - "attention_mask": attention_mask, - "pixel_values": pixel_values, - "moe_token_types": moe_token_types, - "image_grid_thw": image_grid_thw, - "dof_mask": dof_mask, - "agent_pos_mask": agent_pos_mask, - "proprioception": proprioception, - "dataset_names": dataset_names, - } - - # Generate output tokens - predict_output_ids = self.generate( - **batch, - max_new_tokens=100, - eos_token_id=[self.processor.tokenizer.eos_token_id], - use_cache=True, - pad_token_id=self.processor.tokenizer.pad_token_id, - temperature=( - 1.0 if not re_generate else 0.7 - ), # Higher temperature for regeneration - do_sample=( - False if not re_generate else True - ), # Enable sampling for regeneration + moe_token_types = ( + moe_token_types[:, :prefix_length] + if moe_token_types is not None + else None ) + else: + prefix_length = input_ids.shape[1] - # Decode generated and ground truth text - gt_output_text = self.processor.batch_decode( - gt_output_ids, - skip_special_tokens=False, - clean_up_tokenization_spaces=True, - ) - predict_output_text = self.processor.batch_decode( - predict_output_ids, - skip_special_tokens=False, - clean_up_tokenization_spaces=True, - ) - output["gt_output_text"] = gt_output_text - output["predict_output_text"] = predict_output_text - - # Convert tokens to actions for fast prediction mode - if predict_mode == "fast": - action_id = [] - # Extract action tokens from generated sequence - for token_id_i in predict_output_ids[0]: - if ( - token_id_i.item() - >= self.processor.tokenizer.init_kwargs["action_token_start_index"] - ): - action_id.append( - token_id_i.item() - - self.processor.tokenizer.init_kwargs[ - "action_token_start_index" - ] - ) - - predict_action = self.processor.action_processor.decode( - [action_id], time_horizon=pred_horizon, action_dim=action_dim - ) - # Handle action decoding errors - if np.sum(predict_action) == 0: - print("Error in decoding action, predict_action is None") - output["predict_action"] = None - else: - # Convert discrete tokens to continuous actions - predict_action = torch.tensor(predict_action, device=self.device) - dof_mask = dof_mask.to(self.device).to(pixel_values.dtype) - predict_action = ( - self.action_preprocessor.normalizer_action.unnormalize_data( - predict_action, dataset_names, dof_mask - ) - ) - output["predict_action"] = predict_action - - # Process ground truth actions if available - if action_chunk is not None: - # Apply DOF mask and unnormalize action chunk to get ground truth actions - action_chunk = action_chunk[:, :, dof_mask[0, 0, :].bool()] - output["gt_action"] = ( - self.action_preprocessor.normalizer_action.unnormalize_data( - action_chunk, dataset_names, dof_mask - ) - ) - else: - output["gt_action"] = None - - # Handle diffusion-based action prediction - if predict_mode == "diffusion": - # Initialize with random noise - noisy_action = torch.randn( - size=(batch_size, pred_horizon, action_dim), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) - dof_mask = dof_mask.to(inputs_embeds.device).to(inputs_embeds.dtype) - - # Calculate token distribution across MoE expert groups + if start_indices is None or end_indices is None: + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here group_size = torch.zeros( self.config.num_experts, dtype=torch.long, device="cpu" ) @@ -1867,91 +1580,188 @@ class Qwen2_5_VLMoEForAction( start_indices = torch.cumsum(group_size, dim=0) - group_size end_indices = torch.cumsum(group_size, dim=0) - def step(timestep, noisy_action): - """ - Single denoising step for diffusion process. - - Args: - timestep: Current diffusion timestep - noisy_action: Current noisy action estimate - - Returns: - torch.Tensor: Predicted clean action - """ - action_mask = input_ids == self.action_token_id_set["action_token_id"] - assert action_mask.any(), "No action token found in input_ids" - - # Prepare timestep for batch processing - timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) - action_embed, _ = self.action_preprocessor.step( - timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask - ) - action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]) - - # Create temporary copy of embeddings for thread safety - temp_inputs_embeds = inputs_embeds.clone() - temp_inputs_embeds[action_mask] = action_embed.to( - temp_inputs_embeds.dtype - ) - - # Forward pass through transformer - transformer_outputs = self.model( - input_ids=None, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=temp_inputs_embeds, - moe_token_types=moe_token_types, - start_indices=start_indices, - end_indices=end_indices, - use_cache=True, - output_attentions=False, - output_hidden_states=False, - return_dict=True, - ) - - # Extract action predictions from hidden states - hidden_states = transformer_outputs.last_hidden_state - action_mask = input_ids == self.action_token_id_set["action_token_id"] - action_hidden_states = hidden_states[action_mask] - pred = self.action_preprocessor.action_proj_back( - action_hidden_states[ - :, : self.action_preprocessor.action_hidden_size - ] - ) - return pred.reshape(batch_size, pred_horizon, action_dim) - - # Perform ODE integration for diffusion sampling - times = torch.linspace( - 0, - 1, - num_inference_timesteps + 1, - device=inputs_embeds.device, - dtype=inputs_embeds.dtype, + batch = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "moe_token_types": moe_token_types, + "start_indices": start_indices, + "end_indices": end_indices, + "image_grid_thw": image_grid_thw, + "dof_mask": dof_mask, + "agent_pos_mask": agent_pos_mask, + "proprioception": proprioception, + "dataset_names": dataset_names, + # "prefix_length": prefix_length, + } + predict_output_ids = self.generate( + **batch, + max_new_tokens=100, + eos_token_id=[ + self.processor.tokenizer.eos_token_id + ], # set multiple end markers + use_cache=True, + pad_token_id=self.processor.tokenizer.pad_token_id, # explicitly set pad_token_id + temperature=( + 1.0 if not re_generate else 0.7 + ), # use a higher temperature when regenerating + do_sample=( + False if not re_generate else True + ), # use sampling when regenerating + ) + input_text = self.processor.batch_decode( + input_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True + ) + if gt_output_ids is not None and gt_output_ids.shape[-1] != 0: + gt_output_text = self.processor.batch_decode( + gt_output_ids, + skip_special_tokens=False, + clean_up_tokenization_spaces=True, ) - action_trajectory = odeint( - step, - noisy_action.to(torch.float32), - times.to(torch.float32), - method="euler", + else: + gt_output_text = None + predict_output_text = self.processor.batch_decode( + predict_output_ids[:, prefix_length:], + skip_special_tokens=False, + clean_up_tokenization_spaces=True, + ) + output = { + "input_text": input_text, + "gt_output_text": gt_output_text, + "predict_output_text": predict_output_text, + "predict_output_ids": predict_output_ids, + } + return output + + @torch.no_grad() + def generate_ar_action( + self, + input_ids, + action_horizon, + action_dim, + state=None, + dataset_names=None, + action_chunk=None, + dof_mask=None, + unnorm=True, + max_retries=3, # maximum retry count + proprioception=None, # proprioception for v3.1_delta decoding + robot_type_id=None, # robot type ID for v3.1_delta decoding + **kwargs, + ): + # Regeneration mechanism + predict_action = None + output = None + decode_success = False + if dataset_names is None and (unnorm or action_chunk is not None): + raise KeyError( + "generate_ar_action requires dataset_names when normalizing or " + "unnormalizing actions." ) - # Extract final predicted action and unnormalize - predict_action = action_trajectory[-1] - predict_action = ( - self.action_preprocessor.normalizer_action.unnormalize_data( - predict_action, dataset_names - ) - ) - output["predict_action"] = predict_action + for retry_attempt in range(max_retries): + use_re_generate = retry_attempt > 0 - # Process ground truth actions if available - if action_chunk is not None: + if retry_attempt > 0: + logger.info( + "Retrying action generation (attempt %s/%s)...", + retry_attempt + 1, + max_retries, + ) + + output = self.generate_text( + input_ids=input_ids, re_generate=use_re_generate, **kwargs + ) + predict_output_ids = output["predict_output_ids"] + + # Decode uniformly through tokenizer_mixin + predict_action, decode_success = self.tokenizer_mixin.decode_action( + output_ids=predict_output_ids, + action_mapper=self.action_mapper, + action_horizon=action_horizon, + action_dim=action_dim, + device=input_ids.device, + proprioception=proprioception, + dof_mask=dof_mask, + robot_type_id=robot_type_id, + state=state, + ) + + if decode_success: + if retry_attempt > 0: + logger.info( + "Action generation succeeded on attempt %s", + retry_attempt + 1, + ) + break + else: + if retry_attempt < max_retries - 1: + logger.warning( + "Error in decoding action (attempt %s/%s), retrying with re_generate mode...", + retry_attempt + 1, + max_retries, + ) + else: + logger.warning( + "Error in decoding action after %s attempts, returning None", + max_retries, + ) + + # Check whether dof_mask is required through mixin attributes + uses_dof_mask = ( + self.tokenizer_mixin.uses_dof_mask_for_unnorm + if self.tokenizer_mixin is not None + else True + ) + + if action_chunk is not None: + action_chunk = action_chunk.to(input_ids.device).to(torch.bfloat16) + if uses_dof_mask: + action_chunk = action_chunk[:, :, dof_mask[0, 0, :].bool()] output["gt_action"] = ( self.action_preprocessor.normalizer_action.unnormalize_data( - action_chunk, dataset_names + action_chunk, dataset_names, dof_mask ) ) + else: + # Output full dimensions without mask filtering + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names, None + ) + ) + + if not decode_success or predict_action is None: + logger.warning("Error in decoding action, predict_action is None") + output["predict_action"] = None + else: + if unnorm: + if isinstance(predict_action, np.ndarray): + predict_action = torch.tensor( + predict_action, device=input_ids.device + ) + elif predict_action.device != input_ids.device: + predict_action = predict_action.to(input_ids.device) + + # Add the batch dimension + if predict_action.dim() == 2: + predict_action = predict_action.unsqueeze(0) + + if uses_dof_mask: + predict_action = unnormalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + predict_action, + dataset_names, + dof_mask, + ) + else: + # Output full dimensions without mask filtering + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names, None + ) + ) + output["predict_action"] = predict_action return output @@ -1962,8 +1772,7 @@ class Qwen2_5_VLMoEForAction( action_horizon, action_dim, num_inference_timesteps: int = 10, - padding_action: Optional[torch.Tensor] = None, - prefix_length: Optional[int] = None, + prefix_length: Optional[int] = None, # prefix length attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -1990,13 +1799,94 @@ class Qwen2_5_VLMoEForAction( dataset_names: Optional[str] = None, dof_mask: Optional[torch.FloatTensor] = None, agent_pos_mask: Optional[torch.FloatTensor] = None, - unnorm: Optional[bool] = True, **kwargs, ): - total_start_time = time.time() - timing_results = {} + # assert self.config._attn_implementation == "sdpa", "generate_flow_action only support sdpa attn implementation" + ctx = self._prepare_flow_action_inputs( + input_ids=input_ids, + action_horizon=action_horizon, + action_dim=action_dim, + num_inference_timesteps=num_inference_timesteps, + prefix_length=prefix_length, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, + start_indices=start_indices, + end_indices=end_indices, + positional_masks=positional_masks, + labels=labels, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + action_chunk=action_chunk, + proprioception=proprioception, + unnorm_proprioception=unnorm_proprioception, + rope_deltas=rope_deltas, + cache_position=cache_position, + second_per_grid_ts=second_per_grid_ts, + dataset_names=dataset_names, + dof_mask=dof_mask, + agent_pos_mask=agent_pos_mask, + ) + action_trajectory, attention_maps = self._execute_flow_action(ctx) + + output = self._finalize_flow_action_output(action_trajectory, ctx) + if output_attentions: + output["attention_maps"] = attention_maps + return output + + def _execute_flow_action(self, ctx): + """Strategy hook: selects which flow-action backend to use. + + Returns: + action_trajectory: list of action tensors + attention_maps: attention maps or None + """ + return self._generate_flow_action_vanilla(ctx) + + def _prepare_flow_action_inputs( + self, + input_ids, + action_horizon, + action_dim, + num_inference_timesteps, + prefix_length, + attention_mask, + position_ids, + past_key_values, + inputs_embeds, + moe_token_types, + start_indices, + end_indices, + positional_masks, + labels, + use_cache, + output_attentions, + output_hidden_states, + return_dict, + pixel_values, + pixel_values_videos, + image_grid_thw, + video_grid_thw, + action_chunk, + proprioception, + unnorm_proprioception, + rope_deltas, + cache_position, + second_per_grid_ts, + dataset_names, + dof_mask, + agent_pos_mask, + ): batch_size = ( input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] ) @@ -2014,11 +1904,14 @@ class Qwen2_5_VLMoEForAction( return_dict if return_dict is not None else self.config.use_return_dict ) - embed_start_time = time.time() + # Timing: input embedding processing + img_mask = None if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) if pixel_values is not None: pixel_values = pixel_values.type(self.visual.dtype) + # from wall_x.utils.timers import ScopeTimer + # with ScopeTimer("pixel_values.visual"): image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) n_image_tokens = (input_ids == self.config.image_token_id).sum().item() n_image_features = image_embeds.shape[0] @@ -2027,8 +1920,8 @@ class Qwen2_5_VLMoEForAction( f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" ) - mask = input_ids == self.config.image_token_id - mask_unsqueezed = mask.unsqueeze(-1) + img_mask = input_ids == self.config.image_token_id + mask_unsqueezed = img_mask.unsqueeze(-1) mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) image_mask = mask_expanded.to(inputs_embeds.device) @@ -2079,9 +1972,6 @@ class Qwen2_5_VLMoEForAction( if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) - timing_results["embed_processing"] = time.time() - embed_start_time - - position_start_time = time.time() # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme if position_ids is None and ( attention_mask is None or attention_mask.ndim == 2 @@ -2092,12 +1982,17 @@ class Qwen2_5_VLMoEForAction( or self.rope_deltas is None or (past_key_values is None or past_key_values.get_seq_length() == 0) ): - position_ids, rope_deltas = self.get_rope_index( + position_ids, rope_deltas = get_rope_index( input_ids, image_grid_thw, video_grid_thw, second_per_grid_ts, attention_mask, + self.config.vision_config.spatial_merge_size, + self.config.image_token_id, + self.config.video_token_id, + self.config.vision_start_token_id, + self.config.vision_config.tokens_per_second, ) self.rope_deltas = rope_deltas # then use the prev pre-calculated rope-deltas to get the correct position ids @@ -2115,25 +2010,46 @@ class Qwen2_5_VLMoEForAction( position_ids = position_ids.add(delta) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] + + if prefix_length is None: + has_true = flow_action_mask.any(dim=1) + prefix_length = torch.argmax(flow_action_mask.float(), dim=1, keepdim=True) + prefix_length[~has_true] = flow_action_mask.shape[1] + # check if prefix_length is the same for all batch + if not torch.all(prefix_length == prefix_length[0]): + raise ValueError( + "prefix_length differs across batch; batch prompts must align" + ) + prefix_length = int(prefix_length[0].item()) + if start_indices is None or end_indices is None: - # Calculate the start and end positions of each expert group's tokens after permutation (the dataset does not contain `num_expert` information, so this calculation must be done here). - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (moe_token_types == i).sum() + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here + # Cache key: shape + sum is sufficient for stable-state (moe_token_types structure doesn't change) + _moe_key = (moe_token_types.shape, int(moe_token_types.sum().item())) + _moe_cache = self._infer_stable_cache.get("moe_indices") + if _moe_cache is None or _moe_cache[0] != _moe_key: + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() - # Calculate start and end indices for each expert group - start_indices = torch.cumsum(group_size, dim=0) - group_size - end_indices = torch.cumsum(group_size, dim=0) + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + self._infer_stable_cache["moe_indices"] = ( + _moe_key, + start_indices, + end_indices, + ) + else: + _, start_indices, end_indices = _moe_cache - timing_results["position_encoding"] = time.time() - position_start_time - - action_init_start_time = time.time() + # Timing: action initialization if action_chunk is not None: action_chunk = action_chunk.to(inputs_embeds.device).to(torch.float32) - output = {} # Reproduce # torch.manual_seed(0) noise = torch.randn( @@ -2142,18 +2058,16 @@ class Qwen2_5_VLMoEForAction( device=inputs_embeds.device, ) noisy_action = noise.clone() + dof_mask = dof_mask.to(inputs_embeds.device).to(torch.float32) if num_inference_timesteps not in self.times_cache: - self.times_cache[num_inference_timesteps] = torch.linspace( - 0.0, - 1.0, - num_inference_timesteps + 1, - device=inputs_embeds.device, - dtype=torch.float32, + self.times_cache[num_inference_timesteps] = ( + self.action_preprocessor.get_inference_times( + num_inference_timesteps, inputs_embeds.device, torch.float32 + ) ) times = self.times_cache[num_inference_timesteps] - dt = times[1] - times[0] time_0 = times[0].unsqueeze(0).repeat(noisy_action.shape[0]) action_embed, adarms_cond = self.action_preprocessor.step( timestep=time_0, noisy_action=noisy_action, dof_mask=dof_mask @@ -2161,134 +2075,101 @@ class Qwen2_5_VLMoEForAction( action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]).to( inputs_embeds.dtype ) - flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] - inputs_embeds[flow_action_mask] = action_embed - timing_results["action_initialization"] = time.time() - action_init_start_time + # Automatically generate padding actions from dof_mask + if not dof_mask.all(): + padding_action = ( + torch.zeros((1, dof_mask.shape[-1])) + .to(dof_mask.device) + .to(torch.float32) + ) + padding_action = normalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + padding_action, + dataset_names, + dof_mask, + ) + else: + padding_action = None - prefetch_start_time = time.time() + return { + "batch_size": batch_size, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "position_ids": position_ids, + "prefix_length": prefix_length, + "flow_action_mask": flow_action_mask, + "start_indices": start_indices, + "end_indices": end_indices, + "dof_mask": dof_mask, + "noise": noise, + "noisy_action": noisy_action, + "times": times, + "adarms_cond": adarms_cond, + "action_chunk": action_chunk, + "padding_action": padding_action, + "action_horizon": action_horizon, + "action_dim": action_dim, + "input_ids": input_ids, + "moe_token_types": moe_token_types, + "positional_masks": positional_masks, + "dataset_names": dataset_names, + "output_attentions": output_attentions, + "img_mask": img_mask, + } + + def _generate_flow_action_vanilla(self, ctx): + # Timing: prefill forward pass + all_attention_maps = [] + output_attentions = True if ctx["output_attentions"] else False prefetch_output = self.model( input_ids=None, - attention_mask=attention_mask, - position_ids=position_ids, + attention_mask=ctx["attention_mask"], + position_ids=ctx["position_ids"], past_key_values=None, - inputs_embeds=inputs_embeds, - moe_token_types=moe_token_types, - start_indices=start_indices, - end_indices=end_indices, - positional_masks=positional_masks, + inputs_embeds=ctx["inputs_embeds"], + moe_token_types=ctx["moe_token_types"], + start_indices=ctx["start_indices"], + end_indices=ctx["end_indices"], + positional_masks=ctx["positional_masks"], use_cache=True, - output_attentions=False, + output_attentions=output_attentions, output_hidden_states=False, return_dict=True, - adarms_conds=[None, adarms_cond], + adarms_conds=[None, ctx["adarms_cond"]], ) hidden_states = prefetch_output.last_hidden_state prefix_kv_cache = prefetch_output.past_key_values - - action_hidden_states = hidden_states[flow_action_mask].to(torch.float32) + if output_attentions: + all_attention_maps.append(prefetch_output.attentions) + action_hidden_states = hidden_states[ctx["flow_action_mask"]].to(torch.float32) action_pred = self.action_preprocessor.action_proj_back( action_hidden_states[:, : self.action_preprocessor.action_hidden_size] ) if getattr(self.config, "use_x_pred", False): - v_0 = action_pred - noise.reshape(-1, noise.shape[-1]) + v_0 = action_pred - ctx["noise"].reshape(-1, ctx["noise"].shape[-1]) else: v_0 = action_pred + v_0 = v_0.reshape(ctx["batch_size"], ctx["action_horizon"], ctx["action_dim"]) - if (not dof_mask.all()) and (padding_action is not None): - print("use padding action", flush=True) - v_padding = padding_action - noisy_action - v_0 = (v_padding) * (1 - dof_mask) + v_0 * dof_mask + v_padding = None + if (not ctx["dof_mask"].all()) and (ctx["padding_action"] is not None): + v_padding = ctx["padding_action"] - ctx["noisy_action"] + v_0 = (v_padding) * (1 - ctx["dof_mask"]) + v_0 * ctx["dof_mask"] + dt = ctx["times"][1] - ctx["times"][0] + ctx["noisy_action"] = ctx["noisy_action"] + dt * v_0 - noisy_action = noisy_action + dt * v_0.reshape( - batch_size, action_horizon, action_dim - ) - - timing_results["prefetch_forward"] = time.time() - prefetch_start_time - - cache_prep_start_time = time.time() - - if prefix_length is None: - has_true = flow_action_mask.any(dim=1) - prefix_length = torch.argmax(flow_action_mask.float(), dim=1, keepdim=True) - prefix_length[~has_true] = flow_action_mask.shape[1] - prefix_length = prefix_length[0] - - # support different transformers version - if hasattr(prefix_kv_cache, "key_cache"): - for layer_i in range(len(prefix_kv_cache.key_cache)): - prefix_kv_cache.key_cache[layer_i] = prefix_kv_cache.key_cache[layer_i][ - :, :, :prefix_length, : - ] - prefix_kv_cache.value_cache[layer_i] = prefix_kv_cache.value_cache[ - layer_i - ][:, :, :prefix_length, :] - else: - for layer_i in range(len(prefix_kv_cache.layers)): - prefix_kv_cache.layers[layer_i].keys = prefix_kv_cache.layers[ - layer_i - ].keys[:, :, :prefix_length, :] - prefix_kv_cache.layers[layer_i].values = prefix_kv_cache.layers[ - layer_i - ].values[:, :, :prefix_length, :] - - postfix_position_ids = position_ids[:, :, prefix_length:] - postfix_inputs_embeds = inputs_embeds[:, prefix_length:, :] - postfix_attention_mask = attention_mask[:, prefix_length:] - postfix_moe_token_types = moe_token_types[:, prefix_length:] - postfix_input_ids = input_ids[:, prefix_length:] - - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (postfix_moe_token_types == i).sum() - - # Calculate start and end indices for each expert group - postfix_start_indices = torch.cumsum(group_size, dim=0) - group_size - postfix_end_indices = torch.cumsum(group_size, dim=0) - - pad_token_id = self.processor.tokenizer.pad_token_id - padding_mask = input_ids == pad_token_id - - # prefix_length, postfix_length = prefix_indices.shape[0], postfix_indices.shape[0] - - postfix_length = input_ids.shape[-1] - prefix_length - _postfix_attention_mask = torch.ones( - (batch_size, postfix_length, prefix_length + postfix_length), - dtype=torch.bool, - device=postfix_attention_mask.device, - ) - - # Use a padding mask to set the corresponding rows and columns to false. - # Get the padding mask for the postfix portion. - postfix_padding_mask = padding_mask[ - :, prefix_length: - ] # [batch_size, postfix_length] - full_padding_mask = padding_mask # [batch_size, prefix_length + postfix_length] - - # causal mask for postfix attention - if self.config.causal_action_attention_mask: - _postfix_attention_mask[:, :, prefix_length:] = torch.tril( - torch.ones( - (postfix_length, postfix_length), - dtype=torch.bool, - device=postfix_attention_mask.device, - ) - ) - - for batch_idx in range(padding_mask.shape[0]): - # Set the rows corresponding to the padding positions to False (where the query position is padding). - _postfix_attention_mask[batch_idx, postfix_padding_mask[batch_idx], :] = ( - False - ) - # Set the columns corresponding to the padding positions to False (the key position is the padding). - _postfix_attention_mask[batch_idx, :, full_padding_mask[batch_idx]] = False - - timing_results["cache_preprocessing"] = time.time() - cache_prep_start_time - - ode_start_time = time.time() + ( + postfix_position_ids, + postfix_inputs_embeds, + _postfix_attention_mask, + postfix_moe_token_types, + postfix_input_ids, + postfix_start_indices, + postfix_end_indices, + padding_mask, + ) = self._prepare_flow_postfix(ctx, prefix_kv_cache) def step_with_kvcache(timestep, noisy_action): action_mask = ( @@ -2297,10 +2178,11 @@ class Qwen2_5_VLMoEForAction( assert action_mask.any(), "No action token found in input_ids" timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) action_embed, adarms_cond = self.action_preprocessor.step( - timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask + timestep=timestep, noisy_action=noisy_action, dof_mask=ctx["dof_mask"] ) action_embed = action_embed.reshape(-1, postfix_inputs_embeds.shape[-1]) + # Create a temporary inputs_embeds copy for thread safety temp_inputs_embeds = postfix_inputs_embeds.clone() temp_inputs_embeds[action_mask] = action_embed.to(temp_inputs_embeds.dtype) transformer_outputs = self.model( @@ -2313,81 +2195,173 @@ class Qwen2_5_VLMoEForAction( start_indices=postfix_start_indices, end_indices=postfix_end_indices, use_cache=False, - output_attentions=False, + output_attentions=output_attentions, output_hidden_states=False, return_dict=True, adarms_conds=[None, adarms_cond], ) - + if output_attentions: + all_attention_maps.append(transformer_outputs.attentions) hidden_states = transformer_outputs.last_hidden_state action_hidden_states = hidden_states[action_mask].to(torch.float32) action_pred = self.action_preprocessor.action_proj_back( action_hidden_states[:, : self.action_preprocessor.action_hidden_size] ) if getattr(self.config, "use_x_pred", False): - v_t = action_pred - noise.reshape(-1, noise.shape[-1]) + # Align noisy_action (and timestep) to action_pred's shape + B, action_horizon, action_dim = noisy_action.shape + noisy_action_flat = noisy_action.reshape( + -1, action_dim + ) # [B * action_horizon, action_dim] + timestep_expand = timestep.view(-1, 1).repeat_interleave( + action_horizon, dim=0 + ) # [B * action_horizon, 1] + v_t = (action_pred - noisy_action_flat) / torch.clamp( + 1 - timestep_expand, min=0.05 + ) else: v_t = action_pred - return v_t.reshape(batch_size, action_horizon, action_dim) + v_t = v_t.reshape( + ctx["batch_size"], ctx["action_horizon"], ctx["action_dim"] + ) + + if (not ctx["dof_mask"].all()) and (ctx["padding_action"] is not None): + v_t = (v_padding) * (1 - ctx["dof_mask"]) + v_t * ctx["dof_mask"] + + return v_t action_trajectory = odeint( - step_with_kvcache, noisy_action, times[1:], method="euler" + step_with_kvcache, ctx["noisy_action"], ctx["times"][1:], method="euler" ) - timing_results["ode_integration"] = time.time() - ode_start_time + attention_maps = None + if output_attentions and ctx["img_mask"] is not None: + attention_maps = [torch.stack(map, dim=0) for map in all_attention_maps] + attention_maps = torch.stack( + attention_maps, dim=0 + ) # [flow steps, layer depths, action tokens, all tokens] + attention_maps = attention_maps.mean(2) + image_mask_indices = ctx["img_mask"][0].nonzero(as_tuple=True)[0] + attention_maps = attention_maps[:, :, image_mask_indices] + return action_trajectory, attention_maps - postprocess_start_time = time.time() + def _prepare_flow_postfix(self, ctx, prefix_kv_cache=None): + postfix_position_ids = ctx["position_ids"][:, :, ctx["prefix_length"] :] + postfix_inputs_embeds = ctx["inputs_embeds"][:, ctx["prefix_length"] :, :] + postfix_attention_mask = ctx["attention_mask"][:, ctx["prefix_length"] :] + postfix_moe_token_types = ctx["moe_token_types"][:, ctx["prefix_length"] :] + postfix_input_ids = ctx["input_ids"][:, ctx["prefix_length"] :] + + postfix_start_indices = None + postfix_end_indices = None + padding_mask = None + + if prefix_kv_cache is not None: + if ctx["prefix_length"] is None: + has_true = ctx["flow_action_mask"].any(dim=1) + prefix_length = torch.argmax( + ctx["flow_action_mask"].float(), dim=1, keepdim=True + ) + prefix_length[~has_true] = ctx["flow_action_mask"].shape[1] + # check if prefix_length is the same for all batch + if not torch.all(prefix_length == prefix_length[0]): + raise ValueError( + "prefix_length differs across batch; batch prompts must align" + ) + prefix_length = int(prefix_length[0].item()) + ctx["prefix_length"] = prefix_length + + if hasattr(prefix_kv_cache, "key_cache"): + for layer_i in range(len(prefix_kv_cache.key_cache)): + prefix_kv_cache.key_cache[layer_i] = prefix_kv_cache.key_cache[ + layer_i + ][:, :, : ctx["prefix_length"], :] + prefix_kv_cache.value_cache[layer_i] = prefix_kv_cache.value_cache[ + layer_i + ][:, :, : ctx["prefix_length"], :] + else: + for layer_i in range(len(prefix_kv_cache.layers)): + prefix_kv_cache.layers[layer_i].keys = prefix_kv_cache.layers[ + layer_i + ].keys[:, :, : ctx["prefix_length"], :] + prefix_kv_cache.layers[layer_i].values = prefix_kv_cache.layers[ + layer_i + ].values[:, :, : ctx["prefix_length"], :] + + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (postfix_moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + postfix_start_indices = torch.cumsum(group_size, dim=0) - group_size + postfix_end_indices = torch.cumsum(group_size, dim=0) + + pad_token_id = self.processor.tokenizer.pad_token_id + padding_mask = ctx["input_ids"] == pad_token_id + else: + pad_token_id = self.processor.tokenizer.pad_token_id + padding_mask = ctx["input_ids"] == pad_token_id + + postfix_length = ctx["input_ids"].shape[-1] - ctx["prefix_length"] + _postfix_attention_mask = torch.ones( + (ctx["batch_size"], postfix_length, ctx["prefix_length"] + postfix_length), + dtype=torch.bool, + device=postfix_attention_mask.device, + ) + + postfix_padding_mask = padding_mask[:, ctx["prefix_length"] :] + full_padding_mask = padding_mask + + if self.config.causal_action_attention_mask: + _postfix_attention_mask[:, :, ctx["prefix_length"] :] = torch.tril( + torch.ones( + (postfix_length, postfix_length), + dtype=torch.bool, + device=postfix_attention_mask.device, + ) + ) + + # postfix_padding_mask: [B, postfix_length], True where padding (query rows) + # full_padding_mask: [B, prefix_length+postfix_length], True where padding (key cols) + _postfix_attention_mask.masked_fill_(postfix_padding_mask.unsqueeze(2), False) + _postfix_attention_mask.masked_fill_(full_padding_mask.unsqueeze(1), False) + + return ( + postfix_position_ids, + postfix_inputs_embeds, + _postfix_attention_mask, + postfix_moe_token_types, + postfix_input_ids, + postfix_start_indices, + postfix_end_indices, + padding_mask, + ) + + def _finalize_flow_action_output(self, action_trajectory, ctx): predict_action = action_trajectory[-1] - if unnorm: - predict_action = ( - self.action_preprocessor.normalizer_action.unnormalize_data( - predict_action, dataset_names - ) + predict_action = unnormalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + predict_action, + ctx["dataset_names"], + ctx.get("dof_mask"), + ) + output = {"predict_action": predict_action} + if ctx["action_chunk"] is not None: + output["gt_action"] = unnormalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + ctx["action_chunk"], + ctx["dataset_names"], + ctx.get("dof_mask"), ) - output["predict_action"] = predict_action - # normalize action chunk to get gt_action - if action_chunk is not None: - output["gt_action"] = ( - self.action_preprocessor.normalizer_action.unnormalize_data( - action_chunk, dataset_names - ) - ) - - timing_results["postprocessing"] = time.time() - postprocess_start_time - timing_results["total_time"] = time.time() - total_start_time - - output["timing_results"] = timing_results return output - def forward( - self, mode: Optional[str] = None, predict_mode: Optional[str] = "text", **kwargs - ): - """ - Main forward pass dispatcher for different execution modes. - - This method routes execution to appropriate forward functions based on the specified mode: - - No mode (None): Training step with gradient disabled - - 'predict': Prediction/inference mode - - 'train': Training mode with gradients enabled - - 'validate': Validation mode with gradients disabled - - Args: - mode (str, optional): Execution mode. If None, defaults to training step without gradients - predict_mode (str, optional): Prediction mode for 'predict' mode ("text", "fast", or "diffusion") - **kwargs: Additional arguments passed to the selected forward function - - Returns: - Model outputs appropriate for the selected mode - - Todo: - - Add support for distinguishing multi-modal data types in prediction mode - """ + def forward(self, mode: Optional[str] = None, **kwargs): if not mode: with torch.no_grad(): return self.train_step_forward(**kwargs) - elif mode == "predict": - return self.generate_flow_action(predict_mode=predict_mode, **kwargs) elif mode == "train": return self.train_step_forward(use_cache=False, **kwargs) elif mode == "validate": @@ -2403,6 +2377,8 @@ class Qwen2_5_VLMoEForAction( attention_mask=None, inputs_embeds=None, moe_token_types=None, + start_indices=None, + end_indices=None, cache_position=None, position_ids=None, use_cache=True, @@ -2417,95 +2393,60 @@ class Qwen2_5_VLMoEForAction( agent_pos_mask=None, **kwargs, ): - """ - Prepare inputs for autoregressive generation with multi-modal support. + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model - This method handles input preparation for generation, including proper slicing of inputs - based on cache position, MoE token type management, and multi-modal data handling. - Vision inputs are selectively forwarded only when needed during generation. + # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens + # Exception 1: when passing input_embeds, input_ids may be missing entries + # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here + # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case. + # (we can't check exception 3 while compiling) + # Exception 4: If input_embeds are passed then slice it through `cache_position`, to keep only the unprocessed tokens and + # generate the first token for each sequence. Later use the generated Input ids for continuation. - Args: - input_ids: Input token IDs - past_key_values: Cached key-value pairs from previous generation steps - attention_mask: Attention mask for input tokens - inputs_embeds: Pre-computed input embeddings - moe_token_types: Token type assignments for MoE routing - cache_position: Current cache position for generation - position_ids: Position IDs for tokens - use_cache: Whether to use key-value caching - pixel_values: Image pixel values - pixel_values_videos: Video pixel values - image_grid_thw: Image grid dimensions - video_grid_thw: Video grid dimensions - second_per_grid_ts: Time interval per temporal grid - dataset_names: Dataset names for processing - proprioception: Proprioceptive sensor data - dof_mask: Degrees of freedom mask - agent_pos_mask: Agent position mask - **kwargs: Additional arguments - - Returns: - dict: Prepared model inputs for generation step - - Todo: - - Test this function thoroughly with various input configurations - - Note: - This is an overridden method that handles specific cases for multi-modal generation: - - Slices input_ids through cache_position to keep only unprocessed tokens - - Handles special cases for input_embeds, generation methods, and GPU synchronization - - Manages vision inputs to avoid unnecessary forward passes - """ - # Initialize MoE token types if not provided if moe_token_types is None: moe_token_types = torch.zeros_like( input_ids - ) # FIXME: Handle case when input_embeds is used instead + ) # FIXME input_ids or input_embeds else: - # Ensure moe_token_types length matches input_ids + # Ensure moe_token_types has the same length as input_ids if moe_token_types.shape[1] < input_ids.shape[1]: - # Calculate required padding length + # Compute the required padding length pad_length = input_ids.shape[1] - moe_token_types.shape[1] - # Create padding tensor with default token type (0) + # Create the padding tensor using 0 as the default token type pad_tensor = torch.zeros( (moe_token_types.shape[0], pad_length), dtype=moe_token_types.dtype, device=moe_token_types.device, ) - # Concatenate padding to existing moe_token_types + # Append padding after the existing moe_token_types moe_token_types = torch.cat([moe_token_types, pad_tensor], dim=1) - # Handle input slicing based on cache state and special cases if past_key_values is not None: - if ( - inputs_embeds is not None and input_ids.shape[1] == 0 - ): # Exception 4: input_embeds case + if inputs_embeds is not None and input_ids.shape[1] == 0: # Exception 4 inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] moe_token_types = moe_token_types[:, -cache_position.shape[0] :] - elif inputs_embeds is not None or ( # Exception 1: input_embeds provided + elif inputs_embeds is not None or ( # Exception 1 is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1] - ): # Exception 3: GPU sync edge case + ): # Exception 3 input_ids = input_ids[:, -cache_position.shape[0] :] moe_token_types = moe_token_types[:, -cache_position.shape[0] :] elif ( input_ids.shape[1] != cache_position.shape[0] - ): # Default case (Exception 2 is no-op) + ): # Default case (the "else", a no op, is Exception 2) cache_pos = cache_position.clone() input_ids = input_ids[:, cache_pos] moe_token_types = moe_token_types[:, cache_pos] - # Skip vision inputs for continuation steps (not initial generation) if cache_position[0] != 0: pixel_values = None pixel_values_videos = None - # Determine whether to use inputs_embeds or input_ids for this generation step + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]: model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} else: model_inputs = {"input_ids": input_ids, "inputs_embeds": None} - # Prepare 4D causal attention mask for static cache if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2: if model_inputs["inputs_embeds"] is not None: batch_size, sequence_length, _ = inputs_embeds.shape @@ -2528,12 +2469,19 @@ class Qwen2_5_VLMoEForAction( ) ) - # Assemble all model inputs for generation + if ( + self.config._attn_implementation == "flash_attention_2_ki" + and self.config.attention_moe is True + ): + attention_mask = None # Pass None as the mask to use the official flash attention for autoregressive prediction + model_inputs.update( { "position_ids": position_ids, "past_key_values": past_key_values, "moe_token_types": moe_token_types, + "start_indices": start_indices, + "end_indices": end_indices, "use_cache": use_cache, "attention_mask": attention_mask, "pixel_values": pixel_values, @@ -2555,30 +2503,25 @@ class Qwen2_5_VLMoEForAction( input_ids: Optional[torch.LongTensor], ) -> Tuple[torch.Tensor, torch.Tensor]: """ - Get the number of images and videos for each sample to calculate tensor separation lengths. - - These parameters are computed directly from input_ids rather than being passed through - the processor to avoid unpredictable impacts from interface modifications. + Get the number of images and videos for each sample to calculate the separation length of the sample tensor. + These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications. Args: - input_ids (torch.LongTensor): Input token IDs of shape (batch_size, sequence_length) + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Returns: - tuple: - - image_nums (torch.LongTensor): Number of images per sample - - video_nums (torch.LongTensor): Number of videos per sample + image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`) + video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`) """ image_token_id = self.config.image_token_id video_token_id = self.config.video_token_id vision_start_token_id = self.config.vision_start_token_id - # Find vision start tokens and their following tokens vision_start_mask = input_ids == vision_start_token_id vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1) image_mask = input_ids == image_token_id video_mask = input_ids == video_token_id - - # Count images and videos following vision start tokens image_nums = torch.sum(vision_first_mask & image_mask, dim=1) video_nums = torch.sum(vision_first_mask & video_mask, dim=1) @@ -2591,28 +2534,14 @@ class Qwen2_5_VLMoEForAction( input_ids: Optional[torch.LongTensor] = None, **model_kwargs, ) -> Tuple[torch.LongTensor, Dict[str, Any]]: - """ - Expand inputs for generation with support for multi-modal tensors. + # Overwritten -- Support for expanding tensors without a batch size dimension + # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t + # pixel_values.shape[0] is sum(seqlen_images for samples) + # image_grid_thw.shape[0] is sum(num_images for samples) - This is an overridden method that supports expanding tensors without a standard batch - size dimension, specifically for vision-related tensors: - - pixel_values.shape[0] = sum(sequence_lengths for all image samples) - - image_grid_thw.shape[0] = sum(num_images for all samples) - - Similar patterns for video tensors - - Args: - expand_size (int): Factor by which to expand inputs (for beam search, etc.) - is_encoder_decoder (bool): Whether using encoder-decoder architecture - input_ids (torch.LongTensor, optional): Input token IDs - **model_kwargs: Additional model arguments to expand - - Returns: - tuple: (expanded_input_ids, expanded_model_kwargs) - """ if expand_size == 1: return input_ids, model_kwargs - # Define keys for vision-related tensors that need special handling visual_keys = [ "pixel_values", "image_grid_thw", @@ -2622,13 +2551,11 @@ class Qwen2_5_VLMoEForAction( ] def _expand_dict_for_generation_visual(dict_to_expand): - """Expand vision-related tensors based on image/video counts per sample.""" image_grid_thw = model_kwargs.get("image_grid_thw", None) video_grid_thw = model_kwargs.get("video_grid_thw", None) image_nums, video_nums = self._get_image_nums_and_video_nums(input_ids) def _repeat_interleave_samples(x, lengths, repeat_times): - """Split tensor by lengths and repeat each sample.""" samples = torch.split(x, lengths) repeat_args = [repeat_times] + [1] * (x.dim() - 1) result = torch.cat( @@ -2638,33 +2565,31 @@ class Qwen2_5_VLMoEForAction( for key in dict_to_expand: if key == "pixel_values": - # Split images into samples and compute sequence lengths + # split images into samples samples = torch.split(image_grid_thw, list(image_nums)) + # compute the sequence length of images for each sample lengths = [torch.prod(sample, dim=1).sum() for sample in samples] dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "image_grid_thw": - # Expand based on number of images per sample + # get the num of images for each sample lengths = list(image_nums) dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "pixel_values_videos": - # Split videos into samples and compute sequence lengths samples = torch.split(video_grid_thw, list(video_nums)) lengths = [torch.prod(sample, dim=1).sum() for sample in samples] dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "video_grid_thw": - # Expand based on number of videos per sample lengths = list(video_nums) dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "second_per_grid_ts": - # Handle list-type temporal grid data if not isinstance(dict_to_expand[key], list): raise TypeError( f"Expected value for key '{key}' to be a list, but got {type(dict_to_expand[key])} instead." @@ -2678,7 +2603,6 @@ class Qwen2_5_VLMoEForAction( return dict_to_expand def _expand_dict_for_generation(dict_to_expand): - """Expand standard tensors using repeat_interleave.""" for key in dict_to_expand: if ( key != "cache_position" @@ -2691,19 +2615,16 @@ class Qwen2_5_VLMoEForAction( ) return dict_to_expand - # Expand visual inputs only if input_ids is available for counting images/videos - # If input_ids is unavailable, visual inputs won't be used, so no expansion needed + # input_ids is required for expanding visual inputs + # If input_ids is unavailable, visual inputs will not be used; therefore, there is no need to expand visual inputs. if input_ids is not None and input_ids.numel() != 0: model_kwargs = _expand_dict_for_generation_visual(model_kwargs) - # Expand input_ids using standard repeat_interleave if input_ids is not None: input_ids = input_ids.repeat_interleave(expand_size, dim=0) - # Expand all other model arguments model_kwargs = _expand_dict_for_generation(model_kwargs) - # Handle encoder-decoder specific expansion if is_encoder_decoder: if model_kwargs.get("encoder_outputs") is None: raise ValueError( @@ -2714,3 +2635,431 @@ class Qwen2_5_VLMoEForAction( ) return input_ids, model_kwargs + + def rename_vlm_weights_for_vla(self, merged_weights): + renamed = {} + + for key, value in merged_weights.items(): + + if key.startswith("model.layers") and "mlp." in key and self.config.mlp_moe: + layer_num = key.split(".layers.")[1].split(".mlp")[0] + new_key = key.replace( + f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0." + ) + renamed[new_key] = value + continue + + if ( + key.startswith("model.layers") + and "self_attn." in key + and self.config.attention_moe + ): + layer_num = key.split(".layers.")[1].split(".self_attn")[0] + proj_types = ["q_proj", "k_proj", "v_proj", "o_proj"] + for proj in proj_types: + if proj in key: + new_key = key.replace( + f"layers.{layer_num}.self_attn.{proj}", + f"layers.{layer_num}.self_attn.{proj}_experts.0", + ) + renamed[new_key] = value + break + continue + + if self.config.norm_moe and ".input_layernorm." in key: + renamed[key.replace("input_layernorm", "input_layernorms.0")] = value + continue + if self.config.norm_moe and ".post_attention_layernorm." in key: + renamed[ + key.replace( + "post_attention_layernorm", "post_attention_layernorms.0" + ) + ] = value + continue + if self.config.norm_moe and ".norm." in key: + renamed[key.replace("norm", "norms.0")] = value + continue + + renamed[key] = value + + fused = Qwen2_5_VLMoEForAction.convert_to_fused(renamed) + + return fused + + @staticmethod + def fuse_gate_up( + fused, prefix, suffix_gate="gate_proj", suffix_up="up_proj", out="gate_up_proj" + ): + gate_w = fused.get(prefix + f"{suffix_gate}.weight") + up_w = fused.get(prefix + f"{suffix_up}.weight") + gate_b = fused.get(prefix + f"{suffix_gate}.bias") + up_b = fused.get(prefix + f"{suffix_up}.bias") + + # Skip fusion if gate_up_proj already exists + if prefix + f"{out}.weight" in fused: + return + + if gate_w is not None and up_w is not None: + fused[prefix + f"{out}.weight"] = torch.cat([gate_w, up_w], dim=0) + if gate_b is not None and up_b is not None: + fused[prefix + f"{out}.bias"] = torch.cat([gate_b, up_b], dim=0) + + # Remove the old modules + for n in [ + f"{suffix_gate}.weight", + f"{suffix_up}.weight", + f"{suffix_gate}.bias", + f"{suffix_up}.bias", + ]: + full = prefix + n + if full in fused: + del fused[full] + + @staticmethod + def fuse_qkv(fused, prefix, *, experts=False, expert_id=None): + """ + prefix: + - non-MoE: model.layers.X.self_attn. + - MoE: model.layers.X.self_attn. + + When experts=True, expert_id must be provided + """ + + if experts: + assert expert_id is not None, "expert_id must be provided for MoE attention" + + eid = expert_id + + q_w = fused.get(prefix + f"q_proj_experts.{eid}.weight") + k_w = fused.get(prefix + f"k_proj_experts.{eid}.weight") + v_w = fused.get(prefix + f"v_proj_experts.{eid}.weight") + + q_b = fused.get(prefix + f"q_proj_experts.{eid}.bias") + k_b = fused.get(prefix + f"k_proj_experts.{eid}.bias") + v_b = fused.get(prefix + f"v_proj_experts.{eid}.bias") + + out_w = f"qkv_proj_experts.{eid}.weight" + out_b = f"qkv_proj_experts.{eid}.bias" + + remove_list = [ + f"q_proj_experts.{eid}.weight", + f"k_proj_experts.{eid}.weight", + f"v_proj_experts.{eid}.weight", + f"q_proj_experts.{eid}.bias", + f"k_proj_experts.{eid}.bias", + f"v_proj_experts.{eid}.bias", + ] + + else: + q_w = fused.get(prefix + "q_proj.weight") + k_w = fused.get(prefix + "k_proj.weight") + v_w = fused.get(prefix + "v_proj.weight") + + q_b = fused.get(prefix + "q_proj.bias") + k_b = fused.get(prefix + "k_proj.bias") + v_b = fused.get(prefix + "v_proj.bias") + + out_w = "qkv_proj.weight" + out_b = "qkv_proj.bias" + + remove_list = [ + "q_proj.weight", + "k_proj.weight", + "v_proj.weight", + "q_proj.bias", + "k_proj.bias", + "v_proj.bias", + ] + + if q_w is None or k_w is None or v_w is None: + return + + fused[prefix + out_w] = torch.cat([q_w, k_w, v_w], dim=0) + + if q_b is not None and k_b is not None and v_b is not None: + fused[prefix + out_b] = torch.cat([q_b, k_b, v_b], dim=0) + + for n in remove_list: + full = prefix + n + if full in fused: + del fused[full] + + @staticmethod + def is_fused(state_dict): + # Check whether the model is already fused by looking for fused weight markers + return any(".moe.experts.0.gate_up_proj" in key for key in state_dict.keys()) + + @staticmethod + def convert_to_fused(state_dict): + """ + Convert an unfused checkpoint to a fused checkpoint + """ + fused = {} + + # ========================= + # Phase 1: copy everything first + # ========================= + for k, v in state_dict.items(): + fused[k] = v + + # ========================= + # Phase 2: fuse gate + up (MoE aware) + # ========================= + for key in list(fused.keys()): + # language MoE mlp (ANY expert) + m = re.match( + r"(model\.layers\.\d+\.moe\.experts\.\d+\.)gate_proj\.weight", + key, + ) + if m: + prefix = m.group(1) + Qwen2_5_VLMoEForAction.fuse_gate_up(fused, prefix) + continue + + # language non-MoE mlp + if re.match(r"(model\.layers\.\d+\.mlp\.)gate_proj\.weight", key): + prefix = key.replace("gate_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_gate_up(fused, prefix) + continue + + # visual mlp + if re.match(r"(visual\.blocks\.\d+\.mlp\.)gate_proj\.weight", key): + prefix = key.replace("gate_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_gate_up(fused, prefix) + + # ========================= + # Phase 3: fuse attention qkv (MoE aware) + # ========================= + for key in list(fused.keys()): + # language MoE attention: ANY expert + m = re.match( + r"(model\.layers\.\d+\.self_attn\.)q_proj_experts\.(\d+)\.weight", + key, + ) + if m: + prefix = m.group(1) + expert_id = int(m.group(2)) + Qwen2_5_VLMoEForAction.fuse_qkv( + fused, + prefix, + experts=True, + expert_id=expert_id, + ) + continue + + # language non-MoE attention + if re.match(r"(model\.layers\.\d+\.self_attn\.)q_proj\.weight", key): + prefix = key.replace("q_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_qkv( + fused, + prefix, + experts=False, + ) + continue + + # visual attention + if re.match(r"(visual\.blocks\.\d+\.attn\.)q_proj\.weight", key): + prefix = key.replace("q_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_qkv( + fused, + prefix, + experts=False, + ) + + # ========================= + # Phase 4: sanity check(optional) + # ========================= + for k in fused.keys(): + assert not any( + x in k + for x in [ + ".gate_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ] + ), f"Unfused key still exists: {k}" + + return fused + + def convert_to_mix_precision(self): + # Mix Precision + params_to_keep_float32 = [] + for name, _ in self.named_parameters(): + if any( + k in name + for k in [ + "input_layernorm", + "post_attention_layernorm", + "model.norm", + "action_preprocessor", + ] + ): + params_to_keep_float32.append(name) + for name, param in self.named_parameters(): + if name not in params_to_keep_float32: + param.data = param.data.to(torch.bfloat16) + if name in params_to_keep_float32: + param.data = param.data.to(torch.float32) + + def convert_to_fsdp( + self, + *, + mesh, + mp_policy, + offload_policy=None, + reshard_after_forward: bool = True, + use_dmuon: bool = False, + ): + """Wrap with FSDP2: per-decoder-layer + per-vision-block fully_shard + with bf16 ``mp_policy``; layernorms and ``action_preprocessor`` leaves + get their own nested ``fully_shard`` with an fp32 policy so their + forward stays in fp32 (matching the old FSDP1 NO_SHARD + fp32_policy + layout). + + ``mesh`` and ``mp_policy`` are required and produced by + ``FSDPStrategy._build_fsdp2_layout``. The pre-migration FSDP1 wrap + used ``ShardingStrategy.SHARD_GRAD_OP`` regardless of the yaml + config; the new path honors ``fsdp_sharding_strategy`` (default + ``full_shard`` -> ``reshard_after_forward=True``). Active runs that + relied on the implicit SHARD_GRAD_OP throughput should set + ``distributed.fsdp_sharding_strategy: shard_grad_op`` explicitly. + + When ``use_dmuon=True``, DMuon owns selected trainable + matrix weights before the remaining parameters are wrapped by FSDP2. + """ + from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy + + if use_dmuon: + import dmuon + + # DMuon keeps fp32 master weights for dedicated params; symmetric + # FSDP2 params still use MixedPrecisionPolicy for bf16 compute. + for param in self.parameters(): + if param.requires_grad and param.dtype != torch.float32: + param.data = param.data.float() + + action_preprocessor_linear_ids = { + id(module) + for module in self.action_preprocessor.modules() + if isinstance(module, nn.Linear) + } + + def hook_boundary(module: nn.Module) -> bool: + # ActionProcessor helper methods call internal Linear layers + # directly, bypassing ActionProcessor.forward(). Those internal + # Linear layers need their own hooks. + if id(module) in action_preprocessor_linear_ids: + return True + return isinstance( + module, + ( + Qwen2_5_VLDecoderLayer_with_MoE, + Qwen2_5_VLVisionBlock, + Qwen2_5_VLPatchMerger, + type(self.action_preprocessor), + ), + ) + + assignment = dmuon.dedicate_params( + self, + mesh, + predicate=_is_qwen25_dmuon_target_param, + hook_boundary_predicate=hook_boundary, + hook_boundary_strict=True, + compute_dtype=torch.bfloat16, + reshard_after_forward=reshard_after_forward, + ) + logger.info( + "[Qwen2.5 FSDP2 + DMuon] dedicate_params: %d params assigned", + len(assignment), + ) + + fp32_mp = MixedPrecisionPolicy( + param_dtype=torch.float32, + reduce_dtype=torch.float32, + ) + shard_kwargs = dict( + mesh=mesh, + mp_policy=mp_policy, + reshard_after_forward=reshard_after_forward, + ) + if offload_policy is not None: + shard_kwargs["offload_policy"] = offload_policy + fp32_shard_kwargs = dict(shard_kwargs, mp_policy=fp32_mp) + + def fully_shard_fp32_leaf_or_container(name: str, module: nn.Module) -> None: + # FSDP2 cannot shard containers without forward(), e.g. ModuleList. + # When a norm match resolves to a container, shard only norm leaves. + def is_norm_leaf(leaf_name: str, leaf: nn.Module) -> bool: + return ( + "norm" in leaf_name.lower() or "norm" in type(leaf).__name__.lower() + ) + + if isinstance(module, (nn.ModuleList, nn.ModuleDict)): + for leaf_name, leaf in module.named_modules(): + if not leaf_name: + continue + if any(True for _ in leaf.children()): + continue + if not is_norm_leaf(leaf_name, leaf): + continue + logger.info( + "[FSDP2] fully_shard fp32 layernorm: %s.%s", + name, + leaf_name, + ) + fully_shard(leaf, **fp32_shard_kwargs) + return + + logger.info("[FSDP2] fully_shard fp32 layernorm: %s", name) + fully_shard(module, **fp32_shard_kwargs) + + # Layernorm leaves inside decoder layers / vision blocks / model.norm + # - wrap with fp32 mp_policy BEFORE wrapping their parents so the + # nested policy survives. Match-by-FQN mirrors the FSDP1 logic. + for module_name, module in list(self.named_modules()): + for child_name, child in list(module.named_children()): + if any( + k in child_name.lower() + for k in ("input_layernorm", "post_attention_layernorm") + ) or ("norm" in child_name.lower() and module_name.endswith("model")): + fully_shard_fp32_leaf_or_container( + f"{module_name}.{child_name}", child + ) + + # action_preprocessor leaves: same fp32 treatment as layernorms. + for name, module in list(self.named_modules()): + if "action_preprocessor" not in name.lower(): + continue + if any(True for _ in module.children()): + continue # only leaves + logger.info( + "[FSDP2] fully_shard fp32 action_preprocessor leaf: %s", + name, + ) + fully_shard(module, **fp32_shard_kwargs) + + # Each transformer decoder block: bf16 mp_policy from strategy. + for idx, layer in enumerate(self.model.layers): + fully_shard(layer, **shard_kwargs) + if idx == 0: + logger.info( + "[FSDP2] fully_shard model.layers.* (bf16, " + "reshard_after_forward=%s)", + reshard_after_forward, + ) + + # Each vision tower block: same. + if hasattr(self, "visual") and hasattr(self.visual, "blocks"): + for idx, block in enumerate(self.visual.blocks): + fully_shard(block, **shard_kwargs) + if idx == 0: + logger.info("[FSDP2] fully_shard visual.blocks.* (bf16)") + + # Root: covers lm_head, embed_tokens, and any params not in a + # nested fully_shard. Inner fp32 wraps maintain their own policy. + logger.info("[FSDP2] fully_shard root (bf16)") + fully_shard(self, **shard_kwargs) + return self diff --git a/wall_x/model/qact/tokenizer_mixin.py b/wall_x/model/qact/tokenizer_mixin.py new file mode 100644 index 0000000..221bcaa --- /dev/null +++ b/wall_x/model/qact/tokenizer_mixin.py @@ -0,0 +1,962 @@ +"""Action tokenizer mixin for loading tokenizers and action mappings.""" + +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import torch +from transformers import AutoProcessor + +from wall_x.utils.constant import is_action_dataset_name + +# Delay imports so missing optional packages do not fail at import time +try: + from spatial_tokenizer.spatial_tokenizer import SpatialActionTokenizer +except ImportError: + SpatialActionTokenizer = None + + +class ActionTokenizerMixin(ABC): + """Base class for action tokenizers""" + + def __init__(self): + self._tokenizer = None + self.action_normalizer = None + self._tokenizer_type: str = "" + self._action_mapper_cache: Optional[Dict] = None # Cached action_mapper + self.dllm = False + self.input_placeholder_flag = False + + @property + def tokenizer_type(self) -> str: + """Return the tokenizer type identifier""" + return self._tokenizer_type + + @property + def tokenizer(self): + """Return the underlying tokenizer instance""" + return self._tokenizer + + @property + def action_mapper(self) -> Optional[Dict]: + """Return the cached action_mapper""" + return self._action_mapper_cache + + @abstractmethod + def load_tokenizer(self, config: dict, normalizer, device: str = "cpu") -> Any: + """ + Load the tokenizer instance + + Args: + config: Configuration dictionary + normalizer: Normalizer + device: Device, usually "cpu" for training and "cuda" for inference + + Returns: + tokenizer instance + """ + pass + + @abstractmethod + def get_val_tokenizer(self, config: dict) -> Any: + """ + Get the tokenizer used for validation/inference + + Args: + config: Configuration dictionary + + Returns: + validation tokenizer instance + """ + pass + + @abstractmethod + def get_special_tokens(self) -> List[str]: + """ + Return special tokens to add to the vocabulary + + Returns: + list of token strings + """ + pass + + def get_all_special_tokens(self) -> Tuple[List[str], Optional[List[str]]]: + """ + Return all special tokens and keep <|action_token_0|> before AR action tokens + + Returns: + (new_tokens, special_tokens) tuple + """ + tokens, special_tokens = self.get_special_tokens() + if "<|action_token_0|>" not in tokens: + tokens.insert(0, "<|action_token_0|>") + return tokens, special_tokens + + @abstractmethod + def build_action_mapper(self, processor) -> Optional[Dict]: + """ + Build action_mapper + + Args: + processor: HuggingFace processor,used to convert tokens to IDs + + Returns: + action_mapper dictionary; format depends on tokenizer type + """ + pass + + @abstractmethod + def get_action_token_list(self, processor) -> List[int]: + """ + Get the action token ID list + + Args: + processor: HuggingFace processor + + Returns: + action token ID list + """ + pass + + @abstractmethod + def decode_action( + self, + output_ids: torch.Tensor, + action_mapper: Dict, + action_horizon: int, + action_dim: int, + device: torch.device, + proprioception: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_id: Optional[int] = None, + state: Optional[torch.Tensor] = None, + ) -> Tuple[Optional[Union[np.ndarray, torch.Tensor]], bool]: + """ + Unified decoding interface + + Args: + output_ids: model output token IDs [1, seq_len] + action_mapper: action_mapper dictionary + action_horizon: action horizon + action_dim: action dimension + device: Device + proprioception: proprioception, normalized when required by a tokenizer + dof_mask: DOF mask when required by a tokenizer + robot_type_id: robot type ID when required by a tokenizer + state: state when required by fast/spatial tokenizers + + Returns: + (predict_action, decode_success) + - predict_action: decoded action [T, action_dim] or None + - decode_success: whether decoding succeeded + """ + pass + + @abstractmethod + def compute_accuracy( + self, + logits: torch.Tensor, + labels: torch.Tensor, + action_mapper: Dict, + action_token_id_set: Dict, + ) -> Dict[str, torch.Tensor]: + """ + Compute accuracy metrics + + Args: + logits: model output logits + labels: Labels + action_mapper: action_mapper dictionary + action_token_id_set: set of action token IDs + + Returns: + accuracy metric dictionary, such as {"action_accuracy": tensor, ...} + """ + pass + + @abstractmethod + def get_accuracy_keys(self) -> List[str]: + """ + Return accuracy metric keys for logging + + Returns: + list of key strings + """ + pass + + @property + @abstractmethod + def vocab_size(self) -> int: + """Return the vocabulary size""" + pass + + @property + @abstractmethod + def uses_dof_mask_for_unnorm(self) -> bool: + """ + Whether unnormalization requires dof_mask + + Returns: + True: dof_mask is required by fast/spatial tokenizers + False: dof_mask is not required and full dimensions are returned + """ + pass + + @property + @abstractmethod + def needs_action_crop(self) -> bool: + """ + Whether actions must be clipped before encoding + + Returns: + True: clip by chunk_size and dof_mask for fast/spatial tokenizers + False: no clipping; the encoder handles the full sequence internally + """ + pass + + @abstractmethod + def encode_to_tokens( + self, + actions: torch.Tensor, + obs_state: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_ids: Optional[List[int]] = None, + is_train: bool = True, + ) -> List[List[str]]: + """ + Encode actions into token strings for training data processing + + Args: + actions: Normalized actions + - fast/spatial: List[Tensor], each [T, D] (clipped) + - v3.1 delta: Tensor [B, T, D] (full sequence) + obs_state: observation state when required by a tokenizer[B, obs_horizon, D] + dof_mask: DOF mask[B, T, D] + robot_type_ids: robot type ID list when required by a tokenizer + + Returns: + List[List[str]]: token string list for each sample + """ + pass + + def init_inference(self, robot_type: Optional[str] = None) -> None: + """ + Inference initialization hook; subclasses may override + + Args: + robot_type: robot type name + """ + pass + + def prepare_action_for_ar_encoding( + self, + ar_actionchunk: torch.Tensor, + dataset_names: List[str], + agent_pos: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Prepare action data before AR encoding; subclasses may override + + SpatialVLA needs waypoint selection; other tokenizers return normalized_action directly + + Args: + ar_actionchunk: [B, T, D] actions before normalization + dataset_names: [str] dataset name list + agent_pos: [B, T, D] agent positions required by SpatialVLA + + Returns: + prepared action data + """ + if self.action_normalizer is not None: + ar_actionchunk = self.action_normalizer.normalize_data( + ar_actionchunk, dataset_names + ) + return ar_actionchunk + + def get_robot_type_ids( + self, + uids: List[Optional[str]], + dataset_names: List[str], + ) -> Optional[List[int]]: + """ + Get robot_type_ids; subclasses may override + + Only the v3.1 delta tokenizer needs this; other tokenizers return None + + Args: + uids: UID list + dataset_names: dataset name list + + Returns: + robot_type_id list or None + """ + return None + + def modify_inputs_for_dllm( + self, + inputs: Dict[str, torch.Tensor], + processor, + sample_time: torch.Tensor, + dataset_names: List[str], + ): + return NotImplementedError + + +class FastTokenizerMixin(ActionTokenizerMixin): + """Fast tokenizer implementation""" + + def __init__(self): + super().__init__() + self._tokenizer_type = "fast" + + def load_tokenizer(self, config: dict, normalizer, device: str = "cpu") -> Any: + """Load the fast tokenizer""" + self._tokenizer = AutoProcessor.from_pretrained( + config["action_tokenizer_path"], trust_remote_code=True + ) + self.action_normalizer = normalizer + return self._tokenizer + + def get_val_tokenizer(self, config: dict) -> Any: + return self._tokenizer + + def get_special_tokens(self) -> List[str]: + """Return special tokens for the fast tokenizer""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + # tokens = ["<|ar_action|>", "<|ar_pad|>"] #temporary compatibility with existing checkpoints + tokens = [] + for i in range(self._tokenizer.vocab_size): + tokens.append(f"<|action_token_{i}|>") + return tokens, None + + def build_action_mapper(self, processor) -> Dict[int, int]: + """ + Build the cached action_mapper for the fast tokenizer + + Returns: + Dict[token_id, action_idx] + """ + # Return cached value if available + if self._action_mapper_cache is not None: + return self._action_mapper_cache + + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + action_mapper = {} + for i in range(self._tokenizer.vocab_size): + token = f"<|action_token_{i}|>" + token_id = processor.tokenizer.convert_tokens_to_ids(token) + action_mapper[token_id] = i + + self._action_mapper_cache = action_mapper + return action_mapper + + def get_action_token_list(self, processor) -> List[int]: + """Get the action token ID list""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + action_token_list = [] + for i in range(self._tokenizer.vocab_size): + token_id = processor.tokenizer.convert_tokens_to_ids( + f"<|action_token_{i}|>" + ) + action_token_list.append(token_id) + return action_token_list + + def decode_action( + self, + output_ids: torch.Tensor, + action_mapper: Dict, + action_horizon: int, + action_dim: int, + device: torch.device, + proprioception: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_id: Optional[int] = None, + state: Optional[torch.Tensor] = None, + ) -> Tuple[Optional[Union[np.ndarray, torch.Tensor]], bool]: + """Fast tokenizer decoding""" + action_id = [] + for token_id_i in output_ids[0]: + if token_id_i.item() in action_mapper: + action_id.append(action_mapper[token_id_i.item()]) + + if len(action_id) == 0: + return np.zeros((action_horizon, action_dim)), False + + predict_action = self._tokenizer.decode( + [action_id], time_horizon=action_horizon, action_dim=action_dim + ) + + # Check whether decoding succeeded + decode_success = False + if isinstance(predict_action, np.ndarray): + decode_success = np.sum(predict_action) != 0 + elif isinstance(predict_action, torch.Tensor): + decode_success = predict_action.sum().item() != 0 + + return predict_action, decode_success + + def compute_accuracy( + self, + logits: torch.Tensor, + labels: torch.Tensor, + action_mapper: Dict, + action_token_id_set: Dict, + ) -> Dict[str, torch.Tensor]: + """Compute fast tokenizer accuracy""" + result = {} + + if len(action_token_id_set.get("action_token_list", [])) > 0: + shift_logits = logits[..., :-1, :].contiguous() + action_preds = shift_logits.argmax(dim=-1) + shift_labels = labels[..., 1:].contiguous() + action_mask = shift_labels > action_token_id_set["action_token_list"][0] + correct_preds = (action_preds == shift_labels) & action_mask + action_accuracy = correct_preds.sum().float() / action_mask.sum().float() + result["action_accuracy"] = action_accuracy + + return result + + def get_accuracy_keys(self) -> List[str]: + """Return fast tokenizer accuracy keys""" + return ["action_accuracy"] + + @property + def vocab_size(self) -> int: + if self._tokenizer is None: + return 0 + return self._tokenizer.vocab_size + + @property + def uses_dof_mask_for_unnorm(self) -> bool: + """fast tokenizer requires dof_mask""" + return True + + @property + def needs_action_crop(self) -> bool: + return True + + @property + def inference_ar_steps_for_dllm(self) -> int: + return self.max_length + + def encode_to_tokens( + self, + actions: List, + obs_state: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_ids: Optional[List[int]] = None, + is_train: bool = True, + ) -> List[List[str]]: + """ + Fast tokenizer encoding + + Args: + actions: List[Tensor/ndarray], each [T, D] (clipped) + """ + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + all_action_tokens = [] + for i in range(len(actions)): + action = actions[i] + if isinstance(action, torch.Tensor): + action = action.cpu().numpy() + token_id = self._tokenizer(action) + action_tokens = [f"<|action_token_{idx}|>" for idx in token_id[0]] + all_action_tokens.append(action_tokens) + return all_action_tokens + + def modify_inputs_for_dllm( + self, + inputs: Dict[str, torch.Tensor], + processor, + sample_time: torch.Tensor, + dataset_names: List[str], + ): + # Untested + input_ids = inputs["input_ids"] + labels = inputs["labels"] + prefix_length = inputs["prefix_length"] + bs, seqlen = input_ids.shape + + ar_token_length = self.max_length + ar_step_num = ar_token_length + + device = input_ids.device + dtype = input_ids.dtype + + placeholder_ids = torch.tensor( + processor.placeholder_seq, device=device, dtype=dtype + ) + + if not torch.is_tensor(sample_time): + sample_time = torch.tensor(sample_time, device=device, dtype=torch.float32) + else: + sample_time = sample_time.to(device=device, dtype=torch.float32) + + sample_time = sample_time.clamp(0.0, 1.0) + + noisy_steps_per_sample = ( + torch.ceil((1.0 - sample_time) * (ar_step_num + 1)).long() - 1 + ) + noisy_steps_per_sample = noisy_steps_per_sample.clamp(min=0, max=ar_step_num) + start = prefix_length - ar_token_length - 2 + end = prefix_length - 2 + + # Prepare the noise sequence for each sample + noise_seqs = placeholder_ids.repeat(bs, 1) + + ar_len = end - start + if noise_seqs.size(1) != ar_len: + # These should usually match; defensively truncate to ar_len + noise_seqs = noise_seqs[:, :ar_len] + rand = torch.rand(bs, ar_len, device=device) + perm = rand.argsort(dim=-1) + ranks = perm.argsort(dim=-1) + noisy_mask = ranks < noisy_steps_per_sample.view(-1, 1) + + ar_input = input_ids[:, start:end] + ar_labels = labels[:, start + 1 : end + 1] + ar_input[noisy_mask] = noise_seqs[noisy_mask] + ar_labels[~noisy_mask] = -100 + + inputs["input_ids"] = input_ids + inputs["labels"] = labels + return inputs + + def update_placeholder_mask(self, processor, prefix_length, input_ids): + # Untested + ar_action_mask = torch.zeros_like(input_ids) + + inc = torch.arange( + 1, + self.max_length + 1, + ) + inc = inc.unsqueeze(0).expand(ar_action_mask.size(0), -1) # [bs, ar_len] + ar_action_mask[ + :, + prefix_length - self.max_length - 2 : prefix_length - 2, + ] = inc + ar_action_mask[:, prefix_length - 2 : prefix_length] = -1 # eos + + return {"ar_action_mask": ar_action_mask} + + def get_placeholder_for_dllm(self): + placeholder_seq = ["<|ar_action|>"] * self.max_length + return placeholder_seq + + +class SpatialVLATokenizerMixin(ActionTokenizerMixin): + """SpatialVLA tokenizer implementation""" + + def __init__(self): + super().__init__() + self._tokenizer_type = "spatialvla" + + def load_tokenizer(self, config: dict, normalizer, device: str = "cpu") -> Any: + """Load the SpatialVLA tokenizer""" + if SpatialActionTokenizer is None: + raise ImportError( + "SpatialActionTokenizer is not installed. " + "Please install spatial_tokenizer package." + ) + self._tokenizer = SpatialActionTokenizer( + normalizer=normalizer, + augment_ratio=config.get("augment_ratio", 0.0), + max_waypoints=config.get("max_waypoints", 5), + with_gripper=config.get("with_gripper", True), + single_arm=config.get("single_arm", False), + ) + self._val_tokenizer = None + self.config = config + self.dllm = config.get("dllm", False) + self.input_placeholder_flag = config.get("input_placeholder_flag", False) + self.action_normalizer = normalizer + self.with_gripper = config.get("with_gripper", True) + return self._tokenizer + + def get_placeholder_for_dllm(self): + if self.with_gripper: + placeholder_seq = [ + "<|left_xyz|>", + "<|left_rpy|>", + "<|left_gripper|>", + "<|right_xyz|>", + "<|right_rpy|>", + "<|right_gripper|>", + ] + else: + placeholder_seq = [ + "<|left_xyz|>", + "<|left_rpy|>", + "<|right_xyz|>", + "<|right_rpy|>", + ] + if self._tokenizer.single_arm: + placeholder_seq = placeholder_seq[len(placeholder_seq) // 2 :] + placeholder_seq = placeholder_seq * self._tokenizer.max_waypoints + return placeholder_seq + + def get_val_tokenizer(self, config: dict) -> Any: + if self._val_tokenizer: + return self._val_tokenizer + + self._val_tokenizer = SpatialActionTokenizer( + normalizer=self.action_normalizer, + augment_ratio=0, + max_waypoints=self.config.get("max_waypoints", 5), + with_gripper=self.config.get("with_gripper", True), + single_arm=self.config.get("single_arm", False), + ) + return self._val_tokenizer + + def get_special_tokens(self) -> List[str]: + """Return special tokens for the SpatialVLA tokenizer""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + tokens = [ + "<|step|>", + "<|left|>", + "<|right|>", + "<|move|>", + ] # only for compatibility with existing checkpoints + if self.input_placeholder_flag: + special_tokens = [ + "<|left_xyz|>", + "<|left_rpy|>", + "<|left_gripper|>", + "<|right_xyz|>", + "<|right_rpy|>", + "<|right_gripper|>", + ] + tokens += special_tokens + if not self._tokenizer.with_gripper: + indices = [0, 1, 3, 4] + special_tokens = [special_tokens[i] for i in indices] + if self._tokenizer.single_arm: + special_tokens = special_tokens[len(special_tokens) // 2 :] + for i in range(self._tokenizer.vocab_size): + tokens.append(f"<|action_token_{i}|>") + + return tokens, special_tokens + + def build_action_mapper(self, processor) -> Dict[int, int]: + """ + Build the cached action_mapper for the SpatialVLA tokenizer + + Returns: + Dict[token_id, action_idx] + """ + # Return cached value if available + if self._action_mapper_cache is not None: + return self._action_mapper_cache + + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + action_mapper = {} + for i in range(self._tokenizer.vocab_size): + token = f"<|action_token_{i}|>" + token_id = processor.tokenizer.convert_tokens_to_ids(token) + action_mapper[token_id] = i + + self._action_mapper_cache = action_mapper + return action_mapper + + def get_action_token_list(self, processor) -> List[int]: + """Get the action token ID list""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + action_token_list = [] + for i in range(self._tokenizer.vocab_size): + token_id = processor.tokenizer.convert_tokens_to_ids( + f"<|action_token_{i}|>" + ) + action_token_list.append(token_id) + return action_token_list + + def decode_action( + self, + output_ids: torch.Tensor, + action_mapper: Dict, + action_horizon: int, + action_dim: int, + device: torch.device, + proprioception: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_id: Optional[int] = None, + state: Optional[torch.Tensor] = None, + ) -> Tuple[Optional[Union[np.ndarray, torch.Tensor]], bool]: + """SpatialVLA tokenizer decoding""" + action_id = [] + for token_id_i in output_ids[0]: + if token_id_i.item() in action_mapper: + action_id.append(action_mapper[token_id_i.item()]) + + if len(action_id) == 0: + return np.zeros((action_horizon, action_dim)), False + + if state is not None: + predict_action = self._tokenizer.decode( + [action_id], + state=state[0, 0, :action_dim], + time_horizon=action_horizon, + action_dim=action_dim, + ) + else: + predict_action = self._tokenizer.decode( + [action_id], time_horizon=action_horizon, action_dim=action_dim + ) + + # Check whether decoding succeeded + decode_success = False + if isinstance(predict_action, np.ndarray): + decode_success = np.sum(predict_action) != 0 + elif isinstance(predict_action, torch.Tensor): + decode_success = predict_action.sum().item() != 0 + + return predict_action, decode_success + + def compute_accuracy( + self, + logits: torch.Tensor, + labels: torch.Tensor, + action_mapper: Dict, + action_token_id_set: Dict, + ) -> Dict[str, torch.Tensor]: + """Compute SpatialVLA tokenizer accuracy, same as fast""" + result = {} + + if len(action_token_id_set.get("action_token_list", [])) > 0: + shift_logits = logits[..., :-1, :].contiguous() + action_preds = shift_logits.argmax(dim=-1) + shift_labels = labels[..., 1:].contiguous() + action_mask = shift_labels > action_token_id_set["action_token_list"][0] + correct_preds = (action_preds == shift_labels) & action_mask + action_accuracy = correct_preds.sum().float() / action_mask.sum().float() + result["action_accuracy"] = action_accuracy + + return result + + def get_accuracy_keys(self) -> List[str]: + """Return SpatialVLA tokenizer accuracy keys""" + return ["action_accuracy"] + + @property + def vocab_size(self) -> int: + if self._tokenizer is None: + return 0 + return self._tokenizer.vocab_size + + @property + def uses_dof_mask_for_unnorm(self) -> bool: + """SpatialVLA tokenizer requires dof_mask""" + return True + + @property + def needs_action_crop(self) -> bool: + return False + + @property + def inference_ar_steps_for_dllm(self) -> int: + return self._tokenizer.max_waypoints + + def prepare_action_for_ar_encoding( + self, + ar_actionchunk: torch.Tensor, + dataset_names: List[str], + agent_pos: Optional[torch.Tensor] = None, + ) -> List[torch.Tensor]: + step = ar_actionchunk.shape[1] // self._tokenizer.max_waypoints + indices = np.arange(0, ar_actionchunk.shape[1], step)[ + : self._tokenizer.max_waypoints + ] + ar_action = ar_actionchunk[:, indices, :] + ar_action = self.action_normalizer.normalize_data(ar_action, dataset_names) + return ar_action + + def encode_to_tokens( + self, + actions: List, + obs_state: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_ids: Optional[List[int]] = None, + is_train: bool = True, + ) -> List[List[str]]: + """ + SpatialVLA tokenizer encoding + + Args: + actions: List[Tensor/ndarray], each [T, D] (clipped) + """ + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + if isinstance(actions, torch.Tensor): + # Convert torch tensors to numpy arrays + actions = actions.cpu().numpy() + tokenizer = self._tokenizer if is_train else self.get_val_tokenizer({}) + token_ids_group = tokenizer.batch_encode(actions) + all_action_tokens = [] + for i in range(len(token_ids_group)): + token_ids = np.array(token_ids_group[i]).reshape(-1) + action_token = [f"<|action_token_{i}|>" for i in token_ids] + all_action_tokens.append(action_token) + return all_action_tokens + + def modify_inputs_for_dllm( + self, + inputs: Dict[str, torch.Tensor], + processor, + sample_time: torch.Tensor, + dataset_names: List[str], + ): + """ + Prepare AR DLLM inputs by encoding sample_time into input_ids and labels + + Noise injection strategy: + - Split [0, 1] into N equal parts (N = ar_step_num) + - For each sample, compute the number of noisy steps k in [1, N] + k = clamp(ceil((1 - t) * N), 1, N) + -> smaller t means more noise; values closer to 1 mean less noise, with at least one noised step + - seq is the noise sequence formed by concatenating ar_step_num placeholder_seq blocks + Each step maps to len(placeholder_seq) consecutive tokens + - Noise injection overwrites matching tokens in input_ids / labels step by step from seq + """ + + input_ids = inputs["input_ids"] + labels = inputs["labels"] + prefix_length = inputs["prefix_length"] # int or tensor(1,) + bs, seqlen = input_ids.shape + + # Total number of AR steps + ar_step_num = self._tokenizer.max_waypoints # N + step_token_len = len(processor.placeholder_seq) # Number of tokens per step + ar_token_length = ar_step_num * step_token_len # Total AR token length + + device = input_ids.device + dtype = input_ids.dtype + + placeholder_ids = torch.tensor( + processor.placeholder_seq, device=device, dtype=dtype + ) # (step_token_len,) + noise_seq = placeholder_ids.repeat(ar_step_num) # (ar_token_length,) + + if not torch.is_tensor(sample_time): + sample_time = torch.tensor(sample_time, device=device, dtype=torch.float32) + else: + sample_time = sample_time.to(device=device, dtype=torch.float32) + + sample_time = sample_time.clamp(0.0, 1.0) + + noisy_steps_per_sample = ( + torch.ceil((1.0 - sample_time) * (ar_step_num + 1)).long() - 1 + ) + noisy_steps_per_sample = noisy_steps_per_sample.clamp(min=0, max=ar_step_num) + start = prefix_length - ar_token_length - 2 + end = prefix_length - 2 + + action_idx = 0 # action sample counter + for i in range(bs): + if not is_action_dataset_name(dataset_names[i]): + continue + k = noisy_steps_per_sample[action_idx].item() + action_idx += 1 + perm = torch.randperm(ar_step_num, device=device) + chosen_steps = perm[:k] # (k,) + + step_mask = torch.zeros(ar_step_num, dtype=torch.bool, device=device) + step_mask[chosen_steps] = True + + token_mask = step_mask.repeat_interleave( + step_token_len + ) # (ar_token_length,) + ignore_mask = ~token_mask + # Take a view of the current sample's AR span for mask-based replacement + cur_input_view = input_ids[i, start:end] + cur_label_view = labels[ + i, start + 1 : end + 1 + ] # shift placeholders one position to the right + + # Overwrite selected tokens with noise + cur_input_view[token_mask] = noise_seq[token_mask] + cur_label_view[ignore_mask] = -100 + + inputs["input_ids"] = input_ids + inputs["labels"] = labels + return inputs + + def update_positional_masks_for_dllm( + self, positional_masks, inputs, processor, visible_predict_ar_ratio=1 + ): + if self.dllm and self.input_placeholder_flag: + mask = self.update_placeholder_mask( + processor, + inputs["prefix_length"], + inputs["input_ids"], + ) + positional_masks.update(mask) + if ( + np.random.rand() < visible_predict_ar_ratio + ): # FIXME ar_visible is temporarily decided per batch; mixed settings are untested and need optimization. + positional_masks["ar_visible"] = True + if "ar_predict_token_positions" in positional_masks: + del positional_masks["ar_predict_token_positions"] + # positional_masks["ar_predict_token_positions"] = None + else: + positional_masks["ar_visible"] = False + return positional_masks + + def update_placeholder_mask(self, processor, prefix_length, input_ids): + ar_action_mask = torch.zeros_like(input_ids) + current_index = 1 + step_len = len(processor.placeholder_seq) + + for b, seq in enumerate(input_ids): + for i in range(self._tokenizer.max_waypoints): + start_idx = ( + prefix_length - (self._tokenizer.max_waypoints - i) * step_len - 2 + ) + end_idx = start_idx + step_len + ar_action_mask[b, start_idx:end_idx] = current_index + current_index += 1 + ar_action_mask[b, end_idx:prefix_length] = -1 # eos + + return {"ar_action_mask": ar_action_mask} + + +# ============================================================ +# Factory function +# ============================================================ + +_TOKENIZER_REGISTRY: Dict[str, type] = { + "fast": FastTokenizerMixin, + "spatialvla": SpatialVLATokenizerMixin, +} + + +def get_action_tokenizer_mixin(tokenizer_type: str) -> ActionTokenizerMixin: + """ + Get the mixin instance for a tokenizer type + + Args: + tokenizer_type: tokenizer type, supports "fast", "spatialvla" + + Returns: + ActionTokenizerMixin instance + + Raises: + ValueError: Unsupported tokenizer type + """ + if tokenizer_type not in _TOKENIZER_REGISTRY: + raise ValueError( + f"Unsupported action tokenizer type: {tokenizer_type}. " + f"Supported types: {list(_TOKENIZER_REGISTRY.keys())}" + ) + return _TOKENIZER_REGISTRY[tokenizer_type]() diff --git a/wall_x/model/qwen2_5_based/__init__.py b/wall_x/model/qwen2_5_based/__init__.py deleted file mode 100644 index d9ed878..0000000 --- a/wall_x/model/qwen2_5_based/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from .modeling_qwen2_5_vl_act import Qwen2_5_VLMoEModel, Qwen2_5_VLMoEForAction -from .configuration_qwen2_5_vl import Qwen2_5_VLConfig - -__all__ = [ - "Qwen2_5_VLMoEModel", - "Qwen2_5_VLMoEForAction", - "Qwen2_5_VLConfig", -] diff --git a/wall_x/model/registry.py b/wall_x/model/registry.py new file mode 100644 index 0000000..250d8e8 --- /dev/null +++ b/wall_x/model/registry.py @@ -0,0 +1,64 @@ +"""Decorator-based model adapter registry.""" + +from __future__ import annotations + +from typing import Dict, Type + +# Global registry mapping model_type string -> adapter class +_MODEL_REGISTRY: Dict[str, Type] = {} + + +def register_model(model_type: str): + """Decorator to register an adapter class for a given model_type. + + Can be stacked to register the same class for multiple model_types: + + @register_model("qwen2_5") + class QActAdapter: ... + """ + + def decorator(cls): + if model_type in _MODEL_REGISTRY: + existing = _MODEL_REGISTRY[model_type] + if existing is not cls: + raise ValueError( + f"model_type '{model_type}' already registered to " + f"{existing.__name__}, cannot re-register to {cls.__name__}" + ) + _MODEL_REGISTRY[model_type] = cls + return cls + + return decorator + + +def get_adapter(model_type: str, **kwargs): + """Instantiate the adapter registered for model_type. + + Args: + model_type: Registered model type string (e.g. "qwen2_5"). + **kwargs: Passed to the adapter constructor. + + Returns: + An adapter instance. + + Raises: + KeyError: If model_type is not registered. + """ + if model_type not in _MODEL_REGISTRY: + available = ", ".join(sorted(_MODEL_REGISTRY.keys())) or "(none)" + raise KeyError( + f"Unknown model_type '{model_type}'. " + f"Available: {available}. " + f"Did you forget to import the adapter module?" + ) + return _MODEL_REGISTRY[model_type](**kwargs) + + +def list_registered_models() -> list: + """Return sorted list of registered model_type strings.""" + return sorted(_MODEL_REGISTRY.keys()) + + +def clear_registry(): + """Clear all registrations. Intended for testing only.""" + _MODEL_REGISTRY.clear() diff --git a/wall_x/model/vla_mixin.py b/wall_x/model/vla_mixin.py deleted file mode 100644 index 1c98aa2..0000000 --- a/wall_x/model/vla_mixin.py +++ /dev/null @@ -1,987 +0,0 @@ -import torch -import torch.nn as nn -import torch.utils.checkpoint as cp - -from torch.distributed.fsdp import MixedPrecision as MP -from torch.distributed.fsdp import FullyShardedDataParallel as FSDP - -from wall_x.fusions import ops - -from peft import LoraConfig, get_peft_model -from typing import Optional, Union, Dict -from packaging import version - -from transformers import GenerationMixin -from transformers.activations import ACT2FN -from transformers.modeling_utils import AttentionInterface - -from transformers.utils import logging, is_torch_xla_available - -from wall_x.model.action_head import ActionProcessor -from wall_x.model.model_utils import find_first_last_ones - -ALL_ATTENTION_FUNCTIONS: AttentionInterface = AttentionInterface() -logger = logging.get_logger(__name__) - - -X2ROBOT_ATTENTION_FUNCTIONS = [] -ATTENTION_TYPES_WITH_2D_MASK = [ - "sdpa", -] -ATTENTION_TYPES_WITH_FLASH_MASK = [] - - -class TokenTypeRouter(nn.Module): - def __init__(self, num_experts: int): - super().__init__() - self.num_experts = num_experts - - def forward(self, token_types: torch.Tensor) -> torch.Tensor: - """ - Assigns tokens to different experts based on `token_type`. - Args: - token_types (torch.Tensor): A tensor of shape (batch_size, seq_length) representing the type of each token. - - Returns: - experts_indices (torch.Tensor): A tensor of shape (batch_size, seq_length) representing the expert index assigned to each token. - """ - experts_indices = token_types % self.num_experts - return experts_indices - - -class BlockSparseMLP(nn.Module): - def __init__(self, config, use_selective_recompute: bool = False): - super().__init__() - self.hidden_size = config["hidden_size"] - self.intermediate_size = config["intermediate_size"] - self.hidden_act = config["hidden_act"] - - self.use_selective_recompute = use_selective_recompute - - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - - self.act_fn = ACT2FN[self.hidden_act] - - def _full_mlp(self, hidden_state): - gate_out = self.gate_proj(hidden_state) - up_out = self.up_proj(hidden_state) - act_out = self.act_fn(gate_out) * up_out - return self.down_proj(act_out) - - def forward(self, hidden_state): - if self.use_selective_recompute: - # Perform checkpoint recalculation for the entire expert MLP. - return cp.checkpoint( - self._full_mlp, - hidden_state, - use_reentrant=False, - ) - else: - return self._full_mlp(hidden_state) - - -class SparseMoeBlock(nn.Module): - def __init__(self, config, num_experts: int, use_selective_recompute: bool = False): - super().__init__() - self.num_experts = num_experts - self.use_selective_recompute = use_selective_recompute - - # Pass the `use_selective_recompute` parameter to each expert. - self.experts = nn.ModuleList( - [ - BlockSparseMLP( - config.experts[i], use_selective_recompute=use_selective_recompute - ) - for i in range(num_experts) - ] - ) - - if not hasattr(config, "dim_inputs") or not config.dim_inputs: - raise ValueError("Configuration must contain a valid dim_inputs") - - self.dim_inputs = config.dim_inputs - self.permuted = config.mot_opt - - def forward( - self, - hidden_states: torch.Tensor, - experts_indices: torch.Tensor, - start_indices: torch.Tensor, - end_indices: torch.Tensor, - ) -> torch.Tensor: - - if self.permuted: - permuted_inputs = hidden_states - else: - batch_size, seq_length, hidden_dim = hidden_states.shape - - flat_hidden = hidden_states.reshape(-1, hidden_dim) - experts_indices = experts_indices.reshape(-1) - probs = torch.ones_like(experts_indices, dtype=torch.float32).reshape(-1, 1) - permuted_inputs, row_id_map = ops.permute(flat_hidden, experts_indices) - - # buffer - final_output = torch.zeros_like(permuted_inputs) - - # Expert forward contain selective recompute - for expert_idx, expert in enumerate(self.experts): - start, end = start_indices[expert_idx], end_indices[expert_idx] - if start == end: - continue - - dim_input = self.dim_inputs[expert_idx] - expert_input = permuted_inputs[start:end, :dim_input] - - partial_output = expert(expert_input) - final_output[start:end, :dim_input] = partial_output[:, :dim_input] - - if self.permuted: - return final_output - else: - final_output = ops.unpermute(final_output, row_id_map, probs) - return final_output.reshape(batch_size, seq_length, hidden_dim) - - -class ActionModelMixMin: - # config: Qwen2_5_VLConfig - action_preprocessor: ActionProcessor - router: TokenTypeRouter - moe: SparseMoeBlock - - def __init__(self, config, action_preprocessor, router, moe): - self.config = config - self.action_preprocessor = action_preprocessor - self.router = router - self.moe = moe - self._mot_opt_warned = False - - def set_normalizer(self, normalizer_action, normalizer_propri): - if hasattr(self, "action_preprocessor"): - self.action_preprocessor.set_normalizer( - normalizer_action, normalizer_propri - ) - else: - logger.warning( - "ActionModelMixMin.set_normalizer is called but action_preprocessor is not set" - ) - - def _apply_mlp_moe(self, hidden_states, token_types, start_indices, end_indices): - if self.config.mlp_moe: - hidden_states = self.moe( - hidden_states, token_types, start_indices, end_indices - ) - else: - hidden_states = self.mlp(hidden_states) - return hidden_states - - def _apply_norm_moe( - self, - hidden_states, - token_types, - adarms_conds, - norms, # list of norm layers (expert-wise) - norm, # shared norm if not norm_moe - start_indices=None, - end_indices=None, - use_selective_recompute=False, - ): - """ - MoE-aware LayerNorm with optional selective activation recomputation. - - Only activation math is recomputed. No GEMM is recomputed. - Safe for FSDP (use_reentrant=False). - """ - - gate = None - gate_mask = None - - # ------------------------- - # Case 1: norm_moe=True (expert-wise norm) - # ------------------------- - if self.config.norm_moe: - - # --------------------------------------------------------- - # Case 1A: mot_opt=True (segments assigned by start/end) - # --------------------------------------------------------- - if self.config.mot_opt: - new_hidden_states = torch.zeros_like(hidden_states) - - for expert_idx, expert_norm in enumerate(norms): - start = start_indices[expert_idx] - end = end_indices[expert_idx] - if start == end: - continue - - dim_input = self.config.dim_inputs[expert_idx] - selected = hidden_states[start:end] # [K, D] - - # ====== reshape if adarms on flow expert ====== - if self.config.use_adarms and expert_idx == 1: - selected = selected.view( - -1, - self.config.action_horizon_flow, - selected.shape[-1], - ) - input_slice = selected[:, :, :dim_input] - cond = adarms_conds[expert_idx] - else: - input_slice = selected[:, :dim_input] - cond = adarms_conds[expert_idx] - - if use_selective_recompute: - - def norm_chunk(t_x, t_cond, expert_norm=expert_norm): - if t_cond is None or ( - isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0 - ): - out, _ = expert_norm(t_x) - else: - out, _ = expert_norm(t_x, t_cond) - return out - - cond_for_cp = ( - cond - if cond is not None - else torch.empty(0, device=input_slice.device) - ) - processed = cp.checkpoint( - norm_chunk, - input_slice, - cond_for_cp, - use_reentrant=False, - ) - else: - processed, gate = expert_norm(input_slice, cond) - - # reshape back if needed - if self.config.use_adarms and expert_idx == 1: - processed = processed.view(-1, dim_input) - - new_hidden_states[start:end, :dim_input] = processed.to( - hidden_states.dtype - ) - - hidden_states = new_hidden_states - - # --------------------------------------------------------- - # Case 1B: mot_opt=False (token-level mask) - # --------------------------------------------------------- - else: - - new_hidden_states = torch.zeros_like(hidden_states) - B, S, D = hidden_states.shape - - for expert_idx, expert_norm in enumerate(norms): - mask = token_types == expert_idx - if mask.sum() == 0: - continue - - dim_input = self.config.dim_inputs[expert_idx] - selected = hidden_states[mask] # [K, D] - - if self.config.use_adarms and expert_idx == 1: - gate_mask = mask - selected = selected.view( - -1, - self.config.action_horizon_flow, - selected.shape[-1], - ) - input_slice = selected[:, :, :dim_input] - cond = adarms_conds[expert_idx] - else: - input_slice = selected[:, :dim_input] - cond = adarms_conds[expert_idx] - - if use_selective_recompute: - - def norm_chunk(t_x, t_cond, expert_norm=expert_norm): - if t_cond is None or ( - isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0 - ): - out, _ = expert_norm(t_x) - else: - out, _ = expert_norm(t_x, t_cond) - return out - - cond_for_cp = ( - cond - if cond is not None - else torch.empty(0, device=input_slice.device) - ) - - processed = cp.checkpoint( - norm_chunk, - input_slice, - cond_for_cp, - use_reentrant=False, - ) - else: - processed, gate = expert_norm(input_slice, cond) - - if self.config.use_adarms and expert_idx == 1: - processed = processed.view(-1, dim_input) - - # scatter back - b_id, s_id = torch.where(mask) - new_hidden_states[b_id, s_id, :dim_input] = processed.to( - hidden_states.dtype - ) - - hidden_states = new_hidden_states - - # ------------------------- - # Case 2: norm_moe=False (single LN) - # ------------------------- - else: - - def norm_chunk_shared(t_x, dummy, norm_module=norm): - out, _ = norm_module(t_x) - return out - - if use_selective_recompute: - dummy = torch.empty(0, device=hidden_states.device) - hidden_states = cp.checkpoint( - norm_chunk_shared, - hidden_states, - dummy, - use_reentrant=False, - ) - else: - hidden_states, gate = norm(hidden_states) - - return hidden_states, gate, gate_mask - - def _gated_residual(self, x, y, gate, start_indices=None, end_indices=None): - """ - Applies gated residual connection with optional gate parameter. - - Args: - x: Input tensor (residual) - y: Output tensor to be added - gate: Optional gate tensor to modulate the addition - - Returns: - x + y if gate is None, otherwise x + y * gate - """ - if x is None and y is None: - return None - if x is None or y is None: - return x if x is not None else y - if gate is None: - return x + y - - new_y = y.clone() - selected_y = y[start_indices[1] : end_indices[1]] - selected_y = selected_y.view( - -1, self.config.action_horizon_flow, selected_y.shape[-1] - )[:, :, : self.config.dim_inputs[1]] - selected_y = selected_y.to(torch.float32) * gate - new_y[start_indices[1] : end_indices[1], : self.config.dim_inputs[1]] = ( - selected_y.view(-1, self.config.dim_inputs[1]).to(new_y.dtype) - ) - - return x + new_y - - def scatter_proprioception_embeddings( - self, input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask - ): - if ( - proprioception is not None - and not self.config.use_state_string_representation - ): - proprioception = proprioception.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - proprioception = self.action_preprocessor.proprioception_proj( - proprioception, - dataset_names, - agent_pos_mask, - use_history=proprioception.shape[1] > 1, - ) - mask = input_ids == self.action_token_id_set["propri_token_id"] - mask_unsqueezed = mask.unsqueeze(-1) - mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) - proprioception_mask = mask_expanded.to(inputs_embeds.device) - - proprioception = proprioception.to( - inputs_embeds.device, inputs_embeds.dtype - ) - inputs_embeds = inputs_embeds.masked_scatter( - proprioception_mask, proprioception - ) - - return inputs_embeds - - def scatter_flow_action_embeddings( - self, input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask - ): - if not self.config.use_flow_action_expert: - return inputs_embeds, None, None - adarms_cond, flow = None, None - if action_chunk is not None: - action_chunk = action_chunk.to(inputs_embeds.device) - dof_mask = dof_mask.to(inputs_embeds.device) - noisy_action_emb, flow, adarms_cond = self.action_preprocessor( - action_chunk, dataset_names, dof_mask - ) - mask = input_ids == self.action_token_id_set["action_token_id"] - mask_unsqueezed = mask.unsqueeze(-1) - mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) - action_mask = mask_expanded.to(inputs_embeds.device) - - noisy_action_emb = noisy_action_emb.to( - inputs_embeds.device, inputs_embeds.dtype - ) - inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb) - - return inputs_embeds, flow, adarms_cond - - @staticmethod - def _update_position_ids( - position_ids, - moe_token_types, - positional_masks, - ): - if ( - positional_masks is None - or "ar_predict_token_positions" not in positional_masks - ): - return position_ids - - new_position_ids = position_ids.clone() - ar_predict_token_positions = positional_masks["ar_predict_token_positions"] - flow_mask = moe_token_types == 1 - - start_ar_pos, end_ar_pos = find_first_last_ones(ar_predict_token_positions) - start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) - - for bs_i in range(position_ids.shape[1]): - if start_ar_pos[bs_i] != -1 and end_ar_pos[bs_i] != -1: - start_ar_ids = new_position_ids[:, bs_i, start_ar_pos[bs_i]] - start_flow_ids = new_position_ids[:, bs_i, start_flow_pos[bs_i]] - diff = start_flow_ids - start_ar_ids - new_position_ids[:, bs_i, start_flow_pos[bs_i] :] = position_ids[ - :, bs_i, start_flow_pos[bs_i] : - ] - diff.unsqueeze(-1) - - return new_position_ids - - def _update_joint_attention_mask_2d( - self, - attention_mask, - moe_token_types, - positional_masks, - ): - if attention_mask.dim() == 3: # bs, seq_len, seq_len - return attention_mask - - bs, seq_len = moe_token_types.shape[0], moe_token_types.shape[1] - # Create a lower triangular matrix as a causal mask. - causal_mask = torch.tril( - torch.ones( - (seq_len, seq_len), dtype=torch.bfloat16, device=moe_token_types.device - ) - ) - # Extended to the batch dimension. - attention_mask = causal_mask.unsqueeze(0).expand(bs, -1, -1) - - if positional_masks is not None and "padding_positions" in positional_masks: - padding_positions = positional_masks["padding_positions"] - # The padding is set to zero. - attention_mask = torch.where( - padding_positions[:, None, :], - torch.zeros_like(attention_mask), - attention_mask, - ) - # The padding is set to zero. - attention_mask = torch.where( - padding_positions[:, :, None], - torch.zeros_like(attention_mask), - attention_mask, - ) - - # Set all values ​​in the moe1 section to 1, and disable the fast section. - moe1_mask = (moe_token_types[:, :, None]) & (moe_token_types[:, None, :]) - - if ( - not self.config.causal_action_attention_mask - ): # If a causal action attention mask is not used, then all elements in the moe1 section are set to 1. - attention_mask = torch.where( - moe1_mask, torch.ones_like(attention_mask), attention_mask - ) - - if ( - positional_masks is not None - and "ar_predict_token_positions" in positional_masks - ): - ar_predict_token_positions = positional_masks["ar_predict_token_positions"] - moe1_mask = (moe_token_types[:, :, None]) & ( - ar_predict_token_positions[:, None, :] - ) - attention_mask = torch.where( - moe1_mask, torch.zeros_like(attention_mask), attention_mask - ) - - if ( - positional_masks is not None - and "valid_flow_action_positions" in positional_masks - ): - # true in moe_token_types but false in valid_flow_action_positions - nonvalid_flow_action_positions = ( - moe_token_types & ~positional_masks["valid_flow_action_positions"] - ) - attention_mask = torch.where( - nonvalid_flow_action_positions[:, None, :], - torch.zeros_like(attention_mask), - attention_mask, - ) - attention_mask = torch.where( - nonvalid_flow_action_positions[:, :, None], - torch.zeros_like(attention_mask), - attention_mask, - ) - - return attention_mask - - def _update_joint_attention_flash_mask( - self, - attention_mask, - moe_token_types, - positional_masks, - debug=False, - ): - device = moe_token_types.device - B, S = moe_token_types.shape - i32 = torch.int32 - - # ---- Return vector initialization ---- - LTS = torch.ones((B, S), device=device, dtype=i32) * S - UTE = ( - torch.arange(S, device=device, dtype=i32).unsqueeze(0).expand(B, S).clone() - ) - - # Handling padding positions - if positional_masks is not None and "padding_positions" in positional_masks: - padding_positions = positional_masks["padding_positions"] - LTS[padding_positions] = 0 - UTE[padding_positions] = S - - # Handling ar predict tokens - if ( - positional_masks is not None - and "ar_predict_token_positions" in positional_masks - ): - start_ar_pos, end_ar_pos = find_first_last_ones( - positional_masks["ar_predict_token_positions"] - ) - for bs_i in range(B): - if end_ar_pos[bs_i] != -1: - LTS[bs_i, positional_masks["ar_predict_token_positions"][bs_i]] = ( - end_ar_pos[bs_i].to(i32) + 1 - ) - - # Handling flow action bidirectional mask - flow_mask = moe_token_types == 1 - if not self.config.causal_action_attention_mask: - start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) - for bs_i in range(B): - if start_flow_pos[bs_i] != -1: - UTE[bs_i, flow_mask[bs_i]] = start_flow_pos[bs_i].to(i32) - - # Handling validate flow - if ( - positional_masks is not None - and "valid_flow_action_positions" in positional_masks - ): - flow_mask = moe_token_types == 1 - nonvalid_flow_action_positions = ( - flow_mask & ~positional_masks["valid_flow_action_positions"] - ) - if nonvalid_flow_action_positions.any(): - LTS[nonvalid_flow_action_positions] = 0 - UTE[nonvalid_flow_action_positions] = S - - LTS = LTS.unsqueeze(-1) - UTE = UTE.unsqueeze(-1) - - startend_row_indices = torch.cat([LTS, UTE], dim=-1) - # startend_row_indices = LTS - - # add num_heads dimension - startend_row_indices = startend_row_indices.unsqueeze(1) - - return startend_row_indices - - -class ActionGenerationMixin(GenerationMixin): - action_preprocessor: ActionProcessor - - def to_bfloat16_for_selected_params(self, fsdp_plugin=None, accelerator=None): - """ - Keep some model parameters as float32, and convert others to bfloat16. - - If `fsdp_plugin` exists, use FSDP v1's `mixed_precision` wrapper. - - Otherwise, directly modify the parameter dtype. - """ - - def _assign_child(root_module, dotted_name: str, new_child): - parts = dotted_name.split(".") - parent = root_module - for p in parts[:-1]: - parent = getattr(parent, p) - setattr(parent, parts[-1], new_child) - - if fsdp_plugin: - fsdp_version = getattr(fsdp_plugin, "fsdp_version", None) - if fsdp_version != 1: - raise RuntimeError("Only FSDP v1 is supported (fsdp_version=1).") - - device = getattr( - accelerator, "device", torch.device("cuda", torch.cuda.current_device()) - ) - if isinstance(device, torch.device) and device.type == "cuda": - if device.index is not None: - torch.cuda.set_device(device.index) - device_id = device.index - - # move model to device - self = self.to(device) - - # Define the mixed-precision strategy. - bf16_policy = MP( - param_dtype=torch.bfloat16, - reduce_dtype=torch.float32, - buffer_dtype=torch.bfloat16, - cast_forward_inputs=False, - cast_root_forward_inputs=False, - ) - - fp32_policy = MP( - param_dtype=torch.float32, - reduce_dtype=torch.float32, - buffer_dtype=torch.float32, - cast_forward_inputs=False, - cast_root_forward_inputs=False, - ) - - # Step 1️⃣: Identify the top-level ActionProcessor module and wrap it separately with FSDP (FP32). - for name, module in list(self.named_modules()): - if isinstance(module, nn.Module) and any( - k in name.lower() for k in ["action_preprocessor"] - ): - if any(True for _ in module.children()): - continue - if getattr(module, "_fsdp_wrapped", False): - continue - - print(f"[FSDP v1] wrapping module in FP32: {name}") - wrapped = FSDP( - module, - mixed_precision=fp32_policy, - sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, - backward_prefetch="BACKWARD_PRE", - device_id=device_id, - use_orig_params=True, - ) - _assign_child(self, name, wrapped) - setattr(wrapped, "_fsdp_wrapped", True) - - # Step 2️⃣: The outermost layer uses unified FSDP (BF16 strategy). - print("[FSDP v1] wrapping root model with bf16 mixed precision...") - self = FSDP( - self, - mixed_precision=bf16_policy, - sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, - backward_prefetch="BACKWARD_PRE", - device_id=device_id, - use_orig_params=True, - ) - - return self - - # ----------------- Non-FSDP scenarios ----------------- - else: - print("[INFO] Running manual dtype conversion (no FSDP).") - self.to(dtype=torch.float32) - - params_to_keep_float32 = [] - for name, _ in self.named_parameters(): - if any( - k in name - for k in [ - "input_layernorm", - "post_attention_layernorm", - "model.norm", - "action_preprocessor", - ] - ): - params_to_keep_float32.append(name) - - for name, param in self.named_parameters(): - if name not in params_to_keep_float32: - param.data = param.data.to(torch.bfloat16) - - return self - - def define_action_token_id(self): - action_token_list = [] - if self.action_tokenizer_type: - for i in range(self.action_tokenizer.vocab_size): - action_token_id = self.processor.tokenizer.convert_tokens_to_ids( - f"<|action_token_{i}|>" - ) - action_token_list.append(action_token_id) - - action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") - propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>") - self.action_token_id_set = { - "action_token_list": action_token_list, - "propri_token_id": propri_token_id, - "action_token_id": action_token_id, - } - - def add_lora( - self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1 - ): - """Add LoRA adapter""" - config = LoraConfig( - r=r, - lora_alpha=lora_alpha, - target_modules=target_modules, - lora_dropout=lora_dropout, - bias="none", - task_type="CAUSAL_LM", - ) - self.model = get_peft_model(self.model, config) - # Print trainable parameter information. - self.model.print_trainable_parameters() - - def compute_loss( - self, - hidden_states, - logits, - input_ids=None, - dataset_names=None, - labels=None, - action_chunk=None, - dof_mask=None, - flow=None, - flow_loss_mask=None, - **kwargs, - ): - if input_ids is not None: - batch_size, seq_length = input_ids.shape - - loss = 0 - cross_entropy_loss, flow_loss = None, None - - # if dataset_names is not None: - # unique_datasets_name = list(set(dataset_names)) - # channel_loss_dict = { - # dataset_name: torch.tensor(0.0, device=logits.device) - # for dataset_name in _ACTION_DATASET_NAMES + _MULTIMODAL_DATASET_NAMES - # } - # channel_loss_count_dict = { - # dataset_name: torch.tensor(0, device=logits.device) - # for dataset_name in _ACTION_DATASET_NAMES + _MULTIMODAL_DATASET_NAMES - # } - # else: - unique_datasets_name, channel_loss_dict, channel_loss_count_dict = ( - None, - None, - None, - ) - - if labels is not None: - action_accuracy = 0 - - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - non_ignored_mask = shift_labels != -100 - _cross_entropy_loss = self.loss_fct(shift_logits, shift_labels) - cross_entropy_loss = ( - _cross_entropy_loss[non_ignored_mask].mean() - if non_ignored_mask.any() - else torch.tensor(0.0, device=shift_logits.device) - ) - - # compute channel loss - _cross_entropy_loss = _cross_entropy_loss.view(batch_size, seq_length - 1) - non_ignored_mask = non_ignored_mask.view(batch_size, seq_length - 1) - for dataset_name_i in unique_datasets_name: - dataset_mask = torch.tensor( - [name == dataset_name_i for name in dataset_names], - device=logits.device, - ) - combined_mask = dataset_mask.unsqueeze(1) & non_ignored_mask - channel_loss_dict[dataset_name_i] = ( - _cross_entropy_loss[combined_mask].sum() - if combined_mask.any() - else torch.tensor(0.0, device=shift_logits.device) - ) - channel_loss_count_dict[dataset_name_i] += combined_mask.sum() - - if not torch.isnan(cross_entropy_loss): - loss += cross_entropy_loss - else: - with torch.no_grad(): - cross_entropy_loss.detach() - - # compute action token accuracy - if len(self.action_token_id_set["action_token_list"]) > 0: - shift_logits = logits[..., :-1, :].contiguous() - action_preds = shift_logits.argmax(dim=-1) - shift_labels = labels[..., 1:].contiguous() - action_mask = ( - shift_labels > self.action_token_id_set["action_token_list"][0] - ) - correct_preds = (action_preds == shift_labels) & action_mask - action_accuracy = ( - correct_preds.sum().float() / action_mask.sum().float() - ) - channel_loss_dict["action_accuracy"] = action_accuracy - - if action_chunk is not None: - action_mask = input_ids == self.action_token_id_set["action_token_id"] - if action_mask.any(): - action_hidden_states = hidden_states[action_mask].to(torch.float32) - flow = flow.reshape(-1, flow.shape[-1]) - _flow_loss = self.action_preprocessor.flow_loss( - action_hidden_states, flow, action_chunk, dof_mask, flow_loss_mask - ) - if isinstance(_flow_loss, torch.Tensor): - flow_loss = _flow_loss.mean() - loss += flow_loss * self.config.flow_loss_weight - _flow_loss = _flow_loss.view( - dof_mask.shape[0], dof_mask.shape[1], dof_mask.shape[2] - ) - - return ( - loss, - cross_entropy_loss, - flow_loss, - channel_loss_dict, - channel_loss_count_dict, - ) - - -class AttentionsSelectorMixin: - - @classmethod - def _autoset_attn_implementation( - cls, - config, - use_flash_attention_2: bool = False, - torch_dtype: Optional[torch.dtype] = None, - device_map: Optional[Union[str, Dict[str, int]]] = None, - check_device_map: bool = True, - ): - """ - Automatically checks and dispatches to a default attention implementation. In order of priority: - 1. An implementation specified in `config._attn_implementation` (due for example to the argument attn_implementation="sdpa" in from_pretrained). - 2. DEPRECATED: if use_flash_attention_2 is set to `True` and `flash_attn` is available, flash attention. (`LlamaFlashAttention` for example) - 3. SDPA implementation, if available and supported by the model type. (`LlamaSdpaAttention` for example) - 4. The default model's implementation otherwise (`LlamaAttention` for example) . - """ - # Here we use config._attn_implementation_internal to check whether the attention implementation was explicitly set by the user. - # The property `PretrainedConfig._attn_implementation` is never `None`, for backward compatibility (always fall back on "eager"). - # The `hasattr` here is used as some Transformers tests for some reason do not call PretrainedConfig __init__ (e.g. test_no_super_init_config_and_model) - requested_attn_implementation = None - if ( - hasattr(config, "_attn_implementation_internal") - and config._attn_implementation_internal is not None - ): - if ( - config._attn_implementation != "flash_attention_2" - and use_flash_attention_2 - ): - raise ValueError( - f'Both attn_implementation="{config._attn_implementation}" and `use_flash_attention_2=True` were used when loading the model, which are not compatible.' - ' We recommend to just use `attn_implementation="flash_attention_2"` when loading the model.' - ) - - if ( - not isinstance(config._attn_implementation, dict) - and config._attn_implementation - not in ["eager"] - + ALL_ATTENTION_FUNCTIONS.valid_keys() - + X2ROBOT_ATTENTION_FUNCTIONS - ): - message = f'Specified `attn_implementation="{config._attn_implementation}"` is not supported. The only possible arguments are `attn_implementation="eager"` (manual attention implementation)' - if cls._supports_flash_attn_2: - message += ', `"attn_implementation=flash_attention_2"` (implementation using flash attention 2)' - if cls._supports_sdpa: - message += ', `"attn_implementation=sdpa"` (implementation using torch.nn.functional.scaled_dot_product_attention)' - if cls._supports_flex_attn: - message += ', `"attn_implementation=flex_attention"` (implementation using torch\'s flex_attention)' - raise ValueError(message + ".") - - # If a config is passed with a preset attn_implementation, we skip the automatic dispatch and use the user-provided config, with hard checks that the requested attention implementation is available. - requested_attn_implementation = config._attn_implementation_internal - - if use_flash_attention_2: - logger.warning_once( - 'The model was loaded with use_flash_attention_2=True, which is deprecated and may be removed in a future release. Please use `attn_implementation="flash_attention_2"` instead.' - ) - config._attn_implementation = "flash_attention_2" - - if config._attn_implementation == "flash_attention_2": - cls._check_and_enable_flash_attn_2( - config, - torch_dtype=torch_dtype, - device_map=device_map, - hard_check_only=False, - check_device_map=check_device_map, - ) - elif requested_attn_implementation == "flex_attention": - config = cls._check_and_enable_flex_attn(config, hard_check_only=True) - elif ( - requested_attn_implementation in [None, "sdpa"] - and not is_torch_xla_available() - ): - # use_flash_attention_2 takes priority over SDPA, hence SDPA treated in this elif. - config = cls._check_and_enable_sdpa( - config, - hard_check_only=( - False if requested_attn_implementation is None else True - ), - ) - - if ( - torch.version.hip is not None - and config._attn_implementation == "sdpa" - and torch.cuda.device_count() > 1 - and version.parse(torch.__version__) < version.parse("2.4.1") - ): - logger.warning_once( - "Using the `SDPA` attention implementation on multi-gpu setup with ROCM may lead to performance issues due to the FA backend. Disabling it to use alternative backends." - ) - torch.backends.cuda.enable_flash_sdp(False) - elif requested_attn_implementation in ALL_ATTENTION_FUNCTIONS.valid_keys(): - config._attn_implementation = requested_attn_implementation - elif isinstance(requested_attn_implementation, dict): - config._attn_implementation = None - elif config._attn_implementation in X2ROBOT_ATTENTION_FUNCTIONS: - pass - else: - config._attn_implementation = "eager" - - config._attn_implementation_autoset = True - return config - - def _check_and_adjust_attn_implementation( - self, attn_implementation: Optional[str], is_init_check: bool = False - ) -> str: - assert ( - attn_implementation - in ["eager", "flash_attention_2", "sdpa"] + X2ROBOT_ATTENTION_FUNCTIONS - ) - return attn_implementation diff --git a/wall_x/serving/README.md b/wall_x/serving/README.md deleted file mode 100644 index febfe3a..0000000 --- a/wall_x/serving/README.md +++ /dev/null @@ -1,263 +0,0 @@ -# Wall-X Model Serving - -This directory contains scripts for serving Wall-X models via a websocket server, allowing remote clients to connect and get action predictions from observations. - -## Overview - -The serving infrastructure consists of three main components: - -1. **WebsocketPolicyServer** (`wall_x/serving/websocket_policy_server.py`): Generic websocket server that can serve any policy implementing the `BasePolicy` interface -2. **WallXPolicy** (`wall_x/serving/policy/wall_x_policy.py`): Policy wrapper that adapts the Wall-X model to the `BasePolicy` interface -3. **launch_serving.py**: Main script for starting the server with various configurations - -## Quick Start - -### Basic Usage - -Serve a model with default LIBERO configuration: - -```bash -cd /x2robot_v2/vincent/workspace/opensource -python -m wall_x.serving.launch_serving \ - --env libero \ - --model-config.model-path /path/to/libero_model_stuff \ - --model-config.action-tokenizer-path /path/to/fast/ \ - --model-config.train-config-path /path/to/config.yml -``` - -### Specify Environment - -Serve with a specific environment preset: - -```bash -# LIBERO (single arm, 7 DOF) -python -m wall_x.serving.launch_serving --env libero - -# ALOHA (dual arm, 14 DOF) -python -m wall_x.serving.launch_serving --env aloha -``` - -### Custom Configuration - -Serve with custom model paths and settings: - -```bash -python -m wall_x.serving.launch_serving \ - --model-config.model-path /path/to/model \ - --model-config.action-tokenizer-path /path/to/tokenizer \ - --model-config.train-config-path /path/to/train_config.yml \ - --model-config.action-dim 7 \ - --model-config.state-dim 8 \ - --model-config.pred-horizon 32 \ - --model-config.camera-key front_view left_wrist_view \ - --port 8000 -``` - -## Command Line Arguments - -### Basic Arguments - -- `--env {libero,aloha}`: Environment mode (default: libero) -- `--port PORT`: Port to serve on (default: 8000) -- `--host HOST`: Host to bind to (default: 0.0.0.0) -- `--default-prompt TEXT`: Default text prompt if not provided in observation -- `--debug`: Enable debug logging - -### Model Configuration - -All model configuration arguments use the `--model-config.` prefix: - -- `--model-config.model-path PATH`: Path to pretrained model checkpoint (required) -- `--model-config.action-tokenizer-path PATH`: Path to action tokenizer (required) -- `--model-config.train-config-path PATH`: Path to train config YAML file (required) -- `--model-config.action-dim INT`: Action space dimension (default: 7) -- `--model-config.state-dim INT`: Robot state dimension (default: 8) -- `--model-config.pred-horizon INT`: Prediction horizon (default: 32) -- `--model-config.device {cuda,cpu}`: Device to run on (default: cuda) -- `--model-config.dtype {bfloat16,float16,float32}`: Model dtype (default: bfloat16) -- `--model-config.predict-mode {fast,diffusion}`: Prediction mode (default: fast) -- `--model-config.camera-key KEY1 KEY2 ...`: Camera keys for observation images - -### Camera Keys - -The `camera-key` parameter specifies which camera views are expected in the observation dictionary. This is **critical** for proper operation: - -- Keys must match between server configuration and client observations -- Order matters: keys are processed in the order specified -- Common keys: `front_view`, `left_wrist_view`, `right_wrist_view`, `face_view` - -Example: -```bash ---model-config.camera-key front_view left_wrist_view -``` - -Client must send observations with matching keys: -```python -obs = { - "front_view": image1, # Must match camera-key[0] - "left_wrist_view": image2, # Must match camera-key[1] - "prompt": "task description", - "state": robot_state, -} -``` - -## Default Configurations - -### LIBERO (Single Arm) - -```python -ModelConfig( - model_path="/path/to/model", - action_tokenizer_path="/path/to/action_tokenizer", - train_config_path="/path/to/train_config", - state_dim=8, - action_dim=7, - pred_horizon=32, - device="cuda", - dtype="bfloat16", - predict_mode="fast", - camera_key=["front_view", "left_wrist_view"], -) -``` - -### ALOHA (Dual Arm) - -```python -ModelConfig( - model_path="/path/to/model", - action_tokenizer_path="/path/to/action_tokenizer", - train_config_path="/path/to/train_config", - state_dim=14, - action_dim=14, - pred_horizon=32, - device="cuda", - dtype="bfloat16", - predict_mode="fast", - camera_key=["face_view", "left_wrist_view", "right_wrist_view"], -) -``` - -## Server Protocol - -### Connection Flow - -1. Client connects to `ws://host:port` -2. Server sends metadata JSON with policy information -3. Client sends observation (msgpack-encoded) -4. Server responds with action prediction (msgpack-encoded) -5. Repeat steps 3-4 for each inference - -### Observation Format - -Observations must be a dictionary with camera keys matching server configuration: - -```python -obs = { - # Image observations - keys must match server's camera_key configuration - "front_view": np.ndarray, # (H, W, 3) uint8 or float - "left_wrist_view": np.ndarray, # (H, W, 3) uint8 or float - - # Required fields - "prompt": str, # Task description - "dataset_names": List[str], # Dataset/robot name, e.g., ["physical-intelligence/libero"] - "state": np.ndarray, # Robot proprioception state (state_dim,) -} -``` - -**Important**: The image keys (`front_view`, `left_wrist_view`, etc.) must exactly match the `camera_key` parameter configured on the server. - -### Action Response Format - -Actions are returned as a dictionary: - -```python -{ - "action": np.ndarray, # Predicted action [pred_horizon, action_dim] - "server_timing": { - "infer_ms": float, # Inference time in milliseconds - "prev_total_ms": float, # Total time for previous request - } -} -``` - -### Server Metadata - -When connecting, the server sends metadata: - -```python -{ - "action_dim": int, # Action space dimension - "pred_horizon": int, # Number of future actions predicted - "device": str, # Device model runs on - "predict_mode": str, # Prediction mode (fast/diffusion) - "env": str, # Environment name -} -``` - -### Health Check - -HTTP health check endpoint available at: -``` -http://host:port/healthz -``` - -Returns `200 OK` if the server is running. - -## Client Example - -### Synchronous Python Client - -For synchronous usage, see `wall_x/serving/client.py`: - -```python -from wall_x.serving.client import WallXClient - -# Create and connect -client = WallXClient(uri="ws://localhost:8000") -client.connect_sync() - -# Prepare observation -obs = { - "front_view": image1, - "left_wrist_view": image2, - "prompt": "task description", - "state": robot_state, - "dataset_names": ["physical-intelligence/libero"], -} - -# Get prediction -response = client.predict_sync(obs) -action = response["action"] - -# Close connection -client.close_sync() -``` - -## Architecture - -### WebsocketPolicyServer - -Generic websocket server that: -- Handles websocket connections with msgpack serialization -- Tracks inference timing and performance metrics -- Provides health check endpoint -- Handles errors gracefully with proper logging -- Supports concurrent client connections - -### WallXPolicy - -Policy wrapper that: -- Loads and manages the Wall-X model from pretrained checkpoint -- Processes multi-camera observations -- Handles image preprocessing (smart resize, normalization) -- Manages device placement and dtype conversion -- Provides policy metadata to clients -- Supports both fast tokenizer and diffusion prediction modes - -### Image Processing Pipeline - -1. **Camera Key Matching**: Extracts images from observation dict using configured camera keys -2. **Format Conversion**: Converts numpy arrays to PIL Images -3. **Smart Resize**: Applies Qwen's smart resize algorithm based on min/max pixels -4. **Vision Token Formatting**: Inserts vision tokens in text prompt -5. **Batch Preparation**: Creates model-ready BatchFeature input diff --git a/wall_x/serving/client.py b/wall_x/serving/client.py deleted file mode 100644 index 8bb60a9..0000000 --- a/wall_x/serving/client.py +++ /dev/null @@ -1,397 +0,0 @@ -#!/usr/bin/env python3 -""" -Example client for Wall-X model server with sync support. - -This script demonstrates how to connect to a Wall-X server and request -action predictions from observations in both sync and async contexts. -""" - -import asyncio -import logging -from typing import Dict, List -import numpy as np -import threading -import yaml -import torch -import matplotlib.pyplot as plt -import os - -from wall_x.data.utils import update_action_statistics -from wall_x.utils.constant import action_statistic_dof -from wall_x.model.action_head import Normalizer - -try: - import msgpack - import msgpack_numpy as m - - m.patch() -except ImportError: - print("Please install msgpack-numpy: pip install msgpack-numpy") - exit(1) - -try: - import websockets -except ImportError: - print("Please install websockets: pip install websockets") - exit(1) - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - - -class WallXClient: - """Client for connecting to Wall-X model server.""" - - def __init__( - self, - config_path: str, - uri: str = "ws://localhost:8000", - norm_stats_path: str = "x2_norm_stats.json", - ): - """Initialize client. - - Args: - config_path: Path to train config file - uri: WebSocket URI of the server (e.g., ws://localhost:8000) - norm_stats_path: Path to normalization stats file - """ - self.uri = uri - self.websocket = None - self.metadata = None - self._loop = None - self._thread = None - self.norm_stats_path = norm_stats_path - - with open(config_path, "r") as f: - self.train_config = yaml.load(f, Loader=yaml.FullLoader) - - self.init_normalizer(self.train_config) - - async def connect(self): - """Connect to the server and receive metadata.""" - logger.info(f"Connecting to {self.uri}...") - self.websocket = await websockets.connect( - self.uri, - ping_interval=None, - ping_timeout=None, - max_size=None, - ) - - self.metadata = msgpack.unpackb(await self.websocket.recv()) - logger.info(f"Connected! Server metadata: {self.metadata}") - - async def predict(self, obs: Dict) -> Dict: - """Get action prediction from observation. - - Args: - obs: Observation dictionary containing: - - 'image': Image array (H, W, C) - - 'prompt': Optional text prompt - - 'state': Optional robot state - - Returns: - Dictionary with: - - 'action': Predicted action array - - 'server_timing': Timing information - """ - if self.websocket is None: - raise RuntimeError("Not connected. Call connect() first.") - - await self.websocket.send(msgpack.packb(obs)) - response = msgpack.unpackb(await self.websocket.recv()) - return response - - async def close(self): - """Close the connection.""" - if self.websocket: - await self.websocket.close() - logger.info("Connection closed") - - async def reset(self): - """Reset the policy (if supported).""" - pass - - # ============ Synchronous methods (using independent thread event loop) ============ - - def _start_background_loop(self): - """Start event loop in background thread.""" - self._loop = asyncio.new_event_loop() - asyncio.set_event_loop(self._loop) - self._loop.run_forever() - - def _ensure_loop(self): - """Ensure background event loop is running.""" - if self._loop is None or not self._loop.is_running(): - self._thread = threading.Thread( - target=self._start_background_loop, daemon=True - ) - self._thread.start() - # Wait for loop to start - import time - - while self._loop is None: - time.sleep(0.01) - - def _run_async(self, coro): - """Run coroutine in background event loop.""" - self._ensure_loop() - future = asyncio.run_coroutine_threadsafe(coro, self._loop) - return future.result() - - def connect_sync(self): - """Synchronously connect to server.""" - return self._run_async(self.connect()) - - def norm_state( - self, - state: np.ndarray, - dataset_names: List[str], - state_mask: torch.Tensor = None, - ) -> np.ndarray: - """Normalize state.""" - return self.normalizer_propri.normalize_data(state, dataset_names, state_mask) - - def predict_sync(self, obs: Dict) -> Dict: - """Synchronous prediction method. - - Args: - obs: Observation dictionary - - Returns: - Prediction result dictionary - """ - return self._run_async(self.predict(obs)) - - def close_sync(self): - """Synchronously close connection.""" - result = self._run_async(self.close()) - # Stop event loop - if self._loop: - self._loop.call_soon_threadsafe(self._loop.stop) - return result - - def init_normalizer(self, train_config): - # Define default configurations - dof_config = {"biarm_eed_with_base": 20} - - agent_pos_config = {"biarm_eed_with_base": 20} - - update_action_statistics( - action_statistic_dof=action_statistic_dof, - norm_stats_path=self.norm_stats_path, - repo_id="x2", - dof_config=dof_config, - agent_pos_config=agent_pos_config, - ) - - self.normalizer_action = Normalizer(action_statistic_dof, dof_config) - self.normalizer_propri = Normalizer(action_statistic_dof, agent_pos_config) - - print("Normalizer initialized") - - -def prepare_batch_sync(data, normalizer_action, normalizer_propri, dataset_names): - """Synchronous version of prepare_batch.""" - image = (data["image"].permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy() - wrist_image = ( - (data["wrist_image"].permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy() - ) - prompt = data["task"] - - state = data["state"].to("cuda") - if state.dim() == 1: - state = state.unsqueeze(0) - - state_mask = torch.ones([1, 32, 20]).to("cuda") - state_mask[:, :, 8:] = 0 - - state = normalizer_propri.normalize_data(state, dataset_names, state_mask) - state = state.cpu().numpy().astype(np.float32) - - obs = { - "front_view": image, - "left_wrist_view": wrist_image, - "prompt": prompt, - "state": state, - "dataset_names": dataset_names, - } - return obs - - -def init_serving_sample_dataset(train_config): - from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata - - repo_id = train_config["data"]["lerobot_config"]["repo_id"] - - meta_info = LeRobotDatasetMetadata(repo_id) - dataset_fps = meta_info.fps - delta_timestamps = { - "actions": [t / dataset_fps for t in range(32)], - } - dataset = LeRobotDataset( - repo_id, - episodes=[0], - delta_timestamps=delta_timestamps, - video_backend="pyav", - ) - - return dataset, repo_id - - -# ============ Synchronous version of main function ============ - - -def main_sync(args): - """Synchronous version of main function.""" - - # Create client and connect - client = WallXClient( - args.config_path, uri=args.uri, norm_stats_path=args.norm_stats_path - ) - client.connect_sync() - - dataset, repo_id = init_serving_sample_dataset(client.train_config) - - total_frames = len(dataset) - gt_traj = np.zeros((total_frames, args.action_dim)) - pred_traj = np.zeros((total_frames, args.action_dim)) - import torch - - dof_mask = torch.ones([1, 32, 20]).to("cuda") - dof_mask[:, :, args.action_dim :] = 0 - - # Synchronous processing - for idx, data in enumerate(dataset): - if idx % args.pred_horizon == 0 and idx + args.pred_horizon < total_frames: - print(f"Processing frame {idx}") - obs = prepare_batch_sync( - data, - client.normalizer_action, - client.normalizer_propri, - dataset_names=[repo_id], - ) - response = client.predict_sync(obs) - pred_action = response["action"] - pred_traj[idx : idx + args.pred_horizon] = pred_action - gt_traj[idx : idx + args.pred_horizon] = data["actions"] - - # Draw plot - timesteps = gt_traj.shape[0] - fig, axs = plt.subplots( - args.action_dim, 1, figsize=(15, 5 * args.action_dim), sharex=True - ) - fig.suptitle("Action Comparison for lerobot", fontsize=16) - - for i in range(args.action_dim): - axs[i].plot(range(timesteps), gt_traj[:, i], label="Ground Truth") - axs[i].plot(range(timesteps), pred_traj[:, i], label="Prediction") - axs[i].set_ylabel(f"Action Dim {i+1}") - axs[i].legend() - axs[i].grid(True) - - axs[-1].set_xlabel("Timestep") - plt.tight_layout(rect=[0, 0.03, 1, 0.95]) - os.makedirs(args.save_dir, exist_ok=True) - save_path = os.path.join(args.save_dir, "lerobot_comparison_serving.png") - plt.savefig(save_path) - print(f"Saved plot to {save_path}") - plt.close() - - # Close connection - client.close_sync() - - -# ============ Asynchronous version of main function (keep original functionality) ============ - - -async def main(args): - client = WallXClient( - args.config_path, uri=args.uri, norm_stats_path=args.norm_stats_path - ) - await client.connect() - dataset, repo_id = init_serving_sample_dataset(client.train_config) - - total_frames = len(dataset) - gt_traj = np.zeros((total_frames, args.action_dim)) - pred_traj = np.zeros((total_frames, args.action_dim)) - - for idx, data in enumerate(dataset): - if idx % args.pred_horizon == 0 and idx + args.pred_horizon < total_frames: - print(f"Processing frame {idx}") - obs = prepare_batch_sync( - data, - client.normalizer_action, - client.normalizer_propri, - dataset_names=[repo_id], - ) - response = await client.predict(obs) - pred_action = response["action"] - print(pred_action.shape) - pred_traj[idx : idx + args.pred_horizon] = pred_action - gt_traj[idx : idx + args.pred_horizon] = data["actions"] - - timesteps = gt_traj.shape[0] - - fig, axs = plt.subplots( - args.action_dim, 1, figsize=(15, 5 * args.action_dim), sharex=True - ) - fig.suptitle("Action Comparison for lerobot", fontsize=16) - - for i in range(args.action_dim): - axs[i].plot(range(timesteps), gt_traj[:, i], label="Ground Truth") - axs[i].plot(range(timesteps), pred_traj[:, i], label="Prediction") - axs[i].set_ylabel(f"Action Dim {i+1}") - axs[i].legend() - axs[i].grid(True) - - axs[-1].set_xlabel("Timestep") - plt.tight_layout(rect=[0, 0.03, 1, 0.95]) - os.makedirs(args.save_dir, exist_ok=True) - save_path = os.path.join(args.save_dir, "lerobot_comparison_serving.png") - plt.savefig(save_path) - print(f"Saved plot to {save_path}") - plt.close() - - -if __name__ == "__main__": - """Asynchronous version of main function.""" - import argparse - - parser = argparse.ArgumentParser(description="Wall-X client examples") - parser.add_argument( - "--example", - choices=["single", "multiple", "benchmark"], - default="single", - help="Example to run", - ) - parser.add_argument( - "--uri", - default="ws://localhost:8000", - help="Server URI", - ) - parser.add_argument( - "--pred_horizon", type=int, default=32, help="Prediction horizon" - ) - parser.add_argument("--action_dim", type=int, default=7, help="Action dimension") - parser.add_argument( - "--config_path", - default="config_from_qwen_libero.yml", - help="Train config path", - ) - parser.add_argument( - "--save_dir", - default="libero", - help="Save directory", - ) - parser.add_argument( - "--norm_stats_path", - default="x2_norm_stats.json", - help="Normalization stats path", - ) - args = parser.parse_args() - - # Synchronous mode - main_sync(args) - - # Asynchronous mode - # asyncio.run(main(args)) diff --git a/wall_x/serving/launch_serving.py b/wall_x/serving/launch_serving.py deleted file mode 100644 index c675a6a..0000000 --- a/wall_x/serving/launch_serving.py +++ /dev/null @@ -1,220 +0,0 @@ -#!/usr/bin/env python3 -""" -Server script for Wall-X model. - -This script serves a Wall-X model using a websocket server, allowing -clients to connect and get action predictions from observations. - -Based on the OpenPI serve_policy.py script structure. -""" - -import dataclasses -from dataclasses import field -import enum -import logging -import socket -import sys -import yaml -from pathlib import Path -from typing import List - -import tyro - -from wall_x.serving.policy.wall_x_policy import WallXPolicy -from wall_x.serving.websocket_policy_server import WebsocketPolicyServer - -logger = logging.getLogger(__name__) - - -class EnvMode(enum.Enum): - """Supported environments/datasets.""" - - LIBERO = "libero" - ALOHA = "aloha" - - -@dataclasses.dataclass -class ModelConfig: - """Configuration for loading a Wall-X model.""" - - # Path to the pretrained model checkpoint - model_path: str - # Path to the action tokenizer - action_tokenizer_path: str - # Path to train config yaml - train_config_path: str - # Action dimension for the environment - action_dim: int = 7 - # State dimension for the environment - state_dim: int = 8 - # Prediction horizon (number of future actions to predict) - pred_horizon: int = 32 - # Device to run model on - device: str = "cuda" - # Model dtype (bfloat16, float16, float32) - dtype: str = "bfloat16" - # Prediction mode (fast or slow) - predict_mode: str = "fast" - # Camera key for the environment - camera_key: List[str] = field( - default_factory=lambda: ["front_view", "left_wrist_view", "right_wrist_view"] - ) - - -@dataclasses.dataclass -class Args: - """Arguments for the serve_wall_x script.""" - - # Environment mode (used for default configurations) - env: EnvMode = EnvMode.LIBERO - - # Model configuration. If not provided, uses default config for the environment - model_config: ModelConfig | None = None - - # Default text prompt to use if not provided in observation - default_prompt: str | None = None - - # Port to serve the policy on - port: int = 8000 - - # Host to bind the server to - host: str = "0.0.0.0" - - # Enable debug logging - debug: bool = False - - -# Default model configurations for each environment -DEFAULT_CONFIGS: dict[EnvMode, ModelConfig] = { - EnvMode.LIBERO: ModelConfig( - model_path="/path/to/model", - action_tokenizer_path="/path/to/action_tokenizer", - train_config_path="/path/to/train_config", - state_dim=8, - action_dim=7, - pred_horizon=32, - device="cuda", - dtype="bfloat16", - predict_mode="fast", - camera_key=["front_view", "left_wrist_view"], - ), - EnvMode.ALOHA: ModelConfig( - model_path="/path/to/model", - action_tokenizer_path="/path/to/action_tokenizer", - train_config_path="/path/to/train_config", - state_dim=14, - action_dim=14, - pred_horizon=32, - device="cuda", - dtype="bfloat16", - predict_mode="fast", - camera_key=["face_view", "left_wrist_view", "right_wrist_view"], - ), -} - - -def get_model_config(args: Args) -> ModelConfig: - """Get model configuration from args or defaults.""" - if args.model_config is not None: - return args.model_config - - if config := DEFAULT_CONFIGS.get(args.env): - logger.info(f"Using default configuration for {args.env.value}") - return config - - raise ValueError( - f"No default configuration for {args.env.value}. " - f"Please provide --model-config with model_path and action_tokenizer_path." - ) - - -def create_policy(args: Args) -> WallXPolicy: - """Create a Wall-X policy from the given arguments.""" - config = get_model_config(args) - logger.info(f"Creating Wall-X policy with config: {config}") - - # Validate paths - if not Path(config.model_path).exists(): - logger.warning(f"Model path does not exist: {config.model_path}") - - if not Path(config.action_tokenizer_path).exists(): - logger.warning( - f"Action tokenizer path does not exist: {config.action_tokenizer_path}" - ) - - with open(config.train_config_path, "r") as f: - train_config = yaml.load(f, Loader=yaml.FullLoader) - - policy = WallXPolicy( - model_path=config.model_path, - train_config=train_config, - action_tokenizer_path=config.action_tokenizer_path, - action_dim=config.action_dim, - agent_pos_dim=config.state_dim, - pred_horizon=config.pred_horizon, - device=config.device, - dtype=config.dtype, - predict_mode=config.predict_mode, - default_prompt=args.default_prompt, - camera_key=config.camera_key, - ) - - return policy - - -def main(args: Args) -> None: - """Main function to start the Wall-X model server.""" - log_level = logging.DEBUG if args.debug else logging.INFO - logging.basicConfig( - level=log_level, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", - ) - - logger.info("Starting Wall-X model server") - logger.info(f"Environment: {args.env.value}") - logger.info(f"Port: {args.port}") - logger.info(f"Host: {args.host}") - - # Create policy - try: - policy = create_policy(args) - except Exception as e: - logger.error(f"Failed to create policy: {e}") - sys.exit(1) - - # Get policy metadata - policy_metadata = policy.metadata - policy_metadata["env"] = args.env.value - - # Get network info - hostname = socket.gethostname() - try: - local_ip = socket.gethostbyname(hostname) - except Exception: - local_ip = "unknown" - - logger.info(f"Server hostname: {hostname}") - logger.info(f"Server IP: {local_ip}") - logger.info(f"Server will be available at: ws://{args.host}:{args.port}") - logger.info(f"Health check endpoint: http://{args.host}:{args.port}/healthz") - - # Create and start server - server = WebsocketPolicyServer( - policy=policy, - host=args.host, - port=args.port, - metadata=policy_metadata, - ) - - logger.info("Starting server...") - try: - server.serve_forever() - except KeyboardInterrupt: - logger.info("Server stopped by user") - except Exception as e: - logger.error(f"Server error: {e}") - sys.exit(1) - - -if __name__ == "__main__": - main(tyro.cli(Args)) diff --git a/wall_x/serving/policy/utils.py b/wall_x/serving/policy/utils.py deleted file mode 100644 index d13d768..0000000 --- a/wall_x/serving/policy/utils.py +++ /dev/null @@ -1,265 +0,0 @@ -from typing import Dict, List -import logging -import numpy as np -from wall_x.data.utils import preprocesser_call -from qwen_vl_utils.vision_process import smart_resize -import torch -from PIL import Image -from transformers import BatchFeature - -logger = logging.getLogger(__name__) - - -def prepare_batch( - obs: Dict, - processor, - normalizer_propri, - camera_key: List[str], - agent_pos_dim, - action_dim, - pred_horizon, - fixed_action_dim, - max_length, - image_factor: int, - min_pixels: int, - max_pixels: int, - predict_mode: str = "fast", - device: str = "cuda", -) -> BatchFeature: - """Prepare observation into model input format. - - Args: - obs: Dictionary containing: - - 'camera_key_0' : image 0 - - 'camera_key_1' : image 1 - ... - - 'prompt': Text prompt - - 'state': Robot state/proprioception - - 'dataset_names': Dataset names - - Returns: - BatchFeature object ready for model input - """ - # Handle images - can be single image, list of images, or dict of images - images = [] - images = [obs[key] for key in camera_key] - # Convert numpy arrays to PIL Images - processed_images = [] - for img in images: - if isinstance(img, np.ndarray): - # Debug: Log the shape and dtype - logger.debug(f"Image shape: {img.shape}, dtype: {img.dtype}") - - # Handle unexpected dimensions - squeeze if needed - if img.ndim > 3: - logger.warning( - f"Image has {img.ndim} dimensions, squeezing extra dimensions" - ) - img = np.squeeze(img) - - # Verify shape is valid for PIL - if img.ndim == 2: - # Grayscale image - pass - elif img.ndim == 3: - # Check if channel dimension is first or last - if img.shape[0] == 3 or img.shape[0] == 1: - # Channels first, transpose to channels last - img = np.transpose(img, (1, 2, 0)) - elif img.shape[2] == 3 or img.shape[2] == 1: - # Already channels last - pass - else: - raise ValueError( - f"Unexpected image shape: {img.shape}. Expected (H, W, C) or (C, H, W)" - ) - else: - raise ValueError( - f"Invalid image dimensions: {img.ndim}. Expected 2 or 3 dimensions, got shape {img.shape}" - ) - - # Convert to PIL Image - if img.dtype == np.uint8: - img = Image.fromarray(img) - else: - img = Image.fromarray((img * 255).astype(np.uint8)) - processed_images.append(img) - - # print("processed_images:",processed_images) - # Apply smart resize to images - resized_images = process_images( - processed_images, image_factor, min_pixels, max_pixels - ) - - # Handle text prompt - format with vision tokens - instruction = obs["prompt"] - formatted_text = format_text_with_vision_tokens( - instruction, camera_key, predict_mode, pred_horizon - ) - - # Use processor to prepare inputs - inputs = preprocesser_call( - processor=processor, - text=[formatted_text], - images=[resized_images], - videos=None, - padding=True, - truncation=True, - return_tensors="pt", - max_length=max_length, - ) - - action_token_id = processor.tokenizer.convert_tokens_to_ids("<|action|>") - moe_token_types = inputs.input_ids == action_token_id - inputs["moe_token_types"] = torch.tensor(moe_token_types) - - # obs["dataset_names"]="libero_all" - - # Handle robot state/proprioception if available - if "state" in obs: - state = obs["state"] - if isinstance(state, np.ndarray): - state = torch.from_numpy(state).float() - elif not isinstance(state, torch.Tensor): - state = torch.tensor(state, dtype=torch.float32) - - # Add batch dimension if needed - if state.dim() == 1: - state = state.unsqueeze(0) - if state.dim() == 2: - state = state.unsqueeze(1) # [batch, 1, state_dim] - - # Pad to 20 dimensions if needed (same as training) - # if state.shape[-1] < 20: - # padding = torch.zeros(state.shape[0], state.shape[1], 20 - state.shape[-1]) - # state = torch.cat([state, padding], dim=-1) - - # Create mask for valid dimensions - agent_pos_mask = torch.ones_like(state) - if state.shape[-1] > agent_pos_dim: - agent_pos_mask[:, :, agent_pos_dim:] = 0 - - normalizer_propri.normalize_data(state, [obs["dataset_names"]] * state.shape[0]) - - inputs["proprioception"] = state - inputs["agent_pos_mask"] = agent_pos_mask - - # Add dataset name (required by model) - inputs["dataset_names"] = [obs["dataset_names"]] * state.shape[0] - - # Move all tensors to device - for key in inputs: - if isinstance(inputs[key], torch.Tensor): - inputs[key] = inputs[key].to(device) - - dof_mask = torch.ones([state.shape[0], pred_horizon, fixed_action_dim]) - dof_mask[:, :, action_dim:] = 0 - - inputs["dof_mask"] = dof_mask - - # Convert to BatchFeature to maintain consistency with training pipeline - return BatchFeature(data=dict(inputs)).to(device) - - -def process_images( - images: List[Image.Image], image_factor: int, min_pixels: int, max_pixels: int -) -> List[Image.Image]: - """Process images with smart resize following the data loading pattern. - - Args: - images: List of PIL Images - - Returns: - List of resized PIL Images - """ - resized_images = [] - for img_pil in images: - - orig_width, orig_height = img_pil.size - target_size = 256 - if target_size != -1: - # Maintain aspect ratio logic - if orig_width > orig_height: # Landscape image - new_width = target_size - new_height = int(target_size * orig_height / orig_width) - else: # Portrait image - new_height = target_size - new_width = int(target_size * orig_width / orig_height) - img_pil = img_pil.resize((new_width, new_height)) - - # Apply smart scaling (Qwen logic) - current_width, current_height = img_pil.size - resized_height, resized_width = smart_resize( - current_height, - current_width, - factor=image_factor, - min_pixels=min_pixels, - max_pixels=max_pixels, - ) - - resized_img = img_pil.resize((resized_width, resized_height)) - resized_images.append(resized_img) - - return resized_images - - -def format_text_with_vision_tokens( - instruction: str, - camera_key: List[str], - predict_mode: str = "diffusion", - pred_horizon: int = 32, -) -> str: - """Format text prompt with vision tokens for the model. - - Args: - instruction: Task instruction text - camera_key: List of camera names - - Returns: - Formatted text with special tokens - """ - # Special tokens for formatting - role_start_symbol = "<|im_start|>" - role_end_symbol = "<|im_end|>" - vision_start_symbol = "<|vision_start|>" - vision_end_symbol = "<|vision_end|>" - image_pad_symbol = "<|image_pad|>" - propri_symbol = "<|propri|>" - action_symbol = "<|action|>" - action_fast_symbol = "<|action_fast|>" - - # Camera name mapping - camera_name_mapping = { - "front_view": "front view", - "face_view": "front view", - "left_wrist_view": "left wrist view", - "right_wrist_view": "right wrist view", - "top_view": "top view", - "wall_view": "wall view", - } - - # System prologue - prologue = ( - f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n" - ) - - # User request with observation - user_request = f"{role_start_symbol}user\nObservation:" - if camera_key: - for cam_name in camera_key: - view_name = camera_name_mapping.get(cam_name, cam_name) - user_request += f" {view_name}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" - user_request += "\nInstruction:" - - text_prompt = ( - f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n" - ) - user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" - assistant_output = ( - f"{role_start_symbol}assistant\n{action_fast_symbol}{role_end_symbol}\n" - ) - if predict_mode == "diffusion": - assistant_output = f"{role_start_symbol}assistant\n{action_symbol * pred_horizon}{role_end_symbol}\n" - complete_text = prologue + user_message + assistant_output - - return complete_text diff --git a/wall_x/serving/policy/wall_x_policy.py b/wall_x/serving/policy/wall_x_policy.py deleted file mode 100644 index e2d6fa9..0000000 --- a/wall_x/serving/policy/wall_x_policy.py +++ /dev/null @@ -1,182 +0,0 @@ -import logging -from typing import Dict, Any, List -import torch -import copy -import numpy as np -from wall_x.serving.websocket_policy_server import BasePolicy -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction -from wall_x.serving.policy.utils import prepare_batch -from wall_x.model.model_utils import load_wallx_processors, register_normalizers - -logger = logging.getLogger(__name__) - - -class WallXPolicy(BasePolicy): - """Policy wrapper for Wall-X model that implements the BasePolicy interface.""" - - def __init__( - self, - model_path: str, - train_config: dict, - action_tokenizer_path: str, - action_dim: int, - agent_pos_dim: int, - pred_horizon: int, - camera_key: List[str], - device: str = "cuda", - dtype: str = "bfloat16", - predict_mode: str = "diffusion", - default_prompt: str | None = None, - min_pixels: int = 4 * 28 * 28, - max_pixels: int = 16384 * 28 * 28, - image_factor: int = 28, - max_length: int = 2048, - ): - """Initialize the Wall-X policy. - - Args: - model_path: Path to the pretrained model checkpoint - action_tokenizer_path: Path to the action tokenizer - action_dim: Dimension of action space - pred_horizon: Prediction horizon for actions - device: Device to run model on ('cuda' or 'cpu') - dtype: Data type for model ('bfloat16', 'float16', or 'float32') - predict_mode: Prediction mode ('fast' or 'slow') - default_prompt: Default text prompt for the model - min_pixels: Minimum pixels for image resizing - max_pixels: Maximum pixels for image resizing - image_factor: Factor for smart resize - max_length: Maximum sequence length for text - """ - logger.info(f"Loading Wall-X model from {model_path}") - - self.normalizer_action, self.normalizer_propri = register_normalizers( - train_config, model_path - ) - - self.model = Qwen2_5_VLMoEForAction.from_pretrained( - model_path, - train_config=train_config, - action_tokenizer_path=action_tokenizer_path, - ) - self.model.set_normalizer( - copy.deepcopy(self.normalizer_action), copy.deepcopy(self.normalizer_propri) - ) - self.model.eval() - self.model = self.model.to(device) - self.model.to_bfloat16_for_selected_params() - - # hard code the action dim to 20 for align to wall-x configuration - self.fixed_action_dim = action_dim - - self.action_dim = action_dim - self.agent_pos_dim = action_dim - self.pred_horizon = pred_horizon - self.device = device - self.predict_mode = predict_mode - self.default_prompt = default_prompt - self.camera_key = camera_key - - # Image preprocessing config - self.min_pixels = min_pixels - self.max_pixels = max_pixels - self.image_factor = image_factor - self.max_length = max_length - - print("predict_mode", predict_mode) - print("camera_key", camera_key) - - # Load processor - logger.info("Loading processor and tokenizer...") - - processors_dict = load_wallx_processors(train_config) - self.processor = processors_dict["processor"] - - # Action buffer for multi-step predictions - self.action_buffer = [] - self.buffer_index = 0 - - logger.info( - f"Model loaded successfully. Device: {device}, Action dim: {action_dim}, Horizon: {pred_horizon}" - ) - - @property - def metadata(self) -> Dict[str, Any]: - """Return metadata about the policy.""" - return { - "action_dim": self.action_dim, - "pred_horizon": self.pred_horizon, - "device": self.device, - "predict_mode": self.predict_mode, - } - - def reset(self) -> None: - """Reset the policy state.""" - self.action_buffer = [] - self.buffer_index = 0 - logger.debug("Policy reset") - - def infer(self, obs: Dict) -> Dict: - """Infer action from observation. - - Args: - obs: Dictionary containing: - - 'image': Image observation (numpy array or PIL Image) - - 'prompt': Optional text prompt - - 'state': Optional robot state - - Other modality-specific observations - - Returns: - Dictionary containing: - - 'action': Predicted action (numpy array) - - Additional metadata - """ - try: - # Need to predict new actions - input_batch = prepare_batch( - obs, - self.processor, - self.normalizer_propri, - self.camera_key, - self.agent_pos_dim, - self.action_dim, - self.pred_horizon, - self.fixed_action_dim, - self.max_length, - self.image_factor, - self.min_pixels, - self.max_pixels, - self.predict_mode, - self.device, - ) - - with torch.no_grad(): - outputs = self.model( - **input_batch, - action_dim=( - self.action_dim - if self.predict_mode == "fast" - else self.fixed_action_dim - ), - action_horizon=self.pred_horizon, - mode="predict", - predict_mode=self.predict_mode, - ) - - if outputs["predict_action"] is None: - predicted_actions = np.zeros( - [1, self.pred_horizon, self.action_dim] - ).astype(np.float32) - - predicted_actions = ( - outputs["predict_action"][:, :, : self.action_dim] - .detach() - .cpu() - .to(torch.float32) - .numpy() - ) - return {"predict_action": predicted_actions} - - except Exception as e: - logger.error(f"Error during inference: {e}") - raise diff --git a/wall_x/trainer/adapters/__init__.py b/wall_x/trainer/adapters/__init__.py new file mode 100644 index 0000000..af6752d --- /dev/null +++ b/wall_x/trainer/adapters/__init__.py @@ -0,0 +1,128 @@ +"""Public model adapter registry.""" + +from importlib import import_module + +from wall_x.trainer.adapters.base_adapter import ModelAdapter +from wall_x.trainer.adapters.vla_model_adapter import VLAdapter + +_ADAPTER_SPECS = { + "qwen2_5": ("wall_x.model.qact.qwen2_5.adapter", "Qwen2_5Adapter"), +} +_ADAPTER_CLASS_EXPORTS = { + "Qwen2_5Adapter": "qwen2_5", +} +ADAPTER_IMPORT_ERRORS: dict[str, str] = {} +_ADAPTER_LOADING: set[str] = set() + + +class _AdapterRegistry(dict[str, type[ModelAdapter]]): + def get(self, model_type: str, default=None): + _load_public_adapter(model_type) + return dict.get(self, model_type, default) + + def __contains__(self, model_type: object) -> bool: + if isinstance(model_type, str): + _load_public_adapter(model_type) + return dict.__contains__(self, model_type) + + def __getitem__(self, model_type: str) -> type[ModelAdapter]: + _load_public_adapter(model_type) + return dict.__getitem__(self, model_type) + + def __iter__(self): + _load_all_public_adapters() + return dict.__iter__(self) + + def keys(self): + _load_all_public_adapters() + return dict.keys(self) + + def items(self): + _load_all_public_adapters() + return dict.items(self) + + def values(self): + _load_all_public_adapters() + return dict.values(self) + + +ADAPTER_REGISTRY: _AdapterRegistry = _AdapterRegistry() + + +def _register_adapter(model_type: str, adapter_cls: type[ModelAdapter]) -> None: + existing = dict.get(ADAPTER_REGISTRY, model_type) + if existing is not None and existing is not adapter_cls: + raise ValueError( + f"model_type {model_type!r} already registered to " + f"{existing.__name__}, cannot re-register to {adapter_cls.__name__}" + ) + declared_model_type = getattr(adapter_cls, "MODEL_TYPE", model_type) + if declared_model_type and declared_model_type != model_type: + raise ValueError( + f"adapter {adapter_cls.__name__} declares MODEL_TYPE=" + f"{declared_model_type!r}, but is registered as {model_type!r}" + ) + dict.__setitem__(ADAPTER_REGISTRY, model_type, adapter_cls) + + +def _load_public_adapter(model_type: str) -> type[ModelAdapter] | None: + if dict.__contains__(ADAPTER_REGISTRY, model_type): + return dict.__getitem__(ADAPTER_REGISTRY, model_type) + spec = _ADAPTER_SPECS.get(model_type) + if spec is None or model_type in _ADAPTER_LOADING: + return None + module_name, class_name = spec + _ADAPTER_LOADING.add(model_type) + try: + module = import_module(module_name) + adapter_cls = getattr(module, class_name) + except (ImportError, AttributeError) as exc: + ADAPTER_IMPORT_ERRORS[model_type] = f"{type(exc).__name__}: {exc}" + return None + finally: + _ADAPTER_LOADING.discard(model_type) + _register_adapter(model_type, adapter_cls) + return adapter_cls + + +def _load_all_public_adapters() -> None: + for model_type in _ADAPTER_SPECS: + _load_public_adapter(model_type) + + +def format_adapter_error(model_type: str) -> str: + msg = ( + f"Unsupported model type: {model_type}. " + f"Registered: {sorted(ADAPTER_REGISTRY)}" + ) + if model_type in ADAPTER_IMPORT_ERRORS: + msg += f". Import failed: {ADAPTER_IMPORT_ERRORS[model_type]}" + return msg + + +def resolve_adapter(model_type: str) -> type[ModelAdapter]: + """Return a registered adapter class, loading lazily when needed.""" + adapter_cls = ADAPTER_REGISTRY.get(model_type) + if adapter_cls is None: + raise ValueError(format_adapter_error(model_type)) + return adapter_cls + + +def __getattr__(name: str): + model_type = _ADAPTER_CLASS_EXPORTS.get(name) + if model_type is not None: + adapter_cls = _load_public_adapter(model_type) + if adapter_cls is not None: + return adapter_cls + raise AttributeError(name) + + +__all__ = [ + "ModelAdapter", + "VLAdapter", + "Qwen2_5Adapter", + "ADAPTER_REGISTRY", + "ADAPTER_IMPORT_ERRORS", + "format_adapter_error", + "resolve_adapter", +] diff --git a/wall_x/trainer/adapters/base_adapter.py b/wall_x/trainer/adapters/base_adapter.py new file mode 100644 index 0000000..99f0803 --- /dev/null +++ b/wall_x/trainer/adapters/base_adapter.py @@ -0,0 +1,494 @@ +import logging +import os +from abc import ABC, abstractmethod +from collections import defaultdict + +import torch +import torch.distributed as dist +from tqdm import tqdm + +from wall_x.trainer.utils import move_batch_to_device + +logger = logging.getLogger(__name__) + + +def load_trainer_data_config(cfg): + """Return optional backend-specific trainer data config. + + Public backends build datasets directly from TrainConfig. A backend that + needs an additional trainer data config may expose + ``load_trainer_data_config(cfg)``. + """ + from wall_x.data import data_backend + + backend = data_backend() + if backend.supports("load_trainer_data_config"): + return backend.load_trainer_data_config(cfg) + return None + + +class ModelAdapter(ABC): + """Pure strategy that encapsulates model-type-specific behavior. + + The adapter is stateless w.r.t. the trainer: every method declares + its inputs as parameters and communicates results via return values. + The trainer orchestrates calls and manages its own state. + + Abstract methods (subclasses MUST implement): + load_processor -- load processor / tokenizer artifacts + build_model_config -- build model-specific config object + create_model -- instantiate the model + load_weights -- load pretrained weights & set normalizers + get_transformer_layer_cls -- layer classes for FSDP auto-wrap + load_dataset -- build dataset & training dataloader + forward -- run a forward pass + extract_loss -- extract scalar loss from model outputs + + Concrete methods with default behaviour (override when needed): + load_step_and_epoch -- restore global_step / epoch from checkpoint + log_model_info -- log model-specific info (no-op) + collect_output_stats -- collect auxiliary loss / accuracy stats (no-op) + collect_param_norms -- compute per-group param L2 norms (no-op) + get_fwd_flops -- optional forward FLOPs hook + get_output_field -- retrieve a named field from outputs (static) + """ + + def __init__(self, *, cfg=None, logger=None, model_type=None): + self.cfg = cfg + self.logger = logger + + # ---- processor ---- + @abstractmethod + def load_processor(self, action_statistic_dof): + """Load processor and associated tokenizer artifacts. + + Reads model/data config from ``self.cfg``. + + Args: + action_statistic_dof: normalizer statistics dict. + + Returns: + dict with keys: + "processor", "data_config", "tokenizer_mixin", + "train_action_tokenizer", "val_action_tokenizer", + "action_mapper", "num_added_tokens" + """ + ... + + # ---- model ---- + @abstractmethod + def create_model(self, processor, tokenizer_mixin, model_config): + """Instantiate the model (before weight loading / FSDP wrap). + + Reads model config from ``self.cfg``. + """ + ... + + @abstractmethod + def build_model_config(self): + """Build the model-specific config object from ``self.cfg.model``.""" + ... + + @abstractmethod + def load_weights(self, model, normalizer_action, normalizer_propri, **kwargs): + """Load pretrained weights, resize embeddings, set normalizers, etc.""" + ... + + # ---- FSDP wrapping ---- + @abstractmethod + def get_transformer_layer_cls(self): + """Return layer classes for FSDP transformer_auto_wrap_policy.""" + ... + + # ---- dataset ---- + @abstractmethod + def load_dataset(self, data_config, processor, rank, world_size, **kwargs): + """Load dataset and build the training dataloader. + + Args: + data_config: parsed data config from load_processor(). + processor: processor instance. + rank: current process rank. + world_size: total number of processes. + **kwargs: extra model-specific args. + + Returns: + (dataset, train_dataloader, train_num) tuple. + """ + ... + + def load_step_and_epoch(self, checkpoint_path: str, is_incomplete_epoch: bool): + global_step = 0 + start_epoch = 0 + + global_step_path = os.path.join(checkpoint_path, "global_step.pth") + if os.path.exists(global_step_path): + global_step = torch.load(global_step_path)["global_step"] + + current_epoch_path = os.path.join(checkpoint_path, "current_epoch.pth") + if os.path.exists(current_epoch_path): + start_epoch = torch.load(current_epoch_path)["current_epoch"] + if not is_incomplete_epoch: + start_epoch = start_epoch + 1 + + return {"global_step": global_step, "start_epoch": start_epoch} + + # ---- forward / loss ---- + @abstractmethod + def forward(self, model, batch, **kwargs): + """Run a forward pass. + + Args: + model: the (possibly wrapped) model. + batch: dict of tensors already on device. + **kwargs: extra context (e.g. global_step, mode). + + Returns: + Raw model outputs (dict or object). + """ + ... + + @abstractmethod + def extract_loss(self, outputs): + """Extract the scalar training loss from model outputs. + + Returns: + torch.Tensor (scalar). + """ + ... + + # ---- optional hooks (override when needed) ---- + def log_model_info(self, model, log_fn): + """Log model-specific info (e.g. attention implementation). + + Args: + model: the unwrapped model. + log_fn: callable(str) for logging. + """ + pass + + def collect_output_stats( + self, + outputs, + step_stats, + reduce_tensor_fn, + true_gather_fn, + tokenizer_mixin=None, + ): + """Collect model-output-specific statistics into step_stats. + + Extracts auxiliary losses (cross_entropy_loss, flow_loss), + per-dataset channel losses, and action accuracy metrics from + the model outputs and merges them into *step_stats* in-place. + + This is an optional hook - subclasses override it when their + output format carries extra fields. The default implementation + is a no-op. + + Args: + outputs: raw model outputs (dict or object). + step_stats: dict being built by the trainer; mutated in-place. + reduce_tensor_fn: callable(tensor, average=True) -> scalar tensor, + wraps dist.all_reduce for the current parallelism. + true_gather_fn: callable(tensor) -> scalar | None, + gathers a value to rank-0 (returns None on + non-main ranks or when input is None). + tokenizer_mixin: optional ActionTokenizerMixin used to retrieve + extra accuracy keys (e.g. RVQ layer accuracy). + """ + pass + + def collect_param_norms( + self, model, step_stats, device, reduce_tensor_fn, params_sharded=False + ): + """Compute per-group parameter L2 norms and merge into step_stats. + + Groups parameters by name keyword: + - "visual" -> visual_param_norm + - "action" -> action_expert_param_norm + - (others) -> org_vlm_param_norm + - all -> total_param_norm + + When parameters are sharded across ranks (FSDP FULL_SHARD / + HYBRID_SHARD), the local squared-sum is partial and must be + all_reduce'd (sum, not avg) before taking the square root. + When parameters are replicated (DDP, or FSDP with SHARD_GRAD_OP / + NO_SHARD), all_reduce would inflate the result by world_size. + + This is an optional hook - the default implementation is a no-op. + + Args: + model: the (possibly FSDP/DDP-wrapped) model. + step_stats: dict being built by the trainer; mutated in-place. + device: torch.device for tensor allocation. + reduce_tensor_fn: callable(tensor, average=False) -> tensor, + wraps dist.all_reduce for the current parallelism. + params_sharded: True when parameters are actually sharded across + ranks (FSDP FULL_SHARD / HYBRID_SHARD). + """ + pass + + def collect_grad_norms( + self, model, step_stats, device, reduce_tensor_fn, params_sharded=False + ): + """Compute per-component gradient L2 norms and merge into step_stats. + + Must be called while gradients are still available (before + optimizer.zero_grad). Under FSDP each rank holds a gradient shard, + so local squared-sums must be all_reduce'd (sum, not avg) before + taking the square root. + + Default is a no-op; subclasses override for model-specific groupings. + """ + pass + + def get_fwd_flops(self, autocast_context, model, batch, global_step): + """Forward FLOPs profiling is disabled in the public package.""" + return None + + @staticmethod + def get_output_field(outputs, key, default=None): + """Retrieve a named field from model outputs (dict or object).""" + if isinstance(outputs, dict): + return outputs.get(key, default) + return getattr(outputs, key, default) + + # ---- console logging fields ---- + def console_fields(self, tokenizer_mixin=None) -> list: + """Per-step console fields as (stat_key, pretty_label, fmt_spec) triples. + + MetricsLogger renders them in order, skipping any key not present + in step_stats. Public adapters may override this method to add + task-specific fields. + """ + fields = [ + ("video_loss", "vid_loss", ".6f"), + ("action_loss", "act_loss", ".6f"), + ("action_accuracy", "accuracy", ".4f"), + ("flow_loss", "flow_loss", ".6f"), + ] + if tokenizer_mixin is not None: + for key in tokenizer_mixin.get_accuracy_keys(): + if key != "action_accuracy": + short_key = key.replace("action_accuracy_", "acc_") + fields.append((key, short_key, ".4f")) + fields.extend( + [ + ("action_grad_norm", "act_gnorm", ".4f"), + ("video_grad_norm", "vid_gnorm", ".4f"), + ] + ) + return fields + + # ---- distribution hooks ---- + def convert_to_mix_precision_hint( + self, + model, + *, + device, + use_fsdp: bool, + log_fn=None, + ): + """Prepare *model* dtype + placement before FSDP/DDP wrap. + + FSDP path: cast to fp32 on device; MixedPrecision policy handles the + bf16 compute cast at forward time. DDP path: delegate to the model's + ``convert_to_mix_precision`` method (manual bf16 conversion, since + DDP has no equivalent of MixedPrecision), then move to device. + + Adapters that handle placement themselves may override this hook. + """ + _log = log_fn or (lambda _msg: None) + if use_fsdp: + model.to(dtype=torch.float32) + _log("FSDP mode: all params fp32, MixedPrecision handles bf16 compute") + model.to(device) + else: + model.convert_to_mix_precision() + _log("DDP mode: params converted to bf16 (no MixedPrecision policy)") + model.to(device) + + # ---- validation ---- + def run_validation( + self, + *, + model, + val_dataloader, + rank, + world_size, + device, + autocast_context, + reduce_fn, + gather_fn, + logger, + global_step, + output_path, + tokenizer_mixin=None, + ): + """Default validation: forward-pass loss + adapter.collect_output_stats. + + Public VLA adapters may override this for task-specific metrics. + + ``output_path`` is accepted for interface uniformity but not used + in this default. + """ + del output_path + model.eval() + log_dict = defaultdict(float) + pbar = tqdm( + val_dataloader, + desc="Validating", + total=len(val_dataloader), + disable=rank != 0, + ) + + for batch in pbar: + batch = move_batch_to_device(batch, device) + with autocast_context(): + outputs = self.forward(model, batch, global_step=0) + + loss = outputs["total_loss"] if "total_loss" in outputs else outputs["loss"] + log_dict["val_loss"] += reduce_fn(loss.detach()).item() + + val_ce_loss = self.get_output_field(outputs, "cross_entropy_loss") + if val_ce_loss is not None: + log_ce_loss = gather_fn(val_ce_loss) + if log_ce_loss is not None and rank == 0: + log_dict["val_cross_entropy_loss"] += log_ce_loss + val_flow_loss = self.get_output_field(outputs, "flow_loss") + if val_flow_loss is not None: + log_flow_loss = gather_fn(val_flow_loss) + if log_flow_loss is not None and rank == 0: + log_dict["val_flow_loss"] += log_flow_loss + + for k, v in collect_channel_loss_stats( + outputs, + prefix="val_", + tokenizer_mixin=tokenizer_mixin, + ).items(): + log_dict[k] += v + + log_dict = {k: v / len(val_dataloader) for k, v in log_dict.items()} + + if rank == 0: + if self.logger is not None: + self.logger.info(f"[FSDP Val] Step {global_step}: {log_dict}") + if logger is not None and hasattr(logger, "log"): + logger.log(log_dict, step=global_step) + + val_info = f"[Validation] step {global_step}" + if "val_loss" in log_dict: + val_info += f" | val_loss {log_dict['val_loss']:.4f}" + if "val_cross_entropy_loss" in log_dict: + val_info += f" | val_ce_loss {log_dict['val_cross_entropy_loss']:.4f}" + if "val_flow_loss" in log_dict: + val_info += f" | val_flow_loss {log_dict['val_flow_loss']:.6f}" + if "val_action_accuracy" in log_dict: + val_info += f" | val_accuracy {log_dict['val_action_accuracy']:.4f}" + if tokenizer_mixin is not None: + for key in tokenizer_mixin.get_accuracy_keys(): + if key != "action_accuracy": + val_key = f"val_{key}" + if val_key in log_dict: + short_key = key.replace("action_accuracy_", "acc_") + val_info += f" | val_{short_key} {log_dict[val_key]:.4f}" + if self.logger is not None: + self.logger.info(val_info) + + model.train() + return log_dict + + def init_validation(self, normalizer_action, normalizer_propri): + """Optional pre-training-loop setup for validation. Default is a no-op.""" + del normalizer_action, normalizer_propri + + # ---- prediction / inference ---- + def predict( + self, + prediction_type: str, + *, + model, + val_dataloader, + rank, + world_size, + device, + processor, + tokenizer_mixin, + logger, + current_step, + max_iteration=None, + save_dir=None, + max_samples=None, + ): + """Run an inference loop for the given ``prediction_type``. + + Subclasses implement the dispatch table (flow_action / ar_action / + dllm_action / text / ...). Default raises - only VLAdapter has + meaningful predict paths today. + """ + raise NotImplementedError( + f"{type(self).__name__} does not support predict(prediction_type=...)" + ) + + # ---- optimizer-related defaults ---- + @property + def default_action_lr_keywords(self) -> list[str]: + """Keyword list used to identify action-expert parameters. + + The trainer splits parameters into two groups when + ``action_expert_learning_rate`` is set: names containing any of these + keywords go into the action-LR group. Subclasses override this when + the architecture uses different naming conventions. + """ + return [ + "action_preprocessor", + "moe.experts.1", + "qkv_proj_experts.1", + "o_proj_experts.1", + "input_layernorms.1", + "post_attention_layernorms.1", + "model.norms.1", + ] + + +def collect_channel_loss_stats(outputs, prefix="", tokenizer_mixin=None): + """Gather per-dataset channel losses + accuracy metrics via dist.all_reduce. + + Mirrors the legacy fsdp_trainer._collect_channel_loss_stats so training + and validation stats come out with identical keys. Returns {} when the + model output dict lacks ``channel_loss_dict``. + """ + stats = {} + channel_loss_dict = ( + outputs.get("channel_loss_dict") + if isinstance(outputs, dict) + else getattr(outputs, "channel_loss_dict", None) + ) + if channel_loss_dict is None: + return stats + channel_loss_count_dict = ( + outputs.get("channel_loss_count_dict") + if isinstance(outputs, dict) + else getattr(outputs, "channel_loss_count_dict", None) + ) + for dataset_name_i in channel_loss_dict: + count_tensor = channel_loss_count_dict[dataset_name_i].clone() + loss_tensor = channel_loss_dict[dataset_name_i].detach().clone() + dist.all_reduce(count_tensor, op=dist.ReduceOp.SUM) + dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM) + count_sum = count_tensor.item() + if count_sum >= 0.5: + stats[f"{prefix}channel_loss_{dataset_name_i}"] = ( + loss_tensor.item() / count_sum + ) + if "action_accuracy" in channel_loss_dict and tokenizer_mixin is not None: + acc_tensor = channel_loss_dict["action_accuracy"].detach().clone() + dist.all_reduce(acc_tensor, op=dist.ReduceOp.SUM) + world_size = dist.get_world_size() if dist.is_initialized() else 1 + stats[f"{prefix}action_accuracy"] = acc_tensor.item() / world_size + for key in tokenizer_mixin.get_accuracy_keys(): + if key != "action_accuracy" and key in channel_loss_dict: + rvq_acc_tensor = channel_loss_dict[key].detach().clone() + dist.all_reduce(rvq_acc_tensor, op=dist.ReduceOp.SUM) + stats[f"{prefix}{key}"] = rvq_acc_tensor.item() / world_size + return stats diff --git a/wall_x/trainer/adapters/vla_model_adapter.py b/wall_x/trainer/adapters/vla_model_adapter.py new file mode 100644 index 0000000..955c9dc --- /dev/null +++ b/wall_x/trainer/adapters/vla_model_adapter.py @@ -0,0 +1,990 @@ +import os + +import torch +import torch.distributed as dist +from tqdm import tqdm + +from wall_x.trainer.adapters.base_adapter import ModelAdapter +from wall_x.trainer.optimizer.dmuon import is_dmuon_model +from wall_x.trainer.trainer_utils import ( + compute_action_metrics, + load_qwen_pretrain_weight, + load_wallx_processors_from_cfg, + save_text_results_to_file, +) + + +class VLAdapter(ModelAdapter): + """Base adapter for public Qwen2.5 VLA models. + + Concrete subclasses define the model/config classes, FSDP wrap layers, + and optional task-specific logging hooks. + """ + + #: Subclass must set this. Identifies the variant in ``ADAPTER_REGISTRY``. + MODEL_TYPE: str = "" + + def __init__(self, *, cfg=None, logger=None, model_type=None): + # Concrete subclasses set MODEL_TYPE; tests may pass model_type explicitly + # (e.g. to swap a stub class into a different slot). + resolved = model_type or self.MODEL_TYPE + super().__init__(cfg=cfg, logger=logger, model_type=resolved) + if not resolved: + raise ValueError( + f"{type(self).__name__} requires a MODEL_TYPE class attribute " + "or an explicit model_type kwarg; got neither." + ) + self.model_type = resolved + self._dmuon_module = None + self._dmuon_import_checked = False + self._dmuon_named_param_cache = None + + def _get_dmuon_module(self): + if not self._dmuon_import_checked: + try: + import dmuon + except ImportError: + dmuon = None + self._dmuon_module = dmuon + self._dmuon_import_checked = True + return self._dmuon_module + + def _get_named_dmuon_dedicated_params(self, model): + cache = self._dmuon_named_param_cache + model_id = id(model) + if cache is not None and cache[0] == model_id: + return cache[1] + + dmuon = self._get_dmuon_module() + if dmuon is None: + named_params = [] + else: + module_to_name = { + id(module): module_name for module_name, module in model.named_modules() + } + named_params = [] + for dparam in dmuon.get_dedicated_params(model): + prefix = module_to_name.get(id(dparam.module), "") + name = f"{prefix}.{dparam.param_name}" if prefix else dparam.param_name + named_params.append((name, dparam)) + + self._dmuon_named_param_cache = (model_id, named_params) + return named_params + + # ---- variant hooks (subclass overrides) ---- + @classmethod + def model_class(cls): + """Return the training-side model class for this variant.""" + raise NotImplementedError(f"{cls.__name__} must override model_class()") + + @classmethod + def config_class(cls): + """Return the HF PretrainedConfig class for this variant.""" + raise NotImplementedError(f"{cls.__name__} must override config_class()") + + @classmethod + def inference_model_class(cls): + """Return the model class to use for inference.""" + return cls.model_class() + + # ---- helpers ---- + def _get_model_and_config_class(self): + """Resolve (ModelClass, ConfigClass) for this variant via classmethods.""" + cls = type(self) + return cls.model_class(), cls.config_class() + + # ---- processor ---- + def _build_processor_dict(self) -> dict: + """Flat-dict shape expected by legacy ``load_wallx_processors`` / + ``update_model_config`` / prediction loops. + + Derived from typed TrainConfig. Kept as the one-place adapter layer + between typed configs and legacy dict-based APIs; drop me when those + downstream APIs are typed-ified. + """ + import dataclasses + + flat = dataclasses.asdict(self.cfg.model) + flat["model_type"] = self.cfg.model_type + flat["data"] = dict(self.cfg._raw_data or {}) + flat["dof_config"] = self.cfg.task.dof_config + flat["agent_pos_config"] = self.cfg.task.agent_pos_config + if self.cfg.task.ar_dof_config is not None: + flat["ar_dof_config"] = self.cfg.task.ar_dof_config + flat["batch_size_per_gpu"] = self.cfg.hyperparams.batch_size_per_gpu + return flat + + def load_processor(self, action_statistic_dof): + processors_dict = load_wallx_processors_from_cfg( + self.cfg, + normalizer=getattr(self, "normalizer_action", None), + action_statistic_dof=action_statistic_dof, + ) + self.logger.info( + f"processor vocab size: {len(processors_dict['processor'].tokenizer.vocab)}" + ) + self.logger.info( + f"num added tokens to processor: {processors_dict['num_added_tokens']}" + ) + return { + "processor": processors_dict["processor"], + "data_config": self._build_processor_dict(), + "tokenizer_mixin": processors_dict.get("tokenizer_mixin"), + "train_action_tokenizer": processors_dict["train_action_tokenizer"], + "val_action_tokenizer": processors_dict["val_action_tokenizer"], + "action_mapper": processors_dict["action_mapper"], + "num_added_tokens": processors_dict["num_added_tokens"], + } + + # ---- model ---- + def build_model_config(self): + _, ConfigClass = self._get_model_and_config_class() + qwen_vl_act_config_path = self.cfg.model.config_path + if qwen_vl_act_config_path.endswith(".json"): + model_config = ConfigClass.from_json_file(qwen_vl_act_config_path) + else: + model_config = ConfigClass.from_pretrained(qwen_vl_act_config_path) + + assert self.model_type in model_config.model_type, ( + f"Mismatch of model type: model type in config file is " + f"{model_config.model_type}, but the model type is {self.model_type}." + ) + + model_config.update_model_config(self._build_processor_dict()) + return model_config + + def create_model(self, processor, tokenizer_mixin, model_config): + ModelClass, _ = self._get_model_and_config_class() + use_selective_recompute = self.cfg.distributed.use_selective_recompute + return ModelClass( + model_config, + processor, + tokenizer_mixin=tokenizer_mixin, + use_selective_recompute=use_selective_recompute, + ) + + def load_weights(self, model, normalizer_action, normalizer_propri, **kwargs): + import copy + + processor = kwargs.get("processor") + + if self.cfg.model.pretrained_path: + model, err = load_qwen_pretrain_weight( + model, self.cfg.model.pretrained_path + ) + + if processor is not None: + model.resize_token_embeddings(len(processor.tokenizer)) + + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + elif ( + hasattr(model, "get_input_embeddings") + and model.get_input_embeddings() is not None + ): + + def _make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook( + _make_inputs_require_grad + ) + + if hasattr(model, "set_normalizer"): + model.set_normalizer( + copy.deepcopy(normalizer_action), + copy.deepcopy(normalizer_propri), + ) + + return model + + # ---- FSDP wrapping ---- + def get_transformer_layer_cls(self): + """Return the FSDP transformer-wrap layer classes for this variant. + + Subclass override - see the per-variant adapter for the actual + layer classes. + """ + raise NotImplementedError( + f"{type(self).__name__} must override get_transformer_layer_cls()" + ) + + # ---- dataset ---- + def load_dataset(self, data_config, processor, rank, world_size, **kwargs): + """Dispatch to the backend named by ``cfg.data.dataset_type``. + + The backend-specific wiring (resume indices, pool offsets, + processor wrapping) lives inside each backend's ``build()``; + this method only assembles a :class:`BuildContext` and forwards. + """ + import copy + + from wall_x.data import BuildContext, build_data + + resume_state = None + resume_batches = kwargs.get("resume_batches", 0) + indices = kwargs.get("resume_indices") + if indices is None and self.cfg.checkpoint.resume_from: + checkpoint_path = self.cfg.checkpoint.resume_from + if os.path.isfile(checkpoint_path): + checkpoint_path = os.path.dirname(checkpoint_path) + if os.path.isdir(checkpoint_path): + indices = self._load_episode_indices(checkpoint_path, rank) + + if resume_batches or indices or kwargs.get("episode_container_checkpoint_path"): + resume_state = { + "resume_batches": resume_batches, + "indices": indices, + "episode_container_checkpoint_path": kwargs.get( + "episode_container_checkpoint_path" + ), + } + + ctx = BuildContext( + rank=rank, + world_size=world_size, + processor=processor, + tokenizer_mixin=kwargs.get("tokenizer_mixin"), + normalizer_action=copy.deepcopy(kwargs.get("normalizer_action")), + normalizer_propri=copy.deepcopy(kwargs.get("normalizer_propri")), + model_config=kwargs.get("model_config"), + resume_state=resume_state, + ) + + bundle = build_data(self.cfg, ctx) + return bundle.dataset, bundle.train_loader, bundle.train_iters + + # ---- forward / loss ---- + def forward(self, model, batch, **kwargs): + self._last_batch_info = self._extract_batch_info(batch) + mode = kwargs.get("mode", "train") + return model(**batch, mode=mode) + + def extract_loss(self, outputs): + return outputs.loss + + def collect_output_stats( + self, + outputs, + step_stats, + reduce_tensor_fn, + true_gather_fn, + tokenizer_mixin=None, + ): + """Collect VL-specific auxiliary losses and accuracy metrics. + + Handles: + - cross_entropy_loss / flow_loss (all-reduce + true_gather) + - per-dataset channel losses (all-reduce via count-weighted avg) + - action_accuracy and extra RVQ layer accuracies + """ + import torch.distributed as dist + + cross_entropy_loss = self.get_output_field(outputs, "cross_entropy_loss") + if cross_entropy_loss is not None: + step_stats["cross_entropy_loss"] = reduce_tensor_fn( + cross_entropy_loss + ).item() + + flow_loss = self.get_output_field(outputs, "flow_loss") + if flow_loss is not None: + step_stats["flow_loss"] = reduce_tensor_fn(flow_loss).item() + + log_ce_loss = true_gather_fn( + self.get_output_field(outputs, "cross_entropy_loss") + ) + if log_ce_loss is not None: + step_stats["cross_entropy_loss"] = log_ce_loss + + log_flow_loss = true_gather_fn(self.get_output_field(outputs, "flow_loss")) + if log_flow_loss is not None: + step_stats["flow_loss"] = log_flow_loss + + channel_loss_dict = self.get_output_field(outputs, "channel_loss_dict") + channel_loss_count_dict = self.get_output_field( + outputs, "channel_loss_count_dict" + ) + if channel_loss_dict is not None: + for dataset_name_i in channel_loss_dict: + count_tensor = channel_loss_count_dict[dataset_name_i].clone() + loss_tensor = channel_loss_dict[dataset_name_i].detach().clone() + + dist.all_reduce(count_tensor, op=dist.ReduceOp.SUM) + dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM) + + cout_sum = count_tensor.item() + if cout_sum >= 0.5: + step_stats[f"channel_loss_{dataset_name_i}"] = ( + loss_tensor.item() / cout_sum + ) + + if "action_accuracy" in channel_loss_dict and tokenizer_mixin is not None: + acc_tensor = channel_loss_dict["action_accuracy"].detach().clone() + dist.all_reduce(acc_tensor, op=dist.ReduceOp.SUM) + world_size = dist.get_world_size() if dist.is_initialized() else 1 + step_stats["action_accuracy"] = acc_tensor.item() / world_size + + for key in tokenizer_mixin.get_accuracy_keys(): + if key != "action_accuracy" and key in channel_loss_dict: + rvq_acc_tensor = channel_loss_dict[key].detach().clone() + dist.all_reduce(rvq_acc_tensor, op=dist.ReduceOp.SUM) + step_stats[key] = rvq_acc_tensor.item() / world_size + + def collect_param_norms( + self, model, step_stats, device, reduce_tensor_fn, params_sharded=False + ): + """Compute per-group parameter L2 norms for VL models. + + Groups: + "visual" -> visual_param_norm + "action" -> action_expert_param_norm + others -> org_vlm_param_norm + all -> total_param_norm + """ + total_sq = torch.tensor(0.0, dtype=torch.float32, device=device) + visual_sq = torch.tensor(0.0, dtype=torch.float32, device=device) + action_sq = torch.tensor(0.0, dtype=torch.float32, device=device) + org_vlm_sq = torch.tensor(0.0, dtype=torch.float32, device=device) + dedicated_sq = torch.tensor(0.0, dtype=torch.float32, device=device) + dedicated_params = [] + + def add_named_sq(name, sq): + nonlocal total_sq, visual_sq, action_sq, org_vlm_sq + total_sq += sq + if "visual" in name: + visual_sq += sq + elif "action" in name: + action_sq += sq + else: + org_vlm_sq += sq + + with torch.no_grad(): + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + if param.numel() == 0 or hasattr(param, "_dedicated_owner_rank"): + continue + tensor = param.detach() + if hasattr(tensor, "to_local"): + tensor = tensor.to_local() + if tensor.numel() == 0: + continue + sq = torch.sum(tensor.float() ** 2) + add_named_sq(name, sq) + + if is_dmuon_model(model): + dedicated_params = self._get_named_dmuon_dedicated_params(model) + # DMuon replaces dedicated params with placeholders in + # named_parameters(). Count the authoritative dedicated + # storage here so total/category norms cover the full model. + for name, dparam in dedicated_params: + if getattr(dparam, "_dmuon_route", None) == "sharded_adamw": + replicate_group = getattr(dparam, "replicate_group", None) + # replicate_group=None is 1D shard-only mode: each rank + # contributes a distinct shard and the later reduce sums + # them into the global norm. + if ( + replicate_group is not None + and replicate_group.rank() + != getattr(dparam, "owner_replicate", 0) + ): + continue + tensor = getattr(dparam, "_sharded_adamw_data", None) + if tensor is None: + continue + valid_numel = int( + getattr( + dparam, + "_sharded_adamw_valid_numel", + tensor.numel(), + ) + ) + tensor = tensor[:valid_numel] + else: + if not bool(getattr(dparam, "is_owner", False)): + continue + tensor = getattr(dparam, "_owned_data", None) + if tensor is None: + continue + + if tensor.numel() == 0: + continue + sq = torch.sum(tensor.detach().float() ** 2) + dedicated_sq += sq + add_named_sq(name, sq) + + if params_sharded: + total_sq = reduce_tensor_fn(total_sq, average=False) + visual_sq = reduce_tensor_fn(visual_sq, average=False) + action_sq = reduce_tensor_fn(action_sq, average=False) + org_vlm_sq = reduce_tensor_fn(org_vlm_sq, average=False) + dedicated_sq = reduce_tensor_fn(dedicated_sq, average=False) + + step_stats["total_param_norm"] = torch.sqrt(total_sq).item() + step_stats["visual_param_norm"] = torch.sqrt(visual_sq).item() + step_stats["action_expert_param_norm"] = torch.sqrt(action_sq).item() + step_stats["org_vlm_param_norm"] = torch.sqrt(org_vlm_sq).item() + if dedicated_params: + step_stats["dmuon_dedicated_param_norm"] = torch.sqrt(dedicated_sq).item() + + # ---- MFU computation ---- + @staticmethod + def _extract_batch_info(batch): + """Extract token counts from a training batch.""" + info = {} + input_ids = batch.get("input_ids") + if input_ids is not None: + info["batch_size"] = input_ids.shape[0] + info["seq_length"] = input_ids.shape[1] + + moe_token_types = batch.get("moe_token_types") + if moe_token_types is not None: + info["num_lang_tokens"] = ( + int((moe_token_types == 0).sum().item()) // info["batch_size"] + ) + info["num_action_tokens"] = ( + int((moe_token_types == 1).sum().item()) // info["batch_size"] + ) + else: + info["num_lang_tokens"] = info.get("seq_length", 0) + info["num_action_tokens"] = 0 + + pixel_values = batch.get("pixel_values") + info["vision_seq_length"] = ( + pixel_values.shape[0] if pixel_values is not None else 0 + ) + + labels = batch.get("labels") + if labels is not None: + info["num_loss_tokens"] = int((labels[..., 1:] != -100).sum().item()) + else: + info["num_loss_tokens"] = None + + return info + + def _compute_detailed_flops(self, model, config): + """Compute detailed per-module FLOPs for one training step. + + Uses the same formulas as scripts/profile_forward.py:compute_module_flops, + with fwd+bwd multipliers: + - Frozen modules (e.g. ViT if frozen): 1x forward + - Trainable modules: 3x forward (fwd + 2x bwd) + """ + + batch_info = getattr(self, "_last_batch_info", None) + if not batch_info: + return None + + model_config = model.config if hasattr(model, "config") else None + if model_config is None: + return None + + B = batch_info["batch_size"] + S = batch_info["seq_length"] + num_lang = batch_info["num_lang_tokens"] + num_act = batch_info["num_action_tokens"] + N_lang = B * num_lang + N_act = B * num_act + N_total = N_lang + N_act + Nv = batch_info["vision_seq_length"] + num_loss_tokens = batch_info["num_loss_tokens"] + + H = model_config.hidden_size + num_layers = model_config.num_hidden_layers + num_heads = model_config.num_attention_heads + num_kv = model_config.num_key_value_heads + vocab_size = getattr(model_config, "padded_vocab_size", model_config.vocab_size) + + use_mot = getattr(model_config, "attention_moe", False) + use_moe_mlp = getattr(model_config, "mlp_moe", False) + dim_inputs = getattr(model_config, "dim_inputs", (H, H)) + dim_lang, dim_act = dim_inputs + + grad_accum = self.cfg.hyperparams.gradient_accumulation_steps + + vit_frozen = True + if hasattr(model, "visual"): + for p in model.visual.parameters(): + if p.requires_grad: + vit_frozen = False + break + vit_mult = 1.0 if vit_frozen else 3.0 + + vit_fwd_flops = 0 + if hasattr(model_config, "vision_config") and Nv > 0: + vcfg = model_config.vision_config + Hv = vcfg.hidden_size + Iv = vcfg.intermediate_size + out_hidden = vcfg.out_hidden_size + depth_v = vcfg.depth + + F_qkv_v = 6 * Nv * Hv * Hv + F_o_v = 2 * Nv * Hv * Hv + F_mlp_v = 6 * Nv * Hv * Iv + 2 * Nv * Iv + F_linear_per_layer = F_qkv_v + F_o_v + F_mlp_v + + fullatt_set = set(getattr(vcfg, "fullatt_block_indexes", [])) + num_images = max(Nv // 768, 1) + si = Nv // num_images + sum_si_sq = num_images * si * si + sum_wi_sq = sum_si_sq // 16 + + for i in range(depth_v): + if i in fullatt_set: + vit_fwd_flops += F_linear_per_layer + 4 * Hv * sum_si_sq + else: + vit_fwd_flops += F_linear_per_layer + 4 * Hv * sum_wi_sq + + merge_unit = getattr(vcfg, "spatial_merge_size", 2) ** 2 + merger_hidden = Hv * merge_unit + Nv_merged = Nv // merge_unit + vit_fwd_flops += ( + 2 * Nv_merged * merger_hidden * merger_hidden + + 2 * Nv_merged * merger_hidden * out_hidden + ) + + kv_ratio = num_kv / num_heads + F_matmul_fwd = 4 * B * (S**2) * H + + if not use_mot: + F_attn_fwd = ( + 2 * N_total * H * H + + 4 * N_total * H * H * kv_ratio + + 2 * N_total * H * H + + F_matmul_fwd + ) + else: + F_attn_fwd = ( + N_lang * dim_lang * H * (4 + 4 * kv_ratio) + + N_act * dim_act * H * (4 + 4 * kv_ratio) + + F_matmul_fwd + ) + + if not use_moe_mlp: + ffn_hidden = model_config.intermediate_size + F_mlp_fwd = 6 * N_total * H * ffn_hidden + 2 * N_total * ffn_hidden + else: + hid_lang = model_config.experts[0]["intermediate_size"] + hid_act = model_config.experts[1]["intermediate_size"] + F_mlp_fwd = (6 * N_lang * dim_lang * hid_lang + 2 * N_lang * hid_lang) + ( + 6 * N_act * dim_act * hid_act + 2 * N_act * hid_act + ) + + decoder_fwd_flops = num_layers * (F_attn_fwd + F_mlp_fwd) + + N_lm = num_loss_tokens if num_loss_tokens is not None else N_total + lm_head_fwd_flops = 2 * N_lm * H * vocab_size + + total_flops = ( + vit_mult * vit_fwd_flops + 3.0 * decoder_fwd_flops + 3.0 * lm_head_fwd_flops + ) * grad_accum + + return { + "total_flops": total_flops, + "vit_fwd_flops": vit_fwd_flops, + "decoder_fwd_flops": decoder_fwd_flops, + "lm_head_fwd_flops": lm_head_fwd_flops, + "vit_mult": vit_mult, + } + + def compute_mfu(self, model, step_time_seconds): + """Compute detailed Model FLOPs Utilization for VL transformer training.""" + if step_time_seconds <= 0: + return None + try: + config = self._build_processor_dict() + flops_info = self._compute_detailed_flops(model, config) + if flops_info is None: + return None + + gpu_peak_tflops = 312.0 # TODO: move to DebugConfig if needed + peak_flops = gpu_peak_tflops * 1e12 + + # Data-parallel: each GPU independently processes its own micro-batch, + # so per-GPU MFU = per_gpu_flops / (step_time * per_gpu_peak). + # No division by num_gpus needed. + total_flops = flops_info["total_flops"] + mfu = total_flops / (step_time_seconds * peak_flops) + + return { + "mfu": mfu, + "flops_per_step_T": total_flops / 1e12, + } + except Exception: + return None + + # ---- optional hooks ---- + @staticmethod + def log_attention_implementation(logger, model): + """Log the attention implementation name. Variant-specific because the + model layout (model.model vs model.model.language_model) differs.""" + raise NotImplementedError( + "VLAdapter subclasses must override log_attention_implementation()" + ) + + def _load_episode_indices(self, checkpoint_path: str, rank: int): + """Read backend-managed per-rank resume offsets if supported.""" + from wall_x.data import data_backend + + backend = data_backend() + if not backend.supports("load_episode_indices"): + return None + return backend.load_episode_indices(checkpoint_path, rank) + + # ---- prediction / inference ---- + def predict( + self, + prediction_type: str, + *, + model, + val_dataloader, + rank, + world_size, + device, + processor, + tokenizer_mixin, + logger, + current_step, + max_iteration=None, + save_dir=None, + max_samples=None, + ): + """Dispatch VLA inference by prediction_type. + + flow_action / dllm_action return per-rank L1 action metrics; ar_action + is reserved (was unimplemented pre-refactor too); text runs text + generation with optional point-L1 distance scoring. + """ + config = self._build_processor_dict() + if prediction_type == "flow_action": + return self._predict_flow_action( + model=model, + val_dataloader=val_dataloader, + rank=rank, + world_size=world_size, + device=device, + config=config, + logger=logger, + current_step=current_step, + max_iteration=max_iteration, + ) + if prediction_type == "ar_action": + return None # matches pre-refactor behaviour (was pass) + if prediction_type == "dllm_action": + return self._predict_dllm_action( + model=model, + val_dataloader=val_dataloader, + rank=rank, + world_size=world_size, + device=device, + config=config, + processor=processor, + tokenizer_mixin=tokenizer_mixin, + logger=logger, + current_step=current_step, + max_iteration=max_iteration, + ) + if prediction_type == "text": + return self._predict_text( + model=model, + val_dataloader=val_dataloader, + rank=rank, + device=device, + logger=logger, + current_step=current_step, + max_samples=max_samples, + save_dir=save_dir, + ) + raise ValueError(f"Unsupported prediction type: {prediction_type}") + + @torch.no_grad() + def _predict_flow_action( + self, + *, + model, + val_dataloader, + rank, + world_size, + device, + config, + logger, + current_step, + max_iteration, + ): + if dist.is_initialized(): + dist.barrier() + + total_num = len(val_dataloader) + if max_iteration: + total_num = min(max_iteration, total_num) + model.eval() + + all_preds, all_actions = [], [] + pepoch = tqdm( + total=total_num, + desc=f"Predicting ckpt at step {current_step}", + disable=rank != 0, + ) + for batch_idx, batch in enumerate(val_dataloader): + if batch_idx >= total_num: + break + batch = _move_batch(batch, device) + + model_output = model.generate_flow_action( + action_horizon=config["data"]["action_horizon_flow"], + action_dim=model.action_preprocessor.action_dim, + **batch, + ) + pred_action, gt_action = ( + model_output["predict_action"], + model_output["gt_action"], + ) + + pred_list = [torch.zeros_like(pred_action) for _ in range(world_size)] + gt_list = [torch.zeros_like(gt_action) for _ in range(world_size)] + dist.all_gather(pred_list, pred_action) + dist.all_gather(gt_list, gt_action) + + if dist.is_initialized(): + dist.barrier() + if rank == 0: + all_preds.append(torch.cat(pred_list, dim=0).cpu()) + all_actions.append(torch.cat(gt_list, dim=0).cpu()) + pepoch.update(1) + + pepoch.close() + + if rank == 0: + step_log = {} + if all_preds: + all_preds = torch.cat(all_preds, dim=0) + all_actions = torch.cat(all_actions, dim=0) + step_log = compute_action_metrics( + all_preds, all_actions, config=config, step_log=step_log + ) + if logger and step_log: + logger.log(step_log, step=current_step) + if self.logger is not None: + self.logger.info( + f"Step {current_step}, Validation L1 Loss: {step_log.get('val_action_l1', -1)}" + ) + return step_log + return None + + @torch.no_grad() + def _predict_dllm_action( + self, + *, + model, + val_dataloader, + rank, + world_size, + device, + config, + processor, + tokenizer_mixin, + logger, + current_step, + max_iteration, + ): + if dist.is_initialized(): + dist.barrier() + + total_num = len(val_dataloader) + if max_iteration: + total_num = min(max_iteration, total_num) + model.eval() + + all_preds, all_actions = [], [] + pepoch = tqdm( + total=total_num, + desc=f"Predicting ckpt at step {current_step}", + disable=rank != 0, + ) + for batch_idx, batch in enumerate(val_dataloader): + if batch_idx >= total_num: + break + batch = _move_batch(batch, device) + batch = self._preprocess_dllm_batch(batch, processor, tokenizer_mixin) + + if hasattr(model.action_tokenizer, "max_waypoints"): + total_ar_step = model.action_tokenizer.max_waypoints + elif hasattr(model.action_tokenizer, "max_length"): + total_ar_step = model.action_tokenizer.max_length + else: + raise ValueError( + "Unknown action_tokenizer type for dllm action inference" + ) + model_output = model.generate_dllm_action( + action_dim=7, + action_horizon=config["data"]["action_horizon_flow"], + use_ar_action=False, + total_ar_step=total_ar_step, + **batch, + ) + pred_action, gt_action = ( + model_output["predict_action"], + model_output["gt_action"], + ) + + pred_list = [torch.zeros_like(pred_action) for _ in range(world_size)] + gt_list = [torch.zeros_like(gt_action) for _ in range(world_size)] + dist.all_gather(pred_list, pred_action) + dist.all_gather(gt_list, gt_action) + + if dist.is_initialized(): + dist.barrier() + if rank == 0: + all_preds.append(torch.cat(pred_list, dim=0).cpu()) + all_actions.append(torch.cat(gt_list, dim=0).cpu()) + pepoch.update(1) + + pepoch.close() + + if rank == 0: + step_log = {} + if all_preds: + all_preds = torch.cat(all_preds, dim=0) + all_actions = torch.cat(all_actions, dim=0) + step_log = compute_action_metrics( + all_preds, all_actions, config=config, step_log=step_log + ) + if logger and step_log: + logger.log(step_log, step=current_step) + if self.logger is not None: + self.logger.info( + f"Step {current_step}, Validation L1 Loss: {step_log.get('val_action_l1', -1)}" + ) + return step_log + return None + + @staticmethod + def _preprocess_dllm_batch(batch, processor, tokenizer_mixin): + """Patch the dllm action placeholders into input_ids (in place).""" + input_ids = batch["input_ids"] + prefix_length = batch["prefix_length"] + placeholder_str = tokenizer_mixin.get_placeholder_for_dllm() + placeholder_seq = torch.tensor( + processor.tokenizer.convert_tokens_to_ids(placeholder_str) + ) + placeholder_len = placeholder_seq.shape[0] + input_ids[:, prefix_length - placeholder_len - 2 : prefix_length - 2] = ( + placeholder_seq + ) + batch.update({"input_ids": input_ids}) + return batch + + @torch.no_grad() + def _predict_text( + self, + *, + model, + val_dataloader, + rank, + device, + logger, + current_step, + max_samples, + save_dir, + ): + from wall_x._vendor.x2robot_utils.grounding import calculate_point_l1_distance + + if dist.is_initialized(): + dist.barrier() + + total_num = ( + len(val_dataloader) + if max_samples is None + else min(max_samples, len(val_dataloader)) + ) + model.eval() + + all_input_texts, all_gt_texts, all_pred_texts = [], [], [] + all_point_l1_distances = [] + + pepoch = tqdm( + total=total_num, + desc=f"Predicting text at step {current_step}", + disable=rank != 0, + ) + + for batch_idx, batch in enumerate(val_dataloader): + if batch_idx >= total_num: + break + batch = _move_batch(batch, device) + + model_output = model.generate_text( + input_ids=batch.get("input_ids"), + attention_mask=batch.get("attention_mask"), + moe_token_types=batch.get("moe_token_types"), + pixel_values=batch.get("pixel_values"), + image_grid_thw=batch.get("image_grid_thw"), + proprioception=batch.get("proprioception"), + dataset_names=batch.get("dataset_names"), + dof_mask=batch.get("dof_mask"), + agent_pos_mask=batch.get("agent_pos_mask"), + prefix_length=batch.get("prefix_length"), + positional_masks=batch.get("positional_masks"), + re_generate=False, + ) + + input_texts = model_output["input_text"] + gt_texts = model_output["gt_output_text"] + pred_texts = model_output["predict_output_text"] + + if rank == 0: + all_input_texts.extend(input_texts) + all_gt_texts.extend(gt_texts) + all_pred_texts.extend(pred_texts) + for gt_text, pred_text in zip(gt_texts, pred_texts): + gt_clean = gt_text[0] if isinstance(gt_text, list) else gt_text + pred_clean = ( + pred_text[0] if isinstance(pred_text, list) else pred_text + ) + point_l1_dist = calculate_point_l1_distance(gt_clean, pred_clean) + if point_l1_dist is not None: + all_point_l1_distances.append(point_l1_dist) + + pepoch.update(1) + + pepoch.close() + + if rank == 0 and all_pred_texts: + if all_point_l1_distances: + avg_point_l1 = sum(all_point_l1_distances) / len(all_point_l1_distances) + min_point_l1 = min(all_point_l1_distances) + max_point_l1 = max(all_point_l1_distances) + point_stats = { + "text_prediction_point_l1_avg": avg_point_l1, + "text_prediction_point_l1_min": min_point_l1, + "text_prediction_point_l1_max": max_point_l1, + "text_prediction_point_samples_count": len(all_point_l1_distances), + } + if logger is not None: + logger.log(point_stats, step=current_step) + if self.logger is not None: + self.logger.info( + f"Point L1 Statistics - Avg: {avg_point_l1:.4f}, " + f"Min: {min_point_l1:.4f}, Max: {max_point_l1:.4f}, " + f"Samples: {len(all_point_l1_distances)}" + ) + save_text_results_to_file( + current_step, + all_input_texts, + all_gt_texts, + all_pred_texts, + save_dir if save_dir is not None else "./", + ) + + +def _move_batch(batch, device): + """Recursively move dict/list/tensor to device (VLA-local copy to keep + VLAdapter self-contained; base_adapter has a similar helper used by + validation).""" + if isinstance(batch, dict): + return {k: _move_batch(v, device) for k, v in batch.items()} + if isinstance(batch, list): + return [_move_batch(v, device) for v in batch] + if torch.is_tensor(batch): + return batch.to(device) + return batch diff --git a/wall_x/trainer/fsdp_trainer/__init__.py b/wall_x/trainer/fsdp_trainer/__init__.py new file mode 100644 index 0000000..891c586 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/__init__.py @@ -0,0 +1,34 @@ +"""Distributed trainer package exports.""" + +from wall_x.trainer.fsdp_trainer.base_trainer import ( + BaseDistributedTrainer, + DistributedTrainer, + all_gather, + all_reduce, + barrier, + cleanup_distributed, + get_local_rank, + get_rank, + get_world_size, + is_main_process, + setup_distributed, +) +from wall_x.trainer.fsdp_trainer.fsdp_trainer import FSDPTrainer + +__all__ = [ + # Base classes and utilities + "BaseDistributedTrainer", + "DistributedTrainer", + "setup_distributed", + "cleanup_distributed", + "get_rank", + "get_world_size", + "get_local_rank", + "is_main_process", + "barrier", + "all_reduce", + "all_gather", + # Trainers + # "DDPTrainer", + "FSDPTrainer", +] diff --git a/wall_x/trainer/fsdp_trainer/base_trainer.py b/wall_x/trainer/fsdp_trainer/base_trainer.py new file mode 100644 index 0000000..2ad2ea4 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/base_trainer.py @@ -0,0 +1,430 @@ +""" +Base Trainer for Pure PyTorch Distributed Training (DDP/FSDP) +No accelerate dependency - can be launched directly with torchrun +""" + +import contextlib +import gc +import logging +import os +from abc import ABC, abstractmethod +from datetime import timedelta +from typing import Any, Optional + +import torch +import torch.distributed as dist + +from wall_x.config.hyperparams_config import AdamWConfig +from wall_x.config.train_config import TrainConfig +from wall_x.trainer.trainer_utils import seed_all +from wall_x.utils.logger import DistributedLogger +from wall_x.utils.timers import Timers + + +def setup_distributed(): + """Initialize distributed training environment. + + NCCL watchdog default is 10 min - some dataset backends (e.g. lerobot's + per-rank sequential LeRobotDataset indexing) can exceed that during the + first build, causing barrier #2 to time out on the waiting rank. Allow + the timeout to be raised via ``WALLX_DIST_TIMEOUT_MINUTES`` (default 30). + """ + if not dist.is_initialized(): + timeout_min = int(os.environ.get("WALLX_DIST_TIMEOUT_MINUTES", "30")) + dist.init_process_group( + backend="nccl", + timeout=timedelta(minutes=timeout_min), + ) + + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + + return local_rank + + +def cleanup_distributed(): + """Clean up distributed training environment""" + if dist.is_initialized(): + dist.destroy_process_group() + + +def get_rank() -> int: + """Get current process rank""" + if dist.is_initialized(): + return dist.get_rank() + return 0 + + +def get_world_size() -> int: + """Get total number of processes""" + if dist.is_initialized(): + return dist.get_world_size() + return 1 + + +def get_local_rank() -> int: + """Get local rank (GPU index on current node)""" + return int(os.environ.get("LOCAL_RANK", 0)) + + +def is_main_process() -> bool: + """Check if current process is the main process (rank 0)""" + return get_rank() == 0 + + +# Trainer-level collectives should use an explicit global process group. +# FSDP's mesh groups may be 1D or 2D and have sharding-specific semantics, +# while trainer barriers/metric reductions need all-rank semantics. +_trainer_process_group = None + + +def set_trainer_process_group(pg): + global _trainer_process_group + _trainer_process_group = pg + + +def barrier(): + """Synchronize all processes""" + if dist.is_initialized(): + dist.barrier(group=_trainer_process_group) + + +def all_reduce(tensor: torch.Tensor, op=dist.ReduceOp.SUM) -> torch.Tensor: + """All-reduce tensor across all processes""" + if dist.is_initialized(): + dist.all_reduce(tensor, op=op, group=_trainer_process_group) + return tensor + + +def all_gather(tensor: torch.Tensor) -> torch.Tensor: + """All-gather tensor from all processes""" + if not dist.is_initialized(): + return tensor + + group = _trainer_process_group + world_size = dist.get_world_size(group=group) + gathered = [torch.zeros_like(tensor) for _ in range(world_size)] + dist.all_gather(gathered, tensor, group=group) + return torch.stack(gathered) + + +class BaseDistributedTrainer(ABC): + """ + Abstract base class for distributed training without accelerate. + Supports both DDP and FSDP through subclasses. + + Launch with torchrun: + torchrun --nproc_per_node=8 train.py --config config.yaml + """ + + @abstractmethod + def train_loop(self, epoch: int): + """Training loop for one epoch""" + raise NotImplementedError("train_loop must be implemented") + + @abstractmethod + def val_loop(self): + """Validation loop""" + raise NotImplementedError("val_loop must be implemented") + + @abstractmethod + def save_checkpoint(self, epoch: int, step: int = 0): + """Save model checkpoint""" + raise NotImplementedError("save_checkpoint must be implemented") + + @abstractmethod + def load_model(self): + """Load model""" + raise NotImplementedError("load_model must be implemented") + + @abstractmethod + def load_dataset(self): + """Load dataset""" + raise NotImplementedError("load_dataset must be implemented") + + @abstractmethod + def backward(self, loss: torch.Tensor): + """Perform backward pass""" + raise NotImplementedError("backward must be implemented") + + @abstractmethod + def clip_grad_norm(self, max_norm: float) -> torch.Tensor: + """Clip gradient norm""" + raise NotImplementedError("clip_grad_norm must be implemented") + + +class DistributedTrainer(BaseDistributedTrainer): + """ + Concrete base trainer with common functionality for DDP/FSDP. + Subclasses should implement wrap_model, backward, clip_grad_norm, etc. + """ + + def __init__( + self, + train_config: TrainConfig, + wandb_run: Optional[Any] = None, + ): + self.cfg = train_config + + # wandb Run object (metric recorder) - None on non-main ranks and + # when use_wandb is disabled. Distinct from self.logger below. + self.wandb_run = wandb_run + + # Initialize distributed environment + self.local_rank = setup_distributed() + self.rank = get_rank() + self.world_size = get_world_size() + self.device = torch.device(f"cuda:{self.local_rank}") + + # Dataset config: optional backend-provided trainer config. + from wall_x.trainer.adapters.base_adapter import load_trainer_data_config + + self.data_config = load_trainer_data_config(self.cfg) + + # Text logger + self.logger = DistributedLogger( + name=self.__class__.__name__, + save_path=self.cfg.checkpoint.save_path, + ) + + # Training state + self.seed = self.cfg.hyperparams.seed + self.logger.info(f"seed {self.seed}") + + seed_all(self.seed) + self.start_epoch = 0 + self.global_step = 0 + self.micro_step = 0 + self.num_epoch = self.cfg.hyperparams.num_epoch + self.dataset_config_path = getattr(self.cfg.data, "dataset_config_path", None) + self.initial_step = 0 + + # Optimizer hyperparameters (betas / weight_decay / eps) + # used by build_optimizer for both AdamW and Muon paths. + opt = self.cfg.hyperparams.optimizer + if isinstance(opt, AdamWConfig): + self.adamw_betas = tuple(opt.betas) + self.adamw_weight_decay = opt.weight_decay + self.adamw_eps = opt.eps + elif getattr(opt, "optimizer_type", None) == "muon": + self.adamw_betas = (opt.beta_1, opt.beta_2) + self.adamw_weight_decay = opt.weight_decay + self.adamw_eps = opt.eps + else: + self.adamw_betas = (0.9, 0.98) + self.adamw_weight_decay = 1e-8 + self.adamw_eps = 1e-8 + + # Performance options + self.nvtx = self.cfg.debug.nvtx + self.timers = Timers(log_level=0, log_option="minmax") + + # Training hyperparameters + self.max_grad_norm = opt.max_grad_norm + self.grad_accum_steps = self.cfg.hyperparams.gradient_accumulation_steps + self.log_interval = self.cfg.logging.log_interval + self.log_stats_buffer = [] + + # Model and optimizer (to be set by subclass) + self.model = None + self.optimizer = None + self.lr_scheduler = None + + # Mixed precision + dist_cfg = self.cfg.distributed + self.use_amp = dist_cfg.use_amp + self.amp_dtype = torch.bfloat16 if dist_cfg.bf16 else torch.float16 + self.grad_scaler = None + if self.use_amp and self.amp_dtype == torch.float16: + self.grad_scaler = torch.amp.GradScaler("cuda") + + def log( + self, message: str, level: int = logging.INFO, main_process_only: bool = True + ): + """Log a text message via self.logger. + + Kept for backward compat with call sites that still use ``self.log(...)`` + (new code should call ``self.logger.info/warning/error(...)`` directly). + Under the new DistributedLogger the ``main_process_only`` flag is + redundant for stdout (stdout is rank-0-only by construction); it still + controls whether non-zero ranks write the message to their file log. + """ + if main_process_only and not is_main_process(): + return + self.logger.log(message, level=level, main_process_only=main_process_only) + + def fit(self): + """Main training loop""" + barrier() + + # Optional: validate before training + if self.cfg.checkpoint.validate_first: + self.val_loop() + barrier() + + if self.nvtx: + torch.cuda.cudart().cudaProfilerStart() + + # Profiler setup + enable_profiling = self.cfg.debug.profile + profiler = ( + self._setup_profiler() if enable_profiling else contextlib.nullcontext() + ) + + for epoch in range(self.start_epoch, self.num_epoch): + # Training + self.train_loop(epoch, profiler=profiler) + barrier() + num_training_steps = getattr(self, "num_training_steps", 0) + if num_training_steps > 0 and self.global_step >= num_training_steps: + break + + # Save checkpoint + if (epoch + 1) % self.cfg.logging.epoch_save_interval == 0: + self.save_checkpoint(epoch) + + # Validation + self.val_loop() + barrier() + + gc.collect() + torch.cuda.empty_cache() + + if self.nvtx: + torch.cuda.cudart().cudaProfilerStop() + + def _setup_profiler(self): + """Setup PyTorch profiler""" + return torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + schedule=torch.profiler.schedule( + wait=self.cfg.debug.profile_wait_iters, + warmup=self.cfg.debug.profile_warmup_iters, + active=self.cfg.debug.profile_active_iters, + ), + on_trace_ready=torch.profiler.tensorboard_trace_handler( + self.cfg.debug.profile_save_path, + worker_name=f"worker{self.rank}", + ), + record_shapes=True, + profile_memory=True, + with_stack=True, + ) + + def gather_tensor(self, tensor: torch.Tensor) -> torch.Tensor: + """Gather tensor from all processes""" + return all_gather(tensor) + + def reduce_tensor(self, tensor: torch.Tensor, average: bool = True) -> torch.Tensor: + """Reduce tensor across all processes""" + tensor = tensor.clone() + all_reduce(tensor) + if average: + tensor = tensor / self.world_size + return tensor + + def sync_gradients(self) -> bool: + """Check if gradients should be synchronized (for gradient accumulation)""" + return (self.micro_step + 1) % self.grad_accum_steps == 0 + + def optimizer_zero_grad(self): + """Zero optimizer gradients""" + self.optimizer.zero_grad(set_to_none=True) + + def optimizer_step(self): + """Perform optimizer step""" + if self.grad_scaler is not None: + self.grad_scaler.step(self.optimizer) + self.grad_scaler.update() + else: + self.optimizer.step() + + def lr_scheduler_step(self): + """Perform learning rate scheduler step""" + if self.lr_scheduler is not None: + self.lr_scheduler.step() + + def get_lr(self) -> float: + """Get current learning rate""" + if self.lr_scheduler is not None: + return self.lr_scheduler.get_last_lr()[0] + return self.optimizer.param_groups[0]["lr"] + + def autocast_context(self): + """Get autocast context for mixed precision""" + if self.use_amp: + return torch.amp.autocast("cuda", dtype=self.amp_dtype) + return contextlib.nullcontext() + + def training_log( + self, + current_epoch: int, + total_epoch: int, + current_iter: int, + total_iter: int, + loss: torch.Tensor, + lr: float, + time_per_step: float, + show_time_details: bool = False, + ): + """Log training progress""" + if not is_main_process(): + return + + # Prefer the cross-rank reduced loss (train_loss) from + # _current_step_stats so the console line matches what wall_x_2509's + # fsdp_trainer prints and what wandb sees. Falls back to the raw + # rank-0 local tensor when stats haven't been populated yet. + _smoothed = getattr(self, "_current_step_stats", None) or {} + _loss_to_print = _smoothed.get("train_loss", float(loss)) + + log_string = "" + log_string += f" epoch {current_epoch:3d}/{total_epoch:3d} |" + log_string += f" iter {current_iter:6d}/{total_iter:6d} |" + log_string += f" loss {_loss_to_print:.10f} |" + log_string += f" lr {lr:.6f} |" + log_string += f" time {time_per_step:.4f}s |" + + self.log(log_string) + + if show_time_details: + timers_to_log = [ + "interval-time", + "data-load", + "forward-compute", + "backward-compute", + "optimizer", + ] + self.timers.log(timers_to_log, normalizer=1) + + def true_gather(self, value): + """Gather values across all processes and compute mean over non-None values.""" + device = next(self.model.parameters()).device + + # CRITICAL: must clone() here, not just detach(). + # detach() only disconnects the computation graph but still shares the + # same underlying storage. dist.all_reduce() is **in-place** - it + # overwrites that storage with the sum across all ranks. If the caller + # passes a tensor that is also referenced elsewhere (e.g. outputs["loss"] + # and outputs["flow_loss"] pointing to the same scalar_loss), the + # in-place all_reduce silently mutates the original, causing downstream + # readers (like training_log) to see a value multiplied by world_size. + value_to_gather = ( + value.detach().clone() + if value is not None + else torch.tensor(0.0, device=device) + ) + count_to_gather = torch.tensor(1.0 if value is not None else 0.0, device=device) + + # all_reduce is cheaper than all_gather because only the sum is needed. + dist.all_reduce(value_to_gather, op=dist.ReduceOp.SUM) + dist.all_reduce(count_to_gather, op=dist.ReduceOp.SUM) + + total_processes = count_to_gather.item() + if total_processes > 0: + return value_to_gather.item() / total_processes + else: + return None diff --git a/wall_x/trainer/fsdp_trainer/checkpoint_io.py b/wall_x/trainer/fsdp_trainer/checkpoint_io.py new file mode 100644 index 0000000..77755ac --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/checkpoint_io.py @@ -0,0 +1,1033 @@ +"""Checkpoint save/load helpers for distributed training.""" + +from __future__ import annotations + +import gc +import logging +import os +import random +import shutil +import time +from typing import Callable, Optional, Tuple + +import numpy as np +import torch +import torch.distributed as dist +import yaml +from safetensors.torch import load_file, save_file +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_state_dict, + set_model_state_dict, + set_optimizer_state_dict, +) +from torch.nn.parallel import DistributedDataParallel as DDP + +from wall_x.trainer.optimizer.dmuon import is_dmuon_model + + +def _noop_log(_msg: str, **_kw) -> None: + pass + + +def _dict_section(config: dict, key: str) -> dict: + section = config.get(key, {}) + return section if isinstance(section, dict) else {} + + +# ---------------------------------------------------------------------- +# Detectors +# ---------------------------------------------------------------------- + + +def _is_fsdp2_model(model: torch.nn.Module) -> bool: + """Detect FSDP2 models by the presence of DTensor parameters. + + FSDP2's ``fully_shard`` converts parameters in-place to DTensors without + wrapping the module in an outer class. DDP wraps in + ``DistributedDataParallel``; unwrapped models have plain + ``torch.Tensor`` parameters. + """ + if isinstance(model, DDP): + return False + try: + from torch.distributed.tensor import DTensor + except ImportError: + return False + for p in model.parameters(): + if isinstance(p, DTensor): + return True + return False + + +def _detect_legacy_fsdp1_format(checkpoint_path: str) -> bool: + """True when the ckpt directory looks like a pre-migration FSDP1 save. + + Legacy: ``model.safetensors`` is present (rank-0 full state), but + optimizer state lives in per-rank ``optimizer_rank{N}.pt`` files + instead of a single ``optimizer.pt``. After the FSDP1 -> FSDP2 + migration we cannot reshard those flat_param-keyed optim files into + the new DTensor layout, so the legacy loader cold-starts the + optimizer and warns. + """ + if not os.path.isdir(checkpoint_path): + return False + if os.path.exists(os.path.join(checkpoint_path, "optimizer.pt")): + return False + try: + return any( + f.startswith("optimizer_rank") and f.endswith(".pt") + for f in os.listdir(checkpoint_path) + ) + except OSError: + return False + + +# ---------------------------------------------------------------------- +# Save +# ---------------------------------------------------------------------- + + +def save_checkpoint( + *, + ckpt_path: str, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + lr_scheduler, + config: dict, + rank: int, + is_main: bool, + epoch: int, + global_step: int, + seed: int, + normalizer_action, + normalizer_propri, + dataset=None, + grad_scaler=None, + log_fn: Optional[Callable] = None, + frozen_prefixes: Optional[Tuple[str, ...]] = None, +) -> None: + """Save model + optimizer + scheduler + metadata at *ckpt_path*. + + ``frozen_prefixes``: when set, model state-dict keys starting with any + of these prefixes are excluded from the saved file. This is used for + frozen-by-design submodules that should not be duplicated in every + checkpoint. Loaders already use ``strict=False``. + """ + log_fn = log_fn or _noop_log + os.makedirs(ckpt_path, exist_ok=True) + + if is_main: + _save_training_checkpoint_metadata( + ckpt_path=ckpt_path, + config=config, + epoch=epoch, + global_step=global_step, + seed=seed, + world_size=_world_size_for_metadata(), + normalizer_action=normalizer_action, + normalizer_propri=normalizer_propri, + grad_scaler=grad_scaler, + log_fn=log_fn, + ) + + if is_dmuon_model(model): + _save_dmuon_state_dict( + ckpt_path=ckpt_path, + model=model, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + is_main=is_main, + log_fn=log_fn, + frozen_prefixes=frozen_prefixes, + ) + elif _is_fsdp2_model(model): + _save_fsdp2_full_state_dict( + ckpt_path=ckpt_path, + model=model, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + is_main=is_main, + log_fn=log_fn, + frozen_prefixes=frozen_prefixes, + ) + else: + # DDP or unwrapped (fallback). + _save_ddp_state_dict( + ckpt_path=ckpt_path, + model=model, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + is_main=is_main, + log_fn=log_fn, + frozen_prefixes=frozen_prefixes, + ) + + # Per-rank dataset state (only when saving mid-epoch, step != 0). + if dataset is not None and global_step != 0: + _save_dataset_state( + ckpt_path=ckpt_path, dataset=dataset, rank=rank, log_fn=log_fn + ) + + +def _save_training_checkpoint_metadata( + *, + ckpt_path: str, + config: dict, + epoch: int, + global_step: int, + seed: int, + world_size: int, + normalizer_action, + normalizer_propri, + grad_scaler=None, + log_fn: Callable, +) -> None: + torch.save({"seed": seed}, os.path.join(ckpt_path, "seed.pth")) + torch.save({"global_step": global_step}, os.path.join(ckpt_path, "global_step.pth")) + torch.save({"current_epoch": epoch}, os.path.join(ckpt_path, "current_epoch.pth")) + + # world_size: public metadata, used at resume to detect reshard. + # Written unconditionally (previously sharded-mode only). + torch.save({"world_size": world_size}, os.path.join(ckpt_path, "world_size.pth")) + + # RNG state (rank-0 snapshot; every rank restores the same state on + # resume, matching the existing seed_all(seed) convention where every + # rank is seeded identically). + rng_state = { + "torch": torch.get_rng_state(), + "numpy": np.random.get_state(), + "python": random.getstate(), + } + if torch.cuda.is_available(): + rng_state["cuda"] = torch.cuda.get_rng_state() + torch.save(rng_state, os.path.join(ckpt_path, "rng_state.pt")) + + # GradScaler state (fp16 AMP only). + if grad_scaler is not None: + torch.save(grad_scaler.state_dict(), os.path.join(ckpt_path, "grad_scaler.pt")) + + with open(os.path.join(ckpt_path, "config.yml"), "w") as f: + yaml.dump(config, f, default_flow_style=False, allow_unicode=True) + + model_cfg = _dict_section(config, "model") + data_cfg = _dict_section(config, "data") + + # Copy processor files to checkpoint directory. + processor_dir = model_cfg.get("processor_path") or config.get("processor_path") + if processor_dir is None: + # Backward compatibility: fall back to pretrained_qwen_vl_path. + processor_dir = config.get("pretrained_qwen_vl_path") + if processor_dir is not None: + log_fn( + "WARNING: 'pretrained_qwen_vl_path' is deprecated for processor " + "file copying, please use 'processor_path' instead.", + level=logging.WARNING, + ) + + if processor_dir is not None: + for filename in ( + "preprocessor_config.json", + "tokenizer_config.json", + "tokenizer.json", + "tokenizer.model", + "special_tokens_map.json", + "vocab.json", + ): + src = os.path.join(processor_dir, filename) + if os.path.exists(src): + shutil.copy(src, os.path.join(ckpt_path, filename)) + + act_config_path = model_cfg.get("config_path") or config.get( + "qwen_vl_act_config_path" + ) + if act_config_path is not None: + if os.path.exists(act_config_path): + shutil.copy(act_config_path, os.path.join(ckpt_path, "config.json")) + log_fn(f"[Checkpoint] Copied act config to {ckpt_path}/config.json") + else: + log_fn(f"[Checkpoint] WARNING: {act_config_path} not found, skipping.") + + norm_stats_path = data_cfg.get("norm_stats_path") or config.get("norm_stats_path") + if norm_stats_path is not None: + if os.path.exists(norm_stats_path): + shutil.copy(norm_stats_path, os.path.join(ckpt_path, "norm_stats.json")) + log_fn(f"[Checkpoint] Copied norm stats to {ckpt_path}/norm_stats.json") + elif (data_cfg.get("dataset_type") or config.get("dataset_type")) == "lerobot": + log_fn(f"[Checkpoint] WARNING: {norm_stats_path} not found, skipping.") + + torch.save( + normalizer_action.state_dict(), + os.path.join(ckpt_path, "normalizer_action.pth"), + ) + torch.save( + normalizer_propri.state_dict(), + os.path.join(ckpt_path, "normalizer_propri.pth"), + ) + + +def _save_fsdp2_full_state_dict( + *, + ckpt_path: str, + model, + optimizer, + lr_scheduler, + is_main: bool, + log_fn: Callable, + frozen_prefixes: Optional[Tuple[str, ...]] = None, +) -> None: + """Save FSDP2 model + optimizer as rank-0 full state dict. + + ``get_state_dict`` with ``full_state_dict=True, cpu_offload=True`` + gathers DTensors to full CPU tensors on rank 0 (other ranks get empty + dicts). This is the cross-world-size-compatible format: on resume, + ``set_state_dict`` with ``broadcast_from_rank0=True`` re-shards from + rank 0's copy to whatever mesh the new run has. + """ + options = StateDictOptions(full_state_dict=True, cpu_offload=True) + model_sd, optim_sd = get_state_dict(model, optimizer, options=options) + + if is_main: + model_sd = _filter_frozen_prefixes(model_sd, frozen_prefixes, log_fn) + model_sd_out = _make_contiguous_and_clone_shared(model_sd) + save_file(model_sd_out, os.path.join(ckpt_path, "model.safetensors")) + torch.save(optim_sd, os.path.join(ckpt_path, "optimizer.pt")) + torch.save(lr_scheduler.state_dict(), os.path.join(ckpt_path, "scheduler.pt")) + log_fn("[Checkpoint] Saved FSDP2 full state dict (rank 0)") + + # Release the consolidated copy immediately; non-rank-0 already held {}. + del model_sd, optim_sd + gc.collect() + + +def _save_dmuon_state_dict( + *, + ckpt_path: str, + model, + optimizer, + lr_scheduler, + is_main: bool, + log_fn: Callable, + frozen_prefixes: Optional[Tuple[str, ...]] = None, +) -> None: + """Save via DMuon's state-dict helpers (full tensors, HF-compatible keys).""" + import dmuon + + model_sd = dmuon.get_model_state_dict(model, cpu_offload=True, rank0_only=True) + if is_main: + model_sd = _filter_frozen_prefixes(model_sd, frozen_prefixes, log_fn) + model_sd = _make_contiguous_and_clone_shared(model_sd) + save_file(model_sd, os.path.join(ckpt_path, "model.safetensors")) + log_fn("[Checkpoint] Saved DMuon model state dict (full, rank0)") + + optim_sd = dmuon.get_optimizer_state_dict( + model, optimizer, cpu_offload=True, rank0_only=True + ) + if is_main: + torch.save(optim_sd, os.path.join(ckpt_path, "optimizer.pt")) + torch.save(lr_scheduler.state_dict(), os.path.join(ckpt_path, "scheduler.pt")) + + del model_sd, optim_sd + gc.collect() + + +def _save_ddp_state_dict( + *, + ckpt_path: str, + model, + optimizer, + lr_scheduler, + is_main: bool, + log_fn: Callable = _noop_log, + frozen_prefixes: Optional[Tuple[str, ...]] = None, +) -> None: + if not is_main: + return + model_state = ( + model.module.state_dict() if isinstance(model, DDP) else model.state_dict() + ) + model_state = _filter_frozen_prefixes(model_state, frozen_prefixes, log_fn) + model_state_contiguous = _make_contiguous_and_clone_shared(model_state) + save_file(model_state_contiguous, os.path.join(ckpt_path, "model.safetensors")) + torch.save(optimizer.state_dict(), os.path.join(ckpt_path, "optimizer.pt")) + torch.save(lr_scheduler.state_dict(), os.path.join(ckpt_path, "scheduler.pt")) + + +def _save_dataset_state( + *, ckpt_path: str, dataset, rank: int, log_fn: Callable +) -> None: + """Save per-rank dataset resume state if the dataset supports it.""" + if not hasattr(dataset, "save_episode_containers"): + return + ec_path = os.path.join(ckpt_path, f"episode_containers_rank_{rank}.pkl") + dataset.save_episode_containers(ec_path) + + _t0 = time.time() + while not (os.path.exists(ec_path) and os.path.getsize(ec_path) > 0): + time.sleep(0.5) + if time.time() - _t0 > 120: + log_fn( + f"WARNING: episode container checkpoint save timeout: {ec_path}", + ) + break + + +def _filter_frozen_prefixes( + state_dict: dict, + frozen_prefixes: Optional[Tuple[str, ...]], + log_fn: Callable, +) -> dict: + """Drop entries whose key starts with any frozen prefix. + + Used to exclude frozen-by-design submodules from checkpoints. Returns + the input unchanged when no prefixes are configured or the dict is + already empty (non-rank-0 case). + """ + if not frozen_prefixes or not state_dict: + return state_dict + original = len(state_dict) + filtered = { + k: v for k, v in state_dict.items() if not k.startswith(frozen_prefixes) + } + log_fn( + f"[Checkpoint] Filtered state dict: {len(filtered)}/{original} entries " + f"(excluded {original - len(filtered)} frozen entries)" + ) + return filtered + + +def _make_contiguous_and_clone_shared(state_dict: dict) -> dict: + """Make tensors contiguous for safetensors, cloning any that share storage.""" + seen_data_ptrs = {} + out = {} + for k, v in state_dict.items(): + ptr = v.data_ptr() + if ptr in seen_data_ptrs: + v = v.clone() + else: + seen_data_ptrs[ptr] = k + out[k] = v.contiguous() if v.is_floating_point() or v.is_complex() else v + return out + + +def _world_size_for_metadata() -> int: + if dist.is_initialized(): + return dist.get_world_size() + return 1 + + +# ---------------------------------------------------------------------- +# Load +# ---------------------------------------------------------------------- + +_EMBED_WEIGHT_KEYS = ( + "model.embed_tokens.weight", + "model.language_model.embed_tokens.weight", +) + + +def _maybe_resize_token_embeddings_for_load( + model: torch.nn.Module, + state_dict: dict, + log_fn: Optional[Callable] = None, +) -> None: + """Resize model embeddings when checkpoint vocab size differs from the model.""" + log_fn = log_fn or _noop_log + for key in _EMBED_WEIGHT_KEYS: + if key not in state_dict: + continue + ckpt_vocab = state_dict[key].shape[0] + embed = ( + model.get_input_embeddings() + if hasattr(model, "get_input_embeddings") + else None + ) + if embed is None: + return + cur_vocab = embed.weight.shape[0] + if cur_vocab == ckpt_vocab: + return + log_fn( + f"resize_token_embeddings from {cur_vocab} to {ckpt_vocab} " + f"to match checkpoint ({key})" + ) + if hasattr(model, "resize_token_embeddings"): + model.resize_token_embeddings(ckpt_vocab) + elif hasattr(model, "model") and hasattr( + model.model, "resize_token_embeddings" + ): + model.model.resize_token_embeddings(ckpt_vocab) + return + + +def load_weights( + *, + model: torch.nn.Module, + resume_config: dict, + model_class=None, + log_fn: Optional[Callable] = None, +): + """Load ``model`` weights from ``resume_config['ckpt']`` (file-level). + + Supports .safetensors / .pth sources and optional ``try_harder`` shape + matching for action-preprocessor weights when action dims changed. + """ + log_fn = log_fn or _noop_log + src = resume_config["ckpt"] + if src.endswith(".safetensors"): + log_fn(f"Loading model from safetensors: {src}") + state_dict = load_file(src) + elif src.endswith(".pth"): + checkpoint = torch.load(src, map_location="cpu", weights_only=False) + state_dict = checkpoint["model_state_dict"] + else: + raise ValueError(f"Unsupported checkpoint format: {src}") + + if model_class is not None and hasattr(model_class, "is_fused"): + if not model_class.is_fused(state_dict): + log_fn("Converting non-fused weights to fused format...") + state_dict = model_class.convert_to_fused(state_dict) + else: + log_fn("The weights is fused, skipping conversion.") + + filtered_state_dict = _drop_checkpoint_normalizer_state(state_dict, log_fn) + _maybe_resize_token_embeddings_for_load(model, filtered_state_dict, log_fn=log_fn) + + if resume_config.get("try_harder", False): + log_fn("### try harder to squeeze checkpoint weights into new model ###") + new_state_dict = reshape_compatible_state_dict( + filtered_state_dict, model.state_dict(), log_fn=log_fn + ) + err = model.load_state_dict(new_state_dict, strict=False) + else: + err = model.load_state_dict(filtered_state_dict, strict=False) + + log_fn(f"err in load model: {err}") + return model + + +def reshape_compatible_state_dict( + state_dict: dict, model_sd: dict, log_fn: Optional[Callable] = None +) -> dict: + """Pad / slice action-preprocessor weights to match target model shape.""" + log_fn = log_fn or _noop_log + out = {} + for name, param in state_dict.items(): + if name not in model_sd: + log_fn(f"Not used parameter: {name}") + continue + if "action_preprocessor" in name or "action_processor" in name: + if param.size() == model_sd[name].size(): + out[name] = param + continue + size_0 = param.size() + size_1 = model_sd[name].size() + if any(old_dim > new_dim for old_dim, new_dim in zip(size_0, size_1)): + raise ValueError( + f"Shape mismatch for '{name}': checkpoint shape {tuple(size_0)} is " + f"larger than model shape {tuple(size_1)} in at least one dimension. " + "Loading a larger checkpoint into a smaller model is not supported here. " + "If the action dimension has changed, please configure action padding " + "in the dataset processor so checkpoint actions match the new action " + "size." + ) + out[name] = model_sd[name].clone() + slices = [ + slice(0, min(old_dim, new_dim)) + for old_dim, new_dim in zip(size_0, size_1) + ] + out[name][slices] = param[slices] + log_fn( + f"Not match key: {name}, checkpoint shape: {tuple(size_0)}, " + f"model shape: {tuple(size_1)}. Filled checkpoint weights into the first " + f"{[s.stop for s in slices]} dims, remaining dims keep model init weights." + ) + else: + if param.size() == model_sd[name].size(): + out[name] = param + else: + log_fn( + f"Skipping '{name}': checkpoint shape {tuple(param.size())} " + f"!= model shape {tuple(model_sd[name].size())}" + ) + return out + + +def resume_from_checkpoint( + *, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + lr_scheduler, + resume_config: dict, + rank: int, + grad_scaler=None, + model_class=None, + log_fn: Optional[Callable] = None, +) -> None: + """Restore model / optimizer / scheduler / RNG / grad_scaler from ckpt.""" + log_fn = log_fn or _noop_log + checkpoint_path = resume_config["ckpt"] + + is_fsdp2 = _is_fsdp2_model(model) + is_dmuon = is_dmuon_model(model) + + # --- world_size diagnostic --------------------------------------- + ckpt_world_size = _read_ckpt_world_size(checkpoint_path, log_fn) + ws_mismatch = ( + ckpt_world_size is not None and ckpt_world_size != _world_size_for_metadata() + ) + if ws_mismatch: + log_fn( + f"Cross-world-size resume: ckpt_world_size={ckpt_world_size}, " + f"current={_world_size_for_metadata()}." + ) + + # --- single-file path (.safetensors / .pth) ---------------------- + if checkpoint_path.endswith(".safetensors") or checkpoint_path.endswith(".pth"): + _load_weights_into_model( + model=model, + resume_config=resume_config, + is_fsdp2=is_fsdp2, + is_dmuon=is_dmuon, + model_class=model_class, + log_fn=log_fn, + ) + log_fn(f"Resumed weights from single-file checkpoint: {checkpoint_path}") + return + + # --- directory path ---------------------------------------------- + if not os.path.isdir(checkpoint_path): + raise FileNotFoundError(f"Checkpoint path not found: {checkpoint_path}") + + safetensors_path = os.path.join(checkpoint_path, "model.safetensors") + rank_shard_path = os.path.join(checkpoint_path, f"model_rank{rank}.pt") + + if os.path.exists(safetensors_path): + if _detect_legacy_fsdp1_format(checkpoint_path): + # Pre-migration FSDP1 ckpt: per-rank optimizer files cannot be + # resharded into the new DTensor layout. Load model only and + # cold-start the optimizer. + _load_legacy_fsdp1_full( + checkpoint_path=checkpoint_path, + model=model, + is_fsdp2=is_fsdp2, + is_dmuon=is_dmuon, + model_class=model_class, + try_harder=resume_config.get("try_harder", False), + log_fn=log_fn, + ) + elif is_fsdp2 or is_dmuon: + # New-format path: model + optimizer via state_dict helpers, + # with automatic reshard on load via broadcast_from_rank0. + _load_fsdp2_or_dmuon_full( + checkpoint_path=checkpoint_path, + model=model, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + is_dmuon=is_dmuon, + ws_mismatch=ws_mismatch, + log_fn=log_fn, + ) + else: + # DDP / unwrapped: single-file model + single-file optimizer. + inner_resume = { + "ckpt": safetensors_path, + "try_harder": resume_config.get("try_harder", False), + } + _load_weights_into_model( + model=model, + resume_config=inner_resume, + is_fsdp2=False, + is_dmuon=False, + model_class=model_class, + log_fn=log_fn, + ) + _resume_ddp_optimizer_scheduler( + checkpoint_path=checkpoint_path, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + log_fn=log_fn, + ) + elif os.path.exists(rank_shard_path): + # FSDP1 sharded checkpoints are rank-layout-bound; dropping support. + raise RuntimeError( + f"Legacy FSDP1 sharded checkpoint at {checkpoint_path} is no " + f"longer supported. Convert to a single model.safetensors first." + ) + else: + raise FileNotFoundError( + f"No model.safetensors or model_rank*.pt found under {checkpoint_path}" + ) + + # --- auxiliary state (all paths) --------------------------------- + _resume_rng(checkpoint_path, log_fn) + if grad_scaler is not None: + _resume_grad_scaler(checkpoint_path, grad_scaler, log_fn) + log_fn(f"Resumed from checkpoint: {checkpoint_path}") + + +def _read_ckpt_world_size(checkpoint_path: str, log_fn: Callable): + if not os.path.isdir(checkpoint_path): + return None + ws_path = os.path.join(checkpoint_path, "world_size.pth") + if not os.path.exists(ws_path): + return None + try: + return int(torch.load(ws_path, map_location="cpu")["world_size"]) + except (KeyError, ValueError, RuntimeError, TypeError) as e: + log_fn( + f"world_size.pth unreadable ({e}); assuming same ws.", + level=logging.WARNING, + ) + return None + + +def _strip_fused_flags(osd: dict) -> None: + """In-place pop of fused / foreach flags from an optimizer state dict's + param_groups. Old checkpoints may have fused=True which breaks dtype + matching when the optimizer is reconstructed under a different precision.""" + for pg in osd.get("param_groups", []) or []: + pg.pop("fused", None) + pg.pop("foreach", None) + + +def _drop_checkpoint_normalizer_state(state_dict: dict, log_fn: Callable) -> dict: + """Drop saved normalizer buffers so current-run stats stay authoritative.""" + normalizer_state_prefixes = ( + "action_preprocessor.normalizer", + "action_processor.normalizer", + ) + filtered = { + k: v + for k, v in state_dict.items() + if not k.startswith(normalizer_state_prefixes) + } + dropped = len(state_dict) - len(filtered) + if dropped: + log_fn( + f"[Checkpoint] Dropped {dropped} checkpoint normalizer entries; " + "keeping current-run normalizers." + ) + return filtered + + +def _load_fsdp2_or_dmuon_full( + *, + checkpoint_path: str, + model, + optimizer, + lr_scheduler, + is_dmuon: bool, + ws_mismatch: bool, + log_fn: Callable, +) -> None: + """Load FSDP2 (or DMuon) model + optimizer from a full-state-dict ckpt. + + FSDP2 path: rank 0 deserializes, ``broadcast_from_rank0=True`` in + ``set_*_state_dict`` reshards to the current mesh. + + DMuon path: every rank MUST load the full state dict from disk. + ``dmuon.set_model_state_dict`` / ``set_optimizer_state_dict`` iterate + ``fqn_to_dp`` and skip any FQN not present in the provided dict - so + if only rank 0 has the data, new owners on other ranks (after a + cross-ws resume) silently miss their assigned params and + ``_owned_data`` stays at fresh-init (random weights), producing a + massive post-resume loss spike. The ckpt files live on CPFS and are + shared, so per-rank reads are essentially free. + """ + is_main = dist.get_rank() == 0 if dist.is_initialized() else True + safetensors_path = os.path.join(checkpoint_path, "model.safetensors") + optim_path = os.path.join(checkpoint_path, "optimizer.pt") + sched_path = os.path.join(checkpoint_path, "scheduler.pt") + + # --- model --- + if is_dmuon: + if not os.path.exists(safetensors_path): + raise FileNotFoundError( + f"DMuon resume requires {safetensors_path} on every rank." + ) + # Every rank loads the full state dict so new owners (after a + # cross-ws reshard, which re-runs dedicate_params and may assign + # FQNs to different ranks) each see their own FQN in the dict. + # dmuon.set_model_state_dict is FQN-keyed and only the new owner + # writes into _owned_data, so passing the full dict on every + # rank is correct (and required) for the cross-ws case. + model_sd = _drop_checkpoint_normalizer_state( + load_file(safetensors_path), log_fn + ) + import dmuon + + dmuon.set_model_state_dict(model, model_sd) + log_fn("[Checkpoint] DMuon model state loaded.") + else: + if is_main and os.path.exists(safetensors_path): + model_sd = _drop_checkpoint_normalizer_state( + load_file(safetensors_path), log_fn + ) + else: + model_sd = {} + options = StateDictOptions( + full_state_dict=True, + cpu_offload=True, + broadcast_from_rank0=True, + strict=False, + ) + set_model_state_dict( + model, + model_state_dict=model_sd, + options=options, + ) + log_fn("[Checkpoint] FSDP2 model state loaded (reshard via broadcast).") + del model_sd + gc.collect() + + # --- optimizer --- + if os.path.exists(optim_path): + if is_dmuon: + # Same rationale as model_sd: every rank needs full optim state. + optim_sd = torch.load(optim_path, map_location="cpu", weights_only=False) + _strip_fused_flags(optim_sd) + elif is_main: + optim_sd = torch.load(optim_path, map_location="cpu", weights_only=False) + _strip_fused_flags(optim_sd) + else: + optim_sd = {} + + try: + if is_dmuon: + import dmuon + + dmuon.set_optimizer_state_dict(model, optimizer, optim_sd) + log_fn("[Checkpoint] DMuon optimizer state loaded.") + else: + options = StateDictOptions( + full_state_dict=True, + cpu_offload=True, + broadcast_from_rank0=True, + ) + set_optimizer_state_dict( + model, + optimizers=optimizer, + optim_state_dict=optim_sd, + options=options, + ) + log_fn( + "[Checkpoint] FSDP2 optimizer state loaded " + "(reshard via broadcast)." + ) + except (ValueError, RuntimeError, TypeError, KeyError) as e: + # Most likely a legacy FSDP1 consolidated OSD (flat_param-specific). + log_fn( + f"[Checkpoint] Optimizer state incompatible with current " + f"layout ({e!r}); cold-starting optimizer. Model weights are " + f"loaded; momentum resets to zero.", + level=logging.WARNING, + ) + del optim_sd + gc.collect() + else: + log_fn( + "[Checkpoint] optimizer.pt not found; cold-starting optimizer.", + level=logging.WARNING, + ) + + # --- scheduler --- + if os.path.exists(sched_path): + try: + lr_scheduler.load_state_dict( + torch.load(sched_path, map_location="cpu", weights_only=False) + ) + log_fn("[Checkpoint] Scheduler state loaded.") + except (ValueError, RuntimeError, TypeError, KeyError) as e: + log_fn( + f"Scheduler load failed ({e}); keeping fresh state.", + level=logging.WARNING, + ) + + +def _resume_ddp_optimizer_scheduler( + *, + checkpoint_path: str, + optimizer, + lr_scheduler, + log_fn: Callable, +) -> None: + """Restore optimizer / scheduler for DDP layouts (single-file).""" + optimizer_path = os.path.join(checkpoint_path, "optimizer.pt") + if os.path.exists(optimizer_path): + try: + optim_sd = torch.load( + optimizer_path, map_location="cpu", weights_only=False + ) + _strip_fused_flags(optim_sd) + optimizer.load_state_dict(optim_sd) + log_fn("[Checkpoint] Optimizer state loaded (single file).") + except (ValueError, RuntimeError) as e: + log_fn( + f"Failed to load optimizer state dict, " + f"optimizer will be re-initialized. Error: {e}", + level=logging.WARNING, + ) + + sched_path = os.path.join(checkpoint_path, "scheduler.pt") + if os.path.exists(sched_path): + try: + lr_scheduler.load_state_dict( + torch.load(sched_path, map_location="cpu", weights_only=False) + ) + log_fn(f"[Checkpoint] Scheduler state loaded ({sched_path}).") + except (ValueError, RuntimeError, KeyError, TypeError) as e: + log_fn( + f"Scheduler load failed ({e}); keeping fresh state.", + level=logging.WARNING, + ) + + +def _load_legacy_fsdp1_full( + *, + checkpoint_path: str, + model, + is_fsdp2: bool, + is_dmuon: bool, + model_class, + try_harder: bool, + log_fn: Callable, +) -> None: + """Load a pre-migration FSDP1 ckpt: weights only, optimizer cold-starts. + + The per-rank ``optimizer_rank{N}.pt`` files use FSDP1's flat_param + layout, which cannot be resharded into FSDP2's DTensor layout. The + only safe action is to load the rank-0 ``model.safetensors`` (which + is layout-agnostic) and let the optimizer warm up from scratch. + """ + inner_resume = { + "ckpt": os.path.join(checkpoint_path, "model.safetensors"), + "try_harder": try_harder, + } + _load_weights_into_model( + model=model, + resume_config=inner_resume, + is_fsdp2=is_fsdp2, + is_dmuon=is_dmuon, + model_class=model_class, + log_fn=log_fn, + ) + log_fn( + f"Legacy FSDP1 ckpt detected at {checkpoint_path}: model weights " + f"loaded, but per-rank optimizer files cannot reshard into FSDP2. " + f"Optimizer state DROPPED - it will warm up from scratch. For long " + f"resume runs prefer a fresh FSDP2 ckpt; for short runs the loss " + f"bump is usually negligible.", + level=logging.WARNING, + ) + + +def _resume_rng(checkpoint_path: str, log_fn: Callable) -> None: + path = os.path.join(checkpoint_path, "rng_state.pt") + if not os.path.exists(path): + log_fn( + "rng_state.pt not found; RNG stays at seed_all() state.", + level=logging.WARNING, + ) + return + try: + sd = torch.load(path, map_location="cpu", weights_only=False) + if "torch" in sd: + torch.set_rng_state(sd["torch"]) + if "cuda" in sd and torch.cuda.is_available(): + torch.cuda.set_rng_state(sd["cuda"]) + if "numpy" in sd: + np.random.set_state(sd["numpy"]) + if "python" in sd: + random.setstate(sd["python"]) + log_fn("[Checkpoint] RNG state restored.") + except (RuntimeError, ValueError, TypeError, KeyError) as e: + log_fn( + f"RNG load failed ({e}); keeping current RNG state.", + level=logging.WARNING, + ) + + +def _resume_grad_scaler(checkpoint_path: str, grad_scaler, log_fn: Callable) -> None: + path = os.path.join(checkpoint_path, "grad_scaler.pt") + if not os.path.exists(path): + log_fn( + "grad_scaler.pt not found; GradScaler stays fresh.", + level=logging.WARNING, + ) + return + try: + grad_scaler.load_state_dict( + torch.load(path, map_location="cpu", weights_only=False) + ) + log_fn("[Checkpoint] GradScaler state restored.") + except (RuntimeError, ValueError, TypeError, KeyError) as e: + log_fn( + f"GradScaler load failed ({e}); keeping fresh state.", + level=logging.WARNING, + ) + + +def _load_weights_into_model( + *, + model: torch.nn.Module, + resume_config: dict, + is_fsdp2: bool = False, + is_dmuon: bool = False, + model_class, + log_fn: Callable, +) -> None: + """Weights-only load for single .safetensors / .pth sources.""" + if is_dmuon or is_fsdp2: + src = resume_config["ckpt"] + is_main = dist.get_rank() == 0 if dist.is_initialized() else True + # DMuon needs the full state dict on every rank (FQN-keyed lookup; + # see _load_fsdp2_or_dmuon_full above). FSDP2 with + # broadcast_from_rank0=True only needs rank 0 to read. + if is_dmuon or is_main: + if src.endswith(".safetensors"): + state_dict = load_file(src) + else: + state_dict = torch.load(src, map_location="cpu", weights_only=False) + if isinstance(state_dict, dict) and "model_state_dict" in state_dict: + state_dict = state_dict["model_state_dict"] + state_dict = _drop_checkpoint_normalizer_state(state_dict, log_fn) + else: + state_dict = {} + if is_dmuon: + import dmuon + + dmuon.set_model_state_dict(model, state_dict) + log_fn(f"[DMuon] Loaded model weights from {src}") + else: + options = StateDictOptions( + full_state_dict=True, + cpu_offload=True, + broadcast_from_rank0=True, + strict=False, + ) + set_model_state_dict( + model, + model_state_dict=state_dict, + options=options, + ) + log_fn(f"[FSDP2] Loaded model weights from {src}") + else: + # DDP / unwrapped. + unwrapped_model = model.module if isinstance(model, DDP) else model + load_weights( + model=unwrapped_model, + resume_config=resume_config, + model_class=model_class, + log_fn=log_fn, + ) + + +def finalize_save(log_fn: Optional[Callable] = None) -> None: + """Barrier + GC hygiene after a checkpoint write. Trainer calls this.""" + if dist.is_initialized(): + dist.barrier() + gc.collect() + torch.cuda.empty_cache() + if log_fn: + pass # trainer already logs the "Saved checkpoint to X" message diff --git a/wall_x/trainer/fsdp_trainer/distribution_strategy.py b/wall_x/trainer/fsdp_trainer/distribution_strategy.py new file mode 100644 index 0000000..65ad4cd --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/distribution_strategy.py @@ -0,0 +1,400 @@ +"""Distribution strategies for model wrapping and gradient coordination.""" + +from __future__ import annotations + +import logging +import os +from abc import ABC, abstractmethod +from contextlib import contextmanager +from importlib import import_module +from typing import NamedTuple, Optional + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh +from torch.distributed.fsdp import ( + MixedPrecisionPolicy, + OffloadPolicy, +) +from torch.nn.parallel import DistributedDataParallel as DDP + +_logger = logging.getLogger(__name__) + + +class FSDP2Layout(NamedTuple): + """Per-call FSDP2 wrap configuration produced by ``FSDPStrategy``.""" + + mesh: DeviceMesh + dp_process_group: Optional[dist.ProcessGroup] + shard_process_group: Optional[dist.ProcessGroup] + replicate_process_group: Optional[dist.ProcessGroup] + trainer_process_group: Optional[dist.ProcessGroup] + mp_policy: Optional[MixedPrecisionPolicy] + offload_policy: Optional[OffloadPolicy] + reshard_after_forward: bool + + +class DistributionStrategy(ABC): + """Common interface for FSDP2 / DDP wrapping + gradient operations.""" + + _cfg: dict + last_grad_clip_stats: Optional[dict] = None + + def _maybe_enable_grad_ckpt(self, model: torch.nn.Module) -> None: + """Call ``gradient_checkpointing_enable`` if the model exposes it.""" + if not self._cfg.get("use_gradient_checkpointing", False): + return + method = getattr(model, "gradient_checkpointing_enable", None) + if method is None: + _logger.warning( + "use_gradient_checkpointing=true but %s has no " + "gradient_checkpointing_enable method. If this model " + "routes recomputation through its config the request may " + "still be honored; otherwise " + "gradient checkpointing is effectively disabled for " + "this run.", + type(model).__name__, + ) + return + method() + + @abstractmethod + def wrap(self, model: torch.nn.Module) -> torch.nn.Module: + """Wrap *model* with the active distribution strategy.""" + + @abstractmethod + def clip_grad_norm( + self, + model: torch.nn.Module, + max_norm: float, + *, + optimizer=None, + ) -> torch.Tensor: + """Clip gradient L2 norm; return total norm before clipping.""" + + @abstractmethod + def no_sync(self, model: torch.nn.Module): + """Context manager that disables cross-rank gradient sync. + + Used by the trainer for gradient accumulation micro-batches. + FSDP2 toggles ``set_requires_gradient_sync``; DDP yields + ``model.no_sync()``. + """ + + @property + @abstractmethod + def params_sharded(self) -> bool: + """Whether parameters are actually sharded across ranks. + + FSDP2 always shards (every ``fully_shard`` unit is sharded across + the mesh). DDP never shards. Used by adapters to decide whether + per-rank norm contributions need an all-reduce. + """ + + +class FSDPStrategy(DistributionStrategy): + def __init__(self, config: dict): + self._cfg = config + self._layout: Optional[FSDP2Layout] = None + self.last_grad_clip_stats = None + + @property + def trainer_process_group(self) -> Optional[dist.ProcessGroup]: + """Process group for trainer-level collectives. + + Returns `dist.group.WORLD` (the implicit all-rank group created + by `init_process_group`) so that trainer barriers / metric + reductions behave consistently regardless of FSDP mesh shape + (1D for full_shard / shard_grad_op, 2D for hybrid_shard / + _hybrid_shard_zero2). No extra NCCL communicator is allocated. + """ + return self._layout.trainer_process_group if self._layout else None + + @staticmethod + def _build_trainer_process_group( + world_size: int, + ) -> Optional[dist.ProcessGroup]: + if not dist.is_initialized() or world_size <= 1: + return None + return dist.group.WORLD + + def _build_fsdp2_layout(self) -> FSDP2Layout: + """Construct the FSDP2 wrap layout from yaml config. + + Builds the device mesh (1D for full_shard / shard_grad_op, 2D + for hybrid_shard / _hybrid_shard_zero2), explicit FSDP mesh process + groups, a trainer-level global process group, the matching + mixed-precision policy, offload policy, and reshard_after_forward flag. + """ + name = self._cfg.get("fsdp_sharding_strategy", "full_shard") + world_size = dist.get_world_size() if dist.is_initialized() else 1 + dp_pg: Optional[dist.ProcessGroup] = None + shard_pg: Optional[dist.ProcessGroup] = None + replicate_pg: Optional[dist.ProcessGroup] = None + + # 2D mesh for HSDP variants; 1D for full_shard / shard_grad_op. + if name in ("hybrid_shard", "_hybrid_shard_zero2"): + replicate_size = self._cfg.get("fsdp_hsdp_replicate_size") or int( + os.environ.get("LOCAL_WORLD_SIZE", "0") + ) + if replicate_size <= 0: + # Fallback: assume one replica per node, equal-sized shard groups. + # Caller should set fsdp_hsdp_replicate_size explicitly when this + # default is wrong (e.g. uneven node sizes). + replicate_size = max(1, world_size // 8) if world_size >= 8 else 1 + if world_size % replicate_size != 0: + raise ValueError( + f"world_size ({world_size}) not divisible by HSDP replicate " + f"size ({replicate_size}); set fsdp_hsdp_replicate_size." + ) + shard_size = world_size // replicate_size + mesh = init_device_mesh( + "cuda", + (replicate_size, shard_size), + mesh_dim_names=("replicate", "shard"), + ) + replicate_pg = mesh.get_group("replicate") + shard_pg = mesh.get_group("shard") + else: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) + dp_pg = mesh.get_group("dp") + + trainer_pg = self._build_trainer_process_group(world_size) + + reshard_after_forward = name not in ( + "shard_grad_op", + "_hybrid_shard_zero2", + ) + + offload_policy: Optional[OffloadPolicy] = ( + OffloadPolicy(pin_memory=True) + if self._cfg.get("fsdp_cpu_offload", False) + else None + ) + + mp_policy: Optional[MixedPrecisionPolicy] + if self._cfg.get("use_mixed_precision", True): + dtype = torch.bfloat16 if self._cfg.get("bf16", True) else torch.float16 + reduce_dtype = ( + torch.float32 + if self._cfg.get("fsdp_reduce_dtype", "bf16") == "fp32" + else dtype + ) + mp_policy = MixedPrecisionPolicy( + param_dtype=dtype, + reduce_dtype=reduce_dtype, + cast_forward_inputs=False, + ) + else: + mp_policy = None + + assert mp_policy is not None, ( + "FSDP2 requires a MixedPrecisionPolicy in this torch version. " + "Set distributed.use_mixed_precision=true." + ) + + return FSDP2Layout( + mesh=mesh, + dp_process_group=dp_pg, + shard_process_group=shard_pg, + replicate_process_group=replicate_pg, + trainer_process_group=trainer_pg, + mp_policy=mp_policy, + offload_policy=offload_policy, + reshard_after_forward=reshard_after_forward, + ) + + def _wrap_fsdp( + self, + model: torch.nn.Module, + *, + use_dmuon: bool = False, + ) -> torch.nn.Module: + if not hasattr(model, "convert_to_fsdp"): + raise NotImplementedError( + f"Model {model.__class__.__name__} has no convert_to_fsdp method" + ) + self._maybe_enable_grad_ckpt(model) + layout = self._build_fsdp2_layout() + self._layout = layout + wrapped = model.convert_to_fsdp( + mesh=layout.mesh, + mp_policy=layout.mp_policy, + offload_policy=layout.offload_policy, + reshard_after_forward=layout.reshard_after_forward, + use_dmuon=use_dmuon, + ) + torch.cuda.empty_cache() + return wrapped + + def wrap(self, model: torch.nn.Module) -> torch.nn.Module: + return self._wrap_fsdp(model, use_dmuon=False) + + def clip_grad_norm( + self, + model: torch.nn.Module, + max_norm: float, + *, + optimizer=None, + ) -> torch.Tensor: + del optimizer + self.last_grad_clip_stats = None + # Separate DTensor (FSDP2) from regular tensor params - PyTorch's + # foreach_mul_ in clip_grad_norm_ can't mix them. + try: + from torch.distributed.tensor import DTensor + except ImportError: + DTensor = None + dtensor_params = [] + regular_params = [] + for p in model.parameters(): + if p.grad is None: + continue + if DTensor is not None and isinstance(p, DTensor): + dtensor_params.append(p) + else: + regular_params.append(p) + + device = next(iter(p for p in model.parameters())).device + total_sq = torch.tensor(0.0, device=device) + if dtensor_params: + n = torch.nn.utils.clip_grad_norm_(dtensor_params, max_norm) + total_sq = total_sq + n.to(device).pow(2) + if regular_params: + n = torch.nn.utils.clip_grad_norm_(regular_params, max_norm) + total_sq = total_sq + n.to(device).pow(2) + return total_sq.sqrt() + + @contextmanager + def no_sync(self, model: torch.nn.Module): + toggled = hasattr(model, "set_requires_gradient_sync") + if toggled: + model.set_requires_gradient_sync(False) + try: + yield + finally: + if toggled: + model.set_requires_gradient_sync(True) + + @property + def params_sharded(self) -> bool: + return True + + +class DMuonFSDPStrategy(FSDPStrategy): + """FSDP2 plus the DMuon dedicated-parameter runtime.""" + + @staticmethod + def _runtime(): + return import_module("dmuon") + + def wrap(self, model: torch.nn.Module) -> torch.nn.Module: + return self._wrap_fsdp(model, use_dmuon=True) + + def clip_grad_norm( + self, + model: torch.nn.Module, + max_norm: float, + *, + optimizer=None, + ) -> torch.Tensor: + total_norm = super().clip_grad_norm(model, max_norm, optimizer=optimizer) + self.last_grad_clip_stats = None + if optimizer is None: + return total_norm + + clip_stats = self._runtime().clip_grad_norm_(optimizer, max_norm) + self.last_grad_clip_stats = clip_stats.as_dict() + device = next(iter(p for p in model.parameters())).device + total_norm_t = ( + total_norm.to(device) + if torch.is_tensor(total_norm) + else torch.tensor(float(total_norm), device=device) + ) + dedicated_norm_t = torch.tensor(float(clip_stats.total_norm), device=device) + return (total_norm_t.pow(2) + dedicated_norm_t.pow(2)).sqrt() + + @contextmanager + def no_sync(self, model: torch.nn.Module): + with self._runtime().no_sync(model): + yield + + +class DDPStrategy(DistributionStrategy): + def __init__( + self, + config: dict, + *, + device: torch.device, + local_rank: int, + ): + self._cfg = config + self._device = device + self._local_rank = local_rank + self.find_unused_parameters = config.get("find_unused_parameters", False) + self.broadcast_buffers = config.get("broadcast_buffers", True) + self.bucket_cap_mb = config.get("bucket_cap_mb", 25) + self.last_grad_clip_stats = None + + def wrap(self, model: torch.nn.Module) -> torch.nn.Module: + # Adapter.convert_to_mix_precision_hint already placed the model on + # self._device; this .to() is a defensive no-op (same as the original + # wrap_ddp_model). + model = model.to(self._device) + self._maybe_enable_grad_ckpt(model) + return DDP( + model, + device_ids=[self._local_rank], + output_device=self._local_rank, + find_unused_parameters=self.find_unused_parameters, + broadcast_buffers=self.broadcast_buffers, + bucket_cap_mb=self.bucket_cap_mb, + ) + + def clip_grad_norm( + self, + model: torch.nn.Module, + max_norm: float, + *, + optimizer=None, + ) -> torch.Tensor: + del optimizer + self.last_grad_clip_stats = None + return torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) + + @contextmanager + def no_sync(self, model: torch.nn.Module): + with model.no_sync(): + yield + + @property + def params_sharded(self) -> bool: + return False + + +def build_strategy( + config: dict, + *, + device: torch.device, + local_rank: int, + use_dmuon: bool = False, +) -> DistributionStrategy: + """Pick FSDP vs DDP based on ``use_fsdp`` flag. + + ``use_gradient_checkpointing`` is read from the config dict by + each strategy's ``_maybe_enable_grad_ckpt``; not a separate kwarg here. + """ + if config.get("use_fsdp", False): + if use_dmuon: + return DMuonFSDPStrategy(config) + return FSDPStrategy(config) + if use_dmuon: + raise NotImplementedError( + "DMuon dedicated-parameter training requires " + "distributed.use_fsdp=true in this trainer." + ) + return DDPStrategy( + config, + device=device, + local_rank=local_rank, + ) diff --git a/wall_x/trainer/fsdp_trainer/fsdp_trainer.py b/wall_x/trainer/fsdp_trainer/fsdp_trainer.py new file mode 100644 index 0000000..19ffd77 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/fsdp_trainer.py @@ -0,0 +1,963 @@ +"""FSDP trainer implementation.""" + +import contextlib +import dataclasses +import gc +import logging +import os +import time +from typing import Any, Optional, Tuple + +import torch +from torch.nn.parallel import DistributedDataParallel as DDP + +from wall_x.model.core.action.normalizer import create_normalizers +from wall_x.trainer.adapters import ADAPTER_REGISTRY, format_adapter_error +from wall_x.trainer.fsdp_trainer import checkpoint_io as _ckpt_io +from wall_x.trainer.fsdp_trainer.base_trainer import ( + DistributedTrainer, + barrier, + is_main_process, +) +from wall_x.trainer.fsdp_trainer.distribution_strategy import build_strategy +from wall_x.trainer.fsdp_trainer.metrics_logger import MetricsLogger +from wall_x.trainer.optimizer import get_optimizer +from wall_x.trainer.scheduler.scheduler import get_scheduler +from wall_x.trainer.utils import move_batch_to_device + + +class FSDPTrainer(DistributedTrainer): + """ + FSDP-based Trainer using pure PyTorch FSDP2 (``fully_shard``). + + Features: + - No accelerate dependency + - Direct torchrun launch + - Mixed precision training support (bfloat16/float16) + - Gradient accumulation + - Full/Sharded state dict saving + - CPU offload support + - Activation checkpointing + + Launch: + torchrun --nproc_per_node=8 --master_port=29500 train_fsdp.py --config config.yaml + """ + + def __init__( + self, + train_config, + wandb_run: Optional[Any] = None, + ): + if train_config is None: + raise ValueError("train_config is required") + + self.model_type = train_config.model_type + super().__init__(train_config, wandb_run) + + if self.model_type not in ADAPTER_REGISTRY: + raise ValueError(format_adapter_error(self.model_type)) + adapter_cls = ADAPTER_REGISTRY[self.model_type] + self.adapter = adapter_cls( + cfg=self.cfg, + logger=self.logger, + model_type=self.model_type, + ) + + self.strategy = build_strategy( + dataclasses.asdict(self.cfg.distributed), + device=self.device, + local_rank=self.local_rank, + use_dmuon=self.cfg.hyperparams.optimizer.optimizer_type == "dmuon", + ) + + # Model config + self.action_dim = self.cfg.action_dim + self.use_selective_recompute = self.cfg.distributed.use_selective_recompute + self.show_time_details = self.cfg.debug.show_time_details + + # num_training_steps: read from any scheduler that exposes it + # (CosineSchedulerConfig currently; future schedulers may add it). + # Used in train_loop to trigger loss_guard_should_stop - works with + # constant scheduler too once the field is set. + sched = self.cfg.hyperparams.scheduler + self.num_training_steps = int(getattr(sched, "num_training_steps", 0) or 0) + self.metrics_logger = MetricsLogger( + wandb_run=self.wandb_run, + log_interval=self.log_interval, + smooth_window=self.cfg.logging.loss_log_smooth_window, + ) + + # Loss tracking + self.base_l1_loss = None + self.base_l1_loss_detail = {} + + self.load_normalizer() + self.load_processor() + self.load_model() + + # DDP needs to see correct requires_grad at wrap time, so freeze first. + self._freeze_params_if_needed(self.model) + + # Frozen-submodule prefixes for checkpoint filtering. Must be computed + # before FSDP wrapping, while child modules still have stable names. + self._frozen_prefixes = self._compute_frozen_prefixes() + + # Original shapes captured before FSDP flattening. + self.model._orig_param_shapes = { + name: p.shape for name, p in self.model.named_parameters() + } + + if self._resume_from_single_file(): + self.load_state_dict( + self.model, + {"ckpt": self.cfg.checkpoint.resume_from}, + ) + + self._wrap_model(self.model) + self._create_optimizer() + self._create_scheduler() + + if self._resume_from_training_checkpoint(): + self.resume_from_checkpoint() + + self.load_dataset() + + self.adapter.init_validation(self.normalizer_action, self.normalizer_propri) + + def load_normalizer(self): + norm_cfg = self.cfg.data.normalizer_config or {} + custom_stats_path = norm_cfg.get("customized_action_statistic_dof") + if not custom_stats_path and self.cfg.data.dataset_type == "lerobot": + from wall_x.data.backends.lerobot.build import load_lerobot_normalizers + + loaded = load_lerobot_normalizers(self.cfg) + if loaded is not None: + self.normalizer_action, self.normalizer_propri = ( + loaded[0], + loaded[1], + ) + self._action_statistic_dof = None + self.logger.info( + "Loaded LeRobot normalizers from %s with dataset key %s", + loaded[2], + loaded[3], + ) + return + + merged = { + "dof_config": self.cfg.task.dof_config, + "agent_pos_config": self.cfg.task.agent_pos_config, + "customized_action_statistic_dof": custom_stats_path, + "min_key": norm_cfg.get("min_key", "min"), + "delta_key": norm_cfg.get("delta_key", "delta"), + } + self.normalizer_action, self.normalizer_propri, self._action_statistic_dof = ( + create_normalizers(merged) + ) + + def backward(self, loss: torch.Tensor): + """Perform backward pass.""" + if self.grad_scaler is not None: + self.grad_scaler.scale(loss).backward() + else: + loss.backward() + + def clip_grad_norm(self, max_norm: float) -> torch.Tensor: + """Unscale and clip gradient norm via the active distribution strategy.""" + if self.grad_scaler is not None: + self.grad_scaler.unscale_(self.optimizer) + + total_norm = self.strategy.clip_grad_norm( + self.model, + max_norm, + optimizer=self.optimizer, + ) + self._dedicated_param_grad_clip_stats = getattr( + self.strategy, + "last_grad_clip_stats", + None, + ) + return total_norm + + def load_model(self): + """Load and prepare model, optimizer, and scheduler""" + self.model_config = self.adapter.build_model_config() + self.model = self.adapter.create_model( + self.processor, self.tokenizer_mixin, self.model_config + ) + type(self.adapter).log_attention_implementation(self.logger, self.model) + self.adapter.load_weights( + self.model, + self.normalizer_action, + self.normalizer_propri, + processor=self.processor, + ) + + def load_processor(self): + """Load processor and tokenizers""" + self.adapter.normalizer_action = self.normalizer_action + self.adapter.normalizer_propri = self.normalizer_propri + processors_dict = self.adapter.load_processor(self._action_statistic_dof) + self.processor = processors_dict["processor"] + self.train_action_tokenizer = processors_dict["train_action_tokenizer"] + self.val_action_tokenizer = processors_dict["val_action_tokenizer"] + self.action_mapper = processors_dict["action_mapper"] + self.tokenizer_mixin = processors_dict.get("tokenizer_mixin") + + def _freeze_params_if_needed(self, model: torch.nn.Module): + """Freeze non-action parameters when train_action_expert_only is set. + + Must be called BEFORE wrapping with DDP/FSDP so the wrapper sees the + correct requires_grad flags and does not expect gradients for frozen params. + """ + from wall_x.trainer.optimizer.utils import resolve_lr_group_configs + + opt = self.cfg.hyperparams.optimizer + if not opt.train_action_expert_only: + return + lr_groups = resolve_lr_group_configs( + opt, self.adapter.default_action_lr_keywords + ) + if not lr_groups: + self.log( + "WARNING: train_action_expert_only is True but no optimizer LR " + "group is set. No parameters will be frozen.", + level=logging.WARNING, + ) + return + + frozen_count = 0 + grouped_count = 0 + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + matches_group = any( + any(keyword in name for keyword in group.include) for group in lr_groups + ) + if matches_group: + grouped_count += 1 + else: + param.requires_grad = False + frozen_count += 1 + + if grouped_count == 0: + raise ValueError( + "No grouped params found for train_action_expert_only. " + "Please check optimizer.lr_groups or legacy action_lr_keywords." + ) + + self.log( + f"*** train_action_expert_only: frozen {frozen_count} base params, " + f"keeping {grouped_count} grouped params trainable ***" + ) + + def _wrap_model(self, model: torch.nn.Module): + """Prepare model dtype/placement via adapter, then wrap via strategy. + + If the strategy created an explicit trainer process group, register + it so trainer-level barriers and metric collectives keep all-rank + semantics independent of the FSDP mesh topology. + """ + self.adapter.convert_to_mix_precision_hint( + model, + device=self.device, + use_fsdp=self.cfg.distributed.use_fsdp, + log_fn=self.log, + ) + self.model = self.strategy.wrap(model) + + trainer_pg = getattr(self.strategy, "trainer_process_group", None) + if trainer_pg is not None: + from wall_x.trainer.fsdp_trainer.base_trainer import ( + set_trainer_process_group, + ) + + set_trainer_process_group(trainer_pg) + self.log("[FSDP2] routing trainer collectives through trainer PG") + + def _create_optimizer(self): + """Create optimizer for FSDP wrapped model.""" + from wall_x.trainer.optimizer.utils import ( + build_lr_param_groups, + resolve_lr_group_configs, + uses_legacy_action_lr_groups, + ) + + opt_cfg = self.cfg.hyperparams.optimizer + + param_groups = None + lr_group_configs = resolve_lr_group_configs( + opt_cfg, self.adapter.default_action_lr_keywords + ) + if lr_group_configs: + if opt_cfg.optimizer_type not in ("adamw", "muon", "dmuon"): + raise ValueError( + "optimizer.lr_groups are only supported with adamw, muon, " + "or dmuon" + ) + base_group_name = ( + "base_lr_group" if uses_legacy_action_lr_groups(opt_cfg) else "base" + ) + param_groups = build_lr_param_groups( + self.model, + opt_cfg, + lr_group_configs, + base_group_name=base_group_name, + ) + summary = ", ".join( + f"{group.name}={group.lr} include={group.include}" + for group in lr_group_configs + ) + self.log( + f"setting optimizer LR groups: base={opt_cfg.learning_rate}; " + f"{summary} ({opt_cfg.optimizer_type})" + ) + + optimizer_kwargs = { + "opt_cfg": opt_cfg, + "param_groups": param_groups, + } + if opt_cfg.optimizer_type == "dmuon": + optimizer_kwargs["log_fn"] = self.log + + self.optimizer = get_optimizer( + opt_cfg.optimizer_type, + self.model, + **optimizer_kwargs, + ) + + def _create_scheduler(self): + """Create learning rate scheduler""" + from wall_x.config.hyperparams_config import ( + ConstantSchedulerConfig, + CosineSchedulerConfig, + ) + + sched = self.cfg.hyperparams.scheduler + lr = self.cfg.hyperparams.optimizer.learning_rate + + if isinstance(sched, ConstantSchedulerConfig): + self.lr_scheduler = torch.optim.lr_scheduler.ConstantLR( + self.optimizer, factor=1.0, total_iters=0 + ) + elif isinstance(sched, CosineSchedulerConfig): + if sched.num_training_steps <= 0: + raise ValueError( + "num_training_steps must be > 0 for cosine scheduler. " + "Please set it explicitly in the config." + ) + min_lr = sched.min_lr if sched.min_lr is not None else 0.1 * lr + self.lr_scheduler = get_scheduler( + optimizer=self.optimizer, + lr_scheduler_type="cosine", + num_warmup_steps=sched.num_warmup_steps, + num_training_steps=sched.num_training_steps, + peak_lr=lr, + end_lr=min_lr, + ) + else: + raise ValueError(f"Unsupported scheduler type: {type(sched).__name__}") + + def load_dataset(self): + """Load training and validation datasets""" + torch.cuda.empty_cache() + barrier() + # Dispatch to adapter with unified signature (supports model-specific kwargs) + extra_dataset_kwargs = dict( + model_config=self.model_config, + tokenizer_mixin=self.tokenizer_mixin, + normalizer_action=self.normalizer_action, + normalizer_propri=self.normalizer_propri, + ) + + resume_path = self.cfg.checkpoint.resume_from + if resume_path is not None and os.path.isdir(resume_path): + # Step and epoch restore are generic. Backend-owned episode + # metadata is optional and only loaded when the adapter exposes + # the hook. + is_incomplete_epoch = False + if hasattr(self.adapter, "_load_episode_indices"): + resume_indices = self.adapter._load_episode_indices( + resume_path, self.rank + ) + if resume_indices is not None: + is_incomplete_epoch = resume_indices["is_incomplete_epoch"] + extra_dataset_kwargs["resume_indices"] = resume_indices + + train_state = self.adapter.load_step_and_epoch( + resume_path, is_incomplete_epoch + ) + self.global_step = train_state["global_step"] + self.start_epoch = train_state["start_epoch"] + self.log( + f"global_step: {self.global_step}, start_epoch: " + f"{self.start_epoch} from {resume_path}" + ) + # global_step records completed steps, so resume from the next step + # to avoid triggering _should_save_checkpoint() immediately. + self.global_step += 1 + ( + self.dataset, + self.train_dataloader, + self.train_num, + ) = self.adapter.load_dataset( + self.data_config, + self.processor, + self.rank, + self.world_size, + **extra_dataset_kwargs, + ) + + def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): + """Execute training for a single epoch""" + self.model.train() + + # VGDynamicRobotDataset (v1 path) doesn't expose set_epoch; other + # dataset backends (v2, lerobot) do. Dispatch conditionally. + if hasattr(self.dataset, "set_epoch"): + self.dataset.set_epoch(epoch) + total = len(self.train_dataloader) + # Drop last + stop_step = total - total % self.grad_accum_steps + + # Disable automatic GC to prevent random ~200ms stalls during compute + gc_interval = self.cfg.logging.gc_interval_steps + gc.disable() + + t0 = time.time() + with profiler: + self.timers("interval-time", log_level=0).start(barrier=True) + self.timers("data-load", log_level=0).start(barrier=True) + + for i, batch in enumerate(self.train_dataloader, self.initial_step): + self.timers("data-load").stop() + + # Save first batch for offline profiling (rank 0 only) + debug_batch_path = self.cfg.debug.save_debug_batch_path + if debug_batch_path and i == self.initial_step and self.rank == 0: + os.makedirs(os.path.dirname(debug_batch_path) or ".", exist_ok=True) + torch.save({"batch": batch}, debug_batch_path) + self.log(f"Saved debug batch to {debug_batch_path}") + + # Move batch to device + batch = move_batch_to_device(batch, self.device) + + if self.cfg.debug.enable_mfu_profile and self.global_step == 0: + self.log( + "[MFU] Forward FLOPs profiling is not available in the public package.", + level=logging.INFO, + ) + + # Forward pass with autocast + self.timers("forward-compute", log_level=0).start(barrier=False) + with self.autocast_context(): + outputs = self.adapter.forward( + self.model, batch, global_step=self.global_step + ) + self.timers("forward-compute").stop() + + loss = self.adapter.extract_loss(outputs) + + # Check for NaN loss -- replace with 0 instead of `continue` + # to avoid skipping collective ops (backward, all_reduce, etc.) + # which would cause NCCL deadlock across ranks. + nan_loss = torch.isnan(loss) + if nan_loss: + self.log( + f"Warning: nan in loss at epoch: {epoch}, step: {i}", + level=logging.WARNING, + ) + loss = torch.zeros_like(loss) + + # Backward pass + self.timers("backward-compute", log_level=0).start(barrier=True) + + context = ( + contextlib.nullcontext() + if self.sync_gradients() + else self.strategy.no_sync(self.model) + ) + with context: + scaled_loss = loss / self.grad_accum_steps + self.backward(scaled_loss) + self.timers("backward-compute").stop() + + # Gradient sync and optimizer step + if self.sync_gradients(): + self.timers("optimizer", log_level=0).start(barrier=True) + + # Per-component grad norms (before clipping) + self.timers("optimizer-grad-norms", log_level=0).start( + barrier=False + ) + self._component_grad_norms = {} + self.adapter.collect_grad_norms( + self.model, + self._component_grad_norms, + device=self.device, + reduce_tensor_fn=self.reduce_tensor, + params_sharded=self.strategy.params_sharded, + ) + self.timers("optimizer-grad-norms").stop() + + # Clip gradients + self.timers("optimizer-clip", log_level=0).start(barrier=False) + if self.cfg.hyperparams.optimizer.enable_grad_clip: + total_norm = self.clip_grad_norm(self.max_grad_norm) + else: + self._dedicated_param_grad_clip_stats = None + total_norm = 0 + self.timers("optimizer-clip").stop() + + self.timers("optimizer-step", log_level=0).start(barrier=False) + self.optimizer_step() + self.timers("optimizer-step").stop() + + self.timers("optimizer-zero-grad", log_level=0).start(barrier=False) + self.optimizer_zero_grad() + self.timers("optimizer-zero-grad").stop() + self.timers("optimizer").stop() + + # Scheduler step + self.lr_scheduler_step() + + # Logging + self.timers("logging", log_level=0).start(barrier=True) + _t_metrics_start = time.time() + self._log_training_metrics( + epoch, i, total, loss, total_norm, outputs + ) + _t_metrics_ms = (time.time() - _t_metrics_start) * 1000 + + t1 = time.time() + self.training_log( + epoch, + self.num_epoch, + i, + total, + loss, + self.get_lr(), + t1 - t0, + self.show_time_details, + ) + self.timers("logging").stop() + t0 = time.time() + + # Optional: log per-step breakdown to pinpoint spikes (e.g. param_norms every 100 steps) + if ( + self.show_time_details + and _t_metrics_ms > 5000 + and is_main_process() + ): + self.log( + f"[Step time breakdown] _log_training_metrics took {_t_metrics_ms:.0f} ms at global_step={self.global_step}", + level=logging.INFO, + ) + + # Checkpoint saving (FSDP full state_dict can take 10-20s every save_interval steps) + if self._should_save_checkpoint(): + _t_save_start = time.time() + self.save_checkpoint(epoch, self.global_step) + if is_main_process(): + self.log( + f"[Step time breakdown] save_checkpoint took {(time.time() - _t_save_start):.1f} s at global_step={self.global_step}", + level=logging.INFO, + ) + + # Validation + if self._should_validate(): + self.val_loop() + + self.global_step += 1 + self.micro_step = 0 + + # Manual GC outside timing window to avoid random stalls (can add 1-5s every gc_interval_steps) + if self.global_step % gc_interval == 0: + _t_gc_start = time.time() + gc.collect() + if is_main_process() and self.show_time_details: + self.log( + f"[Step time breakdown] gc.collect took {(time.time() - _t_gc_start):.1f} s at global_step={self.global_step}", + level=logging.INFO, + ) + else: + self.micro_step += 1 + + del batch + self.timers("interval-time").stop() + + # Drop last + if i == stop_step: + break + + if not isinstance(profiler, contextlib.nullcontext): + profiler.step() + + if ( + self.num_training_steps > 0 + and self.global_step >= self.num_training_steps + ): + break + + # Setup timers for next iteration + if i < total - 1: + self.timers("interval-time", log_level=0).start(barrier=True) + self.timers("data-load", log_level=0).start(barrier=True) + + # Re-enable automatic GC after training loop + gc.enable() + + # Reset dataloader for next epoch + self.train_dataloader = self.dataset.get_train_dataloader() + + def _log_training_metrics(self, epoch, step, total, loss, total_norm, outputs): + """Collect per-step stats, delegate buffering + wandb emission.""" + lr = self.get_lr() + train_loss = self.reduce_tensor(loss.detach()).item() + + step_stats = { + "lr": lr, + "train_loss": train_loss, + "grad_norm": ( + total_norm.item() if torch.is_tensor(total_norm) else float(total_norm) + ), + } + for idx, group in enumerate(self.optimizer.param_groups): + group_name = group.get("group_name", f"group_{idx}") + step_stats[f"lr_group/{group_name}"] = float(group["lr"]) + + # Model-family-specific auxiliary losses / accuracies. + self.adapter.collect_output_stats( + outputs, + step_stats, + reduce_tensor_fn=self.reduce_tensor, + true_gather_fn=self.true_gather, + tokenizer_mixin=self.tokenizer_mixin, + ) + + # Per-component grad norms (captured before clip & optimizer.zero_grad). + if hasattr(self, "_component_grad_norms"): + step_stats.update(self._component_grad_norms) + + dedicated_clip_stats = getattr(self, "_dedicated_param_grad_clip_stats", None) + if dedicated_clip_stats is not None: + step_stats.update( + { + "muon_grad_norm": dedicated_clip_stats["total_norm"], + "muon_grad_clip_coef": dedicated_clip_stats["clip_coef"], + "muon_grad_clipped": float(dedicated_clip_stats["clipped"]), + } + ) + + if self.global_step % 100 == 0: + self.adapter.collect_param_norms( + self.model, + step_stats, + device=self.device, + reduce_tensor_fn=self.reduce_tensor, + params_sharded=self.strategy.params_sharded, + ) + + self._current_step_raw_stats = step_stats + + # Display-smoothing rolling window (DZ-style) + self._current_step_stats = self.metrics_logger.smooth(step_stats) + + def training_log( + self, + current_epoch, + total_epoch, + current_train_iter, + total_train_iter, + loss, + lr, + time_per_step, + show_time_details=False, + ): + # timers.log() contains all_gather - must run on ALL ranks before the + # is_main_process() guard to avoid NCCL deadlock. + if show_time_details: + self.timers.log( + [ + "interval-time", + "data-load", + "forward-compute", + "backward-compute", + "optimizer", + "optimizer-grad-norms", + "optimizer-clip", + "optimizer-step", + "optimizer-zero-grad", + "logging", + ], + normalizer=1, + ) + + main = is_main_process() + stats = getattr(self, "_current_step_stats", None) or {} + + # MFU is computed after step time is known, then merged into both the + # console stats and the MetricsLogger buffer so wandb records it. + mfu_info = None + if self.cfg.debug.enable_mfu and hasattr(self.adapter, "compute_mfu"): + unwrapped = ( + self.model.module if hasattr(self.model, "module") else self.model + ) + mfu_info = self.adapter.compute_mfu(unwrapped, time_per_step) + if mfu_info is not None: + self.log( + "[MFU-rank] rank={} mfu={:.2f}% step={:.3f}s " + "flops_step={:.3f}T flops_fwd={:.3f}T " + "seq={} latent={} source={}".format( + getattr(self, "rank", 0), + mfu_info["mfu"] * 100.0, + time_per_step, + mfu_info["flops_per_step_T"], + mfu_info.get("flops_forward_T", 0.0) or 0.0, + mfu_info.get("seq_dims"), + mfu_info.get("latent_dims"), + mfu_info.get("latent_source"), + ), + main_process_only=False, + ) + if main and mfu_info is not None and stats is not None: + mfu_stats = { + "mfu": mfu_info["mfu"], + "flops_per_step_T": mfu_info["flops_per_step_T"], + "flops_forward_T": mfu_info.get("flops_forward_T"), + } + profile_fwd_flops = mfu_info.get("profile_fwd_flops_T") + if profile_fwd_flops is not None: + mfu_stats["profile_fwd_flops_T"] = profile_fwd_flops + stats.update(mfu_stats) + raw_stats = getattr(self, "_current_step_raw_stats", None) + if raw_stats is not None: + raw_stats.update(mfu_stats) + + if main: + raw_stats = getattr(self, "_current_step_raw_stats", None) + if raw_stats is not None: + self.metrics_logger.record_step(raw_stats, is_main=True) + avg_stats = self.metrics_logger.flush_if_due(self.global_step) + if avg_stats is not None: + self.logger.info( + f"[FSDP Train] Step {self.global_step}: {avg_stats}" + ) + + if not main: + return + + loss_to_print = stats.get( + "train_loss", loss.item() if torch.is_tensor(loss) else float(loss) + ) + fields = self.adapter.console_fields(tokenizer_mixin=self.tokenizer_mixin) + + line = self.metrics_logger.format_training_line( + epoch=current_epoch, + total_epoch=total_epoch, + current_iter=current_train_iter, + total_iter=total_train_iter, + loss=loss_to_print, + lr=lr, + time_per_step=time_per_step, + stats=stats, + fields=fields, + mfu_info=mfu_info, + ) + self.log(line) + + def _should_save_checkpoint(self) -> bool: + """Check if checkpoint should be saved""" + log = self.cfg.logging + return ( + self.global_step >= log.ignore_until_interval + and self.global_step % log.save_interval == 0 + and self.global_step != 0 + ) + + def _should_validate(self) -> bool: + """Check if validation should be performed""" + log = self.cfg.logging + return ( + self.global_step >= log.ignore_until_interval + and self.global_step % log.val_interval == 0 + and self.global_step > 0 + ) + + @torch.no_grad() + def val_loop(self): + """Delegate the full validation run to the adapter. + + Skips silently when the dataset has no val split - adapter's + run_validation iterates ``val_dataloader`` with ``tqdm(..., + total=len(...))`` and cannot accept None. v2 returns None here + when the YAML only declares a train split. + """ + self.val_dataloader = self.dataset.get_val_dataloader() + if self.val_dataloader is None: + if is_main_process(): + self.logger.info("No val split configured, skipping validation.") + return + save_path = self.cfg.checkpoint.save_path + self.adapter.run_validation( + model=self.model, + val_dataloader=self.val_dataloader, + rank=self.rank, + world_size=self.world_size, + device=self.device, + autocast_context=self.autocast_context, + reduce_fn=self.reduce_tensor, + gather_fn=self.true_gather, + logger=self.wandb_run if is_main_process() else None, + global_step=self.global_step, + output_path=os.path.join(save_path, f"val_rank_{self.rank}"), + tokenizer_mixin=self.tokenizer_mixin, + ) + barrier() + + def _compute_frozen_prefixes(self) -> Optional[Tuple[str, ...]]: + """Return state-dict key prefixes for submodules that are fully frozen. + + For models with a ``pipe`` container, any child not present in + ``cfg.model.trainable_models`` is treated as frozen. The matching key + prefixes let checkpoint_io drop those entries from the saved file. + Empty / missing ``trainable_models`` disables filtering and preserves + "save everything" behavior. + """ + if not hasattr(self.model, "pipe"): + return None + trainable_models = getattr(self.cfg.model, "trainable_models", None) + if trainable_models is None: + return None + if isinstance(trainable_models, str): + trainable_set = { + s.strip() for s in trainable_models.split(",") if s.strip() + } + else: + trainable_set = set(trainable_models) + if not trainable_set: + self.log( + "WARNING: trainable_models is empty; not filtering frozen " + "entries from checkpoint to avoid saving an empty file." + ) + return None + frozen_prefixes = tuple( + f"pipe.{name}." + for name, _ in self.model.pipe.named_children() + if name not in trainable_set + ) + if frozen_prefixes: + self.log( + f"Frozen submodule prefixes excluded from checkpoints: " + f"{list(frozen_prefixes)}" + ) + return frozen_prefixes + return None + + def save_checkpoint(self, epoch: int, step: int = 0): + """Save model checkpoint via checkpoint_io (dispatches on model wrapper type).""" + save_path = self.cfg.checkpoint.save_path + ckpt_path = f"{save_path}/{epoch}_{step}" if step else f"{save_path}/{epoch}" + _ckpt_io.save_checkpoint( + ckpt_path=ckpt_path, + model=self.model, + optimizer=self.optimizer, + lr_scheduler=self.lr_scheduler, + config=dataclasses.asdict(self.cfg), + rank=self.rank, + is_main=is_main_process(), + epoch=epoch, + global_step=self.global_step, + seed=self.seed, + normalizer_action=self.normalizer_action, + normalizer_propri=self.normalizer_propri, + dataset=getattr(self, "dataset", None) if step != 0 else None, + grad_scaler=self.grad_scaler, + log_fn=self.log, + frozen_prefixes=self._frozen_prefixes, + ) + _ckpt_io.finalize_save() + self.log(f"Saved checkpoint to {ckpt_path}") + + def load_state_dict(self, model, resume_config): + """Load state dict with fused-weight conversion + try_harder support.""" + return _ckpt_io.load_weights( + model=model, + resume_config=resume_config, + model_class=self._checkpoint_model_class(), + log_fn=self.log, + ) + + def _checkpoint_model_class(self): + model_class = getattr(self, "ModelClass", None) + if model_class is None and hasattr(type(self.adapter), "model_class"): + model_class = type(self.adapter).model_class() + return model_class + + def _resume_from_single_file(self) -> bool: + path = self.cfg.checkpoint.resume_from + return bool(path) and str(path).endswith((".safetensors", ".pth")) + + def _resume_from_training_checkpoint(self) -> bool: + path = self.cfg.checkpoint.resume_from + return bool(path) and not self._resume_from_single_file() + + def resume_from_checkpoint(self): + """Resume training from checkpoint via checkpoint_io.""" + _ckpt_io.resume_from_checkpoint( + model=self.model, + optimizer=self.optimizer, + lr_scheduler=self.lr_scheduler, + resume_config={"ckpt": self.cfg.checkpoint.resume_from}, + rank=self.rank, + grad_scaler=self.grad_scaler, + model_class=self._checkpoint_model_class(), + log_fn=self.log, + ) + barrier() + + def predict_action_loop( + self, + current_step=0, + max_iteration=None, + prediction_type="flow_action", + mode="acc", + ): + """Delegate action prediction to the adapter.""" + del mode # unused in current dispatch; kept for signature compatibility + return self.adapter.predict( + prediction_type, + model=self._unwrapped_model_for_inference(), + val_dataloader=self.dataset.get_val_dataloader(), + rank=self.rank, + world_size=self.world_size, + device=self.device, + processor=self.processor, + tokenizer_mixin=self.tokenizer_mixin, + logger=self.wandb_run, + current_step=current_step, + max_iteration=max_iteration, + ) + + def predict_text_loop(self, current_step=0, max_samples=None, save_dir="./"): + """Delegate text prediction to the adapter.""" + return self.adapter.predict( + "text", + model=self._unwrapped_model_for_inference(), + val_dataloader=self.dataset.get_val_dataloader(), + rank=self.rank, + world_size=self.world_size, + device=self.device, + processor=self.processor, + tokenizer_mixin=self.tokenizer_mixin, + logger=self.wandb_run, + current_step=current_step, + max_samples=max_samples, + save_dir=save_dir, + ) + + def _unwrapped_model_for_inference(self): + """Return the underlying module for inference: unwrap DDP, pass FSDP through.""" + if isinstance(self.model, DDP): + return self.model.module + return self.model diff --git a/wall_x/trainer/fsdp_trainer/metrics/__init__.py b/wall_x/trainer/fsdp_trainer/metrics/__init__.py new file mode 100644 index 0000000..3123573 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/metrics/__init__.py @@ -0,0 +1,3 @@ +from wall_x.trainer.fsdp_trainer.metrics.norms import DistributedNormAccumulator + +__all__ = ["DistributedNormAccumulator"] diff --git a/wall_x/trainer/fsdp_trainer/metrics/norms.py b/wall_x/trainer/fsdp_trainer/metrics/norms.py new file mode 100644 index 0000000..f518534 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/metrics/norms.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +from collections import defaultdict +from typing import Callable, DefaultDict, Optional + +import torch +from torch.distributed.tensor import DTensor, Replicate + + +class DistributedNormAccumulator: + """Accumulate L2 norms without leaking parallel-layout branches to adapters.""" + + def __init__( + self, + *, + device: torch.device, + reduce_tensor_fn: Optional[Callable[[torch.Tensor, bool], torch.Tensor]], + ): + self.device = device + self.reduce_tensor_fn = reduce_tensor_fn + self._regular: DefaultDict[str, list[torch.Tensor]] = defaultdict(list) + self._dtensor: DefaultDict[str, list[torch.Tensor]] = defaultdict(list) + self._owner_only: DefaultDict[str, list[torch.Tensor]] = defaultdict(list) + + @staticmethod + def _is_dtensor(tensor: torch.Tensor) -> bool: + return isinstance(tensor, DTensor) + + def add(self, group: Optional[str], tensor: Optional[torch.Tensor]) -> None: + if group is None or tensor is None or tensor.numel() == 0: + return + tensor = tensor.detach() + if self._is_dtensor(tensor): + self._dtensor[group].append(tensor) + else: + self._regular[group].append(tensor) + + def add_owner_only( + self, group: Optional[str], tensor: Optional[torch.Tensor] + ) -> None: + if group is None or tensor is None or tensor.numel() == 0: + return + self._owner_only[group].append(tensor.detach()) + + def squared_norm(self, group: str) -> torch.Tensor: + with torch.no_grad(): + total = torch.zeros((), dtype=torch.float32, device=self.device) + total = total + self._regular_squared_norm(self._regular[group]) + total = total + self._owner_only_squared_norm(self._owner_only[group]) + total = total + self._dtensor_squared_norm(self._dtensor[group]) + return total + + def norm(self, group: str) -> torch.Tensor: + return self.squared_norm(group).sqrt() + + def _regular_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor: + if not tensors: + return torch.zeros((), dtype=torch.float32, device=self.device) + per = self._foreach_norm(tensors) + return self._sum_norms_squared(per) + + def _owner_only_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor: + local_sq = self._regular_squared_norm(tensors) + if self.reduce_tensor_fn is None: + return local_sq + return self.reduce_tensor_fn(local_sq, average=False) + + def _dtensor_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor: + if not tensors: + return torch.zeros((), dtype=torch.float32, device=self.device) + + per = self._foreach_norm(tensors) + + # Group by (mesh, placements) so each unique layout costs only one + # collective (redistribute). + groups: DefaultDict[object, list[torch.Tensor]] = defaultdict(list) + for norm in per: + if self._is_dtensor(norm): + key = (id(norm.device_mesh), tuple(norm.placements)) + else: + key = None + groups[key].append(norm) + + parts: list[torch.Tensor] = [] + for group_tensors in groups.values(): + try: + stacked = torch.stack(group_tensors) + local = self._to_local_replicated(stacked).to( + device=self.device, dtype=torch.float32 + ) + parts.append(local.pow(2).sum()) + except (RuntimeError, NotImplementedError): + # Stack/redistribute not supported for this layout in the + # current PyTorch version; fall back to per-tensor for safety. + for norm in group_tensors: + local = self._to_local_replicated(norm).to( + device=self.device, dtype=torch.float32 + ) + parts.append(local.pow(2).sum()) + + return torch.stack(parts).sum() + + def _foreach_norm(self, tensors: list[torch.Tensor]) -> list[torch.Tensor]: + try: + return torch._foreach_norm(tensors, 2.0) + except (RuntimeError, NotImplementedError): + return [torch.linalg.vector_norm(t, 2.0) for t in tensors] + + def _sum_norms_squared(self, norms: list[torch.Tensor]) -> torch.Tensor: + if not norms: + return torch.zeros((), dtype=torch.float32, device=self.device) + local_norms = [ + self._to_local_replicated(norm).to(device=self.device, dtype=torch.float32) + for norm in norms + ] + return torch.stack(local_norms).pow(2).sum() + + def _to_local_replicated(self, tensor: torch.Tensor) -> torch.Tensor: + if not self._is_dtensor(tensor): + return tensor.to(device=self.device) + mesh = tensor.device_mesh + tensor = tensor.redistribute(placements=[Replicate()] * mesh.ndim) + return tensor.to_local().to(device=self.device) diff --git a/wall_x/trainer/fsdp_trainer/metrics_logger.py b/wall_x/trainer/fsdp_trainer/metrics_logger.py new file mode 100644 index 0000000..ee00efe --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/metrics_logger.py @@ -0,0 +1,111 @@ +"""Training metric buffering and wandb emission helpers.""" + +from __future__ import annotations + +from collections import deque +from typing import Any, Deque, Dict, List, Optional, Tuple + + +class MetricsLogger: + """Buffer per-step stats and emit averaged batches to wandb. + + Also formats the per-step console line via ``format_training_line``, + which consumes an ``(stat_key, pretty_label, fmt_spec)`` triple list + from ``adapter.console_fields()``. + + ``smooth_window`` controls a separate rolling-mean buffer used purely + for display smoothing (console + tqdm). It is independent of the + ``log_interval`` wandb-flush buffer; the smoothed dict is returned by + :meth:`smooth` for the caller to assign to ``_current_step_stats``. + 1 (default) preserves historical per-step values. + """ + + def __init__(self, *, wandb_run, log_interval: int, smooth_window: int = 1): + self._wandb_run = wandb_run + self._log_interval = log_interval + self._buffer: List[Dict[str, Any]] = [] + self._smooth_window = max(1, int(smooth_window)) + self._smooth_buffers: Dict[str, Deque[float]] = {} + + def record_step(self, step_stats: Dict[str, Any], *, is_main: bool) -> None: + """Push a per-step stats dict into the rank-0 buffer (no-op off rank-0).""" + if is_main: + self._buffer.append(step_stats) + + def flush_if_due(self, global_step: int) -> Optional[Dict[str, Any]]: + """At log_interval boundaries, average + emit; return avg dict or None.""" + if global_step % self._log_interval != 0 or not self._buffer: + return None + avg = self._average(self._buffer) + if self._wandb_run is not None and hasattr(self._wandb_run, "log"): + self._wandb_run.log(avg, step=global_step) + self._buffer = [] + return avg + + def smooth(self, stats: Dict[str, Any]) -> Dict[str, Any]: + """Return *stats* with numeric entries replaced by their rolling mean. + + Uses one ``deque(maxlen=smooth_window)`` per metric key. Non-numeric + entries pass through unchanged. When ``smooth_window <= 1`` this is + the identity (historical behavior). + """ + if self._smooth_window <= 1: + return stats + out: Dict[str, Any] = {} + for key, val in stats.items(): + if not isinstance(val, (int, float)): + out[key] = val + continue + buf = self._smooth_buffers.get(key) + if buf is None or buf.maxlen != self._smooth_window: + buf = deque(maxlen=self._smooth_window) + self._smooth_buffers[key] = buf + buf.append(float(val)) + out[key] = sum(buf) / len(buf) + return out + + @staticmethod + def _average(buf: List[Dict[str, Any]]) -> Dict[str, Any]: + all_keys: set = set() + for stats in buf: + all_keys.update(stats.keys()) + avg: Dict[str, Any] = {} + for key in all_keys: + values = [s[key] for s in buf if key in s] + if values: + avg[key] = sum(values) / len(values) + return avg + + def format_training_line( + self, + *, + epoch: int, + total_epoch: int, + current_iter: int, + total_iter: int, + loss: float, + lr: float, + time_per_step: float, + stats: Dict[str, Any], + fields: List[Tuple[str, str, str]], + mfu_info: Optional[Dict[str, Any]] = None, + ) -> str: + """Render a line matching pre-refactor training_log output exactly. + + Layout: " epoch E/T | iter I/T | loss L | [adapter fields] | lr L | + time_current_backward_step Ts | [MFU P% |]" + """ + parts = [ + " epoch {:3d}/{:3d} |".format(epoch, total_epoch), + " iter {:6d}/{:6d} |".format(current_iter, total_iter), + " loss {:.6f} |".format(loss), + ] + for key, label, fmt in fields: + value = stats.get(key) + if value is not None: + parts.append(" {} {:{fmt}} |".format(label, value, fmt=fmt)) + parts.append(" lr {:.6f} |".format(lr)) + parts.append(" time_current_backward_step {:.6f}s |".format(time_per_step)) + if mfu_info is not None: + parts.append(" MFU {:.2f}% |".format(mfu_info["mfu"] * 100)) + return "".join(parts) diff --git a/wall_x/trainer/fsdp_trainer/train_fsdp.py b/wall_x/trainer/fsdp_trainer/train_fsdp.py new file mode 100644 index 0000000..e30fdfb --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/train_fsdp.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python3 +""" +FSDP Training Entry Point +Launch with torchrun: + torchrun --nproc_per_node=8 --master_port=29500 train_fsdp.py --config config.yaml +""" +import argparse +import dataclasses +import logging +import os +import sys +from datetime import datetime + +import torch +import wandb + +from wall_x.config.loader import load_config as load_typed_config +from wall_x.trainer.fsdp_trainer import ( + FSDPTrainer, + cleanup_distributed, + is_main_process, +) + +logger = logging.getLogger(__name__) + + +class TeeOutput: + """Write output to both the terminal and a file.""" + + def __init__(self, file_path, mode="w"): + self.terminal = sys.stdout + self.log = open(file_path, mode) + + def write(self, message): + self.terminal.write(message) + self.log.write(message) + self.log.flush() + + def flush(self): + self.terminal.flush() + self.log.flush() + + +def parse_args(): + parser = argparse.ArgumentParser(description="FSDP Training Script") + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to training config YAML file", + ) + # FSDP specific overrides + parser.add_argument( + "--fsdp_sharding_strategy", + type=str, + default=None, + choices=["full_shard", "shard_grad_op", "no_shard", "hybrid_shard"], + help="FSDP sharding strategy (overrides config)", + ) + parser.add_argument( + "--debug", + action="store_true", + default=False, + help="Enable debug mode: reduces buffer size, sets log/save path to debug.", + ) + parser.add_argument( + "--wandb_offline", + type=str, + default=None, + help="Whether to run wandb in offline mode (overrides config).", + ) + parser.add_argument( + "--visualize", + action="store_true", + default=False, + help="Whether to visualize samples during training.", + ) + parser.add_argument( + "--log_to_file", + action="store_true", + default=False, + help="Whether to redirect stdout and stderr to log files.", + ) + return parser.parse_args() + + +def load_config(config_path: str, cli_args=None): + """Load configuration from YAML file into TrainConfig.""" + return load_typed_config(config_path, cli_args=cli_args) + + +def setup_logger(cfg): + """Setup wandb logger if enabled""" + log = cfg.logging + if log.use_wandb and is_main_process(): + logger.info( + "rank %s is initializing wandb , is main process %s", + torch.distributed.get_rank(), + is_main_process(), + ) + wandb_run = wandb.init( + project=log.log_project, + name=log.log_name, + entity=log.log_entity, + config=dataclasses.asdict(cfg), + save_code=False, + force=False, + mode="offline" if log.wandb_offline else "online", + ) + logger.info("Wandb Initialized") + return wandb_run + return None + + +def print_fsdp_config(cfg): + """Print FSDP configuration""" + if is_main_process(): + dist = cfg.distributed + logger.info("%s", "=" * 60) + logger.info("FSDP Configuration:") + logger.info(" use_fsdp: %s", dist.use_fsdp) + logger.info("%s", "=" * 60) + + +def main(): + args = parse_args() + + # Redirect logs to files only on the main process to avoid write races. + if args.log_to_file: + import yaml as _yaml + + with open(args.config, "r") as f: + _tmp = _yaml.safe_load(f) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + # support both old (save_path) and new (checkpoint.save_path) schema + log_dir = _tmp.get("save_path") or _tmp.get("checkpoint", {}).get( + "save_path", "./ckpt" + ) + if args.debug: + log_dir = "./ckpt/debug" + os.makedirs(log_dir, exist_ok=True) + log_file = os.path.join(log_dir, f"training_log_{timestamp}.log") + sys.stdout = TeeOutput(log_file, mode="w") + sys.stderr = TeeOutput(log_file.replace(".log", "_stderr.log"), mode="w") + logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) + logger.info("\n%s", "=" * 80) + logger.info("LOG TO FILE MODE: All output will be saved to:") + logger.info(" STDOUT: %s", log_file) + logger.info(" STDERR: %s", log_file.replace(".log", "_stderr.log")) + logger.info("%s\n", "=" * 80) + else: + logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) + + torch.cuda.init() + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + torch.distributed.init_process_group("nccl", device_id=device) + cfg = load_config(args.config, cli_args=args) + wandb_run = setup_logger(cfg) + print_fsdp_config(cfg) + + try: + trainer = FSDPTrainer( + train_config=cfg, + wandb_run=wandb_run, + ) + trainer.fit() + + except Exception as e: + logger.exception("Training failed with error: %s", e) + raise + + finally: + cleanup_distributed() + if wandb_run is not None: + wandb_run.finish() + + +if __name__ == "__main__": + main() diff --git a/wall_x/trainer/optimizer/__init__.py b/wall_x/trainer/optimizer/__init__.py new file mode 100644 index 0000000..596a6aa --- /dev/null +++ b/wall_x/trainer/optimizer/__init__.py @@ -0,0 +1,9 @@ +from .utils import get_optimizer, register_optimizer + +# DMuon is optional (external package). Skip registration if dmuon isn't +# installed so wall_x still imports; the registry will simply not have +# "dmuon" and get_optimizer("dmuon", ...) will raise a clear KeyError. +try: + from . import dmuon # noqa: F401 - side-effect registers "dmuon" +except ImportError: + pass diff --git a/wall_x/trainer/optimizer/dmuon/__init__.py b/wall_x/trainer/optimizer/dmuon/__init__.py new file mode 100644 index 0000000..b3d637c --- /dev/null +++ b/wall_x/trainer/optimizer/dmuon/__init__.py @@ -0,0 +1,16 @@ +import torch.nn as nn + +from .utils import get_dmuon_optimizer + + +def is_dmuon_model(model: nn.Module) -> bool: + """True if ``dmuon.dedicate_params()`` has been applied to this model. + + Checked via an attribute the external ``dmuon`` package attaches to + the root module, so this predicate works without importing ``dmuon`` + and returns ``False`` for ordinary (non-DMuon) models. + """ + return hasattr(model, "_dedicated_comm_ctx") + + +__all__ = ["get_dmuon_optimizer", "is_dmuon_model"] diff --git a/wall_x/trainer/optimizer/dmuon/utils.py b/wall_x/trainer/optimizer/dmuon/utils.py new file mode 100644 index 0000000..3d03b82 --- /dev/null +++ b/wall_x/trainer/optimizer/dmuon/utils.py @@ -0,0 +1,156 @@ +"""DMuon optimizer builder.""" + +import inspect +import logging + +from ..utils import register_optimizer + +_logger = logging.getLogger(__name__) + + +def _emit(log_fn, message, *args, level=logging.INFO): + if args: + message = message % args + if log_fn is not None: + log_fn(message, level=level) + else: + _logger.log(level, message) + + +def _is_rank0(): + try: + import torch.distributed as dist + + return ( + not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 + ) + except Exception: + return True + + +def _build_ns_backend(dmuon, opt_cfg): + coefficients = getattr(opt_cfg, "ns_coefficients", "default") + if coefficients in (None, "default"): + return opt_cfg.ns_backend + + if coefficients != "wallx_muon": + raise ValueError( + "Unsupported DMuon ns_coefficients=" + f"{coefficients!r}. Supported: 'default', 'wallx_muon'." + ) + if opt_cfg.ns_backend != "direct": + raise ValueError( + "ns_coefficients='wallx_muon' is intended to match Wall-X's " + "direct-space Muon implementation. Set ns_backend='direct'." + ) + + wallx_coefficients = [[3.4445, -4.7750, 2.0315] for _ in range(opt_cfg.ns_steps)] + return dmuon.NewtonSchulz( + backend="direct", + coefficients=wallx_coefficients, + ) + + +def get_dmuon_optimizer(model, *, opt_cfg, param_groups=None, log_fn=None): + """Build dmuon.Muon from a DMuonConfig. + + When ``param_groups`` is provided, Wall-X expects DMuon to preserve the + PyTorch optimizer group semantics and then split each user group into + dedicated/Muon and non-dedicated/AdamW subgroups internally. + """ + from wall_x.config.hyperparams_config import DMuonConfig + + if not isinstance(opt_cfg, DMuonConfig): + raise TypeError( + f"get_dmuon_optimizer expects DMuonConfig, got {type(opt_cfg).__name__}" + ) + import dmuon + + muon_signature = inspect.signature(dmuon.Muon) + supports_param_groups = "param_groups" in muon_signature.parameters + if param_groups is not None and not supports_param_groups: + raise RuntimeError( + "Wall-X built optimizer param_groups for DMuon, but the installed " + "dmuon.Muon does not accept a param_groups= argument. Please update " + "DMuon to the param-group-aware implementation before enabling " + "action_expert_learning_rate with optimizer_type='dmuon'." + ) + + ns_backend = _build_ns_backend(dmuon, opt_cfg) + + _emit( + log_fn, + "DMuon: Muon lr=%s momentum=%s ns_steps=%s; " + "AdamW lr=%s betas=%s wd=%s; " + "ns_backend=%s ns_coefficients=%s nesterov=%s", + opt_cfg.muon_lr, + opt_cfg.momentum, + opt_cfg.ns_steps, + opt_cfg.adamw_lr, + opt_cfg.adamw_betas, + opt_cfg.adamw_weight_decay, + opt_cfg.ns_backend, + opt_cfg.ns_coefficients, + opt_cfg.nesterov, + ) + if param_groups is not None: + _emit( + log_fn, + "DMuon param_groups enabled: %s", + [ + { + "group_name": group.get("group_name", f"group_{idx}"), + "lr": group.get("lr"), + "num_params": len(group.get("params", [])), + } + for idx, group in enumerate(param_groups) + ], + ) + + kwargs = {} + if param_groups is not None: + kwargs["param_groups"] = param_groups + + optimizer = dmuon.Muon( + model, + lr=opt_cfg.muon_lr, + momentum=opt_cfg.momentum, + weight_decay=opt_cfg.muon_weight_decay, + ns_steps=opt_cfg.ns_steps, + adamw_lr=opt_cfg.adamw_lr, + adamw_betas=tuple(opt_cfg.adamw_betas), + adamw_weight_decay=opt_cfg.adamw_weight_decay, + adamw_eps=opt_cfg.adamw_eps, + ns_backend=ns_backend, + nesterov=opt_cfg.nesterov, + **kwargs, + ) + + if param_groups is not None and _is_rank0(): + summarize = getattr(dmuon, "summarize_param_groups", None) + format_summary = getattr(dmuon, "format_param_group_summary", None) + if summarize is None or format_summary is None: + _emit( + log_fn, + "DMuon param_groups are enabled, but the installed DMuon package " + "does not expose param-group diagnostics. Update DMuon if you need " + "startup verification of the Muon/AdamW subgroup split.", + level=logging.WARNING, + ) + else: + try: + summary = summarize(model, optimizer, max_rows=80) + _emit(log_fn, "%s", format_summary(summary)) + except Exception as exc: + _logger.exception("Failed to summarize DMuon param_groups") + _emit( + log_fn, + "Failed to summarize DMuon param_groups: %s", + exc, + level=logging.WARNING, + ) + + return optimizer + + +register_optimizer("dmuon", get_dmuon_optimizer) diff --git a/wall_x/trainer/optimizer/utils.py b/wall_x/trainer/optimizer/utils.py new file mode 100644 index 0000000..eb1c17c --- /dev/null +++ b/wall_x/trainer/optimizer/utils.py @@ -0,0 +1,230 @@ +import inspect + +from torch.optim import AdamW + +from wall_x.config.hyperparams_config import LRGroupConfig + +_OPTIMIZERS = {} + + +def register_optimizer(name, optimizer_fn): + _OPTIMIZERS[name] = optimizer_fn + + +def get_optimizer(name, *args, **kwargs): + if name not in _OPTIMIZERS: + raise KeyError(f"Unknown optimizer '{name}'. Registered: {sorted(_OPTIMIZERS)}") + return _OPTIMIZERS[name](*args, **kwargs) + + +def _group_weight_decay(opt_cfg): + return getattr(opt_cfg, "weight_decay", None) + + +def resolve_lr_group_configs(opt_cfg, default_action_lr_keywords): + """Return structured LR groups, with legacy action config fallback. + + ``optimizer.lr_groups`` is the preferred path. The legacy + ``action_expert_learning_rate`` fields are still converted into a single + action group so older configs keep working. + """ + if opt_cfg.lr_groups: + if ( + opt_cfg.action_expert_learning_rate is not None + or opt_cfg.action_lr_keywords is not None + ): + raise ValueError( + "Use either optimizer.lr_groups or legacy " + "action_expert_learning_rate/action_lr_keywords, not both." + ) + return opt_cfg.lr_groups + + if opt_cfg.action_expert_learning_rate is None: + return [] + + action_lr_keywords = ( + opt_cfg.action_lr_keywords + if opt_cfg.action_lr_keywords is not None + else default_action_lr_keywords + ) + return [ + LRGroupConfig( + name="action_lr_group", + lr=opt_cfg.action_expert_learning_rate, + include=action_lr_keywords, + fail_on_empty=True, + ) + ] + + +def uses_legacy_action_lr_groups(opt_cfg) -> bool: + return not opt_cfg.lr_groups and opt_cfg.action_expert_learning_rate is not None + + +def _make_param_group(name, params, lr, opt_cfg): + group = { + "params": params, + "lr": lr, + "group_name": name, + } + weight_decay = _group_weight_decay(opt_cfg) + if weight_decay is not None: + group["weight_decay"] = weight_decay + return group + + +def _validate_lr_group(group: LRGroupConfig, *, base_group_name: str): + if not group.name: + raise ValueError("optimizer.lr_groups entries must have a non-empty name") + if group.name == base_group_name: + raise ValueError( + f"optimizer.lr_groups name {group.name!r} is reserved for the base group" + ) + if "/" in group.name: + raise ValueError( + f"optimizer.lr_groups name {group.name!r} must not contain '/'. " + "DMuon appends '/muon' and '/adamw' to group names." + ) + if not group.include: + raise ValueError( + f"optimizer.lr_groups.{group.name} must define at least one include keyword" + ) + + +def build_lr_param_groups(model, opt_cfg, lr_groups, *, base_group_name="base"): + """Split trainable params into named LR groups plus a base group. + + ``lr_groups`` is a list of :class:`LRGroupConfig`. Each group matches + parameter names by substring. A parameter may match at most one explicit + group; unmatched trainable parameters remain in the ``base`` group using + ``opt_cfg.learning_rate``. + + Returns a list of torch.optim-compatible param_group dicts. The + ``group_name`` key is non-standard but preserved by torch.optim via + ``setdefault`` in ``add_param_group`` and is consumed downstream for + per-group lr logging. + + For AdamW / native Muon, each returned group includes ``weight_decay``. + For DMuon, weight decay is split between Muon and AdamW defaults, so the + groups only carry lr and metadata; DMuon applies its own per-route defaults. + The caller is expected to gate on ``opt_cfg.optimizer_type`` before calling + this. + """ + if not lr_groups: + return None + + names = [group.name for group in lr_groups] + duplicate_names = sorted({name for name in names if names.count(name) > 1}) + if duplicate_names: + raise ValueError( + f"optimizer.lr_groups contains duplicate names: {duplicate_names}" + ) + + for group in lr_groups: + _validate_lr_group(group, base_group_name=base_group_name) + + base_params = [] + grouped_params = {group.name: [] for group in lr_groups} + ambiguous = [] + + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + matches = [ + group.name + for group in lr_groups + if any(keyword in name for keyword in group.include) + ] + if len(matches) > 1: + ambiguous.append((name, matches)) + continue + if matches: + grouped_params[matches[0]].append(param) + else: + base_params.append(param) + + if ambiguous: + examples = ", ".join(f"{name} -> {matches}" for name, matches in ambiguous[:10]) + raise ValueError( + "Some parameters match multiple optimizer.lr_groups. Make group " + f"include patterns disjoint. Examples: {examples}" + ) + + if opt_cfg.train_action_expert_only: + assert len(base_params) == 0, ( + f"Expected 0 base_params after pre-wrap freeze, got {len(base_params)}. " + "Ensure base params are frozen before building the optimizer." + ) + + param_groups = [] + if len(base_params) > 0: + param_groups.append( + _make_param_group( + base_group_name, base_params, opt_cfg.learning_rate, opt_cfg + ) + ) + + for group in lr_groups: + params = grouped_params[group.name] + if len(params) == 0: + if group.fail_on_empty: + raise ValueError( + f"No params found for optimizer.lr_groups.{group.name}. " + f"Please check include={group.include!r}." + ) + continue + param_groups.append(_make_param_group(group.name, params, group.lr, opt_cfg)) + + return param_groups + + +def build_action_expert_param_groups(model, opt_cfg, action_lr_keywords): + """Compatibility wrapper for the legacy action-expert LR config.""" + return build_lr_param_groups( + model, + opt_cfg, + [ + LRGroupConfig( + name="action_lr_group", + lr=opt_cfg.action_expert_learning_rate, + include=action_lr_keywords, + fail_on_empty=True, + ) + ], + base_group_name="base_lr_group", + ) + + +def get_adamw_optimizer(model, *, opt_cfg, param_groups=None): + """Build AdamW from AdamWConfig.""" + from wall_x.config.hyperparams_config import AdamWConfig + + if not isinstance(opt_cfg, AdamWConfig): + raise TypeError( + f"get_adamw_optimizer expects AdamWConfig, got {type(opt_cfg).__name__}" + ) + + if param_groups is None: + params = [p for p in model.parameters() if p.requires_grad] + else: + # Per-group lr / weight_decay are preserved by torch.optim (via + # setdefault in add_param_group), so top-level values act only as + # defaults. Extra keys like ``group_name`` are kept in-place and used + # downstream for per-group lr logging. + params = param_groups + + kw = { + "lr": opt_cfg.learning_rate, + "weight_decay": opt_cfg.weight_decay, + "betas": tuple(opt_cfg.betas), + "eps": opt_cfg.eps, + } + sig_params = inspect.signature(AdamW.__init__).parameters + if "foreach" in sig_params and opt_cfg.foreach is not None: + kw["foreach"] = opt_cfg.foreach + if "fused" in sig_params: + kw["fused"] = opt_cfg.fused + return AdamW(params, **kw) + + +register_optimizer("adamw", get_adamw_optimizer) diff --git a/wall_x/trainer/qwen_vl_act_trainer.py b/wall_x/trainer/qwen_vl_act_trainer.py deleted file mode 100644 index c981e24..0000000 --- a/wall_x/trainer/qwen_vl_act_trainer.py +++ /dev/null @@ -1,1144 +0,0 @@ -import os -import gc -import time -import yaml -import shutil -import torch -import random -import numpy as np -import torch.nn as nn -import torch.distributed as dist -import json -from tqdm import tqdm -from functools import wraps -from datetime import datetime -from torch.optim import AdamW -from torch.distributed.tensor import distribute_tensor -from accelerate import Accelerator -from safetensors.torch import load_file -from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup -from transformers import AutoProcessor -from wall_x.model.action_head import Normalizer -from wall_x.utils.timers import Timers -from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction, Qwen2_5_VLConfig -from wall_x.utils.constant import action_statistic_dof as default_action_statistic_dof -from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES -from wall_x.data.load_lerobot_dataset import ( - PreprocessedDataset, - get_data_configs, - load_lerobot_data, -) -import copy - - -def timer(func): - """ - Decorator to measure function execution time. - - Args: - func: Function to be timed - - Returns: - Wrapped function with timing functionality - """ - - @wraps(func) - def wrapper(*args, **kwargs): - start_time = time.time() - result = func(*args, **kwargs) - end_time = time.time() - print( - f"\033[92m[current time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Function {func.__name__} took {end_time - start_time:.2f} seconds to execute\033[0m" - ) - return result - - return wrapper - - -def print_rank_last(message): - """ - Print message only on the last rank in distributed training. - - Args: - message (str): Message to print - """ - if torch.distributed.is_initialized(): - if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1): - print(message, flush=True) - else: - print(message, flush=True) - - -def seed_all(seed): - """ - Set random seeds for reproducible training. - - Args: - seed (int): Random seed value - """ - np.random.seed(seed) - random.seed(seed) - torch.manual_seed(seed) - - -def update_model_config(train_config, model_config): - model_config.use_state_string_representation = train_config["data"].get( - "use_state_string_representation", False - ) - model_config.flow_loss_weight = train_config.get("flow_loss_weight", 1.0) - - model_config.dof_config = train_config["dof_config"] - model_config.agent_pos_config = train_config["agent_pos_config"] - - model_config.action_horizon_flow = train_config["data"].get( - "action_horizon_flow", 32 - ) - - if train_config.get("_attn_implementation", None) is not None: - model_config._attn_implementation = train_config["_attn_implementation"] - - if train_config.get("attn_deterministic", None) is not None: - model_config.attn_deterministic = train_config["attn_deterministic"] - model_config.vision_config.attn_deterministic = train_config[ - "attn_deterministic" - ] - print("[DEBUG] Attention is using deterministic kernel for this run!") - else: - model_config.attn_deterministic = False - model_config.vision_config.attn_deterministic = False - - return model_config - - -class QwenVlAct_Trainer: - """ - Vision-Language-Action trainer for Qwen-VL models with robotic action prediction. - - This trainer handles multi-modal learning combining vision, language, and action data - for robotic control applications. It supports distributed training, mixed precision, - gradient accumulation, and various optimization strategies including MoE (Mixture of Experts). - - Features: - - Multi-modal data processing (vision + language + actions) - - Distributed training with Accelerate - - Gradient accumulation and clipping - - Learning rate scheduling with warmup - - Checkpoint saving and resuming - - Comprehensive logging and monitoring - """ - - @timer - def __init__( - self, - config, - logger, - accelerator: Accelerator = None, - seed=42, - data_config_path=None, - ): - """ - Initialize the Vision-Language-Action trainer. - - Args: - config (dict): Training configuration dictionary containing: - - processor_path (str): Path to data preprocessing processor - - qwen_vl_act_config_path (str): Path to model configuration file - - learning_rate (float): Base learning rate for training - - num_epoch (int): Number of training epochs - - pretrained_wallx_path (str): Path to pretrained model - - And other training hyperparameters - logger: Logger instance for tracking metrics - accelerator (Accelerator, optional): Hugging Face Accelerate instance for distributed training - seed (int, optional): Random seed for reproducibility. Defaults to 42. - data_config_path (str, optional): Path to data configuration file - - Raises: - ValueError: If required configuration keys are missing - """ - # Validate required configuration keys - required_keys = ["learning_rate", "num_epoch"] - for key in required_keys: - if key not in config: - raise ValueError(f"Missing required configuration key: {key}") - - self.config = config - self.logger = logger - self.accelerator = accelerator - self.seed = seed - - # Initialize random seeds for reproducibility - seed_all(self.seed) - - # Training state variables - self.start_epoch = 0 - self.global_step = 0 - self.num_epoch = self.config["num_epoch"] - self.initial_step = 0 - - # Data and model configuration - self.dataload_config = get_data_configs(self.config["data"]) - self.data_config_path = data_config_path - self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) - self.use_selective_recompute = self.config.get("use_selective_recompute", False) - - # Load model and initialize training components - self.load_normalizer() - self.load_model() - self.action_dim = sum(self.config["dof_config"].values()) - - # Distributed training setup - self.rank = self.accelerator.process_index - self.world_size = self.accelerator.num_processes - print( - f"rank {self.accelerator.process_index} after load model memory usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", - flush=True, - ) - - # Load training data - self.load_qact_data() - print( - f"rank {self.accelerator.process_index} after load qact data usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", - flush=True, - ) - - # Resume from checkpoint if specified - if "resume" in self.config: - self.resume_from_checkpoint() - - # Initialize special token IDs - self.propri_token_id = self.processor.tokenizer.convert_tokens_to_ids( - "<|propri|>" - ) - self.action_token_id = self.processor.tokenizer.convert_tokens_to_ids( - "<|action|>" - ) - - # Initialize evaluation metrics - self.base_l1_loss = None - self.base_l1_loss_detail = {} - - # Performance monitoring - self.timers = Timers(log_level=0, log_option="minmax") - - # Adjust global step if resuming from checkpoint - if self.initial_step != 0: - self.global_step = self.initial_step // self.config.get( - "gradient_accumulation_steps", 1 - ) - - def load_normalizer(self): - if self.config.get("norm_stats_path", None): - self.print_rank0( - f"loading customized action statistic dof from {self.config['norm_stats_path']}" - ) - action_statistic_dof = json.load(open(self.config["norm_stats_path"], "r")) - else: - self.print_rank0( - "loading default action statistic dof from default_action_statistic_dof" - ) - action_statistic_dof = default_action_statistic_dof - - self.normalizer_action = Normalizer( - action_statistic_dof, - self.config["dof_config"], - min_key=self.config.get("min_key", "min"), - delta_key=self.config.get("delta_key", "delta"), - ) - - print("self.normalizer_action.min: ", self.normalizer_action) - self.normalizer_propri = Normalizer( - action_statistic_dof, - self.config["agent_pos_config"], - min_key=self.config.get("min_key", "min"), - delta_key=self.config.get("delta_key", "delta"), - ) - - def print_rank0(self, msg, flush=True): - """ - Print message only on rank 0 (main process). - - Args: - msg: Message to print - flush (bool): Whether to flush output buffer - """ - if self.accelerator.is_main_process: - print(msg, flush=flush) - - def fit(self): - """ - Main training loop executing multiple epochs with validation. - - Handles the complete training process including: - - Training loop execution - - Validation after each epoch - - Process synchronization - - Memory cleanup - """ - self.accelerator.wait_for_everyone() - # Optional validation before training starts - if self.config.get("resume", None) is not None and self.config["resume"].get( - "validate_first", False - ): - self.val_loop() - self.accelerator.wait_for_everyone() - - # Main training loop - for epoch in range(self.start_epoch, self.num_epoch): - self.train_loop(epoch) - self.accelerator.wait_for_everyone() - - if (epoch + 1) % self.config.get("epoch_save_interval", 1) == 0: - self.save_checkpoint(epoch) - - # Validation after each epoch - # self.val_loop() - self.accelerator.wait_for_everyone() - - # Memory cleanup - gc.collect() - - def train_loop(self, epoch): - """ - Execute training for a single epoch. - - Args: - epoch (int): Current epoch number - - Handles: - - Data loading and batching - - Forward/backward passes - - Gradient accumulation and clipping - - Learning rate scheduling - - Loss logging and monitoring - - Performance profiling (optional) - """ - # Initialize training dataloader for current epoch - if isinstance(self.dataset, PreprocessedDataset): - if getattr(self, "train_dataloader", None) is not None: - self.dataset._train() - self.train_sampler.set_epoch(epoch) - else: - self.train_dataloader, self.train_sampler = ( - self.dataset.get_train_dataloader() - ) - self.train_sampler.set_epoch(epoch) - else: - self.train_dataloader = self.dataset.get_train_dataloader() - - self.model.train() - total = len(self.train_dataloader) - t0 = time.time() - enable_profiling = self.config["profile"] - - # Optional PyTorch profiler for performance analysis - if enable_profiling: - profiler = torch.profiler.profile( - activities=[ - torch.profiler.ProfilerActivity.CPU, - torch.profiler.ProfilerActivity.CUDA, - ], - schedule=torch.profiler.schedule( - wait=self.config["profile_wait_iters"], - warmup=self.config["profile_warmup_iters"], - active=self.config["profile_active_iters"], - ), - on_trace_ready=torch.profiler.tensorboard_trace_handler( - self.config["profile_save_path"], worker_name="worker0" - ), - record_shapes=True, - profile_memory=True, - with_stack=True, - ) - profiler.__enter__() - - try: - - # Setup timers for First iteration - self.timers("interval-time", log_level=0).start(barrier=False) - self.timers("data-load", log_level=0).start(barrier=False) - - for i, batch in enumerate(self.train_dataloader, self.initial_step): - # Move batch to device - if isinstance(self.dataset, PreprocessedDataset): - batch = { - k: ( - v.to(self.accelerator.device, non_blocking=True) - if isinstance(v, torch.Tensor) - else v - ) - for k, v in batch.items() - } - - self.timers("data-load").stop() - - with self.accelerator.accumulate(self.model): - # Forward pass - self.timers("forward-compute", log_level=0).start(barrier=False) - outputs = self.model(**batch, mode="train") - self.timers("forward-compute").stop() - - loss = outputs.loss - # Check for NaN loss - if torch.isnan(loss): - print( - f"Warning: NaN loss detected in epoch: {epoch}, step: {i}", - flush=True, - ) - continue - - # Backward pass - self.timers("backward-compute", log_level=0).start(barrier=False) - self.accelerator.backward(loss) - self.timers("backward-compute").stop() - - # Gradient clipping - total_norm = self.accelerator.clip_grad_norm_( - self.model.parameters(), self.config.get("max_grad_norm", 1.0) - ) - - # Optimizer step - self.timers("optimizer", log_level=0).start(barrier=False) - self.optimizer.step() - self.optimizer.zero_grad() - self.timers("optimizer").stop() - - # Update global step and learning rate after gradient accumulation - if self.accelerator.sync_gradients: - self.lr_scheduler.step() - self.global_step += 1 - lr = self.lr_scheduler.get_last_lr()[0] - - # Gather loss across all processes for logging - train_loss = ( - self.accelerator.gather(loss.detach()).mean().item() - ) - _log_dict = { - "lr": lr, - "train_loss": train_loss, - } - - # Log component losses - if ( - "cross_entropy_loss" in outputs - and outputs.cross_entropy_loss is not None - ): - _log_dict["cross_entropy_loss"] = ( - self.accelerator.gather( - outputs.cross_entropy_loss.detach() - ) - .mean() - .item() - ) - - if "flow_loss" in outputs and outputs.flow_loss is not None: - _log_dict["flow_loss"] = ( - self.accelerator.gather(outputs.flow_loss.detach()) - .mean() - .item() - ) - - # Log per-dataset channel losses - if ( - "channel_loss_dict" in outputs - and outputs.channel_loss_dict is not None - ): - for dataset_name_i in ( - ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES - ): - count_sum = ( - self.accelerator.gather( - outputs.channel_loss_count_dict[dataset_name_i] - ) - .sum() - .item() - ) - if count_sum > 0: - channel_loss = ( - self.accelerator.gather( - outputs.channel_loss_dict[ - dataset_name_i - ].detach() - ) - .sum() - .item() - / count_sum - ) - _log_dict[f"channel_loss_{dataset_name_i}"] = ( - channel_loss - ) - - # Log action accuracy for fast tokenizer - if ( - "action_accuracy" in outputs.channel_loss_dict - and self.use_fast_tokenizer - ): - _log_dict["action_accuracy"] = ( - self.accelerator.gather( - outputs.channel_loss_dict[ - "action_accuracy" - ].detach() - ) - .mean() - .item() - ) - - # Log metrics - if self.logger is not None: - self.logger.log(_log_dict, step=self.global_step) - - # Log gradient norm - if self.logger is not None and self.accelerator.sync_gradients: - self.logger.log( - {"total_norm": total_norm}, step=self.global_step - ) - - self.timers("interval-time").stop() - - # Setup timers for next iteration - if i < len(self.train_dataloader) - 1: - self.timers("interval-time", log_level=0).start(barrier=False) - self.timers("data-load", log_level=0).start(barrier=False) - - # Periodic logging - t1 = time.time() - if i % 1 == 0: - lr = self.lr_scheduler.get_last_lr()[0] - self.training_log( - epoch, self.num_epoch, i, total, loss, lr, t1 - t0 - ) - t0 = time.time() - - if enable_profiling: - profiler.step() - - finally: - if enable_profiling: - profiler.__exit__(None, None, None) - - @torch.no_grad() - def val_loop(self): - """ - Execute validation loop with gradient computation disabled. - - Evaluates model performance on validation set and logs validation loss. - """ - # Initialize validation dataloader - if getattr(self, "val_dataloader", None) is not None: - self.dataset._eval() - self.val_sampler.set_epoch(0) - else: - self.val_dataloader, self.val_sampler = self.dataset.get_val_dataloader() - self.val_sampler.set_epoch(0) - - self.model.eval() - self.val_loss = 0 - - # Validation loop - for i, batch in enumerate( - tqdm( - self.val_dataloader, - desc="Validating", - total=len(self.val_dataloader), - disable=not self.accelerator.is_main_process, - ) - ): - if isinstance(self.dataset, PreprocessedDataset): - batch = { - k: ( - v.to(self.accelerator.device, non_blocking=True) - if isinstance(v, torch.Tensor) - else v - ) - for k, v in batch.items() - } - - with torch.no_grad(): - outputs = self.model(**batch, mode="train") - loss = outputs.loss - self.val_loss += self.accelerator.gather(loss.detach()).mean().item() - - # Calculate average validation loss - self.val_loss /= len(self.val_dataloader) - - # Log validation metrics - if self.logger is not None: - self.logger.log({"val_loss": self.val_loss}, step=self.global_step) - - self.model.train() - - @timer - def load_model(self): - """ - Load and configure the Vision-Language-Action model. - - Handles: - - Model loading from pretrained weights - - Processor initialization - - Optimizer configuration (with support for different learning rates for different components) - - Learning rate scheduler setup - - Model preparation for distributed training - """ - # Load pretrained model - model_type = self.config.get("model_type", "qwen2_5") - assert model_type in ["wall-oss", "qwen2_5"] - if model_type == "wall-oss": - model = Qwen2_5_VLMoEForAction.from_pretrained( - self.config["pretrained_wallx_path"], - train_config=self.config, - action_tokenizer_path=( - self.config["action_tokenizer_path"] - if self.use_fast_tokenizer - else None - ), - ) - self.processor = model.processor - model = model.to(torch.bfloat16) - elif model_type == "qwen2_5": - - model_config = Qwen2_5_VLConfig.from_pretrained( - self.config["qwen_vl_act_config_path"] - ) - flow_loss_weight = self.config.get("flow_loss_weight", 1.0) - self.processor = AutoProcessor.from_pretrained( - self.config["pretrained_wallx_path"], use_fast=True - ) - new_tokens = ["<|propri|>", "<|action|>"] - self.processor.tokenizer.add_tokens(new_tokens) - if self.config.get("use_fast_tokenizer", False): - action_tokenizer_path = self.config["action_tokenizer_path"] - action_tokenizer = AutoProcessor.from_pretrained( - action_tokenizer_path, trust_remote_code=True - ) - # process for use fast - new_tokens = [ - f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size) - ] - self.processor.tokenizer.add_tokens(new_tokens) - begin_idx_token = "<|action_token_0|>" - token_id = self.processor.tokenizer.convert_tokens_to_ids( - begin_idx_token - ) - self.processor.tokenizer.init_kwargs["action_token_start_index"] = ( - token_id - ) - self.processor.tokenizer.init_kwargs["action_token_vocab_size"] = ( - action_tokenizer.vocab_size - ) - self.processor.action_processor = action_tokenizer - - # Set the customized robot configuration to ensure consistency between cross-embodiment - # representations and the Wall-X action dimensionality. - # Qwen2_5_VLMoEForAction._set_customized_config(self.config) - customized_dof_config = self.config["customized_robot_config"][ - "customized_dof_config" - ] - customized_agent_pos_config = self.config["customized_robot_config"][ - "customized_agent_pos_config" - ] - setattr(model_config, "customized_dof_config", customized_dof_config) - setattr( - model_config, "customized_agent_pos_config", customized_agent_pos_config - ) - - model_config = update_model_config(self.config, model_config) - model = Qwen2_5_VLMoEForAction( - model_config, - self.use_fast_tokenizer, - self.processor, - flow_loss_weight=flow_loss_weight, - use_selective_recompute=self.use_selective_recompute, - ) - - model = model.to(torch.bfloat16) - model = self.load_qwen_pretrain_weight( - model, self.config["pretrained_wallx_path"] - ) - model.resize_token_embeddings(len(self.processor.tokenizer)) - model = model.to(torch.bfloat16) - else: - raise NotImplementedError(f"Invalid model type: {model_type}") - - # Configure optimizer based on training strategy - if "freeze_vlm" in self.config and self.config["freeze_vlm"]: - print("Freezing VLM parameters, training only MoE experts", flush=True) - moe_params = [] - for name, param in model.named_parameters(): - if "moe.experts.1." not in name: - param.requires_grad = False - else: - moe_params.append(param) - param_groups = [{"params": moe_params, "lr": self.config["learning_rate"]}] - self.optimizer = AdamW(param_groups, weight_decay=0.1) - - elif "action_expert_learning_rate" in self.config: - # Separate learning rates for VLM and action expert parameters - moe_params = [] - vlm_params = [] - for name, param in model.named_parameters(): - if "moe.experts.1." in name: - moe_params.append(param) - else: - vlm_params.append(param) - - # Configure parameter groups - if self.config.get("train_action_expert_only", False): - self.print_rank0("Training action expert only", flush=True) - param_groups = [ - { - "params": moe_params, - "lr": self.config["action_expert_learning_rate"], - } - ] - else: - param_groups = [ - {"params": vlm_params, "lr": self.config["learning_rate"]}, - { - "params": moe_params, - "lr": self.config["action_expert_learning_rate"], - }, - ] - - self.optimizer = AdamW(param_groups, weight_decay=0.1) - self.print_rank0( - f"Setting MoE learning rate to {self.config['action_expert_learning_rate']}, " - f"VLM learning rate to {self.config['learning_rate']}", - flush=True, - ) - else: - # Standard optimizer configuration - self.optimizer = AdamW( - model.parameters(), - lr=self.config["learning_rate"], - weight_decay=0.1, - ) - - # Configure learning rate scheduler - warmup_steps = self.config.get("num_warmup_steps", 0) - num_training_steps = self.config.get("num_training_steps", 1000000000) - min_lr = self.config.get("min_lr", 0.1 * self.config["learning_rate"]) - self.lr_scheduler = get_cosine_with_min_lr_schedule_with_warmup( - optimizer=self.optimizer, - num_warmup_steps=warmup_steps, - num_training_steps=num_training_steps, - min_lr=min_lr, - ) - - self.model = model - - # Enable gradient computation for embeddings - if hasattr(model, "enable_input_require_grads"): - self.model.enable_input_require_grads() - else: - - def make_inputs_require_grad(module, input, output): - output.requires_grad_(True) - - self.model.get_input_embeddings().register_forward_hook( - make_inputs_require_grad - ) - - # Prepare model, optimizer, and scheduler for distributed training - self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare( - self.model, self.optimizer, self.lr_scheduler - ) - - @timer - def load_qact_data(self): - """ - Load and configure training data for Vision-Language-Action learning. - - Supports LeRobot dataset format and handles distributed data loading - across multiple processes. - """ - print(f"Loading Vision-Language-Action data from {__file__}") - self.accelerator.wait_for_everyone() - - # Load LeRobot dataset - self.dataset, self.train_num = load_lerobot_data( - config=self.config, - lerobot_config=self.dataload_config.get("lerobot_config", {}), - normalizer_action=copy.deepcopy(self.normalizer_action), - normalizer_propri=copy.deepcopy(self.normalizer_propri), - rank=self.rank, - world_size=self.world_size, - ) - - @timer - def load_qwen_pretrain_weight(self, model, pretrain_weight_path): - """ - Load pretrained Qwen weights with MoE adaptation. - - Args: - model: Model instance to load weights into - pretrain_weight_path (str): Path to pretrained weight files - - Returns: - Model with loaded pretrained weights - - Handles weight key renaming for MoE architecture compatibility. - """ - # Load all safetensors files - weight_files = sorted( - [f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")] - ) - merged_weights = {} - - # Merge weights from all files - for weight_file in weight_files: - file_path = os.path.join(pretrain_weight_path, weight_file) - weights = load_file(file_path) - merged_weights.update(weights) - - # Rename weights for MoE compatibility - renamed_weights = {} - for key, value in merged_weights.items(): - if ( - key.startswith("model.layers") - and "mlp." in key - and model.config.mlp_moe - ): - # Rename MLP weights for MoE structure - layer_num = key.split(".layers.")[1].split(".mlp")[0] - new_key = key.replace( - f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0." - ) - renamed_weights[new_key] = value - elif ( - key.startswith("model.layers") - and "self_attn." in key - and model.config.attention_moe - ): - # Rename attention weights for MoE structure - layer_num = key.split(".layers.")[1].split(".self_attn")[0] - proj_types = ["q_proj", "k_proj", "v_proj", "o_proj"] - for proj in proj_types: - if proj in key: - new_key = key.replace( - f"layers.{layer_num}.self_attn.{proj}", - f"layers.{layer_num}.self_attn.{proj}_experts.0", - ) - renamed_weights[new_key] = value - break - else: - renamed_weights[key] = value - - # Load weights into model - # err = model.load_state_dict(renamed_weights, strict=False) - # self.print_rank0(f"Weight loading report: {err}", flush=True) - if self.accelerator.is_main_process: - self.print_rank0(f"Loaded pretrained weights from: {pretrain_weight_path}") - - return model - - def training_log( - self, - current_epoch, - total_epoch, - current_train_iter, - total_train_iter, - loss, - lr, - time_per_step, - ): - """ - Log training progress and performance metrics. - - Args: - current_epoch (int): Current epoch number - total_epoch (int): Total number of epochs - current_train_iter (int): Current training iteration - total_train_iter (int): Total iterations in epoch - loss (torch.Tensor): Current loss value - lr (float): Current learning rate - time_per_step (float): Time taken for current step - """ - timers_to_log = [ - "interval-time", - "data-load", - "forward-compute", - "backward-compute", - "optimizer", - ] - - log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]" - log_string += " epoch {:3d}/{:3d} |".format(current_epoch, total_epoch) - log_string += " iter {:6d}/{:6d} |".format(current_train_iter, total_train_iter) - log_string += " loss {:.6f} |".format(loss) - log_string += " lr {:.6f} |".format(lr) - log_string += " time_per_step_avg {:.6f}s |".format(time_per_step) - - print_rank_last(log_string) - self.timers.log(timers_to_log, normalizer=1) - - def save_checkpoint(self, epoch, step=0): - save_path = self.config["save_path"] - if step == 0: - ckpt_path = f"{save_path}/{epoch}" - else: - ckpt_path = f"{save_path}/{epoch}_{step}" - self.accelerator.save_state(ckpt_path) - - # Save random seed - if self.accelerator.is_main_process: - # FIXME the dataset does not have random seed now. Should the dataset set the random seed? - torch.save( - {"seed": self.seed}, os.path.join(ckpt_path, "seed.pth") - ) # seed is shared by all ranks; seed follows dataset - torch.save( - {"global_step": self.global_step}, - os.path.join(ckpt_path, "global_step.pth"), - ) - torch.save( - {"current_epoch": epoch}, os.path.join(ckpt_path, "current_epoch.pth") - ) - - # Save configuration in YAML format - config_path = os.path.join(ckpt_path, "config.yml") - with open(config_path, "w", encoding="utf-8") as f: - yaml.dump( - self.config, - f, - default_flow_style=False, - allow_unicode=True, - indent=2, - sort_keys=False, - ) - - pretrained_dir = self.config.get("pretrained_qwen_vl_path", None) - if pretrained_dir is not None: - files_to_copy = [ - "preprocessor_config.json", - "tokenizer_config.json", - "tokenizer.json", - "vocab.json", - ] - - for filename in files_to_copy: - src = os.path.join(pretrained_dir, filename) - dst = os.path.join(ckpt_path, filename) - - if os.path.exists(src): - shutil.copy(src, dst) - print(f"[Checkpoint] Copied {filename} to {ckpt_path}") - else: - print(f"[Checkpoint] WARNING: {src} not found, skip copying.") - - act_config_path = self.config.get("qwen_vl_act_config_path", None) - if act_config_path is not None: - dst = os.path.join(ckpt_path, "config.json") - - if os.path.exists(act_config_path): - shutil.copy(act_config_path, dst) - print(f"[Checkpoint] Copied act config to {dst}") - else: - print( - f"[Checkpoint] WARNING: {act_config_path} not found, skipping." - ) - # Save normalizer - torch.save( - self.normalizer_action.state_dict(), - os.path.join(ckpt_path, "normalizer_action.pth"), - ) - torch.save( - self.normalizer_propri.state_dict(), - os.path.join(ckpt_path, "normalizer_propri.pth"), - ) - - # Save current iter steps - if step != 0: # step==0, no need for dataset resume - _rank = self.accelerator.process_index - if self.data_config["multimodal_data_ratio"] != 1: - torch.save( - { - "episode_start_index": self.dataset.primary_pool_start_index.value - }, - os.path.join(ckpt_path, f"episode_start_index_rank_{_rank}.pth"), - ) - torch.save( - { - "multimodal_episode_start_index": self.dataset.secondary_pool_start_index.value - }, - os.path.join( - ckpt_path, f"multimodal_episode_start_index_rank_{_rank}.pth" - ), - ) - else: - torch.save( - { - "episode_start_index": self.dataset.secondary_pool_start_index.value - }, - os.path.join(ckpt_path, f"episode_start_index_rank_{_rank}.pth"), - ) - torch.save( - { - "multimodal_episode_start_index": self.dataset.primary_pool_start_index.value - }, - os.path.join( - ckpt_path, f"multimodal_episode_start_index_rank_{_rank}.pth" - ), - ) - - def resume_from_checkpoint(self): - """ - Resume training from a saved checkpoint. - - Handles both full checkpoint loading and model-only loading based on configuration. - """ - - if self.config.get("resume", {}).get("load_ckpt_only", False): - if self.config.get("FSDP2", False): - self._load_fsdp_state_dict_with_distribute_tensor() - - else: - # Load only model weights - ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors" - state_dict = load_file(ckpt_path, device="cpu") - - # Add module prefix if needed for distributed training - new_state_dict = {} - for key in state_dict: - if not key.startswith("module."): - new_key = "module." + key - new_state_dict[new_key] = state_dict[key] - - self.model.load_state_dict(new_state_dict, strict=False) - else: - # Load full checkpoint including optimizer and scheduler states - - # self.accelerator.load_state(checkpoint_path) - state_dict = load_file( - self.config["resume"]["ckpt"] + "/model.safetensors", device="cpu" - ) - - filtered_state_dict = { - k: v - for k, v in state_dict.items() - if not k.startswith("action_preprocessor.normalizer") - } - - if self.config["resume"].get("try_harder", False): - new_state_dict = {} - for name, param in filtered_state_dict.items(): - if name in self.model.state_dict(): - if param.size() == self.model.state_dict()[name].size(): - new_state_dict[name] = param - else: - size_0 = param.size() - size_1 = self.model.state_dict()[name].size() - new_state_dict[name] = self.model.state_dict()[name] - slices = [ - slice(0, min(old_dim, new_dim)) - for old_dim, new_dim in zip(size_0, size_1) - ] - new_state_dict[name][slices] = param[slices] - self.print_rank0( - f"Not match key: {name}, required shape: {size_1}, loaded shape: {size_0}, new shape: {new_state_dict[name].size()}" - ) - elif "module." + name in self.model.state_dict(): - name = "module." + name - if param.size() == self.model.state_dict()[name].size(): - new_state_dict[name] = param - else: - size_0 = param.size() - size_1 = self.model.state_dict()[name].size() - new_state_dict[name] = self.model.state_dict()[name] - slices = [ - slice(0, min(old_dim, new_dim)) - for old_dim, new_dim in zip(size_0, size_1) - ] - new_state_dict[name][slices] = param[slices] - self.print_rank0( - f"Not match key: {name}, required shape: {size_1}, loaded shape: {size_0}, new shape: {new_state_dict[name].size()}" - ) - else: - self.print_rank0(f"Not used parameter: {name}") - err = self.model.load_state_dict(new_state_dict, strict=False) - else: - err = self.model.load_state_dict(filtered_state_dict, strict=False) - - self.print_rank0(f"err in load model: {err}", err) - - def _load_fsdp_state_dict_with_distribute_tensor(self): - - rank = dist.get_rank() if dist.is_initialized() else 0 - - full_sd = load_file( - self.config["resume"]["ckpt"] + "/model.safetensors", device="cpu" - ) - meta_sharded_sd = self.model.state_dict() - sharded_sd = {} - - def find_matching_key(target_key, available_keys): - if target_key in available_keys: - return target_key - prefixed_key = f"_orig_mod.{target_key}" - if prefixed_key in available_keys: - return prefixed_key - if target_key.startswith("_orig_mod."): - unprefixed_key = target_key[len("_orig_mod.") :] - if unprefixed_key in available_keys: - return unprefixed_key - return None - - for param_name, full_tensor in full_sd.items(): - matching_key = find_matching_key(param_name, meta_sharded_sd.keys()) - if matching_key is None: - if rank == 0: - print( - f"[Rank {rank}] Warning: Parameter not found:", - param_name, - flush=True, - ) - continue - sharded_meta_param = meta_sharded_sd[matching_key] - sharded_tensor = distribute_tensor( - full_tensor, - sharded_meta_param.device_mesh, - sharded_meta_param.placements, - ) - sharded_sd[matching_key] = nn.Parameter(sharded_tensor) - - self.model.load_state_dict(sharded_sd, assign=True, strict=False) - - def log_l1_details(self, all_label, all_pred, all_task, all_dof_mask): - """ - Log detailed L1 loss metrics by degrees of freedom. - - Args: - all_label (torch.Tensor): Ground truth action labels - all_pred (torch.Tensor): Predicted actions - all_task (list): Task identifiers - all_dof_mask (torch.Tensor): Degrees of freedom mask - - Computes and logs L1 loss for each DOF component separately for detailed analysis. - """ - all_task = all_task[: len(all_label)] - - # Apply DOF mask - all_label = all_label * all_dof_mask - all_pred = all_pred * all_dof_mask - - # Compute baseline L1 loss (predict mean action) - if self.base_l1_loss is None: - mean_action = all_label.mean(dim=0) - self.base_l1_loss = nn.functional.l1_loss(all_label, mean_action) - - self.logger.log( - {"base_l1_loss": self.base_l1_loss.item()}, step=self.global_step - ) - - # Log L1 loss for each DOF component - start_idx = 0 - dof_config = self.config.get("dof_config", {}) - for dof in dof_config: - end_idx = start_idx + dof_config[dof] - dof_label = all_label[:, :, start_idx:end_idx] - dof_pred = all_pred[:, :, start_idx:end_idx] - dof_l1 = nn.functional.l1_loss(dof_pred, dof_label) - - self.print_rank0(f"DOF {dof}, L1 loss: {dof_l1.item()}", flush=True) - self.logger.log( - {f"detail/l1_loss_{dof}": dof_l1.item()}, step=self.global_step - ) - - start_idx = end_idx diff --git a/wall_x/trainer/scheduler/scheduler.py b/wall_x/trainer/scheduler/scheduler.py new file mode 100644 index 0000000..2b328bb --- /dev/null +++ b/wall_x/trainer/scheduler/scheduler.py @@ -0,0 +1,71 @@ +import math + +from torch.optim import Optimizer +from torch.optim.lr_scheduler import ConstantLR, LambdaLR + + +def create_cosine_scheduler( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + peak_lr: float | None = None, + end_lr: float | None = None, + last_epoch: int = -1, +): + if peak_lr is None: + peak_lr = float(optimizer.defaults["lr"]) + if end_lr is None: + end_lr = peak_lr * 0.1 + + def lr_lambda(current_step: int): + if current_step < num_warmup_steps: + # Start from peak_lr / (warmup_steps + 1). + init_lr = peak_lr / (num_warmup_steps + 1) + current_lr = init_lr + (peak_lr - init_lr) * current_step / num_warmup_steps + return current_lr / peak_lr # LambdaLR multiplies by base_lr + else: + # Cosine decay + decay_steps = num_training_steps - num_warmup_steps + progress = min(1.0, (current_step - num_warmup_steps) / max(1, decay_steps)) + cos = 0.5 * (1 + math.cos(math.pi * progress)) + current_lr = end_lr + (peak_lr - end_lr) * cos + return current_lr / peak_lr + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def create_step_scheduler( + optimizer: Optimizer, + lr_decay_steps: str, + lr_gamma: float = 0.1, +): + decay_steps = [int(s.strip()) for s in lr_decay_steps.split(",")] + + def lr_lambda(current_step): + factor = 1.0 + for step in decay_steps: + if current_step >= step: + factor *= lr_gamma + return factor + + return LambdaLR(optimizer, lr_lambda) + + +def create_constant_scheduler( + optimizer: Optimizer, + factor: float = 1 / 3, + total_iters: int = 5, + last_epoch: int = -1, +): + return ConstantLR(optimizer, factor, total_iters, last_epoch) + + +def get_scheduler(optimizer: Optimizer, lr_scheduler_type: str, **kwargs): + if lr_scheduler_type == "cosine": + return create_cosine_scheduler(optimizer, **kwargs) + elif lr_scheduler_type == "step": + return create_step_scheduler(optimizer, **kwargs) + elif lr_scheduler_type == "constant": + return create_constant_scheduler(optimizer, **kwargs) + else: + raise ValueError(f"Unsupported lr_scheduler: {lr_scheduler_type}") diff --git a/wall_x/trainer/trainer_utils.py b/wall_x/trainer/trainer_utils.py new file mode 100644 index 0000000..eab59f0 --- /dev/null +++ b/wall_x/trainer/trainer_utils.py @@ -0,0 +1,574 @@ +import logging +import math +import os +import random +import threading + +import numpy as np +import psutil +import torch +import torch.nn.functional as F +from safetensors.torch import load_file +from torch.optim.lr_scheduler import LambdaLR +from transformers import AutoProcessor + +from wall_x.model.core.action.normalizer import Normalizer +from wall_x.model.qact.tokenizer_mixin import get_action_tokenizer_mixin +from wall_x.utils.metrics import dtw_distance, frechet_distance, get_action_accuracy + +_logger = logging.getLogger(__name__) + + +def load_wallx_processors( + config, + normalizer=None, + action_statistic_dof=None, + device: str = "cpu", +): + """ + Load Wall-X processors, including tokenizer and action mapper. + + Args: + config: Configuration dictionary. + device: Tokenizer device. Training usually uses "cpu"; inference + usually uses "cuda". + + Returns: + Dictionary containing: + - processor: HuggingFace processor + - train_action_tokenizer: action tokenizer for training + - val_action_tokenizer: action tokenizer for validation + - action_mapper: action mapper dictionary + - num_added_tokens: number of added tokens + - tokenizer_mixin: ActionTokenizerMixin instance + """ + processor = AutoProcessor.from_pretrained(config["processor_path"], use_fast=True) + # pad side = left + processor.tokenizer.padding_side = "left" + + new_tokens = ["<|propri|>", "<|action|>"] + if config.get("new_special_tokens", None) is not None: + new_tokens.extend(config.get("new_special_tokens")) + + action_tokenizer_type = config.get("action_tokenizer_type", None) + + train_action_tokenizer = None + val_action_tokenizer = None + action_mapper = None + tokenizer_mixin = None + + if action_tokenizer_type: + # Use tokenizer_mixin as the single action-tokenizer interface. + action_tokenizer_config = config.get("action_tokenizer", {}) + # Backward compatibility: read top-level keys as fallback values. + action_tokenizer_config.setdefault( + "action_tokenizer_path", config.get("action_tokenizer_path") + ) + action_tokenizer_config.setdefault( + "action_tokenizer_checkpoint_path", + config.get("action_tokenizer_checkpoint_path"), + ) + action_tokenizer_config.setdefault( + "action_tokenizer_config_dir", config.get("action_tokenizer_config_dir") + ) + # Pass action_horizon_ar to the tokenizer for DLLM. + data_config = config.get("data", {}) + action_tokenizer_config.setdefault( + "action_horizon_ar", data_config.get("action_horizon_ar", 32) + ) + # Fall back to dof_config when ar_dof_config is not provided. + ar_dof_config = config.get("ar_dof_config") or config.get("dof_config") + assert ar_dof_config is not None, "Missing ar_dof_config and dof_config" + if normalizer is None: + if action_statistic_dof is None: + raise ValueError( + "Action tokenizer setup requires an explicit normalizer or " + "action statistics. Public Wall-X builds do not bundle " + "default action statistics." + ) + ar_normalizer = Normalizer(action_statistic_dof, ar_dof_config) + else: + ar_normalizer = normalizer + tokenizer_mixin = get_action_tokenizer_mixin(action_tokenizer_type) + tokenizer_mixin.load_tokenizer( + action_tokenizer_config, ar_normalizer, device=device + ) + + # Collect special tokens. + _new_tokens, special_tokens = tokenizer_mixin.get_all_special_tokens() + new_tokens += _new_tokens + + # Add tokens to the vocabulary. + num_added_tokens = processor.tokenizer.add_tokens(new_tokens) + + # Set placeholder_seq for discrete diffusion. + if special_tokens and action_tokenizer_config.get( + "input_placeholder_flag", False + ): + processor.placeholder_seq = [ + processor.tokenizer.convert_tokens_to_ids(token) + for token in special_tokens + ] + + # Backward compatibility: use the first added token when + # <|action_token_0|> does not exist. + ar_first_token_id = processor.tokenizer.convert_tokens_to_ids( + "<|action_token_0|>" + ) + if ( + ar_first_token_id is None + or ar_first_token_id == processor.tokenizer.unk_token_id + ): + ar_first_token_id = processor.tokenizer.convert_tokens_to_ids( + _new_tokens[0] + ) + processor.ar_first_token = ar_first_token_id + + # Build action_mapper. + action_mapper = tokenizer_mixin.build_action_mapper(processor) + + # Fetch the underlying tokenizer. + train_action_tokenizer = tokenizer_mixin.tokenizer + # Fast tokenizers need a separate validation instance; others share one. + val_action_tokenizer = tokenizer_mixin.get_val_tokenizer(config) + else: + num_added_tokens = processor.tokenizer.add_tokens(new_tokens) + + return { + "processor": processor, + "train_action_tokenizer": train_action_tokenizer, + "val_action_tokenizer": val_action_tokenizer, + "action_mapper": action_mapper, + "num_added_tokens": num_added_tokens, + "tokenizer_mixin": tokenizer_mixin, + } + + +def load_wallx_processors_from_cfg( + cfg, + normalizer=None, + action_statistic_dof=None, + device: str = "cpu", +): + """Typed convenience wrapper around ``load_wallx_processors``. + + Builds the flat dict that the legacy function expects from typed + TrainConfig sub-configs, then delegates. Callers using TrainConfig + can use this directly instead of hand-flattening. + """ + import dataclasses + + flat = dataclasses.asdict(cfg.model) + flat["model_type"] = cfg.model_type + flat["data"] = dict(cfg._raw_data or {}) + flat["dof_config"] = cfg.task.dof_config + flat["agent_pos_config"] = cfg.task.agent_pos_config + if cfg.task.ar_dof_config is not None: + flat["ar_dof_config"] = cfg.task.ar_dof_config + flat["batch_size_per_gpu"] = cfg.hyperparams.batch_size_per_gpu + return load_wallx_processors( + flat, + normalizer=normalizer, + action_statistic_dof=action_statistic_dof, + device=device, + ) + + +def load_qwen_pretrain_weight(model, pretrain_weight_path): + weight_files = sorted( + [f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")] + ) + # Initialize empty dictionary to store merged weights + merged_weights = {} + + # Load and merge each file sequentially + for weight_file in weight_files: + file_path = os.path.join(pretrain_weight_path, weight_file) + weights = load_file(file_path) + merged_weights.update(weights) + + renamed_weights = model.rename_vlm_weights_for_vla(merged_weights) + renamed_weights = { + k: v + for k, v in renamed_weights.items() + if "action_preprocessor.normalizer_" not in k + } # remove normalizer weights + if ( + model.config.model_type == "qwen2_5_vl" + and model.model.embed_tokens.weight.shape[0] + != renamed_weights["model.embed_tokens.weight"].shape[0] + ): + _logger.info( + "resize_token_embeddings from %d to %d", + model.model.embed_tokens.weight.shape[0], + renamed_weights["model.embed_tokens.weight"].shape[0], + ) + model.model.resize_token_embeddings( + renamed_weights["model.embed_tokens.weight"].shape[0] + ) + + err = model.load_state_dict(renamed_weights, strict=False) + + return model, err + + +def update_model_config(train_config, model_config): + model_config.use_state_string_representation = train_config["data"].get( + "use_state_string_representation", False + ) + model_config.ar_loss_weight = train_config.get("ar_loss_weight", 1.0) + + model_config.dof_config = train_config["dof_config"] + model_config.agent_pos_config = train_config["agent_pos_config"] + + model_config.action_horizon_flow = train_config["data"].get( + "action_horizon_flow", 32 + ) + + if train_config.get("_attn_implementation", None) is not None: + model_config._attn_implementation = train_config["_attn_implementation"] + + if train_config.get("attn_deterministic", None) is not None: + model_config.attn_deterministic = train_config["attn_deterministic"] + model_config.vision_config.attn_deterministic = train_config[ + "attn_deterministic" + ] + _logger.info("Attention is using deterministic kernel for this run") + else: + model_config.attn_deterministic = True + model_config.vision_config.attn_deterministic = True + + if train_config.get("noise_scheduler", None) is not None: + model_config.noise_scheduler = train_config["noise_scheduler"] + + return model_config + + +def update_data_config(config): + """Keep the top-level model type aligned with the nested data config.""" + config["data"]["model_type"] = config.get("model_type") + + if config.get("use_state_string_representation", None) is not None: + config["data"]["use_state_string_representation"] = config[ + "use_state_string_representation" + ] + + return config + + +def get_detailed_memory_usage(): + """Return process memory and thread-count diagnostics.""" + process = psutil.Process() + memory_info = process.memory_info() + current_threads = threading.active_count() + return { + "rss": f"{memory_info.rss / 1024 / 1024:.2f}MB ", + "vms": f"{memory_info.vms / 1024 / 1024:.2f}MB ", + "threads_count": current_threads, + } + + +def is_last_rank(): + return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1) + + +def print_rank_last(message): + """If distributed is initialized, log only on last rank.""" + if torch.distributed.is_initialized(): + if is_last_rank(): + _logger.info(message) + else: + _logger.info(message) + + +def seed_all(seed): + np.random.seed(seed) + random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ["PYTHONHASHSEED"] = str(seed) + + +def save_text_results_to_file( + current_step, all_input_texts, all_gt_texts, all_pred_texts, save_dir +): + """Save text-generation results to a local JSON file.""" + import json + + os.makedirs(save_dir, exist_ok=True) + + results = [] + for i, (input_text, gt_text, pred_text) in enumerate( + zip(all_input_texts, all_gt_texts, all_pred_texts) + ): + # Some callers pass one-item lists instead of plain strings. + input_clean = input_text[0] if isinstance(input_text, list) else input_text + gt_clean = gt_text[0] if isinstance(gt_text, list) else gt_text + pred_clean = pred_text[0] if isinstance(pred_text, list) else pred_text + + results.append( + { + "sample_id": i, + "input": input_clean, + "ground_truth": gt_clean, + "prediction": pred_clean, + "step": current_step, + } + ) + + filename = os.path.join(save_dir, f"text_predictions_step_{current_step}.json") + with open(filename, "w", encoding="utf-8") as f: + json.dump(results, f, ensure_ascii=False, indent=2) + + +# Add accuracy metrics. +# Add ade, fde, dtw, and frechet metrics from Ryan. +def compute_action_metrics(all_preds, all_actions, config, step_log={}): + metrics_settings = config.get("metrics_settings", "default") + metrics_available = ["l1", "mse", "accuracy", "ade", "fde", "dtw", "frechet"] + metrics_default = ["l1", "mse", "accuracy"] + if metrics_settings == "all": + metrics_enabled = metrics_available + elif metrics_settings == "default": + metrics_enabled = metrics_default + elif isinstance(metrics_settings, str): + metrics_enabled = [ + m.strip() + for m in metrics_settings.split(" ") + if m.strip() in metrics_available + ] + elif isinstance(metrics_settings, list): + metrics_enabled = [m for m in metrics_settings if m in metrics_available] + else: + metrics_enabled = metrics_default + _logger.warning( + 'Unknown metrics_settings "%s". Using default metrics: %s', + metrics_settings, + metrics_enabled, + ) + + overall_l1 = F.l1_loss(all_preds, all_actions) + overall_mse = F.mse_loss(all_preds, all_actions) + + step_log["val_action_l1"] = overall_l1.item() + step_log["val_action_mse"] = overall_mse.item() + + if "accuracy" in metrics_enabled: + accuracy_thresholds = [0.05, 0.1, 0.2, 0.4] + # The accuracy that all predicted action dimensions are within a certain range of the ground truth. + accuracies = get_action_accuracy( + all_preds, all_actions, thresholds=accuracy_thresholds + ) + for th_idx, threshold in enumerate(accuracy_thresholds): + step_log[f"val_action_acc_thr{threshold}"] = accuracies[th_idx].item() + + start_idx = 0 + dof_config = config["dof_config"] + for dof_key, dof_dim in dof_config.items(): + end_idx = start_idx + dof_dim + # all_preds.shape = (B, T, action_dim) + pred_dof = all_preds[..., start_idx:end_idx] + action_dof = all_actions[..., start_idx:end_idx] + dof_l1 = F.l1_loss(pred_dof, action_dof) + dof_mse = F.mse_loss(pred_dof, action_dof) + step_log[f"val_l1/{dof_key}"] = dof_l1.item() + step_log[f"val_mse/{dof_key}"] = dof_mse.item() + + if "accuracy" in metrics_enabled: + accuracies = get_action_accuracy( + action_dof, pred_dof, thresholds=accuracy_thresholds + ) + for th_idx, threshold in enumerate(accuracy_thresholds): + step_log[f"val_acc/{dof_key}_thr{threshold}"] = accuracies[ + th_idx + ].item() + + if "ee_cartesian_pos" in dof_key: + if "ade" in metrics_enabled: + displacement_error = torch.norm(pred_dof - action_dof, dim=-1) + ade = torch.mean(displacement_error) + step_log[f"val_ade/{dof_key}"] = ade.item() + + if "fde" in metrics_enabled: + final_pred = pred_dof[:, -1, :] + final_gt = action_dof[:, -1, :] + fde = torch.mean(torch.norm(final_pred - final_gt, dim=-1)) + step_log[f"val_fde/{dof_key}"] = fde.item() + + if "dtw" in metrics_enabled or "frechet" in metrics_enabled: + dtw_distances = [] + frechet_distances = [] + batch_size = pred_dof.shape[0] + for i in range(batch_size): + pred_seq = pred_dof[i] # shape: (T, D) + gt_seq = action_dof[i] # shape: (T, D) + dtw_dist = dtw_distance(pred_seq, gt_seq) + dtw_distances.append(dtw_dist) + + frechet_dist = frechet_distance(pred_seq, gt_seq) + frechet_distances.append(frechet_dist) + + avg_dtw = torch.mean(torch.stack(dtw_distances)) + if "dtw" in metrics_enabled: + step_log[f"val_dtw/{dof_key}"] = avg_dtw.item() + + avg_frechet = torch.mean(torch.stack(frechet_distances)) + if "frechet" in metrics_enabled: + step_log[f"val_frechet/{dof_key}"] = avg_frechet.item() + + start_idx = end_idx + + return step_log + + +def get_warmup_cosine_schedule( + optimizer, + num_warmup_steps: int, + num_training_steps: int, + peak_lr: float = None, + end_lr: float = None, + last_epoch: int = -1, +): + """ + Create a schedule with linear warmup followed by cosine decay: + - Warmup: linearly increases from peak_lr/(warmup_steps+1) to peak_lr + - Decay: cosine decay from peak_lr to end_lr + + Args: + optimizer: The optimizer for which to schedule the learning rate. + num_warmup_steps: The number of steps for the warmup phase. + num_training_steps: The total number of training steps. + peak_lr: The peak learning rate. If None, uses optimizer's initial lr. + end_lr: The minimum learning rate at the end. If None, defaults to peak_lr * 0.1. + last_epoch: The index of the last epoch when resuming training. + + Return: + torch.optim.lr_scheduler.LambdaLR with the appropriate schedule. + """ + if peak_lr is None: + peak_lr = optimizer.defaults["lr"] + if end_lr is None: + end_lr = peak_lr * 0.1 + + def lr_lambda(current_step: int): + if current_step < num_warmup_steps: + # Start from peak_lr / (warmup_steps + 1). + init_lr = peak_lr / (num_warmup_steps + 1) + current_lr = init_lr + (peak_lr - init_lr) * current_step / num_warmup_steps + return current_lr / peak_lr # LambdaLR multiplies by base_lr + else: + # Cosine decay + decay_steps = num_training_steps - num_warmup_steps + progress = min(1.0, (current_step - num_warmup_steps) / max(1, decay_steps)) + cos = 0.5 * (1 + math.cos(math.pi * progress)) + current_lr = end_lr + (peak_lr - end_lr) * cos + return current_lr / peak_lr + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def plot_openloop( + action_pred_list, + action_gt_list, + l1_loss, + episode_index, + save_path, + is_static_list=None, +): + """ + Plot openloop action comparison visualization. + + Args: + action_pred_list: List of predicted actions, each with shape (horizon, action_dim) + action_gt_list: List of ground truth actions, each with shape (horizon, action_dim) + l1_loss: L1 loss array with shape (total_frames, action_dim) + episode_index: Index of the episode being visualized + save_path: Directory path to save the plot + is_static_list: Optional list of booleans indicating static frames + """ + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + assert len(action_pred_list) == len( + action_gt_list + ), "Predicted action and ground truth action must have the same shape." + + dim = action_pred_list[0].shape[1] + plt.figure(figsize=(12, 4 * dim)) + + for i in range(dim): + plt.subplot(dim, 1, i + 1) + + has_labeled_static = False + for j in range(len(action_gt_list)): + gt_action = action_gt_list[j] + predict_action = action_pred_list[j] + + x_vals_gt = np.linspace(j, j + 1, len(gt_action)) + x_vals_pred = np.linspace(j, j + 1, len(predict_action)) + + if is_static_list is not None and is_static_list[j]: + label = None + if not has_labeled_static: + label = "Static GT" + has_labeled_static = True + plt.axvspan(j, j + 1, color="gray", alpha=0.2, label=label) + + if j == 0: + plt.plot( + x_vals_gt, + gt_action[:, i], + label="Ground Truth", + color="blue", + linewidth=2, + linestyle="-", + marker="o", + markersize=3, + ) + plt.plot( + x_vals_pred, + predict_action[:, i], + label="Model Output", + color="orange", + linewidth=2, + linestyle="--", + marker="x", + markersize=4, + ) + else: + plt.plot( + x_vals_gt, + gt_action[:, i], + color="blue", + linewidth=2, + linestyle="-", + marker="o", + markersize=3, + ) + plt.plot( + x_vals_pred, + predict_action[:, i], + color="orange", + linewidth=2, + linestyle="--", + marker="x", + markersize=4, + ) + + plt.title(f"Action Dimension {i + 1}, L1 Loss: {l1_loss[:, i].mean():.6f}") + plt.xlabel("Number of Chunk") + plt.ylabel("Action Value") + plt.legend() + + plt.suptitle( + f"Openloop Action Comparison for Episode {episode_index}, L1 Loss: {l1_loss.mean():.6f}" + ) + plt.tight_layout(rect=[0, 0, 1, 0.98]) + os.makedirs(save_path, exist_ok=True) + plt.savefig(f"{save_path}/{episode_index}.png") + plt.close() + _logger.info("Saved openloop plot to %s/%s.png", save_path, episode_index) diff --git a/wall_x/trainer/utils/__init__.py b/wall_x/trainer/utils/__init__.py new file mode 100644 index 0000000..d74f475 --- /dev/null +++ b/wall_x/trainer/utils/__init__.py @@ -0,0 +1,6 @@ +"""Trainer utility helpers.""" + +from .data import move_batch_to_device +from .diagnostics import log_gpu_memory + +__all__ = ["move_batch_to_device", "log_gpu_memory"] diff --git a/wall_x/trainer/utils/data.py b/wall_x/trainer/utils/data.py new file mode 100644 index 0000000..8ca974c --- /dev/null +++ b/wall_x/trainer/utils/data.py @@ -0,0 +1,46 @@ +"""Data-related utility functions used by the trainer main loop.""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def move_batch_to_device( + batch: Any, + device: torch.device, + *, + non_blocking: bool = True, +) -> Any: + """Move every tensor in ``batch`` to ``device``, recursing into dict/list. + + Returns a new structure with the same shape; the input ``batch`` is not + mutated. dict / list containers are rebuilt; tensors are moved via + ``.to(device, non_blocking=...)``; everything else is passed through + by reference. + + Parameters + ---------- + batch : Any + Typically a dict produced by the dataloader, but recursion accepts + dict / list / tensor / arbitrary leaves. + device : torch.device + Target device, typically ``self.device`` on the trainer. + non_blocking : bool + Whether to use pinned-memory async copy. Default True because the + trainer uses pinned loaders; pass False if the dataloader hasn't + pinned memory. + """ + if isinstance(batch, dict): + return { + k: move_batch_to_device(v, device, non_blocking=non_blocking) + for k, v in batch.items() + } + if isinstance(batch, list): + return [ + move_batch_to_device(v, device, non_blocking=non_blocking) for v in batch + ] + if isinstance(batch, torch.Tensor): + return batch.to(device, non_blocking=non_blocking) + return batch diff --git a/wall_x/trainer/utils/diagnostics.py b/wall_x/trainer/utils/diagnostics.py new file mode 100644 index 0000000..b8356fb --- /dev/null +++ b/wall_x/trainer/utils/diagnostics.py @@ -0,0 +1,40 @@ +"""Training-run diagnostic helpers (CUDA memory, etc.).""" + +from __future__ import annotations + +from typing import Callable, Optional + +import torch + + +def log_gpu_memory( + device: torch.device, + rank: int, + *, + stage: str = "", + log_fn: Optional[Callable] = None, +) -> None: + """Log per-rank GPU memory (allocated / reserved / total) via ``log_fn``. + + Calls ``torch.cuda.synchronize`` on ``device`` so the numbers reflect + the actual post-op usage, not pending work. If ``log_fn`` is None this + is a no-op. + """ + if log_fn is None: + return + torch.cuda.synchronize() + allocated = torch.cuda.memory_allocated(device) / 1024**3 + reserved = torch.cuda.memory_reserved(device) / 1024**3 + peak_allocated = torch.cuda.max_memory_allocated(device) / 1024**3 + peak_reserved = torch.cuda.max_memory_reserved(device) / 1024**3 + total = torch.cuda.get_device_properties(device).total_memory / 1024**3 + tag = f"[{stage}] " if stage else "" + log_fn( + f"{tag}GPU memory rank{rank} " + f"| allocated {allocated:.2f} GiB" + f" | reserved {reserved:.2f} GiB" + f" | peak_allocated {peak_allocated:.2f} GiB" + f" | peak_reserved {peak_reserved:.2f} GiB" + f" | total {total:.2f} GiB", + main_process_only=False, + ) diff --git a/wall_x/utils/constant.py b/wall_x/utils/constant.py index 3074591..4c1c71f 100644 --- a/wall_x/utils/constant.py +++ b/wall_x/utils/constant.py @@ -1,362 +1,58 @@ -action_statistic_dof = { - "x2_normal": { - # water flowers - "follow_left_arm_joint_cur": { - "min": [-3.7121], - "delta": [7.6008], - }, - "follow_right_arm_joint_cur": { - "min": [-3.6176], - "delta": [8.5015], - }, - "follow_left_ee_cartesian_pos": { - "min": [-0.036, -0.3241, -0.1245], - "delta": [0.4389, 0.557, 0.479], - }, - "follow_left_ee_rotation": { - "min": [-1.2373, -0.1929, -1.5182], - "delta": [2.2009, 1.5669, 2.0936], - }, - "follow_left_gripper": {"min": [-0.1196], "delta": [4.5226]}, - "follow_right_ee_cartesian_pos": { - "min": [-0.0326, -0.2273, -0.1377], - "delta": [0.4574, 0.5704, 0.4743], - }, - "follow_right_ee_rotation": { - "min": [-1.2201, -0.2611, -0.7427], - "delta": [2.6623, 1.6622, 2.4186], - }, - "follow_right_gripper": {"min": [-0.1208], "delta": [4.5261]}, - "height": {"min": [-0.0001], "delta": [0.5051]}, - "head_actions": {"min": [-1.5000, -1.4167], "delta": [2.5000, 1.8879]}, - "base_velocity": { - "min": [-0.0359, -0.084, -0.0162], - "delta": [0.1539, 0.1848, 0.0322], - }, - }, - "DobbE": { - "follow_right_ee_cartesian_pos": { - "min": [-0.6107, -0.3272, -0.4282], - "delta": [1.2629, 1.5297, 0.8349], - }, - "follow_right_ee_rotation": { - "min": [-1.7378, -1.4597, -1.8712], - "delta": [2.7031, 2.8182, 3.5921], - }, - "follow_right_gripper": {"min": [0.0], "delta": [0.9983]}, - }, - "RH20T": { - "follow_right_ee_cartesian_pos": { - "min": [0.3646, -0.2722, 0.0066], - "delta": [0.3813, 0.5973, 0.3277], - }, - "follow_right_ee_rotation": { - "min": [-1.8716, -0.4398, -3.1414], - "delta": [3.4145, 1.0225, 6.2828], - }, - "follow_right_gripper": {"min": [0.0], "delta": [95.0]}, - }, - "agibotworld_alpha": { - "follow_left_ee_cartesian_pos": { - "min": [0.4954, 0.0166, 0.1729], - "delta": [0.3336, 0.5123, 0.9189], - }, - "follow_left_ee_rotation": { - "min": [-3.1064, -1.2629, -3.1238], - "delta": [6.2127, 2.5923, 6.2496], - }, - "follow_left_gripper": {"min": [34.6222], "delta": [86.1921]}, - "follow_right_ee_cartesian_pos": { - "min": [0.4615, -0.5975, 0.1638], - "delta": [0.3823, 0.5577, 0.8873], - }, - "follow_right_ee_rotation": { - "min": [-3.0891, -1.0739, -2.5091], - "delta": [6.1707, 2.3074, 3.8533], - }, - "follow_right_gripper": {"min": [34.6222], "delta": [85.7635]}, - "height": {"min": [0.0], "delta": [0.4535]}, - "head_actions": {"min": [-0.1746, 0.0523], "delta": [0.2444, 0.4713]}, - }, - "austin_buds": { - "follow_right_ee_cartesian_pos": { - "min": [0.3496, -0.2855, 0.0105], - "delta": [0.3748, 0.492, 0.3116], - }, - "follow_right_ee_rotation": { - "min": [-3.1405, -0.151, -0.0737], - "delta": [6.2813, 0.3218, 0.1536], - }, - "follow_right_gripper": {"min": [0.0076], "delta": [0.0724]}, - }, - "austin_sailor": { - "follow_right_ee_cartesian_pos": { - "min": [0.387, -0.3165, 0.0244], - "delta": [0.2999, 0.5252, 0.2308], - }, - "follow_right_ee_rotation": { - "min": [-3.1402, -0.1618, -1.5918], - "delta": [6.2804, 0.337, 2.9478], - }, - "follow_right_gripper": {"min": [0.0005], "delta": [0.0773]}, - }, - "austin_sirius": { - "follow_right_ee_cartesian_pos": { - "min": [0.0, -0.1182, 0.0], - "delta": [0.5329, 0.3812, 0.2723], - }, - "follow_right_ee_rotation": { - "min": [-3.1407, -0.1243, -1.7434], - "delta": [6.2823, 0.1975, 1.8073], - }, - "follow_right_gripper": {"min": [0.0334], "delta": [0.046]}, - }, - "bc_z": { - "follow_right_ee_cartesian_pos": { - "min": [-0.3883, -0.1116, 0.6113], - "delta": [0.7199, 0.4288, 0.3709], - }, - "follow_right_ee_rotation": { - "min": [-1.056, -1.0587, -2.6295], - "delta": [1.9142, 1.9455, 4.8064], - }, - "follow_right_gripper": {"min": [0.2], "delta": [0.8]}, - }, - "berkeley_autolab_ur5": { - "follow_right_ee_cartesian_pos": { - "min": [0.3018, -0.2129, -0.1888], - "delta": [0.3121, 0.52, 0.3107], - }, - "follow_right_ee_rotation": { - "min": [-3.1396, -0.2278, 1.1413], - "delta": [6.279, 0.454, 0.9841], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "berkeley_cable_routing": { - "follow_right_ee_cartesian_pos": { - "min": [0.4617, -0.28, 0.03], - "delta": [0.1838, 0.5665, 0.1272], - }, - "follow_right_ee_rotation": { - "min": [-3.1413, -0.0299, -0.7665], - "delta": [6.2826, 0.0692, 3.322], - }, - }, - "berkeley_fanuc_manipulation": { - "follow_right_ee_cartesian_pos": { - "min": [0.3718, -0.4072, 0.0184], - "delta": [0.3483, 0.7201, 0.5229], - }, - "follow_right_ee_rotation": { - "min": [-3.1399, -1.0166, -1.6988], - "delta": [6.2802, 1.4498, 3.2074], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "bridge_data_v2": { - "follow_right_ee_cartesian_pos": { - "min": [0.1498, -0.2178, -0.0901], - "delta": [0.3012, 0.469, 0.298], - }, - "follow_right_ee_rotation": { - "min": [-0.3279, -0.6105, -1.0578], - "delta": [0.7378, 1.0353, 2.2552], - }, - "follow_right_gripper": {"min": [0.0692], "delta": [0.9426]}, - }, - "dlr_edan_shared_control": { - "follow_right_ee_cartesian_pos": { - "min": [-0.8387, 0.1473, -0.3934], - "delta": [0.6579, 0.6025, 1.1566], - }, - "follow_right_ee_rotation": { - "min": [-3.1217, -1.5197, -2.2516], - "delta": [6.2505, 1.5594, 4.2831], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "droid": { - "follow_right_ee_cartesian_pos": { - "min": [0.2667, -0.4396, -0.0472], - "delta": [0.5159, 0.8806, 0.8331], - }, - "follow_right_ee_rotation": { - "min": [-3.1374, -1.216, -2.1741], - "delta": [6.2749, 2.1075, 4.2259], - }, - "follow_right_gripper": {"min": [0.0], "delta": [0.9912]}, - }, - "fmb": { - "follow_right_ee_cartesian_pos": { - "min": [0.3554, -0.2844, 0.0354], - "delta": [0.336, 0.4961, 0.2943], - }, - "follow_right_ee_rotation": { - "min": [-3.1404, -0.9302, -0.0599], - "delta": [6.2807, 1.724, 1.8284], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "fractal": { - "follow_right_ee_cartesian_pos": { - "min": [0.3242, -0.2836, 0.1405], - "delta": [0.5518, 0.4963, 0.9328], - }, - "follow_right_ee_rotation": { - "min": [-3.1308, -0.2421, -2.9685], - "delta": [6.2609, 1.7343, 5.819], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "furniture_bench": { - "follow_right_ee_cartesian_pos": { - "min": [0.3691, -0.181, 0.0058], - "delta": [0.2962, 0.3582, 0.1775], - }, - "follow_right_ee_rotation": { - "min": [-3.1394, -0.6121, -1.9958], - "delta": [6.2786, 1.6114, 3.7748], - }, - "follow_right_gripper": {"min": [0.0035], "delta": [0.0762]}, - }, - "jaco_play": { - "follow_right_ee_cartesian_pos": { - "min": [-0.3787, -0.6294, 0.1682], - "delta": [0.5898, 0.3587, 0.2183], - }, - "follow_right_ee_rotation": { - "min": [0.9792, -0.0668, -0.0498], - "delta": [0.0175, 0.1277, 0.0686], - }, - "follow_right_gripper": {"min": [0.0791], "delta": [0.1033]}, - }, - "nyu_rot": { - "follow_right_ee_cartesian_pos": { - "min": [0.25, -1.0, -0.2], - "delta": [0.75, 2.0, 1.2], - }, - "follow_right_ee_rotation": { - "min": [-3.1416, -3.1416, 6.2831], - "delta": [9.4248, 4.1416, 0.0], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "stanford_hydra": { - "follow_right_ee_cartesian_pos": { - "min": [0.2068, -0.274, 0.1317], - "delta": [0.4929, 0.4981, 0.4588], - }, - "follow_right_ee_rotation": { - "min": [-3.1321, -0.7496, -3.0269], - "delta": [6.2658, 1.5176, 5.8261], - }, - "follow_right_gripper": {"min": [0.0], "delta": [0.0811]}, - }, - "stanford_kuka_multimodal": { - "follow_right_ee_cartesian_pos": { - "min": [0.4781, -0.0659, 0.3424], - "delta": [0.0868, 0.0864, 0.1863], - }, - "follow_right_ee_rotation": { - "min": [-3.136, -0.0521, -3.1413], - "delta": [6.2727, 0.1109, 6.2825], - }, - "follow_right_gripper": {"min": [-0.4713], "delta": [0.9485]}, - }, - "taco_play": { - "follow_right_ee_cartesian_pos": { - "min": [0.1375, -0.4291, 0.2052], - "delta": [0.5327, 1.0237, 0.3913], - }, - "follow_right_ee_rotation": { - "min": [-3.1391, -0.6946, -1.2808], - "delta": [6.2784, 0.8196, 3.0856], - }, - "follow_right_gripper": {"min": [0.0001], "delta": [0.0806]}, - }, - "utaustin_mutex": { - "follow_right_ee_cartesian_pos": { - "min": [0.3213, -0.4734, 0.0141], - "delta": [0.2108, 0.8471, 0.5644], - }, - "follow_right_ee_rotation": { - "min": [-3.1404, -0.2202, -1.5489], - "delta": [6.2805, 0.582, 1.9282], - }, - "follow_right_gripper": {"min": [0.0019], "delta": [0.0738]}, - }, - "viola": { - "follow_right_ee_cartesian_pos": { - "min": [0.4011, -0.2521, 0.0103], - "delta": [0.2444, 0.4305, 0.4355], - }, - "follow_right_ee_rotation": { - "min": [-3.1403, -0.2737, -1.8626], - "delta": [6.2804, 0.4901, 2.0618], - }, - "follow_right_gripper": {"min": [0.0002], "delta": [0.0773]}, - }, - "kuka": { - "follow_right_ee_cartesian_pos": { - "min": [0.3914, -0.4901, 0.0175], - "delta": [0.3339, 0.8357, 0.9064], - }, - "follow_right_ee_rotation": { - "min": [-3.1416, -0.9903, -3.1416], - "delta": [6.2832, 2.2421, 6.2832], - }, - "follow_right_gripper": { - "min": [0.0000], - "delta": [1.0000], - }, - }, - "UMI-biarm": { - "follow_left_ee_cartesian_pos": { - "min": [-0.2917, -0.4926, 0.0063], - "delta": [0.9028, 0.8168, 0.3473], - }, - "follow_left_ee_rotation": { - "min": [-2.5309, -1.5706, -1.3309], - "delta": [0.8758, 2.3315, 1.7076], - }, - "follow_left_gripper": { - "min": [0.0029], - "delta": [0.0812], - }, - "follow_right_ee_cartesian_pos": { - "min": [-0.0023, -0.5191, -0.0358], - "delta": [0.7474, 0.8668, 0.351], - }, - "follow_right_ee_rotation": { - "min": [-2.4945, -2.0149, -0.8088], - "delta": [1.0941, 3.2628, 2.0018], - }, - "follow_right_gripper": { - "min": [0.0019], - "delta": [0.0814], - }, - }, - "agibotworld_beta": { - "follow_left_ee_cartesian_pos": { - "min": [0.4954, 0.0166, 0.1729], - "delta": [0.3336, 0.5123, 0.9189], - }, - "follow_left_ee_rotation": { - "min": [-3.1064, -1.2629, -3.1238], - "delta": [6.2127, 2.5923, 6.2496], - }, - "follow_left_gripper": {"min": [34.6222], "delta": [86.1921]}, - "follow_right_ee_cartesian_pos": { - "min": [0.4615, -0.5975, 0.1638], - "delta": [0.3823, 0.5577, 0.8873], - }, - "follow_right_ee_rotation": { - "min": [-3.0891, -1.0739, -2.5091], - "delta": [6.1707, 2.3074, 3.8533], - }, - "follow_right_gripper": {"min": [34.6222], "delta": [85.7635]}, - "height": {"min": [0.0], "delta": [0.4535]}, - "head_actions": {"min": [-0.1746, 0.0523], "delta": [0.2444, 0.4713]}, - }, -} +"""OSS-safe compatibility constants. + +The internal repository keeps canonical robot, dataset, action-key, and action +statistics tables in ``x2robot_utils.constants``. Those tables are intentionally +not bundled in the public export. This shim keeps the public VLA code importable +without exposing private dataset catalogues or default normalization stats. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +ACTION_KEY_FULL_MAPPING: dict[str, str] = {} +_ACTION_KEY_FULL_MAPPING = ACTION_KEY_FULL_MAPPING + +ACTION_DATASET_NAMES: tuple[str, ...] = () +MULTIMODAL_DATASET_NAMES: tuple[str, ...] = () +_ACTION_DATASET_NAMES = ACTION_DATASET_NAMES +_MULTIMODAL_DATASET_NAMES = MULTIMODAL_DATASET_NAMES + +VIEW_SLOT_KEYS = ["view1", "view2", "view3"] +PHYSICAL_VIEW_IDS: dict[str, int] = {} + + +def get_view_slot_keys(num_views: int) -> list[str]: + """Return canonical view slot names for ``num_views`` cameras.""" + return [f"view{i + 1}" for i in range(max(0, int(num_views)))] + + +def is_multimodal_dataset_name(dataset_name: str | None) -> bool: + """Return whether a dataset is marked multimodal in the public shim.""" + return dataset_name in MULTIMODAL_DATASET_NAMES + + +def is_action_dataset_name(dataset_name: str | None) -> bool: + """Return whether a dataset row should be treated as an action sample.""" + return dataset_name is not None and not is_multimodal_dataset_name(dataset_name) + + +def iter_action_dataset_names(dataset_names: Iterable[str]) -> list[str]: + """Filter a sequence down to action dataset names.""" + return [name for name in dataset_names if is_action_dataset_name(name)] + + +__all__ = [ + "ACTION_KEY_FULL_MAPPING", + "_ACTION_KEY_FULL_MAPPING", + "ACTION_DATASET_NAMES", + "MULTIMODAL_DATASET_NAMES", + "_ACTION_DATASET_NAMES", + "_MULTIMODAL_DATASET_NAMES", + "VIEW_SLOT_KEYS", + "PHYSICAL_VIEW_IDS", + "get_view_slot_keys", + "is_multimodal_dataset_name", + "is_action_dataset_name", + "iter_action_dataset_names", +] diff --git a/wall_x/utils/cudagraph_wrapper.py b/wall_x/utils/cudagraph_wrapper.py new file mode 100644 index 0000000..6064e17 --- /dev/null +++ b/wall_x/utils/cudagraph_wrapper.py @@ -0,0 +1,278 @@ +import os + +import torch + +ENABLE_CUDA_GRAPH = os.environ.get("ENABLE_CUDA_GRAPH", "True").lower() == "true" +# When using a single master buffer + multi-bucket graph, max batch must be pre-allocated; +# otherwise resize/reallocation after capture invalidates addresses recorded by the graph. +CUDA_GRAPH_MAX_BS = int(os.environ.get("CUDA_GRAPH_MAX_BS", "128")) +_SHAPE_GUARD_FASTPATH = ( + os.environ.get("CUDAGRAPH_DISABLE_SHAPE_GUARD_FASTPATH", "0") != "1" +) + + +class CUDAGraph_Wrapper: + """ + Bucketed CUDA Graph wrapper: + - When batch size changes, select/reuse a CUDAGraph by bucket_size + - For batch_size < bucket_size: copy inputs into static buffer [:B], return outputs [:B] after replay + """ + + def __init__( + self, + model, + warm_up_times: int = 3, + enable: bool | None = None, + batch_size_key: str | None = None, + ): + self.model = model + self.warm_up_times = warm_up_times + self._enable = ENABLE_CUDA_GRAPH if enable is None else bool(enable) + self._batch_size_key = batch_size_key or "suffix_inputs_embeds" + + # graph cache: bucket_bs -> CUDAGraph + self.graphs: dict[int, torch.cuda.CUDAGraph] = {} + # Per bucket, store views into the master buffer (graph capture depends on these views' address/shape/stride) + self.static_inputs_map: dict[int, dict[str, torch.Tensor]] = {} + self.static_output_tensor: dict[int, torch.Tensor] = {} + self.graph_pool = None + + # Master buffers (shared across all buckets) + self.graph_vars: dict[str, torch.Tensor] = {} + self._max_bs: int | None = None + + # Record non-batch dimension shape signature (globally consistent) to avoid silent errors + self._shape_signature: dict[str, tuple] = {} + # Batch dimension index per input (default 0; e.g. suffix_position_ids batch is at dim=1) + self._batch_dim: dict[str, int] = {} + self._shape_guard_verified: dict[int, bool] = {} + + def _get_batch_size(self, **kwargs) -> int: + key = self._batch_size_key + if key in kwargs: + v = kwargs[key] + if isinstance(v, torch.Tensor) and v.dim() >= 1: + return int(v.shape[0]) + raise ValueError(f"CUDAGraph_Wrapper.forward requires {key} argument") + + def _select_bucket_bs(self, bs: int) -> int: + """ + Bucket selection: + - 1, 2, 4, 8 use pow2 buckets + - >8 align to 16 (16, 32, 48, ...) + """ + if bs <= 1: + return 1 + if bs <= 2: + return 2 + if bs <= 4: + return 4 + if bs <= 8: + return 8 + return int(((bs + 15) // 16) * 16) + + def _batch_dim_for_key(self, key: str, tensor: torch.Tensor) -> int: + # Minimal special-case handling for known inputs only + # suffix_position_ids: shape [3, B, T], batch at dim=1 + if key == "suffix_position_ids" and tensor.dim() >= 2: + return 1 + return 0 + + def _view_for_bucket(self, tensor: torch.Tensor, bucket_bs: int, batch_dim: int): + if tensor.dim() == 0: + return tensor + if batch_dim == 0: + return tensor[:bucket_bs] + if batch_dim == 1: + return tensor[:, :bucket_bs] + raise ValueError(f"Unsupported batch_dim={batch_dim} for cudagraph wrapper") + + def _copy_and_pad( + self, + master: torch.Tensor, + tensor: torch.Tensor, + bs: int, + bucket_bs: int, + batch_dim: int, + ): + if tensor.dim() == 0: + master.copy_(tensor) + return + if batch_dim == 0: + master[:bs].copy_(tensor) + if bs < bucket_bs: + master[bs:bucket_bs].zero_() + return + if batch_dim == 1: + master[:, :bs].copy_(tensor) + if bs < bucket_bs: + master[:, bs:bucket_bs].zero_() + return + raise ValueError(f"Unsupported batch_dim={batch_dim} for cudagraph wrapper") + + def _allocate_master_tensor( + self, key: str, tensor: torch.Tensor, max_bs: int, bs: int + ) -> torch.Tensor: + """ + Allocate a master buffer for one input (single large buffer): + - Extend only the batch dimension to max_bs; other dims unchanged + - Copy the current bs valid range first; zero the rest (avoid stale values on replay) + """ + if tensor.dim() == 0: + return tensor.clone() + batch_dim = self._batch_dim_for_key(key, tensor) + shape = list(tensor.shape) + if batch_dim >= len(shape): + raise ValueError( + f"Invalid batch_dim={batch_dim} for key={key}, tensor.shape={tuple(tensor.shape)}" + ) + shape[batch_dim] = max_bs + buf = torch.zeros(tuple(shape), device=tensor.device, dtype=tensor.dtype) + # Initial copy: pad to current bucket (at least cover the valid bs range) + self._copy_and_pad(buf, tensor, bs=bs, bucket_bs=bs, batch_dim=batch_dim) + return buf + + def _ensure_master_buffers(self, bs: int, **kwargs): + """ + Ensure master buffers are allocated once. + - max_bs from CUDA_GRAPH_MAX_BS; if unset, use the first bucket_bs as max_bs + and disallow exceeding it later (resize would invalidate captured graphs). + """ + if self._max_bs is not None: + return + + bucket_bs = self._select_bucket_bs(bs) + max_bs = CUDA_GRAPH_MAX_BS if CUDA_GRAPH_MAX_BS > 0 else bucket_bs + self._max_bs = int(max_bs) + + for k, v in kwargs.items(): + if not isinstance(v, torch.Tensor): + raise TypeError( + f"CUDAGraph_Wrapper only supports Tensor kwargs, got {k}={type(v)}" + ) + batch_dim = self._batch_dim_for_key(k, v) if v.dim() > 0 else 0 + self._batch_dim[k] = batch_dim + # Record non-batch dimensions (globally consistent) + if v.dim() == 0: + self._shape_signature[k] = tuple() + else: + sig = list(v.shape) + sig.pop(batch_dim) + self._shape_signature[k] = tuple(sig) + self.graph_vars[k] = self._allocate_master_tensor( + k, v, max_bs=self._max_bs, bs=bs + ) + + def _initialize_bucket(self, bucket_bs: int, bs: int, **kwargs): + assert self._max_bs is not None + self._shape_guard_verified[bucket_bs] = False + if bucket_bs > self._max_bs: + raise ValueError( + f"CUDAGraph bucket_bs({bucket_bs}) exceeds master max_bs({self._max_bs}). " + f"Set CUDA_GRAPH_MAX_BS >= {bucket_bs}." + ) + + print(f"[CUDAGraph] Initializing bucket_bs={bucket_bs} (current bs={bs}) ...") + + # Build view map for this bucket (fixed address/shape/stride) + static_map: dict[str, torch.Tensor] = {} + for k, v in kwargs.items(): + # Shape guard: non-batch dims must match + expected = self._shape_signature.get(k) + if v.dim() == 0: + got = tuple() + else: + batch_dim = self._batch_dim.get(k, self._batch_dim_for_key(k, v)) + got_list = list(v.shape) + got_list.pop(batch_dim) + got = tuple(got_list) + if expected is not None and got != expected: + raise ValueError( + f"CUDAGraph bucket init: non-batch dim changed: key={k}, expected={expected}, got={got}" + ) + master = self.graph_vars[k] + batch_dim = self._batch_dim.get(k, 0) + static_map[k] = self._view_for_bucket( + master, bucket_bs=bucket_bs, batch_dim=batch_dim + ) + + # Warmup (stabilize kernels/caches) + out = None + for _ in range(self.warm_up_times): + out = self.forward_naive(**static_map) + assert isinstance(out, torch.Tensor) + + # Output master buffer (allocate once) + if "_outputs" not in self.graph_vars: + out_shape = (self._max_bs,) + tuple(out.shape[1:]) + self.graph_vars["_outputs"] = torch.empty( + out_shape, device=out.device, dtype=out.dtype + ) + + self.static_inputs_map[bucket_bs] = static_map + self.static_output_tensor[bucket_bs] = self.graph_vars["_outputs"][:bucket_bs] + + # capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self.graph_pool): + self.static_output_tensor[bucket_bs].copy_(self.forward_naive(**static_map)) + if self.graph_pool is None: + self.graph_pool = graph.pool() + self.graphs[bucket_bs] = graph + torch.cuda.synchronize() + print(f"[CUDAGraph] Bucket {bucket_bs} captured.") + + def forward_naive(self, **kwargs): + return self.model(**kwargs) + + def forward(self, **kwargs): + if not self._enable: + return self.forward_naive(**kwargs) + bs = self._get_batch_size(**kwargs) + bucket_bs = self._select_bucket_bs(bs) + + # Initialize master buffers (first call) + self._ensure_master_buffers(bs=bs, **kwargs) + assert self._max_bs is not None + if bucket_bs > self._max_bs: + # Do not exceed master max_bs at runtime (resize would invalidate captured graphs) + return self.forward_naive(**kwargs) + + # lazy capture for this bucket + if bucket_bs not in self.graphs: + self._initialize_bucket(bucket_bs=bucket_bs, bs=bs, **kwargs) + + # Shape guard: non-batch dim change is unsafe for cudagraph; fallback + # fast-path: skip guard if already verified for this bucket (stable-state optimization) + if not ( + _SHAPE_GUARD_FASTPATH and self._shape_guard_verified.get(bucket_bs, False) + ): + for k, v in kwargs.items(): + if not isinstance(v, torch.Tensor): + raise TypeError( + f"CUDAGraph_Wrapper only supports Tensor kwargs, got {k}={type(v)}" + ) + expected = self._shape_signature.get(k) + if v.dim() == 0: + got = tuple() + else: + batch_dim = self._batch_dim.get(k, self._batch_dim_for_key(k, v)) + got_list = list(v.shape) + got_list.pop(batch_dim) + got = tuple(got_list) + if expected is not None and got != expected: + return self.forward_naive(**kwargs) + # Mark verified for fast-path on subsequent calls + self._shape_guard_verified[bucket_bs] = True + + # Copy inputs into master buffer (per batch_dim, zero padding) + for k, tensor in kwargs.items(): + master = self.graph_vars[k] + batch_dim = self._batch_dim.get(k, 0) + self._copy_and_pad( + master, tensor, bs=bs, bucket_bs=bucket_bs, batch_dim=batch_dim + ) + + graph = self.graphs[bucket_bs] + graph.replay() + return self.graph_vars["_outputs"][:bs] diff --git a/wall_x/utils/logger.py b/wall_x/utils/logger.py new file mode 100644 index 0000000..c7d2afa --- /dev/null +++ b/wall_x/utils/logger.py @@ -0,0 +1,95 @@ +"""Rank-aware text logger for distributed training.""" + +from __future__ import annotations + +import logging +import os +import sys +from pathlib import Path +from typing import Optional + +from torch.distributed import get_rank, is_initialized + + +class DistributedLogger: + def __init__( + self, + name: str = "wallx", + save_path: Optional[str] = None, + level: int = logging.INFO, + ): + if is_initialized(): + self._rank = get_rank() + else: + self._rank = int(os.environ.get("RANK", 0)) + logging.warning( + "DistributedLogger created before init_process_group; " + "falling back to RANK env var (rank=%d).", + self._rank, + ) + logger = logging.getLogger(f"{name}.rank{self._rank}") + logger.setLevel(level) + logger.propagate = False + for h in logger.handlers[:]: + h.close() + logger.removeHandler(h) + + fmt = logging.Formatter( + f"%(asctime)s - [rank{self._rank}] - %(levelname)s - %(message)s" + ) + + # All ranks write a file log (if a save_path was given). + if save_path: + log_dir = Path(save_path) / "logs" + log_dir.mkdir(parents=True, exist_ok=True) + file_handler = logging.FileHandler(log_dir / f"rank_{self._rank}.log") + file_handler.setFormatter(fmt) + logger.addHandler(file_handler) + + # Only rank 0 writes to stdout. + if self._rank == 0: + stream_handler = logging.StreamHandler(sys.stdout) + stream_handler.setFormatter(fmt) + logger.addHandler(stream_handler) + + self._logger = logger + + # Thin pass-throughs - callers use standard logging verbs. + def info(self, msg, *args, **kwargs): + self._logger.info(msg, *args, **kwargs) + + def warning(self, msg, *args, **kwargs): + self._logger.warning(msg, *args, **kwargs) + + def error(self, msg, *args, **kwargs): + self._logger.error(msg, *args, **kwargs) + + def debug(self, msg, *args, **kwargs): + self._logger.debug(msg, *args, **kwargs) + + @property + def rank(self) -> int: + return self._rank + + # --- backward-compatible shim ---------------------------------------- + # Some legacy call sites still invoke `.log(msg, level=..., + # main_process_only=...)`. Forward those to the standard logger so we + # don't break them during the migration; new code should use .info / + # .warning / .error / .debug directly. + def log( + self, + message, + level: int = logging.INFO, + main_process_only: bool = False, + ): + if main_process_only and self._rank != 0: + return + self._logger.log(level, message) + + # Pytorch's ``accelerate``-style fallback for code that still passes an + # accelerator object to the old constructor; accept and ignore it. + # Older codepaths can be migrated incrementally. + @classmethod + def legacy(cls, name: str, level: int = logging.INFO, accelerator=None): + del accelerator # ignored + return cls(name=name, save_path=None, level=level) diff --git a/wall_x/utils/metrics.py b/wall_x/utils/metrics.py new file mode 100644 index 0000000..e2ffb0f --- /dev/null +++ b/wall_x/utils/metrics.py @@ -0,0 +1,95 @@ +from typing import List + +import torch + + +def get_action_accuracy( + gt: torch.FloatTensor, # [Batch_Size, Horizon, Action_Dim] + pred: torch.FloatTensor, + thresholds: List[float] = [0.1, 0.2], +) -> torch.FloatTensor: + device = gt.device + diff = torch.abs(gt - pred).reshape(-1, gt.shape[-1]) + + # get the percentage of diff lower than threshold for all action dimensions + accuracies = torch.zeros(len(thresholds), device=device) + for idx, threshold in enumerate(thresholds): + accuracy = torch.mean( + (torch.mean((diff < threshold).float(), dim=1) >= 1.0).float() + ) + accuracies[idx] = accuracy + return accuracies + + +def dtw_distance(seq1, seq2): + """ + Compute the Dynamic Time Warping distance between two sequences. + + ``seq1`` and ``seq2`` must have shape ``(T, D)``, where ``T`` is the + number of time steps and ``D`` is the feature dimension. + """ + n, m = seq1.shape[0], seq2.shape[0] + + seq1_d, seq2_d = seq1.double(), seq2.double() + dtw_matrix = torch.full( + (n + 1, m + 1), float("inf"), dtype=torch.float64, device=seq1.device + ) + dtw_matrix[0, 0] = 0 + + for i in range(1, n + 1): + for j in range(1, m + 1): + cost = torch.dist(seq1_d[i - 1], seq2_d[j - 1]) + dtw_matrix[i, j] = cost + torch.min( + torch.stack( + [ + dtw_matrix[i - 1, j], + dtw_matrix[i, j - 1], + dtw_matrix[i - 1, j - 1], + ] + ) + ) + + return dtw_matrix[n, m] + + +def frechet_distance(seq1, seq2): + """ + Compute the discrete Frechet distance between two sequences. + + ``seq1`` and ``seq2`` must have shape ``(T, D)``, where ``T`` is the + number of time steps and ``D`` is the feature dimension. + """ + n, m = seq1.shape[0], seq2.shape[0] + + seq1_d, seq2_d = seq1.double(), seq2.double() + frechet_matrix = torch.full( + (n, m), float("inf"), dtype=torch.float64, device=seq1.device + ) + frechet_matrix[0, 0] = torch.dist(seq1_d[0], seq2_d[0]) + + for j in range(1, m): + frechet_matrix[0, j] = torch.max( + frechet_matrix[0, j - 1], torch.dist(seq1_d[0], seq2_d[j]) + ) + + for i in range(1, n): + frechet_matrix[i, 0] = torch.max( + frechet_matrix[i - 1, 0], torch.dist(seq1_d[i], seq2_d[0]) + ) + + for i in range(1, n): + for j in range(1, m): + frechet_matrix[i, j] = torch.max( + torch.min( + torch.stack( + [ + frechet_matrix[i - 1, j], + frechet_matrix[i, j - 1], + frechet_matrix[i - 1, j - 1], + ] + ) + ), + torch.dist(seq1_d[i], seq2_d[j]), + ) + + return frechet_matrix[n - 1, m - 1] diff --git a/wall_x/utils/timers.py b/wall_x/utils/timers.py index 4233d19..e7a05dc 100644 --- a/wall_x/utils/timers.py +++ b/wall_x/utils/timers.py @@ -1,12 +1,15 @@ +import logging +import os import time -from torch.cuda import nvtx from abc import ABC, abstractmethod +from contextlib import nullcontext +from functools import wraps from typing import List import torch -from functools import wraps -from contextlib import nullcontext -import os +from torch.cuda import nvtx + +logger = logging.getLogger(__name__) ENABLE_PERFORMANCE_TIMING = ( os.environ.get("ENABLE_PERFORMANCE_TIMING", "True").lower() == "true" @@ -32,23 +35,14 @@ class ScopeTimerContext: torch.cuda.synchronize() end_time = time.perf_counter() cost_ms = (end_time - self.start_time) * 1e3 - print(f"\033[92m{self.msg} took {cost_ms:.3f} ms to execute\033[0m") + logger.info("%s took %.3f ms to execute", self.msg, cost_ms) ScopeTimer = ScopeTimerContext if ENABLE_PERFORMANCE_TIMING else nullcontext def timer(func, msg=None): - """ - Decorator to measure function execution time. - - Args: - func: Function to be timed - - Returns: - Wrapped function with timing functionality - """ - + """Decorator to measure function execution time.""" if msg is None: msg = func.__name__ else: @@ -58,44 +52,36 @@ def timer(func, msg=None): def wrapper(*args, **kwargs): with ScopeTimer(msg): result = func(*args, **kwargs) - return result return wrapper -# Helper functions to check for distributed environment def _is_distributed(): - """Checks if the current environment is set up for distributed training.""" return torch.distributed.is_available() and torch.distributed.is_initialized() def _get_world_size(): - """Safely retrieves the world size (number of processes).""" if _is_distributed(): return torch.distributed.get_world_size() return 1 def _get_rank(): - """Safely retrieves the rank of the current process.""" if _is_distributed(): return torch.distributed.get_rank() return 0 def _barrier(group=None): - """Safely executes a distributed barrier to synchronize processes.""" if _is_distributed(): torch.distributed.barrier(group=group) -# Dynamically set the all_gather function if torch.distributed.is_available(): try: dist_all_gather_func = torch.distributed.all_gather_into_tensor except AttributeError: - # Fallback to standard all_gather if all_gather_into_tensor is missing dist_all_gather_func = torch.distributed.all_gather else: dist_all_gather_func = None @@ -109,51 +95,35 @@ class TimerBase(ABC): @abstractmethod def start(self, barrier=False): - """Start the timer. - - Args: - barrier (bool, optional): Synchronizes ranks before starting. Defaults to False. - """ + """Start the timer, optionally syncing all ranks with a barrier first.""" pass @abstractmethod def stop(self, barrier=False): - """Stop the timer. - - Args: - barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False. - """ + """Stop the timer, optionally syncing all ranks with a barrier first.""" pass @abstractmethod def reset(self): - """Reset timer.""" + """Reset accumulated elapsed time to zero.""" pass @abstractmethod def elapsed(self, reset=True, barrier=False): - """Calculates the elapsed time and restarts timer. - - Args: - reset (bool, optional): Resets timer before restarting. Defaults to True. - barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False. - - Returns: - float: Elapsed time. - """ + """Return accumulated elapsed time in seconds; reset if reset=True.""" pass class DummyTimer(TimerBase): - """Dummy Timer.""" + """Dummy Timer - no-op placeholder used when log level exceeds threshold.""" def __init__(self): super().__init__("dummy timer") - def start(self, barrier=False, nvtx_push=False): + def start(self, barrier=False, nvtx_push=False, sync=False, **kwargs): return - def stop(self, barrier=False, nvtx_pop=False): + def stop(self, barrier=False, sync=False, **kwargs): return def reset(self): @@ -166,9 +136,6 @@ class DummyTimer(TimerBase): ) def active_time(self): - """Returns the cumulative duration the timer has been active. - Note: Not supported for DummyTimer. - """ raise Exception( "active timer should not be used to calculate elapsed time, " "check if timer's log_level <= self._log_level." @@ -188,34 +155,18 @@ class Timer(TimerBase): """ def __init__(self, name): - """Initialize Timer. - - Args: - name (str): Name of the timer. - """ super().__init__(name) self._elapsed = 0.0 self._active_time = 0.0 self._started = False - # Note that None will default to the global process group self._barrier_group = None self._start_time = time.time() self.nvtx = False def set_barrier_group(self, barrier_group): - """Sets barrier group. - - Args: - barrier_group (ProcessGroup): Torch ProcessGroup for barrier. - """ self._barrier_group = barrier_group def start(self, barrier=False, nvtx_push=False, sync=False): - """Start the timer. - - Args: - barrier (bool, optional): Synchronizes ranks before starting. Defaults to False. - """ assert not self._started, "timer has already been started" if barrier: _barrier(group=self._barrier_group) @@ -228,11 +179,6 @@ class Timer(TimerBase): self.nvtx = True def stop(self, barrier=False, sync=False): - """Stop the timer. - - Args: - barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False. - """ if self.nvtx: nvtx.range_pop() assert self._started, "timer is not started" @@ -246,37 +192,21 @@ class Timer(TimerBase): self._started = False def reset(self): - """Reset timer.""" - # Don't reset _active_time self._elapsed = 0.0 self._started = False def elapsed(self, reset=True, barrier=False): - """Calculates the elapsed time and restarts timer. - - Args: - reset (bool, optional): Resets timer before restarting. Defaults to True. - barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False. - - Returns: - float: Elapsed time. - """ _started = self._started - # If the timing in progress, end it first. if self._started: self.stop(barrier=barrier) - # Get the elapsed time. _elapsed = self._elapsed - # Reset the elapsed time if reset: self.reset() - # If timing was in progress, set it back. if _started: self.start(barrier=barrier) return _elapsed def active_time(self): - """Calculates the cumulative duration for which the timer has been active""" return self._active_time @@ -284,13 +214,6 @@ class Timers: """Class for a group of Timers.""" def __init__(self, log_level, log_option): - """Initialize group of timers. - - Args: - log_level (int): Log level to control what timers are enabled. - log_option (str): Setting for logging statistics over ranks for all the timers. - Allowed: ['max', 'minmax', 'all']. - """ self._log_level = log_level allowed_log_options = set(["max", "minmax", "all"]) assert ( @@ -305,9 +228,6 @@ class Timers: self._max_log_level = 2 def __call__(self, name, log_level=None): - """Call timer with name and log level.""" - # If the timer has already been set, then check if the log-level - # is provided, it matches the one that the timer was created with. if name in self._timers: if log_level is not None: assert log_level == self._log_levels[name], ( @@ -317,8 +237,6 @@ class Timers: ) ) return self._timers[name] - # If timer does not exist and no log level is provided, - # set it to the max log level which is 2. if log_level is None: log_level = self._max_log_level assert ( @@ -326,38 +244,19 @@ class Timers: ), "log level {} is larger than max supported log level {}".format( log_level, self._max_log_level ) - # Now if the input log level is larger than the one set for - # the timers class, just ignore it and return a dummy timer. if log_level > self._log_level: return self._dummy_timer - # Otherwise, initalize the timer and set the level. self._timers[name] = Timer(name) self._log_levels[name] = log_level return self._timers[name] def _get_elapsed_time_all_ranks(self, names, reset, barrier): - """Returns elapsed times of timers in names. - - For single-node/single-GPU cases, directly returns the time for the current rank. - For distributed cases, maintains the existing all_gather logic. - - Args: - names (List[str]): list of timer names - reset (bool): reset the timer after recording the elapsed time - barrier (bool): if set, do a global barrier before time measurements - - Returns: - torch.tensor: Tensor of size [world_size, len(names)] with times in float. - """ - - # First make sure all the callers are in sync. if barrier: _barrier() world_size = _get_world_size() rank = _get_rank() - # Create device tensor if torch.cuda.is_available(): device = torch.cuda.current_device() else: @@ -367,33 +266,26 @@ class Timers: (world_size, len(names)), dtype=torch.float, device=device ) - # Fill timing data for the current rank for i, name in enumerate(names): if name in self._timers: rank_name_to_time[rank, i] = self._timers[name].elapsed(reset=reset) - # Return directly for single-node; perform all_gather for distributed setup if world_size > 1 and _is_distributed() and dist_all_gather_func is not None: try: dist_all_gather_func( rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1) ) except Exception as e: - # If all_gather fails, print a warning and proceed with single rank timing - print(f"Warning: all_gather failed: {e}. Using single rank timing.") + logger.warning("all_gather failed: %s. Using single rank timing.", e) return rank_name_to_time def _get_global_min_max_time(self, names, reset, barrier, normalizer): - """Report only min and max times across all ranks.""" - rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier) name_to_min_max_time = {} for i, name in enumerate(names): rank_to_time = rank_name_to_time[:, i] - # filter out the ones we did not have any timings for rank_to_time = rank_to_time[rank_to_time > 0.0] - # If the timer exists: if rank_to_time.numel() > 0: name_to_min_max_time[name] = ( rank_to_time.min().item() / normalizer, @@ -404,7 +296,6 @@ class Timers: def _get_global_min_max_time_string( self, names, reset, barrier, normalizer, max_only ): - """Report strings for max/minmax times across all ranks.""" name_to_min_max_time = self._get_global_min_max_time( names, reset, barrier, normalizer ) @@ -413,17 +304,13 @@ class Timers: world_size = _get_world_size() if world_size == 1: - # Simplified output for single-node setup output_string = "time (ms):" for name in name_to_min_max_time: - _, max_time = name_to_min_max_time[ - name - ] # min and max are identical for a single rank + _, max_time = name_to_min_max_time[name] output_string += "\n {}: {:.2f}".format( (name + " ").ljust(48, "."), max_time ) else: - # Maintain original output format for multi-node setup if max_only: output_string = "max time across ranks (ms):" else: @@ -441,7 +328,6 @@ class Timers: return output_string def _get_all_ranks_time_string(self, names, reset, barrier, normalizer): - """Report times across all ranks.""" rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier) world_size = _get_world_size() @@ -474,32 +360,20 @@ class Timers: reset: bool = True, barrier: bool = False, ): - """Returns the output string with logged timer values according to configured options. + """Return a formatted timing string for the given timer names. Args: - names (List[str]): Names of the timers to log. If None, all registered timers are - fetched. Defaults to None. - normalizer (float, optional): Normalizes the timer values by the factor. - Defaults to 1.0. - reset (bool, optional): Whether to reset timer values after logging. Defaults to True. - barrier (bool, optional): Whether to do a global barrier before time measurments. - Defaults to False. - - Raises: - Exception: Raises if log option is invalid. - - Returns: - str: Formatted string with the timer values. + names: Timers to include; defaults to all registered timers. + normalizer: Divide raw seconds by this value (e.g. 1000 for ms output). + reset: Reset each timer after reading its elapsed time. + barrier: Synchronize across ranks before gathering times. """ - - if names is None: # get all registered timers + if names is None: names = list(self._timers.keys()) assert normalizer > 0.0 if self._log_option in ["max", "minmax"]: - max_only = False - if self._log_option == "max": - max_only = True + max_only = self._log_option == "max" output_string = self._get_global_min_max_time_string( names, reset, barrier, normalizer / 1000.0, max_only ) @@ -519,30 +393,23 @@ class Timers: reset: bool = True, barrier: bool = False, ): - """logs the timers passed in names to stdout. Example usage is to log average per step - value for timer 'foo', this function can be called with normalizer factor set to logging - interval. + """Print timing results for the given names to stdout on one rank. Args: - names (List[str]): Names of the timers to log. - rank (int, optional): logs the timers to a specific rank. If set to None, logs to the - last rank. Defaults to None. - normalizer (float, optional): Normalizes the timer values by the factor. - Defaults to 1.0. - reset (bool, optional): Whether to reset timer values after logging. Defaults to True. - barrier (bool, optional): Whether to do a global barrier before time measurments. - Defaults to False. + names: Timer names to log. + rank: Rank that prints; defaults to the last rank (world_size - 1). + normalizer: Divide raw seconds by this value before printing. + reset: Reset each timer after reading. + barrier: Synchronize across ranks first. """ - output_string = self.get_all_timers_string(names, normalizer, reset, barrier) - # If no input rank is provided, log on last rank. world_size = _get_world_size() current_rank = _get_rank() if rank is None: rank = world_size - 1 if rank == current_rank and output_string is not None: - print(output_string, flush=True) + logger.info("%s", output_string) def write( self, @@ -553,22 +420,16 @@ class Timers: reset: bool = True, barrier: bool = False, ): - """Write timers to a tensorboard writer. - Note that we only report maximum time across ranks to tensorboard. + """Write per-timer max times as TensorBoard scalars. Args: - names (List[str]): Names of the timers to log. - writer (SummaryWriter): Tensorboard SummaryWriter object - iteration (int): Current iteration. - normalizer (float, optional): Normalizes the timer values by the factor. - Defaults to 1.0. - reset (bool, optional): Whether to reset timer values after logging. Defaults to True. - barrier (bool, optional): Whether to do a global barrier before time measurments. - Defaults to False. + names: Timer names to write. + writer: TensorBoard SummaryWriter instance. + iteration: Global step value for the scalar. + normalizer: Divide raw seconds by this value. + reset: Reset each timer after reading. + barrier: Synchronize across ranks first. """ - # currently when using add_scalars, - # torch.utils.add_scalars makes each timer its own run, which - # polutes the runs list, so we just add each as a scalar assert normalizer > 0.0 name_to_min_max_time = self._get_global_min_max_time( names, reset, barrier, normalizer diff --git a/workspace/README.md b/workspace/README.md index b6a9677..e4072ca 100644 --- a/workspace/README.md +++ b/workspace/README.md @@ -1,160 +1,387 @@ -# Training Guide +# Wall-X-OSS Usage Guide -This document explains the key configuration parameters and memory requirements for Wall-X training. +This guide explains how to fine-tune, evaluate in simulation, and deploy on real robots with the **Wall-OSS-0.5** pretrained model on LeRobot-format datasets. -## Quick Start Checklist +> All commands below assume you are in the **repository root** (`wall-x/`). -### 🚀 **Step 1: Prepare Model** -Choose one of our pretrained models: -- **WALL-OSS-FLOW**: https://huggingface.co/x-square-robot/wall-oss-flow -- **WALL-OSS-FAST**: https://huggingface.co/x-square-robot/wall-oss-fast -Or from Qwen-2.5-VL -- Download https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct, settings refer to `config_qact_from_vlm.yml` +> **Note:** This open-source release targets **Wall-OSS-0.5**. If you are using **Wall-OSS-FLOW** or **Wall-OSS-FAST** instead, switch back to the previous codebase version: +> +> ```bash +> git checkout 97406f2ab5de414c79b091873f946c112d105c72 +> ``` -### ⚙️ **Step 2: Configure Environment** -- Update `run.sh`: Set `code_dir` and `config_path` to your actual paths -- Set `CUDA_VISIBLE_DEVICES` for your available GPUs +--- -### 📝 **Step 3: Update Configuration Files** -- Replace all `/path/to/` placeholders in `config_qact.yml` with actual paths -- Configure robot settings: `dof_config` and `agent_pos_config` -- Set dataset: Choose appropriate `repo_id` -- Adjust `batch_size_per_gpu` based on your GPU memory +## Environment Setup -### ▶️ **Step 4: Start Training** ```bash -bash ./workspace/lerobot_example/run.sh +conda create --name wallx python=3.10 +conda activate wallx + +pip install -r requirements.txt +pip install "dmuon @ git+https://github.com/X-Square-Robot/dmuon.git" + +git clone https://github.com/huggingface/lerobot.git +cd lerobot +git checkout c66cd401767e60baece16e1cf68da2824227e076 +pip install --no-deps -e . +cd - + +# Optional: only needed for LIBERO simulator evaluation. +pip install -r requirements-libero.txt +mkdir -p third_party +git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git third_party/LIBERO + +# Install wall-x (from repo root) +MAX_JOBS=8 pip install --no-build-isolation -e . ``` -## Enable FAST tokenizer -To fine-tune using the FAST tokenizer, please download the repository and update the `action_tokenizer_path`. Make sure to set `use_fast_tokenizer` to `true` and q01 and q99 to normalize the dataset, refer to `wall-x/scripts/compute_norm_stats.py`: +--- + +## Part 1: Download Wall-OSS-0.5 Weights + +Wall-OSS-0.5 is a VLA foundation model built on Qwen2.5-VL-3B. Fine-tuning requires both the **pretrained weights** and the **VLM processor**. + +### 1.1 Download Wall-OSS-0.5 + ```bash -git clone https://huggingface.co/physical-intelligence/fast +# Option A: huggingface-cli +huggingface-cli download X-Square-Robot/wall-oss-0.5 \ + --local-dir /path/to/wall-oss-0.5 + +# Option B: Python +python -c " +from huggingface_hub import snapshot_download +snapshot_download('X-Square-Robot/wall-oss-0.5', local_dir='/path/to/wall-oss-0.5') +" ``` -## Required Paths (Must Modify) +After download, the directory should contain at least: + +| File | Purpose | +|------|---------| +| `config.json` | Model architecture; maps to `model.config_path` in training YAML | +| `model.safetensors` | Pretrained weights; maps to `checkpoint.resume_from` | +| tokenizer / processor files | Used at inference time | + +HuggingFace: + +### 1.2 Download Qwen2.5-VL-3B-Instruct (processor) + +Set `model.processor_path` and `model.pretrained_path` to the Qwen2.5-VL processor directory: + +```bash +huggingface-cli download Qwen/Qwen2.5-VL-3B-Instruct \ + --local-dir /path/to/Qwen2.5-VL-3B-Instruct +``` + +--- + +## Part 2: Simulation Tasks (LIBERO) + +### 2.1 Download the LIBERO Dataset + +Download the LeRobot-format LIBERO dataset from HuggingFace: + +```bash +huggingface-cli download lerobot/libero \ + --repo-type dataset \ + --local-dir /path/to/libero_all +``` + +Dataset page: + +### 2.2 Edit Your Training Config (`config.yml`) + +Copy an example config and replace every `/path/to/*` placeholder: + +```bash +cp workspace/example/libero.yml /path/to/my_libero_config.yml +``` + +**Required paths:** + ```yaml -pretrained_wallx_path: "/path/to/wallx_model/" # Path to pretrained wallx model -save_path: "/path/to/workspace/" # Path to save training outputs -use_fast_tokenizer: False # True: train FAST, False: train Flow -action_tokenizer_path: "/path/to/fast/" # Must set if use_fast_tokenizer is True -norm_stats_path: "/path/to/stats/" # Must set for normalize dataset +model: + config_path: /path/to/wall-oss-0.5/config.json + processor_path: /path/to/Qwen2.5-VL-3B-Instruct + pretrained_path: /path/to/Qwen2.5-VL-3B-Instruct + +data: + lerobot_config: + repo_id: /path/to/libero_all # local LeRobot dataset root + norm_stats_path: /path/to/libero_all_norm_stats.json + key_mappings: # must match your dataset keys + camera: + observation.images.faceImg: face_view + observation.images.rightImg: right_wrist_view + state: observation.state + action: action + +checkpoint: + save_path: /path/to/libero_training_output + resume_from: /path/to/wall-oss-0.5/model.safetensors ``` -## Customize your robot configuration -Ensure that the sum of the configuration dimensions corresponds to the values specified in norm_stats.json, and that each key is unique. The maximum dimensionality is set to 20, consistent with our robot configuration. + +See `workspace/example/libero.yml` for the full example. LIBERO uses a 7-dim single-arm action; pad to 26 dims with `action_padding` to match the Wall-OSS-0.5 pretraining space (see comments in the YAML). + +### 2.3 Compute LIBERO Normalization Stats + +Generate `norm_stats.json` from the dataset before training: + +```bash +python scripts/compute_norm_stats.py \ + --train_config /path/to/my_libero_config.yml \ + --data_root /path/to/libero_all \ + --output_path /path/to/libero_all_norm_stats.json +``` + +Then set `data.norm_stats_path` to the generated JSON file. + +### 2.4 Start Training + +```bash +# Single GPU +CUDA_VISIBLE_DEVICES=0 \ + python wall_x/trainer/fsdp_trainer/train_fsdp.py \ + --config /path/to/my_libero_config.yml + +# Multi-GPU (recommended) +CUDA_VISIBLE_DEVICES=0,1,2,3 \ + torchrun --nproc_per_node=4 \ + wall_x/trainer/fsdp_trainer/train_fsdp.py \ + --config /path/to/my_libero_config.yml +``` + +Logs and checkpoints are written to `checkpoint.save_path`. If training saved FSDP-sharded checkpoints, merge them before inference: + +```bash +python scripts/merge_sharded_weights.py \ + /path/to/sharded_checkpoint \ + /path/to/merged_checkpoint +``` + +Single-GPU training needs at least **48 GB** VRAM. For multi-GPU runs, enable `distributed.use_fsdp: true`. + +### 2.5 Run Inference (LIBERO Simulation) + +`scripts/run_libero.sh` requires the optional LIBERO simulator stack: + +```bash +pip install -r requirements-libero.txt +mkdir -p third_party +git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git third_party/LIBERO +``` + +The script checks for LIBERO, robosuite, MuJoCo, PyOpenGL, BDDL, Gym, and h5py +before loading the model, so missing simulator dependencies fail fast with +installation instructions. + +Use `scripts/run_libero.sh` for batch evaluation in the LIBERO simulator: + +```bash +CHECKPOINT_PATH=/path/to/checkpoint \ +TRAIN_CONFIG_PATH=/path/to/my_libero_config.yml \ +TASK_SUITE_NAME=libero_spatial \ +NUM_TRIALS_PER_TASK=50 \ +bash scripts/run_libero.sh + +# Quick smoke test (1 trial per task) +SMOKE=1 CHECKPOINT_PATH=/path/to/checkpoint bash scripts/run_libero.sh +``` + +Useful environment variables: + +| Variable | Description | +|----------|-------------| +| `CHECKPOINT_PATH` | Checkpoint directory | +| `TRAIN_CONFIG_PATH` | Training YAML | +| `TASK_SUITE_NAME` | `libero_spatial` / `libero_object` / `libero_goal` / `libero_10` | +| `ALL_SUITES=1` | Run all 4 standard suites sequentially | +| `TASK_INDICES` | Task indices, e.g. `0,1,2` | +| `CUDA_ID` | GPU index | + +--- + +## Part 3: Real-Robot Tasks + +### 3.1 Prepare Your Real-Robot Training Dataset + +Convert your robot data to **LeRobot v3** format and note the local dataset root path (used for `data.lerobot_config.repo_id`). + +Reference config: `workspace/example/maniparena_example.yml` (dual-arm, 448px, 3 cameras). + +### 3.2 Edit Your Training Config (`config.yml`) + +```bash +cp workspace/example/maniparena_example.yml /path/to/my_robot_config.yml +``` + +**Required paths:** + ```yaml - customized_dof_config: - "action_eef": 6 - "action_gripper": 1 +model: + config_path: /path/to/wall-oss-0.5/config.json + processor_path: /path/to/Qwen2.5-VL-3B-Instruct + pretrained_path: /path/to/Qwen2.5-VL-3B-Instruct - customized_agent_pos_config: - "state_eef_with_gripper": 7 +data: + lerobot_config: + repo_id: /path/to/your_robot_dataset + norm_stats_path: /path/to/your_robot_norm_stats.json + key_mappings: # match your dataset camera / state / action keys + camera: + observation.images.faceImg: face_view + # ... + state: observation.state + action: action + +checkpoint: + save_path: /path/to/robot_training_output + resume_from: /path/to/wall-oss-0.5/model.safetensors ``` -## Using Lerobot Dataset -- Each dataset employs distinct keys; please specify the corresponding key mappings as described in `wall-x/wall_x/data/utils.py`. -```python -"lerobot/aloha_mobile_cabinet": { - "camera": { - "observation.images.cam_high": "face_view", - "observation.images.cam_left_wrist": "left_wrist_view", - "observation.images.cam_right_wrist": "right_wrist_view", - }, - "state": "observation.state", - "action": "action", - } -``` +`task.dof_config` defines **predicted action dimensions**; `task.agent_pos_config` defines **observation state dimensions**. The sum of each config must match `norm_stats.json`. If your robot has fewer than 26 DOF, pad with `action_padding` (see comments in the example YAML). + +### 3.3 Compute Dataset Normalization Stats -## Compute stats ```bash - python wall-x/scripts/compute_norm_stats.py +python scripts/compute_norm_stats.py \ + --train_config /path/to/my_robot_config.yml \ + --data_root /path/to/your_robot_dataset \ + --output_path /path/to/your_robot_norm_stats.json ``` -## Configuration Explain -- `agent_pos_config` corresponds to `obs_action_keys` and subsequently to state, while `dof_config` corresponds to `predict_action_keys` and subsequently to action. Note that the state and action may not necessarily share the same set of DoF. +### 3.4 Start Training -## Training Parameters (Commonly Modified) - -### Learning Rate Settings -- `learning_rate`: Initial learning rate (default: 0.00009) -- `min_lr`: Minimum learning rate for scheduler (default: 0.00005) -- `num_warmup_steps`: Number of warmup steps (default: 100) - -### Batch Size and Memory -- `batch_size_per_gpu`: Batch size per GPU - adjust based on GPU memory -- `gradient_accumulation_steps`: Gradient accumulation steps -- `num_training_steps`: Total training steps -- `num_epoch`: Number of training epochs - -### Training Optimization Settings -- `FSDP2`: Enable FSDP2 for distributed training (default: True) - **Recommended for multi-GPU** -- `torch_compile`: Enable PyTorch compilation optimization (default: False) - -**⚠️ Important Note on torch_compile:** -- **Benefits**: Enabling `torch_compile` can significantly improve training efficiency -- **Requirements**: Requires that the data input shape is always consistent throughout training -- **Caution**: If you don't have sufficient understanding of torch compile, please **DO NOT** enable it as it may cause unexpected issues with dynamic input shapes - -## Robot Configuration (Modify for Your Robot) - -### DOF Configuration -Modify `dof_config` to match your robot's action space: -- Add/remove action keys based on your robot's capabilities -- Ensure DOF numbers match your robot's action dimensions - -### Agent Position Configuration -Keep `agent_pos_config` consistent with `dof_config`. - -### Action Keys -- `obs_action_keys`: Actions used as observation context -- `predict_action_keys`: Actions to predict/control - -## Data Configuration - -### Dataset -- `repo_id`: LeRobot dataset identifier -- `train_test_split`: Training/validation split ratio (default: 0.95) -- `action_horizon`: Number of future actions to predict (default: 32) - -### Image Settings -- `resolution`: Image resolution for different camera views -- `download_videos`: Whether to download video files (true/false) - -## Resume Training (Optional) -- `resume.ckpt`: Path to checkpoint for resuming training -- `resume.load_ckpt_only`: Only load model weights, not optimizer state - -## Merge checkpoint -- If FSDP SHARDED_STATE_DICT is used, please run command below to merge checkpoint into a single safetensors ```bash - # refer to accelerate/commands/merge.py - accelerate merge-weights /path/to/sharded_tensors /path/to/model.safetensors - # copy the saved processor files - cp /path/to/saved_processor_dir/* /path/to/model.safetensors - - # In earlier versions of PyTorch, errors may occur. You can use our provided script to address this issue; refer to wall-x/scripts/merge_sharded_weights.py for details. +CUDA_VISIBLE_DEVICES=0,1,2,3 \ + torchrun --nproc_per_node=4 \ + wall_x/trainer/fsdp_trainer/train_fsdp.py \ + --config /path/to/my_robot_config.yml ``` -## Memory Usage +To resume training, point `checkpoint.resume_from` to a checkpoint **directory** (not a single `.safetensors` file). -Below are the memory consumption benchmarks for different training configurations using the `lerobot/aloha_mobile_cabinet` dataset: +### 3.5 Start Inference (WebSocket Server) -| Dataset | Batch Size | FSDP2 | Torch Compile | Num GPUs | Max Allocated Memory | -|---------|------------|--------|---------------|----------|---------------------| -| lerobot/aloha_mobile_cabinet | 1 | ❌ | ❌ | 1 | 40.11G | -| lerobot/aloha_mobile_cabinet | 1 | ❌ | ❌ | 8 | 48.02G | -| lerobot/aloha_mobile_cabinet | 1 | ✅ | ❌ | 2 | 43.70G | -| lerobot/aloha_mobile_cabinet | 1 | ✅ | ❌ | 8 | 24.96G | -| lerobot/aloha_mobile_cabinet | 1 | ✅ | ✅ | 8 | 24.21G | +Use `scripts/run_serving.sh` to launch the inference server for real-robot clients or open-loop evaluation: +```bash +bash scripts/run_serving.sh \ + --checkpoint-path /path/to/checkpoint \ + --train-config-path /path/to/my_robot_config.yml \ + --port 32195 +``` -**Hardware Recommendations:** +By default this wrapper returns raw model action chunks, which is the expected +mode for open-loop evaluation. Pass `--serialize-actions` if your client expects +robot-serialized actions. -- For single GPU training: Ensure at least 48GB VRAM (e.g., RTX 6000 Ada, A6000) -- For multi-GPU training: Enable FSDP2 for optimal memory distribution +Or invoke the serving module directly (adjust parameters as needed): -## Reproduce +```bash +export ENABLE_CUDA_GRAPH=True +export ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=True -Openloop plot `wall-x/workspace/lerobot_example/evaluation/lerobot_openloop.png` +CKPT_PATH=/path/to/checkpoint -To reproduce the results, use the config file wall-x/workspace/lerobot_example/config_qact_from_vlm.yml with a global batch size of 128, adjusted via `gradient_accumulation_steps` and numbers of gpu. +python -m wall_x._vendor.harrix.serving.launch_serving \ + --env X2ROBOT \ + --port 32195 \ + --no-serialize-actions \ + model-config:server-model-config \ + --model-config.checkpoint-path "$CKPT_PATH" \ + --model-config.train-config-path /path/to/my_robot_config.yml \ + --model-config.action-horizon 32 \ + --model-config.robot-action-interpolate-multiplier 1 \ + --model-config.robot-action-end-ratio 1.0 \ + --model-config.robot-type desktop +``` + +Clients connect at `ws://127.0.0.1:32195` to send observations and receive predicted actions. + +### 3.6 Plot Open-Loop Results to Verify the Server + +Open-loop evaluation compares model-predicted action trajectories against dataset ground truth **without executing actions or feeding back state**. + +**Terminal 1:** Keep the inference server from section3.5 running. + +**Terminal 2:** Run the open-loop plotting script: + +```bash +python scripts/draw_openloop_plot.py \ + --uri ws://127.0.0.1:32195 \ + --dataset-root /path/to/your_robot_dataset \ + --train-config /path/to/my_robot_config.yml \ + --episode-indices 0,1,2 \ + --save-dir ./openloop_plots +``` + +The script loads episodes from the LeRobot dataset, queries the WebSocket server frame by frame, and saves **predicted vs. ground-truth** comparison plots (PNG) under `--save-dir`. +`--dataset-root` and `--train-config` are both required; pass the same training +config used for the checkpoint so dataset layouts and action dimensions match. + +Common options: + +| Option | Description | +|--------|-------------| +| `--episode-indices` | Comma-separated episode indices to evaluate | +| `--start-ratio` | Start position as a fraction of episode length (0.0 = from the beginning) | +| `--stride` | Frames between inference calls (default: `action_horizon`) | +| `--max-inferences` | Maximum inference requests per episode | + +--- + +## Appendix + +### Example Config Files + +| File | Scenario | +|------|----------| +| `workspace/example/libero.yml` | LIBERO single-arm fine-tuning | +| `workspace/example/maniparena_example.yml` | Real-robot dual-arm fine-tuning | + +### Helper Scripts + +| Script | Purpose | +|--------|---------| +| `scripts/compute_norm_stats.py` | Compute dataset normalization statistics | +| `scripts/fake_inference.py` | Inference smoke test | +| `scripts/run_libero.sh` | Batch LIBERO simulation evaluation | +| `scripts/run_serving.sh` | Launch WebSocket inference server | +| `scripts/draw_openloop_plot.py` | Open-loop evaluation and plotting | +| `scripts/merge_sharded_weights.py` | Merge FSDP sharded checkpoints | + +See [`scripts/README.md`](../scripts/README.md) for more script details. + +### Training Hyperparameters + +| Parameter | Description | Default (libero example) | +|-----------|-------------|--------------------------| +| `hyperparams.batch_size_per_gpu` | Batch size per GPU | 4 | +| `hyperparams.gradient_accumulation_steps` | Gradient accumulation steps | 4 | +| `hyperparams.optimizer.learning_rate` | Learning rate | 5e-5 | +| `hyperparams.num_epoch` | Number of training epochs | 100 | +| `distributed.use_fsdp` | Multi-GPU FSDP training | true | +| `logging.save_interval` | Checkpoint save interval (steps) | 2000 | + +### Quick Start Checklist + +**Simulation (LIBERO)** + +- [ ] Downloaded `wall-oss-0.5` and `Qwen2.5-VL-3B-Instruct` +- [ ] Downloaded the `lerobot/libero` dataset +- [ ] Wrote `config.yml` and replaced all `/path/to/*` placeholders +- [ ] Ran `compute_norm_stats.py` to generate norm stats +- [ ] Launched `train_fsdp.py` and completed fine-tuning +- [ ] Ran simulation evaluation with `run_libero.sh` + +**Real Robot** + +- [ ] Prepared a LeRobot v3 real-robot dataset +- [ ] Wrote `config.yml` and configured `key_mappings` / DOF +- [ ] Ran `compute_norm_stats.py` to generate norm stats +- [ ] Launched `train_fsdp.py` and completed fine-tuning +- [ ] Started the inference server with `run_serving.sh` +- [ ] Verified server output with `draw_openloop_plot.py` diff --git a/workspace/example/lerobot/qwen2_5_lerobot_template.yml b/workspace/example/lerobot/qwen2_5_lerobot_template.yml new file mode 100644 index 0000000..e44d62a --- /dev/null +++ b/workspace/example/lerobot/qwen2_5_lerobot_template.yml @@ -0,0 +1,93 @@ +# Public Wall-X Qwen2.5 + LeRobot training template. +# +# Replace all /path/to/... values before launching training. +model_type: qwen2_5 + +task: + dof_config: + master_right_ee_cartesian_pos: 3 + master_right_ee_rotation: 3 + master_right_gripper: 1 + agent_pos_config: + follow_right_ee_cartesian_pos: 3 + follow_right_ee_rotation: 3 + follow_right_gripper: 2 + action_horizon: 32 + action_horizon_flow: 32 + use_state_string_representation: false + +model: + backbone: qwen2_5 + config_path: workspace/models_config/qwen2_5_moe_flash.json + processor_path: Qwen/Qwen2.5-VL-3B-Instruct + pretrained_path: Qwen/Qwen2.5-VL-3B-Instruct + attn_deterministic: true + use_ema: false + flow_loss_weight: 1.0 + enable_customized_robot_config: true + customized_robot_config: + name: your_robot + customized_dof_config: + action: 7 + customized_agent_pos_config: + observation.state: 8 + +hyperparams: + num_epoch: 100 + batch_size_per_gpu: 8 + gradient_accumulation_steps: 2 + seed: 10233 + optimizer: + optimizer_type: dmuon + learning_rate: 0.0001 + max_grad_norm: 1.0 + enable_grad_clip: true + betas: [0.9, 0.95] + weight_decay: 1.0e-8 + eps: 1.0e-8 + scheduler: + scheduler_type: cosine + num_warmup_steps: 1000 + num_training_steps: 64000000 + min_lr: 1.0e-5 + +distributed: + use_fsdp: true + use_mixed_precision: true + bf16: true + +data: + dataset_type: lerobot + lerobot_config: + repo_id: /path/to/your/lerobot_dataset + root: null + key_mappings: + camera: + observation.images.image: face_view + observation.images.image2: right_wrist_view + state: observation.state + action: action + norm_stats_path: /path/to/your/norm_stats.json + train_test_split: 0.95 + num_workers: 4 + resolution: + face_view: 256 + left_wrist_view: 256 + right_wrist_view: 256 + +logging: + log_name: wallx_lerobot_train + log_project: wallx_public + log_entity: null + use_wandb: false + log_interval: 10 + save_interval: 1000000 + val_interval: 1000000 + epoch_save_interval: 1 + +checkpoint: + save_path: /path/to/output/checkpoints + +debug: + profile: false + nvtx: false diff --git a/workspace/example/libero.yml b/workspace/example/libero.yml new file mode 100644 index 0000000..5a57e32 --- /dev/null +++ b/workspace/example/libero.yml @@ -0,0 +1,125 @@ +# LIBERO single-arm finetune example (Euler delta action, 256px, 2 cameras). +# +# Replace every /path/to/* placeholder before training: +# +# model.config_path -> model architecture JSON (mot_flash_mask_causal_xloss.json) +# model.processor_path -> Qwen2.5-VL-3B-Instruct directory +# model.pretrained_path -> same as processor_path, or HuggingFace cache path +# data.lerobot_config.repo_id -> local LeRobot dataset root (libero_all) +# data.norm_stats_path -> q01/q99 normalization JSON +# checkpoint.save_path -> writable directory for training checkpoints +# checkpoint.resume_from -> Wall-OSS-0.5 pretrained .safetensors or checkpoint directory +# +# Compute norm stats first: +# +# python scripts/compute_norm_stats.py \ +# --train_config workspace/example/libero.yml \ +# --data_root /path/to/libero_all \ +# --output_path /path/to/libero_all_norm_stats.json +# +# Launch training (from repo root): +# +# torchrun --nproc_per_node= wall_x/trainer/fsdp_trainer/train_fsdp.py \ +# --config workspace/example/libero.yml +# +# Strategy: keep dof / agent_pos totals at 26 to match the pretraining action space +# via ``action_padding``. The lerobot collator right-pads libero's 7-dim action / +# 8-dim state with zeros; loss does not flow through the padded tail. + +model_type: qwen2_5 + +task: + # Libero delta action: pos3 + rot3 + gripper1 = 7, plus action_padding(19) = 26. + dof_config: + master_right_ee_cartesian_pos: 3 # delta position + master_right_ee_rotation: 3 # delta rotation (ZYX euler) + master_right_gripper: 1 + action_padding: 19 + ar_dof_config: + master_right_ee_cartesian_pos: 3 + master_right_ee_rotation: 3 + master_right_gripper: 1 + action_padding: 19 + # State: pos3 + rot3 + gripper2 = 8, plus action_padding(18) = 26. + agent_pos_config: + follow_right_ee_cartesian_pos: 3 + follow_right_ee_rotation: 3 + follow_right_gripper: 2 + action_padding: 18 + action_horizon: 10 + action_horizon_flow: 10 + use_state_string_representation: false + +model: + backbone: qwen2_5 + config_path: /path/to/wall-oss-0.5/config.json + processor_path: /path/to/Qwen2.5-VL-3B-Instruct + pretrained_path: /path/to/Qwen2.5-VL-3B-Instruct + attn_deterministic: true + use_ema: false + flow_loss_weight: 1.0 + ar_loss_weight: 0.01 + +hyperparams: + num_epoch: 100 + batch_size_per_gpu: 4 + gradient_accumulation_steps: 4 + seed: 10222 + optimizer: + optimizer_type: adamw + learning_rate: 5.0e-05 + max_grad_norm: 1.0 + enable_grad_clip: true + betas: [0.9, 0.95] + weight_decay: 1.0e-8 + eps: 1.0e-8 + scheduler: + scheduler_type: cosine + num_warmup_steps: 1000 + num_training_steps: 200000 + min_lr: 1.0e-6 + +distributed: + use_fsdp: true + use_mixed_precision: true + bf16: true + +data: + dataset_type: lerobot + lerobot_config: + repo_id: /path/to/libero_all + root: null + key_mappings: + # libero_all v3.0 only has faceImg + rightImg (no leftImg). + camera: + observation.images.faceImg: face_view + observation.images.rightImg: right_wrist_view + state: observation.state + action: action + norm_stats_path: /path/to/libero_all_norm_stats.json + train_test_split: 0.95 + num_workers: 4 + max_length: 1024 + resolution: + face_view: 256 + right_wrist_view: 256 + +logging: + log_name: libero_ft + log_project: lerobot_libero_ft + log_entity: your_wandb_entity + use_wandb: true + log_interval: 10 + save_interval: 2000 + val_interval: 1000000 + epoch_save_interval: 1 + +checkpoint: + save_path: /path/to/libero + # Single-file .safetensors loads as pretrain weights before FSDP wrapping. + # Use a checkpoint directory for full resume (optimizer / scheduler / RNG). + resume_from: /path/to/wall-oss-0.5/model.safetensors + +debug: + profile: false + nvtx: false diff --git a/workspace/example/maniparena_example.yml b/workspace/example/maniparena_example.yml new file mode 100644 index 0000000..7afe8dd --- /dev/null +++ b/workspace/example/maniparena_example.yml @@ -0,0 +1,134 @@ +# ManipArena / CVPR dual-arm finetune example (6D relative action, 448px, 3 cameras). +# +# Replace every /path/to/* placeholder before training: +# +# model.config_path -> model architecture JSON (mot_flash_mask_causal_xloss.json) +# model.processor_path -> Qwen2.5-VL-3B-Instruct directory +# model.pretrained_path -> same as processor_path, or HuggingFace cache path +# data.lerobot_config.repo_id -> local LeRobot dataset root (cvpr_4tasks) +# data.norm_stats_path -> q01/q99 normalization JSON +# checkpoint.save_path -> writable directory for training checkpoints +# checkpoint.resume_from -> Wall-OSS-0.5 pretrained .safetensors or checkpoint directory +# +# Compute norm stats first: +# +# python scripts/compute_norm_stats.py \ +# --train_config workspace/example/maniparena_example.yml \ +# --data_root /path/to/cvpr_4tasks \ +# --output_path /path/to/cvpr_4tasks_norm_stats.json +# +# Launch training (from repo root): +# +# torchrun --nproc_per_node= wall_x/trainer/fsdp_trainer/train_fsdp.py \ +# --config workspace/example/maniparena_example.yml +# +# Strategy: keep dof / agent_pos totals at 26 to match the pretraining action space +# via ``action_padding``. The lerobot collator right-pads real action / state dims; +# loss does not flow through the padded tail. + +model_type: qwen2_5 + +task: + # Dual-arm 6D relative action: 10 + 10 = 20, plus action_padding(6) = 26. + dof_config: + follow_left_ee_cartesian_pos_relative: 3 + follow_left_ee_rotation_6D_relative: 6 + follow_left_gripper: 1 + follow_right_ee_cartesian_pos_relative: 3 + follow_right_ee_rotation_6D_relative: 6 + follow_right_gripper: 1 + action_padding: 6 + ar_dof_config: + follow_left_ee_cartesian_pos_relative: 3 + follow_left_ee_rotation_6D_relative: 6 + follow_left_gripper: 1 + follow_right_ee_cartesian_pos_relative: 3 + follow_right_ee_rotation_6D_relative: 6 + follow_right_gripper: 1 + action_padding: 6 + agent_pos_config: + follow_left_ee_cartesian_pos: 3 + follow_left_ee_rotation_6D: 6 + follow_left_gripper: 1 + follow_right_ee_cartesian_pos: 3 + follow_right_ee_rotation_6D: 6 + follow_right_gripper: 1 + action_padding: 6 + action_horizon: 32 + action_horizon_flow: 32 + use_state_string_representation: false + +model: + backbone: qwen2_5 + config_path: /path/to/wall-oss-0.5/config.json + processor_path: /path/to/Qwen2.5-VL-3B-Instruct + pretrained_path: /path/to/Qwen2.5-VL-3B-Instruct + attn_deterministic: true + use_ema: false + flow_loss_weight: 1.0 + ar_loss_weight: 0.01 + +hyperparams: + num_epoch: 100 + batch_size_per_gpu: 4 + gradient_accumulation_steps: 4 + seed: 10222 + optimizer: + optimizer_type: adamw + learning_rate: 5.0e-05 + max_grad_norm: 1.0 + enable_grad_clip: true + betas: [0.9, 0.95] + weight_decay: 1.0e-8 + eps: 1.0e-8 + scheduler: + scheduler_type: cosine + num_warmup_steps: 1000 + num_training_steps: 200000 + min_lr: 1.0e-6 + +distributed: + use_fsdp: true + use_mixed_precision: true + bf16: true + +data: + dataset_type: lerobot + lerobot_config: + repo_id: /path/to/cvpr_4tasks + root: null + key_mappings: + camera: + observation.images.faceImg: face_view + observation.images.leftImg: left_wrist_view + observation.images.rightImg: right_wrist_view + state: observation.state + action: action + norm_stats_path: /path/to/cvpr_4tasks_norm_stats.json + train_test_split: 0.95 + num_workers: 4 + max_length: 1024 + resolution: + face_view: 448 + left_wrist_view: 448 + right_wrist_view: 448 + +logging: + log_name: maniparena_ft + log_project: lerobot_maniparena_ft + log_entity: your_wandb_entity + use_wandb: true + log_interval: 10 + save_interval: 2000 + val_interval: 1000000 + epoch_save_interval: 1 + +checkpoint: + save_path: /path/to/cvpr_4tasks + # Single-file .safetensors loads as pretrain weights before FSDP wrapping. + # Use a checkpoint directory for full resume (optimizer / scheduler / RNG). + resume_from: /path/to/wall-oss-0.5/model.safetensors + +debug: + profile: false + nvtx: false diff --git a/workspace/lerobot_example/config_qact.yml b/workspace/lerobot_example/config_qact.yml deleted file mode 100644 index 1248fa2..0000000 --- a/workspace/lerobot_example/config_qact.yml +++ /dev/null @@ -1,151 +0,0 @@ -# Training Configuration for Wall-X Robotic Multi-Modal Learning -# This configuration supports multi-modal learning with vision, language, and action data - -# Model and paths configuration -log_name: "robotic_training" -log_project: "vla_training" -model_type: wall-oss -pretrained_wallx_path: "/path/to/wallx_model/" # Must set -save_path: "/path/to/workspace/" # Must set -use_fast_tokenizer: False # True: train FAST, False: train Flow -action_tokenizer_path: "/path/to/fast/" # Must set if use_fast_tokenizer is true - -# Torch Profile -profile: False -profile_save_path: /path/to/profile/ -profile_wait_iters: 10 -profile_warmup_iters: 5 -profile_active_iters: 2 - -# Training hyperparameters -num_warmup_steps: 100 -num_training_steps: 64000000 -learning_rate: 0.00005 -min_lr: 0.00005 -num_epoch: 100 -gradient_accumulation_steps: 32 -batch_size_per_gpu: 8 -padding_side: left -epoch_save_interval: 10 - -# Training optimization settings -FSDP2: True -torch_compile: False - -# Robot configuration - Define degrees of freedom for each component -dof_config: - follow_left_ee_cartesian_pos: 3 # Left end-effector Cartesian position - follow_left_ee_rotation: 3 # Left end-effector rotation - follow_left_gripper: 1 # Left gripper control - follow_right_ee_cartesian_pos: 3 # Right end-effector Cartesian position - follow_right_ee_rotation: 3 # Right end-effector rotation - follow_right_gripper: 1 # Right gripper control - head_actions: 2 # Head/camera movement - height: 1 # Mobile base height control - car_pose: 3 # Mobile base pose (x, y, theta) - -# Agent proprioception configuration (typically matches DOF config) -agent_pos_config: - follow_left_ee_cartesian_pos: 3 - follow_left_ee_rotation: 3 - follow_left_gripper: 1 - follow_right_ee_cartesian_pos: 3 - follow_right_ee_rotation: 3 - follow_right_gripper: 1 - head_actions: 2 - height: 1 - car_pose: 3 - -# # Checkpoint resuming configuration -# resume: -# ckpt: "/path/to/resume_model/" -# load_ckpt_only: true - -norm_stats_path: "/path/to/norm_stats.json" - -enable_customized_robot_config: true -customized_robot_config: - name: "lerobot/aloha_mobile_cabinet" - customized_dof_config: - "action_left_shoulder" : 1 - "action_left_elbow" : 1 - "action_left_forearm_roll" : 1 - "action_left_wrist_angle" : 1 - "action_left_wrist_rotate" : 1 - "action_left_gripper" : 1 - "action_right_waist" : 1 - "action_right_shoulder" : 1 - "action_right_elbow" : 1 - "action_right_forearm_roll" : 1 - "action_right_wrist_angle" : 1 - "action_right_wrist_rotate" : 1 - "action_right_gripper" : 1 - - customized_agent_pos_config: - "state_left_shoulder" : 1 - "state_left_elbow" : 1 - "state_left_forearm_roll" : 1 - "state_left_wrist_angle" : 1 - "state_left_wrist_rotate" : 1 - "state_left_gripper" : 1 - "state_right_waist" : 1 - "state_right_shoulder" : 1 - "state_right_elbow" : 1 - "state_right_forearm_roll" : 1 - "state_right_wrist_angle" : 1 - "state_right_wrist_rotate" : 1 - "state_right_gripper" : 1 - -# Data configuration -data: - use_lerobot: true - - # LeRobot dataset configuration - lerobot_config: - repo_id: "lerobot/aloha_mobile_cabinet" - root: null - episodes: null - image_transforms: null - delta_timestamps: null - tolerance_s: 1e-4 - revision: null - force_cache_sync: false - download_videos: true - video_backend: null - - action_horizon: 32 - train_test_split: 0.95 - - # Action keys for observation and prediction - obs_action_keys: - - follow_left_ee_cartesian_pos - - follow_left_ee_rotation - - follow_left_gripper - - follow_right_ee_cartesian_pos - - follow_right_ee_rotation - - follow_right_gripper - - head_actions - - height - - car_pose - - predict_action_keys: - - follow_left_ee_cartesian_pos - - follow_left_ee_rotation - - follow_left_gripper - - follow_right_ee_cartesian_pos - - follow_right_ee_rotation - - follow_right_gripper - - head_actions - - height - - car_pose - - # Image resolution configuration for different camera views - resolution: - face_view: 256 - left_wrist_view: 256 - right_wrist_view: 256 - move1_view: 256 - move2_view: 256 - top_view: 256 - wall_view: 256 - multi_modal: 256 diff --git a/workspace/lerobot_example/config_qact_from_vlm.yml b/workspace/lerobot_example/config_qact_from_vlm.yml deleted file mode 100644 index e872311..0000000 --- a/workspace/lerobot_example/config_qact_from_vlm.yml +++ /dev/null @@ -1,152 +0,0 @@ -# Train from Qwen-2.5-VL - -# Model and paths configuration -log_name: "robotic_training" -log_project: "vla_training" -model_type: qwen2_5 -pretrained_wallx_path: "/path/to/wallx_model/" # Must set -save_path: "/path/to/workspace/" # Must set -use_fast_tokenizer: True # True: train FAST, False: train Flow -action_tokenizer_path: "/path/to/fast/" # Must set if use_fast_tokenizer is true -qwen_vl_act_config_path: "wall-x/workspace/lerobot_example/qwen25_config.json" - - -# Torch Profile -profile: False -profile_save_path: /path/to/profile/ -profile_wait_iters: 10 -profile_warmup_iters: 5 -profile_active_iters: 2 - -# Training hyperparameters -num_warmup_steps: 100 -num_training_steps: 64000000 -learning_rate: 0.00009 -min_lr: 0.00005 -num_epoch: 100 -gradient_accumulation_steps: 1 -batch_size_per_gpu: 8 -padding_side: left -epoch_save_interval: 10 - -# Training optimization settings -FSDP2: True -torch_compile: False - -# Robot configuration - Define degrees of freedom for each component -dof_config: - follow_left_ee_cartesian_pos: 3 # Left end-effector Cartesian position - follow_left_ee_rotation: 3 # Left end-effector rotation - follow_left_gripper: 1 # Left gripper control - follow_right_ee_cartesian_pos: 3 # Right end-effector Cartesian position - follow_right_ee_rotation: 3 # Right end-effector rotation - follow_right_gripper: 1 # Right gripper control - head_actions: 2 # Head/camera movement - height: 1 # Mobile base height control - car_pose: 3 # Mobile base pose (x, y, theta) - -# Agent proprioception configuration (typically matches DOF config) -agent_pos_config: - follow_left_ee_cartesian_pos: 3 - follow_left_ee_rotation: 3 - follow_left_gripper: 1 - follow_right_ee_cartesian_pos: 3 - follow_right_ee_rotation: 3 - follow_right_gripper: 1 - head_actions: 2 - height: 1 - car_pose: 3 - -# # Checkpoint resuming configuration -# resume: -# ckpt: "/path/to/resume_model/" -# load_ckpt_only: true - -norm_stats_path: "/path/to/norm_stats.json" - -enable_customized_robot_config: true -customized_robot_config: - name: "physical-intelligence/libero" - customized_dof_config: - "action_left_shoulder" : 1 - "action_left_elbow" : 1 - "action_left_forearm_roll" : 1 - "action_left_wrist_angle" : 1 - "action_left_wrist_rotate" : 1 - "action_left_gripper" : 1 - "action_right_waist" : 1 - "action_right_shoulder" : 1 - "action_right_elbow" : 1 - "action_right_forearm_roll" : 1 - "action_right_wrist_angle" : 1 - "action_right_wrist_rotate" : 1 - "action_right_gripper" : 1 - - customized_agent_pos_config: - "state_left_shoulder" : 1 - "state_left_elbow" : 1 - "state_left_forearm_roll" : 1 - "state_left_wrist_angle" : 1 - "state_left_wrist_rotate" : 1 - "state_left_gripper" : 1 - "state_right_waist" : 1 - "state_right_shoulder" : 1 - "state_right_elbow" : 1 - "state_right_forearm_roll" : 1 - "state_right_wrist_angle" : 1 - "state_right_wrist_rotate" : 1 - "state_right_gripper" : 1 - -# Data configuration -data: - use_lerobot: true - - # LeRobot dataset configuration - lerobot_config: - repo_id: "lerobot/aloha_mobile_cabinet" - root: null - episodes: null - image_transforms: null - delta_timestamps: null - tolerance_s: 1e-4 - revision: null - force_cache_sync: false - download_videos: true - video_backend: null - - action_horizon: 32 - train_test_split: 0.95 - - # Action keys for observation and prediction - obs_action_keys: - - follow_left_ee_cartesian_pos - - follow_left_ee_rotation - - follow_left_gripper - - follow_right_ee_cartesian_pos - - follow_right_ee_rotation - - follow_right_gripper - - head_actions - - height - - car_pose - - predict_action_keys: - - follow_left_ee_cartesian_pos - - follow_left_ee_rotation - - follow_left_gripper - - follow_right_ee_cartesian_pos - - follow_right_ee_rotation - - follow_right_gripper - - head_actions - - height - - car_pose - - # Image resolution configuration for different camera views - resolution: - face_view: 256 - left_wrist_view: 256 - right_wrist_view: 256 - move1_view: 256 - move2_view: 256 - top_view: 256 - wall_view: 256 - multi_modal: 256 diff --git a/workspace/lerobot_example/evaluation/lerobot_openloop.png b/workspace/lerobot_example/evaluation/lerobot_openloop.png deleted file mode 100644 index dfc72e7..0000000 --- a/workspace/lerobot_example/evaluation/lerobot_openloop.png +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:aae8646e566b79b64a669956d6d2c779d72809010c802f75b2623ee371444b47 -size 985979 diff --git a/workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml b/workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml deleted file mode 100644 index 21198e4..0000000 --- a/workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml +++ /dev/null @@ -1,100 +0,0 @@ -# Training Configuration for Wall-X Robotic Multi-Modal Learning -# This configuration supports multi-modal learning with vision, language, and action data - -# Model and paths configuration -log_name: "opensource_training" -log_project: "libero" -model_type: qwen2_5 -use_fast_tokenizer: true -pretrained_wallx_path: "/path/to/qwen/" -action_tokenizer_path: "/path/to/fast/" -qwen_vl_act_config_path: "/path/to/qwen25_config.json" - -save_path: "/path/to/save" -# Torch Profile -profile: False -profile_save_path: /path/to/profile/ -profile_wait_iters: 10 -profile_warmup_iters: 5 -profile_active_iters: 2 - -# Training hyperparameters -num_warmup_steps: 100 -num_training_steps: 64000000 -learning_rate: 0.00005 -min_lr: 0.00005 -num_epoch: 100 -gradient_accumulation_steps: 1 -batch_size_per_gpu: 8 -padding_side: left -epoch_save_interval: 1 - -# Robot configuration - Define degrees of freedom for each component -dof_config: - master_right_ee_cartesian_pos: 3 # Right end-effector Cartesian position - master_right_ee_rotation: 3 # Right end-effector rotation - master_right_gripper: 1 # Right gripper control - -# Agent proprioception configuration (typically matches DOF config) -agent_pos_config: - follow_right_ee_cartesian_pos: 3 - follow_right_ee_rotation: 3 - follow_right_gripper: 1 - -norm_stats_path: "/path/to/libero_norm_stats.json" - -enable_customized_robot_config: true -customized_robot_config: - name: "libero_all" - customized_dof_config: - "panda_action_eef_with_gripper": 7 - - customized_agent_pos_config: - "panda_state_eef_with_gripper": 7 - -# Checkpoint resuming configuration -# resume: -# ckpt: "/path/to/ckpt" -# load_ckpt_only: false - -# Data configuration -data: - use_lerobot: true - - # LeRobot dataset configuration - lerobot_config: - repo_id: "libero_all" - root: null - episodes: null - image_transforms: null - delta_timestamps: null - tolerance_s: 1e-4 - revision: null - force_cache_sync: false - download_videos: true - video_backend: null - - action_horizon: 10 - train_test_split: 0.95 - - # Action keys for observation and prediction - obs_action_keys: - - follow_right_ee_cartesian_pos - - follow_right_ee_rotation - - follow_right_gripper - - predict_action_keys: - - master_right_ee_cartesian_pos - - master_right_ee_rotation - - master_right_gripper - - # Image resolution configuration for different camera views - resolution: - face_view: 256 - left_wrist_view: 256 - right_wrist_view: 256 - move1_view: 256 - move2_view: 256 - top_view: 256 - wall_view: 256 - multi_modal: 256 diff --git a/workspace/lerobot_example/qwen25_config.json b/workspace/lerobot_example/qwen25_config.json deleted file mode 100644 index c92cc20..0000000 --- a/workspace/lerobot_example/qwen25_config.json +++ /dev/null @@ -1,106 +0,0 @@ -{ - "architectures": [ - "Qwen2_5_VLForConditionalGeneration" - ], - "attention_dropout": 0.0, - "bos_token_id": 151643, - "eos_token_id": 151645, - "vision_start_token_id": 151652, - "vision_end_token_id": 151653, - "vision_token_id": 151654, - "image_token_id": 151655, - "video_token_id": 151656, - "hidden_act": "silu", - "hidden_size": 2048, - "initializer_range": 0.02, - "intermediate_size": 11008, - "max_position_embeddings": 128000, - "max_window_layers": 70, - "model_type": "qwen2_5_vl", - "num_attention_heads": 16, - "num_hidden_layers": 36, - "num_key_value_heads": 2, - "rms_norm_eps": 1e-06, - "rope_theta": 1000000.0, - "sliding_window": 32768, - "tie_word_embeddings": true, - "torch_dtype": "bfloat16", - "transformers_version": "4.41.2", - "_attn_implementation": "flash_attention_2", - "use_cache": true, - "use_sliding_window": false, - "vision_config": { - "depth": 32, - "hidden_act": "silu", - "hidden_size": 1280, - "intermediate_size": 3420, - "num_heads": 16, - "in_chans": 3, - "out_hidden_size": 2048, - "patch_size": 14, - "spatial_merge_size": 2, - "spatial_patch_size": 14, - "window_size": 112, - "fullatt_block_indexes": [ - 7, - 15, - 23, - 31 - ], - "tokens_per_second": 2, - "temporal_patch_size": 2 - }, - "rope_scaling": { - "type": "mrope", - "mrope_section": [ - 16, - 24, - 24 - ] - }, - "vocab_size": 151936, - "num_experts": 2, - "experts":[ - { - "hidden_size": 2048, - "intermediate_size": 11008, - "hidden_act": "silu" - }, - { - "hidden_size": 2048, - "intermediate_size": 2048, - "hidden_act": "silu" - } - ], - "dof_config": { - "follow_left_ee_cartesian_pos": 3, - "follow_left_ee_rotation": 3, - "follow_left_gripper": 1, - "follow_right_ee_cartesian_pos": 3, - "follow_right_ee_rotation": 3, - "follow_right_gripper": 1, - "head_actions": 2, - "height": 1, - "car_pose": 3 - }, - "agent_pos_config": { - "follow_left_ee_cartesian_pos": 3, - "follow_left_ee_rotation": 3, - "follow_left_gripper": 1, - "follow_right_ee_cartesian_pos": 3, - "follow_right_ee_rotation": 3, - "follow_right_gripper": 1, - "head_actions": 2, - "height": 1, - "car_pose": 3 - }, - "noise_scheduler": { - "beta_alpha": 1.5, - "beta_beta": 1.0, - "s": 0.999, - "num_inference_timesteps": 5 - }, - "dim_inputs": [2048,2048], - "attention_moe": false, - "mlp_moe": true - } diff --git a/workspace/lerobot_example/run.sh b/workspace/lerobot_example/run.sh deleted file mode 100644 index 7d18106..0000000 --- a/workspace/lerobot_example/run.sh +++ /dev/null @@ -1,23 +0,0 @@ -#!/bin/bash -export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 -NUM_GPUS=$(echo $CUDA_VISIBLE_DEVICES | tr ',' '\n' | wc -l) - -# print current time -echo "[current time: $(date +'%Y-%m-%d %H:%M:%S')]" - -code_dir="/path/to/wall-x" -config_path="/path/to/wall-x/workspace/lerobot_example" - -# Use a fixed port instead of a random one -export PORT=$((21000 + $RANDOM % 30000)) - -MASTER_PORT=10239 # use 5 digits ports - -export LAUNCHER="accelerate launch --num_processes=$NUM_GPUS --main_process_port=$PORT" - -export SCRIPT="${code_dir}/train_qact.py" -export SCRIPT_ARGS="--config ${config_path}/config_qact.yml --seed $MASTER_PORT" - -echo "Running command: $LAUNCHER $SCRIPT $SCRIPT_ARGS" - -$LAUNCHER $SCRIPT $SCRIPT_ARGS diff --git a/workspace/models_config/qwen2_5_moe_flash.json b/workspace/models_config/qwen2_5_moe_flash.json new file mode 100644 index 0000000..01b5e01 --- /dev/null +++ b/workspace/models_config/qwen2_5_moe_flash.json @@ -0,0 +1,108 @@ +{ + "architectures": [ + "Qwen2_5_VLForConditionalGeneration" + ], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "vision_start_token_id": 151652, + "vision_end_token_id": 151653, + "vision_token_id": 151654, + "image_token_id": 151655, + "video_token_id": 151656, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 11008, + "max_position_embeddings": 128000, + "max_window_layers": 70, + "model_type": "qwen2_5_vl", + "num_attention_heads": 16, + "num_hidden_layers": 36, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "transformers_version": "4.41.2", + "_attn_implementation": "flash_attention_2", + "use_cache": true, + "use_sliding_window": false, + "vision_config": { + "depth": 32, + "hidden_act": "silu", + "hidden_size": 1280, + "intermediate_size": 3420, + "num_heads": 16, + "in_chans": 3, + "out_hidden_size": 2048, + "patch_size": 14, + "spatial_merge_size": 2, + "spatial_patch_size": 14, + "window_size": 112, + "fullatt_block_indexes": [ + 7, + 15, + 23, + 31 + ], + "tokens_per_second": 2, + "temporal_patch_size": 2 + }, + "rope_scaling": { + "type": "mrope", + "mrope_section": [ + 16, + 24, + 24 + ] + }, + "vocab_size": 151936, + "num_experts": 2, + "experts": [ + { + "hidden_size": 2048, + "intermediate_size": 11008, + "hidden_act": "silu" + }, + { + "hidden_size": 2048, + "intermediate_size": 2048, + "hidden_act": "silu" + } + ], + "dof_config": { + "follow_left_ee_cartesian_pos": 3, + "follow_left_ee_rotation": 3, + "follow_left_gripper": 1, + "follow_right_ee_cartesian_pos": 3, + "follow_right_ee_rotation": 3, + "follow_right_gripper": 1 + }, + "agent_pos_config": { + "follow_left_ee_cartesian_pos": 3, + "follow_left_ee_rotation": 3, + "follow_left_gripper": 1, + "follow_right_ee_cartesian_pos": 3, + "follow_right_ee_rotation": 3, + "follow_right_gripper": 1, + "follow_left_arm_joint_cur": 1, + "follow_right_arm_joint_cur": 1 + }, + "noise_scheduler": { + "beta_alpha": 1.5, + "beta_beta": 1.0, + "t_eps": 0.001, + "s": 0.999, + "num_inference_timesteps": 10 + }, + "dim_inputs": [ + 2048, + 2048 + ], + "attention_moe": false, + "mlp_moe": true, + "ar_loss_weight": 1, + "causal_action_attention_mask": true +}