commit 24dbdbd24b8b6cb3404697407e60becaab049fe3 Author: Starrick <27157812@qq.com> Date: Sun Sep 7 14:59:17 2025 +0800 Init diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..142aa08 --- /dev/null +++ b/.gitignore @@ -0,0 +1,270 @@ +# ============================================================================= +# 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 +__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/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# 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 +*.png +*.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 +wandb/ +_wandb/ + +# TensorBoard +runs/ +tensorboard/ + +# MLflow +mlruns/ + +# ============================================================================= +# TEMPORARY AND CACHE FILES +# ============================================================================= + +# Temporary files +*.tmp +*.temp +*.bak +*.backup + +# Cache directories +.cache/ +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 diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..741dda4 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "3rdparty/cutlass"] + path = 3rdparty/cutlass + url = https://github.com/NVIDIA/cutlass.git diff --git a/3rdparty/cutlass b/3rdparty/cutlass new file mode 160000 index 0000000..e51efbf --- /dev/null +++ b/3rdparty/cutlass @@ -0,0 +1 @@ +Subproject commit e51efbfe18fe4f4cbb66ab814c55bf4aa0185491 diff --git a/README.md b/README.md new file mode 100644 index 0000000..957c891 --- /dev/null +++ b/README.md @@ -0,0 +1,69 @@ +# Wall-X + +## Overview + +Wall-X is a multimodal foundation model designed for robotics applications, combining vision, language, and action capabilities. The model architecture is built upon Qwen2.5-3B-VL with specialized adaptations for robotic control tasks. + +## Environment Setup + +Create and activate 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: +```bash +git clone https://github.com/huggingface/lerobot.git +cd lerobot +pip install -e . +``` + +Install wall_x: +```bash +git submodule update --init --recursive +MAX_JOBS=4 pip install --no-build-isolation --verbose . +``` + +## Training + +### Finetune on LeRobot Datasets + +Before training, please refer to `workspace/README.md` for detailed configuration instructions including: + +Training script path configuration + +- GPU setup +- Model and data paths +- Robot DOF configuration +- Training hyperparameters + +```bash +bash ./workspace/lerobot_example/run.sh +``` + +## Inference + +For model inference, please refer to: + +```bash +python ./scripts/fake_inference.py +``` + +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 + +To generate an open-loop comparison plot, please follow: + +```bash +python ./scripts/draw_openloop_plot.py +``` \ No newline at end of file diff --git a/csrc/README.md b/csrc/README.md new file mode 100644 index 0000000..d73f8d4 --- /dev/null +++ b/csrc/README.md @@ -0,0 +1,23 @@ +# Fusion Operators (CSRC) + +High-performance CUDA kernels for accelerating model training. + +## Operators + +### Asymmetric Dual Expert GEMM +- `asym_dual_gmm`: Simultaneous matrix multiplication for two experts +- Supports all transpose combinations (NN, TN, NT, TT) + +### Token Permutation +- `permute`: Token permutation for MoE routing +- `unpermute`: Token recovery after expert computation +- `unpermute_bwd`: Backward pass for token recovery + +### Multimodal RoPE +- `rope`: Rotary Position Embedding forward pass +- `rope_bwd`: RoPE backward pass +- Support for multimodal inputs with configurable sections + +## Acknowledgments + +The `permute` and `unpermute` operators are adapted from [fanshiqing/grouped_gemm](https://github.com/fanshiqing/grouped_gemm). Thanks for their open-source contributions. \ No newline at end of file diff --git a/csrc/dual_asym_grouped_gemm.cu b/csrc/dual_asym_grouped_gemm.cu new file mode 100644 index 0000000..031cc58 --- /dev/null +++ b/csrc/dual_asym_grouped_gemm.cu @@ -0,0 +1,366 @@ +#include "dual_asym_grouped_gemm.h" + +#include +#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 new file mode 100644 index 0000000..4d7abde --- /dev/null +++ b/csrc/dual_asym_grouped_gemm.h @@ -0,0 +1,10 @@ +#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 new file mode 100644 index 0000000..059c77e --- /dev/null +++ b/csrc/ops.cu @@ -0,0 +1,16 @@ +#include "dual_asym_grouped_gemm.h" +#include "permute.h" +#include "rope.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"); +} + diff --git a/csrc/permute.cu b/csrc/permute.cu new file mode 100644 index 0000000..7eb7c26 --- /dev/null +++ b/csrc/permute.cu @@ -0,0 +1,942 @@ +/************************************************************************* + * 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 new file mode 100644 index 0000000..45a8965 --- /dev/null +++ b/csrc/permute.h @@ -0,0 +1,31 @@ +/************************************************************************* + * 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 new file mode 100644 index 0000000..09c1a11 --- /dev/null +++ b/csrc/rope.cu @@ -0,0 +1,552 @@ +#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) + { + rotate_val = -rotate_val; // First half: negate second half + } + + // 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; + } + int *d_mrope_section_doubled; + cudaMalloc(&d_mrope_section_doubled, 3 * sizeof(int)); + cudaMemcpyAsync(d_mrope_section_doubled, mrope_section_doubled.data(), 3 * sizeof(int), + cudaMemcpyHostToDevice, stream); + + 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; + } + int *d_mrope_section_doubled; + cudaMalloc(&d_mrope_section_doubled, 3 * sizeof(int)); + cudaMemcpyAsync(d_mrope_section_doubled, mrope_section_doubled.data(), 3 * sizeof(int), + cudaMemcpyHostToDevice, stream); + + 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 new file mode 100644 index 0000000..930c83e --- /dev/null +++ b/csrc/rope.h @@ -0,0 +1,14 @@ +#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/requirements.txt b/requirements.txt new file mode 100644 index 0000000..3f88666 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,9 @@ +torch==2.6.0 +torchvision==0.21.0 +torchaudio==2.6.0 +transformers==4.49.0 +accelerate==1.10.1 +peft==0.17.1 +scipy==1.15.3 +torchdiffeq==0.2.5 +qwen_vl_utils==0.0.11 \ No newline at end of file diff --git a/scripts/draw_openloop_plot.py b/scripts/draw_openloop_plot.py new file mode 100644 index 0000000..a649392 --- /dev/null +++ b/scripts/draw_openloop_plot.py @@ -0,0 +1,76 @@ +import os +import yaml +import torch +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 + + +model_path = "path/to/model" +action_tokenizer_path = "path/to/action_tokenizer" +save_dir = "path/to/plot" +model = Qwen2_5_VLMoEForAction.from_pretrained(model_path, action_tokenizer_path=action_tokenizer_path) +model.eval() +model = model.to("cuda") +model = model.bfloat16() + +def load_config(config_path): + """Load configuration from YAML file.""" + with open(config_path, "r") as f: + config = yaml.load(f, Loader=yaml.FullLoader) + + config["data"]["model_type"] = config.get("model_type") + + return config + +# get test dataloader +path = "path/to/config" +config = load_config(path) +dataload_config = get_data_configs(config["data"]) +lerobot_config = dataload_config.get("lerobot_config", {}) +dataset = load_test_dataset(config, lerobot_config, seed=42) +dataloader = dataset.get_dataloader() + +total_frames = len(dataloader) + +pred_horizon = 32 +action_dim = 14 +gt_traj = torch.zeros((total_frames, action_dim)) +pred_traj = torch.zeros((total_frames, action_dim)) + +for idx, batch in enumerate(dataloader): + gt_traj[idx] = batch['action_chunk'][0, 0,:action_dim] + if idx % 32 ==0 and idx + 32 < total_frames: + batch = batch.to("cuda") + with torch.no_grad(): + outputs = model( + **batch, + action_dim=action_dim, + pred_horizon=pred_horizon, + mode="predict", + predict_mode="fast" + ) + pred_traj[idx : idx + pred_horizon] = outputs['predict_action'].detach().cpu() + + +gt_traj_np = gt_traj.numpy() +pred_traj_np = pred_traj.numpy() + +timesteps = gt_traj.shape[0] + +import matplotlib.pyplot as plt + +fig, axs = plt.subplots(action_dim, 1, figsize=(15, 5 * action_dim), sharex=True) +fig.suptitle(f'Action Comparison for lerobot', fontsize=16) + +for i in range(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) +plt.savefig(os.path.join(save_dir, f"lerobot_comparison.png")) +plt.close() diff --git a/scripts/fake_inference.py b/scripts/fake_inference.py new file mode 100644 index 0000000..8c2cdf6 --- /dev/null +++ b/scripts/fake_inference.py @@ -0,0 +1,80 @@ +import torch +from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction + +model_path = "/path/to/model" +model = Qwen2_5_VLMoEForAction.from_pretrained(model_path) +model.eval() + +# Gen Fake data +batch_size = 1 +seq_length = 50 + +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"] + + +device = "cuda" + +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() + +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" + ) + + 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}") + + # 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") + + if not torch.isinf(outputs.logits).any(): + print("✅ Output contains no infinity values") + else: + print("❌ Output contains infinity values") + + print(f"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}") + +except Exception as e: + print(f"❌ Fake inference test failed: {e}") + import traceback + traceback.print_exc() \ No newline at end of file diff --git a/scripts/merge_tokenizer.py b/scripts/merge_tokenizer.py new file mode 100644 index 0000000..2ec453e --- /dev/null +++ b/scripts/merge_tokenizer.py @@ -0,0 +1,25 @@ +from transformers import AutoProcessor +import os + +processor_path = "/path/to/Qwen2.5-VL-3B-Instruct" +action_tokenizer_path = "/path/to/fast" +use_fast_tokenizer = True + +processor = AutoProcessor.from_pretrained(processor_path, use_fast=True) +processor.tokenizer.padding_side = "left" + +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) + +begin_idx_token = f"<|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) + diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..e13c518 --- /dev/null +++ b/setup.py @@ -0,0 +1,60 @@ +import os +import torch +from pathlib import Path +from setuptools import setup, find_packages +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + +cwd = Path(os.path.dirname(os.path.abspath(__file__))) + +nvcc_flags = [ + "-std=c++17", # NOTE: CUTLASS requires c++17 + "-DENABLE_BF16", # Enable BF16 for cuda_version >= 11 +] + +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]}" + +if device_capability: + nvcc_flags.extend( + [ + f"--generate-code=arch=compute_{device_capability},code=sm_{device_capability}", + f"-DGROUPED_GEMM_DEVICE_CAPABILITY={device_capability}", + ] + ) + +ext_modules = [ + CUDAExtension( + "wallx_csrc", + [ + "csrc/ops.cu", + "csrc/dual_asym_grouped_gemm.cu", + "csrc/permute.cu", + "csrc/rope.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.0", + author="X2Robot Team", + classifiers=[ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: BSD License", + "Operating System :: Unix", + ], + packages=find_packages(), + ext_modules=ext_modules, + cmdclass={"build_ext": BuildExtension}, +) diff --git a/train_qact.py b/train_qact.py new file mode 100755 index 0000000..e0281f7 --- /dev/null +++ b/train_qact.py @@ -0,0 +1,106 @@ +import os +import json +import time +import yaml +import wandb +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) + + accelerator = Accelerator( + kwargs_handlers=[ddp_kwargs], + mixed_precision="bf16", + 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) \ No newline at end of file diff --git a/wall_x/__init__.py b/wall_x/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/wall_x/data/__init__.py b/wall_x/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/wall_x/data/config.py b/wall_x/data/config.py new file mode 100644 index 0000000..240110c --- /dev/null +++ b/wall_x/data/config.py @@ -0,0 +1,95 @@ +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" +] + +# 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 \ No newline at end of file diff --git a/wall_x/data/load_lerobot_dataset.py b/wall_x/data/load_lerobot_dataset.py new file mode 100644 index 0000000..d502417 --- /dev/null +++ b/wall_x/data/load_lerobot_dataset.py @@ -0,0 +1,496 @@ +""" +LeRobot Dataset Loader - Distributed Version +""" + +import torch +from torch.utils.data import DistributedSampler +from lerobot.datasets.lerobot_dataset import LeRobotDataset +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 + +T_co = TypeVar("T_co", covariant=True) + +CAMERA_KEY_MAPPINGS = { + "lerobot/aloha_mobile_cabinet": { + "observation.images.cam_high": "face_view", + "observation.images.cam_left_wrist": "left_wrist_view", + "observation.images.cam_right_wrist": "right_wrist_view", + }, +} + + +# 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, seed=42, rank=0, world_size=1): + self._dataset = dataset + 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.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 = CAMERA_KEY_MAPPINGS[self._dataset.meta.repo_id] + + def _vision_preprocess(self, frames): + processed_frames = [] + for key in self._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["observation.state"] + action = data["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, + 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 get_train_dataloader(self): + """ + Get distributed training dataloader + + Args: + rank: Current process rank + world_size: Total number of processes + seed: Random seed for reproducibility + """ + + 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._dataset.meta.stats), + 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) + """ + + 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._dataset.meta.stats), + 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, stats): + self.config = config + self.dataload_config = dataload_config + self.stats = stats + self.min_stat = stats["action"]["min"] + self.max_stat = stats["action"]["max"] + self.delta = self.max_stat - self.min_stat + self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) + self.load_processor() + + def load_processor(self): + processor_path = self.config["processor_path"] + action_tokenizer_path = self.config["action_tokenizer_path"] + + # Use cached processors if available + if processor_path not in self._processor_cache: + self._processor_cache[processor_path] = AutoProcessor.from_pretrained(processor_path, use_fast=True) + if self.config.get("padding_side", "left") == "left": + self._processor_cache[processor_path].tokenizer.padding_side = "left" + + if 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) + + self.processor = self._processor_cache[processor_path] + self.val_processor = self._processor_cache[processor_path] + self.train_action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path] + self.val_action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path] + + new_tokens = ["<|propri|>", "<|action|>"] + + new_tokens += [f"<|action_token_{i}|>" for i in range(self.train_action_tokenizer.vocab_size)] + if not self.use_fast_tokenizer: + self.train_action_tokenizer = None + self.val_action_tokenizer = None + + # Only add tokens if not already added + if "<|propri|>" not in self.processor.tokenizer.get_vocab(): + num_added_tokens = self.processor.tokenizer.add_tokens(new_tokens) + self.val_processor.tokenizer.add_tokens(new_tokens) + + if self.use_fast_tokenizer: + self.action_mapper = {} + for i in range(self.train_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 + else: + self.action_mapper = None + + @classmethod + def _normalize(cls, action, min_stat, 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() + agent_pos.nan_to_num_(nan=0.0) + agent_pos = self._normalize(agent_pos, self.min_stat, self.delta) + 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 + ) + 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 = self._normalize(action, self.min_stat, self.delta) + 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) + 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, + ["x2_normal"] * 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"] = ["x2_normal"] * inputs["action_chunk"].shape[0] + + return inputs + + +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) + + dataset_fps = 50 + dataload_config = get_data_configs(config["data"]) + + delta_timestamps = { + # action chunk + "action": [t / dataset_fps for t in range(dataload_config.get("action_horizon", 32) - 1)], + } + batch_size = config.get("batch_size_per_gpu", 8) + + # repo_id = "lerobot/aloha_mobile_cabinet" + repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet") + dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps, video_backend="pyav") + + if rank == 0: + 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 = PreprocessedDataset(dataset, config, dataload_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, seed=42): + super().__init__(dataset, config, dataload_config, seed=seed, rank=0, world_size=1) + + 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._dataset.meta.stats), + ) + + return dataloader + +def load_test_dataset( + config, + lerobot_config, + 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) + + dataset_fps = 50 + dataload_config = get_data_configs(config["data"]) + + delta_timestamps = { + # action chunk + "action": [t / dataset_fps for t in range(dataload_config.get("action_horizon", 32) - 1)], + } + + repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet") + dataset = LeRobotDataset(repo_id, episodes=[episode], delta_timestamps=delta_timestamps, video_backend="pyav") + + 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, seed=seed) + + return dataset \ No newline at end of file diff --git a/wall_x/data/utils.py b/wall_x/data/utils.py new file mode 100644 index 0000000..cc3517f --- /dev/null +++ b/wall_x/data/utils.py @@ -0,0 +1,560 @@ +""" +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 re +import torch +import random +from collections import OrderedDict +from typing import List, Dict, Any, Optional, Union, Tuple +from transformers import BatchFeature + + +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", +} + + +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", +] + +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 preprocesser_call( + processor, + images: Optional[Union[List, Any]] = None, + text: Optional[Union[str, List[str]]] = None, + videos: Optional[Union[List, Any]] = None, + padding: Union[bool, str] = False, + truncation: Optional[bool] = None, + max_length: Optional[int] = None, + return_tensors: str = "pt", +) -> BatchFeature: + """Unified preprocessing function for Wall-X model handling text, image and video inputs. + + Processes inputs into format suitable for multimodal transformer models, including: + - Text tokenization and special token handling + - Image/video processing through image processor + - Attention mask and label generation + - Padding and truncation handling + + Args: + processor: Multimodal processor containing tokenizer and image processor + images: Input images (PIL, numpy arrays, or torch tensors) + text: Text or list of texts to tokenize + videos: Input videos (numpy arrays or torch tensors) + padding: Whether to pad sequences to same length + truncation: Whether to truncate sequences longer than max_length + max_length: Maximum length for truncation/padding + return_tensors: Format for returned tensors ('pt', 'np', etc.) + + Returns: + BatchFeature containing processed inputs with keys: + - input_ids: Tokenized text + - attention_mask: Attention mask for text + - pixel_values: Processed image pixels + - pixel_values_videos: Processed video frames + - image_grid_thw: Image grid dimensions for LLM + - video_grid_thw: Video grid dimensions for LLM + - labels: Training labels with masking + """ + # Process image inputs + if images is not None and len(images) > 0: + image_inputs = processor.image_processor( + images=images, videos=None, 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 + ) + video_grid_thw = videos_inputs["video_grid_thw"] + else: + videos_inputs = {} + video_grid_thw = None + + # Ensure text input is in list format + if not isinstance(text, list): + text = [text] + + # Process image placeholder tokens in text + if image_grid_thw is not None: + merge_length = processor.image_processor.merge_size ** 2 + index = 0 + for i in range(len(text)): + 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") + break + # Replace image placeholder with actual token count + token_count = image_grid_thw[index].prod() // merge_length + text[i] = text[i].replace( + "<|image_pad|>", "<|placeholder|>" * token_count, 1 + ) + index += 1 + text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>") + + # Process video placeholder tokens in text + if video_grid_thw is not None: + merge_length = processor.image_processor.merge_size ** 2 + index = 0 + for i in range(len(text)): + while "<|video_pad|>" in text[i]: + # Replace video placeholder with actual token count + token_count = video_grid_thw[index].prod() // merge_length + text[i] = text[i].replace( + "<|video_pad|>", "<|placeholder|>" * token_count, 1 + ) + index += 1 + text[i] = text[i].replace("<|placeholder|>", "<|video_pad|>") + + # Tokenize complete input text + text_inputs = processor.tokenizer( + text, + return_tensors=return_tensors, + padding=padding, + truncation=truncation, + max_length=max_length + ) + + # Get pad token ID for label generation + pad_token_id = processor.tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = processor.tokenizer.eos_token_id + + # Generate labels for multi-turn dialogue, keeping only assistant response loss + labels = torch.full_like(text_inputs.input_ids, -100) + assistant_marker = "<|im_start|>assistant\n" + im_end_token_id = processor.tokenizer.convert_tokens_to_ids("<|im_end|>") + assistant_tokens = processor.tokenizer( + "<|im_start|>assistant\n", add_special_tokens=False + ).input_ids + + for i in range(len(text)): + assistant_regions = [] + parts = text[i].split(assistant_marker) + + # Process each part to determine which tokens belong to assistant responses + # Count left padding tokens + num_left_pads = 0 + for token_id in text_inputs.input_ids[i]: + if token_id == pad_token_id: + num_left_pads += 1 + else: + break + current_pos = num_left_pads + + for j, part in enumerate(parts): + part_tokens = processor.tokenizer(part, add_special_tokens=False).input_ids + if j == 0: + # First part is system prompt or user question, all labels are -100 + current_pos += len(part_tokens) + continue + + # From second part onwards, each part starts with assistant response + for k in range(current_pos + 1, len(text_inputs.input_ids[i])): + if text_inputs.input_ids[i][k] == im_end_token_id: + assistant_regions.append(( + current_pos + len(assistant_tokens), k + 2 + )) + break + current_pos += len(part_tokens) + 3 + + # Set labels for assistant response regions + for start, end in assistant_regions: + labels[i][start:end] = text_inputs.input_ids[i][start:end] + + # Mask special action tokens in labels + action_token_id = processor.tokenizer.encode("<|action|>")[0] + propri_token_id = processor.tokenizer.encode("<|propri|>")[0] + labels[labels == action_token_id] = -100 + labels[labels == propri_token_id] = -100 + labels[labels == processor.tokenizer.pad_token_id] = -100 + + # Set labels to None if all are invalid to skip cross entropy loss + if (labels != -100).any().item(): + text_inputs["labels"] = labels + else: + 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 diff --git a/wall_x/fusions/__init__.py b/wall_x/fusions/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/wall_x/fusions/backend.py b/wall_x/fusions/backend.py new file mode 100644 index 0000000..b386be6 --- /dev/null +++ b/wall_x/fusions/backend.py @@ -0,0 +1,270 @@ +""" +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) \ No newline at end of file diff --git a/wall_x/fusions/ops.py b/wall_x/fusions/ops.py new file mode 100644 index 0000000..defb39c --- /dev/null +++ b/wall_x/fusions/ops.py @@ -0,0 +1,397 @@ +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!") + expert_for_rows = expert_for_rows.cuda() + + # 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) diff --git a/wall_x/model/__init__.py b/wall_x/model/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/wall_x/model/action_head.py b/wall_x/model/action_head.py new file mode 100755 index 0000000..5a29631 --- /dev/null +++ b/wall_x/model/action_head.py @@ -0,0 +1,400 @@ + +import math +import torch +import torch.nn as nn +from torch.distributions import Beta +from wall_x.utils.constant import action_statistic_dof + +class Normalizer(nn.Module): + """ + Action data normalizer for multi-robot systems. + + This module handles normalization and denormalization of action data for different robot + configurations. It maintains per-robot statistics (min values and deltas) and applies + normalization to map actions to the [-1, 1] range. + """ + + def __init__(self, action_statistic_dof, dof_config): + """ + Initialize the normalizer with robot-specific action statistics. + + Args: + action_statistic_dof (dict): Statistical data for each robot's degrees of freedom + dof_config (dict): Configuration mapping for degrees of freedom per robot + """ + super(Normalizer, self).__init__() + + action_statistic = {} + + # Process statistics for each robot + for robot_name in action_statistic_dof.keys(): + action_statistic[robot_name] = {} + all_dof_min = [] + all_dof_delta = [] + + # Collect min and delta values for all DOFs + for k in dof_config: + if k in action_statistic_dof[robot_name]: + all_dof_min.extend(action_statistic_dof[robot_name][k]["min"]) + all_dof_delta.extend(action_statistic_dof[robot_name][k]["delta"]) + else: + # Use default values if statistics not available + 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"] = all_dof_min + action_statistic[robot_name]["delta"] = all_dof_delta + + # Register statistics as non-trainable parameters + self.min = nn.ParameterDict({ + k: nn.Parameter(action_statistic[k]["min"], requires_grad=False) + for k in action_statistic.keys() + }) + self.delta = nn.ParameterDict({ + k: nn.Parameter(action_statistic[k]["delta"], requires_grad=False) + for k in action_statistic.keys() + }) + + def normalize_data(self, xs, dataset_names): + """ + Normalize action data to [-1, 1] range using robot-specific statistics. + + Args: + xs: Input action data tensors + dataset_names: List of dataset/robot names corresponding to each tensor + + Returns: + torch.Tensor: Normalized action data in [-1, 1] range + """ + new_xs = [] + # Filter out multimodal dataset entries + dataset_names = [name for name in dataset_names if name != "x2_multimodal"] + + for x, dataset_name in zip(xs, dataset_names): + # Apply min-max normalization + x = (x - self.min[dataset_name]) / (self.delta[dataset_name]) + # Scale to [-1, 1] range + x = x * 2 - 1 + # Clamp to ensure bounds + 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): + """ + Convert normalized data back to original action space. + + Args: + xs: Normalized action data in [-1, 1] range + dataset_names: List of dataset/robot names + dof_mask: Optional mask to select specific degrees of freedom + + Returns: + torch.Tensor: Denormalized action data in original scale + """ + new_xs = [] + # Filter out multimodal dataset entries + 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): + # Convert from [-1, 1] to [0, 1] range + x = (x + 1) / 2 + + # Apply DOF mask if provided + 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] + + # Scale back to original range + 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): + """ + Sinusoidal positional embedding for diffusion timesteps. + + Generates sinusoidal embeddings commonly used in diffusion models to encode + timestep information with different frequencies. + """ + + def __init__(self, dim): + """ + Initialize sinusoidal positional embedding. + + Args: + dim (int): Embedding dimension (must be even) + """ + super().__init__() + self.dim = dim + + def forward(self, x): + """ + Generate sinusoidal embeddings for input timesteps. + + Args: + x (torch.Tensor): Input timesteps + + Returns: + torch.Tensor: Sinusoidal embeddings of shape (..., dim) + """ + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device) * -emb) + emb = x[:, None] * emb[None, :] + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + + +class ActionProcessor(nn.Module): + """ + Action sequence processor for robotic control with flow matching. + + This module handles action sequence processing for robotic systems with the following capabilities: + 1. Adds controlled noise to action sequences using Beta distribution scheduling + 2. Generates temporal embeddings for timestep conditioning + 3. Projects actions to model hidden space for transformer processing + 4. Supports proprioceptive data integration and multi-robot configurations + + The Beta distribution provides more flexible noise injection strategies compared to + traditional linear schedules, allowing better control over the noise scheduling process. + """ + + def __init__(self, config): + """ + Initialize the action processor with multi-robot support. + + Args: + config: Configuration object containing: + - dof_config (dict): Degrees of freedom configuration per robot type + - agent_pos_config (dict): Agent position/proprioception configuration + - hidden_size (int): Model hidden layer dimension + - noise_scheduler (dict): Noise scheduler configuration with Beta parameters + """ + super().__init__() + + # Calculate action and proprioception dimensions from configuration + 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()]) + + # Log configuration details for debugging + print("ActionProcessor Configuration:", flush=True) + print(f" Action dimension: {self.action_dim}", flush=True) + print(f" Proprioception dimension: {self.propri_dim}", flush=True) + print(" DOF configuration:", flush=True) + for key, value in self.dof_config.items(): + print(f" {key}: {value}", flush=True) + print(" Agent position configuration:", flush=True) + for key, value in self.agent_pos_config.items(): + print(f" {key}: {value}", flush=True) + + self.hidden_size = config.hidden_size + + # Initialize data normalizers for actions and proprioception + self.normalizer_action = Normalizer(action_statistic_dof, config.dof_config) + self.normalizer_propri = Normalizer(action_statistic_dof, config.agent_pos_config) + + # Proprioception projection layer (includes history/current state) + self.propri_proj = nn.Linear(self.propri_dim * 2, self.hidden_size, bias=False) + + # Beta distribution noise scheduler configuration + noise_scheduler_config = config.noise_scheduler + self.beta_alpha = noise_scheduler_config.get('beta_alpha', 1.5) # Beta distribution α parameter + self.beta_beta = noise_scheduler_config.get('beta_beta', 1.0) # Beta distribution β parameter + self.s = noise_scheduler_config.get('s', 0.999) # Scaling factor + + # Initialize Beta distribution for noise scheduling + 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) + + # Sinusoidal positional embedding for timesteps + self.time_embed = SinusoidalPosEmb(config.hidden_size) + + # Action embedding network: project to hidden space + self.w1 = nn.Linear(self.action_dim * 2, self.hidden_size, bias=False) # *2 for action + DOF mask + self.w2 = nn.Linear(self.hidden_size * 2, self.hidden_size, bias=False) # *2 for action + time embeddings + self.w3 = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.act_fn = nn.SiLU() + + # Project back to action space for flow matching loss + self.action_proj_back = nn.Linear(self.hidden_size, self.action_dim, bias=False) + self.mse_loss = nn.MSELoss(reduction='none') + + def sample_time(self, batch_size, device, dtype): + """ + Sample timesteps using Beta distribution for noise scheduling. + + Generates random timesteps in [0,1] range using Beta distribution, then scales them. + This provides more flexible control over the noise injection schedule compared to + uniform sampling. + + Args: + batch_size (int): Number of timesteps to sample + device: Target device for tensors + dtype: Target data type for tensors + + Returns: + torch.Tensor: Sampled timesteps of shape [batch_size] + """ + sample = self.beta_dist.sample([batch_size]).to(device=device, dtype=dtype) + time = (self.s - sample) / self.s + return time + + def proprioception_proj(self, proprioception, dataset_names=None, dof_mask=None, use_history=False): + """ + Project proprioceptive data (joint positions, orientations) to hidden space. + + Args: + proprioception (torch.Tensor): Proprioceptive data of shape [batch_size, seq_len, propri_dim] + dataset_names (list, optional): Dataset names for normalization. Defaults to None. + dof_mask (torch.Tensor, optional): DOF mask of shape [batch_size, propri_dim]. Defaults to None. + use_history (bool, optional): Whether to use historical proprioceptive data. Defaults to False. + + Returns: + torch.Tensor: Projected proprioceptive features of shape [batch_size, seq_len, hidden_size] + """ + # Ensure proper device and dtype alignment + proprioception = proprioception.to(device=self.propri_proj.weight.device).to(dtype=self.propri_proj.weight.dtype) + + if dof_mask is not None: + # Concatenate proprioception with DOF mask + # TODO: Use variable-based dimension checking for better flexibility + if use_history: + proprioception = torch.cat([proprioception, dof_mask], dim=-1) + else: + proprioception = torch.cat([proprioception, dof_mask], dim=-1) + + proprioception = proprioception.to(device=self.propri_proj.weight.device).to(dtype=self.propri_proj.weight.dtype) + return self.propri_proj(proprioception) + + def forward(self, action_chunk, dataset_names, dof_mask=None): + """ + Process action sequences with noise injection and temporal embedding. + + This method implements the forward pass for flow matching training: + 1. Adds Beta-distributed noise to action sequences + 2. Generates sinusoidal timestep embeddings + 3. Projects noisy actions to hidden space + 4. Combines action and temporal features + + Args: + action_chunk (torch.Tensor): Action sequences of shape [batch_size, seq_len, action_dim] + dataset_names (list): Dataset names for normalization + dof_mask (torch.Tensor, optional): DOF mask of shape [batch_size, seq_len, action_dim]. + Defaults to None. + + Returns: + tuple: (action_embeddings, flow_target) where: + - action_embeddings: Processed action features of shape [batch_size, seq_len, hidden_size] + - flow_target: Flow matching target (action_chunk - noise) for loss computation + """ + batch_size = action_chunk.shape[0] + device = action_chunk.device + dtype = action_chunk.dtype + + # 1. Add noise to action sequences using flow matching + noise = torch.randn_like(action_chunk) + time = self.sample_time(batch_size, device, dtype) + t = time.unsqueeze(-1).unsqueeze(-1) # Broadcast to match action dimensions + + # Linear interpolation between noise and action (flow matching) + noisy_action = (1 - t) * noise + t * action_chunk + flow = action_chunk - noise # Flow target for loss computation + + # 2. Generate sinusoidal positional encoding for timesteps + time_embed = self.time_embed(time) + + # 3. Project noisy actions with DOF mask to hidden space + 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) + + # Repeat time embedding for each sequence position + time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1).to(dtype=self.w2.weight.dtype) + + # Combine action and temporal embeddings + concat_embed = torch.cat([action_embed, time_embed], dim=-1) + concat_embed = self.w2(concat_embed) + embed = self.w3(self.act_fn(concat_embed)) + + return embed, flow + + def step(self, timestep, noisy_action, dof_mask=None): + """ + Single denoising step for diffusion inference. + + Processes noisy actions at a specific timestep for iterative denoising during inference. + + Args: + timestep (torch.Tensor): Current timesteps of shape [batch_size] + noisy_action (torch.Tensor): Noisy actions of shape [batch_size, seq_len, action_dim] + dof_mask (torch.Tensor, optional): DOF mask for action space. Defaults to None. + + Returns: + torch.Tensor: Processed action embeddings of shape [batch_size, seq_len, hidden_size] + """ + # Concatenate noisy action with DOF mask if provided + if dof_mask is not None: + noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) + + # Generate timestep embeddings + time_embed = self.time_embed(timestep) # [batch_size, hidden_size] + + # Project noisy actions + action_embed = self.w1(noisy_action) + + # Broadcast time embeddings to sequence length + 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) + + # Combine embeddings and process through MLP + concat_embed = torch.cat([action_embed, time_embed], dim=-1) + concat_embed = self.w2(concat_embed) + embed = self.w3(self.act_fn(concat_embed)) + + return embed + + def flow_loss(self, action_hidden_states, flow, dof_mask=None): + """ + Compute flow matching loss between predicted and target actions. + + Args: + action_hidden_states (torch.Tensor): Hidden states from transformer + flow (torch.Tensor): Target flow (action - noise) for matching + dof_mask (torch.Tensor, optional): DOF mask to weight loss per dimension. Defaults to None. + + Returns: + torch.Tensor: Flow matching loss (no reduction for channel loss computation) + """ + # Project hidden states back to action space + action_pred = self.action_proj_back(action_hidden_states) + + # Compute MSE loss between predicted and target flow + loss = self.mse_loss(action_pred, flow) + + # Apply DOF mask if provided + if dof_mask is not None: + dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]) + loss = loss * dof_mask + + # Return loss without reduction for channel-wise loss computation + return loss \ No newline at end of file diff --git a/wall_x/model/qwen2_5_based/__init__.py b/wall_x/model/qwen2_5_based/__init__.py new file mode 100644 index 0000000..dd5b053 --- /dev/null +++ b/wall_x/model/qwen2_5_based/__init__.py @@ -0,0 +1,2 @@ +from .modeling_qwen2_5_vl_act import Qwen2_5_VLMoEModel,Qwen2_5_VLMoEForAction +from .configuration_qwen2_5_vl import Qwen2_5_VLConfig \ No newline at end of file diff --git a/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py b/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py new file mode 100644 index 0000000..8420b80 --- /dev/null +++ b/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py @@ -0,0 +1,248 @@ +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_rope_utils import rope_config_validation + + +class Qwen2_5_VLVisionConfig(PretrainedConfig): + model_type = "qwen2_5_vl" + base_config_key = "vision_config" + + def __init__( + self, + depth=32, + hidden_size=3584, + hidden_act="silu", + intermediate_size=3420, + num_heads=16, + in_channels=3, + patch_size=14, + spatial_merge_size=2, + temporal_patch_size=2, + tokens_per_second=4, + window_size=112, + out_hidden_size=3584, + fullatt_block_indexes=[7, 15, 23, 31], + **kwargs, + ): + super().__init__(**kwargs) + + self.depth = depth + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.intermediate_size = intermediate_size + self.num_heads = num_heads + self.in_channels = in_channels + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + self.temporal_patch_size = temporal_patch_size + self.tokens_per_second = tokens_per_second + self.window_size = window_size + self.fullatt_block_indexes = fullatt_block_indexes + self.out_hidden_size = out_hidden_size + + +class Qwen2_5_VLConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`Qwen2_5_VLModel`]. It is used to instantiate a + Qwen2-VL model according to the specified arguments, defining the model architecture. Instantiating a configuration + with the defaults will yield a similar configuration to that of + Qwen2-VL-7B-Instruct [Qwen/Qwen2-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct). + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 152064): + Vocabulary size of the Qwen2_5_VL model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`Qwen2_5_VLModel`] + hidden_size (`int`, *optional*, defaults to 8192): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 29568): + Dimension of the MLP representations. + num_hidden_layers (`int`, *optional*, defaults to 80): + Number of hidden layers in the Transformer encoder. + num_attention_heads (`int`, *optional*, defaults to 64): + Number of attention heads for each attention layer in the Transformer encoder. + num_key_value_heads (`int`, *optional*, defaults to 8): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 32768): + The maximum sequence length that this model might ever be used with. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + rms_norm_eps (`float`, *optional*, defaults to 1e-05): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether the model's input and output word embeddings should be tied. + rope_theta (`float`, *optional*, defaults to 1000000.0): + The base period of the RoPE embeddings. + use_sliding_window (`bool`, *optional*, defaults to `False`): + Whether to use sliding window attention. + sliding_window (`int`, *optional*, defaults to 4096): + Sliding window attention (SWA) window size. If not specified, will default to `4096`. + max_window_layers (`int`, *optional*, defaults to 80): + The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + vision_config (`Dict`, *optional*): + The config for the visual encoder initialization. + rope_scaling (`Dict`, *optional*): + Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type + and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value + accordingly. + Expected contents: + `rope_type` (`str`): + The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', + 'llama3'], with 'default' being the original RoPE implementation. + `factor` (`float`, *optional*): + Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In + most scaling types, a `factor` of x will enable the model to handle sequences of length x * + original maximum pre-trained length. + `original_max_position_embeddings` (`int`, *optional*): + Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during + pretraining. + `attention_factor` (`float`, *optional*): + Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention + computation. If unspecified, it defaults to value recommended by the implementation, using the + `factor` field to infer the suggested value. + `beta_fast` (`float`, *optional*): + Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear + ramp function. If unspecified, it defaults to 32. + `beta_slow` (`float`, *optional*): + Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear + ramp function. If unspecified, it defaults to 1. + `short_factor` (`List[float]`, *optional*): + Only used with 'longrope'. The scaling factor to be applied to short contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `long_factor` (`List[float]`, *optional*): + Only used with 'longrope'. The scaling factor to be applied to long contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `low_freq_factor` (`float`, *optional*): + Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE + `high_freq_factor` (`float`, *optional*): + Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE + + ```python + >>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig + + >>> # Initializing a Qwen2_5_VL style configuration + >>> configuration = Qwen2_5_VLConfig() + + >>> # Initializing a model from the Qwen2-VL-7B style configuration + >>> model = Qwen2_5_VLForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "qwen2_5_vl" + sub_configs = {"vision_config": Qwen2_5_VLVisionConfig} + keys_to_ignore_at_inference = ["past_key_values"] + # Default tensor parallel plan for base model `Qwen2_5_VL` + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + def __init__( + self, + vocab_size=152064, + hidden_size=8192, + intermediate_size=29568, + num_hidden_layers=80, + num_attention_heads=64, + num_key_value_heads=8, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-05, + use_cache=True, + tie_word_embeddings=False, + rope_theta=1000000.0, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=80, + attention_dropout=0.0, + vision_config=None, + rope_scaling=None, + num_experts=4, + experts=None, + dof_config=None, + noise_scheduler=None, + dim_inputs=(1536,1536), + attention_moe=False, + mlp_moe=False, + **kwargs, + ): + if isinstance(vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**vision_config) + elif vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.use_sliding_window = use_sliding_window + self.sliding_window = sliding_window + self.max_window_layers = max_window_layers + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.attention_dropout = attention_dropout + self.rope_scaling = rope_scaling + + self.num_experts = num_experts + self.experts = experts + self.dof_config = dof_config + self.noise_scheduler = noise_scheduler + self.dim_inputs = tuple(dim_inputs) + self.attention_moe = attention_moe + self.mlp_moe = mlp_moe + + # 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 + # one can set it to "linear"/"dynamic" etc. to have scaled RoPE + # TODO: @raushan update config in the hub + if self.rope_scaling is not None and "type" in self.rope_scaling: + if self.rope_scaling["type"] == "mrope": + self.rope_scaling["type"] = "default" + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + rope_config_validation(self, ignore_keys={"mrope_section"}) + + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + + +__all__ = ["Qwen2_5_VLConfig"] diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py new file mode 100644 index 0000000..e0d891f --- /dev/null +++ b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py @@ -0,0 +1,2078 @@ + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from dataclasses import dataclass +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 +from transformers.generation import GenerationMixin +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + is_torchdynamo_compiling, + logging, + replace_return_docstrings, +) +from .configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig +from wall_x.fusions import ops + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_varlen_func + from flash_attn.layers.rotary import apply_rotary_emb +else: + flash_attn_varlen_func = None + apply_rotary_emb = None + + +if is_flash_attn_2_available(): + from transformers.modeling_flash_attention_utils import _flash_attention_forward +else: + flash_attn_varlen_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): + 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) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_state): + return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) + + +class Qwen2_5_VisionPatchEmbed(nn.Module): + def __init__( + self, + patch_size: int = 14, + temporal_patch_size: int = 2, + in_channels: int = 3, + embed_dim: int = 1152, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.temporal_patch_size = temporal_patch_size + self.in_channels = in_channels + self.embed_dim = embed_dim + + kernel_size = [temporal_patch_size, patch_size, patch_size] + self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_dtype = self.proj.weight.dtype + hidden_states = hidden_states.view( + -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) + return hidden_states + + +class Qwen2_5_VisionRotaryEmbedding(nn.Module): + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, seqlen: int) -> torch.Tensor: + seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + freqs = torch.outer(seq, self.inv_freq) + return freqs + + +class Qwen2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Qwen2RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + 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 + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class Qwen2_5_VLPatchMerger(nn.Module): + def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: + super().__init__() + self.hidden_size = context_dim * (spatial_merge_size**2) + self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6) + self.mlp = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.mlp(self.ln_q(x).view(-1, self.hidden_size)) + return x + + +def apply_rotary_pos_emb_flashatt( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + cos = cos.chunk(2, dim=-1)[0].contiguous() + sin = sin.chunk(2, dim=-1)[0].contiguous() + q_embed = apply_rotary_emb(q.float(), cos.float(), sin.float()).type_as(q) + k_embed = apply_rotary_emb(k.float(), cos.float(), sin.float()).type_as(k) + return q_embed, k_embed + + +class Qwen2_5_VLVisionFlashAttention2(nn.Module): + def __init__(self, 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) + + def forward( + 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: + seq_length = hidden_states.shape[0] + q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `rotary_pos_emb` (2D tensor of RoPE theta values), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be " + "removed and `position_embeddings` will be mandatory." + ) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + cos = emb.cos().float() + sin = emb.sin().float() + else: + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb_flashatt(q.unsqueeze(0), k.unsqueeze(0), cos, sin) + q = q.squeeze(0) + k = k.squeeze(0) + + 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).reshape( + seq_length, -1 + ) + attn_output = self.proj(attn_output) + return attn_output + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb_vision( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + orig_q_dtype = q.dtype + orig_k_dtype = k.dtype + q, k = q.float(), k.float() + cos, sin = cos.unsqueeze(-2), sin.unsqueeze(-2) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + q_embed = q_embed.to(orig_q_dtype) + k_embed = k_embed.to(orig_k_dtype) + return q_embed, k_embed + + +class Qwen2_5_VLVisionAttention(nn.Module): + def __init__(self, dim: int, num_heads: int = 16) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.qkv = nn.Linear(dim, dim * 3, bias=True) + self.proj = nn.Linear(dim, dim) + + def forward( + 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: + seq_length = hidden_states.shape[0] + q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `rotary_pos_emb` (2D tensor of RoPE theta values), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be " + "removed and `position_embeddings` will be mandatory." + ) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + cos = emb.cos().float() + sin = emb.sin().float() + else: + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) + + attention_mask = torch.full( + [1, seq_length, seq_length], torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype + ) + for i in range(1, len(cu_seqlens)): + attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0 + + q = q.transpose(0, 1) + k = k.transpose(0, 1) + v = v.transpose(0, 1) + attn_weights = torch.matmul(q, k.transpose(1, 2)) / math.sqrt(self.head_dim) + attn_weights = attn_weights + attention_mask + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype) + attn_output = torch.matmul(attn_weights, v) + attn_output = attn_output.transpose(0, 1) + attn_output = attn_output.reshape(seq_length, -1) + attn_output = self.proj(attn_output) + return attn_output + + +class Qwen2_5_VLVisionSdpaAttention(nn.Module): + def __init__(self, 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) + + def forward( + 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: + seq_length = hidden_states.shape[0] + q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `rotary_pos_emb` (2D tensor of RoPE theta values), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be " + "removed and `position_embeddings` will be mandatory." + ) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + cos = emb.cos().float() + sin = emb.sin().float() + else: + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) + + attention_mask = torch.zeros([1, seq_length, seq_length], device=q.device, dtype=torch.bool) + for i in range(1, len(cu_seqlens)): + attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = True + q = q.transpose(0, 1) + k = k.transpose(0, 1) + v = v.transpose(0, 1) + attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0) + attn_output = attn_output.transpose(0, 1) + attn_output = attn_output.reshape(seq_length, -1) + attn_output = self.proj(attn_output) + return attn_output + + +QWEN2_5_VL_VISION_ATTENTION_CLASSES = { + "eager": Qwen2_5_VLVisionAttention, + "flash_attention_2": Qwen2_5_VLVisionFlashAttention2, + "sdpa": Qwen2_5_VLVisionSdpaAttention, +} + + +class Qwen2_5_VLVisionBlock(nn.Module): + def __init__(self, config, attn_implementation: str = "sdpa") -> None: + super().__init__() + 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 + ) + self.mlp = Qwen2_5_VLMLP(config, bias=True) + + def forward( + 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: + hidden_states = hidden_states + self.attn( + self.norm1(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)) + return hidden_states + + +Qwen2_5_VL_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`Qwen2_5_VLConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "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): + config_class = Qwen2_5_VLConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + _supports_static_cache = False # TODO (joao): fix. torch.compile failing probably due to `cache_positions` + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, (nn.Linear, nn.Conv3d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +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: + super().__init__(config, *inputs, **kwargs) + self.spatial_merge_size = config.spatial_merge_size + self.patch_size = config.patch_size + self.fullatt_block_indexes = config.fullatt_block_indexes + self.window_size = config.window_size + self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + + self.patch_embed = Qwen2_5_VisionPatchEmbed( + patch_size=config.patch_size, + temporal_patch_size=config.temporal_patch_size, + in_channels=config.in_channels, + embed_dim=config.hidden_size, + ) + + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList( + [Qwen2_5_VLVisionBlock(config, config._attn_implementation) for _ in range(config.depth)] + ) + self.merger = Qwen2_5_VLPatchMerger( + dim=config.out_hidden_size, + context_dim=config.hidden_size, + spatial_merge_size=config.spatial_merge_size, + ) + self.gradient_checkpointing = False + + def rot_pos_emb(self, grid_thw): + pos_ids = [] + for t, h, w in grid_thw: + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + hpos_ids = hpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + hpos_ids = hpos_ids.permute(0, 2, 1, 3) + hpos_ids = hpos_ids.flatten() + + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + wpos_ids = wpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + wpos_ids = wpos_ids.permute(0, 2, 1, 3) + wpos_ids = wpos_ids.flatten() + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = grid_thw[:, 1:].max() + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb + + def get_window_index(self, grid_thw): + window_index: list = [] + cu_window_seqlens: list = [0] + window_index_id = 0 + vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size + + for grid_t, grid_h, grid_w in grid_thw: + llm_grid_h, llm_grid_w = ( + grid_h // self.spatial_merge_size, + grid_w // self.spatial_merge_size, + ) + index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w) + pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size + pad_w = vit_merger_window_size - llm_grid_w % 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", -100) + index_padded = index_padded.reshape( + grid_t, + num_windows_h, + vit_merger_window_size, + num_windows_w, + vit_merger_window_size, + ) + index_padded = index_padded.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 != -100).sum([2, 3]).reshape(-1) + index_padded = index_padded.reshape(-1) + index_new = index_padded[index_padded != -100] + window_index.append(index_new + window_index_id) + cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1] + cu_window_seqlens.extend(cu_seqlens_tmp.tolist()) + window_index_id += (grid_t * llm_grid_h * llm_grid_w).item() + window_index = torch.cat(window_index, dim=0) + + return window_index, cu_window_seqlens + + def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + """ + Args: + hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): + The final hidden states of the model. + grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): + The temporal, height and width of feature shape of each image in LLM. + + Returns: + `torch.Tensor`: hidden_states. + """ + hidden_states = self.patch_embed(hidden_states) + rotary_pos_emb = self.rot_pos_emb(grid_thw) + window_index, cu_window_seqlens = self.get_window_index(grid_thw) + window_index = window_index.to(hidden_states.device) + cu_window_seqlens = torch.tensor( + cu_window_seqlens, + device=hidden_states.device, + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) + + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + hidden_states = hidden_states[window_index, :, :] + hidden_states = hidden_states.reshape(seq_len, -1) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + rotary_pos_emb = rotary_pos_emb[window_index, :, :] + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, + # Select dtype based on the following factors: + # - FA2 requires that cu_seqlens_q must have dtype int32 + # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw + # See https://github.com/huggingface/transformers/pull/34852 for more information + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + max_seqlen_full = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + max_seqlen_window = (cu_window_seqlens[1:] - cu_window_seqlens[:-1]).max().item() + + for layer_num, blk in enumerate(self.blocks): + if layer_num in self.fullatt_block_indexes: + cu_seqlens_now = cu_seqlens + max_seqlen_now = max_seqlen_full + else: + cu_seqlens_now = cu_window_seqlens + max_seqlen_now = max_seqlen_window + if self.gradient_checkpointing and self.training: + hidden_states = self._gradient_checkpointing_func( + blk.__call__, hidden_states, cu_seqlens_now, None, position_embeddings + ) + else: + hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens_now, max_seqlen=max_seqlen_now, position_embeddings=position_embeddings) + + hidden_states = self.merger(hidden_states) + reverse_indices = torch.argsort(window_index) + hidden_states = hidden_states[reverse_indices, :] + + return hidden_states + + +class Qwen2_5_VLRotaryEmbedding(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, device=None): + super().__init__() + # 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")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + 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 + + def _dynamic_frequency_update(self, position_ids, device): + """ + dynamic RoPE layers should recompute `inv_freq` in the following situations: + 1 - growing beyond the cached sequence length (allow scaling) + 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) + """ + seq_len = torch.max(position_ids) + 1 + if seq_len > self.max_seq_len_cached: # growth + inv_freq, self.attention_scaling = self.rope_init_fn( + self.config, device, seq_len=seq_len, **self.rope_kwargs + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation + self.max_seq_len_cached = seq_len + + if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + self.max_seq_len_cached = self.original_max_seq_len + + @torch.no_grad() + def forward(self, x, position_ids): + if "dynamic" in self.rope_type: + self._dynamic_frequency_update(position_ids, device=x.device) + + # Core RoPE block. In contrast to other models, Qwen2_5_VL has different position ids for thw grids + # So we expand the inv_freq to shape (3, ...) + inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) + position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions) + # Force float32 (see https://github.com/huggingface/transformers/pull/29285) + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention + cos = cos * self.attention_scaling + sin = sin * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Qwen2MLP(nn.Module): + def __init__(self, config): + super().__init__() + 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.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)) + return down_proj + + +def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1): + """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/). + + Explanation: + Multimodal 3D rotary position embedding is an extension to 1D rotary position embedding. The input embedding + sequence contains vision (images / videos) embedding and text embedding or just contains text embedding. For + vision embedding part, we apply rotary position embedding on temporal, height and width dimension seperately. + Here we split the channel dimension to 3 chunks for the temporal, height and width rotary position embedding. + For text embedding part, we just apply 1D rotary position embedding. The three rotary position index (temporal, + height and width) of text embedding is always the same, so the text embedding rotary position embedding has no + difference with modern LLMs. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + mrope_section(`List(int)`): + Multimodal rope section is for channel dimension of temporal, height and width in rope calculation. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + mrope_section = mrope_section * 2 + cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze( + unsqueeze_dim + ) + sin = 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) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class Qwen2_5_VLAttention(nn.Module): + """ + Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer + and "Generating Long Sequences with Sparse Transformers". + """ + + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.is_causal = True + self.attention_dropout = config.attention_dropout + self.rope_scaling = config.rope_scaling + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + 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) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) + + 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 + ) -> 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) + + 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) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + # Fix precision issues in Qwen2-VL float16 inference + # Replace inf values with zeros in attention weights to prevent NaN propagation + if query_states.dtype == torch.float16: + attn_weights = torch.where(torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights) + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + +from flash_attn import flash_attn_func + +class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): + """ + Qwen2_5_VL flash attention module, following Qwen2_5_VL attention module. This module inherits from `Qwen2_5_VLAttention` + as the weights of the module stays untouched. The only required change would be on the forward pass + where it needs to correctly call the public API of flash attention and deal with padding tokens + in case the input contains any of them. Additionally, for sliding window attention, we apply SWA only to the bottom + config.max_window_layers layers. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # 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() + + 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() + + 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) + + # 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"]) + 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.q_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) + + attn_output = flash_attn_func( + query_states, key_states, value_states, dropout_rate, softmax_scale=None, causal=self.is_causal + ) + + 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 + + +class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): + """ + Qwen2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `Qwen2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from Qwen2Attention.forward + 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 + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if output_attentions: + # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. + logger.warning_once( + "Qwen2_5_VLModel is using Qwen2_5_VLSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " + 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + 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) + + 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) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + causal_mask = attention_mask + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if query_states.device.type == "cuda" and attention_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal = True if causal_mask is None and q_len > 1 else False + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=causal_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + return attn_output, None, past_key_value + + +QWEN2_5_VL_ATTENTION_CLASSES = { + "eager": Qwen2_5_VLAttention, + "flash_attention_2": Qwen2_5_VLFlashAttention2, + "sdpa": Qwen2_5_VLSdpaAttention, +} + + +class Qwen2_5_VLDecoderLayer(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + if config.use_sliding_window and config._attn_implementation != "flash_attention_2": + logger.warning_once( + f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " + "unexpected results may be encountered." + ) + self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx) + + self.mlp = Qwen2MLP(config) + self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. + position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +@add_start_docstrings( + "The bare Qwen2_5_VL Model outputting raw hidden-states without any specific head on top.", + Qwen2_5_VL_START_DOCSTRING, +) +class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): + def __init__(self, config: Qwen2_5_VLConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self._attn_implementation = config._attn_implementation + self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: torch.LongTensor = None, + 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, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + 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 + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # torch.jit.trace() doesn't support cache objects in the output + if use_cache and past_key_values is None and not torch.jit.is_tracing(): + past_key_values = DynamicCache() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + # the hard coded `3` is for temporal, height and width. + if position_ids is None: + position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) + elif position_ids.dim() == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and past_key_values is not None: + is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0] + if is_padding_right: + raise ValueError( + "You are attempting to perform batched generation with padding_side='right'" + " this may lead to unexpected behaviour for Flash Attention version of Qwen2_5_VL. Make sure to " + " call `tokenizer.padding_side = 'left'` before tokenizing the input. " + ) + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if ( + self.config._attn_implementation == "sdpa" + and not (using_static_cache or using_sliding_window_cache) + and not output_attentions + ): + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + sliding_window=self.config.sliding_window, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + # SlidingWindowCache or StaticCache + if using_sliding_window_cache or using_static_cache: + target_length = past_key_values.get_max_cache_shape() + # DynamicCache or no cache + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=target_length, + dtype=dtype, + device=device, + cache_position=cache_position, + batch_size=input_tensor.shape[0], + config=self.config, + past_key_values=past_key_values, + ) + + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type in ["cuda", "xpu"] + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + @staticmethod + def _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask: torch.Tensor, + sequence_length: int, + target_length: int, + dtype: torch.dtype, + device: torch.device, + cache_position: torch.Tensor, + batch_size: int, + config: Qwen2_5_VLConfig, + past_key_values: Cache, + ): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. + + Args: + attention_mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. + sequence_length (`int`): + The sequence length being processed. + target_length (`int`): + The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. + dtype (`torch.dtype`): + The dtype to use for the 4D attention mask. + device (`torch.device`): + The device to plcae the 4D attention mask on. + cache_position (`torch.Tensor`): + Indices depicting the position of the input sequence tokens in the sequence. + batch_size (`torch.Tensor`): + Batch size. + config (`Qwen2_5_VLConfig`): + The model's configuration class + past_key_values (`Cache`): + The cache class that is being used currently to generate + """ + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + causal_mask = attention_mask + else: + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + if config.sliding_window is not None: + # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also + # the check is needed to verify is current checkpoint was trained with sliding window or not + if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length: + sliding_attend_mask = torch.arange(target_length, device=device) <= ( + cache_position.reshape(-1, 1) - config.sliding_window + ) + diagonal_attend_mask.bitwise_or_(sliding_attend_mask) + causal_mask *= diagonal_attend_mask + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + if attention_mask.shape[-1] > target_length: + attention_mask = attention_mask[:, :target_length] + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( + causal_mask.device + ) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + return causal_mask + + +@dataclass +class Qwen2_5_VLCausalLMOutputWithPast(ModelOutput): + """ + Base class for Qwen2_5_VL causal language model (or autoregressive) outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + """ + + loss: Optional[torch.FloatTensor] = None + logits: torch.FloatTensor = None + past_key_values: Optional[List[torch.FloatTensor]] = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + rope_deltas: Optional[torch.LongTensor] = None + + +QWEN2_5_VL_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids) + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + pixel_values (`torch.FloatTensor` of shape `(seq_length, num_channels * image_size * image_size)): + The tensors corresponding to the input images. Pixel values can be obtained using + [`AutoImageProcessor`]. See [`Qwen2_5_VLImageProcessor.__call__`] for details. [`Qwen2_5_VLProcessor`] uses + [`Qwen2_5_VLImageProcessor`] for processing images. + pixel_values_videos (`torch.FloatTensor` of shape `(seq_length, num_channels * temporal_size * image_size * image_size)): + The tensors corresponding to the input videos. Pixel values can be obtained using + [`AutoImageProcessor`]. See [`Qwen2_5_VLImageProcessor.__call__`] for details. [`Qwen2_5_VLProcessor`] uses + [`Qwen2_5_VLImageProcessor`] for processing videos. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. +""" + + +class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + config_class = Qwen2_5_VLConfig + _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] + + def __init__(self, config): + super().__init__(config) + self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config) + self.model = Qwen2_5_VLModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.rope_deltas = None # cache rope_deltas here + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def get_rope_index( + self, + input_ids: Optional[torch.LongTensor] = None, + 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, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Calculate the 3D rope index based on image and video's temporal, height and width in LLM. + + Explanation: + Each embedding sequence contains vision embedding and text embedding or just contains text embedding. + + For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. + Examples: + input_ids: [T T T T T], here T is for text. + temporal position_ids: [0, 1, 2, 3, 4] + height position_ids: [0, 1, 2, 3, 4] + width position_ids: [0, 1, 2, 3, 4] + + For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part + and 1D rotary position embeddin for text part. + Examples: + Temporal (Time): 3 patches, representing different segments of the video in time. + Height: 2 patches, dividing each frame vertically. + Width: 2 patches, dividing each frame horizontally. + We also have some important parameters: + fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. + tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. + temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. + interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. + input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. + vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] + vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] + vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] + text temporal position_ids: [101, 102, 103, 104, 105] + text height position_ids: [101, 102, 103, 104, 105] + text width position_ids: [101, 102, 103, 104, 105] + Here we calculate the text start position_ids as the max vision position_ids plus 1. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ + 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 + mrope_position_deltas = [] + if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): + 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 in enumerate(total_input_ids): + input_ids = input_ids[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens 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 input_tokens 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: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + second_per_grid_t = 0 + image_index += 1 + remain_images -= 1 + ed = ed_image + + else: + 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 = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // 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).view(1, -1).expand(3, -1) + st_idx) + + range_tensor = torch.arange(llm_grid_t).view(-1, 1) + expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) + + time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second + + time_tensor_long = time_tensor.long() + t_index = time_tensor_long.flatten() + + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).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 + + 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).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=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + 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 + + @add_start_docstrings_to_model_forward(QWEN2_5_VL_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=Qwen2_5_VLCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + 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, + 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, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + ) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + >>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + + >>> messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": "What is shown in this image?"}, + ], + }, + ] + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos]) + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..." + ```""" + + 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 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) + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + + loss = None + if labels is not None: + # Upcast to float if we need to compute the loss to avoid potential precision issues + logits = logits.float() + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + 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) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return Qwen2_5_VLCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=self.rope_deltas, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + position_ids=None, + use_cache=True, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + second_per_grid_ts=None, + **kwargs, + ): + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model + + # 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. + if past_key_values is not None: + if inputs_embeds is not None and input_ids.shape[1] == 0: # Exception 4 + inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] + elif ( + inputs_embeds is not None # Exception 1 + or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1]) # Exception 3 + ): + input_ids = input_ids[:, -cache_position.shape[0] :] + elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) + input_ids = input_ids[:, cache_position] + + if cache_position[0] != 0: + pixel_values = None + pixel_values_videos = None + + # 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} + + 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 + device = inputs_embeds.device + else: + batch_size, sequence_length = input_ids.shape + device = input_ids.device + + attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=past_key_values.get_max_cache_shape(), + dtype=self.lm_head.weight.dtype, + device=device, + cache_position=cache_position, + batch_size=batch_size, + config=self.config, + past_key_values=past_key_values, + ) + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": use_cache, + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "pixel_values_videos": pixel_values_videos, + "image_grid_thw": image_grid_thw, + "video_grid_thw": video_grid_thw, + "cache_position": cache_position, + "second_per_grid_ts": second_per_grid_ts, + } + ) + return model_inputs + + def _get_image_nums_and_video_nums( + self, + input_ids: Optional[torch.LongTensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + 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` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Returns: + 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 + + 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 + image_nums = torch.sum(vision_first_mask & image_mask, dim=1) + video_nums = torch.sum(vision_first_mask & video_mask, dim=1) + + return image_nums, video_nums + + def _expand_inputs_for_generation( + self, + expand_size: int = 1, + is_encoder_decoder: bool = False, + input_ids: Optional[torch.LongTensor] = None, + **model_kwargs, + ) -> Tuple[torch.LongTensor, Dict[str, Any]]: + # 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) + + if expand_size == 1: + return input_ids, model_kwargs + + visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"] + + def _expand_dict_for_generation_visual(dict_to_expand): + 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): + samples = torch.split(x, lengths) + repeat_args = [repeat_times] + [1] * (x.dim() - 1) + result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0) + return result + + for key in dict_to_expand: + if key == "pixel_values": + # 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": + # 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": + 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": + 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": + 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." + ) + tensor = torch.tensor(dict_to_expand[key]) + lengths = list(video_nums) + tensor = _repeat_interleave_samples(tensor, lengths=lengths, repeat_times=expand_size) + dict_to_expand[key] = tensor.tolist() + return dict_to_expand + + def _expand_dict_for_generation(dict_to_expand): + for key in dict_to_expand: + if ( + key != "cache_position" + and dict_to_expand[key] is not None + and isinstance(dict_to_expand[key], torch.Tensor) + and key not in visual_keys + ): + dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0) + return dict_to_expand + + # 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) + + if input_ids is not None: + input_ids = input_ids.repeat_interleave(expand_size, dim=0) + + model_kwargs = _expand_dict_for_generation(model_kwargs) + + if is_encoder_decoder: + if model_kwargs.get("encoder_outputs") is None: + raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.") + model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"]) + + return input_ids, model_kwargs + + +__all__ = ["Qwen2_5_VLForConditionalGeneration", "Qwen2_5_VLModel", "Qwen2_5_VLPreTrainedModel"] diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py new file mode 100644 index 0000000..1a0d65b --- /dev/null +++ b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py @@ -0,0 +1,2007 @@ +import os +import torch +import numpy as np +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 + +from transformers import AutoConfig, AutoProcessor +from transformers.activations import ACT2FN +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 transformers.models.qwen2_vl.modeling_qwen2_vl import ( + Qwen2RMSNorm, +) +from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VLMLP, + Qwen2_5_VLRotaryEmbedding, + Qwen2_5_VLPreTrainedModel, + 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.qwen2_5_based.modeling_qwen2_5_vl import Qwen2_5_VisionTransformerPretrainedModel, Qwen2_5_VLAttention, Qwen2_5_VLFlashAttention2, Qwen2_5_VLSdpaAttention +from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES + + +logger = logging.get_logger(__name__) + + +@dataclass +class Qwen2_5_VLACausalLMOutputWithPast(ModelOutput): + loss: Optional[torch.FloatTensor] = None + flow_loss: Optional[torch.FloatTensor] = None + cross_entropy_loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[List[torch.FloatTensor]] = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + rope_deltas: Optional[torch.LongTensor] = None + + channel_loss_dict: Optional[dict[torch.FloatTensor]] = None + channel_loss_count_dict: Optional[dict[torch.FloatTensor]] = None + + +class BlockSparseMLP(nn.Module): + def __init__(self, config): + super().__init__() + + self.hidden_size = config["hidden_size"] + self.intermediate_size = config["intermediate_size"] + self.hidden_act = config["hidden_act"] + 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 forward(self, hidden_state): + return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) + +class SparseMoeBlock(nn.Module): + """Sparse Mixture of Experts (MoE) Block optimized with Grouped GEMM. + + This module implements a sparse MoE layer where tokens are dynamically routed + to different expert networks. Uses grouped GEMM operations for efficient + computation across multiple experts. + + Args: + config: Configuration object containing expert specifications + num_experts: Number of expert networks in the MoE block + """ + + def __init__(self, config, num_experts: int): + super().__init__() + self.num_experts = num_experts + # Initialize expert networks based on configuration + self.experts = nn.ModuleList([ + BlockSparseMLP(config.experts[i]) for i in range(num_experts) + ]) + + def forward( + self, + hidden_states: torch.Tensor, + experts_indices: torch.Tensor, + start_indices: torch.Tensor, + end_indices: torch.Tensor + ) -> torch.Tensor: + """Forward pass through the Sparse MoE block. + + Routes different hidden states to their corresponding expert networks + for processing. Uses efficient grouped operations to minimize overhead. + + Args: + hidden_states: Input tensor of shape (batch_size, seq_length, hidden_dim) + experts_indices: Expert assignment indices of shape (batch_size, seq_length) + indicating which expert each token should be routed to + start_indices: Starting indices for each expert's assigned tokens + end_indices: Ending indices for each expert's assigned tokens + + Returns: + output: Processed tensor of shape (batch_size, seq_length, hidden_dim) + after expert processing and token reordering + """ + batch_size, seq_length, hidden_dim = hidden_states.size() + + # Flatten inputs for efficient grouped processing + hidden_states = hidden_states.view(-1, hidden_dim) # [total_tokens, hidden_dim] + experts_indices = experts_indices.view(-1) # [total_tokens] + + # Create uniform probabilities for all tokens (can be modified for weighted routing) + probs = torch.ones_like(experts_indices, dtype=torch.float32).view(-1, 1) + + # Permute inputs to group tokens by expert assignment + permuted_inputs, row_id_map = ops.permute(hidden_states, experts_indices) + final_output = torch.zeros_like(permuted_inputs) + + # Process tokens through their assigned experts + for expert_idx, expert in enumerate(self.experts): + # Skip experts with no assigned tokens + if start_indices[expert_idx] == end_indices[expert_idx]: + continue + + # Process tokens assigned to this expert + expert_input = permuted_inputs[start_indices[expert_idx]:end_indices[expert_idx]] + expert_output = expert(expert_input) + final_output[start_indices[expert_idx]:end_indices[expert_idx]] = expert_output + + # Restore original token ordering + unpermuted_outputs = ops.unpermute(final_output, row_id_map, probs) + + # Reshape back to original dimensions + output = unpermuted_outputs.view(batch_size, seq_length, hidden_dim) + + return output + +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): + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int, num_experts: int): + super().__init__() + self.hidden_size = config.hidden_size + + if config.use_sliding_window and config._attn_implementation != "flash_attention_2": + logger.warning_once( + f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " + "unexpected results may be encountered." + ) + + self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx) + + self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + if config.mlp_moe: + self.moe = SparseMoeBlock(config, num_experts=num_experts) + self.mlp = None + else: + self.mlp = Qwen2_5_VLMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + token_types=None, + start_indices=None, + end_indices=None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. + position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + if self.mlp is None: # using moe mlp + hidden_states = self.moe(hidden_states, token_types, start_indices, end_indices) + else: + hidden_states = self.mlp(hidden_states) + + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + if use_cache: + outputs += (present_key_value,) + return outputs + +class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): + """Qwen2.5-VL model with Mixture of Experts (MoE) architecture. + + This model extends the base Qwen2.5-VL model by incorporating MoE layers + for improved scalability and specialization across different token types. + """ + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str, + num_experts: Optional[int] = None, + *args, + **kwargs + ): + """Load a pretrained model with optional MoE configuration. + + Args: + pretrained_model_name_or_path: Path or name of the pretrained model + num_experts: Number of experts for MoE layers (if not in config) + *args: Additional arguments passed to parent class + **kwargs: Additional keyword arguments passed to parent class + + Returns: + Initialized model instance with MoE configuration + """ + config = kwargs.get("config", None) + if config is None: + config = AutoConfig.from_pretrained(pretrained_model_name_or_path) + + # Override number of experts if specified + if num_experts is not None: + config.num_experts = num_experts + + kwargs["config"] = config + return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) + + def __init__(self, config: Qwen2_5_VLConfig): + """Initialize the Qwen2.5-VL MoE model. + + Args: + config: Model configuration containing architecture parameters + """ + super().__init__(config) + + # Basic model parameters + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + # Model components + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) + + # Decoder layers with MoE support + self.layers = nn.ModuleList([ + Qwen2_5_VLDecoderLayer_with_MoE(config, layer_idx, config.num_experts) + for layer_idx in range(config.num_hidden_layers) + ]) + + # Model configuration + self._attn_implementation = config._attn_implementation + self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Embedding: + """Get the input embedding layer. + + Returns: + The token embedding layer + """ + return self.embed_tokens + + def set_input_embeddings(self, value: nn.Embedding) -> None: + """Set the input embedding layer. + + Args: + value: New embedding layer to use + """ + self.embed_tokens = value + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + 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, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Union[Tuple, BaseModelOutputWithPast]: + # Set default output options + 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 + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Validate inputs + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("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") + + # Handle gradient checkpointing compatibility + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # Initialize cache if needed + if use_cache and past_key_values is None and not torch.jit.is_tracing(): + past_key_values = DynamicCache() + + # Get input embeddings + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # Set up cache position + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device + ) + + # Set up position IDs (hardcoded 3 dimensions for temporal, height, width) + if position_ids is None: + position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) + elif position_ids.dim() == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + + # Create causal attention mask + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions, moe_token_types + ) + + hidden_states = inputs_embeds + + # Create position embeddings to be shared across decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # Initialize output collections + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + # Process through decoder layers + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + # Use gradient checkpointing during training + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + moe_token_types, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + # Regular forward pass + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + token_types=moe_token_types, + start_indices=start_indices, + end_indices=end_indices, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + hidden_states = layer_outputs[0] + + # Update cache if using it + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + # Collect attention weights if requested + if output_attentions: + all_self_attns += (layer_outputs[1],) + + # Apply final layer normalization + hidden_states = self.norm(hidden_states) + + # Add final hidden states if collecting all states + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + + # Return outputs in requested format + if not return_dict: + return tuple( + v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] + if v is not None + ) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + moe_token_types: Optional[torch.LongTensor] = None, + ): + """Update causal attention mask with support for bidirectional attention for specific token types. + + This method creates and modifies attention masks to support different attention patterns: + - Standard causal (unidirectional) attention for most tokens + - Bidirectional attention for specific token types (e.g., MoE routing tokens) + + Args: + attention_mask: Input attention mask to avoid attending to padding tokens + input_tensor: Input embeddings tensor for shape and device information + cache_position: Position indices for caching mechanisms + past_key_values: Cached key-value pairs from previous forward passes + output_attentions: Whether attention weights will be returned + moe_token_types: Optional tensor indicating token types for MoE routing + (type 1 tokens will use bidirectional attention) + + Returns: + Updated causal attention mask, or None if using Flash Attention 2 + """ + # Flash Attention 2 handles masking internally + if self.config._attn_implementation == "flash_attention_2": + return None + + # Calculate sequence lengths for cache management + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) + + # For SDPA (Scaled Dot Product Attention), use `is_causal` argument when possible + # instead of explicit attention mask to enable Flash Attention 2 dispatch + # Note: This optimization is not compatible with static cache + if ( + self.config._attn_implementation == "sdpa" + and not (using_static_cache or using_sliding_window_cache) + and not output_attentions + ): + # Check if we can ignore the causal mask and rely on SDPA's internal handling + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + sliding_window=self.config.sliding_window, + is_training=self.training, + ): + return None + + # Extract tensor properties for mask creation + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + + # Determine target length based on cache type + if using_sliding_window_cache or using_static_cache: + # Use maximum cache shape for sliding window or static caches + target_length = past_key_values.get_max_cache_shape() + else: + # For dynamic cache or no cache, calculate based on attention mask or sequence length + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + # Generate 4D causal attention mask from 2D input mask if provided + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=target_length, + dtype=dtype, + device=device, + cache_position=cache_position, + batch_size=input_tensor.shape[0], + config=self.config, + past_key_values=past_key_values, + ) + + # Modify mask to support bidirectional attention for specific token types + if moe_token_types is not None: + # Identify positions of type 1 tokens (MoE routing tokens) + type1_tokens = (moe_token_types == 1).unsqueeze(1).unsqueeze(2) # Shape: [B, 1, 1, S] + + # Create bidirectional attention region for type 1 tokens + # This allows type 1 tokens to attend to each other bidirectionally + type1_mask = torch.zeros_like(causal_mask) # Shape: [B, num_heads, S, S] + type1_region = type1_tokens & type1_tokens.transpose(-1, -2) # Shape: [B, 1, S, S] + type1_mask = type1_mask.masked_fill(type1_region, 1.0).to(torch.bool) + + # Apply bidirectional attention: zero out causal constraints in type 1 regions + causal_mask = torch.where( + type1_mask, # Where type 1 tokens interact with each other + torch.zeros_like(causal_mask), # Remove causal masking (allow bidirectional) + causal_mask # Keep original causal masking for other regions + ) + + # Handle special case for SDPA with CUDA/XPU devices + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type in ["cuda", "xpu"] + and not output_attentions + ): + # Ensure attention to all tokens in fully masked rows for memory-efficient attention + # This is required for F.scaled_dot_product_attention's memory-efficient path + # when using left padding. See: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + @staticmethod + def _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask: torch.Tensor, + sequence_length: int, + target_length: int, + dtype: torch.dtype, + device: torch.device, + cache_position: torch.Tensor, + batch_size: int, + config: Qwen2_5_VLConfig, + past_key_values: Cache, + ): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. + + Args: + attention_mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. + sequence_length (`int`): + The sequence length being processed. + target_length (`int`): + The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. + dtype (`torch.dtype`): + The dtype to use for the 4D attention mask. + device (`torch.device`): + The device to plcae the 4D attention mask on. + cache_position (`torch.Tensor`): + Indices depicting the position of the input sequence tokens in the sequence. + batch_size (`torch.Tensor`): + Batch size. + config (`Qwen2_5_VLConfig`): + The model's configuration class + past_key_values (`Cache`): + The cache class that is being used currently to generate + """ + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + causal_mask = attention_mask + else: + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + if config.sliding_window is not None: + # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also + # the check is needed to verify is current checkpoint was trained with sliding window or not + if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length: + sliding_attend_mask = torch.arange(target_length, device=device) <= ( + cache_position.reshape(-1, 1) - config.sliding_window + ) + diagonal_attend_mask.bitwise_or_(sliding_attend_mask) + causal_mask *= diagonal_attend_mask + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + if attention_mask.shape[-1] > target_length: + attention_mask = attention_mask[:, :target_length] + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( + causal_mask.device + ) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + return causal_mask + +class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): + """ + 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"] + config_class = Qwen2_5_VLConfig + _no_split_modules = ["Qwen2_5_VLDecoderLayer_with_MoE", "Qwen2_5_VLVisionBlock"] + + @classmethod + def from_pretrained(cls, pretrained_model_path, config_path=None, processor_path=None, action_tokenizer_path=None, **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 + model_path = os.path.join(pretrained_model_path, "model.safetensors") + config_path = os.path.join(pretrained_model_path, "config.json") + config = cls.config_class.from_pretrained(config_path) + processor = AutoProcessor.from_pretrained(pretrained_model_path, use_fast=True) + if action_tokenizer_path is not None: + processor.action_processor = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True) + + # Initialize model with configuration and processor + model = cls( + 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 + state_dict = load_file(model_path, device="cpu") + msg = model.load_state_dict(state_dict, strict=False) + + return model + + def __init__(self, config, use_fast_tokenizer=False, processor=None, action_tokenizer=None, action_mapper=None, flow_loss_weight=1.0): + """ + 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.model = Qwen2_5_VLMoEModel(config) + 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.processor = processor + + # Define action token IDs + self.define_action_token_id() + + # Cache for rope deltas + self.rope_deltas = None + + # Initialize action preprocessor + self.action_preprocessor = ActionProcessor(config) + + # Apply LoRA if specified in configuration + if hasattr(config, "use_lora") and config.use_lora: + self.add_lora( + r=config.lora_r, + lora_alpha=config.lora_alpha, + target_modules=config.lora_target_modules, + lora_dropout=config.lora_dropout + ) + + # Initialize weights and apply final processing + self.post_init() + + 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 = [] + 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(f"<|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" + ) + self.model = get_peft_model(self.model, config) + + # Print information about trainable parameters + self.model.print_trainable_parameters() + + 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( + self, + input_ids: Optional[torch.LongTensor] = None, + 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, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Calculate 3D RoPE (Rotary Position Embedding) indices for vision and text tokens. + + This method computes position embeddings that account for the temporal, height, and width + dimensions of vision tokens (images/videos) while maintaining standard 1D position embeddings + for text tokens. + + For vision tokens, 3D position embeddings are calculated based on: + - Temporal dimension: Time patches in videos + - Height dimension: Vertical patches in images/video frames + - Width dimension: Horizontal patches in images/video frames + + For text tokens, standard 1D position embeddings are used, continuing from the maximum + vision position ID plus 1. + + Args: + input_ids (torch.LongTensor, optional): Input token IDs of shape (batch_size, sequence_length) + image_grid_thw (torch.LongTensor, optional): Image grid dimensions (num_images, 3) for [temporal, height, width] + video_grid_thw (torch.LongTensor, optional): Video grid dimensions (num_videos, 3) for [temporal, height, width] + second_per_grid_ts (torch.Tensor, optional): Time interval per temporal grid (num_videos,) + attention_mask (torch.Tensor, optional): Attention mask (batch_size, sequence_length) + + Returns: + tuple: + - position_ids (torch.LongTensor): 3D position IDs of shape (3, batch_size, sequence_length) + - mrope_position_deltas (torch.Tensor): Position deltas for mRoPE of shape (batch_size, 1) + """ + 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 + mrope_position_deltas = [] + + if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + + # Initialize 3D position IDs tensor + 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) + + # Process each sequence in the batch + for i, input_ids in enumerate(total_input_ids): + input_ids = input_ids[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + + # Find vision tokens and count images/videos + vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + + # Process each vision token (image or video) + for _ in range(image_nums + video_nums): + # Find next image or video token + if image_token_id in input_tokens 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 input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + + # Determine if processing image or video token + if ed_image < ed_video: + # Process image token + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + second_per_grid_t = 0 + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + # Process video token + 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 + + # Calculate grid dimensions after spatial merging + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + # Add position IDs for text tokens before vision token + 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).view(1, -1).expand(3, -1) + st_idx) + + # Calculate 3D position embeddings for vision tokens + range_tensor = torch.arange(llm_grid_t).view(-1, 1) + expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) + + # Calculate temporal position IDs with time scaling + time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second + time_tensor_long = time_tensor.long() + t_index = time_tensor_long.flatten() + + # Calculate spatial position IDs + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() + + # Add 3D position IDs for vision tokens + 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 + + # Add position IDs for remaining text tokens + 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).view(1, -1).expand(3, -1) + st_idx) + + # Concatenate all position IDs for this sequence + 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=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + # Handle case without vision tokens - use standard 1D position embeddings + 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 + + def train_step_forward( + self, + input_ids: torch.LongTensor = None, + 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 type assignments + 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, # Action trajectory chunks + proprioception: Optional[torch.FloatTensor] = None, # Joint position/orientation data + 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, + ) -> Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]: + """ + Forward pass for training with multi-modal inputs including vision, text, and action data. + + This method handles the complete forward pass during training, processing various input modalities + including images, videos, text, proprioceptive data, and action sequences. It computes losses + for both language modeling and action prediction using flow matching. + + Args: + 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 for generation + 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 loss computation + 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 (temporal, height, width) + video_grid_thw (torch.LongTensor, optional): Video grid dimensions (temporal, height, width) + action_chunk (torch.FloatTensor, optional): Action trajectory data chunks + proprioception (torch.FloatTensor, optional): Proprioceptive sensor data (joint positions, etc.) + 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 + dataset_names (str, optional): Names of datasets in the current batch + dof_mask (torch.FloatTensor, optional): Degrees of freedom mask for action tokens + agent_pos_mask (torch.FloatTensor, optional): Agent position mask for proprioceptive data + **kwargs: Additional keyword arguments + + Returns: + Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]: Model outputs including losses, logits, + and auxiliary information, or tuple if return_dict=False + """ + batch_size, seq_length = input_ids.shape + + # 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 + + # 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 = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + ) + self.rope_deltas = rope_deltas + # Use previously calculated rope deltas to get correct position IDs + else: + delta = ( + (cache_position[0] + self.rope_deltas).to(self.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=self.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) + + # Calculate token distribution across MoE expert groups + 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) + + # 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) + 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 (joint positions, orientations, etc.) + 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) + 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) + elif self.training: + # Dummy forward pass to ensure gradient registration in DDP + # This handles cases where one process has proprioception data while another doesn't + # Without this, DDP would hang waiting for a gradient that will never be computed + dummy_output = sum(p.sum() for p in self.action_preprocessor.proprioception_proj.parameters()) + dummy_input = torch.randn(2, self.action_preprocessor.propri_dim*2, device=inputs_embeds.device) + dummy_forward = self.action_preprocessor.proprioception_proj(dummy_input) + dummy_loss = sum(p.sum() for p in dummy_forward) + inputs_embeds = inputs_embeds + 0 * dummy_loss + + # Process action chunk data + if action_chunk is not None: + action_chunk = action_chunk.to(inputs_embeds.device).to(inputs_embeds.dtype) + dof_mask = dof_mask.to(inputs_embeds.device).to(inputs_embeds.dtype) + noisy_action_emb, flow = 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) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + # Forward pass through the main model + 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 for MoE routing + start_indices=start_indices, + end_indices=end_indices, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + + # Initialize loss computation variables + loss = None + cross_entropy_loss, flow_loss = None, None + channel_loss_dict = None + channel_loss_count_dict = None + + # Compute losses if labels are provided + if labels is not None: + loss = 0 + action_accuracy = 0 + unique_datasets_name = list(set(dataset_names)) + + # Initialize per-dataset loss tracking dictionaries + 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 + } + + # Compute standard cross-entropy loss for language modeling + 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 by moving labels to correct device + 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 per-dataset channel losses + _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() + + # Add cross-entropy loss to total loss if valid + if not torch.isnan(cross_entropy_loss): + loss += cross_entropy_loss + else: + with torch.no_grad(): + cross_entropy_loss.detach() + + # Compute action token prediction accuracy + shift_logits = logits[..., :-1, :].contiguous() + action_preds = shift_logits.argmax(dim=-1) + shift_labels = labels[..., 1:].contiguous() + if self.use_fast_tokenizer: + action_mask = shift_labels > self.action_token_id_set['fast_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] + flow = flow.reshape(-1, flow.shape[-1]) + _flow_loss = self.action_preprocessor.flow_loss(action_hidden_states, flow, dof_mask) + if isinstance(_flow_loss, torch.Tensor): + flow_loss = _flow_loss.mean() + if loss is not None: + loss += self.flow_loss_weight * flow_loss + else: + loss = self.flow_loss_weight * flow_loss + _flow_loss = _flow_loss.view(dof_mask.shape[0], dof_mask.shape[1], dof_mask.shape[2]) + + # Return outputs based on return_dict setting + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return Qwen2_5_VLACausalLMOutputWithPast( + loss=loss, + cross_entropy_loss=cross_entropy_loss.clone() if cross_entropy_loss is not None else None, + flow_loss=flow_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=self.rope_deltas, + channel_loss_dict=channel_loss_dict, + channel_loss_count_dict=channel_loss_count_dict, + ) + + def predict_action(self, predict_mode: str, **kwargs): + """ + Predict actions using specified prediction mode. + + 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 + ) + + return output['predict_action'], output.get('gt_action', None) + + @torch.no_grad() + def predict( + self, + predict_mode: str, + pred_horizon: Optional[int] = None, + action_dim: Optional[int] = None, + input_ids: torch.LongTensor = None, + 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, + 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, + 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, + re_generate: bool = False, + **kwargs, + ): + """ + Multi-modal prediction method supporting text generation, fast action prediction, and diffusion-based action prediction. + + 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]) + + 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 = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + ) + 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) + 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() + # 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) + 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 + ) + + # 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) + + 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 + + # 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, + 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) + return pred.reshape(batch_size, pred_horizon, action_dim) + + # Perform ODE integration for diffusion sampling + times = torch.linspace(0, 1, num_inference_timesteps, + device=inputs_embeds.device, dtype=inputs_embeds.dtype) + action_trajectory = odeint(step, noisy_action, times, method="euler") + + # 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 + + # Process ground truth actions if available + if action_chunk is not None: + output['gt_action'] = self.action_preprocessor.normalizer_action.unnormalize_data(action_chunk, dataset_names) + + 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 + """ + if not mode: + with torch.no_grad(): + return self.train_step_forward(**kwargs) + elif mode == 'predict': + return self.predict(predict_mode=predict_mode, **kwargs) + elif mode == 'train': + return self.train_step_forward(use_cache=False, **kwargs) + elif mode == 'validate': + with torch.no_grad(): + return self.train_step_forward(use_cache=False, **kwargs) + else: + raise NotImplementedError('invalid key') + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + moe_token_types=None, + cache_position=None, + position_ids=None, + use_cache=True, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + second_per_grid_ts=None, + dataset_names=None, + proprioception=None, + dof_mask=None, + agent_pos_mask=None, + **kwargs, + ): + """ + Prepare inputs for autoregressive generation with multi-modal support. + + 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. + + 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 + else: + # Ensure moe_token_types length matches input_ids + if moe_token_types.shape[1] < input_ids.shape[1]: + # Calculate required padding length + pad_length = input_ids.shape[1] - moe_token_types.shape[1] + # Create padding tensor with default token type (0) + 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 + 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 + inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] + moe_token_types = moe_token_types[:, -cache_position.shape[0] :] + elif ( + inputs_embeds is not None # Exception 1: input_embeds provided + or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1]) # Exception 3: GPU sync edge case + ): + 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) + 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 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 + device = inputs_embeds.device + else: + batch_size, sequence_length = input_ids.shape + device = input_ids.device + + attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=past_key_values.get_max_cache_shape(), + dtype=self.lm_head.weight.dtype, + device=device, + cache_position=cache_position, + batch_size=batch_size, + config=self.config, + past_key_values=past_key_values, + ) + + # Assemble all model inputs for generation + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "moe_token_types": moe_token_types, + "use_cache": use_cache, + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "pixel_values_videos": pixel_values_videos, + "image_grid_thw": image_grid_thw, + "video_grid_thw": video_grid_thw, + "cache_position": cache_position, + "second_per_grid_ts": second_per_grid_ts, + "proprioception": proprioception, + "dataset_names": dataset_names, + "dof_mask": dof_mask, + "agent_pos_mask": agent_pos_mask, + } + ) + return model_inputs + + def _get_image_nums_and_video_nums( + self, + 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. + + Args: + input_ids (torch.LongTensor): Input token IDs of shape (batch_size, sequence_length) + + Returns: + tuple: + - image_nums (torch.LongTensor): Number of images per sample + - video_nums (torch.LongTensor): Number of videos per 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) + + return image_nums, video_nums + + def _expand_inputs_for_generation( + self, + expand_size: int = 1, + is_encoder_decoder: bool = False, + input_ids: Optional[torch.LongTensor] = None, + **model_kwargs, + ) -> Tuple[torch.LongTensor, Dict[str, Any]]: + """ + Expand inputs for generation with support for multi-modal tensors. + + 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", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"] + + 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([sample.repeat(*repeat_args) for sample in samples], dim=0) + return result + + for key in dict_to_expand: + if key == "pixel_values": + # Split images into samples and compute sequence lengths + samples = torch.split(image_grid_thw, list(image_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 == "image_grid_thw": + # Expand based on number of images per 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." + ) + tensor = torch.tensor(dict_to_expand[key]) + lengths = list(video_nums) + tensor = _repeat_interleave_samples(tensor, lengths=lengths, repeat_times=expand_size) + dict_to_expand[key] = tensor.tolist() + 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" + and dict_to_expand[key] is not None + and isinstance(dict_to_expand[key], torch.Tensor) + and key not in visual_keys + ): + dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0) + 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 + 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("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.") + model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"]) + + return input_ids, model_kwargs \ No newline at end of file diff --git a/wall_x/trainer/__init__.py b/wall_x/trainer/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/wall_x/trainer/qwen_vl_act_trainer.py b/wall_x/trainer/qwen_vl_act_trainer.py new file mode 100644 index 0000000..f3301ea --- /dev/null +++ b/wall_x/trainer/qwen_vl_act_trainer.py @@ -0,0 +1,671 @@ +import os +import gc +import time +import torch +import random +import numpy as np +import torch.nn as nn + +from tqdm import tqdm +from functools import wraps +from datetime import datetime +from torch.optim import AdamW +from accelerate import Accelerator +from safetensors.torch import load_file +from accelerate.utils import DistributedType +from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup + +from wall_x.utils.timers import Timers +from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction +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 + + +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) + + +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_qwen_vl_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 = ["processor_path", "qwen_vl_act_config_path", "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) + + # Load model and initialize training components + 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 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() + + # 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.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() + grad_accum_steps = self.config.get("gradient_accumulation_steps", 1) + 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: + 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 (i + 1) % grad_accum_steps == 0: + 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.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 = Qwen2_5_VLMoEForAction.from_pretrained( + self.config["pretrained_qwen_vl_path"], + **{"use_fast_tokenizer": self.use_fast_tokenizer} + ) + self.processor = model.processor + model = model.to(torch.bfloat16) + + # 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( + self.config, + self.dataload_config.get("lerobot_config", {}), + 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 training checkpoint. + + Args: + epoch (int): Current epoch number + step (int, optional): Current step number. Defaults to 0. + + Saves model state, optimizer state, and training progress information. + """ + 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 current iteration steps for dataset resuming + if step != 0: + _rank = self.accelerator.process_index + if isinstance(self.dataset, PreprocessedDataset): + torch.save( + {"epoch": epoch, "step": step}, + os.path.join(ckpt_path, f"epoch_{epoch}_step_{step}_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. + """ + checkpoint_path = self.config["resume"]["ckpt"] + + if self.config.get("resume", {}).get("load_ckpt_only", False): + # 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] + + err = 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) + + self.print_rank0(f"Resumed from checkpoint: {checkpoint_path}") + + 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 \ No newline at end of file diff --git a/wall_x/utils/__init__.py b/wall_x/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/wall_x/utils/constant.py b/wall_x/utils/constant.py new file mode 100755 index 0000000..9da35fd --- /dev/null +++ b/wall_x/utils/constant.py @@ -0,0 +1,191 @@ +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]}, + }, +} diff --git a/wall_x/utils/timers.py b/wall_x/utils/timers.py new file mode 100644 index 0000000..0a67fb3 --- /dev/null +++ b/wall_x/utils/timers.py @@ -0,0 +1,483 @@ +import time +import torch +from torch.cuda import nvtx +from abc import ABC, abstractmethod +from typing import List + +def _is_distributed(): + return torch.distributed.is_available() and torch.distributed.is_initialized() + +def _get_world_size(): + if _is_distributed(): + return torch.distributed.get_world_size() + return 1 + +def _get_rank(): + if _is_distributed(): + return torch.distributed.get_rank() + return 0 + +def _barrier(group=None): + if _is_distributed(): + torch.distributed.barrier(group=group) + +if torch.distributed.is_available(): + try: + dist_all_gather_func = torch.distributed.all_gather_into_tensor + except AttributeError: + dist_all_gather_func = torch.distributed.all_gather +else: + dist_all_gather_func = None + +class TimerBase(ABC): + """Timer base class.""" + + def __init__(self, name): + self.name = name + + @abstractmethod + def start(self, barrier=False): + """Start the timer. + + Args: + barrier (bool, optional): Synchronizes ranks before starting. Defaults to False. + """ + pass + + @abstractmethod + def stop(self, barrier=False): + """Stop the timer. + + Args: + barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False. + """ + pass + + @abstractmethod + def reset(self): + """Reset timer.""" + 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. + """ + pass + + +class DummyTimer(TimerBase): + """Dummy Timer.""" + + def __init__(self): + super().__init__('dummy timer') + + def start(self, barrier=False, nvtx_push=False): + return + + def stop(self, barrier=False, nvtx_pop=False): + return + + def reset(self): + return + + def elapsed(self, reset=True, barrier=False): + raise Exception( + 'dummy timer should not be used to calculate elapsed time, ' + 'check if timer\'s log_level <= self._log_level.' + ) + + 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.' + ) + + +class Timer(TimerBase): + """ + Timer class with ability to start/stop. + + Comment on using `barrier`: If this flag is passed, then all + the caller processes will wait till all reach the timing routine. + It is up to the user to make sure all the ranks in `barrier_group` + call it otherwise, it will result in a hang. + Comment on `barrier_group`: By default it is set to None which + in torch distributed land, it will result in the global communicator. + """ + + 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): + """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) + if torch.cuda.is_available(): + torch.cuda.synchronize() + self._start_time = time.time() + self._started = True + if nvtx_push: + nvtx.range_push("{}".format(self.name)) + 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' + if barrier: + _barrier(group=self._barrier_group) + if torch.cuda.is_available() and sync: + torch.cuda.synchronize() + elapsed = time.time() - self._start_time + self._elapsed += elapsed + self._active_time += elapsed + 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 + + +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 ( + log_option in allowed_log_options + ), 'input log option {} is invalid. It must be one of {}'.format( + log_option, allowed_log_options + ) + self._log_option = log_option + self._timers = {} + self._log_levels = {} + self._dummy_timer = DummyTimer() + 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], ( + 'input log level {} does not match already existing ' + 'log level {} for {} timer'.format(log_level, self._log_levels[name], name) + ) + 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 ( + log_level <= self._max_log_level + ), '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. + + 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 measurments + + 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() + + if torch.cuda.is_available(): + device = torch.cuda.current_device() + else: + device = torch.device('cpu') + + rank_name_to_time = torch.zeros( + (world_size, len(names)), dtype=torch.float, device=device + ) + + for i, name in enumerate(names): + if name in self._timers: + rank_name_to_time[rank, i] = self._timers[name].elapsed(reset=reset) + + 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: + print(f"Warning: all_gather failed: {e}. Using single rank timing.") + + 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, + rank_to_time.max().item() / normalizer, + ) + return name_to_min_max_time + + 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) + if not name_to_min_max_time: + return None + + world_size = _get_world_size() + if world_size == 1: + output_string = 'time (ms):' + for name in name_to_min_max_time: + _, max_time = name_to_min_max_time[name] + output_string += '\n {}: {:.2f}'.format((name + ' ').ljust(48, '.'), max_time) + else: + if max_only: + output_string = 'max time across ranks (ms):' + else: + output_string = '(min, max) time across ranks (ms):' + for name in name_to_min_max_time: + min_time, max_time = name_to_min_max_time[name] + if max_only: + output_string += '\n {}: {:.2f}'.format((name + ' ').ljust(48, '.'), max_time) + else: + output_string += '\n {}: ({:.2f}, {:.2f})'.format( + (name + ' ').ljust(48, '.'), min_time, max_time + ) + 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() + + output_string = 'times across ranks (ms):' + no_reported_timing = True + for i, name in enumerate(names): + not_yet_found = True + for rank in range(world_size): + if rank_name_to_time[rank, i] > 0: + no_reported_timing = False + if not_yet_found: + not_yet_found = False + output_string += '\n {}:'.format(name) + if world_size == 1: + output_string += '\n {:.2f}'.format( + rank_name_to_time[rank, i] / normalizer + ) + else: + output_string += '\n rank {:2d}: {:.2f}'.format( + rank, rank_name_to_time[rank, i] / normalizer + ) + if no_reported_timing: + return None + return output_string + + def get_all_timers_string( + self, + names: List[str] = None, + normalizer: float = 1.0, + reset: bool = True, + barrier: bool = False, + ): + """Returns the output string with logged timer values according to configured options. + + 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. + """ + + if names == None: # get all registered timers + 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 + output_string = self._get_global_min_max_time_string( + names, reset, barrier, normalizer / 1000.0, max_only + ) + elif self._log_option == 'all': + output_string = self._get_all_ranks_time_string( + names, reset, barrier, normalizer / 1000.0 + ) + else: + raise Exception('unknown timing log option {}'.format(self._log_option)) + return output_string + + def log( + self, + names: List[str], + rank: int = None, + normalizer: float = 1.0, + 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. + + 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. + """ + + 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) + + def write( + self, + names: List[str], + writer, + iteration: int, + normalizer: float = 1.0, + reset: bool = True, + barrier: bool = False, + ): + """Write timers to a tensorboard writer. + Note that we only report maximum time across ranks to tensorboard. + + 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. + """ + # 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) + if writer is not None: + for name in name_to_min_max_time: + _, max_time = name_to_min_max_time[name] + writer.add_scalar(name + '-time', max_time, iteration) \ No newline at end of file diff --git a/workspace/README.md b/workspace/README.md new file mode 100644 index 0000000..ccc7a7a --- /dev/null +++ b/workspace/README.md @@ -0,0 +1,67 @@ +# Training Configuration Guide + +This document explains the key configuration parameters that can be modified for Wall-X training. + +## Quick Start Checklist +1. **Update run.sh**: Set `code_dir` and `config_path` to your actual paths +2. **Configure GPUs**: Set `CUDA_VISIBLE_DEVICES` for your available GPUs +3. **Update config paths**: Replace all `/path/to/` placeholders in config.yml with actual paths +4. **Configure robot**: Set `dof_config` and `agent_pos_config` for your robot +5. **Set dataset**: Choose appropriate `repo_id` for your dataset +6. **Adjust batch size**: Set `batch_size_per_gpu` based on GPU memory +7. **Run training**: Execute `bash ./workspace/lerobot_example/run.sh` + +## Required Paths (Must Modify) +```yaml +processor_path: "/path/to/model/" # Path to model processor +pretrained_qwen_vl_path: "/path/to/qwen_vl_model/" # Path to pretrained Qwen VL model +qwen_vl_act_config_path: "/path/to/config.json" # Path to model config file +action_tokenizer_path: "/path/to/fast/" # Path to action tokenizer +save_path: "/path/to/workspace/" # Path to save training outputs +``` + +## 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 + +## 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 + +## Performance Settings (Optional) +- `profile`: Enable PyTorch profiling (true/false) +- `padding_side`: Token padding side (left/right) \ No newline at end of file diff --git a/workspace/lerobot_example/config_qact.yml b/workspace/lerobot_example/config_qact.yml new file mode 100644 index 0000000..4fa62cd --- /dev/null +++ b/workspace/lerobot_example/config_qact.yml @@ -0,0 +1,112 @@ +# 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: qwen2_5 +processor_path: "/path/to/model/" +pretrained_qwen_vl_path: "/path/to/qwen_vl_model/" +qwen_vl_act_config_path: "/path/to/config.json" +action_tokenizer_path: "/path/to/fast/" +save_path: "/path/to/workspace/" + +# 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: 32 +batch_size_per_gpu: 8 +padding_side: left + +# 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 + +# 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/run.sh b/workspace/lerobot_example/run.sh new file mode 100644 index 0000000..8baeb43 --- /dev/null +++ b/workspace/lerobot_example/run.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# export CUDA_VISIBLE_DEVICES=4,5,6,7 +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 \ No newline at end of file