diff --git a/CMakeLists.txt b/CMakeLists.txt index 708004d..11fece1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -144,6 +144,7 @@ add_library(vla_core src/models/vla_adapter.cpp src/models/openvla_oft.cpp src/models/vla_jepa.cpp + src/models/turbovla.cpp ) target_include_directories(vla_core PUBLIC diff --git a/README.md b/README.md index 7978201..69a9434 100644 --- a/README.md +++ b/README.md @@ -326,6 +326,7 @@ supported (released and benchmarked), `~` = in progress, `-` = planned. | [OpenVLA-OFT](https://hf.co/vrfai/openvla-oft-libero-gguf) | Y | Y | - | Y | Y | - | | [VLA-JEPA](https://hf.co/vrfai/vla-jepa-libero) | Y | Y | - | Y | Y | - | | [Octo-Small](https://hf.co/vrfai/octo-small-libero-gguf) | Y | Y | Y | Y | - | - | +| [TurboVLA](https://hf.co/vrfai/turbovla-libero-gguf) | Y | Y | Y | Y | Y | - | --- @@ -364,6 +365,7 @@ Supported VLA models: - [GR00T N1.x](https://github.com/NVIDIA/Isaac-GR00T) - NVIDIA Isaac. - [VLA-JEPA](https://github.com/ginwind/VLA-JEPA) - Jingwen Sun et al. - [Octo](https://github.com/octo-models/octo) - Octo Model Team, UC Berkeley RAIL. +- [TurboVLA](https://github.com/H-EmbodVis/TurboVLA) - Hengyi Xie et al. Built on: diff --git a/eval/client/run_sim_client_direct.py b/eval/client/run_sim_client_direct.py index 6d65a2c..0bb62f5 100644 --- a/eval/client/run_sim_client_direct.py +++ b/eval/client/run_sim_client_direct.py @@ -88,7 +88,8 @@ parser.add_argument( "--stats-json", type=str, default=None, - help="[bitvla/gr00t_n1_6/gr00t_n1_7] path to dataset_statistics.json. Default for bitvla: " + help="[turbovla/bitvla/gr00t_n1_6/gr00t_n1_7] path to dataset statistics. TurboVLA defaults " + "to the H-EmbodVis/TurboVLA release stats. Default for bitvla: " "/dataset_statistics.json. For gr00t_n1_{6,7}, pass the canonical " "/experiment_cfg/dataset_statistics.json explicitly. Without " "it, gr00t_n1_7 returns the raw normalized [40, 132] chunk and gr00t_n1_6 " diff --git a/eval/client/vla_cpp_client.py b/eval/client/vla_cpp_client.py index 6d81559..b295905 100644 --- a/eval/client/vla_cpp_client.py +++ b/eval/client/vla_cpp_client.py @@ -22,8 +22,9 @@ import sys import tempfile from collections import deque +from collections.abc import Sequence from pathlib import Path -from typing import Any, Sequence +from typing import Any import numpy as np import torch @@ -49,6 +50,12 @@ "max_state_dim": 8}, "vla_jepa": {"image_size": 256, "tokenizer": "Qwen/Qwen3-VL-2B-Instruct", "max_state_dim": 8, "use_processor": True}, + "turbovla": { + "image_size": 256, + "tokenizer": "bert-base-uncased", + "max_state_dim": 8, + "max_length": 21, + }, "gr00t_n1_7": {"image_size": 256, "tokenizer": "nvidia/Cosmos-Reason2-2B", "max_state_dim": 132}, "gr00t_n1_5": {"image_size": 224, "tokenizer": "lerobot/eagle2hg-processor-groot-n1p5", @@ -285,6 +292,65 @@ def _oft_norm(x, q01=q01, q99=q99, mask=mask): f"{stats_path}::observation.state ({self._pi05_state_q01.shape[0]}-D)", flush=True) + self._turbovla_state_norm = None + self._turbovla_action_unnorm = None + if arch == "turbovla": + if stats_json: + stats_path = Path(stats_json) + else: + from huggingface_hub import hf_hub_download + stats_path = Path(hf_hub_download( + repo_id="H-EmbodVis/TurboVLA", filename="libero_all4_stats.json")) + if not stats_path.exists(): + raise FileNotFoundError( + f"TurboVLA stats not found at {stats_path}. Pass --stats-json " + ".") + blob = json.loads(stats_path.read_text()) + try: + if "libero_all4_no_noops" in blob: + blob = blob["libero_all4_no_noops"] + state = blob["proprio"] if "proprio" in blob else blob["observation.state"] + action = blob["action"] + state_mean = np.asarray(state["mean"], dtype=np.float32) + state_std = np.asarray(state["std"], dtype=np.float32) + action_min = np.asarray(action["min"], dtype=np.float32) + action_max = np.asarray(action["max"], dtype=np.float32) + except (KeyError, TypeError) as exc: + raise ValueError( + f"{stats_path} must contain observation.state mean/std and action min/max") from exc + if ( + state_mean.shape != (8,) + or state_std.shape != (8,) + or action_min.shape != (7,) + or action_max.shape != (7,) + ): + raise ValueError( + f"{stats_path} does not contain TurboVLA's 8-D state and 7-D action stats" + ) + if np.any(state_std <= 0) or np.any(action_max <= action_min): + raise ValueError(f"{stats_path} contains invalid TurboVLA normalization ranges") + + def _state_norm(state_8d, mean=state_mean, std=state_std): + return ((state_8d - mean) / (std + 1e-6)).astype(np.float32) + + def _action_unnorm(chunk, mn=action_min, mx=action_max): + normalized = np.clip(chunk.astype(np.float32), -1.0, 1.0) + action = normalized.copy() + action[..., :6] = ((normalized[..., :6] + 1.0) * 0.5 + * (mx[:6] - mn[:6]) + mn[:6]) + # TurboVLA thresholds the normalized gripper instead of applying + # continuous min/max denormalization; exactly 0 maps to +1. See + # H-EmbodVis/TurboVLA@b29ab142, turbovla/evaluation/policy.py:205-221. + action[..., 6] = np.where(normalized[..., 6] < 0.0, -1.0, 1.0) + return action.astype(np.float32) + + self._turbovla_state_norm = _state_norm + self._turbovla_action_unnorm = _action_unnorm + print( + "vla-cpp-direct[arch=turbovla]: release MEAN_STD state / MIN_MAX " + f"arm stats via {stats_path}", + flush=True, + ) self._gr00t_action_unnorm = None self._gr00t_state_norm = None @@ -581,6 +647,8 @@ def get_action(self, observations: dict[str, Any]) -> np.ndarray: def _predict_chunk(self, observations: dict[str, Any]) -> np.ndarray: + if self.arch == "turbovla": + return self._predict_chunk_turbovla(observations) if self.arch == "evo1": return self._predict_chunk_evo1(observations) if self.arch == "gr00t_n1_7": @@ -650,6 +718,68 @@ def _predict_chunk(self, observations: dict[str, Any]) -> np.ndarray: return (np.array(resp.action_chunk, dtype=np.float32) .reshape(resp.chunk_size, resp.action_dim)) + def _predict_chunk_turbovla(self, observations: dict[str, Any]) -> np.ndarray: + if self._turbovla_state_norm is None or self._turbovla_action_unnorm is None: + raise RuntimeError("TurboVLA statistics were not initialized") + images_f32: list[np.ndarray] = [] + for key in self.image_keys: + if key not in observations: + raise KeyError( + f"TurboVLA image key {key!r} missing; got {list(observations)}" + ) + image = observations[key] + if isinstance(image, torch.Tensor): + image = image.numpy() + image = np.asarray(image, dtype=np.float32) + if image.ndim != 3 or image.shape[0] != 3: + raise ValueError(f"{key}: expected CHW float32 [3,H,W], got {image.shape}") + image = _resize_with_pad(image, self.image_size, self.image_size) + images_f32.append( + np.ascontiguousarray(np.transpose(image, (1, 2, 0)), dtype=np.float32) + ) + + state = observations["observation.state"] + if isinstance(state, torch.Tensor): + state = state.numpy() + state = np.asarray(state, dtype=np.float32).reshape(-1) + if state.shape != (8,): + raise ValueError(f"TurboVLA requires an 8-D state, got {state.shape}") + + task = observations.get("task", "") + if isinstance(task, bytes): + task = task.decode() + if not isinstance(task, str) or not task.strip(): + raise ValueError("TurboVLA requires a non-empty task instruction") + lang = self.tok( + task, + padding=False, + truncation=True, + max_length=self.max_length, + return_tensors="np", + )["input_ids"][0].astype(np.int32) + + req = self.pb.PredictRequest() + req.request_id = self._step + self._step += 1 + for image in images_f32: + proto_image = req.images.add() + proto_image.encoding = self.pb.Image.F32_RGB_01 + proto_image.height, proto_image.width = image.shape[:2] + proto_image.data = image.tobytes() + req.lang_tokens.extend(lang.tolist()) + req.state.extend(self._turbovla_state_norm(state).tolist()) + + self.sock.send(req.SerializeToString()) + resp = self.pb.PredictResponse() + resp.ParseFromString(self.sock.recv()) + if resp.error: + raise RuntimeError(f"vla-server error: {resp.error}") + self._last_response = resp + chunk = np.array(resp.action_chunk, dtype=np.float32).reshape( + resp.chunk_size, resp.action_dim + ) + return self._turbovla_action_unnorm(chunk) + _VJ_PS, _VJ_TPS, _VJ_MERGE, _VJ_SIDE = 16, 2, 2, 256 @classmethod diff --git a/scripts/convert_turbovla_to_gguf.py b/scripts/convert_turbovla_to_gguf.py new file mode 100644 index 0000000..e910827 --- /dev/null +++ b/scripts/convert_turbovla_to_gguf.py @@ -0,0 +1,665 @@ +#!/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. + +"""Convert a TurboVLA checkpoint (H-EmbodVis/TurboVLA) to GGUF. + +Reads the released .pth ({model_state_dict, model_config}, e.g. +checkpoints/libero/turbovla_libero.pth) or a directory holding +model.safetensors and config.json. Dimensions come from tensor shapes. + +Usage: + python convert_turbovla_to_gguf.py turbovla_libero.pth -o turbovla-libero.gguf [--verify] +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch +from safetensors import safe_open + +import gguf + +F32 = gguf.GGMLQuantizationType.F32 +BF16 = gguf.GGMLQuantizationType.BF16 + +ARCH = "turbovla" + + +# Prefixes used by official TurboVLA checkpoints (without a "model." prefix). +PREFIX_VIT = "vision_encoder.backbone" +PREFIX_TEXT = "text_encoder.bert" +PREFIX_VIT_PROJ = "vision_projection" +PREFIX_FUSION = "vision_language_interaction.fusion_layers" +PREFIX_VL_TEXT = "vision_language_interaction.text_layers" +PREFIX_ACT_DEC = "action_head.decoder" +PREFIX_STATE_PROJ = "action_head.state_projection" +KEY_VIEW_EMB = "view_embedding" +KEY_TEXT_PROJ = "text_encoder.text_projection" + + +def kv_prefix(name: str) -> str: + return f"{ARCH}.{name}" + + +def bf16_u16(t: torch.Tensor) -> np.ndarray: + return t.contiguous().view(torch.uint16).cpu().numpy() + + +def add_tensor(writer: gguf.GGUFWriter, name: str, t: torch.Tensor) -> None: + """Add tensor with preserved dtype.""" + if t.dtype == torch.float32: + writer.add_tensor(name, t.contiguous().cpu().numpy(), raw_dtype=F32) + elif t.dtype == torch.bfloat16: + writer.add_tensor(name, bf16_u16(t), raw_shape=list(t.shape), raw_dtype=BF16) + elif t.dtype == torch.float16: + writer.add_tensor(name, t.contiguous().cpu().numpy().astype(np.float32), raw_dtype=F32) + else: + raise NotImplementedError(f"unsupported dtype {t.dtype} for {name}") + + +def max_layer(keys: set[str], pfx: str) -> int: + """Count number of layers with given prefix.""" + 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 + + +# facebook/dinov3-vitb16-pretrain-lvd1689m is gated, but only these architecture +# values are needed: the fine-tuned weights ship inside the TurboVLA checkpoint. +DINOV3_VITB16 = {"rope_theta": 100.0, "num_register_tokens": 4} + +# Tensors the runtime never reads: DINOv3's final norm (TurboVLA taps +# hidden_states[-1], before it), BERT's pooler, and the MAE mask token. +UNUSED = (f"{PREFIX_VIT}.norm.", f"{PREFIX_TEXT}.pooler.", f"{PREFIX_VIT}.embeddings.mask_token") + + +class TrackedTensors(dict): + """A state dict that records which keys the writers read.""" + + def __init__(self, *a, **k): + super().__init__(*a, **k) + self.read = set() + + def __getitem__(self, key): + self.read.add(key) + return super().__getitem__(key) + + +def load_checkpoint(ckpt: Path) -> tuple[TrackedTensors, dict]: + """Return (state dict, TurboVLA model_config) from a .pth or a directory.""" + if ckpt.is_file(): + blob = torch.load(str(ckpt), map_location="cpu", weights_only=False) + if not isinstance(blob, dict) or "model_state_dict" not in blob: + raise SystemExit(f"{ckpt} is not a TurboVLA checkpoint (no model_state_dict)") + state = blob["model_state_dict"] + cfg = blob.get("model_config") or {} + else: + st_path = ckpt / "model.safetensors" + if not st_path.exists(): + raise SystemExit(f"model.safetensors not found in {ckpt}") + with safe_open(str(st_path), framework="pt", device="cpu") as f: + state = {k: f.get_tensor(k) for k in f.keys()} + cfg_path = ckpt / "config.json" + cfg = json.loads(cfg_path.read_text()) if cfg_path.exists() else {} + state = {k.removeprefix("module."): v for k, v in state.items()} + return TrackedTensors(state), cfg + + +class TurboVLADims: + """Auto-detect all TurboVLA dimensions from checkpoint tensor shapes.""" + + def __init__(self, tensors: dict[str, torch.Tensor], keys: set[str], cfg_json: dict | None = None, + dinov3_cfg: dict | None = None): + self.tensors = tensors + self.keys = keys + self.cfg = cfg_json or {} + self.dinov3_cfg = dinov3_cfg or {} + + q0 = self._get(f"{PREFIX_VIT}.layer.0.attention.q_proj.weight") + self.vit_dim = int(q0.shape[0]) + self.vit_layers = max_layer(keys, f"{PREFIX_VIT}.layer.") + + q0 = self._get(f"{PREFIX_TEXT}.encoder.layer.0.attention.self.query.weight") + self.text_dim = int(q0.shape[0]) + self.text_layers = max_layer(keys, f"{PREFIX_TEXT}.encoder.layer.") + + self.num_fusion_layers = max_layer(keys, f"{PREFIX_FUSION}.") + self.num_text_layers = max_layer(keys, f"{PREFIX_VL_TEXT}.") + + # BiMultiHeadAttention projects to interaction.enhancer_inner_dim, whereas + # the text enhancer and ACT decoder operate at interaction.hidden_dim. + interaction_cfg = self.cfg.get("interaction", {}) + self.hidden_dim = int(interaction_cfg.get("hidden_dim", 256)) + raw_nheads = int(interaction_cfg.get("nheads", 8)) + self.fusion_heads = max(1, raw_nheads // 2) + self.fusion_dim = int(interaction_cfg.get("enhancer_inner_dim", 0)) + if self.fusion_dim <= 0 or self.fusion_dim % self.fusion_heads: + raise SystemExit("interaction.enhancer_inner_dim must be divisible by fusion_heads") + self.fusion_head_dim = self.fusion_dim // self.fusion_heads + + self.text_enhancer_heads = max(1, raw_nheads // 2) + self.text_enhancer_head_dim = self.hidden_dim // self.text_enhancer_heads + + # TurboVLA constructs the ACT decoder with interaction.nheads. + self.action_heads = raw_nheads + self.action_head_dim = self.hidden_dim // self.action_heads + + self.num_state_tokens = int(self.cfg.get("action", {}).get("num_state_tokens", 2)) + + self.num_action_decoder_layers = max_layer(keys, f"{PREFIX_ACT_DEC}.decoder.layers.") + action_q = self._get(f"{PREFIX_ACT_DEC}.action_queries.weight") + self.action_horizon = int(action_q.shape[0]) + act_proj_2 = self._get(f"{PREFIX_ACT_DEC}.action_projection.layers.2.weight") + self.action_dim = int(act_proj_2.shape[0]) + + state_weight = self._get(f"{PREFIX_STATE_PROJ}.net.1.weight") + self.state_dim = int(state_weight.shape[1]) + + view_emb = self._get(KEY_VIEW_EMB) + if view_emb.ndim == 3: + if view_emb.shape[0] != 1: + raise SystemExit(f"Unexpected view_embedding shape: {tuple(view_emb.shape)}") + self.num_views = int(view_emb.shape[1]) + elif view_emb.ndim == 2: + self.num_views = int(view_emb.shape[0]) + else: + raise SystemExit(f"Unexpected view_embedding shape: {tuple(view_emb.shape)}") + + cfg_image_size = self.cfg.get("vision", {}).get("image_size", self.cfg.get("image_size")) + if cfg_image_size is not None: + if isinstance(cfg_image_size, (list, tuple)): + if not cfg_image_size: + raise SystemExit("empty image_size in config") + self.image_size = int(cfg_image_size[0]) + else: + self.image_size = int(cfg_image_size) + else: + inferred = self._infer_image_size() + if inferred is None: + raise SystemExit( + "Cannot determine image_size. " + "Checkpoint does not encode enough information " + "and config.json has no image_size." + ) + self.image_size = inferred + + patch_weight = self._get(f"{PREFIX_VIT}.embeddings.patch_embeddings.weight") + if len(patch_weight.shape) == 4: + self.patch_size = patch_weight.shape[2] + else: + raise SystemExit("Cannot determine patch_size from patch_embeddings") + + reg_key = f"{PREFIX_VIT}.embeddings.register_tokens" + if reg_key in tensors: + reg = tensors[reg_key] + self.num_register_tokens = int(reg.shape[1]) if reg.ndim == 3 else 0 + else: + self.num_register_tokens = 0 + + # DINOv3 RoPE is architectural data, not a shape-derived default. The + # checkpoint's TurboVLA config names the backbone, while the matching + # DINO config supplies rope_theta and coordinate normalization. + if "rope_theta" not in self.dinov3_cfg: + backbone = self.cfg.get("vision", {}).get("model_name_or_path", "") + if "dinov3-vitb16" not in backbone: + raise SystemExit(f"unknown DINOv3 backbone {backbone!r}; pass --dinov3-config") + self.dinov3_cfg = DINOV3_VITB16 + self.rope_theta = float(self.dinov3_cfg["rope_theta"]) + self.rope_normalize_coords = str(self.dinov3_cfg.get("rope_normalize_coords", "separate")) + if self.rope_normalize_coords != "separate": + raise SystemExit( + f"unsupported DINOv3 rope coordinate normalization: {self.rope_normalize_coords!r}" + ) + if int(self.dinov3_cfg.get("num_register_tokens", self.num_register_tokens)) != self.num_register_tokens: + raise SystemExit("DINOv3 config num_register_tokens disagrees with checkpoint weights") + + word_emb = self._get(f"{PREFIX_TEXT}.embeddings.word_embeddings.weight") + self.vocab_size = int(word_emb.shape[0]) + + + def _get(self, key: str) -> torch.Tensor: + if key not in self.tensors: + raise KeyError(f"Required tensor not found: {key}") + return self.tensors[key] + + def _infer_image_size(self) -> int | None: + """Infer image size from tensor shapes if config not available. + + Note: DINOv3 uses RoPE (rotary position embeddings) instead of learned + position embeddings, so position_embeddings.weight may not exist or may + have different semantics. Use patch_size * grid_size from config instead. + """ + if self.dinov3_cfg: + image_size = self.dinov3_cfg.get("image_size") + if image_size: + return int(image_size) + + pos_key = f"{PREFIX_VIT}.embeddings.position_embeddings.weight" + if pos_key in self.tensors: + pos = self.tensors[pos_key] + seq_len = pos.shape[0] if pos.ndim == 2 else pos.shape[1] + num_patches = seq_len - 1 + patch = self.tensors.get(f"{PREFIX_VIT}.embeddings.patch_embeddings.weight") + if patch is not None and len(patch.shape) == 4: + patch_size = patch.shape[2] + grid_size = int(num_patches ** 0.5) + if grid_size * grid_size == num_patches: + return grid_size * patch_size + + return None + + def __repr__(self): + return ( + f"TurboVLADims(vit={self.vit_dim}d×{self.vit_layers}L, " + f"text={self.text_dim}d×{self.text_layers}L, " + f"hidden={self.hidden_dim}, " + f"fusion={self.num_fusion_layers}L, " + f"decoder={self.num_action_decoder_layers}L, " + f"action={self.action_dim}d×{self.action_horizon}horizon, " + f"state={self.state_dim}d, views={self.num_views})" + ) + + +def write_vision_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADims) -> None: + """Write DINOv3 ViT vision encoder. + + Tensor naming follows TurboVLA naming convention with attn_q/k/v/o. + LayerScale is baked into weights. + """ + root = PREFIX_VIT + + # PyTorch patch_embedding weight: [H, IC, KH, KW] = [768, 3, 16, 16] + # For matmul-based patch embedding: flatten weight to [H, IC*KH*KW] = [768, 768] + patch_weight = tensors[f"{root}.embeddings.patch_embeddings.weight"] + patch_weight = patch_weight.reshape(dims.vit_dim, -1) + add_tensor(writer, "vit.cls_token", tensors[f"{root}.embeddings.cls_token"].squeeze(0)) + add_tensor(writer, "vit.patch_embed.weight", patch_weight) + add_tensor(writer, "vit.patch_embed.bias", tensors[f"{root}.embeddings.patch_embeddings.bias"]) + + reg_key = f"{root}.embeddings.register_tokens" + if reg_key in tensors: + add_tensor(writer, "vit.register_tokens", tensors[reg_key].squeeze(0)) + + for i in range(dims.vit_layers): + lr = f"{root}.layer.{i}" + + ls1 = tensors[f"{lr}.layer_scale1.lambda1"].clone() + ls2 = tensors[f"{lr}.layer_scale2.lambda1"].clone() + + w_q = tensors[f"{lr}.attention.q_proj.weight"] + b_q = tensors[f"{lr}.attention.q_proj.bias"] + w_k = tensors[f"{lr}.attention.k_proj.weight"] + k_bias_key = f"{lr}.attention.k_proj.bias" + b_k = tensors[k_bias_key] if k_bias_key in tensors else torch.zeros(dims.vit_dim, dtype=torch.float32) + w_v = tensors[f"{lr}.attention.v_proj.weight"] + b_v = tensors[f"{lr}.attention.v_proj.bias"] + w_o = tensors[f"{lr}.attention.o_proj.weight"] * ls1.view(-1, 1) + b_o = tensors[f"{lr}.attention.o_proj.bias"] * ls1 + + add_tensor(writer, f"vit.blk.{i}.attn_q.weight", w_q) + add_tensor(writer, f"vit.blk.{i}.attn_q.bias", b_q) + add_tensor(writer, f"vit.blk.{i}.attn_k.weight", w_k) + add_tensor(writer, f"vit.blk.{i}.attn_k.bias", b_k) + add_tensor(writer, f"vit.blk.{i}.attn_v.weight", w_v) + add_tensor(writer, f"vit.blk.{i}.attn_v.bias", b_v) + add_tensor(writer, f"vit.blk.{i}.attn_o.weight", w_o) + add_tensor(writer, f"vit.blk.{i}.attn_o.bias", b_o) + + add_tensor(writer, f"vit.blk.{i}.ln1.weight", tensors[f"{lr}.norm1.weight"]) + add_tensor(writer, f"vit.blk.{i}.ln1.bias", tensors[f"{lr}.norm1.bias"]) + add_tensor(writer, f"vit.blk.{i}.ln2.weight", tensors[f"{lr}.norm2.weight"]) + add_tensor(writer, f"vit.blk.{i}.ln2.bias", tensors[f"{lr}.norm2.bias"]) + + w_fc1 = tensors[f"{lr}.mlp.up_proj.weight"] + b_fc1 = tensors[f"{lr}.mlp.up_proj.bias"] + w_fc2 = tensors[f"{lr}.mlp.down_proj.weight"] * ls2.view(-1, 1) + b_fc2 = tensors[f"{lr}.mlp.down_proj.bias"] * ls2 + + add_tensor(writer, f"vit.blk.{i}.fc1.weight", w_fc1) + add_tensor(writer, f"vit.blk.{i}.fc1.bias", b_fc1) + add_tensor(writer, f"vit.blk.{i}.fc2.weight", w_fc2) + add_tensor(writer, f"vit.blk.{i}.fc2.bias", b_fc2) + + +def write_text_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADims) -> None: + """Write BERT text encoder.""" + root = PREFIX_TEXT + + add_tensor(writer, "text.embed.word_embeddings", tensors[f"{root}.embeddings.word_embeddings.weight"]) + add_tensor(writer, "text.embed.position_embeddings", tensors[f"{root}.embeddings.position_embeddings.weight"]) + add_tensor(writer, "text.embed.token_type_embeddings", tensors[f"{root}.embeddings.token_type_embeddings.weight"]) + add_tensor(writer, "text.embed.LayerNorm.weight", tensors[f"{root}.embeddings.LayerNorm.weight"]) + add_tensor(writer, "text.embed.LayerNorm.bias", tensors[f"{root}.embeddings.LayerNorm.bias"]) + + for i in range(dims.text_layers): + lr = f"{root}.encoder.layer.{i}" + + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.query.weight", tensors[f"{lr}.attention.self.query.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.query.bias", tensors[f"{lr}.attention.self.query.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.key.weight", tensors[f"{lr}.attention.self.key.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.key.bias", tensors[f"{lr}.attention.self.key.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.value.weight", tensors[f"{lr}.attention.self.value.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.value.bias", tensors[f"{lr}.attention.self.value.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.output.dense.weight", tensors[f"{lr}.attention.output.dense.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.output.dense.bias", tensors[f"{lr}.attention.output.dense.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.output.LayerNorm.weight", tensors[f"{lr}.attention.output.LayerNorm.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.output.LayerNorm.bias", tensors[f"{lr}.attention.output.LayerNorm.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.intermediate.dense.weight", tensors[f"{lr}.intermediate.dense.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.intermediate.dense.bias", tensors[f"{lr}.intermediate.dense.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.output.dense.weight", tensors[f"{lr}.output.dense.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.output.dense.bias", tensors[f"{lr}.output.dense.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.output.LayerNorm.weight", tensors[f"{lr}.output.LayerNorm.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.output.LayerNorm.bias", tensors[f"{lr}.output.LayerNorm.bias"]) + + add_tensor(writer, "text_proj.weight", tensors[KEY_TEXT_PROJ + ".weight"]) + add_tensor(writer, "text_proj.bias", tensors[KEY_TEXT_PROJ + ".bias"]) + + +def write_vision_projection(writer: gguf.GGUFWriter, tensors: dict) -> None: + """Write vision projection MLP.""" + root = PREFIX_VIT_PROJ + + add_tensor(writer, "vit_proj.input_norm.weight", tensors[f"{root}.input_norm.weight"]) + add_tensor(writer, "vit_proj.input_norm.bias", tensors[f"{root}.input_norm.bias"]) + add_tensor(writer, "vit_proj.mlp.0.weight", tensors[f"{root}.mlp.0.weight"]) + add_tensor(writer, "vit_proj.mlp.0.bias", tensors[f"{root}.mlp.0.bias"]) + add_tensor(writer, "vit_proj.mlp.3.weight", tensors[f"{root}.mlp.3.weight"]) + add_tensor(writer, "vit_proj.mlp.3.bias", tensors[f"{root}.mlp.3.bias"]) + add_tensor(writer, "vit_proj.skip.weight", tensors[f"{root}.skip.weight"]) + add_tensor(writer, "vit_proj.output_norm.weight", tensors[f"{root}.output_norm.weight"]) + add_tensor(writer, "vit_proj.output_norm.bias", tensors[f"{root}.output_norm.bias"]) + + +def write_vision_language_interaction(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADims) -> None: + """Write VL fusion and text self-attention layers.""" + fusion_root = PREFIX_FUSION + text_root = PREFIX_VL_TEXT + + for i in range(dims.num_fusion_layers): + lr = f"{fusion_root}.{i}" + + add_tensor(writer, f"vl_fusion.{i}.v_proj.weight", tensors[f"{lr}.attn.v_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.v_proj.bias", tensors[f"{lr}.attn.v_proj.bias"]) + add_tensor(writer, f"vl_fusion.{i}.l_proj.weight", tensors[f"{lr}.attn.l_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.l_proj.bias", tensors[f"{lr}.attn.l_proj.bias"]) + add_tensor(writer, f"vl_fusion.{i}.values_v.weight", tensors[f"{lr}.attn.values_v_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.values_v.bias", tensors[f"{lr}.attn.values_v_proj.bias"]) + add_tensor(writer, f"vl_fusion.{i}.values_l.weight", tensors[f"{lr}.attn.values_l_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.values_l.bias", tensors[f"{lr}.attn.values_l_proj.bias"]) + add_tensor(writer, f"vl_fusion.{i}.out_v.weight", tensors[f"{lr}.attn.out_v_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.out_v.bias", tensors[f"{lr}.attn.out_v_proj.bias"]) + add_tensor(writer, f"vl_fusion.{i}.out_l.weight", tensors[f"{lr}.attn.out_l_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.out_l.bias", tensors[f"{lr}.attn.out_l_proj.bias"]) + + add_tensor(writer, f"vl_fusion.{i}.norm_v.weight", tensors[f"{lr}.layer_norm_v.weight"]) + add_tensor(writer, f"vl_fusion.{i}.norm_v.bias", tensors[f"{lr}.layer_norm_v.bias"]) + add_tensor(writer, f"vl_fusion.{i}.norm_l.weight", tensors[f"{lr}.layer_norm_l.weight"]) + add_tensor(writer, f"vl_fusion.{i}.norm_l.bias", tensors[f"{lr}.layer_norm_l.bias"]) + + add_tensor(writer, f"vl_fusion.{i}.gamma_v", tensors[f"{lr}.gamma_v"]) + add_tensor(writer, f"vl_fusion.{i}.gamma_l", tensors[f"{lr}.gamma_l"]) + + for i in range(dims.num_text_layers): + lr = f"{text_root}.{i}" + + add_tensor(writer, f"vl_text.{i}.attn_qkv.weight", tensors[f"{lr}.self_attn.in_proj_weight"]) + add_tensor(writer, f"vl_text.{i}.attn_qkv.bias", tensors[f"{lr}.self_attn.in_proj_bias"]) + add_tensor(writer, f"vl_text.{i}.attn_o.weight", tensors[f"{lr}.self_attn.out_proj.weight"]) + add_tensor(writer, f"vl_text.{i}.attn_o.bias", tensors[f"{lr}.self_attn.out_proj.bias"]) + + add_tensor(writer, f"vl_text.{i}.ln1.weight", tensors[f"{lr}.norm1.weight"]) + add_tensor(writer, f"vl_text.{i}.ln1.bias", tensors[f"{lr}.norm1.bias"]) + + add_tensor(writer, f"vl_text.{i}.fc1.weight", tensors[f"{lr}.linear1.weight"]) + add_tensor(writer, f"vl_text.{i}.fc1.bias", tensors[f"{lr}.linear1.bias"]) + add_tensor(writer, f"vl_text.{i}.fc2.weight", tensors[f"{lr}.linear2.weight"]) + add_tensor(writer, f"vl_text.{i}.fc2.bias", tensors[f"{lr}.linear2.bias"]) + add_tensor(writer, f"vl_text.{i}.ln2.weight", tensors[f"{lr}.norm2.weight"]) + add_tensor(writer, f"vl_text.{i}.ln2.bias", tensors[f"{lr}.norm2.bias"]) + + +def write_action_decoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADims) -> None: + """Write ACT-style action decoder.""" + root = PREFIX_ACT_DEC + + add_tensor(writer, "act.q.weight", tensors[f"{root}.action_queries.weight"]) + + for i in range(dims.num_action_decoder_layers): + lr = f"{root}.decoder.layers.{i}" + + add_tensor(writer, f"act.dec.{i}.self_qkv.weight", tensors[f"{lr}.self_attn.in_proj_weight"]) + add_tensor(writer, f"act.dec.{i}.self_qkv.bias", tensors[f"{lr}.self_attn.in_proj_bias"]) + add_tensor(writer, f"act.dec.{i}.self_out.weight", tensors[f"{lr}.self_attn.out_proj.weight"]) + add_tensor(writer, f"act.dec.{i}.self_out.bias", tensors[f"{lr}.self_attn.out_proj.bias"]) + + add_tensor(writer, f"act.dec.{i}.cross_qkv.weight", tensors[f"{lr}.multihead_attn.in_proj_weight"]) + add_tensor(writer, f"act.dec.{i}.cross_qkv.bias", tensors[f"{lr}.multihead_attn.in_proj_bias"]) + add_tensor(writer, f"act.dec.{i}.cross_out.weight", tensors[f"{lr}.multihead_attn.out_proj.weight"]) + add_tensor(writer, f"act.dec.{i}.cross_out.bias", tensors[f"{lr}.multihead_attn.out_proj.bias"]) + + add_tensor(writer, f"act.dec.{i}.ln1.weight", tensors[f"{lr}.norm1.weight"]) + add_tensor(writer, f"act.dec.{i}.ln1.bias", tensors[f"{lr}.norm1.bias"]) + add_tensor(writer, f"act.dec.{i}.ln2.weight", tensors[f"{lr}.norm2.weight"]) + add_tensor(writer, f"act.dec.{i}.ln2.bias", tensors[f"{lr}.norm2.bias"]) + add_tensor(writer, f"act.dec.{i}.ln3.weight", tensors[f"{lr}.norm3.weight"]) + add_tensor(writer, f"act.dec.{i}.ln3.bias", tensors[f"{lr}.norm3.bias"]) + + add_tensor(writer, f"act.dec.{i}.fc1.weight", tensors[f"{lr}.linear1.weight"]) + add_tensor(writer, f"act.dec.{i}.fc1.bias", tensors[f"{lr}.linear1.bias"]) + add_tensor(writer, f"act.dec.{i}.fc2.weight", tensors[f"{lr}.linear2.weight"]) + add_tensor(writer, f"act.dec.{i}.fc2.bias", tensors[f"{lr}.linear2.bias"]) + + act_proj_root = f"{root}.action_projection" + add_tensor(writer, "act.proj.0.weight", tensors[f"{act_proj_root}.layers.0.weight"]) + add_tensor(writer, "act.proj.0.bias", tensors[f"{act_proj_root}.layers.0.bias"]) + add_tensor(writer, "act.proj.1.weight", tensors[f"{act_proj_root}.layers.1.weight"]) + add_tensor(writer, "act.proj.1.bias", tensors[f"{act_proj_root}.layers.1.bias"]) + add_tensor(writer, "act.proj.2.weight", tensors[f"{act_proj_root}.layers.2.weight"]) + add_tensor(writer, "act.proj.2.bias", tensors[f"{act_proj_root}.layers.2.bias"]) + + +def write_state_projection(writer: gguf.GGUFWriter, tensors: dict) -> None: + """Write state projection MLP.""" + root = PREFIX_STATE_PROJ + + add_tensor(writer, "state.proj.0.weight", tensors[f"{root}.net.0.weight"]) + add_tensor(writer, "state.proj.0.bias", tensors[f"{root}.net.0.bias"]) + add_tensor(writer, "state.proj.1.weight", tensors[f"{root}.net.1.weight"]) + add_tensor(writer, "state.proj.1.bias", tensors[f"{root}.net.1.bias"]) + add_tensor(writer, "state.proj.4.weight", tensors[f"{root}.net.4.weight"]) + add_tensor(writer, "state.proj.4.bias", tensors[f"{root}.net.4.bias"]) + + add_tensor(writer, "state.proj.output_norm.weight", tensors[f"{root}.output_norm.weight"]) + add_tensor(writer, "state.proj.output_norm.bias", tensors[f"{root}.output_norm.bias"]) + add_tensor(writer, "state.proj.position", tensors[f"{root}.position"]) + + +def write_view_embeddings(writer: gguf.GGUFWriter, tensors: dict) -> None: + """Write view embeddings.""" + add_tensor(writer, "view_emb", tensors[KEY_VIEW_EMB]) + + +def write_text_groups(writer: gguf.GGUFWriter, text_cfg: dict) -> None: + """Store each training instruction's BERT tokens and padded length. + + The ACT decoder attends the padded text rows too, so the padded length is + part of the model's output. TurboVLA pads each instruction to the length + its training batch used (padding_length_by_instruction); the runtime only + sees token ids, so the table is keyed by them. + """ + groups = text_cfg.get("padding_length_by_instruction") or {} + writer.add_uint32(kv_prefix("text_groups.count"), len(groups)) + if not groups: + return + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(text_cfg.get("model_name_or_path", "bert-base-uncased")) + ids, lengths, pad_to = [], [], [] + for instruction, length in sorted(groups.items()): + seq = tok(str(instruction), truncation=True, max_length=int(length))["input_ids"] + ids += [int(t) for t in seq] + lengths.append(len(seq)) + pad_to.append(int(length)) + writer.add_array(kv_prefix("text_groups.tokens"), ids) + writer.add_array(kv_prefix("text_groups.lengths"), lengths) + writer.add_array(kv_prefix("text_groups.pad_to"), pad_to) + + +def verify_consumed_tensors(tensors: TrackedTensors) -> None: + """Fail on any checkpoint tensor that no writer read.""" + left = sorted(k for k in tensors if k not in tensors.read and not k.startswith(UNUSED)) + if left: + raise SystemExit(f"{len(left)} checkpoint tensors not converted: {left[:20]}") + print(f" All {len(tensors.read)} used tensors consumed.") + + +def main() -> int: + import argparse + parser = argparse.ArgumentParser( + description=f"Convert TurboVLA checkpoint to GGUF ({ARCH})" + ) + parser.add_argument("ckpt", type=Path, + help="TurboVLA .pth, or a directory with model.safetensors + config.json") + parser.add_argument("-o", "--out", type=Path, default=None, help="Output GGUF path") + parser.add_argument("--dinov3-config", type=Path, default=None, + help="config.json of the DINOv3 backbone (default: built-in ViT-B/16 values)") + parser.add_argument("--verify", action="store_true", help="Fail if a checkpoint tensor is left unconverted") + args = parser.parse_args() + + ckpt = args.ckpt.resolve() + out = args.out or (ckpt.with_suffix(".gguf") if ckpt.is_file() else ckpt / f"{ARCH}.gguf") + + print(f"Loading checkpoint from {ckpt}...") + tensors, cfg_json = load_checkpoint(ckpt) + keys = set(tensors.keys()) + print(f" Loaded {len(tensors)} tensors") + + dinov3_cfg = {} + if args.dinov3_config: + dinov3_cfg = json.loads(args.dinov3_config.read_text()) + print(f" DINOv3 config: {args.dinov3_config}") + + print(" Inferring dimensions from checkpoint...") + dims = TurboVLADims(tensors, keys, cfg_json, dinov3_cfg) + print(f" Detected: {dims}") + + print(f"Writing GGUF to {out}...") + out.parent.mkdir(parents=True, exist_ok=True) + writer = gguf.GGUFWriter(str(out), ARCH) + + kv = kv_prefix + writer.add_string(kv("architecture"), ARCH) + writer.add_uint32(kv("hidden"), dims.hidden_dim) + writer.add_uint32(kv("vit_dim"), dims.vit_dim) + writer.add_uint32(kv("vit_layers"), dims.vit_layers) + writer.add_uint32(kv("vit_head_dim"), 64) + writer.add_uint32(kv("vit_heads"), 12) + writer.add_uint32(kv("text_dim"), dims.text_dim) + writer.add_uint32(kv("text_layers"), dims.text_layers) + writer.add_uint32(kv("text_head_dim"), 64) + writer.add_uint32(kv("text_heads"), 12) + writer.add_uint32(kv("num_fusion_layers"), dims.num_fusion_layers) + # Fusion attention splits the configured head count in half. + writer.add_uint32(kv("fusion_heads"), dims.fusion_heads) + writer.add_uint32(kv("fusion_head_dim"), dims.fusion_head_dim) + writer.add_uint32(kv("num_text_layers"), dims.num_text_layers) + writer.add_uint32(kv("text_enhancer_heads"), dims.text_enhancer_heads) + writer.add_uint32(kv("text_enhancer_head_dim"), dims.text_enhancer_head_dim) + writer.add_uint32(kv("action_heads"), dims.action_heads) + writer.add_uint32(kv("action_head_dim"), dims.action_head_dim) + writer.add_uint32(kv("num_action_decoder_layers"), dims.num_action_decoder_layers) + writer.add_uint32(kv("action_dim"), dims.action_dim) + writer.add_uint32(kv("state_dim"), dims.state_dim) + writer.add_uint32(kv("num_state_tokens"), dims.num_state_tokens) + writer.add_uint32(kv("action_horizon"), dims.action_horizon) + writer.add_uint32(kv("image_size"), dims.image_size) + writer.add_uint32(kv("num_views"), dims.num_views) + writer.add_uint32(kv("patch_size"), dims.patch_size) + writer.add_uint32(kv("num_register_tokens"), dims.num_register_tokens) + writer.add_float32(kv("rope_theta"), float(dims.rope_theta)) + writer.add_uint32(kv("vocab_size"), dims.vocab_size) + # The runtime pads to a fixed length; a checkpoint without text.padding_length + # pads each batch to its longest instruction, which it does not implement. + padding_length = cfg_json.get("text", {}).get("padding_length") + if padding_length is None: + raise SystemExit("checkpoint config has no text.padding_length; padding to the longest " + "instruction is not supported") + writer.add_uint32(kv("max_text_length"), int(padding_length)) + # TurboVLA's BERT wrapper uses these exact punctuation IDs when it creates + # sub-sentence attention masks. Persist them so the GGUF runtime does not + # silently depend on a tokenizer installation. + text_cfg = cfg_json.get("text", {}) + model_name = text_cfg.get("model_name_or_path", "bert-base-uncased") + if model_name not in ("bert-base-uncased", "google-bert/bert-base-uncased"): + raise SystemExit( + "unsupported TurboVLA text tokenizer for deterministic mask conversion: " + f"{model_name!r}" + ) + writer.add_uint32(kv("pad_token_id"), 0) + writer.add_uint32(kv("cls_token_id"), 101) + writer.add_uint32(kv("sep_token_id"), 102) + writer.add_uint32(kv("period_token_id"), 1012) + writer.add_uint32(kv("question_token_id"), 1029) + write_text_groups(writer, text_cfg) + + print(" Writing vision encoder...") + write_vision_encoder(writer, tensors, dims) + + print(" Writing text encoder...") + write_text_encoder(writer, tensors, dims) + + print(" Writing vision projection...") + write_vision_projection(writer, tensors) + + print(" Writing VL interaction...") + write_vision_language_interaction(writer, tensors, dims) + + print(" Writing action decoder...") + write_action_decoder(writer, tensors, dims) + + print(" Writing state projection...") + write_state_projection(writer, tensors) + + print(" Writing view embeddings...") + write_view_embeddings(writer, tensors) + + if cfg_json: + writer.add_string(kv("config_json"), json.dumps(cfg_json)) + + if args.verify: + print(" Verifying tensor consumption...") + verify_consumed_tensors(tensors) + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + size_mb = out.stat().st_size / (1024 * 1024) + print(f" Done! Output: {out} ({size_mb:.1f} MiB)") + return 0 + + +if __name__ == "__main__": + import sys + sys.exit(main()) diff --git a/scripts/tokenize_prompt.py b/scripts/tokenize_prompt.py index 4968224..aa604d6 100644 --- a/scripts/tokenize_prompt.py +++ b/scripts/tokenize_prompt.py @@ -37,6 +37,7 @@ "vla_jepa": "Qwen/Qwen3-VL-2B-Instruct", "gr00t_n1_5": "lerobot/eagle2hg-processor-groot-n1p5", "gr00t_n1_7": "nvidia/Cosmos-Reason2-2B", + "turbovla": "bert-base-uncased", } TRUST_REMOTE_CODE = {"evo1", "gr00t_n1_5"} diff --git a/src/arch.h b/src/arch.h index 517f4cb..47be355 100644 --- a/src/arch.h +++ b/src/arch.h @@ -71,6 +71,7 @@ enum class Arch { VLA_ADAPTER,// OpenHelix VLA-Adapter DINOv2 + SigLIP + Bridge-Attention. OPENVLA_OFT,// DINOv2-L/14-reg4 + SigLIP-so400m/14 +Llama-2-7B + MLPResNet. VLA_JEPA, // LeRobot Qwen3-VL-2B-Instruct+V-JEPÀ+DiT-B FM. + TURBOVLA, // TurboVLA (DINOv3 + BERT + VL Fusion + ACT decoder). }; /** @@ -217,6 +218,15 @@ std::unique_ptr vla_jepa_create(const std::string& mmproj_path, const std::string& config_path, const Options& opts); +/** + * @brief Build a TurboVLA model. Vision is baked into @p ckpt_path. + * @copydetails smolvla_create + */ +std::unique_ptr turbovla_create(const std::string& mmproj_path, + const std::string& ckpt_path, + const std::string& config_path, + const Options& opts); + /** * @brief Inspect a GGUF and identify the architecture tag. * diff --git a/src/layers/ffn.h b/src/layers/ffn.h index f514d93..2e26690 100644 --- a/src/layers/ffn.h +++ b/src/layers/ffn.h @@ -26,6 +26,11 @@ inline ggml_tensor * ffn_gelu(ggml_context * C, ggml_tensor * W1, ggml_tensor * return linear(C, W2, b2, ggml_gelu(C, linear(C, W1, b1, x))); } +inline ggml_tensor * ffn_relu(ggml_context * C, ggml_tensor * W1, ggml_tensor * b1, + ggml_tensor * W2, ggml_tensor * b2, ggml_tensor * x) { + return linear(C, W2, b2, ggml_relu(C, linear(C, W1, b1, x))); +} + // DINOv2 and SigLIP-so400m were trained with the exact erf form. inline ggml_tensor * ffn_gelu_erf(ggml_context * C, ggml_tensor * W1, ggml_tensor * b1, ggml_tensor * W2, ggml_tensor * b2, ggml_tensor * x) { diff --git a/src/loader.cpp b/src/loader.cpp index 8e231c0..0137731 100644 --- a/src/loader.cpp +++ b/src/loader.cpp @@ -154,8 +154,9 @@ bool WeightLoader::upload(ggml_backend_t backend, ggml_backend_buffer_t * out_bu std::vector bytes = g_.read_convert(name, t->type, gn); if (bytes.empty() || bytes.size() != ggml_nbytes(t)) { - std::fprintf(stderr, "vla(%s): failed to load %s (%zu vs %zu bytes)\n", - arch_, name, bytes.size(), ggml_nbytes(t)); + std::fprintf(stderr, + "vla(%s): failed to load %s (got %zu bytes, expected %zu, type=%d)\n", + arch_, name, bytes.size(), ggml_nbytes(t), (int) t->type); return false; } ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); diff --git a/src/model.cpp b/src/model.cpp index 5242efb..1b14754 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -76,7 +76,8 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { try_str("bitvla.architecture", arch_str) || try_str("openvla_oft.architecture", arch_str) || try_str("vla_jepa.architecture", arch_str) || - try_str("vla_adapter.architecture", arch_str)) { + try_str("vla_adapter.architecture", arch_str) || + try_str("turbovla.architecture", arch_str)) { if (arch_str == "smolvla") { *out = Arch::SMOLVLA; ok = true; @@ -127,6 +128,10 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { *out = Arch::VLA_JEPA; ok = true; } + else if (arch_str == "turbovla") { + *out = Arch::TURBOVLA; + ok = true; + } } gguf_free(gctx); @@ -280,6 +285,10 @@ Model* model_load(const std::string& mmproj_path, const std::string& ckpt_path, std::printf("vla: arch = vla_jepa\n"); impl = vla_jepa_create(mmproj_path, ckpt_path, config_path, opts); break; + case Arch::TURBOVLA: + std::printf("vla: arch = turbovla\n"); + impl = turbovla_create(mmproj_path, ckpt_path, config_path, opts); + break; } if (!impl) return nullptr; diff --git a/src/models/turbovla.cpp b/src/models/turbovla.cpp new file mode 100644 index 0000000..236aa75 --- /dev/null +++ b/src/models/turbovla.cpp @@ -0,0 +1,851 @@ +// 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. + +// TurboVLA (H-EmbodVis/TurboVLA@b29ab142): DINOv3 ViT-B/16 over each camera, +// BERT-base over the instruction, six Grounding-DINO bi-attention fusion layers +// each followed by a text enhancer, and a three-layer ACT decoder whose twelve +// learned queries read the fused tokens plus two state tokens. One forward pass, +// no denoising loop, so the whole model is one cached graph. + +#include "arch.h" +#include "backend.h" +#include "gguf.h" +#include "gguf_reader.h" +#include "loader.h" +#include "model.h" +#include "options.h" +#include "scratch_ctx.h" +#include "layers/attn.h" +#include "layers/ffn.h" +#include "layers/linear.h" +#include "layers/norm.h" +#include "modules/preprocess.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vla { +namespace { + +constexpr float kImagenetMean[3] = {0.485f, 0.456f, 0.406f}; +constexpr float kImagenetStd[3] = {0.229f, 0.224f, 0.225f}; + +// Epsilons are the module defaults upstream: nn.LayerNorm everywhere except +// BERT, whose config sets 1e-12. +constexpr float kLnEps = 1e-5f; +constexpr float kBertEps = 1e-12f; + +struct VitLayerW { + ggml_tensor *ln1_w, *ln1_b, *qkv_w, *qkv_b, *o_w, *o_b; + ggml_tensor *ln2_w, *ln2_b, *fc1_w, *fc1_b, *fc2_w, *fc2_b; +}; + +struct BertLayerW { + ggml_tensor *qkv_w, *qkv_b, *o_w, *o_b, *ln1_w, *ln1_b; + ggml_tensor *fc1_w, *fc1_b, *fc2_w, *fc2_b, *ln2_w, *ln2_b; +}; + +// v_proj/values_v and l_proj/values_l are fused per side: v_proj(v) is the +// vision query and also the key the language side attends, l_proj(l) the +// other way round, so each side needs one GEMM for both of its projections. +struct FusionLayerW { + ggml_tensor *norm_v_w, *norm_v_b, *norm_l_w, *norm_l_b; + ggml_tensor *v_w, *v_b, *l_w, *l_b; + ggml_tensor *out_v_w, *out_v_b, *out_l_w, *out_l_b; + ggml_tensor *gamma_v, *gamma_l; +}; + +struct EnhancerLayerW { + ggml_tensor *qkv_w, *qkv_b, *o_w, *o_b, *ln1_w, *ln1_b; + ggml_tensor *fc1_w, *fc1_b, *fc2_w, *fc2_b, *ln2_w, *ln2_b; +}; + +struct DecoderLayerW { + ggml_tensor *ln1_w, *ln1_b, *self_qkv_w, *self_qkv_b, *self_o_w, *self_o_b; + ggml_tensor *ln2_w, *ln2_b, *cross_qkv_w, *cross_qkv_b, *cross_o_w, *cross_o_b; + ggml_tensor *ln3_w, *ln3_b, *fc1_w, *fc1_b, *fc2_w, *fc2_b; +}; + +// A training instruction's BERT token ids and the length its batch was padded +// to. See text_len() for why the length is part of the model. +struct TextGroup { + std::vector ids; + int64_t pad_to = 0; +}; + +} // namespace + +struct TurboVlaModelArch : public ModelArchBase { + TurboVlaModelArch() : ModelArchBase(Arch::TURBOVLA) {} + ~TurboVlaModelArch() override { + graph.release(); + if (const_buf) ggml_backend_buffer_free(const_buf); + if (ctx_const) ggml_free(ctx_const); + if (weight_buf) ggml_backend_buffer_free(weight_buf); + if (ctx_weights) ggml_free(ctx_weights); + if (backend) ggml_backend_free(backend); + } + + ggml_backend_t backend = nullptr; + int n_threads = default_cpu_threads(); + ggml_context * ctx_weights = nullptr; + ggml_backend_buffer_t weight_buf = nullptr; + ggml_context * ctx_const = nullptr; + ggml_backend_buffer_t const_buf = nullptr; + ggml_type mt = GGML_TYPE_F32; + + int64_t hidden = 256, n_views = 2, image_size = 256, patch = 16, n_reg = 4; + int64_t vit_dim = 768, vit_layers = 12, vit_heads = 12; + int64_t bert_dim = 768, bert_layers = 12, bert_heads = 12, vocab = 30522, bert_max_pos = 512; + int64_t fusion_layers = 6, fusion_heads = 4, fusion_dim = 1024, enh_heads = 4; + int64_t dec_layers = 3, dec_heads = 8, horizon = 12, action_dim = 7, state_dim = 8, n_state_tok = 2; + int64_t text_len_max = 21; + float rope_theta = 100.0f; + int32_t pad_id = 0, cls_id = 101, sep_id = 102, period_id = 1012, question_id = 1029; + std::vector text_groups; + + ggml_tensor *patch_w = nullptr, *patch_b = nullptr, *cls_tok = nullptr, *reg_tok = nullptr; + std::vector vit; + ggml_tensor *vp_in_w = nullptr, *vp_in_b = nullptr, *vp_fc1_w = nullptr, *vp_fc1_b = nullptr; + ggml_tensor *vp_fc2_w = nullptr, *vp_fc2_b = nullptr, *vp_skip_w = nullptr; + ggml_tensor *vp_out_w = nullptr, *vp_out_b = nullptr, *view_emb = nullptr; + ggml_tensor *word_emb = nullptr, *pos_emb = nullptr, *type_emb = nullptr; + ggml_tensor *emb_ln_w = nullptr, *emb_ln_b = nullptr; + std::vector bert; + ggml_tensor *text_proj_w = nullptr, *text_proj_b = nullptr; + std::vector fusion; + std::vector enhancer; + ggml_tensor *st_ln_w = nullptr, *st_ln_b = nullptr, *st_fc1_w = nullptr, *st_fc1_b = nullptr; + ggml_tensor *st_fc2_w = nullptr, *st_fc2_b = nullptr, *st_pos = nullptr, *st_out_w = nullptr, *st_out_b = nullptr; + ggml_tensor *act_q = nullptr; + std::vector dec; + ggml_tensor *ap0_w = nullptr, *ap0_b = nullptr, *ap1_w = nullptr, *ap1_b = nullptr; + ggml_tensor *ap2_w = nullptr, *ap2_b = nullptr; + + // DINOv3's RoPE depends only on the patch grid, so its tables are computed + // once and live next to the weights, shaped [head_dim, 1, T] to broadcast + // over heads. Rows for the CLS and register tokens are the identity (cos 1, + // sin 0), which lets the rotation cover the whole sequence instead of + // splitting off the prefix in every layer. rope_sin carries rotate_half's + // sign: negative on the first half of the head. + ggml_tensor *rope_cos = nullptr, *rope_sin = nullptr; + + struct Key { + int64_t bert_len = -1; + bool operator==(const Key & o) const { return bert_len == o.bert_len; } + }; + struct IO { + ggml_tensor *patches = nullptr, *ids = nullptr, *pos = nullptr, *bert_mask = nullptr; + ggml_tensor *enh_mask = nullptr, *fus_mask = nullptr, *state = nullptr, *actions = nullptr; + }; + graph_cache graph; + + int64_t grid() const { return image_size / patch; } + int64_t n_patches() const { return grid() * grid(); } + int64_t vit_seq() const { return 1 + n_reg + n_patches(); } + + int64_t text_len(const int32_t * ids, int64_t n) const; + ggml_cgraph * build(ggml_context * C, IO & io, int64_t bert_len) const; + bool upload_rope_tables(); + + std::vector predict(const Inputs& in) override; +}; + +namespace { + +ggml_tensor * split_heads(ggml_context * C, ggml_tensor * qkv, int64_t hd, int64_t heads, + int64_t T, int64_t nv, int part) { + const size_t es = ggml_element_size(qkv); + return ggml_view_4d(C, qkv, hd, heads, T, nv, (size_t) hd*es, qkv->nb[1], qkv->nb[2], + (size_t) part*hd*heads*es); +} + +// Fused-QKV self-attention over [3*dim, T, nv]; returns [dim, T, nv]. +ggml_tensor * self_attention(ggml_context * C, ggml_tensor * qkv, ggml_tensor * mask, + int64_t hd, int64_t heads, int64_t T, int64_t nv) { + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, split_heads(C, qkv, hd, heads, T, nv, 0), 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, split_heads(C, qkv, hd, heads, T, nv, 1), 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, split_heads(C, qkv, hd, heads, T, nv, 2), 1, 2, 0, 3)); + return attention(C, Q, K, V, mask, 1.0f / std::sqrt((float) hd), hd*heads, T, nv); +} + +// Half-split RoPE applied in the projection's own [hd, heads, T, nv] layout, so +// no head permute is needed first: rotate_half(x) is [-x1, x0], and the sign +// lives in the sin table, leaving one concat of two views. +ggml_tensor * rope_heads(ggml_context * C, ggml_tensor * x, ggml_tensor * cos_t, ggml_tensor * sin_signed) { + const int64_t half = x->ne[0] / 2; + const size_t es = ggml_element_size(x); + ggml_tensor * x0 = ggml_view_4d(C, x, half, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], 0); + ggml_tensor * x1 = ggml_view_4d(C, x, half, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], half*es); + return ggml_add(C, ggml_mul(C, x, cos_t), ggml_mul(C, ggml_concat(C, x1, x0, 0), sin_signed)); +} + +// DINOv3 attention. Flash attention takes Q and K as permuted views and casts +// K and V itself, so the rotated projections are never copied again. +ggml_tensor * vit_attention(ggml_context * C, ggml_tensor * qkv, ggml_tensor * cos_t, ggml_tensor * sin_signed, + int64_t hd, int64_t heads, int64_t T, int64_t nv, bool flash) { + const float scale = 1.0f / std::sqrt((float) hd); + ggml_tensor * q = rope_heads(C, split_heads(C, qkv, hd, heads, T, nv, 0), cos_t, sin_signed); + ggml_tensor * k = rope_heads(C, split_heads(C, qkv, hd, heads, T, nv, 1), cos_t, sin_signed); + ggml_tensor * v = split_heads(C, qkv, hd, heads, T, nv, 2); + if (flash) { + ggml_tensor * o = flash_attention(C, ggml_permute(C, q, 0, 2, 1, 3), ggml_permute(C, k, 0, 2, 1, 3), + ggml_permute(C, v, 0, 2, 1, 3), nullptr, scale); + return ggml_reshape_3d(C, o, hd*heads, T, nv); + } + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, q, 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, k, 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, v, 1, 2, 0, 3)); + return attention(C, Q, K, V, nullptr, scale, hd*heads, T, nv); +} + +// The released BERT wrapper (text/bert.py, from Grounding DINO) splits the +// instruction at [CLS], [SEP], '.' and '?': a token attends only inside its +// sub-sentence and its position restarts there. Tokens outside any span (the +// padding, and a special token at either end) see only themselves at position 0. +void special_token_blocks(const std::vector & ids, const TurboVlaModelArch & m, + std::vector & allowed, std::vector & pos) { + const int64_t n = (int64_t) ids.size(); + allowed.assign((size_t) n*n, 0); + pos.assign((size_t) n, 0); + for (int64_t i = 0; i < n; ++i) + allowed[(size_t) i*n + i] = 1; + int64_t prev = 0; + for (int64_t col = 0; col < n; ++col) { + const int32_t id = ids[(size_t) col]; + if (id != m.cls_id && id != m.sep_id && id != m.period_id && id != m.question_id) + continue; + if (col != 0 && col != n - 1) { + for (int64_t q = prev + 1; q <= col; ++q) { + pos[(size_t) q] = (int32_t) (q - prev - 1); + for (int64_t k = prev + 1; k <= col; ++k) + allowed[(size_t) q*n + k] = 1; + } + } + prev = col; + } +} + +} // namespace + +// The ACT decoder attends every text row, padding included (encode_condition +// concatenates them unmasked), so the padded length is part of the output. The +// checkpoint pins it per training instruction (11, 14 or 21 for LIBERO): BERT +// runs over that many tokens and the rows after it, up to text_len_max, are +// zeros that the text projection turns into its bias. Instructions the +// checkpoint does not list are padded to text_len_max. A caller that sends +// trailing [PAD] tokens has chosen the length itself. +int64_t TurboVlaModelArch::text_len(const int32_t * ids, int64_t n) const { + int64_t core = n; + while (core > 0 && ids[core - 1] == pad_id) + --core; + if (core < n) + return n; + for (const TextGroup & g : text_groups) + if ((int64_t) g.ids.size() == n && std::equal(g.ids.begin(), g.ids.end(), ids)) + return g.pad_to; + return text_len_max; +} + +bool TurboVlaModelArch::upload_rope_tables() { + // DINOv3 axial RoPE (HF modeling_dinov3_vit.py): patch centres normalized to + // [-1, 1] per axis, head_dim/4 periods base^(4i/head_dim), angles ordered + // [h, w, h, w] across the head, rotated half-split. See rope_heads. + const int64_t hd = vit_dim / vit_heads, q = hd / 4, T = vit_seq(), g = grid(), pre = 1 + n_reg; + std::vector c((size_t) hd*T, 1.0f), s((size_t) hd*T, 0.0f); + for (int64_t p = 0; p < n_patches(); ++p) { + const float h = 2.0f * ((float) (p / g) + 0.5f) / (float) g - 1.0f; + const float w = 2.0f * ((float) (p % g) + 0.5f) / (float) g - 1.0f; + float * cr = c.data() + (size_t) (pre + p)*hd; + float * sr = s.data() + (size_t) (pre + p)*hd; + for (int64_t i = 0; i < q; ++i) { + const float period = std::pow(rope_theta, 4.0f * (float) i / (float) hd); + const float ah = 6.28318530717958647692f * h / period; + const float aw = 6.28318530717958647692f * w / period; + for (int64_t blk = 0; blk < 4; ++blk) { + const float a = (blk % 2 == 0) ? ah : aw; + cr[blk*q + i] = std::cos(a); + sr[blk*q + i] = blk < 2 ? -std::sin(a) : std::sin(a); + } + } + } + + ggml_init_params p = { 2*ggml_tensor_overhead(), nullptr, true }; + ctx_const = ggml_init(p); + if (!ctx_const) + return false; + rope_cos = ggml_new_tensor_3d(ctx_const, GGML_TYPE_F32, hd, 1, T); + rope_sin = ggml_new_tensor_3d(ctx_const, GGML_TYPE_F32, hd, 1, T); + const_buf = alloc_weights(ctx_const, backend); + if (!const_buf) + return false; + ggml_backend_tensor_set(rope_cos, c.data(), 0, ggml_nbytes(rope_cos)); + ggml_backend_tensor_set(rope_sin, s.data(), 0, ggml_nbytes(rope_sin)); + return true; +} + +ggml_cgraph * TurboVlaModelArch::build(ggml_context * C, IO & io, int64_t bert_len) const { + const int64_t NP = n_patches(), T = vit_seq(), nv = n_views, VS = nv*NP, LT = text_len_max; + const int64_t pdim = 3*patch*patch; + + // DINOv3. Patches arrive flattened on the host; views are the batch dim. + io.patches = ggml_new_tensor_3d(C, GGML_TYPE_F32, pdim, NP, nv); + ggml_set_input(io.patches); + ggml_tensor * prefix = ggml_reshape_2d(C, cls_tok, vit_dim, 1); + if (reg_tok) + prefix = ggml_concat(C, prefix, reg_tok, 1); + prefix = ggml_repeat_4d(C, prefix, vit_dim, 1 + n_reg, nv, 1); + ggml_tensor * x = ggml_concat(C, prefix, linear(C, patch_w, patch_b, io.patches), 1); + for (const VitLayerW & l : vit) { + ggml_tensor * qkv = linear(C, l.qkv_w, l.qkv_b, layer_norm(C, x, l.ln1_w, l.ln1_b, kLnEps)); + ggml_tensor * att = vit_attention(C, qkv, rope_cos, rope_sin, vit_dim / vit_heads, vit_heads, T, nv, + flash_attn_enabled()); + x = ggml_add(C, x, linear(C, l.o_w, l.o_b, att)); + x = ggml_add(C, x, ffn_gelu_erf(C, l.fc1_w, l.fc1_b, l.fc2_w, l.fc2_b, + layer_norm(C, x, l.ln2_w, l.ln2_b, kLnEps))); + } + // TurboVLA reads hidden_states[-1], which is before DINOv3's final norm + // (models/vision_encoder.py:109), then drops the CLS and register tokens. + x = ggml_view_3d(C, x, vit_dim, NP, nv, x->nb[1], x->nb[2], (size_t) (1 + n_reg)*x->nb[1]); + x = ggml_reshape_2d(C, ggml_cont(C, x), vit_dim, VS); + + // VisionProjection, then one learned embedding per camera. + ggml_tensor * mlp = linear(C, vp_fc1_w, vp_fc1_b, layer_norm(C, x, vp_in_w, vp_in_b, kLnEps)); + mlp = linear(C, vp_fc2_w, vp_fc2_b, ggml_gelu_erf(C, mlp)); + ggml_tensor * v = layer_norm(C, ggml_add(C, ggml_mul_mat(C, vp_skip_w, x), mlp), vp_out_w, vp_out_b, kLnEps); + v = ggml_reshape_2d(C, ggml_add(C, ggml_reshape_3d(C, v, hidden, NP, nv), + ggml_reshape_3d(C, view_emb, hidden, 1, nv)), hidden, VS); + + // BERT over the checkpoint's padded length, token type 0 throughout. + io.ids = ggml_new_tensor_1d(C, GGML_TYPE_I32, bert_len); + io.pos = ggml_new_tensor_1d(C, GGML_TYPE_I32, bert_len); + io.bert_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, bert_len, bert_len); + ggml_set_input(io.ids); + ggml_set_input(io.pos); + ggml_set_input(io.bert_mask); + ggml_tensor * t = ggml_add(C, ggml_get_rows(C, word_emb, io.ids), ggml_get_rows(C, pos_emb, io.pos)); + t = ggml_add(C, t, ggml_view_1d(C, type_emb, bert_dim, 0)); + t = layer_norm(C, t, emb_ln_w, emb_ln_b, kBertEps); + for (const BertLayerW & l : bert) { + ggml_tensor * att = self_attention(C, linear(C, l.qkv_w, l.qkv_b, t), io.bert_mask, + bert_dim / bert_heads, bert_heads, bert_len, 1); + t = layer_norm(C, ggml_add(C, t, linear(C, l.o_w, l.o_b, att)), l.ln1_w, l.ln1_b, kBertEps); + t = layer_norm(C, ggml_add(C, t, ffn_gelu_erf(C, l.fc1_w, l.fc1_b, l.fc2_w, l.fc2_b, t)), + l.ln2_w, l.ln2_b, kBertEps); + } + ggml_tensor * lang = linear(C, text_proj_w, text_proj_b, t); + if (bert_len < LT) + lang = ggml_concat(C, lang, ggml_repeat_4d(C, text_proj_b, hidden, LT - bert_len, 1, 1), 1); + + // Fusion. One score matrix serves both directions upstream (the language + // side softmaxes its transpose); two small products are cheaper than a + // transpose here. Only vision queries mask padded language keys. + io.enh_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, LT, LT); + io.fus_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, LT, VS); + ggml_set_input(io.enh_mask); + ggml_set_input(io.fus_mask); + const int64_t fhd = fusion_dim / fusion_heads, ehd = hidden / enh_heads; + const float fscale = 1.0f / std::sqrt((float) fhd); + for (int64_t i = 0; i < fusion_layers; ++i) { + const FusionLayerW & f = fusion[(size_t) i]; + ggml_tensor * vn = layer_norm(C, v, f.norm_v_w, f.norm_v_b, kLnEps); + ggml_tensor * ln = layer_norm(C, lang, f.norm_l_w, f.norm_l_b, kLnEps); + ggml_tensor * vp = linear(C, f.v_w, f.v_b, vn); // [q_v | values_v] + ggml_tensor * lp = linear(C, f.l_w, f.l_b, ln); // [k_l | values_l] + auto heads = [&](ggml_tensor * p, int64_t L, int part) { + return ggml_view_3d(C, p, fhd, fusion_heads, L, (size_t) fhd*ggml_element_size(p), p->nb[1], + (size_t) part*fusion_dim*ggml_element_size(p)); + }; + ggml_tensor * qv = ggml_cont(C, ggml_permute(C, heads(vp, VS, 0), 0, 2, 1, 3)); + ggml_tensor * kl = ggml_cont(C, ggml_permute(C, heads(lp, LT, 0), 0, 2, 1, 3)); + ggml_tensor * vv = ggml_cont(C, ggml_permute(C, heads(vp, VS, 1), 1, 2, 0, 3)); + ggml_tensor * vl = ggml_cont(C, ggml_permute(C, heads(lp, LT, 1), 1, 2, 0, 3)); + ggml_tensor * dv = attention(C, qv, kl, vl, io.fus_mask, fscale, fusion_dim, VS); + ggml_tensor * dl = attention(C, kl, qv, vv, nullptr, fscale, fusion_dim, LT); + // residual_style "normalized": the gated deltas land on the normed tokens. + v = ggml_add(C, vn, ggml_mul(C, linear(C, f.out_v_w, f.out_v_b, dv), f.gamma_v)); + lang = ggml_add(C, ln, ggml_mul(C, linear(C, f.out_l_w, f.out_l_b, dl), f.gamma_l)); + + // Text enhancer: post-norm, ReLU, the sub-sentence mask and no + // key-padding mask (the key_padding_mask argument is commented out + // upstream, models/components/transformer.py). + const EnhancerLayerW & e = enhancer[(size_t) i]; + ggml_tensor * att = self_attention(C, linear(C, e.qkv_w, e.qkv_b, lang), io.enh_mask, ehd, enh_heads, LT, 1); + lang = layer_norm(C, ggml_add(C, lang, linear(C, e.o_w, e.o_b, att)), e.ln1_w, e.ln1_b, kLnEps); + lang = layer_norm(C, ggml_add(C, lang, ffn_relu(C, e.fc1_w, e.fc1_b, e.fc2_w, e.fc2_b, lang)), + e.ln2_w, e.ln2_b, kLnEps); + } + + // State tokens: LayerNorm, MLP to num_state_tokens*hidden, + position, norm. + io.state = ggml_new_tensor_1d(C, GGML_TYPE_F32, state_dim); + ggml_set_input(io.state); + ggml_tensor * s = linear(C, st_fc1_w, st_fc1_b, layer_norm(C, io.state, st_ln_w, st_ln_b, kLnEps)); + s = ggml_reshape_2d(C, linear(C, st_fc2_w, st_fc2_b, ggml_gelu_erf(C, s)), hidden, n_state_tok); + s = layer_norm(C, ggml_add(C, s, ggml_reshape_2d(C, st_pos, hidden, n_state_tok)), st_out_w, st_out_b, kLnEps); + + ggml_tensor * mem = ggml_concat(C, ggml_concat(C, v, lang, 1), s, 1); + const int64_t M = VS + LT + n_state_tok; + + // ACT decoder: nn.TransformerDecoderLayer, norm_first, ReLU, no final norm. + const int64_t dhd = hidden / dec_heads; + const float dscale = 1.0f / std::sqrt((float) dhd); + ggml_tensor * a = act_q; + for (const DecoderLayerW & l : dec) { + ggml_tensor * qkv = linear(C, l.self_qkv_w, l.self_qkv_b, layer_norm(C, a, l.ln1_w, l.ln1_b, kLnEps)); + ggml_tensor * att = self_attention(C, qkv, nullptr, dhd, dec_heads, horizon, 1); + a = ggml_add(C, a, linear(C, l.self_o_w, l.self_o_b, att)); + + // in_proj is [W_q; W_k; W_v]: the queries keep the Q block of their + // projection, the memory its K and V blocks. Slicing the weight instead + // saves little and breaks ggml-openvino once the weights are BF16. + ggml_tensor * q = linear(C, l.cross_qkv_w, l.cross_qkv_b, layer_norm(C, a, l.ln2_w, l.ln2_b, kLnEps)); + ggml_tensor * kv = linear(C, l.cross_qkv_w, l.cross_qkv_b, mem); + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, split_heads(C, q, dhd, dec_heads, horizon, 1, 0), 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, split_heads(C, kv, dhd, dec_heads, M, 1, 1), 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, split_heads(C, kv, dhd, dec_heads, M, 1, 2), 1, 2, 0, 3)); + att = attention(C, Q, K, V, nullptr, dscale, hidden, horizon); + a = ggml_add(C, a, linear(C, l.cross_o_w, l.cross_o_b, att)); + a = ggml_add(C, a, ffn_relu(C, l.fc1_w, l.fc1_b, l.fc2_w, l.fc2_b, + layer_norm(C, a, l.ln3_w, l.ln3_b, kLnEps))); + } + a = ggml_relu(C, linear(C, ap0_w, ap0_b, a)); + a = ggml_relu(C, linear(C, ap1_w, ap1_b, a)); + io.actions = ggml_tanh(C, linear(C, ap2_w, ap2_b, a)); + ggml_set_output(io.actions); + + ggml_cgraph * gf = ggml_new_graph_custom(C, 8192, false); + ggml_build_forward_expand(gf, io.actions); + return gf; +} + +namespace { + +bool load_config(const gguf_reader & g, TurboVlaModelArch & m) { + auto U = [&](const char * k, int64_t & dst) { + char key[96]; + std::snprintf(key, sizeof(key), "turbovla.%s", k); + if (!g.has(key)) + return; + dst = (int64_t) g.u32(key); + }; + auto I = [&](const char * k, int32_t & dst) { + int64_t v = dst; + U(k, v); + dst = (int32_t) v; + }; + U("hidden", m.hidden); U("num_views", m.n_views); + U("image_size", m.image_size); U("patch_size", m.patch); + U("num_register_tokens", m.n_reg); + U("vit_dim", m.vit_dim); U("vit_layers", m.vit_layers); U("vit_heads", m.vit_heads); + U("text_dim", m.bert_dim); U("text_layers", m.bert_layers); U("text_heads", m.bert_heads); + U("vocab_size", m.vocab); + U("num_fusion_layers", m.fusion_layers); U("fusion_heads", m.fusion_heads); + U("text_enhancer_heads", m.enh_heads); + U("num_action_decoder_layers", m.dec_layers); U("action_heads", m.dec_heads); + U("action_horizon", m.horizon); U("action_dim", m.action_dim); + U("state_dim", m.state_dim); U("num_state_tokens", m.n_state_tok); + U("max_text_length", m.text_len_max); + I("pad_token_id", m.pad_id); I("cls_token_id", m.cls_id); I("sep_token_id", m.sep_id); + I("period_token_id", m.period_id); I("question_token_id", m.question_id); + if (g.has("turbovla.rope_theta")) + m.rope_theta = g.f32("turbovla.rope_theta"); + int64_t fhd = m.fusion_dim / m.fusion_heads; + U("fusion_head_dim", fhd); + m.fusion_dim = fhd * m.fusion_heads; + + if (m.image_size % m.patch || m.vit_dim % m.vit_heads || (m.vit_dim / m.vit_heads) % 4 || + m.bert_dim % m.bert_heads || m.hidden % m.enh_heads || m.hidden % m.dec_heads) { + std::fprintf(stderr, "vla(turbovla): inconsistent dimensions in GGUF metadata\n"); + return false; + } + + // The padding table (see text_len). A GGUF without it predates the table and + // would pad every instruction to max_text_length, which is wrong for the + // shorter LIBERO groups. + if (!g.has("turbovla.text_groups.count")) { + std::fprintf(stderr, "vla(turbovla): GGUF has no text padding table; re-convert it with " + "scripts/convert_turbovla_to_gguf.py\n"); + return false; + } + const size_t n = g.u32("turbovla.text_groups.count"); + if (n == 0) + return true; + const int64_t kt = gguf_find_key(g.gctx, "turbovla.text_groups.tokens"); + const int64_t kl = gguf_find_key(g.gctx, "turbovla.text_groups.lengths"); + const int64_t kp = gguf_find_key(g.gctx, "turbovla.text_groups.pad_to"); + if (kt < 0 || kl < 0 || kp < 0) { + std::fprintf(stderr, "vla(turbovla): malformed turbovla.text_groups metadata\n"); + return false; + } + if (gguf_get_arr_type(g.gctx, kt) != GGUF_TYPE_INT32 || gguf_get_arr_type(g.gctx, kl) != GGUF_TYPE_INT32 || + gguf_get_arr_type(g.gctx, kp) != GGUF_TYPE_INT32 || gguf_get_arr_n(g.gctx, kl) != n || + gguf_get_arr_n(g.gctx, kp) != n) { + std::fprintf(stderr, "vla(turbovla): malformed turbovla.text_groups metadata\n"); + return false; + } + const int32_t * tok = (const int32_t *) gguf_get_arr_data(g.gctx, kt); + const int32_t * len = (const int32_t *) gguf_get_arr_data(g.gctx, kl); + const int32_t * pad = (const int32_t *) gguf_get_arr_data(g.gctx, kp); + const size_t n_tok = gguf_get_arr_n(g.gctx, kt); + size_t off = 0; + for (size_t i = 0; i < n; ++i) { + if (len[i] <= 0 || off + (size_t) len[i] > n_tok || pad[i] < len[i] || pad[i] > m.text_len_max) { + std::fprintf(stderr, "vla(turbovla): malformed turbovla.text_groups entry %zu\n", i); + return false; + } + m.text_groups.push_back({ std::vector(tok + off, tok + off + len[i]), pad[i] }); + off += (size_t) len[i]; + } + return true; +} + +bool load_weights(TurboVlaModelArch & m, gguf_reader & g) { + ggml_init_params wp = { (size_t) 16*1024*1024, nullptr, true }; + m.ctx_weights = ggml_init(wp); + if (!m.ctx_weights) + return false; + WeightLoader L("turbovla", g, m.ctx_weights, m.mt); + char b[128]; + auto N = [&](const char * fmt, int64_t i, const char * s) { + std::snprintf(b, sizeof(b), fmt, (long long) i, s); + return std::string(b); + }; + + m.cls_tok = L.f32("vit.cls_token"); + if (m.n_reg > 0) + m.reg_tok = L.f32("vit.register_tokens"); + m.patch_w = L.gemm("vit.patch_embed.weight"); + m.patch_b = L.f32("vit.patch_embed.bias"); + m.vit.resize((size_t) m.vit_layers); + for (int64_t i = 0; i < m.vit_layers; ++i) { + VitLayerW & w = m.vit[(size_t) i]; + const char * f = "vit.blk.%lld.%s"; + w.ln1_w = L.f32("%s", N(f, i, "ln1.weight").c_str()); + w.ln1_b = L.f32("%s", N(f, i, "ln1.bias").c_str()); + w.qkv_w = L.fuse_gemm(N(f, i, "attn_qkv.weight").c_str(), + { N(f, i, "attn_q.weight"), N(f, i, "attn_k.weight"), N(f, i, "attn_v.weight") }); + w.qkv_b = L.fuse_f32(N(f, i, "attn_qkv.bias").c_str(), + { N(f, i, "attn_q.bias"), N(f, i, "attn_k.bias"), N(f, i, "attn_v.bias") }); + w.o_w = L.gemm("%s", N(f, i, "attn_o.weight").c_str()); + w.o_b = L.f32("%s", N(f, i, "attn_o.bias").c_str()); + w.ln2_w = L.f32("%s", N(f, i, "ln2.weight").c_str()); + w.ln2_b = L.f32("%s", N(f, i, "ln2.bias").c_str()); + w.fc1_w = L.gemm("%s", N(f, i, "fc1.weight").c_str()); + w.fc1_b = L.f32("%s", N(f, i, "fc1.bias").c_str()); + w.fc2_w = L.gemm("%s", N(f, i, "fc2.weight").c_str()); + w.fc2_b = L.f32("%s", N(f, i, "fc2.bias").c_str()); + } + + m.vp_in_w = L.f32("vit_proj.input_norm.weight"); + m.vp_in_b = L.f32("vit_proj.input_norm.bias"); + m.vp_fc1_w = L.gemm("vit_proj.mlp.0.weight"); + m.vp_fc1_b = L.f32("vit_proj.mlp.0.bias"); + m.vp_fc2_w = L.gemm("vit_proj.mlp.3.weight"); + m.vp_fc2_b = L.f32("vit_proj.mlp.3.bias"); + m.vp_skip_w = L.gemm("vit_proj.skip.weight"); + m.vp_out_w = L.f32("vit_proj.output_norm.weight"); + m.vp_out_b = L.f32("vit_proj.output_norm.bias"); + m.view_emb = L.f32("view_emb"); + + m.word_emb = L.f32("text.embed.word_embeddings"); + m.pos_emb = L.f32("text.embed.position_embeddings"); + m.type_emb = L.f32("text.embed.token_type_embeddings"); + m.emb_ln_w = L.f32("text.embed.LayerNorm.weight"); + m.emb_ln_b = L.f32("text.embed.LayerNorm.bias"); + m.bert.resize((size_t) m.bert_layers); + for (int64_t i = 0; i < m.bert_layers; ++i) { + BertLayerW & w = m.bert[(size_t) i]; + const char * f = "text.encoder.layer.%lld.%s"; + w.qkv_w = L.fuse_gemm(N(f, i, "attention.self.qkv.weight").c_str(), + { N(f, i, "attention.self.query.weight"), N(f, i, "attention.self.key.weight"), + N(f, i, "attention.self.value.weight") }); + w.qkv_b = L.fuse_f32(N(f, i, "attention.self.qkv.bias").c_str(), + { N(f, i, "attention.self.query.bias"), N(f, i, "attention.self.key.bias"), + N(f, i, "attention.self.value.bias") }); + w.o_w = L.gemm("%s", N(f, i, "attention.output.dense.weight").c_str()); + w.o_b = L.f32("%s", N(f, i, "attention.output.dense.bias").c_str()); + w.ln1_w = L.f32("%s", N(f, i, "attention.output.LayerNorm.weight").c_str()); + w.ln1_b = L.f32("%s", N(f, i, "attention.output.LayerNorm.bias").c_str()); + w.fc1_w = L.gemm("%s", N(f, i, "intermediate.dense.weight").c_str()); + w.fc1_b = L.f32("%s", N(f, i, "intermediate.dense.bias").c_str()); + w.fc2_w = L.gemm("%s", N(f, i, "output.dense.weight").c_str()); + w.fc2_b = L.f32("%s", N(f, i, "output.dense.bias").c_str()); + w.ln2_w = L.f32("%s", N(f, i, "output.LayerNorm.weight").c_str()); + w.ln2_b = L.f32("%s", N(f, i, "output.LayerNorm.bias").c_str()); + } + m.text_proj_w = L.gemm("text_proj.weight"); + m.text_proj_b = L.f32("text_proj.bias"); + + m.fusion.resize((size_t) m.fusion_layers); + m.enhancer.resize((size_t) m.fusion_layers); + for (int64_t i = 0; i < m.fusion_layers; ++i) { + FusionLayerW & w = m.fusion[(size_t) i]; + const char * f = "vl_fusion.%lld.%s"; + w.norm_v_w = L.f32("%s", N(f, i, "norm_v.weight").c_str()); + w.norm_v_b = L.f32("%s", N(f, i, "norm_v.bias").c_str()); + w.norm_l_w = L.f32("%s", N(f, i, "norm_l.weight").c_str()); + w.norm_l_b = L.f32("%s", N(f, i, "norm_l.bias").c_str()); + w.v_w = L.fuse_gemm(N(f, i, "v.weight").c_str(), { N(f, i, "v_proj.weight"), N(f, i, "values_v.weight") }); + w.v_b = L.fuse_f32(N(f, i, "v.bias").c_str(), { N(f, i, "v_proj.bias"), N(f, i, "values_v.bias") }); + w.l_w = L.fuse_gemm(N(f, i, "l.weight").c_str(), { N(f, i, "l_proj.weight"), N(f, i, "values_l.weight") }); + w.l_b = L.fuse_f32(N(f, i, "l.bias").c_str(), { N(f, i, "l_proj.bias"), N(f, i, "values_l.bias") }); + w.out_v_w = L.gemm("%s", N(f, i, "out_v.weight").c_str()); + w.out_v_b = L.f32("%s", N(f, i, "out_v.bias").c_str()); + w.out_l_w = L.gemm("%s", N(f, i, "out_l.weight").c_str()); + w.out_l_b = L.f32("%s", N(f, i, "out_l.bias").c_str()); + w.gamma_v = L.f32("%s", N(f, i, "gamma_v").c_str()); + w.gamma_l = L.f32("%s", N(f, i, "gamma_l").c_str()); + + EnhancerLayerW & e = m.enhancer[(size_t) i]; + const char * t = "vl_text.%lld.%s"; + e.qkv_w = L.gemm("%s", N(t, i, "attn_qkv.weight").c_str()); + e.qkv_b = L.f32("%s", N(t, i, "attn_qkv.bias").c_str()); + e.o_w = L.gemm("%s", N(t, i, "attn_o.weight").c_str()); + e.o_b = L.f32("%s", N(t, i, "attn_o.bias").c_str()); + e.ln1_w = L.f32("%s", N(t, i, "ln1.weight").c_str()); + e.ln1_b = L.f32("%s", N(t, i, "ln1.bias").c_str()); + e.fc1_w = L.gemm("%s", N(t, i, "fc1.weight").c_str()); + e.fc1_b = L.f32("%s", N(t, i, "fc1.bias").c_str()); + e.fc2_w = L.gemm("%s", N(t, i, "fc2.weight").c_str()); + e.fc2_b = L.f32("%s", N(t, i, "fc2.bias").c_str()); + e.ln2_w = L.f32("%s", N(t, i, "ln2.weight").c_str()); + e.ln2_b = L.f32("%s", N(t, i, "ln2.bias").c_str()); + } + + m.st_ln_w = L.f32("state.proj.0.weight"); + m.st_ln_b = L.f32("state.proj.0.bias"); + m.st_fc1_w = L.gemm("state.proj.1.weight"); + m.st_fc1_b = L.f32("state.proj.1.bias"); + m.st_fc2_w = L.gemm("state.proj.4.weight"); + m.st_fc2_b = L.f32("state.proj.4.bias"); + m.st_pos = L.f32("state.proj.position"); + m.st_out_w = L.f32("state.proj.output_norm.weight"); + m.st_out_b = L.f32("state.proj.output_norm.bias"); + + m.act_q = L.f32("act.q.weight"); + m.dec.resize((size_t) m.dec_layers); + for (int64_t i = 0; i < m.dec_layers; ++i) { + DecoderLayerW & w = m.dec[(size_t) i]; + const char * f = "act.dec.%lld.%s"; + w.ln1_w = L.f32("%s", N(f, i, "ln1.weight").c_str()); + w.ln1_b = L.f32("%s", N(f, i, "ln1.bias").c_str()); + w.self_qkv_w = L.gemm("%s", N(f, i, "self_qkv.weight").c_str()); + w.self_qkv_b = L.f32("%s", N(f, i, "self_qkv.bias").c_str()); + w.self_o_w = L.gemm("%s", N(f, i, "self_out.weight").c_str()); + w.self_o_b = L.f32("%s", N(f, i, "self_out.bias").c_str()); + w.ln2_w = L.f32("%s", N(f, i, "ln2.weight").c_str()); + w.ln2_b = L.f32("%s", N(f, i, "ln2.bias").c_str()); + w.cross_qkv_w = L.gemm("%s", N(f, i, "cross_qkv.weight").c_str()); + w.cross_qkv_b = L.f32("%s", N(f, i, "cross_qkv.bias").c_str()); + w.cross_o_w = L.gemm("%s", N(f, i, "cross_out.weight").c_str()); + w.cross_o_b = L.f32("%s", N(f, i, "cross_out.bias").c_str()); + w.ln3_w = L.f32("%s", N(f, i, "ln3.weight").c_str()); + w.ln3_b = L.f32("%s", N(f, i, "ln3.bias").c_str()); + w.fc1_w = L.gemm("%s", N(f, i, "fc1.weight").c_str()); + w.fc1_b = L.f32("%s", N(f, i, "fc1.bias").c_str()); + w.fc2_w = L.gemm("%s", N(f, i, "fc2.weight").c_str()); + w.fc2_b = L.f32("%s", N(f, i, "fc2.bias").c_str()); + } + m.ap0_w = L.gemm("act.proj.0.weight"); + m.ap0_b = L.f32("act.proj.0.bias"); + m.ap1_w = L.gemm("act.proj.1.weight"); + m.ap1_b = L.f32("act.proj.1.bias"); + m.ap2_w = L.gemm("act.proj.2.weight"); + m.ap2_b = L.f32("act.proj.2.bias"); + + if (!L.upload(m.backend, &m.weight_buf)) + return false; + + m.bert_max_pos = m.pos_emb->ne[1]; + if (m.vocab != m.word_emb->ne[1] || m.text_len_max > m.bert_max_pos || + m.patch_w->ne[0] != 3*m.patch*m.patch || (m.reg_tok && m.reg_tok->ne[1] != m.n_reg) || + ggml_nelements(m.view_emb) != m.hidden*m.n_views || m.act_q->ne[1] != m.horizon || + m.dec[0].cross_qkv_w->ne[1] != 3*m.hidden) { + std::fprintf(stderr, "vla(turbovla): tensor shapes disagree with GGUF metadata\n"); + return false; + } + return true; +} + +// Flattens one view into DINOv3's patch-embedding input: column p is patch p +// (row-major over the grid), rows follow the conv weight's (c, ky, kx) order. +void patchify(const ImageView & v, int64_t side, int64_t P, float * out) { + const int64_t g = side / P, pdim = 3*P*P; + for (int64_t p = 0; p < g*g; ++p) { + const int64_t y0 = (p / g)*P, x0 = (p % g)*P; + float * col = out + p*pdim; + for (int64_t c = 0; c < 3; ++c) + for (int64_t ky = 0; ky < P; ++ky) + for (int64_t kx = 0; kx < P; ++kx) { + const size_t idx = (size_t) ((y0 + ky)*side + x0 + kx)*3 + c; + const float px = v.format == PixelFormat::U8 ? ((const uint8_t *) v.data)[idx] / 255.0f + : ((const float *) v.data)[idx]; + col[(c*P + ky)*P + kx] = (px - kImagenetMean[c]) / kImagenetStd[c]; + } + } +} + +} // namespace + +std::unique_ptr turbovla_create(const std::string& mmproj_path, + const std::string& ckpt_path, + const std::string&, + const Options& opts) { + if (!mmproj_path.empty()) + std::printf("vla(turbovla): note - mmproj '%s' is ignored (vision is baked into the GGUF)\n", + mmproj_path.c_str()); + + auto m = std::make_unique(); + m->mt = opts.weight_dtype.value_or(GGML_TYPE_F32); + + gguf_reader g("turbovla"); + if (!g.open(ckpt_path)) + return nullptr; + if (!g.has("turbovla.architecture")) { + std::fprintf(stderr, "vla(turbovla): %s is not a TurboVLA GGUF\n", ckpt_path.c_str()); + return nullptr; + } + if (!load_config(g, *m)) + return nullptr; + + const Backend b = backend_init("vla(turbovla)", m->n_threads); + if (!b.handle) + return nullptr; + m->backend = b.handle; + + if (!load_weights(*m, g) || !m->upload_rope_tables()) + return nullptr; + + m->cfg.n_img = m->n_views * m->n_patches(); + m->cfg.n_lang = m->text_len_max; + m->cfg.n_state = m->n_state_tok; + m->cfg.n_suffix = m->horizon; + m->cfg.hidden = m->hidden; + m->cfg.max_state_dim = m->state_dim; + m->cfg.real_state_dim = m->state_dim; + m->cfg.max_action_dim = m->action_dim; + m->cfg.real_action_dim = m->action_dim; + + std::printf("vla(turbovla): weights resident %.2f GiB (%s) - DINOv3 ViT-B/16 x%lld views + BERT + " + "%lld fusion layers + ACT decoder, horizon %lld, %zu text groups\n", + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), dtype_name(m->mt), + (long long) m->n_views, (long long) m->fusion_layers, (long long) m->horizon, + m->text_groups.size()); + return m; +} + +std::vector TurboVlaModelArch::predict(const Inputs& in) { + using clock = std::chrono::steady_clock; + const auto t0 = clock::now(); + stats = Stats{}; + + if (in.precomputed_img_emb) { + std::fprintf(stderr, "vla(turbovla): precomputed_img_emb is not supported; pass raw images\n"); + return {}; + } + if (!in.images || in.n_images != n_views) { + std::fprintf(stderr, "vla(turbovla): expected %lld views (agentview, wrist), got %d\n", + (long long) n_views, in.n_images); + return {}; + } + for (int64_t i = 0; i < n_views; ++i) { + const ImageView & iv = in.images[i]; + if (!view_is_side(iv.data, iv.w, iv.h, image_size)) { + std::fprintf(stderr, "vla(turbovla): view %lld is %dx%d, expected %lldx%lld\n", + (long long) i, iv.w, iv.h, (long long) image_size, (long long) image_size); + return {}; + } + } + if (!in.lang_tokens || in.n_lang < 1 || in.n_lang > text_len_max) { + std::fprintf(stderr, "vla(turbovla): %d language tokens; need 1..%lld (truncate when tokenizing)\n", + in.n_lang, (long long) text_len_max); + return {}; + } + // ggml_get_rows does not bound-check. + for (int i = 0; i < in.n_lang; ++i) + if (in.lang_tokens[i] < 0 || in.lang_tokens[i] >= vocab) { + std::fprintf(stderr, "vla(turbovla): token %d out of vocab\n", in.lang_tokens[i]); + return {}; + } + + const int64_t LT = text_len_max, VS = n_views*n_patches(); + const int64_t bert_len = text_len(in.lang_tokens, in.n_lang); + std::vector ids((size_t) bert_len, pad_id); + std::copy(in.lang_tokens, in.lang_tokens + in.n_lang, ids.begin()); + + std::vector allowed; + std::vector pos; + special_token_blocks(ids, *this, allowed, pos); + const float NEG = -INFINITY; + std::vector bert_mask((size_t) bert_len*bert_len), enh_mask((size_t) LT*LT, NEG), fus_mask((size_t) LT*VS); + for (int64_t q = 0; q < bert_len; ++q) + for (int64_t k = 0; k < bert_len; ++k) { + bert_mask[(size_t) q*bert_len + k] = allowed[(size_t) q*bert_len + k] ? 0.0f : NEG; + enh_mask[(size_t) q*LT + k] = bert_mask[(size_t) q*bert_len + k]; + } + for (int64_t q = bert_len; q < LT; ++q) + enh_mask[(size_t) q*LT + q] = 0.0f; + for (int64_t k = 0; k < LT; ++k) { + const float val = (k < bert_len && ids[(size_t) k] != pad_id) ? 0.0f : NEG; + for (int64_t q = 0; q < VS; ++q) + fus_mask[(size_t) q*LT + k] = val; + } + + std::vector patches((size_t) 3*patch*patch*VS); + for (int64_t i = 0; i < n_views; ++i) + patchify(in.images[i], image_size, patch, patches.data() + (size_t) i*3*patch*patch*n_patches()); + std::vector state((size_t) state_dim, 0.0f); + if (in.state) + std::copy(in.state, in.state + state_dim, state.begin()); + + const size_t arena = ggml_tensor_overhead()*8192 + ggml_graph_overhead_custom(8192, false); + if (!graph.ensure(backend, Key{bert_len}, arena, + [&](ggml_context * C, IO & io) { return build(C, io, bert_len); })) { + std::fprintf(stderr, "vla(turbovla): graph build/alloc failed\n"); + return {}; + } + IO & io = graph.io(); + ggml_backend_tensor_set(io.patches, patches.data(), 0, ggml_nbytes(io.patches)); + ggml_backend_tensor_set(io.ids, ids.data(), 0, ggml_nbytes(io.ids)); + ggml_backend_tensor_set(io.pos, pos.data(), 0, ggml_nbytes(io.pos)); + ggml_backend_tensor_set(io.bert_mask, bert_mask.data(), 0, ggml_nbytes(io.bert_mask)); + ggml_backend_tensor_set(io.enh_mask, enh_mask.data(), 0, ggml_nbytes(io.enh_mask)); + ggml_backend_tensor_set(io.fus_mask, fus_mask.data(), 0, ggml_nbytes(io.fus_mask)); + ggml_backend_tensor_set(io.state, state.data(), 0, ggml_nbytes(io.state)); + + const auto tc = clock::now(); + graph_unique_names(graph.graph()); + if (ggml_backend_graph_compute(backend, graph.graph()) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(turbovla): compute failed\n"); + return {}; + } + std::vector out((size_t) (horizon*action_dim)); + ggml_backend_tensor_get(io.actions, out.data(), 0, out.size()*sizeof(float)); + + stats.ms_inference = std::chrono::duration(clock::now() - tc).count(); + stats.ms_total = std::chrono::duration(clock::now() - t0).count(); + return out; +} + +} // namespace vla diff --git a/src/serving/vla-cli.cpp b/src/serving/vla-cli.cpp index 7d31060..ec48e13 100644 --- a/src/serving/vla-cli.cpp +++ b/src/serving/vla-cli.cpp @@ -130,6 +130,7 @@ const char * arch_slug(Arch a) { case Arch::OPENVLA_OFT: return "openvla_oft"; case Arch::VLA_JEPA: return "vla_jepa"; case Arch::OCTO: return "octo"; + case Arch::TURBOVLA: return "turbovla"; } return ""; } diff --git a/tests/py/test_converters.py b/tests/py/test_converters.py index a9467b7..ba403e2 100644 --- a/tests/py/test_converters.py +++ b/tests/py/test_converters.py @@ -61,6 +61,9 @@ def numpy(self): return None def float(self): return _F32Tensor() def to(self, *a): return self def reshape(self, *a): return self + def squeeze(self, *a): return self + def clone(self): return self + def __mul__(self, other): return self class _F32Tensor(_Tensor): @@ -259,6 +262,92 @@ def test_prismatic_tower_and_lm(): assert qsrcs[2] == "language_model.model.layers.0.self_attn.q_proj.weight" +class _SourceTensors(dict): + """Return a tensor stub while recording every converter source lookup.""" + + def __init__(self): + self.keys_read = [] + + def __getitem__(self, key): + self.keys_read.append(key) + return _Tensor() + + def __contains__(self, key): + return True + + +def test_turbovla_converter_remap(): + import importlib + + T = importlib.import_module("convert_turbovla_to_gguf") + tensors = _SourceTensors() + dims = types.SimpleNamespace( + vit_dim=1, + vit_layers=1, + text_layers=1, + num_fusion_layers=1, + num_text_layers=1, + num_action_decoder_layers=1, + ) + writer = _Writer() + + T.write_vision_encoder(writer, tensors, dims) + T.write_text_encoder(writer, tensors, dims) + T.write_vision_projection(writer, tensors) + T.write_vision_language_interaction(writer, tensors, dims) + T.write_action_decoder(writer, tensors, dims) + T.write_state_projection(writer, tensors) + T.write_view_embeddings(writer, tensors) + + assert writer.names == [ + "vit.cls_token", "vit.patch_embed.weight", "vit.patch_embed.bias", "vit.register_tokens", + "vit.blk.0.attn_q.weight", "vit.blk.0.attn_q.bias", "vit.blk.0.attn_k.weight", "vit.blk.0.attn_k.bias", + "vit.blk.0.attn_v.weight", "vit.blk.0.attn_v.bias", "vit.blk.0.attn_o.weight", "vit.blk.0.attn_o.bias", + "vit.blk.0.ln1.weight", "vit.blk.0.ln1.bias", "vit.blk.0.ln2.weight", "vit.blk.0.ln2.bias", + "vit.blk.0.fc1.weight", "vit.blk.0.fc1.bias", "vit.blk.0.fc2.weight", "vit.blk.0.fc2.bias", + "text.embed.word_embeddings", "text.embed.position_embeddings", "text.embed.token_type_embeddings", + "text.embed.LayerNorm.weight", "text.embed.LayerNorm.bias", + "text.encoder.layer.0.attention.self.query.weight", "text.encoder.layer.0.attention.self.query.bias", + "text.encoder.layer.0.attention.self.key.weight", "text.encoder.layer.0.attention.self.key.bias", + "text.encoder.layer.0.attention.self.value.weight", "text.encoder.layer.0.attention.self.value.bias", + "text.encoder.layer.0.attention.output.dense.weight", "text.encoder.layer.0.attention.output.dense.bias", + "text.encoder.layer.0.attention.output.LayerNorm.weight", "text.encoder.layer.0.attention.output.LayerNorm.bias", + "text.encoder.layer.0.intermediate.dense.weight", "text.encoder.layer.0.intermediate.dense.bias", + "text.encoder.layer.0.output.dense.weight", "text.encoder.layer.0.output.dense.bias", + "text.encoder.layer.0.output.LayerNorm.weight", "text.encoder.layer.0.output.LayerNorm.bias", + "text_proj.weight", "text_proj.bias", + "vit_proj.input_norm.weight", "vit_proj.input_norm.bias", "vit_proj.mlp.0.weight", "vit_proj.mlp.0.bias", + "vit_proj.mlp.3.weight", "vit_proj.mlp.3.bias", "vit_proj.skip.weight", "vit_proj.output_norm.weight", + "vit_proj.output_norm.bias", + "vl_fusion.0.v_proj.weight", "vl_fusion.0.v_proj.bias", "vl_fusion.0.l_proj.weight", "vl_fusion.0.l_proj.bias", + "vl_fusion.0.values_v.weight", "vl_fusion.0.values_v.bias", "vl_fusion.0.values_l.weight", "vl_fusion.0.values_l.bias", + "vl_fusion.0.out_v.weight", "vl_fusion.0.out_v.bias", "vl_fusion.0.out_l.weight", "vl_fusion.0.out_l.bias", + "vl_fusion.0.norm_v.weight", "vl_fusion.0.norm_v.bias", "vl_fusion.0.norm_l.weight", "vl_fusion.0.norm_l.bias", + "vl_fusion.0.gamma_v", "vl_fusion.0.gamma_l", + "vl_text.0.attn_qkv.weight", "vl_text.0.attn_qkv.bias", "vl_text.0.attn_o.weight", "vl_text.0.attn_o.bias", + "vl_text.0.ln1.weight", "vl_text.0.ln1.bias", "vl_text.0.fc1.weight", "vl_text.0.fc1.bias", + "vl_text.0.fc2.weight", "vl_text.0.fc2.bias", "vl_text.0.ln2.weight", "vl_text.0.ln2.bias", + "act.q.weight", "act.dec.0.self_qkv.weight", "act.dec.0.self_qkv.bias", "act.dec.0.self_out.weight", + "act.dec.0.self_out.bias", "act.dec.0.cross_qkv.weight", "act.dec.0.cross_qkv.bias", + "act.dec.0.cross_out.weight", "act.dec.0.cross_out.bias", "act.dec.0.ln1.weight", "act.dec.0.ln1.bias", + "act.dec.0.ln2.weight", "act.dec.0.ln2.bias", "act.dec.0.ln3.weight", "act.dec.0.ln3.bias", + "act.dec.0.fc1.weight", "act.dec.0.fc1.bias", "act.dec.0.fc2.weight", "act.dec.0.fc2.bias", + "act.proj.0.weight", "act.proj.0.bias", "act.proj.1.weight", "act.proj.1.bias", "act.proj.2.weight", + "act.proj.2.bias", "state.proj.0.weight", "state.proj.0.bias", "state.proj.1.weight", "state.proj.1.bias", + "state.proj.4.weight", "state.proj.4.bias", "state.proj.output_norm.weight", "state.proj.output_norm.bias", + "state.proj.position", "view_emb", + ] + assert { + "vision_encoder.backbone.embeddings.patch_embeddings.weight", + "text_encoder.bert.encoder.layer.0.attention.self.query.weight", + "vision_language_interaction.fusion_layers.0.attn.v_proj.weight", + "vision_language_interaction.text_layers.0.self_attn.in_proj_weight", + "action_head.decoder.decoder.layers.0.multihead_attn.in_proj_weight", + "action_head.state_projection.net.4.weight", + "view_embedding", + } <= set(tensors.keys_read) + + def test_every_converter_imports(): # every converter must resolve against the two shared modules import importlib