diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..9354afc --- /dev/null +++ b/.dockerignore @@ -0,0 +1,10 @@ +# Keep the build context small: llama.cpp is fetched inside the image, models +# are mounted at runtime, and build/output dirs are host artifacts. +.git/ +models/ +build/ +build-*/ +outputs/ +third_party/ +eval/sim/ +*.gguf diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml new file mode 100644 index 0000000..1838000 --- /dev/null +++ b/.github/workflows/build.yml @@ -0,0 +1,52 @@ +# Fast build + unit gate on GitHub-hosted runners. The heavier cross-platform +# LIBERO regression sweep lives in vla-ci.yml (self-hosted). +name: build + +on: + push: + branches: [main, dev] + pull_request: + branches: [main] + +concurrency: + group: build-${{ github.ref }} + cancel-in-progress: true + +jobs: + cpp-unit: + runs-on: ubuntu-24.04 + steps: + - uses: actions/checkout@v4 + - name: pixel-shuffle channel order + run: | + g++ -std=c++17 -Isrc -Wall -Wextra tests/test_vision_common.cpp -o /tmp/test_vision_common + /tmp/test_vision_common + + py-tooling: + runs-on: ubuntu-24.04 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - name: converter remap + run: python tests/py/test_converters.py + + build-gate: + runs-on: ubuntu-24.04 + steps: + - uses: actions/checkout@v4 + - name: deps + run: | + sudo apt-get update -qq + sudo apt-get install -y -qq --no-install-recommends \ + build-essential cmake git ca-certificates pkg-config \ + libzmq3-dev cppzmq-dev libprotobuf-dev protobuf-compiler + - uses: actions/cache@v4 + with: + path: build/_deps + key: llama-b9866-${{ runner.os }} + - name: build vla-server + vlm-server + vla-cli (CPU, -Wall -Wextra) + run: | + cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=OFF + cmake --build build -j"$(nproc)" --target vla-server vlm-server vla-cli diff --git a/.gitignore b/.gitignore index 148882d..8894aa1 100644 --- a/.gitignore +++ b/.gitignore @@ -53,3 +53,4 @@ eval/sim/libero/LIBERO/ eval/sim/libero/libero_uv/ eval/sim/simpler/SimplerEnv/ eval/sim/simpler/simpler_uv/ +/models/ diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..a043125 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,44 @@ +# Changelog + +Notable changes to vla.cpp. Format loosely follows [Keep a Changelog](https://keepachangelog.com). + +## [0.1.1] - 2026-07-04 + +### Added +- `vla-cli`: one-shot inference from the command line (image + tokens to action), no server needed. +- `scripts/quantize_gguf.py`: repack LM weights to Q8_0/Q4_0. The loader runs quantized GGUFs directly (Q8_0 roughly halves the LM, near-lossless). + +### Changed +- One shared GGUF reader across the model loaders, replacing the per-arch copies. +- Cap inbound message size (256 MiB) and image count (16) on `vla-server` and `vlm-server`. +- Scale CPU threads to the machine core count across all loaders instead of a fixed 4. +- Read GGUF file offsets as 64-bit and reject non-float embedding tensors in row-fetch. + +### Fixed +- Reject out-of-range language tokens in OpenVLA-OFT and VLA-Adapter. +- Zero the padded action dimensions so only real action dims carry values. +- Reject images that do not match the model input size in VLA-Adapter and OpenVLA-OFT (out-of-bounds read on a smaller view). +- Validate Evo-1 action dims at load so a client-supplied noise buffer cannot underrun. +- Only enable the BitVLA CUDA path once every device buffer allocates. + +## [0.1.0] - 2026-07-03 + +First tagged release. One self-contained GGUF per model (vision tower + LM + action +expert + dataset stats), CPU or CUDA, no external mmproj and no patch to llama.cpp. + +### Added +- Seven VLA policies auto-detected from the GGUF: SmolVLA, pi0, BitVLA, Evo-1, GR00T N1.5/N1.6/N1.7. +- In-tree vision towers (SigLIP, BitSigLIP, InternViT, RADIO) on stable public ggml/llama APIs. +- ZeroMQ + protobuf `vla-server`; a separate `vlm-server` for VLM chat. +- BitVLA 1.58-bit custom ternary CUDA kernels. +- Per-arch HuggingFace -> GGUF converters and mmproj-merge helpers (`scripts/`). +- Robot eval harness for LIBERO, SimplerEnv, and ALOHA (`eval/`), with device benchmark reports. +- Minimal CI: pixel-shuffle unit test, converter-remap test, CPU build gate. +- `pyproject.toml` for the Python tooling and a CUDA `Dockerfile` for `vla-server`. + +### Changed +- llama.cpp is fetched + pinned via CMake `FetchContent` (tag `b9866`); bumping is a + one-line `GIT_TAG` change. Removed the `patches/` fetch script. + +[0.1.1]: https://github.com/VinRobotics/vla.cpp/releases/tag/v0.1.1 +[0.1.0]: https://github.com/VinRobotics/vla.cpp/releases/tag/v0.1.0 diff --git a/CMakeLists.txt b/CMakeLists.txt index d4d493a..7e757e1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -16,7 +16,14 @@ set(LLAMA_BUILD_TOOLS ON CACHE BOOL "" FORCE) set(LLAMA_BUILD_SERVER OFF CACHE BOOL "" FORCE) set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE) set(LLAMA_BUILD_TESTS OFF CACHE BOOL "" FORCE) -add_subdirectory(third_party/llama.cpp) +# llama.cpp fetched + pinned at configure; bump = one-line GIT_TAG change. +include(FetchContent) +FetchContent_Declare(llama + GIT_REPOSITORY https://github.com/ggml-org/llama.cpp + GIT_TAG b9866 + GIT_SHALLOW TRUE +) +FetchContent_MakeAvailable(llama) add_library(vla_core src/model.cpp @@ -35,10 +42,9 @@ add_library(vla_core target_include_directories(vla_core PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/src - ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/tools/mtmd - ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/vendor + ${llama_SOURCE_DIR}/vendor ) -target_link_libraries(vla_core PUBLIC llama ggml mtmd) +target_link_libraries(vla_core PUBLIC llama ggml) if(GGML_CUDA) target_compile_definitions(vla_core PUBLIC GGML_USE_CUDA) @@ -46,7 +52,21 @@ if(GGML_CUDA) enable_language(CUDA) find_package(CUDAToolkit REQUIRED) if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) - set(CMAKE_CUDA_ARCHITECTURES 80 86 87 89 90 100 120 CACHE STRING "" FORCE) + # Emit SASS per supported GPU (real) and PTX (virtual) only for the newest + # arch, so future cards can JIT without shipping PTX for every target. + # Blackwell sm_100/sm_120 require CUDA >= 12.8. + set(_vla_cuda_archs 80-real 86-real 87-real 89-real 90-real) + set(_vla_cuda_ptx 90-virtual) + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.8) + list(APPEND _vla_cuda_archs 100-real 120-real) + set(_vla_cuda_ptx 120-virtual) + else() + message(WARNING + "CUDA ${CUDAToolkit_VERSION} < 12.8: omitting Blackwell sm_100/sm_120. " + "Pass -DCMAKE_CUDA_ARCHITECTURES= explicitly to override.") + endif() + set(CMAKE_CUDA_ARCHITECTURES ${_vla_cuda_archs} ${_vla_cuda_ptx} + CACHE STRING "" FORCE) endif() add_library(bitvla_cuda_kernels STATIC src/kernels/bitvla/bitnet_kernels.cu @@ -79,9 +99,9 @@ target_include_directories(vlm_core PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/src PRIVATE - ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/common - ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/tools/mtmd - ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/vendor + ${llama_SOURCE_DIR}/common + ${llama_SOURCE_DIR}/tools/mtmd + ${llama_SOURCE_DIR}/vendor ) target_link_libraries(vlm_core PUBLIC llama ggml mtmd llama-common) @@ -89,6 +109,15 @@ find_package(Protobuf REQUIRED) find_package(PkgConfig REQUIRED) pkg_check_modules(ZeroMQ REQUIRED IMPORTED_TARGET libzmq) +# cppzmq (zmq.hpp) is the header-only C++ binding, packaged separately from the +# libzmq C library (Debian/Ubuntu: cppzmq-dev). The servers #include . +find_path(CPPZMQ_INCLUDE_DIR NAMES zmq.hpp) +if(NOT CPPZMQ_INCLUDE_DIR) + message(FATAL_ERROR + "cppzmq (zmq.hpp) not found. Install 'cppzmq-dev' (Debian/Ubuntu) or set " + "CMAKE_PREFIX_PATH to your cppzmq install prefix.") +endif() + set(VLA_PROTO_SRC ${CMAKE_CURRENT_SOURCE_DIR}/src/serving/vla.proto) set(VLA_PROTO_GEN_DIR ${CMAKE_CURRENT_BINARY_DIR}/proto-gen/serving) file(MAKE_DIRECTORY ${VLA_PROTO_GEN_DIR}) @@ -110,7 +139,8 @@ add_executable(vla-server ) target_include_directories(vla-server PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/proto-gen - ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/vendor/stb + ${llama_SOURCE_DIR}/vendor/stb + ${CPPZMQ_INCLUDE_DIR} ) target_link_libraries(vla-server PRIVATE vla_core @@ -139,9 +169,44 @@ add_executable(vlm-server ) target_include_directories(vlm-server PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/proto-gen + ${CPPZMQ_INCLUDE_DIR} ) target_link_libraries(vlm-server PRIVATE vlm_core protobuf::libprotobuf PkgConfig::ZeroMQ ) + +# One-shot inference CLI: image + tokens -> action, no server or simulator. +add_executable(vla-cli + src/serving/vla-cli.cpp +) +target_include_directories(vla-cli PRIVATE + ${llama_SOURCE_DIR}/vendor/stb +) +target_link_libraries(vla-cli PRIVATE vla_core) + +# --- First-party build hygiene (never applied to the vendored llama.cpp subtree) -- +set(VLA_FIRST_PARTY_TARGETS vla_core vlm_core vla-server vlm-server vla-cli) + +foreach(tgt IN LISTS VLA_FIRST_PARTY_TARGETS) + # Warn on our own C++ only; nvcc device code keeps its own diagnostics. + target_compile_options(${tgt} PRIVATE $<$:-Wall -Wextra>) +endforeach() + +# Link-time optimization on Release builds when the toolchain supports it. +include(CheckIPOSupported) +check_ipo_supported(RESULT VLA_IPO_OK OUTPUT VLA_IPO_MSG) +if(VLA_IPO_OK AND CMAKE_BUILD_TYPE STREQUAL "Release") + foreach(tgt IN LISTS VLA_FIRST_PARTY_TARGETS) + set_property(TARGET ${tgt} PROPERTY INTERPROCEDURAL_OPTIMIZATION ON) + endforeach() +else() + message(STATUS "vla: IPO/LTO disabled (${VLA_IPO_MSG})") +endif() + +option(VLA_BUILD_TESTS "Build vla.cpp test harnesses and unit tests" OFF) +if(VLA_BUILD_TESTS) + enable_testing() + add_subdirectory(tests) +endif() diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..c399342 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,39 @@ +# vla-server, CPU or CUDA. cmake fetches llama.cpp at build time. +# GPU sm_120 (default): docker build -t vla-cpp . +# GPU other arch: --build-arg CUDA_ARCH=89 (89=RTX40 90=H100 87=Orin; sm_120 needs CUDA>=12.8) +# older card: --build-arg BASE_IMAGE=nvidia/cuda:12.4.1-devel-ubuntu24.04 --build-arg CUDA_ARCH=86 +# CPU: --build-arg BACKEND=cpu --build-arg BASE_IMAGE=ubuntu:24.04 -t vla-cpp-cpu +# run: docker run --gpus all -p5555:5555 -v $PWD/models:/models vla-cpp --bind tcp://*:5555 /models/M.gguf +# (CDI hosts use --device nvidia.com/gpu=all) + +ARG BASE_IMAGE=nvidia/cuda:12.9.1-devel-ubuntu24.04 +FROM ${BASE_IMAGE} + +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + build-essential cmake git ca-certificates pkg-config \ + libzmq3-dev cppzmq-dev libprotobuf-dev protobuf-compiler \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /src +COPY . . + +ARG BACKEND=cuda +ARG CUDA_ARCH=120 +# nvcc can segfault on the flash-attn kernels under high -j; lower JOBS if so. +ARG JOBS= +# CUDA: -devel ships only a libcuda stub (real driver injected at runtime), so +# link ggml-cuda's driver calls against it. CPU build skips this. +RUN set -eux; \ + if [ "$BACKEND" = "cuda" ]; then \ + export LIBRARY_PATH=/usr/local/cuda/lib64/stubs; \ + cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=ON -DGGML_CUDA_GRAPHS=ON \ + -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCH}" \ + -DCMAKE_SHARED_LINKER_FLAGS="-lcuda" -DCMAKE_EXE_LINKER_FLAGS="-lcuda"; \ + else \ + cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=OFF; \ + fi; \ + cmake --build build -j"${JOBS:-$(nproc)}" --target vla-server vla-cli; \ + cp build/vla-server build/vla-cli /usr/local/bin/ + +EXPOSE 5555 +ENTRYPOINT ["vla-server"] diff --git a/README.md b/README.md index e26a1b4..51e1a81 100644 --- a/README.md +++ b/README.md @@ -42,10 +42,9 @@ Identify your machine CUDA architecture: Configure and build the source: -```bash -# Fetch llama.cpp at pinned tag and apply local patch -bash ./patches/patch.sh +CMake fetches and pins `llama.cpp` automatically (no patch, no submodule): +```bash # CPU build: cmake -B build -DCMAKE_BUILD_TYPE=Release cmake --build build -j$(nproc) @@ -68,6 +67,23 @@ export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH On Apple Silicon (e.g. Mac Mini M4), Metal is enabled by default and runs both the transformer and vision tower on the GPU. See [docs/backend/metal.md](docs/backend/metal.md) for building `vla.cpp` on macOS. +## Quickstart + +Once the binaries are built, run one CPU prediction without a server or simulator: + +```bash +pip install -U "huggingface_hub[cli]" gguf +hf download vrfai/smolvla-libero-gguf --local-dir models/smolvla + +# front.jpg must already be the model input size (512x512 for this checkpoint). +./build/vla-cli --ckpt models/smolvla/smolvla-libero.gguf \ + --image front.jpg --tokens 1,100,200,2 --pretty +``` + +`--tokens` are language token ids from the client tokenizer. For the design overview +see [docs/ARCHITECTURE.md](docs/ARCHITECTURE.md); for the server path and other models +see [Roadmap](#roadmap). + ## Install simulators The eval scaffold under [`eval/`](eval/) supports two simulators end-to-end. Each setup script bootstraps an isolated Python 3.10 `uv` venv next to itself and clones the upstream sim repo. Both require [`uv`](https://github.com/astral-sh/uv) on `PATH`. @@ -86,17 +102,13 @@ Clones LIBERO into [`eval/sim/libero/LIBERO/`](eval/sim/libero/LIBERO), creates bash eval/sim/simpler/setup_SimplerEnv.sh ``` -Clones SimplerEnv (and its nested `ManiSkill2_real2sim`) into [`eval/sim/simpler/SimplerEnv/`](eval/sim/simpler/SimplerEnv), creates `eval/sim/simpler/simpler_uv/.venv/`, and pins gymnasium 0.29.1, numpy 1.26.4, transformers 4.51.3, plus the ManiSkill2 + SimplerEnv editable installs. +Clones SimplerEnv (and its nested `ManiSkill2_real2sim`) into [`eval/sim/simpler/SimplerEnv/`](eval/sim/simpler/SimplerEnv), creates `eval/sim/simpler/simpler_uv/.venv/`. ## Running the server `vla-server` loads the model once at startup and answers ZeroMQ REQ/REP requests synchronously. ```bash -# SmolVLA / π0 (mmproj + ckpt): -./build/vla-server "$VLA_MMPROJ" "$VLA_GGUF" - -# Evo-1 / BitVLA / GR00T-N1.{5,6,7} (vision baked into the ckpt - omit mmproj): ./build/vla-server "$VLA_GGUF" ``` @@ -108,12 +120,25 @@ vla-server: bound to tcp://*:5555. ready. Bound address and port can be configured by `--bind` flag. Stop server with `Ctrl-C`. +## One-shot CLI + +`vla-cli` runs a single prediction without a server or simulator: give it a model, +an image, and the tokenized instruction, and it prints the action chunk. Handy for +smoke-testing a GGUF or scripting a quick inference. + +```bash +./build/vla-cli --ckpt "$VLA_GGUF" \ + --image front.jpg --image wrist.jpg --tokens 1,100,200,2 --pretty +``` + +Tokenization stays in the Python client, so the instruction is passed as token ids. +`--pretty` prints one action row per line; `--state` sets proprioception (defaults to zeros). ## Running the client [`eval/client/`](eval/client/) ships an end-to-end LIBERO benchmark runner that drives `vla-server` directly over the protobuf protocol. Make sure the LIBERO venv from [Install simulators](#install-simulators) is set up first. -### Run an episode (LIBERO) +### LIBERO With `vla-server` already running: @@ -125,13 +150,11 @@ python eval/client/run_sim_client_direct.py \ --arch "$VLA_ARCH" ``` -Note: +The GR00T models needs: + - `--stats-json /path/to/dataset_statistics.json` in client side. + - `VLA_GR00T_EMBODIMENT` (`new_embodiment` for N1.5, `libero_panda` for N1.6, `libero_sim` for N1.7) and `VLA_GR00T_BF16_WEIGHTS=1` (to fit the 8 GB card) in server side. -- Use proper `--arch` flag (see [Models](#models)) to match the GGUF that `vla-server` is serving. -- Pi0 uses the gated `google/paligemma-3b-pt-224` tokenizer (`huggingface-cli login` + accept the licence, or point `--tokenizer` at a local copy) -- The GR00T arches need `--stats-json /dataset_statistics.json` (action/state un-normalisation) and an embodiment selected server-side via `VLA_GR00T_EMBODIMENT` (`new_embodiment` for N1.5, `libero_panda` for N1.6, `libero_sim` for N1.7), with `VLA_GR00T_BF16_WEIGHTS=1` to fit the 8 GB card. - -### Run an episode (SimplerEnv) +### SimplerEnv So far only **GR00T-N1.6** is wired (the `gr00t-n1d6-bridge` checkpoint with the `oxe_widowx` embodiment). Start `vla-server` on port 5566 with `oxe_widowx` embodiment: @@ -151,30 +174,18 @@ python eval/client/run_simpler_client_direct.py \ --stats-json "$VLA_STATS_JSON" ``` -`$VLA_STATS_JSON` is the `statistics.json` shipped beside the bridge GGUF. The default 224-px GGUF mis-localises on WidowX (≈20% success) - the 252-px build is required. - ## Models -Each model ships a combined VLA GGUF (LM + action expert + dataset stats + arch config) and, where applicable, a matching mmproj GGUF (vision tower). SmolVLA, π0 and π0.5 ship a separate mmproj; BitVLA, Evo-1, VLA-Adapter and OpenVLA-OFT bake their vision tower into the combined GGUF, so no mmproj file is needed. - -| Model | Converted GGUF | Source ckpt | Client `--arch` flag | -|---|---|---|---| -| SmolVLA | [`smolvla-libero`](https://huggingface.co/vrfai/smolvla-libero-gguf) | [link](https://huggingface.co/HuggingFaceVLA/smolvla_libero) | `smolvla` | -| π0 | [`pi0-libero`](https://huggingface.co/vrfai/pi0-libero-finetuned-v044-gguf) | [link](https://huggingface.co/lerobot/pi0_libero_finetuned_v044) | `pi0` | -| BitVLA | [`bitvla-libero`](https://huggingface.co/vrfai/bitvla-libero-gguf) | [link](https://huggingface.co/hongyuw/ft-bitvla-bitsiglipL-224px-libero_object-bf16) | `bitvla` | -| OpenVLA-OFT | [`openvla-oft-libero`](https://huggingface.co/vrfai/openvla-oft-libero-gguf) | [link](https://huggingface.co/moojink/openvla-7b-oft-finetuned-libero-spatial-object-goal-10) | `openvla_oft` | -| Groot-N1.7 | [`gr00tn1d7-libero`](https://huggingface.co/vrfai/gr00tn1d7-libero-gguf) | [link](https://huggingface.co/nvidia/GR00T-N1.7-LIBERO) | `gr00t_n1_7` | +### Conversion -`vla.cpp` also supports `gr00t_n1_5`, `gr00t_n1_6`, `evo1`, `vla_adapter`, `vla_jepa` and `pi05`. +Each model ships as a single self-contained GGUF. If you would rather convert a HuggingFace safetensors checkpoint yourself, [`scripts/`](scripts/) provides per-arch GGUF converters. Set up a venv for converter by: ```bash -# Assume third_party/llama.cpp has been cloned and patched python3 -m venv .venv-converter source .venv-converter/bin/activate -pip install -r third_party/llama.cpp/requirements/requirements-convert_hf_to_gguf.txt -pip install safetensors +pip install -e ".[convert]" ``` Then run any of the per-arch converters (`--help` for the full flag list): @@ -185,28 +196,36 @@ python scripts/convert_smolvla_to_gguf.py \ --out /path/to/smolvla-libero-bf16.gguf ``` +### Quantization + +The shipped GGUFs are bf16. `scripts/quantize_gguf.py` repacks the LM-backbone weight +matrices to a smaller type and copies everything else unchanged; the loader keeps the +packed weights and lets `ggml_mul_mat` dequantize at compute, so the file just loads and +runs like the bf16 one. + +```bash +python scripts/quantize_gguf.py --in model-bf16.gguf --out model-q8_0.gguf --type Q8_0 +``` + +`Q8_0` is near-lossless and roughly halves the LM. `Q4_0` is 4-bit for a bigger cut. +Embeddings, the output head, norms and the action expert stay float; pass `--vision` to +pack the vision tower too (smaller, but more accuracy loss). + ## Benchmarks -Full `libero_object` sweep - all 10 tasks × 20 episodes (200 episodes per arch), -run via `vla-server` + `eval/client/run_sim_client_direct.py` across three deployment -targets: an **RTX 3060** (sm_86), an **NVIDIA Jetson AGX Orin** (sm_87, Jetson-class -deployment hardware), and an **NVIDIA Jetson Orin Nano (8 GB)** (sm_87, the cheapest -Jetson and the project's primary deployment target). On the Orin Nano, the 8 GB budget -forces the ~6 GB servers (`gr00t_n1_5`, `pi0`) to run split - server on the Nano, client -(sim) on the RTX 3060 - while lighter models run co-located; GR00T-N1.6 and GR00T-N1.7 -could not be loaded on 8 GB and are omitted there. - -Columns report `client/call (ms)` and `Peak RAM (MiB)` for each hardware target. - -| Model | 3060 call (ms) | 3060 RAM (MiB) | AGX Orin call (ms) | AGX Orin RAM (MiB) | Orin Nano call (ms) | Orin Nano RAM (MiB) | -|---|---:|---:|---:|---:|---:|---:| -| `smolvla` | 113 | 1410 | 262 | 689.4 | 567 | 2031.2 | -| `bitvla` | 303 | 1312 | 809 | 1148.8 | 2845 | 2199.0 | -| `evo1` | 509 | 1564 | 1048 | 637.5 | 3671 | 2135.0 | -| `pi0` | 312 | 5548 | 893 | 640.4 | 1955 | 6067.7 | -| `gr00t_n1_5` | 227 | 4866 | 461 | 1331.3 | 1356 | 5974.9 | -| `gr00t_n1_6` | 165 | 6048 | 427 | 1340.5 | - | - | -| `gr00t_n1_7` | 164 | 6302 | 429 | 1316.5 | - | - | +Latency (ms, inference time + transport overhead) measured at client sides +across four deployment targets: an **RTX 3090**, an **NVIDIA Jetson AGX +Orin**, an **NVIDIA Jetson Orin Nano (8 GB)**, and an **Apple M4**. + + +| Model | 3090 call (ms) | AGX Orin call (ms) | Orin Nano call (ms) | M4 call (ms) | +|---|---:|---:|---:|---:| +| `smolvla` | 113 | 262 | 567 | 888 | +| `pi0` | 312 | 893 | 1955 | 1135 | +| `gr00t_n1_5` | 227 | 461 | 1356 | - | +| `gr00t_n1_7` | 164 | 429 | - | 755 | +| `bitvla` | 303 | 809 | 2845 | - | +| `evo1` | 509 | 1048 | 3671 | - | ## Roadmap @@ -215,17 +234,17 @@ supported (released and benchmarked), `~` = in progress, `-` = planned. | Model | CPU (x86-64 / ARM) | CUDA | Metal | OpenVINO | Hexagon | |---|:--:|:--:|:--:|:--:|:--:| -| SmolVLA | Y | Y | Y | - | - | -| π0 | Y | Y | Y | - | - | -| π0.5 | Y | Y | ~ | - | - | -| GR00T N1.5 | Y | Y | ~ | - | - | -| GR00T N1.6 | Y | Y | ~ | - | - | -| GR00T N1.7 | Y | Y | Y | - | - | -| BitVLA | Y | Y | ~ | - | - | -| Evo-1 | Y | Y | ~ | - | - | -| VLA-Adapter | Y | Y | ~ | - | - | -| OpenVLA-OFT | Y | Y | ~ | - | - | -| VLA-JEPA | Y | Y | ~ | - | - | +| [SmolVLA](https://hf.co/vrfai/smolvla-libero-gguf) | Y | Y | Y | - | - | +| [π0](https://hf.co/vrfai/pi0-libero-finetuned-v044-gguf) | Y | Y | Y | - | - | +| [π0.5](https://hf.co/vrfai/pi05-libero-gguf) | Y | Y | ~ | - | - | +| [GR00T N1.5](https://hf.co/vrfai/gr00tn1d5-libero-object-gguf) | Y | Y | ~ | - | - | +| [GR00T N1.6](https://hf.co/vrfai/gr00tn1d6-libero-gguf) | Y | Y | ~ | - | - | +| [GR00T N1.7](https://hf.co/vrfai/gr00tn1d7-libero-gguf) | Y | Y | Y | - | - | +| [BitVLA](https://hf.co/vrfai/bitvla-libero-gguf) | Y | Y | ~ | - | - | +| [Evo-1](https://hf.co/vrfai/evo1-libero-gguf) | Y | Y | ~ | - | - | +| [VLA-Adapter](https://hf.co/vrfai/vla-adapter-libero-gguf) | Y | Y | ~ | - | - | +| [OpenVLA-OFT](https://hf.co/vrfai/openvla-oft-libero-gguf) | Y | Y | ~ | - | - | +| [VLA-JEPA](https://hf.co/vrfai/vla-jepa-libero) | Y | Y | ~ | - | - | Looking ahead, we will support more models, more platforms, and continue to optimize the framework. diff --git a/ci/README.md b/ci/README.md index 9845450..ae36dfe 100644 --- a/ci/README.md +++ b/ci/README.md @@ -29,8 +29,8 @@ platform; each platform is a remote server reached over the LAN. | Platform | Models | Suites | |---|---|---| -| `rtx3090` | all 7 (smolvla, pi0, bitvla, evo1, gr00t_n1_5/6/7) | `libero_object`; **+ spatial/object/goal/10** for bitvla & gr00t_n1_7 | -| `orin` | all except gr00t_n1_6/_7 | `libero_object` | +| `rtx3090` | all 10 (smolvla, pi0, pi05, bitvla, evo1, vla_adapter, openvla_oft, gr00t_n1_5/6/7) | `libero_object`; **+ spatial/object/goal/10** for bitvla & gr00t_n1_7 | +| `orin` | smolvla, pi0, bitvla, evo1, gr00t_n1_5 | `libero_object` | | `m4` | smolvla, pi0, gr00t_n1_7 | `libero_object` | Every cell is **10 tasks × 1 episode**. diff --git a/ci/config/matrix.env b/ci/config/matrix.env index 463ec76..94da008 100644 --- a/ci/config/matrix.env +++ b/ci/config/matrix.env @@ -43,7 +43,7 @@ models_for() { local v="MODELS_$1"; echo "${!v}"; } MULTISUITE_MODELS_rtx3090="bitvla gr00t_n1_7" MULTISUITE_SUITES="libero_spatial libero_object libero_goal libero_10" -multisuite_models_for() { local v="MULTISUITE_MODELS_$1"; echo "${!v:-}"; } +multisuite_models_for() { local v=KNOWN_ISSUES"MULTISUITE_MODELS_$1"; echo "${!v:-}"; } # Suites a given (platform, model) should run = DEFAULT_SUITE, plus the # multisuite set if the model is listed for that platform. diff --git a/ci/lib/common.sh b/ci/lib/common.sh index eeb4efa..f2bff2c 100755 --- a/ci/lib/common.sh +++ b/ci/lib/common.sh @@ -77,17 +77,14 @@ server_args_for() { local arch="$1" root="$2" suite="${3:-${DEFAULT_SUITE:-libero_object}}" sd tok case "$arch" in pi0) - echo "${root}/pi0-libero-finetuned-v044-gguf/mmproj-pi0-libero-finetuned-v044.gguf" echo "${root}/pi0-libero-finetuned-v044-gguf/pi0-libero-finetuned-v044.gguf" ;; pi05) - echo "${root}/pi05-libero-gguf/mmproj-pi05-libero.gguf" echo "${root}/pi05-libero-gguf/pi05-libero.gguf" ;; vla_adapter) echo "${root}/vla-adapter-libero-object-gguf/libero_object/vla-adapter-libero-object.gguf" ;; openvla_oft) echo "${root}/openvla-oft-libero-gguf/openvla-oft-libero.gguf" ;; smolvla) - echo "${root}/smolvla-libero-gguf/mmproj-smolvla-libero.gguf" echo "${root}/smolvla-libero-gguf/smolvla-libero.gguf" ;; evo1) echo "${root}/evo1-libero-gguf/evo1-libero.gguf" ;; diff --git a/docs/ADOPTION.md b/docs/ADOPTION.md new file mode 100644 index 0000000..b6db733 --- /dev/null +++ b/docs/ADOPTION.md @@ -0,0 +1,29 @@ +# Adoption notes + +vla.cpp is technically solid (7 architectures, self-contained GGUFs, CUDA + Jetson, +real benchmarks). The gap to llama.cpp-style reach is mostly distribution, not code. +Ordered by leverage: + +1. **Publish on GitHub.** The repo lives on Bitbucket (`bitbucket.org/vinrobotics/vla.cpp`) + while the README already links a `github.com/VinRobotics/vla.cpp` URL. llama.cpp's reach + came from GitHub visibility, issues, and PRs. A public GitHub mirror is the single biggest + lever; nothing else here matters as much. + +2. **Ship the models.** All seven GGUFs are already published under + [`vrfai`](https://huggingface.co/vrfai) on the Hub - the README's "coming soon" rows are + stale (now fixed). Keep the model table pointing at the real repos so the policies are + one `hf download` away. + +3. **Cut releases.** `v0.1.0` is the first tag (see `CHANGELOG.md`). Tagged releases + + changelog give users something to pin and cite. + +4. **Rotate the committed credential.** The local `.git/config` remote URL embeds an + access token (`https://@bitbucket.org/...`). It is never pushed (git config is not + tracked), so this is hygiene, not a live leak - but rotate it and use a credential helper + or SSH remote instead of an inline token. + +5. **Lower the build bar (optional).** A CUDA `Dockerfile` now exists; publishing a prebuilt + image (and, later, macOS/Metal or ROCm backends) removes the from-source step that stops + most drive-by users. + +None of these change inference behaviour; they change who can find and run it. diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md new file mode 100644 index 0000000..3ad8cda --- /dev/null +++ b/docs/ARCHITECTURE.md @@ -0,0 +1,67 @@ +# Architecture + +vla.cpp runs Vision-Language-Action (VLA) policies on the ggml/llama.cpp runtime. +Every model is a self-contained GGUF (or a checkpoint plus a vision mmproj) that the +engine loads, detects, and drives on CPU, CUDA, or Metal. This page is the map; the +source is the detail. + +## Layers + +- `src/model.h` - public API: `model_load`, `predict`, `model_config`, `last_stats`. +- `src/arch.h` - the `Arch` enum, the `ModelArchBase` interface, and one `*_create` + factory per architecture. +- `src/model.cpp` - loads a checkpoint, detects the architecture from its GGUF keys + (or safetensors namespace), and dispatches to the matching factory. +- `src/models/*.cpp` - one translation unit per architecture. Each owns its ggml + contexts, vision tower, weights, and compute graph. +- `src/models/gguf_reader.h` - the shared GGUF reader (metadata, tensor bytes, + on-demand embedding rows). +- `src/models/vision_common.h` - small pure vision helpers (pixel-shuffle, view checks). +- `src/serving/` - `vla-server` (ZeroMQ + protobuf, action prediction), `vlm-server` + (chat), and `vla-cli` (one-shot inference). +- `src/kernels/bitvla/` - custom 1.58-bit ternary CUDA kernels for BitVLA. + +## The prediction path + +A forward pass has two stages that most architectures share. + +1. **Prefix.** Camera views go through a vision tower, language tokens through the + embedding table, and proprioception through a small projection. Concatenated, they + form the prefix that the language backbone attends over (bidirectionally, minus any + padded language tokens). + +2. **Action head.** A smaller expert reads the prefix and produces an action chunk of + shape `[num_steps, max_action_dim]`, where only the first `real_action_dim` columns + carry values and the rest are zero padding. The head comes in three flavours: + - **Flow-matching expert** (SmolVLA, pi0, pi0.5): integrates a velocity field with + Euler steps from noise at `t=1` to the action at `t=0`. SmolVLA alternates + self-attention among action tokens with cross-attention to the prefix. + - **DiT** (GR00T N1.5/1.6/1.7, Evo-1, VLA-JEPA): a diffusion transformer head. + - **Parallel decode** (BitVLA, OpenVLA-OFT, VLA-Adapter): OpenVLA-OFT-style + bidirectional decode of the action tokens in a single pass. + +Actions leave `predict` normalised to the training statistics; the caller +un-normalises into world units. + +## Vision deployment + +Two patterns, chosen per architecture: + +- **Baked-in tower** (BitVLA, Evo-1, GR00T, OpenVLA-OFT, VLA-Adapter, VLA-JEPA): the + vision tower ships inside the combined GGUF, so a single file is enough. +- **Separate mmproj** (SmolVLA, pi0, pi0.5): the SigLIP tower ships as a second GGUF; + merge it into the checkpoint with `scripts/merge_*_mmproj_to_gguf.py` before loading. + +## Backends and packaging + +llama.cpp is fetched and pinned by CMake `FetchContent`; a bump is a one-line +`GIT_TAG` change. Weights are bf16 by default and can be repacked to Q8_0/Q4_0 with +`scripts/quantize_gguf.py`; the loader runs quantized GGUFs directly and lets +`ggml_mul_mat` dequantize at compute. CPU thread count scales to the machine core +count; CUDA and Metal run the towers and the transformer on the GPU. + +## Adding an architecture + +Extend the `Arch` enum, declare a `*_create` factory in `arch.h`, implement it under +`src/models/`, wire detection and dispatch in `src/model.cpp`, and add a converter in +`scripts/`. Reuse `gguf_reader.h` and `vision_common.h` rather than copying them. diff --git a/docs/KNOWN_ISSUES.md b/docs/KNOWN_ISSUES.md new file mode 100644 index 0000000..c071692 --- /dev/null +++ b/docs/KNOWN_ISSUES.md @@ -0,0 +1,31 @@ +# Known issues + +## Evo-1 scores 0% on LIBERO despite the 94.5% benchmark (NEEDS DOUBLE-CHECK) + +The released `vrfai/evo1-libero-gguf` scores **0%** on `libero_object` (checked on +task_0/1/2, RTX 5090) while the README benchmark reports **94.5%**. + +Evidence this is a real, pre-existing regression, not a harness or config artifact: +- **smolvla scores 100%** on the same harness, build, and task, so the eval path works. +- evo1 `predict_check` is bit-identical on llama.cpp b9860 and b9866, so the bump did + not change evo1 numerics. +- The model loads and emits finite, non-degenerate actions (~13 ms/step); episodes + just never complete the task. +- Client wiring looks right: image_size 448, InternVL3 tokenizer, a dedicated evo1 + path, and a token-count guard that fails loudly on mismatch rather than scoring 0. + +Most likely the uploaded `evo1-libero.gguf` diverged from the checkpoint that produced +the 94.5% number, or the evo1 loader/converter regressed before it was benchmarked. +Do not blind-fix the flow-matching base noise or attention masking: the base is +N(0,1), and switching it to the reference's uniform[-1,1] also gives 0%. + +Double-check: diff the released GGUF's tensors and config against a fresh +`convert_evo1_to_gguf.py` from `MINT-SJTU/Evo1_LIBERO`, and compare one denoise step's +`v_t` against the Evo-1 reference on identical inputs. + +Status: the N(0,1) base noise matches the reference and is in the tree. The DiT +cross-attention is left unmasked to match the Evo-1 reference (its `flow_matching.py` +attends over the full padded context with no key mask). Both masked and unmasked score +0/3 on `libero_object` task_0 here, so the attention mask is not the cause. The released +GGUF is the remaining suspect; the tensor/config diff above is the open step and needs +the upstream safetensors checkpoint, which is not bundled here. diff --git a/eval/README.md b/eval/README.md new file mode 100644 index 0000000..0aba3d4 --- /dev/null +++ b/eval/README.md @@ -0,0 +1,86 @@ +# Evaluation + +The eval scaffold drives `vla-server` against two simulators end-to-end over the ZeroMQ + +protobuf protocol. The C++ server does all model inference on CPU/GPU; the Python client only runs +the simulator and the per-arch normalisation, so it stays on CPU. + +## Install simulators + +Each setup script bootstraps an isolated Python 3.10 [`uv`](https://github.com/astral-sh/uv) venv +next to itself and clones the upstream sim. Requires `uv` on `PATH`. + +### LIBERO + +```bash +bash eval/sim/libero/setup_libero.sh +``` + +Clones LIBERO into `eval/sim/libero/LIBERO/`, creates `eval/sim/libero/libero_uv/.venv/`, pins +compatible torch / lerobot / transformers / gymnasium (and `mujoco==2.3.2`, required by +robosuite 1.4.0), and seeds `~/.libero/config.yaml` non-interactively. + +### SimplerEnv + +```bash +bash eval/sim/simpler/setup_SimplerEnv.sh +``` + +Clones SimplerEnv (and its nested `ManiSkill2_real2sim`) into `eval/sim/simpler/SimplerEnv/`, +creates `eval/sim/simpler/simpler_uv/.venv/`, and pins its ManiSkill2 + SimplerEnv editable +installs. + +## Run an episode (LIBERO) + +Start the server, then drive it from the LIBERO venv: + +```bash +./build/vla-server "$VLA_GGUF" # terminal 1 + +# terminal 2 +MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 \ +eval/sim/libero/libero_uv/.venv/bin/python eval/client/run_sim_client_direct.py \ + --arch "$VLA_ARCH" \ + --task libero_object --task-id 0 --n-episodes 1 \ + --output-dir /tmp/libero_outputs +``` + +Notes: + +- `--arch` must match the served GGUF (see the model table in the top-level README). +- **Rendering + torch on one box**: run the client with `MUJOCO_GL=egl` and + `CUDA_VISIBLE_DEVICES=0`. robosuite's EGL renderer needs a valid device index (an empty + `CUDA_VISIBLE_DEVICES` breaks it); the client's torch stays on CPU, so it does not need a + torch build matching the GPU's compute capability. The server (a separate process) uses the + GPU for inference. +- **π0** uses the gated `google/paligemma-3b-pt-224` tokenizer. Run `huggingface-cli login` and + accept the licence, or point `--tokenizer` at a local copy. +- **GR00T** arches need `--stats-json /dataset_statistics.json` and an embodiment selected + server-side via `VLA_GR00T_EMBODIMENT` (`new_embodiment` for N1.5, `libero_panda` for N1.6, + `libero_sim` for N1.7), plus `VLA_GR00T_BF16_WEIGHTS=1` to fit an 8 GB card. + +To sweep every model over `libero_object` tasks 0–9, use `eval/run_libero.sh -i `. + +## Run an episode (SimplerEnv) + +So far only **GR00T-N1.6** is wired (the `gr00t-n1d6-bridge` checkpoint with the `oxe_widowx` +embodiment). Serve it, then drive from the SimplerEnv venv: + +```bash +VLA_GR00T_BF16_WEIGHTS=1 VLA_GR00T_EMBODIMENT=oxe_widowx \ + ./build/vla-server "$GR00T_N1D6_GGUF" + +eval/sim/simpler/simpler_uv/.venv/bin/python eval/client/run_simpler_client_direct.py \ + --arch gr00t_n1_6 \ + --task-id oxe_widowx/widowx_spoon_on_towel --n-episodes 1 \ + --stats-json "$VLA_STATS_JSON" \ + --embodiment oxe_widowx --image-size 252 +``` + +`$VLA_STATS_JSON` is the `statistics.json` beside the bridge GGUF. The default 224-px GGUF +mis-localises on WidowX (≈20% success); the 252-px build is required. + +## Reports + +`eval/collect_libero_results.py` / `collect_simpler_results.py` aggregate per-episode outputs into +the markdown reports under [`reports/`](reports/); `scripts/print_versions.sh` emits the +reproducibility block (host, toolchain, GGUF hashes) for each. diff --git a/eval/collect_libero_results.py b/eval/collect_libero_results.py index 8ec132e..6352fa2 100644 --- a/eval/collect_libero_results.py +++ b/eval/collect_libero_results.py @@ -342,7 +342,7 @@ def main() -> int: help="do not write a markdown report") ap.add_argument("--gguf", type=Path, nargs="*", default=[], metavar="PATH", help="GGUF files to checksum in the ## Reproducibility block " - "(combined ckpt + mmproj where applicable)") + "(the self-contained ckpt GGUF)") ap.add_argument("--libero-venv", type=Path, default=REPO_ROOT / "eval" / "sim" / "libero" / "libero_uv" / ".venv", help="LIBERO uv venv whose torch/transformers/lerobot/numpy pins are " diff --git a/eval/run_libero.sh b/eval/run_libero.sh index aa32b17..63cc9b3 100644 --- a/eval/run_libero.sh +++ b/eval/run_libero.sh @@ -38,7 +38,7 @@ Usage: $(basename "$0") -i [-o ] [-n ] [- (default: all) -h show this help -Env overrides: BIND_ADDR, CLIENT_ADDR, BITVLA_TOKENIZER, +Env overrides: BIND_ADDR, CLIENT_ADDR, BITVLA_TOKENIZER, GR00T_N1_6_TOKENIZER, GR00T_N1_5_STATS, GR00T_N1_6_STATS, GR00T_N1_7_STATS EOF } @@ -98,6 +98,10 @@ TASK_SUITE="libero_object" # the Hub. Optional override (offline): BITVLA_TOKENIZER=/path/to/bitvla-ckpt-dir BITVLA_TOKENIZER="${BITVLA_TOKENIZER:-}" +# gr00t_n1_6 has no HF-default tokenizer; its Eagle tokenizer is vendored in the +# model dir. Defaults to the model dir; override with GR00T_N1_6_TOKENIZER. +GR00T_N1_6_TOKENIZER="${GR00T_N1_6_TOKENIZER:-}" + # GR00T-N1.5 / N1.6 / N1.7 need a dataset_statistics.json GR00T_N1_5_STATS="${GR00T_N1_5_STATS:-}" GR00T_N1_6_STATS="${GR00T_N1_6_STATS:-}" @@ -273,7 +277,7 @@ run_model() { local n_action_steps="$3" local stats_json="$4" # optional dataset_statistics.json (gr00t_n1_{6,7}); "" to skip shift 4 - local server_args=("$@") # bitvla/evo1/gr00t_n1_{6,7}: just ckpt; smolvla/pi0: mmproj + ckpt + local server_args=("$@") local client_extra=() # bitvla auto-loads tokenizer + dataset_statistics.json from the GGUF repo on @@ -281,6 +285,11 @@ run_model() { if [[ "${arch}" == "bitvla" && -n "${BITVLA_TOKENIZER}" ]]; then client_extra+=(--tokenizer "${BITVLA_TOKENIZER}") fi + # gr00t_n1_6 has no HF-default tokenizer; point the client at the vendored + # Eagle tokenizer in the model dir (override via GR00T_N1_6_TOKENIZER). + if [[ "${arch}" == "gr00t_n1_6" ]]; then + client_extra+=(--tokenizer "${GR00T_N1_6_TOKENIZER:-${model_dir}}") + fi if [[ -n "${stats_json}" ]]; then client_extra+=(--stats-json "${stats_json}") fi @@ -348,30 +357,24 @@ if should_run pi0; then "${MODELS_ROOT}/pi0-libero-finetuned-v044-gguf" \ "${N_ACTION_STEPS_PI0}" \ "" \ - "${MODELS_ROOT}/pi0-libero-finetuned-v044-gguf/mmproj-pi0-libero-finetuned-v044.gguf" \ "${MODELS_ROOT}/pi0-libero-finetuned-v044-gguf/pi0-libero-finetuned-v044.gguf" fi -# pi05: PaliGemma + Gemma-300m adaRMS expert; mmproj + ckpt (same layout as pi0). -# State quantiles auto-load from the lerobot/libero dataset; the pi05 client preset -# sets max_length=200 automatically. Pass --stats-json via PI05_STATS to pin a -# local LIBERO meta/stats.json (offline). if should_run pi05; then run_model pi05 \ "${MODELS_ROOT}/pi05-libero-gguf" \ "${N_ACTION_STEPS_PI05}" \ "${PI05_STATS:-}" \ - "${MODELS_ROOT}/pi05-libero-gguf/mmproj-pi05-libero.gguf" \ "${MODELS_ROOT}/pi05-libero-gguf/pi05-libero.gguf" fi +# smolvla: vision baked into the ckpt (no mmproj). if should_run smol; then run_model smolvla \ - "${MODELS_ROOT}/smolvla-libero-bf16-gguf" \ + "${MODELS_ROOT}/smolvla-libero-gguf" \ "${N_ACTION_STEPS_SMOL}" \ "" \ - "${MODELS_ROOT}/smolvla-libero-bf16-gguf/mmproj-smolvla-libero.gguf" \ - "${MODELS_ROOT}/smolvla-libero-bf16-gguf/smolvla-libero.gguf" + "${MODELS_ROOT}/smolvla-libero-gguf/smolvla-libero.gguf" fi # evo1: vision baked into ckpt (no mmproj) @@ -387,10 +390,10 @@ fi # GGUF repo on the Hub. Set BITVLA_TOKENIZER= to override (offline). if should_run bit; then run_model bitvla \ - "${MODELS_ROOT}/bitvla-libero-object-gguf" \ + "${MODELS_ROOT}/bitvla-libero-gguf/libero_object" \ "${N_ACTION_STEPS_BIT}" \ "" \ - "${MODELS_ROOT}/bitvla-libero-object-gguf/bitvla-libero-object-int2.gguf" + "${MODELS_ROOT}/bitvla-libero-gguf/libero_object/bitvla-libero-object.gguf" fi # vla_adapter: Qwen2.5-0.5B + Bridge-Attention; vision baked in (no mmproj), @@ -438,17 +441,17 @@ fi # gr00t_n1_6: needs GR00T_N1_6_STATS dataset_statistics.json if should_run gr00t_n1_6; then - g6_stats_default="${MODELS_ROOT}/gr00t-n1d6-libero-gguf/dataset_statistics.json" + g6_stats_default="${MODELS_ROOT}/gr00tn1d6-libero-gguf/dataset_statistics.json" g6_stats="${GR00T_N1_6_STATS:-}" if [[ -z "${g6_stats}" && -f "${g6_stats_default}" ]]; then g6_stats="${g6_stats_default}" fi if [[ -n "${g6_stats}" && -f "${g6_stats}" ]]; then run_model gr00t_n1_6 \ - "${MODELS_ROOT}/gr00t-n1d6-libero-gguf" \ + "${MODELS_ROOT}/gr00tn1d6-libero-gguf" \ "${N_ACTION_STEPS_GR00T_N1_6}" \ "${g6_stats}" \ - "${MODELS_ROOT}/gr00t-n1d6-libero-gguf/gr00t-n1d6-libero.gguf" + "${MODELS_ROOT}/gr00tn1d6-libero-gguf/gr00tn1d6-libero.gguf" else echo "[skip] gr00t_n1_6: set GR00T_N1_6_STATS= to enable" fi @@ -456,14 +459,14 @@ fi # gr00t_n1_7: needs dataset_statistics.json if should_run gr00t_n1_7; then - g7_stats_default="${MODELS_ROOT}/gr00t-n1d7-libero-object-gguf/dataset_statistics.json" + g7_stats_default="${MODELS_ROOT}/gr00tn1d7-libero-gguf/libero_object/dataset_statistics.json" g7_stats="${GR00T_N1_7_STATS:-${g7_stats_default}}" if [[ -f "${g7_stats}" ]]; then run_model gr00t_n1_7 \ - "${MODELS_ROOT}/gr00t-n1d7-libero-object-gguf" \ + "${MODELS_ROOT}/gr00tn1d7-libero-gguf/libero_object" \ "${N_ACTION_STEPS_GR00T_N1_7}" \ "${g7_stats}" \ - "${MODELS_ROOT}/gr00t-n1d7-libero-object-gguf/gr00t-n1d7-libero-object.gguf" + "${MODELS_ROOT}/gr00tn1d7-libero-gguf/libero_object/gr00tn1d7-libero-object.gguf" else echo "[skip] gr00t_n1_7: dataset_statistics.json not found at ${g7_stats}; set GR00T_N1_7_STATS to override" fi diff --git a/eval/run_libero_client.sh b/eval/run_libero_client.sh index 02fb41c..140bda9 100644 --- a/eval/run_libero_client.sh +++ b/eval/run_libero_client.sh @@ -137,7 +137,6 @@ print_server_cmd() { pi0) cat < just ckpt. Weights live here on -# the server, NOT on the client. (dataset_statistics.json is a CLIENT-side arg.) +# pi0 / gr00t_*: vision baked into the ckpt -> just the self-contained ckpt GGUF. +# Weights live here on the server, NOT on the client. (dataset_statistics.json is a +# CLIENT-side arg.) SERVER_ARGS=() case "${MODEL}" in pi0) SERVER_ARGS=( - "${MODELS_ROOT}/pi0-libero-finetuned-v044-gguf/mmproj-pi0-libero-finetuned-v044.gguf" "${MODELS_ROOT}/pi0-libero-finetuned-v044-gguf/pi0-libero-finetuned-v044.gguf" ) ;; diff --git a/eval/sim/libero/setup_libero.sh b/eval/sim/libero/setup_libero.sh index 718eb5a..88dc9f0 100644 --- a/eval/sim/libero/setup_libero.sh +++ b/eval/sim/libero/setup_libero.sh @@ -15,6 +15,10 @@ set -euxo pipefail +# egl-probe and friends request cmake_minimum_required < 3.5, which CMake >= 4.0 +# refuses; this lets them configure anyway (env var honored since CMake 3.31). +export CMAKE_POLICY_VERSION_MINIMUM=3.5 + # Get the directory where this script is located SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" @@ -39,7 +43,7 @@ else git clone "$LIBERO_GIT_URL" "$LIBERO_REPO" fi -rm -rf "$LIBERO_UV_ENV"deac +rm -rf "$LIBERO_UV_ENV" mkdir -p "$LIBERO_UV_ENV" uv venv "$LIBERO_UV_ENV/.venv" --python 3.10 source "$LIBERO_UV_ENV/.venv/bin/activate" @@ -52,5 +56,11 @@ uv pip install pandas==2.0.3 uv pip install pyarrow==12.0.1 uv pip install diffusers==0.30.1 uv pip install numpy==1.26.4 +# robosuite 1.4.0 calls the mujoco 2.3 mj_fullM(m, dst, M) signature; mujoco 3.x +# changed it, so a transitive dep pulling 3.x breaks the OSC controller at reset. +uv pip install mujoco==2.3.2 +# LIBERO prompts for a dataset path on first import if ~/.libero/config.yaml is +# missing, which hangs a headless eval; clear stale config and seed defaults (N). rm -rf "$HOME/.libero" +echo "N" | MUJOCO_GL=egl python -c "import libero.libero" || true diff --git a/patches/llama.cpp-vla.patch b/patches/llama.cpp-vla.patch deleted file mode 100644 index 91257fa..0000000 --- a/patches/llama.cpp-vla.patch +++ /dev/null @@ -1,592 +0,0 @@ -diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py -index 7f4a4018f..6be188287 100755 ---- a/convert_hf_to_gguf.py -+++ b/convert_hf_to_gguf.py -@@ -7323,6 +7323,32 @@ class Gemma3VisionModel(MmprojModel): - return # skip other tensors - - -+# vla.cpp: PaliGemma vision tower + projector (π0). See patches/PATCH.md. -+@ModelBase.register("PaliGemmaForConditionalGeneration") -+class PaliGemmaVisionModel(MmprojModel): -+ def set_gguf_parameters(self): -+ super().set_gguf_parameters() -+ hparams = self.hparams -+ self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.PALIGEMMA) -+ self.gguf_writer.add_vision_attention_layernorm_eps(hparams.get("layer_norm_eps", 1e-6)) -+ self.gguf_writer.add_vision_use_gelu(True) -+ -+ def tensor_force_quant(self, name, new_name, bid, n_dims): -+ if ".embeddings." in name: -+ return gguf.GGMLQuantizationType.F32 -+ return super().tensor_force_quant(name, new_name, bid, n_dims) -+ -+ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: -+ if name.startswith("multi_modal_projector.") or name.startswith("vision_tower."): -+ if name.startswith("multi_modal_projector.linear."): -+ name = name.replace( -+ "multi_modal_projector.linear.", -+ "multi_modal_projector.mm_input_projection.", -+ ) -+ yield from super().modify_tensors(data_torch, name, bid) -+ return -+ -+ - class ConformerAudioModel(MmprojModel): - _batch_norm_tensors: list[dict[str, Tensor]] | None = None - -diff --git a/ggml/src/ggml-cuda/fattn-tile.cu b/ggml/src/ggml-cuda/fattn-tile.cu -index d60634cc0..83bcd3f8c 100644 ---- a/ggml/src/ggml-cuda/fattn-tile.cu -+++ b/ggml/src/ggml-cuda/fattn-tile.cu -@@ -10,6 +10,10 @@ void ggml_cuda_flash_attn_ext_tile(ggml_backend_cuda_context & ctx, ggml_tensor - GGML_ASSERT(V->ne[0] == K->ne[0]); - ggml_cuda_flash_attn_ext_tile_case< 40, 40>(ctx, dst); - } break; -+ case 48: { -+ GGML_ASSERT(V->ne[0] == K->ne[0]); -+ ggml_cuda_flash_attn_ext_tile_case< 48, 48>(ctx, dst); -+ } break; - case 64: { - GGML_ASSERT(V->ne[0] == K->ne[0]); - ggml_cuda_flash_attn_ext_tile_case< 64, 64>(ctx, dst); -diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh -index 585f2c228..afbc26f13 100644 ---- a/ggml/src/ggml-cuda/fattn-tile.cuh -+++ b/ggml/src/ggml-cuda/fattn-tile.cuh -@@ -26,6 +26,12 @@ static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_nv - GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40, 40, 16, 256, 2, 64, 40) - GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40, 40, 32, 256, 2, 64, 40) - -+ GGML_CUDA_FATTN_TILE_CONFIG_CASE( 48, 48, 2, 64, 2, 64, 48) -+ GGML_CUDA_FATTN_TILE_CONFIG_CASE( 48, 48, 4, 128, 2, 64, 48) -+ GGML_CUDA_FATTN_TILE_CONFIG_CASE( 48, 48, 8, 256, 2, 64, 48) -+ GGML_CUDA_FATTN_TILE_CONFIG_CASE( 48, 48, 16, 256, 2, 64, 48) -+ GGML_CUDA_FATTN_TILE_CONFIG_CASE( 48, 48, 32, 256, 2, 64, 48) -+ - GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64, 64, 2, 64, 2, 64, 64) - GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64, 64, 4, 128, 2, 64, 64) - GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64, 64, 8, 256, 2, 64, 64) -@@ -88,6 +94,12 @@ static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_nv - GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40, 40, 16, 256, 2, 32, 40) - GGML_CUDA_FATTN_TILE_CONFIG_CASE( 40, 40, 32, 256, 2, 32, 40) - -+ GGML_CUDA_FATTN_TILE_CONFIG_CASE( 48, 48, 2, 64, 2, 32, 48) -+ GGML_CUDA_FATTN_TILE_CONFIG_CASE( 48, 48, 4, 128, 2, 32, 48) -+ GGML_CUDA_FATTN_TILE_CONFIG_CASE( 48, 48, 8, 256, 2, 32, 48) -+ GGML_CUDA_FATTN_TILE_CONFIG_CASE( 48, 48, 16, 256, 2, 32, 48) -+ GGML_CUDA_FATTN_TILE_CONFIG_CASE( 48, 48, 32, 256, 2, 32, 48) -+ - GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64, 64, 2, 128, 3, 64, 64) - GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64, 64, 4, 128, 3, 32, 64) - GGML_CUDA_FATTN_TILE_CONFIG_CASE( 64, 64, 8, 128, 3, 32, 64) -diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu -index 8256591b2..dca4f0b8d 100644 ---- a/ggml/src/ggml-cuda/fattn.cu -+++ b/ggml/src/ggml-cuda/fattn.cu -@@ -357,6 +357,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const - - switch (K->ne[0]) { - case 40: -+ case 48: - case 64: - case 72: - case 80: -@@ -428,7 +429,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const - const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && K->ne[1] % FATTN_KQ_STRIDE == 0; - - // If Turing tensor cores are available, use them: -- if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { -+ if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 48 && Q->ne[0] != 72) { - if (can_use_vector_kernel) { - if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) { - if (cc >= GGML_CUDA_CC_ADA_LOVELACE && Q->ne[1] == 1 && Q->ne[3] == 1 && !(gqa_ratio > 4 && K->ne[1] >= 8192)) { -@@ -452,7 +453,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const - return BEST_FATTN_KERNEL_MMA_F16; - } - -- if (volta_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { -+ if (volta_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 48 && Q->ne[0] != 72) { - int gqa_ratio_eff = 1; - const int ncols2_max = Q->ne[0] == 576 ? 16 : 8; - while (gqa_ratio % (2*gqa_ratio_eff) == 0 && gqa_ratio_eff < ncols2_max) { -@@ -468,14 +469,14 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const - } - - // Use the WMMA kernel if possible: -- if (ggml_cuda_should_use_wmma_fattn(cc) && K->ne[1] % FATTN_KQ_STRIDE == 0 && Q->ne[0] != 40 && Q->ne[0] != 72 && Q->ne[0] != 512 && Q->ne[0] != 576) { -+ if (ggml_cuda_should_use_wmma_fattn(cc) && K->ne[1] % FATTN_KQ_STRIDE == 0 && Q->ne[0] != 40 && Q->ne[0] != 48 && Q->ne[0] != 72 && Q->ne[0] != 512 && Q->ne[0] != 576) { - if (can_use_vector_kernel && Q->ne[1] <= 2) { - return BEST_FATTN_KERNEL_VEC; - } - return BEST_FATTN_KERNEL_WMMA_F16; - } - -- if (amd_wmma_available(cc) && GGML_CUDA_CC_IS_RDNA4(cc) && gqa_opt_applies && Q->ne[0] <= 128 && Q->ne[0] != 40 && Q->ne[0] != 72) { -+ if (amd_wmma_available(cc) && GGML_CUDA_CC_IS_RDNA4(cc) && gqa_opt_applies && Q->ne[0] <= 128 && Q->ne[0] != 40 && Q->ne[0] != 48 && Q->ne[0] != 72) { - if (can_use_vector_kernel) { - if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) { - if (Q->ne[1] == 1) { -@@ -501,7 +502,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const - } - - // Use MFMA flash attention for CDNA (MI100+): -- if (amd_mfma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72 && Q->ne[0] != 256 && Q->ne[0] != 512 && Q->ne[0] != 576) { -+ if (amd_mfma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 48 && Q->ne[0] != 72 && Q->ne[0] != 256 && Q->ne[0] != 512 && Q->ne[0] != 576) { - const int64_t eff_nq = Q->ne[1] * (gqa_opt_applies ? gqa_ratio : 1); - // MMA vs tile crossover benchmarked on MI300X @ d32768: - // hsk=64 (gqa=4): MMA wins at eff >= 128 (+11%) -diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py -index 83ae51ce9..85209b51a 100644 ---- a/gguf-py/gguf/constants.py -+++ b/gguf-py/gguf/constants.py -@@ -4158,6 +4158,7 @@ class VisionProjectorType: - NEMOTRON_V2_VL = "nemotron_v2_vl" - HUNYUANOCR = "hunyuanocr" - HUNYUANVL = "hunyuanvl" -+ PALIGEMMA = "paligemma" # vla.cpp - - - # Items here are (block size, type size) -diff --git a/tools/mtmd/clip-graph.h b/tools/mtmd/clip-graph.h -index d3e7b1ed0..759246612 100644 ---- a/tools/mtmd/clip-graph.h -+++ b/tools/mtmd/clip-graph.h -@@ -32,6 +32,8 @@ struct clip_graph { - float kq_scale; // TODO: maybe move this to hparams - const clip_flash_attn_type flash_attn_type; - -+ bool use_bf16_runtime = false; -+ int batch_size = 1; - ggml_context_ptr ctx0_ptr; - ggml_context * ctx0; - ggml_cgraph * gf; -diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h -index 7d6484eea..7be74331f 100644 ---- a/tools/mtmd/clip-impl.h -+++ b/tools/mtmd/clip-impl.h -@@ -108,6 +108,7 @@ - #define TN_MM_INP_NORM "mm.input_norm.weight" - #define TN_MM_INP_NORM_B "mm.input_norm.bias" - #define TN_MM_INP_PROJ "mm.input_projection.weight" // gemma3 -+#define TN_MM_INP_PROJ_B "mm.input_projection.bias" // paligemma - #define TN_MM_SOFT_EMB_N "mm.soft_emb_norm.weight" // gemma3 - #define TN_MM_PROJECTOR "mm.model.fc.%s" // idefics3, deepseekocr - #define TN_MM_PATCH_MERGER "mm.patch_merger.%s" // mistral small 3.1, glm4v -@@ -304,6 +305,7 @@ enum projector_type { - PROJECTOR_TYPE_NEMOTRON_V2_VL, - PROJECTOR_TYPE_HUNYUANOCR, - PROJECTOR_TYPE_HUNYUANVL, -+ PROJECTOR_TYPE_PALIGEMMA, // vla.cpp - PROJECTOR_TYPE_UNKNOWN, - }; - -@@ -351,6 +353,7 @@ static std::map PROJECTOR_TYPE_NAMES = { - { PROJECTOR_TYPE_NEMOTRON_V2_VL, "nemotron_v2_vl"}, - { PROJECTOR_TYPE_HUNYUANOCR, "hunyuanocr"}, - { PROJECTOR_TYPE_HUNYUANVL, "hunyuanvl"}, -+ { PROJECTOR_TYPE_PALIGEMMA, "paligemma"}, // vla.cpp - }; - - static projector_type clip_projector_type_from_string(const std::string & str) { -diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h -index bf8031b55..c39f749fb 100644 ---- a/tools/mtmd/clip-model.h -+++ b/tools/mtmd/clip-model.h -@@ -407,6 +407,8 @@ struct clip_model { - ggml_tensor * mm_input_proj_w = nullptr; - ggml_tensor * mm_soft_emb_norm_w = nullptr; - -+ ggml_tensor * mm_input_proj_b = nullptr; // paligemma (vla.cpp) -+ - // mobilenetv5 for gemma3n - std::vector mobilenet_blocks; - std::vector mobilenet_stage_ends; -diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp -index 45e39898d..4c582d52c 100644 ---- a/tools/mtmd/clip.cpp -+++ b/tools/mtmd/clip.cpp -@@ -162,6 +162,7 @@ struct clip_ctx { - - bool debug_output_embeddings = false; - -+ bool use_bf16_runtime = false; - clip_ctx(clip_context_params & ctx_params) { - flash_attn_type = ctx_params.flash_attn_type; - backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); -@@ -254,6 +255,13 @@ clip_graph::clip_graph(clip_ctx * ctx, const clip_image_f32 & img) : - ctx0_ptr.reset(ggml_init(params)); - ctx0 = ctx0_ptr.get(); - gf = ggml_new_graph_custom(ctx0, ctx->max_nodes, false); -+ use_bf16_runtime = ctx->use_bf16_runtime; -+ static bool logged_once = false; -+ if (!logged_once) { -+ LOG_INF("%s: vision-encoder activation precision = %s\n", -+ __func__, use_bf16_runtime ? "BF16 (HF-matched)" : "F32"); -+ logged_once = true; -+ } - } - - ggml_tensor * clip_graph::build_mm(ggml_tensor * w, ggml_tensor * x) const { -@@ -394,9 +402,15 @@ ggml_tensor * clip_graph::build_vit( - } - } - -- Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); -- Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos); -- Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos); -+ if (batch_size == 1) { -+ Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); -+ Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos); -+ Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos); -+ } else { -+ Qcur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, n_pos, batch_size); -+ Kcur = ggml_reshape_4d(ctx0, Kcur, d_head, n_head, n_pos, batch_size); -+ Vcur = ggml_reshape_4d(ctx0, Vcur, d_head, n_head, n_pos, batch_size); -+ } - - if (norm_per_head) { - if (layer.q_norm) { -@@ -505,7 +519,11 @@ ggml_tensor * clip_graph::build_vit( - ggml_tensor * clip_graph::build_inp() { - ggml_tensor * inp_raw = build_inp_raw(); - ggml_tensor * inp = ggml_conv_2d(ctx0, model.patch_embeddings_0, inp_raw, patch_size, patch_size, 0, 0, 1, 1); -- inp = ggml_reshape_2d(ctx0, inp, n_patches, n_embd); -+ if (batch_size == 1) { -+ inp = ggml_reshape_2d(ctx0, inp, n_patches, n_embd); -+ } else { -+ inp = ggml_reshape_3d(ctx0, inp, n_patches, n_embd, batch_size); -+ } - inp = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); - if (model.patch_bias) { - inp = ggml_add(ctx0, inp, model.patch_bias); -@@ -515,7 +533,9 @@ ggml_tensor * clip_graph::build_inp() { - } - - ggml_tensor * clip_graph::build_inp_raw(int channels) { -- ggml_tensor * inp_raw = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, img.nx, img.ny, channels); -+ ggml_tensor * inp_raw = (batch_size == 1) -+ ? ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, img.nx, img.ny, channels) -+ : ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, img.nx, img.ny, channels, batch_size); - ggml_set_name(inp_raw, "inp_raw"); - ggml_set_input(inp_raw); - return inp_raw; -@@ -666,7 +686,11 @@ ggml_tensor * clip_graph::build_attn( - cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, 0.0f, 0.0f); - ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32); - -- cur = ggml_reshape_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]); -+ if (batch_size == 1) { -+ cur = ggml_reshape_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]); -+ } else { -+ cur = ggml_reshape_3d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2], cur->ne[3]); -+ } - - } else { - ggml_tensor * v = ggml_permute(ctx0, v_cur, 1, 2, 0, 3); -@@ -680,7 +704,11 @@ ggml_tensor * clip_graph::build_attn( - - ggml_tensor * kqv = ggml_mul_mat(ctx0, v, kq); - cur = ggml_permute(ctx0, kqv, 0, 2, 1, 3); -- cur = ggml_cont_2d(ctx0, cur, cur->ne[0] * cur->ne[1], cur->ne[2] * cur->ne[3]); -+ if (batch_size == 1) { -+ cur = ggml_cont_2d(ctx0, cur, cur->ne[0] * cur->ne[1], cur->ne[2] * cur->ne[3]); -+ } else { -+ cur = ggml_cont_3d(ctx0, cur, cur->ne[0] * cur->ne[1], cur->ne[2], cur->ne[3]); -+ } - } - - cb(cur, "kqv_out", il); -@@ -804,35 +832,48 @@ ggml_tensor * clip_graph::build_patch_merge_permute(ggml_tensor * cur, int scale - // pad width and height to factor - const int64_t pad_width = CLIP_ALIGN(width, scale_factor) - width; - const int64_t pad_height = CLIP_ALIGN(height, scale_factor) - height; -- cur = ggml_reshape_3d(ctx0, cur, n_embd, width, height); -+ if (batch_size == 1) { -+ cur = ggml_reshape_3d(ctx0, cur, n_embd, width, height); -+ } else { -+ cur = ggml_reshape_4d(ctx0, cur, n_embd, width, height, batch_size); -+ } - if (pad_width || pad_height) { - cur = ggml_pad(ctx0, cur, 0, pad_width, pad_height, 0); - width += pad_width; - height += pad_height; - } - -- // unshuffle h -- cur = ggml_reshape_3d(ctx0, cur, n_embd * scale_factor, width / scale_factor, height); -- cur = ggml_permute(ctx0, cur, 0, 2, 1, 3); -+ if (batch_size == 1) { -+ cur = ggml_reshape_3d(ctx0, cur, n_embd * scale_factor, width / scale_factor, height); -+ cur = ggml_permute(ctx0, cur, 0, 2, 1, 3); -+ cur = ggml_cont_3d(ctx0, cur, n_embd * scale_factor * scale_factor, height / scale_factor, width / scale_factor); -+ cur = ggml_permute(ctx0, cur, 0, 2, 1, 3); - -- // unshuffle w -- cur = ggml_cont_3d(ctx0, cur, n_embd * scale_factor * scale_factor, height / scale_factor, width / scale_factor); -- cur = ggml_permute(ctx0, cur, 0, 2, 1, 3); -+ cur = ggml_cont_2d(ctx0, cur, cur->ne[0], cur->ne[1] * cur->ne[2]); -+ } else { -+ cur = ggml_reshape_4d(ctx0, cur, n_embd * scale_factor, width / scale_factor, height, batch_size); -+ cur = ggml_permute(ctx0, cur, 0, 2, 1, 3); -+ cur = ggml_cont_4d(ctx0, cur, n_embd * scale_factor * scale_factor, height / scale_factor, width / scale_factor, batch_size); -+ cur = ggml_permute(ctx0, cur, 0, 2, 1, 3); - -- cur = ggml_cont_2d(ctx0, cur, cur->ne[0], cur->ne[1] * cur->ne[2]); -+ cur = ggml_cont_3d(ctx0, cur, cur->ne[0], cur->ne[1] * cur->ne[2], cur->ne[3]); -+ } - cb(cur, "pixel_shuffle", -1); - - return cur; - } - - static ggml_cgraph * clip_image_build_graph(clip_ctx * ctx, const clip_image_f32_batch & imgs) { -- GGML_ASSERT(imgs.entries.size() == 1 && "n_batch > 1 is not supported"); -+ if (ctx->proj_type() != PROJECTOR_TYPE_IDEFICS3) { -+ GGML_ASSERT(imgs.entries.size() == 1 && "n_batch > 1 is not supported"); -+ } - - const clip_image_f32 & img = *imgs.entries[0]; - std::unique_ptr builder; - - switch (ctx->proj_type()) { - case PROJECTOR_TYPE_GEMMA3: -+ case PROJECTOR_TYPE_PALIGEMMA: // vla.cpp - case PROJECTOR_TYPE_IDEFICS3: - case PROJECTOR_TYPE_LFM2: - case PROJECTOR_TYPE_JANUS_PRO: -@@ -956,6 +997,7 @@ static ggml_cgraph * clip_image_build_graph(clip_ctx * ctx, const clip_image_f32 - GGML_ABORT("missing cgraph builder"); - } - -+ builder->batch_size = static_cast(imgs.entries.size()); - return builder->build(); - } - -@@ -1331,6 +1373,12 @@ struct clip_model_loader { - get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false); - } break; - -+ case PROJECTOR_TYPE_PALIGEMMA: // vla.cpp -+ { -+ hparams.n_merge = 1; -+ hparams.image_resize_algo = RESIZE_ALGO_BILINEAR; -+ } break; -+ - case PROJECTOR_TYPE_GEMMA4V: - { - hparams.rope_theta = 100.0f; -@@ -1932,6 +1980,11 @@ struct clip_model_loader { - model.mm_input_proj_w = get_tensor(TN_MM_INP_PROJ); - model.mm_soft_emb_norm_w = get_tensor(TN_MM_SOFT_EMB_N); - } break; -+ case PROJECTOR_TYPE_PALIGEMMA: // vla.cpp -+ { -+ model.mm_input_proj_w = get_tensor(TN_MM_INP_PROJ); -+ model.mm_input_proj_b = get_tensor(TN_MM_INP_PROJ_B, false); -+ } break; - case PROJECTOR_TYPE_GEMMA4V: - { - model.mm_input_proj_w = get_tensor(TN_MM_INP_PROJ); -@@ -2719,6 +2772,14 @@ struct clip_init_result clip_init(const char * fname, struct clip_context_params - ctx_vision = new clip_ctx(ctx_params); - loader.load_hparams(ctx_vision->model, CLIP_MODALITY_VISION); - loader.load_tensors(*ctx_vision); -+ ctx_vision->use_bf16_runtime = -+ ctx_vision->model.proj_type == PROJECTOR_TYPE_IDEFICS3 -+ && !ctx_vision->model.layers.empty() -+ && ctx_vision->model.layers[0].o_w -+ && ctx_vision->model.layers[0].o_w->type == GGML_TYPE_BF16; -+ if (const char * env = std::getenv("MTMD_VISION_BF16")) { -+ ctx_vision->use_bf16_runtime = !(env[0] == '0' && env[1] == '\0'); -+ } - if (ctx_params.warmup) { - loader.warmup(*ctx_vision); - } -@@ -2978,6 +3039,7 @@ int clip_n_output_tokens(const struct clip_ctx * ctx, struct clip_image_f32 * im - n_patches = x_patch * y_patch; - } break; - case PROJECTOR_TYPE_GEMMA3: -+ case PROJECTOR_TYPE_PALIGEMMA: // vla.cpp - case PROJECTOR_TYPE_GEMMA4V: - case PROJECTOR_TYPE_IDEFICS3: - case PROJECTOR_TYPE_INTERNVL: -@@ -3125,10 +3187,18 @@ bool clip_image_batch_encode(clip_ctx * ctx, const int n_threads, const clip_ima - const clip_image_f32_batch & imgs = *imgs_c_ptr; - int batch_size = imgs.entries.size(); - -- // TODO @ngxson : implement batch size > 1 as a loop -- // we don't need true batching support because the cgraph will gonna be big anyway -- if (batch_size != 1) { -- return false; // only support batch size of 1 -+ if (batch_size != 1 && ctx->proj_type() != PROJECTOR_TYPE_IDEFICS3) { -+ return false; -+ } -+ if (batch_size > 1) { -+ const int nx0 = imgs.entries[0]->nx; -+ const int ny0 = imgs.entries[0]->ny; -+ for (int b = 1; b < batch_size; ++b) { -+ if (imgs.entries[b]->nx != nx0 || imgs.entries[b]->ny != ny0) { -+ LOG_ERR("%s: batched encode requires identical (nx, ny) across images\n", __func__); -+ return false; -+ } -+ } - } - - // if buffers are not allocated, we need to do a warmup run to allocate them -@@ -3199,21 +3269,18 @@ bool clip_image_batch_encode(clip_ctx * ctx, const int n_threads, const clip_ima - // └─────┘ │ - // ──────┘ x B - -- for (size_t i = 0; i < imgs.entries.size(); i++) { -- const int nx = imgs.entries[i]->nx; -- const int ny = imgs.entries[i]->ny; -- const int n = nx * ny; -- -- for (int b = 0; b < batch_size; b++) { -- float * batch_entry = inp_raw.data() + b * (3*n); -- for (int y = 0; y < ny; y++) { -- for (int x = 0; x < nx; x++) { -- size_t base_src = 3*(y * nx + x); // idx of the first channel -- size_t base_dst = y * nx + x; // idx of the first channel -- batch_entry[ base_dst] = imgs.entries[b]->buf[base_src ]; -- batch_entry[1*n + base_dst] = imgs.entries[b]->buf[base_src + 1]; -- batch_entry[2*n + base_dst] = imgs.entries[b]->buf[base_src + 2]; -- } -+ for (int b = 0; b < batch_size; ++b) { -+ const int nx = imgs.entries[b]->nx; -+ const int ny = imgs.entries[b]->ny; -+ const int n = nx * ny; -+ float * dst = inp_raw.data() + (size_t)b * 3 * n; -+ for (int y = 0; y < ny; y++) { -+ for (int x = 0; x < nx; x++) { -+ const size_t base_src = 3 * (y * nx + x); -+ const size_t base_dst = y * nx + x; -+ dst[ base_dst] = imgs.entries[b]->buf[base_src ]; -+ dst[1*n + base_dst] = imgs.entries[b]->buf[base_src + 1]; -+ dst[2*n + base_dst] = imgs.entries[b]->buf[base_src + 2]; - } - } - } -@@ -3276,6 +3343,29 @@ bool clip_image_batch_encode(clip_ctx * ctx, const int n_threads, const clip_ima - } - set_input_f32("omega", omega); - } break; -+ case PROJECTOR_TYPE_IDEFICS3: -+ { -+ const int n = ctx->model.hparams.image_size / patch_size; -+ std::vector positions(pos_h * pos_w); -+ if (ctx->use_bf16_runtime) { -+ for (int i = 0, id = 0; i < pos_h; i++) { -+ for (int j = 0; j < pos_w; j++) { -+ positions[id++] = i * n + j; -+ } -+ } -+ } else { -+ std::vector bucket(n); -+ for (int i = 0; i < n; i++) { -+ bucket[i] = i == 0 ? 0 : i - 1; -+ } -+ for (int i = 0, id = 0; i < pos_h; i++) { -+ for (int j = 0; j < pos_w; j++) { -+ positions[id++] = bucket[i] * n + bucket[j]; -+ } -+ } -+ } -+ set_input_i32("positions", positions); -+ } break; - case PROJECTOR_TYPE_QWEN2VL: - case PROJECTOR_TYPE_QWEN3VL: - case PROJECTOR_TYPE_GLM4V: -@@ -3529,8 +3619,8 @@ bool clip_image_batch_encode(clip_ctx * ctx, const int n_threads, const clip_ima - set_input_i32("rel_pos_indices_global", rel_pos_indices_global); - } break; - case PROJECTOR_TYPE_GEMMA3: -+ case PROJECTOR_TYPE_PALIGEMMA: // vla.cpp - case PROJECTOR_TYPE_GEMMA3NV: -- case PROJECTOR_TYPE_IDEFICS3: - case PROJECTOR_TYPE_INTERNVL: - case PROJECTOR_TYPE_NEMOTRON_V2_VL: - case PROJECTOR_TYPE_QWEN2A: -@@ -3812,6 +3902,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { - case PROJECTOR_TYPE_GEMMA3: - case PROJECTOR_TYPE_GEMMA3NV: - return ctx->model.mm_input_proj_w->ne[0]; -+ case PROJECTOR_TYPE_PALIGEMMA: // vla.cpp: ne(vision, text) → output width = ne[1] -+ return ctx->model.mm_input_proj_w->ne[1]; - case PROJECTOR_TYPE_GEMMA4V: - return ctx->model.mm_input_proj_w->ne[1]; - case PROJECTOR_TYPE_IDEFICS3: -@@ -3909,6 +4001,26 @@ bool clip_encode_float_image (struct clip_ctx * ctx, int n_threads, float * img, - return true; - } - -+bool clip_encode_float_images(struct clip_ctx * ctx, int n_threads, -+ const float * img, int h, int w, int n_images, -+ float * vec) { -+ if (n_images < 1) { -+ return false; -+ } -+ if (n_images == 1) { -+ return clip_encode_float_image(ctx, n_threads, const_cast(img), h, w, vec); -+ } -+ clip_image_f32_batch batch; -+ const size_t per_image = (size_t)h * (size_t)w * 3; -+ for (int b = 0; b < n_images; ++b) { -+ clip_image_f32_ptr p(clip_image_f32_init()); -+ p->nx = w; -+ p->ny = h; -+ p->buf.assign(img + b * per_image, img + (b + 1) * per_image); -+ batch.entries.push_back(std::move(p)); -+ } -+ return clip_image_batch_encode(ctx, n_threads, &batch, vec); -+} - // - // API used internally with mtmd - // -diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h -index a859b3865..dbe3a438a 100644 ---- a/tools/mtmd/clip.h -+++ b/tools/mtmd/clip.h -@@ -110,6 +110,9 @@ bool clip_is_llava(const struct clip_ctx * ctx); - - bool clip_encode_float_image (struct clip_ctx * ctx, int n_threads, float * img, int h, int w, float * vec); - -+bool clip_encode_float_images(struct clip_ctx * ctx, int n_threads, -+ const float * img, int h, int w, int n_images, -+ float * vec); - // use by audio input - void clip_image_f32_batch_add_mel(struct clip_image_f32_batch * batch, int n_mel, int n_frames, float * mel); - -diff --git a/tools/mtmd/models/siglip.cpp b/tools/mtmd/models/siglip.cpp -index 7ef98eed0..890314699 100644 ---- a/tools/mtmd/models/siglip.cpp -+++ b/tools/mtmd/models/siglip.cpp -@@ -7,6 +7,12 @@ ggml_cgraph * clip_graph_siglip::build() { - if (proj_type == PROJECTOR_TYPE_LFM2 || proj_type == PROJECTOR_TYPE_PHI4) { - learned_pos_embd = resize_position_embeddings(); - } -+ if (proj_type == PROJECTOR_TYPE_IDEFICS3) { -+ ggml_tensor * positions = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_patches); -+ ggml_set_name(positions, "positions"); -+ ggml_set_input(positions); -+ learned_pos_embd = ggml_get_rows(ctx0, model.position_embeddings, positions); -+ } - - ggml_tensor * cur = build_vit( - inp, n_patches, -@@ -38,6 +44,14 @@ ggml_cgraph * clip_graph_siglip::build() { - ggml_cont(ctx0, ggml_transpose(ctx0, model.mm_input_proj_w)), - cur); - -+ } else if (proj_type == PROJECTOR_TYPE_PALIGEMMA) { -+ // vla.cpp: π0's PaliGemma projector - see patches/PATCH.md. -+ cur = ggml_mul_mat(ctx0, model.mm_input_proj_w, cur); -+ if (model.mm_input_proj_b) { -+ cur = ggml_add(ctx0, cur, model.mm_input_proj_b); -+ } -+ cur = ggml_scale(ctx0, cur, 1.0f / sqrtf((float) cur->ne[0])); -+ - } else if (proj_type == PROJECTOR_TYPE_IDEFICS3) { - // pixel_shuffle - // https://github.com/huggingface/transformers/blob/0a950e0bbe1ed58d5401a6b547af19f15f0c195e/src/transformers/models/idefics3/modeling_idefics3.py#L578 diff --git a/patches/patch.sh b/patches/patch.sh deleted file mode 100755 index 83db471..0000000 --- a/patches/patch.sh +++ /dev/null @@ -1,54 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2026 VinRobotics -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# patches/patch.sh - fetch the pinned llama.cpp and apply vla.cpp's local patch. -# -# vla.cpp builds against llama.cpp at tag b9016 with a small set of local fixes -# shipped as patches/llama.cpp-vla.patch. Run this once from the repo root after -# cloning, before configuring CMake: -# -# bash ./patches/patch.sh -# -# Re-running is safe: an existing checkout is left in place and an -# already-applied patch is detected and skipped. - -set -euo pipefail - -# Repo root is one level up from this script (which lives in patches/). -ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -LLAMA_DIR="$ROOT/third_party/llama.cpp" -LLAMA_URL="https://github.com/ggml-org/llama.cpp.git" -LLAMA_TAG="b9016" -PATCH="$ROOT/patches/llama.cpp-vla.patch" - -[[ -f "$PATCH" ]] || { echo "patch not found: $PATCH" >&2; exit 1; } - -if [[ -e "$LLAMA_DIR/.git" ]]; then - echo ">> third_party/llama.cpp already present - skipping clone" -else - echo ">> cloning llama.cpp @ $LLAMA_TAG" - git clone "$LLAMA_URL" "$LLAMA_DIR" - git -C "$LLAMA_DIR" checkout "$LLAMA_TAG" -fi - -if git -C "$LLAMA_DIR" apply --reverse --check "$PATCH" 2>/dev/null; then - echo ">> patch already applied - nothing to do" -elif git -C "$LLAMA_DIR" apply --check "$PATCH" 2>/dev/null; then - git -C "$LLAMA_DIR" apply "$PATCH" - echo ">> patched - third_party/llama.cpp is ready to build" -else - echo "ERROR: patch does not apply cleanly to $LLAMA_DIR (expected tag $LLAMA_TAG)" >&2 - exit 1 -fi diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..85ec419 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,32 @@ +[project] +name = "vla-cpp-tooling" +version = "0.1.1" +description = "Python tooling for vla.cpp: HuggingFace -> GGUF converters and the ZeroMQ eval client." +readme = "README.md" +requires-python = ">=3.10" +license = { text = "Apache-2.0" } +dependencies = ["numpy>=1.24"] + +[project.optional-dependencies] +# scripts/convert_*_to_gguf.py, scripts/merge_*_mmproj_to_gguf.py, scripts/bitvla_int2_pack.py +convert = [ + "torch>=2.5", + "safetensors>=0.4", + "gguf>=0.10", + "numpy>=1.24", +] +# eval/client/* talks to vla-server over ZeroMQ. The simulators (LIBERO, SimplerEnv, +# ALOHA) keep their own pinned venvs - see eval/sim/*/setup_*.sh - and are not packaged here. +client = [ + "pyzmq>=25", + "msgpack>=1.1", + "msgpack-numpy>=0.4.8", + "pillow>=10", + "torch>=2.5", + "transformers>=4.51,<4.52", + "numpy>=1.24", +] + +# Deps-only project: the code lives in scripts/ and eval/, not an importable package. +[tool.uv] +package = false diff --git a/scripts/convert_bitvla_to_gguf.py b/scripts/convert_bitvla_to_gguf.py index f830a3e..ef14f7c 100644 --- a/scripts/convert_bitvla_to_gguf.py +++ b/scripts/convert_bitvla_to_gguf.py @@ -13,10 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. + from __future__ import annotations import argparse import json +import re import sys from pathlib import Path @@ -70,9 +72,7 @@ def weight_quant(W: torch.Tensor) -> torch.Tensor: return Wq.to(dtype) def _bf16_u16(t: torch.Tensor) -> np.ndarray: - if t.dtype != torch.bfloat16: - print(f" warn: casting non-BF16 tensor (dtype={t.dtype}) to BF16 for storage") - t = t.to(torch.bfloat16) + assert t.dtype == torch.bfloat16, t.dtype return t.contiguous().view(torch.uint16).cpu().numpy() def _add(writer: gguf.GGUFWriter, name: str, t: torch.Tensor) -> None: @@ -165,10 +165,18 @@ def main() -> int: W = _load_sharded(ckpt) print(f" {len(W)} main tensors") - ah_path = ckpt / "action_head--10000_checkpoint.pt" - pp_path = ckpt / "proprio_projector--10000_checkpoint.pt" - if not ah_path.exists() or not pp_path.exists(): + def _find_sidecar(stem: str) -> Path | None: + cands = sorted( + ckpt.glob(f"{stem}--*_checkpoint.pt"), + key=lambda p: int(m.group(1)) if (m := re.search(r"--(\d+)_checkpoint\.pt$", p.name)) else -1, + ) + return cands[-1] if cands else None + + ah_path = _find_sidecar("action_head") + pp_path = _find_sidecar("proprio_projector") + if ah_path is None or pp_path is None: raise SystemExit(f"missing action_head/proprio_projector sidecars in {ckpt}") + print(f" sidecars: {ah_path.name}, {pp_path.name}") AH = _load_pt_sidecar(ah_path) PP = _load_pt_sidecar(pp_path) print(f" +{len(AH)} action_head tensors, +{len(PP)} proprio_projector tensors") diff --git a/scripts/convert_pi05_to_gguf.py b/scripts/convert_pi05_to_gguf.py new file mode 100644 index 0000000..ff87a34 --- /dev/null +++ b/scripts/convert_pi05_to_gguf.py @@ -0,0 +1,384 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Optional + +import numpy as np +import torch +from safetensors import safe_open + +import gguf + +ARCH = "pi05" +KV = lambda name: f"{ARCH}.{name}" + +PFX_VLM = "paligemma_with_expert.paligemma.model.language_model" +PFX_VLM_HEAD = "paligemma_with_expert.paligemma.lm_head.weight" +PFX_AEX = "paligemma_with_expert.gemma_expert.model" + +PFX_VIS_CANDIDATES = [ + "paligemma_with_expert.paligemma.model.vision_tower.vision_model", + "paligemma_with_expert.paligemma.vision_tower.vision_model", +] +PFX_MMP_CANDIDATES = [ + "paligemma_with_expert.paligemma.model.multi_modal_projector", + "paligemma_with_expert.paligemma.multi_modal_projector", +] + +GEMMA_2B = dict(hidden=2048, n_q_heads=8, n_kv_heads=1, head_dim=256, intermediate=16384) +GEMMA_300M = dict(expert_h=1024, expert_inter=4096) + +ROPE_THETA = 10000.0 +RMS_NORM_EPS = 1e-6 + +def _bf16_to_u16_bytes(t: torch.Tensor) -> np.ndarray: + assert t.dtype == torch.bfloat16, t.dtype + return t.view(torch.uint16).contiguous().cpu().numpy() + +def _f32_np(t: torch.Tensor) -> np.ndarray: + return t.float().contiguous().cpu().numpy() + +def _add_one_tensor(writer: gguf.GGUFWriter, dst_name: str, t: torch.Tensor) -> None: + if t.dtype == torch.float32: + writer.add_tensor(dst_name, _f32_np(t), raw_dtype=gguf.GGMLQuantizationType.F32) + elif t.dtype == torch.bfloat16: + writer.add_tensor(dst_name, _bf16_to_u16_bytes(t), + raw_shape=list(t.shape), raw_dtype=gguf.GGMLQuantizationType.BF16) + else: + raise NotImplementedError(f"unsupported dtype {t.dtype} for {dst_name}") + +def _stream_vlm_block(writer, sf, n_layers: int) -> None: + + suffix_map = [ + ("input_layernorm.weight", "attn_norm.weight"), + ("self_attn.q_proj.weight", "attn_q.weight"), + ("self_attn.k_proj.weight", "attn_k.weight"), + ("self_attn.v_proj.weight", "attn_v.weight"), + ("self_attn.o_proj.weight", "attn_o.weight"), + ("post_attention_layernorm.weight", "ffn_norm.weight"), + ("mlp.gate_proj.weight", "ffn_gate.weight"), + ("mlp.up_proj.weight", "ffn_up.weight"), + ("mlp.down_proj.weight", "ffn_down.weight"), + ] + for i in range(n_layers): + for src_suf, dst_suf in suffix_map: + t = sf.get_tensor(f"{PFX_VLM}.layers.{i}.{src_suf}") + _add_one_tensor(writer, f"vlm.blk.{i}.{dst_suf}", t) + +def _stream_aex_block(writer, sf, n_layers: int) -> None: + + norm_map = [ + ("input_layernorm.dense.weight", "attn_norm.weight"), + ("input_layernorm.dense.bias", "attn_norm.bias"), + ("post_attention_layernorm.dense.weight", "ffn_norm.weight"), + ("post_attention_layernorm.dense.bias", "ffn_norm.bias"), + ] + mat_map = [ + ("self_attn.q_proj.weight", "attn_q.weight"), + ("self_attn.k_proj.weight", "attn_k.weight"), + ("self_attn.v_proj.weight", "attn_v.weight"), + ("self_attn.o_proj.weight", "attn_o.weight"), + ("mlp.gate_proj.weight", "ffn_gate.weight"), + ("mlp.up_proj.weight", "ffn_up.weight"), + ("mlp.down_proj.weight", "ffn_down.weight"), + ] + for i in range(n_layers): + for src_suf, dst_suf in norm_map + mat_map: + t = sf.get_tensor(f"{PFX_AEX}.layers.{i}.{src_suf}") + _add_one_tensor(writer, f"aex.blk.{i}.{dst_suf}", t) + +def _probe_vision(sf, keys, cfg_json) -> dict: + vis = next((p for p in PFX_VIS_CANDIDATES + if f"{p}.embeddings.patch_embedding.weight" in keys), None) + mmp = next((p for p in PFX_MMP_CANDIDATES if f"{p}.linear.weight" in keys), None) + if vis is None or mmp is None: + raise SystemExit("vision_tower / multi_modal_projector not found in checkpoint; " + f"tried {PFX_VIS_CANDIDATES} and {PFX_MMP_CANDIDATES}") + conv = sf.get_slice(f"{vis}.embeddings.patch_embedding.weight").get_shape() + pos = sf.get_slice(f"{vis}.embeddings.position_embedding.weight").get_shape() + n_vit = 0 + while f"{vis}.encoder.layers.{n_vit}.layer_norm1.weight" in keys: + n_vit += 1 + vcfg = cfg_json.get("vision_config") or {} + patch = int(conv[2]); ntok = int(pos[0]); grid = int(round(ntok ** 0.5)) + return dict(prefix=vis, mmp=mmp, vit_hidden=int(conv[0]), vit_layers=n_vit, + vit_heads=int(vcfg.get("num_attention_heads", 16)), + image_size=grid * patch, patch_size=patch, n_img_tokens=ntok, + vit_ln_eps=float(vcfg.get("layer_norm_eps", 1e-6))) + +def _add_vision_tensors(writer: gguf.GGUFWriter, sf, v: dict) -> None: + vis, mmp = v["prefix"], v["mmp"] + af32 = lambda dst, src: _add_one_tensor(writer, dst, sf.get_tensor(src).float()) + akeep = lambda dst, src: _add_one_tensor(writer, dst, sf.get_tensor(src)) + af32("vit.patch_embd.weight", f"{vis}.embeddings.patch_embedding.weight") + af32("vit.patch_embd.bias", f"{vis}.embeddings.patch_embedding.bias") + af32("vit.pos_embd", f"{vis}.embeddings.position_embedding.weight") + for i in range(v["vit_layers"]): + L = f"{vis}.encoder.layers.{i}." + af32(f"vit.blk.{i}.ln1.weight", L + "layer_norm1.weight"); af32(f"vit.blk.{i}.ln1.bias", L + "layer_norm1.bias") + af32(f"vit.blk.{i}.ln2.weight", L + "layer_norm2.weight"); af32(f"vit.blk.{i}.ln2.bias", L + "layer_norm2.bias") + for q in ("q", "k", "v"): + akeep(f"vit.blk.{i}.attn_{q}.weight", L + f"self_attn.{q}_proj.weight") + af32 (f"vit.blk.{i}.attn_{q}.bias", L + f"self_attn.{q}_proj.bias") + akeep(f"vit.blk.{i}.attn_o.weight", L + "self_attn.out_proj.weight"); af32(f"vit.blk.{i}.attn_o.bias", L + "self_attn.out_proj.bias") + akeep(f"vit.blk.{i}.fc1.weight", L + "mlp.fc1.weight"); af32(f"vit.blk.{i}.fc1.bias", L + "mlp.fc1.bias") + akeep(f"vit.blk.{i}.fc2.weight", L + "mlp.fc2.weight"); af32(f"vit.blk.{i}.fc2.bias", L + "mlp.fc2.bias") + af32("vit.post_ln.weight", f"{vis}.post_layernorm.weight"); af32("vit.post_ln.bias", f"{vis}.post_layernorm.bias") + akeep("mm.proj.weight", f"{mmp}.linear.weight") + af32 ("mm.proj.bias", f"{mmp}.linear.bias") + +def _read_norm_eps(meta_path: Path) -> Optional[float]: + if not meta_path.exists(): + return None + meta = json.loads(meta_path.read_text()) + for step in meta.get("steps", []): + if step.get("registry_name") in ("normalizer_processor", "unnormalizer_processor"): + cfg = step.get("config", {}) + if "eps" in cfg: + return float(cfg["eps"]) + return None + +def _load_dataset_stats(stats_json: Optional[Path], dataset_repo: Optional[str], + real_state_dim: int, real_action_dim: int) -> dict[str, np.ndarray]: + + def _from_json(path: Path) -> dict[str, np.ndarray]: + d = json.loads(path.read_text()) + + def grab(feat, dim): + s = d[feat] + mean = np.asarray(s["mean"], dtype=np.float32).reshape(-1) + std = np.asarray(s["std"], dtype=np.float32).reshape(-1) + if mean.size < dim or std.size < dim: + raise SystemExit(f"stats {feat} dim {mean.size} < expected {dim}") + out = {"mean": mean[:dim], "std": std[:dim]} + + for q in ("q01", "q99"): + if q not in s: + raise SystemExit( + f"stats {feat} missing {q} (QUANTILES normalization needs it). " + f"Use a meta/stats.json with quantile stats.") + out[q] = np.asarray(s[q], dtype=np.float32).reshape(-1)[:dim] + return out + st = grab("observation.state", real_state_dim) + ac = grab("action", real_action_dim) + return {"state_mean": st["mean"], "state_std": st["std"], + "state_q01": st["q01"], "state_q99": st["q99"], + "action_mean": ac["mean"], "action_std": ac["std"], + "action_q01": ac["q01"], "action_q99": ac["q99"]} + + if stats_json is not None: + print(f" stats: loading from {stats_json}") + return _from_json(stats_json) + + if dataset_repo is not None: + from huggingface_hub import hf_hub_download + print(f" stats: fetching meta/stats.json from dataset repo {dataset_repo}") + p = hf_hub_download(repo_id=dataset_repo, filename="meta/stats.json", repo_type="dataset") + return _from_json(Path(p)) + + raise SystemExit( + "no stats source: π0.5 checkpoints carry no normalizer stats — pass " + "--dataset-stats or --dataset-repo lerobot/libero. " + "Refusing to bake identity stats (would make the policy miss; skill footgun #1).") + +def _add_kv(writer: gguf.GGUFWriter, cfg: dict) -> None: + writer.add_string (KV("architecture"), ARCH) + writer.add_string (KV("paligemma_variant"), cfg["paligemma_variant"]) + writer.add_string (KV("action_expert_variant"), cfg["action_expert_variant"]) + writer.add_uint32 (KV("hidden"), cfg["hidden"]) + writer.add_uint32 (KV("intermediate"), cfg["intermediate"]) + writer.add_uint32 (KV("n_q_heads"), cfg["n_q_heads"]) + writer.add_uint32 (KV("n_kv_heads"), cfg["n_kv_heads"]) + writer.add_uint32 (KV("head_dim"), cfg["head_dim"]) + writer.add_uint32 (KV("n_layers"), cfg["n_layers"]) + writer.add_uint32 (KV("vocab_size"), cfg["vocab_size"]) + writer.add_uint32 (KV("expert_h"), cfg["expert_h"]) + writer.add_uint32 (KV("expert_inter"), cfg["expert_inter"]) + writer.add_uint32 (KV("chunk_size"), cfg["chunk_size"]) + writer.add_uint32 (KV("num_steps"), cfg["num_steps"]) + writer.add_uint32 (KV("n_action_steps"), cfg["n_action_steps"]) + writer.add_uint32 (KV("max_state_dim"), cfg["max_state_dim"]) + writer.add_uint32 (KV("max_action_dim"), cfg["max_action_dim"]) + writer.add_uint32 (KV("real_state_dim"), cfg["real_state_dim"]) + writer.add_uint32 (KV("real_action_dim"), cfg["real_action_dim"]) + writer.add_uint32 (KV("tokenizer_max_length"), cfg["tokenizer_max_length"]) + writer.add_bool (KV("use_adarms_expert"), True) + writer.add_uint32 (KV("adarms_cond_dim"), cfg["expert_h"]) + writer.add_string (KV("norm_mode"), "quantiles") + writer.add_float64 (KV("min_period"), cfg["min_period"]) + writer.add_float64 (KV("max_period"), cfg["max_period"]) + writer.add_float64 (KV("rope_theta"), cfg["rope_theta"]) + writer.add_float32 (KV("rms_norm_eps"), cfg["rms_norm_eps"]) + writer.add_float32 (KV("norm_eps"), cfg["norm_eps"]) + if "vit" in cfg: + v = cfg["vit"] + writer.add_uint32 (KV("vit_hidden"), v["vit_hidden"]) + writer.add_uint32 (KV("vit_layers"), v["vit_layers"]) + writer.add_uint32 (KV("vit_heads"), v["vit_heads"]) + writer.add_uint32 (KV("image_size"), v["image_size"]) + writer.add_uint32 (KV("patch_size"), v["patch_size"]) + writer.add_uint32 (KV("n_img_tokens"), v["n_img_tokens"]) + writer.add_float32(KV("vit_ln_eps"), v["vit_ln_eps"]) + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--ckpt", type=Path, required=True, + help="lerobot π0.5 checkpoint dir (model.safetensors + config.json + policy_*processor.json)") + ap.add_argument("--out", type=Path, default=None, help="Output GGUF (default /pi05.gguf)") + ap.add_argument("--dataset-stats", type=Path, default=None, + help="Path to a LIBERO dataset meta/stats.json for MEAN_STD norm stats") + ap.add_argument("--dataset-repo", type=str, default=None, + help="HF dataset repo to fetch meta/stats.json from (e.g. lerobot/libero)") + args = ap.parse_args() + + ckpt = args.ckpt.resolve() + out = (args.out or ckpt / "pi05.gguf").resolve() + sf_path = ckpt / "model.safetensors" + cfg_path = ckpt / "config.json" + if not sf_path.exists(): raise SystemExit(f"missing {sf_path}") + if not cfg_path.exists(): raise SystemExit(f"missing {cfg_path}") + + cfg_json = json.loads(cfg_path.read_text()) + if cfg_json.get("type") != "pi05": + raise SystemExit(f"config.json type is {cfg_json.get('type')!r}, expected 'pi05'") + + cfg = dict(GEMMA_2B, **GEMMA_300M) + cfg["paligemma_variant"] = str(cfg_json.get("paligemma_variant", "gemma_2b")) + cfg["action_expert_variant"] = str(cfg_json.get("action_expert_variant", "gemma_300m")) + cfg["chunk_size"] = int(cfg_json["chunk_size"]) + cfg["num_steps"] = int(cfg_json["num_inference_steps"]) + cfg["n_action_steps"] = int(cfg_json["n_action_steps"]) + cfg["max_state_dim"] = int(cfg_json["max_state_dim"]) + cfg["max_action_dim"] = int(cfg_json["max_action_dim"]) + cfg["min_period"] = float(cfg_json["min_period"]) + cfg["max_period"] = float(cfg_json["max_period"]) + cfg["tokenizer_max_length"] = int(cfg_json["tokenizer_max_length"]) + cfg["real_state_dim"] = int(cfg_json["input_features"]["observation.state"]["shape"][0]) + cfg["real_action_dim"] = int(cfg_json["output_features"]["action"]["shape"][0]) + cfg["rope_theta"] = ROPE_THETA + cfg["rms_norm_eps"] = RMS_NORM_EPS + + norm_eps = _read_norm_eps(ckpt / "policy_preprocessor.json") + if norm_eps is None: + norm_eps = _read_norm_eps(ckpt / "policy_postprocessor.json") + cfg["norm_eps"] = float(norm_eps if norm_eps is not None else 1e-8) + + print(f"opening {sf_path}") + sf = safe_open(sf_path, framework="pt") + keys = set(sf.keys()) + + def _maxlayer(pfx: str) -> int: + m = -1 + for k in keys: + if k.startswith(pfx): + try: m = max(m, int(k[len(pfx):].split(".", 1)[0])) + except ValueError: pass + return m + 1 + + n_layers_vlm = _maxlayer(f"{PFX_VLM}.layers.") + n_layers_aex = _maxlayer(f"{PFX_AEX}.layers.") + if n_layers_vlm <= 0: + raise SystemExit("cannot find PaliGemma language-model layers in checkpoint") + if n_layers_aex != n_layers_vlm: + raise SystemExit(f"layer count mismatch: VLM={n_layers_vlm} expert={n_layers_aex}") + cfg["n_layers"] = n_layers_vlm + + q0 = sf.get_slice(f"{PFX_VLM}.layers.0.self_attn.q_proj.weight").get_shape() + kv0 = sf.get_slice(f"{PFX_VLM}.layers.0.self_attn.k_proj.weight").get_shape() + gate0 = sf.get_slice(f"{PFX_VLM}.layers.0.mlp.gate_proj.weight").get_shape() + if q0[1] != cfg["hidden"]: + raise SystemExit(f"hidden mismatch: cfg={cfg['hidden']} ckpt={q0[1]}") + if q0[0] != cfg["n_q_heads"] * cfg["head_dim"]: + raise SystemExit(f"q_proj rows {q0[0]} != n_q_heads*head_dim") + if kv0[0] != cfg["n_kv_heads"] * cfg["head_dim"]: + raise SystemExit(f"k_proj rows {kv0[0]} != n_kv_heads*head_dim") + if gate0[0] != cfg["intermediate"]: + raise SystemExit(f"intermediate mismatch: cfg={cfg['intermediate']} ckpt={gate0[0]}") + + ada0 = sf.get_slice(f"{PFX_AEX}.layers.0.input_layernorm.dense.weight").get_shape() + if ada0 != [3 * cfg["expert_h"], cfg["expert_h"]]: + raise SystemExit(f"expert adaRMS dense shape {ada0} != [3*expert_h, expert_h] " + f"{[3*cfg['expert_h'], cfg['expert_h']]}") + aex_o0 = sf.get_slice(f"{PFX_AEX}.layers.0.self_attn.o_proj.weight").get_shape() + if aex_o0 != [cfg["expert_h"], cfg["n_q_heads"] * cfg["head_dim"]]: + raise SystemExit(f"expert o_proj shape {aex_o0} unexpected") + + head_w = sf.get_slice(PFX_VLM_HEAD).get_shape() + cfg["vocab_size"] = int(head_w[0]) + + print(f"resolved cfg: hidden={cfg['hidden']} n_layers={cfg['n_layers']} " + f"expert_h={cfg['expert_h']} vocab={cfg['vocab_size']} chunk={cfg['chunk_size']} " + f"steps={cfg['num_steps']} real_state={cfg['real_state_dim']} " + f"real_action={cfg['real_action_dim']} max_len={cfg['tokenizer_max_length']} " + f"norm_eps={cfg['norm_eps']:g}") + + vit = _probe_vision(sf, keys, cfg_json) + cfg["vit"] = vit + print(f"vision: SigLIP hidden={vit['vit_hidden']} layers={vit['vit_layers']} " + f"heads={vit['vit_heads']} image={vit['image_size']} patch={vit['patch_size']} " + f"tokens={vit['n_img_tokens']} ln_eps={vit['vit_ln_eps']:g}") + + print("loading dataset normalizer stats...") + stats = _load_dataset_stats(args.dataset_stats, args.dataset_repo, + cfg["real_state_dim"], cfg["real_action_dim"]) + print(f" state_q01[:3]={stats['state_q01'][:3]} state_q99[:3]={stats['state_q99'][:3]}") + print(f" action_q01[:3]={stats['action_q01'][:3]} action_q99[:3]={stats['action_q99'][:3]} (QUANTILES)") + + out.parent.mkdir(parents=True, exist_ok=True) + print(f"writing {out}") + writer = gguf.GGUFWriter(str(out), arch=ARCH) + _add_kv(writer, cfg) + + _add_one_tensor(writer, "token_embd.weight", sf.get_tensor(PFX_VLM_HEAD)) + _add_one_tensor(writer, "vlm.output_norm.weight", sf.get_tensor(f"{PFX_VLM}.norm.weight")) + _stream_vlm_block(writer, sf, cfg["n_layers"]) + + _add_one_tensor(writer, "aex.output_norm.weight", sf.get_tensor(f"{PFX_AEX}.norm.dense.weight")) + _add_one_tensor(writer, "aex.output_norm.bias", sf.get_tensor(f"{PFX_AEX}.norm.dense.bias")) + _stream_aex_block(writer, sf, cfg["n_layers"]) + + for suf in [ + "action_in_proj.weight", "action_in_proj.bias", + "time_mlp_in.weight", "time_mlp_in.bias", + "time_mlp_out.weight", "time_mlp_out.bias", + "action_out_proj.weight", "action_out_proj.bias", + ]: + _add_one_tensor(writer, suf, sf.get_tensor(suf)) + + _add_vision_tensors(writer, sf, vit) + + for k, v in stats.items(): + writer.add_tensor(k, v.astype(np.float32, copy=False), + raw_dtype=gguf.GGMLQuantizationType.F32) + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + print(f"done. {out} ({out.stat().st_size / (1024*1024):.1f} MiB)") + print("note: self-contained GGUF — SigLIP vision tower + PaliGemma projector are baked in; " + "no separate mmproj is needed.") + return 0 + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/convert_pi0_to_gguf.py b/scripts/convert_pi0_to_gguf.py index 15b1fe7..0546fe6 100644 --- a/scripts/convert_pi0_to_gguf.py +++ b/scripts/convert_pi0_to_gguf.py @@ -34,6 +34,15 @@ PFX_AEX = "model.paligemma_with_expert.gemma_expert.model" PFX_PROJ = "model" +PFX_VIS_CANDIDATES = [ + "model.paligemma_with_expert.paligemma.model.vision_tower.vision_model", + "model.paligemma_with_expert.paligemma.vision_tower.vision_model", +] +PFX_MMP_CANDIDATES = [ + "model.paligemma_with_expert.paligemma.model.multi_modal_projector", + "model.paligemma_with_expert.paligemma.multi_modal_projector", +] + GEMMA_2B = dict(hidden=2048, n_q_heads=8, n_kv_heads=1, head_dim=256, intermediate=16384) GEMMA_300M = dict(expert_h=1024, expert_inter=4096) @@ -80,6 +89,48 @@ def _stream_block(writer: gguf.GGUFWriter, sf, src_pfx: str, dst_pfx: str, n_lay t = sf.get_tensor(f"{src_pfx}.layers.{i}.{src_suf}") _add_one_tensor(writer, f"{dst_pfx}.blk.{i}.{dst_suf}", t) +def _probe_vision(sf, keys, cfg_json) -> Optional[dict]: + vis = next((p for p in PFX_VIS_CANDIDATES + if f"{p}.embeddings.patch_embedding.weight" in keys), None) + mmp = next((p for p in PFX_MMP_CANDIDATES if f"{p}.linear.weight" in keys), None) + if vis is None or mmp is None: + raise SystemExit("vision_tower / multi_modal_projector not found in checkpoint; " + f"tried {PFX_VIS_CANDIDATES} and {PFX_MMP_CANDIDATES}") + conv = sf.get_slice(f"{vis}.embeddings.patch_embedding.weight").get_shape() + pos = sf.get_slice(f"{vis}.embeddings.position_embedding.weight").get_shape() + n_vit = 0 + while f"{vis}.encoder.layers.{n_vit}.layer_norm1.weight" in keys: + n_vit += 1 + vcfg = cfg_json.get("vision_config") or {} + patch = int(conv[2]); ntok = int(pos[0]); grid = int(round(ntok ** 0.5)) + return dict(prefix=vis, mmp=mmp, vit_hidden=int(conv[0]), vit_layers=n_vit, + vit_heads=int(vcfg.get("num_attention_heads", 16)), + image_size=grid * patch, patch_size=patch, n_img_tokens=ntok, + vit_ln_eps=float(vcfg.get("layer_norm_eps", 1e-6))) + + +def _add_vision_tensors(writer: gguf.GGUFWriter, sf, v: dict) -> None: + vis, mmp = v["prefix"], v["mmp"] + af32 = lambda dst, src: _add_one_tensor(writer, dst, sf.get_tensor(src).float()) + akeep = lambda dst, src: _add_one_tensor(writer, dst, sf.get_tensor(src)) + af32("vit.patch_embd.weight", f"{vis}.embeddings.patch_embedding.weight") + af32("vit.patch_embd.bias", f"{vis}.embeddings.patch_embedding.bias") + af32("vit.pos_embd", f"{vis}.embeddings.position_embedding.weight") + for i in range(v["vit_layers"]): + L = f"{vis}.encoder.layers.{i}." + af32(f"vit.blk.{i}.ln1.weight", L + "layer_norm1.weight"); af32(f"vit.blk.{i}.ln1.bias", L + "layer_norm1.bias") + af32(f"vit.blk.{i}.ln2.weight", L + "layer_norm2.weight"); af32(f"vit.blk.{i}.ln2.bias", L + "layer_norm2.bias") + for q in ("q", "k", "v"): + akeep(f"vit.blk.{i}.attn_{q}.weight", L + f"self_attn.{q}_proj.weight") + af32 (f"vit.blk.{i}.attn_{q}.bias", L + f"self_attn.{q}_proj.bias") + akeep(f"vit.blk.{i}.attn_o.weight", L + "self_attn.out_proj.weight"); af32(f"vit.blk.{i}.attn_o.bias", L + "self_attn.out_proj.bias") + akeep(f"vit.blk.{i}.fc1.weight", L + "mlp.fc1.weight"); af32(f"vit.blk.{i}.fc1.bias", L + "mlp.fc1.bias") + akeep(f"vit.blk.{i}.fc2.weight", L + "mlp.fc2.weight"); af32(f"vit.blk.{i}.fc2.bias", L + "mlp.fc2.bias") + af32("vit.post_ln.weight", f"{vis}.post_layernorm.weight"); af32("vit.post_ln.bias", f"{vis}.post_layernorm.bias") + akeep("mm.proj.weight", f"{mmp}.linear.weight") + af32 ("mm.proj.bias", f"{mmp}.linear.bias") + + def _read_norm_eps(meta_path: Path) -> Optional[float]: if not meta_path.exists(): @@ -200,6 +251,15 @@ def _add_kv(writer: gguf.GGUFWriter, cfg: dict) -> None: writer.add_float64 (KV("rope_theta"), cfg["rope_theta"]) writer.add_float32 (KV("rms_norm_eps"), cfg["rms_norm_eps"]) writer.add_float32 (KV("norm_eps"), cfg["norm_eps"]) + if "vit" in cfg: + v = cfg["vit"] + writer.add_uint32 (KV("vit_hidden"), v["vit_hidden"]) + writer.add_uint32 (KV("vit_layers"), v["vit_layers"]) + writer.add_uint32 (KV("vit_heads"), v["vit_heads"]) + writer.add_uint32 (KV("image_size"), v["image_size"]) + writer.add_uint32 (KV("patch_size"), v["patch_size"]) + writer.add_uint32 (KV("n_img_tokens"), v["n_img_tokens"]) + writer.add_float32(KV("vit_ln_eps"), v["vit_ln_eps"]) def main() -> int: ap = argparse.ArgumentParser(description=__doc__, @@ -305,6 +365,12 @@ def _maxlayer(pfx: str) -> int: print("loading normalizer stats...") stats = _load_stats(sf, ckpt, cfg["real_state_dim"], cfg["real_action_dim"]) + vit = _probe_vision(sf, keys, cfg_json) + cfg["vit"] = vit + print(f"vision: SigLIP hidden={vit['vit_hidden']} layers={vit['vit_layers']} " + f"heads={vit['vit_heads']} image={vit['image_size']} patch={vit['patch_size']} " + f"tokens={vit['n_img_tokens']} ln_eps={vit['vit_ln_eps']:g}") + out.parent.mkdir(parents=True, exist_ok=True) print(f"writing {out}") writer = gguf.GGUFWriter(str(out), arch=ARCH) @@ -326,6 +392,8 @@ def _maxlayer(pfx: str) -> int: ]: _add_one_tensor(writer, suf, sf.get_tensor(f"{PFX_PROJ}.{suf}")) + _add_vision_tensors(writer, sf, vit) + for k, v in stats.items(): writer.add_tensor(k, v, raw_dtype=gguf.GGMLQuantizationType.F32) @@ -335,8 +403,7 @@ def _maxlayer(pfx: str) -> int: writer.close() print(f"done. {out} ({out.stat().st_size / (1024*1024):.1f} MiB)") - print("note: the SigLIP vision tower + multi_modal_projector are NOT in this file - " - "produce the mmproj GGUF separately (see tests/pi0/img_emb_mtmd/prepare_mmproj.py).") + print("self-contained: SigLIP vision tower + projector are bundled; no separate mmproj needed.") return 0 if __name__ == "__main__": diff --git a/scripts/convert_smolvla_to_gguf.py b/scripts/convert_smolvla_to_gguf.py index 02b747a..b53af11 100644 --- a/scripts/convert_smolvla_to_gguf.py +++ b/scripts/convert_smolvla_to_gguf.py @@ -105,6 +105,55 @@ def _load(meta_json: str, registry: str, key_prefix: str, "action", "action_mean", "action_std", real_action_dim) return out +PFX_VIS = "model.vlm_with_expert.vlm.model.vision_model" +PFX_CONN = "model.vlm_with_expert.vlm.model.connector.modality_projection.proj.weight" + + +def _probe_vision(sf, keys, cfg_json) -> dict: + if f"{PFX_VIS}.embeddings.patch_embedding.weight" not in keys or PFX_CONN not in keys: + raise SystemExit(f"vision_model / connector not found under {PFX_VIS}") + conv = sf.get_slice(f"{PFX_VIS}.embeddings.patch_embedding.weight").get_shape() + pos = sf.get_slice(f"{PFX_VIS}.embeddings.position_embedding.weight").get_shape() + fc1 = sf.get_slice(f"{PFX_VIS}.encoder.layers.0.mlp.fc1.weight").get_shape() + n_vit = 0 + while f"{PFX_VIS}.encoder.layers.{n_vit}.layer_norm1.weight" in keys: + n_vit += 1 + vcfg = cfg_json.get("vision_config") or {} + scale = int(cfg_json.get("scale_factor", vcfg.get("scale_factor", 4))) + patch = int(conv[2]); grid = int(round(int(pos[0]) ** 0.5)) + return dict(vit_hidden=int(conv[0]), vit_layers=n_vit, + vit_heads=int(vcfg.get("num_attention_heads", 12)), vit_inter=int(fc1[0]), + image_size=grid * patch, patch_size=patch, vit_pixel_shuffle=scale, + n_img_tokens=(grid // scale) ** 2, + vit_ln_eps=float(vcfg.get("layer_norm_eps", 1e-6))) + + +def _add_vision_tensors(writer: gguf.GGUFWriter, sf) -> None: + af32 = lambda dst, src: _add_one_tensor(writer, dst, sf.get_tensor(src).float()) + akeep = lambda dst, src: _add_one_tensor(writer, dst, sf.get_tensor(src)) + af32("vit.patch_embd.weight", f"{PFX_VIS}.embeddings.patch_embedding.weight") + af32("vit.patch_embd.bias", f"{PFX_VIS}.embeddings.patch_embedding.bias") + af32("vit.pos_embd", f"{PFX_VIS}.embeddings.position_embedding.weight") + i = 0 + while True: + L = f"{PFX_VIS}.encoder.layers.{i}." + try: + sf.get_slice(L + "layer_norm1.weight") + except Exception: + break + af32(f"vit.blk.{i}.ln1.weight", L + "layer_norm1.weight"); af32(f"vit.blk.{i}.ln1.bias", L + "layer_norm1.bias") + af32(f"vit.blk.{i}.ln2.weight", L + "layer_norm2.weight"); af32(f"vit.blk.{i}.ln2.bias", L + "layer_norm2.bias") + for q in ("q", "k", "v"): + akeep(f"vit.blk.{i}.attn_{q}.weight", L + f"self_attn.{q}_proj.weight") + af32 (f"vit.blk.{i}.attn_{q}.bias", L + f"self_attn.{q}_proj.bias") + akeep(f"vit.blk.{i}.attn_o.weight", L + "self_attn.out_proj.weight"); af32(f"vit.blk.{i}.attn_o.bias", L + "self_attn.out_proj.bias") + akeep(f"vit.blk.{i}.fc1.weight", L + "mlp.fc1.weight"); af32(f"vit.blk.{i}.fc1.bias", L + "mlp.fc1.bias") + akeep(f"vit.blk.{i}.fc2.weight", L + "mlp.fc2.weight"); af32(f"vit.blk.{i}.fc2.bias", L + "mlp.fc2.bias") + i += 1 + af32("vit.post_ln.weight", f"{PFX_VIS}.post_layernorm.weight"); af32("vit.post_ln.bias", f"{PFX_VIS}.post_layernorm.bias") + akeep("mm.fc.weight", PFX_CONN) + + def _add_kv(writer: gguf.GGUFWriter, cfg: dict) -> None: writer.add_string (KV("architecture"), ARCH) @@ -129,6 +178,17 @@ def _add_kv(writer: gguf.GGUFWriter, cfg: dict) -> None: writer.add_float64 (KV("max_period"), cfg["max_period"]) writer.add_float64 (KV("expert_width_multiplier"), cfg["expert_width_multiplier"]) writer.add_float32 (KV("norm_eps"), cfg["norm_eps"]) + if "vit" in cfg: + v = cfg["vit"] + writer.add_uint32 (KV("vit_hidden"), v["vit_hidden"]) + writer.add_uint32 (KV("vit_layers"), v["vit_layers"]) + writer.add_uint32 (KV("vit_heads"), v["vit_heads"]) + writer.add_uint32 (KV("vit_inter"), v["vit_inter"]) + writer.add_uint32 (KV("image_size"), v["image_size"]) + writer.add_uint32 (KV("patch_size"), v["patch_size"]) + writer.add_uint32 (KV("vit_pixel_shuffle"), v["vit_pixel_shuffle"]) + writer.add_uint32 (KV("n_img_tokens"), v["n_img_tokens"]) + writer.add_float32(KV("vit_ln_eps"), v["vit_ln_eps"]) def _add_one_tensor(writer: gguf.GGUFWriter, dst_name: str, t: torch.Tensor) -> None: @@ -242,6 +302,11 @@ def main() -> int: print("loading normalizer stats...") stats = _try_load_stats(ckpt, cfg["real_state_dim"], cfg["real_action_dim"]) + cfg["vit"] = _probe_vision(sf, set(sf.keys()), cfg_json) + print(f"vision: SigLIP hidden={cfg['vit']['vit_hidden']} layers={cfg['vit']['vit_layers']} " + f"heads={cfg['vit']['vit_heads']} inter={cfg['vit']['vit_inter']} image={cfg['vit']['image_size']} " + f"patch={cfg['vit']['patch_size']} shuffle={cfg['vit']['vit_pixel_shuffle']} tokens={cfg['vit']['n_img_tokens']}") + out.parent.mkdir(parents=True, exist_ok=True) print(f"writing {out}") writer = gguf.GGUFWriter(str(out), arch=ARCH) @@ -272,6 +337,8 @@ def main() -> int: ]: _add_one_tensor(writer, dst_name, sf.get_tensor(f"{PFX_PROJ}.{src_suf}")) + _add_vision_tensors(writer, sf) + for k, v in stats.items(): writer.add_tensor(k, v, raw_dtype=gguf.GGMLQuantizationType.F32) diff --git a/scripts/merge_pi0_mmproj_to_gguf.py b/scripts/merge_pi0_mmproj_to_gguf.py new file mode 100644 index 0000000..f2afe59 --- /dev/null +++ b/scripts/merge_pi0_mmproj_to_gguf.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +"""Fold a PaliGemma mmproj GGUF into a pi0 checkpoint GGUF, producing one +self-contained file whose SigLIP tower is built in-tree by src/models/pi0.cpp +(no llama.cpp clip.cpp). Use this when you only have the released pi0 + mmproj +GGUFs; convert_pi0_to_gguf.py bakes the same tensors straight from safetensors. + +The clip tensor names (v.blk.*, mm.input_projection.*) are remapped to the +in-tree vit.* / mm.proj.* scheme, and the SigLIP geometry is copied into pi0.* KV. +""" +import argparse +import sys + +import gguf + +# clip.cpp mmproj tensor name -> in-tree pi0 name +FIXED = { + "v.patch_embd.weight": "vit.patch_embd.weight", + "v.patch_embd.bias": "vit.patch_embd.bias", + "v.position_embd.weight": "vit.pos_embd", + "v.post_ln.weight": "vit.post_ln.weight", + "v.post_ln.bias": "vit.post_ln.bias", + "mm.input_projection.weight": "mm.proj.weight", + "mm.input_projection.bias": "mm.proj.bias", +} + + +def remap(name: str): + if name in FIXED: + return FIXED[name] + if name.startswith("v.blk."): + idx, sub = name[len("v.blk."):].split(".", 1) + sub = sub.replace("attn_out", "attn_o").replace("ffn_up", "fc1").replace("ffn_down", "fc2") + return f"vit.blk.{idx}.{sub}" + return None # audio / unrelated -> drop + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("ckpt", help="pi0-*.gguf produced by convert_pi0_to_gguf.py") + ap.add_argument("mmproj", help="mmproj-*.gguf (PaliGemma SigLIP tower)") + ap.add_argument("out", help="output combined .gguf") + args = ap.parse_args() + + ck = gguf.GGUFReader(args.ckpt) + mm = gguf.GGUFReader(args.mmproj) + + if not ck.get_field("pi0.architecture"): + sys.exit(f"{args.ckpt} is not a pi0 GGUF (pi0.architecture missing)") + + w = gguf.GGUFWriter(args.out, arch="pi0") + + # copy every ckpt KV verbatim (skip virtual + the arch GGUFWriter re-adds) + for f in ck.fields.values(): + if f.name == gguf.Keys.General.ARCHITECTURE or f.name.startswith("GGUF."): + continue + w.add_key_value(f.name, f.contents(), f.types[0], + sub_type=f.types[1] if len(f.types) > 1 else None) + + # SigLIP geometry from the mmproj -> pi0.vit_* KV + g = lambda k: mm.get_field(k).contents() + vh, vl, vhd = g("clip.vision.embedding_length"), g("clip.vision.block_count"), g("clip.vision.attention.head_count") + img, patch = g("clip.vision.image_size"), g("clip.vision.patch_size") + eps = g("clip.vision.attention.layer_norm_epsilon") + ntok = (img // patch) ** 2 + w.add_uint32("pi0.vit_hidden", int(vh)) + w.add_uint32("pi0.vit_layers", int(vl)) + w.add_uint32("pi0.vit_heads", int(vhd)) + w.add_uint32("pi0.image_size", int(img)) + w.add_uint32("pi0.patch_size", int(patch)) + w.add_uint32("pi0.n_img_tokens", int(ntok)) + w.add_float32("pi0.vit_ln_eps", float(eps)) + print(f"vision: hidden={vh} layers={vl} heads={vhd} image={img} patch={patch} tokens={ntok} eps={eps}") + + # ckpt tensors as-is + for t in ck.tensors: + w.add_tensor(t.name, t.data, raw_dtype=t.tensor_type) + + # remapped vision tensors + n_vis = 0 + for t in mm.tensors: + nn = remap(t.name) + if nn is None: + continue + w.add_tensor(nn, t.data, raw_dtype=t.tensor_type) + n_vis += 1 + print(f"copied {len(ck.tensors)} ckpt tensors + {n_vis} vision tensors") + + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/print_versions.sh b/scripts/print_versions.sh index 9820fe1..4681af5 100644 --- a/scripts/print_versions.sh +++ b/scripts/print_versions.sh @@ -26,8 +26,8 @@ set -u ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -LLAMA_DIR="$ROOT/third_party/llama.cpp" -PATCH="$ROOT/patches/llama.cpp-vla.patch" +# llama.cpp lives in the build tree (cmake FetchContent); set VLA_BUILD_DIR to override. +LLAMA_DIR="${VLA_BUILD_DIR:-$ROOT/build}/_deps/llama-src" sha256() { if command -v sha256sum >/dev/null 2>&1; then @@ -76,11 +76,6 @@ echo VLA_HEAD=$(git_head_or_dash "$ROOT") VLA_DESC=$(git_describe_or_dash "$ROOT") echo "- repo HEAD: \`${VLA_HEAD}\` (\`${VLA_DESC}\`)" -if [[ -f "$PATCH" ]]; then - echo "- patch \`patches/llama.cpp-vla.patch\` sha256: \`$(sha256 "$PATCH")\`" -else - echo "- patch \`patches/llama.cpp-vla.patch\`: **missing**" -fi # ---- llama.cpp ---- echo @@ -89,7 +84,8 @@ echo LLAMA_HEAD=$(git_head_or_dash "$LLAMA_DIR") LLAMA_DESC=$(git_describe_or_dash "$LLAMA_DIR") echo "- HEAD: \`${LLAMA_HEAD}\` (\`${LLAMA_DESC}\`)" -echo "- expected pinned tag (from \`patches/patch.sh\`): \`b9016\`" +LLAMA_TAG=$(grep -m1 'GIT_TAG' "$ROOT/CMakeLists.txt" | grep -oE 'b[0-9]+' || echo '?') +echo "- expected pinned tag (from \`CMakeLists.txt\`): \`${LLAMA_TAG}\`" # ---- GGUFs ---- echo diff --git a/scripts/quantize_gguf.py b/scripts/quantize_gguf.py new file mode 100644 index 0000000..3ebd0e0 --- /dev/null +++ b/scripts/quantize_gguf.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Requantize a vla.cpp GGUF: pack large weight matrices to Q8_0/Q4_0 and copy +everything else unchanged. The loader keeps quantized weights packed and lets +ggml_mul_mat dequantize at compute, so a Q8_0 file is about half the size of the +bf16 one with near-identical actions. Embeddings, the output head, norms, conv +patch embeddings and position tables stay float (row-fetch and small tensors do +not benefit and can lose accuracy). + + python scripts/quantize_gguf.py --in model-bf16.gguf --out model-q8_0.gguf --type Q8_0 +""" + +import argparse +import numpy as np +import gguf + +# Substrings that keep a tensor at its source precision. Embeddings, the output +# head, norms, conv, position tables and the action expert stay float. The vision +# tower stays float too by default; add --vision to pack it as well. +SKIP = ("token_embd", "output.weight", "patch_embd", "norm", "pos", "embed", + "cls", "action", "state", "expert", "dit", "adaln", "ada_", "time") +SKIP_VISION = ("vit", "vision") + +# Block size per row (ne0 must divide this). Only the types the gguf writer can +# pack are offered; Q8_0 is near-lossless, Q4_0/Q4_1 are 4-bit. +QK = {"Q8_0": 32, "Q4_0": 32, "Q5_0": 32, "Q4_1": 32, "Q5_1": 32} + + +def bf16_to_f32(u8: np.ndarray) -> np.ndarray: + u16 = u8.view(np.uint16).astype(np.uint32) + return (u16 << 16).view(np.float32) + + +def to_f32(t) -> np.ndarray: + # C-order array with the ggml ne axes reversed (ne0 is the last axis). + shape = tuple(int(x) for x in t.shape[::-1]) + if t.tensor_type == gguf.GGMLQuantizationType.F32: + return t.data.astype(np.float32, copy=False).reshape(shape) + if t.tensor_type == gguf.GGMLQuantizationType.BF16: + return bf16_to_f32(np.ascontiguousarray(t.data).reshape(-1)).reshape(shape) + return None # already packed / unsupported source + + +def eligible(name: str, shape, qtype: str, skip) -> bool: + if len(shape) != 2 or not name.endswith(".weight"): + return False + if any(s in name for s in skip): + return False + return int(shape[0]) % QK[qtype] == 0 # ne0 is the quantization axis + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--in", dest="src", required=True) + ap.add_argument("--out", dest="dst", required=True) + ap.add_argument("--type", default="Q8_0", choices=sorted(QK)) + ap.add_argument("--vision", action="store_true", help="also pack the vision tower") + args = ap.parse_args() + skip = SKIP if args.vision else SKIP + SKIP_VISION + + r = gguf.GGUFReader(args.src) + arch = r.fields["general.architecture"].contents() + w = gguf.GGUFWriter(args.dst, arch) + + meta = {"GGUF.version", "GGUF.tensor_count", "GGUF.kv_count", "general.architecture"} + for name, f in r.fields.items(): + if name in meta: + continue + if f.types and f.types[0] == gguf.GGUFValueType.ARRAY: + w.add_array(name, f.contents()) + else: + w.add_key_value(name, f.contents(), f.types[0]) + + qtype = getattr(gguf.GGMLQuantizationType, args.type) + F32, BF16 = gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.BF16 + n_q = 0 + bytes_in = bytes_out = 0 + for t in r.tensors: + src_bytes = int(t.data.nbytes) + bytes_in += src_bytes + f32 = to_f32(t) if eligible(t.name, t.shape, args.type, skip) else None + if f32 is not None: + packed = gguf.quants.quantize(f32, qtype) # rows are byte-sized; shape inferred + w.add_tensor(t.name, packed, raw_dtype=qtype) + bytes_out += int(packed.nbytes) + n_q += 1 + else: + # Pass copies in their natural dtype so the writer keeps the size right. + data = np.ascontiguousarray(t.data) + if t.tensor_type == BF16: + data = data.view(np.uint16) + elif t.tensor_type == F32: + data = data.astype(np.float32, copy=False) + w.add_tensor(t.name, data, raw_dtype=t.tensor_type) + bytes_out += src_bytes + + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + print(f"{args.type}: quantized {n_q}/{len(r.tensors)} tensors, " + f"weights {bytes_in/1e9:.2f} GB -> {bytes_out/1e9:.2f} GB " + f"({100*bytes_out/bytes_in:.0f}%)") + + +if __name__ == "__main__": + main() diff --git a/src/arch.h b/src/arch.h index d556ebc..633e8ff 100644 --- a/src/arch.h +++ b/src/arch.h @@ -27,11 +27,20 @@ #include "model.h" +#include #include #include +#include namespace vla { +// Default CPU thread count for the in-tree loaders: all cores, capped at 8, +// with a safe fallback when the hardware count is unknown. +inline int default_cpu_threads() { + const unsigned hw = std::thread::hardware_concurrency(); + return hw == 0 ? 4 : (int) std::min(hw, 8u); +} + /** * @brief Identifier for every VLA architecture the engine can serve. * diff --git a/src/kernels/bitvla/bitvla_lm_cuda.cu b/src/kernels/bitvla/bitvla_lm_cuda.cu index a2ffad0..0e4c455 100644 --- a/src/kernels/bitvla/bitvla_lm_cuda.cu +++ b/src/kernels/bitvla/bitvla_lm_cuda.cu @@ -410,26 +410,12 @@ static int run_layer(bitvla_lm_cuda_ctx* ctx, int L, int seq, cudaStream_t strea ctx->d_act_s, lr.v_ws, seq, hkv, hidden, stream); l0_dump("L0_02_qkv_proj", ctx->d_qkv, (size_t)seq * (hq + 2*hkv)); - const int q_stride_bytes = hq * (int)sizeof(__nv_bfloat16); - const int k_stride_bytes = hkv * (int)sizeof(__nv_bfloat16); - for (int h = 0; h < n_q; ++h) { - CUDA_OK(cudaMemcpy2DAsync( - ctx->d_q_HShd + (size_t)h * seq * hd, hd * sizeof(__nv_bfloat16), - q_dense + (size_t)h * hd, q_stride_bytes, - hd * sizeof(__nv_bfloat16), seq, cudaMemcpyDeviceToDevice, stream)); - } - for (int h = 0; h < n_kv; ++h) { - CUDA_OK(cudaMemcpy2DAsync( - ctx->d_k_HShd + (size_t)h * seq * hd, hd * sizeof(__nv_bfloat16), - k_dense + (size_t)h * hd, k_stride_bytes, - hd * sizeof(__nv_bfloat16), seq, cudaMemcpyDeviceToDevice, stream)); - } - for (int h = 0; h < n_kv; ++h) { - CUDA_OK(cudaMemcpy2DAsync( - ctx->d_v_HShd + (size_t)h * seq * hd, hd * sizeof(__nv_bfloat16), - v_dense + (size_t)h * hd, k_stride_bytes, - hd * sizeof(__nv_bfloat16), seq, cudaMemcpyDeviceToDevice, stream)); - } + // Split the interleaved [seq, H*hd] projections into head-major [H, seq, hd] + // with one kernel per tensor instead of a cudaMemcpy2DAsync per head (30 tiny + // launches per layer). Same transpose the ViT head-split already uses. + bitvla_transpose_sNhd_to_NshHd_bf16(q_dense, ctx->d_q_HShd, seq, n_q, hd, stream); + bitvla_transpose_sNhd_to_NshHd_bf16(k_dense, ctx->d_k_HShd, seq, n_kv, hd, stream); + bitvla_transpose_sNhd_to_NshHd_bf16(v_dense, ctx->d_v_HShd, seq, n_kv, hd, stream); bitvla_rope_neox_bf16(ctx->d_q_HShd, ctx->d_cos, ctx->d_sin, n_q, seq, hd, stream); bitvla_rope_neox_bf16(ctx->d_k_HShd, ctx->d_cos, ctx->d_sin, n_kv, seq, hd, stream); diff --git a/src/model.cpp b/src/model.cpp index 73c9198..b99ba3c 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -47,6 +47,9 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { auto try_str = [&](const char * key, std::string& val) -> bool { const int64_t kid = gguf_find_key(gctx, key); if (kid < 0) return false; + // gguf_get_val_str asserts (aborts) if the key is not a string, so a + // malformed GGUF would kill the process here. Fail the probe instead. + if (gguf_get_kv_type(gctx, kid) != GGUF_TYPE_STRING) return false; const char * s = gguf_get_val_str(gctx, kid); if (!s) return false; val = s; @@ -199,6 +202,7 @@ const Stats& last_stats(const Model* m) { } std::vector predict(Model* m, const Inputs& in) { + if (!m || !m->impl) return {}; return m->impl->predict(in); } diff --git a/src/model.h b/src/model.h index 93957e5..9440f86 100644 --- a/src/model.h +++ b/src/model.h @@ -38,7 +38,8 @@ namespace vla { * Populated from the GGUF metadata at load time. Sequence counts are in * tokens; @c hidden / @c expert_h / @c intermediate refer to model widths. * The action chunk returned by @ref predict has shape - * @c [num_steps, real_action_dim] in row-major order. + * @c [num_steps, max_action_dim] in row-major order; only the first + * @c real_action_dim columns are valid, the rest are zero padding. */ struct Config { int64_t n_img; ///< Image-token count fed to the LM. @@ -127,10 +128,12 @@ struct Inputs { const int32_t* lang_tokens; ///< Tokenised language instruction. int n_lang; ///< Length of @ref lang_tokens. - const float* state; ///< Proprioception, length @c real_state_dim. + const float* state; ///< Proprioception, length @c max_state_dim + /// (pad @c real_state_dim..max with zeros). const float* noise; ///< Initial noise for the action expert. - /// Optional override for the language attention mask (per-token). + /// Optional per-token language attention mask. Only Evo-1 honors this; the + /// other architectures derive their own mask and ignore it. const int32_t* attention_mask = nullptr; int attention_mask_n = 0; ///< Length of @ref attention_mask. @@ -176,8 +179,8 @@ const Config& model_config(const Model* m); * * @param m A handle from @ref model_load. * @param in Filled-in @ref Inputs struct. - * @return Action chunk of length @c num_steps * real_action_dim, in - * row-major order. + * @return Action chunk of length @c num_steps * max_action_dim, in + * row-major order (first @c real_action_dim columns valid). */ std::vector predict(Model* m, const Inputs& in); diff --git a/src/models/bitvla.cpp b/src/models/bitvla.cpp index bfd7a7a..57f23fe 100644 --- a/src/models/bitvla.cpp +++ b/src/models/bitvla.cpp @@ -22,6 +22,7 @@ #include "ggml-cuda.h" #endif #include "gguf.h" +#include "models/gguf_reader.h" #ifdef VLA_BITVLA_CUDA_KERNELS #include "kernels/bitvla/bitvla_lm_cuda.h" @@ -48,91 +49,6 @@ namespace vla { namespace { -struct gguf_reader { - gguf_context * gctx = nullptr; - ggml_context * meta_ctx = nullptr; - FILE * fp = nullptr; - size_t data_off = 0; - - bool open(const std::string & path) { - gguf_init_params p{}; - p.no_alloc = true; - p.ctx = &meta_ctx; - gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) { std::fprintf(stderr, "vla(bitvla): gguf_init_from_file failed for %s\n", path.c_str()); return false; } - fp = std::fopen(path.c_str(), "rb"); - if (!fp) { std::fprintf(stderr, "vla(bitvla): fopen failed for %s\n", path.c_str()); return false; } - data_off = gguf_get_data_offset(gctx); - return true; - } - ~gguf_reader() { - if (fp) std::fclose(fp); - if (gctx) gguf_free(gctx); - if (meta_ctx) ggml_free(meta_ctx); - } - gguf_reader() = default; - gguf_reader(const gguf_reader &) = delete; - gguf_reader & operator=(const gguf_reader &) = delete; - - bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { const int64_t id = gguf_find_key(gctx, k); return id < 0 ? std::string() : std::string(gguf_get_val_str(gctx, id)); } - const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } - - bool read_raw(const char * name, void * buf) { - const int64_t id = gguf_find_tensor(gctx, name); - if (id < 0) { std::fprintf(stderr, "vla(bitvla): missing tensor %s\n", name); return false; } - const size_t off = data_off + gguf_get_tensor_offset(gctx, id); - const size_t nb = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) off, SEEK_SET) != 0) return false; - return std::fread(buf, 1, nb, fp) == nb; - } - std::vector read_f32(const char * name) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(bitvla): missing tensor %s\n", name); return {}; } - const int64_t n = ggml_nelements(t); - std::vector out(n); - if (t->type == GGML_TYPE_F32) { if (!read_raw(name, out.data())) return {}; } - else if (t->type == GGML_TYPE_BF16) { std::vector tmp(n); if (!read_raw(name, tmp.data())) return {}; ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); } - else { std::fprintf(stderr, "vla(bitvla): tensor %s unsupported type %d\n", name, (int) t->type); return {}; } - return out; - } - std::vector read_convert(const char * name, ggml_type target) { - if (target == GGML_TYPE_I8) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(bitvla): missing tensor %s\n", name); return {}; } - std::vector o(ggml_nbytes(t)); - if (!read_raw(name, o.data())) return {}; - return o; - } - std::vector f = read_f32(name); - if (f.empty()) return {}; - const int64_t n = (int64_t) f.size(); - if (target == GGML_TYPE_F32) { std::vector o(n * 4); std::memcpy(o.data(), f.data(), o.size()); return o; } - if (target == GGML_TYPE_BF16) { std::vector o(n * 2); ggml_fp32_to_bf16_row(f.data(), reinterpret_cast(o.data()), n); return o; } - std::fprintf(stderr, "vla(bitvla): unsupported resident type %d for %s\n", (int) target, name); return {}; - } - bool fetch_rows_f32(const char * name, const std::vector & row_ids, float * dst, int64_t cols) { - const ggml_tensor * t = meta(name); - if (!t || t->ne[0] != cols) { std::fprintf(stderr, "vla(bitvla): %s shape unfit for row-fetch\n", name); return false; } - const int64_t rows = t->ne[1]; - const int64_t id = gguf_find_tensor(gctx, name); - const size_t base = data_off + gguf_get_tensor_offset(gctx, id); - const size_t elsz = (t->type == GGML_TYPE_F32) ? 4u : 2u; - const size_t rb = (size_t) cols * elsz; - std::vector row(rb); - for (size_t k = 0; k < row_ids.size(); ++k) { - const int32_t r = row_ids[k]; - if (r < 0 || r >= rows) { std::fprintf(stderr, "vla(bitvla): row %d out of range for %s\n", r, name); return false; } - if (std::fseek(fp, (long) (base + (size_t) r * rb), SEEK_SET) != 0) return false; - if (std::fread(row.data(), 1, rb, fp) != rb) return false; - if (elsz == 4) std::memcpy(dst + k * cols, row.data(), rb); - else ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst + k * cols, cols); - } - return true; - } -}; struct VitLayerW { ggml_tensor *ln1w, *ln1b, *ln2w, *ln2b; @@ -178,9 +94,11 @@ struct BitvlaModelArch : public ModelArchBase { ~BitvlaModelArch() override; std::string gguf_path; + gguf_reader emb_reader{"bitvla"}; // stays open for per-step token-embedding row fetches + std::vector stop_embed; // cached constant stop-token embedding row ggml_backend_t backend = nullptr; bool is_cuda = false; - int n_threads = 4; + int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; @@ -461,12 +379,14 @@ namespace { static void recover_ternary_and_scale(const float* W, int64_t n, std::vector& ternary, float& absmean) { - - float amax = 0.0f; - for (int64_t i = 0; i < n; ++i) amax = std::max(amax, std::fabs(W[i])); - if (amax < 1e-5f) amax = 1e-5f; - absmean = amax; - const float inv = 1.0f / amax; + // Per-tensor absmean scale (1/mean|W|), matching scripts/bitvla_int2_pack.py; + // the int2-packed path bakes the same scale. + double s = 0.0; + for (int64_t i = 0; i < n; ++i) s += std::fabs((double) W[i]); + float mean = n > 0 ? (float) (s / (double) n) : 0.0f; + if (mean < 1e-5f) mean = 1e-5f; + absmean = mean; + const float inv = 1.0f / mean; ternary.resize(n); for (int64_t i = 0; i < n; ++i) { float q = std::nearbyintf(W[i] * inv); @@ -598,12 +518,23 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, m->gguf_path = ckpt_path; m->matmul_type = std::getenv("VLA_BITVLA_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; - gguf_reader g; + gguf_reader g("bitvla"); if (!g.open(ckpt_path)) return nullptr; if (!g.has("bitvla.architecture") && !g.has("general.architecture")) { std::fprintf(stderr, "vla(bitvla): %s is not a bitvla GGUF\n", ckpt_path.c_str()); return nullptr; } if (!load_config(g, *m, m->cfg)) return nullptr; + + // Keep one reader open for the per-step token-embedding fetches (token_embd + // stays on disk under int2 packing) and cache the constant stop-token row, + // so predict() no longer re-opens and re-parses the GGUF twice per call. + if (!m->emb_reader.open(ckpt_path)) return nullptr; + m->stop_embed.resize((size_t) m->lm_hidden); + { + std::vector sid{ m->stop_id }; + if (!m->emb_reader.fetch_rows_f32("token_embd.weight", sid, + m->stop_embed.data(), m->lm_hidden)) return nullptr; + } std::printf("vla(bitvla): vit=BitSigLIP-L %lldd×%lldL×%lldh@%lld ⇒ %lld patches mm=%lld→%lld " "lm=BitNet %lldd×%lldL (%lldq/%lldkv×%lld) rope=%g chunk×dim=%lld×%lld vocab=%lld resident=%s\n", (long long) m->vit_hidden, (long long) m->vit_layers, (long long) m->vit_heads, (long long) m->image_size, @@ -624,7 +555,7 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); if (!gt) { std::fprintf(stderr, "vla(bitvla): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, type, ggml_n_dims(gt), gt->ne); + ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); ggml_set_name(t, name); return t; }; @@ -784,25 +715,29 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, __nv_bfloat16* onorm = upload_bf16_from_f32((const float*) m->lm_output_norm->data, m->lm_hidden, m->cuda_devptrs); bitvla_lm_cuda_set_output_norm(m->lm_cuda_ctx, onorm); - cudaMalloc(&m->d_inputs_embeds, (size_t) max_seq * m->lm_hidden * sizeof(__nv_bfloat16)); - cudaMalloc(&m->d_last_hidden, (size_t) max_seq * m->lm_hidden * sizeof(__nv_bfloat16)); - cudaMalloc(&m->d_action_hidden, (size_t) (m->num_actions_chunk * m->action_dim) * m->lm_hidden * sizeof(__nv_bfloat16)); - cudaMalloc(&m->d_action_ids, (size_t) (m->num_actions_chunk * m->action_dim) * sizeof(int32_t)); - m->cuda_lm_ready = true; - size_t packed_bytes = 0; - for (void* p : m->cuda_devptrs) { - (void) p; - + cudaError_t lm_ce = cudaMalloc(&m->d_inputs_embeds, (size_t) max_seq * m->lm_hidden * sizeof(__nv_bfloat16)); + if (lm_ce == cudaSuccess) lm_ce = cudaMalloc(&m->d_last_hidden, (size_t) max_seq * m->lm_hidden * sizeof(__nv_bfloat16)); + if (lm_ce == cudaSuccess) lm_ce = cudaMalloc(&m->d_action_hidden, (size_t) (m->num_actions_chunk * m->action_dim) * m->lm_hidden * sizeof(__nv_bfloat16)); + if (lm_ce == cudaSuccess) lm_ce = cudaMalloc(&m->d_action_ids, (size_t) (m->num_actions_chunk * m->action_dim) * sizeof(int32_t)); + // only enable the CUDA LM once every work buffer is really allocated. + if (lm_ce != cudaSuccess) { + std::fprintf(stderr, "vla(bitvla): CUDA LM buffer alloc failed (%s); using CPU LM\n", + cudaGetErrorString(lm_ce)); + cudaFree(m->d_inputs_embeds); cudaFree(m->d_last_hidden); + cudaFree(m->d_action_hidden); cudaFree(m->d_action_ids); + m->d_inputs_embeds = nullptr; m->d_last_hidden = nullptr; + m->d_action_hidden = nullptr; m->d_action_ids = nullptr; + } else { + m->cuda_lm_ready = true; + const size_t packed_bytes = (size_t) m->lm_layers * ( + (size_t)(m->lm_q + 2*m->lm_kv) * m->lm_head_dim * m->lm_hidden / 4 + + (size_t) m->lm_hidden * m->lm_hidden / 4 + + (size_t) 2 * m->lm_inter * m->lm_hidden / 4 + + (size_t) m->lm_hidden * m->lm_inter / 4); + std::printf("vla(bitvla): CUDA LM forward ENABLED - packed int2 weights = %.2f MiB, max_seq=%d\n", + packed_bytes / (1024.0 * 1024.0), max_seq); } - packed_bytes = (size_t) m->lm_layers * ( - (size_t)(m->lm_q + 2*m->lm_kv) * m->lm_head_dim * m->lm_hidden / 4 + - (size_t) m->lm_hidden * m->lm_hidden / 4 + - (size_t) 2 * m->lm_inter * m->lm_hidden / 4 + - (size_t) m->lm_hidden * m->lm_inter / 4); - std::printf("vla(bitvla): CUDA LM forward ENABLED - packed int2 weights = %.2f MiB, max_seq=%d\n", - packed_bytes / (1024.0 * 1024.0), max_seq); - const int patch_flat = 3 * m->patch_size * m->patch_size; const int mm_out = (int) m->lm_hidden; const int ffn_pad = ((m->vit_inter + 127) / 128) * 128; @@ -1195,9 +1130,7 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { std::fprintf(stderr, "vla(bitvla): prompt token %d out of vocab\n", id); return {}; } } - gguf_reader g_re; - if (!g_re.open(gguf_path)) return {}; - if (!g_re.fetch_rows_f32("token_embd.weight", ids, inputs_embeds.data(), hidden_l)) return {}; + if (!emb_reader.fetch_rows_f32("token_embd.weight", ids, inputs_embeds.data(), hidden_l)) return {}; int64_t k_img = 0; for (int64_t i = 0; i < n_lang_in; ++i) { @@ -1214,11 +1147,8 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { } } - gguf_reader g_re2; - if (!g_re2.open(gguf_path)) return {}; - std::vector stop_id_v{ stop_id }; - if (!g_re2.fetch_rows_f32("token_embd.weight", stop_id_v, - inputs_embeds.data() + (size_t) (seq - 1) * hidden_l, hidden_l)) return {}; + std::memcpy(inputs_embeds.data() + (size_t) (seq - 1) * hidden_l, + stop_embed.data(), (size_t) hidden_l * sizeof(float)); } else { const int64_t n_prompt = n_lang_in; @@ -1229,19 +1159,14 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { std::fprintf(stderr, "vla(bitvla): prompt token %d out of vocab\n", id); return {}; } } - gguf_reader g_re; - if (!g_re.open(gguf_path)) return {}; - if (!g_re.fetch_rows_f32("token_embd.weight", ids, inputs_embeds.data(), hidden_l)) return {}; + if (!emb_reader.fetch_rows_f32("token_embd.weight", ids, inputs_embeds.data(), hidden_l)) return {}; } std::memcpy(inputs_embeds.data() + (size_t) n_prompt * hidden_l, img_embeds_host.data(), (size_t) n_img_tok * hidden_l * sizeof(float)); std::memcpy(inputs_embeds.data() + (size_t) (n_prompt + n_img_tok) * hidden_l, proprio_embed_host.data(), (size_t) hidden_l * sizeof(float)); - gguf_reader g_re2; - if (!g_re2.open(gguf_path)) return {}; - std::vector stop_id_v{ stop_id }; - if (!g_re2.fetch_rows_f32("token_embd.weight", stop_id_v, - inputs_embeds.data() + (size_t) (seq - 1) * hidden_l, hidden_l)) return {}; + std::memcpy(inputs_embeds.data() + (size_t) (seq - 1) * hidden_l, + stop_embed.data(), (size_t) hidden_l * sizeof(float)); } _dump_bin("inputs_embeds", inputs_embeds.data(), inputs_embeds.size()); diff --git a/src/models/evo1.cpp b/src/models/evo1.cpp index a0e4a9a..77e169c 100644 --- a/src/models/evo1.cpp +++ b/src/models/evo1.cpp @@ -25,6 +25,7 @@ #include "ggml-metal.h" #endif #include "gguf.h" +#include "models/gguf_reader.h" #include #include @@ -40,105 +41,6 @@ namespace vla { namespace { -struct gguf_reader { - gguf_context * gctx = nullptr; - ggml_context * meta_ctx = nullptr; - FILE * fp = nullptr; - size_t data_off = 0; - - bool open(const std::string & path) { - gguf_init_params p{}; - p.no_alloc = true; - p.ctx = &meta_ctx; - gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) { std::fprintf(stderr, "vla(evo1): gguf_init_from_file failed for %s\n", path.c_str()); return false; } - fp = std::fopen(path.c_str(), "rb"); - if (!fp) { std::fprintf(stderr, "vla(evo1): fopen failed for %s\n", path.c_str()); return false; } - data_off = gguf_get_data_offset(gctx); - return true; - } - ~gguf_reader() { - if (fp) std::fclose(fp); - if (gctx) gguf_free(gctx); - if (meta_ctx) ggml_free(meta_ctx); - } - gguf_reader() = default; - gguf_reader(const gguf_reader &) = delete; - gguf_reader & operator=(const gguf_reader &) = delete; - - bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - double f64(const char * k) const { return gguf_get_val_f64(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { return gguf_get_val_str(gctx, gguf_find_key(gctx, k)); } - const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } - - bool read_raw(const char * name, void * buf) { - const int64_t id = gguf_find_tensor(gctx, name); - if (id < 0) { std::fprintf(stderr, "vla(evo1): missing tensor %s\n", name); return false; } - const size_t off = data_off + gguf_get_tensor_offset(gctx, id); - const size_t nb = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) off, SEEK_SET) != 0) return false; - return std::fread(buf, 1, nb, fp) == nb; - } - - std::vector read_f32(const char * name) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(evo1): missing tensor %s\n", name); return {}; } - const int64_t n = ggml_nelements(t); - std::vector out(n); - if (t->type == GGML_TYPE_F32) { - if (!read_raw(name, out.data())) return {}; - } else if (t->type == GGML_TYPE_BF16) { - std::vector tmp(n); - if (!read_raw(name, tmp.data())) return {}; - ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); - } else { - std::fprintf(stderr, "vla(evo1): tensor %s unsupported type %d\n", name, (int) t->type); return {}; - } - return out; - } - - std::vector read_convert(const char * name, ggml_type target) { - std::vector f32 = read_f32(name); - if (f32.empty()) return {}; - const int64_t n = (int64_t) f32.size(); - if (target == GGML_TYPE_F32) { - std::vector out(n * sizeof(float)); - std::memcpy(out.data(), f32.data(), out.size()); - return out; - } - if (target == GGML_TYPE_BF16) { - std::vector out(n * sizeof(ggml_bf16_t)); - ggml_fp32_to_bf16_row(f32.data(), reinterpret_cast(out.data()), n); - return out; - } - std::fprintf(stderr, "vla(evo1): unsupported resident type %d for %s\n", (int) target, name); - return {}; - } - - bool fetch_rows_f32(const char * name, const std::vector & row_ids, float * dst, int64_t cols) { - const ggml_tensor * t = meta(name); - if (!t || t->ne[0] != cols || t->ne[2] != 1 || t->ne[3] != 1) { - std::fprintf(stderr, "vla(evo1): %s shape unfit for row-fetch\n", name); return false; - } - const int64_t rows = t->ne[1]; - const int64_t id = gguf_find_tensor(gctx, name); - const size_t base = data_off + gguf_get_tensor_offset(gctx, id); - const size_t elsz = (t->type == GGML_TYPE_F32) ? 4u : 2u; - const size_t rb = (size_t) cols * elsz; - std::vector row(rb); - for (size_t k = 0; k < row_ids.size(); ++k) { - const int32_t r = row_ids[k]; - if (r < 0 || r >= rows) { std::fprintf(stderr, "vla(evo1): row %d out of range for %s\n", r, name); return false; } - if (std::fseek(fp, (long) (base + (size_t) r * rb), SEEK_SET) != 0) return false; - if (std::fread(row.data(), 1, rb, fp) != rb) return false; - if (elsz == 4) std::memcpy(dst + k * cols, row.data(), rb); - else ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst + k * cols, cols); - } - return true; - } -}; struct Qwen2LayerW { ggml_tensor *attn_norm, *Wq, *bq, *Wk, *bk, *Wv, *bv, *Wo, *ffn_norm, *Wgate, *Wup, *Wdown; @@ -160,7 +62,7 @@ struct Evo1ModelArch : public ModelArchBase { ggml_backend_t backend = nullptr; bool is_cuda = false; bool is_gpu = false; - int n_threads = 4; + int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_BF16; @@ -321,6 +223,12 @@ bool load_config(const gguf_reader & g, Evo1ModelArch & m, Config & cfg) { std::fprintf(stderr, "vla(evo1): embed_dim (%lld) != lm_hidden (%lld) - not handled\n", (long long) m.embed_dim, (long long) m.lm_hidden); return false; } + // predict() reads action_dim noise floats; the server sizes noise as + // horizon*per_action_dim, so they must agree or a client noise buffer underruns. + if (m.action_dim != m.horizon * m.per_a) { + std::fprintf(stderr, "vla(evo1): action_dim (%lld) != horizon (%lld) * per_action_dim (%lld)\n", + (long long) m.action_dim, (long long) m.horizon, (long long) m.per_a); return false; + } cfg = Config{}; cfg.n_suffix = m.horizon; @@ -367,7 +275,7 @@ std::unique_ptr evo1_create(const std::string& mmproj_path, m->gguf_path = ckpt_path; m->matmul_type = std::getenv("VLA_EVO1_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; - gguf_reader g; + gguf_reader g("evo1"); if (!g.open(ckpt_path)) return nullptr; if (!g.has("evo1.architecture")) { std::fprintf(stderr, "vla(evo1): %s is not an evo1 GGUF (no evo1.architecture KV)\n", ckpt_path.c_str()); return nullptr; @@ -405,7 +313,7 @@ std::unique_ptr evo1_create(const std::string& mmproj_path, auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); if (!gt) { std::fprintf(stderr, "vla(evo1): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, type, ggml_n_dims(gt), gt->ne); + ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); ggml_set_name(t, name); return t; }; @@ -517,7 +425,11 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { const float * img_emb_ptr = nullptr; std::vector img_emb_host; if (in.precomputed_img_emb) { - n_views = in.n_img_views > 0 ? in.n_img_views : n_images; + if (in.n_img_views < 1) { + std::fprintf(stderr, "vla(evo1): precomputed_img_emb set but n_img_views=%d; cannot infer view count\n", in.n_img_views); + return {}; + } + n_views = in.n_img_views; img_emb_ptr = in.precomputed_img_emb; } else if (in.images && in.n_images > 0) { if (!have_vision) { @@ -537,7 +449,9 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { ggml_gallocr_t vga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); if (!vga || !ggml_gallocr_alloc_graph(vga, vg)) { std::fprintf(stderr, "vla(evo1): vision ggml_gallocr_alloc_graph failed\n"); - if (vga) ggml_gallocr_free(vga); ggml_free(VC); return {}; + if (vga) ggml_gallocr_free(vga); + ggml_free(VC); + return {}; } img_emb_host.assign((size_t) n_views * num_image_token * lm_hidden, 0.0f); std::vector chw; @@ -587,7 +501,7 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { input_ids.resize(max_text_length, pad_id); const int64_t SEQ = max_text_length; - gguf_reader g; + gguf_reader g("evo1"); if (!g.open(gguf_path)) return {}; std::vector inputs_embeds((size_t) SEQ * lm_hidden); if (!g.fetch_rows_f32("token_embd.weight", input_ids, inputs_embeds.data(), lm_hidden)) return {}; @@ -601,7 +515,10 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { ++img_idx; } } - if (img_idx != n_img_tokens) std::fprintf(stderr, "vla(evo1): warning - spliced %lld of %lld ViT embeds\n", (long long) img_idx, (long long) n_img_tokens); + if (img_idx != n_img_tokens) { + std::fprintf(stderr, "vla(evo1): spliced %lld of %lld ViT embeds - IMG_CTX slot count does not match the images\n", (long long) img_idx, (long long) n_img_tokens); + return {}; + } } std::vector attn_ok(SEQ, 0); @@ -620,7 +537,8 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { for (int64_t i = 0; i < per_a; ++i) { const float lo = state_min[i], hi = state_max[i]; float xn = 2.0f * (in.state[i] - lo) / (hi - lo + norm_eps_denom) - 1.0f; - if (xn < -1.0f) xn = -1.0f; if (xn > 1.0f) xn = 1.0f; + if (xn < -1.0f) xn = -1.0f; + if (xn > 1.0f) xn = 1.0f; state_norm[i] = xn; } @@ -628,16 +546,19 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { if (in.noise) { std::memcpy(x_init.data(), in.noise, x_init.size() * sizeof(float)); } else { - + // Evo1 is trained with uniform[-1,1] noise uint64_t s = 0xE701ACE5ULL ^ (uint64_t) std::chrono::steady_clock::now().time_since_epoch().count(); - for (auto & v : x_init) { s = s * 6364136223846793005ULL + 1442695040888963407ULL; v = ((float) (uint32_t)(s >> 32) / 2147483648.0f) - 1.0f; } + for (auto & v : x_init) { + s = s * 6364136223846793005ULL + 1442695040888963407ULL; + v = ((float) (uint32_t) (s >> 32) / 2147483648.0f) - 1.0f; // uniform[-1,1) + } } ggml_init_params cp = { (size_t) 96 * 1024 * 1024, nullptr, true }; ggml_context * C = ggml_init(cp); if (!C) { std::fprintf(stderr, "vla(evo1): ggml_init(ctx_compute) failed\n"); return {}; } - const int64_t E = embed_dim, hd_dit = E / dit_heads, inner_dit = 4 * E; + const int64_t E = embed_dim, hd_dit = E / dit_heads; const float scale_dit = 1.0f / std::sqrt((float) hd_dit); const int64_t Nctx = SEQ + 1; @@ -686,6 +607,8 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { ggml_tensor * qp = ggml_add(C, ggml_mul_mat(C, c.Wq, x_q), c.bq); ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, qp, hd_dit, dit_heads, horizon), 0, 2, 1, 3)); ggml_tensor * kq = ggml_mul_mat(C, c.K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + // The Evo-1 reference cross-attends over the full padded context + // (no key mask), so the action queries see every LM position. ggml_tensor * aw = ggml_soft_max_ext(C, kq, nullptr, scale_dit, 0.0f); ggml_tensor * kqv = ggml_mul_mat(C, c.V, aw); ggml_tensor * att_pre = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, kqv, 0, 2, 1, 3)), E, horizon); @@ -723,7 +646,9 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); if (!galloc || !ggml_gallocr_alloc_graph(galloc, gf)) { std::fprintf(stderr, "vla(evo1): ggml_gallocr_alloc_graph failed\n"); - if (galloc) ggml_gallocr_free(galloc); ggml_free(C); return {}; + if (galloc) ggml_gallocr_free(galloc); + ggml_free(C); + return {}; } ggml_backend_tensor_set(t_embeds, inputs_embeds.data(), 0, ggml_nbytes(t_embeds)); { std::vector pp(SEQ); for (int64_t i = 0; i < SEQ; ++i) pp[i] = (int32_t) i; ggml_backend_tensor_set(t_pos, pp.data(), 0, ggml_nbytes(t_pos)); } diff --git a/src/models/gguf_reader.h b/src/models/gguf_reader.h new file mode 100644 index 0000000..3b3f739 --- /dev/null +++ b/src/models/gguf_reader.h @@ -0,0 +1,137 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Shared GGUF reader for the in-tree model loaders. Reads a tensor as F32/BF16 +// (dequantizing BF16), passes packed bytes through when the resident type is not +// float, and fetches token-embedding rows on demand. The arch name only labels +// error messages. + +#pragma once + +#include "ggml.h" +#include "gguf.h" + +#include +#include +#include +#include +#include + +namespace vla { + +struct gguf_reader { + const char * arch = "vla"; + gguf_context * gctx = nullptr; + ggml_context * meta_ctx = nullptr; + FILE * fp = nullptr; + size_t data_off = 0; + + explicit gguf_reader(const char * arch_ = "vla") : arch(arch_) {} + ~gguf_reader() { + if (fp) std::fclose(fp); + if (gctx) gguf_free(gctx); + if (meta_ctx) ggml_free(meta_ctx); + } + gguf_reader(const gguf_reader &) = delete; + gguf_reader & operator=(const gguf_reader &) = delete; + + bool open(const std::string & path) { + gguf_init_params p{}; + p.no_alloc = true; + p.ctx = &meta_ctx; + gctx = gguf_init_from_file(path.c_str(), p); + if (!gctx) { std::fprintf(stderr, "vla(%s): gguf_init_from_file failed for %s\n", arch, path.c_str()); return false; } + fp = std::fopen(path.c_str(), "rb"); + if (!fp) { std::fprintf(stderr, "vla(%s): fopen failed for %s\n", arch, path.c_str()); return false; } + data_off = gguf_get_data_offset(gctx); + return true; + } + + bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } + uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } + float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } + double f64(const char * k) const { return gguf_get_val_f64(gctx, gguf_find_key(gctx, k)); } + std::string str(const char * k) const { const int64_t id = gguf_find_key(gctx, k); return id < 0 ? std::string() : std::string(gguf_get_val_str(gctx, id)); } + const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } + + // Resident type for a weight: keep a quantized source type (Q8_0, Q4_0, ...) + // so it stays packed and ggml_mul_mat dequantizes at compute; otherwise use + // prefer (the model's F32/BF16 matmul type, or an explicit packed type). + ggml_type resident_type(const ggml_tensor * src, ggml_type prefer) const { + return (src && ggml_is_quantized(src->type)) ? src->type : prefer; + } + + bool read_raw(const char * name, void * buf) { + const int64_t id = gguf_find_tensor(gctx, name); + if (id < 0) { std::fprintf(stderr, "vla(%s): missing tensor %s\n", arch, name); return false; } + const size_t off = data_off + gguf_get_tensor_offset(gctx, id); + const size_t nb = gguf_get_tensor_size(gctx, id); + if (fseeko(fp, (off_t) off, SEEK_SET) != 0) return false; + return std::fread(buf, 1, nb, fp) == nb; + } + + std::vector read_f32(const char * name) { + const ggml_tensor * t = meta(name); + if (!t) { std::fprintf(stderr, "vla(%s): missing tensor %s\n", arch, name); return {}; } + const int64_t n = ggml_nelements(t); + std::vector out(n); + if (t->type == GGML_TYPE_F32) { if (!read_raw(name, out.data())) return {}; } + else if (t->type == GGML_TYPE_BF16) { std::vector tmp(n); if (!read_raw(name, tmp.data())) return {}; ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); } + else { std::fprintf(stderr, "vla(%s): tensor %s unsupported type %d\n", arch, name, (int) t->type); return {}; } + return out; + } + + // F32/BF16 targets dequantize to that resident type. Any other target (I8, + // Q8_0, Q4_0, ...) is stored raw so ggml_mul_mat dequantizes at compute. + // gemma_norm adds 1.0 per weight. + std::vector read_convert(const char * name, ggml_type target, bool gemma_norm = false) { + if (target != GGML_TYPE_F32 && target != GGML_TYPE_BF16) { + const ggml_tensor * t = meta(name); + if (!t) { std::fprintf(stderr, "vla(%s): missing tensor %s\n", arch, name); return {}; } + std::vector o(ggml_nbytes(t)); + if (!read_raw(name, o.data())) return {}; + return o; + } + std::vector f = read_f32(name); + if (f.empty()) return {}; + const int64_t n = (int64_t) f.size(); + if (gemma_norm) for (int64_t i = 0; i < n; ++i) f[i] += 1.0f; + if (target == GGML_TYPE_F32) { std::vector o(n * sizeof(float)); std::memcpy(o.data(), f.data(), o.size()); return o; } + if (target == GGML_TYPE_BF16) { std::vector o(n * sizeof(ggml_bf16_t)); ggml_fp32_to_bf16_row(f.data(), reinterpret_cast(o.data()), n); return o; } + std::fprintf(stderr, "vla(%s): unsupported resident type %d for %s\n", arch, (int) target, name); return {}; + } + + bool fetch_rows_f32(const char * name, const std::vector & row_ids, float * dst, int64_t cols) { + const ggml_tensor * t = meta(name); + if (!t || t->ne[0] != cols || t->ne[2] != 1 || t->ne[3] != 1) { std::fprintf(stderr, "vla(%s): %s shape unfit for row-fetch\n", arch, name); return false; } + if (t->type != GGML_TYPE_F32 && t->type != GGML_TYPE_BF16) { std::fprintf(stderr, "vla(%s): %s type %d not f32/bf16 for row-fetch\n", arch, name, (int) t->type); return false; } + const int64_t rows = t->ne[1]; + const int64_t id = gguf_find_tensor(gctx, name); + const size_t base = data_off + gguf_get_tensor_offset(gctx, id); + const size_t elsz = (t->type == GGML_TYPE_F32) ? 4u : 2u; + const size_t rb = (size_t) cols * elsz; + std::vector row(rb); + for (size_t k = 0; k < row_ids.size(); ++k) { + const int32_t r = row_ids[k]; + if (r < 0 || r >= rows) { std::fprintf(stderr, "vla(%s): row %d out of range for %s\n", arch, r, name); return false; } + if (fseeko(fp, (off_t) (base + (size_t) r * rb), SEEK_SET) != 0) return false; + if (std::fread(row.data(), 1, rb, fp) != rb) return false; + if (elsz == 4) std::memcpy(dst + k * cols, row.data(), rb); + else ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst + k * cols, cols); + } + return true; + } +}; + +} // namespace vla diff --git a/src/models/gr00tn1d5.cpp b/src/models/gr00tn1d5.cpp index 0c7b654..6c20589 100644 --- a/src/models/gr00tn1d5.cpp +++ b/src/models/gr00tn1d5.cpp @@ -25,6 +25,7 @@ #include "ggml-metal.h" #endif #include "gguf.h" +#include "models/gguf_reader.h" #include #include @@ -41,86 +42,6 @@ namespace vla { namespace { -struct gguf_reader { - gguf_context * gctx = nullptr; - ggml_context * meta_ctx = nullptr; - FILE * fp = nullptr; - size_t data_off = 0; - - bool open(const std::string & path) { - gguf_init_params p{}; - p.no_alloc = true; - p.ctx = &meta_ctx; - gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) { std::fprintf(stderr, "vla(gr00tn1d5): gguf_init_from_file failed for %s\n", path.c_str()); return false; } - fp = std::fopen(path.c_str(), "rb"); - if (!fp) { std::fprintf(stderr, "vla(gr00tn1d5): fopen failed for %s\n", path.c_str()); return false; } - data_off = gguf_get_data_offset(gctx); - return true; - } - ~gguf_reader() { - if (fp) std::fclose(fp); - if (gctx) gguf_free(gctx); - if (meta_ctx) ggml_free(meta_ctx); - } - gguf_reader() = default; - gguf_reader(const gguf_reader &) = delete; - gguf_reader & operator=(const gguf_reader &) = delete; - - bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - double f64(const char * k) const { return gguf_get_val_f64(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { const int64_t id = gguf_find_key(gctx, k); return id < 0 ? std::string() : std::string(gguf_get_val_str(gctx, id)); } - const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } - - bool read_raw(const char * name, void * buf) { - const int64_t id = gguf_find_tensor(gctx, name); - if (id < 0) { std::fprintf(stderr, "vla(gr00tn1d5): missing tensor %s\n", name); return false; } - const size_t off = data_off + gguf_get_tensor_offset(gctx, id); - const size_t nb = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) off, SEEK_SET) != 0) return false; - return std::fread(buf, 1, nb, fp) == nb; - } - std::vector read_f32(const char * name) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(gr00tn1d5): missing tensor %s\n", name); return {}; } - const int64_t n = ggml_nelements(t); - std::vector out(n); - if (t->type == GGML_TYPE_F32) { if (!read_raw(name, out.data())) return {}; } - else if (t->type == GGML_TYPE_BF16) { std::vector tmp(n); if (!read_raw(name, tmp.data())) return {}; ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); } - else { std::fprintf(stderr, "vla(gr00tn1d5): tensor %s unsupported type %d\n", name, (int) t->type); return {}; } - return out; - } - std::vector read_convert(const char * name, ggml_type target) { - std::vector f = read_f32(name); - if (f.empty()) return {}; - const int64_t n = (int64_t) f.size(); - if (target == GGML_TYPE_F32) { std::vector o(n * 4); std::memcpy(o.data(), f.data(), o.size()); return o; } - if (target == GGML_TYPE_BF16) { std::vector o(n * 2); ggml_fp32_to_bf16_row(f.data(), reinterpret_cast(o.data()), n); return o; } - std::fprintf(stderr, "vla(gr00tn1d5): unsupported resident type %d for %s\n", (int) target, name); return {}; - } - - bool fetch_rows_f32(const char * name, const std::vector & row_ids, float * dst, int64_t cols) { - const ggml_tensor * t = meta(name); - if (!t || t->ne[0] != cols || t->ne[2] != 1 || t->ne[3] != 1) { std::fprintf(stderr, "vla(gr00tn1d5): %s shape unfit for row-fetch\n", name); return false; } - const int64_t rows = t->ne[1]; - const int64_t id = gguf_find_tensor(gctx, name); - const size_t base = data_off + gguf_get_tensor_offset(gctx, id); - const size_t elsz = (t->type == GGML_TYPE_F32) ? 4u : 2u; - const size_t rb = (size_t) cols * elsz; - std::vector row(rb); - for (size_t k = 0; k < row_ids.size(); ++k) { - const int32_t r = row_ids[k]; - if (r < 0 || r >= rows) { std::fprintf(stderr, "vla(gr00tn1d5): row %d out of range for %s\n", r, name); return false; } - if (std::fseek(fp, (long) (base + (size_t) r * rb), SEEK_SET) != 0) return false; - if (std::fread(row.data(), 1, rb, fp) != rb) return false; - if (elsz == 4) std::memcpy(dst + k * cols, row.data(), rb); - else ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst + k * cols, cols); - } - return true; - } -}; struct SigLipLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; struct Qwen3LayerW { ggml_tensor *attn_norm,*Wq,*Wk,*Wv,*Wo,*q_norm,*k_norm,*ffn_norm,*Wgate,*Wup,*Wdown; }; @@ -137,7 +58,7 @@ struct Gr00tN1d5ModelArch : public ModelArchBase { ggml_backend_t backend = nullptr; bool is_cuda = false; bool is_gpu = false; - int n_threads = 4; + int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; @@ -369,7 +290,7 @@ std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, m->gguf_path = ckpt_path; m->matmul_type = std::getenv("VLA_GR00T_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; - gguf_reader g; + gguf_reader g("gr00tn1d5"); if (!g.open(ckpt_path)) return nullptr; if (!g.has("gr00t_n1_5.architecture")) { std::fprintf(stderr, "vla(gr00tn1d5): %s is not a gr00t_n1_5 GGUF\n", ckpt_path.c_str()); return nullptr; } if (!load_config(g, *m, m->cfg)) return nullptr; @@ -405,7 +326,7 @@ std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); if (!gt) { std::fprintf(stderr, "vla(gr00tn1d5): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, type, ggml_n_dims(gt), gt->ne); + ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); ggml_set_name(t, name); return t; }; auto mk_mm = [&](const char * name) { return mk(name, m->matmul_type); }; @@ -554,7 +475,7 @@ std::vector Gr00tN1d5ModelArch::predict(const Inputs& in) { const int64_t SEQ = (int64_t) input_ids.size(); if (SEQ > max_seq_len) { std::fprintf(stderr, "vla(gr00tn1d5): prompt too long (%lld > %lld)\n", (long long) SEQ, (long long) max_seq_len); return {}; } - gguf_reader g; + gguf_reader g("gr00tn1d5"); if (!g.open(gguf_path)) return {}; std::vector inputs_embeds((size_t) SEQ * H); if (!g.fetch_rows_f32("token_embd.weight", input_ids, inputs_embeds.data(), H)) return {}; diff --git a/src/models/gr00tn1d6.cpp b/src/models/gr00tn1d6.cpp index bcb8a9a..7017d9a 100644 --- a/src/models/gr00tn1d6.cpp +++ b/src/models/gr00tn1d6.cpp @@ -25,6 +25,7 @@ #include "ggml-metal.h" #endif #include "gguf.h" +#include "models/gguf_reader.h" #include #include @@ -41,86 +42,6 @@ namespace vla { namespace { -struct gguf_reader { - gguf_context * gctx = nullptr; - ggml_context * meta_ctx = nullptr; - FILE * fp = nullptr; - size_t data_off = 0; - - bool open(const std::string & path) { - gguf_init_params p{}; - p.no_alloc = true; - p.ctx = &meta_ctx; - gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) { std::fprintf(stderr, "vla(gr00tn1d6): gguf_init_from_file failed for %s\n", path.c_str()); return false; } - fp = std::fopen(path.c_str(), "rb"); - if (!fp) { std::fprintf(stderr, "vla(gr00tn1d6): fopen failed for %s\n", path.c_str()); return false; } - data_off = gguf_get_data_offset(gctx); - return true; - } - ~gguf_reader() { - if (fp) std::fclose(fp); - if (gctx) gguf_free(gctx); - if (meta_ctx) ggml_free(meta_ctx); - } - gguf_reader() = default; - gguf_reader(const gguf_reader &) = delete; - gguf_reader & operator=(const gguf_reader &) = delete; - - bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - double f64(const char * k) const { return gguf_get_val_f64(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { const int64_t id = gguf_find_key(gctx, k); return id < 0 ? std::string() : std::string(gguf_get_val_str(gctx, id)); } - const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } - - bool read_raw(const char * name, void * buf) { - const int64_t id = gguf_find_tensor(gctx, name); - if (id < 0) { std::fprintf(stderr, "vla(gr00tn1d6): missing tensor %s\n", name); return false; } - const size_t off = data_off + gguf_get_tensor_offset(gctx, id); - const size_t nb = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) off, SEEK_SET) != 0) return false; - return std::fread(buf, 1, nb, fp) == nb; - } - std::vector read_f32(const char * name) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(gr00tn1d6): missing tensor %s\n", name); return {}; } - const int64_t n = ggml_nelements(t); - std::vector out(n); - if (t->type == GGML_TYPE_F32) { if (!read_raw(name, out.data())) return {}; } - else if (t->type == GGML_TYPE_BF16) { std::vector tmp(n); if (!read_raw(name, tmp.data())) return {}; ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); } - else { std::fprintf(stderr, "vla(gr00tn1d6): tensor %s unsupported type %d\n", name, (int) t->type); return {}; } - return out; - } - std::vector read_convert(const char * name, ggml_type target) { - std::vector f = read_f32(name); - if (f.empty()) return {}; - const int64_t n = (int64_t) f.size(); - if (target == GGML_TYPE_F32) { std::vector o(n * 4); std::memcpy(o.data(), f.data(), o.size()); return o; } - if (target == GGML_TYPE_BF16) { std::vector o(n * 2); ggml_fp32_to_bf16_row(f.data(), reinterpret_cast(o.data()), n); return o; } - std::fprintf(stderr, "vla(gr00tn1d6): unsupported resident type %d for %s\n", (int) target, name); return {}; - } - - bool fetch_rows_f32(const char * name, const std::vector & row_ids, float * dst, int64_t cols) { - const ggml_tensor * t = meta(name); - if (!t || t->ne[0] != cols || t->ne[2] != 1 || t->ne[3] != 1) { std::fprintf(stderr, "vla(gr00tn1d6): %s shape unfit for row-fetch\n", name); return false; } - const int64_t rows = t->ne[1]; - const int64_t id = gguf_find_tensor(gctx, name); - const size_t base = data_off + gguf_get_tensor_offset(gctx, id); - const size_t elsz = (t->type == GGML_TYPE_F32) ? 4u : 2u; - const size_t rb = (size_t) cols * elsz; - std::vector row(rb); - for (size_t k = 0; k < row_ids.size(); ++k) { - const int32_t r = row_ids[k]; - if (r < 0 || r >= rows) { std::fprintf(stderr, "vla(gr00tn1d6): row %d out of range for %s\n", r, name); return false; } - if (std::fseek(fp, (long) (base + (size_t) r * rb), SEEK_SET) != 0) return false; - if (std::fread(row.data(), 1, rb, fp) != rb) return false; - if (elsz == 4) std::memcpy(dst + k * cols, row.data(), rb); - else ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst + k * cols, cols); - } - return true; - } -}; struct SigLipLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; struct Qwen3LayerW { ggml_tensor *attn_norm,*Wq,*Wk,*Wv,*Wo,*q_norm,*k_norm,*ffn_norm,*Wgate,*Wup,*Wdown; }; @@ -136,7 +57,7 @@ struct Gr00tN1d6ModelArch : public ModelArchBase { ggml_backend_t backend = nullptr; bool is_cuda = false; bool is_gpu = false; - int n_threads = 4; + int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; @@ -374,7 +295,7 @@ std::unique_ptr gr00t_n1_6_create(const std::string& mmproj_path, m->gguf_path = ckpt_path; m->matmul_type = std::getenv("VLA_GR00T_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; - gguf_reader g; + gguf_reader g("gr00tn1d6"); if (!g.open(ckpt_path)) return nullptr; if (!g.has("gr00t_n1_6.architecture")) { std::fprintf(stderr, "vla(gr00tn1d6): %s is not a gr00t_n1_6 GGUF\n", ckpt_path.c_str()); return nullptr; } if (!load_config(g, *m, m->cfg)) return nullptr; @@ -410,7 +331,7 @@ std::unique_ptr gr00t_n1_6_create(const std::string& mmproj_path, auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); if (!gt) { std::fprintf(stderr, "vla(gr00tn1d6): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, type, ggml_n_dims(gt), gt->ne); + ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); ggml_set_name(t, name); return t; }; auto mk_mm = [&](const char * name) { return mk(name, m->matmul_type); }; @@ -582,7 +503,7 @@ std::vector Gr00tN1d6ModelArch::predict(const Inputs& in) { const int64_t SEQ = (int64_t) input_ids.size(); if (SEQ > max_seq_len) { std::fprintf(stderr, "vla(gr00tn1d6): prompt too long (%lld > %lld)\n", (long long) SEQ, (long long) max_seq_len); return {}; } - gguf_reader g; + gguf_reader g("gr00tn1d6"); if (!g.open(gguf_path)) return {}; std::vector inputs_embeds((size_t) SEQ * H); if (!g.fetch_rows_f32("token_embd.weight", input_ids, inputs_embeds.data(), H)) return {}; diff --git a/src/models/gr00tn1d7.cpp b/src/models/gr00tn1d7.cpp index 397ab17..ca107db 100644 --- a/src/models/gr00tn1d7.cpp +++ b/src/models/gr00tn1d7.cpp @@ -25,6 +25,7 @@ #include "ggml-metal.h" #endif #include "gguf.h" +#include "models/gguf_reader.h" #include #include @@ -42,86 +43,6 @@ namespace vla { namespace { -struct gguf_reader { - gguf_context * gctx = nullptr; - ggml_context * meta_ctx = nullptr; - FILE * fp = nullptr; - size_t data_off = 0; - - bool open(const std::string & path) { - gguf_init_params p{}; - p.no_alloc = true; - p.ctx = &meta_ctx; - gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) { std::fprintf(stderr, "vla(gr00tn1d7): gguf_init_from_file failed for %s\n", path.c_str()); return false; } - fp = std::fopen(path.c_str(), "rb"); - if (!fp) { std::fprintf(stderr, "vla(gr00tn1d7): fopen failed for %s\n", path.c_str()); return false; } - data_off = gguf_get_data_offset(gctx); - return true; - } - ~gguf_reader() { - if (fp) std::fclose(fp); - if (gctx) gguf_free(gctx); - if (meta_ctx) ggml_free(meta_ctx); - } - gguf_reader() = default; - gguf_reader(const gguf_reader &) = delete; - gguf_reader & operator=(const gguf_reader &) = delete; - - bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - double f64(const char * k) const { return gguf_get_val_f64(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { const int64_t id = gguf_find_key(gctx, k); return id < 0 ? std::string() : std::string(gguf_get_val_str(gctx, id)); } - const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } - - bool read_raw(const char * name, void * buf) { - const int64_t id = gguf_find_tensor(gctx, name); - if (id < 0) { std::fprintf(stderr, "vla(gr00tn1d7): missing tensor %s\n", name); return false; } - const size_t off = data_off + gguf_get_tensor_offset(gctx, id); - const size_t nb = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) off, SEEK_SET) != 0) return false; - return std::fread(buf, 1, nb, fp) == nb; - } - std::vector read_f32(const char * name) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(gr00tn1d7): missing tensor %s\n", name); return {}; } - const int64_t n = ggml_nelements(t); - std::vector out(n); - if (t->type == GGML_TYPE_F32) { if (!read_raw(name, out.data())) return {}; } - else if (t->type == GGML_TYPE_BF16) { std::vector tmp(n); if (!read_raw(name, tmp.data())) return {}; ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); } - else { std::fprintf(stderr, "vla(gr00tn1d7): tensor %s unsupported type %d\n", name, (int) t->type); return {}; } - return out; - } - std::vector read_convert(const char * name, ggml_type target) { - std::vector f = read_f32(name); - if (f.empty()) return {}; - const int64_t n = (int64_t) f.size(); - if (target == GGML_TYPE_F32) { std::vector o(n * 4); std::memcpy(o.data(), f.data(), o.size()); return o; } - if (target == GGML_TYPE_BF16) { std::vector o(n * 2); ggml_fp32_to_bf16_row(f.data(), reinterpret_cast(o.data()), n); return o; } - std::fprintf(stderr, "vla(gr00tn1d7): unsupported resident type %d for %s\n", (int) target, name); return {}; - } - - bool fetch_rows_f32(const char * name, const std::vector & row_ids, float * dst, int64_t cols) { - const ggml_tensor * t = meta(name); - if (!t || t->ne[0] != cols || t->ne[2] != 1 || t->ne[3] != 1) { std::fprintf(stderr, "vla(gr00tn1d7): %s shape unfit for row-fetch\n", name); return false; } - const int64_t rows = t->ne[1]; - const int64_t id = gguf_find_tensor(gctx, name); - const size_t base = data_off + gguf_get_tensor_offset(gctx, id); - const size_t elsz = (t->type == GGML_TYPE_F32) ? 4u : 2u; - const size_t rb = (size_t) cols * elsz; - std::vector row(rb); - for (size_t k = 0; k < row_ids.size(); ++k) { - const int32_t r = row_ids[k]; - if (r < 0 || r >= rows) { std::fprintf(stderr, "vla(gr00tn1d7): row %d out of range for %s\n", r, name); return false; } - if (std::fseek(fp, (long) (base + (size_t) r * rb), SEEK_SET) != 0) return false; - if (std::fread(row.data(), 1, rb, fp) != rb) return false; - if (elsz == 4) std::memcpy(dst + k * cols, row.data(), rb); - else ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst + k * cols, cols); - } - return true; - } -}; constexpr float CLIP_MEAN[3] = {0.5f, 0.5f, 0.5f}; constexpr float CLIP_STD [3] = {0.5f, 0.5f, 0.5f}; @@ -143,7 +64,7 @@ struct Gr00tN1d7ModelArch : public ModelArchBase { ggml_backend_t backend = nullptr; bool is_cuda = false; bool is_gpu = false; - int n_threads = 4; + int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; @@ -234,6 +155,7 @@ ggml_tensor * head_view(ggml_context * C, ggml_tensor * proj, int64_t hd, int64_ ggml_tensor * flash_attn(ggml_context * C, ggml_tensor * q, ggml_tensor * k, ggml_tensor * v, ggml_tensor * mask, float scale, int64_t hidden) { + (void) hidden; ggml_tensor * kf = (k->type == GGML_TYPE_F16) ? k : ggml_cast(C, k, GGML_TYPE_F16); ggml_tensor * vf = (v->type == GGML_TYPE_F16) ? v : ggml_cast(C, v, GGML_TYPE_F16); ggml_tensor * o = ggml_flash_attn_ext(C, q, kf, vf, mask, scale, 0.0f, 0.0f); @@ -533,7 +455,7 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, m->gguf_path = ckpt_path; m->matmul_type = std::getenv("VLA_GR00T_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; - gguf_reader g; + gguf_reader g("gr00tn1d7"); if (!g.open(ckpt_path)) return nullptr; if (!g.has("gr00t_n1_7.architecture")) { std::fprintf(stderr, "vla(gr00tn1d7): %s is not a gr00t_n1_7 GGUF\n", ckpt_path.c_str()); return nullptr; } if (!load_config(g, *m, m->cfg)) return nullptr; @@ -571,7 +493,7 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); if (!gt) { std::fprintf(stderr, "vla(gr00tn1d7): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, type, ggml_n_dims(gt), gt->ne); + ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); ggml_set_name(t, name); return t; }; auto mk_mm = [&](const char * name) { return mk(name, m->matmul_type); }; diff --git a/src/models/openvla_oft.cpp b/src/models/openvla_oft.cpp index b76549f..c91d6ca 100644 --- a/src/models/openvla_oft.cpp +++ b/src/models/openvla_oft.cpp @@ -14,6 +14,7 @@ #include "arch.h" #include "model.h" +#include "vision_common.h" #include "ggml.h" #include "ggml-cpu.h" @@ -25,6 +26,7 @@ #include "ggml-metal.h" #endif #include "gguf.h" +#include "models/gguf_reader.h" #include #include @@ -40,59 +42,6 @@ namespace vla { namespace { -struct gguf_reader { - gguf_context * gctx = nullptr; - ggml_context * meta_ctx = nullptr; - FILE * fp = nullptr; - size_t data_off = 0; - - bool open(const std::string & path) { - gguf_init_params p{}; p.no_alloc = true; p.ctx = &meta_ctx; - gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) { std::fprintf(stderr, "vla(openvla_oft): gguf_init_from_file failed for %s\n", path.c_str()); return false; } - fp = std::fopen(path.c_str(), "rb"); - if (!fp) { std::fprintf(stderr, "vla(openvla_oft): fopen failed for %s\n", path.c_str()); return false; } - data_off = gguf_get_data_offset(gctx); - return true; - } - ~gguf_reader() { if (fp) std::fclose(fp); if (gctx) gguf_free(gctx); if (meta_ctx) ggml_free(meta_ctx); } - gguf_reader() = default; - gguf_reader(const gguf_reader &) = delete; - gguf_reader & operator=(const gguf_reader &) = delete; - - bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { return gguf_get_val_str(gctx, gguf_find_key(gctx, k)); } - const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } - - bool read_raw(const char * name, void * buf) { - const int64_t id = gguf_find_tensor(gctx, name); - if (id < 0) { std::fprintf(stderr, "vla(openvla_oft): missing tensor %s\n", name); return false; } - const size_t off = data_off + gguf_get_tensor_offset(gctx, id); - const size_t nb = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) off, SEEK_SET) != 0) return false; - return std::fread(buf, 1, nb, fp) == nb; - } - std::vector read_f32(const char * name) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(openvla_oft): missing tensor %s\n", name); return {}; } - const int64_t n = ggml_nelements(t); - std::vector out(n); - if (t->type == GGML_TYPE_F32) { if (!read_raw(name, out.data())) return {}; } - else if (t->type == GGML_TYPE_BF16) { std::vector tmp(n); if (!read_raw(name, tmp.data())) return {}; ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); } - else { std::fprintf(stderr, "vla(openvla_oft): tensor %s unsupported type %d\n", name, (int) t->type); return {}; } - return out; - } - std::vector read_convert(const char * name, ggml_type target) { - std::vector f = read_f32(name); - if (f.empty()) return {}; - const int64_t n = (int64_t) f.size(); - if (target == GGML_TYPE_F32) { std::vector o(n * sizeof(float)); std::memcpy(o.data(), f.data(), o.size()); return o; } - if (target == GGML_TYPE_BF16) { std::vector o(n * sizeof(ggml_bf16_t)); ggml_fp32_to_bf16_row(f.data(), reinterpret_cast(o.data()), n); return o; } - std::fprintf(stderr, "vla(openvla_oft): unsupported resident type %d for %s\n", (int) target, name); return {}; - } -}; bool parse_stats(const std::string & js, int64_t want, std::vector & q01, std::vector & q99, std::vector & mask, std::string & suite) { @@ -147,7 +96,7 @@ struct OpenVlaOftModelArch : public ModelArchBase { } ggml_backend_t backend = nullptr; - bool is_gpu = false; int n_threads = 4; + bool is_gpu = false; int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; ggml_backend_buffer_t weight_buf = nullptr; ggml_type mt = GGML_TYPE_BF16; @@ -225,7 +174,7 @@ std::unique_ptr openvla_oft_create(const std::string& mmproj_path auto m = std::make_unique(); m->mt = std::getenv("VLA_OPENVLA_OFT_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; - gguf_reader g; + gguf_reader g("openvla_oft"); if (!g.open(ckpt_path)) return nullptr; if (!g.has("openvla_oft.architecture")) { std::fprintf(stderr, "vla(openvla_oft): not an openvla_oft GGUF\n"); return nullptr; } @@ -274,7 +223,7 @@ std::unique_ptr openvla_oft_create(const std::string& mmproj_path bool ok = true; auto mk=[&](const char*name, ggml_type ty)->ggml_tensor*{ const ggml_tensor*gt=g.meta(name); if(!gt){ std::fprintf(stderr,"vla(openvla_oft): missing %s\n",name); ok=false; return nullptr; } - ggml_tensor*t=ggml_new_tensor(W,ty,ggml_n_dims(gt),gt->ne); ggml_set_name(t,name); return t; }; + ggml_tensor*t=ggml_new_tensor(W,g.resident_type(gt,ty),ggml_n_dims(gt),gt->ne); ggml_set_name(t,name); return t; }; auto mm=[&](const char*n){ return mk(n,m->mt); }; auto f32=[&](const char*n){ return mk(n,GGML_TYPE_F32); }; @@ -357,7 +306,18 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { stats = Stats{}; const int64_t S=image_size, NP=n_patches, HC=lm_hidden; const int64_t n_views = in.n_images; + if (in.precomputed_img_emb) { std::fprintf(stderr, "vla(openvla_oft): precomputed_img_emb is not supported; the DINOv2+SigLIP tower is baked into the GGUF, pass raw images\n"); return {}; } if (n_views < 1) { std::fprintf(stderr, "vla(openvla_oft): need >=1 image view\n"); return {}; } + if (!in.images) { std::fprintf(stderr, "vla(openvla_oft): n_images=%d but the images pointer is null\n", in.n_images); return {}; } + // towers read S*S*3 per view; reject any view that is not exactly SxS. + for (int64_t v = 0; v < n_views; ++v) { + const ImageView& iv = in.images[v]; + if (!view_is_side(iv.data, iv.w, iv.h, S)) { + std::fprintf(stderr, "vla(openvla_oft): image view %lld is %dx%d, expected %lldx%lld\n", + (long long) v, iv.w, iv.h, (long long) S, (long long) S); + return {}; + } + } static const float DMEAN[3]={0.484375f,0.455078125f,0.40625f}, DSTD[3]={0.228515625f,0.2236328125f,0.224609375f}; static const float SMEAN[3]={0.5f,0.5f,0.5f}, SSTD[3]={0.5f,0.5f,0.5f}; @@ -393,6 +353,16 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { } const int64_t L = in.n_lang; + // ggml_get_rows does not bound-check, so reject out-of-range tokens here. + for (int64_t i = 0; i < L; ++i) + if (in.lang_tokens[i] < 0 || in.lang_tokens[i] >= vocab) { + std::fprintf(stderr, "vla(openvla_oft): token %d out of vocab\n", in.lang_tokens[i]); + return {}; + } + if ((int64_t) stop_id < 0 || (int64_t) stop_id >= vocab) { + std::fprintf(stderr, "vla(openvla_oft): stop_id %lld out of vocab\n", (long long) stop_id); + return {}; + } const int64_t n_act = chunk * action_dim; const int64_t NUM_PATCHES = NPATCH + 1; const int64_t NUM_PROMPT_TOKENS = L - 1; diff --git a/src/models/pi0.cpp b/src/models/pi0.cpp index 53a941f..0c38116 100644 --- a/src/models/pi0.cpp +++ b/src/models/pi0.cpp @@ -15,8 +15,6 @@ #include "arch.h" #include "model.h" -#include "clip.h" -#include "mtmd.h" #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" @@ -28,6 +26,7 @@ #include "ggml-metal.h" #endif #include "gguf.h" +#include "models/gguf_reader.h" #include #include @@ -46,113 +45,6 @@ namespace vla { namespace { -struct gguf_reader { - gguf_context * gctx = nullptr; - ggml_context * meta_ctx = nullptr; - FILE * fp = nullptr; - size_t data_off = 0; - - bool open(const std::string & path) { - gguf_init_params p{}; - p.no_alloc = true; - p.ctx = &meta_ctx; - gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) { - std::fprintf(stderr, "vla(pi0): gguf_init_from_file failed for %s\n", path.c_str()); - return false; - } - fp = std::fopen(path.c_str(), "rb"); - if (!fp) { - std::fprintf(stderr, "vla(pi0): fopen failed for %s\n", path.c_str()); - return false; - } - data_off = gguf_get_data_offset(gctx); - return true; - } - ~gguf_reader() { - if (fp) std::fclose(fp); - if (gctx) gguf_free(gctx); - if (meta_ctx) ggml_free(meta_ctx); - } - gguf_reader() = default; - gguf_reader(const gguf_reader &) = delete; - gguf_reader & operator=(const gguf_reader &) = delete; - - bool has_key(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - double f64(const char * k) const { return gguf_get_val_f64(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { return gguf_get_val_str(gctx, gguf_find_key(gctx, k)); } - - const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } - - bool read_raw(const char * name, void * buf) { - const int64_t id = gguf_find_tensor(gctx, name); - if (id < 0) { std::fprintf(stderr, "vla(pi0): missing tensor %s\n", name); return false; } - const size_t off = data_off + gguf_get_tensor_offset(gctx, id); - const size_t nb = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) off, SEEK_SET) != 0) return false; - return std::fread(buf, 1, nb, fp) == nb; - } - - std::vector read_convert(const char * name, ggml_type target, bool gemma_norm) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(pi0): missing tensor %s\n", name); return {}; } - const int64_t n = ggml_nelements(t); - - std::vector f32(n); - if (t->type == GGML_TYPE_F32) { - if (!read_raw(name, f32.data())) return {}; - } else if (t->type == GGML_TYPE_BF16) { - std::vector tmp(n); - if (!read_raw(name, tmp.data())) return {}; - ggml_bf16_to_fp32_row(tmp.data(), f32.data(), n); - } else { - std::fprintf(stderr, "vla(pi0): tensor %s has unsupported type %d\n", name, (int) t->type); - return {}; - } - if (gemma_norm) for (int64_t i = 0; i < n; ++i) f32[i] += 1.0f; - - if (target == GGML_TYPE_F32) { - std::vector out(n * sizeof(float)); - std::memcpy(out.data(), f32.data(), out.size()); - return out; - } - if (target == GGML_TYPE_BF16) { - std::vector out(n * sizeof(ggml_bf16_t)); - ggml_fp32_to_bf16_row(f32.data(), reinterpret_cast(out.data()), n); - return out; - } - std::fprintf(stderr, "vla(pi0): unsupported resident type %d for %s\n", (int) target, name); - return {}; - } - - bool fetch_rows_f32(const char * name, const std::vector & row_ids, - float * dst, int64_t cols) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(pi0): missing tensor %s\n", name); return false; } - if (t->ne[0] != cols || t->ne[2] != 1 || t->ne[3] != 1) { - std::fprintf(stderr, "vla(pi0): %s shape unfit for row-fetch\n", name); return false; - } - const int64_t rows = t->ne[1]; - const int64_t id = gguf_find_tensor(gctx, name); - const size_t base = data_off + gguf_get_tensor_offset(gctx, id); - const size_t elsz = (t->type == GGML_TYPE_F32) ? 4u : 2u; - const size_t rb = (size_t) cols * elsz; - std::vector row(rb); - for (size_t k = 0; k < row_ids.size(); ++k) { - const int32_t r = row_ids[k]; - if (r < 0 || r >= rows) { - std::fprintf(stderr, "vla(pi0): row %d out of range for %s\n", r, name); return false; - } - if (std::fseek(fp, (long) (base + (size_t) r * rb), SEEK_SET) != 0) return false; - if (std::fread(row.data(), 1, rb, fp) != rb) return false; - if (elsz == 4) std::memcpy(dst + k * cols, row.data(), rb); - else ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst + k * cols, cols); - } - return true; - } -}; struct GemmaLayerW { ggml_tensor * ln_in = nullptr; @@ -166,7 +58,15 @@ struct GemmaLayerW { ggml_tensor * Wdown = nullptr; }; -bool is_gemma_norm(const std::string & name) { return name.find("norm.weight") != std::string::npos; } +// SigLIP-So400m vision block weights (PaliGemma tower, built in-tree like gr00tn1d5). +struct SigLipLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; + +// The +1 RMSNorm fixup is a Gemma quirk that must only touch the language towers, +// never the SigLIP LayerNorms (whose names would otherwise never match anyway). +bool is_gemma_norm(const std::string & name) { + const bool lm = name.rfind("vlm.", 0) == 0 || name.rfind("aex.", 0) == 0; + return lm && name.find("norm.weight") != std::string::npos; +} std::vector sinusoidal_time_emb(double t, int64_t dim, double min_p, double max_p) { const int64_t half = dim / 2; @@ -194,7 +94,6 @@ struct Pi0ModelArch : public ModelArchBase { std::vector predict(const Inputs& in) override; - clip_ctx * cctx = nullptr; ggml_backend_t backend = nullptr; bool is_cuda = false; bool is_gpu = false; @@ -203,6 +102,15 @@ struct Pi0ModelArch : public ModelArchBase { std::string ckpt_path_; ggml_type matmul_type = GGML_TYPE_BF16; + // In-tree SigLIP-So400m/14 vision tower (was llama.cpp clip.cpp mmproj). + int64_t vit_hidden = 1152, vit_layers = 27, vit_heads = 16; + int64_t vit_image_size = 224, vit_patch_size = 14, vit_n_tokens = 256; + float vit_ln_eps = 1e-6f; + ggml_tensor * vit_patch_w = nullptr, * vit_patch_b = nullptr, * vit_pos = nullptr; + ggml_tensor * vit_post_ln_w = nullptr, * vit_post_ln_b = nullptr; + std::vector vit; + ggml_tensor * mm_proj_w = nullptr, * mm_proj_b = nullptr; + std::vector pl_layers; std::vector ex_layers; @@ -222,6 +130,47 @@ struct Pi0ModelArch : public ModelArchBase { namespace { +// One pre-norm SigLIP encoder block, identical to gr00tn1d5's in-tree tower +// (the PaliGemma vision tower is the same SigLIP-So400m/14). Bidirectional +// attention (nullptr mask), F32 score accumulation, tanh GELU FFN. +ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_tensor * x, + int64_t seq, int64_t heads, int64_t head_dim, int64_t hidden, float ln_eps) { + const float scale = 1.0f / std::sqrt((float) head_dim); + ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); + ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n1), w.bq); + ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, n1), w.bk); + ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, n1), w.bv); + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, head_dim, heads, seq), 1, 2, 0, 3)); + ggml_tensor * kq = ggml_mul_mat(C, K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + ggml_tensor * aw = ggml_soft_max_ext(C, kq, nullptr, scale, 0.0f); + ggml_tensor * att = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, ggml_mul_mat(C, V, aw), 0, 2, 1, 3)), hidden, seq); + ggml_tensor * h1 = ggml_add(C, x, ggml_add(C, ggml_mul_mat(C, w.Wo, att), w.bo)); + ggml_tensor * n2 = ggml_add(C, ggml_mul(C, ggml_norm(C, h1, ln_eps), w.ln2w), w.ln2b); + ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wfc2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wfc1, n2), w.bfc1))), w.bfc2); + return ggml_add(C, h1, ff); +} + +// CHW-planar float image in [-1,1] for ggml_conv_2d (SigLIP mean/std 0.5). +bool preprocess_image_chw(const ImageView & v, int64_t side, std::vector & out) { + if (v.w != (int) side || v.h != (int) side || !v.data) { + std::fprintf(stderr, "vla(pi0): image view is %dx%d, expected %lldx%lld\n", + v.w, v.h, (long long) side, (long long) side); + return false; + } + out.assign((size_t) 3 * side * side, 0.0f); + for (int64_t h = 0; h < side; ++h) + for (int64_t w = 0; w < side; ++w) + for (int64_t c = 0; c < 3; ++c) { + float px; + if (v.format == PixelFormat::U8) px = ((const uint8_t *) v.data)[(h * side + w) * 3 + c] / 255.0f; + else px = ((const float *) v.data)[(h * side + w) * 3 + c]; + out[c * side * side + h * side + w] = px * 2.0f - 1.0f; + } + return true; +} + ggml_tensor * build_gemma_layer( ggml_context * ctx, const GemmaLayerW & w, ggml_tensor * x_in, ggml_tensor * positions, @@ -299,7 +248,7 @@ ggml_tensor * build_embed_suffix(ggml_context * ctx, const Pi0ModelArch & m, bool load_config(const gguf_reader & g, Config & cfg) { auto need = [&](const char * k) { - if (!g.has_key(k)) { std::fprintf(stderr, "vla(pi0): gguf missing key %s\n", k); return false; } + if (!g.has(k)) { std::fprintf(stderr, "vla(pi0): gguf missing key %s\n", k); return false; } return true; }; for (const char * k : {"pi0.hidden", "pi0.intermediate", "pi0.n_q_heads", "pi0.n_kv_heads", @@ -333,11 +282,11 @@ bool load_config(const gguf_reader & g, Config & cfg) { cfg.q_full_dim = cfg.n_q_heads * cfg.head_dim; cfg.kv_full_dim = cfg.n_kv_heads * cfg.head_dim; cfg.self_attn_every_n = 0; - cfg.rms_eps = g.has_key("pi0.rms_norm_eps") ? g.f32("pi0.rms_norm_eps") : 1e-6f; - cfg.norm_eps = g.has_key("pi0.norm_eps") ? g.f32("pi0.norm_eps") : 1e-8f; + cfg.rms_eps = g.has("pi0.rms_norm_eps") ? g.f32("pi0.rms_norm_eps") : 1e-6f; + cfg.norm_eps = g.has("pi0.norm_eps") ? g.f32("pi0.norm_eps") : 1e-8f; cfg.rope_mode = GGML_ROPE_TYPE_NEOX; cfg.rope_n_dims = (int) cfg.head_dim; - cfg.rope_freq_base = g.has_key("pi0.rope_theta") ? (float) g.f64("pi0.rope_theta") : 10000.f; + cfg.rope_freq_base = g.has("pi0.rope_theta") ? (float) g.f64("pi0.rope_theta") : 10000.f; cfg.n_prefix = 0; cfg.n_full = 0; return true; @@ -368,7 +317,6 @@ Pi0ModelArch::~Pi0ModelArch() { if (weight_buf) ggml_backend_buffer_free(weight_buf); if (ctx_weights) ggml_free(ctx_weights); if (backend) ggml_backend_free(backend); - if (cctx) clip_free(cctx); } std::unique_ptr pi0_create(const std::string& mmproj_path, @@ -388,9 +336,9 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, m->ckpt_path_ = ckpt_path; m->matmul_type = std::getenv("VLA_PI0_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; - gguf_reader g; + gguf_reader g("pi0"); if (!g.open(ckpt_path)) return nullptr; - if (!g.has_key("pi0.architecture") || g.str("pi0.architecture") != "pi0") { + if (!g.has("pi0.architecture") || g.str("pi0.architecture") != "pi0") { std::fprintf(stderr, "vla(pi0): '%s' is not a π₀ GGUF (pi0.architecture missing/wrong)\n", ckpt_path.c_str()); return nullptr; @@ -426,26 +374,18 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, std::printf("vla(pi0): backend = CPU (%d threads)\n", m->n_threads); } + // The SigLIP tower is now bundled in the ckpt GGUF; mmproj_path is ignored. + (void) mmproj_path; { - clip_context_params cp = {}; - cp.use_gpu = m->is_gpu; - cp.flash_attn_type = m->is_gpu ? CLIP_FLASH_ATTN_TYPE_AUTO : CLIP_FLASH_ATTN_TYPE_DISABLED; - cp.image_min_tokens = -1; - cp.image_max_tokens = -1; - cp.warmup = m->is_gpu; - cp.cb_eval = nullptr; - cp.cb_eval_user_data = nullptr; - clip_init_result r = clip_init(mmproj_path.c_str(), cp); - if (!r.ctx_v) { - std::fprintf(stderr, "vla(pi0): clip_init failed for %s\n", mmproj_path.c_str()); - return nullptr; - } - m->cctx = r.ctx_v; - const int img_sz = clip_get_image_size(m->cctx); - const int mm_embd = clip_n_mmproj_embd(m->cctx); - if (img_sz != 224 || mm_embd != (int) cfg.hidden) { - std::fprintf(stderr, "vla(pi0): mmproj mismatch (image_size=%d mmproj_embd=%d; want 224 / %lld)\n", - img_sz, mm_embd, (long long) cfg.hidden); + auto vu = [&](const char * k, int64_t & d) { if (g.has(k)) d = (int64_t) g.u32(k); }; + vu("pi0.vit_hidden", m->vit_hidden); vu("pi0.vit_layers", m->vit_layers); + vu("pi0.vit_heads", m->vit_heads); vu("pi0.image_size", m->vit_image_size); + vu("pi0.patch_size", m->vit_patch_size); vu("pi0.n_img_tokens", m->vit_n_tokens); + if (g.has("pi0.vit_ln_eps")) m->vit_ln_eps = g.f32("pi0.vit_ln_eps"); + const int64_t grid = m->vit_image_size / m->vit_patch_size; + if (grid * grid != m->vit_n_tokens || m->vit_n_tokens != cfg.n_img) { + std::fprintf(stderr, "vla(pi0): vit geometry mismatch (grid^2=%lld n_img_tokens=%lld cfg.n_img=%lld)\n", + (long long) (grid * grid), (long long) m->vit_n_tokens, (long long) cfg.n_img); return nullptr; } } @@ -461,7 +401,7 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, auto mk = [&](const char * name, ggml_type type, int n_dims, const int64_t * ne) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); if (!gt) { std::fprintf(stderr, "vla(pi0): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, type, n_dims, ne); + ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), n_dims, ne); ggml_set_name(t, name); weights.push_back(t); return t; @@ -493,6 +433,25 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, return lw.ln_in && lw.Wq && lw.Wk && lw.Wv && lw.Wo && lw.ln_post && lw.Wgate && lw.Wup && lw.Wdown; }; + // Vision tower weights (SigLIP-So400m + PaliGemma projector), bundled in the ckpt GGUF. + m->vit_patch_w = mk_f32("vit.patch_embd.weight"); + m->vit_patch_b = mk_f32("vit.patch_embd.bias"); + m->vit_pos = mk_f32("vit.pos_embd"); + m->vit_post_ln_w = mk_f32("vit.post_ln.weight"); + m->vit_post_ln_b = mk_f32("vit.post_ln.bias"); + m->vit.resize(m->vit_layers); + for (int64_t i = 0; i < m->vit_layers; ++i) { + char p[64]; + auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.blk.%lld.%s", (long long) i, s); return (const char *) p; }; + auto & w = m->vit[i]; + w.ln1w=mk_f32(N("ln1.weight")); w.ln1b=mk_f32(N("ln1.bias")); w.ln2w=mk_f32(N("ln2.weight")); w.ln2b=mk_f32(N("ln2.bias")); + w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); + w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); + w.Wfc1=mk_mm(N("fc1.weight")); w.bfc1=mk_f32(N("fc1.bias")); w.Wfc2=mk_mm(N("fc2.weight")); w.bfc2=mk_f32(N("fc2.bias")); + } + m->mm_proj_w = mk_mm("mm.proj.weight"); + m->mm_proj_b = g.meta("mm.proj.bias") ? mk_f32("mm.proj.bias") : nullptr; // PaliGemma projector bias (optional) + m->pl_layers.resize(cfg.n_layers); m->ex_layers.resize(cfg.n_layers); for (int64_t i = 0; i < cfg.n_layers; ++i) { @@ -558,35 +517,48 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { std::fprintf(stderr, "vla(pi0): predict: no images and no precomputed_img_emb\n"); return {}; } - const int img_sz = clip_get_image_size(cctx); - const size_t per_pix = (size_t) 3 * img_sz * img_sz; - const size_t per_out = clip_embd_nbytes_by_img(cctx, img_sz, img_sz) / sizeof(float); - img_emb_host.resize(per_out * (size_t) in.n_images); - std::vector hwc(per_pix); + const int64_t K = vit_n_tokens, H = hidden_pl, grid = vit_image_size / vit_patch_size; + n_img_tokens = (int64_t) in.n_images * K; + img_emb_host.assign((size_t) in.n_images * K * H, 0.0f); + + ggml_init_params vp = { (size_t) 128 * 1024 * 1024, nullptr, true }; + ggml_context * VC = ggml_init(vp); + if (!VC) { std::fprintf(stderr, "vla(pi0): ggml_init(vision ctx) failed\n"); return {}; } + ggml_tensor * t_px = ggml_new_tensor_3d(VC, GGML_TYPE_F32, vit_image_size, vit_image_size, 3); ggml_set_input(t_px); + ggml_tensor * conv = ggml_conv_2d(VC, vit_patch_w, t_px, (int) vit_patch_size, (int) vit_patch_size, 0, 0, 1, 1); + ggml_tensor * patches = ggml_cont(VC, ggml_transpose(VC, ggml_reshape_2d(VC, conv, grid * grid, vit_hidden))); + ggml_tensor * h = ggml_add(VC, ggml_add(VC, patches, vit_patch_b), vit_pos); + for (int64_t i = 0; i < vit_layers; ++i) + h = build_siglip_layer(VC, vit[i], h, K, vit_heads, vit_hidden / vit_heads, vit_hidden, vit_ln_eps); + h = ggml_add(VC, ggml_mul(VC, ggml_norm(VC, h, vit_ln_eps), vit_post_ln_w), vit_post_ln_b); + // PaliGemma projector: linear (+ optional bias), then 1/sqrt(hidden) scale (matches clip.cpp siglip.cpp). + ggml_tensor * proj = ggml_mul_mat(VC, mm_proj_w, h); + if (mm_proj_b) proj = ggml_add(VC, proj, mm_proj_b); + ggml_tensor * vit_emb = ggml_scale(VC, proj, 1.0f / std::sqrt((float) proj->ne[0])); + ggml_set_output(vit_emb); + + ggml_cgraph * vg = ggml_new_graph_custom(VC, 8192, false); + ggml_build_forward_expand(vg, vit_emb); + ggml_gallocr_t vga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!vga || !ggml_gallocr_alloc_graph(vga, vg)) { + std::fprintf(stderr, "vla(pi0): vision gallocr alloc failed\n"); + if (vga) ggml_gallocr_free(vga); + ggml_free(VC); + return {}; + } const auto tv0 = clk::now(); + std::vector chw; for (int v = 0; v < in.n_images; ++v) { - const ImageView & view = in.images[v]; - if (view.w != img_sz || view.h != img_sz) { - std::fprintf(stderr, "vla(pi0): image[%d] is %dx%d; π₀ requires %dx%d\n", - v, view.w, view.h, img_sz, img_sz); - return {}; - } - if (view.format == PixelFormat::U8) { - const uint8_t * src = static_cast(view.data); - for (size_t i = 0; i < per_pix; ++i) hwc[i] = (float) src[i] / 127.5f - 1.0f; - } else { - const float * src = static_cast(view.data); - for (size_t i = 0; i < per_pix; ++i) hwc[i] = src[i] * 2.0f - 1.0f; - } - - if (!clip_encode_float_image(cctx, n_threads, hwc.data(), img_sz, img_sz, - img_emb_host.data() + (size_t) v * per_out)) { - std::fprintf(stderr, "vla(pi0): clip_encode_float_image failed (view %d)\n", v); - return {}; + if (!preprocess_image_chw(in.images[v], vit_image_size, chw)) { ggml_gallocr_free(vga); ggml_free(VC); return {}; } + ggml_backend_tensor_set(t_px, chw.data(), 0, ggml_nbytes(t_px)); + if (ggml_backend_graph_compute(backend, vg) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(pi0): vision compute failed (view %d)\n", v); + ggml_gallocr_free(vga); ggml_free(VC); return {}; } + ggml_backend_tensor_get(vit_emb, img_emb_host.data() + (size_t) v * K * H, 0, ggml_nbytes(vit_emb)); } stats.ms_vision = std::chrono::duration(clk::now() - tv0).count(); - n_img_tokens = (int64_t) ((per_out / (size_t) hidden_pl) * (size_t) in.n_images); + ggml_gallocr_free(vga); ggml_free(VC); } if (in.n_lang < 1 || !in.lang_tokens) { @@ -600,7 +572,7 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { std::vector lang_ids(in.lang_tokens, in.lang_tokens + n_lang); std::vector lang_rows((size_t) n_lang * hidden_pl); { - gguf_reader g; + gguf_reader g("pi0"); if (!g.open(ckpt_path_)) return {}; if (!g.fetch_rows_f32("token_embd.weight", lang_ids, lang_rows.data(), hidden_pl)) return {}; } @@ -724,8 +696,8 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { ggml_free(C); for (int64_t t = 0; t < chunk; ++t) { float * row = out.data() + (size_t) t * max_ad; - for (int64_t j = 0; j < cfg.real_action_dim && j < max_ad; ++j) - row[j] = row[j] * (action_std[j] + cfg.norm_eps) + action_mean[j]; + for (int64_t j = 0; j < max_ad; ++j) + row[j] = j < cfg.real_action_dim ? row[j] * (action_std[j] + cfg.norm_eps) + action_mean[j] : 0.0f; } stats.ms_total = std::chrono::duration(clk::now() - t0).count(); diff --git a/src/models/pi05.cpp b/src/models/pi05.cpp index 2f7e49b..9a30087 100644 --- a/src/models/pi05.cpp +++ b/src/models/pi05.cpp @@ -15,8 +15,6 @@ #include "arch.h" #include "model.h" -#include "clip.h" -#include "mtmd.h" #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" @@ -28,6 +26,7 @@ #include "ggml-metal.h" #endif #include "gguf.h" +#include "models/gguf_reader.h" #include #include @@ -46,113 +45,6 @@ namespace vla { namespace { -struct gguf_reader { - gguf_context * gctx = nullptr; - ggml_context * meta_ctx = nullptr; - FILE * fp = nullptr; - size_t data_off = 0; - - bool open(const std::string & path) { - gguf_init_params p{}; - p.no_alloc = true; - p.ctx = &meta_ctx; - gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) { - std::fprintf(stderr, "vla(pi05): gguf_init_from_file failed for %s\n", path.c_str()); - return false; - } - fp = std::fopen(path.c_str(), "rb"); - if (!fp) { - std::fprintf(stderr, "vla(pi05): fopen failed for %s\n", path.c_str()); - return false; - } - data_off = gguf_get_data_offset(gctx); - return true; - } - ~gguf_reader() { - if (fp) std::fclose(fp); - if (gctx) gguf_free(gctx); - if (meta_ctx) ggml_free(meta_ctx); - } - gguf_reader() = default; - gguf_reader(const gguf_reader &) = delete; - gguf_reader & operator=(const gguf_reader &) = delete; - - bool has_key(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - double f64(const char * k) const { return gguf_get_val_f64(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { return gguf_get_val_str(gctx, gguf_find_key(gctx, k)); } - - const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } - - bool read_raw(const char * name, void * buf) { - const int64_t id = gguf_find_tensor(gctx, name); - if (id < 0) { std::fprintf(stderr, "vla(pi05): missing tensor %s\n", name); return false; } - const size_t off = data_off + gguf_get_tensor_offset(gctx, id); - const size_t nb = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) off, SEEK_SET) != 0) return false; - return std::fread(buf, 1, nb, fp) == nb; - } - - std::vector read_convert(const char * name, ggml_type target, bool gemma_norm) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(pi05): missing tensor %s\n", name); return {}; } - const int64_t n = ggml_nelements(t); - - std::vector f32(n); - if (t->type == GGML_TYPE_F32) { - if (!read_raw(name, f32.data())) return {}; - } else if (t->type == GGML_TYPE_BF16) { - std::vector tmp(n); - if (!read_raw(name, tmp.data())) return {}; - ggml_bf16_to_fp32_row(tmp.data(), f32.data(), n); - } else { - std::fprintf(stderr, "vla(pi05): tensor %s has unsupported type %d\n", name, (int) t->type); - return {}; - } - if (gemma_norm) for (int64_t i = 0; i < n; ++i) f32[i] += 1.0f; - - if (target == GGML_TYPE_F32) { - std::vector out(n * sizeof(float)); - std::memcpy(out.data(), f32.data(), out.size()); - return out; - } - if (target == GGML_TYPE_BF16) { - std::vector out(n * sizeof(ggml_bf16_t)); - ggml_fp32_to_bf16_row(f32.data(), reinterpret_cast(out.data()), n); - return out; - } - std::fprintf(stderr, "vla(pi05): unsupported resident type %d for %s\n", (int) target, name); - return {}; - } - - bool fetch_rows_f32(const char * name, const std::vector & row_ids, - float * dst, int64_t cols) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(pi05): missing tensor %s\n", name); return false; } - if (t->ne[0] != cols || t->ne[2] != 1 || t->ne[3] != 1) { - std::fprintf(stderr, "vla(pi05): %s shape unfit for row-fetch\n", name); return false; - } - const int64_t rows = t->ne[1]; - const int64_t id = gguf_find_tensor(gctx, name); - const size_t base = data_off + gguf_get_tensor_offset(gctx, id); - const size_t elsz = (t->type == GGML_TYPE_F32) ? 4u : 2u; - const size_t rb = (size_t) cols * elsz; - std::vector row(rb); - for (size_t k = 0; k < row_ids.size(); ++k) { - const int32_t r = row_ids[k]; - if (r < 0 || r >= rows) { - std::fprintf(stderr, "vla(pi05): row %d out of range for %s\n", r, name); return false; - } - if (std::fseek(fp, (long) (base + (size_t) r * rb), SEEK_SET) != 0) return false; - if (std::fread(row.data(), 1, rb, fp) != rb) return false; - if (elsz == 4) std::memcpy(dst + k * cols, row.data(), rb); - else ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst + k * cols, cols); - } - return true; - } -}; struct VlmLayerW { ggml_tensor * ln_in = nullptr; @@ -180,6 +72,9 @@ struct ExpertLayerW { ggml_tensor * Wdown = nullptr; }; +// SigLIP-So400m vision block weights (PaliGemma tower, built in-tree like gr00tn1d5). +struct SigLipLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; + std::vector sinusoidal_time_emb(double t, int64_t dim, double min_p, double max_p) { const int64_t half = dim / 2; std::vector out(dim); @@ -214,7 +109,6 @@ struct Pi05ModelArch : public ModelArchBase { std::vector predict(const Inputs& in) override; - clip_ctx * cctx = nullptr; ggml_backend_t backend = nullptr; bool is_cuda = false; bool is_gpu = false; @@ -224,6 +118,15 @@ struct Pi05ModelArch : public ModelArchBase { ggml_type matmul_type = GGML_TYPE_BF16; int64_t adarms_cond_dim = 0; + // In-tree SigLIP-So400m/14 vision tower (was llama.cpp clip.cpp mmproj). + int64_t vit_hidden = 1152, vit_layers = 27, vit_heads = 16; + int64_t vit_image_size = 224, vit_patch_size = 14, vit_n_tokens = 256; + float vit_ln_eps = 1e-6f; + ggml_tensor * vit_patch_w = nullptr, * vit_patch_b = nullptr, * vit_pos = nullptr; + ggml_tensor * vit_post_ln_w = nullptr, * vit_post_ln_b = nullptr; + std::vector vit; + ggml_tensor * mm_proj_w = nullptr, * mm_proj_b = nullptr; + std::vector pl_layers; std::vector ex_layers; @@ -245,6 +148,47 @@ struct Pi05ModelArch : public ModelArchBase { namespace { +// One pre-norm SigLIP encoder block, identical to gr00tn1d5's in-tree tower +// (the PaliGemma vision tower is the same SigLIP-So400m/14). Bidirectional +// attention (nullptr mask), F32 score accumulation, tanh GELU FFN. +ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_tensor * x, + int64_t seq, int64_t heads, int64_t head_dim, int64_t hidden, float ln_eps) { + const float scale = 1.0f / std::sqrt((float) head_dim); + ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); + ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n1), w.bq); + ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, n1), w.bk); + ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, n1), w.bv); + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, head_dim, heads, seq), 1, 2, 0, 3)); + ggml_tensor * kq = ggml_mul_mat(C, K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + ggml_tensor * aw = ggml_soft_max_ext(C, kq, nullptr, scale, 0.0f); + ggml_tensor * att = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, ggml_mul_mat(C, V, aw), 0, 2, 1, 3)), hidden, seq); + ggml_tensor * h1 = ggml_add(C, x, ggml_add(C, ggml_mul_mat(C, w.Wo, att), w.bo)); + ggml_tensor * n2 = ggml_add(C, ggml_mul(C, ggml_norm(C, h1, ln_eps), w.ln2w), w.ln2b); + ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wfc2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wfc1, n2), w.bfc1))), w.bfc2); + return ggml_add(C, h1, ff); +} + +// CHW-planar float image in [-1,1] for ggml_conv_2d (SigLIP mean/std 0.5). +bool preprocess_image_chw(const ImageView & v, int64_t side, std::vector & out) { + if (v.w != (int) side || v.h != (int) side || !v.data) { + std::fprintf(stderr, "vla(pi05): image view is %dx%d, expected %lldx%lld\n", + v.w, v.h, (long long) side, (long long) side); + return false; + } + out.assign((size_t) 3 * side * side, 0.0f); + for (int64_t h = 0; h < side; ++h) + for (int64_t w = 0; w < side; ++w) + for (int64_t c = 0; c < 3; ++c) { + float px; + if (v.format == PixelFormat::U8) px = ((const uint8_t *) v.data)[(h * side + w) * 3 + c] / 255.0f; + else px = ((const float *) v.data)[(h * side + w) * 3 + c]; + out[c * side * side + h * side + w] = px * 2.0f - 1.0f; + } + return true; +} + ggml_tensor * build_vlm_layer( ggml_context * ctx, const VlmLayerW & w, ggml_tensor * x_in, ggml_tensor * positions, @@ -376,7 +320,7 @@ ggml_tensor * build_expert_layer( bool load_config(const gguf_reader & g, Config & cfg) { auto need = [&](const char * k) { - if (!g.has_key(k)) { std::fprintf(stderr, "vla(pi05): gguf missing key %s\n", k); return false; } + if (!g.has(k)) { std::fprintf(stderr, "vla(pi05): gguf missing key %s\n", k); return false; } return true; }; for (const char * k : {"pi05.hidden", "pi05.intermediate", "pi05.n_q_heads", "pi05.n_kv_heads", @@ -410,11 +354,11 @@ bool load_config(const gguf_reader & g, Config & cfg) { cfg.q_full_dim = cfg.n_q_heads * cfg.head_dim; cfg.kv_full_dim = cfg.n_kv_heads * cfg.head_dim; cfg.self_attn_every_n = 0; - cfg.rms_eps = g.has_key("pi05.rms_norm_eps") ? g.f32("pi05.rms_norm_eps") : 1e-6f; - cfg.norm_eps = g.has_key("pi05.norm_eps") ? g.f32("pi05.norm_eps") : 1e-8f; + cfg.rms_eps = g.has("pi05.rms_norm_eps") ? g.f32("pi05.rms_norm_eps") : 1e-6f; + cfg.norm_eps = g.has("pi05.norm_eps") ? g.f32("pi05.norm_eps") : 1e-8f; cfg.rope_mode = GGML_ROPE_TYPE_NEOX; cfg.rope_n_dims = (int) cfg.head_dim; - cfg.rope_freq_base = g.has_key("pi05.rope_theta") ? (float) g.f64("pi05.rope_theta") : 10000.f; + cfg.rope_freq_base = g.has("pi05.rope_theta") ? (float) g.f64("pi05.rope_theta") : 10000.f; cfg.n_prefix = 0; cfg.n_full = 0; return true; @@ -452,7 +396,6 @@ Pi05ModelArch::~Pi05ModelArch() { if (weight_buf) ggml_backend_buffer_free(weight_buf); if (ctx_weights) ggml_free(ctx_weights); if (backend) ggml_backend_free(backend); - if (cctx) clip_free(cctx); } std::unique_ptr pi05_create(const std::string& mmproj_path, @@ -471,17 +414,17 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, m->ckpt_path_ = ckpt_path; m->matmul_type = std::getenv("VLA_PI05_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; - gguf_reader g; + gguf_reader g("pi05"); if (!g.open(ckpt_path)) return nullptr; - if (!g.has_key("pi05.architecture") || g.str("pi05.architecture") != "pi05") { + if (!g.has("pi05.architecture") || g.str("pi05.architecture") != "pi05") { std::fprintf(stderr, "vla(pi05): '%s' is not a π0.5 GGUF (pi05.architecture missing/wrong)\n", ckpt_path.c_str()); return nullptr; } if (!load_config(g, m->cfg)) return nullptr; const Config & cfg = m->cfg; - m->adarms_cond_dim = g.has_key("pi05.adarms_cond_dim") ? g.u32("pi05.adarms_cond_dim") : cfg.expert_h; - m->quantile_norm = g.has_key("pi05.norm_mode") && g.str("pi05.norm_mode") == "quantiles"; + m->adarms_cond_dim = g.has("pi05.adarms_cond_dim") ? g.u32("pi05.adarms_cond_dim") : cfg.expert_h; + m->quantile_norm = g.has("pi05.norm_mode") && g.str("pi05.norm_mode") == "quantiles"; std::printf("vla(pi05): hidden=%lld inter=%lld heads=%lldq/%lldkv x%lld n_layers=%lld " "expert_h=%lld expert_inter=%lld chunk=%lld steps=%d real_state=%lld real_action=%lld " "max_len=%lld adarms_cond=%lld matmul_weights=%s\n", @@ -512,26 +455,18 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, std::printf("vla(pi05): backend = CPU (%d threads)\n", m->n_threads); } + // The SigLIP tower is now bundled in the ckpt GGUF; mmproj_path is ignored. + (void) mmproj_path; { - clip_context_params cp = {}; - cp.use_gpu = m->is_cuda; - cp.flash_attn_type = m->is_cuda ? CLIP_FLASH_ATTN_TYPE_AUTO : CLIP_FLASH_ATTN_TYPE_DISABLED; - cp.image_min_tokens = -1; - cp.image_max_tokens = -1; - cp.warmup = m->is_cuda; - cp.cb_eval = nullptr; - cp.cb_eval_user_data = nullptr; - clip_init_result r = clip_init(mmproj_path.c_str(), cp); - if (!r.ctx_v) { - std::fprintf(stderr, "vla(pi05): clip_init failed for %s\n", mmproj_path.c_str()); - return nullptr; - } - m->cctx = r.ctx_v; - const int img_sz = clip_get_image_size(m->cctx); - const int mm_embd = clip_n_mmproj_embd(m->cctx); - if (img_sz != 224 || mm_embd != (int) cfg.hidden) { - std::fprintf(stderr, "vla(pi05): mmproj mismatch (image_size=%d mmproj_embd=%d; want 224 / %lld)\n", - img_sz, mm_embd, (long long) cfg.hidden); + auto vu = [&](const char * k, int64_t & d) { if (g.has(k)) d = (int64_t) g.u32(k); }; + vu("pi05.vit_hidden", m->vit_hidden); vu("pi05.vit_layers", m->vit_layers); + vu("pi05.vit_heads", m->vit_heads); vu("pi05.image_size", m->vit_image_size); + vu("pi05.patch_size", m->vit_patch_size); vu("pi05.n_img_tokens", m->vit_n_tokens); + if (g.has("pi05.vit_ln_eps")) m->vit_ln_eps = g.f32("pi05.vit_ln_eps"); + const int64_t grid = m->vit_image_size / m->vit_patch_size; + if (grid * grid != m->vit_n_tokens || m->vit_n_tokens != cfg.n_img) { + std::fprintf(stderr, "vla(pi05): vit geometry mismatch (grid^2=%lld n_img_tokens=%lld cfg.n_img=%lld)\n", + (long long) (grid * grid), (long long) m->vit_n_tokens, (long long) cfg.n_img); return nullptr; } } @@ -547,7 +482,7 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, auto mk = [&](const char * name, ggml_type type, int n_dims, const int64_t * ne) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); if (!gt) { std::fprintf(stderr, "vla(pi05): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, type, n_dims, ne); + ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), n_dims, ne); ggml_set_name(t, name); weights.push_back(t); return t; @@ -595,6 +530,25 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, lw.ada_post_w && lw.ada_post_b && lw.Wgate && lw.Wup && lw.Wdown; }; + // Vision tower weights (SigLIP-So400m + PaliGemma projector), bundled in the ckpt GGUF. + m->vit_patch_w = mk_f32("vit.patch_embd.weight"); + m->vit_patch_b = mk_f32("vit.patch_embd.bias"); + m->vit_pos = mk_f32("vit.pos_embd"); + m->vit_post_ln_w = mk_f32("vit.post_ln.weight"); + m->vit_post_ln_b = mk_f32("vit.post_ln.bias"); + m->vit.resize(m->vit_layers); + for (int64_t i = 0; i < m->vit_layers; ++i) { + char p[64]; + auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.blk.%lld.%s", (long long) i, s); return (const char *) p; }; + auto & w = m->vit[i]; + w.ln1w=mk_f32(N("ln1.weight")); w.ln1b=mk_f32(N("ln1.bias")); w.ln2w=mk_f32(N("ln2.weight")); w.ln2b=mk_f32(N("ln2.bias")); + w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); + w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); + w.Wfc1=mk_mm(N("fc1.weight")); w.bfc1=mk_f32(N("fc1.bias")); w.Wfc2=mk_mm(N("fc2.weight")); w.bfc2=mk_f32(N("fc2.bias")); + } + m->mm_proj_w = mk_mm("mm.proj.weight"); + m->mm_proj_b = g.meta("mm.proj.bias") ? mk_f32("mm.proj.bias") : nullptr; // PaliGemma projector bias (optional) + m->pl_layers.resize(cfg.n_layers); m->ex_layers.resize(cfg.n_layers); for (int64_t i = 0; i < cfg.n_layers; ++i) { @@ -659,35 +613,51 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { std::fprintf(stderr, "vla(pi05): predict: no images and no precomputed_img_emb\n"); return {}; } - const int img_sz = clip_get_image_size(cctx); - const size_t per_pix = (size_t) 3 * img_sz * img_sz; - const size_t per_out = clip_embd_nbytes_by_img(cctx, img_sz, img_sz) / sizeof(float); - img_emb_host.resize(per_out * (size_t) in.n_images); - std::vector hwc(per_pix); + const int64_t K = vit_n_tokens, H = hidden_pl, grid = vit_image_size / vit_patch_size; + n_img_tokens = (int64_t) in.n_images * K; + img_emb_host.assign((size_t) in.n_images * K * H, 0.0f); + + ggml_init_params vp = { (size_t) 128 * 1024 * 1024, nullptr, true }; + ggml_context * VC = ggml_init(vp); + if (!VC) { std::fprintf(stderr, "vla(pi05): ggml_init(vision ctx) failed\n"); return {}; } + ggml_tensor * t_px = ggml_new_tensor_3d(VC, GGML_TYPE_F32, vit_image_size, vit_image_size, 3); ggml_set_input(t_px); + ggml_tensor * conv = ggml_conv_2d(VC, vit_patch_w, t_px, (int) vit_patch_size, (int) vit_patch_size, 0, 0, 1, 1); + ggml_tensor * patches = ggml_cont(VC, ggml_transpose(VC, ggml_reshape_2d(VC, conv, grid * grid, vit_hidden))); + ggml_tensor * h = ggml_add(VC, ggml_add(VC, patches, vit_patch_b), vit_pos); + for (int64_t i = 0; i < vit_layers; ++i) + h = build_siglip_layer(VC, vit[i], h, K, vit_heads, vit_hidden / vit_heads, vit_hidden, vit_ln_eps); + h = ggml_add(VC, ggml_mul(VC, ggml_norm(VC, h, vit_ln_eps), vit_post_ln_w), vit_post_ln_b); + // PaliGemma projector: linear (+ optional bias), then 1/sqrt(hidden) scale (matches clip.cpp siglip.cpp). + ggml_tensor * proj = ggml_mul_mat(VC, mm_proj_w, h); + if (mm_proj_b) proj = ggml_add(VC, proj, mm_proj_b); + ggml_tensor * vit_emb = ggml_scale(VC, proj, 1.0f / std::sqrt((float) proj->ne[0])); + ggml_set_output(vit_emb); + + ggml_cgraph * vg = ggml_new_graph_custom(VC, 8192, false); + ggml_build_forward_expand(vg, vit_emb); + ggml_gallocr_t vga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!vga || !ggml_gallocr_alloc_graph(vga, vg)) { + std::fprintf(stderr, "vla(pi05): vision gallocr alloc failed\n"); + if (vga) ggml_gallocr_free(vga); + ggml_free(VC); + return {}; + } const auto tv0 = clk::now(); + std::vector chw; for (int v = 0; v < in.n_images; ++v) { - const ImageView & view = in.images[v]; - if (view.w != img_sz || view.h != img_sz) { - std::fprintf(stderr, "vla(pi05): image[%d] is %dx%d; π0.5 requires %dx%d\n", - v, view.w, view.h, img_sz, img_sz); - return {}; - } - if (view.format == PixelFormat::U8) { - const uint8_t * src = static_cast(view.data); - for (size_t i = 0; i < per_pix; ++i) hwc[i] = (float) src[i] / 127.5f - 1.0f; - } else { - const float * src = static_cast(view.data); - for (size_t i = 0; i < per_pix; ++i) hwc[i] = src[i] * 2.0f - 1.0f; - } - if (!clip_encode_float_image(cctx, n_threads, hwc.data(), img_sz, img_sz, - img_emb_host.data() + (size_t) v * per_out)) { - std::fprintf(stderr, "vla(pi05): clip_encode_float_image failed (view %d)\n", v); - return {}; + if (!preprocess_image_chw(in.images[v], vit_image_size, chw)) { ggml_gallocr_free(vga); ggml_free(VC); return {}; } + ggml_backend_tensor_set(t_px, chw.data(), 0, ggml_nbytes(t_px)); + if (ggml_backend_graph_compute(backend, vg) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(pi05): vision compute failed (view %d)\n", v); + ggml_gallocr_free(vga); ggml_free(VC); return {}; } + ggml_backend_tensor_get(vit_emb, img_emb_host.data() + (size_t) v * K * H, 0, ggml_nbytes(vit_emb)); } stats.ms_vision = std::chrono::duration(clk::now() - tv0).count(); - n_img_tokens = (int64_t) ((per_out / (size_t) hidden_pl) * (size_t) in.n_images); + ggml_gallocr_free(vga); ggml_free(VC); + // π0.5's image tokens are the raw PaliGemma projector features: this undoes + // the 1/sqrt(hidden) the shared vision graph applies (π0 keeps them scaled). const float img_scale = (float) std::sqrt((double) hidden_pl); for (float & x : img_emb_host) x *= img_scale; } @@ -698,12 +668,11 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { } const int64_t n_lang = in.n_lang; const int64_t n_prefix = n_img_tokens + n_lang; - const int64_t n_total = n_prefix + n_suf; std::vector lang_ids(in.lang_tokens, in.lang_tokens + n_lang); std::vector lang_rows((size_t) n_lang * hidden_pl); { - gguf_reader g; + gguf_reader g("pi05"); if (!g.open(ckpt_path_)) return {}; if (!g.fetch_rows_f32("token_embd.weight", lang_ids, lang_rows.data(), hidden_pl)) return {}; } @@ -805,8 +774,10 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { if (!std::getenv("VLA_PI05_SKIP_UNNORM")) { for (int64_t t = 0; t < chunk; ++t) { float * row = out.data() + (size_t) t * max_ad; - for (int64_t j = 0; j < cfg.real_action_dim && j < max_ad; ++j) { - if (quantile_norm) + for (int64_t j = 0; j < max_ad; ++j) { + if (j >= cfg.real_action_dim) + row[j] = 0.0f; + else if (quantile_norm) row[j] = (row[j] + 1.0f) * (action_q99[j] - action_q01[j]) * 0.5f + action_q01[j]; else row[j] = row[j] * (action_std[j] + cfg.norm_eps) + action_mean[j]; diff --git a/src/models/smolvla.cpp b/src/models/smolvla.cpp index 3e99e73..357e00d 100644 --- a/src/models/smolvla.cpp +++ b/src/models/smolvla.cpp @@ -17,9 +17,8 @@ #include "arch.h" #include "model.h" +#include "vision_common.h" -#include "clip.h" -#include "mtmd.h" #include "ggml.h" #include "ggml-backend.h" #include "ggml-cpu.h" @@ -270,6 +269,9 @@ struct ExpertLayerW { ggml_tensor * Wdown; }; +// SigLIP-B/16 vision block weights (SmolVLM2 tower, built in-tree). +struct SigLipLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; + } struct SmolVLAModelArch : public ModelArchBase { @@ -278,7 +280,13 @@ struct SmolVLAModelArch : public ModelArchBase { std::vector predict(const Inputs& in) override; - clip_ctx * cctx = nullptr; + // In-tree SigLIP-B/16 vision tower (was llama.cpp clip.cpp mmproj). + int64_t vit_hidden = 768, vit_layers = 12, vit_heads = 12, vit_inter = 3072; + int64_t vit_patch = 16, vit_image = 512, vit_scale = 4, vit_n_tokens = 64; + float vit_ln_eps = 1e-6f; + ggml_tensor * vit_patch_w = nullptr, * vit_patch_b = nullptr, * vit_pos = nullptr; + ggml_tensor * vit_post_ln_w = nullptr, * vit_post_ln_b = nullptr, * mm_fc = nullptr; + std::vector vit; ggml_backend_t backend = nullptr; ggml_backend_buffer_t weight_buf = nullptr; @@ -337,6 +345,46 @@ struct SmolVLAModelArch : public ModelArchBase { namespace { +// One pre-norm SigLIP encoder block (SmolVLM2 tower), same graph as the other +// in-tree models. Bidirectional attention, F32 score accumulation, tanh GELU. +ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_tensor * x, + int64_t seq, int64_t heads, int64_t head_dim, int64_t hidden, float ln_eps) { + const float scale = 1.0f / std::sqrt((float) head_dim); + ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); + ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n1), w.bq); + ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, n1), w.bk); + ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, n1), w.bv); + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, head_dim, heads, seq), 1, 2, 0, 3)); + ggml_tensor * kq = ggml_mul_mat(C, K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + ggml_tensor * aw = ggml_soft_max_ext(C, kq, nullptr, scale, 0.0f); + ggml_tensor * att = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, ggml_mul_mat(C, V, aw), 0, 2, 1, 3)), hidden, seq); + ggml_tensor * h1 = ggml_add(C, x, ggml_add(C, ggml_mul_mat(C, w.Wo, att), w.bo)); + ggml_tensor * n2 = ggml_add(C, ggml_mul(C, ggml_norm(C, h1, ln_eps), w.ln2w), w.ln2b); + ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wfc2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wfc1, n2), w.bfc1))), w.bfc2); + return ggml_add(C, h1, ff); +} + +// CHW-planar float image in [-1,1] for ggml_conv_2d (SigLIP mean/std 0.5). +bool preprocess_image_chw(const ImageView & v, int64_t side, std::vector & out) { + if (v.w != (int) side || v.h != (int) side || !v.data) { + std::fprintf(stderr, "vla(smolvla): image view is %dx%d, expected %lldx%lld\n", + v.w, v.h, (long long) side, (long long) side); + return false; + } + out.assign((size_t) 3 * side * side, 0.0f); + for (int64_t h = 0; h < side; ++h) + for (int64_t w = 0; w < side; ++w) + for (int64_t c = 0; c < 3; ++c) { + float px; + if (v.format == PixelFormat::U8) px = ((const uint8_t *) v.data)[(h * side + w) * 3 + c] / 255.0f; + else px = ((const float *) v.data)[(h * side + w) * 3 + c]; + out[c * side * side + h * side + w] = px * 2.0f - 1.0f; + } + return true; +} + bool load_config_from_json(const std::string & path, Config & cfg) { std::ifstream f(path); if (!f) { @@ -599,6 +647,45 @@ std::string hf_to_gguf(const std::string & n) { if (starts_with(n, AEX_LAYER_PFX)) { return layer_translate(n.substr(std::strlen(AEX_LAYER_PFX)), "aex"); } + + static const char * VIS_PFX = "model.vlm_with_expert.vlm.model.vision_model."; + if (n == "model.vlm_with_expert.vlm.model.connector.modality_projection.proj.weight") + return "mm.fc.weight"; + if (starts_with(n, VIS_PFX)) { + const std::string rest = n.substr(std::strlen(VIS_PFX)); + if (rest == "embeddings.patch_embedding.weight") return "vit.patch_embd.weight"; + if (rest == "embeddings.patch_embedding.bias") return "vit.patch_embd.bias"; + if (rest == "embeddings.position_embedding.weight") return "vit.pos_embd"; + if (rest == "post_layernorm.weight") return "vit.post_ln.weight"; + if (rest == "post_layernorm.bias") return "vit.post_ln.bias"; + if (starts_with(rest, "encoder.layers.")) { + const size_t e = rest.find('.', 15); + if (e == std::string::npos) return n; + const std::string idx = rest.substr(15, e - 15); + const std::string suf = rest.substr(e + 1); + std::string ds; + if (suf == "layer_norm1.weight") ds = "ln1.weight"; + else if (suf == "layer_norm1.bias") ds = "ln1.bias"; + else if (suf == "layer_norm2.weight") ds = "ln2.weight"; + else if (suf == "layer_norm2.bias") ds = "ln2.bias"; + else if (suf == "self_attn.q_proj.weight") ds = "attn_q.weight"; + else if (suf == "self_attn.q_proj.bias") ds = "attn_q.bias"; + else if (suf == "self_attn.k_proj.weight") ds = "attn_k.weight"; + else if (suf == "self_attn.k_proj.bias") ds = "attn_k.bias"; + else if (suf == "self_attn.v_proj.weight") ds = "attn_v.weight"; + else if (suf == "self_attn.v_proj.bias") ds = "attn_v.bias"; + else if (suf == "self_attn.out_proj.weight") ds = "attn_o.weight"; + else if (suf == "self_attn.out_proj.bias") ds = "attn_o.bias"; + else if (suf == "mlp.fc1.weight") ds = "fc1.weight"; + else if (suf == "mlp.fc1.bias") ds = "fc1.bias"; + else if (suf == "mlp.fc2.weight") ds = "fc2.weight"; + else if (suf == "mlp.fc2.bias") ds = "fc2.bias"; + else return n; + return "vit.blk." + idx + "." + ds; + } + return n; + } + if (starts_with(n, MODEL_PFX)) { return n.substr(std::strlen(MODEL_PFX)); @@ -713,9 +800,20 @@ ggml_tensor * build_expert_self_attn_layer( return ggml_add(ctx, h1, mm_w(ctx, w.Wdown, inter)); } +// Reproject the constant prefix cache into a cross-attn layer's K/V. Depends only +// on the prefix, so it is built once and shared across all denoise steps. +void expert_cross_kv(ggml_context * ctx, const ExpertLayerW & w, + ggml_tensor * cached_K, ggml_tensor * cached_V, const Config & cfg, + ggml_tensor ** K_repro, ggml_tensor ** V_repro) { + ggml_tensor * cK_flat = ggml_reshape_2d(ctx, cached_K, cfg.kv_full_dim, cfg.n_prefix); + ggml_tensor * cV_flat = ggml_reshape_2d(ctx, cached_V, cfg.kv_full_dim, cfg.n_prefix); + *K_repro = ggml_reshape_3d(ctx, mm_w(ctx, w.Wk, cK_flat), cfg.head_dim, cfg.n_kv_heads, cfg.n_prefix); + *V_repro = ggml_reshape_3d(ctx, mm_w(ctx, w.Wv, cV_flat), cfg.head_dim, cfg.n_kv_heads, cfg.n_prefix); +} + ggml_tensor * build_expert_cross_attn_layer( ggml_context * ctx, const ExpertLayerW & w, - ggml_tensor * x_in, ggml_tensor * cached_K, ggml_tensor * cached_V, + ggml_tensor * x_in, ggml_tensor * K_repro, ggml_tensor * V_repro, ggml_tensor * positions_rebased, ggml_tensor * mask_prefix_only, const Config & cfg) { @@ -725,13 +823,6 @@ ggml_tensor * build_expert_cross_attn_layer( ggml_tensor * q_h = ggml_reshape_3d(ctx, q_proj, cfg.head_dim, cfg.n_q_heads, cfg.n_suffix); ggml_tensor * q_rope = rope_q_or_k(ctx, q_h, positions_rebased, cfg); - ggml_tensor * cK_flat = ggml_reshape_2d(ctx, cached_K, cfg.kv_full_dim, cfg.n_prefix); - ggml_tensor * cV_flat = ggml_reshape_2d(ctx, cached_V, cfg.kv_full_dim, cfg.n_prefix); - ggml_tensor * K_repro = ggml_reshape_3d(ctx, mm_w(ctx, w.Wk, cK_flat), - cfg.head_dim, cfg.n_kv_heads, cfg.n_prefix); - ggml_tensor * V_repro = ggml_reshape_3d(ctx, mm_w(ctx, w.Wv, cV_flat), - cfg.head_dim, cfg.n_kv_heads, cfg.n_prefix); - ggml_tensor * Q = ggml_permute(ctx, q_rope, 0, 2, 1, 3); ggml_tensor * Kp = ggml_permute(ctx, K_repro, 0, 2, 1, 3); ggml_tensor * Vp = ggml_permute(ctx, V_repro, 0, 2, 1, 3); @@ -861,54 +952,35 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, delete m; return nullptr; } - ggml_backend_cpu_set_n_threads(m->backend, 4); - std::printf("vla: backend = CPU (4 threads)\n"); + ggml_backend_cpu_set_n_threads(m->backend, default_cpu_threads()); + std::printf("vla: backend = CPU (%d threads)\n", default_cpu_threads()); } vram_probe(m->backend, "after backend init"); m->weight_dtype = resolve_weight_dtype(); std::printf("vla: tower weights resident as %s\n", ggml_type_name(m->weight_dtype)); - clip_context_params cparams = {}; - cparams.use_gpu = m->is_gpu; - - cparams.flash_attn_type = CLIP_FLASH_ATTN_TYPE_AUTO; - cparams.image_min_tokens = -1; - cparams.image_max_tokens = -1; - cparams.warmup = true; - cparams.cb_eval = nullptr; - cparams.cb_eval_user_data= nullptr; - auto init_res = clip_init(mmproj_path.c_str(), cparams); - if (!init_res.ctx_v) { - std::fprintf(stderr, "vla: clip_init failed for %s\n", mmproj_path.c_str()); - ggml_backend_free(m->backend); - delete m; - return nullptr; - } - m->cctx = init_res.ctx_v; - vram_probe(m->backend, "after clip init"); - - if (clip_n_mmproj_embd(m->cctx) != static_cast(m->cfg.hidden)) { - std::fprintf(stderr, "vla: clip hidden dim mismatch\n"); - clip_free(m->cctx); - ggml_backend_free(m->backend); - delete m; - return nullptr; + // Vision tower geometry: from gguf KV (self-contained ckpt), else SmolVLM2-500M defaults. + (void) mmproj_path; + if (use_gguf) { + auto vu = [&](const char * k, int64_t & d) { if (gst.has_key(k)) d = (int64_t) gst.get_u32(k); }; + vu("smolvla.vit_hidden", m->vit_hidden); vu("smolvla.vit_layers", m->vit_layers); + vu("smolvla.vit_heads", m->vit_heads); vu("smolvla.patch_size", m->vit_patch); + vu("smolvla.image_size", m->vit_image); vu("smolvla.vit_pixel_shuffle", m->vit_scale); + vu("smolvla.n_img_tokens", m->vit_n_tokens); vu("smolvla.vit_inter", m->vit_inter); + if (gst.has_key("smolvla.vit_ln_eps")) m->vit_ln_eps = gst.get_f32("smolvla.vit_ln_eps"); } - { - const int img_size = clip_get_image_size(m->cctx); - const size_t img_emb_n = - clip_embd_nbytes_by_img(m->cctx, img_size, img_size) / sizeof(float); - if (img_emb_n % size_t(m->cfg.hidden) != 0) { - std::fprintf(stderr, "vla: clip embedding %zu not a multiple of hidden %lld\n", - img_emb_n, (long long) m->cfg.hidden); - clip_free(m->cctx); + const int64_t grid = m->vit_image / m->vit_patch; + const int64_t k = grid / m->vit_scale; + if (k * k != m->vit_n_tokens) { + std::fprintf(stderr, "vla: smolvla vit geometry mismatch (grid=%lld scale=%lld -> %lld tokens, KV says %lld)\n", + (long long) grid, (long long) m->vit_scale, (long long) (k * k), (long long) m->vit_n_tokens); ggml_backend_free(m->backend); delete m; return nullptr; } - m->cfg.n_img = static_cast(img_emb_n / size_t(m->cfg.hidden)); + m->cfg.n_img = m->vit_n_tokens; } m->cfg.n_prefix = m->cfg.n_img + m->cfg.n_lang + m->cfg.n_state; m->cfg.n_full = m->cfg.n_prefix + m->cfg.n_suffix; @@ -916,7 +988,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, if (!use_gguf) { if (!st.open(ckpt_path)) { std::fprintf(stderr, "vla: failed to open %s\n", ckpt_path.c_str()); - clip_free(m->cctx); ggml_backend_free(m->backend); delete m; return nullptr; @@ -932,7 +1003,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, } if (max_layer < 0) { std::fprintf(stderr, "vla: cannot infer n_layers from %s\n", ckpt_path.c_str()); - clip_free(m->cctx); ggml_backend_free(m->backend); delete m; return nullptr; @@ -943,7 +1013,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, const auto it = st.tensors.find("model.vlm_with_expert.lm_expert.layers.0.mlp.gate_proj.weight"); if (it == st.tensors.end() || it->second.shape.size() != 2) { std::fprintf(stderr, "vla: missing/malformed expert gate_proj for shape derivation\n"); - clip_free(m->cctx); ggml_backend_free(m->backend); delete m; return nullptr; @@ -954,7 +1023,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, std::fprintf(stderr, "vla: expert_h mismatch - config implies %lld, " "checkpoint gate_proj has %lld\n", (long long) m->cfg.expert_h, (long long) it->second.shape[1]); - clip_free(m->cctx); ggml_backend_free(m->backend); delete m; return nullptr; @@ -975,7 +1043,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, if (use_gguf) { if (!load_normalizer_stats_from_gguf(gst, *m)) { std::fprintf(stderr, "vla: failed to load normalizer stats from gguf\n"); - clip_free(m->cctx); ggml_backend_free(m->backend); delete m; return nullptr; @@ -997,7 +1064,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, m->ctx_weights = ggml_init(gparams); if (!m->ctx_weights) { std::fprintf(stderr, "vla: ggml_init (weights) failed\n"); - clip_free(m->cctx); ggml_backend_free(m->backend); delete m; return nullptr; @@ -1021,6 +1087,49 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, pending_f32.push_back({"model.state_proj.weight", m->Wstate, {cfg.hidden, cfg.max_state_dim}}); pending_f32.push_back({"model.state_proj.bias", m->bstate, {cfg.hidden}}); + // Vision tower weights (SigLIP-B/16 encoder + single-linear pixel-shuffle connector). + { + const int64_t H = m->vit_hidden, FF = m->vit_inter, P = m->vit_patch; + const int64_t grid = m->vit_image / P, n_patches = grid * grid; + const int64_t c4 = H * m->vit_scale * m->vit_scale; + const char * VP = "model.vlm_with_expert.vlm.model.vision_model."; + m->vit_patch_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, P, P, 3, H); + m->vit_patch_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + m->vit_pos = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, H, n_patches); + m->vit_post_ln_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + m->vit_post_ln_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + pending_f32.push_back({std::string(VP) + "embeddings.patch_embedding.weight", m->vit_patch_w, {H, 3, P, P}}); + pending_f32.push_back({std::string(VP) + "embeddings.patch_embedding.bias", m->vit_patch_b, {H}}); + pending_f32.push_back({std::string(VP) + "embeddings.position_embedding.weight", m->vit_pos, {n_patches, H}}); + pending_f32.push_back({std::string(VP) + "post_layernorm.weight", m->vit_post_ln_w, {H}}); + pending_f32.push_back({std::string(VP) + "post_layernorm.bias", m->vit_post_ln_b, {H}}); + m->vit.resize(m->vit_layers); + for (int64_t i = 0; i < m->vit_layers; ++i) { + SigLipLayerW & w = m->vit[i]; + char pb[256]; std::snprintf(pb, sizeof(pb), "%sencoder.layers.%lld.", VP, (long long) i); + const std::string pf = pb; + w.ln1w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); w.ln1b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + w.ln2w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); w.ln2b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + w.Wq = ggml_new_tensor_2d(ctx, wdt, H, H); w.bq = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + w.Wk = ggml_new_tensor_2d(ctx, wdt, H, H); w.bk = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + w.Wv = ggml_new_tensor_2d(ctx, wdt, H, H); w.bv = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + w.Wo = ggml_new_tensor_2d(ctx, wdt, H, H); w.bo = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + w.Wfc1 = ggml_new_tensor_2d(ctx, wdt, H, FF); w.bfc1 = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, FF); + w.Wfc2 = ggml_new_tensor_2d(ctx, wdt, FF, H); w.bfc2 = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + pending_f32.push_back({pf + "layer_norm1.weight", w.ln1w, {H}}); pending_f32.push_back({pf + "layer_norm1.bias", w.ln1b, {H}}); + pending_f32.push_back({pf + "layer_norm2.weight", w.ln2w, {H}}); pending_f32.push_back({pf + "layer_norm2.bias", w.ln2b, {H}}); + pending_f32.push_back({pf + "self_attn.q_proj.weight", w.Wq, {H, H}}); pending_f32.push_back({pf + "self_attn.q_proj.bias", w.bq, {H}}); + pending_f32.push_back({pf + "self_attn.k_proj.weight", w.Wk, {H, H}}); pending_f32.push_back({pf + "self_attn.k_proj.bias", w.bk, {H}}); + pending_f32.push_back({pf + "self_attn.v_proj.weight", w.Wv, {H, H}}); pending_f32.push_back({pf + "self_attn.v_proj.bias", w.bv, {H}}); + pending_f32.push_back({pf + "self_attn.out_proj.weight", w.Wo, {H, H}}); pending_f32.push_back({pf + "self_attn.out_proj.bias", w.bo, {H}}); + pending_f32.push_back({pf + "mlp.fc1.weight", w.Wfc1, {FF, H}}); pending_f32.push_back({pf + "mlp.fc1.bias", w.bfc1, {FF}}); + pending_f32.push_back({pf + "mlp.fc2.weight", w.Wfc2, {H, FF}}); pending_f32.push_back({pf + "mlp.fc2.bias", w.bfc2, {H}}); + } + m->mm_fc = ggml_new_tensor_2d(ctx, wdt, c4, cfg.hidden); + pending_f32.push_back({"model.vlm_with_expert.vlm.model.connector.modality_projection.proj.weight", + m->mm_fc, {cfg.hidden, c4}}); + } + m->vlm_layers.resize(cfg.n_layers); for (int i = 0; i < cfg.n_layers; ++i) { VlmLayerW & w = m->vlm_layers[i]; @@ -1250,6 +1359,15 @@ bool build_compute_graph(SmolVLAModelArch* m, int n_views) { } } + // Reproject each cross-attn layer's prefix K/V once; reused every denoise step. + std::vector xk_cache(cfg.n_layers, nullptr); + std::vector xv_cache(cfg.n_layers, nullptr); + for (int li = 0; li < cfg.n_layers; ++li) { + if (!m->expert_layers[li].is_self_attn) + expert_cross_kv(ctx, m->expert_layers[li], k_cache[li], v_cache[li], + cfg_built, &xk_cache[li], &xv_cache[li]); + } + const float dt = -1.f / static_cast(cfg.num_steps); ggml_tensor * x_t = x0; @@ -1269,7 +1387,7 @@ bool build_compute_graph(SmolVLAModelArch* m, int n_views) { pos_full, mask_full_f16, cfg_built); } else { h = build_expert_cross_attn_layer(ctx, m->expert_layers[li], h, - k_cache[li], v_cache[li], + xk_cache[li], xv_cache[li], pos_rebased, mask_pfx_only_f16, cfg_built); } } @@ -1327,7 +1445,6 @@ SmolVLAModelArch::~SmolVLAModelArch() { if (weight_buf) ggml_backend_buffer_free(weight_buf); if (ctx_weights) ggml_free(ctx_weights); if (backend) ggml_backend_free(backend); - if (cctx) clip_free(cctx); } namespace { @@ -1339,7 +1456,6 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { Config cfg = m->cfg; - const int img_size = clip_get_image_size(m->cctx); const size_t per_view_n = size_t(m->cfg.n_img * cfg.hidden); int n_views = 0; @@ -1361,49 +1477,72 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { std::fprintf(stderr, "vla: at least one image is required\n"); return {}; } - const size_t per_view_clip = - clip_embd_nbytes_by_img(m->cctx, img_size, img_size) / sizeof(float); - if (per_view_clip != per_view_n) { - std::fprintf(stderr, "vla: clip per-view embedding size mismatch\n"); - return {}; - } n_views = in.n_images; img_emb_n = per_view_n * size_t(n_views); img_emb_pre.resize(img_emb_n); - const size_t per_image = static_cast(img_size) * img_size * 3; - std::vector img_f32_all(per_image * size_t(n_views)); - for (int v = 0; v < in.n_images; ++v) { - const ImageView & view = in.images[v]; - if (view.w != img_size || view.h != img_size) { - std::fprintf(stderr, "vla: image[%d] size %dx%d != model image_size %dx%d\n", - v, view.w, view.h, img_size, img_size); - return {}; - } + const int64_t H = m->vit_hidden, grid = m->vit_image / m->vit_patch, n_patches = grid * grid; + const int64_t s = m->vit_scale, c4 = H * s * s, K = m->vit_n_tokens; + const auto t_vision_begin = clock::now(); - float * dst = img_f32_all.data() + size_t(v) * per_image; - if (view.format == PixelFormat::U8) { - const uint8_t * src = static_cast(view.data); - for (size_t i = 0; i < per_image; ++i) { - dst[i] = static_cast(src[i]) / 127.5f - 1.0f; - } - } else { - const float * src = static_cast(view.data); - for (size_t i = 0; i < per_image; ++i) { - dst[i] = src[i] * 2.0f - 1.0f; - } - } + // Graph A: SigLIP ViT (conv patch-embed -> +pos -> layers -> post_ln), plain sequential positions. + ggml_init_params vpA = { size_t(256) * 1024 * 1024, nullptr, true }; + ggml_context * VC = ggml_init(vpA); + if (!VC) { std::fprintf(stderr, "vla(smolvla): ggml_init(vision ctx) failed\n"); return {}; } + ggml_tensor * t_px = ggml_new_tensor_3d(VC, GGML_TYPE_F32, m->vit_image, m->vit_image, 3); ggml_set_input(t_px); + ggml_tensor * conv = ggml_conv_2d(VC, m->vit_patch_w, t_px, (int) m->vit_patch, (int) m->vit_patch, 0, 0, 1, 1); + ggml_tensor * patches = ggml_cont(VC, ggml_transpose(VC, ggml_reshape_2d(VC, conv, n_patches, H))); + ggml_tensor * hv = ggml_add(VC, ggml_add(VC, patches, m->vit_patch_b), m->vit_pos); + for (int64_t i = 0; i < m->vit_layers; ++i) + hv = build_siglip_layer(VC, m->vit[i], hv, n_patches, m->vit_heads, H / m->vit_heads, H, m->vit_ln_eps); + ggml_tensor * post_ln = ggml_add(VC, ggml_mul(VC, ggml_norm(VC, hv, m->vit_ln_eps), m->vit_post_ln_w), m->vit_post_ln_b); + ggml_set_output(post_ln); + ggml_cgraph * gA = ggml_new_graph_custom(VC, 8192, false); + ggml_build_forward_expand(gA, post_ln); + ggml_gallocr_t vgA = ggml_gallocr_new(ggml_backend_get_default_buffer_type(m->backend)); + if (!vgA || !ggml_gallocr_alloc_graph(vgA, gA)) { + std::fprintf(stderr, "vla(smolvla): vision gallocr A alloc failed\n"); + if (vgA) ggml_gallocr_free(vgA); + ggml_free(VC); + return {}; } - const auto t_vision_begin = clock::now(); - if (!clip_encode_float_images(m->cctx, 4, - img_f32_all.data(), img_size, img_size, - n_views, img_emb_pre.data())) { - std::fprintf(stderr, "vla: clip_encode_float_images failed (n_views=%d)\n", - n_views); + + // Graph B: pixel-shuffle connector, a single bias-free matmul (c4 -> hidden). + ggml_init_params vpB = { size_t(64) * 1024 * 1024, nullptr, true }; + ggml_context * MC = ggml_init(vpB); + if (!MC) { std::fprintf(stderr, "vla(smolvla): ggml_init(connector ctx) failed\n"); ggml_gallocr_free(vgA); ggml_free(VC); return {}; } + ggml_tensor * t_shuf = ggml_new_tensor_2d(MC, GGML_TYPE_F32, c4, K); ggml_set_input(t_shuf); + ggml_tensor * img_embeds = ggml_mul_mat(MC, m->mm_fc, t_shuf); + ggml_set_output(img_embeds); + ggml_cgraph * gB = ggml_new_graph(MC); + ggml_build_forward_expand(gB, img_embeds); + ggml_gallocr_t vgB = ggml_gallocr_new(ggml_backend_get_default_buffer_type(m->backend)); + if (!vgB || !ggml_gallocr_alloc_graph(vgB, gB)) { + std::fprintf(stderr, "vla(smolvla): vision gallocr B alloc failed\n"); + if (vgB) ggml_gallocr_free(vgB); + ggml_gallocr_free(vgA); ggml_free(MC); ggml_free(VC); return {}; } - m->stats.ms_vision = std::chrono::duration( - clock::now() - t_vision_begin).count(); + + std::vector chw, post_host((size_t) H * n_patches), shuf_host((size_t) c4 * K); + bool vok = true; + for (int v = 0; v < n_views && vok; ++v) { + if (!preprocess_image_chw(in.images[v], m->vit_image, chw)) { vok = false; break; } + ggml_backend_tensor_set(t_px, chw.data(), 0, ggml_nbytes(t_px)); + if (ggml_backend_graph_compute(m->backend, gA) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(smolvla): vision compute A failed (view %d)\n", v); vok = false; break; + } + ggml_backend_tensor_get(post_ln, post_host.data(), 0, ggml_nbytes(post_ln)); + pixel_shuffle_hf(post_host.data(), shuf_host.data(), H, grid, s); + ggml_backend_tensor_set(t_shuf, shuf_host.data(), 0, ggml_nbytes(t_shuf)); + if (ggml_backend_graph_compute(m->backend, gB) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(smolvla): connector compute failed (view %d)\n", v); vok = false; break; + } + ggml_backend_tensor_get(img_embeds, img_emb_pre.data() + size_t(v) * per_view_n, 0, ggml_nbytes(img_embeds)); + } + ggml_gallocr_free(vgB); ggml_free(MC); ggml_gallocr_free(vgA); ggml_free(VC); + if (!vok) return {}; + m->stats.ms_vision = std::chrono::duration(clock::now() - t_vision_begin).count(); } cfg.n_img = m->cfg.n_img * int64_t(n_views); @@ -1416,6 +1555,18 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { return {}; } + // The language tokens index E_lang via ggml_get_rows, which does not bound + // its indices. Reject any token id outside the embedding table before the + // gather so an out-of-range id cannot read past the weights. + const int64_t vocab_rows = m->E_lang ? m->E_lang->ne[1] : 0; + for (int i = 0; i < in.n_lang; ++i) { + if (in.lang_tokens[i] < 0 || in.lang_tokens[i] >= vocab_rows) { + std::fprintf(stderr, "vla: lang_tokens[%d]=%d out of vocab range [0, %lld)\n", + i, in.lang_tokens[i], (long long) vocab_rows); + return {}; + } + } + if (in.timing_detail == TimingDetail::NONE) { if (m->gf_cached == nullptr || m->cached_n_views != n_views) { @@ -1521,8 +1672,8 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { ggml_backend_tensor_get(m->out_x_t, out.data(), 0, out.size() * sizeof(float)); for (int64_t r = 0; r < cfg.n_suffix; ++r) { float * row = out.data() + r * cfg.max_action_dim; - for (int64_t j = 0; j < cfg.real_action_dim; ++j) { - row[j] = row[j] * (m->action_std[j] + cfg.norm_eps) + m->action_mean[j]; + for (int64_t j = 0; j < cfg.max_action_dim; ++j) { + row[j] = j < cfg.real_action_dim ? row[j] * (m->action_std[j] + cfg.norm_eps) + m->action_mean[j] : 0.0f; } } @@ -1648,6 +1799,13 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { std::vector time_bcasts(cfg.num_steps, nullptr); std::vector> time_host (cfg.num_steps); + std::vector xk_cache(cfg.n_layers, nullptr); + std::vector xv_cache(cfg.n_layers, nullptr); + for (int li = 0; li < cfg.n_layers; ++li) { + if (!m->expert_layers[li].is_self_attn) + expert_cross_kv(ctx, m->expert_layers[li], K_ref[li], V_ref[li], cfg, &xk_cache[li], &xv_cache[li]); + } + for (int step = 0; step < cfg.num_steps; ++step) { const double time = 1.0 + static_cast(step) * static_cast(dt); time_host[step] = sinusoidal_time_emb(time, cfg.expert_h, cfg.min_period, cfg.max_period); @@ -1670,7 +1828,7 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { pos_full, mask_full_f16, cfg); } else { h = build_expert_cross_attn_layer(ctx, m->expert_layers[li], h, - K_ref[li], V_ref[li], + xk_cache[li], xv_cache[li], pos_rebased, mask_prefix_only_f16, cfg); } } @@ -1751,8 +1909,8 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { for (int64_t r = 0; r < cfg.n_suffix; ++r) { float * row = out.data() + r * cfg.max_action_dim; - for (int64_t j = 0; j < cfg.real_action_dim; ++j) { - row[j] = row[j] * (m->action_std[j] + cfg.norm_eps) + m->action_mean[j]; + for (int64_t j = 0; j < cfg.max_action_dim; ++j) { + row[j] = j < cfg.real_action_dim ? row[j] * (m->action_std[j] + cfg.norm_eps) + m->action_mean[j] : 0.0f; } } diff --git a/src/models/vision_common.h b/src/models/vision_common.h new file mode 100644 index 0000000..35ad863 --- /dev/null +++ b/src/models/vision_common.h @@ -0,0 +1,49 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Small pure vision helpers shared by the in-tree towers, split out so they can +// be unit-tested without a model or a GPU. + +#pragma once + +#include +#include + +namespace vla { + +// A tower that reads side*side*3 from a view needs the view to be exactly that +// size with real data, else it runs past the buffer. +inline bool view_is_side(const void * data, int w, int h, int64_t side) { + return data != nullptr && (int64_t) w == side && (int64_t) h == side; +} + +// IDEFICS3/SmolVLM pixel-shuffle (space-to-depth), c-innermost channel order. +// src [embed, n_patches] row-major (patch p, channel e) -> dst [embed*s^2, (grid/s)^2]. +inline void pixel_shuffle_hf(const float * src, float * dst, + int64_t embed, int64_t grid, int64_t s) { + const int64_t g2 = grid / s, c4 = embed * s * s; + for (int64_t h2 = 0; h2 < g2; ++h2) + for (int64_t w2 = 0; w2 < g2; ++w2) { + const int64_t t = h2 * g2 + w2; + for (int64_t hs = 0; hs < s; ++hs) + for (int64_t ws = 0; ws < s; ++ws) { + const int64_t p = (h2 * s + hs) * grid + (w2 * s + ws); + const int64_t base = (hs * s + ws) * embed; + std::memcpy(dst + t * c4 + base, src + p * embed, + (size_t) embed * sizeof(float)); + } + } +} + +} // namespace vla diff --git a/src/models/vla_adapter.cpp b/src/models/vla_adapter.cpp index 156df17..fca8d0d 100644 --- a/src/models/vla_adapter.cpp +++ b/src/models/vla_adapter.cpp @@ -14,6 +14,7 @@ #include "arch.h" #include "model.h" +#include "vision_common.h" #include "ggml.h" #include "ggml-cpu.h" @@ -25,6 +26,7 @@ #include "ggml-metal.h" #endif #include "gguf.h" +#include "models/gguf_reader.h" #include #include @@ -40,59 +42,6 @@ namespace vla { namespace { -struct gguf_reader { - gguf_context * gctx = nullptr; - ggml_context * meta_ctx = nullptr; - FILE * fp = nullptr; - size_t data_off = 0; - - bool open(const std::string & path) { - gguf_init_params p{}; p.no_alloc = true; p.ctx = &meta_ctx; - gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) { std::fprintf(stderr, "vla(vla_adapter): gguf_init_from_file failed for %s\n", path.c_str()); return false; } - fp = std::fopen(path.c_str(), "rb"); - if (!fp) { std::fprintf(stderr, "vla(vla_adapter): fopen failed for %s\n", path.c_str()); return false; } - data_off = gguf_get_data_offset(gctx); - return true; - } - ~gguf_reader() { if (fp) std::fclose(fp); if (gctx) gguf_free(gctx); if (meta_ctx) ggml_free(meta_ctx); } - gguf_reader() = default; - gguf_reader(const gguf_reader &) = delete; - gguf_reader & operator=(const gguf_reader &) = delete; - - bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { return gguf_get_val_str(gctx, gguf_find_key(gctx, k)); } - const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } - - bool read_raw(const char * name, void * buf) { - const int64_t id = gguf_find_tensor(gctx, name); - if (id < 0) { std::fprintf(stderr, "vla(vla_adapter): missing tensor %s\n", name); return false; } - const size_t off = data_off + gguf_get_tensor_offset(gctx, id); - const size_t nb = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) off, SEEK_SET) != 0) return false; - return std::fread(buf, 1, nb, fp) == nb; - } - std::vector read_f32(const char * name) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(vla_adapter): missing tensor %s\n", name); return {}; } - const int64_t n = ggml_nelements(t); - std::vector out(n); - if (t->type == GGML_TYPE_F32) { if (!read_raw(name, out.data())) return {}; } - else if (t->type == GGML_TYPE_BF16) { std::vector tmp(n); if (!read_raw(name, tmp.data())) return {}; ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); } - else { std::fprintf(stderr, "vla(vla_adapter): tensor %s unsupported type %d\n", name, (int) t->type); return {}; } - return out; - } - std::vector read_convert(const char * name, ggml_type target) { - std::vector f = read_f32(name); - if (f.empty()) return {}; - const int64_t n = (int64_t) f.size(); - if (target == GGML_TYPE_F32) { std::vector o(n * sizeof(float)); std::memcpy(o.data(), f.data(), o.size()); return o; } - if (target == GGML_TYPE_BF16) { std::vector o(n * sizeof(ggml_bf16_t)); ggml_fp32_to_bf16_row(f.data(), reinterpret_cast(o.data()), n); return o; } - std::fprintf(stderr, "vla(vla_adapter): unsupported resident type %d for %s\n", (int) target, name); return {}; - } -}; bool parse_stats(const std::string & js, int64_t want, std::vector & q01, std::vector & q99, std::vector & mask, std::string & suite) { @@ -148,7 +97,7 @@ struct VlaAdapterModelArch : public ModelArchBase { } ggml_backend_t backend = nullptr; - bool is_gpu = false; int n_threads = 4; + bool is_gpu = false; int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; ggml_backend_buffer_t weight_buf = nullptr; ggml_type mt = GGML_TYPE_BF16; @@ -239,7 +188,7 @@ std::unique_ptr vla_adapter_create(const std::string& mmproj_path auto m = std::make_unique(); m->mt = std::getenv("VLA_ADAPTER_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; - gguf_reader g; + gguf_reader g("vla_adapter"); if (!g.open(ckpt_path)) return nullptr; if (!g.has("vla_adapter.architecture")) { std::fprintf(stderr, "vla(vla_adapter): not a vla_adapter GGUF\n"); return nullptr; } @@ -288,7 +237,7 @@ std::unique_ptr vla_adapter_create(const std::string& mmproj_path bool ok = true; auto mk=[&](const char*name, ggml_type ty)->ggml_tensor*{ const ggml_tensor*gt=g.meta(name); if(!gt){ std::fprintf(stderr,"vla(vla_adapter): missing %s\n",name); ok=false; return nullptr; } - ggml_tensor*t=ggml_new_tensor(W,ty,ggml_n_dims(gt),gt->ne); ggml_set_name(t,name); return t; }; + ggml_tensor*t=ggml_new_tensor(W,g.resident_type(gt,ty),ggml_n_dims(gt),gt->ne); ggml_set_name(t,name); return t; }; auto mm=[&](const char*n){ return mk(n,m->mt); }; auto f32=[&](const char*n){ return mk(n,GGML_TYPE_F32); }; char nm[96]; @@ -379,7 +328,18 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { stats = Stats{}; const int64_t S=image_size, NP=n_patches, HC=lm_hidden, HD=head_dim, NH=head_heads; const int64_t n_views = in.n_images; + if (in.precomputed_img_emb) { std::fprintf(stderr, "vla(vla_adapter): precomputed_img_emb is not supported; the DINOv2+SigLIP tower is baked into the GGUF, pass raw images\n"); return {}; } if (n_views < 1) { std::fprintf(stderr, "vla(vla_adapter): need >=1 image view\n"); return {}; } + if (!in.images) { std::fprintf(stderr, "vla(vla_adapter): n_images=%d but the images pointer is null\n", in.n_images); return {}; } + // towers read S*S*3 per view; reject any view that is not exactly SxS. + for (int64_t v = 0; v < n_views; ++v) { + const ImageView& iv = in.images[v]; + if (!view_is_side(iv.data, iv.w, iv.h, S)) { + std::fprintf(stderr, "vla(vla_adapter): image view %lld is %dx%d, expected %lldx%lld\n", + (long long) v, iv.w, iv.h, (long long) S, (long long) S); + return {}; + } + } static const float DMEAN[3]={0.484375f,0.455078125f,0.40625f}, DSTD[3]={0.228515625f,0.2236328125f,0.224609375f}; static const float SMEAN[3]={0.5f,0.5f,0.5f}, SSTD[3]={0.5f,0.5f,0.5f}; @@ -415,6 +375,16 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { } const int64_t NPROMPT = in.n_lang; + // ggml_get_rows does not bound-check, so reject out-of-range tokens here. + for (int64_t i = 0; i < NPROMPT; ++i) + if (in.lang_tokens[i] < 0 || in.lang_tokens[i] >= vocab) { + std::fprintf(stderr, "vla(vla_adapter): token %d out of vocab\n", in.lang_tokens[i]); + return {}; + } + if ((int64_t) stop_id < 0 || (int64_t) stop_id >= vocab) { + std::fprintf(stderr, "vla(vla_adapter): stop_id %lld out of vocab\n", (long long) stop_id); + return {}; + } const int64_t NUM_PROMPT_TOKENS = NPROMPT - 1; const int64_t NPATCH = NP * n_views; const int64_t SEQ = 1 + NPATCH + (NPROMPT-1) + num_tokens + 1; @@ -510,7 +480,8 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { { std::vector ids(NPROMPT+num_tokens+1); for(int64_t i=0;i pp(SEQ); for(int64_t i=0;i #include @@ -42,86 +43,6 @@ namespace vla { namespace { -struct gguf_reader { - gguf_context * gctx = nullptr; - ggml_context * meta_ctx = nullptr; - FILE * fp = nullptr; - size_t data_off = 0; - - bool open(const std::string & path) { - gguf_init_params p{}; - p.no_alloc = true; - p.ctx = &meta_ctx; - gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) { std::fprintf(stderr, "vla(vla_jepa): gguf_init_from_file failed for %s\n", path.c_str()); return false; } - fp = std::fopen(path.c_str(), "rb"); - if (!fp) { std::fprintf(stderr, "vla(vla_jepa): fopen failed for %s\n", path.c_str()); return false; } - data_off = gguf_get_data_offset(gctx); - return true; - } - ~gguf_reader() { - if (fp) std::fclose(fp); - if (gctx) gguf_free(gctx); - if (meta_ctx) ggml_free(meta_ctx); - } - gguf_reader() = default; - gguf_reader(const gguf_reader &) = delete; - gguf_reader & operator=(const gguf_reader &) = delete; - - bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - double f64(const char * k) const { return gguf_get_val_f64(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { const int64_t id = gguf_find_key(gctx, k); return id < 0 ? std::string() : std::string(gguf_get_val_str(gctx, id)); } - const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } - - bool read_raw(const char * name, void * buf) { - const int64_t id = gguf_find_tensor(gctx, name); - if (id < 0) { std::fprintf(stderr, "vla(vla_jepa): missing tensor %s\n", name); return false; } - const size_t off = data_off + gguf_get_tensor_offset(gctx, id); - const size_t nb = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) off, SEEK_SET) != 0) return false; - return std::fread(buf, 1, nb, fp) == nb; - } - std::vector read_f32(const char * name) { - const ggml_tensor * t = meta(name); - if (!t) { std::fprintf(stderr, "vla(vla_jepa): missing tensor %s\n", name); return {}; } - const int64_t n = ggml_nelements(t); - std::vector out(n); - if (t->type == GGML_TYPE_F32) { if (!read_raw(name, out.data())) return {}; } - else if (t->type == GGML_TYPE_BF16) { std::vector tmp(n); if (!read_raw(name, tmp.data())) return {}; ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); } - else { std::fprintf(stderr, "vla(vla_jepa): tensor %s unsupported type %d\n", name, (int) t->type); return {}; } - return out; - } - std::vector read_convert(const char * name, ggml_type target) { - std::vector f = read_f32(name); - if (f.empty()) return {}; - const int64_t n = (int64_t) f.size(); - if (target == GGML_TYPE_F32) { std::vector o(n * 4); std::memcpy(o.data(), f.data(), o.size()); return o; } - if (target == GGML_TYPE_BF16) { std::vector o(n * 2); ggml_fp32_to_bf16_row(f.data(), reinterpret_cast(o.data()), n); return o; } - std::fprintf(stderr, "vla(vla_jepa): unsupported resident type %d for %s\n", (int) target, name); return {}; - } - - bool fetch_rows_f32(const char * name, const std::vector & row_ids, float * dst, int64_t cols) { - const ggml_tensor * t = meta(name); - if (!t || t->ne[0] != cols || t->ne[2] != 1 || t->ne[3] != 1) { std::fprintf(stderr, "vla(vla_jepa): %s shape unfit for row-fetch\n", name); return false; } - const int64_t rows = t->ne[1]; - const int64_t id = gguf_find_tensor(gctx, name); - const size_t base = data_off + gguf_get_tensor_offset(gctx, id); - const size_t elsz = (t->type == GGML_TYPE_F32) ? 4u : 2u; - const size_t rb = (size_t) cols * elsz; - std::vector row(rb); - for (size_t k = 0; k < row_ids.size(); ++k) { - const int32_t r = row_ids[k]; - if (r < 0 || r >= rows) { std::fprintf(stderr, "vla(vla_jepa): row %d out of range for %s\n", r, name); return false; } - if (std::fseek(fp, (long) (base + (size_t) r * rb), SEEK_SET) != 0) return false; - if (std::fread(row.data(), 1, rb, fp) != rb) return false; - if (elsz == 4) std::memcpy(dst + k * cols, row.data(), rb); - else ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst + k * cols, cols); - } - return true; - } -}; constexpr float CLIP_MEAN[3] = {0.5f, 0.5f, 0.5f}; constexpr float CLIP_STD [3] = {0.5f, 0.5f, 0.5f}; @@ -141,7 +62,7 @@ struct VlaJepaModelArch : public ModelArchBase { ggml_backend_t backend = nullptr; bool is_cuda = false; bool is_gpu = false; - int n_threads = 4; + int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; @@ -416,7 +337,7 @@ std::unique_ptr vla_jepa_create(const std::string& mmproj_path, m->gguf_path = ckpt_path; m->matmul_type = std::getenv("VLA_JEPA_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; - gguf_reader g; + gguf_reader g("vla_jepa"); if (!g.open(ckpt_path)) return nullptr; if (!g.has("vla_jepa.architecture")) { std::fprintf(stderr, "vla(vla_jepa): %s is not a vla_jepa GGUF\n", ckpt_path.c_str()); return nullptr; } if (!load_config(g, *m, m->cfg)) return nullptr; @@ -450,7 +371,7 @@ std::unique_ptr vla_jepa_create(const std::string& mmproj_path, auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); if (!gt) { std::fprintf(stderr, "vla(vla_jepa): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, type, ggml_n_dims(gt), gt->ne); + ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); ggml_set_name(t, name); return t; }; auto mk_mm = [&](const char * name) { return mk(name, m->matmul_type); }; @@ -592,8 +513,12 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { std::fclose(fp); std::printf("vla(vla_jepa): conditioning injected from %s (action-head isolation)\n", cond_file); } else { - - int64_t n_views = in.n_images > 0 ? in.n_images : in.n_img_views; + if (in.precomputed_img_emb) { + std::fprintf(stderr, "vla(vla_jepa): precomputed_img_emb is not supported. The V-JEPA tower also " + "emits deepstack features that a single embedding buffer cannot carry; pass raw images.\n"); + return {}; + } + int64_t n_views = in.n_images; if (n_views <= 0) { std::fprintf(stderr, "vla(vla_jepa): no images in the request\n"); return {}; } std::vector img_emb_host((size_t) n_views * K * H), ds_host[3]; for (int j = 0; j < 3; ++j) ds_host[j].assign((size_t) n_views * K * H, 0.0f); @@ -607,6 +532,7 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { std::fclose(fp); std::printf("vla(vla_jepa): pixel_values injected from %s\n", patches_file); } + if (inj_patches.empty() && !in.images) { std::fprintf(stderr, "vla(vla_jepa): n_images=%d but the images pointer is null\n", in.n_images); return {}; } ggml_init_params vp = { (size_t) 512 * 1024 * 1024, nullptr, true }; ggml_context * VC = ggml_init(vp); diff --git a/src/serving/server.cpp b/src/serving/server.cpp index 55d7d8d..1cecaad 100644 --- a/src/serving/server.cpp +++ b/src/serving/server.cpp @@ -17,11 +17,17 @@ #define STB_IMAGE_IMPLEMENTATION #define STB_IMAGE_STATIC +// stb ships its full implementation here; silence its unused-function noise so +// our own -Wall -Wextra output stays meaningful. +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wunused-function" #include "stb_image.h" +#pragma GCC diagnostic pop #include #include +#include #include #include #include @@ -35,6 +41,10 @@ std::atomic g_shutdown{false}; void on_signal(int) { g_shutdown.store(true, std::memory_order_relaxed); } +// Reject absurd image dimensions before any size arithmetic, so an untrusted +// width or height cannot overflow size_t or truncate to a negative int. +constexpr unsigned kMaxImageDim = 8192; + bool decode_image(const vla::Image & img, std::vector & u8, std::vector & f32, @@ -42,6 +52,11 @@ bool decode_image(const vla::Image & img, if (img.encoding() == vla::Image::JPEG) { int w = 0, h = 0, ch = 0; const auto & data = img.data(); + if (data.size() > size_t(INT_MAX)) { + std::fprintf(stderr, "vla-server: JPEG payload too large (%zu bytes)\n", + data.size()); + return false; + } unsigned char * px = stbi_load_from_memory( reinterpret_cast(data.data()), static_cast(data.size()), @@ -51,11 +66,23 @@ bool decode_image(const vla::Image & img, stbi_failure_reason()); return false; } + if (w <= 0 || h <= 0 || w > int(kMaxImageDim) || h > int(kMaxImageDim)) { + std::fprintf(stderr, "vla-server: decoded JPEG dims %dx%d out of range (max %u)\n", + w, h, kMaxImageDim); + stbi_image_free(px); + return false; + } u8.assign(px, px + size_t(3) * w * h); stbi_image_free(px); view = { u8.data(), w, h, vla::PixelFormat::U8 }; return true; } else if (img.encoding() == vla::Image::RGB_U8) { + if (img.width() == 0 || img.height() == 0 || + img.width() > kMaxImageDim || img.height() > kMaxImageDim) { + std::fprintf(stderr, "vla-server: RGB_U8 dims %ux%u out of range (max %u)\n", + img.width(), img.height(), kMaxImageDim); + return false; + } const size_t expected = size_t(3) * img.width() * img.height(); if (img.data().size() != expected) { std::fprintf(stderr, "vla-server: RGB_U8 size %zu != 3*%u*%u = %zu\n", @@ -67,6 +94,12 @@ bool decode_image(const vla::Image & img, view = { u8.data(), int(img.width()), int(img.height()), vla::PixelFormat::U8 }; return true; } else if (img.encoding() == vla::Image::F32_RGB_01) { + if (img.width() == 0 || img.height() == 0 || + img.width() > kMaxImageDim || img.height() > kMaxImageDim) { + std::fprintf(stderr, "vla-server: F32_RGB_01 dims %ux%u out of range (max %u)\n", + img.width(), img.height(), kMaxImageDim); + return false; + } const size_t pixels = size_t(3) * img.width() * img.height(); const size_t expected = pixels * sizeof(float); if (img.data().size() != expected) { @@ -197,9 +230,32 @@ int main(int argc, char ** argv) { zmq::context_t zctx( 1); zmq::socket_t sock(zctx, zmq::socket_type::rep); sock.set(zmq::sockopt::linger, 0); + // cap inbound messages so one oversized request cannot exhaust memory. + sock.set(zmq::sockopt::maxmsgsize, int64_t(256) * 1024 * 1024); sock.bind(bind_addr); std::printf("vla-server: bound to %s. ready.\n", bind_addr.c_str()); + if (bind_addr.find("127.0.0.1") == std::string::npos && + bind_addr.find("localhost") == std::string::npos) { + std::fprintf(stderr, + "vla-server: WARNING: bound to %s with no authentication. Any host that\n" + " can reach this address may submit requests. Restrict access to a\n" + " trusted network, or bind tcp://127.0.0.1:PORT for local use.\n", + bind_addr.c_str()); + } + + // A failed reply desyncs the REP recv/send lockstep, so this socket cannot + // recover; log and shut down cleanly rather than crash (the old unguarded + // send) or spin forever on the wedged socket. + auto send_reply = [&sock](const std::string & body) { + try { + sock.send(zmq::buffer(body), zmq::send_flags::none); + } catch (const zmq::error_t & e) { + std::fprintf(stderr, "vla-server: reply send failed (%s); shutting down\n", e.what()); + g_shutdown.store(true, std::memory_order_relaxed); + } + }; + std::signal(SIGINT, on_signal); std::signal(SIGTERM, on_signal); @@ -212,7 +268,9 @@ int main(int argc, char ** argv) { zmq::poll(poll, 1, std::chrono::milliseconds(200)); } catch (const zmq::error_t & e) { if (e.num() == EINTR) continue; - throw; + if (e.num() == ETERM) break; + std::fprintf(stderr, "vla-server: zmq error: %s\n", e.what()); + continue; } if (!(poll[0].revents & ZMQ_POLLIN)) continue; @@ -222,7 +280,9 @@ int main(int argc, char ** argv) { if (!rr) continue; } catch (const zmq::error_t & e) { if (e.num() == EINTR) continue; - throw; + if (e.num() == ETERM) break; + std::fprintf(stderr, "vla-server: zmq error: %s\n", e.what()); + continue; } vla::PredictRequest req; @@ -230,7 +290,7 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "vla-server: PredictRequest parse failed (size=%zu)\n", req_msg.size()); const std::string body = make_error_response(0, "request parse failed"); - sock.send(zmq::buffer(body), zmq::send_flags::none); + send_reply(body); continue; } @@ -239,14 +299,18 @@ int main(int argc, char ** argv) { if (req.images_size() < 1 && req.precomputed_img_emb_size() == 0) { const std::string body = make_error_response(rid, "PredictRequest must contain images or precomputed_img_emb"); - sock.send(zmq::buffer(body), zmq::send_flags::none); + send_reply(body); + continue; + } + if (req.images_size() > 16) { + send_reply(make_error_response(rid, "too many image views (max 16)")); continue; } if (req.lang_tokens_size() < 1 || req.lang_tokens_size() > int(cfg.n_lang)) { char buf[128]; std::snprintf(buf, sizeof(buf), "lang_tokens length %d out of range [1, %lld]", req.lang_tokens_size(), (long long) cfg.n_lang); - sock.send(zmq::buffer(make_error_response(rid, buf)), zmq::send_flags::none); + send_reply(make_error_response(rid, buf)); continue; } { @@ -255,7 +319,7 @@ int main(int argc, char ** argv) { if (req.lang_tokens(t) < 0) { char buf[128]; std::snprintf(buf, sizeof(buf), "lang_tokens[%d] = %d is negative", t, req.lang_tokens(t)); - sock.send(zmq::buffer(make_error_response(rid, buf)), zmq::send_flags::none); + send_reply(make_error_response(rid, buf)); tokens_ok = false; break; } @@ -266,7 +330,7 @@ int main(int argc, char ** argv) { char buf[128]; std::snprintf(buf, sizeof(buf), "state length %d != expected %lld", req.state_size(), (long long) cfg.max_state_dim); - sock.send(zmq::buffer(make_error_response(rid, buf)), zmq::send_flags::none); + send_reply(make_error_response(rid, buf)); continue; } const int expected_noise_n = int(cfg.n_suffix * cfg.max_action_dim); @@ -274,7 +338,7 @@ int main(int argc, char ** argv) { char buf[128]; std::snprintf(buf, sizeof(buf), "noise length %d != 0 or %d (chunk_size * action_dim)", req.noise_size(), expected_noise_n); - sock.send(zmq::buffer(make_error_response(rid, buf)), zmq::send_flags::none); + send_reply(make_error_response(rid, buf)); continue; } @@ -289,15 +353,15 @@ int main(int argc, char ** argv) { if (use_precomputed) { precomputed_n_views = static_cast(req.precomputed_img_emb_n_views()); - const int per_view = int(cfg.n_img * cfg.hidden); - const int expected = per_view * precomputed_n_views; - if (precomputed_n_views < 1 || req.precomputed_img_emb_size() != expected) { + const int64_t per_view = cfg.n_img * cfg.hidden; + const int64_t expected = per_view * static_cast(precomputed_n_views); + if (precomputed_n_views < 1 || + static_cast(req.precomputed_img_emb_size()) != expected) { char buf[160]; std::snprintf(buf, sizeof(buf), - "precomputed_img_emb size %d != %d (n_views=%d * n_img_per_view=%lld * hidden=%lld)", - req.precomputed_img_emb_size(), expected, precomputed_n_views, + "precomputed_img_emb size %d != %lld (n_views=%d * n_img_per_view=%lld * hidden=%lld)", + req.precomputed_img_emb_size(), (long long) expected, precomputed_n_views, (long long) cfg.n_img, (long long) cfg.hidden); - sock.send(zmq::buffer(make_error_response(rid, buf)), - zmq::send_flags::none); + send_reply(make_error_response(rid, buf)); continue; } precomputed_emb.assign(req.precomputed_img_emb().begin(), @@ -308,7 +372,7 @@ int main(int argc, char ** argv) { char buf[128]; std::snprintf(buf, sizeof(buf), "precomputed_img_emb[%d] = %g is not finite (NaN/Inf)", bad, precomputed_emb[bad]); - sock.send(zmq::buffer(make_error_response(rid, buf)), zmq::send_flags::none); + send_reply(make_error_response(rid, buf)); continue; } } else { @@ -317,8 +381,7 @@ int main(int argc, char ** argv) { for (int v = 0; v < n_views; ++v) { if (!decode_image(req.images(v), u8_bufs[v], f32_bufs[v], img_views[v])) { char buf[64]; std::snprintf(buf, sizeof(buf), "image[%d] decode failed", v); - sock.send(zmq::buffer(make_error_response(rid, buf)), - zmq::send_flags::none); + send_reply(make_error_response(rid, buf)); decode_ok = false; break; } @@ -333,7 +396,7 @@ int main(int argc, char ** argv) { if (bad >= 0) { char buf[128]; std::snprintf(buf, sizeof(buf), "state[%d] = %g is not finite (NaN/Inf)", bad, state_vec[bad]); - sock.send(zmq::buffer(make_error_response(rid, buf)), zmq::send_flags::none); + send_reply(make_error_response(rid, buf)); continue; } } @@ -344,7 +407,7 @@ int main(int argc, char ** argv) { if (bad >= 0) { char buf[128]; std::snprintf(buf, sizeof(buf), "noise[%d] = %g is not finite (NaN/Inf)", bad, noise_vec[bad]); - sock.send(zmq::buffer(make_error_response(rid, buf)), zmq::send_flags::none); + send_reply(make_error_response(rid, buf)); continue; } } @@ -378,8 +441,7 @@ int main(int argc, char ** argv) { const auto & st = vla::last_stats(model); if (action_chunk.empty()) { - sock.send(zmq::buffer(make_error_response(rid, "predict failed")), - zmq::send_flags::none); + send_reply(make_error_response(rid, "predict failed")); continue; } @@ -396,7 +458,7 @@ int main(int argc, char ** argv) { resp.set_latency_ms_denoise(st.ms_denoise); const std::string body = resp.SerializeAsString(); - sock.send(zmq::buffer(body), zmq::send_flags::none); + send_reply(body); ++served; if (served % 10 == 1) { diff --git a/src/serving/vla-cli.cpp b/src/serving/vla-cli.cpp new file mode 100644 index 0000000..597a010 --- /dev/null +++ b/src/serving/vla-cli.cpp @@ -0,0 +1,178 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// One-shot action prediction from the command line. Loads a model, decodes an +// image plus an already-tokenized instruction, runs one predict(), and prints +// the action chunk. No server, no simulator. Tokenization stays in the Python +// client, so language is passed as token ids here. +// +// vla-cli [--mmproj m.gguf] --ckpt c.gguf --image img.jpg [--image img2.jpg] +// --tokens id,id,... [--state f,f,...] [--pretty] + +#include "model.h" + +#define STB_IMAGE_IMPLEMENTATION +#define STB_IMAGE_STATIC +// stb pulls in its full implementation here; keep its unused-function noise out +// of our -Wall -Wextra output. +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wunused-function" +#include "stb_image.h" +#pragma GCC diagnostic pop + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace vla; + +namespace { + +// Comma/space separated ints. Returns false on junk or out-of-int32 values. +bool parse_ints(const std::string & s, std::vector & out) { + out.clear(); + size_t i = 0; + while (i < s.size()) { + while (i < s.size() && (s[i] == ',' || s[i] == ' ')) ++i; + if (i >= s.size()) break; + errno = 0; + char * e = nullptr; + long long x = std::strtoll(s.c_str() + i, &e, 10); + if (e == s.c_str() + i) { std::fprintf(stderr, "vla-cli: bad token near '%s'\n", s.c_str() + i); return false; } + if (errno == ERANGE || x < INT32_MIN || x > INT32_MAX) { std::fprintf(stderr, "vla-cli: token %lld out of int32 range\n", x); return false; } + out.push_back((int32_t) x); + i = (size_t) (e - s.c_str()); + } + return true; +} + +// Comma/space separated floats. Returns false on junk or non-finite values. +bool parse_floats(const std::string & s, std::vector & out) { + out.clear(); + size_t i = 0; + while (i < s.size()) { + while (i < s.size() && (s[i] == ',' || s[i] == ' ')) ++i; + if (i >= s.size()) break; + char * e = nullptr; + float x = std::strtof(s.c_str() + i, &e); + if (e == s.c_str() + i) { std::fprintf(stderr, "vla-cli: bad number near '%s'\n", s.c_str() + i); return false; } + if (!std::isfinite(x)) { std::fprintf(stderr, "vla-cli: non-finite value in --state\n"); return false; } + out.push_back(x); + i = (size_t) (e - s.c_str()); + } + return true; +} + +// Decode an image file to interleaved RGB8. buf must outlive the ImageView. +bool load_image(const char * path, std::vector & buf, int & w, int & h) { + int ch = 0; + unsigned char * px = stbi_load(path, &w, &h, &ch, 3); + if (!px) { + std::fprintf(stderr, "vla-cli: cannot load image %s: %s\n", path, stbi_failure_reason()); + return false; + } + buf.assign(px, px + size_t(3) * w * h); + stbi_image_free(px); + return true; +} + +void usage(const char * prog) { + std::fprintf(stderr, + "usage: %s [--mmproj m.gguf] --ckpt c.gguf --image img.jpg [--image ...]\n" + " --tokens id,id,... [--state f,f,...] [--pretty]\n" + " --mmproj vision-tower GGUF (SmolVLA/pi0/pi0.5); omit for baked-vision archs\n" + " --ckpt model checkpoint GGUF\n" + " --image image file, repeat for multi-view (decoded via stb_image)\n" + " --tokens language token ids, comma-separated (tokenize in the client)\n" + " --state proprioception floats, comma-separated (default zeros)\n" + " --pretty print one action row (max_action_dim values) per line\n", + prog); +} + +} // namespace + +int main(int argc, char ** argv) { + std::string mmproj, ckpt, tokens_s, state_s; + std::vector image_paths; + bool pretty = false; + + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + auto need = [&](const char * name) -> const char * { + if (i + 1 >= argc) { std::fprintf(stderr, "vla-cli: %s needs a value\n", name); std::exit(1); } + return argv[++i]; + }; + if (a == "--mmproj") mmproj = need("--mmproj"); + else if (a == "--ckpt") ckpt = need("--ckpt"); + else if (a == "--image") image_paths.push_back(need("--image")); + else if (a == "--tokens") tokens_s = need("--tokens"); + else if (a == "--state") state_s = need("--state"); + else if (a == "--pretty") pretty = true; + else if (a == "-h" || a == "--help") { usage(argv[0]); return 0; } + else { std::fprintf(stderr, "vla-cli: unknown argument %s\n", a.c_str()); usage(argv[0]); return 1; } + } + if (ckpt.empty() || image_paths.empty() || tokens_s.empty()) { usage(argv[0]); return 1; } + + // Validate the cheap args before loading the model. + std::vector lang; + std::vector state; + if (!parse_ints(tokens_s, lang) || !parse_floats(state_s, state)) return 1; + if (lang.empty()) { std::fprintf(stderr, "vla-cli: --tokens parsed to nothing\n"); return 1; } + + Model * m = model_load(mmproj, ckpt, ""); + if (!m) { std::fprintf(stderr, "vla-cli: model_load failed\n"); return 1; } + const Config & cfg = model_config(m); + + if (!state.empty() && (int64_t) state.size() != cfg.max_state_dim) + std::fprintf(stderr, "vla-cli: --state has %zu values, model expects %lld; padding or truncating\n", + state.size(), (long long) cfg.max_state_dim); + state.resize((size_t) cfg.max_state_dim, 0.0f); + + std::vector> imgbuf(image_paths.size()); + std::vector views(image_paths.size()); + for (size_t v = 0; v < image_paths.size(); ++v) { + int w = 0, h = 0; + if (!load_image(image_paths[v].c_str(), imgbuf[v], w, h)) { model_free(m); return 1; } + views[v] = ImageView{ imgbuf[v].data(), w, h, PixelFormat::U8 }; + } + + Inputs in{}; + in.images = views.data(); + in.n_images = (int) views.size(); + in.lang_tokens = lang.data(); + in.n_lang = (int) lang.size(); + in.state = state.data(); + in.noise = nullptr; // predict() samples N(0,1) when omitted + + std::vector act = predict(m, in); + if (act.empty()) { std::fprintf(stderr, "vla-cli: predict failed\n"); model_free(m); return 2; } + + const int64_t cols = cfg.max_action_dim > 0 ? cfg.max_action_dim : 1; + if (pretty) { + for (size_t i = 0; i < act.size(); ++i) + std::printf("%.6g%c", act[i], ((int64_t) (i + 1) % cols == 0) ? '\n' : ' '); + } else { + std::printf("action_len=%zu\n", act.size()); + for (float x : act) std::printf("%.9g\n", x); + } + std::fflush(stdout); + + model_free(m); + return 0; +} diff --git a/src/serving/vlm-server.cpp b/src/serving/vlm-server.cpp index 1547d56..3fa3e5d 100644 --- a/src/serving/vlm-server.cpp +++ b/src/serving/vlm-server.cpp @@ -17,6 +17,7 @@ #include +#include #include #include #include @@ -29,6 +30,9 @@ namespace { std::atomic g_shutdown{false}; void on_signal(int) { g_shutdown.store(true, std::memory_order_relaxed); } +// Reject absurd image dimensions before any size arithmetic on untrusted input. +constexpr unsigned kMaxImageDim = 8192; + std::string make_error_stream(uint64_t request_id, const std::string & msg) { vlm_chat::StreamMessage sm; vlm_chat::ChatResponse * resp = sm.mutable_final(); @@ -97,9 +101,20 @@ int main(int argc, char ** argv) { zmq::context_t zctx( 1); zmq::socket_t sock(zctx, zmq::socket_type::router); sock.set(zmq::sockopt::linger, 0); + // cap inbound messages so one oversized request cannot exhaust memory. + sock.set(zmq::sockopt::maxmsgsize, int64_t(256) * 1024 * 1024); sock.bind(bind_addr); std::printf("vlm-server: bound to %s. ready.\n", bind_addr.c_str()); + if (bind_addr.find("127.0.0.1") == std::string::npos && + bind_addr.find("localhost") == std::string::npos) { + std::fprintf(stderr, + "vlm-server: WARNING: bound to %s with no authentication. Any host that can\n" + " reach this address may submit requests. Restrict to a trusted\n" + " network, or bind tcp://127.0.0.1:PORT for local use.\n", + bind_addr.c_str()); + } + std::signal(SIGINT, on_signal); std::signal(SIGTERM, on_signal); @@ -111,7 +126,9 @@ int main(int argc, char ** argv) { zmq::poll(poll, 1, std::chrono::milliseconds(200)); } catch (const zmq::error_t & e) { if (e.num() == EINTR) continue; - throw; + if (e.num() == ETERM) break; + std::fprintf(stderr, "vlm-server: zmq error: %s\n", e.what()); + continue; } if (!(poll[0].revents & ZMQ_POLLIN)) continue; @@ -124,8 +141,9 @@ int main(int argc, char ** argv) { auto rr = sock.recv(part, zmq::recv_flags::none); if (!rr) { recv_ok = false; break; } } catch (const zmq::error_t & e) { - if (e.num() == EINTR) { recv_ok = false; break; } - throw; + if (e.num() != EINTR) + std::fprintf(stderr, "vlm-server: zmq recv error: %s\n", e.what()); + recv_ok = false; break; } if (sock.get(zmq::sockopt::rcvmore)) { env.emplace_back(static_cast(part.data()), part.size()); @@ -138,10 +156,17 @@ int main(int argc, char ** argv) { if (!recv_ok || !have_payload || env.empty()) continue; auto send_reply = [&](const std::string & body) { - for (const auto & e : env) { - sock.send(zmq::buffer(e), zmq::send_flags::sndmore); + try { + for (const auto & e : env) { + sock.send(zmq::buffer(e), zmq::send_flags::sndmore); + } + sock.send(zmq::buffer(body), zmq::send_flags::none); + } catch (const zmq::error_t & e) { + // A throw mid-multipart leaves the ROUTER frame open, which would + // corrupt the next reply; shut down cleanly instead of crashing. + std::fprintf(stderr, "vlm-server: reply send failed (%s); shutting down\n", e.what()); + g_shutdown.store(true, std::memory_order_relaxed); } - sock.send(zmq::buffer(body), zmq::send_flags::none); }; vlm_chat::ChatRequest req; @@ -156,6 +181,10 @@ int main(int argc, char ** argv) { send_reply(make_error_stream(rid, "ChatRequest has no messages")); continue; } + if (req.images_size() > 16) { + send_reply(make_error_stream(rid, "too many image views (max 16)")); + continue; + } std::vector images; images.reserve(req.images_size()); @@ -172,6 +201,14 @@ int main(int argc, char ** argv) { decode_ok = false; break; } } else if (im.encoding() == vlm_chat::Image::RGB_U8) { + if (im.width() == 0 || im.height() == 0 || + im.width() > kMaxImageDim || im.height() > kMaxImageDim) { + char buf[96]; std::snprintf(buf, sizeof(buf), + "image[%d] RGB_U8 dims %ux%u out of range (max %u)", v, + im.width(), im.height(), kMaxImageDim); + send_reply(make_error_stream(rid, buf)); + decode_ok = false; break; + } const size_t expected = size_t(3) * im.width() * im.height(); if (im.data().size() != expected) { char buf[96]; std::snprintf(buf, sizeof(buf), @@ -204,7 +241,11 @@ int main(int argc, char ** argv) { sp.temperature = s.temperature(); sp.top_p = s.top_p(); sp.top_k = s.top_k(); - sp.max_tokens = s.max_tokens() > 0 ? s.max_tokens() : 256; + // Cap client-requested length at the context window so a huge + // max_tokens cannot monopolize the single-threaded decode loop. + const uint32_t cap = uint32_t(lp.n_ctx > 0 ? lp.n_ctx : 4096); + const uint32_t want = s.max_tokens() > 0 ? s.max_tokens() : 256; + sp.max_tokens = int(std::min(want, cap)); sp.seed = s.seed() != 0 ? s.seed() : 0xFFFFFFFFu; } diff --git a/src/vlm/engine.cpp b/src/vlm/engine.cpp index 7f090c8..8f5e81f 100644 --- a/src/vlm/engine.cpp +++ b/src/vlm/engine.cpp @@ -133,7 +133,7 @@ bool Engine::decode_image_file(const std::string & path, Image & out) const { if (!loaded()) { return false; } - mtmd::bitmap bmp(mtmd_helper_bitmap_init_from_file(impl_->vision.get(), path.c_str())); + mtmd::bitmap bmp(mtmd_helper_bitmap_init_from_file(impl_->vision.get(), path.c_str(), false).bitmap); return bitmap_to_image(bmp, out); } @@ -141,7 +141,7 @@ bool Engine::decode_image_buf(const uint8_t * data, size_t len, Image & out) con if (!loaded() || !data || len == 0) { return false; } - mtmd::bitmap bmp(mtmd_helper_bitmap_init_from_buf(impl_->vision.get(), data, len)); + mtmd::bitmap bmp(mtmd_helper_bitmap_init_from_buf(impl_->vision.get(), data, len, false).bitmap); return bitmap_to_image(bmp, out); } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt new file mode 100644 index 0000000..3e20804 --- /dev/null +++ b/tests/CMakeLists.txt @@ -0,0 +1,11 @@ +# vla.cpp tests, enabled with -DVLA_BUILD_TESTS=ON. predict_check needs a real +# GGUF so it is built but not run by ctest; the unit tests below are pure. + +add_executable(vla_predict_check predict_check.cpp) +target_link_libraries(vla_predict_check PRIVATE vla_core) +target_compile_options(vla_predict_check PRIVATE -Wall -Wextra) + +add_executable(test_vision_common test_vision_common.cpp) +target_include_directories(test_vision_common PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_compile_options(test_vision_common PRIVATE -Wall -Wextra) +add_test(NAME vision_common COMMAND test_vision_common) diff --git a/tests/predict_check.cpp b/tests/predict_check.cpp new file mode 100644 index 0000000..1d71659 --- /dev/null +++ b/tests/predict_check.cpp @@ -0,0 +1,113 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Deterministic output-equality + timing harness for vla::predict. Feeds fixed +// images/language/state/noise and prints the action chunk; diff stdout before and +// after a change to confirm numerics are unchanged (bitvla ignores noise, the +// others read it, so fixed noise is reproducible for all archs). +// +// predict_check [mmproj.gguf] [n_images] +// env: VLA_IMG_SIZE (square input, default 224), VLA_BENCH_ITERS (>0 = time it) + +#include "model.h" + +#include +#include +#include +#include +#include + +using namespace vla; + +int main(int argc, char** argv) { + if (argc < 2) { + std::fprintf(stderr, "usage: %s [mmproj.gguf] [n_images]\n", argv[0]); + return 1; + } + const char* ckpt = argv[1]; + const char* mmproj = (argc > 2 && argv[2][0] && argv[2][0] != '-') ? argv[2] : ""; + const int n_images = argc > 3 ? std::atoi(argv[3]) : 2; + + Model* m = model_load(mmproj, ckpt, ""); + if (!m) { + std::fprintf(stderr, "model_load failed\n"); + return 1; + } + const Config& cfg = model_config(m); + std::fprintf(stderr, + "loaded: real_state=%lld max_state=%lld real_action=%lld max_action=%lld " + "n_suffix=%lld num_steps=%d\n", + (long long)cfg.real_state_dim, (long long)cfg.max_state_dim, + (long long)cfg.real_action_dim, (long long)cfg.max_action_dim, + (long long)cfg.n_suffix, cfg.num_steps); + + const char* isz = std::getenv("VLA_IMG_SIZE"); + const int W = isz ? std::atoi(isz) : 224, H = W; + std::vector> imgbuf(n_images, std::vector((size_t)3 * W * H)); + std::vector views(n_images); + for (int v = 0; v < n_images; ++v) { + for (int y = 0; y < H; ++y) + for (int x = 0; x < W; ++x) + for (int c = 0; c < 3; ++c) + imgbuf[v][((size_t)y * W + x) * 3 + c] = (uint8_t)((x + 2 * y + 40 * c + 17 * v) & 0xFF); + views[v] = ImageView{ imgbuf[v].data(), W, H, PixelFormat::U8 }; + } + + std::vector lang = {1, 100, 200, 300, 400, 2}; + std::vector state((size_t)cfg.max_state_dim, 0.0f); + for (int i = 0; i < (int)cfg.real_state_dim; ++i) state[i] = 0.01f * (float)(i + 1); + + const size_t noise_n = (size_t)cfg.max_action_dim * (size_t)cfg.n_suffix; + std::vector noise(noise_n); + for (size_t i = 0; i < noise_n; ++i) + noise[i] = 0.001f * (float)((i * 2654435761u) % 1000) - 0.5f; + + Inputs in{}; + in.images = views.data(); + in.n_images = n_images; + in.lang_tokens = lang.data(); + in.n_lang = (int)lang.size(); + in.state = state.data(); + in.noise = noise.data(); + in.timing_detail = TimingDetail::NONE; + + std::vector act = predict(m, in); + std::printf("action_len=%zu\n", act.size()); + for (size_t i = 0; i < act.size(); ++i) std::printf("%.9g\n", (double)act[i]); + std::fflush(stdout); + + const char* it_env = std::getenv("VLA_BENCH_ITERS"); + const int iters = it_env ? std::atoi(it_env) : 0; + if (iters > 0) { + for (int w = 0; w < 3; ++w) (void)predict(m, in); + double best = 1e30, sum = 0.0; + for (int i = 0; i < iters; ++i) { + struct timespec t0, t1; + clock_gettime(CLOCK_MONOTONIC, &t0); + std::vector a = predict(m, in); + clock_gettime(CLOCK_MONOTONIC, &t1); + double ms = (t1.tv_sec - t0.tv_sec) * 1e3 + (t1.tv_nsec - t0.tv_nsec) / 1e6; + best = ms < best ? ms : best; + sum += ms; + } + const Stats& st = last_stats(m); + std::fprintf(stderr, "predict() over %d iters: min=%.3f ms avg=%.3f ms\n", + iters, best, sum / iters); + std::fprintf(stderr, "last split: vision=%.2f inference=%.2f total=%.2f ms\n", + st.ms_vision, st.ms_inference, st.ms_total); + } + + model_free(m); + return act.empty() ? 2 : 0; +} diff --git a/tests/py/test_converters.py b/tests/py/test_converters.py new file mode 100644 index 0000000..9200f83 --- /dev/null +++ b/tests/py/test_converters.py @@ -0,0 +1,59 @@ +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +"""Unit tests for the clip->vit tensor-name remap in the mmproj merge scripts. +A wrong mapping silently breaks the in-tree tower. remap() is pure, so the heavy +`gguf` import is stubbed.""" + +import importlib.util +import pathlib +import sys +import types + + +def _load(name): + sys.modules.setdefault("gguf", types.ModuleType("gguf")) + path = pathlib.Path(__file__).resolve().parents[2] / "scripts" / f"{name}.py" + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def _check(remap): + # clip block names -> in-tree vit block names, with the attn/ffn sub-renames. + assert remap("v.blk.5.attn_out.weight") == "vit.blk.5.attn_o.weight" + assert remap("v.blk.0.ffn_up.weight") == "vit.blk.0.fc1.weight" + assert remap("v.blk.11.ffn_down.bias") == "vit.blk.11.fc2.bias" + # a passthrough sub-name keeps its suffix + assert remap("v.blk.2.attn_q.weight") == "vit.blk.2.attn_q.weight" + # fixed top-level names + assert remap("v.patch_embd.weight") == "vit.patch_embd.weight" + assert remap("v.position_embd.weight") == "vit.pos_embd" + assert remap("v.post_ln.bias") == "vit.post_ln.bias" + # unrelated tensors are dropped + assert remap("some.audio.tensor") is None + + +def test_pi0_remap(): + m = _load("merge_pi0_mmproj_to_gguf") + assert m.remap("mm.input_projection.weight") == "mm.proj.weight" + _check(m.remap) + + +# NB: SmolVLA has no merge script - convert_smolvla_to_gguf.py bakes the vision +# tower straight from safetensors, so there is no clip->vit remap to test here. + + +if __name__ == "__main__": + test_pi0_remap() + print("test_converters: remap OK") diff --git a/tests/test_vision_common.cpp b/tests/test_vision_common.cpp new file mode 100644 index 0000000..5dc5ceb --- /dev/null +++ b/tests/test_vision_common.cpp @@ -0,0 +1,79 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Unit test for the IDEFICS3/SmolVLM pixel-shuffle channel order: a c-outermost +// (wrong) arrangement fails the hard-coded expectations below. + +#include "models/vision_common.h" + +#undef NDEBUG // keep assert() live even in Release builds +#include +#include +#include +#include + +int main() { + // embed=2, grid=4, s=2 -> g2=2, n_patches=16, c4=8, out tokens=4. + // src[p*embed + e] = 1000*p + e, unique per (patch, channel). + const int64_t embed = 2, grid = 4, s = 2; + const int64_t n_patches = grid * grid, c4 = embed * s * s, tokens = (grid / s) * (grid / s); + std::vector src((size_t) embed * n_patches); + for (int64_t p = 0; p < n_patches; ++p) + for (int64_t e = 0; e < embed; ++e) + src[p * embed + e] = 1000.0f * (float) p + (float) e; + + std::vector dst((size_t) c4 * tokens, -1.0f); + vla::pixel_shuffle_hf(src.data(), dst.data(), embed, grid, s); + + // Hand-computed token 0 (h2=0,w2=0): inner patches {p0,p1,p4,p5}, channel + // innermost. A c-outermost layout would order these differently. + const float expect0[8] = {0, 1, 1000, 1001, 4000, 4001, 5000, 5001}; + for (int i = 0; i < 8; ++i) { + if (dst[i] != expect0[i]) { + std::fprintf(stderr, "pixel_shuffle_hf token0[%d]=%g, expected %g (c-order wrong?)\n", + i, (double) dst[i], (double) expect0[i]); + return 1; + } + } + if (dst[8] != 2000.0f) { // token 1 (w2=1), inner patch p2, channel 0 + std::fprintf(stderr, "pixel_shuffle_hf token1[0]=%g, expected 2000\n", (double) dst[8]); + return 1; + } + + // Full mapping check against the reference index formula. + const int64_t g2 = grid / s; + for (int64_t h2 = 0; h2 < g2; ++h2) + for (int64_t w2 = 0; w2 < g2; ++w2) + for (int64_t hs = 0; hs < s; ++hs) + for (int64_t ws = 0; ws < s; ++ws) + for (int64_t e = 0; e < embed; ++e) { + const int64_t t = h2 * g2 + w2; + const int64_t p = (h2 * s + hs) * grid + (w2 * s + ws); + const float got = dst[t * c4 + (hs * s + ws) * embed + e]; + const float want = 1000.0f * (float) p + (float) e; + assert(got == want); + (void) got; (void) want; + } + + // view_is_side: only an exact side x side view with real data passes. + int dummy = 0; + assert(vla::view_is_side(&dummy, 224, 224, 224)); + assert(!vla::view_is_side(nullptr, 224, 224, 224)); + assert(!vla::view_is_side(&dummy, 32, 32, 224)); + assert(!vla::view_is_side(&dummy, 224, 32, 224)); + assert(!vla::view_is_side(&dummy, 32, 224, 224)); + + std::printf("test_vision_common: pixel_shuffle_hf c-innermost order OK\n"); + return 0; +}